Replace speculative fold operators at version 1

docs/ContribOperators.md

@@ -2104,117 +2104,129 @@ This version of the operator has been available since version 1 of the 'com.micr
 
 ### <a name="com.microsoft.FoldMid2"></a><a name="com.microsoft.foldmid2">**com.microsoft.FoldMid2**</a>
 
-  Combines speculative-token acceptance, recurrent-state commit bookkeeping, and head rotary-cache lookup for a
-  two-token verification step. Updated control and committed context are returned explicitly; CUDA may reuse
-  their input buffers. When ctl[3] is nonzero, state is preserved unless validation fails.
-  The pointer tensor contains the device-accessible result-ring address and its positive record capacity;
-  ctl[4] must be a nonnegative record counter below INT64_MAX. The caller must supply an aligned device-accessible
-  allocation holding capacity * 8 int64 elements. Both rotary positions ctl[1]+1 and ctl[1]+2 must be in range.
-  Invalid head-cache length, rotary positions, or ring metadata set error bits 4, 8, or 16 respectively in
-  updated_ctl[5], set updated_ctl[3]=1 (muted), and return sel=-1 with zero head/rotary outputs, without committing
-  context or accessing the ring. Ring metadata is not accessed on muted calls.
+  Commits a batch-one speculative prefix and prepares rotary/head inputs for two through four head rows. Control and committed-context outputs must alias their inputs on CUDA. This is stateful: callers must serialize access to these buffers. The result-ring address must identify a device-accessible, aligned allocation of capacity*8 int64 words. The mid stage writes the payload; FoldTail publishes its sequence word after downstream validation. Muted prefill calls prepare inputs without committing state or touching ring metadata. Invalid device metadata sets error bits and mutes execution.
 
 #### Version
 
 This version of the operator has been available since version 1 of the 'com.microsoft' operator set.
 
-#### Inputs
+#### Attributes
+
+<dl>
+<dt><tt>head_capacity</tt> : int (default is 33028)</dt>
+<dd>Maximum head-cache token capacity.</dd>
+<dt><tt>head_rows</tt> : int (default is 2)</dt>
+<dd>Fixed number of head rows, 2 through 4.</dd>
+</dl>
+
+#### Inputs (10 - 11)
 
 <dl>
 <dt><tt>next_tokens</tt> : I64</dt>
-<dd>Two predicted token IDs.</dd>
+<dd>Verifier predictions, batch one.</dd>
 <dt><tt>input_ids</tt> : I64</dt>
-<dd>Two input token IDs; input_ids[1] is the draft token.</dd>
+<dd>Verifier input IDs.</dd>
 <dt><tt>status</tt> : I32</dt>
-<dd>Status values aggregated into the result record.</dd>
+<dd>Verifier status words.</dd>
 <dt><tt>ctl</tt> : I64</dt>
-<dd>Eight-element control state.</dd>
+<dd>Mutable control state: at least 8 words for S2, 16 for S3/S4.</dd>
 <dt><tt>ctx_commit</tt> : I64</dt>
-<dd>Two-element committed context.</dd>
-<dt><tt>ctx_a</tt> : I64</dt>
-<dd>Context selected when the draft is rejected.</dd>
-<dt><tt>ctx_b</tt> : I64</dt>
-<dd>Context selected when the draft is accepted.</dd>
+<dd>Mutable committed context, two words.</dd>
+<dt><tt>prefixes</tt> : I64</dt>
+<dd>Two-word context for each accepted-prefix length; S2 needs only the first.</dd>
+<dt><tt>prefix2</tt> : I64</dt>
+<dd>Two-word context for acceptance length one.</dd>
 <dt><tt>ptrs</tt> : I64</dt>
-<dd>Device pointers containing the result-ring address and capacity.</dd>
+<dd>Result ring address and positive record capacity.</dd>
 <dt><tt>cos_cache</tt> : B</dt>
-<dd>BFloat16 rotary cosine cache with shape [positions, 32].</dd>
+<dd>Rotary cosine cache [positions, width].</dd>
 <dt><tt>sin_cache</tt> : B</dt>
-<dd>BFloat16 rotary sine cache with shape [positions, 32].</dd>
+<dd>Rotary sine cache [positions, width].</dd>
+<dt><tt>hidden</tt> (optional) : B</dt>
+<dd>Optional hidden states [1, verifier_rows, ...].</dd>
 </dl>
 
 #### Outputs
 
 <dl>
 <dt><tt>head_ids</tt> : I64</dt>
-<dd>Two token IDs consumed by the head.</dd>
+<dd>Head token IDs [1, head_rows].</dd>
 <dt><tt>head_c</tt> : I32</dt>
-<dd>Current head-cache length.</dd>
+<dd>Head committed length [1].</dd>
 <dt><tt>slot_ids</tt> : I64</dt>
-<dd>Two consecutive head-cache slots.</dd>
+<dd>Head slots [1, head_rows].</dd>
 <dt><tt>seqlens</tt> : I32</dt>
-<dd>Head sequence length.</dd>
+<dd>Head sequence length [1].</dd>
 <dt><tt>cos_out</tt> : B</dt>
-<dd>Two rows of rotary cosine values.</dd>
+<dd>Gathered rotary cosine [1, head_rows, width].</dd>
 <dt><tt>sin_out</tt> : B</dt>
-<dd>Two rows of rotary sine values.</dd>
+<dd>Gathered rotary sine [1, head_rows, width].</dd>
 <dt><tt>sel</tt> : I64</dt>
-<dd>Selected row: one when accepted, zero when rejected, or -1 on invalid metadata.</dd>
+<dd>Accepted-prefix length, or -1 on invalid metadata.</dd>
+<dt><tt>head_hidden</tt> (optional) : B</dt>
+<dd>Hidden rows, repeating the last source row when necessary.</dd>
 <dt><tt>updated_ctl</tt> : I64</dt>
-<dd>Updated control state with the same shape as ctl.</dd>
+<dd>Updated control state, aliases ctl.</dd>
 <dt><tt>updated_ctx_commit</tt> : I64</dt>
-<dd>Updated committed context with the same shape as ctx_commit.</dd>
+<dd>Updated context, aliases ctx_commit.</dd>
 </dl>
 
 #### Type Constraints
 
 <dl>
 <dt><tt>I64</tt> : tensor(int64)</dt>
-<dd>Constrain control and token tensors to int64.</dd>
+<dd>Token and control words.</dd>
 <dt><tt>I32</tt> : tensor(int32)</dt>
-<dd>Constrain status and sequence tensors to int32.</dd>
+<dd>Status and sequence lengths.</dd>
 <dt><tt>B</tt> : tensor(bfloat16)</dt>
-<dd>Constrain rotary cache tensors to bfloat16.</dd>
+<dd>Rotary and hidden tensors.</dd>
 </dl>
 
 
 ### <a name="com.microsoft.FoldTail"></a><a name="com.microsoft.foldtail">**com.microsoft.FoldTail**</a>
 
-  Updates the next speculative draft token and aggregate error bit from the head result. The eight-element control
-  tensor is returned as updated_ctl, with supported CUDA input-buffer reuse. When ctl[3] is zero, the draft token
-  is updated and nonzero status sets error bit 2; muted calls return done=0 and preserve the control state.
+  Updates the selected draft-control slot and publishes the result-ring sequence after validation. updated_ctl must alias ctl on CUDA. Muted calls do not access the ring. Invalid ring metadata sets error bit 16 and mutes execution. done is a dependency output initialized to zero.
 
 #### Version
 
 This version of the operator has been available since version 1 of the 'com.microsoft' operator set.
 
+#### Attributes
+
+<dl>
+<dt><tt>ctl_slot</tt> : int (default is 2)</dt>
+<dd>Draft-control slot: 2, 8 or 9.</dd>
+</dl>
+
 #### Inputs
 
 <dl>
 <dt><tt>next_token</tt> : I64</dt>
-<dd>Next draft token ID.</dd>
+<dd>Next draft token (or ignored selector in DFlash).</dd>
 <dt><tt>status</tt> : I32</dt>
-<dd>Head status values.</dd>
+<dd>Downstream validation status words.</dd>
 <dt><tt>ctl</tt> : I64</dt>
-<dd>Eight-element control state.</dd>
+<dd>Mutable control state.</dd>
+<dt><tt>ptrs</tt> : I64</dt>
+<dd>Result ring address and record capacity.</dd>
 </dl>
 
 #### Outputs
 
 <dl>
 <dt><tt>done</tt> : I32</dt>
-<dd>One when the control state was updated; zero for muted calls.</dd>
+<dd>Dependency output [1].</dd>
 <dt><tt>updated_ctl</tt> : I64</dt>
-<dd>Updated control state with the same shape as ctl.</dd>
+<dd>Updated control state, aliases ctl.</dd>
 </dl>
 
 #### Type Constraints
 
 <dl>
 <dt><tt>I64</tt> : tensor(int64)</dt>
-<dd>Constrain token and control tensors to int64.</dd>
+<dd>Token and control words.</dd>
 <dt><tt>I32</tt> : tensor(int32)</dt>
-<dd>Constrain status and output tensors to int32.</dd>
+<dd>Status and dependency output.</dd>
 </dl>
 
 

docs/OperatorKernels.md

@@ -1103,8 +1103,8 @@ The **OpSet Version** column uses the following notation:
 |EmbedLayerNormalization|*in* input_ids:**T1**<br> *in* segment_ids:**T1**<br> *in* word_embedding:**T**<br> *in* position_embedding:**T**<br> *in* segment_embedding:**T**<br> *in* gamma:**T**<br> *in* beta:**T**<br> *in* mask:**T1**<br> *in* position_ids:**T1**<br> *out* output:**T**<br> *out* mask_index:**T1**<br> *out* embedding_sum:**T**|1+|**T** = tensor(float), tensor(float16)|
 |EngramGate|*in* key:**T**<br> *in* query:**T**<br> *in* value:**T**<br> *in* key_norm_scale:**T**<br> *in* query_norm_scale:**T**<br> *in* conv_norm_scale:**T**<br> *out* output:**T**<br> *out* gated_value_normed:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)|
 |FastGelu|*in* X:**T**<br> *in* bias:**T**<br> *out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)|
-|FoldMid2|*in* next_tokens:**I64**<br> *in* input_ids:**I64**<br> *in* status:**I32**<br> *in* ctl:**I64**<br> *in* ctx_commit:**I64**<br> *in* ctx_a:**I64**<br> *in* ctx_b:**I64**<br> *in* ptrs:**I64**<br> *in* cos_cache:**B**<br> *in* sin_cache:**B**<br> *out* head_ids:**I64**<br> *out* head_c:**I32**<br> *out* slot_ids:**I64**<br> *out* seqlens:**I32**<br> *out* cos_out:**B**<br> *out* sin_out:**B**<br> *out* sel:**I64**<br> *out* updated_ctl:**I64**<br> *out* updated_ctx_commit:**I64**|1+|**B** = tensor(bfloat16)<br/> **I32** = tensor(int32)<br/> **I64** = tensor(int64)|
-|FoldTail|*in* next_token:**I64**<br> *in* status:**I32**<br> *in* ctl:**I64**<br> *out* done:**I32**<br> *out* updated_ctl:**I64**|1+|**I32** = tensor(int32)<br/> **I64** = tensor(int64)|
+|FoldMid2|*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* prefix2:**I64**<br> *in* ptrs:**I64**<br> *in* cos_cache:**B**<br> *in* sin_cache:**B**<br> *in* hidden:**B**<br> *out* head_ids:**I64**<br> *out* head_c:**I32**<br> *out* slot_ids:**I64**<br> *out* seqlens:**I32**<br> *out* cos_out:**B**<br> *out* sin_out:**B**<br> *out* sel:**I64**<br> *out* head_hidden:**B**<br> *out* updated_ctl:**I64**<br> *out* updated_ctx_commit:**I64**|1+|**B** = tensor(bfloat16)<br/> **I32** = tensor(int32)<br/> **I64** = tensor(int64)|
+|FoldTail|*in* next_token:**I64**<br> *in* status:**I32**<br> *in* ctl:**I64**<br> *in* ptrs:**I64**<br> *out* done:**I32**<br> *out* updated_ctl:**I64**|1+|**I32** = tensor(int32)<br/> **I64** = tensor(int64)|
 |FusedConv|*in* X:**T**<br> *in* W:**T**<br> *in* B:**T**<br> *in* Z:**T**<br> *out* Y:**T**|1+|**T** = tensor(float)|
 |FusedMatMul|*in* A:**T**<br> *in* B:**T**<br> *out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)|
 |GatedAdd|*in* X:**T**<br> *in* Y:**T**<br> *in* gate:**T**<br> *out* output:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)|

docs/contrib_ops/cuda/speculative_fold.md

@@ -0,0 +1,53 @@
+# CUDA speculative fold operators
+
+`com.microsoft::FoldMid2` and `FoldTail` schema version 1 implement batch-one
+speculative acceptance and result publication without a custom-op library.
+This replaces the earlier two-row schema and supports two through four head
+rows, variable rotary width, configurable cache capacity, optional hidden-state
+padding, and deferred publication after downstream validation.
+
+`FoldMid2` takes verifier token predictions, input token IDs, status words,
+control state, committed context, prefix contexts, result-ring metadata,
+rotary caches, and optional hidden rows. Its outputs are head token IDs,
+committed length, slot IDs, sequence length, gathered cosine/sine rows, accepted
+prefix length, optional padded hidden rows, updated control, and updated context.
+The last two outputs **alias** inputs 3 and 4. The caller must serialize access
+to those device buffers. For four head rows, control has at least 16 int64 words
+and prefix contexts contain four pairs of int64 words.
+
+`FoldTail` takes a token, status words, the mid-stage control output, and
+result-ring metadata. It returns a dependency token and updated control.
+The latter aliases input 2. `ctl_slot` selects word 2, 8, or 9. The tail writes
+the ring's sequence word only after a system fence, so host polling cannot
+observe a partial record.
+
+The ring metadata is `[device_accessible_address, positive_record_capacity]`.
+Each record has eight int64 words. Its address must be aligned and its allocation
+must hold the declared number of records; allocation ownership remains with
+the caller. Invalid length, rotary position, or ring metadata sets error bits
+4, 8, or 16, mutes control, and avoids unsafe ring access. Muted prefill uses
+zero positions/cache length and does not commit contexts or access ring metadata.
+Muted prefill may contain more verifier rows than `head_rows`; non-muted
+verification is limited to `head_rows`. The tail unions its downstream status
+with the verifier status stored by the mid stage. The kernels do not allocate
+or synchronize the stream and support CUDA graph replay.
+
+## Migrating legacy packages
+
+Use GenAI's `tools/qwen38_graph/convert_fold_to_contrib.py` on the target ONNX
+graph. It adds explicit state outputs and connects the tail to the mid control
+output, preserving external-weight paths. Remove `custom_ops_library` and
+`custom_ops_libraries` from model session options only after all graph operations
+are available in the rebuilt ORT stack. Keep the required GenAI data-plane
+library: it is a runtime helper, not a custom-op registration library.
+
+Build ORT with CUDA and contrib operators enabled. The plugin build registers
+the schemas in the core and the kernels in the CUDA provider. Rebuild both,
+not just the provider. Ordinary upstream wheels do not contain these additions.
+
+`FoldContribTest.*` in `onnxruntime_provider_test` covers all S1-S4 prefix
+lengths, hidden-row padding, deferred/wrapped publication, muted prefill, invalid
+metadata, verifier/downstream status union, tail slots, and captured-graph replay.
+`SpeculativeFoldTest.InferenceSessionLoadsVersionOneAndAcceptsMutedPrefillRows`
+loads and runs the version-1 schemas through `InferenceSession`, including
+muted prefill batches larger than `head_rows`.