Remove unused mask operand from Cadence quantized softmax

backends/cadence/aot/ops_registrations.py

@@ -1,3 +1,3 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 #
@@ -490,16 +490,16 @@
 )
 
 lib.define(
-    "quantized_softmax(Tensor input, Tensor mask, int dim, int mask_type, Tensor pos, Tensor in_scale, Tensor in_zero_point, Tensor out_scale, Tensor out_zero_point) -> (Tensor out)"
+    "quantized_softmax(Tensor input, int dim, int mask_type, Tensor pos, Tensor in_scale, Tensor in_zero_point, Tensor out_scale, Tensor out_zero_point) -> (Tensor out)"
 )
 lib.define(
-    "quantized_softmax.per_tensor(Tensor input, Tensor mask, int dim, int mask_type, Tensor pos, float in_scale, int in_zero_point, float out_scale, int out_zero_point) -> (Tensor out)"
+    "quantized_softmax.per_tensor(Tensor input, int dim, int mask_type, Tensor pos, float in_scale, int in_zero_point, float out_scale, int out_zero_point) -> (Tensor out)"
 )
 lib.define(
-    "quantized_softmax.out(Tensor input, Tensor mask, int dim, int mask_type, Tensor pos, Tensor in_scale, Tensor in_zero_point, Tensor out_scale, Tensor out_zero_point, *, Tensor(a!) out) -> Tensor (a!)"
+    "quantized_softmax.out(Tensor input, int dim, int mask_type, Tensor pos, Tensor in_scale, Tensor in_zero_point, Tensor out_scale, Tensor out_zero_point, *, Tensor(a!) out) -> Tensor (a!)"
 )
 lib.define(
-    "quantized_softmax.per_tensor_out(Tensor input, Tensor mask, int dim, int mask_type, Tensor pos, float in_scale, int in_zero_point, float out_scale, int out_zero_point, *, Tensor(a!) out) -> Tensor (a!)"
+    "quantized_softmax.per_tensor_out(Tensor input, int dim, int mask_type, Tensor pos, float in_scale, int in_zero_point, float out_scale, int out_zero_point, *, Tensor(a!) out) -> Tensor (a!)"
 )
 
 # pack float/bool mask tensor into a bitmask of type uint8 (each element holding 8 bool mask elements)
@@ -3205,10 +3205,37 @@
     return input_tensor.new_empty(input_tensor.size(), dtype=torch.float32)
 
 
+def _validate_quantized_softmax_args(
+    input: torch.Tensor,
+    dim: int,
+    mask_type: int,
+    pos: torch.Tensor,
+) -> None:
+    assert input.dtype in (
+        torch.int8,
+        torch.uint8,
+        torch.int16,
+    ), "input must be int8, uint8, or int16"
+    assert input.dim() > 0, "input must have at least one dimension"
+    normalized_dim = dim if dim >= 0 else dim + input.dim()
+    assert normalized_dim == input.dim() - 1, "dim must be the last dimension"
+    assert mask_type in (0, 1), "mask_type must be 0 or 1"
+    assert pos.dtype in (torch.int16, torch.int64), "pos must be int16 or int64"
+    assert pos.numel() == 1, "pos must contain exactly one element"
+
+
+def _validate_quantized_softmax_qparam(
+    value: torch.Tensor,
+    name: str,
+    dtype: torch.dtype,
+) -> None:
+    assert value.dtype == dtype, f"{name} must have dtype {dtype}"
+    assert value.numel() == 1, f"{name} must contain exactly one element"
+
+
 @register_fake("cadence::quantized_softmax")
 def quantized_softmax_meta(
     input: torch.Tensor,
-    mask: torch.Tensor,
     dim: int,
     mask_type: int,
     pos: torch.Tensor,
@@ -3217,13 +3244,21 @@
     out_scale: torch.Tensor,
     out_zero_point: torch.Tensor,
 ) -> torch.Tensor:
+    _validate_quantized_softmax_args(input, dim, mask_type, pos)
+    _validate_quantized_softmax_qparam(in_scale, "in_scale", torch.float32)
+    _validate_quantized_softmax_qparam(
+        in_zero_point, "in_zero_point", torch.int64
+    )
+    _validate_quantized_softmax_qparam(out_scale, "out_scale", torch.float32)
+    _validate_quantized_softmax_qparam(
+        out_zero_point, "out_zero_point", torch.int64
+    )
     return input.new_empty(input.size(), dtype=input.dtype)
 
 
 @register_fake("cadence::quantized_softmax.per_tensor")
 def quantized_softmax_per_tensor_meta(
     input: torch.Tensor,
-    mask: torch.Tensor,
     dim: int,
     mask_type: int,
     pos: torch.Tensor,
@@ -3232,6 +3267,7 @@
     out_scale: float,
     out_zero_point: int,
 ) -> torch.Tensor:
+    _validate_quantized_softmax_args(input, dim, mask_type, pos)
     return input.new_empty(input.size(), dtype=input.dtype)
 
 

backends/cadence/aot/quantizer/patterns.py

@@ -12,7 +12,7 @@
 from typing import List, Optional, Tuple, Union
 
 import torch
 from executorch.backends.cadence.aot.compiler_utils import get_shape
 from executorch.backends.cadence.aot.pass_utils import get_arg, replace_with_op
 from executorch.backends.cadence.aot.quantizer.pattern_utils import (
     DQ_PER_TENSOR,
@@ -1198,21 +1198,6 @@
         if quant_node is None:
             return None
         input_q = get_arg(dq_input, "input", fx.Node)
-        quant_input = get_arg(quant_node, "input", fx.Node)
-        mask_shape = get_shape(gm, quant_input)
-        if not mask_shape:
-            return None
-        mask_shape = list(mask_shape)
-        # Softmax mask is packed 16 elements per int32 word.
-        mask_shape[-1] = mask_shape[-1] // 16
-        mask_tensor = insert_node_with_meta(
-            gm,
-            torch.ops.aten.full.default,
-            (mask_shape, 0.0),
-            {"dtype": torch.int32},
-            anchor_node,
-            input_q,
-        )
         # Initial position for streaming softmax (unused, set to 0).
         pos_tensor = insert_node_with_meta(
             gm,
@@ -1224,7 +1209,6 @@
         )
         args = (
             input_q,
-            mask_tensor,
             get_arg(anchor_node, "dim", int),
             0,
             pos_tensor,

backends/cadence/aot/ref_implementations.py

@@ -1,3 +1,3 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 #
@@ -2632,7 +2632,6 @@
 
 def quantized_softmax_per_tensor_common(
     input_tensor: torch.Tensor,
-    mask: torch.Tensor | None,
     dim: int,
     mask_type: int,
     pos: torch.Tensor,
@@ -2646,7 +2645,6 @@
 
     Args:
         - input_tensor (Tensor): The quantized input tensor
-        - mask (Tensor): Mask tensor
         - dim (int): The dimension along which softmax is computed
         - mask_type (int): Masking strategy (0=none, 1=position-based causal)
         - pos (Tensor): Position tensor for causal masking
@@ -2655,16 +2653,17 @@
         - out_scale (float): The scale of the output quantization
         - out_zero_point (int): The zero point of the output quantization
     """
-    # TODO: T228751479 - Add support for mask parameter in softmax
-    assert mask is None
-    assert (
-        mask_type == 0
-    ), f"Only mask_type=0 (no masking) is supported, got {mask_type}"
-    supported_dtypes = [torch.int8, torch.uint8, torch.int16]
-    if input_tensor.dtype not in supported_dtypes:
-        raise ValueError(
-            f"Input dtype must be one of {supported_dtypes}. Got {input_tensor.dtype}"
-        )
+    assert input_tensor.dtype in (
+        torch.int8,
+        torch.uint8,
+        torch.int16,
+    ), "input must be int8, uint8, or int16"
+    assert input_tensor.dim() > 0, "input must have at least one dimension"
+    normalized_dim = dim if dim >= 0 else dim + input_tensor.dim()
+    assert normalized_dim == input_tensor.dim() - 1, "dim must be the last dimension"
+    assert mask_type in (0, 1), "mask_type must be 0 or 1"
+    assert pos.dtype in (torch.int16, torch.int64), "pos must be int16 or int64"
+    assert pos.numel() == 1, "pos must contain exactly one element"
 
     float_input_tensor = dequantize_per_tensor(
         input_tensor,
@@ -2675,7 +2674,28 @@
         input_tensor.dtype,
     )
 
-    softmax_output = torch.nn.functional.softmax(float_input_tensor, dim=dim)
+    if mask_type == 1:
+        row_width = input_tensor.shape[-1]
+        rows = float_input_tensor.reshape(-1, row_width)
+        base_pos = int(pos.reshape(-1)[0].item())
+        if base_pos < 0:
+            softmax_output = torch.zeros_like(float_input_tensor)
+        else:
+            row_positions = base_pos + torch.arange(
+                rows.shape[0], device=rows.device
+            ).unsqueeze(1)
+            column_positions = torch.arange(row_width, device=rows.device).unsqueeze(
+                0
+            )
+            causal_mask = column_positions > row_positions
+            softmax_output = torch.ops.aten._masked_softmax.default(
+                float_input_tensor,
+                causal_mask.reshape_as(float_input_tensor),
+                dim,
+                2,
+            )
+    else:
+        softmax_output = torch.nn.functional.softmax(float_input_tensor, dim=dim)
 
     return quantize_per_tensor(
         softmax_output,
@@ -2690,7 +2710,6 @@
 @impl_tracked(m, "quantized_softmax.per_tensor")
 def quantized_softmax_per_tensor(
     input_tensor: torch.Tensor,
-    mask: torch.Tensor | None,
     dim: int,
     mask_type: int,
     pos: torch.Tensor,
@@ -2701,7 +2720,6 @@
 ) -> torch.Tensor:
     return quantized_softmax_per_tensor_common(
         input_tensor,
-        mask,
         dim,
         mask_type,
         pos,
@@ -2715,7 +2733,6 @@
 @impl_tracked(m, "quantized_softmax")
 def quantized_softmax(
     input_tensor: torch.Tensor,
-    mask: torch.Tensor | None,
     dim: int,
     mask_type: int,
     pos: torch.Tensor,
@@ -2726,7 +2743,6 @@
 ) -> torch.Tensor:
     return quantized_softmax_per_tensor_common(
         input_tensor,
-        mask,
         dim,
         mask_type,
         pos,