diff --git a/cmake/onnxruntime_unittests.cmake b/cmake/onnxruntime_unittests.cmake
index 149c1cd74f239..99842093ed143 100644
--- a/cmake/onnxruntime_unittests.cmake
+++ b/cmake/onnxruntime_unittests.cmake
@@ -1125,6 +1125,7 @@ if (onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS AND onnxruntime_BUILD_CUDA_EP_AS_P
"${TEST_SRC_DIR}/contrib_ops/cuda_kernels/qmoe_fp4_to_fp8_kernel_test.cc"
"${TEST_SRC_DIR}/contrib_ops/cuda_kernels/qmoe_fp8_compaction_test.cc"
"${TEST_SRC_DIR}/contrib_ops/cuda_kernels/softmax_topk_kernel_test.cc"
+ "${TEST_SRC_DIR}/contrib_ops/cuda_kernels/speculative_fold_test.cc"
)
endif()
@@ -1158,6 +1159,7 @@ if (onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS AND onnxruntime_BUILD_CUDA_EP_AS_P
"${ONNXRUNTIME_ROOT}/contrib_ops/cuda/llm/fpA_intB_gemm_adaptor.cu"
"${ONNXRUNTIME_ROOT}/contrib_ops/cuda/llm/fpA_intB_gemm_preprocessors_impl.cu"
"${ONNXRUNTIME_ROOT}/contrib_ops/cuda/llm/moe_gemm/moe_gemm_kernels_fp16_uint2.cu"
+ "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/llm/speculative_fold.cu"
"${ONNXRUNTIME_ROOT}/contrib_ops/cuda/moe/qmoe_kernels.cu"
)
if (onnxruntime_USE_FP4_QMOE)
diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md
index f84d7593141f2..d77423fa7af28 100644
--- a/docs/ContribOperators.md
+++ b/docs/ContribOperators.md
@@ -131,6 +131,7 @@ Do not modify directly.*
* com.microsoft.SparseAttentionIndexer
* com.microsoft.SparsePagedAttention
* com.microsoft.SparseToDenseMatMul
+ * com.microsoft.SpeculativePrefixCommit
* com.microsoft.Tokenizer
* com.microsoft.TorchEmbedding
* com.microsoft.TransposeMatMul
@@ -8171,6 +8172,61 @@ This version of the operator has been available since version 1 of the 'com.micr
+### **com.microsoft.SpeculativePrefixCommit**
+
+ Commits the accepted prefix of a speculative verification step for each request in the batch and publishes a result-ring record for the host. Acceptance compares the verifier predictions with the draft tokens in input_ids. The operator is stateful: the control state and committed context are read from ctl and ctx_commit and the updated state is written to updated_ctl and updated_ctx_commit, which are required outputs. An output may alias its input (same device buffer) or be disjoint from it; partial overlap is not supported. Callers that keep the state in persistent device buffers must bind the outputs to those buffers, otherwise the updated state is lost. Callers must serialize access to the state buffers. The result ring is described by ptrs and holds ptrs[1] records of eight int64 words for each batch item, laid out consecutively per batch item starting at ptrs[0]; it must be a device-accessible, 8-byte aligned allocation. Each record is published by writing its sequence word last, after a system-scope fence. A request whose ctl[3] muted flag is set (for example during prefill) does not commit state or touch the ring, and may use any number of rows. Invalid device metadata sets error bits in ctl[5] (4: invalid row count or state, 16: invalid result ring) and mutes the request. The status words are OR-ed together and a non-zero value sets bit 2 of ctl[5].
+
+#### Version
+
+This version of the operator has been available since version 1 of the 'com.microsoft' operator set.
+
+#### Attributes
+
+
+- head_rows : int
+- Number of verifier rows, 3 through 8. Non-muted requests must provide exactly this many rows. Default is 8.
+
+
+#### Inputs
+
+
+- next_tokens : I64
+- Verifier predictions with shape [batch, rows], or [rows] for batch one.
+- input_ids : I64
+- Verifier input IDs; the draft tokens are input_ids[b, 1:rows]. The number of elements must be a multiple of the batch size and at least batch * rows.
+- status : I32
+- Verifier status words, shared by all batch items.
+- ctl : I64
+- Control state, 16 int64 words per batch item: [0] pending token, [1] length, [3] muted flag, [4] ring sequence number, [5] error flags, [6] accepted count, [7] cache length. Other words are preserved.
+- ctx_commit : I64
+- Committed context, two int64 words per batch item.
+- prefixes : I64
+- Context pairs per batch item: pair k (words 2k and 2k+1) is committed when k tokens are accepted. At least 2 * head_rows words per batch item.
+- ptrs : I64
+- Result ring address and number of records per batch item (must be positive).
+
+
+#### Outputs
+
+
+- head_c : I32
+- Cache length before the commit [batch]; zero for muted or invalid requests.
+- updated_ctl : I64
+- Updated control state, same shape as ctl. May alias ctl.
+- updated_ctx_commit : I64
+- Updated committed context, same shape as ctx_commit. May alias ctx_commit.
+
+
+#### Type Constraints
+
+
+- I64 : tensor(int64)
+- Token, control and pointer words.
+- I32 : tensor(int32)
+- Status words and cache lengths.
+
+
+
### **com.microsoft.Tokenizer**
Tokenizer divides each string in X into a vector of strings along the last axis. Allowed input shapes are [C] and [N, C].
diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md
index 2497a7e6fd31a..378609f83c9a4 100644
--- a/docs/OperatorKernels.md
+++ b/docs/OperatorKernels.md
@@ -1177,6 +1177,7 @@ The **OpSet Version** column uses the following notation:
|SparseAttention|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* past_key:**T**
*in* past_value:**T**
*in* block_row_indices:**M**
*in* block_col_indices:**M**
*in* total_sequence_length:**M**
*in* key_total_sequence_lengths:**M**
*in* cos_cache:**T**
*in* sin_cache:**T**
*out* output:**T**
*out* present_key:**T**
*out* present_value:**T**|1+|**M** = tensor(int32)
**T** = tensor(bfloat16), tensor(float16)|
|SparseAttentionIndexer|*in* query:**T**
*in* key:**T**
*in* query_norm_weight:**T**
*in* key_norm_weight:**T**
*in* cos_cache:**T**
*in* sin_cache:**T**
*in* mask:**TB**
*in* past_key:**T**
*in* gate:**T**
*in* position_bias:**T**
*in* head_weights:**T**
*in* position_ids:**I**
*in* past_sequence_length:**M**
*in* past_proj_buffer:**T**
*out* selected_indices:**M**
*out* present_key:**T**
*out* present_proj_buffer:**T**|1+|**I** = tensor(int64)
**M** = tensor(int32)
**T** = tensor(bfloat16), tensor(float), tensor(float16)
**TB** = tensor(int64)|
|SparsePagedAttention|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* key_cache:**T_CACHE**
*in* value_cache:**T_CACHE**
*in* cumulative_sequence_length:**S**
*in* past_seqlens:**S**
*in* block_table:**S**
*in* slot_mapping:**S**
*in* selected_indices:**S**
*in* selected_counts:**S**
*in* auxiliary_key:**T**
*in* auxiliary_value:**T**
*in* auxiliary_lengths:**S**
*in* cos_cache:**T**
*in* sin_cache:**T**
*in* head_sink:**T**
*in* q_norm_weight:**T**
*in* k_norm_weight:**T**
*in* k_scale:**T_KV_SCALE**
*in* v_scale:**T_KV_SCALE**
*in* attention_metadata:**S**
*out* output:**T**
*out* key_cache_out:**T_CACHE**
*out* value_cache_out:**T_CACHE**|1+|**S** = tensor(int32)
**T** = tensor(bfloat16), tensor(float16)
**T_CACHE** = tensor(bfloat16), tensor(float16), tensor(int8)
**T_KV_SCALE** = tensor(float)|
+|SpeculativePrefixCommit|*in* next_tokens:**I64**
*in* input_ids:**I64**
*in* status:**I32**
*in* ctl:**I64**
*in* ctx_commit:**I64**
*in* prefixes:**I64**
*in* ptrs:**I64**
*out* head_c:**I32**
*out* updated_ctl:**I64**
*out* updated_ctx_commit:**I64**|1+|**I32** = tensor(int32)
**I64** = tensor(int64)|
|TransposeMatMul|*in* A:**T**
*in* B:**T**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)|
|Trilu|*in* X:**T**
*in* k:**tensor(int64)**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)|
|UnfoldTensor|*in* input:**T**
*out* output:**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)|
diff --git a/docs/contrib_ops/cuda/speculative_fold.md b/docs/contrib_ops/cuda/speculative_fold.md
new file mode 100644
index 0000000000000..05480720096ee
--- /dev/null
+++ b/docs/contrib_ops/cuda/speculative_fold.md
@@ -0,0 +1,86 @@
+# CUDA SpeculativePrefixCommit operator
+
+`com.microsoft::SpeculativePrefixCommit` (schema version 1, CUDA EP only) is the commit step of a speculative
+decoding verification pass. For every request in the batch it accepts the longest prefix of draft tokens that the
+verifier confirmed, advances the persistent control state, selects the committed context for the accepted length, and
+publishes a record to a host-visible result ring that the runtime polls. It runs entirely on the device and can be
+captured in a CUDA graph (no allocation, no synchronization).
+
+## Interface
+
+| Kind | Name | Type / shape | Meaning |
+|---|---|---|---|
+| attribute | `head_rows` | int, 3..8, default 8 | Number of verifier rows. Non-muted requests must provide exactly this many rows. |
+| input 0 | `next_tokens` | int64 `[batch, rows]` or `[rows]` (batch one) | Verifier predictions. Any other rank is rejected. |
+| input 1 | `input_ids` | int64, `batch * stride` elements, `stride >= rows` | Verifier input; the draft tokens are `input_ids[b, 1:rows]`. |
+| input 2 | `status` | int32, any size | Verifier status words, OR-ed together and shared by all batch items. |
+| input 3 | `ctl` | int64, 16 words per item | Control state (see below). |
+| input 4 | `ctx_commit` | int64, 2 words per item | Committed context pair. |
+| input 5 | `prefixes` | int64, `>= 2 * head_rows` words per item | Context pair `k` (words `2k`, `2k+1`) is committed when `k` drafts are accepted. |
+| input 6 | `ptrs` | int64 `[2]` | `{result ring address, records per batch item}`. |
+| output 0 | `head_c` | int32 `[batch]` | Cache length before the commit; zero for muted or invalid requests. |
+| output 1 | `updated_ctl` | int64, shape of `ctl` | Updated control state. **Required.** |
+| output 2 | `updated_ctx_commit` | int64, shape of `ctx_commit` | Updated committed context. **Required.** |
+
+`ctl` words used by the operator: `[0]` pending token, `[1]` length, `[3]` muted flag, `[4]` ring sequence number,
+`[5]` error flags, `[6]` accepted count, `[7]` cache length. All other words are copied unchanged.
+
+## State handling
+
+The operator is stateful but never mutates its inputs through hidden aliasing. It reads `ctl` / `ctx_commit` and
+always writes the new state to `updated_ctl` / `updated_ctx_commit`. Each output may be the same device buffer as its
+input (in-place, which the kernel declares via `MayInplace`) or a disjoint buffer; partially overlapping buffers are
+rejected. The kernel copies the unchanged words when the buffers differ, and each batch item is handled by a single
+thread, so aliasing is safe.
+
+A caller that keeps the state in persistent device buffers must bind the two outputs to those buffers (for example with
+`IOBinding`, passing the same device pointer as the input). An output that ORT allocates itself is a separate buffer,
+so the update would be lost unless the caller feeds the returned tensors into the next run. Callers must serialize
+access to the state buffers.
+
+## Result ring
+
+`ptrs[0]` is the device-accessible address of an 8-byte aligned allocation that holds `batch * ptrs[1]` records of
+eight int64 words; batch item `b` uses records `[b * ptrs[1], (b + 1) * ptrs[1])` and writes sequence `s` to record
+`(s - 1) % ptrs[1]`. Allocation ownership stays with the caller. Words (`low | high << 32` where two values share a
+word):
+
+| Word | Contents |
+|---|---|
+| 0 | sequence number (`ctl[4]` after the commit), written last after `__threadfence_system()` |
+| 1 | `tokens[0]` \| `tokens[4]` |
+| 2 | `tokens[1]` \| `tokens[5]` |
+| 3 | accepted count \| `tokens[6]` |
+| 4 | OR of the status words |
+| 5 | pending token \| `tokens[7]` |
+| 6 | new length |
+| 7 | `tokens[2]` \| `tokens[3]` |
+
+Tokens beyond `rows` are stored as zero. A host poller that reads word 0 first never observes a partial record.
+
+## Muted requests and errors
+
+A request with `ctl[3] != 0` (for example during prefill) does not commit state, does not access the ring, reports
+`head_c = 0`, and may use any number of `next_tokens` rows. Invalid device metadata mutes the request and sets error
+flags in `ctl[5]`: `4` for a negative length or cache length, a non-muted row count different from `head_rows`, and `16`
+for an invalid ring (null or misaligned address, non-positive or overflowing record count, bad sequence number). A
+non-zero OR of the status words sets flag `2`.
+
+## Using it in a graph
+
+Replace the previous fold node as follows: connect the seven inputs above, drop any rotary-cache inputs and the
+duplicate `prefix2` input, and expose `updated_ctl` and `updated_ctx_commit` as graph outputs (or otherwise bind them)
+so the runtime can bind them to its persistent control and context buffers. Only `head_c` is typically consumed as a
+regular graph value.
+
+Build ORT with CUDA and contrib operators enabled. The schema lives in the core library and the kernel in the CUDA
+provider, so both must be rebuilt; ordinary upstream wheels do not contain this operator.
+
+## Tests
+
+- `FoldContribTest.*` (`onnxruntime_provider_test`, kernel level) covers every accepted-prefix length for 3, 4 and 8
+ rows, in-place and separate state outputs, batched per-item state and ring, muted prefill, invalid metadata, the
+ status union, and captured-graph replay.
+- `SpeculativeFoldTest.*` (session level) covers batched S8 commits through `InferenceSession`, carrying the returned
+ state into later runs (the sequence counter advances), muted prefill with more rows than `head_rows`, shape inference
+ for known and unknown `next_tokens` shapes, and rejection of ranks other than 1 and 2.
diff --git a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc
index e8f435f99b4d7..8b3b80eabae94 100644
--- a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc
+++ b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc
@@ -100,6 +100,7 @@ class CUDA_ONNX_OP_TYPED_CLASS_NAME(1, MLFloat16, Crop);
class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, MoE);
class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, MoE);
class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, MoE);
+class CUDA_MS_OP_CLASS_NAME(1, SpeculativePrefixCommit);
class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, QMoE);
class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, QMoE);
class CUDA_MS_OP_TYPED_CLASS_NAME(1, float_float, MultiHeadAttention);
@@ -467,6 +468,7 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) {
BuildKernelCreateInfo,
BuildKernelCreateInfo,
BuildKernelCreateInfo,
+ BuildKernelCreateInfo,
BuildKernelCreateInfo,
BuildKernelCreateInfo,
BuildKernelCreateInfo,
diff --git a/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cc b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cc
new file mode 100644
index 0000000000000..3ccbdec16f350
--- /dev/null
+++ b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cc
@@ -0,0 +1,88 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "contrib_ops/cuda/llm/speculative_fold.h"
+#include "core/providers/cuda/cuda_kernel.h"
+#include "core/common/safeint.h"
+#include
+
+namespace onnxruntime::contrib::cuda {
+using namespace onnxruntime::cuda;
+namespace {
+
+constexpr int64_t kCtlWords = 16;
+
+// Aliased (identical) buffers are supported; partially overlapping ones are not.
+bool IdenticalOrDisjoint(const Tensor& a, const Tensor& b) {
+ const auto* a_begin = static_cast(a.DataRaw());
+ const auto* b_begin = static_cast(b.DataRaw());
+ return a_begin == b_begin || a_begin + a.SizeInBytes() <= b_begin || b_begin + b.SizeInBytes() <= a_begin;
+}
+
+class SpeculativePrefixCommit final : public CudaKernel {
+ public:
+ explicit SpeculativePrefixCommit(const OpKernelInfo& info) : CudaKernel(info) {
+ head_rows_ = SafeInt(info.GetAttrOrDefault("head_rows", 8));
+ ORT_ENFORCE(head_rows_ >= 3 && head_rows_ <= 8, "SpeculativePrefixCommit: head_rows must be in [3, 8]");
+ }
+
+ Status ComputeInternal(OpKernelContext* context) const override {
+ const auto* tokens = context->Input(0);
+ const auto* ids = context->Input(1);
+ const auto* status = context->Input(2);
+ const auto* ctl = context->Input(3);
+ const auto* committed = context->Input(4);
+ const auto* prefixes = context->Input(5);
+ const auto* ptrs = context->Input(6);
+ ORT_RETURN_IF_NOT(tokens && ids && status && ctl && committed && prefixes && ptrs,
+ "SpeculativePrefixCommit: missing required input");
+ const auto& tok_shape = tokens->Shape();
+ const size_t tok_rank = tok_shape.NumDimensions();
+ ORT_RETURN_IF_NOT(tok_rank == 1 || tok_rank == 2,
+ "SpeculativePrefixCommit: next_tokens must have rank 1 ([rows]) or 2 ([batch, rows])");
+ const int64_t batch = tok_rank == 2 ? tok_shape[0] : int64_t{1};
+ const int64_t rows = tok_shape[tok_rank - 1];
+ ORT_RETURN_IF_NOT(batch > 0 && batch <= INT_MAX && rows > 0 && rows <= INT_MAX &&
+ ids->Shape().Size() >= batch * rows && ids->Shape().Size() % batch == 0 &&
+ ids->Shape().Size() / batch <= INT_MAX &&
+ status->Shape().Size() <= INT_MAX,
+ "SpeculativePrefixCommit: invalid token/status shape");
+ ORT_RETURN_IF_NOT(ctl->Shape().Size() >= batch * kCtlWords && ctl->Shape().Size() % batch == 0 &&
+ ctl->Shape().Size() / batch <= INT_MAX &&
+ committed->Shape().Size() == 2 * batch &&
+ prefixes->Shape().Size() >= batch * 2 * head_rows_ && prefixes->Shape().Size() % batch == 0 &&
+ prefixes->Shape().Size() / batch <= INT_MAX &&
+ ptrs->Shape().Size() == 2,
+ "SpeculativePrefixCommit: invalid control/context shape");
+ auto* head_c = context->Output(0, TensorShape({batch}));
+ auto* updated_ctl = context->Output(1, ctl->Shape());
+ auto* updated_context = context->Output(2, committed->Shape());
+ ORT_RETURN_IF_NOT(head_c && updated_ctl && updated_context,
+ "SpeculativePrefixCommit: updated_ctl and updated_ctx_commit outputs are required");
+ ORT_RETURN_IF_NOT(IdenticalOrDisjoint(*ctl, *updated_ctl) && IdenticalOrDisjoint(*committed, *updated_context),
+ "SpeculativePrefixCommit: state outputs must alias or be disjoint from their inputs");
+ return speculative_fold::LaunchPrefixCommit(
+ tokens->Data(), ids->Data(), status->Data(),
+ static_cast(status->Shape().Size()), ctl->Data(), updated_ctl->MutableData(),
+ committed->Data(), updated_context->MutableData(), prefixes->Data(),
+ ptrs->Data(), head_c->MutableData(), static_cast(rows), head_rows_,
+ static_cast(batch), static_cast(ids->Shape().Size() / batch),
+ static_cast(ctl->Shape().Size() / batch), static_cast(prefixes->Shape().Size() / batch),
+ Stream(context));
+ }
+
+ private:
+ int head_rows_;
+};
+} // namespace
+
+ONNX_OPERATOR_KERNEL_EX(
+ SpeculativePrefixCommit, kMSDomain, 1, kCudaExecutionProvider,
+ (*KernelDefBuilder::Create())
+ .MayInplace(3, 1)
+ .MayInplace(4, 2)
+ .TypeConstraint("I64", DataTypeImpl::GetTensorType())
+ .TypeConstraint("I32", DataTypeImpl::GetTensorType()),
+ SpeculativePrefixCommit);
+
+} // namespace onnxruntime::contrib::cuda
diff --git a/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cu b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cu
new file mode 100644
index 0000000000000..2c4c84c7ad938
--- /dev/null
+++ b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cu
@@ -0,0 +1,127 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "contrib_ops/cuda/llm/speculative_fold.h"
+#include "core/common/common.h"
+#include
+
+namespace onnxruntime::contrib::cuda::speculative_fold {
+namespace {
+
+constexpr int64_t kRecordWords = 8;
+constexpr int64_t kRecordBytes = kRecordWords * static_cast(sizeof(int64_t));
+
+// Error flags stored in ctl[5].
+constexpr int64_t kStatusFlag = 2;
+constexpr int64_t kInvalidStateFlag = 4;
+constexpr int64_t kInvalidRingFlag = 16;
+
+__device__ bool RingBase(const int64_t* ptrs, int b, int batch_size, int64_t* base) {
+ const int64_t address = ptrs[0];
+ const int64_t records = ptrs[1];
+ if (address <= 0 || address % alignof(int64_t) != 0 || records <= 0 ||
+ records > LLONG_MAX / (kRecordBytes * batch_size)) {
+ return false;
+ }
+ if (address > LLONG_MAX - records * kRecordBytes * batch_size) return false;
+ *base = address + static_cast(b) * records * kRecordBytes;
+ return true;
+}
+
+__device__ int64_t Pack(int64_t low, int64_t high) {
+ return static_cast(static_cast(static_cast(low)) |
+ (static_cast(static_cast(high)) << 32));
+}
+
+// One thread per batch item: the work is a short sequential compare/commit and a ring write.
+__global__ void PrefixCommitKernel(const int64_t* tokens, const int64_t* ids, const int32_t* status, int status_count,
+ const int64_t* ctl, int64_t* ctl_out, const int64_t* ctx, int64_t* ctx_out,
+ const int64_t* prefixes, const int64_t* ptrs, int32_t* head_c,
+ int rows, int head_rows, int batch_size, int ids_stride, int ctl_stride,
+ int prefix_stride) {
+ const int b = static_cast(blockIdx.x);
+ if (b >= batch_size) return;
+ const int64_t* tokens_b = tokens + static_cast(b) * rows;
+ const int64_t* ids_b = ids + static_cast(b) * ids_stride;
+ const int64_t* prefixes_b = prefixes + static_cast(b) * prefix_stride;
+ const int64_t* ctl_in_b = ctl + static_cast(b) * ctl_stride;
+ int64_t* ctl_b = ctl_out + static_cast(b) * ctl_stride;
+ const int64_t* ctx_in_b = ctx + static_cast(b) * 2;
+ int64_t* ctx_b = ctx_out + static_cast(b) * 2;
+
+ if (ctl_b != ctl_in_b) {
+ for (int i = 0; i < ctl_stride; ++i) ctl_b[i] = ctl_in_b[i];
+ }
+ if (ctx_b != ctx_in_b) {
+ ctx_b[0] = ctx_in_b[0];
+ ctx_b[1] = ctx_in_b[1];
+ }
+
+ const bool muted = ctl_b[3] != 0;
+ const int64_t length = muted ? 0 : ctl_b[1];
+ const int64_t cache_length = muted ? 0 : ctl_b[7];
+ int64_t errors = 0;
+ if (length < 0 || cache_length < 0 || cache_length > INT_MAX || (!muted && rows != head_rows)) {
+ errors |= kInvalidStateFlag;
+ }
+ int64_t ring = 0;
+ if (!muted && (!RingBase(ptrs, b, batch_size, &ring) || ctl_b[4] < 0 || ctl_b[4] == LLONG_MAX)) {
+ errors |= kInvalidRingFlag;
+ }
+ if (errors) {
+ ctl_b[5] |= errors;
+ ctl_b[3] = 1;
+ head_c[b] = 0;
+ return;
+ }
+ head_c[b] = static_cast(cache_length);
+ if (muted) return;
+
+ int accepted = 0;
+ while (accepted + 1 < rows && tokens_b[accepted] == ids_b[accepted + 1]) ++accepted;
+ int64_t validation = 0;
+ for (int i = 0; i < status_count; ++i) validation |= status[i];
+ const int64_t pending = tokens_b[accepted];
+ ctx_b[0] = prefixes_b[2 * accepted];
+ ctx_b[1] = prefixes_b[2 * accepted + 1];
+ ctl_b[0] = pending;
+ ctl_b[1] = length + accepted + 1;
+ ctl_b[7] = cache_length + accepted + 1;
+ ++ctl_b[4];
+ ctl_b[5] |= validation ? kStatusFlag : 0;
+ ctl_b[6] = accepted;
+
+ // Record words (low | high << 32): [1] pred0|pred4, [2] pred1|pred5, [3] accepted|pred6, [5] pending|pred7,
+ // [7] pred2|pred3, [4] validation, [6] new length, [0] sequence (published last).
+ volatile int64_t* record =
+ reinterpret_cast(ring) + kRecordWords * ((ctl_b[4] - 1) % ptrs[1]);
+ record[1] = Pack(tokens_b[0], rows > 4 ? tokens_b[4] : 0);
+ record[2] = Pack(tokens_b[1], rows > 5 ? tokens_b[5] : 0);
+ record[3] = Pack(accepted, rows > 6 ? tokens_b[6] : 0);
+ record[4] = validation;
+ record[5] = Pack(pending, rows > 7 ? tokens_b[7] : 0);
+ record[6] = ctl_b[1];
+ record[7] = Pack(tokens_b[2], rows > 3 ? tokens_b[3] : 0);
+ __threadfence_system();
+ record[0] = ctl_b[4];
+}
+
+Status LaunchStatus(cudaError_t error) {
+ return error == cudaSuccess ? Status::OK()
+ : Status(common::StatusCategory::SYSTEM, static_cast(error), cudaGetErrorString(error));
+}
+} // namespace
+
+Status LaunchPrefixCommit(const int64_t* tokens, const int64_t* ids, const int32_t* status, int status_count,
+ const int64_t* ctl, int64_t* ctl_out, const int64_t* ctx, int64_t* ctx_out,
+ const int64_t* prefixes, const int64_t* ptrs, int32_t* head_c,
+ int rows, int head_rows, int batch_size, int ids_stride, int ctl_stride,
+ int prefix_stride, cudaStream_t stream) {
+ if (batch_size <= 0) return Status::OK();
+ PrefixCommitKernel<<(batch_size), 1, 0, stream>>>(
+ tokens, ids, status, status_count, ctl, ctl_out, ctx, ctx_out, prefixes, ptrs, head_c,
+ rows, head_rows, batch_size, ids_stride, ctl_stride, prefix_stride);
+ return LaunchStatus(cudaGetLastError());
+}
+
+} // namespace onnxruntime::contrib::cuda::speculative_fold
diff --git a/onnxruntime/contrib_ops/cuda/llm/speculative_fold.h b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.h
new file mode 100644
index 0000000000000..e36a715286185
--- /dev/null
+++ b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.h
@@ -0,0 +1,23 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include
+#include
+#include "core/common/status.h"
+
+namespace onnxruntime::contrib::cuda::speculative_fold {
+
+// ctl_out/ctx_out must either be the same buffers as ctl/ctx (in-place) or disjoint from them. The kernel
+// copies the unchanged words when they differ, so callers always read the updated state from the outputs.
+// ctl has ctl_stride (>= 16) words and prefixes has prefix_stride (>= 2 * head_rows) words per batch item.
+// ptrs is {ring address, records per batch item}; batch item b uses the ring at
+// address + b * records * 8 * sizeof(int64_t).
+Status LaunchPrefixCommit(const int64_t* tokens, const int64_t* ids, const int32_t* status, int status_count,
+ const int64_t* ctl, int64_t* ctl_out, const int64_t* ctx, int64_t* ctx_out,
+ const int64_t* prefixes, const int64_t* ptrs, int32_t* head_c,
+ int rows, int head_rows, int batch_size, int ids_stride, int ctl_stride,
+ int prefix_stride, cudaStream_t stream);
+
+} // namespace onnxruntime::contrib::cuda::speculative_fold
diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc
index b18836cf0b764..4514c4b61aad6 100644
--- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc
+++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc
@@ -34,6 +34,72 @@ void MatMulShapeInference(
namespace onnxruntime {
namespace contrib {
+ONNX_MS_OPERATOR_SET_SCHEMA(
+ SpeculativePrefixCommit, 1,
+ OpSchema()
+ .SetDoc("Commits the accepted prefix of a speculative verification step for each request in the batch and "
+ "publishes a result-ring record for the host. Acceptance compares the verifier predictions with the "
+ "draft tokens in input_ids. The operator is stateful: the control state and committed context are "
+ "read from ctl and ctx_commit and the updated state is written to updated_ctl and updated_ctx_commit, "
+ "which are required outputs. An output may alias its input (same device buffer) or be disjoint from it; "
+ "partial overlap is not supported. Callers that keep the state in persistent device buffers must bind "
+ "the outputs to those buffers, otherwise the updated state is lost. Callers must serialize access to "
+ "the state buffers. The result ring is described by ptrs and holds ptrs[1] records of eight int64 "
+ "words for each batch item, laid out consecutively per batch item starting at ptrs[0]; it must be a "
+ "device-accessible, 8-byte aligned allocation. Each record is published by writing its sequence word "
+ "last, after a system-scope fence. A request whose ctl[3] muted flag is set (for example during "
+ "prefill) does not commit state or touch the ring, and may use any number of rows. Invalid device "
+ "metadata sets error bits in ctl[5] (4: invalid row count or state, 16: invalid result ring) and mutes "
+ "the request. The status words are OR-ed together and a non-zero value sets bit 2 of ctl[5].")
+ .Attr("head_rows",
+ "Number of verifier rows, 3 through 8. Non-muted requests must provide exactly this many "
+ "rows. Default is 8.",
+ AttributeProto::INT, static_cast(8))
+ .Input(0, "next_tokens", "Verifier predictions with shape [batch, rows], or [rows] for batch one.", "I64")
+ .Input(1, "input_ids",
+ "Verifier input IDs; the draft tokens are input_ids[b, 1:rows]. The number of "
+ "elements must be a multiple of the batch size and at least batch * rows.",
+ "I64")
+ .Input(2, "status", "Verifier status words, shared by all batch items.", "I32")
+ .Input(3, "ctl",
+ "Control state, 16 int64 words per batch item: [0] pending token, [1] length, [3] muted flag, "
+ "[4] ring sequence number, [5] error flags, [6] accepted count, [7] cache length. Other words are "
+ "preserved.",
+ "I64")
+ .Input(4, "ctx_commit", "Committed context, two int64 words per batch item.", "I64")
+ .Input(5, "prefixes",
+ "Context pairs per batch item: pair k (words 2k and 2k+1) is committed when k tokens are accepted. "
+ "At least 2 * head_rows words per batch item.",
+ "I64")
+ .Input(6, "ptrs", "Result ring address and number of records per batch item (must be positive).", "I64")
+ .Output(0, "head_c", "Cache length before the commit [batch]; zero for muted or invalid requests.", "I32")
+ .Output(1, "updated_ctl", "Updated control state, same shape as ctl. May alias ctl.", "I64")
+ .Output(2, "updated_ctx_commit", "Updated committed context, same shape as ctx_commit. May alias ctx_commit.",
+ "I64")
+ .TypeConstraint("I64", {"tensor(int64)"}, "Token, control and pointer words.")
+ .TypeConstraint("I32", {"tensor(int32)"}, "Status words and cache lengths.")
+ .TypeAndShapeInferenceFunction([](InferenceContext& ctx) {
+ propagateElemTypeFromInputToOutput(ctx, 2, 0);
+ propagateElemTypeFromInputToOutput(ctx, 3, 1);
+ propagateElemTypeFromInputToOutput(ctx, 4, 2);
+ const auto rows = getAttribute(ctx, "head_rows", 8);
+ if (rows < 3 || rows > 8) fail_shape_inference("SpeculativePrefixCommit: head_rows must be in [3, 8]");
+ if (hasInputShape(ctx, 3)) propagateShapeFromInputToOutput(ctx, 3, 1);
+ if (hasInputShape(ctx, 4)) propagateShapeFromInputToOutput(ctx, 4, 2);
+ if (hasInputShape(ctx, 0)) {
+ const auto& tokens_shape = getInputShape(ctx, 0);
+ if (tokens_shape.dim_size() != 1 && tokens_shape.dim_size() != 2) {
+ fail_shape_inference("SpeculativePrefixCommit: next_tokens must have rank 1 or 2");
+ }
+ auto* head_c_shape = getOutputShape(ctx, 0);
+ if (tokens_shape.dim_size() == 2) {
+ *head_c_shape->add_dim() = tokens_shape.dim(0);
+ } else {
+ head_c_shape->add_dim()->set_dim_value(1);
+ }
+ }
+ }));
+
void DecoderAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& ctx) {
// Type inference
ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 0, 0);
diff --git a/onnxruntime/core/graph/contrib_ops/ms_opset.h b/onnxruntime/core/graph/contrib_ops/ms_opset.h
index bec1ef0fde030..1d0dd9e5581a4 100644
--- a/onnxruntime/core/graph/contrib_ops/ms_opset.h
+++ b/onnxruntime/core/graph/contrib_ops/ms_opset.h
@@ -84,6 +84,7 @@ class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, MatMulFpQ4);
#endif
class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, MaxpoolWithMask);
class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, MoE);
+class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, SpeculativePrefixCommit);
class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, QMoE);
class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, MultiHeadAttention);
class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, GroupQueryAttention);
@@ -220,6 +221,7 @@ class OpSet_Microsoft_ver1 {
#endif
fn(GetOpSchema());
fn(GetOpSchema());
+ fn(GetOpSchema());
fn(GetOpSchema());
fn(GetOpSchema());
fn(GetOpSchema());
diff --git a/onnxruntime/test/contrib_ops/cuda_kernels/speculative_fold_test.cc b/onnxruntime/test/contrib_ops/cuda_kernels/speculative_fold_test.cc
new file mode 100644
index 0000000000000..24b5b101c0917
--- /dev/null
+++ b/onnxruntime/test/contrib_ops/cuda_kernels/speculative_fold_test.cc
@@ -0,0 +1,295 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+#include "contrib_ops/cuda/llm/speculative_fold.h"
+
+namespace onnxruntime::test {
+namespace {
+namespace fold = contrib::cuda::speculative_fold;
+
+constexpr int kMaxRows = 8;
+constexpr int kCtlWords = 16;
+constexpr int kRecords = 2;
+constexpr int64_t kUntouched = -7;
+
+void Check(cudaError_t status) {
+ if (status != cudaSuccess) throw std::runtime_error(cudaGetErrorString(status));
+}
+
+template
+struct Buffer {
+ explicit Buffer(size_t count) : count(count) { Check(cudaMalloc(&data, count * sizeof(T))); }
+ ~Buffer() {
+ if (cudaFree(data) != cudaSuccess) std::fprintf(stderr, "Fold test cudaFree failed\n");
+ }
+ Buffer(const Buffer&) = delete;
+ Buffer& operator=(const Buffer&) = delete;
+ void Upload(const std::vector& values) {
+ ASSERT_LE(values.size(), count);
+ Check(cudaMemcpy(data, values.data(), values.size() * sizeof(T), cudaMemcpyHostToDevice));
+ }
+ std::vector Read() const {
+ std::vector values(count);
+ Check(cudaMemcpy(values.data(), data, count * sizeof(T), cudaMemcpyDeviceToHost));
+ return values;
+ }
+ T* data{};
+ size_t count;
+};
+
+int64_t Prediction(int b, int k) { return 1000 * b + 10 * (k + 1); }
+int64_t PrefixWord(int b, int k, int word) { return 1000 * b + 100 * (k + 1) + word; }
+int32_t Low(int64_t word) { return static_cast(static_cast(word) & 0xffffffffu); }
+int32_t High(int64_t word) { return static_cast(static_cast(word) >> 32); }
+
+// Batch item b has length 3 + b, cache length 2 + b and ring sequence `sequence + b`. All ring records start as -7.
+struct Fixture {
+ explicit Fixture(int batch = 1)
+ : batch(batch),
+ tokens(batch * kMaxRows),
+ ids(batch * (kMaxRows + 4)),
+ ctl(batch * kCtlWords),
+ ctl_out(batch * kCtlWords),
+ context(batch * 2),
+ context_out(batch * 2),
+ prefixes(batch * 2 * kMaxRows),
+ ptrs(2),
+ ring(batch * kRecords * 8),
+ status(2),
+ head_c(batch) {}
+
+ void Reset(int rows, const std::vector& accepted, const std::vector& muted, int64_t sequence = 0,
+ int ids_padding = 0) {
+ ids_stride = rows + ids_padding;
+ std::vector predictions, input, control(batch * kCtlWords), contexts, prefix_values;
+ for (int b = 0; b < batch; ++b) {
+ for (int k = 0; k < rows; ++k) predictions.push_back(Prediction(b, k));
+ std::vector item(ids_stride, -1);
+ item[0] = 99;
+ for (int i = 0; i < accepted[b]; ++i) item[i + 1] = Prediction(b, i);
+ input.insert(input.end(), item.begin(), item.end());
+ for (int w = 0; w < kCtlWords; ++w) control[b * kCtlWords + w] = 5000 + w;
+ control[b * kCtlWords + 1] = 3 + b;
+ control[b * kCtlWords + 3] = muted[b] ? 1 : 0;
+ control[b * kCtlWords + 4] = sequence + b;
+ control[b * kCtlWords + 5] = 0;
+ control[b * kCtlWords + 6] = 0;
+ control[b * kCtlWords + 7] = 2 + b;
+ contexts.insert(contexts.end(), {-1 - b, -2 - b});
+ for (int k = 0; k < kMaxRows; ++k) prefix_values.insert(prefix_values.end(), {PrefixWord(b, k, 0), PrefixWord(b, k, 1)});
+ }
+ tokens.Upload(predictions);
+ ids.Upload(input);
+ ctl.Upload(control);
+ ctl_out.Upload(std::vector(batch * kCtlWords, -99));
+ context.Upload(contexts);
+ context_out.Upload(std::vector(batch * 2, -99));
+ prefixes.Upload(prefix_values);
+ ptrs.Upload({reinterpret_cast(ring.data), kRecords});
+ ring.Upload(std::vector(batch * kRecords * 8, kUntouched));
+ status.Upload({0, 0});
+ }
+
+ void Reset(int rows, int accepted, bool muted = false, int64_t sequence = 0) {
+ Reset(rows, std::vector(batch, accepted), std::vector(batch, muted), sequence);
+ }
+
+ Status Run(int rows, int head_rows, bool in_place = true, cudaStream_t stream = nullptr) {
+ in_place_ = in_place;
+ return fold::LaunchPrefixCommit(tokens.data, ids.data, status.data, 2, ctl.data,
+ in_place ? ctl.data : ctl_out.data, context.data,
+ in_place ? context.data : context_out.data, prefixes.data, ptrs.data,
+ head_c.data, rows, head_rows, batch, ids_stride, kCtlWords, 2 * kMaxRows, stream);
+ }
+
+ std::vector Ctl() const { return in_place_ ? ctl.Read() : ctl_out.Read(); }
+ std::vector Context() const { return in_place_ ? context.Read() : context_out.Read(); }
+
+ int batch;
+ int ids_stride{kMaxRows};
+ bool in_place_{true};
+ Buffer tokens, ids, ctl, ctl_out, context, context_out, prefixes, ptrs, ring;
+ Buffer status, head_c;
+};
+
+class FoldContribTest : public testing::Test {
+ protected:
+ void SetUp() override {
+ int devices = 0;
+ if (cudaGetDeviceCount(&devices) != cudaSuccess || devices == 0) GTEST_SKIP() << "CUDA is required";
+ }
+};
+
+void ExpectCommitted(const Fixture& f, int rows, int b, int accepted, int64_t sequence) {
+ const auto ctl = f.Ctl();
+ const int64_t* c = ctl.data() + b * kCtlWords;
+ EXPECT_EQ(c[0], Prediction(b, accepted));
+ EXPECT_EQ(c[1], 3 + b + accepted + 1);
+ EXPECT_EQ(c[3], 0);
+ EXPECT_EQ(c[4], sequence + b + 1);
+ EXPECT_EQ(c[5], 0);
+ EXPECT_EQ(c[6], accepted);
+ EXPECT_EQ(c[7], 2 + b + accepted + 1);
+ for (int w : {2, 8, 9, 10, 11, 12, 13, 14, 15}) EXPECT_EQ(c[w], 5000 + w) << "word " << w;
+ const auto context = f.Context();
+ EXPECT_EQ(context[2 * b], PrefixWord(b, accepted, 0));
+ EXPECT_EQ(context[2 * b + 1], PrefixWord(b, accepted, 1));
+ EXPECT_EQ(f.head_c.Read()[b], 2 + b);
+
+ const auto ring = f.ring.Read();
+ const int64_t* record = ring.data() + (b * kRecords + (sequence + b) % kRecords) * 8;
+ EXPECT_EQ(record[0], sequence + b + 1);
+ EXPECT_EQ(Low(record[1]), Prediction(b, 0));
+ EXPECT_EQ(Low(record[2]), Prediction(b, 1));
+ EXPECT_EQ(Low(record[3]), accepted);
+ EXPECT_EQ(record[4], 0);
+ EXPECT_EQ(Low(record[5]), Prediction(b, accepted));
+ EXPECT_EQ(record[6], 3 + b + accepted + 1);
+ EXPECT_EQ(Low(record[7]), Prediction(b, 2));
+ EXPECT_EQ(High(record[7]), rows > 3 ? Prediction(b, 3) : 0);
+ EXPECT_EQ(High(record[1]), rows > 4 ? Prediction(b, 4) : 0);
+ EXPECT_EQ(High(record[2]), rows > 5 ? Prediction(b, 5) : 0);
+ EXPECT_EQ(High(record[3]), rows > 6 ? Prediction(b, 6) : 0);
+ EXPECT_EQ(High(record[5]), rows > 7 ? Prediction(b, 7) : 0);
+}
+
+TEST_F(FoldContribTest, AllPrefixLengthsAndDirectRingPublication) {
+ Fixture f;
+ for (bool in_place : {true, false}) {
+ for (int rows : {3, 4, 8}) {
+ for (int accepted = 0; accepted < rows; ++accepted) {
+ SCOPED_TRACE(testing::Message() << "in_place=" << in_place << " rows=" << rows << " accepted=" << accepted);
+ f.Reset(rows, accepted, false, 3);
+ ASSERT_TRUE(f.Run(rows, rows, in_place).IsOK());
+ Check(cudaDeviceSynchronize());
+ ExpectCommitted(f, rows, 0, accepted, 3);
+ }
+ }
+ }
+}
+
+TEST_F(FoldContribTest, SeparateOutputsLeaveInputStateUntouched) {
+ Fixture f;
+ f.Reset(8, 5);
+ const auto ctl_before = f.ctl.Read();
+ const auto context_before = f.context.Read();
+ ASSERT_TRUE(f.Run(8, 8, false).IsOK());
+ Check(cudaDeviceSynchronize());
+ EXPECT_EQ(f.ctl.Read(), ctl_before);
+ EXPECT_EQ(f.context.Read(), context_before);
+ ExpectCommitted(f, 8, 0, 5, 0);
+}
+
+TEST_F(FoldContribTest, BatchedRequestsUsePerItemStateAndRing) {
+ constexpr int kBatch = 3;
+ const std::vector accepted{3, 0, 1};
+ const std::vector muted{false, true, false};
+ Fixture f(kBatch);
+ for (bool in_place : {true, false}) {
+ f.Reset(kMaxRows, accepted, muted, 1, /*ids_padding=*/3);
+ ASSERT_TRUE(f.Run(kMaxRows, kMaxRows, in_place).IsOK());
+ Check(cudaDeviceSynchronize());
+ ExpectCommitted(f, kMaxRows, 0, accepted[0], 1);
+ ExpectCommitted(f, kMaxRows, 2, accepted[2], 1);
+ const auto ctl = f.Ctl();
+ EXPECT_EQ(ctl[kCtlWords + 4], 2);
+ EXPECT_EQ(ctl[kCtlWords + 0], 5000);
+ EXPECT_EQ(f.Context()[2], -2);
+ EXPECT_EQ(f.Context()[3], -3);
+ EXPECT_EQ(f.head_c.Read()[1], 0);
+ const auto ring = f.ring.Read();
+ for (int r = 0; r < kRecords * 8; ++r) EXPECT_EQ(ring[kRecords * 8 + r], kUntouched);
+ }
+}
+
+TEST_F(FoldContribTest, MutedPrefillDoesNotAccessRingOrCommitContext) {
+ Fixture f;
+ for (bool in_place : {true, false}) {
+ f.Reset(8, 0, true);
+ f.ptrs.Upload({0, 0});
+ // Muted prefill may use any number of rows; the kernel never reads the tokens.
+ ASSERT_TRUE(f.Run(64, 8, in_place).IsOK());
+ Check(cudaDeviceSynchronize());
+ EXPECT_EQ(f.Context(), (std::vector{-1, -2}));
+ const auto ctl = f.Ctl();
+ EXPECT_EQ(ctl[1], 3);
+ EXPECT_EQ(ctl[3], 1);
+ EXPECT_EQ(ctl[4], 0);
+ EXPECT_EQ(ctl[5], 0);
+ EXPECT_EQ(ctl[7], 2);
+ EXPECT_EQ(f.head_c.Read()[0], 0);
+ EXPECT_EQ(f.ring.Read(), std::vector(kRecords * 8, kUntouched));
+ }
+}
+
+TEST_F(FoldContribTest, InvalidDeviceMetadataMutesBeforeRingAccess) {
+ Fixture f;
+ for (int kind = 0; kind < 6; ++kind) {
+ SCOPED_TRACE(kind);
+ f.Reset(4, 3);
+ auto ctl = f.ctl.Read();
+ int64_t bit = 16;
+ int head_rows = 4;
+ if (kind == 0) f.ptrs.Upload({0, 0});
+ if (kind == 1) ctl[4] = INT64_MAX;
+ if (kind == 2) f.ptrs.Upload({reinterpret_cast(f.ring.data), INT64_MAX});
+ if (kind >= 3) bit = 4;
+ if (kind == 3) ctl[7] = -1;
+ if (kind == 4) ctl[1] = -1;
+ if (kind == 5) head_rows = 5;
+ f.ctl.Upload(ctl);
+ ASSERT_TRUE(f.Run(4, head_rows).IsOK());
+ Check(cudaDeviceSynchronize());
+ EXPECT_EQ(f.head_c.Read()[0], 0);
+ EXPECT_EQ(f.Ctl()[5] & bit, bit);
+ EXPECT_EQ(f.Ctl()[3], 1);
+ EXPECT_EQ(f.Context(), (std::vector{-1, -2}));
+ EXPECT_EQ(f.ring.Read(), std::vector(kRecords * 8, kUntouched));
+ }
+}
+
+TEST_F(FoldContribTest, ValidationStatusIsPublished) {
+ Fixture f;
+ for (const auto& statuses : {std::vector{0, 0}, std::vector{1, 4}}) {
+ f.Reset(4, 2);
+ f.status.Upload(statuses);
+ ASSERT_TRUE(f.Run(4, 4).IsOK());
+ Check(cudaDeviceSynchronize());
+ const int32_t combined = statuses[0] | statuses[1];
+ EXPECT_EQ(f.Ctl()[5], combined ? 2 : 0);
+ EXPECT_EQ(f.ring.Read()[4], combined);
+ EXPECT_EQ(f.ring.Read()[0], 1);
+ }
+}
+
+TEST_F(FoldContribTest, CapturedGraphReplaysWithFreshAcceptance) {
+ Fixture f;
+ cudaStream_t stream{};
+ cudaGraph_t graph{};
+ cudaGraphExec_t executable{};
+ Check(cudaStreamCreate(&stream));
+ f.Reset(4, 0);
+ Check(cudaStreamBeginCapture(stream, cudaStreamCaptureModeThreadLocal));
+ ASSERT_TRUE(f.Run(4, 4, true, stream).IsOK());
+ Check(cudaStreamEndCapture(stream, &graph));
+ Check(cudaGraphInstantiate(&executable, graph, nullptr, nullptr, 0));
+ for (int accepted : {0, 3, 1, 2}) {
+ f.Reset(4, accepted);
+ Check(cudaGraphLaunch(executable, stream));
+ Check(cudaStreamSynchronize(stream));
+ ExpectCommitted(f, 4, 0, accepted, 0);
+ }
+ Check(cudaGraphExecDestroy(executable));
+ Check(cudaGraphDestroy(graph));
+ Check(cudaStreamDestroy(stream));
+}
+} // namespace
+} // namespace onnxruntime::test
diff --git a/onnxruntime/test/contrib_ops/speculative_fold_op_test.cc b/onnxruntime/test/contrib_ops/speculative_fold_op_test.cc
new file mode 100644
index 0000000000000..63501115ae5c2
--- /dev/null
+++ b/onnxruntime/test/contrib_ops/speculative_fold_op_test.cc
@@ -0,0 +1,374 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#if defined(USE_CUDA) && !defined(ORT_MINIMAL_BUILD)
+
+#include
+
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+
+#include "core/graph/model.h"
+#include "core/session/inference_session.h"
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+#include "test/common/cuda_op_test_utils.h"
+#include "test/common/tensor_op_test_utils.h"
+#include "test/unittest_util/framework_test_utils.h"
+#include "test/util/include/asserts.h"
+#include "test/util/include/default_providers.h"
+#include "test/util/include/test_environment.h"
+
+namespace onnxruntime {
+namespace test {
+namespace {
+
+constexpr int kCtlWords = 16;
+
+// Host-mapped result ring that the kernel can write through its device address.
+class FoldResultRing {
+ public:
+ explicit FoldResultRing(size_t words) : words_(words) {
+ const cudaError_t error = cudaHostAlloc(reinterpret_cast(&host_), words * sizeof(int64_t), cudaHostAllocMapped);
+ ORT_ENFORCE(error == cudaSuccess, cudaGetErrorString(error));
+ const cudaError_t pointer_error = cudaHostGetDevicePointer(reinterpret_cast(&device_), host_, 0);
+ if (pointer_error != cudaSuccess) {
+ cudaFreeHost(host_);
+ ORT_THROW(cudaGetErrorString(pointer_error));
+ }
+ std::fill_n(host_, words, int64_t{-1});
+ }
+
+ ~FoldResultRing() { cudaFreeHost(host_); }
+
+ ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(FoldResultRing);
+
+ int64_t Address() const { return static_cast(reinterpret_cast(device_)); }
+ const int64_t* Host() const { return host_; }
+ size_t Words() const { return words_; }
+
+ private:
+ size_t words_;
+ int64_t* host_ = nullptr;
+ int64_t* device_ = nullptr;
+};
+
+// A one-node SpeculativePrefixCommit graph. Outputs are left without types so that inference provides them.
+// token_dims == nullopt leaves next_tokens without a shape; negative dims are symbolic.
+class FoldModel {
+ public:
+ FoldModel(std::optional> token_dims, int batch, int head_rows = 8)
+ : model_("speculative_fold", false, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(),
+ {{kMSDomain, 1}}, {}, DefaultLoggingManager().DefaultLogger()) {
+ auto& graph = model_.MainGraph();
+ constexpr int i64 = ONNX_NAMESPACE::TensorProto_DataType_INT64;
+ constexpr int i32 = ONNX_NAMESPACE::TensorProto_DataType_INT32;
+ auto arg = [&](const char* name, int element_type, const std::optional>& dims) {
+ ONNX_NAMESPACE::TypeProto type;
+ auto* tensor = type.mutable_tensor_type();
+ tensor->set_elem_type(element_type);
+ if (dims) {
+ int axis = 0;
+ for (const auto dim : *dims) {
+ auto* shape_dim = tensor->mutable_shape()->add_dim();
+ if (dim < 0) {
+ shape_dim->set_dim_param(axis == 0 ? "batch" : "rows");
+ } else {
+ shape_dim->set_dim_value(dim);
+ }
+ ++axis;
+ }
+ }
+ return &graph.GetOrCreateNodeArg(name, &type);
+ };
+ auto* tokens = arg("tokens", i64, token_dims);
+ auto* ids = arg("ids", i64, std::vector{-1});
+ auto* status = arg("status", i32, std::vector{1});
+ auto* ctl = arg("ctl", i64, std::vector{batch * kCtlWords});
+ auto* context = arg("context", i64, std::vector{batch * 2});
+ auto* prefixes = arg("prefixes", i64, std::vector{batch * kCtlWords});
+ auto* ptrs = arg("ptrs", i64, std::vector{2});
+ head_c_ = &graph.GetOrCreateNodeArg("head_c", nullptr);
+ auto* updated_ctl = &graph.GetOrCreateNodeArg("updated_ctl", nullptr);
+ auto* updated_context = &graph.GetOrCreateNodeArg("updated_context", nullptr);
+
+ ONNX_NAMESPACE::AttributeProto head_rows_attr;
+ head_rows_attr.set_name("head_rows");
+ head_rows_attr.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_INT);
+ head_rows_attr.set_i(head_rows);
+ NodeAttributes attrs{{"head_rows", head_rows_attr}};
+ graph.AddNode("commit", "SpeculativePrefixCommit", "",
+ {tokens, ids, status, ctl, context, prefixes, ptrs},
+ {head_c_, updated_ctl, updated_context}, &attrs, kMSDomain)
+ .SetExecutionProviderType(kCudaExecutionProvider);
+ graph.SetOutputs({head_c_, updated_ctl, updated_context});
+ resolve_status_ = graph.Resolve();
+ }
+
+ const Status& ResolveStatus() const { return resolve_status_; }
+ const NodeArg* HeadC() const { return head_c_; }
+
+ Status CreateSession(std::unique_ptr& session) {
+ auto provider = DefaultCudaExecutionProvider();
+ if (!provider) return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "CUDA provider is unavailable");
+ std::stringstream stream;
+ if (!model_.ToProto().SerializeToOstream(&stream)) return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "serialize failed");
+ session = std::make_unique(SessionOptions{}, GetEnvironment());
+ ORT_RETURN_IF_ERROR(session->RegisterExecutionProvider(std::move(provider)));
+ ORT_RETURN_IF_ERROR(session->Load(stream));
+ return session->Initialize();
+ }
+
+ private:
+ Model model_;
+ Status resolve_status_;
+ NodeArg* head_c_ = nullptr;
+};
+
+struct FoldInputs {
+ TensorShapeVector token_shape;
+ std::vector tokens;
+ std::vector ids;
+ std::vector ctl;
+ std::vector context;
+ std::vector prefixes;
+ std::vector ptrs;
+};
+
+struct FoldOutputs {
+ std::vector head_c;
+ std::vector ctl;
+ std::vector context;
+};
+
+// Feeds CPU tensors, so ORT copies the state to temporary CUDA tensors: the updated state is only visible
+// through the returned outputs.
+Status RunFold(InferenceSession& session, const FoldInputs& in, FoldOutputs* out) {
+ auto allocator = TestCPUExecutionProvider()->CreatePreferredAllocators()[0];
+ NameMLValMap feeds;
+ auto feed = [&](const char* name, const TensorShapeVector& shape, const auto& values) {
+ using T = typename std::decay_t::value_type;
+ OrtValue value;
+ CreateMLValue(allocator, shape, values, &value);
+ feeds[name] = std::move(value);
+ };
+ feed("tokens", in.token_shape, in.tokens);
+ feed("ids", {static_cast(in.ids.size())}, in.ids);
+ feed("status", {1}, std::vector{0});
+ feed("ctl", {static_cast(in.ctl.size())}, in.ctl);
+ feed("context", {static_cast(in.context.size())}, in.context);
+ feed("prefixes", {static_cast(in.prefixes.size())}, in.prefixes);
+ feed("ptrs", {2}, in.ptrs);
+ std::vector outputs;
+ const std::vector output_names{"head_c", "updated_ctl", "updated_context"};
+ ORT_RETURN_IF_ERROR(session.Run(RunOptions{}, feeds, output_names, &outputs));
+ auto copy = [](const OrtValue& value, auto* dest) {
+ const auto& tensor = value.Get();
+ dest->assign(tensor.DataAsSpan::value_type>().begin(),
+ tensor.DataAsSpan::value_type>().end());
+ };
+ copy(outputs[0], &out->head_c);
+ copy(outputs[1], &out->ctl);
+ copy(outputs[2], &out->context);
+ return Status::OK();
+}
+
+TEST(SpeculativeFoldTest, SupportsBatchedSpeculativePrefixCommitS8) {
+ if (!HasCudaEnvironment(800)) GTEST_SKIP() << "Requires CUDA SM80 or newer";
+ constexpr int kBatch = 3;
+ constexpr int kRows = 8;
+ constexpr int kRingCap = 4;
+ FoldResultRing ring(kBatch * kRingCap * 8);
+
+ FoldModel model(std::vector{-1, kRows}, kBatch);
+ ASSERT_STATUS_OK(model.ResolveStatus());
+ const auto* head_c_shape = model.HeadC()->Shape();
+ ASSERT_NE(head_c_shape, nullptr);
+ ASSERT_EQ(head_c_shape->dim_size(), 1);
+ EXPECT_EQ(head_c_shape->dim(0).dim_param(), "batch");
+ std::unique_ptr session;
+ ASSERT_STATUS_OK(model.CreateSession(session));
+
+ FoldInputs in;
+ in.token_shape = {kBatch, kRows};
+ // Batch 0: accepts 3 drafts (tokens[0..2] == ids[1..3], tokens[3] != ids[4])
+ // Batch 1: muted (ctl[3] = 1)
+ // Batch 2: accepts 1 draft (tokens[0] == ids[1], tokens[1] != ids[2])
+ in.tokens = {10, 11, 12, 99, 14, 15, 16, 17,
+ 20, 21, 22, 23, 24, 25, 26, 27,
+ 30, 88, 32, 33, 34, 35, 36, 37};
+ in.ids = {1, 10, 11, 12, 13, 14, 15, 16,
+ 2, 20, 21, 22, 23, 24, 25, 26,
+ 3, 30, 31, 32, 33, 34, 35, 36};
+ in.ctl.assign(kBatch * kCtlWords, 0);
+ // Batch 0: length=10, muted=0, seq=0, cache_len=10
+ in.ctl[0 * kCtlWords + 1] = 10;
+ in.ctl[0 * kCtlWords + 7] = 10;
+ // Batch 1: length=20, muted=1, seq=5, cache_len=20
+ in.ctl[1 * kCtlWords + 1] = 20;
+ in.ctl[1 * kCtlWords + 3] = 1;
+ in.ctl[1 * kCtlWords + 4] = 5;
+ in.ctl[1 * kCtlWords + 7] = 20;
+ // Batch 2: length=30, muted=0, seq=2, cache_len=30
+ in.ctl[2 * kCtlWords + 1] = 30;
+ in.ctl[2 * kCtlWords + 4] = 2;
+ in.ctl[2 * kCtlWords + 7] = 30;
+ in.ctl[2 * kCtlWords + 2] = 77; // untouched word
+ in.context = {1, 2, 3, 4, 5, 6};
+ in.prefixes.resize(kBatch * kCtlWords);
+ for (int b = 0; b < kBatch; ++b) {
+ for (int r = 0; r < kRows; ++r) {
+ in.prefixes[b * kCtlWords + 2 * r + 0] = 1000 * (b + 1) + 10 * r + 0;
+ in.prefixes[b * kCtlWords + 2 * r + 1] = 1000 * (b + 1) + 10 * r + 1;
+ }
+ }
+ in.ptrs = {ring.Address(), kRingCap};
+
+ FoldOutputs out;
+ ASSERT_STATUS_OK(RunFold(*session, in, &out));
+
+ EXPECT_EQ(out.head_c, (std::vector{10, 0, 30}));
+
+ EXPECT_EQ(out.ctl[0 * kCtlWords + 0], 99); // pending = tokens[0, 3]
+ EXPECT_EQ(out.ctl[0 * kCtlWords + 1], 10 + 4); // length = 10 + 3 + 1
+ EXPECT_EQ(out.ctl[0 * kCtlWords + 4], 1); // seq incremented to 1
+ EXPECT_EQ(out.ctl[0 * kCtlWords + 6], 3); // accepted
+ EXPECT_EQ(out.ctl[0 * kCtlWords + 7], 10 + 4); // cache length
+
+ EXPECT_EQ(out.ctl[1 * kCtlWords + 4], 5); // muted: seq unchanged
+ EXPECT_EQ(out.ctl[1 * kCtlWords + 1], 20);
+
+ EXPECT_EQ(out.ctl[2 * kCtlWords + 0], 88); // pending = tokens[2, 1]
+ EXPECT_EQ(out.ctl[2 * kCtlWords + 1], 30 + 2); // length = 30 + 1 + 1
+ EXPECT_EQ(out.ctl[2 * kCtlWords + 2], 77);
+ EXPECT_EQ(out.ctl[2 * kCtlWords + 4], 3); // seq incremented from 2 to 3
+ EXPECT_EQ(out.ctl[2 * kCtlWords + 6], 1);
+
+ // Batch 0 accepted=3 -> prefixes[0, 6..7]; batch 1 muted; batch 2 accepted=1 -> prefixes[2, 2..3]
+ EXPECT_EQ(out.context, (std::vector{1030, 1031, 3, 4, 3010, 3011}));
+
+ // Ring records published for batch 0 (record 0) and batch 2 (record 2 % kRingCap); batch 1 is untouched.
+ const int64_t* host = ring.Host();
+ EXPECT_EQ(host[(0 * kRingCap + 0) * 8 + 0], 1);
+ for (int word = 0; word < kRingCap * 8; ++word) EXPECT_EQ(host[kRingCap * 8 + word], -1);
+ EXPECT_EQ(host[(2 * kRingCap + 2) * 8 + 0], 3);
+}
+
+TEST(SpeculativeFoldTest, CarriesUpdatedStateIntoNextRun) {
+ if (!HasCudaEnvironment(800)) GTEST_SKIP() << "Requires CUDA SM80 or newer";
+ constexpr int kRows = 8;
+ constexpr int kRingCap = 2;
+ FoldResultRing ring(kRingCap * 8);
+
+ FoldModel model(std::vector{1, kRows}, 1);
+ ASSERT_STATUS_OK(model.ResolveStatus());
+ std::unique_ptr session;
+ ASSERT_STATUS_OK(model.CreateSession(session));
+
+ FoldInputs in;
+ in.token_shape = {1, kRows};
+ in.tokens = {10, 11, 12, 13, 14, 15, 16, 17};
+ in.ids = {1, 10, 11, 99, 99, 99, 99, 99}; // accepts two drafts
+ in.ctl.assign(kCtlWords, 0);
+ in.ctl[1] = 10;
+ in.ctl[7] = 10;
+ in.context = {1, 2};
+ in.prefixes.resize(kCtlWords);
+ for (int i = 0; i < kCtlWords; ++i) in.prefixes[i] = 100 + i;
+ in.ptrs = {ring.Address(), kRingCap};
+
+ for (int run = 1; run <= 3; ++run) {
+ SCOPED_TRACE(run);
+ FoldOutputs out;
+ ASSERT_STATUS_OK(RunFold(*session, in, &out));
+ EXPECT_EQ(out.ctl[4], run); // the sequence counter advances across runs
+ EXPECT_EQ(out.ctl[1], 10 + 3 * run);
+ EXPECT_EQ(out.ctl[7], 10 + 3 * run);
+ EXPECT_EQ(out.head_c[0], 10 + 3 * (run - 1));
+ EXPECT_EQ(out.context, (std::vector{104, 105}));
+ EXPECT_EQ(ring.Host()[((run - 1) % kRingCap) * 8], run);
+ in.ctl = out.ctl;
+ in.context = out.context;
+ }
+ // Three records were published into a two-record ring: record 0 was overwritten by sequence 3, record 1 holds 2.
+ EXPECT_EQ(ring.Host()[0], 3);
+ EXPECT_EQ(ring.Host()[8], 2);
+}
+
+TEST(SpeculativeFoldTest, MutedPrefillAcceptsMoreRowsThanHeadRows) {
+ if (!HasCudaEnvironment(800)) GTEST_SKIP() << "Requires CUDA SM80 or newer";
+ constexpr int kRows = 20;
+ FoldResultRing ring(2 * 8);
+
+ FoldModel model(std::vector{kRows}, 1);
+ ASSERT_STATUS_OK(model.ResolveStatus());
+ const auto* head_c_shape = model.HeadC()->Shape();
+ ASSERT_NE(head_c_shape, nullptr); // a rank-1 next_tokens is batch one
+ ASSERT_EQ(head_c_shape->dim_size(), 1);
+ EXPECT_EQ(head_c_shape->dim(0).dim_value(), 1);
+ std::unique_ptr session;
+ ASSERT_STATUS_OK(model.CreateSession(session));
+
+ FoldInputs in;
+ in.token_shape = {kRows};
+ in.tokens.assign(kRows, 10);
+ in.ids.assign(kRows, 10);
+ in.ctl.assign(kCtlWords, 0);
+ in.ctl[1] = 3;
+ in.ctl[3] = 1;
+ in.ctl[7] = 2;
+ in.context = {1, 2};
+ in.prefixes.assign(kCtlWords, 100);
+ in.ptrs = {0, 0};
+
+ FoldOutputs out;
+ ASSERT_STATUS_OK(RunFold(*session, in, &out));
+ EXPECT_EQ(out.head_c, std::vector{0});
+ EXPECT_EQ(out.ctl, in.ctl);
+ EXPECT_EQ(out.context, in.context);
+ EXPECT_EQ(std::count(ring.Host(), ring.Host() + ring.Words(), int64_t{-1}), static_cast(ring.Words()));
+}
+
+TEST(SpeculativeFoldTest, RejectsNextTokensRankOutsideOneAndTwo) {
+ if (!HasCudaEnvironment(800)) GTEST_SKIP() << "Requires CUDA SM80 or newer";
+ {
+ // Schema inference rejects a statically known rank-3 next_tokens.
+ FoldModel model(std::vector{2, 3, 8}, 2);
+ EXPECT_FALSE(model.ResolveStatus().IsOK());
+ EXPECT_THAT(model.ResolveStatus().ErrorMessage(), testing::HasSubstr("next_tokens must have rank 1 or 2"));
+ }
+ {
+ // Without a declared shape, inference declares nothing (no batch 1) and the kernel rejects the runtime shape.
+ FoldModel model(std::nullopt, 2);
+ ASSERT_STATUS_OK(model.ResolveStatus());
+ EXPECT_EQ(model.HeadC()->Shape(), nullptr);
+ std::unique_ptr session;
+ ASSERT_STATUS_OK(model.CreateSession(session));
+
+ FoldInputs in;
+ in.token_shape = {2, 3, 8};
+ in.tokens.assign(2 * 3 * 8, 10);
+ in.ids.assign(2 * 3 * 8, 10);
+ in.ctl.assign(2 * kCtlWords, 0);
+ in.context = {1, 2, 3, 4};
+ in.prefixes.assign(2 * kCtlWords, 100);
+ in.ptrs = {0, 0};
+ FoldOutputs out;
+ const Status status = RunFold(*session, in, &out);
+ EXPECT_FALSE(status.IsOK());
+ EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("next_tokens must have rank 1 ([rows]) or 2 ([batch, rows])"));
+ }
+}
+
+} // namespace
+} // namespace test
+} // namespace onnxruntime
+
+#endif