From 7e4780fa026c5cbe18a4939eaac561ed2730ab95 Mon Sep 17 00:00:00 2001 From: Andrew Grebenisan Date: Wed, 5 Aug 2026 16:11:14 -0700 Subject: [PATCH 1/2] Add mask support to cadence softmax (#21596) Summary: The kernel supports masking, need to support it in the reference as well. Reviewed By: aliafzal Differential Revision: D114671838 --- backends/cadence/aot/ops_registrations.py | 38 +++ backends/cadence/aot/ref_implementations.py | 55 +++-- .../aot/tests/test_ref_implementations.py | 227 ++++++++++++++++-- 3 files changed, 286 insertions(+), 34 deletions(-) diff --git a/backends/cadence/aot/ops_registrations.py b/backends/cadence/aot/ops_registrations.py index f3e73028169..b0cb19b3d7a 100644 --- a/backends/cadence/aot/ops_registrations.py +++ b/backends/cadence/aot/ops_registrations.py @@ -3205,6 +3205,34 @@ def softmax_f32_f32_meta( 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, @@ -3217,6 +3245,15 @@ def quantized_softmax_meta( 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) @@ -3232,6 +3269,7 @@ def quantized_softmax_per_tensor_meta( 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) diff --git a/backends/cadence/aot/ref_implementations.py b/backends/cadence/aot/ref_implementations.py index d3a5c853a4a..26453dccfd2 100644 --- a/backends/cadence/aot/ref_implementations.py +++ b/backends/cadence/aot/ref_implementations.py @@ -2632,7 +2632,7 @@ def softmax_f32_f32( def quantized_softmax_per_tensor_common( input_tensor: torch.Tensor, - mask: torch.Tensor | None, + mask: torch.Tensor, dim: int, mask_type: int, pos: torch.Tensor, @@ -2646,7 +2646,9 @@ def quantized_softmax_per_tensor_common( Args: - input_tensor (Tensor): The quantized input tensor - - mask (Tensor): Mask tensor + - mask (Tensor): Currently ignored. Causal masking is still supported when + mask_type=1 and is derived from pos; this argument is reserved for future + mask types. - 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 +2657,18 @@ def quantized_softmax_per_tensor_common( - 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}" - ) + del mask + 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 +2679,28 @@ def quantized_softmax_per_tensor_common( 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 +2715,7 @@ def quantized_softmax_per_tensor_common( @impl_tracked(m, "quantized_softmax.per_tensor") def quantized_softmax_per_tensor( input_tensor: torch.Tensor, - mask: torch.Tensor | None, + mask: torch.Tensor, dim: int, mask_type: int, pos: torch.Tensor, @@ -2715,7 +2740,7 @@ def quantized_softmax_per_tensor( @impl_tracked(m, "quantized_softmax") def quantized_softmax( input_tensor: torch.Tensor, - mask: torch.Tensor | None, + mask: torch.Tensor, dim: int, mask_type: int, pos: torch.Tensor, diff --git a/backends/cadence/aot/tests/test_ref_implementations.py b/backends/cadence/aot/tests/test_ref_implementations.py index 005bf9d85bd..0353f041c47 100644 --- a/backends/cadence/aot/tests/test_ref_implementations.py +++ b/backends/cadence/aot/tests/test_ref_implementations.py @@ -8,7 +8,7 @@ import typing import unittest -import executorch.backends.cadence.aot.ops_registrations # noqa +import executorch.backends.cadence.aot.ops_registrations as ops_registrations import executorch.backends.cadence.aot.ref_implementations # noqa import numpy as np @@ -3152,7 +3152,7 @@ def test_quantized_w8a32_gru_invalid_hidden_dim(self) -> None: ( "basic_int8_dim_1", torch.tensor([[10, 20, 30]], dtype=torch.int8), - None, + torch.empty(0, dtype=torch.int8), 1, 0.1, 0, @@ -3164,7 +3164,7 @@ def test_quantized_w8a32_gru_invalid_hidden_dim(self) -> None: ( "uint8_with_zero_points", torch.tensor([[128, 130, 132]], dtype=torch.uint8), - None, + torch.empty(0, dtype=torch.int8), 1, 0.1, 128, @@ -3176,7 +3176,7 @@ def test_quantized_w8a32_gru_invalid_hidden_dim(self) -> None: ( "basic_int16", torch.tensor([[100, 200, 300]], dtype=torch.int16), - None, + torch.empty(0, dtype=torch.int8), 1, 0.01, 0, @@ -3188,7 +3188,7 @@ def test_quantized_w8a32_gru_invalid_hidden_dim(self) -> None: ( "multi_row_int8", torch.tensor([[10, 20, 30], [5, 10, 15]], dtype=torch.int8), - None, + torch.empty(0, dtype=torch.int8), 1, 0.1, 0, @@ -3197,25 +3197,13 @@ def test_quantized_w8a32_gru_invalid_hidden_dim(self) -> None: torch.int8, torch.tensor([[23, 61, 127], [47, 77, 127]], dtype=torch.int8), ), - ( - "softmax_dim_0", - torch.tensor([[10, 20], [30, 40]], dtype=torch.int8), - None, - 0, - 0.1, - 0, - 0.004, - 0, - torch.int8, - torch.tensor([[30, 30], [127, 127]], dtype=torch.int8), - ), ] ) def test_quantized_softmax_per_tensor( self, name: str, input_tensor: torch.Tensor, - mask: torch.Tensor | None, + mask: torch.Tensor, dim: int, in_scale: float, in_zero_point: int, @@ -3265,7 +3253,7 @@ def test_quantized_softmax(self) -> None: in_zero_point = torch.tensor([0]) output = torch.ops.cadence.quantized_softmax( input_tensor, - None, # mask + torch.empty(0, dtype=torch.int8), # unused mask 1, # dim 0, # mask_type (no masking) torch.zeros(1, dtype=torch.int64), # pos @@ -3283,6 +3271,207 @@ def test_quantized_softmax(self) -> None: "Output shape should match input shape", ) + @expand( + [ + ( + "input_dtype", + torch.ones((2, 4), dtype=torch.float32), + -1, + 0, + torch.zeros(1, dtype=torch.int64), + "input must be int8, uint8, or int16", + ), + ( + "input_rank", + torch.tensor(1, dtype=torch.int8), + 0, + 0, + torch.zeros(1, dtype=torch.int64), + "input must have at least one dimension", + ), + ( + "dim", + torch.ones((2, 4), dtype=torch.int8), + 0, + 0, + torch.zeros(1, dtype=torch.int64), + "dim must be the last dimension", + ), + ( + "mask_type", + torch.ones((2, 4), dtype=torch.int8), + -1, + 2, + torch.zeros(1, dtype=torch.int64), + "mask_type must be 0 or 1", + ), + ( + "pos_dtype", + torch.ones((2, 4), dtype=torch.int8), + -1, + 1, + torch.zeros(1, dtype=torch.int32), + "pos must be int16 or int64", + ), + ( + "pos_shape", + torch.ones((2, 4), dtype=torch.int8), + -1, + 1, + torch.zeros(2, dtype=torch.int64), + "pos must contain exactly one element", + ), + ] + ) + def test_quantized_softmax_meta_rejects_invalid_arguments( + self, + name: str, + input_tensor: torch.Tensor, + dim: int, + mask_type: int, + pos: torch.Tensor, + error: str, + ) -> None: + with self.assertRaisesRegex(AssertionError, error, msg=name): + ops_registrations.quantized_softmax_per_tensor_meta( + input_tensor, + torch.empty(0, dtype=torch.int8), + dim, + mask_type, + pos, + 1.0, + 0, + 1.0, + 0, + ) + + def test_quantized_softmax_meta_rejects_non_scalar_qparams(self) -> None: + with self.assertRaisesRegex( + AssertionError, "in_scale must contain exactly one element" + ): + ops_registrations.quantized_softmax_meta( + torch.ones((2, 4), dtype=torch.int8), + torch.empty(0, dtype=torch.int8), + -1, + 0, + torch.zeros(1, dtype=torch.int64), + torch.ones(2, dtype=torch.float32), + torch.zeros(1, dtype=torch.int64), + torch.ones(1, dtype=torch.float32), + torch.zeros(1, dtype=torch.int64), + ) + + def test_quantized_softmax_rejects_non_scalar_pos(self) -> None: + with self.assertRaisesRegex( + AssertionError, "pos must contain exactly one element" + ): + torch.ops.cadence.quantized_softmax.per_tensor( + torch.ones((2, 4), dtype=torch.int8), + torch.empty(0, dtype=torch.int8), + -1, + 1, + torch.zeros(2, dtype=torch.int64), + 1.0, + 0, + 1.0, + 0, + ) + + def test_quantized_softmax_per_tensor_causal(self) -> None: + output = torch.ops.cadence.quantized_softmax.per_tensor( + torch.zeros((2, 4), dtype=torch.int8), + torch.zeros((2, 1), dtype=torch.int32), + -1, + 1, + torch.zeros(1, dtype=torch.int64), + 1.0, + 0, + 1.0 / 128, + 0, + ) + + self.assertTrue( + torch.equal( + output, + torch.tensor([[127, 0, 0, 0], [64, 64, 0, 0]], dtype=torch.int8), + ) + ) + + def test_quantized_softmax_negative_base_pos_returns_quantized_zero(self) -> None: + out_zero_point = 17 + output = torch.ops.cadence.quantized_softmax.per_tensor( + torch.ones((2, 4), dtype=torch.int8), + torch.empty(0, dtype=torch.int8), + -1, + 1, + torch.tensor([-1], dtype=torch.int64), + 1.0, + 0, + 1.0, + out_zero_point, + ) + + self.assertTrue( + torch.equal(output, torch.full_like(output, out_zero_point)) + ) + + def test_quantized_softmax_causal_position_advances_across_leading_dims( + self, + ) -> None: + output = torch.ops.cadence.quantized_softmax.per_tensor( + torch.zeros((2, 2, 4), dtype=torch.int8), + torch.empty(0, dtype=torch.int8), + -1, + 1, + torch.zeros(1, dtype=torch.int64), + 1.0, + 0, + 1.0 / 128, + 0, + ) + + self.assertTrue( + torch.equal( + output != 0, + torch.tensor( + [ + [[True, False, False, False], [True, True, False, False]], + [[True, True, True, False], [True, True, True, True]], + ] + ), + ) + ) + + def test_quantized_softmax_rejects_invalid_mask_type(self) -> None: + with self.assertRaisesRegex(AssertionError, "mask_type must be 0 or 1"): + torch.ops.cadence.quantized_softmax.per_tensor( + torch.ones((2, 4), dtype=torch.int8), + torch.empty(0, dtype=torch.int8), + -1, + 2, + torch.zeros(1, dtype=torch.int64), + 1.0, + 0, + 1.0, + 0, + ) + + def test_quantized_softmax_rejects_non_last_dim(self) -> None: + with self.assertRaisesRegex( + AssertionError, "dim must be the last dimension" + ): + torch.ops.cadence.quantized_softmax.per_tensor( + torch.ones((2, 4), dtype=torch.int8), + torch.empty(0, dtype=torch.int8), + 0, + 1, + torch.zeros(1, dtype=torch.int64), + 1.0, + 0, + 1.0, + 0, + ) + @expand( [ # Basic 1D slice_scatter tests From d841c77363bcb3e922e9b4efcc1010a9599cf529 Mon Sep 17 00:00:00 2001 From: Andrew Grebenisan Date: Wed, 5 Aug 2026 16:11:14 -0700 Subject: [PATCH 2/2] Remove unused `mask` operand from Cadence quantized softmax Summary: Remove the unused `mask` operand from the quantized softmax schemas, fake and reference implementations, generic and DLA kernel APIs, and tests. Stop materializing a placeholder packed mask in the softmax quantizer pattern. Preserve position-based causal masking through `mask_type` and `pos`, and remove the dead DLA mask buffer. Differential Revision: D114967226 --- backends/cadence/aot/ops_registrations.py | 10 ++++------ backends/cadence/aot/quantizer/patterns.py | 16 ---------------- backends/cadence/aot/ref_implementations.py | 9 --------- .../aot/tests/test_ref_implementations.py | 15 --------------- .../generic/operators/op_quantized_softmax.cpp | 7 ------- .../generic/operators/op_quantized_softmax.h | 3 --- 6 files changed, 4 insertions(+), 56 deletions(-) diff --git a/backends/cadence/aot/ops_registrations.py b/backends/cadence/aot/ops_registrations.py index b0cb19b3d7a..16739e35385 100644 --- a/backends/cadence/aot/ops_registrations.py +++ b/backends/cadence/aot/ops_registrations.py @@ -490,16 +490,16 @@ def register_fake( ) 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) @@ -3236,7 +3236,6 @@ def _validate_quantized_softmax_qparam( @register_fake("cadence::quantized_softmax") def quantized_softmax_meta( input: torch.Tensor, - mask: torch.Tensor, dim: int, mask_type: int, pos: torch.Tensor, @@ -3260,7 +3259,6 @@ def quantized_softmax_meta( @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, diff --git a/backends/cadence/aot/quantizer/patterns.py b/backends/cadence/aot/quantizer/patterns.py index e3dc7afd0cf..f263531ccdf 100644 --- a/backends/cadence/aot/quantizer/patterns.py +++ b/backends/cadence/aot/quantizer/patterns.py @@ -1198,21 +1198,6 @@ def fuse(self, gm: fx.GraphModule, anchor_node: fx.Node) -> fx.Node | None: 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 @@ def fuse(self, gm: fx.GraphModule, anchor_node: fx.Node) -> fx.Node | None: ) args = ( input_q, - mask_tensor, get_arg(anchor_node, "dim", int), 0, pos_tensor, diff --git a/backends/cadence/aot/ref_implementations.py b/backends/cadence/aot/ref_implementations.py index 26453dccfd2..694273f24a1 100644 --- a/backends/cadence/aot/ref_implementations.py +++ b/backends/cadence/aot/ref_implementations.py @@ -2632,7 +2632,6 @@ def softmax_f32_f32( def quantized_softmax_per_tensor_common( input_tensor: torch.Tensor, - mask: torch.Tensor, dim: int, mask_type: int, pos: torch.Tensor, @@ -2646,9 +2645,6 @@ def quantized_softmax_per_tensor_common( Args: - input_tensor (Tensor): The quantized input tensor - - mask (Tensor): Currently ignored. Causal masking is still supported when - mask_type=1 and is derived from pos; this argument is reserved for future - mask types. - 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 @@ -2657,7 +2653,6 @@ def quantized_softmax_per_tensor_common( - out_scale (float): The scale of the output quantization - out_zero_point (int): The zero point of the output quantization """ - del mask assert input_tensor.dtype in ( torch.int8, torch.uint8, @@ -2715,7 +2710,6 @@ def quantized_softmax_per_tensor_common( @impl_tracked(m, "quantized_softmax.per_tensor") def quantized_softmax_per_tensor( input_tensor: torch.Tensor, - mask: torch.Tensor, dim: int, mask_type: int, pos: torch.Tensor, @@ -2726,7 +2720,6 @@ def quantized_softmax_per_tensor( ) -> torch.Tensor: return quantized_softmax_per_tensor_common( input_tensor, - mask, dim, mask_type, pos, @@ -2740,7 +2733,6 @@ def quantized_softmax_per_tensor( @impl_tracked(m, "quantized_softmax") def quantized_softmax( input_tensor: torch.Tensor, - mask: torch.Tensor, dim: int, mask_type: int, pos: torch.Tensor, @@ -2751,7 +2743,6 @@ def quantized_softmax( ) -> torch.Tensor: return quantized_softmax_per_tensor_common( input_tensor, - mask, dim, mask_type, pos, diff --git a/backends/cadence/aot/tests/test_ref_implementations.py b/backends/cadence/aot/tests/test_ref_implementations.py index 0353f041c47..ae938ec77aa 100644 --- a/backends/cadence/aot/tests/test_ref_implementations.py +++ b/backends/cadence/aot/tests/test_ref_implementations.py @@ -3152,7 +3152,6 @@ def test_quantized_w8a32_gru_invalid_hidden_dim(self) -> None: ( "basic_int8_dim_1", torch.tensor([[10, 20, 30]], dtype=torch.int8), - torch.empty(0, dtype=torch.int8), 1, 0.1, 0, @@ -3164,7 +3163,6 @@ def test_quantized_w8a32_gru_invalid_hidden_dim(self) -> None: ( "uint8_with_zero_points", torch.tensor([[128, 130, 132]], dtype=torch.uint8), - torch.empty(0, dtype=torch.int8), 1, 0.1, 128, @@ -3176,7 +3174,6 @@ def test_quantized_w8a32_gru_invalid_hidden_dim(self) -> None: ( "basic_int16", torch.tensor([[100, 200, 300]], dtype=torch.int16), - torch.empty(0, dtype=torch.int8), 1, 0.01, 0, @@ -3188,7 +3185,6 @@ def test_quantized_w8a32_gru_invalid_hidden_dim(self) -> None: ( "multi_row_int8", torch.tensor([[10, 20, 30], [5, 10, 15]], dtype=torch.int8), - torch.empty(0, dtype=torch.int8), 1, 0.1, 0, @@ -3203,7 +3199,6 @@ def test_quantized_softmax_per_tensor( self, name: str, input_tensor: torch.Tensor, - mask: torch.Tensor, dim: int, in_scale: float, in_zero_point: int, @@ -3214,7 +3209,6 @@ def test_quantized_softmax_per_tensor( ) -> None: output = torch.ops.cadence.quantized_softmax.per_tensor( input_tensor, - mask, dim, 0, # mask_type (no masking) torch.zeros(1, dtype=torch.int64), # pos @@ -3253,7 +3247,6 @@ def test_quantized_softmax(self) -> None: in_zero_point = torch.tensor([0]) output = torch.ops.cadence.quantized_softmax( input_tensor, - torch.empty(0, dtype=torch.int8), # unused mask 1, # dim 0, # mask_type (no masking) torch.zeros(1, dtype=torch.int64), # pos @@ -3335,7 +3328,6 @@ def test_quantized_softmax_meta_rejects_invalid_arguments( with self.assertRaisesRegex(AssertionError, error, msg=name): ops_registrations.quantized_softmax_per_tensor_meta( input_tensor, - torch.empty(0, dtype=torch.int8), dim, mask_type, pos, @@ -3351,7 +3343,6 @@ def test_quantized_softmax_meta_rejects_non_scalar_qparams(self) -> None: ): ops_registrations.quantized_softmax_meta( torch.ones((2, 4), dtype=torch.int8), - torch.empty(0, dtype=torch.int8), -1, 0, torch.zeros(1, dtype=torch.int64), @@ -3367,7 +3358,6 @@ def test_quantized_softmax_rejects_non_scalar_pos(self) -> None: ): torch.ops.cadence.quantized_softmax.per_tensor( torch.ones((2, 4), dtype=torch.int8), - torch.empty(0, dtype=torch.int8), -1, 1, torch.zeros(2, dtype=torch.int64), @@ -3380,7 +3370,6 @@ def test_quantized_softmax_rejects_non_scalar_pos(self) -> None: def test_quantized_softmax_per_tensor_causal(self) -> None: output = torch.ops.cadence.quantized_softmax.per_tensor( torch.zeros((2, 4), dtype=torch.int8), - torch.zeros((2, 1), dtype=torch.int32), -1, 1, torch.zeros(1, dtype=torch.int64), @@ -3401,7 +3390,6 @@ def test_quantized_softmax_negative_base_pos_returns_quantized_zero(self) -> Non out_zero_point = 17 output = torch.ops.cadence.quantized_softmax.per_tensor( torch.ones((2, 4), dtype=torch.int8), - torch.empty(0, dtype=torch.int8), -1, 1, torch.tensor([-1], dtype=torch.int64), @@ -3420,7 +3408,6 @@ def test_quantized_softmax_causal_position_advances_across_leading_dims( ) -> None: output = torch.ops.cadence.quantized_softmax.per_tensor( torch.zeros((2, 2, 4), dtype=torch.int8), - torch.empty(0, dtype=torch.int8), -1, 1, torch.zeros(1, dtype=torch.int64), @@ -3446,7 +3433,6 @@ def test_quantized_softmax_rejects_invalid_mask_type(self) -> None: with self.assertRaisesRegex(AssertionError, "mask_type must be 0 or 1"): torch.ops.cadence.quantized_softmax.per_tensor( torch.ones((2, 4), dtype=torch.int8), - torch.empty(0, dtype=torch.int8), -1, 2, torch.zeros(1, dtype=torch.int64), @@ -3462,7 +3448,6 @@ def test_quantized_softmax_rejects_non_last_dim(self) -> None: ): torch.ops.cadence.quantized_softmax.per_tensor( torch.ones((2, 4), dtype=torch.int8), - torch.empty(0, dtype=torch.int8), 0, 1, torch.zeros(1, dtype=torch.int64), diff --git a/backends/cadence/generic/operators/op_quantized_softmax.cpp b/backends/cadence/generic/operators/op_quantized_softmax.cpp index a1d78810372..8b809790fd4 100644 --- a/backends/cadence/generic/operators/op_quantized_softmax.cpp +++ b/backends/cadence/generic/operators/op_quantized_softmax.cpp @@ -118,7 +118,6 @@ void updatePositionMaskIncremental( template void quantized_softmax_per_tensor_( const Tensor& input, - ET_UNUSED const Tensor& mask, int64_t dim, int64_t mask_type, const Tensor& pos, @@ -280,7 +279,6 @@ void quantized_softmax_per_tensor_( template void quantized_softmax_( const Tensor& input, - const Tensor& mask, const int64_t dim, int64_t mask_type, const Tensor& pos, @@ -296,7 +294,6 @@ void quantized_softmax_( int64_t output_zero_point = out_zero_point.const_data_ptr()[0]; quantized_softmax_per_tensor_( input, - mask, dim, mask_type, pos, @@ -312,7 +309,6 @@ void quantized_softmax_( Tensor& quantized_softmax_out( ET_UNUSED KernelRuntimeContext& ctx, const Tensor& input, - const Tensor& mask, int64_t dim, int64_t mask_type, const Tensor& pos, @@ -325,7 +321,6 @@ Tensor& quantized_softmax_out( case ScalarType::dtype: { \ quantized_softmax_( \ input, \ - mask, \ dim, \ mask_type, \ pos, \ @@ -352,7 +347,6 @@ Tensor& quantized_softmax_out( Tensor& quantized_softmax_per_tensor_out( ET_UNUSED KernelRuntimeContext& ctx, const Tensor& input, - const Tensor& mask, int64_t dim, int64_t mask_type, const Tensor& pos, @@ -365,7 +359,6 @@ Tensor& quantized_softmax_per_tensor_out( case ScalarType::dtype: { \ quantized_softmax_per_tensor_( \ input, \ - mask, \ dim, \ mask_type, \ pos, \ diff --git a/backends/cadence/generic/operators/op_quantized_softmax.h b/backends/cadence/generic/operators/op_quantized_softmax.h index 4ce6a966f77..12e33db4507 100644 --- a/backends/cadence/generic/operators/op_quantized_softmax.h +++ b/backends/cadence/generic/operators/op_quantized_softmax.h @@ -23,7 +23,6 @@ namespace native { * * @param ctx Kernel runtime context (unused) * @param input Input quantized tensor - * @param mask Mask tensor (currently unused, reserved for future mask types) * @param dim Dimension along which to compute softmax. Only the last dimension * is currently supported (dim == -1 or dim == input.dim() - 1) * @param mask_type Masking strategy to use: @@ -49,7 +48,6 @@ namespace native { ::executorch::aten::Tensor& quantized_softmax_out( __ET_UNUSED ::executorch::runtime::KernelRuntimeContext& ctx, const ::executorch::aten::Tensor& input, - const ::executorch::aten::Tensor& mask, int64_t dim, int64_t mask_type, const ::executorch::aten::Tensor& pos, @@ -62,7 +60,6 @@ ::executorch::aten::Tensor& quantized_softmax_out( ::executorch::aten::Tensor& quantized_softmax_per_tensor_out( __ET_UNUSED ::executorch::runtime::KernelRuntimeContext& ctx, const ::executorch::aten::Tensor& input, - const ::executorch::aten::Tensor& mask, int64_t dim, int64_t mask_type, const ::executorch::aten::Tensor& pos,