diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md
index c04af0abd903f..4bd2bba789ce5 100644
--- a/docs/ContribOperators.md
+++ b/docs/ContribOperators.md
@@ -2104,117 +2104,129 @@ This version of the operator has been available since version 1 of the 'com.micr
### **com.microsoft.FoldMid2**
- 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
+
+
+- head_capacity : int (default is 33028)
+- Maximum head-cache token capacity.
+- head_rows : int (default is 2)
+- Fixed number of head rows, 2 through 4.
+
+
+#### Inputs (10 - 11)
- next_tokens : I64
-- Two predicted token IDs.
+- Verifier predictions, batch one.
- input_ids : I64
-- Two input token IDs; input_ids[1] is the draft token.
+- Verifier input IDs.
- status : I32
-- Status values aggregated into the result record.
+- Verifier status words.
- ctl : I64
-- Eight-element control state.
+- Mutable control state: at least 8 words for S2, 16 for S3/S4.
- ctx_commit : I64
-- Two-element committed context.
-- ctx_a : I64
-- Context selected when the draft is rejected.
-- ctx_b : I64
-- Context selected when the draft is accepted.
+- Mutable committed context, two words.
+- prefixes : I64
+- Two-word context for each accepted-prefix length; S2 needs only the first.
+- prefix2 : I64
+- Two-word context for acceptance length one.
- ptrs : I64
-- Device pointers containing the result-ring address and capacity.
+- Result ring address and positive record capacity.
- cos_cache : B
-- BFloat16 rotary cosine cache with shape [positions, 32].
+- Rotary cosine cache [positions, width].
- sin_cache : B
-- BFloat16 rotary sine cache with shape [positions, 32].
+- Rotary sine cache [positions, width].
+- hidden (optional) : B
+- Optional hidden states [1, verifier_rows, ...].
#### Outputs
- head_ids : I64
-- Two token IDs consumed by the head.
+- Head token IDs [1, head_rows].
- head_c : I32
-- Current head-cache length.
+- Head committed length [1].
- slot_ids : I64
-- Two consecutive head-cache slots.
+- Head slots [1, head_rows].
- seqlens : I32
-- Head sequence length.
+- Head sequence length [1].
- cos_out : B
-- Two rows of rotary cosine values.
+- Gathered rotary cosine [1, head_rows, width].
- sin_out : B
-- Two rows of rotary sine values.
+- Gathered rotary sine [1, head_rows, width].
- sel : I64
-- Selected row: one when accepted, zero when rejected, or -1 on invalid metadata.
+- Accepted-prefix length, or -1 on invalid metadata.
+- head_hidden (optional) : B
+- Hidden rows, repeating the last source row when necessary.
- updated_ctl : I64
-- Updated control state with the same shape as ctl.
+- Updated control state, aliases ctl.
- updated_ctx_commit : I64
-- Updated committed context with the same shape as ctx_commit.
+- Updated context, aliases ctx_commit.
#### Type Constraints
- I64 : tensor(int64)
-- Constrain control and token tensors to int64.
+- Token and control words.
- I32 : tensor(int32)
-- Constrain status and sequence tensors to int32.
+- Status and sequence lengths.
- B : tensor(bfloat16)
-- Constrain rotary cache tensors to bfloat16.
+- Rotary and hidden tensors.
### **com.microsoft.FoldTail**
- 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
+
+
+- ctl_slot : int (default is 2)
+- Draft-control slot: 2, 8 or 9.
+
+
#### Inputs
- next_token : I64
-- Next draft token ID.
+- Next draft token (or ignored selector in DFlash).
- status : I32
-- Head status values.
+- Downstream validation status words.
- ctl : I64
-- Eight-element control state.
+- Mutable control state.
+- ptrs : I64
+- Result ring address and record capacity.
#### Outputs
- done : I32
-- One when the control state was updated; zero for muted calls.
+- Dependency output [1].
- updated_ctl : I64
-- Updated control state with the same shape as ctl.
+- Updated control state, aliases ctl.
#### Type Constraints
- I64 : tensor(int64)
-- Constrain token and control tensors to int64.
+- Token and control words.
- I32 : tensor(int32)
-- Constrain status and output tensors to int32.
+- Status and dependency output.
diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md
index cb6d5e79d8cd8..4b9b7c8504cb7 100644
--- a/docs/OperatorKernels.md
+++ b/docs/OperatorKernels.md
@@ -1103,8 +1103,8 @@ The **OpSet Version** column uses the following notation:
|EmbedLayerNormalization|*in* input_ids:**T1**
*in* segment_ids:**T1**
*in* word_embedding:**T**
*in* position_embedding:**T**
*in* segment_embedding:**T**
*in* gamma:**T**
*in* beta:**T**
*in* mask:**T1**
*in* position_ids:**T1**
*out* output:**T**
*out* mask_index:**T1**
*out* embedding_sum:**T**|1+|**T** = tensor(float), tensor(float16)|
|EngramGate|*in* key:**T**
*in* query:**T**
*in* value:**T**
*in* key_norm_scale:**T**
*in* query_norm_scale:**T**
*in* conv_norm_scale:**T**
*out* output:**T**
*out* gated_value_normed:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)|
|FastGelu|*in* X:**T**
*in* bias:**T**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)|
-|FoldMid2|*in* next_tokens:**I64**
*in* input_ids:**I64**
*in* status:**I32**
*in* ctl:**I64**
*in* ctx_commit:**I64**
*in* ctx_a:**I64**
*in* ctx_b:**I64**
*in* ptrs:**I64**
*in* cos_cache:**B**
*in* sin_cache:**B**
*out* head_ids:**I64**
*out* head_c:**I32**
*out* slot_ids:**I64**
*out* seqlens:**I32**
*out* cos_out:**B**
*out* sin_out:**B**
*out* sel:**I64**
*out* updated_ctl:**I64**
*out* updated_ctx_commit:**I64**|1+|**B** = tensor(bfloat16)
**I32** = tensor(int32)
**I64** = tensor(int64)|
-|FoldTail|*in* next_token:**I64**
*in* status:**I32**
*in* ctl:**I64**
*out* done:**I32**
*out* updated_ctl:**I64**|1+|**I32** = tensor(int32)
**I64** = tensor(int64)|
+|FoldMid2|*in* next_tokens:**I64**
*in* input_ids:**I64**
*in* status:**I32**
*in* ctl:**I64**
*in* ctx_commit:**I64**
*in* prefixes:**I64**
*in* prefix2:**I64**
*in* ptrs:**I64**
*in* cos_cache:**B**
*in* sin_cache:**B**
*in* hidden:**B**
*out* head_ids:**I64**
*out* head_c:**I32**
*out* slot_ids:**I64**
*out* seqlens:**I32**
*out* cos_out:**B**
*out* sin_out:**B**
*out* sel:**I64**
*out* head_hidden:**B**
*out* updated_ctl:**I64**
*out* updated_ctx_commit:**I64**|1+|**B** = tensor(bfloat16)
**I32** = tensor(int32)
**I64** = tensor(int64)|
+|FoldTail|*in* next_token:**I64**
*in* status:**I32**
*in* ctl:**I64**
*in* ptrs:**I64**
*out* done:**I32**
*out* updated_ctl:**I64**|1+|**I32** = tensor(int32)
**I64** = tensor(int64)|
|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)|
diff --git a/docs/contrib_ops/cuda/speculative_fold.md b/docs/contrib_ops/cuda/speculative_fold.md
new file mode 100644
index 0000000000000..447e59edacef7
--- /dev/null
+++ b/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`.
diff --git a/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cc b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cc
index c5d4c6a06f893..d35eac055bbf0 100644
--- a/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cc
+++ b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cc
@@ -2,123 +2,129 @@
// Licensed under the MIT License.
#include "contrib_ops/cuda/llm/speculative_fold.h"
-
-#include
-
-#include "core/common/common.h"
#include "core/providers/cuda/cuda_kernel.h"
+#include "core/common/safeint.h"
+#include
-namespace onnxruntime {
-namespace contrib {
-namespace cuda {
-
+namespace onnxruntime::contrib::cuda {
using namespace onnxruntime::cuda;
-
namespace {
-bool HasElementCount(const Tensor* tensor, int64_t count) {
- return tensor != nullptr && tensor->Shape().Size() == count;
+bool Count(const Tensor* tensor, int64_t count) {
+ return tensor && tensor->Shape().Size() == count;
}
class FoldMid2 final : public CudaKernel {
public:
- explicit FoldMid2(const OpKernelInfo& info) : CudaKernel(info) {}
+ explicit FoldMid2(const OpKernelInfo& info) : CudaKernel(info) {
+ head_rows_ = SafeInt(info.GetAttrOrDefault("head_rows", 2));
+ capacity_ = SafeInt(info.GetAttrOrDefault("head_capacity", 33028));
+ ORT_ENFORCE(head_rows_ >= 2 && head_rows_ <= 4 && capacity_ >= head_rows_,
+ "FoldMid2: invalid head_rows/head_capacity");
+ }
Status ComputeInternal(OpKernelContext* context) const override {
- const Tensor* next_tokens = context->Input(0);
- const Tensor* input_ids = context->Input(1);
- const Tensor* status = context->Input(2);
- const Tensor* ctl = context->Input(3);
- const Tensor* ctx_commit = context->Input(4);
- const Tensor* ctx_a = context->Input(5);
- const Tensor* ctx_b = context->Input(6);
- const Tensor* ptrs = context->Input(7);
- const Tensor* cos_cache = context->Input(8);
- const Tensor* sin_cache = context->Input(9);
- ORT_ENFORCE(next_tokens != nullptr && input_ids != nullptr && status != nullptr &&
- ctl != nullptr && ctx_commit != nullptr && ctx_a != nullptr &&
- ctx_b != nullptr && ptrs != nullptr && cos_cache != nullptr &&
- sin_cache != nullptr);
-
- ORT_RETURN_IF_NOT(HasElementCount(next_tokens, 2) &&
- HasElementCount(input_ids, 2) &&
- HasElementCount(ctl, 8) &&
- HasElementCount(ctx_commit, 2) &&
- HasElementCount(ctx_a, 2) &&
- HasElementCount(ctx_b, 2) &&
- HasElementCount(ptrs, 2),
- "FoldMid2 received an invalid control tensor shape");
- ORT_RETURN_IF_NOT(cos_cache->Shape().NumDimensions() == 2 &&
- sin_cache->Shape().NumDimensions() == 2 &&
- cos_cache->Shape()[1] == 32 &&
- sin_cache->Shape()[1] == 32 &&
- cos_cache->Shape()[0] == sin_cache->Shape()[0],
- "FoldMid2 rotary caches must have matching shape [positions, 32]");
- const int64_t status_count = status->Shape().Size();
- ORT_RETURN_IF_NOT(status_count <= std::numeric_limits::max(),
- "FoldMid2 status tensor is too large");
-
- Tensor* head_ids = context->Output(0, TensorShape({1, 2}));
- Tensor* head_c = context->Output(1, TensorShape({1}));
- Tensor* slot_ids = context->Output(2, TensorShape({1, 2}));
- Tensor* seqlens = context->Output(3, TensorShape({1}));
- Tensor* cos_out = context->Output(4, TensorShape({1, 2, 32}));
- Tensor* sin_out = context->Output(5, TensorShape({1, 2, 32}));
- Tensor* sel = context->Output(6, TensorShape({1}));
- Tensor* updated_ctl = context->Output(7, ctl->Shape());
- Tensor* updated_ctx_commit = context->Output(8, ctx_commit->Shape());
- ORT_ENFORCE(head_ids != nullptr && head_c != nullptr && slot_ids != nullptr &&
- seqlens != nullptr && cos_out != nullptr && sin_out != nullptr &&
- sel != nullptr && updated_ctl != nullptr && updated_ctx_commit != nullptr);
-
- return speculative_fold::LaunchFoldMid2(
- next_tokens->Data(), input_ids->Data(),
- status->Data(), static_cast(status_count),
- ctl->Data(), ctx_commit->Data(),
- ctx_a->Data(), ctx_b->Data(), ptrs->Data(),
- reinterpret_cast(cos_cache->Data()),
- reinterpret_cast(sin_cache->Data()),
- head_ids->MutableData(), head_c->MutableData(),
- slot_ids->MutableData(), seqlens->MutableData(),
- reinterpret_cast<__nv_bfloat16*>(cos_out->MutableData()),
- reinterpret_cast<__nv_bfloat16*>(sin_out->MutableData()),
- sel->MutableData(), updated_ctl->MutableData(),
- updated_ctx_commit->MutableData(), cos_cache->Shape()[0], Stream(context));
+ 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* prefix2 = context->Input(6);
+ const auto* ptrs = context->Input(7);
+ const auto* cosine = context->Input(8);
+ const auto* sine = context->Input(9);
+ const auto* hidden = context->Input(10);
+ ORT_RETURN_IF_NOT(tokens && ids && status && ctl && committed && prefixes && prefix2 && ptrs && cosine && sine,
+ "FoldMid2: missing required input");
+ const int64_t rows = tokens->Shape().Size();
+ ORT_RETURN_IF_NOT(rows > 0 && rows <= INT_MAX && ids->Shape().Size() >= rows &&
+ status->Shape().Size() <= INT_MAX,
+ "FoldMid2: invalid token/status shape");
+ ORT_RETURN_IF_NOT(ctl->Shape().Size() >= (head_rows_ > 2 ? 16 : 8) &&
+ Count(committed, 2) && prefixes->Shape().Size() >= (head_rows_ > 2 ? 2 * head_rows_ : 2) &&
+ Count(prefix2, 2) && Count(ptrs, 2),
+ "FoldMid2: invalid control/context shape");
+ const auto& rope = cosine->Shape();
+ ORT_RETURN_IF_NOT(rope.NumDimensions() == 2 && rope[0] > 0 && rope[0] <= INT_MAX &&
+ rope[1] > 0 && rope[1] <= INT_MAX && sine->Shape() == rope,
+ "FoldMid2: rotary caches must have matching [positions, width] shapes");
+ Tensor* head_hidden = nullptr;
+ int64_t width = 0;
+ if (hidden) {
+ auto shape = hidden->Shape().AsShapeVector();
+ ORT_RETURN_IF_NOT(shape.size() >= 3 && shape[0] == 1 && shape[1] == rows,
+ "FoldMid2: hidden state must match token rows");
+ width = hidden->Shape().Size() / rows;
+ ORT_RETURN_IF_NOT(width > 0 && width <= INT64_MAX / head_rows_, "FoldMid2: invalid hidden width");
+ shape[1] = head_rows_;
+ head_hidden = context->Output(7, TensorShape(shape));
+ } else {
+ ORT_RETURN_IF_NOT(!context->Output(7, TensorShape({0})), "FoldMid2: hidden output requires hidden input");
+ }
+ auto* out_ids = context->Output(0, TensorShape({1, head_rows_}));
+ auto* head_c = context->Output(1, TensorShape({1}));
+ auto* slots = context->Output(2, TensorShape({1, head_rows_}));
+ auto* lengths = context->Output(3, TensorShape({1}));
+ auto* out_cos = context->Output(4, TensorShape({1, head_rows_, rope[1]}));
+ auto* out_sin = context->Output(5, TensorShape({1, head_rows_, rope[1]}));
+ auto* selection = context->Output(6, TensorShape({1}));
+ auto* updated_ctl = context->Output(8, ctl->Shape());
+ auto* updated_context = context->Output(9, committed->Shape());
+ ORT_ENFORCE(updated_ctl->DataRaw() == ctl->DataRaw() && updated_context->DataRaw() == committed->DataRaw(),
+ "FoldMid2 requires aliased control/context outputs");
+ return speculative_fold::LaunchMid(
+ tokens->Data(), ids->Data(), status->Data(), static_cast(status->Shape().Size()),
+ updated_ctl->MutableData(), updated_context->MutableData(), prefixes->Data(),
+ prefix2->Data(), ptrs->Data(),
+ reinterpret_cast(cosine->Data()),
+ reinterpret_cast(sine->Data()),
+ out_ids->MutableData(), head_c->MutableData(), slots->MutableData(),
+ lengths->MutableData(), reinterpret_cast(out_cos->MutableData()),
+ reinterpret_cast(out_sin->MutableData()), selection->MutableData(),
+ hidden ? reinterpret_cast(hidden->Data()) : nullptr,
+ head_hidden ? reinterpret_cast(head_hidden->MutableData()) : nullptr, width,
+ static_cast(rows), head_rows_, static_cast(rope[0]), static_cast(rope[1]), capacity_, Stream(context));
}
+
+ private:
+ int head_rows_;
+ int capacity_;
};
class FoldTail final : public CudaKernel {
public:
- explicit FoldTail(const OpKernelInfo& info) : CudaKernel(info) {}
+ explicit FoldTail(const OpKernelInfo& info) : CudaKernel(info) {
+ slot_ = SafeInt(info.GetAttrOrDefault("ctl_slot", 2));
+ ORT_ENFORCE(slot_ == 2 || slot_ == 8 || slot_ == 9, "FoldTail: ctl_slot must be 2, 8 or 9");
+ }
Status ComputeInternal(OpKernelContext* context) const override {
- const Tensor* next_token = context->Input(0);
- const Tensor* status = context->Input(1);
- const Tensor* ctl = context->Input(2);
- ORT_ENFORCE(next_token != nullptr && status != nullptr && ctl != nullptr);
- ORT_RETURN_IF_NOT(HasElementCount(next_token, 1) && HasElementCount(ctl, 8),
- "FoldTail received an invalid control tensor shape");
- const int64_t status_count = status->Shape().Size();
- ORT_RETURN_IF_NOT(status_count <= std::numeric_limits::max(),
- "FoldTail status tensor is too large");
- Tensor* done = context->Output(0, TensorShape({1}));
- Tensor* updated_ctl = context->Output(1, ctl->Shape());
- ORT_ENFORCE(done != nullptr && updated_ctl != nullptr);
-
- return speculative_fold::LaunchFoldTail(
- next_token->Data(), status->Data(),
- static_cast(status_count), ctl->Data(),
- done->MutableData(), updated_ctl->MutableData(), Stream(context));
+ const auto* token = context->Input(0);
+ const auto* status = context->Input(1);
+ const auto* ctl = context->Input(2);
+ const auto* ptrs = context->Input(3);
+ ORT_RETURN_IF_NOT(Count(token, 1) && status && status->Shape().Size() <= INT_MAX &&
+ ctl && ctl->Shape().Size() >= 8 && ctl->Shape().Size() > slot_ && Count(ptrs, 2),
+ "FoldTail: invalid token/status/control/pointer shape");
+ auto* done = context->Output(0, TensorShape({1}));
+ auto* updated_ctl = context->Output(1, ctl->Shape());
+ ORT_ENFORCE(updated_ctl->DataRaw() == ctl->DataRaw(), "FoldTail requires aliased control output");
+ return speculative_fold::LaunchTail(token->Data(), status->Data(),
+ static_cast(status->Shape().Size()), updated_ctl->MutableData(),
+ ptrs->Data(), done->MutableData(), slot_, Stream(context));
}
-};
+ private:
+ int slot_;
+};
} // namespace
ONNX_OPERATOR_KERNEL_EX(
FoldMid2, kMSDomain, 1, kCudaExecutionProvider,
(*KernelDefBuilder::Create())
- .MayInplace(3, 7)
- .MayInplace(4, 8)
+ .Alias(3, 8)
+ .Alias(4, 9)
.TypeConstraint("I64", DataTypeImpl::GetTensorType())
.TypeConstraint("I32", DataTypeImpl::GetTensorType())
.TypeConstraint("B", DataTypeImpl::GetTensorType()),
@@ -127,11 +133,9 @@ ONNX_OPERATOR_KERNEL_EX(
ONNX_OPERATOR_KERNEL_EX(
FoldTail, kMSDomain, 1, kCudaExecutionProvider,
(*KernelDefBuilder::Create())
- .MayInplace(2, 1)
+ .Alias(2, 1)
.TypeConstraint("I64", DataTypeImpl::GetTensorType())
.TypeConstraint("I32", DataTypeImpl::GetTensorType()),
FoldTail);
-} // namespace cuda
-} // namespace contrib
-} // namespace onnxruntime
+} // namespace onnxruntime::contrib::cuda
diff --git a/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cu b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cu
index 7a15d1bbb8dcc..8b90ebe9f255d 100644
--- a/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cu
+++ b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cu
@@ -2,189 +2,149 @@
// Licensed under the MIT License.
#include "contrib_ops/cuda/llm/speculative_fold.h"
-
+#include "core/common/common.h"
#include
-namespace onnxruntime {
-namespace contrib {
-namespace cuda {
-namespace speculative_fold {
+namespace onnxruntime::contrib::cuda::speculative_fold {
namespace {
-constexpr int64_t kHeadCapacity = 33028;
+__device__ bool ValidRing(const int64_t* ptrs, int64_t sequence) {
+ return ptrs[0] != 0 && ptrs[0] % alignof(int64_t) == 0 &&
+ ptrs[1] > 0 && ptrs[1] <= LLONG_MAX / (8 * static_cast(sizeof(int64_t))) &&
+ sequence >= 0 && sequence < LLONG_MAX;
+}
-__global__ void FoldMid2Kernel(
- const int64_t* next_tokens, const int64_t* input_ids,
- const int32_t* status, int status_count, const int64_t* input_ctl,
- const int64_t* input_ctx_commit, const int64_t* ctx_a, const int64_t* ctx_b,
- const int64_t* ptrs, const __nv_bfloat16* cos_cache,
- const __nv_bfloat16* sin_cache, int64_t* head_ids,
- int32_t* head_c, int64_t* slot_ids, int32_t* seqlens,
- __nv_bfloat16* cos_out, __nv_bfloat16* sin_out, int64_t* sel,
- int64_t* ctl, int64_t* ctx_commit, int64_t cache_positions) {
- __shared__ int64_t committed_length;
+__global__ void MidKernel(const int64_t* tokens, const int64_t* ids, const int32_t* status, int status_count,
+ int64_t* ctl, int64_t* context, const int64_t* prefixes, const int64_t* prefix2,
+ const int64_t* ptrs, const uint16_t* cosine, const uint16_t* sine,
+ int64_t* head_ids, int32_t* head_c, int64_t* slots, int32_t* lengths,
+ uint16_t* head_cos, uint16_t* head_sin, int64_t* selection,
+ int rows, int head_rows, int rope_rows, int rope_width, int capacity) {
+ __shared__ int64_t length;
__shared__ int64_t cache_length;
- __shared__ int64_t validation_bits;
+ __shared__ int64_t errors;
+ __shared__ bool muted;
if (threadIdx.x == 0) {
- for (int i = 0; i < 8; ++i) ctl[i] = input_ctl[i];
- for (int i = 0; i < 2; ++i) ctx_commit[i] = input_ctx_commit[i];
- committed_length = ctl[1];
- cache_length = ctl[7];
- validation_bits = 0;
- if (cache_positions < 3 || committed_length < 0 || committed_length > cache_positions - 3) {
- validation_bits |= 8;
- }
- if (cache_length < 0 || cache_length > kHeadCapacity - 2) {
- validation_bits |= 4;
- }
- if (ctl[3] == 0 &&
- (ptrs[0] == 0 || ptrs[0] % alignof(int64_t) != 0 ||
- ptrs[1] <= 0 || ptrs[1] > LLONG_MAX / (8 * static_cast(sizeof(int64_t))) ||
- ctl[4] < 0 || ctl[4] == LLONG_MAX)) {
- validation_bits |= 16;
- }
- if (validation_bits != 0) {
- ctl[5] |= validation_bits;
+ muted = ctl[3] != 0;
+ length = muted ? 0 : ctl[1];
+ cache_length = muted ? 0 : ctl[7];
+ errors = 0;
+ if (length < 0 || length > static_cast(rope_rows) - head_rows - 1) errors |= 8;
+ if (cache_length < 0 || cache_length > capacity - head_rows) errors |= 4;
+ if (!muted && !ValidRing(ptrs, ctl[4])) errors |= 16;
+ if (rows > head_rows && !muted) errors |= 4;
+ if (errors) {
+ ctl[5] |= errors;
ctl[3] = 1;
- head_ids[0] = head_ids[1] = 0;
- head_c[0] = 0;
- slot_ids[0] = slot_ids[1] = 0;
- seqlens[0] = 0;
- sel[0] = -1;
}
}
__syncthreads();
-
- if (threadIdx.x >= 32) {
- const int j = threadIdx.x - 32;
- if (j < 64) {
- if (validation_bits == 0) {
- const int64_t position = committed_length + 1 + (j >> 5);
- const int feature = j & 31;
- cos_out[j] = cos_cache[position * 32 + feature];
- sin_out[j] = sin_cache[position * 32 + feature];
- } else {
- cos_out[j] = sin_out[j] = __float2bfloat16(0.0f);
- }
- }
- return;
+ for (int64_t i = threadIdx.x; i < static_cast(head_rows) * rope_width; i += blockDim.x) {
+ const int64_t position = length + 1 + i / rope_width;
+ head_cos[i] = !errors ? cosine[position * rope_width + i % rope_width] : 0;
+ head_sin[i] = !errors ? sine[position * rope_width + i % rope_width] : 0;
}
- if (threadIdx.x != 0 || validation_bits != 0) {
+ if (threadIdx.x != 0) return;
+ if (errors) {
+ for (int row = 0; row < head_rows; ++row) head_ids[row] = slots[row] = 0;
+ head_c[0] = lengths[0] = 0;
+ selection[0] = -1;
return;
}
-
- const int64_t pred0 = next_tokens[0];
- const int64_t pred1 = next_tokens[1];
- const int64_t draft = input_ids[1];
- const bool accept = pred0 == draft;
- int64_t status_bits = 0;
- for (int i = 0; i < status_count; ++i) {
- status_bits |= static_cast(status[i]);
+ int accepted = 0;
+ if (!muted) {
+ while (accepted + 1 < rows && accepted + 1 < head_rows &&
+ tokens[accepted] == ids[accepted + 1]) ++accepted;
+ }
+ for (int row = 0; row < head_rows; ++row) {
+ head_ids[row] = row < accepted ? ids[row + 1] : tokens[row < rows ? row : rows - 1];
+ slots[row] = cache_length + row;
}
-
- head_ids[0] = accept ? draft : pred0;
- head_ids[1] = pred1;
head_c[0] = static_cast(cache_length);
- slot_ids[0] = cache_length;
- slot_ids[1] = cache_length + 1;
- seqlens[0] = static_cast(cache_length + 1);
- sel[0] = accept ? 1 : 0;
+ lengths[0] = static_cast(cache_length + head_rows - 1);
+ selection[0] = accepted;
+ if (muted) return;
+ int64_t validation = 0;
+ for (int i = 0; i < status_count; ++i) validation |= status[i];
+ const int64_t* prefix = accepted == 1 ? prefix2 : prefixes + 2 * accepted;
+ context[0] = prefix[0];
+ context[1] = prefix[1];
+ const int64_t pending = tokens[accepted];
+ ctl[0] = pending;
+ ctl[1] = length + accepted + 1;
+ ctl[7] = cache_length + accepted + 1;
+ ++ctl[4];
+ ctl[5] |= validation ? 2 : 0;
+ ctl[6] = accepted;
+ volatile int64_t* record = reinterpret_cast(ptrs[0]) + 8 * ((ctl[4] - 1) % ptrs[1]);
+ record[1] = tokens[0];
+ record[2] = rows > 1 ? tokens[1] : tokens[0];
+ record[3] = accepted;
+ record[4] = validation;
+ record[5] = pending;
+ record[6] = ctl[1];
+ record[7] = rows > 2 ? static_cast(
+ static_cast(static_cast(tokens[2])) |
+ (rows > 3 ? static_cast(static_cast(tokens[3])) << 32 : 0))
+ : 0;
+ // The tail publishes record[0] only after all downstream validation has completed.
+}
- if (ctl[3] != 0) {
- return;
+__global__ void HiddenKernel(const uint16_t* source, uint16_t* output, int64_t width, int rows, int head_rows) {
+ for (int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x;
+ i < width * head_rows; i += static_cast(gridDim.x) * blockDim.x) {
+ const int64_t row = i / width;
+ output[i] = source[(row < rows ? row : rows - 1) * width + i % width];
}
-
- const int64_t step = accept ? 2 : 1;
- const int64_t* committed_context = accept ? ctx_b : ctx_a;
- ctx_commit[0] = committed_context[0];
- ctx_commit[1] = committed_context[1];
-
- int64_t* result_ring = reinterpret_cast(ptrs[0]);
- const int64_t ring_size = ptrs[1];
- volatile int64_t* record = result_ring + 8 * (ctl[4] % ring_size);
- ctl[0] = accept ? pred1 : pred0;
- ctl[1] = committed_length + step;
- ctl[7] = cache_length + step;
- ctl[4] += 1;
- ctl[5] |= status_bits != 0 ? 2 : 0;
- ctl[6] = accept ? 1 : 0;
- record[1] = pred0;
- record[2] = pred1;
- record[3] = accept ? 1 : 0;
- record[4] = status_bits;
- record[5] = accept ? pred1 : pred0;
- record[6] = committed_length + step;
- __threadfence_system();
- record[0] = ctl[4];
}
-__global__ void FoldTailKernel(const int64_t* next_token,
- const int32_t* status, int status_count,
- const int64_t* input_ctl, int32_t* done, int64_t* ctl) {
- if (threadIdx.x != 0) {
+__global__ void TailKernel(const int64_t* token, const int32_t* status, int count, int64_t* ctl,
+ const int64_t* ptrs, int32_t* done, int slot) {
+ if (threadIdx.x != 0) return;
+ done[0] = 0;
+ if (ctl[3] != 0) return;
+ if (ctl[4] <= 0 || !ValidRing(ptrs, ctl[4] - 1)) {
+ ctl[5] |= 16;
+ ctl[3] = 1;
return;
}
- for (int i = 0; i < 8; ++i) ctl[i] = input_ctl[i];
- if (ctl[3] != 0) {
- done[0] = 0;
- return;
- }
-
- int64_t status_bits = 0;
- for (int i = 0; i < status_count; ++i) {
- status_bits |= static_cast(status[i]);
- }
- ctl[2] = next_token[0];
- if (status_bits != 0) {
- ctl[5] |= 2;
- }
- done[0] = 1;
+ int32_t error = 0;
+ for (int i = 0; i < count; ++i) error |= status[i];
+ ctl[slot] = token[0];
+ if (error) ctl[5] |= 2;
+ volatile int64_t* record = reinterpret_cast(ptrs[0]) + 8 * ((ctl[4] - 1) % ptrs[1]);
+ record[4] = record[4] | error | ctl[5];
+ __threadfence_system();
+ record[0] = ctl[4];
}
Status LaunchStatus(cudaError_t error) {
- return error == cudaSuccess
- ? Status::OK()
- : Status(common::StatusCategory::SYSTEM,
- static_cast(error), cudaGetErrorString(error));
+ return error == cudaSuccess ? Status::OK()
+ : Status(common::StatusCategory::SYSTEM, static_cast(error), cudaGetErrorString(error));
}
-
} // namespace
-Status LaunchFoldMid2(const int64_t* next_tokens, const int64_t* input_ids,
- const int32_t* status, int status_count, const int64_t* ctl,
- const int64_t* ctx_commit, const int64_t* ctx_a, const int64_t* ctx_b,
- const int64_t* ptrs, const __nv_bfloat16* cos_cache,
- const __nv_bfloat16* sin_cache, int64_t* head_ids,
- int32_t* head_c, int64_t* slot_ids, int32_t* seqlens,
- __nv_bfloat16* cos_out, __nv_bfloat16* sin_out,
- int64_t* sel, int64_t* updated_ctl, int64_t* updated_ctx_commit,
- int64_t cache_positions, cudaStream_t stream) {
- if (status_count < 0) {
- return Status(common::StatusCategory::ONNXRUNTIME,
- common::StatusCode::INVALID_ARGUMENT,
- "FoldMid2 status count must not be negative");
+Status LaunchMid(const int64_t* tokens, const int64_t* ids, const int32_t* status, int status_count,
+ int64_t* ctl, int64_t* context, const int64_t* prefixes, const int64_t* prefix2,
+ const int64_t* ptrs, const uint16_t* cosine, const uint16_t* sine,
+ int64_t* head_ids, int32_t* head_c, int64_t* slots, int32_t* lengths,
+ uint16_t* head_cos, uint16_t* head_sin, int64_t* selection,
+ const uint16_t* hidden, uint16_t* head_hidden, int64_t hidden_width,
+ int rows, int head_rows, int rope_rows, int rope_width, int capacity, cudaStream_t stream) {
+ MidKernel<<<1, 256, 0, stream>>>(tokens, ids, status, status_count, ctl, context, prefixes, prefix2, ptrs,
+ cosine, sine, head_ids, head_c, slots, lengths, head_cos, head_sin, selection,
+ rows, head_rows, rope_rows, rope_width, capacity);
+ ORT_RETURN_IF_ERROR(LaunchStatus(cudaGetLastError()));
+ if (head_hidden) {
+ HiddenKernel<<<32, 256, 0, stream>>>(hidden, head_hidden, hidden_width, rows, head_rows);
+ ORT_RETURN_IF_ERROR(LaunchStatus(cudaGetLastError()));
}
- FoldMid2Kernel<<<1, 96, 0, stream>>>(
- next_tokens, input_ids, status, status_count, ctl, ctx_commit, ctx_a,
- ctx_b, ptrs, cos_cache, sin_cache, head_ids, head_c, slot_ids, seqlens,
- cos_out, sin_out, sel, updated_ctl, updated_ctx_commit, cache_positions);
- return LaunchStatus(cudaGetLastError());
+ return Status::OK();
}
-Status LaunchFoldTail(const int64_t* next_token, const int32_t* status,
- int status_count, const int64_t* ctl, int32_t* done, int64_t* updated_ctl,
- cudaStream_t stream) {
- if (status_count < 0) {
- return Status(common::StatusCategory::ONNXRUNTIME,
- common::StatusCode::INVALID_ARGUMENT,
- "FoldTail status count must not be negative");
- }
- FoldTailKernel<<<1, 32, 0, stream>>>(next_token, status, status_count, ctl,
- done, updated_ctl);
+Status LaunchTail(const int64_t* token, const int32_t* status, int status_count, int64_t* ctl,
+ const int64_t* ptrs, int32_t* done, int slot, cudaStream_t stream) {
+ TailKernel<<<1, 32, 0, stream>>>(token, status, status_count, ctl, ptrs, done, slot);
return LaunchStatus(cudaGetLastError());
}
-
-} // namespace speculative_fold
-} // namespace cuda
-} // namespace contrib
-} // namespace onnxruntime
+} // 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
index 6af434004d952..c21c0c98f9393 100644
--- a/onnxruntime/contrib_ops/cuda/llm/speculative_fold.h
+++ b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.h
@@ -3,33 +3,21 @@
#pragma once
-#include
#include
-
#include
-
#include "core/common/status.h"
-namespace onnxruntime {
-namespace contrib {
-namespace cuda {
-namespace speculative_fold {
+namespace onnxruntime::contrib::cuda::speculative_fold {
-Status LaunchFoldMid2(const int64_t* next_tokens, const int64_t* input_ids,
- const int32_t* status, int status_count, const int64_t* ctl,
- const int64_t* ctx_commit, const int64_t* ctx_a, const int64_t* ctx_b,
- const int64_t* ptrs, const __nv_bfloat16* cos_cache,
- const __nv_bfloat16* sin_cache, int64_t* head_ids,
- int32_t* head_c, int64_t* slot_ids, int32_t* seqlens,
- __nv_bfloat16* cos_out, __nv_bfloat16* sin_out,
- int64_t* sel, int64_t* updated_ctl, int64_t* updated_ctx_commit,
- int64_t cache_positions, cudaStream_t stream);
+Status LaunchMid(const int64_t* tokens, const int64_t* ids, const int32_t* status, int status_count,
+ int64_t* ctl, int64_t* context, const int64_t* prefixes, const int64_t* prefix2,
+ const int64_t* ptrs, const uint16_t* cosine, const uint16_t* sine,
+ int64_t* head_ids, int32_t* head_c, int64_t* slots, int32_t* lengths,
+ uint16_t* head_cos, uint16_t* head_sin, int64_t* selection,
+ const uint16_t* hidden, uint16_t* head_hidden, int64_t hidden_width,
+ int rows, int head_rows, int rope_rows, int rope_width, int capacity, cudaStream_t stream);
-Status LaunchFoldTail(const int64_t* next_token, const int32_t* status,
- int status_count, const int64_t* ctl, int32_t* done, int64_t* updated_ctl,
- cudaStream_t stream);
+Status LaunchTail(const int64_t* token, const int32_t* status, int status_count, int64_t* ctl,
+ const int64_t* ptrs, int32_t* done, int slot, cudaStream_t stream);
-} // namespace speculative_fold
-} // namespace cuda
-} // namespace contrib
-} // namespace onnxruntime
+} // 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 fecac2f17f523..c2b1e39a68546 100644
--- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc
+++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc
@@ -34,6 +34,100 @@ void MatMulShapeInference(
namespace onnxruntime {
namespace contrib {
+ONNX_MS_OPERATOR_SET_SCHEMA(
+ FoldMid2, 1,
+ OpSchema()
+ .SetDoc("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.")
+ .Attr("head_rows", "Fixed number of head rows, 2 through 4.", AttributeProto::INT, int64_t{2})
+ .Attr("head_capacity", "Maximum head-cache token capacity.", AttributeProto::INT, int64_t{33028})
+ .Input(0, "next_tokens", "Verifier predictions, batch one.", "I64")
+ .Input(1, "input_ids", "Verifier input IDs.", "I64")
+ .Input(2, "status", "Verifier status words.", "I32")
+ .Input(3, "ctl", "Mutable control state: at least 8 words for S2, 16 for S3/S4.", "I64")
+ .Input(4, "ctx_commit", "Mutable committed context, two words.", "I64")
+ .Input(5, "prefixes", "Two-word context for each accepted-prefix length; S2 needs only the first.", "I64")
+ .Input(6, "prefix2", "Two-word context for acceptance length one.", "I64")
+ .Input(7, "ptrs", "Result ring address and positive record capacity.", "I64")
+ .Input(8, "cos_cache", "Rotary cosine cache [positions, width].", "B")
+ .Input(9, "sin_cache", "Rotary sine cache [positions, width].", "B")
+ .Input(10, "hidden", "Optional hidden states [1, verifier_rows, ...].", "B", OpSchema::Optional)
+ .Output(0, "head_ids", "Head token IDs [1, head_rows].", "I64")
+ .Output(1, "head_c", "Head committed length [1].", "I32")
+ .Output(2, "slot_ids", "Head slots [1, head_rows].", "I64")
+ .Output(3, "seqlens", "Head sequence length [1].", "I32")
+ .Output(4, "cos_out", "Gathered rotary cosine [1, head_rows, width].", "B")
+ .Output(5, "sin_out", "Gathered rotary sine [1, head_rows, width].", "B")
+ .Output(6, "sel", "Accepted-prefix length, or -1 on invalid metadata.", "I64")
+ .Output(7, "head_hidden", "Hidden rows, repeating the last source row when necessary.", "B", OpSchema::Optional)
+ .Output(8, "updated_ctl", "Updated control state, aliases ctl.", "I64")
+ .Output(9, "updated_ctx_commit", "Updated context, aliases ctx_commit.", "I64")
+ .TypeConstraint("I64", {"tensor(int64)"}, "Token and control words.")
+ .TypeConstraint("I32", {"tensor(int32)"}, "Status and sequence lengths.")
+ .TypeConstraint("B", {"tensor(bfloat16)"}, "Rotary and hidden tensors.")
+ .TypeAndShapeInferenceFunction([](InferenceContext& ctx) {
+ for (int output : {0, 2, 6, 8, 9}) propagateElemTypeFromInputToOutput(ctx, 0, output);
+ for (int output : {1, 3}) propagateElemTypeFromInputToOutput(ctx, 2, output);
+ for (int output : {4, 5}) propagateElemTypeFromInputToOutput(ctx, 8, output);
+ if (hasInputShape(ctx, 10)) propagateElemTypeFromInputToOutput(ctx, 10, 7);
+ const auto rows = getAttribute(ctx, "head_rows", 2);
+ const auto capacity = getAttribute(ctx, "head_capacity", 33028);
+ if (rows < 2 || rows > 4 || capacity < rows || capacity > std::numeric_limits::max())
+ fail_shape_inference("FoldMid2: invalid head_rows/head_capacity");
+ for (int output : {0, 2}) {
+ auto* shape = getOutputShape(ctx, output);
+ shape->add_dim()->set_dim_value(1);
+ shape->add_dim()->set_dim_value(rows);
+ }
+ for (int output : {1, 3, 6}) getOutputShape(ctx, output)->add_dim()->set_dim_value(1);
+ if (hasInputShape(ctx, 8)) {
+ const auto& rope = getInputShape(ctx, 8);
+ if (rope.dim_size() != 2) fail_shape_inference("FoldMid2: rotary cache must have rank 2");
+ TensorShapeProto shape;
+ shape.add_dim()->set_dim_value(1);
+ shape.add_dim()->set_dim_value(rows);
+ *shape.add_dim() = rope.dim(1);
+ updateOutputShape(ctx, 4, shape);
+ updateOutputShape(ctx, 5, shape);
+ }
+ if (hasInputShape(ctx, 10)) {
+ auto shape = getInputShape(ctx, 10);
+ if (shape.dim_size() < 3) fail_shape_inference("FoldMid2: hidden must have rank at least 3");
+ shape.mutable_dim(1)->set_dim_value(rows);
+ updateOutputShape(ctx, 7, shape);
+ }
+ propagateShapeFromInputToOutput(ctx, 3, 8);
+ propagateShapeFromInputToOutput(ctx, 4, 9);
+ }));
+
+ONNX_MS_OPERATOR_SET_SCHEMA(
+ FoldTail, 1,
+ OpSchema()
+ .SetDoc("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.")
+ .Attr("ctl_slot", "Draft-control slot: 2, 8 or 9.", AttributeProto::INT, int64_t{2})
+ .Input(0, "next_token", "Next draft token (or ignored selector in DFlash).", "I64")
+ .Input(1, "status", "Downstream validation status words.", "I32")
+ .Input(2, "ctl", "Mutable control state.", "I64")
+ .Input(3, "ptrs", "Result ring address and record capacity.", "I64")
+ .Output(0, "done", "Dependency output [1].", "I32")
+ .Output(1, "updated_ctl", "Updated control state, aliases ctl.", "I64")
+ .TypeConstraint("I64", {"tensor(int64)"}, "Token and control words.")
+ .TypeConstraint("I32", {"tensor(int32)"}, "Status and dependency output.")
+ .TypeAndShapeInferenceFunction([](InferenceContext& ctx) {
+ propagateElemTypeFromInputToOutput(ctx, 1, 0);
+ propagateElemTypeFromInputToOutput(ctx, 2, 1);
+ const auto slot = getAttribute(ctx, "ctl_slot", 2);
+ if (slot != 2 && slot != 8 && slot != 9) fail_shape_inference("FoldTail: invalid ctl_slot");
+ getOutputShape(ctx, 0)->add_dim()->set_dim_value(1);
+ propagateShapeFromInputToOutput(ctx, 2, 1);
+ }));
+
void DecoderAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& ctx) {
// Type inference
ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 0, 0);
diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc
index cd42a5ba75bb9..9823c47403d35 100644
--- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc
+++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc
@@ -1500,73 +1500,6 @@ ONNX_MS_OPERATOR_SET_SCHEMA(
}
}));
-constexpr const char* FoldMid2_ver1_doc = R"DOC(
-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.
-)DOC";
-
-ONNX_MS_OPERATOR_SET_SCHEMA(
- FoldMid2, 1,
- OpSchema()
- .SetDoc(FoldMid2_ver1_doc)
- .Input(0, "next_tokens", "Two predicted token IDs.", "I64")
- .Input(1, "input_ids", "Two input token IDs; input_ids[1] is the draft token.", "I64")
- .Input(2, "status", "Status values aggregated into the result record.", "I32")
- .Input(3, "ctl", "Eight-element control state.", "I64")
- .Input(4, "ctx_commit", "Two-element committed context.", "I64")
- .Input(5, "ctx_a", "Context selected when the draft is rejected.", "I64")
- .Input(6, "ctx_b", "Context selected when the draft is accepted.", "I64")
- .Input(7, "ptrs", "Device pointers containing the result-ring address and capacity.", "I64")
- .Input(8, "cos_cache", "BFloat16 rotary cosine cache with shape [positions, 32].", "B")
- .Input(9, "sin_cache", "BFloat16 rotary sine cache with shape [positions, 32].", "B")
- .Output(0, "head_ids", "Two token IDs consumed by the head.", "I64")
- .Output(1, "head_c", "Current head-cache length.", "I32")
- .Output(2, "slot_ids", "Two consecutive head-cache slots.", "I64")
- .Output(3, "seqlens", "Head sequence length.", "I32")
- .Output(4, "cos_out", "Two rows of rotary cosine values.", "B")
- .Output(5, "sin_out", "Two rows of rotary sine values.", "B")
- .Output(6, "sel", "Selected row: one when accepted, zero when rejected, or -1 on invalid metadata.", "I64")
- .Output(7, "updated_ctl", "Updated control state with the same shape as ctl.", "I64")
- .Output(8, "updated_ctx_commit", "Updated committed context with the same shape as ctx_commit.", "I64")
- .TypeConstraint("I64", {"tensor(int64)"}, "Constrain control and token tensors to int64.")
- .TypeConstraint("I32", {"tensor(int32)"}, "Constrain status and sequence tensors to int32.")
- .TypeConstraint("B", {"tensor(bfloat16)"}, "Constrain rotary cache tensors to bfloat16.")
- .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) {
- ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 3, 7);
- ONNX_NAMESPACE::propagateShapeFromInputToOutput(ctx, 3, 7);
- ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 4, 8);
- ONNX_NAMESPACE::propagateShapeFromInputToOutput(ctx, 4, 8);
- }));
-
-constexpr const char* FoldTail_ver1_doc = R"DOC(
-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.
-)DOC";
-
-ONNX_MS_OPERATOR_SET_SCHEMA(
- FoldTail, 1,
- OpSchema()
- .SetDoc(FoldTail_ver1_doc)
- .Input(0, "next_token", "Next draft token ID.", "I64")
- .Input(1, "status", "Head status values.", "I32")
- .Input(2, "ctl", "Eight-element control state.", "I64")
- .Output(0, "done", "One when the control state was updated; zero for muted calls.", "I32")
- .Output(1, "updated_ctl", "Updated control state with the same shape as ctl.", "I64")
- .TypeConstraint("I64", {"tensor(int64)"}, "Constrain token and control tensors to int64.")
- .TypeConstraint("I32", {"tensor(int32)"}, "Constrain status and output tensors to int32.")
- .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) {
- ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 2, 1);
- ONNX_NAMESPACE::propagateShapeFromInputToOutput(ctx, 2, 1);
- }));
-
constexpr const char* qMoE_ver1_doc = R"DOC(
Quantized mixture of experts (MoE).
A 2D input is the packed token-major form used by continuous-batching engines: tokens from
diff --git a/onnxruntime/test/contrib_ops/cuda_kernels/speculative_fold_test.cc b/onnxruntime/test/contrib_ops/cuda_kernels/speculative_fold_test.cc
index 3d17acc140fc1..f19f629946634 100644
--- a/onnxruntime/test/contrib_ops/cuda_kernels/speculative_fold_test.cc
+++ b/onnxruntime/test/contrib_ops/cuda_kernels/speculative_fold_test.cc
@@ -1,329 +1,222 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
-#include
#include
#include
-
#include
-#include
#include
+#include
+#include
#include
#include
-
#include "contrib_ops/cuda/llm/speculative_fold.h"
-namespace onnxruntime {
-namespace test {
+namespace onnxruntime::test {
namespace {
+namespace fold = contrib::cuda::speculative_fold;
-void CheckCuda(cudaError_t error) {
- if (error != cudaSuccess) {
- throw std::runtime_error(cudaGetErrorString(error));
- }
+void Check(cudaError_t status) {
+ if (status != cudaSuccess) throw std::runtime_error(cudaGetErrorString(status));
}
-class DeviceBuffer {
- public:
- explicit DeviceBuffer(size_t size) {
- CheckCuda(cudaMalloc(&data_, size));
+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");
}
- ~DeviceBuffer() {
- if (data_ != nullptr) {
- cudaFree(data_);
- }
+ Buffer(const Buffer&) = delete;
+ Buffer& operator=(const Buffer&) = delete;
+ void Upload(const std::vector& values) {
+ ASSERT_EQ(values.size(), count);
+ Check(cudaMemcpy(data, values.data(), count * sizeof(T), cudaMemcpyHostToDevice));
}
- DeviceBuffer(const DeviceBuffer&) = delete;
- DeviceBuffer& operator=(const DeviceBuffer&) = delete;
-
- template
- T* As() const {
- return reinterpret_cast(data_);
+ std::vector Read() const {
+ std::vector values(count);
+ Check(cudaMemcpy(values.data(), data, count * sizeof(T), cudaMemcpyDeviceToHost));
+ return values;
}
- void Upload(const void* source, size_t size) {
- CheckCuda(cudaMemcpy(data_, source, size, cudaMemcpyHostToDevice));
- }
- void Download(void* destination, size_t size) const {
- CheckCuda(cudaMemcpy(destination, data_, size, cudaMemcpyDeviceToHost));
- }
-
- private:
- void* data_ = nullptr;
+ T* data{};
+ size_t count;
};
-class MappedRing {
- public:
- explicit MappedRing(size_t element_count) : element_count_(element_count) {
- CheckCuda(cudaHostAlloc(reinterpret_cast(&host_data_),
- element_count * sizeof(int64_t),
- cudaHostAllocMapped));
- CheckCuda(cudaHostGetDevicePointer(reinterpret_cast(&device_data_),
- host_data_, 0));
- std::fill(host_data_, host_data_ + element_count_, int64_t{0});
- }
- ~MappedRing() {
- if (host_data_ != nullptr) {
- cudaFreeHost(host_data_);
+struct Fixture {
+ Buffer tokens{8}, ids{8}, ctl{16}, context{2}, prefixes{8}, prefix2{2}, ptrs{2}, ring{16};
+ Buffer status{2}, head_c{1}, lengths{1}, done{1};
+ Buffer head_ids{4}, slots{4}, selection{1};
+ Buffer cosine{16 * 3}, sine{16 * 3}, head_cos{12}, head_sin{12}, hidden{24}, out_hidden{12};
+
+ void Reset(int rows, int accepted, bool muted = false, int64_t sequence = 0) {
+ std::vector predictions{10, 20, 30, 40, 50, 60, 70, 80};
+ std::vector input(8, -1);
+ input[0] = 99;
+ for (int i = 0; i < accepted; ++i) input[i + 1] = predictions[i];
+ tokens.Upload(predictions);
+ ids.Upload(input);
+ std::vector control(16, 0);
+ control[1] = 3;
+ control[3] = muted ? 1 : 0;
+ control[4] = sequence;
+ control[7] = 2;
+ ctl.Upload(control);
+ context.Upload({-1, -2});
+ prefixes.Upload({100, 101, 200, 201, 300, 301, 400, 401});
+ prefix2.Upload({200, 201});
+ ptrs.Upload({reinterpret_cast(ring.data), 2});
+ ring.Upload(std::vector(16, -7));
+ status.Upload({0, 0});
+ std::vector cos(48), sin(48), h(24);
+ for (size_t i = 0; i < cos.size(); ++i) {
+ cos[i] = static_cast(i + 1);
+ sin[i] = static_cast(i + 100);
}
+ for (size_t i = 0; i < h.size(); ++i) h[i] = static_cast(i + 200);
+ cosine.Upload(cos);
+ sine.Upload(sin);
+ hidden.Upload(h);
+ (void)rows;
}
- MappedRing(const MappedRing&) = delete;
- MappedRing& operator=(const MappedRing&) = delete;
- int64_t* DeviceData() const { return device_data_; }
- int64_t* HostData() const { return host_data_; }
+ Status Mid(int rows, int heads, cudaStream_t stream = nullptr) {
+ return fold::LaunchMid(tokens.data, ids.data, status.data, 2, ctl.data, context.data,
+ prefixes.data, prefix2.data, ptrs.data, cosine.data, sine.data,
+ head_ids.data, head_c.data, slots.data, lengths.data, head_cos.data, head_sin.data,
+ selection.data, hidden.data, out_hidden.data, 3, rows, heads, 16, 3, 16, stream);
+ }
+ Status Tail(int slot = 2, cudaStream_t stream = nullptr) {
+ return fold::LaunchTail(tokens.data, status.data, 2, ctl.data, ptrs.data, done.data, slot, stream);
+ }
+};
- private:
- size_t element_count_;
- int64_t* host_data_ = nullptr;
- int64_t* device_data_ = nullptr;
+class FoldContribTest : public testing::Test {
+ protected:
+ void SetUp() override {
+ int devices = 0;
+ if (cudaGetDeviceCount(&devices) != cudaSuccess || devices == 0) GTEST_SKIP() << "CUDA is required";
+ }
};
-std::vector<__nv_bfloat16> MakeCache(int positions, float base) {
- std::vector<__nv_bfloat16> cache(static_cast(positions) * 32);
- for (int position = 0; position < positions; ++position) {
- for (int feature = 0; feature < 32; ++feature) {
- cache[static_cast(position) * 32 + feature] =
- __float2bfloat16(base + static_cast(position * 32 + feature));
+TEST_F(FoldContribTest, AllPrefixLengthsAndDeferredRingPublication) {
+ Fixture f;
+ for (int rows : {1, 2, 3, 4}) {
+ const int heads = std::max(2, rows);
+ for (int accepted = 0; accepted < rows; ++accepted) {
+ SCOPED_TRACE(testing::Message() << rows << "/" << accepted);
+ f.Reset(rows, accepted, false, 3);
+ ASSERT_TRUE(f.Mid(rows, heads).IsOK());
+ Check(cudaDeviceSynchronize());
+ EXPECT_EQ(f.selection.Read()[0], accepted);
+ EXPECT_EQ(f.context.Read(), (std::vector{100 * (accepted + 1), 100 * (accepted + 1) + 1}));
+ const auto ctl = f.ctl.Read();
+ EXPECT_EQ(ctl[0], 10 * (accepted + 1));
+ EXPECT_EQ(ctl[1], 4 + accepted);
+ EXPECT_EQ(ctl[7], 3 + accepted);
+ EXPECT_EQ(ctl[4], 4);
+ EXPECT_EQ(ctl[6], accepted);
+ auto record = f.ring.Read();
+ EXPECT_EQ(record[8], -7);
+ EXPECT_EQ(record[11], accepted);
+ EXPECT_EQ(record[13], 10 * (accepted + 1));
+ EXPECT_EQ(record[14], 4 + accepted);
+ if (rows > 2) {
+ const uint64_t packed = 30 | (rows > 3 ? uint64_t{40} << 32 : 0);
+ EXPECT_EQ(static_cast(record[15]), packed);
+ }
+ ASSERT_TRUE(f.Tail().IsOK());
+ Check(cudaDeviceSynchronize());
+ EXPECT_EQ(f.ring.Read()[8], 4);
+ EXPECT_EQ(f.done.Read()[0], 0);
+ const auto cos = f.head_cos.Read();
+ EXPECT_EQ(cos[0], 13);
+ EXPECT_EQ(cos[3 * heads - 1], 12 + 3 * heads);
+ const auto hidden = f.out_hidden.Read();
+ for (int row = 0; row < heads; ++row)
+ for (int col = 0; col < 3; ++col)
+ EXPECT_EQ(hidden[row * 3 + col], 200 + std::min(row, rows - 1) * 3 + col);
}
}
- return cache;
}
-TEST(CUDA_EP_Unittest, FoldMid2AcceptCommitsStateAndPublishesRingRecord) {
- using onnxruntime::contrib::cuda::speculative_fold::LaunchFoldMid2;
-
- constexpr int64_t kRingSize = 4;
- const std::array next_tokens = {7, 8};
- const std::array input_ids = {0, 7};
- const std::array status = {0, 4};
- const std::array ctl = {0, 10, 99, 0, 2, 0, 0, 5};
- const std::array ctx_commit = {1, 2};
- const std::array ctx_a = {11, 12};
- const std::array ctx_b = {21, 22};
- MappedRing ring(kRingSize * 8);
- const std::array ptrs = {
- static_cast(reinterpret_cast(ring.DeviceData())),
- kRingSize};
- const auto cos_cache = MakeCache(16, 1000.0f);
- const auto sin_cache = MakeCache(16, 2000.0f);
-
- DeviceBuffer d_next_tokens(sizeof(next_tokens));
- DeviceBuffer d_input_ids(sizeof(input_ids));
- DeviceBuffer d_status(sizeof(status));
- DeviceBuffer d_ctl(sizeof(ctl));
- DeviceBuffer d_ctx_commit(sizeof(ctx_commit));
- DeviceBuffer d_ctx_a(sizeof(ctx_a));
- DeviceBuffer d_ctx_b(sizeof(ctx_b));
- DeviceBuffer d_ptrs(sizeof(ptrs));
- DeviceBuffer d_cos_cache(cos_cache.size() * sizeof(__nv_bfloat16));
- DeviceBuffer d_sin_cache(sin_cache.size() * sizeof(__nv_bfloat16));
- DeviceBuffer d_head_ids(2 * sizeof(int64_t));
- DeviceBuffer d_head_c(sizeof(int32_t));
- DeviceBuffer d_slot_ids(2 * sizeof(int64_t));
- DeviceBuffer d_seqlens(sizeof(int32_t));
- DeviceBuffer d_cos_out(2 * 32 * sizeof(__nv_bfloat16));
- DeviceBuffer d_sin_out(2 * 32 * sizeof(__nv_bfloat16));
- DeviceBuffer d_sel(sizeof(int64_t));
-
- d_next_tokens.Upload(next_tokens.data(), sizeof(next_tokens));
- d_input_ids.Upload(input_ids.data(), sizeof(input_ids));
- d_status.Upload(status.data(), sizeof(status));
- d_ctl.Upload(ctl.data(), sizeof(ctl));
- d_ctx_commit.Upload(ctx_commit.data(), sizeof(ctx_commit));
- d_ctx_a.Upload(ctx_a.data(), sizeof(ctx_a));
- d_ctx_b.Upload(ctx_b.data(), sizeof(ctx_b));
- d_ptrs.Upload(ptrs.data(), sizeof(ptrs));
- d_cos_cache.Upload(cos_cache.data(), cos_cache.size() * sizeof(__nv_bfloat16));
- d_sin_cache.Upload(sin_cache.data(), sin_cache.size() * sizeof(__nv_bfloat16));
-
- const Status launch_status = LaunchFoldMid2(
- d_next_tokens.As(), d_input_ids.As(),
- d_status.As(), static_cast(status.size()),
- d_ctl.As(), d_ctx_commit.As(),
- d_ctx_a.As(), d_ctx_b.As(), d_ptrs.As(),
- d_cos_cache.As<__nv_bfloat16>(), d_sin_cache.As<__nv_bfloat16>(),
- d_head_ids.As(), d_head_c.As(),
- d_slot_ids.As(), d_seqlens.As(),
- d_cos_out.As<__nv_bfloat16>(), d_sin_out.As<__nv_bfloat16>(),
- d_sel.As(), d_ctl.As(), d_ctx_commit.As(), 16, nullptr);
- ASSERT_TRUE(launch_status.IsOK()) << launch_status.ErrorMessage();
- CheckCuda(cudaStreamSynchronize(nullptr));
-
- std::array actual_ctl{};
- std::array actual_ctx_commit{};
- std::array actual_head_ids{};
- std::array actual_slot_ids{};
- std::array actual_sel{};
- std::array actual_head_c{};
- std::array actual_seqlens{};
- std::array<__nv_bfloat16, 64> actual_cos{};
- std::array<__nv_bfloat16, 64> actual_sin{};
- d_ctl.Download(actual_ctl.data(), sizeof(actual_ctl));
- d_ctx_commit.Download(actual_ctx_commit.data(), sizeof(actual_ctx_commit));
- d_head_ids.Download(actual_head_ids.data(), sizeof(actual_head_ids));
- d_head_c.Download(actual_head_c.data(), sizeof(actual_head_c));
- d_slot_ids.Download(actual_slot_ids.data(), sizeof(actual_slot_ids));
- d_seqlens.Download(actual_seqlens.data(), sizeof(actual_seqlens));
- d_cos_out.Download(actual_cos.data(), sizeof(actual_cos));
- d_sin_out.Download(actual_sin.data(), sizeof(actual_sin));
- d_sel.Download(actual_sel.data(), sizeof(actual_sel));
-
- EXPECT_EQ(actual_head_ids, (std::array{7, 8}));
- EXPECT_EQ(actual_head_c[0], 5);
- EXPECT_EQ(actual_slot_ids, (std::array{5, 6}));
- EXPECT_EQ(actual_seqlens[0], 6);
- EXPECT_EQ(actual_sel[0], 1);
- EXPECT_EQ(actual_ctx_commit, (std::array{21, 22}));
- EXPECT_EQ(actual_ctl[0], 8);
- EXPECT_EQ(actual_ctl[1], 12);
- EXPECT_EQ(actual_ctl[4], 3);
- EXPECT_EQ(actual_ctl[5], 2);
- EXPECT_EQ(actual_ctl[6], 1);
- EXPECT_EQ(actual_ctl[7], 7);
+TEST_F(FoldContribTest, MutedPrefillDoesNotAccessRingOrCommitContext) {
+ Fixture f;
+ f.Reset(8, 0, true);
+ f.ptrs.Upload({0, 0});
+ ASSERT_TRUE(f.Mid(8, 4).IsOK());
+ ASSERT_TRUE(f.Tail().IsOK());
+ Check(cudaDeviceSynchronize());
+ EXPECT_EQ(f.context.Read(), (std::vector{-1, -2}));
+ EXPECT_EQ(f.ctl.Read()[1], 3);
+ EXPECT_EQ(f.ctl.Read()[4], 0);
+ EXPECT_EQ(f.head_c.Read()[0], 0);
+ EXPECT_EQ(f.head_cos.Read()[0], 4);
+ EXPECT_EQ(f.selection.Read()[0], 0);
+}
- for (int row = 0; row < 2; ++row) {
- for (int feature = 0; feature < 32; ++feature) {
- const size_t index = static_cast(row) * 32 + feature;
- EXPECT_EQ(__bfloat162float(actual_cos[index]),
- __bfloat162float(cos_cache[static_cast(11 + row) * 32 + feature]));
- EXPECT_EQ(__bfloat162float(actual_sin[index]),
- __bfloat162float(sin_cache[static_cast(11 + row) * 32 + feature]));
- }
+TEST_F(FoldContribTest, InvalidDeviceMetadataMutesBeforeRingAccess) {
+ Fixture f;
+ for (int kind = 0; kind < 4; ++kind) {
+ f.Reset(4, 3);
+ auto ctl = f.ctl.Read();
+ const int bit = kind < 2 ? 16 : (kind == 2 ? 8 : 4);
+ if (kind == 0) f.ptrs.Upload({0, 0});
+ if (kind == 1) ctl[4] = INT64_MAX;
+ if (kind == 2) ctl[1] = INT64_MAX;
+ if (kind == 3) ctl[7] = -1;
+ f.ctl.Upload(ctl);
+ ASSERT_TRUE(f.Mid(4, 4).IsOK());
+ ASSERT_TRUE(f.Tail().IsOK());
+ Check(cudaDeviceSynchronize());
+ EXPECT_EQ(f.selection.Read()[0], -1);
+ EXPECT_EQ(f.ctl.Read()[5] & bit, bit);
+ EXPECT_EQ(f.ctl.Read()[3], 1);
+ EXPECT_EQ(f.context.Read(), (std::vector{-1, -2}));
+ EXPECT_EQ(f.ring.Read(), std::vector(16, -7));
}
-
- const int64_t* record = ring.HostData() + 8 * 2;
- EXPECT_EQ(record[0], 3);
- EXPECT_EQ(record[1], 7);
- EXPECT_EQ(record[2], 8);
- EXPECT_EQ(record[3], 1);
- EXPECT_EQ(record[4], 4);
- EXPECT_EQ(record[5], 8);
- EXPECT_EQ(record[6], 12);
}
-TEST(CUDA_EP_Unittest, FoldMid2MutedRejectDoesNotCommitState) {
- using onnxruntime::contrib::cuda::speculative_fold::LaunchFoldMid2;
-
- const std::array next_tokens = {5, 6};
- const std::array input_ids = {0, 7};
- const std::array status = {0};
- const std::array ctl = {0, 4, 9, 1, 1, 0, 0, 2};
- const std::array ctx_commit = {31, 32};
- const std::array ctx_a = {41, 42};
- const std::array ctx_b = {51, 52};
- MappedRing ring(8 * 8);
- std::fill(ring.HostData(), ring.HostData() + 8 * 8, int64_t{-1});
- const std::array ptrs = {
- static_cast(reinterpret_cast(ring.DeviceData())), 8};
- const auto cos_cache = MakeCache(16, 1000.0f);
- const auto sin_cache = MakeCache(16, 2000.0f);
-
- DeviceBuffer d_next_tokens(sizeof(next_tokens));
- DeviceBuffer d_input_ids(sizeof(input_ids));
- DeviceBuffer d_status(sizeof(status));
- DeviceBuffer d_ctl(sizeof(ctl));
- DeviceBuffer d_ctx_commit(sizeof(ctx_commit));
- DeviceBuffer d_ctx_a(sizeof(ctx_a));
- DeviceBuffer d_ctx_b(sizeof(ctx_b));
- DeviceBuffer d_ptrs(sizeof(ptrs));
- DeviceBuffer d_cos_cache(cos_cache.size() * sizeof(__nv_bfloat16));
- DeviceBuffer d_sin_cache(sin_cache.size() * sizeof(__nv_bfloat16));
- DeviceBuffer d_head_ids(2 * sizeof(int64_t));
- DeviceBuffer d_head_c(sizeof(int32_t));
- DeviceBuffer d_slot_ids(2 * sizeof(int64_t));
- DeviceBuffer d_seqlens(sizeof(int32_t));
- DeviceBuffer d_cos_out(2 * 32 * sizeof(__nv_bfloat16));
- DeviceBuffer d_sin_out(2 * 32 * sizeof(__nv_bfloat16));
- DeviceBuffer d_sel(sizeof(int64_t));
-
- d_next_tokens.Upload(next_tokens.data(), sizeof(next_tokens));
- d_input_ids.Upload(input_ids.data(), sizeof(input_ids));
- d_status.Upload(status.data(), sizeof(status));
- d_ctl.Upload(ctl.data(), sizeof(ctl));
- d_ctx_commit.Upload(ctx_commit.data(), sizeof(ctx_commit));
- d_ctx_a.Upload(ctx_a.data(), sizeof(ctx_a));
- d_ctx_b.Upload(ctx_b.data(), sizeof(ctx_b));
- d_ptrs.Upload(ptrs.data(), sizeof(ptrs));
- d_cos_cache.Upload(cos_cache.data(), cos_cache.size() * sizeof(__nv_bfloat16));
- d_sin_cache.Upload(sin_cache.data(), sin_cache.size() * sizeof(__nv_bfloat16));
-
- const Status launch_status = LaunchFoldMid2(
- d_next_tokens.As(), d_input_ids.As(),
- d_status.As(), static_cast(status.size()),
- d_ctl.As(), d_ctx_commit.As(),
- d_ctx_a.As(), d_ctx_b.As(), d_ptrs.As(),
- d_cos_cache.As<__nv_bfloat16>(), d_sin_cache.As<__nv_bfloat16>(),
- d_head_ids.As(), d_head_c.As(),
- d_slot_ids.As(), d_seqlens.As(),
- d_cos_out.As<__nv_bfloat16>(), d_sin_out.As<__nv_bfloat16>(),
- d_sel.As(), d_ctl.As(), d_ctx_commit.As(), 16, nullptr);
- ASSERT_TRUE(launch_status.IsOK()) << launch_status.ErrorMessage();
- CheckCuda(cudaStreamSynchronize(nullptr));
-
- std::array actual_ctl{};
- std::array actual_ctx_commit{};
- std::array actual_head_ids{};
- std::array actual_slot_ids{};
- std::array actual_sel{};
- d_ctl.Download(actual_ctl.data(), sizeof(actual_ctl));
- d_ctx_commit.Download(actual_ctx_commit.data(), sizeof(actual_ctx_commit));
- d_head_ids.Download(actual_head_ids.data(), sizeof(actual_head_ids));
- d_slot_ids.Download(actual_slot_ids.data(), sizeof(actual_slot_ids));
- d_sel.Download(actual_sel.data(), sizeof(actual_sel));
-
- EXPECT_EQ(actual_ctl, ctl);
- EXPECT_EQ(actual_ctx_commit, ctx_commit);
- EXPECT_EQ(actual_head_ids, (std::array{5, 6}));
- EXPECT_EQ(actual_slot_ids, (std::array{2, 3}));
- EXPECT_EQ(actual_sel[0], 0);
- EXPECT_TRUE(std::all_of(ring.HostData(), ring.HostData() + 8 * 8,
- [](int64_t value) { return value == -1; }));
+TEST_F(FoldContribTest, TailSlotsAndValidationArePublished) {
+ Fixture f;
+ for (int slot : {2, 8, 9}) {
+ for (const auto& statuses : {std::pair{std::vector{0, 0}, std::vector{4, 8}},
+ std::pair{std::vector{1, 4}, std::vector{0, 8}}}) {
+ f.Reset(4, 2);
+ f.status.Upload(statuses.first);
+ ASSERT_TRUE(f.Mid(4, 4).IsOK());
+ Check(cudaDeviceSynchronize());
+ f.status.Upload(statuses.second);
+ ASSERT_TRUE(f.Tail(slot).IsOK());
+ Check(cudaDeviceSynchronize());
+ const int32_t combined = statuses.first[0] | statuses.first[1] | statuses.second[0] | statuses.second[1];
+ EXPECT_EQ(f.ctl.Read()[slot], 10);
+ EXPECT_EQ(f.ctl.Read()[5], combined ? 2 : 0);
+ EXPECT_EQ(f.ring.Read()[4], combined | (combined ? 2 : 0));
+ EXPECT_EQ(f.ring.Read()[0], 1);
+ }
+ }
}
-TEST(CUDA_EP_Unittest, FoldTailUpdatesDraftAndHonorsMute) {
- using onnxruntime::contrib::cuda::speculative_fold::LaunchFoldTail;
-
- const int64_t next_token = 42;
- const std::array status = {0, 4};
- std::array ctl = {0, 10, 9, 0, 3, 1, 0, 5};
- DeviceBuffer d_next_token(sizeof(next_token));
- DeviceBuffer d_status(sizeof(status));
- DeviceBuffer d_ctl(sizeof(ctl));
- DeviceBuffer d_done(sizeof(int32_t));
- d_next_token.Upload(&next_token, sizeof(next_token));
- d_status.Upload(status.data(), sizeof(status));
- d_ctl.Upload(ctl.data(), sizeof(ctl));
-
- Status launch_status = LaunchFoldTail(
- d_next_token.As(), d_status.As(),
- static_cast(status.size()), d_ctl.As(),
- d_done.As