[OP] SelectiveSSM and PagedSelectiveSSM operations

Author: mmikolajcz

src/core/dev_api/dev_headers.cmake

@@ -28,7 +28,9 @@ set(DEV_HEADERS
     ${CMAKE_CURRENT_LIST_DIR}/openvino/op/paged_attention.hpp
     ${CMAKE_CURRENT_LIST_DIR}/openvino/op/paged_causal_conv1d.hpp
     ${CMAKE_CURRENT_LIST_DIR}/openvino/op/paged_gated_delta_net.hpp
+    ${CMAKE_CURRENT_LIST_DIR}/openvino/op/paged_selective_ssm.hpp
     ${CMAKE_CURRENT_LIST_DIR}/openvino/op/rms_norm.hpp
+    ${CMAKE_CURRENT_LIST_DIR}/openvino/op/selective_ssm.hpp
     ${CMAKE_CURRENT_LIST_DIR}/openvino/op/util/node_util.hpp
     ${CMAKE_CURRENT_LIST_DIR}/openvino/op/util/slice_plan.hpp
     ${CMAKE_CURRENT_LIST_DIR}/openvino/opsets/opset10_decl.hpp

src/core/dev_api/openvino/op/paged_selective_ssm.hpp

@@ -0,0 +1,74 @@
+// Copyright (C) 2018-2026 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+#pragma once
+
+#include "openvino/op/op.hpp"
+
+namespace ov::op::internal {
+/// \note PagedSelectiveSSM op class is under development and subject to change
+///
+/// \brief Operator performing paged SelectiveSSM computation for continuous batching.
+///
+/// Paged variant of the SelectiveSSM (Mamba2 selective state-space, arXiv:2405.21060) recurrence. Processes
+/// tokens from multiple sequences packed into a single batch and manages the recurrent SSM state (a per-head
+/// ``[head_dim, state_size]`` matrix) using a paged block table, enabling non-contiguous memory allocation
+/// across sequences. Per token: ``dA = exp(A * dt)``, ``dtB = dt * B``, ``dBx = dtB (x) x``,
+/// ``state = state * dA + dBx``, ``output = sum(state * C)``.
+///
+/// ``B`` and ``C`` are grouped and shared across heads: each head ``h`` reads group
+/// ``g = h / heads_per_group``, where ``heads_per_group = num_heads / num_groups``.
+///
+/// ``recurrent_state_table`` is updated in place. For sequence ``s``, the state is cached every
+/// ``cache_interval[s]`` tokens into the blocks addressed by
+/// ``la_block_indices[la_block_indices_begins[s] : la_block_indices_begins[s+1]]``;
+/// ``cache_interval[s] <= 0`` disables caching for that sequence. ``num_processed_tokens[s]`` gives the
+/// count of previously processed tokens, used to resume from the correct cached state and block offset.
+/// \ingroup ov_ops_cpp_api
+class OPENVINO_API PagedSelectiveSSM : public ov::op::Op {
+public:
+    OPENVINO_OP("PagedSelectiveSSM");
+
+    PagedSelectiveSSM() = default;
+    /// \brief Constructs a PagedSelectiveSSM operation.
+    ///
+    /// \param A (Negative) log-decay rates per head [num_heads].
+    /// \param dt Per-token, per-head time steps used for discretization [batch_size_in_tokens, num_heads].
+    /// \param B Grouped input projection [batch_size_in_tokens, num_groups, state_size].
+    /// \param x Input hidden states [batch_size_in_tokens, num_heads, head_dim].
+    /// \param C Grouped output projection [batch_size_in_tokens, num_groups, state_size].
+    /// \param recurrent_state_table Paged table of recurrent state snapshots, updated in place
+    ///        [num_blocks, num_heads, head_dim, state_size]; an all-zeros tensor before any tokens are cached.
+    /// \param subsequence_begins Start indices of each sequence's tokens in the flattened token batch
+    ///        [batch_size_in_sequences + 1], element type i32 or i64.
+    /// \param la_block_indices Physical block row indices into recurrent_state_table, concatenated across
+    ///        all sequences [num_blocks], element type i32 or i64.
+    /// \param la_block_indices_begins Splits la_block_indices among sequences
+    ///        [batch_size_in_sequences + 1], element type i32 or i64.
+    /// \param num_processed_tokens Number of tokens already processed for each sequence
+    ///        [batch_size_in_sequences], element type i32 or i64.
+    /// \param cache_interval Interval (in tokens) at which the recurrent state is cached for each sequence;
+    ///        a value <= 0 disables caching [batch_size_in_sequences], element type i32 or i64.
+    PagedSelectiveSSM(const Output<Node>& A,
+                      const Output<Node>& dt,
+                      const Output<Node>& B,
+                      const Output<Node>& x,
+                      const Output<Node>& C,
+                      const Output<Node>& recurrent_state_table,
+                      const Output<Node>& subsequence_begins,
+                      const Output<Node>& la_block_indices,
+                      const Output<Node>& la_block_indices_begins,
+                      const Output<Node>& num_processed_tokens,
+                      const Output<Node>& cache_interval);
+
+    /// \brief Constructs a PagedSelectiveSSM operation from input vector.
+    ///
+    /// \param args Input tensor vector (11 inputs in order listed above).
+    explicit PagedSelectiveSSM(const ov::OutputVector& args);
+
+    void validate_and_infer_types() override;
+    bool visit_attributes(AttributeVisitor& visitor) override;
+    std::shared_ptr<ov::Node> clone_with_new_inputs(const ov::OutputVector& new_args) const override;
+};
+
+}  // namespace ov::op::internal

src/core/dev_api/openvino/op/selective_ssm.hpp

@@ -0,0 +1,51 @@
+// Copyright (C) 2018-2026 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+#pragma once
+
+#include "openvino/op/op.hpp"
+
+namespace ov::op::internal {
+/// \note SelectiveSSM op class is under development and subject to change
+///
+/// \brief Operator performing the Mamba2 selective state-space recurrence (arXiv:2405.21060).
+///
+/// Discretizes ``A`` (log-decay rates) and ``B`` (input projection) with ``dt`` (time steps) ahead of the
+/// recurrence: ``dA = exp(A * dt)``, ``dtB = dt * B``. Then, for each token ``t``:
+/// ``dBx_t = dtB_t (x) x_t``, ``state_t = state_{t-1} * dA_t + dBx_t``, ``y_t = sum(state_t * C_t)``.
+///
+/// ``B`` and ``C`` are grouped and shared across heads: each head ``h`` reads group
+/// ``g = h / heads_per_group``, where ``heads_per_group = num_heads / num_groups``.
+/// \ingroup ov_ops_cpp_api
+class OPENVINO_API SelectiveSSM : public ov::op::Op {
+public:
+    OPENVINO_OP("SelectiveSSM");
+
+    SelectiveSSM() = default;
+    /// \brief Constructs a SelectiveSSM operation.
+    ///
+    /// \param A (Negative) log-decay rates per head [num_heads].
+    /// \param dt Per-token, per-head time steps used for discretization [batch_size, seq_len, num_heads].
+    /// \param B Grouped input projection [batch_size, seq_len, num_groups, state_size].
+    /// \param x Input hidden states [batch_size, seq_len, num_heads, head_dim].
+    /// \param C Grouped output projection [batch_size, seq_len, num_groups, state_size].
+    /// \param recurrent_state Initial SSM hidden state [batch_size, num_heads, head_dim, state_size];
+    ///        an all-zeros tensor for a fresh sequence.
+    SelectiveSSM(const Output<Node>& A,
+                 const Output<Node>& dt,
+                 const Output<Node>& B,
+                 const Output<Node>& x,
+                 const Output<Node>& C,
+                 const Output<Node>& recurrent_state);
+
+    /// \brief Constructs a SelectiveSSM operation from input vector.
+    ///
+    /// \param args Input tensor vector in order: A, dt, B, x, C, recurrent_state.
+    explicit SelectiveSSM(const ov::OutputVector& args);
+
+    void validate_and_infer_types() override;
+    bool visit_attributes(AttributeVisitor& visitor) override;
+    std::shared_ptr<ov::Node> clone_with_new_inputs(const ov::OutputVector& new_args) const override;
+};
+
+}  // namespace ov::op::internal

src/core/shape_inference/CMakeLists.txt

@@ -100,6 +100,7 @@ target_sources(${TARGET_NAME}
     ${SHAPE_INFER_INCLUDE_DIR}/paged_attention_shape_inference.hpp
     ${SHAPE_INFER_INCLUDE_DIR}/paged_causal_conv1d_shape_inference.hpp
     ${SHAPE_INFER_INCLUDE_DIR}/paged_gated_delta_net_shape_inference.hpp
+    ${SHAPE_INFER_INCLUDE_DIR}/paged_selective_ssm_shape_inference.hpp
     ${SHAPE_INFER_INCLUDE_DIR}/pooling_shape_inference_util.hpp
     ${SHAPE_INFER_INCLUDE_DIR}/prior_box_clustered_shape_inference.hpp
     ${SHAPE_INFER_INCLUDE_DIR}/prior_box_shape_inference.hpp
@@ -130,6 +131,7 @@ target_sources(${TARGET_NAME}
     ${SHAPE_INFER_INCLUDE_DIR}/search_sorted_shape_inference.hpp
     ${SHAPE_INFER_INCLUDE_DIR}/segment_max_shape_inference.hpp
     ${SHAPE_INFER_INCLUDE_DIR}/select_shape_inference.hpp
+    ${SHAPE_INFER_INCLUDE_DIR}/selective_ssm_shape_inference.hpp
     ${SHAPE_INFER_INCLUDE_DIR}/sequence_generator.hpp
     ${SHAPE_INFER_INCLUDE_DIR}/shape_infer_type_utils.hpp
     ${SHAPE_INFER_INCLUDE_DIR}/shape_nodes.hpp

src/core/shape_inference/include/paged_selective_ssm_shape_inference.hpp

@@ -0,0 +1,156 @@
+// Copyright (C) 2018-2026 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include "openvino/op/paged_selective_ssm.hpp"
+#include "utils.hpp"
+
+namespace ov::op::internal {
+
+template <class T, class TRShape = result_shape_t<T>>
+std::vector<TRShape> shape_infer(const PagedSelectiveSSM* op, const std::vector<T>& input_shapes) {
+    NODE_VALIDATION_CHECK(op, input_shapes.size() == 11);
+
+    NODE_SHAPE_INFER_CHECK(op, input_shapes, input_shapes[0].rank().compatible(1));
+    NODE_SHAPE_INFER_CHECK(op, input_shapes, input_shapes[1].rank().compatible(2));
+    NODE_SHAPE_INFER_CHECK(op, input_shapes, input_shapes[2].rank().compatible(3));
+    NODE_SHAPE_INFER_CHECK(op, input_shapes, input_shapes[3].rank().compatible(3));
+    NODE_SHAPE_INFER_CHECK(op, input_shapes, input_shapes[4].rank().compatible(3));
+    NODE_SHAPE_INFER_CHECK(op, input_shapes, input_shapes[5].rank().compatible(4));
+    for (size_t i = 6; i < 11; ++i) {
+        NODE_SHAPE_INFER_CHECK(op, input_shapes, input_shapes[i].rank().compatible(1));
+    }
+
+    const auto& A_ps = input_shapes[0];
+    const auto& dt_ps = input_shapes[1];
+    const auto& B_ps = input_shapes[2];
+    const auto& x_ps = input_shapes[3];
+    const auto& C_ps = input_shapes[4];
+    const auto& state_ps = input_shapes[5];
+    const auto& subsequence_begins_ps = input_shapes[6];
+    const auto& la_block_indices_ps = input_shapes[7];
+    const auto& la_block_indices_begins_ps = input_shapes[8];
+    const auto& num_processed_tokens_ps = input_shapes[9];
+    const auto& cache_interval_ps = input_shapes[10];
+
+    const auto A_rank_is_static = A_ps.rank().is_static();
+    const auto dt_rank_is_static = dt_ps.rank().is_static();
+    const auto B_rank_is_static = B_ps.rank().is_static();
+    const auto x_rank_is_static = x_ps.rank().is_static();
+    const auto C_rank_is_static = C_ps.rank().is_static();
+    const auto state_rank_is_static = state_ps.rank().is_static();
+    const auto subsequence_begins_rank_is_static = subsequence_begins_ps.rank().is_static();
+    const auto la_block_indices_rank_is_static = la_block_indices_ps.rank().is_static();
+    const auto la_block_indices_begins_rank_is_static = la_block_indices_begins_ps.rank().is_static();
+    const auto num_processed_tokens_rank_is_static = num_processed_tokens_ps.rank().is_static();
+    const auto cache_interval_rank_is_static = cache_interval_ps.rank().is_static();
+
+    Dimension token_dim, num_heads_dim, head_dim_dim, state_size_dim;
+
+    bool token_ok = true;
+    if (x_rank_is_static)
+        token_ok &= Dimension::merge(token_dim, token_dim, x_ps[0]);
+    if (dt_rank_is_static)
+        token_ok &= Dimension::merge(token_dim, token_dim, dt_ps[0]);
+    if (B_rank_is_static)
+        token_ok &= Dimension::merge(token_dim, token_dim, B_ps[0]);
+    if (C_rank_is_static)
+        token_ok &= Dimension::merge(token_dim, token_dim, C_ps[0]);
+    NODE_SHAPE_INFER_CHECK(op,
+                           input_shapes,
+                           token_ok,
+                           "The token dimension of `dt`, `B`, `x` and `C` should be the same.");
+
+    bool num_heads_ok = true;
+    if (x_rank_is_static)
+        num_heads_ok &= Dimension::merge(num_heads_dim, num_heads_dim, x_ps[1]);
+    if (A_rank_is_static)
+        num_heads_ok &= Dimension::merge(num_heads_dim, num_heads_dim, A_ps[0]);
+    if (dt_rank_is_static)
+        num_heads_ok &= Dimension::merge(num_heads_dim, num_heads_dim, dt_ps[1]);
+    if (state_rank_is_static)
+        num_heads_ok &= Dimension::merge(num_heads_dim, num_heads_dim, state_ps[1]);
+    NODE_SHAPE_INFER_CHECK(op,
+                           input_shapes,
+                           num_heads_ok,
+                           "The number of heads of `A`, `dt`, `x` and `recurrent_state_table` should be the same.");
+
+    if (B_rank_is_static && C_rank_is_static) {
+        NODE_SHAPE_INFER_CHECK(op,
+                               input_shapes,
+                               C_ps[1].compatible(B_ps[1]),
+                               "The number of groups of `B` and `C` should be the same.");
+    }
+
+    bool head_dim_ok = true;
+    if (x_rank_is_static)
+        head_dim_ok &= Dimension::merge(head_dim_dim, head_dim_dim, x_ps[2]);
+    if (state_rank_is_static)
+        head_dim_ok &= Dimension::merge(head_dim_dim, head_dim_dim, state_ps[2]);
+    NODE_SHAPE_INFER_CHECK(op,
+                           input_shapes,
+                           head_dim_ok,
+                           "The head dimension of `x` and `recurrent_state_table` should be the same.");
+
+    bool state_size_ok = true;
+    if (state_rank_is_static)
+        state_size_ok &= Dimension::merge(state_size_dim, state_size_dim, state_ps[3]);
+    if (B_rank_is_static)
+        state_size_ok &= Dimension::merge(state_size_dim, state_size_dim, B_ps[2]);
+    if (C_rank_is_static)
+        state_size_ok &= Dimension::merge(state_size_dim, state_size_dim, C_ps[2]);
+    NODE_SHAPE_INFER_CHECK(op,
+                           input_shapes,
+                           state_size_ok,
+                           "The state size of `B`, `C` and `recurrent_state_table` should be the same.");
+
+    if (num_heads_dim.is_static() && B_rank_is_static) {
+        const auto& num_groups = B_ps[1];
+        if (num_groups.is_static()) {
+            NODE_SHAPE_INFER_CHECK(
+                op,
+                input_shapes,
+                num_groups.get_length() != 0 && num_heads_dim.get_length() % num_groups.get_length() == 0,
+                "The number of heads should be divisible by the number of groups.");
+        }
+    }
+
+    if (la_block_indices_rank_is_static && state_rank_is_static) {
+        NODE_SHAPE_INFER_CHECK(op,
+                               input_shapes,
+                               la_block_indices_ps[0].compatible(state_ps[0]),
+                               "The number of blocks of `la_block_indices` and `recurrent_state_table` should "
+                               "be the same.");
+    }
+    if (subsequence_begins_rank_is_static && la_block_indices_begins_rank_is_static) {
+        NODE_SHAPE_INFER_CHECK(op,
+                               input_shapes,
+                               subsequence_begins_ps[0].compatible(la_block_indices_begins_ps[0]),
+                               "The number of sequences of `subsequence_begins` and `la_block_indices_begins` "
+                               "should be the same.");
+    }
+    if (num_processed_tokens_rank_is_static && cache_interval_rank_is_static) {
+        NODE_SHAPE_INFER_CHECK(op,
+                               input_shapes,
+                               num_processed_tokens_ps[0].compatible(cache_interval_ps[0]),
+                               "The number of sequences of `num_processed_tokens` and `cache_interval` should "
+                               "be the same.");
+    }
+    if (subsequence_begins_rank_is_static && num_processed_tokens_rank_is_static) {
+        NODE_SHAPE_INFER_CHECK(op,
+                               input_shapes,
+                               (subsequence_begins_ps[0] - Dimension(1)).compatible(num_processed_tokens_ps[0]),
+                               "The number of sequences of `subsequence_begins` and `num_processed_tokens` "
+                               "should be the same.");
+    }
+
+    auto output_shapes = std::vector<TRShape>{x_ps};
+    if (x_rank_is_static) {
+        output_shapes[0] = TRShape{token_dim, num_heads_dim, head_dim_dim};
+    }
+    return output_shapes;
+}
+
+}  // namespace ov::op::internal