From 6e528e7d0bbed6c8a4d06d2afcf3a79ae498f00a Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Thu, 22 Oct 2026 22:12:16 +0000 Subject: [PATCH 1/2] Add native speculative fold v2 for multi-row DFlash verification Preserve FoldMid2/FoldTail v1, add explicit aliased state outputs and deferred ring publication, and cover S1-S4 acceptance and captured replay. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- cmake/onnxruntime_unittests.cmake | 2 + docs/contrib_ops/cuda/speculative_fold.md | 47 ++++ .../contrib_ops/cuda/cuda_contrib_kernels.cc | 4 + .../cuda/llm/speculative_fold_v2.cc | 141 ++++++++++++ .../cuda/llm/speculative_fold_v2.cu | 150 ++++++++++++ .../cuda/llm/speculative_fold_v2.h | 23 ++ .../core/graph/contrib_ops/bert_defs.cc | 94 ++++++++ onnxruntime/core/graph/contrib_ops/ms_opset.h | 4 + .../cuda_kernels/speculative_fold_v2_test.cc | 215 ++++++++++++++++++ 9 files changed, 680 insertions(+) create mode 100644 docs/contrib_ops/cuda/speculative_fold.md create mode 100644 onnxruntime/contrib_ops/cuda/llm/speculative_fold_v2.cc create mode 100644 onnxruntime/contrib_ops/cuda/llm/speculative_fold_v2.cu create mode 100644 onnxruntime/contrib_ops/cuda/llm/speculative_fold_v2.h create mode 100644 onnxruntime/test/contrib_ops/cuda_kernels/speculative_fold_v2_test.cc diff --git a/cmake/onnxruntime_unittests.cmake b/cmake/onnxruntime_unittests.cmake index 7c5d8de4dd80d..75bb7b3c7209b 100644 --- a/cmake/onnxruntime_unittests.cmake +++ b/cmake/onnxruntime_unittests.cmake @@ -1127,6 +1127,7 @@ if (onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS AND onnxruntime_BUILD_CUDA_EP_AS_P "${TEST_SRC_DIR}/contrib_ops/cuda_kernels/softmax_topk_kernel_test.cc" "${TEST_SRC_DIR}/contrib_ops/cuda_kernels/mtp_moe_test.cc" "${TEST_SRC_DIR}/contrib_ops/cuda_kernels/speculative_fold_test.cc" + "${TEST_SRC_DIR}/contrib_ops/cuda_kernels/speculative_fold_v2_test.cc" ) endif() @@ -1162,6 +1163,7 @@ if (onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS AND onnxruntime_BUILD_CUDA_EP_AS_P "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/llm/mtp_moe.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/llm/speculative_fold_v2.cu" "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/moe/qmoe_kernels.cu" ) if (onnxruntime_USE_FP4_QMOE) diff --git a/docs/contrib_ops/cuda/speculative_fold.md b/docs/contrib_ops/cuda/speculative_fold.md new file mode 100644 index 0000000000000..bc893e87b39f6 --- /dev/null +++ b/docs/contrib_ops/cuda/speculative_fold.md @@ -0,0 +1,47 @@ +# CUDA speculative fold operators (version 2) + +`com.microsoft::FoldMid2` and `FoldTail` version 2 implement batch-one +speculative acceptance and result publication without a custom-op library. +They extend the two-row version-1 design for 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. This differs from version 1, which published in mid. + +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. +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 +both schemas in the core and both kernels in the CUDA provider. Rebuild both, +not just the provider. Ordinary upstream wheels do not contain these additions. + +`FoldContribV2Test.*` in `onnxruntime_provider_test` covers all S1-S4 prefix +lengths, hidden-row padding, deferred/wrapped publication, muted prefill, invalid +metadata, status aggregation, tail slots, and captured-graph replay. diff --git a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc index 93b62d9fff569..9e52891f8bdbc 100644 --- a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc @@ -103,6 +103,8 @@ class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, MoE); class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, MtpMoe); class CUDA_MS_OP_CLASS_NAME(1, FoldMid2); class CUDA_MS_OP_CLASS_NAME(1, FoldTail); +class CUDA_MS_OP_CLASS_NAME(2, FoldMid2); +class CUDA_MS_OP_CLASS_NAME(2, FoldTail); 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); @@ -470,6 +472,8 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/onnxruntime/contrib_ops/cuda/llm/speculative_fold_v2.cc b/onnxruntime/contrib_ops/cuda/llm/speculative_fold_v2.cc new file mode 100644 index 0000000000000..1a501a1155802 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/llm/speculative_fold_v2.cc @@ -0,0 +1,141 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "contrib_ops/cuda/llm/speculative_fold_v2.h" +#include "core/providers/cuda/cuda_kernel.h" +#include "core/common/safeint.h" +#include + +namespace onnxruntime::contrib::cuda { +using namespace onnxruntime::cuda; +namespace { + +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) { + 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 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 && + (rows <= head_rows_ || rows > 4) && 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_v2::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) { + 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 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_v2::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, 2, kCudaExecutionProvider, + (*KernelDefBuilder::Create()) + .Alias(3, 8) + .Alias(4, 9) + .TypeConstraint("I64", DataTypeImpl::GetTensorType()) + .TypeConstraint("I32", DataTypeImpl::GetTensorType()) + .TypeConstraint("B", DataTypeImpl::GetTensorType()), + FoldMid2); + +ONNX_OPERATOR_KERNEL_EX( + FoldTail, kMSDomain, 2, kCudaExecutionProvider, + (*KernelDefBuilder::Create()) + .Alias(2, 1) + .TypeConstraint("I64", DataTypeImpl::GetTensorType()) + .TypeConstraint("I32", DataTypeImpl::GetTensorType()), + FoldTail); + +} // namespace onnxruntime::contrib::cuda diff --git a/onnxruntime/contrib_ops/cuda/llm/speculative_fold_v2.cu b/onnxruntime/contrib_ops/cuda/llm/speculative_fold_v2.cu new file mode 100644 index 0000000000000..2995d9b1841d9 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/llm/speculative_fold_v2.cu @@ -0,0 +1,150 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "contrib_ops/cuda/llm/speculative_fold_v2.h" +#include "core/common/common.h" +#include + +namespace onnxruntime::contrib::cuda::speculative_fold_v2 { +namespace { + +__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 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 errors; + __shared__ bool muted; + if (threadIdx.x == 0) { + 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 > 4 && !muted) errors |= 4; + if (errors) { + ctl[5] |= errors; + ctl[3] = 1; + } + } + __syncthreads(); + 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) 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; + } + 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_c[0] = static_cast(cache_length); + 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. +} + +__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]; + } +} + +__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; + } + 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)); +} +} // namespace + +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())); + } + return Status::OK(); +} + +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 onnxruntime::contrib::cuda::speculative_fold_v2 diff --git a/onnxruntime/contrib_ops/cuda/llm/speculative_fold_v2.h b/onnxruntime/contrib_ops/cuda/llm/speculative_fold_v2.h new file mode 100644 index 0000000000000..22c9374276ffb --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/llm/speculative_fold_v2.h @@ -0,0 +1,23 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include +#include "core/common/status.h" + +namespace onnxruntime::contrib::cuda::speculative_fold_v2 { + +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 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 onnxruntime::contrib::cuda::speculative_fold_v2 diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index fecac2f17f523..e0d7b9c428c4e 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, 2, + 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, 2, + 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/ms_opset.h b/onnxruntime/core/graph/contrib_ops/ms_opset.h index 71b3c719f9906..c9ef934a3edf8 100644 --- a/onnxruntime/core/graph/contrib_ops/ms_opset.h +++ b/onnxruntime/core/graph/contrib_ops/ms_opset.h @@ -87,6 +87,8 @@ class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, MoE); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, MtpMoe); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, FoldMid2); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, FoldTail); +class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 2, FoldMid2); +class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 2, FoldTail); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, QMoE); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, MultiHeadAttention); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, GroupQueryAttention); @@ -224,6 +226,8 @@ class OpSet_Microsoft_ver1 { fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); + fn(GetOpSchema()); + fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); diff --git a/onnxruntime/test/contrib_ops/cuda_kernels/speculative_fold_v2_test.cc b/onnxruntime/test/contrib_ops/cuda_kernels/speculative_fold_v2_test.cc new file mode 100644 index 0000000000000..ea1532ab6156e --- /dev/null +++ b/onnxruntime/test/contrib_ops/cuda_kernels/speculative_fold_v2_test.cc @@ -0,0 +1,215 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include +#include +#include +#include +#include +#include +#include "contrib_ops/cuda/llm/speculative_fold_v2.h" + +namespace onnxruntime::test { +namespace { +namespace fold = contrib::cuda::speculative_fold_v2; + +void Check(cudaError_t status) { + if (status != cudaSuccess) throw std::runtime_error(cudaGetErrorString(status)); +} + +template +struct Buffer { + explicit Buffer(size_t count) : count(count) { Check(cudaMalloc(&data, count * sizeof(T))); } + ~Buffer() { + if (cudaFree(data) != cudaSuccess) std::fprintf(stderr, "Fold test cudaFree failed\n"); + } + Buffer(const Buffer&) = delete; + Buffer& operator=(const Buffer&) = delete; + void Upload(const std::vector& values) { + ASSERT_EQ(values.size(), count); + Check(cudaMemcpy(data, values.data(), count * sizeof(T), cudaMemcpyHostToDevice)); + } + std::vector Read() const { + std::vector values(count); + Check(cudaMemcpy(values.data(), data, count * sizeof(T), cudaMemcpyDeviceToHost)); + return values; + } + T* data{}; + size_t count; +}; + +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; + } + + 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); + } +}; + +class FoldContribV2Test : public testing::Test { + protected: + void SetUp() override { + int devices = 0; + if (cudaGetDeviceCount(&devices) != cudaSuccess || devices == 0) GTEST_SKIP() << "CUDA is required"; + } +}; + +TEST_F(FoldContribV2Test, 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); + } + } +} + +TEST_F(FoldContribV2Test, 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); +} + +TEST_F(FoldContribV2Test, 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)); + } +} + +TEST_F(FoldContribV2Test, TailSlotsAndValidationArePublished) { + Fixture f; + for (int slot : {2, 8, 9}) { + f.Reset(4, 2); + f.status.Upload({1, 4}); + ASSERT_TRUE(f.Mid(4, 4).IsOK()); + ASSERT_TRUE(f.Tail(slot).IsOK()); + Check(cudaDeviceSynchronize()); + EXPECT_EQ(f.ctl.Read()[slot], 10); + EXPECT_EQ(f.ctl.Read()[5], 2); + EXPECT_EQ(f.ring.Read()[4], 7); + EXPECT_EQ(f.ring.Read()[0], 1); + } +} + +TEST_F(FoldContribV2Test, CapturedGraphReplaysWithFreshAcceptance) { + Fixture f; + cudaStream_t stream{}; + cudaGraph_t graph{}; + cudaGraphExec_t executable{}; + Check(cudaStreamCreate(&stream)); + f.Reset(4, 0); + Check(cudaStreamBeginCapture(stream, cudaStreamCaptureModeThreadLocal)); + ASSERT_TRUE(f.Mid(4, 4, stream).IsOK()); + ASSERT_TRUE(f.Tail(2, stream).IsOK()); + Check(cudaStreamEndCapture(stream, &graph)); + Check(cudaGraphInstantiate(&executable, graph, nullptr, nullptr, 0)); + for (int accepted : {0, 3, 1, 2}) { + f.Reset(4, accepted); + Check(cudaGraphLaunch(executable, stream)); + Check(cudaStreamSynchronize(stream)); + EXPECT_EQ(f.selection.Read()[0], accepted); + EXPECT_EQ(f.ring.Read()[0], 1); + } + Check(cudaGraphExecDestroy(executable)); + Check(cudaGraphDestroy(graph)); + Check(cudaStreamDestroy(stream)); +} +} // namespace +} // namespace onnxruntime::test From 8bdbc5f00e1f246fee9b28dcc32e469a504a30b5 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Fri, 23 Oct 2026 20:50:46 +0000 Subject: [PATCH 2/2] Replace speculative fold operators at version 1 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- cmake/onnxruntime_unittests.cmake | 2 - docs/ContribOperators.md | 100 ++-- docs/OperatorKernels.md | 4 +- docs/contrib_ops/cuda/speculative_fold.md | 24 +- .../contrib_ops/cuda/cuda_contrib_kernels.cc | 4 - .../contrib_ops/cuda/llm/speculative_fold.cc | 196 ++++---- .../contrib_ops/cuda/llm/speculative_fold.cu | 266 +++++----- .../contrib_ops/cuda/llm/speculative_fold.h | 34 +- .../cuda/llm/speculative_fold_v2.cc | 141 ------ .../cuda/llm/speculative_fold_v2.cu | 150 ------ .../cuda/llm/speculative_fold_v2.h | 23 - .../core/graph/contrib_ops/bert_defs.cc | 4 +- .../core/graph/contrib_ops/contrib_defs.cc | 67 --- onnxruntime/core/graph/contrib_ops/ms_opset.h | 4 - .../cuda_kernels/speculative_fold_test.cc | 475 +++++++----------- .../cuda_kernels/speculative_fold_v2_test.cc | 215 -------- .../contrib_ops/speculative_fold_op_test.cc | 219 +++----- 17 files changed, 565 insertions(+), 1363 deletions(-) delete mode 100644 onnxruntime/contrib_ops/cuda/llm/speculative_fold_v2.cc delete mode 100644 onnxruntime/contrib_ops/cuda/llm/speculative_fold_v2.cu delete mode 100644 onnxruntime/contrib_ops/cuda/llm/speculative_fold_v2.h delete mode 100644 onnxruntime/test/contrib_ops/cuda_kernels/speculative_fold_v2_test.cc diff --git a/cmake/onnxruntime_unittests.cmake b/cmake/onnxruntime_unittests.cmake index 75bb7b3c7209b..7c5d8de4dd80d 100644 --- a/cmake/onnxruntime_unittests.cmake +++ b/cmake/onnxruntime_unittests.cmake @@ -1127,7 +1127,6 @@ if (onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS AND onnxruntime_BUILD_CUDA_EP_AS_P "${TEST_SRC_DIR}/contrib_ops/cuda_kernels/softmax_topk_kernel_test.cc" "${TEST_SRC_DIR}/contrib_ops/cuda_kernels/mtp_moe_test.cc" "${TEST_SRC_DIR}/contrib_ops/cuda_kernels/speculative_fold_test.cc" - "${TEST_SRC_DIR}/contrib_ops/cuda_kernels/speculative_fold_v2_test.cc" ) endif() @@ -1163,7 +1162,6 @@ if (onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS AND onnxruntime_BUILD_CUDA_EP_AS_P "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/llm/mtp_moe.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/llm/speculative_fold_v2.cu" "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/moe/qmoe_kernels.cu" ) if (onnxruntime_USE_FP4_QMOE) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 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 index bc893e87b39f6..447e59edacef7 100644 --- a/docs/contrib_ops/cuda/speculative_fold.md +++ b/docs/contrib_ops/cuda/speculative_fold.md @@ -1,9 +1,9 @@ -# CUDA speculative fold operators (version 2) +# CUDA speculative fold operators -`com.microsoft::FoldMid2` and `FoldTail` version 2 implement batch-one +`com.microsoft::FoldMid2` and `FoldTail` schema version 1 implement batch-one speculative acceptance and result publication without a custom-op library. -They extend the two-row version-1 design for two through four head rows, -variable rotary width, configurable cache capacity, optional hidden-state +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, @@ -19,7 +19,7 @@ and prefix contexts contain four pairs of int64 words. 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. This differs from version 1, which published in mid. +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 @@ -27,7 +27,10 @@ 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. -The kernels do not allocate or synchronize the stream and support CUDA graph replay. +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 @@ -39,9 +42,12 @@ 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 -both schemas in the core and both kernels in the CUDA provider. Rebuild both, +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. -`FoldContribV2Test.*` in `onnxruntime_provider_test` covers all S1-S4 prefix +`FoldContribTest.*` in `onnxruntime_provider_test` covers all S1-S4 prefix lengths, hidden-row padding, deferred/wrapped publication, muted prefill, invalid -metadata, status aggregation, tail slots, and captured-graph replay. +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/cuda_contrib_kernels.cc b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc index 9e52891f8bdbc..93b62d9fff569 100644 --- a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc @@ -103,8 +103,6 @@ class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, MoE); class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, MtpMoe); class CUDA_MS_OP_CLASS_NAME(1, FoldMid2); class CUDA_MS_OP_CLASS_NAME(1, FoldTail); -class CUDA_MS_OP_CLASS_NAME(2, FoldMid2); -class CUDA_MS_OP_CLASS_NAME(2, FoldTail); 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); @@ -472,8 +470,6 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cc b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cc 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/contrib_ops/cuda/llm/speculative_fold_v2.cc b/onnxruntime/contrib_ops/cuda/llm/speculative_fold_v2.cc deleted file mode 100644 index 1a501a1155802..0000000000000 --- a/onnxruntime/contrib_ops/cuda/llm/speculative_fold_v2.cc +++ /dev/null @@ -1,141 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -#include "contrib_ops/cuda/llm/speculative_fold_v2.h" -#include "core/providers/cuda/cuda_kernel.h" -#include "core/common/safeint.h" -#include - -namespace onnxruntime::contrib::cuda { -using namespace onnxruntime::cuda; -namespace { - -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) { - 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 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 && - (rows <= head_rows_ || rows > 4) && 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_v2::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) { - 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 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_v2::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, 2, kCudaExecutionProvider, - (*KernelDefBuilder::Create()) - .Alias(3, 8) - .Alias(4, 9) - .TypeConstraint("I64", DataTypeImpl::GetTensorType()) - .TypeConstraint("I32", DataTypeImpl::GetTensorType()) - .TypeConstraint("B", DataTypeImpl::GetTensorType()), - FoldMid2); - -ONNX_OPERATOR_KERNEL_EX( - FoldTail, kMSDomain, 2, kCudaExecutionProvider, - (*KernelDefBuilder::Create()) - .Alias(2, 1) - .TypeConstraint("I64", DataTypeImpl::GetTensorType()) - .TypeConstraint("I32", DataTypeImpl::GetTensorType()), - FoldTail); - -} // namespace onnxruntime::contrib::cuda diff --git a/onnxruntime/contrib_ops/cuda/llm/speculative_fold_v2.cu b/onnxruntime/contrib_ops/cuda/llm/speculative_fold_v2.cu deleted file mode 100644 index 2995d9b1841d9..0000000000000 --- a/onnxruntime/contrib_ops/cuda/llm/speculative_fold_v2.cu +++ /dev/null @@ -1,150 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -#include "contrib_ops/cuda/llm/speculative_fold_v2.h" -#include "core/common/common.h" -#include - -namespace onnxruntime::contrib::cuda::speculative_fold_v2 { -namespace { - -__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 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 errors; - __shared__ bool muted; - if (threadIdx.x == 0) { - 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 > 4 && !muted) errors |= 4; - if (errors) { - ctl[5] |= errors; - ctl[3] = 1; - } - } - __syncthreads(); - 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) 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; - } - 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_c[0] = static_cast(cache_length); - 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. -} - -__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]; - } -} - -__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; - } - 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)); -} -} // namespace - -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())); - } - return Status::OK(); -} - -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 onnxruntime::contrib::cuda::speculative_fold_v2 diff --git a/onnxruntime/contrib_ops/cuda/llm/speculative_fold_v2.h b/onnxruntime/contrib_ops/cuda/llm/speculative_fold_v2.h deleted file mode 100644 index 22c9374276ffb..0000000000000 --- a/onnxruntime/contrib_ops/cuda/llm/speculative_fold_v2.h +++ /dev/null @@ -1,23 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -#pragma once - -#include -#include -#include "core/common/status.h" - -namespace onnxruntime::contrib::cuda::speculative_fold_v2 { - -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 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 onnxruntime::contrib::cuda::speculative_fold_v2 diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index e0d7b9c428c4e..c2b1e39a68546 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -35,7 +35,7 @@ namespace onnxruntime { namespace contrib { ONNX_MS_OPERATOR_SET_SCHEMA( - FoldMid2, 2, + 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 " @@ -105,7 +105,7 @@ ONNX_MS_OPERATOR_SET_SCHEMA( })); ONNX_MS_OPERATOR_SET_SCHEMA( - FoldTail, 2, + 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 " 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/core/graph/contrib_ops/ms_opset.h b/onnxruntime/core/graph/contrib_ops/ms_opset.h index c9ef934a3edf8..71b3c719f9906 100644 --- a/onnxruntime/core/graph/contrib_ops/ms_opset.h +++ b/onnxruntime/core/graph/contrib_ops/ms_opset.h @@ -87,8 +87,6 @@ class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, MoE); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, MtpMoe); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, FoldMid2); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, FoldTail); -class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 2, FoldMid2); -class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 2, FoldTail); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, QMoE); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, MultiHeadAttention); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, GroupQueryAttention); @@ -226,8 +224,6 @@ class OpSet_Microsoft_ver1 { fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); - fn(GetOpSchema()); - fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); diff --git a/onnxruntime/test/contrib_ops/cuda_kernels/speculative_fold_test.cc b/onnxruntime/test/contrib_ops/cuda_kernels/speculative_fold_test.cc 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(), d_ctl.As(), nullptr); - ASSERT_TRUE(launch_status.IsOK()) << launch_status.ErrorMessage(); - CheckCuda(cudaStreamSynchronize(nullptr)); - - std::array actual_ctl{}; - int32_t done = 0; - d_ctl.Download(actual_ctl.data(), sizeof(actual_ctl)); - d_done.Download(&done, sizeof(done)); - EXPECT_EQ(done, 1); - EXPECT_EQ(actual_ctl[2], 42); - EXPECT_EQ(actual_ctl[5], 3); - - ctl[2] = 17; - ctl[3] = 1; - d_ctl.Upload(ctl.data(), sizeof(ctl)); - launch_status = LaunchFoldTail( - d_next_token.As(), d_status.As(), - static_cast(status.size()), d_ctl.As(), - d_done.As(), d_ctl.As(), nullptr); - ASSERT_TRUE(launch_status.IsOK()) << launch_status.ErrorMessage(); - CheckCuda(cudaStreamSynchronize(nullptr)); - d_ctl.Download(actual_ctl.data(), sizeof(actual_ctl)); - d_done.Download(&done, sizeof(done)); - EXPECT_EQ(done, 0); - EXPECT_EQ(actual_ctl, ctl); +TEST_F(FoldContribTest, CapturedGraphReplaysWithFreshAcceptance) { + Fixture f; + cudaStream_t stream{}; + cudaGraph_t graph{}; + cudaGraphExec_t executable{}; + Check(cudaStreamCreate(&stream)); + f.Reset(4, 0); + Check(cudaStreamBeginCapture(stream, cudaStreamCaptureModeThreadLocal)); + ASSERT_TRUE(f.Mid(4, 4, stream).IsOK()); + ASSERT_TRUE(f.Tail(2, stream).IsOK()); + Check(cudaStreamEndCapture(stream, &graph)); + Check(cudaGraphInstantiate(&executable, graph, nullptr, nullptr, 0)); + for (int accepted : {0, 3, 1, 2}) { + f.Reset(4, accepted); + Check(cudaGraphLaunch(executable, stream)); + Check(cudaStreamSynchronize(stream)); + EXPECT_EQ(f.selection.Read()[0], accepted); + EXPECT_EQ(f.ring.Read()[0], 1); + } + Check(cudaGraphExecDestroy(executable)); + Check(cudaGraphDestroy(graph)); + Check(cudaStreamDestroy(stream)); } - } // namespace -} // namespace test -} // namespace onnxruntime +} // namespace onnxruntime::test diff --git a/onnxruntime/test/contrib_ops/cuda_kernels/speculative_fold_v2_test.cc b/onnxruntime/test/contrib_ops/cuda_kernels/speculative_fold_v2_test.cc deleted file mode 100644 index ea1532ab6156e..0000000000000 --- a/onnxruntime/test/contrib_ops/cuda_kernels/speculative_fold_v2_test.cc +++ /dev/null @@ -1,215 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -#include -#include -#include -#include -#include -#include -#include -#include "contrib_ops/cuda/llm/speculative_fold_v2.h" - -namespace onnxruntime::test { -namespace { -namespace fold = contrib::cuda::speculative_fold_v2; - -void Check(cudaError_t status) { - if (status != cudaSuccess) throw std::runtime_error(cudaGetErrorString(status)); -} - -template -struct Buffer { - explicit Buffer(size_t count) : count(count) { Check(cudaMalloc(&data, count * sizeof(T))); } - ~Buffer() { - if (cudaFree(data) != cudaSuccess) std::fprintf(stderr, "Fold test cudaFree failed\n"); - } - Buffer(const Buffer&) = delete; - Buffer& operator=(const Buffer&) = delete; - void Upload(const std::vector& values) { - ASSERT_EQ(values.size(), count); - Check(cudaMemcpy(data, values.data(), count * sizeof(T), cudaMemcpyHostToDevice)); - } - std::vector Read() const { - std::vector values(count); - Check(cudaMemcpy(values.data(), data, count * sizeof(T), cudaMemcpyDeviceToHost)); - return values; - } - T* data{}; - size_t count; -}; - -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; - } - - 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); - } -}; - -class FoldContribV2Test : public testing::Test { - protected: - void SetUp() override { - int devices = 0; - if (cudaGetDeviceCount(&devices) != cudaSuccess || devices == 0) GTEST_SKIP() << "CUDA is required"; - } -}; - -TEST_F(FoldContribV2Test, 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); - } - } -} - -TEST_F(FoldContribV2Test, 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); -} - -TEST_F(FoldContribV2Test, 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)); - } -} - -TEST_F(FoldContribV2Test, TailSlotsAndValidationArePublished) { - Fixture f; - for (int slot : {2, 8, 9}) { - f.Reset(4, 2); - f.status.Upload({1, 4}); - ASSERT_TRUE(f.Mid(4, 4).IsOK()); - ASSERT_TRUE(f.Tail(slot).IsOK()); - Check(cudaDeviceSynchronize()); - EXPECT_EQ(f.ctl.Read()[slot], 10); - EXPECT_EQ(f.ctl.Read()[5], 2); - EXPECT_EQ(f.ring.Read()[4], 7); - EXPECT_EQ(f.ring.Read()[0], 1); - } -} - -TEST_F(FoldContribV2Test, CapturedGraphReplaysWithFreshAcceptance) { - Fixture f; - cudaStream_t stream{}; - cudaGraph_t graph{}; - cudaGraphExec_t executable{}; - Check(cudaStreamCreate(&stream)); - f.Reset(4, 0); - Check(cudaStreamBeginCapture(stream, cudaStreamCaptureModeThreadLocal)); - ASSERT_TRUE(f.Mid(4, 4, stream).IsOK()); - ASSERT_TRUE(f.Tail(2, stream).IsOK()); - Check(cudaStreamEndCapture(stream, &graph)); - Check(cudaGraphInstantiate(&executable, graph, nullptr, nullptr, 0)); - for (int accepted : {0, 3, 1, 2}) { - f.Reset(4, accepted); - Check(cudaGraphLaunch(executable, stream)); - Check(cudaStreamSynchronize(stream)); - EXPECT_EQ(f.selection.Read()[0], accepted); - EXPECT_EQ(f.ring.Read()[0], 1); - } - Check(cudaGraphExecDestroy(executable)); - Check(cudaGraphDestroy(graph)); - Check(cudaStreamDestroy(stream)); -} -} // namespace -} // namespace onnxruntime::test diff --git a/onnxruntime/test/contrib_ops/speculative_fold_op_test.cc b/onnxruntime/test/contrib_ops/speculative_fold_op_test.cc index 058ab6eb4bd06..5b1a973347e41 100644 --- a/onnxruntime/test/contrib_ops/speculative_fold_op_test.cc +++ b/onnxruntime/test/contrib_ops/speculative_fold_op_test.cc @@ -7,20 +7,18 @@ #include #include -#include +#include #include #include #include +#include #include -#include "core/common/logging/logging.h" #include "core/graph/model.h" #include "core/session/inference_session.h" #include "gtest/gtest.h" #include "test/common/cuda_op_test_utils.h" #include "test/common/tensor_op_test_utils.h" -#include "test/unittest_util/framework_test_utils.h" -#include "test/util/include/asserts.h" #include "test/util/include/default_providers.h" #include "test/util/include/test_environment.h" @@ -31,75 +29,93 @@ namespace { class FoldResultRing { public: FoldResultRing() { - cudaError_t error = cudaHostAlloc(reinterpret_cast(&host_), 32 * sizeof(int64_t), cudaHostAllocMapped); + const cudaError_t error = cudaHostAlloc(reinterpret_cast(&host_), 32 * sizeof(int64_t), cudaHostAllocMapped); ORT_ENFORCE(error == cudaSuccess, cudaGetErrorString(error)); - error = cudaHostGetDevicePointer(reinterpret_cast(&device_), host_, 0); - if (error != cudaSuccess) { + const cudaError_t pointer_error = cudaHostGetDevicePointer(reinterpret_cast(&device_), host_, 0); + if (pointer_error != cudaSuccess) { cudaFreeHost(host_); - ORT_THROW(cudaGetErrorString(error)); + ORT_THROW(cudaGetErrorString(pointer_error)); } std::fill_n(host_, 32, int64_t{-1}); } + ~FoldResultRing() { cudaFreeHost(host_); } + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(FoldResultRing); - int64_t* Host() const { return host_; } + int64_t Address() const { return static_cast(reinterpret_cast(device_)); } + int64_t* Host() const { return host_; } private: int64_t* host_ = nullptr; int64_t* device_ = nullptr; }; -TEST(SpeculativeFoldTest, CpuFeedsPropagateStateAcrossRunsAndRejectInvalidMetadata) { +TEST(SpeculativeFoldTest, InferenceSessionLoadsVersionOneAndAcceptsMutedPrefillRows) { if (!HasCudaEnvironment(800)) GTEST_SKIP() << "Requires CUDA SM80 or newer"; auto provider = DefaultCudaExecutionProvider(); ASSERT_NE(provider, nullptr); FoldResultRing ring; - Model model("fold_state_outputs", false, ModelMetaData(), PathString(), + + Model model("speculative_fold", false, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(), {{kMSDomain, 1}}, {}, DefaultLoggingManager().DefaultLogger()); auto& graph = model.MainGraph(); - auto arg = [&](const char* name, int element_type, std::initializer_list shape) -> NodeArg* { + auto arg = [&](const char* name, int element_type, std::initializer_list dims) -> NodeArg* { ONNX_NAMESPACE::TypeProto type; auto* tensor = type.mutable_tensor_type(); tensor->set_elem_type(element_type); - for (auto dim : shape) tensor->mutable_shape()->add_dim()->set_dim_value(dim); + for (const auto dim : dims) { + auto* shape_dim = tensor->mutable_shape()->add_dim(); + if (dim < 0) { + shape_dim->set_dim_param("rows"); + } else { + shape_dim->set_dim_value(dim); + } + } return &graph.GetOrCreateNodeArg(name, &type); }; + constexpr int i64 = ONNX_NAMESPACE::TensorProto_DataType_INT64; constexpr int i32 = ONNX_NAMESPACE::TensorProto_DataType_INT32; constexpr int bf16 = ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16; - auto* next_tokens = arg("next_tokens", i64, {2}); - auto* input_ids = arg("input_ids", i64, {2}); + auto* tokens = arg("tokens", i64, {-1}); + auto* ids = arg("ids", i64, {-1}); auto* status = arg("status", i32, {1}); auto* ctl = arg("ctl", i64, {8}); - auto* ctx = arg("ctx_commit", i64, {2}); - auto* ctx_a = arg("ctx_a", i64, {2}); - auto* ctx_b = arg("ctx_b", i64, {2}); + auto* context = arg("context", i64, {2}); + auto* prefixes = arg("prefixes", i64, {8}); + auto* prefix2 = arg("prefix2", i64, {2}); auto* ptrs = arg("ptrs", i64, {2}); - auto* cos = arg("cos_cache", bf16, {16, 32}); - auto* sin = arg("sin_cache", bf16, {16, 32}); + auto* cosine = arg("cosine", bf16, {16, 3}); + auto* sine = arg("sine", bf16, {16, 3}); + auto* hidden = arg("hidden", bf16, {1, -1, 3}); auto* head_ids = arg("head_ids", i64, {1, 2}); auto* head_c = arg("head_c", i32, {1}); - auto* slots = arg("slot_ids", i64, {1, 2}); - auto* lengths = arg("seqlens", i32, {1}); - auto* cos_out = arg("cos_out", bf16, {1, 2, 32}); - auto* sin_out = arg("sin_out", bf16, {1, 2, 32}); - auto* selected = arg("sel", i64, {1}); + auto* slots = arg("slots", i64, {1, 2}); + auto* lengths = arg("lengths", i32, {1}); + auto* cos_out = arg("cos_out", bf16, {1, 2, 3}); + auto* sin_out = arg("sin_out", bf16, {1, 2, 3}); + auto* selected = arg("selected", i64, {1}); + auto* head_hidden = arg("head_hidden", bf16, {1, 2, 3}); auto* mid_ctl = arg("mid_ctl", i64, {8}); - auto* updated_ctx = arg("updated_ctx", i64, {2}); + auto* updated_context = arg("updated_context", i64, {2}); auto* draft = arg("draft", i64, {1}); + auto* downstream_status = arg("downstream_status", i32, {1}); auto* done = arg("done", i32, {1}); auto* updated_ctl = arg("updated_ctl", i64, {8}); graph.AddNode("mid", "FoldMid2", "", - {next_tokens, input_ids, status, ctl, ctx, ctx_a, ctx_b, ptrs, cos, sin}, - {head_ids, head_c, slots, lengths, cos_out, sin_out, selected, mid_ctl, updated_ctx}, + {tokens, ids, status, ctl, context, prefixes, prefix2, ptrs, cosine, sine, hidden}, + {head_ids, head_c, slots, lengths, cos_out, sin_out, selected, head_hidden, mid_ctl, updated_context}, nullptr, kMSDomain) .SetExecutionProviderType(kCudaExecutionProvider); - graph.AddNode("tail", "FoldTail", "", {draft, status, mid_ctl}, {done, updated_ctl}, nullptr, kMSDomain) + graph.AddNode("tail", "FoldTail", "", {draft, downstream_status, mid_ctl, ptrs}, {done, updated_ctl}, nullptr, + kMSDomain) .SetExecutionProviderType(kCudaExecutionProvider); - graph.SetOutputs({updated_ctl, updated_ctx, selected, cos_out, sin_out, done}); + graph.SetOutputs({head_ids, head_c, slots, lengths, cos_out, sin_out, selected, head_hidden, + updated_ctl, updated_context, done}); ASSERT_STATUS_OK(graph.Resolve()); + std::stringstream stream; ASSERT_TRUE(model.ToProto().SerializeToOstream(&stream)); SessionOptions options; @@ -107,124 +123,53 @@ TEST(SpeculativeFoldTest, CpuFeedsPropagateStateAcrossRunsAndRejectInvalidMetada ASSERT_STATUS_OK(session.RegisterExecutionProvider(std::move(provider))); ASSERT_STATUS_OK(session.Load(stream)); ASSERT_STATUS_OK(session.Initialize()); + auto allocator = TestCPUExecutionProvider()->CreatePreferredAllocators()[0]; NameMLValMap feeds; - const std::vector output_names{ - "updated_ctl", "updated_ctx", "sel", "cos_out", "sin_out", "done"}; auto feed = [&](const char* name, const TensorShapeVector& shape, const auto& values) { using T = typename std::decay_t::value_type; OrtValue value; CreateMLValue(allocator, shape, values, &value); feeds[name] = std::move(value); }; - feed("next_tokens", {2}, std::vector{7, 8}); - feed("input_ids", {2}, std::vector{0, 7}); - feed("status", {1}, std::vector{4}); - feed("ctx_a", {2}, std::vector{11, 12}); - feed("ctx_b", {2}, std::vector{21, 22}); - feed("ptrs", {2}, std::vector{ring.Address(), 4}); - feed("draft", {1}, std::vector{42}); - std::vector cache(16 * 32); - for (size_t i = 0; i < cache.size(); ++i) cache[i] = BFloat16(static_cast(i)); - feed("cos_cache", {16, 32}, cache); - feed("sin_cache", {16, 32}, cache); - std::vector control{0, 10, 99, 0, 0, 0, 0, 5}; - std::vector context{1, 2}; - for (int run = 0; run < 2; ++run) { - feed("ctl", {8}, control); - feed("ctx_commit", {2}, context); - std::vector outputs; - ASSERT_STATUS_OK(session.Run(RunOptions{}, feeds, output_names, &outputs)); - ASSERT_EQ(outputs.size(), 6u); - const auto* result = outputs[0].Get().Data(); - EXPECT_EQ(result[0], 8); - EXPECT_EQ(result[1], control[1] + 2); - EXPECT_EQ(result[2], 42); - EXPECT_EQ(result[4], control[4] + 1); - EXPECT_EQ(result[5], 2); - EXPECT_EQ(result[7], control[7] + 2); - EXPECT_EQ(outputs[2].Get().Data()[0], 1); - EXPECT_EQ(outputs[5].Get().Data()[0], 1); - for (int output_index : {3, 4}) { - const auto* rotary = outputs[output_index].Get().Data(); - for (int i = 0; i < 64; ++i) { - EXPECT_EQ(rotary[i], cache[static_cast(control[1] + 1) * 32 + i]); - } - } - const auto* committed = outputs[1].Get().Data(); - EXPECT_EQ(committed[0], 21); - EXPECT_EQ(committed[1], 22); - EXPECT_TRUE(std::equal(control.begin(), control.end(), feeds.at("ctl").Get().Data())); - EXPECT_TRUE(std::equal(context.begin(), context.end(), feeds.at("ctx_commit").Get().Data())); - control.assign(result, result + 8); - context.assign(committed, committed + 2); - } - EXPECT_EQ(control[1], 14); - EXPECT_EQ(ring.Host()[0], 1); - EXPECT_EQ(ring.Host()[8], 2); - - // Rejection at the last valid pair still commits one token; a muted run preserves both states. - feed("input_ids", {2}, std::vector{0, 9}); - for (int64_t muted : {int64_t{0}, int64_t{1}}) { - control = {0, 13, 99, muted, 2, 0, 0, 5}; - context = {1, 2}; - feed("ctl", {8}, control); - feed("ctx_commit", {2}, context); - feed("ptrs", {2}, std::vector{muted ? 0 : ring.Address(), muted ? 0 : 4}); - std::vector outputs; - ASSERT_STATUS_OK(session.Run(RunOptions{}, feeds, output_names, &outputs)); - const auto* result = outputs[0].Get().Data(); - EXPECT_EQ(outputs[2].Get().Data()[0], 0); - EXPECT_EQ(outputs[5].Get().Data()[0], muted ? 0 : 1); - EXPECT_EQ(result[1], muted ? 13 : 14); - EXPECT_EQ(result[2], muted ? 99 : 42); - EXPECT_EQ(result[4], muted ? 2 : 3); - EXPECT_EQ(result[5], muted ? 0 : 2); - EXPECT_EQ(result[7], muted ? 5 : 6); - const auto* committed = outputs[1].Get().Data(); - EXPECT_EQ(committed[0], muted ? 1 : 11); - EXPECT_EQ(committed[1], muted ? 2 : 12); - for (int output_index : {3, 4}) { - const auto* rotary = outputs[output_index].Get().Data(); - for (int i = 0; i < 64; ++i) EXPECT_EQ(rotary[i], cache[14 * 32 + i]); - } - } + std::vector rope(16 * 3); + for (size_t i = 0; i < rope.size(); ++i) rope[i] = BFloat16(static_cast(i)); - struct InvalidCase { - int64_t position; - int64_t muted; - int64_t address; - int64_t capacity; - int64_t counter; - int64_t error; - int64_t cache_length = 5; - }; - const std::vector invalid_cases{ - {14, 0, ring.Address(), 4, 0, 8}, {14, 1, ring.Address(), 4, 0, 8}, {-1, 0, ring.Address(), 4, 0, 8}, {std::numeric_limits::max(), 0, ring.Address(), 4, 0, 8}, {10, 0, 0, 4, 0, 16}, {10, 0, ring.Address(), 0, 0, 16}, {10, 0, ring.Address(), -1, 0, 16}, {10, 0, ring.Address(), 4, -1, 16}, {10, 0, ring.Address(), 4, std::numeric_limits::max(), 16}, {10, 0, ring.Address() + 1, 4, 0, 16}, {10, 0, ring.Address(), std::numeric_limits::max(), 0, 16}, {10, 0, ring.Address(), 4, 0, 4, -1}, {10, 0, ring.Address(), 4, 0, 4, 33027}}; - const std::vector preserved_ring(ring.Host(), ring.Host() + 32); - for (const auto& invalid : invalid_cases) { - SCOPED_TRACE(invalid.position); - SCOPED_TRACE(invalid.capacity); - SCOPED_TRACE(invalid.counter); - control = {0, invalid.position, 99, invalid.muted, invalid.counter, 0, 0, invalid.cache_length}; - feed("ctl", {8}, control); - feed("ctx_commit", {2}, context); - feed("ptrs", {2}, std::vector{invalid.address, invalid.capacity}); + auto run = [&](int rows, bool muted) { + std::vector token_values(rows); + std::vector id_values(rows, -1); + std::vector hidden_rows(rows * 3); + for (int row = 0; row < rows; ++row) token_values[row] = row + 10; + id_values[0] = 0; + if (rows > 1) id_values[1] = token_values[0]; + feed("tokens", {rows}, token_values); + feed("ids", {rows}, id_values); + feed("status", {1}, std::vector{0}); + feed("ctl", {8}, std::vector{0, 3, 0, muted ? 1 : 0, 0, 0, 0, 2}); + feed("context", {2}, std::vector{1, 2}); + feed("prefixes", {8}, std::vector{100, 101, 200, 201, 300, 301, 400, 401}); + feed("prefix2", {2}, std::vector{200, 201}); + feed("ptrs", {2}, muted ? std::vector{0, 0} : std::vector{ring.Address(), 4}); + feed("cosine", {16, 3}, rope); + feed("sine", {16, 3}, rope); + feed("hidden", {1, rows, 3}, hidden_rows); + feed("draft", {1}, std::vector{42}); + feed("downstream_status", {1}, std::vector{0}); std::vector outputs; - ASSERT_STATUS_OK(session.Run(RunOptions{}, feeds, output_names, &outputs)); - auto expected = control; - expected[3] = 1; - expected[5] = invalid.error; - EXPECT_TRUE(std::equal(expected.begin(), expected.end(), outputs[0].Get().Data())); - EXPECT_TRUE(std::equal(context.begin(), context.end(), outputs[1].Get().Data())); - EXPECT_EQ(outputs[2].Get().Data()[0], -1); - EXPECT_EQ(outputs[5].Get().Data()[0], 0); - for (int output_index : {3, 4}) { - const auto* rotary = outputs[output_index].Get().Data(); - for (int i = 0; i < 64; ++i) EXPECT_EQ(rotary[i].ToFloat(), 0.0f); + const std::vector output_names{ + "head_ids", "head_c", "slots", "lengths", "cos_out", "sin_out", + "selected", "head_hidden", "updated_ctl", "updated_context", "done"}; + const Status status = session.Run(RunOptions{}, feeds, output_names, &outputs); + EXPECT_TRUE(status.IsOK()) << status; + if (status.IsOK()) { + EXPECT_EQ(outputs[6].Get().Data()[0], muted ? 0 : 1); + EXPECT_EQ(outputs[10].Get().Data()[0], muted ? 0 : 1); } - EXPECT_TRUE(std::equal(preserved_ring.begin(), preserved_ring.end(), ring.Host())); - } + }; + + run(2, false); + run(3, true); + run(4, true); } } // namespace