Skip to content

Commit 29d255c

Browse files
mmikolajczallnes
mmikolajcz
and
allnes
committed
Adress review feedback:
- Unify the element type across all float inputs. The recurrent state is a regular input rather than an in-place cache, so both SelectiveSSM outputs use the common type. - Restrict float inputs to f32, f16 and bf16, matching the other recurrent ops. - Merge the metadata element type across the paged scheduling inputs. - Make shape inference generic over the dimension type so it instantiates under StaticShape. - Merge the group dimension from B and C before validating non-zero and divisibility, and reject a zero state size. - Drop the la_block_indices / recurrent_state_table block count check. Logical slots and physical rows are independent under aliasing. - Document the paged cache protocol: input slot, write-slot formula, alias rules, required slot count and cross-sequence ownership. - Correct the outer product order to x_t (x) dtB_t. Co-authored-by: Alexander Nesterov <alexander.nesterov@intel.com>
1 parent ce291f6 commit 29d255c

8 files changed

Lines changed: 419 additions & 255 deletions

File tree

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

Lines changed: 22 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -13,17 +13,31 @@ namespace ov::op::internal {
1313
/// Paged variant of the SelectiveSSM (Mamba2 selective state-space, arXiv:2405.21060) recurrence. Processes
1414
/// tokens from multiple sequences packed into a single batch and manages the recurrent SSM state (a per-head
1515
/// ``[head_dim, state_size]`` matrix) using a paged block table, enabling non-contiguous memory allocation
16-
/// across sequences. Per token: ``dA = exp(A * dt)``, ``dtB = dt * B``, ``dBx = dtB (x) x``,
16+
/// across sequences. Per token: ``dA = exp(A * dt)``, ``dtB = dt * B``, ``dBx = x (x) dtB``,
1717
/// ``state = state * dA + dBx``, ``output = sum(state * C)``.
1818
///
1919
/// ``B`` and ``C`` are grouped and shared across heads: each head ``h`` reads group
2020
/// ``g = h / heads_per_group``, where ``heads_per_group = num_heads / num_groups``.
2121
///
22-
/// ``recurrent_state_table`` is updated in place. For sequence ``s``, the state is cached every
23-
/// ``cache_interval[s]`` tokens into the blocks addressed by
24-
/// ``la_block_indices[la_block_indices_begins[s] : la_block_indices_begins[s+1]]``;
25-
/// ``cache_interval[s] <= 0`` disables caching for that sequence. ``num_processed_tokens[s]`` gives the
26-
/// count of previously processed tokens, used to resume from the correct cached state and block offset.
22+
/// ``recurrent_state_table`` is updated in place through the logical slots of sequence ``s``, that is
23+
/// ``la_block_indices[la_block_indices_begins[s] : la_block_indices_begins[s+1]]``. Slot 0 is the input slot
24+
/// holding the state after exactly ``num_processed_tokens[s]`` tokens; it is read before any slot is written.
25+
/// Let ``current = subsequence_begins[s+1] - subsequence_begins[s]``. When ``cache_interval[s] > 0``, with
26+
/// ``interval = cache_interval[s]`` and ``past = num_processed_tokens[s] % interval``, the state is written in
27+
/// order to slots ``1 .. write_count``, where ``write_count = (past + current - 1) / interval + 1``: once each
28+
/// time ``past + t`` reaches a multiple of ``interval`` for the ``t``-th token of the call, and once after the
29+
/// last token unless that token already is such a boundary. The sequence therefore needs at least
30+
/// ``write_count + 1`` slots; surplus slots and unreferenced table rows are ignored. ``cache_interval[s] <= 0``
31+
/// disables caching: only the input slot is required and the table is left unmodified. A sequence with no
32+
/// tokens reads and writes nothing and needs no slots.
33+
///
34+
/// Slot 0 may alias slot 1 to update the state in place; the input state is read first, so the result matches
35+
/// a non-aliased copy. Across sequences the write set must be disjoint from every other sequence's read and
36+
/// write sets, while a read-only slot may be shared. The caller owns slot 0: it must hold a valid state before
37+
/// execution, including when ``num_processed_tokens[s]`` is 0, and zero-initialization and recycled-page
38+
/// synchronization are caller responsibilities. Metadata values are trusted: both begins arrays start at 0, are
39+
/// non-decreasing and end at the token and logical-slot counts, ``num_processed_tokens`` is non-negative, and
40+
/// every block index is below ``num_physical_blocks``.
2741
/// \ingroup ov_ops_cpp_api
2842
class OPENVINO_API PagedSelectiveSSM : public ov::op::Op {
2943
public:
@@ -38,11 +52,11 @@ class OPENVINO_API PagedSelectiveSSM : public ov::op::Op {
3852
/// \param x Input hidden states [batch_size_in_tokens, num_heads, head_dim].
3953
/// \param C Grouped output projection [batch_size_in_tokens, num_groups, state_size].
4054
/// \param recurrent_state_table Paged table of recurrent state snapshots, updated in place
41-
/// [num_blocks, num_heads, head_dim, state_size]; an all-zeros tensor before any tokens are cached.
55+
/// [num_physical_blocks, num_heads, head_dim, state_size]; an all-zeros tensor before any tokens are cached.
4256
/// \param subsequence_begins Start indices of each sequence's tokens in the flattened token batch
4357
/// [batch_size_in_sequences + 1], element type i32 or i64.
4458
/// \param la_block_indices Physical block row indices into recurrent_state_table, concatenated across
45-
/// all sequences [num_blocks], element type i32 or i64.
59+
/// all sequences [num_logical_blocks], element type i32 or i64.
4660
/// \param la_block_indices_begins Splits la_block_indices among sequences
4761
/// [batch_size_in_sequences + 1], element type i32 or i64.
4862
/// \param num_processed_tokens Number of tokens already processed for each sequence

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ namespace ov::op::internal {
1212
///
1313
/// Discretizes ``A`` (log-decay rates) and ``B`` (input projection) with ``dt`` (time steps) ahead of the
1414
/// recurrence: ``dA = exp(A * dt)``, ``dtB = dt * B``. Then, for each token ``t``:
15-
/// ``dBx_t = dtB_t (x) x_t``, ``state_t = state_{t-1} * dA_t + dBx_t``, ``y_t = sum(state_t * C_t)``.
15+
/// ``dBx_t = x_t (x) dtB_t``, ``state_t = state_{t-1} * dA_t + dBx_t``, ``y_t = sum(state_t * C_t)``.
1616
///
1717
/// ``B`` and ``C`` are grouped and shared across heads: each head ``h`` reads group
1818
/// ``g = h / heads_per_group``, where ``heads_per_group = num_heads / num_groups``.

‎src/core/shape_inference/include/paged_selective_ssm_shape_inference.hpp‎

Lines changed: 98 additions & 92 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,16 @@
99

1010
namespace ov::op::internal {
1111

12+
template <typename TDim>
13+
bool merge_paged_selective_ssm_dim(TDim& destination, bool& initialized, const TDim& source) {
14+
if (!initialized) {
15+
destination = source;
16+
initialized = true;
17+
return true;
18+
}
19+
return TDim::merge(destination, destination, source);
20+
}
21+
1222
template <class T, class TRShape = result_shape_t<T>>
1323
std::vector<TRShape> shape_infer(const PagedSelectiveSSM* op, const std::vector<T>& input_shapes) {
1424
NODE_VALIDATION_CHECK(op, input_shapes.size() == 11);
@@ -19,136 +29,132 @@ std::vector<TRShape> shape_infer(const PagedSelectiveSSM* op, const std::vector<
1929
NODE_SHAPE_INFER_CHECK(op, input_shapes, input_shapes[3].rank().compatible(3));
2030
NODE_SHAPE_INFER_CHECK(op, input_shapes, input_shapes[4].rank().compatible(3));
2131
NODE_SHAPE_INFER_CHECK(op, input_shapes, input_shapes[5].rank().compatible(4));
22-
for (size_t i = 6; i < 11; ++i) {
23-
NODE_SHAPE_INFER_CHECK(op, input_shapes, input_shapes[i].rank().compatible(1));
32+
for (size_t input = 6; input < 11; ++input) {
33+
NODE_SHAPE_INFER_CHECK(op, input_shapes, input_shapes[input].rank().compatible(1));
2434
}
2535

26-
const auto& A_ps = input_shapes[0];
27-
const auto& dt_ps = input_shapes[1];
28-
const auto& B_ps = input_shapes[2];
29-
const auto& x_ps = input_shapes[3];
30-
const auto& C_ps = input_shapes[4];
31-
const auto& state_ps = input_shapes[5];
32-
const auto& subsequence_begins_ps = input_shapes[6];
33-
const auto& la_block_indices_ps = input_shapes[7];
34-
const auto& la_block_indices_begins_ps = input_shapes[8];
35-
const auto& num_processed_tokens_ps = input_shapes[9];
36-
const auto& cache_interval_ps = input_shapes[10];
37-
38-
const auto A_rank_is_static = A_ps.rank().is_static();
39-
const auto dt_rank_is_static = dt_ps.rank().is_static();
40-
const auto B_rank_is_static = B_ps.rank().is_static();
41-
const auto x_rank_is_static = x_ps.rank().is_static();
42-
const auto C_rank_is_static = C_ps.rank().is_static();
43-
const auto state_rank_is_static = state_ps.rank().is_static();
44-
const auto subsequence_begins_rank_is_static = subsequence_begins_ps.rank().is_static();
45-
const auto la_block_indices_rank_is_static = la_block_indices_ps.rank().is_static();
46-
const auto la_block_indices_begins_rank_is_static = la_block_indices_begins_ps.rank().is_static();
47-
const auto num_processed_tokens_rank_is_static = num_processed_tokens_ps.rank().is_static();
48-
const auto cache_interval_rank_is_static = cache_interval_ps.rank().is_static();
49-
50-
Dimension token_dim, num_heads_dim, head_dim_dim, state_size_dim;
51-
36+
const auto& A_shape = input_shapes[0];
37+
const auto& dt_shape = input_shapes[1];
38+
const auto& B_shape = input_shapes[2];
39+
const auto& x_shape = input_shapes[3];
40+
const auto& C_shape = input_shapes[4];
41+
const auto& state_shape = input_shapes[5];
42+
const auto& subsequence_shape = input_shapes[6];
43+
const auto& block_indices_begins_shape = input_shapes[8];
44+
const auto& processed_shape = input_shapes[9];
45+
const auto& interval_shape = input_shapes[10];
46+
using DimType = typename T::value_type;
47+
48+
DimType token_dim{};
49+
DimType heads_dim{};
50+
DimType head_dim{};
51+
DimType groups_dim{};
52+
DimType state_size_dim{};
53+
54+
bool token_initialized = false;
5255
bool token_ok = true;
53-
if (x_rank_is_static)
54-
token_ok &= Dimension::merge(token_dim, token_dim, x_ps[0]);
55-
if (dt_rank_is_static)
56-
token_ok &= Dimension::merge(token_dim, token_dim, dt_ps[0]);
57-
if (B_rank_is_static)
58-
token_ok &= Dimension::merge(token_dim, token_dim, B_ps[0]);
59-
if (C_rank_is_static)
60-
token_ok &= Dimension::merge(token_dim, token_dim, C_ps[0]);
56+
if (x_shape.rank().is_static())
57+
token_ok &= merge_paged_selective_ssm_dim(token_dim, token_initialized, x_shape[0]);
58+
if (dt_shape.rank().is_static())
59+
token_ok &= merge_paged_selective_ssm_dim(token_dim, token_initialized, dt_shape[0]);
60+
if (B_shape.rank().is_static())
61+
token_ok &= merge_paged_selective_ssm_dim(token_dim, token_initialized, B_shape[0]);
62+
if (C_shape.rank().is_static())
63+
token_ok &= merge_paged_selective_ssm_dim(token_dim, token_initialized, C_shape[0]);
6164
NODE_SHAPE_INFER_CHECK(op,
6265
input_shapes,
6366
token_ok,
6467
"The token dimension of `dt`, `B`, `x` and `C` should be the same.");
6568

66-
bool num_heads_ok = true;
67-
if (x_rank_is_static)
68-
num_heads_ok &= Dimension::merge(num_heads_dim, num_heads_dim, x_ps[1]);
69-
if (A_rank_is_static)
70-
num_heads_ok &= Dimension::merge(num_heads_dim, num_heads_dim, A_ps[0]);
71-
if (dt_rank_is_static)
72-
num_heads_ok &= Dimension::merge(num_heads_dim, num_heads_dim, dt_ps[1]);
73-
if (state_rank_is_static)
74-
num_heads_ok &= Dimension::merge(num_heads_dim, num_heads_dim, state_ps[1]);
69+
bool heads_initialized = false;
70+
bool heads_ok = true;
71+
if (x_shape.rank().is_static())
72+
heads_ok &= merge_paged_selective_ssm_dim(heads_dim, heads_initialized, x_shape[1]);
73+
if (A_shape.rank().is_static())
74+
heads_ok &= merge_paged_selective_ssm_dim(heads_dim, heads_initialized, A_shape[0]);
75+
if (dt_shape.rank().is_static())
76+
heads_ok &= merge_paged_selective_ssm_dim(heads_dim, heads_initialized, dt_shape[1]);
77+
if (state_shape.rank().is_static())
78+
heads_ok &= merge_paged_selective_ssm_dim(heads_dim, heads_initialized, state_shape[1]);
7579
NODE_SHAPE_INFER_CHECK(op,
7680
input_shapes,
77-
num_heads_ok,
81+
heads_ok,
7882
"The number of heads of `A`, `dt`, `x` and `recurrent_state_table` should be the same.");
7983

80-
if (B_rank_is_static && C_rank_is_static) {
81-
NODE_SHAPE_INFER_CHECK(op,
82-
input_shapes,
83-
C_ps[1].compatible(B_ps[1]),
84-
"The number of groups of `B` and `C` should be the same.");
85-
}
86-
84+
bool head_dim_initialized = false;
8785
bool head_dim_ok = true;
88-
if (x_rank_is_static)
89-
head_dim_ok &= Dimension::merge(head_dim_dim, head_dim_dim, x_ps[2]);
90-
if (state_rank_is_static)
91-
head_dim_ok &= Dimension::merge(head_dim_dim, head_dim_dim, state_ps[2]);
86+
if (x_shape.rank().is_static())
87+
head_dim_ok &= merge_paged_selective_ssm_dim(head_dim, head_dim_initialized, x_shape[2]);
88+
if (state_shape.rank().is_static())
89+
head_dim_ok &= merge_paged_selective_ssm_dim(head_dim, head_dim_initialized, state_shape[2]);
9290
NODE_SHAPE_INFER_CHECK(op,
9391
input_shapes,
9492
head_dim_ok,
9593
"The head dimension of `x` and `recurrent_state_table` should be the same.");
9694

95+
bool groups_initialized = false;
96+
bool groups_ok = true;
97+
if (B_shape.rank().is_static())
98+
groups_ok &= merge_paged_selective_ssm_dim(groups_dim, groups_initialized, B_shape[1]);
99+
if (C_shape.rank().is_static())
100+
groups_ok &= merge_paged_selective_ssm_dim(groups_dim, groups_initialized, C_shape[1]);
101+
NODE_SHAPE_INFER_CHECK(op, input_shapes, groups_ok, "The number of groups of `B` and `C` should be the same.");
102+
NODE_SHAPE_INFER_CHECK(op,
103+
input_shapes,
104+
groups_dim.is_dynamic() || groups_dim.get_length() > 0,
105+
"The number of groups must be greater than zero.");
106+
107+
bool state_size_initialized = false;
97108
bool state_size_ok = true;
98-
if (state_rank_is_static)
99-
state_size_ok &= Dimension::merge(state_size_dim, state_size_dim, state_ps[3]);
100-
if (B_rank_is_static)
101-
state_size_ok &= Dimension::merge(state_size_dim, state_size_dim, B_ps[2]);
102-
if (C_rank_is_static)
103-
state_size_ok &= Dimension::merge(state_size_dim, state_size_dim, C_ps[2]);
109+
if (state_shape.rank().is_static())
110+
state_size_ok &= merge_paged_selective_ssm_dim(state_size_dim, state_size_initialized, state_shape[3]);
111+
if (B_shape.rank().is_static())
112+
state_size_ok &= merge_paged_selective_ssm_dim(state_size_dim, state_size_initialized, B_shape[2]);
113+
if (C_shape.rank().is_static())
114+
state_size_ok &= merge_paged_selective_ssm_dim(state_size_dim, state_size_initialized, C_shape[2]);
104115
NODE_SHAPE_INFER_CHECK(op,
105116
input_shapes,
106117
state_size_ok,
107118
"The state size of `B`, `C` and `recurrent_state_table` should be the same.");
119+
NODE_SHAPE_INFER_CHECK(op,
120+
input_shapes,
121+
state_size_dim.is_dynamic() || state_size_dim.get_length() > 0,
122+
"The state size must be greater than zero.");
108123

109-
if (num_heads_dim.is_static() && B_rank_is_static) {
110-
const auto& num_groups = B_ps[1];
111-
if (num_groups.is_static()) {
112-
NODE_SHAPE_INFER_CHECK(
113-
op,
114-
input_shapes,
115-
num_groups.get_length() != 0 && num_heads_dim.get_length() % num_groups.get_length() == 0,
116-
"The number of heads should be divisible by the number of groups.");
117-
}
124+
if (heads_dim.is_static() && groups_dim.is_static()) {
125+
NODE_SHAPE_INFER_CHECK(op,
126+
input_shapes,
127+
groups_dim.get_length() != 0 && heads_dim.get_length() % groups_dim.get_length() == 0,
128+
"The number of heads should be divisible by the number of groups.");
118129
}
119130

120-
if (la_block_indices_rank_is_static && state_rank_is_static) {
131+
if (subsequence_shape.rank().is_static() && block_indices_begins_shape.rank().is_static()) {
121132
NODE_SHAPE_INFER_CHECK(op,
122133
input_shapes,
123-
la_block_indices_ps[0].compatible(state_ps[0]),
124-
"The number of blocks of `la_block_indices` and `recurrent_state_table` should "
125-
"be the same.");
134+
subsequence_shape[0].compatible(block_indices_begins_shape[0]),
135+
"The sizes of `subsequence_begins` and `la_block_indices_begins` should be the same.");
126136
}
127-
if (subsequence_begins_rank_is_static && la_block_indices_begins_rank_is_static) {
137+
if (processed_shape.rank().is_static() && interval_shape.rank().is_static()) {
128138
NODE_SHAPE_INFER_CHECK(op,
129139
input_shapes,
130-
subsequence_begins_ps[0].compatible(la_block_indices_begins_ps[0]),
131-
"The number of sequences of `subsequence_begins` and `la_block_indices_begins` "
132-
"should be the same.");
140+
processed_shape[0].compatible(interval_shape[0]),
141+
"The sizes of `num_processed_tokens` and `cache_interval` should be the same.");
133142
}
134-
if (num_processed_tokens_rank_is_static && cache_interval_rank_is_static) {
143+
if (subsequence_shape.rank().is_static() && processed_shape.rank().is_static()) {
135144
NODE_SHAPE_INFER_CHECK(op,
136145
input_shapes,
137-
num_processed_tokens_ps[0].compatible(cache_interval_ps[0]),
138-
"The number of sequences of `num_processed_tokens` and `cache_interval` should "
139-
"be the same.");
140-
}
141-
if (subsequence_begins_rank_is_static && num_processed_tokens_rank_is_static) {
146+
subsequence_shape[0].is_dynamic() || subsequence_shape[0].get_length() >= 1,
147+
"The size of `subsequence_begins` must be at least one.");
142148
NODE_SHAPE_INFER_CHECK(op,
143149
input_shapes,
144-
(subsequence_begins_ps[0] - Dimension(1)).compatible(num_processed_tokens_ps[0]),
145-
"The number of sequences of `subsequence_begins` and `num_processed_tokens` "
146-
"should be the same.");
150+
(subsequence_shape[0] - DimType(1)).compatible(processed_shape[0]),
151+
"The size of `subsequence_begins` should be one larger than "
152+
"`num_processed_tokens`.");
147153
}
148154

149-
auto output_shapes = std::vector<TRShape>{x_ps};
150-
if (x_rank_is_static) {
151-
output_shapes[0] = TRShape{token_dim, num_heads_dim, head_dim_dim};
155+
auto output_shapes = std::vector<TRShape>{x_shape};
156+
if (x_shape.rank().is_static()) {
157+
output_shapes[0] = TRShape{token_dim, heads_dim, head_dim};
152158
}
153159
return output_shapes;
154160
}

0 commit comments

Comments
 (0)