diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml
index dca6a88b093d8..e5fb682992f2a 100644
--- a/.github/workflows/lint.yml
+++ b/.github/workflows/lint.yml
@@ -52,8 +52,11 @@ jobs:
- name: Setup Python
uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6.3.0
with:
- # Use the version configured in target-version of [tool.black] section in pyproject.toml.
- python-version: "3.10"
+ # Use a version pre-installed in the runner pool's read-only tool cache; requesting an
+ # uncached version makes setup-python fail trying to write to /opt/hostedtoolcache.
+ # This is only the interpreter lintrunner runs on. The Python syntax the linters target
+ # is set by target-version under [tool.ruff] in pyproject.toml, independently of this.
+ python-version: "3.12"
- name: Setup Rust
uses: actions-rs/toolchain@16499b5e05bf2e26879000db0c1d13f7e13fa3af # v1.0.7
with:
diff --git a/.github/workflows/pr_checks.yml b/.github/workflows/pr_checks.yml
index abb3ca86596b2..c7b755150233d 100644
--- a/.github/workflows/pr_checks.yml
+++ b/.github/workflows/pr_checks.yml
@@ -33,7 +33,8 @@ jobs:
- name: Setup Python
uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6.3.0
with:
- python-version: "3.10"
+ # Keep in sync with lint.yml; must be a version cached on the runner pool.
+ python-version: "3.12"
- name: Setup Rust
uses: actions-rs/toolchain@16499b5e05bf2e26879000db0c1d13f7e13fa3af # v1.0.7
with:
diff --git a/cmake/onnxruntime_providers_coreml.cmake b/cmake/onnxruntime_providers_coreml.cmake
index bf46a73e43839..f3425b9272102 100644
--- a/cmake/onnxruntime_providers_coreml.cmake
+++ b/cmake/onnxruntime_providers_coreml.cmake
@@ -13,9 +13,9 @@ add_compile_definitions(COREML_ENABLE_MLPROGRAM=1)
if(LINUX)
find_library(LibUUID_LIBRARY NAMES uuid)
find_path(LibUUID_INCLUDE_DIR NAMES uuid/uuid.h)
- if (NOT LibUUID_INCLUDE_DIR)
- message(FATAL "uuid/uuid.h was not found as is required for ML Program support. "
- "Run `sudo apt install uuid-dev` if you need to test ML Program related CoreML EP code. ")
+ if (NOT LibUUID_INCLUDE_DIR OR NOT LibUUID_LIBRARY)
+ message(FATAL_ERROR "libuuid (uuid/uuid.h) was not found and is required for ML Program support. "
+ "Run `sudo apt install uuid-dev`, or build with `--use_vcpkg` so the libuuid port is used. ")
endif()
endif()
@@ -194,7 +194,8 @@ target_include_directories(onnxruntime_providers_coreml PRIVATE
)
if (LINUX)
- target_link_libraries(onnxruntime_providers_coreml PRIVATE uuid)
+ target_include_directories(onnxruntime_providers_coreml PRIVATE ${LibUUID_INCLUDE_DIR})
+ target_link_libraries(onnxruntime_providers_coreml PRIVATE ${LibUUID_LIBRARY})
endif()
diff --git a/cmake/vcpkg.json b/cmake/vcpkg.json
index 7c58604b9c97a..429e07aafa772 100644
--- a/cmake/vcpkg.json
+++ b/cmake/vcpkg.json
@@ -86,7 +86,13 @@
},
"coreml-ep": {
"description": "Build with CoreML EP",
- "dependencies": ["fp16"]
+ "dependencies": [
+ "fp16",
+ {
+ "name": "libuuid",
+ "platform": "linux"
+ }
+ ]
},
"dml-ep": {
"description": "Build with DirectML EP",
diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md
index 5b23a58c865c2..308546d783298 100644
--- a/docs/ContribOperators.md
+++ b/docs/ContribOperators.md
@@ -31,6 +31,7 @@ Do not modify directly.*
* com.microsoft.DynamicTimeWarping
* com.microsoft.EPContext
* com.microsoft.EmbedLayerNormalization
+ * com.microsoft.EngramGate
* com.microsoft.ExpandDims
* com.microsoft.FastGelu
* com.microsoft.FusedConv
@@ -71,6 +72,7 @@ Do not modify directly.*
* com.microsoft.MulInteger
* com.microsoft.MultiHeadAttention
* com.microsoft.MurmurHash3
+ * com.microsoft.NGramHashMapping
* com.microsoft.NGramRepeatBlock
* com.microsoft.NhwcConv
* com.microsoft.NhwcFusedConv
@@ -930,6 +932,23 @@ This version of the operator has been available since version 1 of the 'com.micr
enforced on the last spatial dimension only.
The optional activation attribute supports fused SiLU/Swish activation.
+
+ The dilation attribute spaces the kernel taps along the causal axis: output position t reads
+ input positions t - (k_1 - 1 - j) * dilation for tap j. The receptive field therefore spans
+ (k_1 - 1) * dilation positions before the current one, and the carry state grows to match:
+ past_state and present_state hold (k_1 - 1) * dilation positions instead of k_1 - 1. Dilation 1
+ (the default) is the undilated case and keeps the original state length, so models exported
+ before the attribute existed are unaffected.
+
+ The channels_last attribute selects a sequence-major layout for the activations and the carry
+ state, so a model that already produces channels-last activations does not have to transpose into
+ and out of the channels-first layout. With channels_last = 1 and ndim = 1, input and output are
+ (batch_size, sequence_length, d_1, ..., d_n) and the state tensors are
+ (batch_size, state_length, d_1, ..., d_n), where channels = d_1 * ... * d_n. Any number of trailing
+ channel axes is accepted, so an activation that keeps hyper-connections and hidden size as separate
+ axes needs no reshape either. weight and bias keep their channels-first (channels, 1, k_1) and
+ (channels) shapes because they have no sequence axis. The computed values are identical to the
+ channels-first layout; only the memory layout differs.
#### Version
@@ -940,23 +959,27 @@ This version of the operator has been available since version 1 of the 'com.micr
- activation : string
- Fused activation function. One of: 'silu', 'swish', 'none'. Default is 'none'.
+- channels_last : int
+- When 1, input, output, past_state and present_state use a sequence-major, channels-last layout: input and output are (batch_size, sequence_length, d_1, ..., d_n) and the state tensors are (batch_size, state_length, d_1, ..., d_n), where channels = d_1 * ... * d_n. weight and bias keep their channels-first shapes. Requires ndim = 1. Default is 0 (channels-first).
+- dilation : int
+- Spacing between kernel taps along the causal (last spatial) axis. The receptive field spans (k_1 - 1) * dilation positions before the current one, and past_state / present_state hold that many positions. Must be >= 1. Default is 1 (undilated).
- ndim : int
- Spatial dimensionality: 1, 2, or 3. Default is 1.
- state_window : int
-- Number of trailing per-position carry states held by past_state and present_state. When 0 (default) the state tensors have no window axis and hold only the state after the last position, i.e. the backward-compatible (batch_size, channels, k_1 - 1). When W > 0 both gain a LEADING axis of extent W, right-aligned: slot j is the state after position (seq_len - W + j), so slot W-1 is always the state after the last position (identical to the W = 0 tensor) and is the slot past_state is read from. The window axis leads the batch axis so that each slot is one contiguous (batch_size, channels, k_1 - 1) block. Slots below max(0, W - seq_len) hold no position from this call and are filled with zeros. A window lets a speculative decoder roll the state back to an accepted prefix without replaying the forward. Valid range is [0, 8].
+- Number of trailing per-position carry states held by past_state and present_state. When 0 (default) the state tensors have no window axis and hold only the state after the last position, i.e. the backward-compatible (batch_size, channels, state_length) where state_length = (k_1 - 1) * dilation. When W > 0 both gain a LEADING axis of extent W, right-aligned: slot j is the state after position (seq_len - W + j), so slot W-1 is always the state after the last position (identical to the W = 0 tensor) and is the slot past_state is read from. The window axis leads the batch axis so that each slot is one contiguous (batch_size, channels, state_length) block. Slots below max(0, W - seq_len) hold no position from this call and are filled with zeros. A window lets a speculative decoder roll the state back to an accepted prefix without replaying the forward. Valid range is [0, 8].
#### Inputs (2 - 4)
- input : T
-- Input tensor with shape (batch_size, channels, ...). Channels-first layout. Spatial dims: 1D: (L,); 2D: (H, W); 3D: (D, H, W).
+- Input tensor with shape (batch_size, channels, ...) in the default channels-first layout. Spatial dims: 1D: (L,); 2D: (H, W); 3D: (D, H, W). When channels_last = 1 the shape is (batch_size, sequence_length, d_1, ..., d_n) instead.
- weight : T
- Depthwise convolution kernel with shape (channels, 1, k_1, ...). Spatial kernel sizes: (k_1, ..., k_ndim).
- bias (optional) : T
- Optional per-channel bias with shape (channels).
- past_state (optional) : T
-- Carry state from previous step. For ndim=1: (batch_size, channels, k_1 - 1), or (W, batch_size, channels, k_1 - 1) when state_window = W > 0, in which case only slot W-1 is read. If not provided, padding is zero.
+- Carry state from previous step. For ndim=1: (batch_size, channels, state_length), or (W, batch_size, channels, state_length) when state_window = W > 0, in which case only slot W-1 is read, where state_length = (k_1 - 1) * dilation. When channels_last = 1 each slot is (batch_size, state_length, d_1, ..., d_n) instead. If not provided, padding is zero.
#### Outputs
@@ -965,7 +988,7 @@ This version of the operator has been available since version 1 of the 'com.micr
output : T
Convolution output with same shape as input.
present_state : T
-Updated carry state. For ndim=1: (batch_size, channels, k_1 - 1), or (W, batch_size, channels, k_1 - 1) when state_window = W > 0. Slot W-1 contains the last (k-1) values from the virtual input along the causal axis; slot j contains the same for the prefix ending at position (seq_len - W + j).
+Updated carry state. For ndim=1: (batch_size, channels, state_length), or (W, batch_size, channels, state_length) when state_window = W > 0, and (batch_size, state_length, d_1, ..., d_n) per slot when channels_last = 1. Slot W-1 contains the last state_length values from the virtual input along the causal axis; slot j contains the same for the prefix ending at position (seq_len - W + j).
#### Type Constraints
@@ -1772,6 +1795,66 @@ This version of the operator has been available since version 1 of the 'com.micr
+### **com.microsoft.EngramGate**
+
+ Fuses the Engram gate.
+
+ The op consumes already projected keys in (batch_size, sequence_length, hc_mult, hidden_size) layout,
+ the hidden-state queries in the same layout, an already projected value in
+ (batch_size, sequence_length, hidden_size) layout that is shared by every hyper-connection, and the two
+ RMSNorm scales. The key and value projections stay outside the op so they can run on the execution
+ provider's tuned MatMul (weight prepacking, tensor cores, quantized weights) and so the value
+ projection is computed once per token instead of once per hyper-connection.
+
+ It computes the Engram gate:
+
+ gate = sigmoid(sign(dot) * sqrt(max(abs(dot), 1e-6))) where
+ dot = sum(RMSNorm(key) * RMSNorm(query)) / sqrt(hidden_size).
+
+ The output is gate * value, broadcast across the hyper-connections. The final Engram residual
+ value + short_conv(value) is then expressed with RMSNorm, CausalConvWithState and Add.
+
+#### Version
+
+This version of the operator has been available since version 1 of the 'com.microsoft' operator set.
+
+#### Attributes
+
+
+- epsilon : float
+- Epsilon used by both RMS normalization steps. Default is 1e-5.
+
+
+#### Inputs
+
+
+- key : T
+- Projected Engram keys with shape (batch_size, sequence_length, hc_mult, hidden_size).
+- query : T
+- Hidden-state queries with shape (batch_size, sequence_length, hc_mult, hidden_size).
+- value : T
+- Projected Engram value shared by every hyper-connection, with shape (batch_size, sequence_length, hidden_size).
+- key_norm_scale : T
+- RMSNorm scale for keys with shape (hc_mult, hidden_size).
+- query_norm_scale : T
+- RMSNorm scale for queries with shape (hc_mult, hidden_size).
+
+
+#### Outputs
+
+
+- output : T
+- Gated value tensor with shape (batch_size, sequence_length, hc_mult, hidden_size).
+
+
+#### Type Constraints
+
+
+- T : tensor(float), tensor(float16), tensor(bfloat16)
+- Constrain input and output types to float tensors.
+
+
+
### **com.microsoft.ExpandDims**
ExpandDims echo operator.
@@ -4182,6 +4265,73 @@ This version of the operator has been available since version 1 of the 'com.micr
+### **com.microsoft.NGramHashMapping**
+
+ Computes Engram n-gram hash ids from pre-compressed tokenizer ids.
+
+ For n in [2, max_ngram_size], the op creates causal shifts of input_ids, padding positions before the
+ sequence with pad_id, and computes
+ mix = shifted_0 * multipliers[0] xor ... xor shifted_(n-1) * multipliers[n-1].
+ For every head of that n-gram order it emits mix modulo the corresponding head vocabulary size.
+ The output layout is (batch_size, sequence_length, (max_ngram_size - 1) * n_head_per_ngram), with
+ heads for n=2 first, then n=3, and so on.
+
+ An n-gram window reaches max_ngram_size - 1 positions before the current token. To keep the op causal
+ across invocations (chunked prefill or autoregressive decode), the optional past_ids input carries
+ those preceding ids and present_ids returns the ids to pass to the next call. Both have shape
+ (batch_size, max_ngram_size - 1) and are right-aligned, so the last slot is the most recent id.
+ Positions before the start of the whole sequence use pad_id. Running the op once over a full sequence
+ and running it over consecutive chunks while threading present_ids into past_ids produce identical
+ hash ids. When past_ids is omitted the missing history is pad_id, which matches a fresh sequence.
+ past_ids and present_ids may use the same allocation. Such in-place execution is transaction-safe
+ only when the whole operator call is unconditionally committed; a caller that may select a prefix or
+ roll back must preserve past_ids.
+
+#### Version
+
+This version of the operator has been available since version 1 of the 'com.microsoft' operator set.
+
+#### Attributes
+
+
+- max_ngram_size : int (required)
+- Maximum n-gram order. Must be at least 2.
+- n_head_per_ngram : int (required)
+- Number of hash heads emitted for each n-gram order.
+- pad_id : int (required)
+- Compressed tokenizer id used to pad causal shifts before the beginning of a sequence.
+
+
+#### Inputs (3 - 4)
+
+
+- input_ids : M
+- Compressed tokenizer ids with shape (batch_size, sequence_length).
+- multipliers : M
+- Per-shift hash multipliers with shape (max_ngram_size). Conventionally odd, but any value is accepted.
+- vocab_sizes : M
+- Per-output-head vocabulary sizes, conventionally prime, with shape ((max_ngram_size - 1) * n_head_per_ngram). Every entry must be strictly positive. The CPU implementation rejects a non-positive entry; GPU implementations guard the modulo to avoid a device-side division by zero and emit a hash id of 0 for that head.
+- past_ids (optional) : M
+- Optional compressed tokenizer ids for the max_ngram_size - 1 positions that precede this call, with shape (batch_size, max_ngram_size - 1). Right-aligned, so the last slot is the most recent id. If omitted the history is pad_id.
+
+
+#### Outputs (1 - 2)
+
+
+- hash_ids : M
+- Hash ids with shape (batch_size, sequence_length, (max_ngram_size - 1) * n_head_per_ngram).
+- present_ids (optional) : M
+- Trailing max_ngram_size - 1 ids of past_ids followed by input_ids, with shape (batch_size, max_ngram_size - 1). Feed this back as past_ids on the next call.
+
+
+#### Type Constraints
+
+
+- M : tensor(int32), tensor(int64)
+- Constrain ids, multipliers, vocabulary sizes, and output ids to integer tensors.
+
+
+
### **com.microsoft.NGramRepeatBlock**
Enforce no repetition of n-grams. Scores are set to `-inf` for tokens that form a repeated n-gram if added to the back of the input_ids.
@@ -7165,7 +7315,8 @@ This version of the operator has been available since version 1 of the 'com.micr
at least one token. weight has shape (channels, 1, kernel_size), and optional bias has shape
(channels). The convolution never reads across a sequence boundary.
- initial_state is required and has shape (batch_size, channels, kernel_size - 1). It contains
+ initial_state is required and has shape (batch_size, channels, state_length), where
+ state_length = (kernel_size - 1) * dilation. It contains
the committed raw activation samples immediately preceding this call. final_state has the same
shape and type and is fully written with the state after each sequence's final token. State
uses the activation type because it stores raw samples, not accumulated convolution values.
@@ -7187,6 +7338,14 @@ This version of the operator has been available since version 1 of the 'com.micr
This device-side containment is not a synchronous validation or rejection mechanism.
The optional activation attribute supports none, SiLU, and Swish.
+
+ The dilation attribute spaces the kernel taps along the sequence axis: local token t of a request
+ reads that request's local positions t - (kernel_size - 1 - j) * dilation for tap j, and positions
+ before the request's first token come from the carry state. The carry state therefore holds
+ state_length = (kernel_size - 1) * dilation positions per request instead of kernel_size - 1.
+ Dilation 1 (the default) is the undilated case and keeps the original state length, so models
+ exported before the attribute existed are unaffected. input and output are already token-major
+ (sequence-major, channels-last), so this op needs no separate layout attribute.
#### Version
@@ -7197,6 +7356,8 @@ This version of the operator has been available since version 1 of the 'com.micr
- activation : string
- Fused activation function. One of: 'silu', 'swish', 'none'. Default is 'none'.
+- dilation : int
+- Spacing between kernel taps along the sequence axis. The receptive field spans (kernel_size - 1) * dilation positions before the current token, and initial_state / final_state hold that many positions per request. Must be >= 1. Default is 1 (undilated).
- state_update_capacity : int
- Static number of compact contiguous-prefix transition values to expose per request. Valid range is [0, 8]. capture_count is required exactly when this is positive.
@@ -7213,7 +7374,7 @@ This version of the operator has been available since version 1 of the 'com.micr
bias (optional) : T
Optional per-channel bias with shape (channels). Because the following initial_state input is required, an omitted bias must still occupy this position as an empty input name so initial_state stays at input index 4.
initial_state : T
-Required committed carry state with shape (batch_size, channels, kernel_size - 1).
+Required committed carry state with shape (batch_size, channels, (kernel_size - 1) * dilation).
capture_count (optional) : M
Optional device int32 tensor with shape (batch_size). For each request, captures that many local tokens from the contiguous prefix, clamped to the sequence length and state_update_capacity. Required exactly when state_update_capacity is positive.
@@ -7224,7 +7385,7 @@ This version of the operator has been available since version 1 of the 'com.micr
output : T
Token-major convolution output with the same shape as input.
final_state : T
-Fully written state after each sequence's final token, with shape (batch_size, channels, kernel_size - 1).
+Fully written state after each sequence's final token, with shape (batch_size, channels, (kernel_size - 1) * dilation).
state_update (optional) : T
Optional compact transition values with shape (batch_size, state_update_capacity, channels). Inactive slots are zero.
diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md
index 96a6c76f63c09..a49a82a073b26 100644
--- a/docs/OperatorKernels.md
+++ b/docs/OperatorKernels.md
@@ -582,6 +582,7 @@ The **OpSet Version** column uses the following notation:
|DynamicQuantizeMatMul|*in* A:**T1**
*in* B:**T2**
*in* b_scale:**T1**
*in* b_zero_point:**T2**
*in* bias:**T1**
*out* Y:**T1**|1+|**T1** = tensor(float)
**T2** = tensor(int8), tensor(uint8)|
|DynamicTimeWarping|*in* input:**F**
*out* output:**I**|1+|**F** = tensor(float)
**I** = tensor(int32)|
|EmbedLayerNormalization|*in* input_ids:**T1**
*in* segment_ids:**T1**
*in* word_embedding:**T**
*in* position_embedding:**T**
*in* segment_embedding:**T**
*in* gamma:**T**
*in* beta:**T**
*in* mask:**T1**
*in* position_ids:**T1**
*out* output:**T**
*out* mask_index:**T1**
*out* embedding_sum:**T**|1+|**T** = tensor(float)|
+|EngramGate|*in* key:**T**
*in* query:**T**
*in* value:**T**
*in* key_norm_scale:**T**
*in* query_norm_scale:**T**
*out* output:**T**|1+|**T** = tensor(float), tensor(float16)|
|ExpandDims|*in* X:**T**
*in* axis:**tensor(int32)**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(string), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)
**axis** = tensor(int32)|
|FastGelu|*in* X:**T**
*in* bias:**T**
*out* Y:**T**|1+|**T** = tensor(float)|
|FusedConv|*in* X:**T**
*in* W:**T**
*in* B:**T**
*in* Z:**T**
*out* Y:**T**|1+|**T** = tensor(float)|
@@ -608,6 +609,7 @@ The **OpSet Version** column uses the following notation:
|MoE|*in* input:**T**
*in* router_probs:**T**
*in* fc1_experts_weights:**T**
*in* fc1_experts_bias:**T**
*in* fc2_experts_weights:**T**
*in* fc2_experts_bias:**T**
*in* fc3_experts_weights:**T**
*in* fc3_experts_bias:**T**
*out* output:**T**|1+|**T** = tensor(float)|
|MultiHeadAttention|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* bias:**T**
*in* key_padding_mask:**M**
*in* attention_bias:**T**
*in* past_key:**T**
*in* past_value:**T**
*in* past_sequence_length:**M**
*in* cache_indirection:**M**
*out* output:**T**
*out* present_key:**T**
*out* present_value:**T**
*out* qk:**QK**|1+|**T** = tensor(float)|
|MurmurHash3|*in* X:**T1**
*out* Y:**T2**|1+|**T1** = tensor(double), tensor(float), tensor(int32), tensor(int64), tensor(string), tensor(uint32), tensor(uint64)
**T2** = tensor(int32), tensor(uint32)|
+|NGramHashMapping|*in* input_ids:**M**
*in* multipliers:**M**
*in* vocab_sizes:**M**
*in* past_ids:**M**
*out* hash_ids:**M**
*out* present_ids:**M**|1+|**M** = tensor(int32), tensor(int64)|
|NGramRepeatBlock|*in* input_ids:**Tid**
*in* scores:**T**
*out* scores_out:**T**|1+|**T** = tensor(float)
**Tid** = tensor(int64)|
|NhwcMaxPool|*in* x:**T**
*out* y:**T**|1+|**T** = tensor(int8), tensor(uint8)|
|Pad|*in* data:**T**
*in* pads:**tensor(int64)**
*in* value:**T**
*out* output:**T**|1+|**T** = tensor(float)|
@@ -1087,6 +1089,7 @@ The **OpSet Version** column uses the following notation:
|DequantizeWithOrder|*in* input:**Q**
*in* scale_input:**S**
*out* output:**F**|1+|**F** = tensor(float), tensor(float16)
**Q** = tensor(int8)
**S** = tensor(float)|
|DynamicTimeWarping|*in* input:**F**
*out* output:**I**|1+|**F** = tensor(float)
**I** = tensor(int32)|
|EmbedLayerNormalization|*in* input_ids:**T1**
*in* segment_ids:**T1**
*in* word_embedding:**T**
*in* position_embedding:**T**
*in* segment_embedding:**T**
*in* gamma:**T**
*in* beta:**T**
*in* mask:**T1**
*in* position_ids:**T1**
*out* output:**T**
*out* mask_index:**T1**
*out* embedding_sum:**T**|1+|**T** = tensor(float), tensor(float16)|
+|EngramGate|*in* key:**T**
*in* query:**T**
*in* value:**T**
*in* key_norm_scale:**T**
*in* query_norm_scale:**T**
*out* output:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)|
|FastGelu|*in* X:**T**
*in* bias:**T**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)|
|FusedConv|*in* X:**T**
*in* W:**T**
*in* B:**T**
*in* Z:**T**
*out* Y:**T**|1+|**T** = tensor(float)|
|FusedMatMul|*in* A:**T**
*in* B:**T**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)|
@@ -1114,6 +1117,7 @@ The **OpSet Version** column uses the following notation:
|MatMulNBits|*in* A:**T1**
*in* B:**T2**
*in* scales:**T1**
*in* zero_points:**T3**
*in* g_idx:**T4**
*in* bias:**T1**
*out* Y:**T1**|1+|**T1** = tensor(bfloat16), tensor(float), tensor(float16)
**T2** = tensor(uint8)
**T3** = tensor(bfloat16), tensor(float), tensor(float16), tensor(uint8)|
|MoE|*in* input:**T**
*in* router_probs:**T**
*in* fc1_experts_weights:**T**
*in* fc1_experts_bias:**T**
*in* fc2_experts_weights:**T**
*in* fc2_experts_bias:**T**
*in* fc3_experts_weights:**T**
*in* fc3_experts_bias:**T**
*out* output:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)|
|MultiHeadAttention|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* bias:**T**
*in* key_padding_mask:**M**
*in* attention_bias:**T**
*in* past_key:**T**
*in* past_value:**T**
*in* past_sequence_length:**M**
*in* cache_indirection:**M**
*out* output:**T**
*out* present_key:**T**
*out* present_value:**T**
*out* qk:**QK**|1+|**QK** = tensor(bfloat16), tensor(float), tensor(float16)
**T** = tensor(bfloat16), tensor(float), tensor(float16)|
+|NGramHashMapping|*in* input_ids:**M**
*in* multipliers:**M**
*in* vocab_sizes:**M**
*in* past_ids:**M**
*out* hash_ids:**M**
*out* present_ids:**M**|1+|**M** = tensor(int32), tensor(int64)|
|NGramRepeatBlock|*in* input_ids:**Tid**
*in* scores:**T**
*out* scores_out:**T**|1+|**T** = tensor(float)
**Tid** = tensor(int64)|
|NhwcConv|*in* X:**T**
*in* W:**T**
*in* B:**T**
*out* Y:**T**|1+|**T** = tensor(float), tensor(float16)|
|PackedAttention|*in* input:**T**
*in* weights:**T**
*in* bias:**T**
*in* token_offset:**M**
*in* cumulative_sequence_length:**M**
*in* attention_bias:**T**
*out* output:**T**|1+|**T** = tensor(float), tensor(float16)|
diff --git a/onnxruntime/contrib_ops/cpu/bert/causal_conv_with_state.cc b/onnxruntime/contrib_ops/cpu/bert/causal_conv_with_state.cc
index d48ba7baa46c2..7585993088805 100644
--- a/onnxruntime/contrib_ops/cpu/bert/causal_conv_with_state.cc
+++ b/onnxruntime/contrib_ops/cpu/bert/causal_conv_with_state.cc
@@ -49,6 +49,10 @@ CausalConvWithState::CausalConvWithState(const OpKernelInfo& info) : OpKernel
ORT_ENFORCE(activation_ == "none" || activation_ == "silu" || activation_ == "swish",
"activation must be one of: none, silu, swish");
+ ORT_THROW_IF_ERROR(causal_conv_with_state_helper::ParseDilation(info, dilation_));
+ ORT_THROW_IF_ERROR(causal_conv_with_state_helper::ParseChannelsLast(info, channels_last_));
+ ORT_ENFORCE(!channels_last_ || ndim_ == 1, "channels_last requires ndim = 1");
+
ORT_ENFORCE(info.GetAttrOrDefault("state_window", 0) == 0,
"CPU CausalConvWithState does not support state_window > 0 (CUDA EP only)");
state_window_ = 0;
@@ -97,84 +101,121 @@ inline void ProcessChannelDecodeFixedK(
}
}
+// A channel's positions are contiguous in the channels-first layout (stride 1) and `channels`
+// apart in the channels-last one, so every per-channel walk goes through these two helpers.
+inline void GatherStrided(float* dst, const float* src, int64_t stride, int64_t count) {
+ if (stride == 1) {
+ std::memcpy(dst, src, static_cast(count) * sizeof(float));
+ return;
+ }
+ for (int64_t i = 0; i < count; ++i) {
+ dst[i] = src[i * stride];
+ }
+}
+
+inline void ScatterStrided(float* dst, int64_t stride, const float* src, int64_t count) {
+ if (stride == 1) {
+ std::memcpy(dst, src, static_cast(count) * sizeof(float));
+ return;
+ }
+ for (int64_t i = 0; i < count; ++i) {
+ dst[i * stride] = src[i];
+ }
+}
+
// Decode fast-path: L=1, no padded buffer needed.
-// The "visible window" for position 0 is [past_state(K-1 values), input(1 value)] = K values.
-// Compute dot(weight, window), shift state left by 1, append new input.
+// The "visible window" is [past_state(pad values), input(1 value)] with pad = (K-1)*dilation.
+// Tap k reads window position k*dilation, so the last tap is the current input. The state is then
+// shifted left by one position and the new input appended.
void ProcessChannelDecode(
- const float* past_row, // past_state for this (b,c): [K-1] or nullptr
- const float* input_val, // &input[b,c,0] — single value
+ const float* past_row, // past_state for this (b,c): [pad] strided, or nullptr
+ const float* input_val, // input for this (b,c) — single value
const float* w, // weight for this channel: [K]
float bias_val,
bool apply_silu,
- float* out_val, // &output[b,c,0] — single value
- float* present_row, // present_state for this (b,c): [K-1]
- int64_t K) {
- int64_t pad = K - 1;
+ float* out_val, // output for this (b,c) — single value
+ float* present_row, // present_state for this (b,c): [pad] strided
+ int64_t state_stride,
+ int64_t K,
+ int64_t dilation) {
+ const int64_t pad = (K - 1) * dilation;
// Dot product over the window: [past_state..., input]
float sum = bias_val;
- // First K-1 elements come from past_state
+ // The first K-1 taps land in past_state, spaced `dilation` apart.
if (past_row != nullptr) {
- for (int64_t k = 0; k < pad; ++k) {
- sum += w[k] * past_row[k];
+ for (int64_t k = 0; k < K - 1; ++k) {
+ sum += w[k] * past_row[k * dilation * state_stride];
}
}
- // Last element is the current input
- sum += w[pad] * input_val[0];
+ // Last tap is the current input
+ sum += w[K - 1] * input_val[0];
if (apply_silu) {
sum = ApplySilu(sum);
}
out_val[0] = sum;
- // Update present_state: shift past_state left by 1, append input
+ // Update present_state: shift past_state left by 1, append input. The copy runs forward from a
+ // higher source index, so it stays correct even if past_state and present_state are the same
+ // buffer.
if (pad > 0) {
- if (past_row != nullptr && pad > 1) {
- std::memcpy(present_row, past_row + 1, static_cast(pad - 1) * sizeof(float));
- } else if (pad > 1) {
- std::memset(present_row, 0, static_cast(pad - 1) * sizeof(float));
+ for (int64_t s = 0; s < pad - 1; ++s) {
+ present_row[s * state_stride] =
+ (past_row != nullptr) ? past_row[(s + 1) * state_stride] : 0.0f;
}
- present_row[pad - 1] = input_val[0];
+ present_row[(pad - 1) * state_stride] = input_val[0];
}
}
// Prefill path: L>1, uses padded buffer for the convolution window.
void ProcessChannelPrefill(
- const float* past_row, // past_state for this (b,c): [K-1] or nullptr
- const float* in_row, // input for this (b,c): [L]
+ const float* past_row, // past_state for this (b,c): [pad] strided, or nullptr
+ const float* in_row, // input for this (b,c): [L] strided
const float* w, // weight for this channel: [K]
float bias_val,
bool apply_silu,
- float* out_row, // output for this (b,c): [L]
- float* present_row, // present_state for this (b,c): [K-1]
- float* padded_row, // scratch buffer: [K-1 + L]
+ float* out_row, // output for this (b,c): [L] strided
+ float* present_row, // present_state for this (b,c): [pad] strided
+ float* padded_row, // scratch buffer: [pad + L]
+ float* out_buf, // scratch buffer: [L], only used when act_stride != 1
+ int64_t act_stride,
+ int64_t state_stride,
int64_t L,
- int64_t K) {
- int64_t pad = K - 1;
- int64_t padded_len = pad + L;
+ int64_t K,
+ int64_t dilation) {
+ const int64_t pad = (K - 1) * dilation;
+ const int64_t padded_len = pad + L;
// Build padded window: [past_state | input]
if (past_row != nullptr) {
- std::memcpy(padded_row, past_row, static_cast(pad) * sizeof(float));
+ GatherStrided(padded_row, past_row, state_stride, pad);
} else {
std::memset(padded_row, 0, static_cast(pad) * sizeof(float));
}
- std::memcpy(padded_row + pad, in_row, static_cast(L) * sizeof(float));
+ GatherStrided(padded_row + pad, in_row, act_stride, L);
- // Depthwise 1D convolution
+ // Depthwise 1D convolution. Tap k of output position l reads padded_row[l + k*dilation]; at
+ // k = K-1 that is padded_row[l + pad], i.e. the current input position.
+ // A contiguous output row is written in place; only a strided one needs the scratch round trip,
+ // so the channels-first layout keeps the exact write pattern it had before strides existed.
+ float* conv_dst = act_stride == 1 ? out_row : out_buf;
for (int64_t l = 0; l < L; ++l) {
float sum = bias_val;
for (int64_t k = 0; k < K; ++k) {
- sum += w[k] * padded_row[l + k];
+ sum += w[k] * padded_row[l + k * dilation];
}
if (apply_silu) {
sum = ApplySilu(sum);
}
- out_row[l] = sum;
+ conv_dst[l] = sum;
+ }
+ if (act_stride != 1) {
+ ScatterStrided(out_row, act_stride, out_buf, L);
}
- // Save present_state: last K-1 elements of (past_state | input)
- std::memcpy(present_row, padded_row + padded_len - pad, static_cast(pad) * sizeof(float));
+ // Save present_state: last pad elements of (past_state | input)
+ ScatterStrided(present_row, state_stride, padded_row + padded_len - pad, pad);
}
} // anonymous namespace
@@ -192,13 +233,22 @@ Status CausalConvWithState::Compute(OpKernelContext* context) const {
const auto& input_shape = input_tensor->Shape();
const auto& weight_shape = weight_tensor->Shape();
- ORT_RETURN_IF_NOT(static_cast(input_shape.NumDimensions()) == 2 + ndim_,
- "input must have ", 2 + ndim_, " dimensions for ndim=", ndim_);
+ if (channels_last_) {
+ // (batch_size, sequence_length, d_1, ..., d_n): any number of trailing channel axes, so a
+ // caller that keeps hyper-connections and hidden size separate needs no reshape.
+ ORT_RETURN_IF_NOT(input_shape.NumDimensions() >= 3,
+ "input must have at least 3 dimensions when channels_last = 1");
+ } else {
+ ORT_RETURN_IF_NOT(static_cast(input_shape.NumDimensions()) == 2 + ndim_,
+ "input must have ", 2 + ndim_, " dimensions for ndim=", ndim_);
+ }
ORT_RETURN_IF_NOT(static_cast(weight_shape.NumDimensions()) == 2 + ndim_,
"weight must have ", 2 + ndim_, " dimensions for ndim=", ndim_);
const int64_t batch_size = input_shape[0];
- const int64_t channels = input_shape[1];
+ const int64_t channels = channels_last_
+ ? input_shape.SizeFromDimension(2)
+ : input_shape[1];
ORT_RETURN_IF_NOT(weight_shape[0] == channels, "weight channels must match input channels");
ORT_RETURN_IF_NOT(weight_shape[1] == 1, "weight must be depthwise (group=1)");
@@ -211,28 +261,26 @@ Status CausalConvWithState::Compute(OpKernelContext* context) const {
// ==== ndim=1 implementation: (B, C, L) with kernel (C, 1, K) ====
if (ndim_ == 1) {
- const int64_t L = input_shape[2];
+ const int64_t L = channels_last_ ? input_shape[1] : input_shape[2];
const int64_t K = weight_shape[2];
- const int64_t pad = K - 1;
-
- if (past_state_tensor != nullptr) {
- const auto& ps_shape = past_state_tensor->Shape();
- ORT_RETURN_IF_NOT(ps_shape.NumDimensions() == 3 &&
- ps_shape[0] == batch_size &&
- ps_shape[1] == channels &&
- ps_shape[2] == pad,
- "past_state must be (B, C, K-1)");
- }
+ const int64_t dilation = dilation_;
+ const int64_t pad = (K - 1) * dilation;
// ==== Allocate outputs ====
Tensor* output_tensor = context->Output(0, input_shape);
float* output_data = output_tensor->MutableData();
- // state_window_ is always 0 on CPU, so this is the legacy (B, C, K-1) shape.
+ // state_window_ is always 0 on CPU, so the state has no leading window axis.
TensorShape state_shape;
- ORT_RETURN_IF_ERROR(causal_conv_with_state_helper::CheckInputs(
- state_window_, static_cast(batch_size), static_cast(channels),
- static_cast(pad), past_state_tensor, state_shape, "CausalConvWithState"));
+ if (channels_last_) {
+ ORT_RETURN_IF_ERROR(causal_conv_with_state_helper::CheckInputsChannelsLast(
+ state_window_, input_shape, static_cast(pad), past_state_tensor, state_shape,
+ "CausalConvWithState"));
+ } else {
+ ORT_RETURN_IF_ERROR(causal_conv_with_state_helper::CheckInputs(
+ state_window_, static_cast(batch_size), static_cast(channels),
+ static_cast(pad), past_state_tensor, state_shape, "CausalConvWithState"));
+ }
Tensor* present_state_tensor = context->Output(1, state_shape);
float* present_data = present_state_tensor->MutableData();
@@ -242,6 +290,11 @@ Status CausalConvWithState::Compute(OpKernelContext* context) const {
const float* past_data = past_state_tensor ? past_state_tensor->Data() : nullptr;
bool apply_silu = (activation_ == "silu" || activation_ == "swish");
+ // Both layouts are dense, so one strided (batch, position, channel) view covers them.
+ const auto act_layout = causal_conv_with_state_helper::MakeLayout(channels_last_, channels, L);
+ const auto state_layout =
+ causal_conv_with_state_helper::MakeLayout(channels_last_, channels, pad);
+
// ==== Thread-parallel over (batch, channel) pairs ====
// Depthwise conv: each channel is fully independent.
int64_t total_tasks = batch_size * channels;
@@ -260,15 +313,19 @@ Status CausalConvWithState::Compute(OpKernelContext* context) const {
int64_t b = task / channels;
int64_t c = task % channels;
- const float* past_row = past_data
- ? past_data + (b * channels + c) * pad
- : nullptr;
- const float* input_val = input_data + (b * channels + c) * L;
+ const int64_t act_offset = act_layout.Offset(b, 0, c);
+ const int64_t state_offset = state_layout.Offset(b, 0, c);
+
+ const float* past_row = past_data ? past_data + state_offset : nullptr;
+ const float* input_val = input_data + act_offset;
const float* w = weight_data + c * K;
float bias_val = bias_data ? bias_data[c] : 0.0f;
- float* out_val = output_data + (b * channels + c) * L;
- float* present_row = present_data + (b * channels + c) * pad;
- switch (K) {
+ float* out_val = output_data + act_offset;
+ float* present_row = present_data + state_offset;
+ // ProcessChannelDecodeFixedK assumes pad == K - 1 and contiguous state, so it only
+ // applies to the undilated channels-first case.
+ const int64_t fixed_k = (dilation == 1 && state_layout.pos_stride == 1) ? K : 0;
+ switch (fixed_k) {
case 2:
ProcessChannelDecodeFixedK<2>(past_row, input_val, w, bias_val, apply_silu,
out_val, present_row);
@@ -287,7 +344,7 @@ Status CausalConvWithState::Compute(OpKernelContext* context) const {
break;
default:
ProcessChannelDecode(past_row, input_val, w, bias_val, apply_silu,
- out_val, present_row, K);
+ out_val, present_row, state_layout.pos_stride, K, dilation);
break;
}
}
@@ -299,24 +356,28 @@ Status CausalConvWithState::Compute(OpKernelContext* context) const {
static_cast(total_tasks),
cost_per_task,
[&](std::ptrdiff_t first, std::ptrdiff_t last) {
- // Per-thread scratch buffer for padded input
+ // Per-thread scratch buffers for the padded input window and, only when the output row
+ // is strided, the contiguous convolution result that is then scattered into it.
std::vector padded_buf(static_cast(pad + L));
+ std::vector out_buf(act_layout.pos_stride == 1 ? 0 : static_cast(L));
for (std::ptrdiff_t task = first; task < last; ++task) {
int64_t b = task / channels;
int64_t c = task % channels;
- const float* past_row = past_data
- ? past_data + (b * channels + c) * pad
- : nullptr;
- const float* in_row = input_data + (b * channels + c) * L;
+ const int64_t act_offset = act_layout.Offset(b, 0, c);
+ const int64_t state_offset = state_layout.Offset(b, 0, c);
+
+ const float* past_row = past_data ? past_data + state_offset : nullptr;
+ const float* in_row = input_data + act_offset;
const float* w = weight_data + c * K;
float bias_val = bias_data ? bias_data[c] : 0.0f;
- float* out_row = output_data + (b * channels + c) * L;
- float* present_row = present_data + (b * channels + c) * pad;
+ float* out_row = output_data + act_offset;
+ float* present_row = present_data + state_offset;
ProcessChannelPrefill(past_row, in_row, w, bias_val, apply_silu,
- out_row, present_row, padded_buf.data(), L, K);
+ out_row, present_row, padded_buf.data(), out_buf.data(),
+ act_layout.pos_stride, state_layout.pos_stride, L, K, dilation);
}
});
}
diff --git a/onnxruntime/contrib_ops/cpu/bert/causal_conv_with_state.h b/onnxruntime/contrib_ops/cpu/bert/causal_conv_with_state.h
index 0e552e7bd27dd..19467ac50a918 100644
--- a/onnxruntime/contrib_ops/cpu/bert/causal_conv_with_state.h
+++ b/onnxruntime/contrib_ops/cpu/bert/causal_conv_with_state.h
@@ -19,6 +19,8 @@ class CausalConvWithState final : public OpKernel {
private:
int ndim_;
+ int dilation_;
+ bool channels_last_;
std::string activation_;
// Always 0 on CPU (a state window is CUDA-only), but kept so the shared shape helper in
// causal_conv_with_state_helper.h is driven the same way on every EP.
diff --git a/onnxruntime/contrib_ops/cpu/bert/causal_conv_with_state_helper.h b/onnxruntime/contrib_ops/cpu/bert/causal_conv_with_state_helper.h
index 9e5c737dabaec..2d1c96dc8175b 100644
--- a/onnxruntime/contrib_ops/cpu/bert/causal_conv_with_state_helper.h
+++ b/onnxruntime/contrib_ops/cpu/bert/causal_conv_with_state_helper.h
@@ -30,8 +30,91 @@ Status ParseStateWindow(const TKernelInfo& info, int& state_window) {
return Status::OK();
}
+// Reads and validates the optional `dilation` attribute.
+//
+// 1 (the default, i.e. attribute absent) is the undilated case, which is what every model exported
+// before the attribute existed uses, so this must stay the default.
+template
+Status ParseDilation(const TKernelInfo& info, int& dilation) {
+ const int64_t value = info.template GetAttrOrDefault("dilation", 1);
+ if (value < 1) {
+ return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "dilation must be >= 1, got ", value);
+ }
+ dilation = static_cast(value);
+ return Status::OK();
+}
+
+// Reads and validates the optional `channels_last` attribute.
+//
+// 0 (the default, i.e. attribute absent) is the channels-first (batch_size, channels, seq_len)
+// layout that every model exported before the attribute existed uses.
+template
+Status ParseChannelsLast(const TKernelInfo& info, bool& channels_last) {
+ const int64_t value = info.template GetAttrOrDefault("channels_last", 0);
+ if (value != 0 && value != 1) {
+ return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "channels_last must be 0 or 1, got ",
+ value);
+ }
+ channels_last = (value == 1);
+ return Status::OK();
+}
+
+// Element strides of a (batch, position, channel) view over an activation or state tensor.
+// Both supported layouts are dense, so a single strided view covers them and every kernel can be
+// written once against (batch, position, channel) coordinates.
+struct Layout {
+ int64_t batch_stride;
+ int64_t pos_stride;
+ int64_t chan_stride;
+
+ int64_t Offset(int64_t b, int64_t pos, int64_t c) const {
+ return b * batch_stride + pos * pos_stride + c * chan_stride;
+ }
+};
+
+// `length` is the extent of the position axis: seq_len for activations, state_length for state.
+constexpr Layout MakeLayout(bool channels_last, int64_t channels, int64_t length) {
+ return channels_last ? Layout{channels * length, channels, 1}
+ : Layout{channels * length, 1, length};
+}
+
+// Derives the expected channels-last past_state / present_state shape,
+// (batch_size, state_length, d_1, ..., d_n), from the input shape and validates past_state.
+// The trailing channel axes are copied verbatim from the input, so a caller that keeps
+// hyper-connections and hidden size as separate axes gets the same split back and needs no reshape.
+template
+Status CheckInputsChannelsLast(int state_window,
+ const TensorShape& input_shape,
+ int state_length,
+ const T* past_state,
+ TensorShape& state_shape,
+ std::string_view op_name) {
+ TensorShapeVector dims;
+ if (state_window > 0) {
+ dims.push_back(state_window);
+ }
+ dims.push_back(input_shape[0]);
+ dims.push_back(state_length);
+ for (size_t i = 2; i < input_shape.NumDimensions(); ++i) {
+ dims.push_back(input_shape[i]);
+ }
+ state_shape = TensorShape(dims);
+
+ if (past_state != nullptr && past_state->Shape() != state_shape) {
+ return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+ "Input 'past_state' is expected to have shape ", state_shape.ToString(),
+ ", got ", past_state->Shape().ToString(),
+ ". ", op_name,
+ " with channels_last = 1 uses "
+ "(batch_size, (kernel_size - 1) * dilation, d_1, ..., d_n), optionally "
+ "led by a state_window axis.");
+ }
+
+ return Status::OK();
+}
+
// Derives the expected past_state / present_state shape and validates past_state against it.
-// `state_length` is the carry length along the causal axis, i.e. kernel_size - 1.
+// `state_length` is the carry length along the causal axis, i.e. (kernel_size - 1) * dilation.
//
// state_window == 0 -> (batch_size, channels, state_length). A single state with no window axis:
// the backward-compatible layout that models exported before the attribute existed use.
@@ -62,9 +145,10 @@ Status CheckInputs(int state_window,
"Input 'past_state' is expected to have shape ", state_shape.ToString(),
", got ", past_state->Shape().ToString(),
". ", op_name,
- " uses (batch_size, channels, kernel_size - 1) when "
+ " uses (batch_size, channels, (kernel_size - 1) * dilation) when "
"the state_window attribute is absent or 0, and "
- "(state_window, batch_size, channels, kernel_size - 1) otherwise.");
+ "(state_window, batch_size, channels, (kernel_size - 1) * dilation) "
+ "otherwise.");
}
return Status::OK();
diff --git a/onnxruntime/contrib_ops/cpu/bert/engram_gate.cc b/onnxruntime/contrib_ops/cpu/bert/engram_gate.cc
new file mode 100644
index 0000000000000..bd7d9da8c241b
--- /dev/null
+++ b/onnxruntime/contrib_ops/cpu/bert/engram_gate.cc
@@ -0,0 +1,125 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "contrib_ops/cpu/bert/engram_gate.h"
+
+#include
+
+#include "contrib_ops/cpu/bert/engram_helper.h"
+#include "core/common/narrow.h"
+#include "core/platform/threadpool.h"
+
+using onnxruntime::concurrency::ThreadPool;
+
+namespace onnxruntime {
+namespace contrib {
+
+#define REGISTER_ENGRAM_GATE_TYPED(T) \
+ ONNX_OPERATOR_TYPED_KERNEL_EX( \
+ EngramGate, \
+ kMSDomain, \
+ 1, \
+ T, \
+ kCpuExecutionProvider, \
+ KernelDefBuilder() \
+ .TypeConstraint("T", DataTypeImpl::GetTensorType()), \
+ EngramGate);
+
+REGISTER_ENGRAM_GATE_TYPED(float)
+REGISTER_ENGRAM_GATE_TYPED(MLFloat16)
+
+#undef REGISTER_ENGRAM_GATE_TYPED
+
+template
+EngramGate::EngramGate(const OpKernelInfo& info) : OpKernel(info) {
+ epsilon_ = info.GetAttrOrDefault("epsilon", 1.0e-5f);
+}
+
+template
+Status EngramGate::Compute(OpKernelContext* context) const {
+ const Tensor* key = context->Input(0);
+ const Tensor* query = context->Input(1);
+ const Tensor* value = context->Input(2);
+ const Tensor* key_norm_scale = context->Input(3);
+ const Tensor* query_norm_scale = context->Input(4);
+
+ const TensorShape& key_shape = key->Shape();
+ ORT_RETURN_IF_NOT(key_shape.NumDimensions() == 4,
+ "key must have shape (batch_size, sequence_length, hc_mult, hidden_size)");
+ const int64_t batch_size = key_shape[0];
+ const int64_t sequence_length = key_shape[1];
+ const int64_t hc_mult = key_shape[2];
+ const int64_t hidden_size = key_shape[3];
+
+ ORT_RETURN_IF_NOT(query->Shape() == key_shape, "query must have the same shape as key");
+ ORT_RETURN_IF_NOT(value->Shape() == TensorShape({batch_size, sequence_length, hidden_size}),
+ "value must have shape (batch_size, sequence_length, hidden_size)");
+ ORT_RETURN_IF_NOT(key_norm_scale->Shape() == TensorShape({hc_mult, hidden_size}),
+ "key_norm_scale must have shape (hc_mult, hidden_size)");
+ ORT_RETURN_IF_NOT(query_norm_scale->Shape() == TensorShape({hc_mult, hidden_size}),
+ "query_norm_scale must have shape (hc_mult, hidden_size)");
+
+ Tensor* output = context->Output(0, key_shape);
+ if (key_shape.Size() == 0) {
+ return Status::OK();
+ }
+
+ const T* key_data = key->Data();
+ const T* query_data = query->Data();
+ const T* value_data = value->Data();
+ const T* key_scale_data = key_norm_scale->Data();
+ const T* query_scale_data = query_norm_scale->Data();
+ T* output_data = output->MutableData();
+
+ const int64_t rows = batch_size * sequence_length * hc_mult;
+ ThreadPool::TryParallelFor(
+ // Each row makes one fused reduction pass and one output pass over hidden_size, plus a
+ // handful of scalar transcendentals. Costing it as a single pass would over-partition.
+ context->GetOperatorThreadPool(), narrow(rows),
+ static_cast(2 * hidden_size + 32),
+ [&](ptrdiff_t begin, ptrdiff_t end) {
+ for (int64_t row = begin; row < end; ++row) {
+ const int64_t g = row % hc_mult;
+ const int64_t token = row / hc_mult;
+ const T* key_row = key_data + row * hidden_size;
+ const T* query_row = query_data + row * hidden_size;
+ const T* value_row = value_data + token * hidden_size;
+
+ // Both inverse RMS factors are scalars, so they can be pulled out of the dot product and
+ // applied afterwards. That folds the two reductions into one pass over key_row and
+ // query_row, which is what the CUDA and WGSL kernels already do.
+ const T* key_scale_row = key_scale_data + g * hidden_size;
+ const T* query_scale_row = query_scale_data + g * hidden_size;
+ float key_sum_sq = 0.0f;
+ float query_sum_sq = 0.0f;
+ float dot_numerator = 0.0f;
+ for (int64_t c = 0; c < hidden_size; ++c) {
+ const float key_value = static_cast(key_row[c]);
+ const float query_value = static_cast(query_row[c]);
+ key_sum_sq += key_value * key_value;
+ query_sum_sq += query_value * query_value;
+ dot_numerator += key_value * static_cast(key_scale_row[c]) *
+ query_value * static_cast(query_scale_row[c]);
+ }
+
+ const float key_inv_rms = 1.0f / std::sqrt(key_sum_sq / static_cast(hidden_size) + epsilon_);
+ const float query_inv_rms = 1.0f / std::sqrt(query_sum_sq / static_cast(hidden_size) + epsilon_);
+ const float dot =
+ dot_numerator * key_inv_rms * query_inv_rms / std::sqrt(static_cast(hidden_size));
+ const float gate = engram_helper::SigmoidFloat(engram_helper::EngramGateArg(dot));
+
+ T* output_row = output_data + row * hidden_size;
+ for (int64_t c = 0; c < hidden_size; ++c) {
+ output_row[c] = static_cast(gate * static_cast(value_row[c]));
+ }
+ }
+ });
+
+ return Status::OK();
+}
+
+template class EngramGate;
+template class EngramGate;
+
+} // namespace contrib
+} // namespace onnxruntime
diff --git a/onnxruntime/contrib_ops/cpu/bert/engram_gate.h b/onnxruntime/contrib_ops/cpu/bert/engram_gate.h
new file mode 100644
index 0000000000000..ec92da027dde4
--- /dev/null
+++ b/onnxruntime/contrib_ops/cpu/bert/engram_gate.h
@@ -0,0 +1,23 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include "core/common/common.h"
+#include "core/framework/op_kernel.h"
+
+namespace onnxruntime {
+namespace contrib {
+
+template
+class EngramGate final : public OpKernel {
+ public:
+ explicit EngramGate(const OpKernelInfo& info);
+ Status Compute(OpKernelContext* context) const override;
+
+ private:
+ float epsilon_;
+};
+
+} // namespace contrib
+} // namespace onnxruntime
diff --git a/onnxruntime/contrib_ops/cpu/bert/engram_helper.h b/onnxruntime/contrib_ops/cpu/bert/engram_helper.h
new file mode 100644
index 0000000000000..72abdf3ee2368
--- /dev/null
+++ b/onnxruntime/contrib_ops/cpu/bert/engram_helper.h
@@ -0,0 +1,54 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include
+#include
+#include
+
+namespace onnxruntime {
+namespace contrib {
+namespace engram_helper {
+
+// Numerically stable logistic function.
+inline float SigmoidFloat(float x) {
+ if (x > 0.0f) {
+ return 1.0f / (1.0f + std::exp(-x));
+ }
+ const float exp_x = std::exp(x);
+ return exp_x / (1.0f + exp_x);
+}
+
+// Engram gate pre-activation: sign(dot) * sqrt(max(abs(dot), 1e-6)).
+// std::copysign cannot be used here because it maps a zero dot product to +sqrt(1e-6) instead of
+// zero, which would disagree with the schema formula and with the other execution providers.
+inline float EngramGateArg(float dot) {
+ if (dot == 0.0f) {
+ return 0.0f;
+ }
+ const float magnitude = std::sqrt(std::max(std::abs(dot), 1.0e-6f));
+ return dot < 0.0f ? -magnitude : magnitude;
+}
+
+// Euclidean modulo: the result always has the sign of `mod`, which must be positive.
+template
+inline T PositiveMod(T value, T mod) {
+ T result = static_cast(value % mod);
+ if (result < 0) {
+ result = static_cast(result + mod);
+ }
+ return result;
+}
+
+// Multiplies through the unsigned counterpart of T so that overflow wraps around instead of
+// being undefined behavior.
+template
+inline T WrappedMultiply(T a, T b) {
+ using UnsignedT = typename std::make_unsigned::type;
+ return static_cast(static_cast(a) * static_cast(b));
+}
+
+} // namespace engram_helper
+} // namespace contrib
+} // namespace onnxruntime
diff --git a/onnxruntime/contrib_ops/cpu/bert/ngram_hash_mapping.cc b/onnxruntime/contrib_ops/cpu/bert/ngram_hash_mapping.cc
new file mode 100644
index 0000000000000..2b8250e9df817
--- /dev/null
+++ b/onnxruntime/contrib_ops/cpu/bert/ngram_hash_mapping.cc
@@ -0,0 +1,161 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "contrib_ops/cpu/bert/ngram_hash_mapping.h"
+
+#include
+#include
+
+#include "contrib_ops/cpu/bert/engram_helper.h"
+#include "core/common/narrow.h"
+#include "core/platform/threadpool.h"
+
+using onnxruntime::concurrency::ThreadPool;
+
+namespace onnxruntime {
+namespace contrib {
+
+#define REGISTER_NGRAM_HASH_TYPED(T) \
+ ONNX_OPERATOR_TYPED_KERNEL_EX( \
+ NGramHashMapping, \
+ kMSDomain, \
+ 1, \
+ T, \
+ kCpuExecutionProvider, \
+ KernelDefBuilder() \
+ .MayInplace(3, 1) \
+ .TypeConstraint("M", DataTypeImpl::GetTensorType()), \
+ NGramHashMapping);
+
+REGISTER_NGRAM_HASH_TYPED(int32_t)
+REGISTER_NGRAM_HASH_TYPED(int64_t)
+
+#undef REGISTER_NGRAM_HASH_TYPED
+
+template
+NGramHashMapping::NGramHashMapping(const OpKernelInfo& info) : OpKernel(info) {
+ ORT_ENFORCE(info.GetAttr("max_ngram_size", &max_ngram_size_).IsOK(),
+ "max_ngram_size attribute is required");
+ ORT_ENFORCE(info.GetAttr("n_head_per_ngram", &n_head_per_ngram_).IsOK(),
+ "n_head_per_ngram attribute is required");
+ int64_t pad_id = 0;
+ ORT_ENFORCE(info.GetAttr("pad_id", &pad_id).IsOK(), "pad_id attribute is required");
+ ORT_ENFORCE(max_ngram_size_ >= 2, "max_ngram_size must be at least 2");
+ ORT_ENFORCE(n_head_per_ngram_ >= 1, "n_head_per_ngram must be positive");
+ ORT_ENFORCE(pad_id >= static_cast(std::numeric_limits::min()) &&
+ pad_id <= static_cast(std::numeric_limits::max()),
+ "pad_id is out of range for the input id type");
+ pad_id_ = static_cast(pad_id);
+}
+
+// Reads the id at right-aligned history slot `slot` of past_ids. Slots outside the provided history
+// (or a missing past_ids) are positions before the start of the whole sequence, so they use pad_id.
+template
+T NGramHashMapping::HistoryId(const T* past_data, int64_t b, int64_t slot, int64_t state_length) const {
+ if (past_data == nullptr || slot < 0 || slot >= state_length) {
+ return pad_id_;
+ }
+ return past_data[b * state_length + slot];
+}
+
+template
+Status NGramHashMapping::Compute(OpKernelContext* context) const {
+ const Tensor* input_ids = context->Input(0);
+ const Tensor* multipliers = context->Input(1);
+ const Tensor* vocab_sizes = context->Input(2);
+ const Tensor* past_ids = context->Input(3);
+
+ const TensorShape& input_shape = input_ids->Shape();
+ ORT_RETURN_IF_NOT(input_shape.NumDimensions() == 2, "input_ids must have rank 2");
+ ORT_RETURN_IF_NOT(multipliers->Shape().NumDimensions() == 1 &&
+ multipliers->Shape()[0] == max_ngram_size_,
+ "multipliers must have shape (max_ngram_size)");
+ const int64_t num_heads = (max_ngram_size_ - 1) * n_head_per_ngram_;
+ ORT_RETURN_IF_NOT(vocab_sizes->Shape().NumDimensions() == 1 && vocab_sizes->Shape()[0] == num_heads,
+ "vocab_sizes must have shape ((max_ngram_size - 1) * n_head_per_ngram)");
+
+ const int64_t batch_size = input_shape[0];
+ const int64_t sequence_length = input_shape[1];
+ // An n-gram window reaches this many positions before the current token.
+ const int64_t state_length = max_ngram_size_ - 1;
+ if (past_ids != nullptr) {
+ ORT_RETURN_IF_NOT(past_ids->Shape() == TensorShape({batch_size, state_length}),
+ "past_ids must have shape (batch_size, max_ngram_size - 1)");
+ }
+
+ Tensor* output = context->Output(0, TensorShape({batch_size, sequence_length, num_heads}));
+ Tensor* present_ids = context->Output(1, TensorShape({batch_size, state_length}));
+
+ const T* input_data = input_ids->Data();
+ const T* multiplier_data = multipliers->Data();
+ const T* vocab_data = vocab_sizes->Data();
+ const T* past_data = past_ids == nullptr ? nullptr : past_ids->Data();
+
+ // A non-positive head vocabulary size has no meaningful modulo. Every EP guards the division to
+ // avoid a device-side divide-by-zero, which turns the mistake into a constant hash id of 0 for that
+ // head rather than a crash. That is a silent wrong answer, so validate it here where vocab_sizes is
+ // already resident on the host and the check costs one pass over a tiny tensor.
+ for (int64_t h = 0; h < num_heads; ++h) {
+ ORT_RETURN_IF_NOT(vocab_data[h] > 0,
+ "vocab_sizes must be positive; entry ", h, " is ", static_cast(vocab_data[h]));
+ }
+
+ if (input_shape.Size() != 0) {
+ T* output_data = output->MutableData();
+ const int64_t total = batch_size * sequence_length;
+ ThreadPool::TryParallelFor(
+ context->GetOperatorThreadPool(), narrow(total), static_cast(max_ngram_size_ * n_head_per_ngram_),
+ [&](ptrdiff_t begin, ptrdiff_t end) {
+ for (int64_t linear = begin; linear < end; ++linear) {
+ const int64_t t = linear % sequence_length;
+ const int64_t b = linear / sequence_length;
+ const int64_t input_base = b * sequence_length;
+ const int64_t output_base = linear * num_heads;
+
+ for (int64_t n = 2; n <= max_ngram_size_; ++n) {
+ T mix = 0;
+ for (int64_t k = 0; k < n; ++k) {
+ const int64_t source_t = t - k;
+ const T token = source_t >= 0 ? input_data[input_base + source_t]
+ : HistoryId(past_data, b, state_length + source_t, state_length);
+ const T product = engram_helper::WrappedMultiply(token, multiplier_data[k]);
+ mix = k == 0 ? product : static_cast(mix ^ product);
+ }
+
+ const int64_t ngram_offset = (n - 2) * n_head_per_ngram_;
+ for (int64_t h = 0; h < n_head_per_ngram_; ++h) {
+ const int64_t out_h = ngram_offset + h;
+ // vocab_sizes was validated to be positive above, so the modulo is always well defined.
+ output_data[output_base + out_h] = engram_helper::PositiveMod(mix, vocab_data[out_h]);
+ }
+ }
+ }
+ });
+ }
+
+ // present_ids is the right-aligned trailing window of (past_ids ++ input_ids), so it is well defined
+ // even when this call is shorter than the window. It is written last because past_ids may share
+ // its allocation, and the hash loop above still needs the original history. Within this loop the
+ // aliased case is safe too: slot j writes index j and reads index j + sequence_length, so the walk
+ // is strictly ahead of itself.
+ if (present_ids != nullptr) {
+ T* present_data = present_ids->MutableData();
+ for (int64_t b = 0; b < batch_size; ++b) {
+ for (int64_t j = 0; j < state_length; ++j) {
+ // Virtual position of slot j relative to the end of input_ids.
+ const int64_t source_t = sequence_length - state_length + j;
+ present_data[b * state_length + j] =
+ source_t >= 0 ? input_data[b * sequence_length + source_t]
+ : HistoryId(past_data, b, state_length + source_t, state_length);
+ }
+ }
+ }
+
+ return Status::OK();
+}
+
+template class NGramHashMapping;
+template class NGramHashMapping;
+
+} // namespace contrib
+} // namespace onnxruntime
diff --git a/onnxruntime/contrib_ops/cpu/bert/ngram_hash_mapping.h b/onnxruntime/contrib_ops/cpu/bert/ngram_hash_mapping.h
new file mode 100644
index 0000000000000..77b9c1cd524fe
--- /dev/null
+++ b/onnxruntime/contrib_ops/cpu/bert/ngram_hash_mapping.h
@@ -0,0 +1,27 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include "core/common/common.h"
+#include "core/framework/op_kernel.h"
+
+namespace onnxruntime {
+namespace contrib {
+
+template
+class NGramHashMapping final : public OpKernel {
+ public:
+ explicit NGramHashMapping(const OpKernelInfo& info);
+ Status Compute(OpKernelContext* context) const override;
+
+ private:
+ T HistoryId(const T* past_data, int64_t b, int64_t slot, int64_t state_length) const;
+
+ int64_t max_ngram_size_;
+ int64_t n_head_per_ngram_;
+ T pad_id_;
+};
+
+} // namespace contrib
+} // namespace onnxruntime
diff --git a/onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc b/onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc
index 1d323b18af6fe..a3ce50516fae2 100644
--- a/onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc
+++ b/onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc
@@ -39,6 +39,10 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1,
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, MLFloat16, LinearAttentionGate);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, GatedRMSNorm);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, MLFloat16, GatedRMSNorm);
+class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, EngramGate);
+class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, MLFloat16, EngramGate);
+class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, int32_t, NGramHashMapping);
+class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, int64_t, NGramHashMapping);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, CausalConvWithState);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, RotaryEmbedding);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, MLFloat16, RotaryEmbedding);
@@ -343,6 +347,10 @@ Status RegisterCpuContribKernels(KernelRegistry& kernel_registry) {
BuildKernelCreateInfo,
BuildKernelCreateInfo,
BuildKernelCreateInfo,
+ BuildKernelCreateInfo,
+ BuildKernelCreateInfo,
+ BuildKernelCreateInfo,
+ BuildKernelCreateInfo,
BuildKernelCreateInfo,
BuildKernelCreateInfo,
BuildKernelCreateInfo,
diff --git a/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state.cc b/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state.cc
index 817097ac34ad4..c79859f735c7f 100644
--- a/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state.cc
+++ b/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state.cc
@@ -38,6 +38,10 @@ CausalConvWithState::CausalConvWithState(const OpKernelInfo& info) : CudaKern
ORT_ENFORCE(activation_ == "none" || activation_ == "silu" || activation_ == "swish",
"activation must be one of: none, silu, swish");
+ ORT_THROW_IF_ERROR(causal_conv_with_state_helper::ParseDilation(info, dilation_));
+ ORT_THROW_IF_ERROR(causal_conv_with_state_helper::ParseChannelsLast(info, channels_last_));
+ ORT_ENFORCE(!channels_last_ || ndim_ == 1, "channels_last requires ndim = 1");
+
// See LinearAttention: only the trailing per-position states are ever consumed, so a window
// caps the allocation and the write traffic for long prompts. 0 keeps the plain single state.
ORT_THROW_IF_ERROR(causal_conv_with_state_helper::ParseStateWindow(info, state_window_));
@@ -57,16 +61,26 @@ Status CausalConvWithState::ComputeInternal(OpKernelContext* context) const {
const auto& weight_shape = weight_tensor->Shape();
// Validate input rank and weight rank
- ORT_RETURN_IF_NOT(input_shape.NumDimensions() == 3,
- "input must be rank 3 (batch, channels, length), got rank ", input_shape.NumDimensions());
+ if (channels_last_) {
+ // (batch_size, sequence_length, d_1, ..., d_n): any number of trailing channel axes, so a
+ // caller that keeps hyper-connections and hidden size separate needs no reshape.
+ ORT_RETURN_IF_NOT(input_shape.NumDimensions() >= 3,
+ "input must have rank >= 3 (batch, length, ...channels) when "
+ "channels_last = 1, got rank ",
+ input_shape.NumDimensions());
+ } else {
+ ORT_RETURN_IF_NOT(input_shape.NumDimensions() == 3,
+ "input must be rank 3 (batch, channels, length), got rank ", input_shape.NumDimensions());
+ }
ORT_RETURN_IF_NOT(weight_shape.NumDimensions() == 3,
"weight must be rank 3 (channels, 1, kernel_size), got rank ", weight_shape.NumDimensions());
const int batch_size = static_cast(input_shape[0]);
- const int channels = static_cast(input_shape[1]);
- const int L = static_cast(input_shape[2]);
+ const int channels = static_cast(channels_last_ ? input_shape.SizeFromDimension(2)
+ : input_shape[1]);
+ const int L = static_cast(channels_last_ ? input_shape[1] : input_shape[2]);
const int K = static_cast(weight_shape[2]);
- const int pad = K - 1;
+ const int pad = (K - 1) * dilation_;
ORT_RETURN_IF_NOT(L > 0, "input length must be positive, got ", L);
@@ -83,13 +97,19 @@ Status CausalConvWithState::ComputeInternal(OpKernelContext* context) const {
"bias must have shape (", channels, "), got ", bias_shape.ToString());
}
- // past_state / present_state are [B, C, K-1], or [W, B, C, K-1] when state_window_ = W > 0.
+ // past_state / present_state are [B, C, pad], or [W, B, C, pad] when state_window_ = W > 0,
+ // where pad = (K-1)*dilation.
// Right-aligned: token t lands in slot t + W - L, so slot W-1 always holds the state after the
// last token (and is the slot past_state is read from).
const int state_slots = state_window_ > 0 ? state_window_ : 1;
TensorShape state_shape;
- ORT_RETURN_IF_ERROR(causal_conv_with_state_helper::CheckInputs(
- state_window_, batch_size, channels, pad, past_state_tensor, state_shape, "CausalConvWithState"));
+ if (channels_last_) {
+ ORT_RETURN_IF_ERROR(causal_conv_with_state_helper::CheckInputsChannelsLast(
+ state_window_, input_shape, pad, past_state_tensor, state_shape, "CausalConvWithState"));
+ } else {
+ ORT_RETURN_IF_ERROR(causal_conv_with_state_helper::CheckInputs(
+ state_window_, batch_size, channels, pad, past_state_tensor, state_shape, "CausalConvWithState"));
+ }
// Allocate outputs
Tensor* output_tensor = context->Output(0, input_shape);
@@ -125,6 +145,9 @@ Status CausalConvWithState::ComputeInternal(OpKernelContext* context) const {
channels,
L,
K,
+ dilation_,
+ MakeCausalConvLayout(channels_last_, channels, L),
+ MakeCausalConvLayout(channels_last_, channels, pad),
apply_silu,
GetDeviceProp().maxThreadsPerBlock,
state_slots);
diff --git a/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state.h b/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state.h
index f0fb66e8485b9..0a1f23f3abfd3 100644
--- a/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state.h
+++ b/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state.h
@@ -20,6 +20,8 @@ class CausalConvWithState final : public onnxruntime::cuda::CudaKernel {
private:
int ndim_;
+ int dilation_;
+ bool channels_last_;
std::string activation_;
// Leading (axis-0) extent of past_state / present_state; 0 means no window axis (single state).
int state_window_;
diff --git a/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state_impl.cu b/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state_impl.cu
index 8c34e672a4797..0c0e1da836521 100644
--- a/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state_impl.cu
+++ b/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state_impl.cu
@@ -5,8 +5,9 @@
//
// Design: One thread block per (batch, channel). Two execution paths:
//
-// 1. Decode (L=1): The convolution window is [past_state(K-1), input(1)].
-// Load K values into registers, compute a single dot product, shift state.
+// 1. Decode (L=1): The convolution window is [past_state(pad), input(1)] with
+// pad = (K-1)*dilation. Load K values into registers, compute a single dot product,
+// shift state.
// One thread block does the entire operation — zero shared memory needed.
//
// 2. Prefill (L>1): Load past_state + input into shared memory as a padded buffer,
@@ -47,6 +48,9 @@ __global__ void CausalConvDecodeKernel(
int batch_channels, // = batch_size * channels (actual element count)
int channels,
int kernel_size,
+ int dilation,
+ CausalConvLayout act_layout,
+ CausalConvLayout state_layout,
bool apply_silu,
int state_window) { // W: axis-0 extent of past_state / present_state (>= 1)
const int bc = blockIdx.x * blockDim.x + threadIdx.x;
@@ -54,14 +58,17 @@ __global__ void CausalConvDecodeKernel(
const int b = bc / channels;
const int c = bc % channels;
- const int pad = kernel_size - 1;
+ const int pad = (kernel_size - 1) * dilation;
+ const int64_t state_pos_stride = state_layout.pos_stride;
// Cache input value in register — avoids redundant global reads
- const float input_val = to_float(input[(int64_t)b * channels + c]);
+ const int64_t act_offset = act_layout.Offset(b, 0, c);
+ const float input_val = to_float(input[act_offset]);
// seq_len == 1, so the single position is window slot W-1 for both the read and the write.
- // Window-major [W, B, C, K-1]: slot stride is batch_channels*pad and (b, c) flattens to bc.
- const int64_t state_offset = (int64_t)(state_window - 1) * batch_channels * pad + (int64_t)bc * pad;
+ // Window-major: one slot is a whole [B, ...] block of batch_channels*pad elements.
+ const int64_t state_offset =
+ (int64_t)(state_window - 1) * batch_channels * pad + state_layout.Offset(b, 0, c);
// Cache past_state base pointer for this (b, c)
const T* ps_in = (past_state != nullptr) ? past_state + state_offset : nullptr;
@@ -70,27 +77,28 @@ __global__ void CausalConvDecodeKernel(
// weight layout: [C, 1, K], so channel c starts at c * K
float sum = (bias != nullptr) ? to_float(bias[c]) : 0.0f;
- // Convolution window: [past_state[0..K-2], input[0]]
- for (int k = 0; k < pad; ++k) {
+ // Convolution window: [past_state(pad), input[0]]. Tap k reads window slot k*dilation.
+ for (int k = 0; k < kernel_size - 1; ++k) {
float wk = to_float(weight[c * kernel_size + k]);
- float xk = (ps_in != nullptr) ? to_float(ps_in[k]) : 0.0f;
+ float xk = (ps_in != nullptr) ? to_float(ps_in[k * dilation * state_pos_stride]) : 0.0f;
sum += wk * xk;
}
- // Last element of window is current input
- sum += to_float(weight[c * kernel_size + pad]) * input_val;
+ // Last tap is the current input
+ sum += to_float(weight[c * kernel_size + kernel_size - 1]) * input_val;
if (apply_silu) {
sum = silu_fn(sum);
}
- output[(int64_t)b * channels + c] = from_float(sum);
+ output[act_offset] = from_float(sum);
// Update present_state: shift left by 1, append input
T* ps_out = present_state + state_offset;
for (int k = 0; k < pad - 1; ++k) {
- ps_out[k] = (ps_in != nullptr) ? ps_in[k + 1] : from_float(0.0f);
+ ps_out[k * state_pos_stride] =
+ (ps_in != nullptr) ? ps_in[(k + 1) * state_pos_stride] : from_float(0.0f);
}
if (pad > 0) {
- ps_out[pad - 1] = from_float(input_val);
+ ps_out[(pad - 1) * state_pos_stride] = from_float(input_val);
}
}
@@ -162,6 +170,9 @@ __global__ void CausalConvPrefillKernel(
int seq_len,
int channels,
int kernel_size,
+ int dilation,
+ CausalConvLayout act_layout,
+ CausalConvLayout state_layout,
bool apply_silu,
int batch_size,
int state_window) { // W: axis-0 extent of past_state / present_state (>= 1)
@@ -169,14 +180,15 @@ __global__ void CausalConvPrefillKernel(
const int c = blockIdx.y;
const int tid = threadIdx.x;
- const int pad = kernel_size - 1;
+ const int pad = (kernel_size - 1) * dilation;
const int padded_len = pad + seq_len;
+ const int64_t state_pos_stride = state_layout.pos_stride;
// Slot W-1 holds the state after the last token; that is what past_state is read from.
// Window-major, so one slot spans the whole batch.
const int64_t slot_stride = (int64_t)batch_size * channels * pad;
const int64_t last_slot_offset =
- (int64_t)(state_window - 1) * slot_stride + ((int64_t)b * channels + c) * pad;
+ (int64_t)(state_window - 1) * slot_stride + state_layout.Offset(b, 0, c);
// Shared memory: padded input [pad + L] floats + weight [K] floats
extern __shared__ float smem[];
@@ -187,14 +199,14 @@ __global__ void CausalConvPrefillKernel(
// Past state portion: [0..pad-1]
for (int i = tid; i < pad; i += blockDim.x) {
if (past_state != nullptr) {
- s_padded[i] = to_float(past_state[last_slot_offset + i]);
+ s_padded[i] = to_float(past_state[last_slot_offset + (int64_t)i * state_pos_stride]);
} else {
s_padded[i] = 0.0f;
}
}
// Current input portion: [pad..pad+L-1]
for (int i = tid; i < seq_len; i += blockDim.x) {
- s_padded[pad + i] = to_float(input[((int64_t)b * channels + c) * seq_len + i]);
+ s_padded[pad + i] = to_float(input[act_layout.Offset(b, i, c)]);
}
// Load weight into shared memory
for (int i = tid; i < kernel_size; i += blockDim.x) {
@@ -207,12 +219,12 @@ __global__ void CausalConvPrefillKernel(
for (int l = tid; l < seq_len; l += blockDim.x) {
float sum = bias_val;
for (int k = 0; k < kernel_size; ++k) {
- sum += s_weight[k] * s_padded[l + k];
+ sum += s_weight[k] * s_padded[l + k * dilation];
}
if (apply_silu) {
sum = silu_fn(sum);
}
- output[((int64_t)b * channels + c) * seq_len + l] = from_float(sum);
+ output[act_layout.Offset(b, l, c)] = from_float(sum);
}
// Save present_state. The carry state after token t is the pad-length window ending at position
@@ -223,9 +235,9 @@ __global__ void CausalConvPrefillKernel(
const int first = seq_len > state_window ? seq_len - state_window : 0;
for (int t = first + tid; t < seq_len; t += blockDim.x) {
T* ps = present_state + (int64_t)(t + state_window - seq_len) * slot_stride +
- ((int64_t)b * channels + c) * pad;
+ state_layout.Offset(b, 0, c);
for (int p = 0; p < pad; ++p) {
- ps[p] = from_float(s_padded[t + 1 + p]);
+ ps[(int64_t)p * state_pos_stride] = from_float(s_padded[t + 1 + p]);
}
}
}
@@ -251,6 +263,9 @@ __global__ void CausalConvPrefillKernelBatched(
int seq_len,
int channels,
int kernel_size,
+ int dilation,
+ CausalConvLayout act_layout,
+ CausalConvLayout state_layout,
bool apply_silu,
int batch_size,
int state_window) { // W: axis-0 extent of past_state / present_state (>= 1)
@@ -258,10 +273,11 @@ __global__ void CausalConvPrefillKernelBatched(
const int c_base = blockIdx.y * CPB;
const int tid = threadIdx.x;
- const int pad = kernel_size - 1;
+ const int pad = (kernel_size - 1) * dilation;
const int padded_len = pad + seq_len;
+ const int64_t state_pos_stride = state_layout.pos_stride;
- // Window-major [W, B, C, K-1]: one slot spans the whole batch.
+ // Window-major: one slot spans the whole batch.
const int64_t slot_stride = (int64_t)batch_size * channels * pad;
// Which channel within this block's CPB group does this thread serve?
@@ -279,17 +295,17 @@ __global__ void CausalConvPrefillKernelBatched(
if (c < channels) {
// Load past state from window slot W-1 (the state after the last token of the previous step)
const int64_t last_slot_offset =
- (int64_t)(state_window - 1) * slot_stride + ((int64_t)b * channels + c) * pad;
+ (int64_t)(state_window - 1) * slot_stride + state_layout.Offset(b, 0, c);
for (int i = local_tid; i < pad; i += threads_per_channel) {
if (past_state != nullptr) {
- s_padded[i] = to_float(past_state[last_slot_offset + i]);
+ s_padded[i] = to_float(past_state[last_slot_offset + (int64_t)i * state_pos_stride]);
} else {
s_padded[i] = 0.0f;
}
}
// Load input
for (int i = local_tid; i < seq_len; i += threads_per_channel) {
- s_padded[pad + i] = to_float(input[((int64_t)b * channels + c) * seq_len + i]);
+ s_padded[pad + i] = to_float(input[act_layout.Offset(b, i, c)]);
}
// Load weight
for (int i = local_tid; i < kernel_size; i += threads_per_channel) {
@@ -303,12 +319,12 @@ __global__ void CausalConvPrefillKernelBatched(
for (int l = local_tid; l < seq_len; l += threads_per_channel) {
float sum = bias_val;
for (int k = 0; k < kernel_size; ++k) {
- sum += s_weight[k] * s_padded[l + k];
+ sum += s_weight[k] * s_padded[l + k * dilation];
}
if (apply_silu) {
sum = silu_fn(sum);
}
- output[((int64_t)b * channels + c) * seq_len + l] = from_float(sum);
+ output[act_layout.Offset(b, l, c)] = from_float(sum);
}
}
@@ -323,14 +339,93 @@ __global__ void CausalConvPrefillKernelBatched(
const int first = seq_len > state_window ? seq_len - state_window : 0;
for (int t = first + local_tid; t < seq_len; t += threads_per_channel) {
T* ps = present_state + (int64_t)(t + state_window - seq_len) * slot_stride +
- ((int64_t)b * channels + c) * pad;
+ state_layout.Offset(b, 0, c);
for (int p = 0; p < pad; ++p) {
- ps[p] = from_float(s_padded[t + 1 + p]);
+ ps[(int64_t)p * state_pos_stride] = from_float(s_padded[t + 1 + p]);
}
}
}
}
+// =============================================================================
+// Channels-last prefill kernel: L>1, one thread per (batch, position, channel)
+// with the channel as the fastest-moving thread axis.
+//
+// Grid: (ceil(channels / threads), seq_len, batch_size)
+// Block: (threads, 1, 1)
+// Shared memory: none
+//
+// The shared-memory kernels above stage one channel per block and walk positions, which is
+// coalesced only when positions are contiguous. Under channels_last the contiguous axis is the
+// channel, so that access pattern turns every load and store into a strided gather. Here adjacent
+// threads hold adjacent channels, so each convolution tap, each state read and every store is a
+// single contiguous transaction. Staging is unnecessary because the overlapping taps of
+// neighbouring positions are served by L1/L2 rather than shared memory.
+// =============================================================================
+template
+__global__ void CausalConvPrefillKernelChannelsLast(
+ const T* __restrict__ input, // [B, L, C]
+ const T* __restrict__ weight, // [C, 1, K]
+ const T* __restrict__ bias, // [C] or nullptr
+ const T* __restrict__ past_state, // [W, B, K-1, C] or nullptr
+ T* __restrict__ output, // [B, L, C]
+ T* __restrict__ present_state, // [W, B, K-1, C]
+ int seq_len,
+ int channels,
+ int kernel_size,
+ int dilation,
+ CausalConvLayout act_layout,
+ CausalConvLayout state_layout,
+ bool apply_silu,
+ int batch_size,
+ int state_window) { // W: axis-0 extent of past_state / present_state (>= 1)
+ const int c = blockIdx.x * blockDim.x + threadIdx.x;
+ if (c >= channels) {
+ return;
+ }
+ const int l = blockIdx.y;
+ const int b = blockIdx.z;
+
+ const int pad = (kernel_size - 1) * dilation;
+ const int64_t state_pos_stride = state_layout.pos_stride;
+ // Window-major: one slot spans the whole batch. Slot W-1 holds the state after the last token of
+ // the previous step, which is the only slot past_state is read from.
+ const int64_t slot_stride = (int64_t)batch_size * channels * pad;
+ const int64_t last_slot_offset =
+ (int64_t)(state_window - 1) * slot_stride + state_layout.Offset(b, 0, c);
+
+ // Reads the virtual stream [past_state (pad samples), input (seq_len samples)] at index `vp`.
+ auto sample = [&](int vp) -> float {
+ if (vp >= pad) {
+ return to_float(input[act_layout.Offset(b, vp - pad, c)]);
+ }
+ return past_state != nullptr
+ ? to_float(past_state[last_slot_offset + (int64_t)vp * state_pos_stride])
+ : 0.0f;
+ };
+
+ float sum = (bias != nullptr) ? to_float(bias[c]) : 0.0f;
+ for (int k = 0; k < kernel_size; ++k) {
+ sum += to_float(weight[(int64_t)c * kernel_size + k]) * sample(l + k * dilation);
+ }
+ if (apply_silu) {
+ sum = silu_fn(sum);
+ }
+ output[act_layout.Offset(b, l, c)] = from_float(sum);
+
+ // The carry state after token l is the pad-length window ending at that token, i.e. virtual
+ // stream positions [l + 1, l + pad]. It goes into the right-aligned slot l + W - seq_len;
+ // earlier tokens fall outside the window. The last token always maps to slot W-1.
+ const int first = seq_len > state_window ? seq_len - state_window : 0;
+ if (l >= first) {
+ T* ps = present_state + (int64_t)(l + state_window - seq_len) * slot_stride +
+ state_layout.Offset(b, 0, c);
+ for (int p = 0; p < pad; ++p) {
+ ps[(int64_t)p * state_pos_stride] = from_float(sample(l + 1 + p));
+ }
+ }
+}
+
} // anonymous namespace
template
@@ -346,6 +441,9 @@ Status LaunchCausalConvWithStateKernel(
int channels,
int seq_len,
int kernel_size,
+ int dilation,
+ CausalConvLayout act_layout,
+ CausalConvLayout state_layout,
bool apply_silu,
int max_threads_per_block,
int state_window) {
@@ -354,7 +452,10 @@ Status LaunchCausalConvWithStateKernel(
int total = batch_size * channels;
int threads = 256;
int blocks = (total + threads - 1) / threads;
- switch (kernel_size) {
+ // The fixed-K decode kernels hard-code pad == K - 1 and a contiguous state row, so they only
+ // apply to the undilated channels-first case (pos_stride is 1 there, and `channels` in the
+ // channels-last one).
+ switch ((dilation == 1 && state_layout.pos_stride == 1) ? kernel_size : 0) {
case 2:
CausalConvDecodeKernelFixedK<<>>(
input, weight, bias, past_state, output, present_state,
@@ -378,12 +479,29 @@ Status LaunchCausalConvWithStateKernel(
default:
CausalConvDecodeKernel<<>>(
input, weight, bias, past_state, output, present_state,
- total, channels, kernel_size, apply_silu, state_window);
+ total, channels, kernel_size, dilation, act_layout, state_layout, apply_silu,
+ state_window);
break;
}
} else {
// Prefill path: choose between batched (short seq) or single-channel (long seq) kernel
- int pad = kernel_size - 1;
+ int pad = (kernel_size - 1) * dilation;
+
+ // Under channels_last the contiguous axis is the channel, so use the kernel whose fastest
+ // thread axis is the channel; the shared-memory kernels below would gather every access.
+ // gridDim.y is capped at 65535; longer sequences fall through to the shared-memory kernels,
+ // which handle either layout correctly (just less efficiently).
+ constexpr int kMaxGridDimY = 65535;
+ if (act_layout.chan_stride == 1 && seq_len <= kMaxGridDimY) {
+ int threads = channels >= 256 ? 256 : ((channels + 31) / 32) * 32;
+ threads = std::min(threads, max_threads_per_block);
+ const dim3 grid((channels + threads - 1) / threads, seq_len, batch_size);
+ CausalConvPrefillKernelChannelsLast<<>>(
+ input, weight, bias, past_state, output, present_state,
+ seq_len, channels, kernel_size, dilation, act_layout, state_layout, apply_silu,
+ batch_size, state_window);
+ return CUDA_CALL(cudaGetLastError());
+ }
// For short sequences, batch multiple channels per block to improve occupancy.
// CPB=4: each block handles 4 channels, reducing block count by 4x.
@@ -416,7 +534,8 @@ Status LaunchCausalConvWithStateKernel(
CausalConvPrefillKernelBatched<<>>(
input, weight, bias, past_state, output, present_state,
- seq_len, channels, kernel_size, apply_silu, batch_size, state_window);
+ seq_len, channels, kernel_size, dilation, act_layout, state_layout, apply_silu,
+ batch_size, state_window);
} else {
// Original single-channel-per-block path for long sequences
const dim3 grid(batch_size, channels, 1);
@@ -440,7 +559,8 @@ Status LaunchCausalConvWithStateKernel(
CausalConvPrefillKernel<<>>(
input, weight, bias, past_state, output, present_state,
- seq_len, channels, kernel_size, apply_silu, batch_size, state_window);
+ seq_len, channels, kernel_size, dilation, act_layout, state_layout, apply_silu,
+ batch_size, state_window);
}
}
@@ -450,16 +570,17 @@ Status LaunchCausalConvWithStateKernel(
// Explicit instantiations
template Status LaunchCausalConvWithStateKernel(
cudaStream_t, const float*, const float*, const float*, const float*,
- float*, float*, int, int, int, int, bool, int, int);
+ float*, float*, int, int, int, int, int, CausalConvLayout, CausalConvLayout, bool, int, int);
template Status LaunchCausalConvWithStateKernel(
cudaStream_t, const half*, const half*, const half*, const half*,
- half*, half*, int, int, int, int, bool, int, int);
+ half*, half*, int, int, int, int, int, CausalConvLayout, CausalConvLayout, bool, int, int);
#if __CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)
template Status LaunchCausalConvWithStateKernel<__nv_bfloat16>(
cudaStream_t, const __nv_bfloat16*, const __nv_bfloat16*, const __nv_bfloat16*, const __nv_bfloat16*,
- __nv_bfloat16*, __nv_bfloat16*, int, int, int, int, bool, int, int);
+ __nv_bfloat16*, __nv_bfloat16*, int, int, int, int, int, CausalConvLayout, CausalConvLayout,
+ bool, int, int);
#endif
} // namespace cuda
diff --git a/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state_impl.h b/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state_impl.h
index b07730e6c0d98..968cb9507339d 100644
--- a/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state_impl.h
+++ b/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state_impl.h
@@ -10,6 +10,27 @@ namespace onnxruntime {
namespace contrib {
namespace cuda {
+// Element strides of a (batch, position, channel) view over an activation or state tensor.
+// Both supported layouts are dense, so a single strided view covers them and one kernel body
+// serves the channels-first (batch, channels, length) and channels-last (batch, length, channels)
+// layouts. For the state tensors these are the strides *within* one state_window slot.
+struct CausalConvLayout {
+ int64_t batch_stride;
+ int64_t pos_stride;
+ int64_t chan_stride;
+
+ __host__ __device__ int64_t Offset(int b, int pos, int c) const {
+ return static_cast(b) * batch_stride + static_cast(pos) * pos_stride +
+ static_cast(c) * chan_stride;
+ }
+};
+
+// `length` is the extent of the position axis: seq_len for activations, state_length for state.
+inline CausalConvLayout MakeCausalConvLayout(bool channels_last, int64_t channels, int64_t length) {
+ return channels_last ? CausalConvLayout{channels * length, channels, 1}
+ : CausalConvLayout{channels * length, 1, length};
+}
+
// Fused causal depthwise conv1d + activation + state management.
// One thread block per (batch, channel). For decode (L=1), this is a simple
// dot product from shared memory. For prefill (L>1), each thread handles
@@ -20,17 +41,20 @@ Status LaunchCausalConvWithStateKernel(
const T* input, // [B, C, L]
const T* weight, // [C, 1, K]
const T* bias, // [C] or nullptr
- const T* past_state, // [W, B, C, K-1] or nullptr
+ const T* past_state, // [W, B, C, (K-1)*dilation] or nullptr
T* output, // [B, C, L]
- T* present_state, // [W, B, C, K-1]
+ T* present_state, // [W, B, C, (K-1)*dilation]
int batch_size,
int channels,
int seq_len,
int kernel_size,
+ int dilation, // spacing between kernel taps along the causal axis (>= 1)
+ CausalConvLayout act_layout, // strides of input / output
+ CausalConvLayout state_layout, // strides within one past_state / present_state slot
bool apply_silu,
int max_threads_per_block,
// Axis-0 extent W of past_state / present_state (>= 1). The window axis leads the batch axis
- // so that a slot is one contiguous [B, C, K-1] block. Right-aligned: token t writes slot
+ // so that a slot is one contiguous [B, C, (K-1)*dilation] block. Right-aligned: token t writes slot
// t + W - seq_len and negative slots are skipped, so slot W-1 always holds the state after the
// last token and is the slot past_state is read from. Pass 1 for a plain single-state tensor
// with no window axis.
diff --git a/onnxruntime/contrib_ops/cuda/bert/engram_gate.cc b/onnxruntime/contrib_ops/cuda/bert/engram_gate.cc
new file mode 100644
index 0000000000000..a108ab3b9ae46
--- /dev/null
+++ b/onnxruntime/contrib_ops/cuda/bert/engram_gate.cc
@@ -0,0 +1,88 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "contrib_ops/cuda/bert/engram_gate.h"
+#include "contrib_ops/cuda/bert/engram_gate_impl.h"
+#include "core/providers/cuda/cuda_common.h"
+#include "core/providers/cuda/cuda_type_conversion.h"
+
+namespace onnxruntime {
+namespace contrib {
+namespace cuda {
+
+using namespace onnxruntime::cuda;
+
+#define REGISTER_KERNEL_TYPED(T) \
+ ONNX_OPERATOR_TYPED_KERNEL_EX( \
+ EngramGate, \
+ kMSDomain, \
+ 1, \
+ T, \
+ kCudaExecutionProvider, \
+ (*KernelDefBuilder::Create()) \
+ .TypeConstraint("T", DataTypeImpl::GetTensorType()), \
+ EngramGate);
+
+REGISTER_KERNEL_TYPED(float)
+REGISTER_KERNEL_TYPED(MLFloat16)
+REGISTER_KERNEL_TYPED(BFloat16)
+
+#undef REGISTER_KERNEL_TYPED
+
+template
+EngramGate::EngramGate(const OpKernelInfo& info) : CudaKernel(info) {
+ epsilon_ = info.GetAttrOrDefault("epsilon", 1.0e-5f);
+}
+
+template
+Status EngramGate::ComputeInternal(OpKernelContext* context) const {
+ using CudaT = typename OrtToCudaType::type;
+ const Tensor* key = context->Input(0);
+ const Tensor* query = context->Input(1);
+ const Tensor* value = context->Input(2);
+ const Tensor* key_norm_scale = context->Input(3);
+ const Tensor* query_norm_scale = context->Input(4);
+
+ const TensorShape& key_shape = key->Shape();
+ ORT_RETURN_IF_NOT(key_shape.NumDimensions() == 4,
+ "key must have shape (batch_size, sequence_length, hc_mult, hidden_size)");
+ const int64_t batch_size = key_shape[0];
+ const int64_t sequence_length = key_shape[1];
+ const int64_t hc_mult = key_shape[2];
+ const int64_t hidden_size = key_shape[3];
+
+ ORT_RETURN_IF_NOT(query->Shape() == key_shape, "query must have the same shape as key");
+ ORT_RETURN_IF_NOT(value->Shape() == TensorShape({batch_size, sequence_length, hidden_size}),
+ "value must have shape (batch_size, sequence_length, hidden_size)");
+ ORT_RETURN_IF_NOT(key_norm_scale->Shape() == TensorShape({hc_mult, hidden_size}),
+ "key_norm_scale must have shape (hc_mult, hidden_size)");
+ ORT_RETURN_IF_NOT(query_norm_scale->Shape() == TensorShape({hc_mult, hidden_size}),
+ "query_norm_scale must have shape (hc_mult, hidden_size)");
+
+ Tensor* output = context->Output(0, key_shape);
+ if (key_shape.Size() == 0) {
+ return Status::OK();
+ }
+
+ return LaunchEngramGateKernel(
+ Stream(context),
+ reinterpret_cast(key->Data()),
+ reinterpret_cast(query->Data()),
+ reinterpret_cast(value->Data()),
+ reinterpret_cast(key_norm_scale->Data()),
+ reinterpret_cast(query_norm_scale->Data()),
+ reinterpret_cast(output->MutableData()),
+ batch_size,
+ sequence_length,
+ hc_mult,
+ hidden_size,
+ epsilon_);
+}
+
+template class EngramGate;
+template class EngramGate;
+template class EngramGate;
+
+} // namespace cuda
+} // namespace contrib
+} // namespace onnxruntime
diff --git a/onnxruntime/contrib_ops/cuda/bert/engram_gate.h b/onnxruntime/contrib_ops/cuda/bert/engram_gate.h
new file mode 100644
index 0000000000000..09aa7018008f2
--- /dev/null
+++ b/onnxruntime/contrib_ops/cuda/bert/engram_gate.h
@@ -0,0 +1,25 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include "core/common/common.h"
+#include "core/providers/cuda/cuda_kernel.h"
+
+namespace onnxruntime {
+namespace contrib {
+namespace cuda {
+
+template
+class EngramGate final : public onnxruntime::cuda::CudaKernel {
+ public:
+ explicit EngramGate(const OpKernelInfo& info);
+ Status ComputeInternal(OpKernelContext* context) const override;
+
+ private:
+ float epsilon_;
+};
+
+} // namespace cuda
+} // namespace contrib
+} // namespace onnxruntime
diff --git a/onnxruntime/contrib_ops/cuda/bert/engram_gate_impl.cu b/onnxruntime/contrib_ops/cuda/bert/engram_gate_impl.cu
new file mode 100644
index 0000000000000..e0dd2a7e31a56
--- /dev/null
+++ b/onnxruntime/contrib_ops/cuda/bert/engram_gate_impl.cu
@@ -0,0 +1,113 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "contrib_ops/cuda/bert/engram_gate_impl.h"
+
+#include
+#include
+#include
+
+#include
+
+#include "contrib_ops/cuda/bert/engram_helper.cuh"
+#include "core/providers/cuda/cu_inc/cuda_type_helper.cuh"
+
+namespace onnxruntime {
+namespace contrib {
+namespace cuda {
+
+namespace {
+
+// One block per (token, g) row. The gate is a scalar for the whole row, so it is reduced once by the
+// block and then broadcast over the value channels.
+template
+__global__ void EngramGateKernel(
+ const T* key,
+ const T* query,
+ const T* value,
+ const T* key_norm_scale,
+ const T* query_norm_scale,
+ T* output,
+ int64_t rows,
+ int64_t hc_mult,
+ int64_t hidden_size,
+ float epsilon) {
+ extern __shared__ float shared[];
+
+ for (int64_t row = blockIdx.x; row < rows; row += gridDim.x) {
+ const int64_t g = row % hc_mult;
+ const int64_t token = row / hc_mult;
+ const T* key_row = key + row * hidden_size;
+ const T* query_row = query + row * hidden_size;
+ const T* value_row = value + token * hidden_size;
+ const T* key_scale_g = key_norm_scale + g * hidden_size;
+ const T* query_scale_g = query_norm_scale + g * hidden_size;
+
+ float key_sum_sq = 0.0f;
+ float query_sum_sq = 0.0f;
+ float dot_numerator = 0.0f;
+
+ for (int64_t d = threadIdx.x; d < hidden_size; d += blockDim.x) {
+ const float key_value = to_float(key_row[d]);
+ const float query_value = to_float(query_row[d]);
+ key_sum_sq += key_value * key_value;
+ query_sum_sq += query_value * query_value;
+ dot_numerator += key_value * to_float(key_scale_g[d]) * query_value * to_float(query_scale_g[d]);
+ }
+
+ // The three partials are independent and available at the same point, so fuse them into one tree
+ // reduction instead of paying three sets of barriers per row.
+ engram_helper::BlockSum3(&key_sum_sq, &query_sum_sq, &dot_numerator, shared);
+
+ const float key_inv_rms = rsqrtf(key_sum_sq / static_cast(hidden_size) + epsilon);
+ const float query_inv_rms = rsqrtf(query_sum_sq / static_cast(hidden_size) + epsilon);
+ const float dot = dot_numerator * key_inv_rms * query_inv_rms / sqrtf(static_cast(hidden_size));
+ const float gate = engram_helper::SigmoidFloat(engram_helper::EngramGateArg(dot));
+
+ T* output_row = output + row * hidden_size;
+ for (int64_t c = threadIdx.x; c < hidden_size; c += blockDim.x) {
+ output_row[c] = from_float(gate * to_float(value_row[c]));
+ }
+ }
+}
+
+} // namespace
+
+template
+Status LaunchEngramGateKernel(
+ cudaStream_t stream,
+ const T* key,
+ const T* query,
+ const T* value,
+ const T* key_norm_scale,
+ const T* query_norm_scale,
+ T* output,
+ int64_t batch_size,
+ int64_t sequence_length,
+ int64_t hc_mult,
+ int64_t hidden_size,
+ float epsilon) {
+ const int64_t rows = batch_size * sequence_length * hc_mult;
+ if (rows == 0 || hidden_size == 0) {
+ return Status::OK();
+ }
+ const int blocks = static_cast(std::min(rows, engram_helper::kMaxGridDimX));
+ const size_t shared_bytes = 3 * static_cast(engram_helper::kThreads) * sizeof(float);
+ EngramGateKernel<<>>(
+ key, query, value, key_norm_scale, query_norm_scale, output, rows, hc_mult, hidden_size, epsilon);
+ return CUDA_CALL(cudaGetLastError());
+}
+
+#define INSTANTIATE_ENGRAM_GATE(T) \
+ template Status LaunchEngramGateKernel(cudaStream_t, const T*, const T*, const T*, const T*, \
+ const T*, T*, int64_t, int64_t, int64_t, int64_t, float);
+
+INSTANTIATE_ENGRAM_GATE(float)
+INSTANTIATE_ENGRAM_GATE(half)
+INSTANTIATE_ENGRAM_GATE(__nv_bfloat16)
+
+#undef INSTANTIATE_ENGRAM_GATE
+
+} // namespace cuda
+} // namespace contrib
+} // namespace onnxruntime
diff --git a/onnxruntime/contrib_ops/cuda/bert/engram_gate_impl.h b/onnxruntime/contrib_ops/cuda/bert/engram_gate_impl.h
new file mode 100644
index 0000000000000..83e2ec8fe77ac
--- /dev/null
+++ b/onnxruntime/contrib_ops/cuda/bert/engram_gate_impl.h
@@ -0,0 +1,30 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include
+#include "core/providers/cuda/cuda_common.h"
+
+namespace onnxruntime {
+namespace contrib {
+namespace cuda {
+
+template
+Status LaunchEngramGateKernel(
+ cudaStream_t stream,
+ const T* key,
+ const T* query,
+ const T* value,
+ const T* key_norm_scale,
+ const T* query_norm_scale,
+ T* output,
+ int64_t batch_size,
+ int64_t sequence_length,
+ int64_t hc_mult,
+ int64_t hidden_size,
+ float epsilon);
+
+} // namespace cuda
+} // namespace contrib
+} // namespace onnxruntime
diff --git a/onnxruntime/contrib_ops/cuda/bert/engram_helper.cuh b/onnxruntime/contrib_ops/cuda/bert/engram_helper.cuh
new file mode 100644
index 0000000000000..872e6955aaf4f
--- /dev/null
+++ b/onnxruntime/contrib_ops/cuda/bert/engram_helper.cuh
@@ -0,0 +1,93 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include
+#include
+#include
+
+namespace onnxruntime {
+namespace contrib {
+namespace cuda {
+namespace engram_helper {
+
+constexpr int kThreads = 256;
+// grid.x is limited to 2^31 - 1 since compute capability 3.0 (the 65535 limit applies to grid.y and
+// grid.z only). All kernels launched through GridSize() use a grid-stride loop, so the clamp only
+// bounds the launch; correctness does not depend on it.
+constexpr int64_t kMaxGridDimX = 2147483647;
+
+// Number of blocks for a grid-stride loop over `count` elements, clamped to the maximum grid size.
+inline int GridSize(int64_t count) {
+ const int64_t blocks = (count + kThreads - 1) / kThreads;
+ return static_cast(std::min(blocks, kMaxGridDimX));
+}
+
+// Sums three independent per-thread partials across the block in a single tree reduction, so a row
+// that needs three reductions pays one set of barriers instead of three. `shared` must point to at
+// least 3 * blockDim.x floats, and blockDim.x must be a power of two. All threads must call this.
+__device__ __forceinline__ void BlockSum3(float* a, float* b, float* c, float* shared) {
+ float* shared_a = shared;
+ float* shared_b = shared + blockDim.x;
+ float* shared_c = shared + 2 * blockDim.x;
+ shared_a[threadIdx.x] = *a;
+ shared_b[threadIdx.x] = *b;
+ shared_c[threadIdx.x] = *c;
+ __syncthreads();
+ for (unsigned int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
+ if (threadIdx.x < stride) {
+ shared_a[threadIdx.x] += shared_a[threadIdx.x + stride];
+ shared_b[threadIdx.x] += shared_b[threadIdx.x + stride];
+ shared_c[threadIdx.x] += shared_c[threadIdx.x + stride];
+ }
+ __syncthreads();
+ }
+ *a = shared_a[0];
+ *b = shared_b[0];
+ *c = shared_c[0];
+ __syncthreads();
+}
+
+// Numerically stable logistic function.
+__device__ __forceinline__ float SigmoidFloat(float x) {
+ return x > 0.0f ? 1.0f / (1.0f + expf(-x)) : expf(x) / (1.0f + expf(x));
+}
+
+// Engram gate pre-activation: sign(dot) * sqrt(max(abs(dot), 1e-6)).
+// copysignf cannot be used here because it maps a zero dot product to +sqrt(1e-6) instead of zero,
+// which would disagree with the schema formula and with the other execution providers.
+__device__ __forceinline__ float EngramGateArg(float dot) {
+ if (dot == 0.0f) {
+ return 0.0f;
+ }
+ const float magnitude = sqrtf(fmaxf(fabsf(dot), 1.0e-6f));
+ return dot < 0.0f ? -magnitude : magnitude;
+}
+
+// Euclidean modulo: the result always has the sign of `mod`, which must be positive.
+template
+__device__ __forceinline__ T PositiveMod(T value, T mod) {
+ const T result = value % mod;
+ return result < 0 ? static_cast(result + mod) : result;
+}
+
+// Multiplies through the unsigned counterpart of T so that overflow wraps around instead of
+// being undefined behavior.
+template
+__device__ __forceinline__ T WrappedMultiply(T a, T b);
+
+template <>
+__device__ __forceinline__ int32_t WrappedMultiply(int32_t a, int32_t b) {
+ return static_cast(static_cast(a) * static_cast(b));
+}
+
+template <>
+__device__ __forceinline__ int64_t WrappedMultiply(int64_t a, int64_t b) {
+ return static_cast(static_cast(a) * static_cast(b));
+}
+
+} // namespace engram_helper
+} // namespace cuda
+} // namespace contrib
+} // namespace onnxruntime
diff --git a/onnxruntime/contrib_ops/cuda/bert/ngram_hash_mapping.cc b/onnxruntime/contrib_ops/cuda/bert/ngram_hash_mapping.cc
new file mode 100644
index 0000000000000..36ef5f6c40ba1
--- /dev/null
+++ b/onnxruntime/contrib_ops/cuda/bert/ngram_hash_mapping.cc
@@ -0,0 +1,95 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "contrib_ops/cuda/bert/ngram_hash_mapping.h"
+#include "contrib_ops/cuda/bert/ngram_hash_mapping_impl.h"
+#include "core/providers/cuda/cuda_common.h"
+
+#include
+
+namespace onnxruntime {
+namespace contrib {
+namespace cuda {
+
+using namespace onnxruntime::cuda;
+
+#define REGISTER_KERNEL_TYPED(T) \
+ ONNX_OPERATOR_TYPED_KERNEL_EX( \
+ NGramHashMapping, \
+ kMSDomain, \
+ 1, \
+ T, \
+ kCudaExecutionProvider, \
+ (*KernelDefBuilder::Create()) \
+ .MayInplace(3, 1) \
+ .TypeConstraint("M", DataTypeImpl::GetTensorType()), \
+ NGramHashMapping);
+
+REGISTER_KERNEL_TYPED(int32_t)
+REGISTER_KERNEL_TYPED(int64_t)
+
+#undef REGISTER_KERNEL_TYPED
+
+template
+NGramHashMapping::NGramHashMapping(const OpKernelInfo& info) : CudaKernel(info) {
+ ORT_ENFORCE(info.GetAttr("max_ngram_size", &max_ngram_size_).IsOK(),
+ "max_ngram_size attribute is required");
+ ORT_ENFORCE(info.GetAttr("n_head_per_ngram", &n_head_per_ngram_).IsOK(),
+ "n_head_per_ngram attribute is required");
+ int64_t pad_id = 0;
+ ORT_ENFORCE(info.GetAttr("pad_id", &pad_id).IsOK(), "pad_id attribute is required");
+ ORT_ENFORCE(max_ngram_size_ >= 2, "max_ngram_size must be at least 2");
+ ORT_ENFORCE(n_head_per_ngram_ >= 1, "n_head_per_ngram must be positive");
+ ORT_ENFORCE(pad_id >= static_cast(std::numeric_limits::min()) &&
+ pad_id <= static_cast(std::numeric_limits::max()),
+ "pad_id is out of range for the input id type");
+ pad_id_ = static_cast(pad_id);
+}
+
+template
+Status NGramHashMapping