diff --git a/backends/cadence/aot/ops_registrations.py b/backends/cadence/aot/ops_registrations.py
index f3e73028169..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)
@@ -3205,10 +3205,37 @@ 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,
-    mask: torch.Tensor,
     dim: int,
     mask_type: int,
     pos: torch.Tensor,
@@ -3217,13 +3244,21 @@ 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)
 
 
 @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 @@ 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/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 d3a5c853a4a..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 | None,
     dim: int,
     mask_type: int,
     pos: torch.Tensor,
@@ -2646,7 +2645,6 @@ def quantized_softmax_per_tensor_common(
 
     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 @@ 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}"
-        )
+    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 @@ 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 +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 | None,
     dim: int,
     mask_type: int,
     pos: torch.Tensor,
@@ -2701,7 +2720,6 @@ def quantized_softmax_per_tensor(
 ) -> torch.Tensor:
     return quantized_softmax_per_tensor_common(
         input_tensor,
-        mask,
         dim,
         mask_type,
         pos,
@@ -2715,7 +2733,6 @@ def quantized_softmax_per_tensor(
 @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 @@ 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 005bf9d85bd..ae938ec77aa 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,6 @@ def test_quantized_w8a32_gru_invalid_hidden_dim(self) -> None:
             (
                 "basic_int8_dim_1",
                 torch.tensor([[10, 20, 30]], dtype=torch.int8),
-                None,
                 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),
-                None,
                 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),
-                None,
                 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),
-                None,
                 1,
                 0.1,
                 0,
@@ -3197,25 +3193,12 @@ 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,
         dim: int,
         in_scale: float,
         in_zero_point: int,
@@ -3226,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
@@ -3265,7 +3247,6 @@ def test_quantized_softmax(self) -> None:
         in_zero_point = torch.tensor([0])
         output = torch.ops.cadence.quantized_softmax(
             input_tensor,
-            None,  # mask
             1,  # dim
             0,  # mask_type (no masking)
             torch.zeros(1, dtype=torch.int64),  # pos
@@ -3283,6 +3264,199 @@ 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,
+                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),
+                -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),
+                -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),
+            -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),
+            -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),
+            -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),
+                -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),
+                0,
+                1,
+                torch.zeros(1, dtype=torch.int64),
+                1.0,
+                0,
+                1.0,
+                0,
+            )
+
     @expand(
         [
             # Basic 1D slice_scatter tests
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 <typename T>
 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 <typename T>
 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<int64_t>()[0];
   quantized_softmax_per_tensor_<T>(
       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_<ctype>(                \
         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_<ctype>(     \
         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,