[CUDA] Add SpeculativePrefixCommit contrib operator

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)

docs/ContribOperators.md

@@ -131,6 +131,7 @@ Do not modify directly.*
   * <a href="#com.microsoft.SparseAttentionIndexer">com.microsoft.SparseAttentionIndexer</a>
   * <a href="#com.microsoft.SparsePagedAttention">com.microsoft.SparsePagedAttention</a>
   * <a href="#com.microsoft.SparseToDenseMatMul">com.microsoft.SparseToDenseMatMul</a>
+  * <a href="#com.microsoft.SpeculativePrefixCommit">com.microsoft.SpeculativePrefixCommit</a>
   * <a href="#com.microsoft.Tokenizer">com.microsoft.Tokenizer</a>
   * <a href="#com.microsoft.TorchEmbedding">com.microsoft.TorchEmbedding</a>
   * <a href="#com.microsoft.TransposeMatMul">com.microsoft.TransposeMatMul</a>
@@ -8171,6 +8172,61 @@ This version of the operator has been available since version 1 of the 'com.micr
 </dl>
 
 
+### <a name="com.microsoft.SpeculativePrefixCommit"></a><a name="com.microsoft.speculativeprefixcommit">**com.microsoft.SpeculativePrefixCommit**</a>
+
+  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
+
+<dl>
+<dt><tt>head_rows</tt> : int</dt>
+<dd>Number of verifier rows, 3 through 8. Non-muted requests must provide exactly this many rows. Default is 8.</dd>
+</dl>
+
+#### Inputs
+
+<dl>
+<dt><tt>next_tokens</tt> : I64</dt>
+<dd>Verifier predictions with shape [batch, rows], or [rows] for batch one.</dd>
+<dt><tt>input_ids</tt> : I64</dt>
+<dd>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.</dd>
+<dt><tt>status</tt> : I32</dt>
+<dd>Verifier status words, shared by all batch items.</dd>
+<dt><tt>ctl</tt> : I64</dt>
+<dd>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.</dd>
+<dt><tt>ctx_commit</tt> : I64</dt>
+<dd>Committed context, two int64 words per batch item.</dd>
+<dt><tt>prefixes</tt> : I64</dt>
+<dd>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.</dd>
+<dt><tt>ptrs</tt> : I64</dt>
+<dd>Result ring address and number of records per batch item (must be positive).</dd>
+</dl>
+
+#### Outputs
+
+<dl>
+<dt><tt>head_c</tt> : I32</dt>
+<dd>Cache length before the commit [batch]; zero for muted or invalid requests.</dd>
+<dt><tt>updated_ctl</tt> : I64</dt>
+<dd>Updated control state, same shape as ctl. May alias ctl.</dd>
+<dt><tt>updated_ctx_commit</tt> : I64</dt>
+<dd>Updated committed context, same shape as ctx_commit. May alias ctx_commit.</dd>
+</dl>
+
+#### Type Constraints
+
+<dl>
+<dt><tt>I64</tt> : tensor(int64)</dt>
+<dd>Token, control and pointer words.</dd>
+<dt><tt>I32</tt> : tensor(int32)</dt>
+<dd>Status words and cache lengths.</dd>
+</dl>
+
+
 ### <a name="com.microsoft.Tokenizer"></a><a name="com.microsoft.tokenizer">**com.microsoft.Tokenizer**</a>
 
   Tokenizer divides each string in X into a vector of strings along the last axis. Allowed input shapes are [C] and [N, C].

docs/OperatorKernels.md

@@ -1177,6 +1177,7 @@ The **OpSet Version** column uses the following notation:
 |SparseAttention|*in* query:**T**<br> *in* key:**T**<br> *in* value:**T**<br> *in* past_key:**T**<br> *in* past_value:**T**<br> *in* block_row_indices:**M**<br> *in* block_col_indices:**M**<br> *in* total_sequence_length:**M**<br> *in* key_total_sequence_lengths:**M**<br> *in* cos_cache:**T**<br> *in* sin_cache:**T**<br> *out* output:**T**<br> *out* present_key:**T**<br> *out* present_value:**T**|1+|**M** = tensor(int32)<br/> **T** = tensor(bfloat16), tensor(float16)|
 |SparseAttentionIndexer|*in* query:**T**<br> *in* key:**T**<br> *in* query_norm_weight:**T**<br> *in* key_norm_weight:**T**<br> *in* cos_cache:**T**<br> *in* sin_cache:**T**<br> *in* mask:**TB**<br> *in* past_key:**T**<br> *in* gate:**T**<br> *in* position_bias:**T**<br> *in* head_weights:**T**<br> *in* position_ids:**I**<br> *in* past_sequence_length:**M**<br> *in* past_proj_buffer:**T**<br> *out* selected_indices:**M**<br> *out* present_key:**T**<br> *out* present_proj_buffer:**T**|1+|**I** = tensor(int64)<br/> **M** = tensor(int32)<br/> **T** = tensor(bfloat16), tensor(float), tensor(float16)<br/> **TB** = tensor(int64)|
 |SparsePagedAttention|*in* query:**T**<br> *in* key:**T**<br> *in* value:**T**<br> *in* key_cache:**T_CACHE**<br> *in* value_cache:**T_CACHE**<br> *in* cumulative_sequence_length:**S**<br> *in* past_seqlens:**S**<br> *in* block_table:**S**<br> *in* slot_mapping:**S**<br> *in* selected_indices:**S**<br> *in* selected_counts:**S**<br> *in* auxiliary_key:**T**<br> *in* auxiliary_value:**T**<br> *in* auxiliary_lengths:**S**<br> *in* cos_cache:**T**<br> *in* sin_cache:**T**<br> *in* head_sink:**T**<br> *in* q_norm_weight:**T**<br> *in* k_norm_weight:**T**<br> *in* k_scale:**T_KV_SCALE**<br> *in* v_scale:**T_KV_SCALE**<br> *in* attention_metadata:**S**<br> *out* output:**T**<br> *out* key_cache_out:**T_CACHE**<br> *out* value_cache_out:**T_CACHE**|1+|**S** = tensor(int32)<br/> **T** = tensor(bfloat16), tensor(float16)<br/> **T_CACHE** = tensor(bfloat16), tensor(float16), tensor(int8)<br/> **T_KV_SCALE** = tensor(float)|
+|SpeculativePrefixCommit|*in* next_tokens:**I64**<br> *in* input_ids:**I64**<br> *in* status:**I32**<br> *in* ctl:**I64**<br> *in* ctx_commit:**I64**<br> *in* prefixes:**I64**<br> *in* ptrs:**I64**<br> *out* head_c:**I32**<br> *out* updated_ctl:**I64**<br> *out* updated_ctx_commit:**I64**|1+|**I32** = tensor(int32)<br/> **I64** = tensor(int64)|
 |TransposeMatMul|*in* A:**T**<br> *in* B:**T**<br> *out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)|
 |Trilu|*in* X:**T**<br> *in* k:**tensor(int64)**<br> *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**<br> *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)|

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.

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<CUDA_MS_OP_TYPED_CLASS_NAME(1, float, MoE)>,
       BuildKernelCreateInfo<CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, MoE)>,
       BuildKernelCreateInfo<CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, MoE)>,
+      BuildKernelCreateInfo<CUDA_MS_OP_CLASS_NAME(1, SpeculativePrefixCommit)>,
       BuildKernelCreateInfo<CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, QMoE)>,
       BuildKernelCreateInfo<CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, QMoE)>,
       BuildKernelCreateInfo<CUDA_MS_OP_TYPED_CLASS_NAME(1, float_float, MultiHeadAttention)>,

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 <climits>
+
+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<const char*>(a.DataRaw());
+  const auto* b_begin = static_cast<const char*>(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<int>(info.GetAttrOrDefault<int64_t>("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<Tensor>(0);
+    const auto* ids = context->Input<Tensor>(1);
+    const auto* status = context->Input<Tensor>(2);
+    const auto* ctl = context->Input<Tensor>(3);
+    const auto* committed = context->Input<Tensor>(4);
+    const auto* prefixes = context->Input<Tensor>(5);
+    const auto* ptrs = context->Input<Tensor>(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<int64_t>(), ids->Data<int64_t>(), status->Data<int32_t>(),
+        static_cast<int>(status->Shape().Size()), ctl->Data<int64_t>(), updated_ctl->MutableData<int64_t>(),
+        committed->Data<int64_t>(), updated_context->MutableData<int64_t>(), prefixes->Data<int64_t>(),
+        ptrs->Data<int64_t>(), head_c->MutableData<int32_t>(), static_cast<int>(rows), head_rows_,
+        static_cast<int>(batch), static_cast<int>(ids->Shape().Size() / batch),
+        static_cast<int>(ctl->Shape().Size() / batch), static_cast<int>(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<int64_t>())
+        .TypeConstraint("I32", DataTypeImpl::GetTensorType<int32_t>()),
+    SpeculativePrefixCommit);
+
+}  // namespace onnxruntime::contrib::cuda