diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md
index 6dcb758a234fe..c811c293b12c2 100644
--- a/docs/ContribOperators.md
+++ b/docs/ContribOperators.md
@@ -1025,8 +1025,18 @@ This version of the operator has been available since version 1 of the 'com.micr
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).
+input_channel_offset : int
+Starting channel of the channels-last input slice (CUDA only).
+input_channels : int
+CUDA channels-last input channel slice width; 0 selects all channels. Nondefault extended attributes require rank-3 input and state_window=0.
ndim : int
Spatial dimensionality: 1, 2, or 3. Default is 1.
+state_channels_last : int
+State layout override (CUDA only): -1 follows channels_last, 0 uses [B,C,state_length], 1 uses [B,state_length,C].
+state_length : int
+Carry history length (CUDA only); 0 uses (kernel_size-1)*dilation. An explicit length must cover the receptive field; convolution reads its trailing history.
+state_update_capacity : int
+Compact speculative state-update capture capacity (CUDA extended path only, 0..8); when > 0, output 2 has shape (batch_size, state_update_capacity, channels) and records the first min(sequence_length, state_update_capacity) projected input rows and any remaining rows are zero.
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, 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].
@@ -1035,22 +1045,24 @@ This version of the operator has been available since version 1 of the 'com.micr
- input : T
-- 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.
+- 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. Extended channel selection uses rank-3 (batch_size, sequence_length, input_width).
- 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, 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.
+- 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. The state_length and state_channels_last attributes may override the length and layout for the CUDA extended path. If not provided, padding is zero.
-#### Outputs
+#### Outputs (2 - 3)
- output : T
-- Convolution output with same shape as input.
+- Convolution output with same shape as input, except the final dimension is input_channels when extended channel selection is enabled.
- present_state : T
-- 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).
+- 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). Extended attributes override the state length and layout as for past_state.
+- state_update (optional) : T
+- Optional compact state-update capture with shape (batch_size, state_update_capacity, channels) when state_update_capacity > 0. CUDA extended path only; requires rank-3 channels-last input and state_window=0.
#### Type Constraints
@@ -2493,7 +2505,7 @@ This version of the operator has been available since version 1 of the 'com.micr
per-chunk system (I + M) is arbitrarily ill-conditioned and the recurrence diverges. Either
normalize upstream or set `qk_l2_norm=1` to have the operator do it.
- Fused activations. `gate_activation='qwen'` computes the effective decay in float32 from the
+ Fused activations. `gate_activation='softplus'` computes the effective decay in float32 from the
raw projection carried by `decay`:
g = -exp(a_log) * Softplus(decay + dt_bias)
@@ -2501,6 +2513,9 @@ This version of the operator has been available since version 1 of the 'com.micr
`beta_activation='sigmoid'` applies a sigmoid to `beta`, and `qk_l2_norm=1` L2-normalizes each
query and key head vector. Folding these in avoids materializing the intermediates and keeps
the gate arithmetic in float32 independent of the input type.
+ Decay and beta may be float or the query element type (both use the same type). CUDA widens
+ FP16/BF16 gate inputs to float before any gate arithmetic, avoiding standalone lossless casts.
+ State, a_log, and dt_bias remain float. WebGPU currently accepts float decay and beta only.
#### Version
@@ -2515,7 +2530,7 @@ This version of the operator has been available since version 1 of the 'com.micr
chunk_size : int
Tuning hint for the chunk-parallel prefill algorithm. On CUDA, 32 pins the narrow chunk and any other value lets the implementation select the widest chunk the device can hold. On WebGPU, 16 opts into the chunkwise prefill route on supported Apple silicon; default 64 retains recurrent execution. Default 64.
gate_activation : string
-'none' (default) treats `decay` as the effective log-space decay. 'qwen' computes -exp(a_log) * Softplus(decay + dt_bias) in float32.
+'none' (default) treats `decay` as the effective log-space decay. 'softplus' computes -exp(a_log) * Softplus(decay + dt_bias) in float32. 'qwen' is accepted as a legacy alias of 'softplus' with identical behavior.
qk_l2_norm : int
When 1, L2-normalize each query and key head vector before the recurrence. Default 0.
scale : float
@@ -2537,16 +2552,16 @@ This version of the operator has been available since version 1 of the 'com.micr
Value, shape (total_tokens, num_heads_v, head_size_v)
cu_seqlens (optional) : TI
Exclusive prefix sums of the per-request token counts, shape (batch_size + 1). Absent means uniform packing.
-decay (optional) : TS
-Log-space decay, shape (total_tokens, num_heads_v) for a scalar per-head decay or (total_tokens, num_heads_v, head_size_qk) for a per-key-dimension decay.
-beta (optional) : TS
-Update rate, shape (total_tokens, num_heads_v)
+decay (optional) : TG
+Log-space decay, shape (total_tokens, num_heads_v) for a scalar per-head decay or (total_tokens, num_heads_v, head_size_qk) for a per-key-dimension decay. Float or the query element type; converted to float before gate arithmetic.
+beta (optional) : TG
+Update rate, shape (total_tokens, num_heads_v). Same element type as decay.
initial_state (optional) : TS
Recurrent state, shape (batch_size, num_heads_v, head_size_v, head_size_qk), V-major. Required for packed QKV. May alias final_state.
a_log (optional) : TS
-Per-head A_log, shape (num_heads_v). Requires gate_activation=qwen.
+Per-head A_log, shape (num_heads_v). Requires gate_activation=softplus.
dt_bias (optional) : TS
-Per-head gate bias, shape (num_heads_v). Requires gate_activation=qwen.
+Per-head gate bias, shape (num_heads_v). Requires gate_activation=softplus.
capture_count (optional) : TI
Number of leading local token transitions to capture for each request, shape (batch_size). Clamped on device to [0, min(state_update_capacity, sequence_length)]. Required exactly when state_update_capacity is positive.
state_update_active (optional) : TI
@@ -2570,7 +2585,9 @@ This version of the operator has been available since version 1 of the 'com.micr
T : tensor(float), tensor(float16), tensor(bfloat16)
Constrain query/key/value/output types.
TS : tensor(float)
-State, gate, beta and compact state-update tensors are always float.
+State, gate parameters and compact state-update tensors are always float.
+TG : tensor(float), tensor(float16), tensor(bfloat16)
+Decay and beta are float or match the query element type. Gate arithmetic is float.
TI : tensor(int32)
Constrain index and count tensors to int32.
diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md
index 6d084f06a97dc..6393bcce36f34 100644
--- a/docs/OperatorKernels.md
+++ b/docs/OperatorKernels.md
@@ -585,7 +585,7 @@ The **OpSet Version** column uses the following notation:
|BifurcationDetector|*in* src_tokens:**T**
*in* cur_tokens:**T**
*in* prev_suffix_match_idx:**T**
*in* pred_tokens:**T**
*out* tokens:**T**
*out* suffix_match_idx:**T**|1+|**T** = tensor(int64)|
|BranchwiseRMSNorm|*in* X:**T**
*in* scale:**M**
*out* Y:**T**|1+|**M** = tensor(bfloat16), tensor(float), tensor(float16)
**T** = tensor(bfloat16), tensor(float), tensor(float16)|
|CDist|*in* A:**T**
*in* B:**T**
*out* C:**T**|1+|**T** = tensor(double), tensor(float)|
-|CausalConvWithState|*in* input:**T**
*in* weight:**T**
*in* bias:**T**
*in* past_state:**T**
*out* output:**T**
*out* present_state:**T**|1+|**T** = tensor(float)|
+|CausalConvWithState|*in* input:**T**
*in* weight:**T**
*in* bias:**T**
*in* past_state:**T**
*out* output:**T**
*out* present_state:**T**
*out* state_update:**T**|1+|**T** = tensor(float)|
|ConvTransposeWithDynamicPads|*in* X:**T**
*in* W:**T**
*in* Pads:**tensor(int64)**
*in* B:**T**
*out* Y:**T**|1+|**T** = tensor(float)|
|CropAndResize|*in* X:**T1**
*in* rois:**T1**
*in* batch_indices:**T2**
*in* crop_size:**T2**
*out* Y:**T1**|1+|**T1** = tensor(float)
**T2** = tensor(int32)|
|DecoderMaskedMultiHeadAttention|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* mask_index:**M**
*in* attention_bias:**T**
*in* past_key:**T**
*in* past_value:**T**
*in* past_sequence_length:**M**
*in* beam_width:**M**
*in* cache_indirection:**M**
*in* bias:**T**
*out* output:**T**
*out* present_key:**T**
*out* present_value:**T**
*out* qk:**QK**|1+|**T** = tensor(float)|
@@ -1103,7 +1103,7 @@ The **OpSet Version** column uses the following notation:
|BitmaskBiasDropout|*in* data:**T**
*in* bias:**T**
*in* residual:**T**
*in* ratio:**T1**
*in* training_mode:**T2**
*out* output:**T**
*out* mask:**T3**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)
**T1** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)
**T2** = tensor(bool)
**T3** = tensor(uint32)|
|BitmaskDropout|*in* data:**T**
*in* ratio:**T1**
*in* training_mode:**T2**
*out* output:**T**
*out* mask:**T3**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)
**T1** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)
**T2** = tensor(bool)
**T3** = tensor(uint32)|
|BranchwiseRMSNorm|*in* X:**T**
*in* scale:**M**
*out* Y:**T**|1+|**M** = tensor(bfloat16), tensor(float), tensor(float16)
**T** = tensor(bfloat16), tensor(float), tensor(float16)|
-|CausalConvWithState|*in* input:**T**
*in* weight:**T**
*in* bias:**T**
*in* past_state:**T**
*out* output:**T**
*out* present_state:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)|
+|CausalConvWithState|*in* input:**T**
*in* weight:**T**
*in* bias:**T**
*in* past_state:**T**
*out* output:**T**
*out* present_state:**T**
*out* state_update:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)|
|ComplexMul|*in* A:**T**
*in* B:**T**
*out* C:**T**|1+|**T** = tensor(float), tensor(float16)|
|ComplexMulConj|*in* A:**T**
*in* B:**T**
*out* C:**T**|1+|**T** = tensor(float), tensor(float16)|
|ConvTransposeWithDynamicPads|*in* X:**T**
*in* W:**T**
*in* Pads:**tensor(int64)**
*in* B:**T**
*out* Y:**T**|1+|**T** = tensor(float)|
@@ -1120,7 +1120,7 @@ The **OpSet Version** column uses the following notation:
|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)|
|GatedAdd|*in* X:**T**
*in* Y:**T**
*in* gate:**T**
*out* output:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)|
-|GatedDeltaNet|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* cu_seqlens:**TI**
*in* decay:**TS**
*in* beta:**TS**
*in* initial_state:**TS**
*in* a_log:**TS**
*in* dt_bias:**TS**
*in* capture_count:**TI**
*in* state_update_active:**TI**
*out* output:**T**
*out* final_state:**TS**
*out* state_update:**TS**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)
**TI** = tensor(int32)
**TS** = tensor(float)|
+|GatedDeltaNet|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* cu_seqlens:**TI**
*in* decay:**TG**
*in* beta:**TG**
*in* initial_state:**TS**
*in* a_log:**TS**
*in* dt_bias:**TS**
*in* capture_count:**TI**
*in* state_update_active:**TI**
*out* output:**T**
*out* final_state:**TS**
*out* state_update:**TS**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)
**TG** = tensor(bfloat16), tensor(float), tensor(float16)
**TI** = tensor(int32)
**TS** = tensor(float)|
|GatedRMSNorm|*in* X:**T**
*in* scale:**T**
*in* gate:**T**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)|
|GatedRelativePositionBias|*in* query_layer:**T**
*in* query_bias:**T**
*in* rel_pos:**T**
*in* weight:**T**
*in* bias:**T**
*in* eco_a:**T**
*in* token_offset:**M**
*out* output:**T**|1+|**T** = tensor(float), tensor(float16)|
|GatherBlockQuantized|*in* data:**T1**
*in* indices:**Tind**
*in* scales:**T2**
*in* zero_points:**T1**
*out* output:**T2**|1+|**T1** = tensor(float4e2m1), tensor(float8e4m3fn), tensor(float8e4m3fnuz), tensor(float8e5m2), tensor(float8e5m2fnuz), tensor(int4), tensor(uint4), tensor(uint8)
**T2** = tensor(bfloat16), tensor(float), tensor(float16)
**Tind** = tensor(int32), tensor(int64)|
diff --git a/docs/contrib_ops/cuda/gdn.md b/docs/contrib_ops/cuda/gdn.md
index 367acb820175f..2cc86db3a951b 100644
--- a/docs/contrib_ops/cuda/gdn.md
+++ b/docs/contrib_ops/cuda/gdn.md
@@ -1,5 +1,95 @@
# GatedDeltaNet — Operator Documentation
+## Experimental short-sequence CUDA dispatch
+
+Set process-level tuning options before the first applicable invocation:
+
+- `ORT_GDN_S2_SPECIALIZE=1` specializes packed QKV for a single request with
+ 1 to 8 tokens, 16 query/key heads, 48 value heads, and 128-wide heads.
+ It additionally requires Qwen softplus-decay gates, sigmoid beta,
+ gated-delta updates, Q/K L2 normalization, uniform non-packed sequence
+ metadata, and no per-key decay, state-update tail pass, or short-row tail
+ pass.
+- `ORT_CAUSAL_CONV_S2_SPECIALIZE=1` specializes the companion
+ `CausalConvWithState` extended (packed-channel) projection path for kernel size
+ four, state length four and dilation one. Lengths 1 to 8 use the per-position
+ specialized kernel and lengths of 16 or more use the tiled kernel (eight
+ positions per CTA); other lengths keep the generic kernel. The `activation`
+ attribute (`none`, `silu` or `swish`; the last two are identical) is honored
+ as configured and is not part of the dispatch condition. Layout, optional
+ bias, and first-token state semantics are unchanged.
+
+Both specialization switches default to `1` and accept only `0` or `1`.
+Unsupported geometry retains generic dispatch. These options do not alter
+operator schemas, state layouts, or the default package configuration.
+
+### Benchmark
+
+Op-level GPU kernel time per call on NVIDIA H200 (SM90, driver 580.159.04, CUDA 13.0, ORT 1.32.0 integration build
+of this branch), measured 2026-10-27 at the Qwen3.8-Flash-Next rc3 production geometry
+(taken from the rc3 DFlash2 graph): `GatedDeltaNet` with packed BF16 QKV `(B, S, 10240)`, 16 q/k heads, 48 v heads,
+head size 128, FP32 state `(B, 48, 128, 128)`, BF16 decay/beta, `gated_delta`/`softplus`/`sigmoid`,
+`qk_l2_norm=1`, `state_update_capacity=7`; `CausalConvWithState` with BF16 `(B, S, 10240)` channels-last input,
+K=4, `state_length=4`, `dilation=1`, `silu`, `state_update_capacity=7`. Median kernel duration in microseconds
+(lower is better), specialized (`=1`, default) versus generic (`=0`) in a fresh process per setting. Outputs are
+bit-identical between the two settings for every row (max abs diff 0 on every output, including the state and
+`state_update`).
+
+`GatedDeltaNet` (`ORT_GDN_S2_SPECIALIZE`; S is tokens per request, S=8 is the production verification step):
+
+| Batch | S | QwenS2 (us) | Generic (us) | Speedup |
+|---:|---:|---:|---:|---:|
+| 1 | 1 | 4.83 | 6.05 | 1.25x |
+| 1 | 2 | 6.82 | 8.16 | 1.20x |
+| 1 | 4 | 10.94 | 12.74 | 1.16x |
+| 1 | 8 | 18.82 | 21.60 | 1.15x |
+| 2 | 8 | 40.74 | 40.74 | 1.00x (generic kernel; not eligible) |
+| 4 | 8 | 74.88 | 74.88 | 1.00x (generic kernel; not eligible) |
+| 8 | 8 | 142.24 | 142.24 | 1.00x (generic kernel; not eligible) |
+
+`CausalConvWithState` (`ORT_CAUSAL_CONV_S2_SPECIALIZE`):
+
+| Batch | S | Specialized (us) | Generic (us) | Speedup |
+|---:|---:|---:|---:|---:|
+| 1 | 1 | 1.76 | 2.21 | 1.25x |
+| 1 | 2 | 2.02 | 2.46 | 1.22x |
+| 1 | 4 | 2.30 | 2.75 | 1.19x |
+| 1 | 8 | 2.46 | 2.98 | 1.21x |
+| 2 | 8 | 2.94 | 3.49 | 1.18x |
+| 4 | 8 | 4.13 | 5.60 | 1.36x |
+| 8 | 8 | 6.78 | 9.28 | 1.37x |
+| 1 | 16 (tiled) | 3.90 | 3.36 | 0.86x |
+| 1 | 2048 (tiled) | 55.13 | 223.97 | 4.06x |
+
+QwenS2 trims 1.2-2.8 us (13-20%) from each single-request decode call; at S=8 that is 2.8 us x 36 GDN layers, or
+roughly 0.1 ms per verification cycle; the S<=8 convolution kernel saves about 0.5 us per call at batch 1 and
+more at larger batch. The tiled convolution kernel is the main win for prefill (4x at 2048 tokens), but at the
+smallest tiled length (S=16) it is slightly slower than the generic kernel (0.86x, +0.5 us), an absolute cost that is
+negligible next to the S>=2048 gain.
+
+End to end on the rc3 model (INT4 `all_int4`, decode ms per DFlash2 S=8 cycle, H200, median of 5 repetitions):
+
+| Context | Default | `ORT_GDN_S2_SPECIALIZE=0` |
+|---:|---:|---:|
+| 2K | 14.93 | 15.03 |
+| 8K | 15.07 | 15.17 |
+| 32K | 15.30 | 15.45 |
+
+That is +0.10 to +0.15 ms per cycle without QwenS2, i.e. QwenS2 decodes about 0.7-1.0% faster.
+`ORT_CAUSAL_CONV_S2_SPECIALIZE=0` leaves decode unchanged but raises time to first token by 2.9-3.1% (2K / 8K / 32K
+prompts: 0.1987 -> 0.2044 s, 0.8506 -> 0.8791 s, 3.5843 -> 3.6942 s).
+
+Methodology: each op is a single-node ONNX model with random BF16/FP32 inputs (fixed seeds), `state_update`
+capture enabled (`capture_count = min(S, 7)`), 48 copies of the node (distinct state tensors, like the 48 layers) in
+one CUDA-graph-captured session with IOBinding and device-resident inputs; 30 warm-up replays and 300 timed replays
+(14,400 kernel instances per row). Kernel durations come from Nsight Systems (`--cuda-graph-trace=node`), so host
+launch and inter-kernel gaps are excluded; the table gives the median of the three per-process medians (spread across
+processes <= 0.15 us, every row run on the same GPU as its counterpart). Wall-clock back-to-back graph replay
+(slope between 16- and 48-node graphs) agrees for the large kernels (conv S=2048: 54 vs 224 us; GDN S=8: 20.5 vs
+23.6 us, which includes about 1.8 us of per-node launch gap) but is noisy for kernels of a few microseconds, so it is
+not used here. Run `ORT_GDN_S2_SPECIALIZE=0/1` and `ORT_CAUSAL_CONV_S2_SPECIALIZE=0/1` in separate processes because
+the switches are cached on first use.
+
This document describes the `com.microsoft::GatedDeltaNet` contrib operator: its schema and
state contract, the chunked and warp-specialised CUDA engines and how one is selected, the
consumer-Blackwell shared-memory constraint, and measured comparisons against
@@ -56,7 +146,7 @@ Source:
|---|---|---|---|
| `update_rule` | string | `gated_delta` | One of `linear`, `gated`, `delta`, `gated_delta`. Selects which of the decay and delta-retrieval terms are present. |
| `scale` | float | `0.0` | Output scale. `0.0` means `1/sqrt(head_size_qk)`. |
-| `gate_activation` | string | `none` | `none` treats `decay` as the effective log-space decay. `qwen` computes `-exp(a_log) * Softplus(decay + dt_bias)` in float32. |
+| `gate_activation` | string | `none` | `none` treats `decay` as the effective log-space decay. `softplus` computes `-exp(a_log) * Softplus(decay + dt_bias)` in float32. `qwen` is accepted as a legacy alias of `softplus` with identical behavior. |
| `beta_activation` | string | `none` | `none` treats `beta` as the effective update rate; `sigmoid` applies a sigmoid. |
| `qk_l2_norm` | int | `0` | When `1`, L2-normalizes each query and key head vector in-kernel. |
| `chunk_size` | int | `64` | Chunk length for the chunked engine. `32` pins the reduced shared-memory configuration (see §6). |
@@ -514,7 +604,7 @@ support three intermediate precision modes controlled by session option
Measured on NVIDIA H200 (`sm_90a`, CUDA 13.0) at the Qwen3.8-Flash-Next geometry
(`batch = 1`, `num_heads_q = num_heads_k = 16`, `num_heads_v = 32`, `head_size_qk = head_size_v = 128`,
-`update_rule = "gated_delta"`, `gate_activation = "qwen"`, `beta_activation = "sigmoid"`,
+`update_rule = "gated_delta"`, `gate_activation = "softplus"`, `beta_activation = "sigmoid"`,
`qk_l2_norm = 1`, 36 GDN layers) against the sequential `float32` `recurrent` engine reference:
| Sequence Length | Mode | 1-Op Latency (us) | 36-Layer Time (ms) | Output RMSE vs. `recurrent` | State RMSE vs. `recurrent` | Exact BF16 Output Match |
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 7585993088805..9e73c5a75f888 100644
--- a/onnxruntime/contrib_ops/cpu/bert/causal_conv_with_state.cc
+++ b/onnxruntime/contrib_ops/cpu/bert/causal_conv_with_state.cc
@@ -41,6 +41,12 @@ REGISTER_KERNEL_TYPED(float)
template
CausalConvWithState::CausalConvWithState(const OpKernelInfo& info) : OpKernel(info) {
+ ORT_ENFORCE(info.GetAttrOrDefault("input_channels", 0) == 0 &&
+ info.GetAttrOrDefault("input_channel_offset", 0) == 0 &&
+ info.GetAttrOrDefault("state_length", 0) == 0 &&
+ info.GetAttrOrDefault("state_channels_last", -1) == -1 &&
+ !causal_conv_with_state_helper::HasOutput(info, 2),
+ "CPU CausalConvWithState does not support extended channel/state options (CUDA only)");
int64_t ndim = info.GetAttrOrDefault("ndim", 1);
ORT_ENFORCE(ndim == 1, "CPU CausalConvWithState only supports ndim=1");
ndim_ = static_cast(ndim);
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 8f16d8ee71f49..0bc24d8cf9502 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
@@ -16,6 +16,21 @@ namespace onnxruntime {
namespace contrib {
namespace causal_conv_with_state_helper {
+// Returns true if output `output_index` is present and not an omitted optional output.
+//
+// The in-tree onnxruntime::Node exposes NodeArgs through OutputDefs(), while the plugin EP adapter
+// node (ep::adapter::Node) only exposes OutputCount() and OutputExists().
+template
+bool HasOutput(const TKernelInfo& info, size_t output_index) {
+ const auto& node = info.node();
+ if constexpr (requires { node.OutputCount(); node.OutputExists(output_index); }) {
+ return node.OutputCount() > output_index && node.OutputExists(output_index);
+ } else {
+ const auto& output_defs = node.OutputDefs();
+ return output_defs.size() > output_index && output_defs[output_index]->Exists();
+ }
+}
+
// Reads and validates the optional `state_window` attribute.
//
// 0 (the default, i.e. attribute absent) selects the legacy unwindowed state layout. Every model
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 c79859f735c7f..cc2a3552de4dc 100644
--- a/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state.cc
+++ b/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state.cc
@@ -6,6 +6,7 @@
#include "contrib_ops/cpu/bert/causal_conv_with_state_helper.h"
#include "core/providers/cuda/cuda_common.h"
#include "core/providers/cuda/cuda_type_conversion.h"
+#include "core/common/safeint.h"
namespace onnxruntime {
namespace contrib {
@@ -45,6 +46,16 @@ CausalConvWithState::CausalConvWithState(const OpKernelInfo& info) : CudaKern
// 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_));
+ input_channels_ = SafeInt(info.GetAttrOrDefault("input_channels", 0));
+ input_channel_offset_ = SafeInt(info.GetAttrOrDefault("input_channel_offset", 0));
+ state_length_ = SafeInt(info.GetAttrOrDefault("state_length", 0));
+ state_channels_last_ = SafeInt(info.GetAttrOrDefault("state_channels_last", -1));
+ state_update_capacity_ = SafeInt(info.GetAttrOrDefault("state_update_capacity", 0));
+ has_state_update_output_ = causal_conv_with_state_helper::HasOutput(info, 2);
+ ORT_ENFORCE(input_channels_ >= 0 && input_channel_offset_ >= 0 && state_length_ >= 0 &&
+ state_channels_last_ >= -1 && state_channels_last_ <= 1 &&
+ state_update_capacity_ >= 0 && state_update_capacity_ <= 8,
+ "invalid packed-channel or state attributes");
}
template
@@ -60,6 +71,52 @@ Status CausalConvWithState::ComputeInternal(OpKernelContext* context) const {
const auto& input_shape = input_tensor->Shape();
const auto& weight_shape = weight_tensor->Shape();
+ const bool extended = input_channels_ != 0 || input_channel_offset_ != 0 || state_length_ != 0 ||
+ state_channels_last_ != -1 || state_update_capacity_ != 0 || has_state_update_output_;
+ if (extended) {
+ ORT_RETURN_IF_NOT(channels_last_ && state_window_ == 0 && input_shape.NumDimensions() == 3 &&
+ weight_shape.NumDimensions() == 3 && weight_shape[1] == 1,
+ "extended CausalConvWithState requires rank-3 channels-last input, rank-3 weight and state_window=0");
+ ORT_RETURN_IF_NOT(!has_state_update_output_ || state_update_capacity_ > 0,
+ "extended CausalConvWithState state_update output requires state_update_capacity > 0");
+ const int batch = SafeInt(input_shape[0]);
+ const int length = SafeInt(input_shape[1]);
+ const int width = SafeInt(input_shape[2]);
+ const int channels = input_channels_ == 0 ? width : input_channels_;
+ const int kernel = SafeInt(weight_shape[2]);
+ const int pad = SafeInt(kernel - 1) * dilation_;
+ const int state_length = state_length_ == 0 ? pad : state_length_;
+ const bool state_last = state_channels_last_ == -1 ? channels_last_ : state_channels_last_ != 0;
+ ORT_RETURN_IF_NOT(batch > 0 && batch <= 65535 && length > 0 && length <= 65535 &&
+ channels > 0 && kernel > 0 && input_channel_offset_ <= width &&
+ channels <= width - input_channel_offset_ && weight_shape[0] == channels &&
+ state_length >= pad,
+ "invalid extended CausalConvWithState dimensions");
+ ORT_RETURN_IF_NOT(bias_tensor == nullptr || bias_tensor->Shape() == TensorShape({channels}),
+ "bias must match selected input channels");
+ const TensorShape state_shape = state_last ? TensorShape({batch, state_length, channels})
+ : TensorShape({batch, channels, state_length});
+ ORT_RETURN_IF_NOT(past_state_tensor == nullptr || past_state_tensor->Shape() == state_shape,
+ "past_state must match configured state layout and length");
+ Tensor* output = context->Output(0, TensorShape({batch, length, channels}));
+ Tensor* present = context->Output(1, state_shape);
+ Tensor* state_update = state_update_capacity_ > 0 && has_state_update_output_
+ ? context->Output(2, TensorShape({batch, state_update_capacity_, channels}))
+ : nullptr;
+ using CudaT = typename OrtToCudaType::type;
+ return LaunchCausalConvProjectionKernel(
+ Stream(context), reinterpret_cast(input_tensor->Data()),
+ reinterpret_cast(weight_tensor->Data()),
+ bias_tensor ? reinterpret_cast(bias_tensor->Data()) : nullptr,
+ past_state_tensor ? reinterpret_cast(past_state_tensor->Data()) : nullptr,
+ reinterpret_cast(output->MutableData()),
+ reinterpret_cast(present->MutableData()),
+ batch, channels, length, kernel, dilation_, width, input_channel_offset_, state_length,
+ state_last, activation_ != "none",
+ state_update ? reinterpret_cast(state_update->MutableData()) : nullptr,
+ state_update_capacity_);
+ }
+
// Validate input rank and weight rank
if (channels_last_) {
// (batch_size, sequence_length, d_1, ..., d_n): any number of trailing channel axes, so a
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 0a1f23f3abfd3..aec21dfdb4928 100644
--- a/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state.h
+++ b/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state.h
@@ -25,6 +25,12 @@ class CausalConvWithState final : public onnxruntime::cuda::CudaKernel {
std::string activation_;
// Leading (axis-0) extent of past_state / present_state; 0 means no window axis (single state).
int state_window_;
+ int input_channels_;
+ int input_channel_offset_;
+ int state_length_;
+ int state_channels_last_;
+ int state_update_capacity_;
+ bool has_state_update_output_;
};
} // namespace cuda
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 0c0e1da836521..1747dcc2de1b0 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
@@ -19,6 +19,7 @@
#include
#include
#include "contrib_ops/cuda/bert/causal_conv_with_state_impl.h"
+#include "core/platform/env_var_utils.h"
#include "core/providers/cuda/cu_inc/cuda_type_helper.cuh"
namespace onnxruntime {
@@ -27,6 +28,105 @@ namespace cuda {
namespace {
+template
+__global__ void CausalConvProjectionTiledK4Kernel(
+ const T* __restrict__ input, const T* __restrict__ weight, const T* __restrict__ bias,
+ const T* __restrict__ past, T* __restrict__ output, T* __restrict__ present,
+ int channels, int length, int input_width, int channel_offset,
+ CausalConvLayout state_layout, bool apply_silu,
+ T* __restrict__ state_update = nullptr, int state_update_capacity = 0) {
+ const int b = blockIdx.y;
+ const int c = blockIdx.x * blockDim.x + threadIdx.x;
+ if (c >= channels) return;
+ const int pos_base = static_cast(blockIdx.z) * kPosTile;
+
+ const float b_val = bias == nullptr ? 0.0f : to_float(bias[c]);
+ const float w0 = to_float(weight[static_cast(c) * 4 + 0]);
+ const float w1 = to_float(weight[static_cast(c) * 4 + 1]);
+ const float w2 = to_float(weight[static_cast(c) * 4 + 2]);
+ const float w3 = to_float(weight[static_cast(c) * 4 + 3]);
+
+ auto load_src = [&](int src) -> float {
+ if (src >= 0) {
+ return to_float(input[(static_cast(b) * length + src) * input_width + channel_offset + c]);
+ }
+ return past == nullptr ? 0.0f : to_float(past[state_layout.Offset(b, 4 + src, c)]);
+ };
+
+ float v0 = load_src(pos_base - 3);
+ float v1 = load_src(pos_base - 2);
+ float v2 = load_src(pos_base - 1);
+
+#pragma unroll
+ for (int p = 0; p < kPosTile; ++p) {
+ const int position = pos_base + p;
+ if (position >= length) break;
+ const T cur_in = input[(static_cast(b) * length + position) * input_width + channel_offset + c];
+ if (state_update != nullptr && position < state_update_capacity) {
+ state_update[(static_cast(b) * state_update_capacity + position) * channels + c] = cur_in;
+ }
+ const float v3 = to_float(cur_in);
+ float sum = b_val + v0 * w0 + v1 * w1 + v2 * w2 + v3 * w3;
+ if (apply_silu) sum = sum / (1.0f + expf(-sum));
+ output[(static_cast(b) * length + position) * channels + c] = from_float(sum);
+ v0 = v1;
+ v1 = v2;
+ v2 = v3;
+ }
+
+ if (blockIdx.z != 0) return;
+#pragma unroll
+ for (int position = 0; position < 4; ++position) {
+ const int source = length - 4 + position;
+ present[state_layout.Offset(b, position, c)] =
+ source >= 0 ? input[(static_cast(b) * length + source) * input_width + channel_offset + c]
+ : (past == nullptr ? from_float(0.0f) : past[state_layout.Offset(b, 4 + source, c)]);
+ }
+}
+
+template
+__global__ void CausalConvProjectionKernel(
+ const T* input, const T* weight, const T* bias, const T* past,
+ T* output, T* present, int channels, int length,
+ int kernel_size, int dilation, int input_width, int channel_offset,
+ int state_length, CausalConvLayout state_layout, bool apply_silu,
+ T* state_update = nullptr, int state_update_capacity = 0) {
+ if constexpr (S2) {
+ kernel_size = state_length = 4;
+ dilation = 1;
+ }
+ const int b = blockIdx.y;
+ const int c = blockIdx.x * blockDim.x + threadIdx.x;
+ if (c >= channels) return;
+ const int pad = (kernel_size - 1) * dilation;
+ const int position = blockIdx.z;
+ {
+ const T cur_in = input[(static_cast(b) * length + position) * input_width + channel_offset + c];
+ if (state_update != nullptr && position < state_update_capacity) {
+ state_update[(static_cast(b) * state_update_capacity + position) * channels + c] = cur_in;
+ }
+ float sum = bias == nullptr ? 0.0f : to_float(bias[c]);
+ for (int tap = 0; tap < kernel_size; ++tap) {
+ const int source = position - pad + tap * dilation;
+ const float value = source >= 0
+ ? (source == position
+ ? to_float(cur_in)
+ : to_float(input[(static_cast(b) * length + source) * input_width + channel_offset + c]))
+ : (past == nullptr ? 0.0f : to_float(past[state_layout.Offset(b, state_length + source, c)]));
+ sum += value * to_float(weight[static_cast(c) * kernel_size + tap]);
+ }
+ if (apply_silu) sum = sum / (1.0f + expf(-sum));
+ output[(static_cast(b) * length + position) * channels + c] = from_float(sum);
+ }
+ if (position != 0) return;
+ for (int position = 0; position < state_length; ++position) {
+ const int source = length - state_length + position;
+ present[state_layout.Offset(b, position, c)] =
+ source >= 0 ? input[(static_cast(b) * length + source) * input_width + channel_offset + c]
+ : (past == nullptr ? from_float(0.0f) : past[state_layout.Offset(b, state_length + source, c)]);
+ }
+}
+
__device__ __forceinline__ float silu_fn(float x) {
return x / (1.0f + expf(-x));
}
@@ -428,6 +528,60 @@ __global__ void CausalConvPrefillKernelChannelsLast(
} // anonymous namespace
+template
+Status LaunchCausalConvProjectionKernel(
+ cudaStream_t stream, const T* input, const T* weight, const T* bias, const T* past,
+ T* output, T* present, int batch, int channels, int length,
+ int kernel_size, int dilation, int input_width, int channel_offset,
+ int state_length, bool state_channels_last, bool apply_silu,
+ T* state_update, int state_update_capacity) {
+ const dim3 grid((channels - 1) / 128 + 1, batch, length);
+ if (state_update != nullptr && length < state_update_capacity) {
+ // The kernels write only the first min(length, capacity) rows per batch; zero the rest so the
+ // output does not expose stale allocator contents. A strided memset keeps this graph-capturable.
+ const size_t row_bytes = static_cast(channels) * sizeof(T);
+ CUDA_RETURN_IF_ERROR(cudaMemset2DAsync(
+ state_update + static_cast(length) * channels,
+ static_cast(state_update_capacity) * row_bytes, 0,
+ static_cast(state_update_capacity - length) * row_bytes, static_cast(batch), stream));
+ }
+ static const int specialize = ParseEnvironmentVariableWithDefault("ORT_CAUSAL_CONV_S2_SPECIALIZE", 1);
+ ORT_ENFORCE(specialize == 0 || specialize == 1, "ORT_CAUSAL_CONV_S2_SPECIALIZE must be 0 or 1.");
+ if (specialize && length >= 1 && kernel_size == 4 && dilation == 1 && state_length == 4) {
+ if (length >= 16) {
+ constexpr int kPosTile = 8;
+ const dim3 tiled_grid((channels - 1) / 128 + 1, batch, (length + kPosTile - 1) / kPosTile);
+ CausalConvProjectionTiledK4Kernel<<>>(
+ input, weight, bias, past, output, present, channels, length,
+ input_width, channel_offset,
+ MakeCausalConvLayout(state_channels_last, channels, state_length), apply_silu,
+ state_update, state_update_capacity);
+ return CUDA_CALL(cudaGetLastError());
+ }
+ CausalConvProjectionKernel<<>>(
+ input, weight, bias, past, output, present, channels, length,
+ kernel_size, dilation, input_width, channel_offset, state_length,
+ MakeCausalConvLayout(state_channels_last, channels, state_length), apply_silu,
+ state_update, state_update_capacity);
+ return CUDA_CALL(cudaGetLastError());
+ }
+ CausalConvProjectionKernel<<>>(
+ input, weight, bias, past, output, present, channels, length,
+ kernel_size, dilation, input_width, channel_offset, state_length,
+ MakeCausalConvLayout(state_channels_last, channels, state_length), apply_silu,
+ state_update, state_update_capacity);
+ return CUDA_CALL(cudaGetLastError());
+}
+
+#define INSTANTIATE_PROJECTION(T) \
+ template Status LaunchCausalConvProjectionKernel( \
+ cudaStream_t, const T*, const T*, const T*, const T*, T*, T*, \
+ int, int, int, int, int, int, int, int, bool, bool, T*, int);
+INSTANTIATE_PROJECTION(float)
+INSTANTIATE_PROJECTION(half)
+INSTANTIATE_PROJECTION(__nv_bfloat16)
+#undef INSTANTIATE_PROJECTION
+
template
Status LaunchCausalConvWithStateKernel(
cudaStream_t stream,
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 968cb9507339d..8339eefb7efbe 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
@@ -31,6 +31,14 @@ inline CausalConvLayout MakeCausalConvLayout(bool channels_last, int64_t channel
: CausalConvLayout{channels * length, 1, length};
}
+template
+Status LaunchCausalConvProjectionKernel(
+ cudaStream_t stream, const T* input, const T* weight, const T* bias, const T* past,
+ T* output, T* present, int batch, int channels, int length,
+ int kernel_size, int dilation, int input_width, int channel_offset,
+ int state_length, bool state_channels_last, bool apply_silu,
+ T* state_update = nullptr, int state_update_capacity = 0);
+
// 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
diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc
index 51a64a041a2a9..d326e3b83be25 100644
--- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc
+++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc
@@ -22,19 +22,22 @@ using onnxruntime::cuda::CudaKernel;
using onnxruntime::cuda::OrtToCudaType;
namespace gdn = gated_delta_net;
-#define REGISTER_KERNEL_TYPED(T) \
- ONNX_OPERATOR_TYPED_KERNEL_EX( \
- GatedDeltaNet, \
- kMSDomain, \
- 1, \
- T, \
- kCudaExecutionProvider, \
- (*KernelDefBuilder::Create()) \
- .TypeConstraint("T", DataTypeImpl::GetTensorType()) \
- .TypeConstraint("TS", DataTypeImpl::GetTensorType()) \
- .TypeConstraint("TI", DataTypeImpl::GetTensorType()) \
- .InputMemoryType(OrtMemTypeCPUInput, 10) \
- .MayInplace(6, 1), \
+#define REGISTER_KERNEL_TYPED(T) \
+ ONNX_OPERATOR_TYPED_KERNEL_EX( \
+ GatedDeltaNet, \
+ kMSDomain, \
+ 1, \
+ T, \
+ kCudaExecutionProvider, \
+ (*KernelDefBuilder::Create()) \
+ .TypeConstraint("T", DataTypeImpl::GetTensorType()) \
+ .TypeConstraint("TS", DataTypeImpl::GetTensorType()) \
+ .TypeConstraint("TG", {DataTypeImpl::GetTensorType(), \
+ DataTypeImpl::GetTensorType(), \
+ DataTypeImpl::GetTensorType()}) \
+ .TypeConstraint("TI", DataTypeImpl::GetTensorType()) \
+ .InputMemoryType(OrtMemTypeCPUInput, 10) \
+ .MayInplace(6, 1), \
GatedDeltaNet);
REGISTER_KERNEL_TYPED(float)
@@ -59,10 +62,10 @@ GatedDeltaNet::GatedDeltaNet(const OpKernelInfo& info) : CudaKernel(info) {
const std::string gate = info.GetAttrOrDefault("gate_activation", "none");
if (gate == "none") {
gate_activation_ = gdn::GateActivation::kNone;
- } else if (gate == "qwen") {
- gate_activation_ = gdn::GateActivation::kQwen;
+ } else if (gate == "softplus" || gate == "qwen") {
+ gate_activation_ = gdn::GateActivation::kSoftplusDecay;
} else {
- ORT_THROW("gate_activation must be one of: none, qwen");
+ ORT_THROW("gate_activation must be one of: none, softplus (legacy alias: qwen)");
}
const std::string beta = info.GetAttrOrDefault("beta_activation", "none");
@@ -123,6 +126,13 @@ Status GatedDeltaNet::ComputeInternal(OpKernelContext* context) const {
const Tensor* capture_count = context->Input(9);
const Tensor* state_update_active = context->Input(10);
+ ORT_RETURN_IF_NOT(query != nullptr, "query is required");
+ for (const Tensor* gate : {decay, beta}) {
+ ORT_RETURN_IF_NOT(gate == nullptr || gate->IsDataType() ||
+ gate->GetElementType() == query->GetElementType(),
+ "decay and beta must be float or match the query element type");
+ }
+
ORT_RETURN_IF_NOT((state_update_capacity_ > 0) == (capture_count != nullptr),
"capture_count must be present exactly when state_update_capacity is positive");
@@ -137,7 +147,6 @@ Status GatedDeltaNet::ComputeInternal(OpKernelContext* context) const {
ORT_RETURN_IF_NOT(needs_beta == (beta != nullptr),
"beta input presence must match update_rule");
- ORT_RETURN_IF_NOT(query != nullptr, "query is required");
const bool is_packed_qkv = key == nullptr && value == nullptr;
ORT_RETURN_IF_NOT(is_packed_qkv || (key != nullptr && value != nullptr),
"key and value must be both present or both absent");
@@ -277,9 +286,9 @@ Status GatedDeltaNet::ComputeInternal(OpKernelContext* context) const {
b[token_dims] == num_heads_v,
"beta must be [...tokens, num_heads_v]");
}
- if (gate_activation_ == gdn::GateActivation::kQwen) {
+ if (gate_activation_ == gdn::GateActivation::kSoftplusDecay) {
ORT_RETURN_IF_NOT(a_log != nullptr && dt_bias != nullptr,
- "gate_activation=qwen requires a_log and dt_bias");
+ "gate_activation=softplus requires a_log and dt_bias");
ORT_RETURN_IF_NOT(a_log->Shape().NumDimensions() == 1 && a_log->Shape()[0] == num_heads_v,
"a_log must be [num_heads_v]");
ORT_RETURN_IF_NOT(dt_bias->Shape().NumDimensions() == 1 &&
@@ -287,7 +296,7 @@ Status GatedDeltaNet::ComputeInternal(OpKernelContext* context) const {
"dt_bias must be [num_heads_v]");
} else {
ORT_RETURN_IF_NOT(a_log == nullptr && dt_bias == nullptr,
- "a_log and dt_bias require gate_activation=qwen");
+ "a_log and dt_bias require gate_activation=softplus");
}
if (capture_count != nullptr) {
@@ -391,7 +400,9 @@ Status GatedDeltaNet::ComputeInternal(OpKernelContext* context) const {
const size_t query_row_bytes = static_cast(num_heads_q) * head_size_qk * sizeof(T);
const size_t key_row_bytes = static_cast(num_heads_k) * head_size_qk * sizeof(T);
const size_t value_row_bytes = static_cast(num_heads_v) * head_size_v * sizeof(T);
- const size_t unpacked_qkv_bytes = is_packed_qkv
+ const bool direct_packed_decode = is_packed_qkv && plan.engine == gdn::Engine::kRecurrent &&
+ plan.warp_specialized && total_tokens <= 8LL * batch;
+ const size_t unpacked_qkv_bytes = is_packed_qkv && !direct_packed_decode
? static_cast(total_tokens) *
(query_row_bytes + key_row_bytes + value_row_bytes)
: 0;
@@ -405,7 +416,10 @@ Status GatedDeltaNet::ComputeInternal(OpKernelContext* context) const {
const CudaT* query_data = reinterpret_cast(query->Data());
const CudaT* key_data = key != nullptr ? reinterpret_cast(key->Data()) : nullptr;
const CudaT* value_data = value != nullptr ? reinterpret_cast(value->Data()) : nullptr;
- if (is_packed_qkv) {
+ if (direct_packed_decode) {
+ key_data = query_data + query_row_bytes / sizeof(T);
+ value_data = key_data + key_row_bytes / sizeof(T);
+ } else if (is_packed_qkv) {
const size_t packed_row_bytes = query_row_bytes + key_row_bytes + value_row_bytes;
auto* unpacked_query = workspace.get();
auto* unpacked_key = unpacked_query + static_cast(total_tokens) * query_row_bytes;
@@ -430,10 +444,13 @@ Status GatedDeltaNet::ComputeInternal(OpKernelContext* context) const {
pack.query = query_data;
pack.key = key_data;
pack.value = value_data;
+ pack.packed_qkv = direct_packed_decode;
pack.cu_seqlens = cu_seqlens != nullptr ? cu_seqlens->Data() : nullptr;
pack.capture_count = desc.state_update_active ? capture_count->Data() : nullptr;
- pack.decay = decay != nullptr ? decay->Data() : nullptr;
- pack.beta = beta != nullptr ? beta->Data() : nullptr;
+ pack.decay = decay != nullptr ? decay->DataRaw() : nullptr;
+ pack.beta = beta != nullptr ? beta->DataRaw() : nullptr;
+ const Tensor* gate = decay != nullptr ? decay : beta;
+ pack.gates_in_input_type = gate != nullptr && !gate->IsDataType();
pack.initial_state = initial_state != nullptr ? initial_state->Data() : nullptr;
pack.a_log = a_log != nullptr ? a_log->Data() : nullptr;
pack.dt_bias = dt_bias != nullptr ? dt_bias->Data() : nullptr;
diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.h b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.h
index bed13f2cb7722..7d321e932a9c6 100644
--- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.h
+++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.h
@@ -14,7 +14,7 @@ namespace contrib {
namespace cuda {
template
-class GatedDeltaNet final : public onnxruntime::cuda::CudaKernel {
+class GatedDeltaNet : public onnxruntime::cuda::CudaKernel {
public:
GatedDeltaNet(const OpKernelInfo& info);
Status ComputeInternal(OpKernelContext* context) const override;
diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu
index 3512a1efd1b77..8278638389397 100644
--- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu
+++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu
@@ -50,6 +50,10 @@
#include "core/providers/cuda/cu_inc/common.cuh"
#include "core/providers/cuda/cuda_common.h"
+// Include env_var_utils.h after cuda_common.h: the latter transitively pulls in provider_api.h,
+// which defines SHARED_PROVIDER and suppresses env_var_utils.h's own logging.h include.
+#include "core/platform/env_var_utils.h"
+
namespace onnxruntime {
namespace contrib {
namespace cuda {
@@ -375,7 +379,7 @@ __device__ __noinline__ void RunChunkedStepExactBf16Fallback(
// ---------------------------------------------------------------------------
// Chunked engine
// ---------------------------------------------------------------------------
-template
__global__ __launch_bounds__(kChunkedThreads) void GatedDeltaNetChunkedKernel(
VariantPack pack, KernelParams p) {
@@ -483,10 +487,10 @@ __global__ __launch_bounds__(kChunkedThreads) void GatedDeltaNetChunkedKernel(
const int64_t gi = tok * p.num_heads_v + hv;
float g = 0.0f, be = 0.0f;
if (tid < len) {
- g = needs_decay ? EffectiveDecay(pack.decay[gi], pack.a_log, pack.dt_bias, hv,
+ g = needs_decay ? EffectiveDecay(LoadGate(pack.decay, gi, InputTypeGates), pack.a_log, pack.dt_bias, hv,
p.gate_activation)
: 0.0f;
- be = pack.beta != nullptr ? EffectiveBeta(pack.beta[gi], p.beta_activation) : 1.0f;
+ be = pack.beta != nullptr ? EffectiveBeta(LoadGate(pack.beta, gi, InputTypeGates), p.beta_activation) : 1.0f;
}
s.gc[tid] = g;
s.beta[tid] = be;
@@ -557,9 +561,21 @@ __device__ __forceinline__ float WarpReduceK(float v) {
return v;
}
-template
+template
__global__ __launch_bounds__(32 * kWarps) void GatedDeltaNetDecodeWarpKernel(
VariantPack pack, KernelParams p, int d_v) {
+ if constexpr (QwenS2) {
+ // Host dispatch checks every field before replacing it with a compile-time constant.
+ p.num_heads_q = p.num_heads_k = 16;
+ p.num_heads_v = 48;
+ p.gate_activation = GateActivation::kSoftplusDecay;
+ p.beta_activation = BetaActivation::kSigmoid;
+ p.update_rule = UpdateRule::kGatedDelta;
+ p.qk_l2_norm = true;
+ p.decay_per_key_dim_flag = 0;
+ p.state_update_tail_pass = p.short_row_tail_pass = false;
+ d_v = 128;
+ }
constexpr int kRegs = DK / 32; // state elements per lane
const int lane = threadIdx.x;
const int warp = threadIdx.y;
@@ -579,9 +595,12 @@ __global__ __launch_bounds__(32 * kWarps) void GatedDeltaNetDecodeWarpKernel(
}
const int capture_count = CaptureCount(pack, p, b, seq_len);
- const int64_t q_stride = static_cast(p.num_heads_q) * DK;
- const int64_t k_stride = static_cast(p.num_heads_k) * DK;
+ const int64_t packed_stride = (static_cast(p.num_heads_q) + p.num_heads_k) * DK +
+ static_cast(p.num_heads_v) * d_v;
+ const int64_t q_stride = PackedQkv ? packed_stride : static_cast(p.num_heads_q) * DK;
+ const int64_t k_stride = PackedQkv ? packed_stride : static_cast(p.num_heads_k) * DK;
const int64_t v_stride = static_cast(p.num_heads_v) * d_v;
+ const int64_t v_input_stride = PackedQkv ? packed_stride : v_stride;
const int64_t st_off =
(static_cast(b) * p.num_heads_v + hv) * static_cast(d_v) * DK +
static_cast(col) * DK;
@@ -634,14 +653,14 @@ __global__ __launch_bounds__(32 * kWarps) void GatedDeltaNetDecodeWarpKernel(
float decay_scalar = 1.0f;
if (needs_decay && !decay_per_key_dim) {
- decay_scalar = __expf(EffectiveDecay(pack.decay[gi], pack.a_log, pack.dt_bias, hv,
+ decay_scalar = __expf(EffectiveDecay(LoadGate(pack.decay, gi, InputTypeGates), pack.a_log, pack.dt_bias, hv,
p.gate_activation));
}
#pragma unroll
for (int i = 0; i < kRegs; ++i) {
if (needs_decay && decay_per_key_dim) {
const int k = lane + i * 32;
- const float raw = pack.decay[(tok * p.num_heads_v + hv) * DK + k];
+ const float raw = LoadGate(pack.decay, (tok * p.num_heads_v + hv) * DK + k, InputTypeGates);
s[i] *= __expf(EffectiveDecay(raw, pack.a_log, pack.dt_bias, hv, p.gate_activation));
} else {
s[i] *= decay_scalar;
@@ -656,8 +675,8 @@ __global__ __launch_bounds__(32 * kWarps) void GatedDeltaNetDecodeWarpKernel(
}
const float beta_t =
- pack.beta != nullptr ? EffectiveBeta(pack.beta[gi], p.beta_activation) : 1.0f;
- const float vv = static_cast(pack.value[tok * v_stride + hv * d_v + col]);
+ pack.beta != nullptr ? EffectiveBeta(LoadGate(pack.beta, gi, InputTypeGates), p.beta_activation) : 1.0f;
+ const float vv = static_cast(pack.value[tok * v_input_stride + hv * d_v + col]);
const float delta = beta_t * (vv - r);
if (t < capture_count) {
@@ -704,7 +723,7 @@ __global__ __launch_bounds__(32 * kWarps) void GatedDeltaNetDecodeWarpKernel(
// Generic recurrent engine: any head geometry the warp kernel cannot take.
// One CTA per (sequence, v-head); the [DK x DV] state lives in shared memory.
// ---------------------------------------------------------------------------
-template
+template
__global__ __launch_bounds__(kThreads) void GatedDeltaNetRecurrentKernel(VariantPack pack,
KernelParams p, int d_k,
int d_v) {
@@ -763,7 +782,7 @@ __global__ __launch_bounds__(kThreads) void GatedDeltaNetRecurrentKernel(Variant
}
if (needs_decay && decay_per_key_dim) {
for (int i = tid; i < d_k; i += kThreads) {
- const float raw = pack.decay[(tok * p.num_heads_v + hv) * d_k + i];
+ const float raw = LoadGate(pack.decay, (tok * p.num_heads_v + hv) * d_k + i, InputTypeGates);
gbuf[i] = __expf(EffectiveDecay(raw, pack.a_log, pack.dt_bias, hv, p.gate_activation));
}
}
@@ -790,11 +809,11 @@ __global__ __launch_bounds__(kThreads) void GatedDeltaNetRecurrentKernel(Variant
const float g_scalar =
(needs_decay && !decay_per_key_dim)
- ? __expf(EffectiveDecay(pack.decay[gi], pack.a_log, pack.dt_bias, hv,
+ ? __expf(EffectiveDecay(LoadGate(pack.decay, gi, InputTypeGates), pack.a_log, pack.dt_bias, hv,
p.gate_activation))
: 1.0f;
const float beta_t =
- pack.beta != nullptr ? EffectiveBeta(pack.beta[gi], p.beta_activation) : 1.0f;
+ pack.beta != nullptr ? EffectiveBeta(LoadGate(pack.beta, gi, InputTypeGates), p.beta_activation) : 1.0f;
// Decay, then retrieval against the decayed state.
for (int idx = tid; idx < d_k * d_v; idx += kThreads) {
@@ -858,27 +877,53 @@ __global__ __launch_bounds__(kThreads) void GatedDeltaNetRecurrentKernel(Variant
constexpr int kRecurrentThreads = 256;
constexpr int kDecodeWarps = kRecurrentThreads / 32;
-template
+template
+void LaunchDecodeWarpSized(const Descriptor& desc, const VariantPack& pack,
+ const KernelParams& p, cudaStream_t stream) {
+ const dim3 grid(desc.batch, desc.num_heads_v,
+ (desc.head_size_v + Warps - 1) / Warps);
+ const dim3 block(32, Warps, 1);
+ switch (desc.head_size_qk) {
+ case 64:
+ GatedDeltaNetDecodeWarpKernel
+ <<>>(pack, p, desc.head_size_v);
+ break;
+ case 128:
+ if constexpr (PackedQkv) {
+ static const int specialize = ParseEnvironmentVariableWithDefault("ORT_GDN_S2_SPECIALIZE", 1);
+ ORT_ENFORCE(specialize == 0 || specialize == 1, "ORT_GDN_S2_SPECIALIZE must be 0 or 1.");
+ if (specialize && desc.batch == 1 && p.total_tokens == p.uniform_len &&
+ p.total_tokens >= 1 && p.total_tokens <= 8 &&
+ p.num_heads_q == 16 && p.num_heads_k == 16 && p.num_heads_v == 48 &&
+ desc.head_size_v == 128 && pack.cu_seqlens == nullptr &&
+ p.gate_activation == GateActivation::kSoftplusDecay &&
+ p.beta_activation == BetaActivation::kSigmoid && p.update_rule == UpdateRule::kGatedDelta &&
+ p.qk_l2_norm && p.decay_per_key_dim_flag == 0 &&
+ !p.state_update_tail_pass && !p.short_row_tail_pass) {
+ GatedDeltaNetDecodeWarpKernel
+ <<>>(pack, p, desc.head_size_v);
+ break;
+ }
+ }
+ GatedDeltaNetDecodeWarpKernel
+ <<>>(pack, p, desc.head_size_v);
+ break;
+ default:
+ GatedDeltaNetDecodeWarpKernel
+ <<>>(pack, p, desc.head_size_v);
+ break;
+ }
+}
+
+template
Status LaunchRecurrent(const Descriptor& desc, bool warp_specialized, const VariantPack& pack,
const KernelParams& p, size_t max_shared_memory_per_block,
cudaStream_t stream) {
if (warp_specialized) {
- const dim3 grid(desc.batch, desc.num_heads_v,
- (desc.head_size_v + kDecodeWarps - 1) / kDecodeWarps);
- const dim3 block(32, kDecodeWarps, 1);
- switch (desc.head_size_qk) {
- case 64:
- GatedDeltaNetDecodeWarpKernel
- <<>>(pack, p, desc.head_size_v);
- break;
- case 128:
- GatedDeltaNetDecodeWarpKernel
- <<>>(pack, p, desc.head_size_v);
- break;
- default:
- GatedDeltaNetDecodeWarpKernel
- <<>>(pack, p, desc.head_size_v);
- break;
+ if (pack.packed_qkv) {
+ LaunchDecodeWarpSized(desc, pack, p, stream);
+ } else {
+ LaunchDecodeWarpSized(desc, pack, p, stream);
}
return CUDA_CALL(cudaGetLastError());
}
@@ -886,20 +931,20 @@ Status LaunchRecurrent(const Descriptor& desc, bool warp_specialized, const Vari
const size_t smem = RecurrentSmemBytes(desc.head_size_qk, desc.head_size_v);
static DynamicSmemConfig recurrent_smem_config;
ORT_RETURN_IF_ERROR(SetMaxDynamicSmem(
- GatedDeltaNetRecurrentKernel,
+ GatedDeltaNetRecurrentKernel,
max_shared_memory_per_block, recurrent_smem_config));
const dim3 grid(desc.batch, desc.num_heads_v, 1);
- GatedDeltaNetRecurrentKernel
+ GatedDeltaNetRecurrentKernel
<<>>(pack, p, desc.head_size_qk, desc.head_size_v);
return CUDA_CALL(cudaGetLastError());
}
} // namespace
-template
-Status LaunchGatedDeltaNet(const Descriptor& desc, const Plan& plan, const VariantPack& pack,
- float scale, int max_threads_per_block,
- size_t max_shared_memory_per_block, cudaStream_t stream) {
+template
+static Status LaunchGatedDeltaNetTypedGates(const Descriptor& desc, const Plan& plan, const VariantPack& pack,
+ float scale, int max_threads_per_block,
+ size_t max_shared_memory_per_block, cudaStream_t stream) {
KernelParams p{};
p.total_tokens = desc.total_tokens;
p.uniform_len = desc.batch > 0 ? desc.total_tokens / desc.batch : 0;
@@ -932,8 +977,8 @@ Status LaunchGatedDeltaNet(const Descriptor& desc, const Plan& plan, const Varia
auto launch_chunked = [&]() -> Status {
static DynamicSmemConfig smem_config;
ORT_RETURN_IF_ERROR(SetMaxDynamicSmem(
- GatedDeltaNetChunkedKernel, plan.smem_bytes, smem_config));
- GatedDeltaNetChunkedKernel
+ GatedDeltaNetChunkedKernel, plan.smem_bytes, smem_config));
+ GatedDeltaNetChunkedKernel
<<>>(pack, p);
return CUDA_CALL(cudaGetLastError());
};
@@ -963,14 +1008,26 @@ Status LaunchGatedDeltaNet(const Descriptor& desc, const Plan& plan, const Varia
// the two row sets are disjoint, so the merged result is deterministic.
// The chunked engine only accepts head_size_qk == head_size_v == 128, so the warp
// kernel is always the right shape here.
- return LaunchRecurrent(desc, /*warp_specialized=*/true, pack, p,
- max_shared_memory_per_block, stream);
+ return LaunchRecurrent(desc, /*warp_specialized=*/true, pack, p,
+ max_shared_memory_per_block, stream);
}
}
(void)max_threads_per_block;
- return LaunchRecurrent(desc, plan.warp_specialized, pack, p, max_shared_memory_per_block,
- stream);
+ return LaunchRecurrent(desc, plan.warp_specialized, pack, p, max_shared_memory_per_block,
+ stream);
+}
+
+template
+Status LaunchGatedDeltaNet(const Descriptor& desc, const Plan& plan, const VariantPack& pack,
+ float scale, int max_threads_per_block,
+ size_t max_shared_memory_per_block, cudaStream_t stream) {
+ if (pack.gates_in_input_type) {
+ return LaunchGatedDeltaNetTypedGates(
+ desc, plan, pack, scale, max_threads_per_block, max_shared_memory_per_block, stream);
+ }
+ return LaunchGatedDeltaNetTypedGates(
+ desc, plan, pack, scale, max_threads_per_block, max_shared_memory_per_block, stream);
}
template Status LaunchGatedDeltaNet(const Descriptor&, const Plan&, const VariantPack&,
diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.h b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.h
index e1a38ea967475..8f2a0eb3d7acc 100644
--- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.h
+++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.h
@@ -24,8 +24,10 @@ struct VariantPack {
const T* value = nullptr; // [total_tokens, Hv, V]
const int32_t* cu_seqlens = nullptr; // [B+1] or null for uniform packing
const int32_t* capture_count = nullptr; // [B] or null when transition capture is disabled
- const float* decay = nullptr; // [total_tokens, Hv] or [total_tokens, Hv, K]
- const float* beta = nullptr; // [total_tokens, Hv]
+ const void* decay = nullptr; // [total_tokens, Hv] or [total_tokens, Hv, K]
+ const void* beta = nullptr; // [total_tokens, Hv]
+ bool gates_in_input_type = false; // gate buffers contain T instead of float
+ bool packed_qkv = false; // short decode reads packed Q/K/V without unpacking
const float* initial_state = nullptr; // [B, Hv, V, K] V-major, may alias final_state
const float* a_log = nullptr; // [Hv]
const float* dt_bias = nullptr; // [Hv]
@@ -35,6 +37,14 @@ struct VariantPack {
void* workspace = nullptr;
};
+#ifdef __CUDACC__
+template
+__device__ __forceinline__ float LoadGate(const void* data, int64_t index, bool gates_in_input_type) {
+ return gates_in_input_type ? static_cast(static_cast(data)[index])
+ : static_cast(data)[index];
+}
+#endif
+
template
Status LaunchGatedDeltaNet(const Descriptor& desc, const Plan& plan, const VariantPack& pack,
float scale, int max_threads_per_block,
diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_mma.cuh b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_mma.cuh
index 65dce96407920..7c3bf5b4b59d5 100644
--- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_mma.cuh
+++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_mma.cuh
@@ -53,7 +53,7 @@ __device__ __forceinline__ void SequenceRange(const int32_t* cu_seqlens, int b,
// Effective per-token decay and beta, after the optional fused activations.
__device__ __forceinline__ float EffectiveDecay(float raw, const float* a_log, const float* dt_bias,
int h, GateActivation act) {
- if (act == GateActivation::kQwen) {
+ if (act == GateActivation::kSoftplusDecay) {
const float bias = dt_bias != nullptr ? dt_bias[h] : 0.0f;
const float scale = a_log != nullptr ? -__expf(a_log[h]) : -1.0f;
return scale * SoftPlus(raw + bias);
@@ -122,7 +122,7 @@ __device__ __forceinline__ uint32_t PackHalf2(__nv_bfloat16 lo, __nv_bfloat16 hi
// mma.sync.m16n8k16 is Ampere and later, so ptxas rejects it for a pre-Ampere target. The only
// caller is the chunked engine, and SelectPlan requires sm_major >= 8 before choosing it, so the
// pre-Ampere body below exists purely to let those targets compile and is never executed.
-template
+template
__device__ __forceinline__ void MmaM16N8K16(float (&d)[4], const uint32_t (&a)[4],
const uint32_t (&b)[2]);
diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_plan.h b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_plan.h
index 7bcfe894ca514..a73cb5742c279 100644
--- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_plan.h
+++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_plan.h
@@ -26,7 +26,7 @@ enum class UpdateRule : int { kLinear = 0,
kDelta = 2,
kGatedDelta = 3 };
enum class GateActivation : int { kNone = 0,
- kQwen = 1 };
+ kSoftplusDecay = 1 };
enum class BetaActivation : int { kNone = 0,
kSigmoid = 1 };
diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_split_impl.cu b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_split_impl.cu
index 1857583c4ea27..b50f47efb4e32 100644
--- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_split_impl.cu
+++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_split_impl.cu
@@ -363,10 +363,10 @@ __global__ __launch_bounds__(kChunkedThreads) void GatedDeltaNetPrepareKernel(
const int64_t gi = tok * p.num_heads_v + hv;
float g = 0.0f, be = 0.0f;
if (tid < len) {
- g = needs_decay ? EffectiveDecay(pack.decay[gi], pack.a_log, pack.dt_bias, hv,
+ g = needs_decay ? EffectiveDecay(LoadGate(pack.decay, gi, pack.gates_in_input_type), pack.a_log, pack.dt_bias, hv,
p.gate_activation)
: 0.0f;
- be = pack.beta != nullptr ? EffectiveBeta(pack.beta[gi], p.beta_activation) : 1.0f;
+ be = pack.beta != nullptr ? EffectiveBeta(LoadGate(pack.beta, gi, pack.gates_in_input_type), p.beta_activation) : 1.0f;
}
s.gc[tid] = g;
s.beta[tid] = be;
@@ -521,10 +521,13 @@ __device__ __noinline__ void RunScanStepExactBf16Fallback(
const int64_t gi = tok * p.num_heads_v + hv;
float g = 0.0f, be = 0.0f;
if (tid < len) {
- g = needs_decay ? EffectiveDecay(pack.decay[gi], pack.a_log, pack.dt_bias, hv,
- p.gate_activation)
+ g = needs_decay ? EffectiveDecay(LoadGate(pack.decay, gi, pack.gates_in_input_type),
+ pack.a_log, pack.dt_bias, hv, p.gate_activation)
: 0.0f;
- be = pack.beta != nullptr ? EffectiveBeta(pack.beta[gi], p.beta_activation) : 1.0f;
+ be = pack.beta != nullptr
+ ? EffectiveBeta(LoadGate(pack.beta, gi, pack.gates_in_input_type),
+ p.beta_activation)
+ : 1.0f;
}
gc[tid] = g;
beta[tid] = be;
diff --git a/onnxruntime/contrib_ops/webgpu/bert/causal_conv_with_state.cc b/onnxruntime/contrib_ops/webgpu/bert/causal_conv_with_state.cc
index a47b313e482fd..a4042ecd9340c 100644
--- a/onnxruntime/contrib_ops/webgpu/bert/causal_conv_with_state.cc
+++ b/onnxruntime/contrib_ops/webgpu/bert/causal_conv_with_state.cc
@@ -38,6 +38,12 @@ ONNX_OPERATOR_KERNEL_EX(
CausalConvWithState::CausalConvWithState(const OpKernelInfo& info)
: WebGpuKernel(info) {
+ ORT_ENFORCE(info.GetAttrOrDefault("input_channels", 0) == 0 &&
+ info.GetAttrOrDefault("input_channel_offset", 0) == 0 &&
+ info.GetAttrOrDefault("state_length", 0) == 0 &&
+ info.GetAttrOrDefault("state_channels_last", -1) == -1 &&
+ !causal_conv_with_state_helper::HasOutput(info, 2),
+ "WebGPU CausalConvWithState does not support extended channel/state options (CUDA only)");
std::string activation_str = info.GetAttrOrDefault("activation", "none");
activation_ = ParseCausalConvActivation(activation_str);
ORT_ENFORCE(info.GetAttr("ndim", &ndim_).IsOK(), "Attribute 'ndim' is required");
diff --git a/onnxruntime/contrib_ops/webgpu/bert/gated_delta_net.cc b/onnxruntime/contrib_ops/webgpu/bert/gated_delta_net.cc
index 2294389a9728c..e1c9560342a6e 100644
--- a/onnxruntime/contrib_ops/webgpu/bert/gated_delta_net.cc
+++ b/onnxruntime/contrib_ops/webgpu/bert/gated_delta_net.cc
@@ -37,6 +37,7 @@ ONNX_OPERATOR_KERNEL_EX(
(*KernelDefBuilder::Create())
.TypeConstraint("T", WebGpuSupportedFloatTypes())
.TypeConstraint("TS", DataTypeImpl::GetTensorType())
+ .TypeConstraint("TG", DataTypeImpl::GetTensorType())
.TypeConstraint("TI", DataTypeImpl::GetTensorType())
.InputMemoryType(OrtMemTypeCPUInput, 10)
.MayInplace(6, 1),
@@ -47,9 +48,10 @@ GatedDeltaNet::GatedDeltaNet(const OpKernelInfo& info) : WebGpuKernel(info) {
ORT_ENFORCE(update_rule_ != GatedDeltaNetUpdateRule::Invalid,
"update_rule must be one of: linear, gated, delta, gated_delta");
const auto gate_activation = info.GetAttrOrDefault("gate_activation", "none");
- ORT_ENFORCE(gate_activation == "none" || gate_activation == "qwen",
- "gate_activation must be one of: none, qwen");
- qwen_gate_ = gate_activation == "qwen";
+ ORT_ENFORCE(gate_activation == "none" || gate_activation == "softplus" ||
+ gate_activation == "qwen",
+ "gate_activation must be one of: none, softplus (legacy alias: qwen)");
+ qwen_gate_ = gate_activation == "softplus" || gate_activation == "qwen";
const auto beta_activation = info.GetAttrOrDefault("beta_activation", "none");
ORT_ENFORCE(beta_activation == "none" || beta_activation == "sigmoid",
"beta_activation must be one of: none, sigmoid");
@@ -451,9 +453,9 @@ Status GatedDeltaNet::ComputeInternal(onnxruntime::webgpu::ComputeContext& conte
ORT_RETURN_IF_NOT(a_log != nullptr && dt_bias != nullptr &&
a_log->Shape().NumDimensions() == 1 && a_log->Shape()[0] == hv &&
dt_bias->Shape().NumDimensions() == 1 && dt_bias->Shape()[0] == hv,
- "gate_activation=qwen requires a_log and dt_bias with shape [num_heads_v]");
+ "gate_activation=softplus requires a_log and dt_bias with shape [num_heads_v]");
} else {
- ORT_RETURN_IF_NOT(a_log == nullptr && dt_bias == nullptr, "a_log and dt_bias require gate_activation=qwen");
+ ORT_RETURN_IF_NOT(a_log == nullptr && dt_bias == nullptr, "a_log and dt_bias require gate_activation=softplus");
}
TensorShapeVector output_dims(q_shape.GetDims().begin(), q_shape.GetDims().begin() + token_dims);
diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc
index 56d57cec1d700..317b50da02c83 100644
--- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc
+++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc
@@ -4690,6 +4690,26 @@ ONNX_MS_OPERATOR_SET_SCHEMA(
CausalConvWithState, 1,
OpSchema()
.SetDoc(CausalConvWithState_ver1_doc)
+ .Attr("input_channels",
+ "CUDA channels-last input channel slice width; 0 selects all channels. "
+ "Nondefault extended attributes require rank-3 input and state_window=0.",
+ AttributeProto::INT, static_cast(0))
+ .Attr("input_channel_offset", "Starting channel of the channels-last input slice (CUDA only).",
+ AttributeProto::INT, static_cast(0))
+ .Attr("state_length",
+ "Carry history length (CUDA only); 0 uses (kernel_size-1)*dilation. "
+ "An explicit length must cover the receptive field; convolution reads its trailing history.",
+ AttributeProto::INT, static_cast(0))
+ .Attr("state_channels_last",
+ "State layout override (CUDA only): -1 follows channels_last, "
+ "0 uses [B,C,state_length], 1 uses [B,state_length,C].",
+ AttributeProto::INT, static_cast(-1))
+ .Attr("state_update_capacity",
+ "Compact speculative state-update capture capacity (CUDA extended path only, 0..8); "
+ "when > 0, output 2 has shape (batch_size, state_update_capacity, channels) and records "
+ "the first min(sequence_length, state_update_capacity) projected input rows and any "
+ "remaining rows are zero.",
+ AttributeProto::INT, static_cast(0))
.Attr("activation",
"Fused activation function. One of: 'silu', 'swish', 'none'. "
"Default is 'none'.",
@@ -4733,7 +4753,8 @@ ONNX_MS_OPERATOR_SET_SCHEMA(
"input",
"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.",
+ "the shape is (batch_size, sequence_length, d_1, ..., d_n) instead. Extended channel "
+ "selection uses rank-3 (batch_size, sequence_length, input_width).",
"T")
.Input(1,
"weight",
@@ -4751,12 +4772,14 @@ ONNX_MS_OPERATOR_SET_SCHEMA(
"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.",
+ "The state_length and state_channels_last attributes may override the length and "
+ "layout for the CUDA extended path. If not provided, padding is zero.",
"T",
OpSchema::Optional)
.Output(0,
"output",
- "Convolution output with same shape as input.",
+ "Convolution output with same shape as input, except the final dimension is "
+ "input_channels when extended channel selection is enabled.",
"T")
.Output(1,
"present_state",
@@ -4765,8 +4788,14 @@ ONNX_MS_OPERATOR_SET_SCHEMA(
"(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).",
+ "(seq_len - W + j). Extended attributes override the state length and layout "
+ "as for past_state.",
"T")
+ .Output(2, "state_update",
+ "Optional compact state-update capture with shape (batch_size, state_update_capacity, channels) "
+ "when state_update_capacity > 0. CUDA extended path only; requires rank-3 channels-last input "
+ "and state_window=0.",
+ "T", OpSchema::Optional)
.TypeConstraint("T",
{"tensor(float)", "tensor(float16)", "tensor(bfloat16)"},
"Constrain input and output types to float tensors.")
@@ -4775,6 +4804,75 @@ ONNX_MS_OPERATOR_SET_SCHEMA(
propagateElemTypeFromInputToOutput(ctx, 0, 1);
const int64_t state_window = getAttribute(ctx, "state_window", 0);
+ const int64_t input_channels = getAttribute(ctx, "input_channels", 0);
+ const int64_t channel_offset = getAttribute(ctx, "input_channel_offset", 0);
+ const int64_t history = getAttribute(ctx, "state_length", 0);
+ const int64_t state_layout = getAttribute(ctx, "state_channels_last", -1);
+ const int64_t state_update_capacity = getAttribute(ctx, "state_update_capacity", 0);
+ const bool has_state_update_output = ctx.getNumOutputs() > 2 && ctx.hasOutput(2);
+ if (input_channels < 0 || input_channels > std::numeric_limits::max() ||
+ channel_offset < 0 || channel_offset > std::numeric_limits::max() ||
+ history < 0 || history > std::numeric_limits::max() ||
+ state_layout < -1 || state_layout > 1 ||
+ state_update_capacity < 0 || state_update_capacity > 8) {
+ fail_shape_inference("CausalConvWithState: invalid extended attributes");
+ }
+ if (input_channels != 0 || channel_offset != 0 || history != 0 || state_layout != -1 ||
+ state_update_capacity != 0 || has_state_update_output) {
+ if (getAttribute(ctx, "channels_last", 0) != 1 || getAttribute(ctx, "ndim", 1) != 1 ||
+ state_window != 0 || getAttribute(ctx, "dilation", 1) < 1) {
+ fail_shape_inference("CausalConvWithState: extended attributes require channels_last=1, ndim=1, state_window=0");
+ }
+ if (has_state_update_output && state_update_capacity <= 0) {
+ fail_shape_inference("CausalConvWithState: state_update output requires state_update_capacity > 0");
+ }
+ if (has_state_update_output) propagateElemTypeFromInputToOutput(ctx, 0, 2);
+ if (!hasInputShape(ctx, 0)) return;
+ const auto& input = getInputShape(ctx, 0);
+ if (input.dim_size() != 3) fail_shape_inference("CausalConvWithState: extended input must have rank 3");
+ TensorShapeProto output = input;
+ auto* channels = output.mutable_dim(2);
+ if (input_channels > 0) channels->set_dim_value(input_channels);
+ if (input.dim(2).has_dim_value() &&
+ (channel_offset > input.dim(2).dim_value() ||
+ (channels->has_dim_value() && channels->dim_value() > input.dim(2).dim_value() - channel_offset))) {
+ fail_shape_inference("CausalConvWithState: channel slice exceeds input width");
+ }
+ updateOutputShape(ctx, 0, output);
+ TensorShapeProto::Dimension length;
+ if (history > 0) length.set_dim_value(history);
+ if (hasInputShape(ctx, 1)) {
+ const auto& weight = getInputShape(ctx, 1);
+ if (weight.dim_size() != 3) fail_shape_inference("CausalConvWithState: weight must have rank 3");
+ if (weight.dim(0).has_dim_value() && channels->has_dim_value() &&
+ weight.dim(0).dim_value() != channels->dim_value()) {
+ fail_shape_inference("CausalConvWithState: weight channels mismatch");
+ }
+ if (weight.dim(2).has_dim_value()) {
+ const int64_t kernel = weight.dim(2).dim_value();
+ const int64_t dilation = getAttribute(ctx, "dilation", 1);
+ if (kernel <= 0 || kernel - 1 > std::numeric_limits::max() / dilation) {
+ fail_shape_inference("CausalConvWithState: invalid receptive field");
+ }
+ const int64_t pad = (kernel - 1) * dilation;
+ if (history != 0 && history < pad) fail_shape_inference("CausalConvWithState: state_length is too short");
+ if (history == 0) length.set_dim_value(pad);
+ }
+ }
+ TensorShapeProto state;
+ *state.add_dim() = input.dim(0);
+ *state.add_dim() = state_layout == 0 ? *channels : length;
+ *state.add_dim() = state_layout == 0 ? length : *channels;
+ updateOutputShape(ctx, 1, state);
+ if (state_update_capacity > 0 && has_state_update_output) {
+ TensorShapeProto update_shape;
+ *update_shape.add_dim() = input.dim(0);
+ update_shape.add_dim()->set_dim_value(state_update_capacity);
+ *update_shape.add_dim() = *channels;
+ updateOutputShape(ctx, 2, update_shape);
+ }
+ return;
+ }
if (state_window < 0 || state_window > kMaxStateWindow) {
fail_shape_inference("CausalConvWithState: state_window must be in [0, ", kMaxStateWindow,
"], got ", state_window);
@@ -5448,7 +5546,7 @@ The delta family ('delta' and 'gated_delta') requires L2-normalized keys. Withou
per-chunk system (I + M) is arbitrarily ill-conditioned and the recurrence diverges. Either
normalize upstream or set `qk_l2_norm=1` to have the operator do it.
-Fused activations. `gate_activation='qwen'` computes the effective decay in float32 from the
+Fused activations. `gate_activation='softplus'` computes the effective decay in float32 from the
raw projection carried by `decay`:
g = -exp(a_log) * Softplus(decay + dt_bias)
@@ -5456,6 +5554,9 @@ raw projection carried by `decay`:
`beta_activation='sigmoid'` applies a sigmoid to `beta`, and `qk_l2_norm=1` L2-normalizes each
query and key head vector. Folding these in avoids materializing the intermediates and keeps
the gate arithmetic in float32 independent of the input type.
+Decay and beta may be float or the query element type (both use the same type). CUDA widens
+FP16/BF16 gate inputs to float before any gate arithmetic, avoiding standalone lossless casts.
+State, a_log, and dt_bias remain float. WebGPU currently accepts float decay and beta only.
)DOC";
@@ -5470,8 +5571,9 @@ ONNX_MS_OPERATOR_SET_SCHEMA(
"Output scaling factor. When 0.0 (default) uses 1/sqrt(head_size_qk).",
AttributeProto::FLOAT, 0.0f)
.Attr("gate_activation",
- "'none' (default) treats `decay` as the effective log-space decay. 'qwen' computes "
- "-exp(a_log) * Softplus(decay + dt_bias) in float32.",
+ "'none' (default) treats `decay` as the effective log-space decay. 'softplus' computes "
+ "-exp(a_log) * Softplus(decay + dt_bias) in float32. 'qwen' is accepted as a legacy "
+ "alias of 'softplus' with identical behavior.",
AttributeProto::STRING, std::string("none"))
.Attr("beta_activation",
"'none' (default) treats `beta` as the effective update rate. 'sigmoid' applies a "
@@ -5502,20 +5604,21 @@ ONNX_MS_OPERATOR_SET_SCHEMA(
"TI", OpSchema::Optional)
.Input(4, "decay",
"Log-space decay, shape (total_tokens, num_heads_v) for a scalar per-head decay or "
- "(total_tokens, num_heads_v, head_size_qk) for a per-key-dimension decay.",
- "TS", OpSchema::Optional)
- .Input(5, "beta", "Update rate, shape (total_tokens, num_heads_v)", "TS",
+ "(total_tokens, num_heads_v, head_size_qk) for a per-key-dimension decay. "
+ "Float or the query element type; converted to float before gate arithmetic.",
+ "TG", OpSchema::Optional)
+ .Input(5, "beta", "Update rate, shape (total_tokens, num_heads_v). Same element type as decay.", "TG",
OpSchema::Optional)
.Input(6, "initial_state",
"Recurrent state, shape (batch_size, num_heads_v, head_size_v, head_size_qk), "
"V-major. Required for packed QKV. May alias final_state.",
"TS", OpSchema::Optional)
.Input(7, "a_log",
- "Per-head A_log, shape (num_heads_v). Requires gate_activation=qwen.",
+ "Per-head A_log, shape (num_heads_v). Requires gate_activation=softplus.",
"TS", OpSchema::Optional)
.Input(8, "dt_bias",
"Per-head gate bias, shape (num_heads_v). "
- "Requires gate_activation=qwen.",
+ "Requires gate_activation=softplus.",
"TS", OpSchema::Optional)
.Input(9, "capture_count",
"Number of leading local token transitions to capture for each request, shape "
@@ -5540,7 +5643,9 @@ ONNX_MS_OPERATOR_SET_SCHEMA(
.TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"},
"Constrain query/key/value/output types.")
.TypeConstraint("TS", {"tensor(float)"},
- "State, gate, beta and compact state-update tensors are always float.")
+ "State, gate parameters and compact state-update tensors are always float.")
+ .TypeConstraint("TG", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"},
+ "Decay and beta are float or match the query element type. Gate arithmetic is float.")
.TypeConstraint("TI", {"tensor(int32)"}, "Constrain index and count tensors to int32.")
.TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) {
propagateElemTypeFromInputToOutput(ctx, 0, 0);
diff --git a/onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc b/onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc
index 0adc85ac191b5..bb9bcf03de1ad 100644
--- a/onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc
+++ b/onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc
@@ -1,6 +1,7 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
+#include
#include
#include
#include
@@ -319,6 +320,228 @@ static void RunCausalConvWithStateTests(
// Basic tests - simple cases
// =============================================================================
+template
+void TestPackedChannelHistory(int length, bool state_last, bool use_past, bool use_bias, bool silu, int dilation,
+ int capacity = 0, int channels = 7) {
+ constexpr int batch = 2, offset = 3, kernel = 4;
+ const int width = channels + 6;
+ const int history = (kernel - 1) * dilation + 1;
+ auto cuda = DefaultCudaExecutionProvider();
+ if (!cuda) GTEST_SKIP() << "CUDA is required for extended convolution";
+ const auto round = [](float value) { return static_cast(T(value)); };
+ std::vector input(batch * length * width), weight(channels * kernel), bias(channels);
+ std::vector past(batch * channels * history), output(batch * length * channels);
+ std::vector present(past.size());
+ for (size_t i = 0; i < input.size(); ++i) input[i] = T(float(int(i % 17) - 8) / 16);
+ for (size_t i = 0; i < weight.size(); ++i) weight[i] = T(float(int(i % 11) - 5) / 32);
+ for (size_t i = 0; i < bias.size(); ++i) bias[i] = T(float(i) / 64);
+ for (size_t i = 0; i < past.size(); ++i) past[i] = T(float(int(i % 13) - 6) / 8);
+ const auto state_index = [&](int b, int t, int c) {
+ return state_last ? (b * history + t) * channels + c : (b * channels + c) * history + t;
+ };
+ for (int b = 0; b < batch; ++b) {
+ for (int c = 0; c < channels; ++c) {
+ const auto value = [&](int t) {
+ return t >= 0 ? static_cast(input[(b * length + t) * width + offset + c])
+ : (use_past ? static_cast(past[state_index(b, history + t, c)]) : 0.0f);
+ };
+ for (int t = 0; t < length; ++t) {
+ float sum = use_bias ? static_cast(bias[c]) : 0.0f;
+ for (int j = 0; j < kernel; ++j) sum += value(t - (kernel - 1) * dilation + j * dilation) * static_cast(weight[c * kernel + j]);
+ if (silu) sum /= 1.0f + std::exp(-sum);
+ output[(b * length + t) * channels + c] = T(round(sum));
+ }
+ for (int t = 0; t < history; ++t) {
+ present[state_index(b, t, c)] = T(value(length - history + t));
+ }
+ }
+ }
+ // Captured rows are the raw projected input; rows beyond min(length, capacity) are defined as zero.
+ std::vector state_update(static_cast(batch) * capacity * channels, T(0.0f));
+ for (int b = 0; b < batch; ++b) {
+ for (int t = 0; t < std::min(length, capacity); ++t) {
+ for (int c = 0; c < channels; ++c) {
+ state_update[(static_cast(b) * capacity + t) * channels + c] = input[(b * length + t) * width + offset + c];
+ }
+ }
+ }
+ OpTester test("CausalConvWithState", 1, kMSDomain);
+ test.AddAttribute("channels_last", 1);
+ if (capacity > 0) test.AddAttribute("state_update_capacity", capacity);
+ test.AddAttribute("input_channels", channels);
+ test.AddAttribute("input_channel_offset", offset);
+ test.AddAttribute