From a4f693dc9fd6e51fd441e618b1fddac708d28b9d Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 26 Oct 2026 20:05:29 +0000 Subject: [PATCH 1/2] [CUDA] Support batch_size >= 1 in SpeculativePrefixCommit - Propagate batch dimension from next_tokens[0] in SpeculativePrefixCommit type and shape inference. - Launch MidKernel across batch_size blocks with per-slot control blocks, contexts, ring buffers, and muted slot flags. - Add SupportsBatchedSpeculativePrefixCommitS8 unit test. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../contrib_ops/cuda/llm/speculative_fold.cc | 46 +++-- .../contrib_ops/cuda/llm/speculative_fold.cu | 128 ++++++++----- .../contrib_ops/cuda/llm/speculative_fold.h | 3 +- .../core/graph/contrib_ops/bert_defs.cc | 44 +++-- .../contrib_ops/speculative_fold_op_test.cc | 177 ++++++++++++++++++ 5 files changed, 319 insertions(+), 79 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cc b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cc index c087b952c73c3..e7cb836f2b4ab 100644 --- a/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cc +++ b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cc @@ -36,25 +36,39 @@ class SpeculativePrefixCommit final : public CudaKernel { const auto* sine = context->Input(9); ORT_RETURN_IF_NOT(tokens && ids && status && ctl && committed && prefixes && prefix2 && ptrs && cosine && sine, "SpeculativePrefixCommit: missing required input"); - const int64_t rows = tokens->Shape().Size(); - ORT_RETURN_IF_NOT(rows > 0 && rows <= INT_MAX && ids->Shape().Size() >= rows && + const auto& tok_shape = tokens->Shape(); + const int64_t batch = tok_shape.NumDimensions() >= 2 + ? tok_shape.SizeToDimension(tok_shape.NumDimensions() - 1) + : int64_t{1}; + const int64_t rows = tok_shape.NumDimensions() >= 2 + ? tok_shape[tok_shape.NumDimensions() - 1] + : tok_shape.Size(); + ORT_RETURN_IF_NOT(batch > 0 && batch <= INT_MAX && rows > 0 && rows <= INT_MAX && + ids->Shape().Size() >= batch * rows && ids->Shape().Size() % batch == 0 && + ids->Shape().Size() / batch <= INT_MAX && status->Shape().Size() <= INT_MAX, "SpeculativePrefixCommit: invalid token/status shape"); - ORT_RETURN_IF_NOT(ctl->Shape().Size() >= (head_rows_ > 2 ? 16 : 8) && - Count(committed, 2) && prefixes->Shape().Size() >= (head_rows_ > 2 ? 2 * head_rows_ : 2) && - Count(prefix2, 2) && Count(ptrs, 2), + const int64_t min_ctl = (head_rows_ > 2 ? 16 : 8); + const int64_t min_pref = (head_rows_ > 2 ? 2 * head_rows_ : 2); + ORT_RETURN_IF_NOT(ctl->Shape().Size() >= batch * min_ctl && ctl->Shape().Size() % batch == 0 && + ctl->Shape().Size() / batch <= INT_MAX && + Count(committed, 2 * batch) && + prefixes->Shape().Size() >= batch * min_pref && prefixes->Shape().Size() % batch == 0 && + prefixes->Shape().Size() / batch <= INT_MAX && + (Count(prefix2, 2 * batch) || Count(prefix2, 2)) && + (Count(ptrs, 2 * batch) || Count(ptrs, 2)), "SpeculativePrefixCommit: 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, "SpeculativePrefixCommit: rotary caches must have matching [positions, width] shapes"); - 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* out_ids = context->Output(0, TensorShape({batch, head_rows_})); + auto* head_c = context->Output(1, TensorShape({batch})); + auto* slots = context->Output(2, TensorShape({batch, head_rows_})); + auto* lengths = context->Output(3, TensorShape({batch})); + auto* out_cos = context->Output(4, TensorShape({batch, head_rows_, rope[1]})); + auto* out_sin = context->Output(5, TensorShape({batch, head_rows_, rope[1]})); + auto* selection = context->Output(6, TensorShape({batch})); int64_t* ctl_ptr = const_cast(ctl->Data()); int64_t* context_ptr = const_cast(committed->Data()); if (context->OutputCount() > 7) { @@ -69,6 +83,11 @@ class SpeculativePrefixCommit final : public CudaKernel { context_ptr = updated_context->MutableData(); } } + const int ids_stride = static_cast(ids->Shape().Size() / batch); + const int ctl_stride = static_cast(ctl->Shape().Size() / batch); + const int prefix_stride = static_cast(prefixes->Shape().Size() / batch); + const int prefix2_stride = (Count(prefix2, 2 * batch) && batch > 1) ? 2 : 0; + const int ptrs_stride = (Count(ptrs, 2 * batch) && batch > 1) ? 2 : 0; return speculative_fold::LaunchMid( tokens->Data(), ids->Data(), status->Data(), static_cast(status->Shape().Size()), ctl_ptr, context_ptr, prefixes->Data(), @@ -79,7 +98,8 @@ class SpeculativePrefixCommit final : public CudaKernel { lengths->MutableData(), reinterpret_cast(out_cos->MutableData()), reinterpret_cast(out_sin->MutableData()), selection->MutableData(), static_cast(rows), head_rows_, static_cast(rope[0]), static_cast(rope[1]), capacity_, Stream(context), - true, ctl->Data(), static_cast(ctl->Shape().Size()), committed->Data()); + true, ctl->Data(), ctl_stride, committed->Data(), + static_cast(batch), ids_stride, prefix_stride, prefix2_stride, ptrs_stride); } private: diff --git a/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cu b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cu index 42a376291edac..904f2e524f6cf 100644 --- a/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cu +++ b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cu @@ -22,87 +22,119 @@ __global__ void MidKernel(const int64_t* tokens, const int64_t* ids, const int32 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, - bool publish_ring) { + bool publish_ring, int batch_size, int ids_stride, + int prefix_stride, int prefix2_stride, int ptrs_stride) { + const int b = static_cast(blockIdx.x); + if (b >= batch_size) return; + const int eff_ids_stride = ids_stride > 0 ? ids_stride : rows; + const int eff_ctl_stride = ctl_count > 0 ? ctl_count : (head_rows > 2 ? 16 : 8); + const int eff_prefix_stride = prefix_stride > 0 ? prefix_stride : (head_rows > 2 ? 2 * head_rows : 2); + const int64_t* tokens_b = tokens + static_cast(b) * rows; + const int64_t* ids_b = ids + static_cast(b) * eff_ids_stride; + const int64_t* input_ctl_b = input_ctl != nullptr ? input_ctl + static_cast(b) * eff_ctl_stride : nullptr; + int64_t* ctl_b = ctl + static_cast(b) * eff_ctl_stride; + const int64_t* input_context_b = input_context != nullptr ? input_context + static_cast(b) * 2 : nullptr; + int64_t* context_b = context + static_cast(b) * 2; + const int64_t* prefixes_b = prefixes + static_cast(b) * eff_prefix_stride; + const int64_t* prefix2_b = (prefix2_stride > 0 || b == 0) + ? (prefix2 + static_cast(b) * prefix2_stride) + : (eff_prefix_stride >= 4 ? prefixes_b + 2 : prefix2); + int64_t* head_ids_b = head_ids + static_cast(b) * head_rows; + int64_t* slots_b = slots + static_cast(b) * head_rows; + uint16_t* head_cos_b = head_cos + static_cast(b) * head_rows * rope_width; + uint16_t* head_sin_b = head_sin + static_cast(b) * head_rows * rope_width; + __shared__ int64_t length; __shared__ int64_t cache_length; __shared__ int64_t errors; __shared__ bool muted; + __shared__ int64_t local_ptrs[2]; if (threadIdx.x == 0) { - if (input_ctl != nullptr && input_ctl != ctl) { - for (int i = 0; i < ctl_count; ++i) ctl[i] = input_ctl[i]; + if (input_ctl_b != nullptr && input_ctl_b != ctl_b) { + for (int i = 0; i < eff_ctl_stride; ++i) ctl_b[i] = input_ctl_b[i]; + } + if (input_context_b != nullptr && input_context_b != context_b) { + context_b[0] = input_context_b[0]; + context_b[1] = input_context_b[1]; } - if (input_context != nullptr && input_context != context) { - context[0] = input_context[0]; - context[1] = input_context[1]; + if (ptrs_stride > 0) { + local_ptrs[0] = ptrs[2 * b + 0]; + local_ptrs[1] = ptrs[2 * b + 1]; + } else { + local_ptrs[1] = ptrs[1]; + local_ptrs[0] = (ptrs[0] != 0 && ptrs[1] > 0 && + ptrs[1] <= LLONG_MAX / (8 * static_cast(sizeof(int64_t)) * batch_size)) + ? ptrs[0] + static_cast(b) * ptrs[1] * 8 * static_cast(sizeof(int64_t)) + : ptrs[0]; } - muted = ctl[3] != 0; - length = muted ? 0 : ctl[1]; - cache_length = muted ? 0 : ctl[7]; + muted = ctl_b[3] != 0; + length = muted ? 0 : ctl_b[1]; + cache_length = muted ? 0 : ctl_b[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 (!muted && !ValidRing(local_ptrs, ctl_b[4])) errors |= 16; if (rows > head_rows && !muted) errors |= 4; if (errors) { - ctl[5] |= errors; - ctl[3] = 1; + ctl_b[5] |= errors; + ctl_b[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; + head_cos_b[i] = !errors ? cosine[position * rope_width + i % rope_width] : 0; + head_sin_b[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; + for (int row = 0; row < head_rows; ++row) head_ids_b[row] = slots_b[row] = 0; + head_c[b] = lengths[b] = 0; + selection[b] = -1; return; } int accepted = 0; if (!muted) { while (accepted + 1 < rows && accepted + 1 < head_rows && - tokens[accepted] == ids[accepted + 1]) ++accepted; + tokens_b[accepted] == ids_b[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_b[row] = row < accepted ? ids_b[row + 1] : tokens_b[row < rows ? row : rows - 1]; + slots_b[row] = cache_length + row; } - head_c[0] = static_cast(cache_length); - lengths[0] = static_cast(cache_length + head_rows - 1); - selection[0] = accepted; + head_c[b] = static_cast(cache_length); + lengths[b] = static_cast(cache_length + head_rows - 1); + selection[b] = 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]); + const int64_t* prefix = accepted == 1 ? prefix2_b : prefixes_b + 2 * accepted; + context_b[0] = prefix[0]; + context_b[1] = prefix[1]; + const int64_t pending = tokens_b[accepted]; + ctl_b[0] = pending; + ctl_b[1] = length + accepted + 1; + ctl_b[7] = cache_length + accepted + 1; + ++ctl_b[4]; + ctl_b[5] |= validation ? 2 : 0; + ctl_b[6] = accepted; + volatile int64_t* record = reinterpret_cast(local_ptrs[0]) + 8 * ((ctl_b[4] - 1) % local_ptrs[1]); auto pack_lo_hi = [](int64_t lo, int64_t hi, bool has_hi) -> int64_t { if (!has_hi) return lo; return static_cast( static_cast(static_cast(lo)) | (static_cast(static_cast(hi)) << 32)); }; - record[1] = pack_lo_hi(tokens[0], rows > 4 ? tokens[4] : 0, rows > 4); - record[2] = pack_lo_hi(rows > 1 ? tokens[1] : tokens[0], rows > 5 ? tokens[5] : 0, rows > 5); - record[3] = pack_lo_hi(accepted, rows > 6 ? tokens[6] : 0, rows > 6); + record[1] = pack_lo_hi(tokens_b[0], rows > 4 ? tokens_b[4] : 0, rows > 4); + record[2] = pack_lo_hi(rows > 1 ? tokens_b[1] : tokens_b[0], rows > 5 ? tokens_b[5] : 0, rows > 5); + record[3] = pack_lo_hi(accepted, rows > 6 ? tokens_b[6] : 0, rows > 6); record[4] = validation; - record[5] = pack_lo_hi(pending, rows > 7 ? tokens[7] : 0, rows > 7); - record[6] = ctl[1]; - record[7] = rows > 2 ? pack_lo_hi(tokens[2], rows > 3 ? tokens[3] : 0, rows > 3) : 0; + record[5] = pack_lo_hi(pending, rows > 7 ? tokens_b[7] : 0, rows > 7); + record[6] = ctl_b[1]; + record[7] = rows > 2 ? pack_lo_hi(tokens_b[2], rows > 3 ? tokens_b[3] : 0, rows > 3) : 0; if (publish_ring) { __threadfence_system(); - record[0] = ctl[4]; + record[0] = ctl_b[4]; } } @@ -118,11 +150,15 @@ Status LaunchMid(const int64_t* tokens, const int64_t* ids, const int32_t* statu 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, cudaStream_t stream, - bool publish_ring, const int64_t* input_ctl, int ctl_count, const int64_t* input_context) { - MidKernel<<<1, 256, 0, stream>>>(tokens, ids, status, status_count, input_ctl, ctl, ctl_count, - input_context, 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, publish_ring); + bool publish_ring, const int64_t* input_ctl, int ctl_count, const int64_t* input_context, + int batch_size, int ids_stride, int prefix_stride, int prefix2_stride, int ptrs_stride) { + if (batch_size <= 0) return Status::OK(); + MidKernel<<(batch_size), 256, 0, stream>>>( + tokens, ids, status, status_count, input_ctl, ctl, ctl_count, + input_context, 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, publish_ring, + batch_size, ids_stride, prefix_stride, prefix2_stride, ptrs_stride); return LaunchStatus(cudaGetLastError()); } diff --git a/onnxruntime/contrib_ops/cuda/llm/speculative_fold.h b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.h index 5c9a4e4b506d9..4bab583cfd793 100644 --- a/onnxruntime/contrib_ops/cuda/llm/speculative_fold.h +++ b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.h @@ -16,6 +16,7 @@ Status LaunchMid(const int64_t* tokens, const int64_t* ids, const int32_t* statu uint16_t* head_cos, uint16_t* head_sin, int64_t* selection, int rows, int head_rows, int rope_rows, int rope_width, int capacity, cudaStream_t stream, bool publish_ring = true, const int64_t* input_ctl = nullptr, int ctl_count = 0, - const int64_t* input_context = nullptr); + const int64_t* input_context = nullptr, int batch_size = 1, int ids_stride = 0, + int prefix_stride = 0, int prefix2_stride = 0, int ptrs_stride = 0); } // namespace onnxruntime::contrib::cuda::speculative_fold diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index cbc43caef5616..f95d36abd0825 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -37,31 +37,31 @@ namespace contrib { ONNX_MS_OPERATOR_SET_SCHEMA( SpeculativePrefixCommit, 1, OpSchema() - .SetDoc("Commits a batch-one speculative prefix, prepares rotary/head inputs for two through eight head rows, " - "and publishes the result-ring record. Control and committed-context outputs must alias their inputs " - "on CUDA when present. 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. " - "Muted prefill calls prepare inputs without committing state or touching ring metadata. " - "Invalid device metadata sets error bits and mutes execution.") + .SetDoc("Commits a speculative prefix for each request in the batch, prepares rotary/head inputs for two " + "through eight head rows, and publishes the result-ring records. Control and committed-context outputs " + "must alias their inputs on CUDA when present. 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 per batch item. 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 8.", AttributeProto::INT, int64_t{2}) .Attr("head_capacity", "Maximum head-cache token capacity.", AttributeProto::INT, int64_t{262144}) - .Input(0, "next_tokens", "Verifier predictions, batch one.", "I64") + .Input(0, "next_tokens", "Verifier predictions, [batch, rows] or [rows] for 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..S8.", "I64") - .Input(4, "ctx_commit", "Mutable committed context, two words.", "I64") + .Input(3, "ctl", "Mutable control state: at least 8 words per batch item for S2, 16 for S3..S8.", "I64") + .Input(4, "ctx_commit", "Mutable committed context, two words per batch item.", "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") - .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(0, "head_ids", "Head token IDs [batch, head_rows].", "I64") + .Output(1, "head_c", "Head committed length [batch].", "I32") + .Output(2, "slot_ids", "Head slots [batch, head_rows].", "I64") + .Output(3, "seqlens", "Head sequence length [batch].", "I32") + .Output(4, "cos_out", "Gathered rotary cosine [batch, head_rows, width].", "B") + .Output(5, "sin_out", "Gathered rotary sine [batch, head_rows, width].", "B") + .Output(6, "sel", "Accepted-prefix length [batch], or -1 on invalid metadata.", "I64") .Output(7, "updated_ctl", "Optional updated control state, aliases ctl.", "I64", OpSchema::Optional) .Output(8, "updated_ctx_commit", "Optional updated context, aliases ctx_commit.", "I64", OpSchema::Optional) .TypeConstraint("I64", {"tensor(int64)"}, "Token and control words.") @@ -79,17 +79,23 @@ ONNX_MS_OPERATOR_SET_SCHEMA( const auto capacity = getAttribute(ctx, "head_capacity", 262144); if (rows < 2 || rows > 8 || capacity < rows || capacity > std::numeric_limits::max()) fail_shape_inference("SpeculativePrefixCommit: invalid head_rows/head_capacity"); + TensorShapeProto_Dimension batch_dim; + if (hasInputShape(ctx, 0) && getInputShape(ctx, 0).dim_size() >= 2) { + batch_dim = getInputShape(ctx, 0).dim(0); + } else { + batch_dim.set_dim_value(1); + } for (int output : {0, 2}) { auto* shape = getOutputShape(ctx, output); - shape->add_dim()->set_dim_value(1); + *shape->add_dim() = batch_dim; shape->add_dim()->set_dim_value(rows); } - for (int output : {1, 3, 6}) getOutputShape(ctx, output)->add_dim()->set_dim_value(1); + for (int output : {1, 3, 6}) *getOutputShape(ctx, output)->add_dim() = batch_dim; if (hasInputShape(ctx, 8)) { const auto& rope = getInputShape(ctx, 8); if (rope.dim_size() != 2) fail_shape_inference("SpeculativePrefixCommit: rotary cache must have rank 2"); TensorShapeProto shape; - shape.add_dim()->set_dim_value(1); + *shape.add_dim() = batch_dim; shape.add_dim()->set_dim_value(rows); *shape.add_dim() = rope.dim(1); updateOutputShape(ctx, 4, shape); diff --git a/onnxruntime/test/contrib_ops/speculative_fold_op_test.cc b/onnxruntime/test/contrib_ops/speculative_fold_op_test.cc index 7855fe391bc46..5d85277c28ceb 100644 --- a/onnxruntime/test/contrib_ops/speculative_fold_op_test.cc +++ b/onnxruntime/test/contrib_ops/speculative_fold_op_test.cc @@ -161,6 +161,183 @@ TEST(SpeculativeFoldTest, InferenceSessionLoadsVersionOneAndAcceptsMutedPrefillR run(4, true); } +TEST(SpeculativeFoldTest, SupportsBatchedSpeculativePrefixCommitS8) { + if (!HasCudaEnvironment(800)) GTEST_SKIP() << "Requires CUDA SM80 or newer"; + auto provider = DefaultCudaExecutionProvider(); + ASSERT_NE(provider, nullptr); + + constexpr int kBatch = 3; + constexpr int kRows = 8; + constexpr int kRingCap = 4; + int64_t* ring_host = nullptr; + int64_t* ring_dev = nullptr; + ASSERT_EQ(cudaHostAlloc(reinterpret_cast(&ring_host), kBatch * kRingCap * 8 * sizeof(int64_t), cudaHostAllocMapped), + cudaSuccess); + ASSERT_EQ(cudaHostGetDevicePointer(reinterpret_cast(&ring_dev), ring_host, 0), cudaSuccess); + std::fill_n(ring_host, kBatch * kRingCap * 8, int64_t{-1}); + + Model model("speculative_fold_batch", false, ModelMetaData(), PathString(), + IOnnxRuntimeOpSchemaRegistryList(), {{kMSDomain, 1}}, {}, + DefaultLoggingManager().DefaultLogger()); + auto& graph = model.MainGraph(); + 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); + int axis = 0; + for (const auto dim : dims) { + auto* shape_dim = tensor->mutable_shape()->add_dim(); + if (dim < 0) { + shape_dim->set_dim_param(axis == 0 ? "batch" : "rows"); + } else { + shape_dim->set_dim_value(dim); + } + ++axis; + } + return &graph.GetOrCreateNodeArg(name, &type); + }; + + 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* tokens = arg("tokens", i64, {-1, kRows}); + auto* ids = arg("ids", i64, {-1, kRows}); + auto* status = arg("status", i32, {1}); + auto* ctl = arg("ctl", i64, {kBatch * 16}); + auto* context = arg("context", i64, {kBatch * 2}); + auto* prefixes = arg("prefixes", i64, {kBatch * 16}); + auto* prefix2 = arg("prefix2", i64, {2}); + auto* ptrs = arg("ptrs", i64, {2}); + auto* cosine = arg("cosine", bf16, {64, 4}); + auto* sine = arg("sine", bf16, {64, 4}); + auto* head_ids = arg("head_ids", i64, {-1, kRows}); + auto* head_c = arg("head_c", i32, {-1}); + auto* slots = arg("slots", i64, {-1, kRows}); + auto* lengths = arg("lengths", i32, {-1}); + auto* cos_out = arg("cos_out", bf16, {-1, kRows, 4}); + auto* sin_out = arg("sin_out", bf16, {-1, kRows, 4}); + auto* selected = arg("selected", i64, {-1}); + auto* updated_ctl = arg("updated_ctl", i64, {kBatch * 16}); + auto* updated_context = arg("updated_context", i64, {kBatch * 2}); + + ONNX_NAMESPACE::AttributeProto head_rows_attr; + head_rows_attr.set_name("head_rows"); + head_rows_attr.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_INT); + head_rows_attr.set_i(kRows); + NodeAttributes attrs{{"head_rows", head_rows_attr}}; + + graph.AddNode("mid", "SpeculativePrefixCommit", "", + {tokens, ids, status, ctl, context, prefixes, prefix2, ptrs, cosine, sine}, + {head_ids, head_c, slots, lengths, cos_out, sin_out, selected, updated_ctl, updated_context}, + &attrs, kMSDomain) + .SetExecutionProviderType(kCudaExecutionProvider); + graph.SetOutputs({head_ids, head_c, slots, lengths, cos_out, sin_out, selected, + updated_ctl, updated_context}); + ASSERT_STATUS_OK(graph.Resolve()); + + std::stringstream stream; + ASSERT_TRUE(model.ToProto().SerializeToOstream(&stream)); + SessionOptions options; + InferenceSession session(options, GetEnvironment()); + 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; + 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); + }; + std::vector rope(64 * 4); + for (size_t i = 0; i < rope.size(); ++i) rope[i] = BFloat16(static_cast(i)); + + // Batch 0: accepts 3 drafts (tokens[0..2] == ids[1..3], tokens[3] != ids[4]) + // Batch 1: muted (ctl[3] = 1) + // Batch 2: accepts 1 draft (tokens[0] == ids[1], tokens[1] != ids[2]) + std::vector token_values = { + 10, 11, 12, 99, 14, 15, 16, 17, + 20, 21, 22, 23, 24, 25, 26, 27, + 30, 88, 32, 33, 34, 35, 36, 37, + }; + std::vector id_values = { + 1, 10, 11, 12, 13, 14, 15, 16, + 2, 20, 21, 22, 23, 24, 25, 26, + 3, 30, 31, 32, 33, 34, 35, 36, + }; + std::vector ctl_values(kBatch * 16, 0); + // Batch 0: length=10, muted=0, seq=0, cache_len=10 + ctl_values[0 * 16 + 1] = 10; + ctl_values[0 * 16 + 7] = 10; + // Batch 1: length=20, muted=1, seq=5, cache_len=20 + ctl_values[1 * 16 + 1] = 20; + ctl_values[1 * 16 + 3] = 1; + ctl_values[1 * 16 + 4] = 5; + ctl_values[1 * 16 + 7] = 20; + // Batch 2: length=30, muted=0, seq=2, cache_len=30 + ctl_values[2 * 16 + 1] = 30; + ctl_values[2 * 16 + 4] = 2; + ctl_values[2 * 16 + 7] = 30; + + std::vector prefixes_values(kBatch * 16); + for (int b = 0; b < kBatch; ++b) { + for (int r = 0; r < kRows; ++r) { + prefixes_values[b * 16 + 2 * r + 0] = 1000 * (b + 1) + 10 * r + 0; + prefixes_values[b * 16 + 2 * r + 1] = 1000 * (b + 1) + 10 * r + 1; + } + } + + feed("tokens", {kBatch, kRows}, token_values); + feed("ids", {kBatch, kRows}, id_values); + feed("status", {1}, std::vector{0}); + feed("ctl", {kBatch * 16}, ctl_values); + feed("context", {kBatch * 2}, std::vector{1, 2, 3, 4, 5, 6}); + feed("prefixes", {kBatch * 16}, prefixes_values); + feed("prefix2", {2}, std::vector{prefixes_values[2], prefixes_values[3]}); + feed("ptrs", {2}, std::vector{reinterpret_cast(ring_dev), kRingCap}); + feed("cosine", {64, 4}, rope); + feed("sine", {64, 4}, rope); + + std::vector outputs; + const std::vector output_names{ + "head_ids", "head_c", "slots", "lengths", "cos_out", "sin_out", + "selected", "updated_ctl", "updated_context"}; + ASSERT_STATUS_OK(session.Run(RunOptions{}, feeds, output_names, &outputs)); + + const int64_t* sel = outputs[6].Get().Data(); + EXPECT_EQ(sel[0], 3); + EXPECT_EQ(sel[1], 0); // muted + EXPECT_EQ(sel[2], 1); + + const int64_t* out_ctl = outputs[7].Get().Data(); + EXPECT_EQ(out_ctl[0 * 16 + 0], 99); // pending = tokens[0, 3] + EXPECT_EQ(out_ctl[0 * 16 + 1], 10 + 4); // length = 10 + 3 + 1 + EXPECT_EQ(out_ctl[0 * 16 + 4], 1); // seq incremented to 1 + + EXPECT_EQ(out_ctl[1 * 16 + 4], 5); // muted: seq unchanged + + EXPECT_EQ(out_ctl[2 * 16 + 0], 88); // pending = tokens[2, 1] + EXPECT_EQ(out_ctl[2 * 16 + 1], 30 + 2); // length = 30 + 1 + 1 + EXPECT_EQ(out_ctl[2 * 16 + 4], 3); // seq incremented from 2 to 3 + + const int64_t* out_ctx = outputs[8].Get().Data(); + // Batch 0 accepted=3 -> prefixes[0, 6..7] + EXPECT_EQ(out_ctx[0], 1030); + EXPECT_EQ(out_ctx[1], 1031); + // Batch 2 accepted=1 -> prefixes[2, 2..3] + EXPECT_EQ(out_ctx[4], 3010); + EXPECT_EQ(out_ctx[5], 3011); + + // Ring records published for batch 0 (slot 0) and batch 2 (slot 2) + EXPECT_EQ(ring_host[(0 * kRingCap + 0) * 8 + 0], 1); + EXPECT_EQ(ring_host[(1 * kRingCap + 1) * 8 + 0], -1); // muted: untouched + EXPECT_EQ(ring_host[(2 * kRingCap + 2) * 8 + 0], 3); + + cudaFreeHost(ring_host); +} + } // namespace } // namespace test } // namespace onnxruntime From 1dc5b37a558f7af4d68ea509847f70b5a3e822b4 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 26 Oct 2026 20:39:20 +0000 Subject: [PATCH 2/2] lintrunner --- .../contrib_ops/speculative_fold_op_test.cc | 56 ++++++++++++++++--- 1 file changed, 49 insertions(+), 7 deletions(-) diff --git a/onnxruntime/test/contrib_ops/speculative_fold_op_test.cc b/onnxruntime/test/contrib_ops/speculative_fold_op_test.cc index 5d85277c28ceb..84aa79f71f1ed 100644 --- a/onnxruntime/test/contrib_ops/speculative_fold_op_test.cc +++ b/onnxruntime/test/contrib_ops/speculative_fold_op_test.cc @@ -258,14 +258,56 @@ TEST(SpeculativeFoldTest, SupportsBatchedSpeculativePrefixCommitS8) { // Batch 1: muted (ctl[3] = 1) // Batch 2: accepts 1 draft (tokens[0] == ids[1], tokens[1] != ids[2]) std::vector token_values = { - 10, 11, 12, 99, 14, 15, 16, 17, - 20, 21, 22, 23, 24, 25, 26, 27, - 30, 88, 32, 33, 34, 35, 36, 37, + 10, + 11, + 12, + 99, + 14, + 15, + 16, + 17, + 20, + 21, + 22, + 23, + 24, + 25, + 26, + 27, + 30, + 88, + 32, + 33, + 34, + 35, + 36, + 37, }; std::vector id_values = { - 1, 10, 11, 12, 13, 14, 15, 16, - 2, 20, 21, 22, 23, 24, 25, 26, - 3, 30, 31, 32, 33, 34, 35, 36, + 1, + 10, + 11, + 12, + 13, + 14, + 15, + 16, + 2, + 20, + 21, + 22, + 23, + 24, + 25, + 26, + 3, + 30, + 31, + 32, + 33, + 34, + 35, + 36, }; std::vector ctl_values(kBatch * 16, 0); // Batch 0: length=10, muted=0, seq=0, cache_len=10 @@ -316,7 +358,7 @@ TEST(SpeculativeFoldTest, SupportsBatchedSpeculativePrefixCommitS8) { EXPECT_EQ(out_ctl[0 * 16 + 1], 10 + 4); // length = 10 + 3 + 1 EXPECT_EQ(out_ctl[0 * 16 + 4], 1); // seq incremented to 1 - EXPECT_EQ(out_ctl[1 * 16 + 4], 5); // muted: seq unchanged + EXPECT_EQ(out_ctl[1 * 16 + 4], 5); // muted: seq unchanged EXPECT_EQ(out_ctl[2 * 16 + 0], 88); // pending = tokens[2, 1] EXPECT_EQ(out_ctl[2 * 16 + 1], 30 + 2); // length = 30 + 1 + 1