From a4441ce26ccd11a5e10244a4d426566b3682b9d1 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Mon, 19 Oct 2026 19:24:25 +0000 Subject: [PATCH 01/15] Add fused MTP MoE and speculative state fold CUDA operators --- cmake/onnxruntime_unittests.cmake | 4 + docs/ContribOperators.md | 150 ++++++++ .../contrib_ops/cuda/cuda_contrib_kernels.cc | 6 + onnxruntime/contrib_ops/cuda/llm/mtp_moe.cc | 91 +++++ onnxruntime/contrib_ops/cuda/llm/mtp_moe.cu | 294 +++++++++++++++ onnxruntime/contrib_ops/cuda/llm/mtp_moe.h | 36 ++ .../contrib_ops/cuda/llm/speculative_fold.cc | 131 +++++++ .../contrib_ops/cuda/llm/speculative_fold.cu | 157 ++++++++ .../contrib_ops/cuda/llm/speculative_fold.h | 34 ++ .../core/graph/contrib_ops/contrib_defs.cc | 79 ++++ onnxruntime/core/graph/contrib_ops/ms_opset.h | 6 + .../contrib_ops/cuda_kernels/mtp_moe_test.cc | 344 ++++++++++++++++++ .../cuda_kernels/speculative_fold_test.cc | 329 +++++++++++++++++ 13 files changed, 1661 insertions(+) create mode 100644 onnxruntime/contrib_ops/cuda/llm/mtp_moe.cc create mode 100644 onnxruntime/contrib_ops/cuda/llm/mtp_moe.cu create mode 100644 onnxruntime/contrib_ops/cuda/llm/mtp_moe.h create mode 100644 onnxruntime/contrib_ops/cuda/llm/speculative_fold.cc create mode 100644 onnxruntime/contrib_ops/cuda/llm/speculative_fold.cu create mode 100644 onnxruntime/contrib_ops/cuda/llm/speculative_fold.h create mode 100644 onnxruntime/test/contrib_ops/cuda_kernels/mtp_moe_test.cc create mode 100644 onnxruntime/test/contrib_ops/cuda_kernels/speculative_fold_test.cc diff --git a/cmake/onnxruntime_unittests.cmake b/cmake/onnxruntime_unittests.cmake index 4e630b5bb45d2..7c5d8de4dd80d 100644 --- a/cmake/onnxruntime_unittests.cmake +++ b/cmake/onnxruntime_unittests.cmake @@ -1125,6 +1125,8 @@ if (onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS AND onnxruntime_BUILD_CUDA_EP_AS_P "${TEST_SRC_DIR}/contrib_ops/cuda_kernels/qmoe_fp4_to_fp8_kernel_test.cc" "${TEST_SRC_DIR}/contrib_ops/cuda_kernels/qmoe_fp8_compaction_test.cc" "${TEST_SRC_DIR}/contrib_ops/cuda_kernels/softmax_topk_kernel_test.cc" + "${TEST_SRC_DIR}/contrib_ops/cuda_kernels/mtp_moe_test.cc" + "${TEST_SRC_DIR}/contrib_ops/cuda_kernels/speculative_fold_test.cc" ) endif() @@ -1157,7 +1159,9 @@ if (onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS AND onnxruntime_BUILD_CUDA_EP_AS_P "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/llm/gemm_tactic_cache.cc" "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/llm/fpA_intB_gemm_adaptor.cu" "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/llm/fpA_intB_gemm_preprocessors_impl.cu" + "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/llm/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/moe/qmoe_kernels.cu" ) if (onnxruntime_USE_FP4_QMOE) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 2703de7cc1fae..7963064068950 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -3,6 +3,9 @@ Do not modify directly.* * com.microsoft + * com.microsoft.MtpMoe + * com.microsoft.FoldMid2 + * com.microsoft.FoldTail * com.microsoft.Attention * com.microsoft.AttnLSTM * com.microsoft.BeamSearch @@ -8772,3 +8775,150 @@ No versioning maintained for experimental ops.
Constrain input and output types to float32 tensors.
+### **com.microsoft.MtpMoe** + + Fused bfloat16 NVFP4 mixture-of-experts projection for multi-token prediction. + +#### Version + +This version of the operator has been available since version 1 of the 'com.microsoft' operator set. + +#### Attributes + +
+
gate_up_global_scale : float (required)
+
Global scale applied to the gate/up projection weights.
+
down_global_scale : float (required)
+
Global scale applied to the down projection weights.
+
+ +#### Inputs + +
+
input : T
+
Bfloat16 hidden states with shape [1, rows, 2560].
+
router_logits : T
+
Bfloat16 router logits with shape [1, rows, 512].
+
gate_up_weights : T1
+
Packed NVFP4 gate/up weights with shape [512, 1280, 1440].
+
down_weights : T1
+
Packed NVFP4 down weights with shape [512, 2560, 360].
+
+ +#### Outputs + +
+
output : T
+
Bfloat16 output with the same shape as input.
+
scratch : T2
+
Float scratch buffer used by the fused expert projections.
+
+ +#### Type Constraints + +
+
T : tensor(bfloat16)
+
Constrain activations and output to bfloat16.
+
T1 : tensor(uint8)
+
Constrain packed expert weights to uint8.
+
T2 : tensor(float)
+
Constrain scratch to float.
+
+ +### **com.microsoft.FoldMid2** + + Combines speculative-token acceptance, recurrent-state commit bookkeeping, and head rotary-cache lookup for a two-token verification step. When `ctl[3]` is zero, the operator updates `ctl` and `ctx_commit` in place and writes a result record through the device-accessible pointer and capacity in `ptrs`. + +#### Version + +This version of the operator has been available since version 1 of the 'com.microsoft' operator set. + +#### Inputs + +
+
next_tokens : I64
+
Two predicted token IDs.
+
input_ids : I64
+
Two input token IDs; input_ids[1] is the draft token.
+
status : I32
+
Status values aggregated into the result record.
+
ctl : I64
+
Mutable eight-element control state.
+
ctx_commit : I64
+
Mutable two-element committed context.
+
ctx_a : I64
+
Context selected when the draft is rejected.
+
ctx_b : I64
+
Context selected when the draft is accepted.
+
ptrs : I64
+
Device pointers containing the result-ring address and capacity.
+
cos_cache : B
+
Bfloat16 rotary cosine cache with shape [positions, 32].
+
sin_cache : B
+
Bfloat16 rotary sine cache with shape [positions, 32].
+
+ +#### Outputs + +
+
head_ids : I64
+
Two token IDs consumed by the head.
+
head_c : I32
+
Current head-cache length.
+
slot_ids : I64
+
Two consecutive head-cache slots.
+
seqlens : I32
+
Head sequence length.
+
cos_out : B
+
Two rows of rotary cosine values.
+
sin_out : B
+
Two rows of rotary sine values.
+
sel : I64
+
Selected row: one when the draft is accepted, otherwise zero.
+
+ +#### Type Constraints + +
+
B : tensor(bfloat16)
+
Constrain rotary cache tensors to bfloat16.
+
I32 : tensor(int32)
+
Constrain status and sequence tensors to int32.
+
I64 : tensor(int64)
+
Constrain control and token tensors to int64.
+
+ +### **com.microsoft.FoldTail** + + Updates the next speculative draft token and aggregate error bit from the head result. When `ctl[3]` is zero, the operator updates the control tensor in place; muted calls return `done=0` without changing the state. + +#### Version + +This version of the operator has been available since version 1 of the 'com.microsoft' operator set. + +#### Inputs + +
+
next_token : I64
+
Next draft token ID.
+
status : I32
+
Head status values.
+
ctl : I64
+
Mutable eight-element control state.
+
+ +#### Outputs + +
+
done : I32
+
One when the control state was updated; zero for muted calls.
+
+ +#### Type Constraints + +
+
I32 : tensor(int32)
+
Constrain status and output tensors to int32.
+
I64 : tensor(int64)
+
Constrain token and control tensors to int64.
+
diff --git a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc index 546b5b12e178f..93b62d9fff569 100644 --- a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc @@ -100,6 +100,9 @@ class CUDA_ONNX_OP_TYPED_CLASS_NAME(1, MLFloat16, Crop); class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, MoE); class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, MoE); class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, MoE); +class CUDA_MS_OP_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_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); @@ -464,6 +467,9 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/onnxruntime/contrib_ops/cuda/llm/mtp_moe.cc b/onnxruntime/contrib_ops/cuda/llm/mtp_moe.cc new file mode 100644 index 0000000000000..167945464010e --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/llm/mtp_moe.cc @@ -0,0 +1,91 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "contrib_ops/cuda/llm/mtp_moe.h" + +#include "core/common/common.h" +#include "core/providers/cuda/cuda_kernel.h" + +#include + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +using namespace onnxruntime::cuda; + +class MtpMoe final : public CudaKernel { + public: + explicit MtpMoe(const OpKernelInfo& info) : CudaKernel(info) { + ORT_ENFORCE(info.GetAttr("gate_up_global_scale", &gate_up_global_scale_).IsOK(), + "MtpMoe requires gate_up_global_scale"); + ORT_ENFORCE(info.GetAttr("down_global_scale", &down_global_scale_).IsOK(), + "MtpMoe requires down_global_scale"); + } + + Status ComputeInternal(OpKernelContext* context) const override { + const Tensor* input = context->Input(0); + const Tensor* router = context->Input(1); + const Tensor* gate_up = context->Input(2); + const Tensor* down = context->Input(3); + ORT_ENFORCE(input != nullptr && router != nullptr && gate_up != nullptr && down != nullptr); + + const TensorShape& input_shape = input->Shape(); + const TensorShape& router_shape = router->Shape(); + const TensorShape& gate_up_shape = gate_up->Shape(); + const TensorShape& down_shape = down->Shape(); + ORT_RETURN_IF_NOT(input_shape.NumDimensions() == 3 && input_shape[0] == 1 && + input_shape[2] == mtp_moe::kHiddenSize, + "MtpMoe input must have shape [1, R, 2560], got ", input_shape); + const int64_t rows = input_shape[1]; + ORT_RETURN_IF_NOT(rows > 0 && rows <= std::numeric_limits::max() && + router_shape.NumDimensions() == 3 && + router_shape[0] == 1 && router_shape[1] == rows && + router_shape[2] == mtp_moe::kExpertCount, + "MtpMoe router logits must have shape [1, R, 512], got ", router_shape); + ORT_RETURN_IF_NOT(gate_up_shape.NumDimensions() == 3 && + gate_up_shape[0] == mtp_moe::kExpertCount && + gate_up_shape[1] == 2 * mtp_moe::kFfnSize && + gate_up_shape[2] == mtp_moe::kHiddenSize * 9 / 16, + "MtpMoe NVFP4 gate_up weights must have shape [512, 1280, 1440], got ", gate_up_shape); + ORT_RETURN_IF_NOT(down_shape.NumDimensions() == 3 && + down_shape[0] == mtp_moe::kExpertCount && + down_shape[1] == mtp_moe::kHiddenSize && + down_shape[2] == mtp_moe::kFfnSize * 9 / 16, + "MtpMoe NVFP4 down weights must have shape [512, 2560, 360], got ", down_shape); + + Tensor* output = context->Output(0, input_shape); + Tensor* scratch = context->Output( + 1, TensorShape({rows, static_cast(mtp_moe::kScratchStride)})); + ORT_ENFORCE(output != nullptr && scratch != nullptr); + + return mtp_moe::LaunchNvfp4MtpMoe( + reinterpret_cast(input->Data()), + reinterpret_cast(router->Data()), + gate_up->Data(), down->Data(), + gate_up_global_scale_, down_global_scale_, + scratch->MutableData(), + reinterpret_cast<__nv_bfloat16*>(output->MutableData()), + static_cast(rows), Stream(context)); + } + + private: + float gate_up_global_scale_; + float down_global_scale_; +}; + +ONNX_OPERATOR_TYPED_KERNEL_EX( + MtpMoe, + kMSDomain, + 1, + BFloat16, + kCudaExecutionProvider, + (*KernelDefBuilder::Create()) + .TypeConstraint("T", DataTypeImpl::GetTensorType()) + .TypeConstraint("T1", DataTypeImpl::GetTensorType()) + .TypeConstraint("T2", DataTypeImpl::GetTensorType()), + MtpMoe); + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/llm/mtp_moe.cu b/onnxruntime/contrib_ops/cuda/llm/mtp_moe.cu new file mode 100644 index 0000000000000..2bdb95a658717 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/llm/mtp_moe.cu @@ -0,0 +1,294 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "contrib_ops/cuda/llm/mtp_moe.h" + +#include + +#include + +namespace onnxruntime { +namespace contrib { +namespace cuda { +namespace mtp_moe { +namespace { + +constexpr int kChunk = 8; +constexpr int kGateUpRows = 2 * kFfnSize; +constexpr int kWarpsPerBlock = 8; +constexpr int kKChunks = kHiddenSize / kChunk; +constexpr int kPerLane = kKChunks / 32; +constexpr int kDownChunks = kFfnSize / kChunk; +constexpr int kDownPerLane = (kDownChunks + 31) / 32; + +__device__ __forceinline__ float RoundBf16(float value) { + return __bfloat162float(__float2bfloat16(value)); +} + +__device__ __forceinline__ float Sigmoid(float value) { + return 1.0f / (1.0f + expf(-value)); +} + +__device__ __forceinline__ void UnpackBf16x8(const uint4& packed, float* values) { + const uint32_t words[4] = {packed.x, packed.y, packed.z, packed.w}; +#pragma unroll + for (int i = 0; i < 4; ++i) { + values[2 * i] = __uint_as_float(words[i] << 16); + values[2 * i + 1] = __uint_as_float(words[i] & 0xffff0000u); + } +} + +__device__ __forceinline__ float DecodeE4M3(uint8_t bits) { + const __half_raw half_bits = __nv_cvt_fp8_to_halfraw(static_cast<__nv_fp8_storage_t>(bits), __NV_E4M3); + return __half2float(*reinterpret_cast(&half_bits)); +} + +__device__ __forceinline__ float DecodeE2M1(uint32_t code) { + const uint32_t magnitude = code & 7u; + const uint32_t exponent_mantissa = + magnitude == 0u ? 0u : ((magnitude == 1u ? 252u : (magnitude + 252u)) << 22); + return __uint_as_float(((code & 8u) << 28) | exponent_mantissa); +} + +__device__ __forceinline__ void UnpackNvfp4(uint32_t packed, uint8_t scale, float* values) { + const float decoded_scale = DecodeE4M3(scale); +#pragma unroll + for (int i = 0; i < 8; ++i) { + values[i] = DecodeE2M1((packed >> (4 * i)) & 0xfu) * decoded_scale; + } +} + +__device__ __forceinline__ void MergeTop(float& best, int& best_index, + float candidate, int candidate_index) { + if (candidate > best || (candidate == best && candidate_index < best_index)) { + best = candidate; + best_index = candidate_index; + } +} + +// Mirrors the reference custom op's FP32 lane-wise softmax and stable iterative top-10. +__global__ void RouterKernel(const __nv_bfloat16* __restrict__ logits, + float* __restrict__ scratch) { + const int row = blockIdx.x; + const int lane = threadIdx.x; + constexpr int kIterations = kExpertCount / 32; + float logits_per_lane[kIterations]; + float max_logit = -INFINITY; +#pragma unroll + for (int i = 0; i < kIterations; ++i) { + logits_per_lane[i] = __bfloat162float(logits[static_cast(row) * kExpertCount + i * 32 + lane]); + max_logit = fmaxf(max_logit, logits_per_lane[i]); + } + for (int offset = 16; offset > 0; offset >>= 1) { + max_logit = fmaxf(max_logit, __shfl_xor_sync(0xffffffffu, max_logit, offset)); + } + + float exponentials[kIterations]; + float denominator = 0.0f; +#pragma unroll + for (int i = 0; i < kIterations; ++i) { + exponentials[i] = expf(logits_per_lane[i] - max_logit); + denominator += exponentials[i]; + } + for (int offset = 16; offset > 0; offset >>= 1) { + denominator += __shfl_xor_sync(0xffffffffu, denominator, offset); + } + + float probabilities[kIterations]; +#pragma unroll + for (int i = 0; i < kIterations; ++i) { + probabilities[i] = exponentials[i] / denominator; + } + + float top_values[kTopK]; + int top_indices[kTopK]; +#pragma unroll + for (int k = 0; k < kTopK; ++k) { + float best = -1.0f; + int best_index = 0x7fffffff; +#pragma unroll + for (int i = 0; i < kIterations; ++i) { + if (probabilities[i] > best) { + best = probabilities[i]; + best_index = i * 32 + lane; + } + } + for (int offset = 16; offset > 0; offset >>= 1) { + const float other_value = __shfl_xor_sync(0xffffffffu, best, offset); + const int other_index = __shfl_xor_sync(0xffffffffu, best_index, offset); + MergeTop(best, best_index, other_value, other_index); + } + top_values[k] = best; + top_indices[k] = best_index; +#pragma unroll + for (int i = 0; i < kIterations; ++i) { + if (best_index == i * 32 + lane) { + probabilities[i] = -2.0f; + } + } + } + + if (lane == 0) { + float top_sum = 0.0f; +#pragma unroll + for (int k = 0; k < kTopK; ++k) { + top_sum += top_values[k]; + } + float* row_scratch = scratch + static_cast(row) * kScratchStride; +#pragma unroll + for (int k = 0; k < kTopK; ++k) { + row_scratch[k] = __int_as_float(top_indices[k]); + row_scratch[kTopK + k] = RoundBf16(top_values[k] / top_sum); + } + } +} + +__global__ void GateUpKernel(const __nv_bfloat16* __restrict__ input, + const uint8_t* __restrict__ gate_up, + float* __restrict__ scratch, + float global_scale) { + const int row = blockIdx.z; + const int slot = blockIdx.y; + const int warp = threadIdx.x >> 5; + const int lane = threadIdx.x & 31; + float* row_scratch = scratch + static_cast(row) * kScratchStride; + const int expert = __float_as_int(row_scratch[slot]); + const int output_row = blockIdx.x * kWarpsPerBlock + warp; + constexpr int kPackedRowBytes = kHiddenSize * 9 / 16; + const uint4* input_vector = reinterpret_cast(input + static_cast(row) * kHiddenSize); + const uint8_t* gate_row = gate_up + (static_cast(expert) * kGateUpRows + output_row) * kPackedRowBytes; + const uint8_t* up_row = + gate_up + (static_cast(expert) * kGateUpRows + kFfnSize + output_row) * kPackedRowBytes; + + uint32_t packed_gate[kPerLane]; + uint32_t packed_up[kPerLane]; + uint8_t gate_scales[kPerLane]; + uint8_t up_scales[kPerLane]; + uint4 packed_input[kPerLane]; +#pragma unroll + for (int i = 0; i < kPerLane; ++i) { + const int column_chunk = i * 32 + lane; + packed_gate[i] = __ldg(reinterpret_cast(gate_row) + column_chunk); + packed_up[i] = __ldg(reinterpret_cast(up_row) + column_chunk); + gate_scales[i] = __ldg(gate_row + kHiddenSize / 2 + (column_chunk >> 1)); + up_scales[i] = __ldg(up_row + kHiddenSize / 2 + (column_chunk >> 1)); + packed_input[i] = __ldg(input_vector + column_chunk); + } + + float gate_accumulator = 0.0f; + float up_accumulator = 0.0f; +#pragma unroll + for (int i = 0; i < kPerLane; ++i) { + float gate_values[8], up_values[8], input_values[8]; + UnpackNvfp4(packed_gate[i], gate_scales[i], gate_values); + UnpackNvfp4(packed_up[i], up_scales[i], up_values); + UnpackBf16x8(packed_input[i], input_values); +#pragma unroll + for (int k = 0; k < 8; ++k) { + gate_accumulator = fmaf(gate_values[k], input_values[k], gate_accumulator); + up_accumulator = fmaf(up_values[k], input_values[k], up_accumulator); + } + } + for (int offset = 16; offset > 0; offset >>= 1) { + gate_accumulator += __shfl_xor_sync(0xffffffffu, gate_accumulator, offset); + up_accumulator += __shfl_xor_sync(0xffffffffu, up_accumulator, offset); + } + + if (lane == 0) { + gate_accumulator *= global_scale; + up_accumulator *= global_scale; + const float gate = RoundBf16(gate_accumulator); + const float up = RoundBf16(up_accumulator); + const float silu = RoundBf16(gate * Sigmoid(gate)); + auto* hidden = reinterpret_cast<__nv_bfloat16*>(row_scratch + 32); + hidden[slot * kFfnSize + output_row] = __float2bfloat16(silu * up); + } +} + +__global__ void DownKernel(const uint8_t* __restrict__ down, + const float* __restrict__ scratch, + __nv_bfloat16* __restrict__ output, + float global_scale) { + const int row = blockIdx.y; + const int slot = threadIdx.x >> 5; + const int lane = threadIdx.x & 31; + const int output_column = blockIdx.x; + const float* row_scratch = scratch + static_cast(row) * kScratchStride; + const int expert = __float_as_int(row_scratch[slot]); + constexpr int kPackedRowBytes = kFfnSize * 9 / 16; + const uint8_t* weight_row = + down + (static_cast(expert) * kHiddenSize + output_column) * kPackedRowBytes; + const auto* hidden = reinterpret_cast(row_scratch + 32) + slot * kFfnSize; + const uint4* packed_hidden = reinterpret_cast(hidden); + float accumulator = 0.0f; +#pragma unroll + for (int chunk = 0; chunk < kDownPerLane; ++chunk) { + const int column_chunk = chunk * 32 + lane; + if (column_chunk < kDownChunks) { + const uint32_t packed_weight = __ldg(reinterpret_cast(weight_row) + column_chunk); + const uint8_t scale = __ldg(weight_row + kFfnSize / 2 + (column_chunk >> 1)); + const uint4 hidden_chunk = __ldg(packed_hidden + column_chunk); + float weights[8], activations[8]; + UnpackNvfp4(packed_weight, scale, weights); + UnpackBf16x8(hidden_chunk, activations); +#pragma unroll + for (int i = 0; i < 8; ++i) { + accumulator = fmaf(weights[i], activations[i], accumulator); + } + } + } + for (int offset = 16; offset > 0; offset >>= 1) { + accumulator += __shfl_xor_sync(0xffffffffu, accumulator, offset); + } + __shared__ float weighted[kTopK]; + if (lane == 0) { + accumulator *= global_scale; + weighted[slot] = RoundBf16(RoundBf16(accumulator) * row_scratch[kTopK + slot]); + } + __syncthreads(); + if (threadIdx.x == 0) { + float total = 0.0f; +#pragma unroll + for (int expert_slot = 0; expert_slot < kTopK; ++expert_slot) { + total += weighted[expert_slot]; + } + output[static_cast(row) * kHiddenSize + output_column] = __float2bfloat16(total); + } +} + +} // namespace + +Status LaunchRouter(const __nv_bfloat16* router_logits, float* scratch, int rows, cudaStream_t stream) { + if (rows <= 0) { + return Status(common::StatusCategory::ONNXRUNTIME, common::StatusCode::INVALID_ARGUMENT, + "MtpMoe requires at least one row"); + } + RouterKernel<<>>(router_logits, scratch); + const cudaError_t error = cudaGetLastError(); + return error == cudaSuccess + ? Status::OK() + : Status(common::StatusCategory::SYSTEM, static_cast(error), cudaGetErrorString(error)); +} + +Status LaunchNvfp4MtpMoe(const __nv_bfloat16* input, const __nv_bfloat16* router_logits, + const uint8_t* gate_up_weights, const uint8_t* down_weights, + float gate_up_global_scale, float down_global_scale, + float* scratch, __nv_bfloat16* output, int rows, cudaStream_t stream) { + if (rows <= 0) { + return Status(common::StatusCategory::ONNXRUNTIME, common::StatusCode::INVALID_ARGUMENT, + "MtpMoe requires at least one row"); + } + RouterKernel<<>>(router_logits, scratch); + GateUpKernel<<>>( + input, gate_up_weights, scratch, gate_up_global_scale); + DownKernel<<>>(down_weights, scratch, output, down_global_scale); + const cudaError_t error = cudaGetLastError(); + return error == cudaSuccess + ? Status::OK() + : Status(common::StatusCategory::SYSTEM, static_cast(error), cudaGetErrorString(error)); +} + +} // namespace mtp_moe +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/llm/mtp_moe.h b/onnxruntime/contrib_ops/cuda/llm/mtp_moe.h new file mode 100644 index 0000000000000..ecba2b3babc83 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/llm/mtp_moe.h @@ -0,0 +1,36 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include + +#include + +#include "core/common/status.h" + +namespace onnxruntime { +namespace contrib { +namespace cuda { +namespace mtp_moe { + +constexpr int kHiddenSize = 2560; +constexpr int kExpertCount = 512; +constexpr int kTopK = 10; +constexpr int kFfnSize = 640; +constexpr int kScratchStride = 32 + kTopK * kFfnSize; + +Status LaunchRouter(const __nv_bfloat16* router_logits, float* scratch, + int rows, cudaStream_t stream); + +Status LaunchNvfp4MtpMoe(const __nv_bfloat16* input, const __nv_bfloat16* router_logits, + const uint8_t* gate_up_weights, const uint8_t* down_weights, + float gate_up_global_scale, float down_global_scale, + float* scratch, __nv_bfloat16* output, + int rows, cudaStream_t stream); + +} // namespace mtp_moe +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cc b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cc new file mode 100644 index 0000000000000..6965c9d1539c1 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cc @@ -0,0 +1,131 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// 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" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +using namespace onnxruntime::cuda; + +namespace { + +bool HasElementCount(const Tensor* tensor, int64_t count) { + return tensor != nullptr && tensor->Shape().Size() == count; +} + +class FoldMid2 final : public CudaKernel { + public: + explicit FoldMid2(const OpKernelInfo& info) : CudaKernel(info) {} + + 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})); + ORT_ENFORCE(head_ids != nullptr && head_c != nullptr && slot_ids != nullptr && + seqlens != nullptr && cos_out != nullptr && sin_out != nullptr && + sel != nullptr); + + return speculative_fold::LaunchFoldMid2( + next_tokens->Data(), input_ids->Data(), + status->Data(), static_cast(status_count), + const_cast(ctl->Data()), + const_cast(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(), Stream(context)); + } +}; + +class FoldTail final : public CudaKernel { + public: + explicit FoldTail(const OpKernelInfo& info) : CudaKernel(info) {} + + 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})); + ORT_ENFORCE(done != nullptr); + + return speculative_fold::LaunchFoldTail( + next_token->Data(), status->Data(), + static_cast(status_count), const_cast(ctl->Data()), + done->MutableData(), Stream(context)); + } +}; + +} // namespace + +ONNX_OPERATOR_KERNEL_EX( + FoldMid2, kMSDomain, 1, kCudaExecutionProvider, + (*KernelDefBuilder::Create()) + .TypeConstraint("I64", DataTypeImpl::GetTensorType()) + .TypeConstraint("I32", DataTypeImpl::GetTensorType()) + .TypeConstraint("B", DataTypeImpl::GetTensorType()), + FoldMid2); + +ONNX_OPERATOR_KERNEL_EX( + FoldTail, kMSDomain, 1, kCudaExecutionProvider, + (*KernelDefBuilder::Create()) + .TypeConstraint("I64", DataTypeImpl::GetTensorType()) + .TypeConstraint("I32", DataTypeImpl::GetTensorType()), + FoldTail); + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cu b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cu new file mode 100644 index 0000000000000..4e034f204bd02 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cu @@ -0,0 +1,157 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "contrib_ops/cuda/llm/speculative_fold.h" + +namespace onnxruntime { +namespace contrib { +namespace cuda { +namespace speculative_fold { +namespace { + +constexpr int64_t kHeadCapacity = 33028; + +__global__ void FoldMid2Kernel( + const int64_t* next_tokens, const int64_t* input_ids, + const int32_t* status, int status_count, int64_t* ctl, + 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) { + __shared__ int64_t committed_length; + __shared__ int64_t cache_length; + if (threadIdx.x == 0) { + committed_length = ctl[1]; + cache_length = ctl[7]; + } + __syncthreads(); + + if (threadIdx.x >= 32) { + const int j = threadIdx.x - 32; + if (j < 64) { + 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]; + } + return; + } + if (threadIdx.x != 0) { + 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]); + } + + 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; + + if (ctl[3] != 0) { + return; + } + + 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) | + (cache_length + 2 > kHeadCapacity ? 4 : 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, + int64_t* ctl, int32_t* done) { + if (threadIdx.x != 0) { + return; + } + 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; +} + +Status LaunchStatus(cudaError_t 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, int64_t* ctl, + 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, cudaStream_t stream) { + if (status_count < 0) { + return Status(common::StatusCategory::ONNXRUNTIME, + common::StatusCode::INVALID_ARGUMENT, + "FoldMid2 status count must not be negative"); + } + 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); + return LaunchStatus(cudaGetLastError()); +} + +Status LaunchFoldTail(const int64_t* next_token, const int32_t* status, + int status_count, int64_t* ctl, int32_t* done, + 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); + return LaunchStatus(cudaGetLastError()); +} + +} // namespace speculative_fold +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/llm/speculative_fold.h b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.h new file mode 100644 index 0000000000000..3f376bc08eca9 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.h @@ -0,0 +1,34 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include + +#include + +#include "core/common/status.h" + +namespace onnxruntime { +namespace contrib { +namespace cuda { +namespace speculative_fold { + +Status LaunchFoldMid2(const int64_t* next_tokens, const int64_t* input_ids, + const int32_t* status, int status_count, int64_t* ctl, + 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, cudaStream_t stream); + +Status LaunchFoldTail(const int64_t* next_token, const int32_t* status, + int status_count, int64_t* ctl, int32_t* done, + cudaStream_t stream); + +} // namespace speculative_fold +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc index c38c4d1e1936e..0271641588b8d 100644 --- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc @@ -1465,6 +1465,85 @@ ONNX_MS_OPERATOR_SET_SCHEMA( .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain input and output types to float tensors.") .TypeAndShapeInferenceFunction(ONNX_NAMESPACE::propagateShapeAndTypeFromFirstInput)); +constexpr const char* MtpMoe_ver1_doc = R"DOC( +Mixture-of-experts projection for multi-token prediction. The input and router logits are bfloat16. +Expert weights are packed NVFP4 bytes with shape [512, 1280, 1440] for gate/up and [512, 2560, 360] +for down projection. The output has the same shape as input. The second output is a float scratch buffer. +)DOC"; + +ONNX_MS_OPERATOR_SET_SCHEMA( + MtpMoe, 1, + OpSchema() + .SetDoc(MtpMoe_ver1_doc) + .Attr("gate_up_global_scale", "Global scale applied to the gate/up projection weights.", + AttributeProto::FLOAT, true) + .Attr("down_global_scale", "Global scale applied to the down projection weights.", + AttributeProto::FLOAT, true) + .Input(0, "input", "BFloat16 hidden states with shape [1, rows, 2560].", "T") + .Input(1, "router_logits", "BFloat16 router logits with shape [1, rows, 512].", "T") + .Input(2, "gate_up_weights", "Packed NVFP4 gate/up weights with shape [512, 1280, 1440].", "T1") + .Input(3, "down_weights", "Packed NVFP4 down weights with shape [512, 2560, 360].", "T1") + .Output(0, "output", "BFloat16 output with the same shape as input.", "T") + .Output(1, "scratch", "Float scratch buffer used by the fused expert projections.", "T2") + .TypeConstraint("T", {"tensor(bfloat16)"}, "Constrain activations and output to bfloat16.") + .TypeConstraint("T1", {"tensor(uint8)"}, "Constrain packed expert weights to uint8.") + .TypeConstraint("T2", {"tensor(float)"}, "Constrain scratch to float.") + .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 0, 0); + ONNX_NAMESPACE::propagateShapeFromInputToOutput(ctx, 0, 0); + if (auto* scratch = ctx.getOutputType(1)) { + scratch->mutable_tensor_type()->set_elem_type( + ONNX_NAMESPACE::TensorProto::FLOAT); + } + })); + +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. The control tensor and context-commit tensor are updated in place when ctl[3] is +zero. The pointer tensor contains the device-accessible result-ring address and its record capacity. +)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", "Mutable eight-element control state.", "I64") + .Input(4, "ctx_commit", "Mutable 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 the draft is accepted, otherwise zero.", "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.")); + +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 updated in place when ctl[3] is zero; muted calls return done=0 without changing 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", "Mutable eight-element control state.", "I64") + .Output(0, "done", "One when the control state was updated; zero for muted calls.", "I32") + .TypeConstraint("I64", {"tensor(int64)"}, "Constrain token and control tensors to int64.") + .TypeConstraint("I32", {"tensor(int32)"}, "Constrain status and output tensors to int32.")); + 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 d68cdd392666b..71b3c719f9906 100644 --- a/onnxruntime/core/graph/contrib_ops/ms_opset.h +++ b/onnxruntime/core/graph/contrib_ops/ms_opset.h @@ -84,6 +84,9 @@ class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, MatMulFpQ4); #endif class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, MaxpoolWithMask); 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, 1, QMoE); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, MultiHeadAttention); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, GroupQueryAttention); @@ -218,6 +221,9 @@ class OpSet_Microsoft_ver1 { #endif 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/mtp_moe_test.cc b/onnxruntime/test/contrib_ops/cuda_kernels/mtp_moe_test.cc new file mode 100644 index 0000000000000..b9f6f91a75e4b --- /dev/null +++ b/onnxruntime/test/contrib_ops/cuda_kernels/mtp_moe_test.cc @@ -0,0 +1,344 @@ +// 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/mtp_moe.h" + +namespace onnxruntime { +namespace test { +namespace { + +namespace mtp_moe = onnxruntime::contrib::cuda::mtp_moe; + +constexpr float kGateUpGlobalScale = 0.0005173456156626344f; +constexpr float kDownGlobalScale = 0.00031534829759038985f; + +void CheckCuda(cudaError_t error) { + if (error != cudaSuccess) { + throw std::runtime_error(cudaGetErrorString(error)); + } +} + +class DeviceBuffer { + public: + explicit DeviceBuffer(size_t size) { + CheckCuda(cudaMalloc(&data_, size)); + } + ~DeviceBuffer() { + if (data_ != nullptr) { + cudaFree(data_); + } + } + DeviceBuffer(const DeviceBuffer&) = delete; + DeviceBuffer& operator=(const DeviceBuffer&) = delete; + + template + T* As() const { + return reinterpret_cast(data_); + } + 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)); + } + void UploadAt(size_t offset, const void* source, size_t size) { + CheckCuda(cudaMemcpy(static_cast(data_) + offset, source, size, cudaMemcpyHostToDevice)); + } + + private: + void* data_ = nullptr; +}; + +float RoundBf16(float value) { + return __bfloat162float(__float2bfloat16(value)); +} + +float DecodeFp4(uint32_t code) { + constexpr std::array kMagnitude = {0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f}; + const float magnitude = kMagnitude[code & 7u]; + return (code & 8u) != 0 ? -magnitude : magnitude; +} + +float DecodeE4M3Scale(uint8_t encoded) { + // The test's synthetic block scales use only exactly representable 1.0 and 0.5 codes. + EXPECT_TRUE(encoded == 0x38 || encoded == 0x30); + return encoded == 0x38 ? 1.0f : 0.5f; +} + +float DecodePackedWeight(const uint8_t* row, int k, int column) { + const uint8_t packed = row[column >> 1]; + const uint32_t nibble = (column & 1) == 0 ? (packed & 0xfu) : (packed >> 4); + const int scale_index = k / 2 + column / 16; + return DecodeFp4(nibble) * DecodeE4M3Scale(row[scale_index]); +} + +float WarpReduceXor(std::array values) { + for (int offset = 16; offset > 0; offset >>= 1) { + const auto previous = values; + for (int lane = 0; lane < 32; ++lane) { + values[lane] = previous[lane] + previous[lane ^ offset]; + } + } + return values[0]; +} + +std::vector MakeNvfp4Weights(int experts, int rows, int k) { + const int row_bytes = k * 9 / 16; + const int packed_bytes = k / 2; + const int scale_bytes = k / 16; + std::vector weights(static_cast(experts) * rows * row_bytes); + for (int expert = 0; expert < experts; ++expert) { + for (int row = 0; row < rows; ++row) { + uint8_t* weight_row = weights.data() + + (static_cast(expert) * rows + row) * row_bytes; + for (int column = 0; column < k; ++column) { + const uint8_t code = static_cast((expert * 13 + row * 7 + column * 3 + column / 17) & 0xf); + uint8_t& packed = weight_row[column >> 1]; + packed = (column & 1) == 0 + ? static_cast((packed & 0xf0u) | code) + : static_cast((packed & 0x0fu) | (code << 4)); + } + for (int group = 0; group < scale_bytes; ++group) { + weight_row[packed_bytes + group] = ((expert + row + group) % 3 == 0) ? 0x30 : 0x38; + } + } + } + return weights; +} + +float ReferenceNvfp4Dot(const uint8_t* weight_row, const __nv_bfloat16* input, int k) { + const int chunks = k / 8; + std::array lane_sum{}; + for (int lane = 0; lane < 32; ++lane) { + float accumulator = 0.0f; + for (int chunk = lane; chunk < chunks; chunk += 32) { + for (int element = 0; element < 8; ++element) { + const int column = chunk * 8 + element; + const float weight = DecodePackedWeight(weight_row, k, column); + const float activation = __bfloat162float(input[column]); + accumulator = std::fma(weight, activation, accumulator); + } + } + lane_sum[lane] = accumulator; + } + return WarpReduceXor(lane_sum); +} + +struct ReferenceResult { + std::vector<__nv_bfloat16> output; + std::vector<__nv_bfloat16> hidden; + std::array experts{}; + std::array routing_weights{}; +}; + +ReferenceResult ReferenceMtpMoe(const std::vector<__nv_bfloat16>& input, + const std::vector<__nv_bfloat16>& router, + const std::vector& gate_up, + const std::vector& down) { + constexpr int hidden = mtp_moe::kHiddenSize; + constexpr int ffn = mtp_moe::kFfnSize; + constexpr int gate_up_rows = 2 * ffn; + constexpr int gate_row_bytes = hidden * 9 / 16; + constexpr int down_row_bytes = ffn * 9 / 16; + + ReferenceResult result; + std::array probabilities{}; + float maximum = -std::numeric_limits::infinity(); + for (int expert = 0; expert < mtp_moe::kExpertCount; ++expert) { + probabilities[expert] = __bfloat162float(router[expert]); + maximum = std::max(maximum, probabilities[expert]); + } + float denominator = 0.0f; + for (float& probability : probabilities) { + probability = std::exp(probability - maximum); + denominator += probability; + } + for (float& probability : probabilities) { + probability /= denominator; + } + std::array selected{}; + for (int slot = 0; slot < mtp_moe::kTopK; ++slot) { + int best = -1; + for (int expert = 0; expert < mtp_moe::kExpertCount; ++expert) { + if (!selected[expert] && + (best < 0 || probabilities[expert] > probabilities[best] || + (probabilities[expert] == probabilities[best] && expert < best))) { + best = expert; + } + } + selected[best] = true; + result.experts[slot] = best; + } + float top_sum = 0.0f; + for (int slot = 0; slot < mtp_moe::kTopK; ++slot) { + top_sum += probabilities[result.experts[slot]]; + } + for (int slot = 0; slot < mtp_moe::kTopK; ++slot) { + result.routing_weights[slot] = RoundBf16(probabilities[result.experts[slot]] / top_sum); + } + + result.hidden.resize(mtp_moe::kTopK * ffn); + for (int slot = 0; slot < mtp_moe::kTopK; ++slot) { + const int expert = result.experts[slot]; + const size_t expert_offset = static_cast(expert) * gate_up_rows * gate_row_bytes; + for (int output_row = 0; output_row < ffn; ++output_row) { + const uint8_t* gate_row = gate_up.data() + expert_offset + output_row * gate_row_bytes; + const uint8_t* up_row = gate_up.data() + expert_offset + (ffn + output_row) * gate_row_bytes; + const float gate = RoundBf16(ReferenceNvfp4Dot(gate_row, input.data(), hidden) * kGateUpGlobalScale); + const float up = RoundBf16(ReferenceNvfp4Dot(up_row, input.data(), hidden) * kGateUpGlobalScale); + const float silu = RoundBf16(gate / (1.0f + std::exp(-gate))); + result.hidden[slot * ffn + output_row] = __float2bfloat16(silu * up); + } + } + + result.output.resize(hidden); + for (int output_column = 0; output_column < hidden; ++output_column) { + float total = 0.0f; + for (int slot = 0; slot < mtp_moe::kTopK; ++slot) { + const int expert = result.experts[slot]; + const size_t expert_offset = static_cast(expert) * hidden * down_row_bytes; + const uint8_t* weight_row = down.data() + expert_offset + output_column * down_row_bytes; + const float down_sum = RoundBf16(ReferenceNvfp4Dot( + weight_row, result.hidden.data() + slot * ffn, ffn) * + kDownGlobalScale); + const float weighted = RoundBf16(down_sum * result.routing_weights[slot]); + total += weighted; + } + result.output[output_column] = __float2bfloat16(total); + } + return result; +} + +TEST(CUDA_EP_Unittest, MtpMoeRouterSelectsStableTopTenAndBf16Normalizes) { + constexpr int rows = 1; + std::vector<__nv_bfloat16> logits(mtp_moe::kExpertCount, __float2bfloat16(-8.0f)); + for (int expert = 0; expert < mtp_moe::kTopK; ++expert) { + logits[expert] = __float2bfloat16(0.0f); + } + + DeviceBuffer device_logits(logits.size() * sizeof(__nv_bfloat16)); + DeviceBuffer device_scratch(rows * mtp_moe::kScratchStride * sizeof(float)); + device_logits.Upload(logits.data(), logits.size() * sizeof(__nv_bfloat16)); + + cudaStream_t stream = nullptr; + CheckCuda(cudaStreamCreate(&stream)); + const Status status = mtp_moe::LaunchRouter(device_logits.As<__nv_bfloat16>(), + device_scratch.As(), rows, stream); + ASSERT_TRUE(status.IsOK()) << status.ErrorMessage(); + CheckCuda(cudaStreamSynchronize(stream)); + CheckCuda(cudaStreamDestroy(stream)); + + std::array scratch{}; + device_scratch.Download(scratch.data(), sizeof(scratch)); + float weight_sum = 0.0f; + for (int slot = 0; slot < mtp_moe::kTopK; ++slot) { + int expert = -1; + std::memcpy(&expert, &scratch[slot], sizeof(expert)); + EXPECT_EQ(expert, slot); + EXPECT_NEAR(scratch[mtp_moe::kTopK + slot], + __bfloat162float(__float2bfloat16(0.1f)), 0.0f); + weight_sum += scratch[mtp_moe::kTopK + slot]; + } + EXPECT_NEAR(weight_sum, 1.0f, 0.004f); +} + +TEST(CUDA_EP_Unittest, MtpMoeNvfp4MatchesReferenceArithmeticEndToEnd) { + constexpr int rows = 1; + constexpr int experts = mtp_moe::kTopK; + constexpr int hidden = mtp_moe::kHiddenSize; + constexpr int ffn = mtp_moe::kFfnSize; + constexpr int gate_row_bytes = hidden * 9 / 16; + constexpr int down_row_bytes = ffn * 9 / 16; + const size_t full_gate_bytes = static_cast(mtp_moe::kExpertCount) * 2 * ffn * gate_row_bytes; + const size_t full_down_bytes = static_cast(mtp_moe::kExpertCount) * hidden * down_row_bytes; + + size_t free_bytes = 0; + size_t total_bytes = 0; + CheckCuda(cudaMemGetInfo(&free_bytes, &total_bytes)); + constexpr size_t kMemoryReserve = size_t{128} << 20; + if (free_bytes < full_gate_bytes + full_down_bytes + kMemoryReserve) { + GTEST_SKIP() << "Insufficient CUDA memory for exact-shape MtpMoe weights"; + } + + std::vector<__nv_bfloat16> input(hidden); + for (int column = 0; column < hidden; ++column) { + input[column] = __float2bfloat16(static_cast((column % 23) - 11) * 6.0f); + } + std::vector<__nv_bfloat16> router(mtp_moe::kExpertCount, __float2bfloat16(-8.0f)); + for (int expert = 0; expert < experts; ++expert) { + router[expert] = __float2bfloat16(0.0f); + } + const std::vector gate_up = MakeNvfp4Weights(experts, 2 * ffn, hidden); + const std::vector down = MakeNvfp4Weights(experts, hidden, ffn); + const ReferenceResult expected = ReferenceMtpMoe(input, router, gate_up, down); + + DeviceBuffer device_input(input.size() * sizeof(__nv_bfloat16)); + DeviceBuffer device_router(router.size() * sizeof(__nv_bfloat16)); + DeviceBuffer device_gate_up(full_gate_bytes); + DeviceBuffer device_down(full_down_bytes); + DeviceBuffer device_output(hidden * sizeof(__nv_bfloat16)); + DeviceBuffer device_scratch(mtp_moe::kScratchStride * sizeof(float)); + device_input.Upload(input.data(), input.size() * sizeof(__nv_bfloat16)); + device_router.Upload(router.data(), router.size() * sizeof(__nv_bfloat16)); + for (int expert = 0; expert < experts; ++expert) { + const size_t gate_expert_bytes = static_cast(2 * ffn) * gate_row_bytes; + const size_t down_expert_bytes = static_cast(hidden) * down_row_bytes; + device_gate_up.UploadAt(static_cast(expert) * gate_expert_bytes, + gate_up.data() + static_cast(expert) * gate_expert_bytes, + gate_expert_bytes); + device_down.UploadAt(static_cast(expert) * down_expert_bytes, + down.data() + static_cast(expert) * down_expert_bytes, + down_expert_bytes); + } + + const Status status = mtp_moe::LaunchNvfp4MtpMoe( + device_input.As<__nv_bfloat16>(), device_router.As<__nv_bfloat16>(), + device_gate_up.As(), device_down.As(), + kGateUpGlobalScale, kDownGlobalScale, device_scratch.As(), + device_output.As<__nv_bfloat16>(), rows, nullptr); + ASSERT_TRUE(status.IsOK()) << status.ErrorMessage(); + + std::vector<__nv_bfloat16> actual_output(hidden); + std::vector actual_scratch(mtp_moe::kScratchStride * sizeof(float)); + device_output.Download(actual_output.data(), actual_output.size() * sizeof(__nv_bfloat16)); + device_scratch.Download(actual_scratch.data(), actual_scratch.size()); + + for (int slot = 0; slot < mtp_moe::kTopK; ++slot) { + int actual_expert = -1; + std::memcpy(&actual_expert, actual_scratch.data() + slot * sizeof(float), sizeof(actual_expert)); + EXPECT_EQ(actual_expert, expected.experts[slot]) << "routing slot " << slot; + float actual_weight = 0.0f; + std::memcpy(&actual_weight, + actual_scratch.data() + (mtp_moe::kTopK + slot) * sizeof(float), + sizeof(actual_weight)); + EXPECT_EQ(actual_weight, expected.routing_weights[slot]) << "routing weight " << slot; + + const auto* actual_hidden = reinterpret_cast( + actual_scratch.data() + 32 * sizeof(float)); + for (int column = 0; column < ffn; ++column) { + EXPECT_EQ(actual_hidden[slot * ffn + column], expected.hidden[slot * ffn + column]) + << "hidden intermediate slot=" << slot << " column=" << column; + } + } + for (int column = 0; column < hidden; ++column) { + const float actual = __bfloat162float(actual_output[column]); + const float reference = __bfloat162float(expected.output[column]); + EXPECT_NEAR(actual, reference, 0.01f + 0.01f * std::abs(reference)) << "output column " << column; + } +} + +} // namespace +} // namespace test +} // namespace onnxruntime diff --git a/onnxruntime/test/contrib_ops/cuda_kernels/speculative_fold_test.cc b/onnxruntime/test/contrib_ops/cuda_kernels/speculative_fold_test.cc new file mode 100644 index 0000000000000..d7d7d4d4802f4 --- /dev/null +++ b/onnxruntime/test/contrib_ops/cuda_kernels/speculative_fold_test.cc @@ -0,0 +1,329 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include +#include + +#include +#include +#include +#include +#include + +#include "contrib_ops/cuda/llm/speculative_fold.h" + +namespace onnxruntime { +namespace test { +namespace { + +void CheckCuda(cudaError_t error) { + if (error != cudaSuccess) { + throw std::runtime_error(cudaGetErrorString(error)); + } +} + +class DeviceBuffer { + public: + explicit DeviceBuffer(size_t size) { + CheckCuda(cudaMalloc(&data_, size)); + } + ~DeviceBuffer() { + if (data_ != nullptr) { + cudaFree(data_); + } + } + DeviceBuffer(const DeviceBuffer&) = delete; + DeviceBuffer& operator=(const DeviceBuffer&) = delete; + + template + T* As() const { + return reinterpret_cast(data_); + } + 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; +}; + +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_); + } + } + MappedRing(const MappedRing&) = delete; + MappedRing& operator=(const MappedRing&) = delete; + + int64_t* DeviceData() const { return device_data_; } + int64_t* HostData() const { return host_data_; } + + private: + size_t element_count_; + int64_t* host_data_ = nullptr; + int64_t* device_data_ = nullptr; +}; + +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)); + } + } + 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(), 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); + + 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])); + } + } + + 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(), 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(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(), 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(), 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); +} + +} // namespace +} // namespace test +} // namespace onnxruntime From e5902d17d1d9f79791925447a9f1c6144a0ee9ae Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Mon, 19 Oct 2026 20:28:28 +0000 Subject: [PATCH 02/15] Fix MTP scratch initialization and explicit fold state propagation Keep router expert indices valid for non-finite logits, initialize the public scratch buffer, return fold state through declared reusable outputs, and guard rotary and ring metadata before committing. Add session-level and kernel regressions and regenerate operator documentation. --- docs/ContribOperators.md | 330 ++++++++++-------- docs/OperatorKernels.md | 3 + onnxruntime/contrib_ops/cuda/llm/mtp_moe.cu | 25 +- .../contrib_ops/cuda/llm/speculative_fold.cc | 20 +- .../contrib_ops/cuda/llm/speculative_fold.cu | 67 +++- .../contrib_ops/cuda/llm/speculative_fold.h | 9 +- .../core/graph/contrib_ops/contrib_defs.cc | 42 ++- .../contrib_ops/cuda_kernels/mtp_moe_test.cc | 34 ++ .../cuda_kernels/speculative_fold_test.cc | 8 +- .../contrib_ops/speculative_fold_op_test.cc | 233 +++++++++++++ 10 files changed, 574 insertions(+), 197 deletions(-) create mode 100644 onnxruntime/test/contrib_ops/speculative_fold_op_test.cc diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 7963064068950..f18bcc0d95cd7 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -3,9 +3,6 @@ Do not modify directly.* * com.microsoft - * com.microsoft.MtpMoe - * com.microsoft.FoldMid2 - * com.microsoft.FoldTail * com.microsoft.Attention * com.microsoft.AttnLSTM * com.microsoft.BeamSearch @@ -39,6 +36,8 @@ Do not modify directly.* * com.microsoft.EngramGate * com.microsoft.ExpandDims * com.microsoft.FastGelu + * com.microsoft.FoldMid2 + * com.microsoft.FoldTail * com.microsoft.FusedConv * com.microsoft.FusedGemm * com.microsoft.FusedMatMul @@ -62,6 +61,7 @@ Do not modify directly.* * com.microsoft.HyperConnectionPreMix * com.microsoft.Inverse * com.microsoft.Irfft + * experimental com.microsoft.IsAllFinite * com.microsoft.LinearAttention * com.microsoft.LinearAttentionGate * com.microsoft.LongformerAttention @@ -77,6 +77,7 @@ Do not modify directly.* * com.microsoft.MatMulNBitsQkv * com.microsoft.MaxpoolWithMask * com.microsoft.MoE + * com.microsoft.MtpMoe * com.microsoft.MulInteger * com.microsoft.MultiHeadAttention * com.microsoft.MurmurHash3 @@ -92,6 +93,7 @@ Do not modify directly.* * com.microsoft.Pad * com.microsoft.PagedAttention * com.microsoft.QAttention + * experimental com.microsoft.QEmbedLayerNormalization * com.microsoft.QGemm * com.microsoft.QLinearAdd * com.microsoft.QLinearAveragePool @@ -142,9 +144,6 @@ Do not modify directly.* * com.microsoft.VarlenNGramHashMapping * com.microsoft.WhisperBeamSearch * com.microsoft.WordConvEmbedding - * experimental com.microsoft.IsAllFinite - * experimental com.microsoft.QEmbedLayerNormalization - ## com.microsoft ### **com.microsoft.Attention** @@ -2103,6 +2102,122 @@ 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. + +#### Version + +This version of the operator has been available since version 1 of the 'com.microsoft' operator set. + +#### Inputs + +
+
next_tokens : I64
+
Two predicted token IDs.
+
input_ids : I64
+
Two input token IDs; input_ids[1] is the draft token.
+
status : I32
+
Status values aggregated into the result record.
+
ctl : I64
+
Eight-element control state.
+
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.
+
ptrs : I64
+
Device pointers containing the result-ring address and capacity.
+
cos_cache : B
+
BFloat16 rotary cosine cache with shape [positions, 32].
+
sin_cache : B
+
BFloat16 rotary sine cache with shape [positions, 32].
+
+ +#### Outputs + +
+
head_ids : I64
+
Two token IDs consumed by the head.
+
head_c : I32
+
Current head-cache length.
+
slot_ids : I64
+
Two consecutive head-cache slots.
+
seqlens : I32
+
Head sequence length.
+
cos_out : B
+
Two rows of rotary cosine values.
+
sin_out : B
+
Two rows of rotary sine values.
+
sel : I64
+
Selected row: one when accepted, zero when rejected, or -1 on invalid metadata.
+
updated_ctl : I64
+
Updated control state with the same shape as ctl.
+
updated_ctx_commit : I64
+
Updated committed context with the same shape as ctx_commit.
+
+ +#### Type Constraints + +
+
I64 : tensor(int64)
+
Constrain control and token tensors to int64.
+
I32 : tensor(int32)
+
Constrain status and sequence tensors to int32.
+
B : tensor(bfloat16)
+
Constrain rotary cache tensors to bfloat16.
+
+ + +### **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. + +#### Version + +This version of the operator has been available since version 1 of the 'com.microsoft' operator set. + +#### Inputs + +
+
next_token : I64
+
Next draft token ID.
+
status : I32
+
Head status values.
+
ctl : I64
+
Eight-element control state.
+
+ +#### Outputs + +
+
done : I32
+
One when the control state was updated; zero for muted calls.
+
updated_ctl : I64
+
Updated control state with the same shape as ctl.
+
+ +#### Type Constraints + +
+
I64 : tensor(int64)
+
Constrain token and control tensors to int64.
+
I32 : tensor(int32)
+
Constrain status and output tensors to int32.
+
+ + ### **com.microsoft.FusedConv** The fused convolution operator schema is the same as Conv besides it includes an attribute @@ -4503,6 +4618,61 @@ This version of the operator has been available since version 1 of the 'com.micr +### **com.microsoft.MtpMoe** + + Mixture-of-experts projection for multi-token prediction. The input and router logits are bfloat16. + Expert weights are packed NVFP4 bytes with shape [512, 1280, 1440] for gate/up and [512, 2560, 360] + for down projection. The output has the same shape as input. The second output is a fully initialized float + scratch buffer. A row containing non-finite router logits selects valid expert indices 0 through 9 and + reports invalid routing through NaN routing weights and output values. + +#### Version + +This version of the operator has been available since version 1 of the 'com.microsoft' operator set. + +#### Attributes + +
+
down_global_scale : float (required)
+
Global scale applied to the down projection weights.
+
gate_up_global_scale : float (required)
+
Global scale applied to the gate/up projection weights.
+
+ +#### Inputs + +
+
input : T
+
BFloat16 hidden states with shape [1, rows, 2560].
+
router_logits : T
+
BFloat16 router logits with shape [1, rows, 512].
+
gate_up_weights : T1
+
Packed NVFP4 gate/up weights with shape [512, 1280, 1440].
+
down_weights : T1
+
Packed NVFP4 down weights with shape [512, 2560, 360].
+
+ +#### Outputs + +
+
output : T
+
BFloat16 output with the same shape as input.
+
scratch : T2
+
Float scratch buffer used by the fused expert projections.
+
+ +#### Type Constraints + +
+
T : tensor(bfloat16)
+
Constrain activations and output to bfloat16.
+
T1 : tensor(uint8)
+
Constrain packed expert weights to uint8.
+
T2 : tensor(float)
+
Constrain scratch to float.
+
+ + ### **com.microsoft.MulInteger** Performs element-wise binary quantized multiplication (with Numpy-style broadcasting support). @@ -8774,151 +8944,3 @@ No versioning maintained for experimental ops.
T : tensor(float)
Constrain input and output types to float32 tensors.
- -### **com.microsoft.MtpMoe** - - Fused bfloat16 NVFP4 mixture-of-experts projection for multi-token prediction. - -#### Version - -This version of the operator has been available since version 1 of the 'com.microsoft' operator set. - -#### Attributes - -
-
gate_up_global_scale : float (required)
-
Global scale applied to the gate/up projection weights.
-
down_global_scale : float (required)
-
Global scale applied to the down projection weights.
-
- -#### Inputs - -
-
input : T
-
Bfloat16 hidden states with shape [1, rows, 2560].
-
router_logits : T
-
Bfloat16 router logits with shape [1, rows, 512].
-
gate_up_weights : T1
-
Packed NVFP4 gate/up weights with shape [512, 1280, 1440].
-
down_weights : T1
-
Packed NVFP4 down weights with shape [512, 2560, 360].
-
- -#### Outputs - -
-
output : T
-
Bfloat16 output with the same shape as input.
-
scratch : T2
-
Float scratch buffer used by the fused expert projections.
-
- -#### Type Constraints - -
-
T : tensor(bfloat16)
-
Constrain activations and output to bfloat16.
-
T1 : tensor(uint8)
-
Constrain packed expert weights to uint8.
-
T2 : tensor(float)
-
Constrain scratch to float.
-
- -### **com.microsoft.FoldMid2** - - Combines speculative-token acceptance, recurrent-state commit bookkeeping, and head rotary-cache lookup for a two-token verification step. When `ctl[3]` is zero, the operator updates `ctl` and `ctx_commit` in place and writes a result record through the device-accessible pointer and capacity in `ptrs`. - -#### Version - -This version of the operator has been available since version 1 of the 'com.microsoft' operator set. - -#### Inputs - -
-
next_tokens : I64
-
Two predicted token IDs.
-
input_ids : I64
-
Two input token IDs; input_ids[1] is the draft token.
-
status : I32
-
Status values aggregated into the result record.
-
ctl : I64
-
Mutable eight-element control state.
-
ctx_commit : I64
-
Mutable two-element committed context.
-
ctx_a : I64
-
Context selected when the draft is rejected.
-
ctx_b : I64
-
Context selected when the draft is accepted.
-
ptrs : I64
-
Device pointers containing the result-ring address and capacity.
-
cos_cache : B
-
Bfloat16 rotary cosine cache with shape [positions, 32].
-
sin_cache : B
-
Bfloat16 rotary sine cache with shape [positions, 32].
-
- -#### Outputs - -
-
head_ids : I64
-
Two token IDs consumed by the head.
-
head_c : I32
-
Current head-cache length.
-
slot_ids : I64
-
Two consecutive head-cache slots.
-
seqlens : I32
-
Head sequence length.
-
cos_out : B
-
Two rows of rotary cosine values.
-
sin_out : B
-
Two rows of rotary sine values.
-
sel : I64
-
Selected row: one when the draft is accepted, otherwise zero.
-
- -#### Type Constraints - -
-
B : tensor(bfloat16)
-
Constrain rotary cache tensors to bfloat16.
-
I32 : tensor(int32)
-
Constrain status and sequence tensors to int32.
-
I64 : tensor(int64)
-
Constrain control and token tensors to int64.
-
- -### **com.microsoft.FoldTail** - - Updates the next speculative draft token and aggregate error bit from the head result. When `ctl[3]` is zero, the operator updates the control tensor in place; muted calls return `done=0` without changing the state. - -#### Version - -This version of the operator has been available since version 1 of the 'com.microsoft' operator set. - -#### Inputs - -
-
next_token : I64
-
Next draft token ID.
-
status : I32
-
Head status values.
-
ctl : I64
-
Mutable eight-element control state.
-
- -#### Outputs - -
-
done : I32
-
One when the control state was updated; zero for muted calls.
-
- -#### Type Constraints - -
-
I32 : tensor(int32)
-
Constrain status and output tensors to int32.
-
I64 : tensor(int64)
-
Constrain token and control tensors to int64.
-
diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md index 669e3083f6b15..cb6d5e79d8cd8 100644 --- a/docs/OperatorKernels.md +++ b/docs/OperatorKernels.md @@ -1103,6 +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)| |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)| @@ -1130,6 +1132,7 @@ The **OpSet Version** column uses the following notation: |MatMulBnb4|*in* A:**T1**
*in* B:**T2**
*in* absmax:**T1**
*out* Y:**T1**|1+|**T1** = tensor(bfloat16), tensor(float), tensor(float16)
**T2** = tensor(uint8)| |MatMulNBits|*in* A:**T1**
*in* B:**T2**
*in* scales:**T1**
*in* zero_points:**T3**
*in* g_idx:**T4**
*in* bias:**T1**
*out* Y:**T1**|1+|**T1** = tensor(bfloat16), tensor(float), tensor(float16)
**T2** = tensor(uint8)
**T3** = tensor(bfloat16), tensor(float), tensor(float16), tensor(uint8)| |MoE|*in* input:**T**
*in* router_probs:**T**
*in* fc1_experts_weights:**T**
*in* fc1_experts_bias:**T**
*in* fc2_experts_weights:**T**
*in* fc2_experts_bias:**T**
*in* fc3_experts_weights:**T**
*in* fc3_experts_bias:**T**
*out* output:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)| +|MtpMoe|*in* input:**T**
*in* router_logits:**T**
*in* gate_up_weights:**T1**
*in* down_weights:**T1**
*out* output:**T**
*out* scratch:**T2**|1+|**T** = tensor(bfloat16)
**T1** = tensor(uint8)
**T2** = tensor(float)| |MultiHeadAttention|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* bias:**T**
*in* key_padding_mask:**M**
*in* attention_bias:**T**
*in* past_key:**T**
*in* past_value:**T**
*in* past_sequence_length:**M**
*in* cache_indirection:**M**
*out* output:**T**
*out* present_key:**T**
*out* present_value:**T**
*out* qk:**QK**|1+|**QK** = tensor(bfloat16), tensor(float), tensor(float16)
**T** = tensor(bfloat16), tensor(float), tensor(float16)| |NGramHashMapping|*in* input_ids:**M**
*in* multipliers:**M**
*in* vocab_sizes:**M**
*in* past_ids:**M**
*in* head_offsets:**M**
*in* eos_token_id:**M**
*in* segment_ids:**tensor(int32)**
*out* hash_ids:**M**
*out* present_ids:**M**|1+|**M** = tensor(int32), tensor(int64)| |NGramRepeatBlock|*in* input_ids:**Tid**
*in* scores:**T**
*out* scores_out:**T**|1+|**T** = tensor(float)
**Tid** = tensor(int64)| diff --git a/onnxruntime/contrib_ops/cuda/llm/mtp_moe.cu b/onnxruntime/contrib_ops/cuda/llm/mtp_moe.cu index 2bdb95a658717..8aec1bca2731d 100644 --- a/onnxruntime/contrib_ops/cuda/llm/mtp_moe.cu +++ b/onnxruntime/contrib_ops/cuda/llm/mtp_moe.cu @@ -4,9 +4,12 @@ #include "contrib_ops/cuda/llm/mtp_moe.h" #include +#include #include +#include "core/common/common.h" + namespace onnxruntime { namespace contrib { namespace cuda { @@ -82,6 +85,21 @@ __global__ void RouterKernel(const __nv_bfloat16* __restrict__ logits, for (int offset = 16; offset > 0; offset >>= 1) { max_logit = fmaxf(max_logit, __shfl_xor_sync(0xffffffffu, max_logit, offset)); } + bool invalid = false; +#pragma unroll + for (int i = 0; i < kIterations; ++i) { + invalid |= !isfinite(logits_per_lane[i]); + } + if (__any_sync(0xffffffffu, invalid)) { + if (lane == 0) { + float* row_scratch = scratch + static_cast(row) * kScratchStride; + for (int k = 0; k < kTopK; ++k) { + row_scratch[k] = __int_as_float(k); + row_scratch[kTopK + k] = CUDART_NAN_F; + } + } + return; + } float exponentials[kIterations]; float denominator = 0.0f; @@ -278,7 +296,12 @@ Status LaunchNvfp4MtpMoe(const __nv_bfloat16* input, const __nv_bfloat16* router return Status(common::StatusCategory::ONNXRUNTIME, common::StatusCode::INVALID_ARGUMENT, "MtpMoe requires at least one row"); } - RouterKernel<<>>(router_logits, scratch); + const cudaError_t clear_error = cudaMemsetAsync( + scratch, 0, static_cast(rows) * kScratchStride * sizeof(float), stream); + if (clear_error != cudaSuccess) { + return Status(common::StatusCategory::SYSTEM, static_cast(clear_error), cudaGetErrorString(clear_error)); + } + ORT_RETURN_IF_ERROR(LaunchRouter(router_logits, scratch, rows, stream)); GateUpKernel<<>>( input, gate_up_weights, scratch, gate_up_global_scale); DownKernel<<>>(down_weights, scratch, output, down_global_scale); diff --git a/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cc b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cc index 6965c9d1539c1..c5d4c6a06f893 100644 --- a/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cc +++ b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cc @@ -65,15 +65,16 @@ class FoldMid2 final : public CudaKernel { 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); + 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), - const_cast(ctl->Data()), - const_cast(ctx_commit->Data()), + ctl->Data(), ctx_commit->Data(), ctx_a->Data(), ctx_b->Data(), ptrs->Data(), reinterpret_cast(cos_cache->Data()), reinterpret_cast(sin_cache->Data()), @@ -81,7 +82,8 @@ class FoldMid2 final : public CudaKernel { slot_ids->MutableData(), seqlens->MutableData(), reinterpret_cast<__nv_bfloat16*>(cos_out->MutableData()), reinterpret_cast<__nv_bfloat16*>(sin_out->MutableData()), - sel->MutableData(), Stream(context)); + sel->MutableData(), updated_ctl->MutableData(), + updated_ctx_commit->MutableData(), cos_cache->Shape()[0], Stream(context)); } }; @@ -100,12 +102,13 @@ class FoldTail final : public CudaKernel { ORT_RETURN_IF_NOT(status_count <= std::numeric_limits::max(), "FoldTail status tensor is too large"); Tensor* done = context->Output(0, TensorShape({1})); - ORT_ENFORCE(done != nullptr); + 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), const_cast(ctl->Data()), - done->MutableData(), Stream(context)); + static_cast(status_count), ctl->Data(), + done->MutableData(), updated_ctl->MutableData(), Stream(context)); } }; @@ -114,6 +117,8 @@ class FoldTail final : public CudaKernel { ONNX_OPERATOR_KERNEL_EX( FoldMid2, kMSDomain, 1, kCudaExecutionProvider, (*KernelDefBuilder::Create()) + .MayInplace(3, 7) + .MayInplace(4, 8) .TypeConstraint("I64", DataTypeImpl::GetTensorType()) .TypeConstraint("I32", DataTypeImpl::GetTensorType()) .TypeConstraint("B", DataTypeImpl::GetTensorType()), @@ -122,6 +127,7 @@ ONNX_OPERATOR_KERNEL_EX( ONNX_OPERATOR_KERNEL_EX( FoldTail, kMSDomain, 1, kCudaExecutionProvider, (*KernelDefBuilder::Create()) + .MayInplace(2, 1) .TypeConstraint("I64", DataTypeImpl::GetTensorType()) .TypeConstraint("I32", DataTypeImpl::GetTensorType()), FoldTail); diff --git a/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cu b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cu index 4e034f204bd02..7a15d1bbb8dcc 100644 --- a/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cu +++ b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cu @@ -3,6 +3,8 @@ #include "contrib_ops/cuda/llm/speculative_fold.h" +#include + namespace onnxruntime { namespace contrib { namespace cuda { @@ -13,31 +15,61 @@ constexpr int64_t kHeadCapacity = 33028; __global__ void FoldMid2Kernel( const int64_t* next_tokens, const int64_t* input_ids, - const int32_t* status, int status_count, int64_t* ctl, - int64_t* ctx_commit, const int64_t* ctx_a, const int64_t* ctx_b, + 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) { + __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; __shared__ int64_t cache_length; + __shared__ int64_t validation_bits; 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; + 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) { - 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]; + 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; } - if (threadIdx.x != 0) { + if (threadIdx.x != 0 || validation_bits != 0) { return; } @@ -74,8 +106,7 @@ __global__ void FoldMid2Kernel( ctl[1] = committed_length + step; ctl[7] = cache_length + step; ctl[4] += 1; - ctl[5] |= (status_bits != 0 ? 2 : 0) | - (cache_length + 2 > kHeadCapacity ? 4 : 0); + ctl[5] |= status_bits != 0 ? 2 : 0; ctl[6] = accept ? 1 : 0; record[1] = pred0; record[2] = pred1; @@ -89,10 +120,11 @@ __global__ void FoldMid2Kernel( __global__ void FoldTailKernel(const int64_t* next_token, const int32_t* status, int status_count, - int64_t* ctl, int32_t* done) { + const int64_t* input_ctl, int32_t* done, int64_t* ctl) { if (threadIdx.x != 0) { return; } + for (int i = 0; i < 8; ++i) ctl[i] = input_ctl[i]; if (ctl[3] != 0) { done[0] = 0; return; @@ -119,13 +151,14 @@ Status LaunchStatus(cudaError_t error) { } // namespace Status LaunchFoldMid2(const int64_t* next_tokens, const int64_t* input_ids, - const int32_t* status, int status_count, int64_t* ctl, - int64_t* ctx_commit, const int64_t* ctx_a, const int64_t* ctx_b, + 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, cudaStream_t stream) { + 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, @@ -134,12 +167,12 @@ Status LaunchFoldMid2(const int64_t* next_tokens, const int64_t* input_ids, 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); + cos_out, sin_out, sel, updated_ctl, updated_ctx_commit, cache_positions); return LaunchStatus(cudaGetLastError()); } Status LaunchFoldTail(const int64_t* next_token, const int32_t* status, - int status_count, int64_t* ctl, int32_t* done, + 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, @@ -147,7 +180,7 @@ Status LaunchFoldTail(const int64_t* next_token, const int32_t* status, "FoldTail status count must not be negative"); } FoldTailKernel<<<1, 32, 0, stream>>>(next_token, status, status_count, ctl, - done); + done, updated_ctl); return LaunchStatus(cudaGetLastError()); } diff --git a/onnxruntime/contrib_ops/cuda/llm/speculative_fold.h b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.h index 3f376bc08eca9..6af434004d952 100644 --- a/onnxruntime/contrib_ops/cuda/llm/speculative_fold.h +++ b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.h @@ -16,16 +16,17 @@ namespace cuda { namespace speculative_fold { Status LaunchFoldMid2(const int64_t* next_tokens, const int64_t* input_ids, - const int32_t* status, int status_count, int64_t* ctl, - int64_t* ctx_commit, const int64_t* ctx_a, const int64_t* ctx_b, + 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, cudaStream_t stream); + int64_t* sel, int64_t* updated_ctl, int64_t* updated_ctx_commit, + int64_t cache_positions, cudaStream_t stream); Status LaunchFoldTail(const int64_t* next_token, const int32_t* status, - int status_count, int64_t* ctl, int32_t* done, + int status_count, const int64_t* ctl, int32_t* done, int64_t* updated_ctl, cudaStream_t stream); } // namespace speculative_fold diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc index 0271641588b8d..34d7bb74f63d7 100644 --- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc @@ -1468,7 +1468,9 @@ ONNX_MS_OPERATOR_SET_SCHEMA( constexpr const char* MtpMoe_ver1_doc = R"DOC( Mixture-of-experts projection for multi-token prediction. The input and router logits are bfloat16. Expert weights are packed NVFP4 bytes with shape [512, 1280, 1440] for gate/up and [512, 2560, 360] -for down projection. The output has the same shape as input. The second output is a float scratch buffer. +for down projection. The output has the same shape as input. The second output is a fully initialized float +scratch buffer. A row containing non-finite router logits selects valid expert indices 0 through 9 and +reports invalid routing through NaN routing weights and output values. )DOC"; ONNX_MS_OPERATOR_SET_SCHEMA( @@ -1499,8 +1501,14 @@ 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. The control tensor and context-commit tensor are updated in place when ctl[3] is -zero. The pointer tensor contains the device-accessible result-ring address and its record capacity. +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( @@ -1510,8 +1518,8 @@ ONNX_MS_OPERATOR_SET_SCHEMA( .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", "Mutable eight-element control state.", "I64") - .Input(4, "ctx_commit", "Mutable two-element committed context.", "I64") + .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") @@ -1523,14 +1531,23 @@ ONNX_MS_OPERATOR_SET_SCHEMA( .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 the draft is accepted, otherwise zero.", "I64") + .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.")); + .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 updated in place when ctl[3] is zero; muted calls return done=0 without changing the control state. +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( @@ -1539,10 +1556,15 @@ ONNX_MS_OPERATOR_SET_SCHEMA( .SetDoc(FoldTail_ver1_doc) .Input(0, "next_token", "Next draft token ID.", "I64") .Input(1, "status", "Head status values.", "I32") - .Input(2, "ctl", "Mutable eight-element control state.", "I64") + .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.")); + .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). diff --git a/onnxruntime/test/contrib_ops/cuda_kernels/mtp_moe_test.cc b/onnxruntime/test/contrib_ops/cuda_kernels/mtp_moe_test.cc index b9f6f91a75e4b..a8553cf31c85c 100644 --- a/onnxruntime/test/contrib_ops/cuda_kernels/mtp_moe_test.cc +++ b/onnxruntime/test/contrib_ops/cuda_kernels/mtp_moe_test.cc @@ -290,6 +290,8 @@ TEST(CUDA_EP_Unittest, MtpMoeNvfp4MatchesReferenceArithmeticEndToEnd) { DeviceBuffer device_down(full_down_bytes); DeviceBuffer device_output(hidden * sizeof(__nv_bfloat16)); DeviceBuffer device_scratch(mtp_moe::kScratchStride * sizeof(float)); + const std::vector scratch_poison(mtp_moe::kScratchStride, -999.0f); + device_scratch.Upload(scratch_poison.data(), scratch_poison.size() * sizeof(float)); device_input.Upload(input.data(), input.size() * sizeof(__nv_bfloat16)); device_router.Upload(router.data(), router.size() * sizeof(__nv_bfloat16)); for (int expert = 0; expert < experts; ++expert) { @@ -314,6 +316,12 @@ TEST(CUDA_EP_Unittest, MtpMoeNvfp4MatchesReferenceArithmeticEndToEnd) { std::vector actual_scratch(mtp_moe::kScratchStride * sizeof(float)); device_output.Download(actual_output.data(), actual_output.size() * sizeof(__nv_bfloat16)); device_scratch.Download(actual_scratch.data(), actual_scratch.size()); + for (size_t i = 20 * sizeof(float); i < actual_scratch.size(); ++i) { + if (i < 32 * sizeof(float) || + i >= 32 * sizeof(float) + mtp_moe::kTopK * ffn * sizeof(__nv_bfloat16)) { + EXPECT_EQ(actual_scratch[i], 0) << "scratch padding byte " << i; + } + } for (int slot = 0; slot < mtp_moe::kTopK; ++slot) { int actual_expert = -1; @@ -339,6 +347,32 @@ TEST(CUDA_EP_Unittest, MtpMoeNvfp4MatchesReferenceArithmeticEndToEnd) { } } +TEST(CUDA_EP_Unittest, MtpMoeRouterNonFiniteRowsKeepExpertIndicesInRange) { + constexpr int rows = 4; + std::vector<__nv_bfloat16> logits(rows * mtp_moe::kExpertCount, __float2bfloat16(0.0f)); + logits[0] = __float2bfloat16(std::numeric_limits::quiet_NaN()); + logits[mtp_moe::kExpertCount + 17] = __float2bfloat16(std::numeric_limits::infinity()); + std::fill(logits.begin() + 2 * mtp_moe::kExpertCount, logits.begin() + 3 * mtp_moe::kExpertCount, + __float2bfloat16(-std::numeric_limits::infinity())); + logits[3 * mtp_moe::kExpertCount + 511] = __float2bfloat16(-std::numeric_limits::infinity()); + DeviceBuffer device_logits(logits.size() * sizeof(__nv_bfloat16)); + DeviceBuffer device_scratch(rows * mtp_moe::kScratchStride * sizeof(float)); + device_logits.Upload(logits.data(), logits.size() * sizeof(__nv_bfloat16)); + const Status status = mtp_moe::LaunchRouter(device_logits.As<__nv_bfloat16>(), + device_scratch.As(), rows, nullptr); + ASSERT_TRUE(status.IsOK()) << status.ErrorMessage(); + std::vector scratch(rows * mtp_moe::kScratchStride); + device_scratch.Download(scratch.data(), scratch.size() * sizeof(float)); + for (int row = 0; row < rows; ++row) { + for (int slot = 0; slot < mtp_moe::kTopK; ++slot) { + int expert = -1; + std::memcpy(&expert, &scratch[row * mtp_moe::kScratchStride + slot], sizeof(expert)); + EXPECT_EQ(expert, slot); + EXPECT_TRUE(std::isnan(scratch[row * mtp_moe::kScratchStride + mtp_moe::kTopK + slot])); + } + } +} + } // namespace } // namespace test } // namespace onnxruntime 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 d7d7d4d4802f4..3d17acc140fc1 100644 --- a/onnxruntime/test/contrib_ops/cuda_kernels/speculative_fold_test.cc +++ b/onnxruntime/test/contrib_ops/cuda_kernels/speculative_fold_test.cc @@ -145,7 +145,7 @@ TEST(CUDA_EP_Unittest, FoldMid2AcceptCommitsStateAndPublishesRingRecord) { 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(), nullptr); + d_sel.As(), d_ctl.As(), d_ctx_commit.As(), 16, nullptr); ASSERT_TRUE(launch_status.IsOK()) << launch_status.ErrorMessage(); CheckCuda(cudaStreamSynchronize(nullptr)); @@ -256,7 +256,7 @@ TEST(CUDA_EP_Unittest, FoldMid2MutedRejectDoesNotCommitState) { 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(), nullptr); + d_sel.As(), d_ctl.As(), d_ctx_commit.As(), 16, nullptr); ASSERT_TRUE(launch_status.IsOK()) << launch_status.ErrorMessage(); CheckCuda(cudaStreamSynchronize(nullptr)); @@ -297,7 +297,7 @@ TEST(CUDA_EP_Unittest, FoldTailUpdatesDraftAndHonorsMute) { Status launch_status = LaunchFoldTail( d_next_token.As(), d_status.As(), static_cast(status.size()), d_ctl.As(), - d_done.As(), nullptr); + d_done.As(), d_ctl.As(), nullptr); ASSERT_TRUE(launch_status.IsOK()) << launch_status.ErrorMessage(); CheckCuda(cudaStreamSynchronize(nullptr)); @@ -315,7 +315,7 @@ TEST(CUDA_EP_Unittest, FoldTailUpdatesDraftAndHonorsMute) { launch_status = LaunchFoldTail( d_next_token.As(), d_status.As(), static_cast(status.size()), d_ctl.As(), - d_done.As(), nullptr); + 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)); diff --git a/onnxruntime/test/contrib_ops/speculative_fold_op_test.cc b/onnxruntime/test/contrib_ops/speculative_fold_op_test.cc new file mode 100644 index 0000000000000..82eb0b1eb371e --- /dev/null +++ b/onnxruntime/test/contrib_ops/speculative_fold_op_test.cc @@ -0,0 +1,233 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#if defined(USE_CUDA) && !defined(ORT_MINIMAL_BUILD) + +#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/default_providers.h" +#include "test/util/include/test_environment.h" + +namespace onnxruntime { +namespace test { +namespace { + +class FoldResultRing { + public: + FoldResultRing() { + 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) { + cudaFreeHost(host_); + ORT_THROW(cudaGetErrorString(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_)); } + + private: + int64_t* host_ = nullptr; + int64_t* device_ = nullptr; +}; + +TEST(SpeculativeFoldTest, CpuFeedsPropagateStateAcrossRunsAndRejectInvalidMetadata) { + 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(), + IOnnxRuntimeOpSchemaRegistryList(), {{kMSDomain, 1}}, {}, + DefaultLoggingManager().DefaultLogger()); + auto& graph = model.MainGraph(); + auto arg = [&](const char* name, int element_type, std::initializer_list shape) -> 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); + 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* 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* ptrs = arg("ptrs", i64, {2}); + auto* cos = arg("cos_cache", bf16, {16, 32}); + auto* sin = arg("sin_cache", bf16, {16, 32}); + 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* mid_ctl = arg("mid_ctl", i64, {8}); + auto* updated_ctx = arg("updated_ctx", i64, {2}); + auto* draft = arg("draft", i64, {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}, + nullptr, kMSDomain) + .SetExecutionProviderType(kCudaExecutionProvider); + graph.AddNode("tail", "FoldTail", "", {draft, status, mid_ctl}, {done, updated_ctl}, nullptr, kMSDomain) + .SetExecutionProviderType(kCudaExecutionProvider); + graph.SetOutputs({updated_ctl, updated_ctx, selected, cos_out, sin_out, done}); + 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); + }; + 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, + {"updated_ctl", "updated_ctx", "sel", "cos_out", "sin_out", "done"}, &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, + {"updated_ctl", "updated_ctx", "sel", "cos_out", "sin_out", "done"}, &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]); + } + } + + 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}); + std::vector outputs; + ASSERT_STATUS_OK(session.Run(RunOptions{}, feeds, + {"updated_ctl", "updated_ctx", "sel", "cos_out", "sin_out", "done"}, &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); + } + EXPECT_TRUE(std::equal(preserved_ring.begin(), preserved_ring.end(), ring.Host())); + } +} + +} // namespace +} // namespace test +} // namespace onnxruntime + +#endif From 24625bf3d29342960b73d564958a76551c91e437 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Mon, 19 Oct 2026 22:19:05 +0000 Subject: [PATCH 03/15] Fix fold session test calls and validate MTP CUDA row limits --- docs/ContribOperators.md | 5 +++-- onnxruntime/contrib_ops/cuda/llm/mtp_moe.cc | 6 ++---- onnxruntime/contrib_ops/cuda/llm/mtp_moe.cu | 5 +---- onnxruntime/contrib_ops/cuda/llm/mtp_moe.h | 9 +++++++++ onnxruntime/core/graph/contrib_ops/contrib_defs.cc | 1 + .../test/contrib_ops/cuda_kernels/mtp_moe_test.cc | 14 ++++++++++++++ .../test/contrib_ops/speculative_fold_op_test.cc | 13 +++++++------ 7 files changed, 37 insertions(+), 16 deletions(-) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index f18bcc0d95cd7..c04af0abd903f 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -61,7 +61,6 @@ Do not modify directly.* * com.microsoft.HyperConnectionPreMix * com.microsoft.Inverse * com.microsoft.Irfft - * experimental com.microsoft.IsAllFinite * com.microsoft.LinearAttention * com.microsoft.LinearAttentionGate * com.microsoft.LongformerAttention @@ -93,7 +92,6 @@ Do not modify directly.* * com.microsoft.Pad * com.microsoft.PagedAttention * com.microsoft.QAttention - * experimental com.microsoft.QEmbedLayerNormalization * com.microsoft.QGemm * com.microsoft.QLinearAdd * com.microsoft.QLinearAveragePool @@ -144,6 +142,8 @@ Do not modify directly.* * com.microsoft.VarlenNGramHashMapping * com.microsoft.WhisperBeamSearch * com.microsoft.WordConvEmbedding + * experimental com.microsoft.IsAllFinite + * experimental com.microsoft.QEmbedLayerNormalization ## com.microsoft ### **com.microsoft.Attention** @@ -4621,6 +4621,7 @@ This version of the operator has been available since version 1 of the 'com.micr ### **com.microsoft.MtpMoe** Mixture-of-experts projection for multi-token prediction. The input and router logits are bfloat16. + The row count must be between 1 and 65535, inclusive, to fit the CUDA projection grid dimensions. Expert weights are packed NVFP4 bytes with shape [512, 1280, 1440] for gate/up and [512, 2560, 360] for down projection. The output has the same shape as input. The second output is a fully initialized float scratch buffer. A row containing non-finite router logits selects valid expert indices 0 through 9 and diff --git a/onnxruntime/contrib_ops/cuda/llm/mtp_moe.cc b/onnxruntime/contrib_ops/cuda/llm/mtp_moe.cc index 167945464010e..ca29c1bb3519a 100644 --- a/onnxruntime/contrib_ops/cuda/llm/mtp_moe.cc +++ b/onnxruntime/contrib_ops/cuda/llm/mtp_moe.cc @@ -6,8 +6,6 @@ #include "core/common/common.h" #include "core/providers/cuda/cuda_kernel.h" -#include - namespace onnxruntime { namespace contrib { namespace cuda { @@ -38,8 +36,8 @@ class MtpMoe final : public CudaKernel { input_shape[2] == mtp_moe::kHiddenSize, "MtpMoe input must have shape [1, R, 2560], got ", input_shape); const int64_t rows = input_shape[1]; - ORT_RETURN_IF_NOT(rows > 0 && rows <= std::numeric_limits::max() && - router_shape.NumDimensions() == 3 && + ORT_RETURN_IF_ERROR(mtp_moe::ValidateRowCount(rows)); + ORT_RETURN_IF_NOT(router_shape.NumDimensions() == 3 && router_shape[0] == 1 && router_shape[1] == rows && router_shape[2] == mtp_moe::kExpertCount, "MtpMoe router logits must have shape [1, R, 512], got ", router_shape); diff --git a/onnxruntime/contrib_ops/cuda/llm/mtp_moe.cu b/onnxruntime/contrib_ops/cuda/llm/mtp_moe.cu index 8aec1bca2731d..7f3f98c4ed9ee 100644 --- a/onnxruntime/contrib_ops/cuda/llm/mtp_moe.cu +++ b/onnxruntime/contrib_ops/cuda/llm/mtp_moe.cu @@ -292,10 +292,7 @@ Status LaunchNvfp4MtpMoe(const __nv_bfloat16* input, const __nv_bfloat16* router const uint8_t* gate_up_weights, const uint8_t* down_weights, float gate_up_global_scale, float down_global_scale, float* scratch, __nv_bfloat16* output, int rows, cudaStream_t stream) { - if (rows <= 0) { - return Status(common::StatusCategory::ONNXRUNTIME, common::StatusCode::INVALID_ARGUMENT, - "MtpMoe requires at least one row"); - } + ORT_RETURN_IF_ERROR(ValidateRowCount(rows)); const cudaError_t clear_error = cudaMemsetAsync( scratch, 0, static_cast(rows) * kScratchStride * sizeof(float), stream); if (clear_error != cudaSuccess) { diff --git a/onnxruntime/contrib_ops/cuda/llm/mtp_moe.h b/onnxruntime/contrib_ops/cuda/llm/mtp_moe.h index ecba2b3babc83..d94e15cbf2f69 100644 --- a/onnxruntime/contrib_ops/cuda/llm/mtp_moe.h +++ b/onnxruntime/contrib_ops/cuda/llm/mtp_moe.h @@ -20,6 +20,15 @@ constexpr int kExpertCount = 512; constexpr int kTopK = 10; constexpr int kFfnSize = 640; constexpr int kScratchStride = 32 + kTopK * kFfnSize; +constexpr int kMaxRows = 65535; + +inline Status ValidateRowCount(int64_t rows) { + if (rows <= 0 || rows > kMaxRows) { + return Status(common::StatusCategory::ONNXRUNTIME, common::StatusCode::INVALID_ARGUMENT, + "MtpMoe requires between 1 and 65535 rows"); + } + return Status::OK(); +} Status LaunchRouter(const __nv_bfloat16* router_logits, float* scratch, int rows, cudaStream_t stream); diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc index 34d7bb74f63d7..cd42a5ba75bb9 100644 --- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc @@ -1467,6 +1467,7 @@ ONNX_MS_OPERATOR_SET_SCHEMA( constexpr const char* MtpMoe_ver1_doc = R"DOC( Mixture-of-experts projection for multi-token prediction. The input and router logits are bfloat16. +The row count must be between 1 and 65535, inclusive, to fit the CUDA projection grid dimensions. Expert weights are packed NVFP4 bytes with shape [512, 1280, 1440] for gate/up and [512, 2560, 360] for down projection. The output has the same shape as input. The second output is a fully initialized float scratch buffer. A row containing non-finite router logits selects valid expert indices 0 through 9 and diff --git a/onnxruntime/test/contrib_ops/cuda_kernels/mtp_moe_test.cc b/onnxruntime/test/contrib_ops/cuda_kernels/mtp_moe_test.cc index a8553cf31c85c..688d0e9858632 100644 --- a/onnxruntime/test/contrib_ops/cuda_kernels/mtp_moe_test.cc +++ b/onnxruntime/test/contrib_ops/cuda_kernels/mtp_moe_test.cc @@ -24,6 +24,20 @@ namespace mtp_moe = onnxruntime::contrib::cuda::mtp_moe; constexpr float kGateUpGlobalScale = 0.0005173456156626344f; constexpr float kDownGlobalScale = 0.00031534829759038985f; +TEST(MtpMoeTest, RowCountBoundaries) { + EXPECT_TRUE(mtp_moe::ValidateRowCount(1).IsOK()); + EXPECT_TRUE(mtp_moe::ValidateRowCount(65535).IsOK()); + for (int rows : {-1, 0, 65536, std::numeric_limits::max()}) { + SCOPED_TRACE(rows); + const auto status = mtp_moe::ValidateRowCount(rows); + EXPECT_EQ(status.Code(), common::StatusCode::INVALID_ARGUMENT); + const auto launch_status = mtp_moe::LaunchNvfp4MtpMoe( + nullptr, nullptr, nullptr, nullptr, 1.0f, 1.0f, nullptr, nullptr, rows, nullptr); + EXPECT_EQ(launch_status.Code(), common::StatusCode::INVALID_ARGUMENT); + EXPECT_EQ(launch_status.ErrorMessage(), status.ErrorMessage()); + } +} + void CheckCuda(cudaError_t error) { if (error != cudaSuccess) { throw std::runtime_error(cudaGetErrorString(error)); diff --git a/onnxruntime/test/contrib_ops/speculative_fold_op_test.cc b/onnxruntime/test/contrib_ops/speculative_fold_op_test.cc index 82eb0b1eb371e..058ab6eb4bd06 100644 --- a/onnxruntime/test/contrib_ops/speculative_fold_op_test.cc +++ b/onnxruntime/test/contrib_ops/speculative_fold_op_test.cc @@ -9,6 +9,7 @@ #include #include #include +#include #include #include @@ -19,6 +20,7 @@ #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" @@ -107,6 +109,8 @@ TEST(SpeculativeFoldTest, CpuFeedsPropagateStateAcrossRunsAndRejectInvalidMetada 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; @@ -130,8 +134,7 @@ TEST(SpeculativeFoldTest, CpuFeedsPropagateStateAcrossRunsAndRejectInvalidMetada feed("ctl", {8}, control); feed("ctx_commit", {2}, context); std::vector outputs; - ASSERT_STATUS_OK(session.Run(RunOptions{}, feeds, - {"updated_ctl", "updated_ctx", "sel", "cos_out", "sin_out", "done"}, &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); @@ -169,8 +172,7 @@ TEST(SpeculativeFoldTest, CpuFeedsPropagateStateAcrossRunsAndRejectInvalidMetada 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, - {"updated_ctl", "updated_ctx", "sel", "cos_out", "sin_out", "done"}, &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); @@ -209,8 +211,7 @@ TEST(SpeculativeFoldTest, CpuFeedsPropagateStateAcrossRunsAndRejectInvalidMetada feed("ctx_commit", {2}, context); feed("ptrs", {2}, std::vector{invalid.address, invalid.capacity}); std::vector outputs; - ASSERT_STATUS_OK(session.Run(RunOptions{}, feeds, - {"updated_ctl", "updated_ctx", "sel", "cos_out", "sin_out", "done"}, &outputs)); + ASSERT_STATUS_OK(session.Run(RunOptions{}, feeds, output_names, &outputs)); auto expected = control; expected[3] = 1; expected[5] = invalid.error; From 6e528e7d0bbed6c8a4d06d2afcf3a79ae498f00a Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Thu, 22 Oct 2026 22:12:16 +0000 Subject: [PATCH 04/15] 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 05/15] 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 From 30d28558e37cc3c13b65b5009027c119e7343084 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Fri, 23 Oct 2026 23:37:31 +0000 Subject: [PATCH 06/15] [CUDA] Squash #33550 into #33474 (FoldMid2/FoldTail S<=4, 256K capacity, optional in-place ctl/ctx) and remove unused MtpMoe Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- cmake/onnxruntime_unittests.cmake | 2 - docs/ContribOperators.md | 87 +-- docs/OperatorKernels.md | 1 - .../contrib_ops/cuda/cuda_contrib_kernels.cc | 2 - onnxruntime/contrib_ops/cuda/llm/mtp_moe.cc | 89 --- onnxruntime/contrib_ops/cuda/llm/mtp_moe.cu | 314 -------- onnxruntime/contrib_ops/cuda/llm/mtp_moe.h | 45 -- .../contrib_ops/cuda/llm/speculative_fold.cc | 60 +- .../contrib_ops/cuda/llm/speculative_fold.cu | 28 +- .../contrib_ops/cuda/llm/speculative_fold.h | 3 +- .../core/graph/contrib_ops/bert_defs.cc | 49 +- .../core/graph/contrib_ops/contrib_defs.cc | 734 +++++++++++------- onnxruntime/core/graph/contrib_ops/ms_opset.h | 2 - .../contrib_ops/cuda_kernels/mtp_moe_test.cc | 392 ---------- .../cuda_kernels/speculative_fold_test.cc | 2 +- 15 files changed, 558 insertions(+), 1252 deletions(-) delete mode 100644 onnxruntime/contrib_ops/cuda/llm/mtp_moe.cc delete mode 100644 onnxruntime/contrib_ops/cuda/llm/mtp_moe.cu delete mode 100644 onnxruntime/contrib_ops/cuda/llm/mtp_moe.h delete mode 100644 onnxruntime/test/contrib_ops/cuda_kernels/mtp_moe_test.cc diff --git a/cmake/onnxruntime_unittests.cmake b/cmake/onnxruntime_unittests.cmake index 7c5d8de4dd80d..78b6c46826d24 100644 --- a/cmake/onnxruntime_unittests.cmake +++ b/cmake/onnxruntime_unittests.cmake @@ -1125,7 +1125,6 @@ if (onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS AND onnxruntime_BUILD_CUDA_EP_AS_P "${TEST_SRC_DIR}/contrib_ops/cuda_kernels/qmoe_fp4_to_fp8_kernel_test.cc" "${TEST_SRC_DIR}/contrib_ops/cuda_kernels/qmoe_fp8_compaction_test.cc" "${TEST_SRC_DIR}/contrib_ops/cuda_kernels/softmax_topk_kernel_test.cc" - "${TEST_SRC_DIR}/contrib_ops/cuda_kernels/mtp_moe_test.cc" "${TEST_SRC_DIR}/contrib_ops/cuda_kernels/speculative_fold_test.cc" ) endif() @@ -1159,7 +1158,6 @@ if (onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS AND onnxruntime_BUILD_CUDA_EP_AS_P "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/llm/gemm_tactic_cache.cc" "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/llm/fpA_intB_gemm_adaptor.cu" "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/llm/fpA_intB_gemm_preprocessors_impl.cu" - "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/llm/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/moe/qmoe_kernels.cu" diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 4bd2bba789ce5..4a6b704a4c9b7 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -76,7 +76,6 @@ Do not modify directly.* * com.microsoft.MatMulNBitsQkv * com.microsoft.MaxpoolWithMask * com.microsoft.MoE - * com.microsoft.MtpMoe * com.microsoft.MulInteger * com.microsoft.MultiHeadAttention * com.microsoft.MurmurHash3 @@ -2104,7 +2103,7 @@ This version of the operator has been available since version 1 of the 'com.micr ### **com.microsoft.FoldMid2** - 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. + 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 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. The mid stage writes the payload; FoldTail publishes its sequence word after downstream validation when ptrs is passed to FoldTail, or FoldMid2 publishes it directly when updated_ctl/updated_ctx_commit are omitted. Muted prefill calls prepare inputs without committing state or touching ring metadata. Invalid device metadata sets error bits and mutes execution. #### Version @@ -2113,7 +2112,7 @@ This version of the operator has been available since version 1 of the 'com.micr #### Attributes
-
head_capacity : int (default is 33028)
+
head_capacity : int (default is 262144)
Maximum head-cache token capacity.
head_rows : int (default is 2)
Fixed number of head rows, 2 through 4.
@@ -2146,7 +2145,7 @@ This version of the operator has been available since version 1 of the 'com.micr
Optional hidden states [1, verifier_rows, ...].
-#### Outputs +#### Outputs (7 - 10)
head_ids : I64
@@ -2165,10 +2164,10 @@ This version of the operator has been available since version 1 of the 'com.micr
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, aliases ctl.
-
updated_ctx_commit : I64
-
Updated context, aliases ctx_commit.
+
updated_ctl (optional) : I64
+
Optional updated control state, aliases ctl.
+
updated_ctx_commit (optional) : I64
+
Optional updated context, aliases ctx_commit.
#### Type Constraints @@ -2185,7 +2184,7 @@ This version of the operator has been available since version 1 of the 'com.micr ### **com.microsoft.FoldTail** - 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. + Updates the selected draft-control slot and publishes the result-ring sequence after validation. updated_ctl must alias ctl on CUDA when present. Muted calls do not access the ring. Invalid ring metadata sets error bit 16 and mutes execution. done is a dependency output set to 1 when unmuted and 0 when muted. #### Version @@ -2195,10 +2194,10 @@ This version of the operator has been available since version 1 of the 'com.micr
ctl_slot : int (default is 2)
-
Draft-control slot: 2, 8 or 9.
+
Draft-control slot: -1 (none), 2, 8 or 9.
-#### Inputs +#### Inputs (3 - 4)
next_token : I64
@@ -2207,17 +2206,17 @@ This version of the operator has been available since version 1 of the 'com.micr
Downstream validation status words.
ctl : I64
Mutable control state.
-
ptrs : I64
-
Result ring address and record capacity.
+
ptrs (optional) : I64
+
Optional result ring address and record capacity.
-#### Outputs +#### Outputs (1 - 2)
done : I32
Dependency output [1].
-
updated_ctl : I64
-
Updated control state, aliases ctl.
+
updated_ctl (optional) : I64
+
Optional updated control state, aliases ctl.
#### Type Constraints @@ -4630,62 +4629,6 @@ This version of the operator has been available since version 1 of the 'com.micr -### **com.microsoft.MtpMoe** - - Mixture-of-experts projection for multi-token prediction. The input and router logits are bfloat16. - The row count must be between 1 and 65535, inclusive, to fit the CUDA projection grid dimensions. - Expert weights are packed NVFP4 bytes with shape [512, 1280, 1440] for gate/up and [512, 2560, 360] - for down projection. The output has the same shape as input. The second output is a fully initialized float - scratch buffer. A row containing non-finite router logits selects valid expert indices 0 through 9 and - reports invalid routing through NaN routing weights and output values. - -#### Version - -This version of the operator has been available since version 1 of the 'com.microsoft' operator set. - -#### Attributes - -
-
down_global_scale : float (required)
-
Global scale applied to the down projection weights.
-
gate_up_global_scale : float (required)
-
Global scale applied to the gate/up projection weights.
-
- -#### Inputs - -
-
input : T
-
BFloat16 hidden states with shape [1, rows, 2560].
-
router_logits : T
-
BFloat16 router logits with shape [1, rows, 512].
-
gate_up_weights : T1
-
Packed NVFP4 gate/up weights with shape [512, 1280, 1440].
-
down_weights : T1
-
Packed NVFP4 down weights with shape [512, 2560, 360].
-
- -#### Outputs - -
-
output : T
-
BFloat16 output with the same shape as input.
-
scratch : T2
-
Float scratch buffer used by the fused expert projections.
-
- -#### Type Constraints - -
-
T : tensor(bfloat16)
-
Constrain activations and output to bfloat16.
-
T1 : tensor(uint8)
-
Constrain packed expert weights to uint8.
-
T2 : tensor(float)
-
Constrain scratch to float.
-
- - ### **com.microsoft.MulInteger** Performs element-wise binary quantized multiplication (with Numpy-style broadcasting support). diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md index 4b9b7c8504cb7..f4702391f746c 100644 --- a/docs/OperatorKernels.md +++ b/docs/OperatorKernels.md @@ -1132,7 +1132,6 @@ The **OpSet Version** column uses the following notation: |MatMulBnb4|*in* A:**T1**
*in* B:**T2**
*in* absmax:**T1**
*out* Y:**T1**|1+|**T1** = tensor(bfloat16), tensor(float), tensor(float16)
**T2** = tensor(uint8)| |MatMulNBits|*in* A:**T1**
*in* B:**T2**
*in* scales:**T1**
*in* zero_points:**T3**
*in* g_idx:**T4**
*in* bias:**T1**
*out* Y:**T1**|1+|**T1** = tensor(bfloat16), tensor(float), tensor(float16)
**T2** = tensor(uint8)
**T3** = tensor(bfloat16), tensor(float), tensor(float16), tensor(uint8)| |MoE|*in* input:**T**
*in* router_probs:**T**
*in* fc1_experts_weights:**T**
*in* fc1_experts_bias:**T**
*in* fc2_experts_weights:**T**
*in* fc2_experts_bias:**T**
*in* fc3_experts_weights:**T**
*in* fc3_experts_bias:**T**
*out* output:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)| -|MtpMoe|*in* input:**T**
*in* router_logits:**T**
*in* gate_up_weights:**T1**
*in* down_weights:**T1**
*out* output:**T**
*out* scratch:**T2**|1+|**T** = tensor(bfloat16)
**T1** = tensor(uint8)
**T2** = tensor(float)| |MultiHeadAttention|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* bias:**T**
*in* key_padding_mask:**M**
*in* attention_bias:**T**
*in* past_key:**T**
*in* past_value:**T**
*in* past_sequence_length:**M**
*in* cache_indirection:**M**
*out* output:**T**
*out* present_key:**T**
*out* present_value:**T**
*out* qk:**QK**|1+|**QK** = tensor(bfloat16), tensor(float), tensor(float16)
**T** = tensor(bfloat16), tensor(float), tensor(float16)| |NGramHashMapping|*in* input_ids:**M**
*in* multipliers:**M**
*in* vocab_sizes:**M**
*in* past_ids:**M**
*in* head_offsets:**M**
*in* eos_token_id:**M**
*in* segment_ids:**tensor(int32)**
*out* hash_ids:**M**
*out* present_ids:**M**|1+|**M** = tensor(int32), tensor(int64)| |NGramRepeatBlock|*in* input_ids:**Tid**
*in* scores:**T**
*out* scores_out:**T**|1+|**T** = tensor(float)
**Tid** = tensor(int64)| diff --git a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc index 93b62d9fff569..1e09d85141a43 100644 --- a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc @@ -100,7 +100,6 @@ class CUDA_ONNX_OP_TYPED_CLASS_NAME(1, MLFloat16, Crop); class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, MoE); class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, MoE); class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, MoE); -class CUDA_MS_OP_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_TYPED_CLASS_NAME(1, BFloat16, QMoE); @@ -467,7 +466,6 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, - BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/onnxruntime/contrib_ops/cuda/llm/mtp_moe.cc b/onnxruntime/contrib_ops/cuda/llm/mtp_moe.cc deleted file mode 100644 index ca29c1bb3519a..0000000000000 --- a/onnxruntime/contrib_ops/cuda/llm/mtp_moe.cc +++ /dev/null @@ -1,89 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -#include "contrib_ops/cuda/llm/mtp_moe.h" - -#include "core/common/common.h" -#include "core/providers/cuda/cuda_kernel.h" - -namespace onnxruntime { -namespace contrib { -namespace cuda { - -using namespace onnxruntime::cuda; - -class MtpMoe final : public CudaKernel { - public: - explicit MtpMoe(const OpKernelInfo& info) : CudaKernel(info) { - ORT_ENFORCE(info.GetAttr("gate_up_global_scale", &gate_up_global_scale_).IsOK(), - "MtpMoe requires gate_up_global_scale"); - ORT_ENFORCE(info.GetAttr("down_global_scale", &down_global_scale_).IsOK(), - "MtpMoe requires down_global_scale"); - } - - Status ComputeInternal(OpKernelContext* context) const override { - const Tensor* input = context->Input(0); - const Tensor* router = context->Input(1); - const Tensor* gate_up = context->Input(2); - const Tensor* down = context->Input(3); - ORT_ENFORCE(input != nullptr && router != nullptr && gate_up != nullptr && down != nullptr); - - const TensorShape& input_shape = input->Shape(); - const TensorShape& router_shape = router->Shape(); - const TensorShape& gate_up_shape = gate_up->Shape(); - const TensorShape& down_shape = down->Shape(); - ORT_RETURN_IF_NOT(input_shape.NumDimensions() == 3 && input_shape[0] == 1 && - input_shape[2] == mtp_moe::kHiddenSize, - "MtpMoe input must have shape [1, R, 2560], got ", input_shape); - const int64_t rows = input_shape[1]; - ORT_RETURN_IF_ERROR(mtp_moe::ValidateRowCount(rows)); - ORT_RETURN_IF_NOT(router_shape.NumDimensions() == 3 && - router_shape[0] == 1 && router_shape[1] == rows && - router_shape[2] == mtp_moe::kExpertCount, - "MtpMoe router logits must have shape [1, R, 512], got ", router_shape); - ORT_RETURN_IF_NOT(gate_up_shape.NumDimensions() == 3 && - gate_up_shape[0] == mtp_moe::kExpertCount && - gate_up_shape[1] == 2 * mtp_moe::kFfnSize && - gate_up_shape[2] == mtp_moe::kHiddenSize * 9 / 16, - "MtpMoe NVFP4 gate_up weights must have shape [512, 1280, 1440], got ", gate_up_shape); - ORT_RETURN_IF_NOT(down_shape.NumDimensions() == 3 && - down_shape[0] == mtp_moe::kExpertCount && - down_shape[1] == mtp_moe::kHiddenSize && - down_shape[2] == mtp_moe::kFfnSize * 9 / 16, - "MtpMoe NVFP4 down weights must have shape [512, 2560, 360], got ", down_shape); - - Tensor* output = context->Output(0, input_shape); - Tensor* scratch = context->Output( - 1, TensorShape({rows, static_cast(mtp_moe::kScratchStride)})); - ORT_ENFORCE(output != nullptr && scratch != nullptr); - - return mtp_moe::LaunchNvfp4MtpMoe( - reinterpret_cast(input->Data()), - reinterpret_cast(router->Data()), - gate_up->Data(), down->Data(), - gate_up_global_scale_, down_global_scale_, - scratch->MutableData(), - reinterpret_cast<__nv_bfloat16*>(output->MutableData()), - static_cast(rows), Stream(context)); - } - - private: - float gate_up_global_scale_; - float down_global_scale_; -}; - -ONNX_OPERATOR_TYPED_KERNEL_EX( - MtpMoe, - kMSDomain, - 1, - BFloat16, - kCudaExecutionProvider, - (*KernelDefBuilder::Create()) - .TypeConstraint("T", DataTypeImpl::GetTensorType()) - .TypeConstraint("T1", DataTypeImpl::GetTensorType()) - .TypeConstraint("T2", DataTypeImpl::GetTensorType()), - MtpMoe); - -} // namespace cuda -} // namespace contrib -} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/llm/mtp_moe.cu b/onnxruntime/contrib_ops/cuda/llm/mtp_moe.cu deleted file mode 100644 index 7f3f98c4ed9ee..0000000000000 --- a/onnxruntime/contrib_ops/cuda/llm/mtp_moe.cu +++ /dev/null @@ -1,314 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -#include "contrib_ops/cuda/llm/mtp_moe.h" - -#include -#include - -#include - -#include "core/common/common.h" - -namespace onnxruntime { -namespace contrib { -namespace cuda { -namespace mtp_moe { -namespace { - -constexpr int kChunk = 8; -constexpr int kGateUpRows = 2 * kFfnSize; -constexpr int kWarpsPerBlock = 8; -constexpr int kKChunks = kHiddenSize / kChunk; -constexpr int kPerLane = kKChunks / 32; -constexpr int kDownChunks = kFfnSize / kChunk; -constexpr int kDownPerLane = (kDownChunks + 31) / 32; - -__device__ __forceinline__ float RoundBf16(float value) { - return __bfloat162float(__float2bfloat16(value)); -} - -__device__ __forceinline__ float Sigmoid(float value) { - return 1.0f / (1.0f + expf(-value)); -} - -__device__ __forceinline__ void UnpackBf16x8(const uint4& packed, float* values) { - const uint32_t words[4] = {packed.x, packed.y, packed.z, packed.w}; -#pragma unroll - for (int i = 0; i < 4; ++i) { - values[2 * i] = __uint_as_float(words[i] << 16); - values[2 * i + 1] = __uint_as_float(words[i] & 0xffff0000u); - } -} - -__device__ __forceinline__ float DecodeE4M3(uint8_t bits) { - const __half_raw half_bits = __nv_cvt_fp8_to_halfraw(static_cast<__nv_fp8_storage_t>(bits), __NV_E4M3); - return __half2float(*reinterpret_cast(&half_bits)); -} - -__device__ __forceinline__ float DecodeE2M1(uint32_t code) { - const uint32_t magnitude = code & 7u; - const uint32_t exponent_mantissa = - magnitude == 0u ? 0u : ((magnitude == 1u ? 252u : (magnitude + 252u)) << 22); - return __uint_as_float(((code & 8u) << 28) | exponent_mantissa); -} - -__device__ __forceinline__ void UnpackNvfp4(uint32_t packed, uint8_t scale, float* values) { - const float decoded_scale = DecodeE4M3(scale); -#pragma unroll - for (int i = 0; i < 8; ++i) { - values[i] = DecodeE2M1((packed >> (4 * i)) & 0xfu) * decoded_scale; - } -} - -__device__ __forceinline__ void MergeTop(float& best, int& best_index, - float candidate, int candidate_index) { - if (candidate > best || (candidate == best && candidate_index < best_index)) { - best = candidate; - best_index = candidate_index; - } -} - -// Mirrors the reference custom op's FP32 lane-wise softmax and stable iterative top-10. -__global__ void RouterKernel(const __nv_bfloat16* __restrict__ logits, - float* __restrict__ scratch) { - const int row = blockIdx.x; - const int lane = threadIdx.x; - constexpr int kIterations = kExpertCount / 32; - float logits_per_lane[kIterations]; - float max_logit = -INFINITY; -#pragma unroll - for (int i = 0; i < kIterations; ++i) { - logits_per_lane[i] = __bfloat162float(logits[static_cast(row) * kExpertCount + i * 32 + lane]); - max_logit = fmaxf(max_logit, logits_per_lane[i]); - } - for (int offset = 16; offset > 0; offset >>= 1) { - max_logit = fmaxf(max_logit, __shfl_xor_sync(0xffffffffu, max_logit, offset)); - } - bool invalid = false; -#pragma unroll - for (int i = 0; i < kIterations; ++i) { - invalid |= !isfinite(logits_per_lane[i]); - } - if (__any_sync(0xffffffffu, invalid)) { - if (lane == 0) { - float* row_scratch = scratch + static_cast(row) * kScratchStride; - for (int k = 0; k < kTopK; ++k) { - row_scratch[k] = __int_as_float(k); - row_scratch[kTopK + k] = CUDART_NAN_F; - } - } - return; - } - - float exponentials[kIterations]; - float denominator = 0.0f; -#pragma unroll - for (int i = 0; i < kIterations; ++i) { - exponentials[i] = expf(logits_per_lane[i] - max_logit); - denominator += exponentials[i]; - } - for (int offset = 16; offset > 0; offset >>= 1) { - denominator += __shfl_xor_sync(0xffffffffu, denominator, offset); - } - - float probabilities[kIterations]; -#pragma unroll - for (int i = 0; i < kIterations; ++i) { - probabilities[i] = exponentials[i] / denominator; - } - - float top_values[kTopK]; - int top_indices[kTopK]; -#pragma unroll - for (int k = 0; k < kTopK; ++k) { - float best = -1.0f; - int best_index = 0x7fffffff; -#pragma unroll - for (int i = 0; i < kIterations; ++i) { - if (probabilities[i] > best) { - best = probabilities[i]; - best_index = i * 32 + lane; - } - } - for (int offset = 16; offset > 0; offset >>= 1) { - const float other_value = __shfl_xor_sync(0xffffffffu, best, offset); - const int other_index = __shfl_xor_sync(0xffffffffu, best_index, offset); - MergeTop(best, best_index, other_value, other_index); - } - top_values[k] = best; - top_indices[k] = best_index; -#pragma unroll - for (int i = 0; i < kIterations; ++i) { - if (best_index == i * 32 + lane) { - probabilities[i] = -2.0f; - } - } - } - - if (lane == 0) { - float top_sum = 0.0f; -#pragma unroll - for (int k = 0; k < kTopK; ++k) { - top_sum += top_values[k]; - } - float* row_scratch = scratch + static_cast(row) * kScratchStride; -#pragma unroll - for (int k = 0; k < kTopK; ++k) { - row_scratch[k] = __int_as_float(top_indices[k]); - row_scratch[kTopK + k] = RoundBf16(top_values[k] / top_sum); - } - } -} - -__global__ void GateUpKernel(const __nv_bfloat16* __restrict__ input, - const uint8_t* __restrict__ gate_up, - float* __restrict__ scratch, - float global_scale) { - const int row = blockIdx.z; - const int slot = blockIdx.y; - const int warp = threadIdx.x >> 5; - const int lane = threadIdx.x & 31; - float* row_scratch = scratch + static_cast(row) * kScratchStride; - const int expert = __float_as_int(row_scratch[slot]); - const int output_row = blockIdx.x * kWarpsPerBlock + warp; - constexpr int kPackedRowBytes = kHiddenSize * 9 / 16; - const uint4* input_vector = reinterpret_cast(input + static_cast(row) * kHiddenSize); - const uint8_t* gate_row = gate_up + (static_cast(expert) * kGateUpRows + output_row) * kPackedRowBytes; - const uint8_t* up_row = - gate_up + (static_cast(expert) * kGateUpRows + kFfnSize + output_row) * kPackedRowBytes; - - uint32_t packed_gate[kPerLane]; - uint32_t packed_up[kPerLane]; - uint8_t gate_scales[kPerLane]; - uint8_t up_scales[kPerLane]; - uint4 packed_input[kPerLane]; -#pragma unroll - for (int i = 0; i < kPerLane; ++i) { - const int column_chunk = i * 32 + lane; - packed_gate[i] = __ldg(reinterpret_cast(gate_row) + column_chunk); - packed_up[i] = __ldg(reinterpret_cast(up_row) + column_chunk); - gate_scales[i] = __ldg(gate_row + kHiddenSize / 2 + (column_chunk >> 1)); - up_scales[i] = __ldg(up_row + kHiddenSize / 2 + (column_chunk >> 1)); - packed_input[i] = __ldg(input_vector + column_chunk); - } - - float gate_accumulator = 0.0f; - float up_accumulator = 0.0f; -#pragma unroll - for (int i = 0; i < kPerLane; ++i) { - float gate_values[8], up_values[8], input_values[8]; - UnpackNvfp4(packed_gate[i], gate_scales[i], gate_values); - UnpackNvfp4(packed_up[i], up_scales[i], up_values); - UnpackBf16x8(packed_input[i], input_values); -#pragma unroll - for (int k = 0; k < 8; ++k) { - gate_accumulator = fmaf(gate_values[k], input_values[k], gate_accumulator); - up_accumulator = fmaf(up_values[k], input_values[k], up_accumulator); - } - } - for (int offset = 16; offset > 0; offset >>= 1) { - gate_accumulator += __shfl_xor_sync(0xffffffffu, gate_accumulator, offset); - up_accumulator += __shfl_xor_sync(0xffffffffu, up_accumulator, offset); - } - - if (lane == 0) { - gate_accumulator *= global_scale; - up_accumulator *= global_scale; - const float gate = RoundBf16(gate_accumulator); - const float up = RoundBf16(up_accumulator); - const float silu = RoundBf16(gate * Sigmoid(gate)); - auto* hidden = reinterpret_cast<__nv_bfloat16*>(row_scratch + 32); - hidden[slot * kFfnSize + output_row] = __float2bfloat16(silu * up); - } -} - -__global__ void DownKernel(const uint8_t* __restrict__ down, - const float* __restrict__ scratch, - __nv_bfloat16* __restrict__ output, - float global_scale) { - const int row = blockIdx.y; - const int slot = threadIdx.x >> 5; - const int lane = threadIdx.x & 31; - const int output_column = blockIdx.x; - const float* row_scratch = scratch + static_cast(row) * kScratchStride; - const int expert = __float_as_int(row_scratch[slot]); - constexpr int kPackedRowBytes = kFfnSize * 9 / 16; - const uint8_t* weight_row = - down + (static_cast(expert) * kHiddenSize + output_column) * kPackedRowBytes; - const auto* hidden = reinterpret_cast(row_scratch + 32) + slot * kFfnSize; - const uint4* packed_hidden = reinterpret_cast(hidden); - float accumulator = 0.0f; -#pragma unroll - for (int chunk = 0; chunk < kDownPerLane; ++chunk) { - const int column_chunk = chunk * 32 + lane; - if (column_chunk < kDownChunks) { - const uint32_t packed_weight = __ldg(reinterpret_cast(weight_row) + column_chunk); - const uint8_t scale = __ldg(weight_row + kFfnSize / 2 + (column_chunk >> 1)); - const uint4 hidden_chunk = __ldg(packed_hidden + column_chunk); - float weights[8], activations[8]; - UnpackNvfp4(packed_weight, scale, weights); - UnpackBf16x8(hidden_chunk, activations); -#pragma unroll - for (int i = 0; i < 8; ++i) { - accumulator = fmaf(weights[i], activations[i], accumulator); - } - } - } - for (int offset = 16; offset > 0; offset >>= 1) { - accumulator += __shfl_xor_sync(0xffffffffu, accumulator, offset); - } - __shared__ float weighted[kTopK]; - if (lane == 0) { - accumulator *= global_scale; - weighted[slot] = RoundBf16(RoundBf16(accumulator) * row_scratch[kTopK + slot]); - } - __syncthreads(); - if (threadIdx.x == 0) { - float total = 0.0f; -#pragma unroll - for (int expert_slot = 0; expert_slot < kTopK; ++expert_slot) { - total += weighted[expert_slot]; - } - output[static_cast(row) * kHiddenSize + output_column] = __float2bfloat16(total); - } -} - -} // namespace - -Status LaunchRouter(const __nv_bfloat16* router_logits, float* scratch, int rows, cudaStream_t stream) { - if (rows <= 0) { - return Status(common::StatusCategory::ONNXRUNTIME, common::StatusCode::INVALID_ARGUMENT, - "MtpMoe requires at least one row"); - } - RouterKernel<<>>(router_logits, scratch); - const cudaError_t error = cudaGetLastError(); - return error == cudaSuccess - ? Status::OK() - : Status(common::StatusCategory::SYSTEM, static_cast(error), cudaGetErrorString(error)); -} - -Status LaunchNvfp4MtpMoe(const __nv_bfloat16* input, const __nv_bfloat16* router_logits, - const uint8_t* gate_up_weights, const uint8_t* down_weights, - float gate_up_global_scale, float down_global_scale, - float* scratch, __nv_bfloat16* output, int rows, cudaStream_t stream) { - ORT_RETURN_IF_ERROR(ValidateRowCount(rows)); - const cudaError_t clear_error = cudaMemsetAsync( - scratch, 0, static_cast(rows) * kScratchStride * sizeof(float), stream); - if (clear_error != cudaSuccess) { - return Status(common::StatusCategory::SYSTEM, static_cast(clear_error), cudaGetErrorString(clear_error)); - } - ORT_RETURN_IF_ERROR(LaunchRouter(router_logits, scratch, rows, stream)); - GateUpKernel<<>>( - input, gate_up_weights, scratch, gate_up_global_scale); - DownKernel<<>>(down_weights, scratch, output, down_global_scale); - const cudaError_t error = cudaGetLastError(); - return error == cudaSuccess - ? Status::OK() - : Status(common::StatusCategory::SYSTEM, static_cast(error), cudaGetErrorString(error)); -} - -} // namespace mtp_moe -} // namespace cuda -} // namespace contrib -} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/llm/mtp_moe.h b/onnxruntime/contrib_ops/cuda/llm/mtp_moe.h deleted file mode 100644 index d94e15cbf2f69..0000000000000 --- a/onnxruntime/contrib_ops/cuda/llm/mtp_moe.h +++ /dev/null @@ -1,45 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -#pragma once - -#include -#include - -#include - -#include "core/common/status.h" - -namespace onnxruntime { -namespace contrib { -namespace cuda { -namespace mtp_moe { - -constexpr int kHiddenSize = 2560; -constexpr int kExpertCount = 512; -constexpr int kTopK = 10; -constexpr int kFfnSize = 640; -constexpr int kScratchStride = 32 + kTopK * kFfnSize; -constexpr int kMaxRows = 65535; - -inline Status ValidateRowCount(int64_t rows) { - if (rows <= 0 || rows > kMaxRows) { - return Status(common::StatusCategory::ONNXRUNTIME, common::StatusCode::INVALID_ARGUMENT, - "MtpMoe requires between 1 and 65535 rows"); - } - return Status::OK(); -} - -Status LaunchRouter(const __nv_bfloat16* router_logits, float* scratch, - int rows, cudaStream_t stream); - -Status LaunchNvfp4MtpMoe(const __nv_bfloat16* input, const __nv_bfloat16* router_logits, - const uint8_t* gate_up_weights, const uint8_t* down_weights, - float gate_up_global_scale, float down_global_scale, - float* scratch, __nv_bfloat16* output, - int rows, cudaStream_t stream); - -} // namespace mtp_moe -} // namespace cuda -} // namespace contrib -} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cc b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cc index d35eac055bbf0..cb2b408b8011d 100644 --- a/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cc +++ b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cc @@ -18,7 +18,7 @@ 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)); + capacity_ = SafeInt(info.GetAttrOrDefault("head_capacity", 262144)); ORT_ENFORCE(head_rows_ >= 2 && head_rows_ <= 4 && capacity_ >= head_rows_, "FoldMid2: invalid head_rows/head_capacity"); } @@ -34,7 +34,7 @@ class FoldMid2 final : public CudaKernel { const auto* ptrs = context->Input(7); const auto* cosine = context->Input(8); const auto* sine = context->Input(9); - const auto* hidden = context->Input(10); + const auto* hidden = context->InputCount() > 10 ? context->Input(10) : nullptr; 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(); @@ -58,8 +58,10 @@ class FoldMid2 final : public CudaKernel { 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 { + if (context->OutputCount() > 7) { + head_hidden = context->Output(7, TensorShape(shape)); + } + } else if (context->OutputCount() > 7) { 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_})); @@ -69,13 +71,27 @@ class FoldMid2 final : public CudaKernel { 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"); + int64_t* ctl_ptr = const_cast(ctl->Data()); + int64_t* context_ptr = const_cast(committed->Data()); + bool publish_ring = true; + if (context->OutputCount() > 8) { + auto* updated_ctl = context->Output(8, ctl->Shape()); + if (updated_ctl) { + ORT_ENFORCE(updated_ctl->DataRaw() == ctl->DataRaw(), "FoldMid2 requires aliased control output"); + ctl_ptr = updated_ctl->MutableData(); + publish_ring = false; + } + } + if (context->OutputCount() > 9) { + auto* updated_context = context->Output(9, committed->Shape()); + if (updated_context) { + ORT_ENFORCE(updated_context->DataRaw() == committed->DataRaw(), "FoldMid2 requires aliased context output"); + context_ptr = updated_context->MutableData(); + } + } return speculative_fold::LaunchMid( tokens->Data(), ids->Data(), status->Data(), static_cast(status->Shape().Size()), - updated_ctl->MutableData(), updated_context->MutableData(), prefixes->Data(), + ctl_ptr, context_ptr, prefixes->Data(), prefix2->Data(), ptrs->Data(), reinterpret_cast(cosine->Data()), reinterpret_cast(sine->Data()), @@ -84,7 +100,8 @@ class FoldMid2 final : public CudaKernel { 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)); + static_cast(rows), head_rows_, static_cast(rope[0]), static_cast(rope[1]), capacity_, Stream(context), + publish_ring); } private: @@ -96,23 +113,32 @@ 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"); + ORT_ENFORCE(slot_ == -1 || slot_ == 2 || slot_ == 8 || slot_ == 9, + "FoldTail: ctl_slot must be -1, 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); + const auto* ptrs = context->InputCount() > 3 ? context->Input(3) : nullptr; 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), + ctl && ctl->Shape().Size() >= 8 && (slot_ < 0 || ctl->Shape().Size() > slot_) && + (ptrs == nullptr || 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"); + int64_t* ctl_ptr = const_cast(ctl->Data()); + if (context->OutputCount() > 1) { + auto* updated_ctl = context->Output(1, ctl->Shape()); + if (updated_ctl) { + ORT_ENFORCE(updated_ctl->DataRaw() == ctl->DataRaw(), "FoldTail requires aliased control output"); + ctl_ptr = updated_ctl->MutableData(); + } + } return speculative_fold::LaunchTail(token->Data(), status->Data(), - static_cast(status->Shape().Size()), updated_ctl->MutableData(), - ptrs->Data(), done->MutableData(), slot_, Stream(context)); + static_cast(status->Shape().Size()), ctl_ptr, + ptrs ? ptrs->Data() : nullptr, + done->MutableData(), slot_, Stream(context)); } private: diff --git a/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cu b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cu index 8b90ebe9f255d..31cb1a4d088a6 100644 --- a/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cu +++ b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cu @@ -19,7 +19,8 @@ __global__ void MidKernel(const int64_t* tokens, const int64_t* ids, const int32 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) { + int rows, int head_rows, int rope_rows, int rope_width, int capacity, + bool publish_ring) { __shared__ int64_t length; __shared__ int64_t cache_length; __shared__ int64_t errors; @@ -87,7 +88,10 @@ __global__ void MidKernel(const int64_t* tokens, const int64_t* ids, const int32 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 (publish_ring) { + __threadfence_system(); + record[0] = ctl[4]; + } } __global__ void HiddenKernel(const uint16_t* source, uint16_t* output, int64_t width, int rows, int head_rows) { @@ -103,19 +107,22 @@ __global__ void TailKernel(const int64_t* token, const int32_t* status, int coun if (threadIdx.x != 0) return; done[0] = 0; if (ctl[3] != 0) return; - if (ctl[4] <= 0 || !ValidRing(ptrs, ctl[4] - 1)) { + if (ptrs != nullptr && (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 (slot >= 0) 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]; + if (ptrs != nullptr) { + 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]; + } + done[0] = 1; } Status LaunchStatus(cudaError_t error) { @@ -130,10 +137,11 @@ 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, 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) { + int rows, int head_rows, int rope_rows, int rope_width, int capacity, cudaStream_t stream, + bool publish_ring) { 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); + rows, head_rows, rope_rows, rope_width, capacity, publish_ring); ORT_RETURN_IF_ERROR(LaunchStatus(cudaGetLastError())); if (head_hidden) { HiddenKernel<<<32, 256, 0, stream>>>(hidden, head_hidden, hidden_width, rows, head_rows); diff --git a/onnxruntime/contrib_ops/cuda/llm/speculative_fold.h b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.h index c21c0c98f9393..fe8ae64e39156 100644 --- a/onnxruntime/contrib_ops/cuda/llm/speculative_fold.h +++ b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.h @@ -15,7 +15,8 @@ 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, 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); + int rows, int head_rows, int rope_rows, int rope_width, int capacity, cudaStream_t stream, + bool publish_ring = false); 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); diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index c2b1e39a68546..8cfc05c72987a 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -38,13 +38,14 @@ ONNX_MS_OPERATOR_SET_SCHEMA( FoldMid2, 1, OpSchema() .SetDoc("Commits a batch-one speculative prefix and prepares rotary/head inputs for two through four head rows. " - "Control and committed-context outputs must alias their inputs on CUDA. This is stateful: callers " + "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. The mid stage writes the payload; FoldTail publishes " - "its sequence word after downstream validation. Muted prefill calls prepare inputs without committing " + "its sequence word after downstream validation when ptrs is passed to FoldTail, or FoldMid2 publishes it directly when updated_ctl/updated_ctx_commit are omitted. " + "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}) + .Attr("head_capacity", "Maximum head-cache token capacity.", AttributeProto::INT, int64_t{262144}) .Input(0, "next_tokens", "Verifier predictions, batch one.", "I64") .Input(1, "input_ids", "Verifier input IDs.", "I64") .Input(2, "status", "Verifier status words.", "I32") @@ -64,18 +65,24 @@ ONNX_MS_OPERATOR_SET_SCHEMA( .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") + .Output(8, "updated_ctl", "Optional updated control state, aliases ctl.", "I64", OpSchema::Optional) + .Output(9, "updated_ctx_commit", "Optional updated context, aliases ctx_commit.", "I64", OpSchema::Optional) .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 : {0, 2, 6, 8, 9}) { + if (static_cast(output) < ctx.getNumOutputs() && ctx.hasOutput(output)) { + 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); + if (ctx.getNumOutputs() > 7 && ctx.hasOutput(7) && hasInputShape(ctx, 10)) { + propagateElemTypeFromInputToOutput(ctx, 10, 7); + } const auto rows = getAttribute(ctx, "head_rows", 2); - const auto capacity = getAttribute(ctx, "head_capacity", 33028); + const auto capacity = getAttribute(ctx, "head_capacity", 262144); 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}) { @@ -94,38 +101,42 @@ ONNX_MS_OPERATOR_SET_SCHEMA( updateOutputShape(ctx, 4, shape); updateOutputShape(ctx, 5, shape); } - if (hasInputShape(ctx, 10)) { + if (ctx.getNumOutputs() > 7 && ctx.hasOutput(7) && 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); + if (ctx.getNumOutputs() > 8 && ctx.hasOutput(8)) propagateShapeFromInputToOutput(ctx, 3, 8); + if (ctx.getNumOutputs() > 9 && ctx.hasOutput(9)) propagateShapeFromInputToOutput(ctx, 4, 9); })); ONNX_MS_OPERATOR_SET_SCHEMA( FoldTail, 1, OpSchema() .SetDoc("Updates the selected draft-control slot and publishes the result-ring sequence after validation. " - "updated_ctl must alias ctl on CUDA. Muted calls do not access the ring. Invalid ring metadata " - "sets error bit 16 and mutes execution. done is a dependency output initialized to zero.") - .Attr("ctl_slot", "Draft-control slot: 2, 8 or 9.", AttributeProto::INT, int64_t{2}) + "updated_ctl must alias ctl on CUDA when present. Muted calls do not access the ring. Invalid ring metadata " + "sets error bit 16 and mutes execution. done is a dependency output set to 1 when unmuted and 0 when muted.") + .Attr("ctl_slot", "Draft-control slot: -1 (none), 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") + .Input(3, "ptrs", "Optional result ring address and record capacity.", "I64", OpSchema::Optional) .Output(0, "done", "Dependency output [1].", "I32") - .Output(1, "updated_ctl", "Updated control state, aliases ctl.", "I64") + .Output(1, "updated_ctl", "Optional updated control state, aliases ctl.", "I64", OpSchema::Optional) .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); + if (ctx.getNumOutputs() > 1 && ctx.hasOutput(1)) { + 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"); + if (slot != -1 && 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); + if (ctx.getNumOutputs() > 1 && ctx.hasOutput(1)) { + propagateShapeFromInputToOutput(ctx, 2, 1); + } })); void DecoderAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& ctx) { diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc index 9823c47403d35..e05119bd05dfa 100644 --- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc @@ -4,6 +4,7 @@ #include #include +#include #include "core/graph/onnx_protobuf.h" #include "onnx/defs/shape_inference.h" @@ -1465,41 +1466,6 @@ ONNX_MS_OPERATOR_SET_SCHEMA( .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain input and output types to float tensors.") .TypeAndShapeInferenceFunction(ONNX_NAMESPACE::propagateShapeAndTypeFromFirstInput)); -constexpr const char* MtpMoe_ver1_doc = R"DOC( -Mixture-of-experts projection for multi-token prediction. The input and router logits are bfloat16. -The row count must be between 1 and 65535, inclusive, to fit the CUDA projection grid dimensions. -Expert weights are packed NVFP4 bytes with shape [512, 1280, 1440] for gate/up and [512, 2560, 360] -for down projection. The output has the same shape as input. The second output is a fully initialized float -scratch buffer. A row containing non-finite router logits selects valid expert indices 0 through 9 and -reports invalid routing through NaN routing weights and output values. -)DOC"; - -ONNX_MS_OPERATOR_SET_SCHEMA( - MtpMoe, 1, - OpSchema() - .SetDoc(MtpMoe_ver1_doc) - .Attr("gate_up_global_scale", "Global scale applied to the gate/up projection weights.", - AttributeProto::FLOAT, true) - .Attr("down_global_scale", "Global scale applied to the down projection weights.", - AttributeProto::FLOAT, true) - .Input(0, "input", "BFloat16 hidden states with shape [1, rows, 2560].", "T") - .Input(1, "router_logits", "BFloat16 router logits with shape [1, rows, 512].", "T") - .Input(2, "gate_up_weights", "Packed NVFP4 gate/up weights with shape [512, 1280, 1440].", "T1") - .Input(3, "down_weights", "Packed NVFP4 down weights with shape [512, 2560, 360].", "T1") - .Output(0, "output", "BFloat16 output with the same shape as input.", "T") - .Output(1, "scratch", "Float scratch buffer used by the fused expert projections.", "T2") - .TypeConstraint("T", {"tensor(bfloat16)"}, "Constrain activations and output to bfloat16.") - .TypeConstraint("T1", {"tensor(uint8)"}, "Constrain packed expert weights to uint8.") - .TypeConstraint("T2", {"tensor(float)"}, "Constrain scratch to float.") - .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { - ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 0, 0); - ONNX_NAMESPACE::propagateShapeFromInputToOutput(ctx, 0, 0); - if (auto* scratch = ctx.getOutputType(1)) { - scratch->mutable_tensor_type()->set_elem_type( - ONNX_NAMESPACE::TensorProto::FLOAT); - } - })); - 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 @@ -1553,256 +1519,454 @@ constexpr const char* qMoE_ver1_doc = R"DOC( When swiglu_fusion=2, two GEMMs are fused, and g and l are concatenated on each row. )DOC"; -ONNX_MS_OPERATOR_SET_SCHEMA( - QMoE, 1, - OpSchema() - .SetDoc(qMoE_ver1_doc) - .Attr("activation_type", - "Activation function to use. Choose from relu, gelu, silu, swiglu and identity. CUDA also supports relu2 (squared ReLU). Default is relu", - AttributeProto::STRING, - std::string("relu")) - .Attr("k", - "Number of top experts to select from expert pool", - AttributeProto::INT, - static_cast(1)) - .Attr("normalize_routing_weights", - "Whether to normalize routing weights", - AttributeProto::INT, - static_cast(0)) - .Attr("use_sparse_mixer", - "Whether to use sparse mixer", - AttributeProto::INT, - static_cast(0)) - .Attr("expert_weight_bits", - "Number of bits used in quantized weights. Supported values are 2, 4, and 8. Default is 4 bits", - AttributeProto::INT, - static_cast(4)) - .Attr("fc1_expert_weight_bits", - "Optional FC1 override for expert_weight_bits. Inherits expert_weight_bits when omitted.", - AttributeProto::INT, - OPTIONAL_VALUE) - .Attr("fc2_expert_weight_bits", - "Optional FC2 override for expert_weight_bits. Inherits expert_weight_bits when omitted.", - AttributeProto::INT, - OPTIONAL_VALUE) - .Attr("fc3_expert_weight_bits", - "Optional FC3 override for expert_weight_bits. Inherits expert_weight_bits when omitted. " - "For fused SwiGLU, the effective FC3 width must equal the effective FC1 width.", - AttributeProto::INT, - OPTIONAL_VALUE) - .Attr("swiglu_fusion", - "0: not fused, 1: fused and interleaved. 2: fused and not interleaved.", - AttributeProto::INT, - static_cast(0)) - .Attr("swiglu_limit", - "The limit used to clamp inputs in SwiGLU. It is infinite when limit is not provided.", - AttributeProto::FLOAT, - OPTIONAL_VALUE) - .Attr("activation_alpha", - "Alpha parameter used in activation function.", - AttributeProto::FLOAT, 1.0f) - .Attr("activation_beta", - "Beta parameter used in activation function.", - AttributeProto::FLOAT, 0.0f) - .Attr("accuracy_level", - "Minimum accuracy level of the expert GEMMs on CPU, with the MatMulNBits meaning. For block-wise 4-bit " - "experts, 0 (default) or 1 keeps fp32 activations and 4 allows int8 activations (int8 dot-product kernels). " - "Block-wise 8-bit experts have no fp32 kernel and use int8 activations at every level. " - "Other values are treated as 0.", - AttributeProto::INT, static_cast(0)) - .Attr("block_size", - "For integer quantization, size of each quantization block along the K (input feature) dimension. " - "Must be power of two and ≥ 16 (e.g., 16, 32, 64, 128). " - "For integer and FP4 quantization, both hidden_size and inter_size must be divisible by " - "the block size. FP8 with block_size=128 supports partial 128x128 tiles. " - "The FP4 modes always use blocking: MXFP4 ('fp4'/'wfp4afp8') is normalized to block_size 32 " - "and NVFP4 ('nvfp4') to block_size 16, even when block_size is omitted. " - "For FP8 ('fp8'), a positive value instead specifies square blocks along both N and K, " - "with floating-point scales shaped [E, ceil(N/block_size), ceil(K/block_size)]; partial blocks are allowed. " - "Without a positive value, FP8 uses per-expert global scales. " - "For integer quantization ('int'), omitting block_size means there is no blocking " - "and a whole column shares one scaling factor. ", - AttributeProto::INT, - OPTIONAL_VALUE) - .Attr("quant_type", - "Quantization type: 'int' for integer quantization (default), 'fp4' for MXFP4 quantization, " - "'nvfp4' for NVFP4 quantization, 'fp8' for FP8 e4m3 weight-only quantization, " - "or 'wfp4afp8' for MXFP4 weight with FP8 activation. " - "When quant_type is 'fp4' or 'nvfp4', weights are stored in E2M1 FP4 format (2 values per byte), " - "fc*_scales inputs contain the FP4 block scales, and fc*_global_scale inputs must be provided. " - "'fp4' uses Float8E8M0 block scales with block_size 32; 'nvfp4' uses Float8E4M3FN block scales " - "with block_size 16. 'fp8' uses float8e4m3fn weights and, when block_size > 0, " - "float32, float16, or bfloat16 square-block scales in fc*_scales instead of fc*_global_scale.", - AttributeProto::STRING, - std::string("int")) - .Attr("weights_prepacked", - "Expert-weight layout selector with quantization-specific meanings. For quant_type='int', " - "-1 and 1 select execution-provider-specific prepacked int4/int8 fc1/fc2 layouts. " - "0 selects raw [E,N,K/pack] tensors as produced by quantize_matmul_{4,8}bits. " - "For quant_type='nvfp4' on CUDA, -1 (default) and 0 select legacy N-packed [E,K,N/2] bytes; " - "1 selects K-packed row-major [E,N,K/2] bytes while retaining logical [E,K,N/2] dimensions. " - "NVFP4 block scales and global scales retain their original layouts and values. " - "Ignored for other quantization types. Supported values are -1, 0, and 1; defaults to -1.", - AttributeProto::INT, - static_cast(-1)) - .Input(0, - "input", - "2D packed token tensor with shape (total_tokens, hidden_size), where tokens from ragged sequences " - "may be concatenated without padding, or " - "3D tensor with shape (batch_size, sequence_length, hidden_size)", - "T") - .Input(1, - "router_probs", - "2D tensor with shape (total_tokens, num_experts), where total_tokens must match the flattened " - "token count of input", - "T") - .Input(2, - "fc1_experts_weights", - "3D tensor with shape (num_experts, fusion_size * inter_size, " - "hidden_size * effective_fc1_bits / 8). The last dimension must be byte-aligned. " - "The fusion_size is 2 for fused swiglu, or 1 otherwise. effective_fc1_bits is " - "fc1_expert_weight_bits when provided, otherwise expert_weight_bits.", - "T1") - .Input(3, - "fc1_scales", - "Optional weight scales. For quant_type='int', this is a 2D tensor with shape " - "(num_experts, fusion_size * inter_size), or a 3D tensor with shape " - "(num_experts, fusion_size * inter_size, hidden_size / block_size) when block_size is provided. " - "For quant_type='fp4' or 'wfp4afp8', this is a float8e8m0 MXFP block-scale tensor with shape " - "(num_experts, fusion_size * inter_size, hidden_size / 32). " - "For quant_type='nvfp4', this is a float8e4m3fn NVFP4 block-scale tensor with shape " - "(num_experts, fusion_size * inter_size, hidden_size / 16). " - "For quant_type='fp8' and block_size > 0, required float32, float16, or bfloat16 scales with shape " - "(num_experts, ceil(fusion_size * inter_size / block_size), ceil(hidden_size / block_size)).", - "T2", - OpSchema::Optional) - .Input(4, - "fc1_experts_bias", - "2D optional tensor with shape (num_experts, fusion_size * inter_size)", "T", OpSchema::Optional) - .Input(5, - "fc2_experts_weights", - "3D tensor with shape (num_experts, hidden_size, inter_size * effective_fc2_bits / 8). " - "The last dimension must be byte-aligned. effective_fc2_bits is fc2_expert_weight_bits " - "when provided, otherwise expert_weight_bits.", - "T1") - .Input(6, - "fc2_scales", - "Optional weight scales. For quant_type='int', this is a 2D tensor with shape " - "(num_experts, hidden_size), or a 3D tensor with shape " - "(num_experts, hidden_size, inter_size / block_size) when block_size is provided. " - "For quant_type='fp4' or 'wfp4afp8', this is a float8e8m0 MXFP block-scale tensor with shape " - "(num_experts, hidden_size, inter_size / 32). " - "For quant_type='nvfp4', this is a float8e4m3fn NVFP4 block-scale tensor with shape " - "(num_experts, hidden_size, inter_size / 16). " - "For quant_type='fp8' and block_size > 0, required float32, float16, or bfloat16 scales with shape " - "(num_experts, ceil(hidden_size / block_size), ceil(inter_size / block_size)).", - "T2", - OpSchema::Optional) - .Input(7, - "fc2_experts_bias", - "2D optional tensor with shape (num_experts, hidden_size)", - "T", - OpSchema::Optional) - .Input(8, - "fc3_experts_weights", - "3D optional tensor with shape (num_experts, inter_size, hidden_size * effective_fc3_bits / 8). " - "The last dimension must be byte-aligned. effective_fc3_bits is fc3_expert_weight_bits " - "when provided, otherwise expert_weight_bits.", - "T1", - OpSchema::Optional) - .Input(9, - "fc3_scales", - "Optional weight scales. For quant_type='int', this is a 2D tensor with shape " - "(num_experts, inter_size), or a 3D tensor with shape " - "(num_experts, inter_size, hidden_size / block_size) when block_size is provided. " - "For quant_type='fp4' or 'wfp4afp8', this is a float8e8m0 MXFP block-scale tensor with shape " - "(num_experts, inter_size, hidden_size / 32). " - "For quant_type='fp8' and block_size > 0, required when FC3 is present, with floating-point scales shaped " - "(num_experts, ceil(inter_size / block_size), ceil(hidden_size / block_size)).", - "T2", - OpSchema::Optional) - .Input(10, - "fc3_experts_bias", - "2D optional tensor with shape (num_experts, inter_size)", - "T", - OpSchema::Optional) - .Input(11, - "fc1_zero_points", - "2D tensor with shape (num_experts, ceil(fusion_size * inter_size * effective_fc1_bits / 8)), or " - "3D tensor with shape (num_experts, fusion_size * inter_size, " - "ceil((hidden_size / block_size) * effective_fc1_bits / 8)) when block_size is provided.", - "T1", - OpSchema::Optional) - .Input(12, - "fc2_zero_points", - "2D tensor with shape (num_experts, ceil(hidden_size * effective_fc2_bits / 8)), or " - "3D tensor with shape (num_experts, hidden_size, " - "ceil((inter_size / block_size) * effective_fc2_bits / 8)) when block_size is provided.", - "T1", - OpSchema::Optional) - .Input(13, - "fc3_zero_points", - "2D optional tensor with shape (num_experts, ceil(inter_size * effective_fc3_bits / 8)), or " - "3D optional tensor with shape (num_experts, inter_size, " - "ceil((hidden_size / block_size) * effective_fc3_bits / 8)) when block_size is provided.", - "T1", - OpSchema::Optional) - .Input(14, - "router_weights", - "2D optional tensor with shape (num_tokens, num_experts). " - "When provided, router_probs is used only for Top-K expert selection, and router_weights is used " - "for aggregating expert outputs (the values at the selected expert indices are gathered and used as " - "mixing weights). This enables DeepSeek-style noaux_tc routing where different tensors are used for " - "selection and aggregation. When not provided, router_probs is used for both selection and aggregation " - "(backward compatible).", - "T", - OpSchema::Optional) - .Input(15, - "fc1_global_scale", - "1D optional tensor with shape (num_experts,). " - "Per-expert global weight scale for FC1. Required when quant_type is 'fp4', 'nvfp4', or 'wfp4afp8', " - "or 'fp8' with block_size <= 0. Not used for block-scaled FP8.", - "T4", - OpSchema::Optional) - .Input(16, - "fc2_global_scale", - "1D optional tensor with shape (num_experts,). " - "Per-expert global weight scale for FC2. Required when quant_type is 'fp4', 'nvfp4', or 'wfp4afp8', " - "or 'fp8' with block_size <= 0. Not used for block-scaled FP8.", - "T4", - OpSchema::Optional) - .Input(17, - "fc1_act_scale", - "1D optional tensor with shape (1,) or (num_experts,). Activation scale for FC1 FP8 activation modes.", - "T4", - OpSchema::Optional) - .Input(18, - "fc2_act_scale", - "1D optional tensor with shape (1,) or (num_experts,). Activation scale for FC2 FP8 activation modes.", - "T4", - OpSchema::Optional) - .Input(19, - "fc1_act_block_scale", - "3D optional float8e8m0 MXFP activation block-scale tensor for FC1 FP8 activation modes.", - "T2", - OpSchema::Optional) - .Input(20, - "fc2_act_block_scale", - "3D optional float8e8m0 MXFP activation block-scale tensor for FC2 FP8 activation modes.", - "T2", - OpSchema::Optional) - .Output(0, - "output", - "output tensor with same shape of input", - "T") - .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain input and output types to float tensors.") - .TypeConstraint("T1", {"tensor(uint8)", "tensor(float8e4m3fn)"}, - "Constrain quantized weight types. Integer and FP4 weights use uint8. FP8 weights use float8e4m3fn.") - .TypeConstraint("T2", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)", "tensor(float8e8m0)", "tensor(float8e4m3fn)"}, - "Constrain scale types. Float tensors are used for integer quantization and FP8 square-block scales. " - "Float8e8m0 tensors are used for MXFP4 block scales; float8e4m3fn tensors are used for NVFP4 block scales.") - .TypeConstraint("T4", {"tensor(float)"}, "Constrain FP4 global scale type to float32 tensors.") - .TypeAndShapeInferenceFunction(ONNX_NAMESPACE::propagateShapeAndTypeFromFirstInput)); +static void QMoETypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& ctx) { + using namespace ONNX_NAMESPACE; + const int64_t root_bits = getAttribute(ctx, "expert_weight_bits", 4); + const int64_t fc1_bits = getAttribute(ctx, "fc1_expert_weight_bits", root_bits); + const int64_t fc2_bits = getAttribute(ctx, "fc2_expert_weight_bits", root_bits); + const int64_t fc3_bits = getAttribute(ctx, "fc3_expert_weight_bits", root_bits); + propagateShapeAndTypeFromFirstInput(ctx); + if (fc1_bits != 3 && fc2_bits != 3 && fc3_bits != 3) { + return; + } + for (int64_t bits : {root_bits, fc1_bits, fc2_bits, fc3_bits}) { + if (bits != 2 && bits != 3 && bits != 4 && bits != 8) { + fail_shape_inference("QMoE INT3 supports only expert weight widths 2, 3, 4, and 8."); + } + } + + const auto* quant_type = ctx.getAttribute("quant_type"); + if ((quant_type && quant_type->s() != "int") || getAttribute(ctx, "weights_prepacked", -1) != 0) { + fail_shape_inference("QMoE INT3 requires integer quantization and explicit weights_prepacked=0 (raw weights)."); + } + const auto* activation_type_attr = ctx.getAttribute("activation_type"); + if (!activation_type_attr || activation_type_attr->s() != "swiglu" || + getAttribute(ctx, "swiglu_fusion", 0) != 1 || fc3_bits != fc1_bits) { + fail_shape_inference("QMoE INT3 requires activation_type='swiglu', swiglu_fusion=1 and equal effective FC1/FC3 widths after inheritance."); + } + const int64_t block_size = getAttribute(ctx, "block_size", 0); + if (block_size != 32 && block_size != 64 && block_size != 128) { + fail_shape_inference("QMoE INT3 requires block_size 32, 64, or 128."); + } + auto input_type = [&](size_t index) -> const TypeProto* { + return index < ctx.getNumInputs() ? ctx.getInputType(index) : nullptr; + }; + for (size_t index : {8, 9, 10, 13, 15, 16, 17, 18, 19, 20}) { + if (input_type(index)) { + fail_shape_inference("QMoE INT3 fused raw profile does not accept input ", index, "."); + } + } + const auto* activation = input_type(0); + const int activation_type = activation && activation->has_tensor_type() ? activation->tensor_type().elem_type() : 0; + if (activation_type && activation_type != TensorProto::FLOAT16 && activation_type != TensorProto::BFLOAT16) { + fail_type_inference("QMoE INT3 requires FP16 or BF16 activations."); + } + auto check_type = [&](size_t index, int expected, bool required) { + const auto* type = input_type(index); + if (!type) { + if (required) { + fail_type_inference("QMoE INT3 requires input ", index, "."); + } + return; + } + if (!type->has_tensor_type() || + (expected && type->tensor_type().elem_type() && type->tensor_type().elem_type() != expected)) { + fail_type_inference("QMoE INT3 input ", index, " has an incompatible tensor element type."); + } + }; + for (size_t index : {2, 5, 11, 12}) { + check_type(index, TensorProto::UINT8, index == 2 || index == 5); + } + for (size_t index : {1, 3, 4, 6, 7, 14}) { + check_type(index, activation_type, index == 1 || index == 3 || index == 6); + } + auto shape = [&](size_t index, int rank) -> const TensorShapeProto* { + if (!input_type(index) || !hasInputShape(ctx, index)) { + return nullptr; + } + const auto& result = getInputShape(ctx, index); + if (result.dim_size() != rank) { + fail_shape_inference("QMoE INT3 input ", index, " must have rank ", rank, "."); + } + for (const auto& dimension : result.dim()) { + if (dimension.has_dim_value() && dimension.dim_value() < 0) { + fail_shape_inference("QMoE INT3 dimensions must be nonnegative."); + } + } + return &result; + }; + const TensorShapeProto* activation_shape = nullptr; + if (hasInputShape(ctx, 0)) { + const auto& candidate = getInputShape(ctx, 0); + if (candidate.dim_size() != 2 && candidate.dim_size() != 3) { + fail_shape_inference("QMoE INT3 input must have rank 2 or 3."); + } + activation_shape = shape(0, candidate.dim_size()); + } + const auto* fc1 = shape(2, 3); + const auto* fc2 = shape(5, 3); + using Dimension = TensorShapeProto::Dimension; + const Dimension* hidden = activation_shape ? &activation_shape->dim(activation_shape->dim_size() - 1) : nullptr; + Dimension experts; + if (fc1) experts = fc1->dim(0); + Dimension intermediate; + if (fc1 && fc1->dim(1).has_dim_value()) { + if (fc1->dim(1).dim_value() % 2 != 0) { + fail_shape_inference("QMoE INT3 fused FC1 row count must be even."); + } + intermediate.set_dim_value(fc1->dim(1).dim_value() / 2); + } + for (const auto* dimension : {hidden, static_cast(&intermediate), + static_cast(&experts)}) { + if (dimension && dimension->has_dim_value() && dimension->dim_value() <= 0) { + fail_shape_inference("QMoE INT3 hidden, intermediate, and expert dimensions must be positive."); + } + } + for (const auto* dimension : {hidden, static_cast(&intermediate)}) { + if (dimension && dimension->has_dim_value() && dimension->dim_value() % block_size != 0) { + fail_shape_inference("QMoE INT3 hidden and intermediate sizes must be divisible by block_size."); + } + } + auto match = [](const Dimension& actual, const Dimension* expected, const char* name) { + if (expected && expected->has_dim_value() && actual.has_dim_value() && + actual.dim_value() != expected->dim_value()) { + fail_shape_inference("QMoE INT3 incompatible ", name, " dimension."); + } + }; + auto packed_bytes = [](int64_t count, int64_t bits) { + return (count / 8) * bits + ((count % 8) * bits + 7) / 8; + }; + auto match_experts = [&](const Dimension& actual) { + match(actual, &experts, "expert count"); + if (actual.has_dim_value()) { + if (actual.dim_value() <= 0) { + fail_shape_inference("QMoE INT3 expert count must be positive."); + } + experts = actual; + } + }; + auto projection = [&](const TensorShapeProto* weights, size_t scale_index, size_t zp_index, + size_t bias_index, const Dimension* rows, const Dimension* reduction, int64_t bits) { + const auto* scales = shape(scale_index, 3); + const auto* zero_points = shape(zp_index, 3); + const auto* bias = shape(bias_index, 2); + Dimension concrete_rows; + if (rows) concrete_rows = *rows; + for (const auto* tensor : {weights, scales, zero_points, bias}) { + if (tensor) { + match_experts(tensor->dim(0)); + match(tensor->dim(1), &concrete_rows, "projection metadata rows"); + if (tensor->dim(1).has_dim_value()) concrete_rows = tensor->dim(1); + } + } + if (reduction && reduction->has_dim_value()) { + Dimension expected; + expected.set_dim_value(packed_bytes(reduction->dim_value(), bits)); + if (weights) match(weights->dim(2), &expected, "packed weight bytes"); + const int64_t blocks = reduction->dim_value() / block_size; + expected.set_dim_value(blocks); + if (scales) match(scales->dim(2), &expected, "scale blocks"); + expected.set_dim_value(packed_bytes(blocks, bits)); + if (zero_points) match(zero_points->dim(2), &expected, "packed zero-point bytes"); + } + }; + projection(fc1, 3, 11, 4, fc1 ? &fc1->dim(1) : nullptr, hidden, fc1_bits); + projection(fc2, 6, 12, 7, hidden, &intermediate, fc2_bits); + const auto* router = shape(1, 2); + const auto* router_weights = shape(14, 2); + Dimension tokens; + if (activation_shape) { + if (activation_shape->dim_size() == 2) { + tokens = activation_shape->dim(0); + } else if (activation_shape->dim(0).has_dim_value() && activation_shape->dim(1).has_dim_value()) { + const int64_t batch = activation_shape->dim(0).dim_value(); + const int64_t sequence = activation_shape->dim(1).dim_value(); + if (sequence != 0 && batch > std::numeric_limits::max() / sequence) { + fail_shape_inference("QMoE INT3 flattened token count overflows int64."); + } + tokens.set_dim_value(batch * sequence); + } + } + for (const auto* tensor : {router, router_weights}) { + if (tensor) { + match_experts(tensor->dim(1)); + match(tensor->dim(0), &tokens, "token count"); + if (router) match(tensor->dim(0), &router->dim(0), "router token count"); + } + } +} + +constexpr const char* qMoE_int3_doc = R"DOC( + The raw INT3 profile extends version 1; execution support is provider-specific. + For INT3, the requirements below supersede the legacy format descriptions above. + Effective widths are 2, 3, 4, or 8. Models without INT3 retain the legacy format contracts. + With INT3, use quant_type='int' (the default), and explicitly set activation_type='swiglu', weights_prepacked=0, swiglu_fusion=1, + and block_size=32, 64, or 128. Effective FC3 must equal FC1 after independent inheritance. + H comes from input and I from half the FC1 row count; both must be divisible by block_size. + Weights are uint8 [E,N,ceil(K*bits/8)], row-local LSB-first bitstreams, without per-block padding. + Scales are required [E,N,K/block_size], matching FP16/BF16 activations. + INT3 codes dequantize as (u-4)*scale; scales must be finite and nonnegative. + Optional INT3 zero points are uint8 [E,N,ceil(3*(K/block_size)/8)]. Each unpacked 3-bit + zero-point code must equal 4 (binary 100), the unsigned offset representing logical zero; + this does not mean each packed byte is 4. Omission implies the same zero point 4. + Unused high padding bits must be zero. INT2/4/8 projections retain their zero-point semantics. + FC1 rows are interleaved gate/up; FC2 is [E,H,ceil(I*fc2_bits/8)]. + Separate FC3, global scales, and activation scales are not accepted by the INT3 profile. +)DOC"; + +static OpSchema QMoESchema() { + constexpr const char* quant_type_doc = + "Quantization type: 'int' for integer quantization (default), 'fp4' for MXFP4 quantization, " + "'nvfp4' for NVFP4 quantization, 'fp8' for FP8 e4m3 weight-only quantization, " + "or 'wfp4afp8' for MXFP4 weight with FP8 activation. " + "When quant_type is 'fp4' or 'nvfp4', weights are stored in E2M1 FP4 format (2 values per byte), " + "fc*_scales inputs contain the FP4 block scales, and fc*_global_scale inputs must be provided. " + "'fp4' uses Float8E8M0 block scales with block_size 32; 'nvfp4' uses Float8E4M3FN block scales " + "with block_size 16. 'fp8' uses float8e4m3fn weights and, when block_size > 0, " + "float32, float16, or bfloat16 square-block scales in fc*_scales instead of fc*_global_scale. " + "The INT3 profile requires integer quantization."; + auto schema = OpSchema() + .SetDoc(std::string(qMoE_ver1_doc) + qMoE_int3_doc) + .Attr("activation_type", + "Activation function to use. Choose from relu, gelu, silu, swiglu and identity. CUDA also supports relu2 (squared ReLU). Default is relu", + AttributeProto::STRING, + std::string("relu")) + .Attr("k", + "Number of top experts to select from expert pool", + AttributeProto::INT, + static_cast(1)) + .Attr("normalize_routing_weights", + "Whether to normalize routing weights", + AttributeProto::INT, + static_cast(0)) + .Attr("use_sparse_mixer", + "Whether to use sparse mixer", + AttributeProto::INT, + static_cast(0)) + .Attr("expert_weight_bits", + "Number of bits used in quantized weights: 2, 3, 4, or 8. Default is 4 bits.", + AttributeProto::INT, + static_cast(4)) + .Attr("fc1_expert_weight_bits", + "Optional FC1 override for expert_weight_bits. Inherits expert_weight_bits when omitted.", + AttributeProto::INT, + OPTIONAL_VALUE) + .Attr("fc2_expert_weight_bits", + "Optional FC2 override for expert_weight_bits. Inherits expert_weight_bits when omitted.", + AttributeProto::INT, + OPTIONAL_VALUE) + .Attr("fc3_expert_weight_bits", + "Optional FC3 override for expert_weight_bits. Inherits expert_weight_bits when omitted. " + "For fused SwiGLU, the effective FC3 width must equal the effective FC1 width.", + AttributeProto::INT, + OPTIONAL_VALUE) + .Attr("swiglu_fusion", + "0: not fused, 1: fused and interleaved. 2: fused and not interleaved.", + AttributeProto::INT, + static_cast(0)) + .Attr("swiglu_limit", + "The limit used to clamp inputs in SwiGLU. It is infinite when limit is not provided.", + AttributeProto::FLOAT, + OPTIONAL_VALUE) + .Attr("activation_alpha", + "Alpha parameter used in activation function.", + AttributeProto::FLOAT, 1.0f) + .Attr("activation_beta", + "Beta parameter used in activation function.", + AttributeProto::FLOAT, 0.0f) + .Attr("accuracy_level", + "Minimum accuracy level of the expert GEMMs on CPU, with the MatMulNBits meaning. For block-wise 4-bit " + "experts, 0 (default) or 1 keeps fp32 activations and 4 allows int8 activations (int8 dot-product kernels). " + "Block-wise 8-bit experts have no fp32 kernel and use int8 activations at every level. " + "Other values are treated as 0.", + AttributeProto::INT, static_cast(0)) + .Attr("block_size", + "For integer quantization, size of each quantization block along the K (input feature) dimension. " + "Must be power of two and ≥ 16 (e.g., 16, 32, 64, 128). " + "For integer and FP4 quantization, both hidden_size and inter_size must be divisible by " + "the block size. FP8 with block_size=128 supports partial 128x128 tiles. " + "The FP4 modes always use blocking: MXFP4 ('fp4'/'wfp4afp8') is normalized to block_size 32 " + "and NVFP4 ('nvfp4') to block_size 16, even when block_size is omitted. " + "For FP8 ('fp8'), a positive value instead specifies square blocks along both N and K, " + "with floating-point scales shaped [E, ceil(N/block_size), ceil(K/block_size)]; partial blocks are allowed. " + "Without a positive value, FP8 uses per-expert global scales. " + "For integer quantization ('int'), omitting block_size means there is no blocking " + "and a whole column shares one scaling factor. ", + AttributeProto::INT, + OPTIONAL_VALUE) + .Attr("weights_prepacked", + "Expert-weight layout selector with quantization-specific meanings. For quant_type='int', " + "-1 and 1 select execution-provider-specific prepacked int4/int8 fc1/fc2 layouts. " + "0 selects raw [E,N,K/pack] tensors as produced by quantize_matmul_{4,8}bits. " + "For quant_type='nvfp4' on CUDA, -1 (default) and 0 select legacy N-packed [E,K,N/2] bytes; " + "1 selects K-packed row-major [E,N,K/2] bytes while retaining logical [E,K,N/2] dimensions. " + "NVFP4 block scales and global scales retain their original layouts and values. " + "Ignored for other quantization types. Supported values are -1, 0, and 1; defaults to -1.", + AttributeProto::INT, + static_cast(-1)) + .Input(0, + "input", + "2D packed token tensor with shape (total_tokens, hidden_size), where tokens from ragged sequences " + "may be concatenated without padding, or " + "3D tensor with shape (batch_size, sequence_length, hidden_size)", + "T") + .Input(1, + "router_probs", + "2D tensor with shape (total_tokens, num_experts), where total_tokens must match the flattened " + "token count of input", + "T") + .Input(2, + "fc1_experts_weights", + "3D tensor with shape (num_experts, fusion_size * inter_size, " + "hidden_size * effective_fc1_bits / 8). The last dimension must be byte-aligned. " + "The fusion_size is 2 for fused swiglu, or 1 otherwise. effective_fc1_bits is " + "fc1_expert_weight_bits when provided, otherwise expert_weight_bits.", + "T1") + .Input(3, + "fc1_scales", + "Optional weight scales. For quant_type='int', this is a 2D tensor with shape " + "(num_experts, fusion_size * inter_size), or a 3D tensor with shape " + "(num_experts, fusion_size * inter_size, hidden_size / block_size) when block_size is provided. " + "For quant_type='fp4' or 'wfp4afp8', this is a float8e8m0 MXFP block-scale tensor with shape " + "(num_experts, fusion_size * inter_size, hidden_size / 32). " + "For quant_type='nvfp4', this is a float8e4m3fn NVFP4 block-scale tensor with shape " + "(num_experts, fusion_size * inter_size, hidden_size / 16). " + "For quant_type='fp8' and block_size > 0, required float32, float16, or bfloat16 scales with shape " + "(num_experts, ceil(fusion_size * inter_size / block_size), ceil(hidden_size / block_size)).", + "T2", + OpSchema::Optional) + .Input(4, + "fc1_experts_bias", + "2D optional tensor with shape (num_experts, fusion_size * inter_size)", "T", OpSchema::Optional) + .Input(5, + "fc2_experts_weights", + "3D tensor with shape (num_experts, hidden_size, inter_size * effective_fc2_bits / 8). " + "The last dimension must be byte-aligned. effective_fc2_bits is fc2_expert_weight_bits " + "when provided, otherwise expert_weight_bits.", + "T1") + .Input(6, + "fc2_scales", + "Optional weight scales. For quant_type='int', this is a 2D tensor with shape " + "(num_experts, hidden_size), or a 3D tensor with shape " + "(num_experts, hidden_size, inter_size / block_size) when block_size is provided. " + "For quant_type='fp4' or 'wfp4afp8', this is a float8e8m0 MXFP block-scale tensor with shape " + "(num_experts, hidden_size, inter_size / 32). " + "For quant_type='nvfp4', this is a float8e4m3fn NVFP4 block-scale tensor with shape " + "(num_experts, hidden_size, inter_size / 16). " + "For quant_type='fp8' and block_size > 0, required float32, float16, or bfloat16 scales with shape " + "(num_experts, ceil(hidden_size / block_size), ceil(inter_size / block_size)).", + "T2", + OpSchema::Optional) + .Input(7, + "fc2_experts_bias", + "2D optional tensor with shape (num_experts, hidden_size)", + "T", + OpSchema::Optional) + .Input(8, + "fc3_experts_weights", + "3D optional tensor with shape (num_experts, inter_size, hidden_size * effective_fc3_bits / 8). " + "The last dimension must be byte-aligned. effective_fc3_bits is fc3_expert_weight_bits " + "when provided, otherwise expert_weight_bits.", + "T1", + OpSchema::Optional) + .Input(9, + "fc3_scales", + "Optional weight scales. For quant_type='int', this is a 2D tensor with shape " + "(num_experts, inter_size), or a 3D tensor with shape " + "(num_experts, inter_size, hidden_size / block_size) when block_size is provided. " + "For quant_type='fp4' or 'wfp4afp8', this is a float8e8m0 MXFP block-scale tensor with shape " + "(num_experts, inter_size, hidden_size / 32). " + "For quant_type='fp8' and block_size > 0, required when FC3 is present, with floating-point scales shaped " + "(num_experts, ceil(inter_size / block_size), ceil(hidden_size / block_size)).", + "T2", + OpSchema::Optional) + .Input(10, + "fc3_experts_bias", + "2D optional tensor with shape (num_experts, inter_size)", + "T", + OpSchema::Optional) + .Input(11, + "fc1_zero_points", + "2D tensor with shape (num_experts, ceil(fusion_size * inter_size * effective_fc1_bits / 8)), or " + "3D tensor with shape (num_experts, fusion_size * inter_size, " + "ceil((hidden_size / block_size) * effective_fc1_bits / 8)) when block_size is provided.", + "T1", + OpSchema::Optional) + .Input(12, + "fc2_zero_points", + "2D tensor with shape (num_experts, ceil(hidden_size * effective_fc2_bits / 8)), or " + "3D tensor with shape (num_experts, hidden_size, " + "ceil((inter_size / block_size) * effective_fc2_bits / 8)) when block_size is provided.", + "T1", + OpSchema::Optional) + .Input(13, + "fc3_zero_points", + "2D optional tensor with shape (num_experts, ceil(inter_size * effective_fc3_bits / 8)), or " + "3D optional tensor with shape (num_experts, inter_size, " + "ceil((hidden_size / block_size) * effective_fc3_bits / 8)) when block_size is provided.", + "T1", + OpSchema::Optional) + .Input(14, + "router_weights", + "2D optional tensor with shape (num_tokens, num_experts). " + "When provided, router_probs is used only for Top-K expert selection, and router_weights is used " + "for aggregating expert outputs (the values at the selected expert indices are gathered and used as " + "mixing weights). This enables DeepSeek-style noaux_tc routing where different tensors are used for " + "selection and aggregation. When not provided, router_probs is used for both selection and aggregation " + "(backward compatible).", + "T", + OpSchema::Optional) + .Input(15, + "fc1_global_scale", + "1D optional tensor with shape (num_experts,). " + "Per-expert global weight scale for FC1. Required when quant_type is 'fp4', 'nvfp4', or 'wfp4afp8', " + "or 'fp8' with block_size <= 0. Not used for block-scaled FP8.", + "T4", + OpSchema::Optional) + .Input(16, + "fc2_global_scale", + "1D optional tensor with shape (num_experts,). " + "Per-expert global weight scale for FC2. Required when quant_type is 'fp4', 'nvfp4', or 'wfp4afp8', " + "or 'fp8' with block_size <= 0. Not used for block-scaled FP8.", + "T4", + OpSchema::Optional) + .Input(17, + "fc1_act_scale", + "1D optional tensor with shape (1,) or (num_experts,). Activation scale for FC1 FP8 activation modes.", + "T4", + OpSchema::Optional) + .Input(18, + "fc2_act_scale", + "1D optional tensor with shape (1,) or (num_experts,). Activation scale for FC2 FP8 activation modes.", + "T4", + OpSchema::Optional) + .Input(19, + "fc1_act_block_scale", + "3D optional float8e8m0 MXFP activation block-scale tensor for FC1 FP8 activation modes.", + "T2", + OpSchema::Optional) + .Input(20, + "fc2_act_block_scale", + "3D optional float8e8m0 MXFP activation block-scale tensor for FC2 FP8 activation modes.", + "T2", + OpSchema::Optional) + .Output(0, + "output", + "output tensor with same shape of input", + "T") + .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain input and output types to float tensors.") + .TypeConstraint("T1", {"tensor(uint8)", "tensor(float8e4m3fn)"}, + "Constrain quantized weight types. Integer and FP4 weights use uint8. FP8 weights use float8e4m3fn.") + .TypeConstraint("T2", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)", "tensor(float8e8m0)", "tensor(float8e4m3fn)"}, + "Constrain scale types. Float tensors are used for integer quantization and FP8 square-block scales. " + "Float8e8m0 tensors are used for MXFP4 block scales; float8e4m3fn tensors are used for NVFP4 block scales.") + .TypeConstraint("T4", {"tensor(float)"}, "Constrain FP4 global scale type to float32 tensors.") + .TypeAndShapeInferenceFunction(QMoETypeAndShapeInference); + schema.Attr("quant_type", quant_type_doc, AttributeProto::STRING, std::string("int")); + return schema; +} + +ONNX_MS_OPERATOR_SET_SCHEMA(QMoE, 1, QMoESchema()); ONNX_MS_OPERATOR_SET_SCHEMA(SampleOp, 1, OpSchema() diff --git a/onnxruntime/core/graph/contrib_ops/ms_opset.h b/onnxruntime/core/graph/contrib_ops/ms_opset.h index 71b3c719f9906..99ebb142af946 100644 --- a/onnxruntime/core/graph/contrib_ops/ms_opset.h +++ b/onnxruntime/core/graph/contrib_ops/ms_opset.h @@ -84,7 +84,6 @@ class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, MatMulFpQ4); #endif class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, MaxpoolWithMask); 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, 1, QMoE); @@ -221,7 +220,6 @@ class OpSet_Microsoft_ver1 { #endif fn(GetOpSchema()); fn(GetOpSchema()); - fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); diff --git a/onnxruntime/test/contrib_ops/cuda_kernels/mtp_moe_test.cc b/onnxruntime/test/contrib_ops/cuda_kernels/mtp_moe_test.cc deleted file mode 100644 index 688d0e9858632..0000000000000 --- a/onnxruntime/test/contrib_ops/cuda_kernels/mtp_moe_test.cc +++ /dev/null @@ -1,392 +0,0 @@ -// 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/mtp_moe.h" - -namespace onnxruntime { -namespace test { -namespace { - -namespace mtp_moe = onnxruntime::contrib::cuda::mtp_moe; - -constexpr float kGateUpGlobalScale = 0.0005173456156626344f; -constexpr float kDownGlobalScale = 0.00031534829759038985f; - -TEST(MtpMoeTest, RowCountBoundaries) { - EXPECT_TRUE(mtp_moe::ValidateRowCount(1).IsOK()); - EXPECT_TRUE(mtp_moe::ValidateRowCount(65535).IsOK()); - for (int rows : {-1, 0, 65536, std::numeric_limits::max()}) { - SCOPED_TRACE(rows); - const auto status = mtp_moe::ValidateRowCount(rows); - EXPECT_EQ(status.Code(), common::StatusCode::INVALID_ARGUMENT); - const auto launch_status = mtp_moe::LaunchNvfp4MtpMoe( - nullptr, nullptr, nullptr, nullptr, 1.0f, 1.0f, nullptr, nullptr, rows, nullptr); - EXPECT_EQ(launch_status.Code(), common::StatusCode::INVALID_ARGUMENT); - EXPECT_EQ(launch_status.ErrorMessage(), status.ErrorMessage()); - } -} - -void CheckCuda(cudaError_t error) { - if (error != cudaSuccess) { - throw std::runtime_error(cudaGetErrorString(error)); - } -} - -class DeviceBuffer { - public: - explicit DeviceBuffer(size_t size) { - CheckCuda(cudaMalloc(&data_, size)); - } - ~DeviceBuffer() { - if (data_ != nullptr) { - cudaFree(data_); - } - } - DeviceBuffer(const DeviceBuffer&) = delete; - DeviceBuffer& operator=(const DeviceBuffer&) = delete; - - template - T* As() const { - return reinterpret_cast(data_); - } - 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)); - } - void UploadAt(size_t offset, const void* source, size_t size) { - CheckCuda(cudaMemcpy(static_cast(data_) + offset, source, size, cudaMemcpyHostToDevice)); - } - - private: - void* data_ = nullptr; -}; - -float RoundBf16(float value) { - return __bfloat162float(__float2bfloat16(value)); -} - -float DecodeFp4(uint32_t code) { - constexpr std::array kMagnitude = {0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f}; - const float magnitude = kMagnitude[code & 7u]; - return (code & 8u) != 0 ? -magnitude : magnitude; -} - -float DecodeE4M3Scale(uint8_t encoded) { - // The test's synthetic block scales use only exactly representable 1.0 and 0.5 codes. - EXPECT_TRUE(encoded == 0x38 || encoded == 0x30); - return encoded == 0x38 ? 1.0f : 0.5f; -} - -float DecodePackedWeight(const uint8_t* row, int k, int column) { - const uint8_t packed = row[column >> 1]; - const uint32_t nibble = (column & 1) == 0 ? (packed & 0xfu) : (packed >> 4); - const int scale_index = k / 2 + column / 16; - return DecodeFp4(nibble) * DecodeE4M3Scale(row[scale_index]); -} - -float WarpReduceXor(std::array values) { - for (int offset = 16; offset > 0; offset >>= 1) { - const auto previous = values; - for (int lane = 0; lane < 32; ++lane) { - values[lane] = previous[lane] + previous[lane ^ offset]; - } - } - return values[0]; -} - -std::vector MakeNvfp4Weights(int experts, int rows, int k) { - const int row_bytes = k * 9 / 16; - const int packed_bytes = k / 2; - const int scale_bytes = k / 16; - std::vector weights(static_cast(experts) * rows * row_bytes); - for (int expert = 0; expert < experts; ++expert) { - for (int row = 0; row < rows; ++row) { - uint8_t* weight_row = weights.data() + - (static_cast(expert) * rows + row) * row_bytes; - for (int column = 0; column < k; ++column) { - const uint8_t code = static_cast((expert * 13 + row * 7 + column * 3 + column / 17) & 0xf); - uint8_t& packed = weight_row[column >> 1]; - packed = (column & 1) == 0 - ? static_cast((packed & 0xf0u) | code) - : static_cast((packed & 0x0fu) | (code << 4)); - } - for (int group = 0; group < scale_bytes; ++group) { - weight_row[packed_bytes + group] = ((expert + row + group) % 3 == 0) ? 0x30 : 0x38; - } - } - } - return weights; -} - -float ReferenceNvfp4Dot(const uint8_t* weight_row, const __nv_bfloat16* input, int k) { - const int chunks = k / 8; - std::array lane_sum{}; - for (int lane = 0; lane < 32; ++lane) { - float accumulator = 0.0f; - for (int chunk = lane; chunk < chunks; chunk += 32) { - for (int element = 0; element < 8; ++element) { - const int column = chunk * 8 + element; - const float weight = DecodePackedWeight(weight_row, k, column); - const float activation = __bfloat162float(input[column]); - accumulator = std::fma(weight, activation, accumulator); - } - } - lane_sum[lane] = accumulator; - } - return WarpReduceXor(lane_sum); -} - -struct ReferenceResult { - std::vector<__nv_bfloat16> output; - std::vector<__nv_bfloat16> hidden; - std::array experts{}; - std::array routing_weights{}; -}; - -ReferenceResult ReferenceMtpMoe(const std::vector<__nv_bfloat16>& input, - const std::vector<__nv_bfloat16>& router, - const std::vector& gate_up, - const std::vector& down) { - constexpr int hidden = mtp_moe::kHiddenSize; - constexpr int ffn = mtp_moe::kFfnSize; - constexpr int gate_up_rows = 2 * ffn; - constexpr int gate_row_bytes = hidden * 9 / 16; - constexpr int down_row_bytes = ffn * 9 / 16; - - ReferenceResult result; - std::array probabilities{}; - float maximum = -std::numeric_limits::infinity(); - for (int expert = 0; expert < mtp_moe::kExpertCount; ++expert) { - probabilities[expert] = __bfloat162float(router[expert]); - maximum = std::max(maximum, probabilities[expert]); - } - float denominator = 0.0f; - for (float& probability : probabilities) { - probability = std::exp(probability - maximum); - denominator += probability; - } - for (float& probability : probabilities) { - probability /= denominator; - } - std::array selected{}; - for (int slot = 0; slot < mtp_moe::kTopK; ++slot) { - int best = -1; - for (int expert = 0; expert < mtp_moe::kExpertCount; ++expert) { - if (!selected[expert] && - (best < 0 || probabilities[expert] > probabilities[best] || - (probabilities[expert] == probabilities[best] && expert < best))) { - best = expert; - } - } - selected[best] = true; - result.experts[slot] = best; - } - float top_sum = 0.0f; - for (int slot = 0; slot < mtp_moe::kTopK; ++slot) { - top_sum += probabilities[result.experts[slot]]; - } - for (int slot = 0; slot < mtp_moe::kTopK; ++slot) { - result.routing_weights[slot] = RoundBf16(probabilities[result.experts[slot]] / top_sum); - } - - result.hidden.resize(mtp_moe::kTopK * ffn); - for (int slot = 0; slot < mtp_moe::kTopK; ++slot) { - const int expert = result.experts[slot]; - const size_t expert_offset = static_cast(expert) * gate_up_rows * gate_row_bytes; - for (int output_row = 0; output_row < ffn; ++output_row) { - const uint8_t* gate_row = gate_up.data() + expert_offset + output_row * gate_row_bytes; - const uint8_t* up_row = gate_up.data() + expert_offset + (ffn + output_row) * gate_row_bytes; - const float gate = RoundBf16(ReferenceNvfp4Dot(gate_row, input.data(), hidden) * kGateUpGlobalScale); - const float up = RoundBf16(ReferenceNvfp4Dot(up_row, input.data(), hidden) * kGateUpGlobalScale); - const float silu = RoundBf16(gate / (1.0f + std::exp(-gate))); - result.hidden[slot * ffn + output_row] = __float2bfloat16(silu * up); - } - } - - result.output.resize(hidden); - for (int output_column = 0; output_column < hidden; ++output_column) { - float total = 0.0f; - for (int slot = 0; slot < mtp_moe::kTopK; ++slot) { - const int expert = result.experts[slot]; - const size_t expert_offset = static_cast(expert) * hidden * down_row_bytes; - const uint8_t* weight_row = down.data() + expert_offset + output_column * down_row_bytes; - const float down_sum = RoundBf16(ReferenceNvfp4Dot( - weight_row, result.hidden.data() + slot * ffn, ffn) * - kDownGlobalScale); - const float weighted = RoundBf16(down_sum * result.routing_weights[slot]); - total += weighted; - } - result.output[output_column] = __float2bfloat16(total); - } - return result; -} - -TEST(CUDA_EP_Unittest, MtpMoeRouterSelectsStableTopTenAndBf16Normalizes) { - constexpr int rows = 1; - std::vector<__nv_bfloat16> logits(mtp_moe::kExpertCount, __float2bfloat16(-8.0f)); - for (int expert = 0; expert < mtp_moe::kTopK; ++expert) { - logits[expert] = __float2bfloat16(0.0f); - } - - DeviceBuffer device_logits(logits.size() * sizeof(__nv_bfloat16)); - DeviceBuffer device_scratch(rows * mtp_moe::kScratchStride * sizeof(float)); - device_logits.Upload(logits.data(), logits.size() * sizeof(__nv_bfloat16)); - - cudaStream_t stream = nullptr; - CheckCuda(cudaStreamCreate(&stream)); - const Status status = mtp_moe::LaunchRouter(device_logits.As<__nv_bfloat16>(), - device_scratch.As(), rows, stream); - ASSERT_TRUE(status.IsOK()) << status.ErrorMessage(); - CheckCuda(cudaStreamSynchronize(stream)); - CheckCuda(cudaStreamDestroy(stream)); - - std::array scratch{}; - device_scratch.Download(scratch.data(), sizeof(scratch)); - float weight_sum = 0.0f; - for (int slot = 0; slot < mtp_moe::kTopK; ++slot) { - int expert = -1; - std::memcpy(&expert, &scratch[slot], sizeof(expert)); - EXPECT_EQ(expert, slot); - EXPECT_NEAR(scratch[mtp_moe::kTopK + slot], - __bfloat162float(__float2bfloat16(0.1f)), 0.0f); - weight_sum += scratch[mtp_moe::kTopK + slot]; - } - EXPECT_NEAR(weight_sum, 1.0f, 0.004f); -} - -TEST(CUDA_EP_Unittest, MtpMoeNvfp4MatchesReferenceArithmeticEndToEnd) { - constexpr int rows = 1; - constexpr int experts = mtp_moe::kTopK; - constexpr int hidden = mtp_moe::kHiddenSize; - constexpr int ffn = mtp_moe::kFfnSize; - constexpr int gate_row_bytes = hidden * 9 / 16; - constexpr int down_row_bytes = ffn * 9 / 16; - const size_t full_gate_bytes = static_cast(mtp_moe::kExpertCount) * 2 * ffn * gate_row_bytes; - const size_t full_down_bytes = static_cast(mtp_moe::kExpertCount) * hidden * down_row_bytes; - - size_t free_bytes = 0; - size_t total_bytes = 0; - CheckCuda(cudaMemGetInfo(&free_bytes, &total_bytes)); - constexpr size_t kMemoryReserve = size_t{128} << 20; - if (free_bytes < full_gate_bytes + full_down_bytes + kMemoryReserve) { - GTEST_SKIP() << "Insufficient CUDA memory for exact-shape MtpMoe weights"; - } - - std::vector<__nv_bfloat16> input(hidden); - for (int column = 0; column < hidden; ++column) { - input[column] = __float2bfloat16(static_cast((column % 23) - 11) * 6.0f); - } - std::vector<__nv_bfloat16> router(mtp_moe::kExpertCount, __float2bfloat16(-8.0f)); - for (int expert = 0; expert < experts; ++expert) { - router[expert] = __float2bfloat16(0.0f); - } - const std::vector gate_up = MakeNvfp4Weights(experts, 2 * ffn, hidden); - const std::vector down = MakeNvfp4Weights(experts, hidden, ffn); - const ReferenceResult expected = ReferenceMtpMoe(input, router, gate_up, down); - - DeviceBuffer device_input(input.size() * sizeof(__nv_bfloat16)); - DeviceBuffer device_router(router.size() * sizeof(__nv_bfloat16)); - DeviceBuffer device_gate_up(full_gate_bytes); - DeviceBuffer device_down(full_down_bytes); - DeviceBuffer device_output(hidden * sizeof(__nv_bfloat16)); - DeviceBuffer device_scratch(mtp_moe::kScratchStride * sizeof(float)); - const std::vector scratch_poison(mtp_moe::kScratchStride, -999.0f); - device_scratch.Upload(scratch_poison.data(), scratch_poison.size() * sizeof(float)); - device_input.Upload(input.data(), input.size() * sizeof(__nv_bfloat16)); - device_router.Upload(router.data(), router.size() * sizeof(__nv_bfloat16)); - for (int expert = 0; expert < experts; ++expert) { - const size_t gate_expert_bytes = static_cast(2 * ffn) * gate_row_bytes; - const size_t down_expert_bytes = static_cast(hidden) * down_row_bytes; - device_gate_up.UploadAt(static_cast(expert) * gate_expert_bytes, - gate_up.data() + static_cast(expert) * gate_expert_bytes, - gate_expert_bytes); - device_down.UploadAt(static_cast(expert) * down_expert_bytes, - down.data() + static_cast(expert) * down_expert_bytes, - down_expert_bytes); - } - - const Status status = mtp_moe::LaunchNvfp4MtpMoe( - device_input.As<__nv_bfloat16>(), device_router.As<__nv_bfloat16>(), - device_gate_up.As(), device_down.As(), - kGateUpGlobalScale, kDownGlobalScale, device_scratch.As(), - device_output.As<__nv_bfloat16>(), rows, nullptr); - ASSERT_TRUE(status.IsOK()) << status.ErrorMessage(); - - std::vector<__nv_bfloat16> actual_output(hidden); - std::vector actual_scratch(mtp_moe::kScratchStride * sizeof(float)); - device_output.Download(actual_output.data(), actual_output.size() * sizeof(__nv_bfloat16)); - device_scratch.Download(actual_scratch.data(), actual_scratch.size()); - for (size_t i = 20 * sizeof(float); i < actual_scratch.size(); ++i) { - if (i < 32 * sizeof(float) || - i >= 32 * sizeof(float) + mtp_moe::kTopK * ffn * sizeof(__nv_bfloat16)) { - EXPECT_EQ(actual_scratch[i], 0) << "scratch padding byte " << i; - } - } - - for (int slot = 0; slot < mtp_moe::kTopK; ++slot) { - int actual_expert = -1; - std::memcpy(&actual_expert, actual_scratch.data() + slot * sizeof(float), sizeof(actual_expert)); - EXPECT_EQ(actual_expert, expected.experts[slot]) << "routing slot " << slot; - float actual_weight = 0.0f; - std::memcpy(&actual_weight, - actual_scratch.data() + (mtp_moe::kTopK + slot) * sizeof(float), - sizeof(actual_weight)); - EXPECT_EQ(actual_weight, expected.routing_weights[slot]) << "routing weight " << slot; - - const auto* actual_hidden = reinterpret_cast( - actual_scratch.data() + 32 * sizeof(float)); - for (int column = 0; column < ffn; ++column) { - EXPECT_EQ(actual_hidden[slot * ffn + column], expected.hidden[slot * ffn + column]) - << "hidden intermediate slot=" << slot << " column=" << column; - } - } - for (int column = 0; column < hidden; ++column) { - const float actual = __bfloat162float(actual_output[column]); - const float reference = __bfloat162float(expected.output[column]); - EXPECT_NEAR(actual, reference, 0.01f + 0.01f * std::abs(reference)) << "output column " << column; - } -} - -TEST(CUDA_EP_Unittest, MtpMoeRouterNonFiniteRowsKeepExpertIndicesInRange) { - constexpr int rows = 4; - std::vector<__nv_bfloat16> logits(rows * mtp_moe::kExpertCount, __float2bfloat16(0.0f)); - logits[0] = __float2bfloat16(std::numeric_limits::quiet_NaN()); - logits[mtp_moe::kExpertCount + 17] = __float2bfloat16(std::numeric_limits::infinity()); - std::fill(logits.begin() + 2 * mtp_moe::kExpertCount, logits.begin() + 3 * mtp_moe::kExpertCount, - __float2bfloat16(-std::numeric_limits::infinity())); - logits[3 * mtp_moe::kExpertCount + 511] = __float2bfloat16(-std::numeric_limits::infinity()); - DeviceBuffer device_logits(logits.size() * sizeof(__nv_bfloat16)); - DeviceBuffer device_scratch(rows * mtp_moe::kScratchStride * sizeof(float)); - device_logits.Upload(logits.data(), logits.size() * sizeof(__nv_bfloat16)); - const Status status = mtp_moe::LaunchRouter(device_logits.As<__nv_bfloat16>(), - device_scratch.As(), rows, nullptr); - ASSERT_TRUE(status.IsOK()) << status.ErrorMessage(); - std::vector scratch(rows * mtp_moe::kScratchStride); - device_scratch.Download(scratch.data(), scratch.size() * sizeof(float)); - for (int row = 0; row < rows; ++row) { - for (int slot = 0; slot < mtp_moe::kTopK; ++slot) { - int expert = -1; - std::memcpy(&expert, &scratch[row * mtp_moe::kScratchStride + slot], sizeof(expert)); - EXPECT_EQ(expert, slot); - EXPECT_TRUE(std::isnan(scratch[row * mtp_moe::kScratchStride + mtp_moe::kTopK + slot])); - } - } -} - -} // namespace -} // namespace test -} // namespace onnxruntime 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 f19f629946634..c7e12ba539d40 100644 --- a/onnxruntime/test/contrib_ops/cuda_kernels/speculative_fold_test.cc +++ b/onnxruntime/test/contrib_ops/cuda_kernels/speculative_fold_test.cc @@ -125,7 +125,7 @@ TEST_F(FoldContribTest, AllPrefixLengthsAndDeferredRingPublication) { ASSERT_TRUE(f.Tail().IsOK()); Check(cudaDeviceSynchronize()); EXPECT_EQ(f.ring.Read()[8], 4); - EXPECT_EQ(f.done.Read()[0], 0); + EXPECT_EQ(f.done.Read()[0], 1); const auto cos = f.head_cos.Read(); EXPECT_EQ(cos[0], 13); EXPECT_EQ(cos[3 * heads - 1], 12 + 3 * heads); From 8193970b8f85c2c4c3bc0ff8f1190d6cacb8daeb Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Fri, 23 Oct 2026 23:56:00 +0000 Subject: [PATCH 07/15] [CUDA] Support MayInplace for FoldMid2/FoldTail under plugin CUDA EP Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../contrib_ops/cuda/llm/speculative_fold.cc | 14 +++++----- .../contrib_ops/cuda/llm/speculative_fold.cu | 27 ++++++++++++++----- .../contrib_ops/cuda/llm/speculative_fold.h | 6 +++-- 3 files changed, 31 insertions(+), 16 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cc b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cc index cb2b408b8011d..eaced1fb8e806 100644 --- a/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cc +++ b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cc @@ -77,7 +77,6 @@ class FoldMid2 final : public CudaKernel { if (context->OutputCount() > 8) { auto* updated_ctl = context->Output(8, ctl->Shape()); if (updated_ctl) { - ORT_ENFORCE(updated_ctl->DataRaw() == ctl->DataRaw(), "FoldMid2 requires aliased control output"); ctl_ptr = updated_ctl->MutableData(); publish_ring = false; } @@ -85,7 +84,6 @@ class FoldMid2 final : public CudaKernel { if (context->OutputCount() > 9) { auto* updated_context = context->Output(9, committed->Shape()); if (updated_context) { - ORT_ENFORCE(updated_context->DataRaw() == committed->DataRaw(), "FoldMid2 requires aliased context output"); context_ptr = updated_context->MutableData(); } } @@ -101,7 +99,7 @@ class FoldMid2 final : public CudaKernel { 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), - publish_ring); + publish_ring, ctl->Data(), static_cast(ctl->Shape().Size()), committed->Data()); } private: @@ -131,14 +129,14 @@ class FoldTail final : public CudaKernel { if (context->OutputCount() > 1) { auto* updated_ctl = context->Output(1, ctl->Shape()); if (updated_ctl) { - ORT_ENFORCE(updated_ctl->DataRaw() == ctl->DataRaw(), "FoldTail requires aliased control output"); ctl_ptr = updated_ctl->MutableData(); } } return speculative_fold::LaunchTail(token->Data(), status->Data(), static_cast(status->Shape().Size()), ctl_ptr, ptrs ? ptrs->Data() : nullptr, - done->MutableData(), slot_, Stream(context)); + done->MutableData(), slot_, Stream(context), + ctl->Data(), static_cast(ctl->Shape().Size())); } private: @@ -149,8 +147,8 @@ class FoldTail final : public CudaKernel { ONNX_OPERATOR_KERNEL_EX( FoldMid2, kMSDomain, 1, kCudaExecutionProvider, (*KernelDefBuilder::Create()) - .Alias(3, 8) - .Alias(4, 9) + .MayInplace(3, 8) + .MayInplace(4, 9) .TypeConstraint("I64", DataTypeImpl::GetTensorType()) .TypeConstraint("I32", DataTypeImpl::GetTensorType()) .TypeConstraint("B", DataTypeImpl::GetTensorType()), @@ -159,7 +157,7 @@ ONNX_OPERATOR_KERNEL_EX( ONNX_OPERATOR_KERNEL_EX( FoldTail, kMSDomain, 1, kCudaExecutionProvider, (*KernelDefBuilder::Create()) - .Alias(2, 1) + .MayInplace(2, 1) .TypeConstraint("I64", DataTypeImpl::GetTensorType()) .TypeConstraint("I32", DataTypeImpl::GetTensorType()), FoldTail); diff --git a/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cu b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cu index 31cb1a4d088a6..3f797ce05a20b 100644 --- a/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cu +++ b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cu @@ -15,7 +15,9 @@ __device__ bool ValidRing(const int64_t* ptrs, int64_t sequence) { } __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* input_ctl, int64_t* ctl, int ctl_count, + const int64_t* input_context, 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, @@ -26,6 +28,13 @@ __global__ void MidKernel(const int64_t* tokens, const int64_t* ids, const int32 __shared__ int64_t errors; __shared__ bool muted; 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_context != nullptr && input_context != context) { + context[0] = input_context[0]; + context[1] = input_context[1]; + } muted = ctl[3] != 0; length = muted ? 0 : ctl[1]; cache_length = muted ? 0 : ctl[7]; @@ -102,9 +111,13 @@ __global__ void HiddenKernel(const uint16_t* source, uint16_t* output, int64_t w } } -__global__ void TailKernel(const int64_t* token, const int32_t* status, int count, int64_t* ctl, +__global__ void TailKernel(const int64_t* token, const int32_t* status, int count, + const int64_t* input_ctl, int64_t* ctl, int ctl_count, const int64_t* ptrs, int32_t* done, int slot) { if (threadIdx.x != 0) return; + if (input_ctl != nullptr && input_ctl != ctl) { + for (int i = 0; i < ctl_count; ++i) ctl[i] = input_ctl[i]; + } done[0] = 0; if (ctl[3] != 0) return; if (ptrs != nullptr && (ctl[4] <= 0 || !ValidRing(ptrs, ctl[4] - 1))) { @@ -138,8 +151,9 @@ 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, 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, - bool publish_ring) { - MidKernel<<<1, 256, 0, stream>>>(tokens, ids, status, status_count, ctl, context, prefixes, prefix2, ptrs, + 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); ORT_RETURN_IF_ERROR(LaunchStatus(cudaGetLastError())); @@ -151,8 +165,9 @@ Status LaunchMid(const int64_t* tokens, const int64_t* ids, const int32_t* statu } 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); + const int64_t* ptrs, int32_t* done, int slot, cudaStream_t stream, + const int64_t* input_ctl, int ctl_count) { + TailKernel<<<1, 32, 0, stream>>>(token, status, status_count, input_ctl, ctl, ctl_count, ptrs, done, slot); return LaunchStatus(cudaGetLastError()); } } // 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 fe8ae64e39156..0508065c6249d 100644 --- a/onnxruntime/contrib_ops/cuda/llm/speculative_fold.h +++ b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.h @@ -16,9 +16,11 @@ 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, 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, - bool publish_ring = false); + bool publish_ring = false, const int64_t* input_ctl = nullptr, int ctl_count = 0, + const int64_t* input_context = nullptr); 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); + const int64_t* ptrs, int32_t* done, int slot, cudaStream_t stream, + const int64_t* input_ctl = nullptr, int ctl_count = 0); } // namespace onnxruntime::contrib::cuda::speculative_fold From df8004aa861725323d94ab57402e46add3cc8665 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sat, 24 Oct 2026 04:10:05 +0000 Subject: [PATCH 08/15] fix(test,docs): add missing test headers and sync speculative fold schema docs Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/ContribOperators.md | 8 +++++--- onnxruntime/test/contrib_ops/speculative_fold_op_test.cc | 2 ++ 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 170e7fc3baf7c..3e786e986b9f1 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -145,6 +145,7 @@ Do not modify directly.* * com.microsoft.WordConvEmbedding * experimental com.microsoft.IsAllFinite * experimental com.microsoft.QEmbedLayerNormalization + ## com.microsoft ### **com.microsoft.Attention** @@ -2179,9 +2180,9 @@ This version of the operator has been available since version 1 of the 'com.micr #### Attributes
-
head_capacity : int (default is 262144)
+
head_capacity : int
Maximum head-cache token capacity.
-
head_rows : int (default is 2)
+
head_rows : int
Fixed number of head rows, 2 through 4.
@@ -2260,7 +2261,7 @@ This version of the operator has been available since version 1 of the 'com.micr #### Attributes
-
ctl_slot : int (default is 2)
+
ctl_slot : int
Draft-control slot: -1 (none), 2, 8 or 9.
@@ -9078,3 +9079,4 @@ No versioning maintained for experimental ops.
Constrain input and output types to float32 tensors.
+ diff --git a/onnxruntime/test/contrib_ops/speculative_fold_op_test.cc b/onnxruntime/test/contrib_ops/speculative_fold_op_test.cc index 5b1a973347e41..94abdd959ab78 100644 --- a/onnxruntime/test/contrib_ops/speculative_fold_op_test.cc +++ b/onnxruntime/test/contrib_ops/speculative_fold_op_test.cc @@ -19,6 +19,8 @@ #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" From 58e0fb5c20146db231a6505b34acca0dacb1948c Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sat, 24 Oct 2026 15:17:45 +0000 Subject: [PATCH 09/15] feat(cuda): extend FoldMid2 and FoldTail to S<=8 (7-draft DFlash2 verification) - Widen FoldMid2 head_rows attribute and shape inference from 2..4 to 2..8, and FoldTail ctl_slot to {-1, 2, 8..13}. - Pack tokens[4..7] into the upper 32-bit halves of ring record words 1, 2, 3, and 5 while preserving the 8-word (64-byte) ring record layout and S<=4 compatibility. - Extend FoldContribTest to verify 8-row prefix acceptance, upper-32-bit token packing, and ctl_slot=13. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/ContribOperators.md | 8 +++---- docs/contrib_ops/cuda/speculative_fold.md | 10 +++++---- .../contrib_ops/cuda/llm/speculative_fold.cc | 6 ++--- .../contrib_ops/cuda/llm/speculative_fold.cu | 19 +++++++++------- .../core/graph/contrib_ops/bert_defs.cc | 12 +++++----- .../cuda_kernels/speculative_fold_test.cc | 22 ++++++++++++------- 6 files changed, 44 insertions(+), 33 deletions(-) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 3e786e986b9f1..26e0f4a0b6422 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -2171,7 +2171,7 @@ This version of the operator has been available since version 1 of the 'com.micr ### **com.microsoft.FoldMid2** - 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 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. The mid stage writes the payload; FoldTail publishes its sequence word after downstream validation when ptrs is passed to FoldTail, or FoldMid2 publishes it directly when updated_ctl/updated_ctx_commit are omitted. Muted prefill calls prepare inputs without committing state or touching ring metadata. Invalid device metadata sets error bits and mutes execution. + Commits a batch-one speculative prefix and prepares rotary/head inputs for two through eight head rows. 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. The mid stage writes the payload; FoldTail publishes its sequence word after downstream validation when ptrs is passed to FoldTail, or FoldMid2 publishes it directly when updated_ctl/updated_ctx_commit are omitted. Muted prefill calls prepare inputs without committing state or touching ring metadata. Invalid device metadata sets error bits and mutes execution. #### Version @@ -2183,7 +2183,7 @@ This version of the operator has been available since version 1 of the 'com.micr
head_capacity : int
Maximum head-cache token capacity.
head_rows : int
-
Fixed number of head rows, 2 through 4.
+
Fixed number of head rows, 2 through 8.
#### Inputs (10 - 11) @@ -2196,7 +2196,7 @@ This version of the operator has been available since version 1 of the 'com.micr
status : I32
Verifier status words.
ctl : I64
-
Mutable control state: at least 8 words for S2, 16 for S3/S4.
+
Mutable control state: at least 8 words for S2, 16 for S3..S8.
ctx_commit : I64
Mutable committed context, two words.
prefixes : I64
@@ -2262,7 +2262,7 @@ This version of the operator has been available since version 1 of the 'com.micr
ctl_slot : int
-
Draft-control slot: -1 (none), 2, 8 or 9.
+
Draft-control slot: -1 (none), 2, or 8..13.
#### Inputs (3 - 4) diff --git a/docs/contrib_ops/cuda/speculative_fold.md b/docs/contrib_ops/cuda/speculative_fold.md index 447e59edacef7..27c7f3e27b593 100644 --- a/docs/contrib_ops/cuda/speculative_fold.md +++ b/docs/contrib_ops/cuda/speculative_fold.md @@ -2,7 +2,7 @@ `com.microsoft::FoldMid2` and `FoldTail` schema version 1 implement batch-one speculative acceptance and result publication without a custom-op library. -This replaces the earlier two-row schema and supports two through four head +This replaces the earlier two-row schema and supports two through eight head rows, variable rotary width, configurable cache capacity, optional hidden-state padding, and deferred publication after downstream validation. @@ -12,12 +12,14 @@ 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. +to those device buffers. For three through eight head rows, control has at least +16 int64 words and prefix contexts contain `head_rows` pairs of int64 words. +When `rows > 4`, predictions `tokens[4..7]` are packed into the upper 32-bit +halves of ring record words `1, 2, 3, 5`. `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 latter aliases input 2. `ctl_slot` selects word -1 (none), 2, or 8..13. The tail writes the ring's sequence word only after a system fence, so host polling cannot observe a partial record. diff --git a/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cc b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cc index eaced1fb8e806..bd89d82b327c1 100644 --- a/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cc +++ b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cc @@ -19,7 +19,7 @@ class FoldMid2 final : public CudaKernel { explicit FoldMid2(const OpKernelInfo& info) : CudaKernel(info) { head_rows_ = SafeInt(info.GetAttrOrDefault("head_rows", 2)); capacity_ = SafeInt(info.GetAttrOrDefault("head_capacity", 262144)); - ORT_ENFORCE(head_rows_ >= 2 && head_rows_ <= 4 && capacity_ >= head_rows_, + ORT_ENFORCE(head_rows_ >= 2 && head_rows_ <= 8 && capacity_ >= head_rows_, "FoldMid2: invalid head_rows/head_capacity"); } @@ -111,8 +111,8 @@ class FoldTail final : public CudaKernel { public: explicit FoldTail(const OpKernelInfo& info) : CudaKernel(info) { slot_ = SafeInt(info.GetAttrOrDefault("ctl_slot", 2)); - ORT_ENFORCE(slot_ == -1 || slot_ == 2 || slot_ == 8 || slot_ == 9, - "FoldTail: ctl_slot must be -1, 2, 8 or 9"); + ORT_ENFORCE(slot_ == -1 || slot_ == 2 || (slot_ >= 8 && slot_ <= 13), + "FoldTail: ctl_slot must be -1, 2, or 8..13"); } Status ComputeInternal(OpKernelContext* context) const override { diff --git a/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cu b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cu index 3f797ce05a20b..5aa6c3fb8a697 100644 --- a/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cu +++ b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cu @@ -87,16 +87,19 @@ __global__ void MidKernel(const int64_t* tokens, const int64_t* ids, const int32 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; + 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[4] = validation; - record[5] = pending; + record[5] = pack_lo_hi(pending, rows > 7 ? tokens[7] : 0, rows > 7); 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; + record[7] = rows > 2 ? pack_lo_hi(tokens[2], rows > 3 ? tokens[3] : 0, rows > 3) : 0; if (publish_ring) { __threadfence_system(); record[0] = ctl[4]; diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index ed2cc225eba54..9ddbefba320b4 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -37,19 +37,19 @@ namespace contrib { ONNX_MS_OPERATOR_SET_SCHEMA( FoldMid2, 1, OpSchema() - .SetDoc("Commits a batch-one speculative prefix and prepares rotary/head inputs for two through four head rows. " + .SetDoc("Commits a batch-one speculative prefix and prepares rotary/head inputs for two through eight head rows. " "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. The mid stage writes the payload; FoldTail publishes " "its sequence word after downstream validation when ptrs is passed to FoldTail, or FoldMid2 publishes it directly when updated_ctl/updated_ctx_commit are omitted. " "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_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(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(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(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") @@ -83,7 +83,7 @@ ONNX_MS_OPERATOR_SET_SCHEMA( } const auto rows = getAttribute(ctx, "head_rows", 2); const auto capacity = getAttribute(ctx, "head_capacity", 262144); - if (rows < 2 || rows > 4 || capacity < rows || capacity > std::numeric_limits::max()) + if (rows < 2 || rows > 8 || 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); @@ -117,7 +117,7 @@ ONNX_MS_OPERATOR_SET_SCHEMA( .SetDoc("Updates the selected draft-control slot and publishes the result-ring sequence after validation. " "updated_ctl must alias ctl on CUDA when present. Muted calls do not access the ring. Invalid ring metadata " "sets error bit 16 and mutes execution. done is a dependency output set to 1 when unmuted and 0 when muted.") - .Attr("ctl_slot", "Draft-control slot: -1 (none), 2, 8 or 9.", AttributeProto::INT, int64_t{2}) + .Attr("ctl_slot", "Draft-control slot: -1 (none), 2, or 8..13.", 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") @@ -132,7 +132,7 @@ ONNX_MS_OPERATOR_SET_SCHEMA( propagateElemTypeFromInputToOutput(ctx, 2, 1); } const auto slot = getAttribute(ctx, "ctl_slot", 2); - if (slot != -1 && slot != 2 && slot != 8 && slot != 9) fail_shape_inference("FoldTail: invalid ctl_slot"); + if (slot != -1 && slot != 2 && (slot < 8 || slot > 13)) fail_shape_inference("FoldTail: invalid ctl_slot"); getOutputShape(ctx, 0)->add_dim()->set_dim_value(1); if (ctx.getNumOutputs() > 1 && ctx.hasOutput(1)) { propagateShapeFromInputToOutput(ctx, 2, 1); 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 c7e12ba539d40..e7252e4679295 100644 --- a/onnxruntime/test/contrib_ops/cuda_kernels/speculative_fold_test.cc +++ b/onnxruntime/test/contrib_ops/cuda_kernels/speculative_fold_test.cc @@ -41,10 +41,10 @@ struct Buffer { }; struct Fixture { - Buffer tokens{8}, ids{8}, ctl{16}, context{2}, prefixes{8}, prefix2{2}, ptrs{2}, ring{16}; + Buffer tokens{8}, ids{8}, ctl{16}, context{2}, prefixes{16}, 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}; + Buffer head_ids{8}, slots{8}, selection{1}; + Buffer cosine{16 * 3}, sine{16 * 3}, head_cos{24}, head_sin{24}, hidden{24}, out_hidden{24}; void Reset(int rows, int accepted, bool muted = false, int64_t sequence = 0) { std::vector predictions{10, 20, 30, 40, 50, 60, 70, 80}; @@ -60,7 +60,7 @@ struct Fixture { control[7] = 2; ctl.Upload(control); context.Upload({-1, -2}); - prefixes.Upload({100, 101, 200, 201, 300, 301, 400, 401}); + prefixes.Upload({100, 101, 200, 201, 300, 301, 400, 401, 500, 501, 600, 601, 700, 701, 800, 801}); prefix2.Upload({200, 201}); ptrs.Upload({reinterpret_cast(ring.data), 2}); ring.Upload(std::vector(16, -7)); @@ -98,7 +98,7 @@ class FoldContribTest : public testing::Test { TEST_F(FoldContribTest, AllPrefixLengthsAndDeferredRingPublication) { Fixture f; - for (int rows : {1, 2, 3, 4}) { + for (int rows : {1, 2, 3, 4, 8}) { const int heads = std::max(2, rows); for (int accepted = 0; accepted < rows; ++accepted) { SCOPED_TRACE(testing::Message() << rows << "/" << accepted); @@ -115,13 +115,19 @@ TEST_F(FoldContribTest, AllPrefixLengthsAndDeferredRingPublication) { 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(static_cast(static_cast(record[11]) & 0xffffffffu), accepted); + EXPECT_EQ(static_cast(static_cast(record[13]) & 0xffffffffu), 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); } + if (rows == 8) { + EXPECT_EQ(static_cast(static_cast(record[9]) >> 32), 50); + EXPECT_EQ(static_cast(static_cast(record[10]) >> 32), 60); + EXPECT_EQ(static_cast(static_cast(record[11]) >> 32), 70); + EXPECT_EQ(static_cast(static_cast(record[13]) >> 32), 80); + } ASSERT_TRUE(f.Tail().IsOK()); Check(cudaDeviceSynchronize()); EXPECT_EQ(f.ring.Read()[8], 4); @@ -176,7 +182,7 @@ TEST_F(FoldContribTest, InvalidDeviceMetadataMutesBeforeRingAccess) { TEST_F(FoldContribTest, TailSlotsAndValidationArePublished) { Fixture f; - for (int slot : {2, 8, 9}) { + for (int slot : {2, 8, 9, 13}) { 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); From 81bcfcd9717fd91dfe5ff60986e98242f0d634e3 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sun, 25 Oct 2026 00:38:26 +0000 Subject: [PATCH 10/15] Rename FoldMid2 and FoldTail to SpeculativePrefixCommit and SpeculativeStepPublish - Rename com.microsoft::FoldMid2 to com.microsoft::SpeculativePrefixCommit and com.microsoft::FoldTail to com.microsoft::SpeculativeStepPublish so the public schema names describe speculative prefix verification/commit and step publication rather than internal graph topology. - Update ContribOperators.md, OperatorKernels.md, speculative_fold.md, and speculative_fold_op_test.cc accordingly. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/ContribOperators.md | 260 +++++++++--------- docs/OperatorKernels.md | 4 +- docs/contrib_ops/cuda/speculative_fold.md | 8 +- .../contrib_ops/cuda/cuda_contrib_kernels.cc | 8 +- .../contrib_ops/cuda/llm/speculative_fold.cc | 36 +-- .../core/graph/contrib_ops/bert_defs.cc | 16 +- onnxruntime/core/graph/contrib_ops/ms_opset.h | 8 +- .../contrib_ops/speculative_fold_op_test.cc | 4 +- 8 files changed, 172 insertions(+), 172 deletions(-) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 26e0f4a0b6422..afc46cc488fc7 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -37,8 +37,6 @@ Do not modify directly.* * com.microsoft.EngramGate * com.microsoft.ExpandDims * com.microsoft.FastGelu - * com.microsoft.FoldMid2 - * com.microsoft.FoldTail * com.microsoft.FusedConv * com.microsoft.FusedGemm * com.microsoft.FusedMatMul @@ -133,6 +131,8 @@ Do not modify directly.* * com.microsoft.SparseAttentionIndexer * com.microsoft.SparsePagedAttention * com.microsoft.SparseToDenseMatMul + * com.microsoft.SpeculativePrefixCommit + * com.microsoft.SpeculativeStepPublish * com.microsoft.Tokenizer * com.microsoft.TorchEmbedding * com.microsoft.TransposeMatMul @@ -2169,134 +2169,6 @@ This version of the operator has been available since version 1 of the 'com.micr -### **com.microsoft.FoldMid2** - - Commits a batch-one speculative prefix and prepares rotary/head inputs for two through eight head rows. 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. The mid stage writes the payload; FoldTail publishes its sequence word after downstream validation when ptrs is passed to FoldTail, or FoldMid2 publishes it directly when updated_ctl/updated_ctx_commit are omitted. 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. - -#### Attributes - -
-
head_capacity : int
-
Maximum head-cache token capacity.
-
head_rows : int
-
Fixed number of head rows, 2 through 8.
-
- -#### Inputs (10 - 11) - -
-
next_tokens : I64
-
Verifier predictions, batch one.
-
input_ids : I64
-
Verifier input IDs.
-
status : I32
-
Verifier status words.
-
ctl : I64
-
Mutable control state: at least 8 words for S2, 16 for S3..S8.
-
ctx_commit : I64
-
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
-
Result ring address and positive record capacity.
-
cos_cache : B
-
Rotary cosine cache [positions, width].
-
sin_cache : B
-
Rotary sine cache [positions, width].
-
hidden (optional) : B
-
Optional hidden states [1, verifier_rows, ...].
-
- -#### Outputs (7 - 10) - -
-
head_ids : I64
-
Head token IDs [1, head_rows].
-
head_c : I32
-
Head committed length [1].
-
slot_ids : I64
-
Head slots [1, head_rows].
-
seqlens : I32
-
Head sequence length [1].
-
cos_out : B
-
Gathered rotary cosine [1, head_rows, width].
-
sin_out : B
-
Gathered rotary sine [1, head_rows, width].
-
sel : I64
-
Accepted-prefix length, or -1 on invalid metadata.
-
head_hidden (optional) : B
-
Hidden rows, repeating the last source row when necessary.
-
updated_ctl (optional) : I64
-
Optional updated control state, aliases ctl.
-
updated_ctx_commit (optional) : I64
-
Optional updated context, aliases ctx_commit.
-
- -#### Type Constraints - -
-
I64 : tensor(int64)
-
Token and control words.
-
I32 : tensor(int32)
-
Status and sequence lengths.
-
B : tensor(bfloat16)
-
Rotary and hidden tensors.
-
- - -### **com.microsoft.FoldTail** - - Updates the selected draft-control slot and publishes the result-ring sequence after validation. updated_ctl must alias ctl on CUDA when present. Muted calls do not access the ring. Invalid ring metadata sets error bit 16 and mutes execution. done is a dependency output set to 1 when unmuted and 0 when muted. - -#### Version - -This version of the operator has been available since version 1 of the 'com.microsoft' operator set. - -#### Attributes - -
-
ctl_slot : int
-
Draft-control slot: -1 (none), 2, or 8..13.
-
- -#### Inputs (3 - 4) - -
-
next_token : I64
-
Next draft token (or ignored selector in DFlash).
-
status : I32
-
Downstream validation status words.
-
ctl : I64
-
Mutable control state.
-
ptrs (optional) : I64
-
Optional result ring address and record capacity.
-
- -#### Outputs (1 - 2) - -
-
done : I32
-
Dependency output [1].
-
updated_ctl (optional) : I64
-
Optional updated control state, aliases ctl.
-
- -#### Type Constraints - -
-
I64 : tensor(int64)
-
Token and control words.
-
I32 : tensor(int32)
-
Status and dependency output.
-
- - ### **com.microsoft.FusedConv** The fused convolution operator schema is the same as Conv besides it includes an attribute @@ -8301,6 +8173,134 @@ This version of the operator has been available since version 1 of the 'com.micr +### **com.microsoft.SpeculativePrefixCommit** + + Commits a batch-one speculative prefix and prepares rotary/head inputs for two through eight head rows. 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. The mid stage writes the payload; SpeculativeStepPublish publishes its sequence word after downstream validation when ptrs is passed to SpeculativeStepPublish, or SpeculativePrefixCommit publishes it directly when updated_ctl/updated_ctx_commit are omitted. 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. + +#### Attributes + +
+
head_capacity : int
+
Maximum head-cache token capacity.
+
head_rows : int
+
Fixed number of head rows, 2 through 8.
+
+ +#### Inputs (10 - 11) + +
+
next_tokens : I64
+
Verifier predictions, batch one.
+
input_ids : I64
+
Verifier input IDs.
+
status : I32
+
Verifier status words.
+
ctl : I64
+
Mutable control state: at least 8 words for S2, 16 for S3..S8.
+
ctx_commit : I64
+
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
+
Result ring address and positive record capacity.
+
cos_cache : B
+
Rotary cosine cache [positions, width].
+
sin_cache : B
+
Rotary sine cache [positions, width].
+
hidden (optional) : B
+
Optional hidden states [1, verifier_rows, ...].
+
+ +#### Outputs (7 - 10) + +
+
head_ids : I64
+
Head token IDs [1, head_rows].
+
head_c : I32
+
Head committed length [1].
+
slot_ids : I64
+
Head slots [1, head_rows].
+
seqlens : I32
+
Head sequence length [1].
+
cos_out : B
+
Gathered rotary cosine [1, head_rows, width].
+
sin_out : B
+
Gathered rotary sine [1, head_rows, width].
+
sel : I64
+
Accepted-prefix length, or -1 on invalid metadata.
+
head_hidden (optional) : B
+
Hidden rows, repeating the last source row when necessary.
+
updated_ctl (optional) : I64
+
Optional updated control state, aliases ctl.
+
updated_ctx_commit (optional) : I64
+
Optional updated context, aliases ctx_commit.
+
+ +#### Type Constraints + +
+
I64 : tensor(int64)
+
Token and control words.
+
I32 : tensor(int32)
+
Status and sequence lengths.
+
B : tensor(bfloat16)
+
Rotary and hidden tensors.
+
+ + +### **com.microsoft.SpeculativeStepPublish** + + Updates the selected draft-control slot and publishes the result-ring sequence after validation. updated_ctl must alias ctl on CUDA when present. Muted calls do not access the ring. Invalid ring metadata sets error bit 16 and mutes execution. done is a dependency output set to 1 when unmuted and 0 when muted. + +#### Version + +This version of the operator has been available since version 1 of the 'com.microsoft' operator set. + +#### Attributes + +
+
ctl_slot : int
+
Draft-control slot: -1 (none), 2, or 8..13.
+
+ +#### Inputs (3 - 4) + +
+
next_token : I64
+
Next draft token (or ignored selector in DFlash).
+
status : I32
+
Downstream validation status words.
+
ctl : I64
+
Mutable control state.
+
ptrs (optional) : I64
+
Optional result ring address and record capacity.
+
+ +#### Outputs (1 - 2) + +
+
done : I32
+
Dependency output [1].
+
updated_ctl (optional) : I64
+
Optional updated control state, aliases ctl.
+
+ +#### Type Constraints + +
+
I64 : tensor(int64)
+
Token and control words.
+
I32 : tensor(int32)
+
Status and dependency output.
+
+ + ### **com.microsoft.Tokenizer** Tokenizer divides each string in X into a vector of strings along the last axis. Allowed input shapes are [C] and [N, C]. diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md index 980eb941fa8c0..08727616ad142 100644 --- a/docs/OperatorKernels.md +++ b/docs/OperatorKernels.md @@ -1117,8 +1117,6 @@ 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* 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)| @@ -1179,6 +1177,8 @@ The **OpSet Version** column uses the following notation: |SparseAttention|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* past_key:**T**
*in* past_value:**T**
*in* block_row_indices:**M**
*in* block_col_indices:**M**
*in* total_sequence_length:**M**
*in* key_total_sequence_lengths:**M**
*in* cos_cache:**T**
*in* sin_cache:**T**
*out* output:**T**
*out* present_key:**T**
*out* present_value:**T**|1+|**M** = tensor(int32)
**T** = tensor(bfloat16), tensor(float16)| |SparseAttentionIndexer|*in* query:**T**
*in* key:**T**
*in* query_norm_weight:**T**
*in* key_norm_weight:**T**
*in* cos_cache:**T**
*in* sin_cache:**T**
*in* mask:**TB**
*in* past_key:**T**
*in* gate:**T**
*in* position_bias:**T**
*in* head_weights:**T**
*in* position_ids:**I**
*in* past_sequence_length:**M**
*in* past_proj_buffer:**T**
*out* selected_indices:**M**
*out* present_key:**T**
*out* present_proj_buffer:**T**|1+|**I** = tensor(int64)
**M** = tensor(int32)
**T** = tensor(bfloat16), tensor(float), tensor(float16)
**TB** = tensor(int64)| |SparsePagedAttention|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* key_cache:**T_CACHE**
*in* value_cache:**T_CACHE**
*in* cumulative_sequence_length:**S**
*in* past_seqlens:**S**
*in* block_table:**S**
*in* slot_mapping:**S**
*in* selected_indices:**S**
*in* selected_counts:**S**
*in* auxiliary_key:**T**
*in* auxiliary_value:**T**
*in* auxiliary_lengths:**S**
*in* cos_cache:**T**
*in* sin_cache:**T**
*in* head_sink:**T**
*in* q_norm_weight:**T**
*in* k_norm_weight:**T**
*in* k_scale:**T_KV_SCALE**
*in* v_scale:**T_KV_SCALE**
*in* attention_metadata:**S**
*out* output:**T**
*out* key_cache_out:**T_CACHE**
*out* value_cache_out:**T_CACHE**|1+|**S** = tensor(int32)
**T** = tensor(bfloat16), tensor(float16)
**T_CACHE** = tensor(bfloat16), tensor(float16), tensor(int8)
**T_KV_SCALE** = tensor(float)| +|SpeculativePrefixCommit|*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)| +|SpeculativeStepPublish|*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)| |TransposeMatMul|*in* A:**T**
*in* B:**T**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)| |Trilu|*in* X:**T**
*in* k:**tensor(int64)**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| |UnfoldTensor|*in* input:**T**
*out* output:**T**|1+|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(string), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| diff --git a/docs/contrib_ops/cuda/speculative_fold.md b/docs/contrib_ops/cuda/speculative_fold.md index 27c7f3e27b593..5a0e4eb683644 100644 --- a/docs/contrib_ops/cuda/speculative_fold.md +++ b/docs/contrib_ops/cuda/speculative_fold.md @@ -1,12 +1,12 @@ -# CUDA speculative fold operators +# CUDA speculative prefix commit and step publish operators -`com.microsoft::FoldMid2` and `FoldTail` schema version 1 implement batch-one +`com.microsoft::SpeculativePrefixCommit` and `SpeculativeStepPublish` schema version 1 implement batch-one speculative acceptance and result publication without a custom-op library. This replaces the earlier two-row schema and supports two through eight 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, +`SpeculativePrefixCommit` 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 @@ -17,7 +17,7 @@ to those device buffers. For three through eight head rows, control has at least When `rows > 4`, predictions `tokens[4..7]` are packed into the upper 32-bit halves of ring record words `1, 2, 3, 5`. -`FoldTail` takes a token, status words, the mid-stage control output, and +`SpeculativeStepPublish` 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 -1 (none), 2, or 8..13. The tail writes the ring's sequence word only after a system fence, so host polling cannot diff --git a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc index f42c7ccf5cd4e..f185b3fcb0529 100644 --- a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc @@ -100,8 +100,8 @@ class CUDA_ONNX_OP_TYPED_CLASS_NAME(1, MLFloat16, Crop); class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, MoE); class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, MoE); class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, MoE); -class CUDA_MS_OP_CLASS_NAME(1, FoldMid2); -class CUDA_MS_OP_CLASS_NAME(1, FoldTail); +class CUDA_MS_OP_CLASS_NAME(1, SpeculativePrefixCommit); +class CUDA_MS_OP_CLASS_NAME(1, SpeculativeStepPublish); 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); @@ -469,8 +469,8 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, 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 bd89d82b327c1..9af6ddece0bb3 100644 --- a/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cc +++ b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cc @@ -14,13 +14,13 @@ bool Count(const Tensor* tensor, int64_t count) { return tensor && tensor->Shape().Size() == count; } -class FoldMid2 final : public CudaKernel { +class SpeculativePrefixCommit final : public CudaKernel { public: - explicit FoldMid2(const OpKernelInfo& info) : CudaKernel(info) { + explicit SpeculativePrefixCommit(const OpKernelInfo& info) : CudaKernel(info) { head_rows_ = SafeInt(info.GetAttrOrDefault("head_rows", 2)); capacity_ = SafeInt(info.GetAttrOrDefault("head_capacity", 262144)); ORT_ENFORCE(head_rows_ >= 2 && head_rows_ <= 8 && capacity_ >= head_rows_, - "FoldMid2: invalid head_rows/head_capacity"); + "SpeculativePrefixCommit: invalid head_rows/head_capacity"); } Status ComputeInternal(OpKernelContext* context) const override { @@ -36,33 +36,33 @@ class FoldMid2 final : public CudaKernel { const auto* sine = context->Input(9); const auto* hidden = context->InputCount() > 10 ? context->Input(10) : nullptr; ORT_RETURN_IF_NOT(tokens && ids && status && ctl && committed && prefixes && prefix2 && ptrs && cosine && sine, - "FoldMid2: missing required input"); + "SpeculativePrefixCommit: 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"); + "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), - "FoldMid2: invalid control/context shape"); + "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, - "FoldMid2: rotary caches must have matching [positions, width] shapes"); + "SpeculativePrefixCommit: 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"); + "SpeculativePrefixCommit: 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"); + ORT_RETURN_IF_NOT(width > 0 && width <= INT64_MAX / head_rows_, "SpeculativePrefixCommit: invalid hidden width"); shape[1] = head_rows_; if (context->OutputCount() > 7) { head_hidden = context->Output(7, TensorShape(shape)); } } else if (context->OutputCount() > 7) { - ORT_RETURN_IF_NOT(!context->Output(7, TensorShape({0})), "FoldMid2: hidden output requires hidden input"); + ORT_RETURN_IF_NOT(!context->Output(7, TensorShape({0})), "SpeculativePrefixCommit: hidden output requires hidden input"); } auto* out_ids = context->Output(0, TensorShape({1, head_rows_})); auto* head_c = context->Output(1, TensorShape({1})); @@ -107,12 +107,12 @@ class FoldMid2 final : public CudaKernel { int capacity_; }; -class FoldTail final : public CudaKernel { +class SpeculativeStepPublish final : public CudaKernel { public: - explicit FoldTail(const OpKernelInfo& info) : CudaKernel(info) { + explicit SpeculativeStepPublish(const OpKernelInfo& info) : CudaKernel(info) { slot_ = SafeInt(info.GetAttrOrDefault("ctl_slot", 2)); ORT_ENFORCE(slot_ == -1 || slot_ == 2 || (slot_ >= 8 && slot_ <= 13), - "FoldTail: ctl_slot must be -1, 2, or 8..13"); + "SpeculativeStepPublish: ctl_slot must be -1, 2, or 8..13"); } Status ComputeInternal(OpKernelContext* context) const override { @@ -123,7 +123,7 @@ class FoldTail final : public CudaKernel { ORT_RETURN_IF_NOT(Count(token, 1) && status && status->Shape().Size() <= INT_MAX && ctl && ctl->Shape().Size() >= 8 && (slot_ < 0 || ctl->Shape().Size() > slot_) && (ptrs == nullptr || Count(ptrs, 2)), - "FoldTail: invalid token/status/control/pointer shape"); + "SpeculativeStepPublish: invalid token/status/control/pointer shape"); auto* done = context->Output(0, TensorShape({1})); int64_t* ctl_ptr = const_cast(ctl->Data()); if (context->OutputCount() > 1) { @@ -145,21 +145,21 @@ class FoldTail final : public CudaKernel { } // namespace ONNX_OPERATOR_KERNEL_EX( - FoldMid2, kMSDomain, 1, kCudaExecutionProvider, + SpeculativePrefixCommit, kMSDomain, 1, kCudaExecutionProvider, (*KernelDefBuilder::Create()) .MayInplace(3, 8) .MayInplace(4, 9) .TypeConstraint("I64", DataTypeImpl::GetTensorType()) .TypeConstraint("I32", DataTypeImpl::GetTensorType()) .TypeConstraint("B", DataTypeImpl::GetTensorType()), - FoldMid2); + SpeculativePrefixCommit); ONNX_OPERATOR_KERNEL_EX( - FoldTail, kMSDomain, 1, kCudaExecutionProvider, + SpeculativeStepPublish, kMSDomain, 1, kCudaExecutionProvider, (*KernelDefBuilder::Create()) .MayInplace(2, 1) .TypeConstraint("I64", DataTypeImpl::GetTensorType()) .TypeConstraint("I32", DataTypeImpl::GetTensorType()), - FoldTail); + SpeculativeStepPublish); } // namespace onnxruntime::contrib::cuda diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index 9ddbefba320b4..5b35a848b8acc 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -35,13 +35,13 @@ namespace onnxruntime { namespace contrib { ONNX_MS_OPERATOR_SET_SCHEMA( - FoldMid2, 1, + SpeculativePrefixCommit, 1, OpSchema() .SetDoc("Commits a batch-one speculative prefix and prepares rotary/head inputs for two through eight head rows. " "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. The mid stage writes the payload; FoldTail publishes " - "its sequence word after downstream validation when ptrs is passed to FoldTail, or FoldMid2 publishes it directly when updated_ctl/updated_ctx_commit are omitted. " + "aligned allocation of capacity*8 int64 words. The mid stage writes the payload; SpeculativeStepPublish publishes " + "its sequence word after downstream validation when ptrs is passed to SpeculativeStepPublish, or SpeculativePrefixCommit publishes it directly when updated_ctl/updated_ctx_commit are omitted. " "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}) @@ -84,7 +84,7 @@ ONNX_MS_OPERATOR_SET_SCHEMA( const auto rows = getAttribute(ctx, "head_rows", 2); const auto capacity = getAttribute(ctx, "head_capacity", 262144); if (rows < 2 || rows > 8 || capacity < rows || capacity > std::numeric_limits::max()) - fail_shape_inference("FoldMid2: invalid head_rows/head_capacity"); + fail_shape_inference("SpeculativePrefixCommit: invalid head_rows/head_capacity"); for (int output : {0, 2}) { auto* shape = getOutputShape(ctx, output); shape->add_dim()->set_dim_value(1); @@ -93,7 +93,7 @@ ONNX_MS_OPERATOR_SET_SCHEMA( 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"); + 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()->set_dim_value(rows); @@ -103,7 +103,7 @@ ONNX_MS_OPERATOR_SET_SCHEMA( } if (ctx.getNumOutputs() > 7 && ctx.hasOutput(7) && hasInputShape(ctx, 10)) { auto shape = getInputShape(ctx, 10); - if (shape.dim_size() < 3) fail_shape_inference("FoldMid2: hidden must have rank at least 3"); + if (shape.dim_size() < 3) fail_shape_inference("SpeculativePrefixCommit: hidden must have rank at least 3"); shape.mutable_dim(1)->set_dim_value(rows); updateOutputShape(ctx, 7, shape); } @@ -112,7 +112,7 @@ ONNX_MS_OPERATOR_SET_SCHEMA( })); ONNX_MS_OPERATOR_SET_SCHEMA( - FoldTail, 1, + SpeculativeStepPublish, 1, OpSchema() .SetDoc("Updates the selected draft-control slot and publishes the result-ring sequence after validation. " "updated_ctl must alias ctl on CUDA when present. Muted calls do not access the ring. Invalid ring metadata " @@ -132,7 +132,7 @@ ONNX_MS_OPERATOR_SET_SCHEMA( propagateElemTypeFromInputToOutput(ctx, 2, 1); } const auto slot = getAttribute(ctx, "ctl_slot", 2); - if (slot != -1 && slot != 2 && (slot < 8 || slot > 13)) fail_shape_inference("FoldTail: invalid ctl_slot"); + if (slot != -1 && slot != 2 && (slot < 8 || slot > 13)) fail_shape_inference("SpeculativeStepPublish: invalid ctl_slot"); getOutputShape(ctx, 0)->add_dim()->set_dim_value(1); if (ctx.getNumOutputs() > 1 && ctx.hasOutput(1)) { propagateShapeFromInputToOutput(ctx, 2, 1); diff --git a/onnxruntime/core/graph/contrib_ops/ms_opset.h b/onnxruntime/core/graph/contrib_ops/ms_opset.h index 689d1ef74e1ea..b5acb5d6f5e8e 100644 --- a/onnxruntime/core/graph/contrib_ops/ms_opset.h +++ b/onnxruntime/core/graph/contrib_ops/ms_opset.h @@ -84,8 +84,8 @@ class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, MatMulFpQ4); #endif class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, MaxpoolWithMask); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, MoE); -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, 1, SpeculativePrefixCommit); +class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, SpeculativeStepPublish); 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); @@ -222,8 +222,8 @@ class OpSet_Microsoft_ver1 { #endif fn(GetOpSchema()); fn(GetOpSchema()); - fn(GetOpSchema()); - fn(GetOpSchema()); + fn(GetOpSchema()); + fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); diff --git a/onnxruntime/test/contrib_ops/speculative_fold_op_test.cc b/onnxruntime/test/contrib_ops/speculative_fold_op_test.cc index 94abdd959ab78..d27a9193c96f0 100644 --- a/onnxruntime/test/contrib_ops/speculative_fold_op_test.cc +++ b/onnxruntime/test/contrib_ops/speculative_fold_op_test.cc @@ -106,12 +106,12 @@ TEST(SpeculativeFoldTest, InferenceSessionLoadsVersionOneAndAcceptsMutedPrefillR 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", "", + graph.AddNode("mid", "SpeculativePrefixCommit", "", {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, downstream_status, mid_ctl, ptrs}, {done, updated_ctl}, nullptr, + graph.AddNode("tail", "SpeculativeStepPublish", "", {draft, downstream_status, mid_ctl, ptrs}, {done, updated_ctl}, nullptr, kMSDomain) .SetExecutionProviderType(kCudaExecutionProvider); graph.SetOutputs({head_ids, head_c, slots, lengths, cos_out, sin_out, selected, head_hidden, From a30a905fc4df542cf951543548390aaca16a5195 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sun, 25 Oct 2026 05:01:27 +0000 Subject: [PATCH 11/15] Remove unused SpeculativeStepPublish op and hidden-state padding from SpeculativePrefixCommit - Remove unused com.microsoft::SpeculativeStepPublish (formerly FoldTail) operator, TailKernel, and LaunchTail. - Remove unused hidden/head_hidden padding path (HiddenKernel) from SpeculativePrefixCommit and always publish the ring record directly. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/ContribOperators.md | 60 +------------- docs/OperatorKernels.md | 3 +- docs/contrib_ops/cuda/speculative_fold.md | 21 ++--- .../contrib_ops/cuda/cuda_contrib_kernels.cc | 2 - .../contrib_ops/cuda/llm/speculative_fold.cc | 78 ++----------------- .../contrib_ops/cuda/llm/speculative_fold.cu | 49 +----------- .../contrib_ops/cuda/llm/speculative_fold.h | 7 +- .../core/graph/contrib_ops/bert_defs.cc | 64 +++------------ onnxruntime/core/graph/contrib_ops/ms_opset.h | 2 - .../cuda_kernels/speculative_fold_test.cc | 47 ++++------- .../contrib_ops/speculative_fold_op_test.cc | 27 ++----- 11 files changed, 54 insertions(+), 306 deletions(-) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index afc46cc488fc7..d136e7809677a 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -132,7 +132,6 @@ Do not modify directly.* * com.microsoft.SparsePagedAttention * com.microsoft.SparseToDenseMatMul * com.microsoft.SpeculativePrefixCommit - * com.microsoft.SpeculativeStepPublish * com.microsoft.Tokenizer * com.microsoft.TorchEmbedding * com.microsoft.TransposeMatMul @@ -8175,7 +8174,7 @@ This version of the operator has been available since version 1 of the 'com.micr ### **com.microsoft.SpeculativePrefixCommit** - Commits a batch-one speculative prefix and prepares rotary/head inputs for two through eight head rows. 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. The mid stage writes the payload; SpeculativeStepPublish publishes its sequence word after downstream validation when ptrs is passed to SpeculativeStepPublish, or SpeculativePrefixCommit publishes it directly when updated_ctl/updated_ctx_commit are omitted. Muted prefill calls prepare inputs without committing state or touching ring metadata. Invalid device metadata sets error bits and mutes execution. + 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. #### Version @@ -8190,7 +8189,7 @@ This version of the operator has been available since version 1 of the 'com.micr
Fixed number of head rows, 2 through 8.
-#### Inputs (10 - 11) +#### Inputs
next_tokens : I64
@@ -8213,11 +8212,9 @@ This version of the operator has been available since version 1 of the 'com.micr
Rotary cosine cache [positions, width].
sin_cache : B
Rotary sine cache [positions, width].
-
hidden (optional) : B
-
Optional hidden states [1, verifier_rows, ...].
-#### Outputs (7 - 10) +#### Outputs (7 - 9)
head_ids : I64
@@ -8234,8 +8231,6 @@ This version of the operator has been available since version 1 of the 'com.micr
Gathered rotary sine [1, head_rows, width].
sel : I64
Accepted-prefix length, or -1 on invalid metadata.
-
head_hidden (optional) : B
-
Hidden rows, repeating the last source row when necessary.
updated_ctl (optional) : I64
Optional updated control state, aliases ctl.
updated_ctx_commit (optional) : I64
@@ -8250,54 +8245,7 @@ This version of the operator has been available since version 1 of the 'com.micr
I32 : tensor(int32)
Status and sequence lengths.
B : tensor(bfloat16)
-
Rotary and hidden tensors.
-
- - -### **com.microsoft.SpeculativeStepPublish** - - Updates the selected draft-control slot and publishes the result-ring sequence after validation. updated_ctl must alias ctl on CUDA when present. Muted calls do not access the ring. Invalid ring metadata sets error bit 16 and mutes execution. done is a dependency output set to 1 when unmuted and 0 when muted. - -#### Version - -This version of the operator has been available since version 1 of the 'com.microsoft' operator set. - -#### Attributes - -
-
ctl_slot : int
-
Draft-control slot: -1 (none), 2, or 8..13.
-
- -#### Inputs (3 - 4) - -
-
next_token : I64
-
Next draft token (or ignored selector in DFlash).
-
status : I32
-
Downstream validation status words.
-
ctl : I64
-
Mutable control state.
-
ptrs (optional) : I64
-
Optional result ring address and record capacity.
-
- -#### Outputs (1 - 2) - -
-
done : I32
-
Dependency output [1].
-
updated_ctl (optional) : I64
-
Optional updated control state, aliases ctl.
-
- -#### Type Constraints - -
-
I64 : tensor(int64)
-
Token and control words.
-
I32 : tensor(int32)
-
Status and dependency output.
+
Rotary tensors.
diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md index 08727616ad142..c1a6d1160bb78 100644 --- a/docs/OperatorKernels.md +++ b/docs/OperatorKernels.md @@ -1177,8 +1177,7 @@ The **OpSet Version** column uses the following notation: |SparseAttention|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* past_key:**T**
*in* past_value:**T**
*in* block_row_indices:**M**
*in* block_col_indices:**M**
*in* total_sequence_length:**M**
*in* key_total_sequence_lengths:**M**
*in* cos_cache:**T**
*in* sin_cache:**T**
*out* output:**T**
*out* present_key:**T**
*out* present_value:**T**|1+|**M** = tensor(int32)
**T** = tensor(bfloat16), tensor(float16)| |SparseAttentionIndexer|*in* query:**T**
*in* key:**T**
*in* query_norm_weight:**T**
*in* key_norm_weight:**T**
*in* cos_cache:**T**
*in* sin_cache:**T**
*in* mask:**TB**
*in* past_key:**T**
*in* gate:**T**
*in* position_bias:**T**
*in* head_weights:**T**
*in* position_ids:**I**
*in* past_sequence_length:**M**
*in* past_proj_buffer:**T**
*out* selected_indices:**M**
*out* present_key:**T**
*out* present_proj_buffer:**T**|1+|**I** = tensor(int64)
**M** = tensor(int32)
**T** = tensor(bfloat16), tensor(float), tensor(float16)
**TB** = tensor(int64)| |SparsePagedAttention|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* key_cache:**T_CACHE**
*in* value_cache:**T_CACHE**
*in* cumulative_sequence_length:**S**
*in* past_seqlens:**S**
*in* block_table:**S**
*in* slot_mapping:**S**
*in* selected_indices:**S**
*in* selected_counts:**S**
*in* auxiliary_key:**T**
*in* auxiliary_value:**T**
*in* auxiliary_lengths:**S**
*in* cos_cache:**T**
*in* sin_cache:**T**
*in* head_sink:**T**
*in* q_norm_weight:**T**
*in* k_norm_weight:**T**
*in* k_scale:**T_KV_SCALE**
*in* v_scale:**T_KV_SCALE**
*in* attention_metadata:**S**
*out* output:**T**
*out* key_cache_out:**T_CACHE**
*out* value_cache_out:**T_CACHE**|1+|**S** = tensor(int32)
**T** = tensor(bfloat16), tensor(float16)
**T_CACHE** = tensor(bfloat16), tensor(float16), tensor(int8)
**T_KV_SCALE** = tensor(float)| -|SpeculativePrefixCommit|*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)| -|SpeculativeStepPublish|*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)| +|SpeculativePrefixCommit|*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**
*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)| |TransposeMatMul|*in* A:**T**
*in* B:**T**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)| |Trilu|*in* X:**T**
*in* k:**tensor(int64)**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| |UnfoldTensor|*in* input:**T**
*out* output:**T**|1+|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(string), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| diff --git a/docs/contrib_ops/cuda/speculative_fold.md b/docs/contrib_ops/cuda/speculative_fold.md index 5a0e4eb683644..572c8f515b560 100644 --- a/docs/contrib_ops/cuda/speculative_fold.md +++ b/docs/contrib_ops/cuda/speculative_fold.md @@ -1,10 +1,10 @@ -# CUDA speculative prefix commit and step publish operators +# CUDA speculative prefix commit operator -`com.microsoft::SpeculativePrefixCommit` and `SpeculativeStepPublish` schema version 1 implement batch-one +`com.microsoft::SpeculativePrefixCommit` schema version 1 implements batch-one speculative acceptance and result publication without a custom-op library. This replaces the earlier two-row schema and supports two through eight head rows, variable rotary width, configurable cache capacity, optional hidden-state -padding, and deferred publication after downstream validation. +padding, and direct ring publication after verification. `SpeculativePrefixCommit` takes verifier token predictions, input token IDs, status words, control state, committed context, prefix contexts, result-ring metadata, @@ -15,13 +15,9 @@ The last two outputs **alias** inputs 3 and 4. The caller must serialize access to those device buffers. For three through eight head rows, control has at least 16 int64 words and prefix contexts contain `head_rows` pairs of int64 words. When `rows > 4`, predictions `tokens[4..7]` are packed into the upper 32-bit -halves of ring record words `1, 2, 3, 5`. - -`SpeculativeStepPublish` 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 -1 (none), 2, or 8..13. The tail writes -the ring's sequence word only after a system fence, so host polling cannot -observe a partial record. +halves of ring record words `1, 2, 3, 5`. The operator writes the ring's +sequence word only after a system fence, so host polling cannot observe a +partial record. The ring metadata is `[device_accessible_address, positive_record_capacity]`. Each record has eight int64 words. Its address must be aligned and its allocation @@ -30,9 +26,8 @@ the caller. Invalid length, rotary position, or ring metadata sets error bits 4, 8, or 16, mutes control, and avoids unsafe ring access. Muted prefill uses zero positions/cache length and does not commit contexts or access ring metadata. Muted prefill may contain more verifier rows than `head_rows`; non-muted -verification is limited to `head_rows`. The tail unions its downstream status -with the verifier status stored by the mid stage. The kernels do not allocate -or synchronize the stream and support CUDA graph replay. +verification is limited to `head_rows`. The kernel does not allocate +or synchronize the stream and supports CUDA graph replay. ## Migrating legacy packages diff --git a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc index f185b3fcb0529..8b3b80eabae94 100644 --- a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc @@ -101,7 +101,6 @@ class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, MoE); class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, MoE); class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, MoE); class CUDA_MS_OP_CLASS_NAME(1, SpeculativePrefixCommit); -class CUDA_MS_OP_CLASS_NAME(1, SpeculativeStepPublish); 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,7 +469,6 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) { 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 9af6ddece0bb3..c087b952c73c3 100644 --- a/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cc +++ b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cc @@ -34,7 +34,6 @@ class SpeculativePrefixCommit final : public CudaKernel { const auto* ptrs = context->Input(7); const auto* cosine = context->Input(8); const auto* sine = context->Input(9); - const auto* hidden = context->InputCount() > 10 ? context->Input(10) : nullptr; 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(); @@ -49,21 +48,6 @@ class SpeculativePrefixCommit final : public CudaKernel { 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"); - 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, - "SpeculativePrefixCommit: hidden state must match token rows"); - width = hidden->Shape().Size() / rows; - ORT_RETURN_IF_NOT(width > 0 && width <= INT64_MAX / head_rows_, "SpeculativePrefixCommit: invalid hidden width"); - shape[1] = head_rows_; - if (context->OutputCount() > 7) { - head_hidden = context->Output(7, TensorShape(shape)); - } - } else if (context->OutputCount() > 7) { - ORT_RETURN_IF_NOT(!context->Output(7, TensorShape({0})), "SpeculativePrefixCommit: 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_})); @@ -73,16 +57,14 @@ class SpeculativePrefixCommit final : public CudaKernel { auto* selection = context->Output(6, TensorShape({1})); int64_t* ctl_ptr = const_cast(ctl->Data()); int64_t* context_ptr = const_cast(committed->Data()); - bool publish_ring = true; - if (context->OutputCount() > 8) { - auto* updated_ctl = context->Output(8, ctl->Shape()); + if (context->OutputCount() > 7) { + auto* updated_ctl = context->Output(7, ctl->Shape()); if (updated_ctl) { ctl_ptr = updated_ctl->MutableData(); - publish_ring = false; } } - if (context->OutputCount() > 9) { - auto* updated_context = context->Output(9, committed->Shape()); + if (context->OutputCount() > 8) { + auto* updated_context = context->Output(8, committed->Shape()); if (updated_context) { context_ptr = updated_context->MutableData(); } @@ -96,70 +78,24 @@ class SpeculativePrefixCommit final : public CudaKernel { 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), - publish_ring, ctl->Data(), static_cast(ctl->Shape().Size()), committed->Data()); + true, ctl->Data(), static_cast(ctl->Shape().Size()), committed->Data()); } private: int head_rows_; int capacity_; }; - -class SpeculativeStepPublish final : public CudaKernel { - public: - explicit SpeculativeStepPublish(const OpKernelInfo& info) : CudaKernel(info) { - slot_ = SafeInt(info.GetAttrOrDefault("ctl_slot", 2)); - ORT_ENFORCE(slot_ == -1 || slot_ == 2 || (slot_ >= 8 && slot_ <= 13), - "SpeculativeStepPublish: ctl_slot must be -1, 2, or 8..13"); - } - - 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->InputCount() > 3 ? context->Input(3) : nullptr; - ORT_RETURN_IF_NOT(Count(token, 1) && status && status->Shape().Size() <= INT_MAX && - ctl && ctl->Shape().Size() >= 8 && (slot_ < 0 || ctl->Shape().Size() > slot_) && - (ptrs == nullptr || Count(ptrs, 2)), - "SpeculativeStepPublish: invalid token/status/control/pointer shape"); - auto* done = context->Output(0, TensorShape({1})); - int64_t* ctl_ptr = const_cast(ctl->Data()); - if (context->OutputCount() > 1) { - auto* updated_ctl = context->Output(1, ctl->Shape()); - if (updated_ctl) { - ctl_ptr = updated_ctl->MutableData(); - } - } - return speculative_fold::LaunchTail(token->Data(), status->Data(), - static_cast(status->Shape().Size()), ctl_ptr, - ptrs ? ptrs->Data() : nullptr, - done->MutableData(), slot_, Stream(context), - ctl->Data(), static_cast(ctl->Shape().Size())); - } - - private: - int slot_; -}; } // namespace ONNX_OPERATOR_KERNEL_EX( SpeculativePrefixCommit, kMSDomain, 1, kCudaExecutionProvider, (*KernelDefBuilder::Create()) - .MayInplace(3, 8) - .MayInplace(4, 9) + .MayInplace(3, 7) + .MayInplace(4, 8) .TypeConstraint("I64", DataTypeImpl::GetTensorType()) .TypeConstraint("I32", DataTypeImpl::GetTensorType()) .TypeConstraint("B", DataTypeImpl::GetTensorType()), SpeculativePrefixCommit); -ONNX_OPERATOR_KERNEL_EX( - SpeculativeStepPublish, kMSDomain, 1, kCudaExecutionProvider, - (*KernelDefBuilder::Create()) - .MayInplace(2, 1) - .TypeConstraint("I64", DataTypeImpl::GetTensorType()) - .TypeConstraint("I32", DataTypeImpl::GetTensorType()), - SpeculativeStepPublish); - } // 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 5aa6c3fb8a697..42a376291edac 100644 --- a/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cu +++ b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cu @@ -106,41 +106,6 @@ __global__ void MidKernel(const int64_t* tokens, const int64_t* ids, const int32 } } -__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, - const int64_t* input_ctl, int64_t* ctl, int ctl_count, - const int64_t* ptrs, int32_t* done, int slot) { - if (threadIdx.x != 0) return; - if (input_ctl != nullptr && input_ctl != ctl) { - for (int i = 0; i < ctl_count; ++i) ctl[i] = input_ctl[i]; - } - done[0] = 0; - if (ctl[3] != 0) return; - if (ptrs != nullptr && (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]; - if (slot >= 0) ctl[slot] = token[0]; - if (error) ctl[5] |= 2; - if (ptrs != nullptr) { - 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]; - } - done[0] = 1; -} - Status LaunchStatus(cudaError_t error) { return error == cudaSuccess ? Status::OK() : Status(common::StatusCategory::SYSTEM, static_cast(error), cudaGetErrorString(error)); @@ -152,25 +117,13 @@ Status LaunchMid(const int64_t* tokens, const int64_t* ids, const int32_t* statu 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, 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); - 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, - const int64_t* input_ctl, int ctl_count) { - TailKernel<<<1, 32, 0, stream>>>(token, status, status_count, input_ctl, ctl, ctl_count, ptrs, done, slot); return LaunchStatus(cudaGetLastError()); } + } // 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 0508065c6249d..5c9a4e4b506d9 100644 --- a/onnxruntime/contrib_ops/cuda/llm/speculative_fold.h +++ b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.h @@ -14,13 +14,8 @@ Status LaunchMid(const int64_t* tokens, const int64_t* ids, const int32_t* statu 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, - bool publish_ring = false, const int64_t* input_ctl = nullptr, int ctl_count = 0, + bool publish_ring = true, const int64_t* input_ctl = nullptr, int ctl_count = 0, const int64_t* input_context = nullptr); -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, - const int64_t* input_ctl = nullptr, int ctl_count = 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 5b35a848b8acc..cbc43caef5616 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -37,13 +37,12 @@ namespace contrib { ONNX_MS_OPERATOR_SET_SCHEMA( SpeculativePrefixCommit, 1, OpSchema() - .SetDoc("Commits a batch-one speculative prefix and prepares rotary/head inputs for two through eight head rows. " - "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. The mid stage writes the payload; SpeculativeStepPublish publishes " - "its sequence word after downstream validation when ptrs is passed to SpeculativeStepPublish, or SpeculativePrefixCommit publishes it directly when updated_ctl/updated_ctx_commit are omitted. " - "Muted prefill calls prepare inputs without committing " - "state or touching ring metadata. Invalid device metadata sets error bits and mutes execution.") + .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.") .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") @@ -56,7 +55,6 @@ ONNX_MS_OPERATOR_SET_SCHEMA( .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") @@ -64,23 +62,19 @@ ONNX_MS_OPERATOR_SET_SCHEMA( .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", "Optional updated control state, aliases ctl.", "I64", OpSchema::Optional) - .Output(9, "updated_ctx_commit", "Optional updated context, aliases ctx_commit.", "I64", OpSchema::Optional) + .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.") .TypeConstraint("I32", {"tensor(int32)"}, "Status and sequence lengths.") - .TypeConstraint("B", {"tensor(bfloat16)"}, "Rotary and hidden tensors.") + .TypeConstraint("B", {"tensor(bfloat16)"}, "Rotary tensors.") .TypeAndShapeInferenceFunction([](InferenceContext& ctx) { - for (int output : {0, 2, 6, 8, 9}) { + for (int output : {0, 2, 6, 7, 8}) { if (static_cast(output) < ctx.getNumOutputs() && ctx.hasOutput(output)) { propagateElemTypeFromInputToOutput(ctx, 0, output); } } for (int output : {1, 3}) propagateElemTypeFromInputToOutput(ctx, 2, output); for (int output : {4, 5}) propagateElemTypeFromInputToOutput(ctx, 8, output); - if (ctx.getNumOutputs() > 7 && ctx.hasOutput(7) && hasInputShape(ctx, 10)) { - propagateElemTypeFromInputToOutput(ctx, 10, 7); - } const auto rows = getAttribute(ctx, "head_rows", 2); const auto capacity = getAttribute(ctx, "head_capacity", 262144); if (rows < 2 || rows > 8 || capacity < rows || capacity > std::numeric_limits::max()) @@ -101,42 +95,8 @@ ONNX_MS_OPERATOR_SET_SCHEMA( updateOutputShape(ctx, 4, shape); updateOutputShape(ctx, 5, shape); } - if (ctx.getNumOutputs() > 7 && ctx.hasOutput(7) && hasInputShape(ctx, 10)) { - auto shape = getInputShape(ctx, 10); - if (shape.dim_size() < 3) fail_shape_inference("SpeculativePrefixCommit: hidden must have rank at least 3"); - shape.mutable_dim(1)->set_dim_value(rows); - updateOutputShape(ctx, 7, shape); - } - if (ctx.getNumOutputs() > 8 && ctx.hasOutput(8)) propagateShapeFromInputToOutput(ctx, 3, 8); - if (ctx.getNumOutputs() > 9 && ctx.hasOutput(9)) propagateShapeFromInputToOutput(ctx, 4, 9); - })); - -ONNX_MS_OPERATOR_SET_SCHEMA( - SpeculativeStepPublish, 1, - OpSchema() - .SetDoc("Updates the selected draft-control slot and publishes the result-ring sequence after validation. " - "updated_ctl must alias ctl on CUDA when present. Muted calls do not access the ring. Invalid ring metadata " - "sets error bit 16 and mutes execution. done is a dependency output set to 1 when unmuted and 0 when muted.") - .Attr("ctl_slot", "Draft-control slot: -1 (none), 2, or 8..13.", 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", "Optional result ring address and record capacity.", "I64", OpSchema::Optional) - .Output(0, "done", "Dependency output [1].", "I32") - .Output(1, "updated_ctl", "Optional updated control state, aliases ctl.", "I64", OpSchema::Optional) - .TypeConstraint("I64", {"tensor(int64)"}, "Token and control words.") - .TypeConstraint("I32", {"tensor(int32)"}, "Status and dependency output.") - .TypeAndShapeInferenceFunction([](InferenceContext& ctx) { - propagateElemTypeFromInputToOutput(ctx, 1, 0); - if (ctx.getNumOutputs() > 1 && ctx.hasOutput(1)) { - propagateElemTypeFromInputToOutput(ctx, 2, 1); - } - const auto slot = getAttribute(ctx, "ctl_slot", 2); - if (slot != -1 && slot != 2 && (slot < 8 || slot > 13)) fail_shape_inference("SpeculativeStepPublish: invalid ctl_slot"); - getOutputShape(ctx, 0)->add_dim()->set_dim_value(1); - if (ctx.getNumOutputs() > 1 && ctx.hasOutput(1)) { - propagateShapeFromInputToOutput(ctx, 2, 1); - } + if (ctx.getNumOutputs() > 7 && ctx.hasOutput(7)) propagateShapeFromInputToOutput(ctx, 3, 7); + if (ctx.getNumOutputs() > 8 && ctx.hasOutput(8)) propagateShapeFromInputToOutput(ctx, 4, 8); })); void DecoderAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& ctx) { diff --git a/onnxruntime/core/graph/contrib_ops/ms_opset.h b/onnxruntime/core/graph/contrib_ops/ms_opset.h index b5acb5d6f5e8e..1d0dd9e5581a4 100644 --- a/onnxruntime/core/graph/contrib_ops/ms_opset.h +++ b/onnxruntime/core/graph/contrib_ops/ms_opset.h @@ -85,7 +85,6 @@ class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, MatMulFpQ4); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, MaxpoolWithMask); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, MoE); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, SpeculativePrefixCommit); -class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, SpeculativeStepPublish); 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); @@ -223,7 +222,6 @@ class OpSet_Microsoft_ver1 { 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 e7252e4679295..7cc3fec77d675 100644 --- a/onnxruntime/test/contrib_ops/cuda_kernels/speculative_fold_test.cc +++ b/onnxruntime/test/contrib_ops/cuda_kernels/speculative_fold_test.cc @@ -81,10 +81,7 @@ struct Fixture { 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); + selection.data, rows, heads, 16, 3, 16, stream); } }; @@ -96,7 +93,7 @@ class FoldContribTest : public testing::Test { } }; -TEST_F(FoldContribTest, AllPrefixLengthsAndDeferredRingPublication) { +TEST_F(FoldContribTest, AllPrefixLengthsAndDirectRingPublication) { Fixture f; for (int rows : {1, 2, 3, 4, 8}) { const int heads = std::max(2, rows); @@ -114,7 +111,7 @@ TEST_F(FoldContribTest, AllPrefixLengthsAndDeferredRingPublication) { EXPECT_EQ(ctl[4], 4); EXPECT_EQ(ctl[6], accepted); auto record = f.ring.Read(); - EXPECT_EQ(record[8], -7); + EXPECT_EQ(record[8], 4); EXPECT_EQ(static_cast(static_cast(record[11]) & 0xffffffffu), accepted); EXPECT_EQ(static_cast(static_cast(record[13]) & 0xffffffffu), 10 * (accepted + 1)); EXPECT_EQ(record[14], 4 + accepted); @@ -128,17 +125,9 @@ TEST_F(FoldContribTest, AllPrefixLengthsAndDeferredRingPublication) { EXPECT_EQ(static_cast(static_cast(record[11]) >> 32), 70); EXPECT_EQ(static_cast(static_cast(record[13]) >> 32), 80); } - ASSERT_TRUE(f.Tail().IsOK()); - Check(cudaDeviceSynchronize()); - EXPECT_EQ(f.ring.Read()[8], 4); - EXPECT_EQ(f.done.Read()[0], 1); 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); } } } @@ -148,7 +137,6 @@ TEST_F(FoldContribTest, MutedPrefillDoesNotAccessRingOrCommitContext) { 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); @@ -170,7 +158,6 @@ TEST_F(FoldContribTest, InvalidDeviceMetadataMutesBeforeRingAccess) { 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); @@ -180,24 +167,17 @@ TEST_F(FoldContribTest, InvalidDeviceMetadataMutesBeforeRingAccess) { } } -TEST_F(FoldContribTest, TailSlotsAndValidationArePublished) { +TEST_F(FoldContribTest, ValidationStatusIsPublished) { Fixture f; - for (int slot : {2, 8, 9, 13}) { - 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); - } + for (const auto& statuses : {std::vector{0, 0}, std::vector{1, 4}}) { + f.Reset(4, 2); + f.status.Upload(statuses); + ASSERT_TRUE(f.Mid(4, 4).IsOK()); + Check(cudaDeviceSynchronize()); + const int32_t combined = statuses[0] | statuses[1]; + EXPECT_EQ(f.ctl.Read()[5], combined ? 2 : 0); + EXPECT_EQ(f.ring.Read()[4], combined); + EXPECT_EQ(f.ring.Read()[0], 1); } } @@ -210,7 +190,6 @@ TEST_F(FoldContribTest, CapturedGraphReplaysWithFreshAcceptance) { 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}) { diff --git a/onnxruntime/test/contrib_ops/speculative_fold_op_test.cc b/onnxruntime/test/contrib_ops/speculative_fold_op_test.cc index d27a9193c96f0..7855fe391bc46 100644 --- a/onnxruntime/test/contrib_ops/speculative_fold_op_test.cc +++ b/onnxruntime/test/contrib_ops/speculative_fold_op_test.cc @@ -91,7 +91,6 @@ TEST(SpeculativeFoldTest, InferenceSessionLoadsVersionOneAndAcceptsMutedPrefillR auto* ptrs = arg("ptrs", i64, {2}); 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("slots", i64, {1, 2}); @@ -99,23 +98,15 @@ TEST(SpeculativeFoldTest, InferenceSessionLoadsVersionOneAndAcceptsMutedPrefillR 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_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}); + auto* updated_context = arg("updated_context", i64, {2}); graph.AddNode("mid", "SpeculativePrefixCommit", "", - {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}, + {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}, nullptr, kMSDomain) .SetExecutionProviderType(kCudaExecutionProvider); - graph.AddNode("tail", "SpeculativeStepPublish", "", {draft, downstream_status, mid_ctl, ptrs}, {done, updated_ctl}, nullptr, - kMSDomain) - .SetExecutionProviderType(kCudaExecutionProvider); - graph.SetOutputs({head_ids, head_c, slots, lengths, cos_out, sin_out, selected, head_hidden, - updated_ctl, updated_context, done}); + 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; @@ -140,7 +131,6 @@ TEST(SpeculativeFoldTest, InferenceSessionLoadsVersionOneAndAcceptsMutedPrefillR 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]; @@ -154,18 +144,15 @@ TEST(SpeculativeFoldTest, InferenceSessionLoadsVersionOneAndAcceptsMutedPrefillR 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; const std::vector output_names{ "head_ids", "head_c", "slots", "lengths", "cos_out", "sin_out", - "selected", "head_hidden", "updated_ctl", "updated_context", "done"}; + "selected", "updated_ctl", "updated_context"}; 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_EQ(outputs[1].Get().Data()[0], muted ? 0 : 2); } }; From 7ac2cd988cc65d12861d645986732b83f56a9bb9 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 26 Oct 2026 13:39:46 -0700 Subject: [PATCH 12/15] [CUDA] Support batch_size >= 1 in SpeculativePrefixCommit (#33626) ### Summary Extends com.microsoft::SpeculativePrefixCommit (introduced in #33474) to support batch_size >= 1 for batched speculative decoding: - Schema and shape inference (bert_defs.cc): Propagates the batch dimension from next_tokens ([batch, rows] or [rows]) to outputs 0..6 (head_ids, head_c, slot_ids, seqlens, cos_out, sin_out, sel). - CUDA kernel (speculative_fold.h/.cc/.cu): Launches MidKernel<<>> with per-slot strides for ctl, ctx_commit, prefixes, prefix2, per-slot result ring buffers, and per-slot muted flags. - Unit test (speculative_fold_op_test.cc): Adds SupportsBatchedSpeculativePrefixCommitS8 covering batch = 3, head_rows = 8 with mixed acceptance lengths and a muted slot. --------- 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 | 219 ++++++++++++++++++ 5 files changed, 361 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..84aa79f71f1ed 100644 --- a/onnxruntime/test/contrib_ops/speculative_fold_op_test.cc +++ b/onnxruntime/test/contrib_ops/speculative_fold_op_test.cc @@ -161,6 +161,225 @@ 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 661152686c24e292b7b93769cc2a69c2bf1b2e29 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 26 Oct 2026 22:57:13 +0000 Subject: [PATCH 13/15] [CUDA] Prune SpeculativePrefixCommit to the rc3 surface and require updated-state outputs Remove inputs, outputs, attributes and modes that the Qwen3.8-Flash-Next rc3 graph and the GenAI runtime do not consume: - rotary gather path (cos_cache/sin_cache inputs, cos_out/sin_out outputs), head-preparation outputs (head_ids, slot_ids, seqlens, sel) - redundant prefix2 input (duplicates prefixes[2..3]) and the size-2*batch ptrs layout (GenAI always passes {ring address, ring size}) - head_capacity attribute and the cache-capacity/rotary error bits - legacy head_rows 2 layout (8-word ctl, 2-word prefixes), the rows < head_rows clamp and the always-true publish_ring parameter updated_ctl and updated_ctx_commit are now required outputs that the kernel always writes (unchanged words are copied when the output does not alias the input), so state is never mutated through const_cast and a CPU-backed feed can no longer lose the update while a ring record is published. The kernel now runs one thread per batch item. Schema inference leaves the batch dimension unknown unless next_tokens has a known shape, and both inference and the kernel reject next_tokens ranks other than 1 and 2. head_rows defaults to 8 and must be in [3, 8]. Tests cover aliased and separate state outputs, batched per-item ring/state, muted prefill, invalid metadata, state carried across session runs, rank-3 rejection and shape inference. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../contrib_ops/cuda/llm/speculative_fold.cc | 109 ++-- .../contrib_ops/cuda/llm/speculative_fold.cu | 200 +++--- .../contrib_ops/cuda/llm/speculative_fold.h | 19 +- .../core/graph/contrib_ops/bert_defs.cc | 119 ++-- .../cuda_kernels/speculative_fold_test.cc | 272 +++++--- .../contrib_ops/speculative_fold_op_test.cc | 585 +++++++++--------- 6 files changed, 651 insertions(+), 653 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cc b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cc index e7cb836f2b4ab..3ccbdec16f350 100644 --- a/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cc +++ b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cc @@ -10,17 +10,20 @@ namespace onnxruntime::contrib::cuda { using namespace onnxruntime::cuda; namespace { -bool Count(const Tensor* tensor, int64_t count) { - return tensor && tensor->Shape().Size() == count; +constexpr int64_t kCtlWords = 16; + +// Aliased (identical) buffers are supported; partially overlapping ones are not. +bool IdenticalOrDisjoint(const Tensor& a, const Tensor& b) { + const auto* a_begin = static_cast(a.DataRaw()); + const auto* b_begin = static_cast(b.DataRaw()); + return a_begin == b_begin || a_begin + a.SizeInBytes() <= b_begin || b_begin + b.SizeInBytes() <= a_begin; } class SpeculativePrefixCommit final : public CudaKernel { public: explicit SpeculativePrefixCommit(const OpKernelInfo& info) : CudaKernel(info) { - head_rows_ = SafeInt(info.GetAttrOrDefault("head_rows", 2)); - capacity_ = SafeInt(info.GetAttrOrDefault("head_capacity", 262144)); - ORT_ENFORCE(head_rows_ >= 2 && head_rows_ <= 8 && capacity_ >= head_rows_, - "SpeculativePrefixCommit: invalid head_rows/head_capacity"); + head_rows_ = SafeInt(info.GetAttrOrDefault("head_rows", 8)); + ORT_ENFORCE(head_rows_ >= 3 && head_rows_ <= 8, "SpeculativePrefixCommit: head_rows must be in [3, 8]"); } Status ComputeInternal(OpKernelContext* context) const override { @@ -30,92 +33,56 @@ class SpeculativePrefixCommit final : public CudaKernel { 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); - ORT_RETURN_IF_NOT(tokens && ids && status && ctl && committed && prefixes && prefix2 && ptrs && cosine && sine, + const auto* ptrs = context->Input(6); + ORT_RETURN_IF_NOT(tokens && ids && status && ctl && committed && prefixes && ptrs, "SpeculativePrefixCommit: missing required input"); 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(); + const size_t tok_rank = tok_shape.NumDimensions(); + ORT_RETURN_IF_NOT(tok_rank == 1 || tok_rank == 2, + "SpeculativePrefixCommit: next_tokens must have rank 1 ([rows]) or 2 ([batch, rows])"); + const int64_t batch = tok_rank == 2 ? tok_shape[0] : int64_t{1}; + const int64_t rows = tok_shape[tok_rank - 1]; ORT_RETURN_IF_NOT(batch > 0 && batch <= INT_MAX && rows > 0 && rows <= INT_MAX && ids->Shape().Size() >= batch * rows && ids->Shape().Size() % batch == 0 && ids->Shape().Size() / batch <= INT_MAX && status->Shape().Size() <= INT_MAX, "SpeculativePrefixCommit: invalid token/status shape"); - 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 && + ORT_RETURN_IF_NOT(ctl->Shape().Size() >= batch * kCtlWords && 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 && + committed->Shape().Size() == 2 * batch && + prefixes->Shape().Size() >= batch * 2 * head_rows_ && 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)), + ptrs->Shape().Size() == 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({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) { - auto* updated_ctl = context->Output(7, ctl->Shape()); - if (updated_ctl) { - ctl_ptr = updated_ctl->MutableData(); - } - } - if (context->OutputCount() > 8) { - auto* updated_context = context->Output(8, committed->Shape()); - if (updated_context) { - 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(), - 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(), - static_cast(rows), head_rows_, static_cast(rope[0]), static_cast(rope[1]), capacity_, Stream(context), - true, ctl->Data(), ctl_stride, committed->Data(), - static_cast(batch), ids_stride, prefix_stride, prefix2_stride, ptrs_stride); + auto* head_c = context->Output(0, TensorShape({batch})); + auto* updated_ctl = context->Output(1, ctl->Shape()); + auto* updated_context = context->Output(2, committed->Shape()); + ORT_RETURN_IF_NOT(head_c && updated_ctl && updated_context, + "SpeculativePrefixCommit: updated_ctl and updated_ctx_commit outputs are required"); + ORT_RETURN_IF_NOT(IdenticalOrDisjoint(*ctl, *updated_ctl) && IdenticalOrDisjoint(*committed, *updated_context), + "SpeculativePrefixCommit: state outputs must alias or be disjoint from their inputs"); + return speculative_fold::LaunchPrefixCommit( + tokens->Data(), ids->Data(), status->Data(), + static_cast(status->Shape().Size()), ctl->Data(), updated_ctl->MutableData(), + committed->Data(), updated_context->MutableData(), prefixes->Data(), + ptrs->Data(), head_c->MutableData(), static_cast(rows), head_rows_, + static_cast(batch), static_cast(ids->Shape().Size() / batch), + static_cast(ctl->Shape().Size() / batch), static_cast(prefixes->Shape().Size() / batch), + Stream(context)); } private: int head_rows_; - int capacity_; }; } // namespace ONNX_OPERATOR_KERNEL_EX( SpeculativePrefixCommit, kMSDomain, 1, kCudaExecutionProvider, (*KernelDefBuilder::Create()) - .MayInplace(3, 7) - .MayInplace(4, 8) + .MayInplace(3, 1) + .MayInplace(4, 2) .TypeConstraint("I64", DataTypeImpl::GetTensorType()) - .TypeConstraint("I32", DataTypeImpl::GetTensorType()) - .TypeConstraint("B", DataTypeImpl::GetTensorType()), + .TypeConstraint("I32", DataTypeImpl::GetTensorType()), SpeculativePrefixCommit); } // 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 904f2e524f6cf..2c4c84c7ad938 100644 --- a/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cu +++ b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cu @@ -8,134 +8,102 @@ namespace onnxruntime::contrib::cuda::speculative_fold { 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; +constexpr int64_t kRecordWords = 8; +constexpr int64_t kRecordBytes = kRecordWords * static_cast(sizeof(int64_t)); + +// Error flags stored in ctl[5]. +constexpr int64_t kStatusFlag = 2; +constexpr int64_t kInvalidStateFlag = 4; +constexpr int64_t kInvalidRingFlag = 16; + +__device__ bool RingBase(const int64_t* ptrs, int b, int batch_size, int64_t* base) { + const int64_t address = ptrs[0]; + const int64_t records = ptrs[1]; + if (address <= 0 || address % alignof(int64_t) != 0 || records <= 0 || + records > LLONG_MAX / (kRecordBytes * batch_size)) { + return false; + } + if (address > LLONG_MAX - records * kRecordBytes * batch_size) return false; + *base = address + static_cast(b) * records * kRecordBytes; + return true; +} + +__device__ int64_t Pack(int64_t low, int64_t high) { + return static_cast(static_cast(static_cast(low)) | + (static_cast(static_cast(high)) << 32)); } -__global__ void MidKernel(const int64_t* tokens, const int64_t* ids, const int32_t* status, int status_count, - const int64_t* input_ctl, int64_t* ctl, int ctl_count, - const int64_t* input_context, 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, - bool publish_ring, int batch_size, int ids_stride, - int prefix_stride, int prefix2_stride, int ptrs_stride) { +// One thread per batch item: the work is a short sequential compare/commit and a ring write. +__global__ void PrefixCommitKernel(const int64_t* tokens, const int64_t* ids, const int32_t* status, int status_count, + const int64_t* ctl, int64_t* ctl_out, const int64_t* ctx, int64_t* ctx_out, + const int64_t* prefixes, const int64_t* ptrs, int32_t* head_c, + int rows, int head_rows, int batch_size, int ids_stride, int ctl_stride, + int prefix_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; + const int64_t* ids_b = ids + static_cast(b) * ids_stride; + const int64_t* prefixes_b = prefixes + static_cast(b) * prefix_stride; + const int64_t* ctl_in_b = ctl + static_cast(b) * ctl_stride; + int64_t* ctl_b = ctl_out + static_cast(b) * ctl_stride; + const int64_t* ctx_in_b = ctx + static_cast(b) * 2; + int64_t* ctx_b = ctx_out + static_cast(b) * 2; - __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_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 (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_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(local_ptrs, ctl_b[4])) errors |= 16; - if (rows > head_rows && !muted) errors |= 4; - if (errors) { - ctl_b[5] |= errors; - ctl_b[3] = 1; - } + if (ctl_b != ctl_in_b) { + for (int i = 0; i < ctl_stride; ++i) ctl_b[i] = ctl_in_b[i]; } - __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_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 (ctx_b != ctx_in_b) { + ctx_b[0] = ctx_in_b[0]; + ctx_b[1] = ctx_in_b[1]; } - if (threadIdx.x != 0) return; - if (errors) { - 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; + + const bool muted = ctl_b[3] != 0; + const int64_t length = muted ? 0 : ctl_b[1]; + const int64_t cache_length = muted ? 0 : ctl_b[7]; + int64_t errors = 0; + if (length < 0 || cache_length < 0 || cache_length > INT_MAX || (!muted && rows != head_rows)) { + errors |= kInvalidStateFlag; } - int accepted = 0; - if (!muted) { - while (accepted + 1 < rows && accepted + 1 < head_rows && - tokens_b[accepted] == ids_b[accepted + 1]) ++accepted; + int64_t ring = 0; + if (!muted && (!RingBase(ptrs, b, batch_size, &ring) || ctl_b[4] < 0 || ctl_b[4] == LLONG_MAX)) { + errors |= kInvalidRingFlag; } - for (int row = 0; row < head_rows; ++row) { - head_ids_b[row] = row < accepted ? ids_b[row + 1] : tokens_b[row < rows ? row : rows - 1]; - slots_b[row] = cache_length + row; + if (errors) { + ctl_b[5] |= errors; + ctl_b[3] = 1; + head_c[b] = 0; + return; } head_c[b] = static_cast(cache_length); - lengths[b] = static_cast(cache_length + head_rows - 1); - selection[b] = accepted; if (muted) return; + + int accepted = 0; + while (accepted + 1 < rows && tokens_b[accepted] == ids_b[accepted + 1]) ++accepted; int64_t validation = 0; for (int i = 0; i < status_count; ++i) validation |= status[i]; - 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]; + ctx_b[0] = prefixes_b[2 * accepted]; + ctx_b[1] = prefixes_b[2 * accepted + 1]; 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[5] |= validation ? kStatusFlag : 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_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 words (low | high << 32): [1] pred0|pred4, [2] pred1|pred5, [3] accepted|pred6, [5] pending|pred7, + // [7] pred2|pred3, [4] validation, [6] new length, [0] sequence (published last). + volatile int64_t* record = + reinterpret_cast(ring) + kRecordWords * ((ctl_b[4] - 1) % ptrs[1]); + record[1] = Pack(tokens_b[0], rows > 4 ? tokens_b[4] : 0); + record[2] = Pack(tokens_b[1], rows > 5 ? tokens_b[5] : 0); + record[3] = Pack(accepted, rows > 6 ? tokens_b[6] : 0); record[4] = validation; - record[5] = pack_lo_hi(pending, rows > 7 ? tokens_b[7] : 0, rows > 7); + record[5] = Pack(pending, rows > 7 ? tokens_b[7] : 0); 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_b[4]; - } + record[7] = Pack(tokens_b[2], rows > 3 ? tokens_b[3] : 0); + __threadfence_system(); + record[0] = ctl_b[4]; } Status LaunchStatus(cudaError_t error) { @@ -144,21 +112,15 @@ Status LaunchStatus(cudaError_t 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, - 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, - int batch_size, int ids_stride, int prefix_stride, int prefix2_stride, int ptrs_stride) { +Status LaunchPrefixCommit(const int64_t* tokens, const int64_t* ids, const int32_t* status, int status_count, + const int64_t* ctl, int64_t* ctl_out, const int64_t* ctx, int64_t* ctx_out, + const int64_t* prefixes, const int64_t* ptrs, int32_t* head_c, + int rows, int head_rows, int batch_size, int ids_stride, int ctl_stride, + int prefix_stride, cudaStream_t stream) { 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); + PrefixCommitKernel<<(batch_size), 1, 0, stream>>>( + tokens, ids, status, status_count, ctl, ctl_out, ctx, ctx_out, prefixes, ptrs, head_c, + rows, head_rows, batch_size, ids_stride, ctl_stride, prefix_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 4bab583cfd793..e36a715286185 100644 --- a/onnxruntime/contrib_ops/cuda/llm/speculative_fold.h +++ b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.h @@ -9,14 +9,15 @@ namespace onnxruntime::contrib::cuda::speculative_fold { -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, - 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, int batch_size = 1, int ids_stride = 0, - int prefix_stride = 0, int prefix2_stride = 0, int ptrs_stride = 0); +// ctl_out/ctx_out must either be the same buffers as ctl/ctx (in-place) or disjoint from them. The kernel +// copies the unchanged words when they differ, so callers always read the updated state from the outputs. +// ctl has ctl_stride (>= 16) words and prefixes has prefix_stride (>= 2 * head_rows) words per batch item. +// ptrs is {ring address, records per batch item}; batch item b uses the ring at +// address + b * records * 8 * sizeof(int64_t). +Status LaunchPrefixCommit(const int64_t* tokens, const int64_t* ids, const int32_t* status, int status_count, + const int64_t* ctl, int64_t* ctl_out, const int64_t* ctx, int64_t* ctx_out, + const int64_t* prefixes, const int64_t* ptrs, int32_t* head_c, + int rows, int head_rows, int batch_size, int ids_stride, int ctl_stride, + int prefix_stride, cudaStream_t stream); } // 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 f95d36abd0825..694416a232d88 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -37,72 +37,65 @@ namespace contrib { ONNX_MS_OPERATOR_SET_SCHEMA( SpeculativePrefixCommit, 1, OpSchema() - .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, 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 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 [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.") - .TypeConstraint("I32", {"tensor(int32)"}, "Status and sequence lengths.") - .TypeConstraint("B", {"tensor(bfloat16)"}, "Rotary tensors.") + .SetDoc("Commits the accepted prefix of a speculative verification step for each request in the batch and " + "publishes a result-ring record for the host. Acceptance compares the verifier predictions with the " + "draft tokens in input_ids. The operator is stateful: the control state and committed context are " + "read from ctl and ctx_commit and the updated state is written to updated_ctl and updated_ctx_commit, " + "which are required outputs. An output may alias its input (same device buffer) or be disjoint from it; " + "partial overlap is not supported. Callers that keep the state in persistent device buffers must bind " + "the outputs to those buffers, otherwise the updated state is lost. Callers must serialize access to " + "the state buffers. The result ring is described by ptrs and holds ptrs[1] records of eight int64 " + "words for each batch item, laid out consecutively per batch item starting at ptrs[0]; it must be a " + "device-accessible, 8-byte aligned allocation. Each record is published by writing its sequence word " + "last, after a system-scope fence. A request whose ctl[3] muted flag is set (for example during " + "prefill) does not commit state or touch the ring, and may use any number of rows. Invalid device " + "metadata sets error bits in ctl[5] (4: invalid row count or state, 16: invalid result ring) and mutes " + "the request. The status words are OR-ed together and a non-zero value sets bit 2 of ctl[5].") + .Attr("head_rows", "Number of verifier rows, 3 through 8. Non-muted requests must provide exactly this many " + "rows. Default is 8.", + AttributeProto::INT, static_cast(8)) + .Input(0, "next_tokens", "Verifier predictions with shape [batch, rows], or [rows] for batch one.", "I64") + .Input(1, "input_ids", "Verifier input IDs; the draft tokens are input_ids[b, 1:rows]. The number of " + "elements must be a multiple of the batch size and at least batch * rows.", + "I64") + .Input(2, "status", "Verifier status words, shared by all batch items.", "I32") + .Input(3, "ctl", + "Control state, 16 int64 words per batch item: [0] pending token, [1] length, [3] muted flag, " + "[4] ring sequence number, [5] error flags, [6] accepted count, [7] cache length. Other words are " + "preserved.", + "I64") + .Input(4, "ctx_commit", "Committed context, two int64 words per batch item.", "I64") + .Input(5, "prefixes", + "Context pairs per batch item: pair k (words 2k and 2k+1) is committed when k tokens are accepted. " + "At least 2 * head_rows words per batch item.", + "I64") + .Input(6, "ptrs", "Result ring address and number of records per batch item (must be positive).", "I64") + .Output(0, "head_c", "Cache length before the commit [batch]; zero for muted or invalid requests.", "I32") + .Output(1, "updated_ctl", "Updated control state, same shape as ctl. May alias ctl.", "I64") + .Output(2, "updated_ctx_commit", "Updated committed context, same shape as ctx_commit. May alias ctx_commit.", + "I64") + .TypeConstraint("I64", {"tensor(int64)"}, "Token, control and pointer words.") + .TypeConstraint("I32", {"tensor(int32)"}, "Status words and cache lengths.") .TypeAndShapeInferenceFunction([](InferenceContext& ctx) { - for (int output : {0, 2, 6, 7, 8}) { - if (static_cast(output) < ctx.getNumOutputs() && ctx.hasOutput(output)) { - propagateElemTypeFromInputToOutput(ctx, 0, output); + propagateElemTypeFromInputToOutput(ctx, 2, 0); + propagateElemTypeFromInputToOutput(ctx, 3, 1); + propagateElemTypeFromInputToOutput(ctx, 4, 2); + const auto rows = getAttribute(ctx, "head_rows", 8); + if (rows < 3 || rows > 8) fail_shape_inference("SpeculativePrefixCommit: head_rows must be in [3, 8]"); + if (hasInputShape(ctx, 3)) propagateShapeFromInputToOutput(ctx, 3, 1); + if (hasInputShape(ctx, 4)) propagateShapeFromInputToOutput(ctx, 4, 2); + if (hasInputShape(ctx, 0)) { + const auto& tokens_shape = getInputShape(ctx, 0); + if (tokens_shape.dim_size() != 1 && tokens_shape.dim_size() != 2) { + fail_shape_inference("SpeculativePrefixCommit: next_tokens must have rank 1 or 2"); + } + auto* head_c_shape = getOutputShape(ctx, 0); + if (tokens_shape.dim_size() == 2) { + *head_c_shape->add_dim() = tokens_shape.dim(0); + } else { + head_c_shape->add_dim()->set_dim_value(1); } } - for (int output : {1, 3}) propagateElemTypeFromInputToOutput(ctx, 2, output); - for (int output : {4, 5}) propagateElemTypeFromInputToOutput(ctx, 8, output); - const auto rows = getAttribute(ctx, "head_rows", 2); - 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() = batch_dim; - shape->add_dim()->set_dim_value(rows); - } - 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() = batch_dim; - shape.add_dim()->set_dim_value(rows); - *shape.add_dim() = rope.dim(1); - updateOutputShape(ctx, 4, shape); - updateOutputShape(ctx, 5, shape); - } - if (ctx.getNumOutputs() > 7 && ctx.hasOutput(7)) propagateShapeFromInputToOutput(ctx, 3, 7); - if (ctx.getNumOutputs() > 8 && ctx.hasOutput(8)) propagateShapeFromInputToOutput(ctx, 4, 8); })); void DecoderAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& ctx) { 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 7cc3fec77d675..24b5b101c0917 100644 --- a/onnxruntime/test/contrib_ops/cuda_kernels/speculative_fold_test.cc +++ b/onnxruntime/test/contrib_ops/cuda_kernels/speculative_fold_test.cc @@ -15,6 +15,11 @@ namespace onnxruntime::test { namespace { namespace fold = contrib::cuda::speculative_fold; +constexpr int kMaxRows = 8; +constexpr int kCtlWords = 16; +constexpr int kRecords = 2; +constexpr int64_t kUntouched = -7; + void Check(cudaError_t status) { if (status != cudaSuccess) throw std::runtime_error(cudaGetErrorString(status)); } @@ -28,8 +33,8 @@ struct Buffer { 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)); + ASSERT_LE(values.size(), count); + Check(cudaMemcpy(data, values.data(), values.size() * sizeof(T), cudaMemcpyHostToDevice)); } std::vector Read() const { std::vector values(count); @@ -40,49 +45,79 @@ struct Buffer { size_t count; }; +int64_t Prediction(int b, int k) { return 1000 * b + 10 * (k + 1); } +int64_t PrefixWord(int b, int k, int word) { return 1000 * b + 100 * (k + 1) + word; } +int32_t Low(int64_t word) { return static_cast(static_cast(word) & 0xffffffffu); } +int32_t High(int64_t word) { return static_cast(static_cast(word) >> 32); } + +// Batch item b has length 3 + b, cache length 2 + b and ring sequence `sequence + b`. All ring records start as -7. struct Fixture { - Buffer tokens{8}, ids{8}, ctl{16}, context{2}, prefixes{16}, prefix2{2}, ptrs{2}, ring{16}; - Buffer status{2}, head_c{1}, lengths{1}, done{1}; - Buffer head_ids{8}, slots{8}, selection{1}; - Buffer cosine{16 * 3}, sine{16 * 3}, head_cos{24}, head_sin{24}, hidden{24}, out_hidden{24}; + explicit Fixture(int batch = 1) + : batch(batch), + tokens(batch * kMaxRows), + ids(batch * (kMaxRows + 4)), + ctl(batch * kCtlWords), + ctl_out(batch * kCtlWords), + context(batch * 2), + context_out(batch * 2), + prefixes(batch * 2 * kMaxRows), + ptrs(2), + ring(batch * kRecords * 8), + status(2), + head_c(batch) {} - 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]; + void Reset(int rows, const std::vector& accepted, const std::vector& muted, int64_t sequence = 0, + int ids_padding = 0) { + ids_stride = rows + ids_padding; + std::vector predictions, input, control(batch * kCtlWords), contexts, prefix_values; + for (int b = 0; b < batch; ++b) { + for (int k = 0; k < rows; ++k) predictions.push_back(Prediction(b, k)); + std::vector item(ids_stride, -1); + item[0] = 99; + for (int i = 0; i < accepted[b]; ++i) item[i + 1] = Prediction(b, i); + input.insert(input.end(), item.begin(), item.end()); + for (int w = 0; w < kCtlWords; ++w) control[b * kCtlWords + w] = 5000 + w; + control[b * kCtlWords + 1] = 3 + b; + control[b * kCtlWords + 3] = muted[b] ? 1 : 0; + control[b * kCtlWords + 4] = sequence + b; + control[b * kCtlWords + 5] = 0; + control[b * kCtlWords + 6] = 0; + control[b * kCtlWords + 7] = 2 + b; + contexts.insert(contexts.end(), {-1 - b, -2 - b}); + for (int k = 0; k < kMaxRows; ++k) prefix_values.insert(prefix_values.end(), {PrefixWord(b, k, 0), PrefixWord(b, k, 1)}); + } 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, 500, 501, 600, 601, 700, 701, 800, 801}); - prefix2.Upload({200, 201}); - ptrs.Upload({reinterpret_cast(ring.data), 2}); - ring.Upload(std::vector(16, -7)); + ctl_out.Upload(std::vector(batch * kCtlWords, -99)); + context.Upload(contexts); + context_out.Upload(std::vector(batch * 2, -99)); + prefixes.Upload(prefix_values); + ptrs.Upload({reinterpret_cast(ring.data), kRecords}); + ring.Upload(std::vector(batch * kRecords * 8, kUntouched)); 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, rows, heads, 16, 3, 16, stream); + void Reset(int rows, int accepted, bool muted = false, int64_t sequence = 0) { + Reset(rows, std::vector(batch, accepted), std::vector(batch, muted), sequence); } + + Status Run(int rows, int head_rows, bool in_place = true, cudaStream_t stream = nullptr) { + in_place_ = in_place; + return fold::LaunchPrefixCommit(tokens.data, ids.data, status.data, 2, ctl.data, + in_place ? ctl.data : ctl_out.data, context.data, + in_place ? context.data : context_out.data, prefixes.data, ptrs.data, + head_c.data, rows, head_rows, batch, ids_stride, kCtlWords, 2 * kMaxRows, stream); + } + + std::vector Ctl() const { return in_place_ ? ctl.Read() : ctl_out.Read(); } + std::vector Context() const { return in_place_ ? context.Read() : context_out.Read(); } + + int batch; + int ids_stride{kMaxRows}; + bool in_place_{true}; + Buffer tokens, ids, ctl, ctl_out, context, context_out, prefixes, ptrs, ring; + Buffer status, head_c; }; class FoldContribTest : public testing::Test { @@ -93,77 +128,131 @@ class FoldContribTest : public testing::Test { } }; +void ExpectCommitted(const Fixture& f, int rows, int b, int accepted, int64_t sequence) { + const auto ctl = f.Ctl(); + const int64_t* c = ctl.data() + b * kCtlWords; + EXPECT_EQ(c[0], Prediction(b, accepted)); + EXPECT_EQ(c[1], 3 + b + accepted + 1); + EXPECT_EQ(c[3], 0); + EXPECT_EQ(c[4], sequence + b + 1); + EXPECT_EQ(c[5], 0); + EXPECT_EQ(c[6], accepted); + EXPECT_EQ(c[7], 2 + b + accepted + 1); + for (int w : {2, 8, 9, 10, 11, 12, 13, 14, 15}) EXPECT_EQ(c[w], 5000 + w) << "word " << w; + const auto context = f.Context(); + EXPECT_EQ(context[2 * b], PrefixWord(b, accepted, 0)); + EXPECT_EQ(context[2 * b + 1], PrefixWord(b, accepted, 1)); + EXPECT_EQ(f.head_c.Read()[b], 2 + b); + + const auto ring = f.ring.Read(); + const int64_t* record = ring.data() + (b * kRecords + (sequence + b) % kRecords) * 8; + EXPECT_EQ(record[0], sequence + b + 1); + EXPECT_EQ(Low(record[1]), Prediction(b, 0)); + EXPECT_EQ(Low(record[2]), Prediction(b, 1)); + EXPECT_EQ(Low(record[3]), accepted); + EXPECT_EQ(record[4], 0); + EXPECT_EQ(Low(record[5]), Prediction(b, accepted)); + EXPECT_EQ(record[6], 3 + b + accepted + 1); + EXPECT_EQ(Low(record[7]), Prediction(b, 2)); + EXPECT_EQ(High(record[7]), rows > 3 ? Prediction(b, 3) : 0); + EXPECT_EQ(High(record[1]), rows > 4 ? Prediction(b, 4) : 0); + EXPECT_EQ(High(record[2]), rows > 5 ? Prediction(b, 5) : 0); + EXPECT_EQ(High(record[3]), rows > 6 ? Prediction(b, 6) : 0); + EXPECT_EQ(High(record[5]), rows > 7 ? Prediction(b, 7) : 0); +} + TEST_F(FoldContribTest, AllPrefixLengthsAndDirectRingPublication) { Fixture f; - for (int rows : {1, 2, 3, 4, 8}) { - 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], 4); - EXPECT_EQ(static_cast(static_cast(record[11]) & 0xffffffffu), accepted); - EXPECT_EQ(static_cast(static_cast(record[13]) & 0xffffffffu), 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); + for (bool in_place : {true, false}) { + for (int rows : {3, 4, 8}) { + for (int accepted = 0; accepted < rows; ++accepted) { + SCOPED_TRACE(testing::Message() << "in_place=" << in_place << " rows=" << rows << " accepted=" << accepted); + f.Reset(rows, accepted, false, 3); + ASSERT_TRUE(f.Run(rows, rows, in_place).IsOK()); + Check(cudaDeviceSynchronize()); + ExpectCommitted(f, rows, 0, accepted, 3); } - if (rows == 8) { - EXPECT_EQ(static_cast(static_cast(record[9]) >> 32), 50); - EXPECT_EQ(static_cast(static_cast(record[10]) >> 32), 60); - EXPECT_EQ(static_cast(static_cast(record[11]) >> 32), 70); - EXPECT_EQ(static_cast(static_cast(record[13]) >> 32), 80); - } - const auto cos = f.head_cos.Read(); - EXPECT_EQ(cos[0], 13); - EXPECT_EQ(cos[3 * heads - 1], 12 + 3 * heads); } } } -TEST_F(FoldContribTest, MutedPrefillDoesNotAccessRingOrCommitContext) { +TEST_F(FoldContribTest, SeparateOutputsLeaveInputStateUntouched) { Fixture f; - f.Reset(8, 0, true); - f.ptrs.Upload({0, 0}); - ASSERT_TRUE(f.Mid(8, 4).IsOK()); + f.Reset(8, 5); + const auto ctl_before = f.ctl.Read(); + const auto context_before = f.context.Read(); + ASSERT_TRUE(f.Run(8, 8, false).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); + EXPECT_EQ(f.ctl.Read(), ctl_before); + EXPECT_EQ(f.context.Read(), context_before); + ExpectCommitted(f, 8, 0, 5, 0); +} + +TEST_F(FoldContribTest, BatchedRequestsUsePerItemStateAndRing) { + constexpr int kBatch = 3; + const std::vector accepted{3, 0, 1}; + const std::vector muted{false, true, false}; + Fixture f(kBatch); + for (bool in_place : {true, false}) { + f.Reset(kMaxRows, accepted, muted, 1, /*ids_padding=*/3); + ASSERT_TRUE(f.Run(kMaxRows, kMaxRows, in_place).IsOK()); + Check(cudaDeviceSynchronize()); + ExpectCommitted(f, kMaxRows, 0, accepted[0], 1); + ExpectCommitted(f, kMaxRows, 2, accepted[2], 1); + const auto ctl = f.Ctl(); + EXPECT_EQ(ctl[kCtlWords + 4], 2); + EXPECT_EQ(ctl[kCtlWords + 0], 5000); + EXPECT_EQ(f.Context()[2], -2); + EXPECT_EQ(f.Context()[3], -3); + EXPECT_EQ(f.head_c.Read()[1], 0); + const auto ring = f.ring.Read(); + for (int r = 0; r < kRecords * 8; ++r) EXPECT_EQ(ring[kRecords * 8 + r], kUntouched); + } +} + +TEST_F(FoldContribTest, MutedPrefillDoesNotAccessRingOrCommitContext) { + Fixture f; + for (bool in_place : {true, false}) { + f.Reset(8, 0, true); + f.ptrs.Upload({0, 0}); + // Muted prefill may use any number of rows; the kernel never reads the tokens. + ASSERT_TRUE(f.Run(64, 8, in_place).IsOK()); + Check(cudaDeviceSynchronize()); + EXPECT_EQ(f.Context(), (std::vector{-1, -2})); + const auto ctl = f.Ctl(); + EXPECT_EQ(ctl[1], 3); + EXPECT_EQ(ctl[3], 1); + EXPECT_EQ(ctl[4], 0); + EXPECT_EQ(ctl[5], 0); + EXPECT_EQ(ctl[7], 2); + EXPECT_EQ(f.head_c.Read()[0], 0); + EXPECT_EQ(f.ring.Read(), std::vector(kRecords * 8, kUntouched)); + } } TEST_F(FoldContribTest, InvalidDeviceMetadataMutesBeforeRingAccess) { Fixture f; - for (int kind = 0; kind < 4; ++kind) { + for (int kind = 0; kind < 6; ++kind) { + SCOPED_TRACE(kind); f.Reset(4, 3); auto ctl = f.ctl.Read(); - const int bit = kind < 2 ? 16 : (kind == 2 ? 8 : 4); + int64_t bit = 16; + int head_rows = 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 == 2) f.ptrs.Upload({reinterpret_cast(f.ring.data), INT64_MAX}); + if (kind >= 3) bit = 4; if (kind == 3) ctl[7] = -1; + if (kind == 4) ctl[1] = -1; + if (kind == 5) head_rows = 5; f.ctl.Upload(ctl); - ASSERT_TRUE(f.Mid(4, 4).IsOK()); + ASSERT_TRUE(f.Run(4, head_rows).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)); + EXPECT_EQ(f.head_c.Read()[0], 0); + EXPECT_EQ(f.Ctl()[5] & bit, bit); + EXPECT_EQ(f.Ctl()[3], 1); + EXPECT_EQ(f.Context(), (std::vector{-1, -2})); + EXPECT_EQ(f.ring.Read(), std::vector(kRecords * 8, kUntouched)); } } @@ -172,10 +261,10 @@ TEST_F(FoldContribTest, ValidationStatusIsPublished) { for (const auto& statuses : {std::vector{0, 0}, std::vector{1, 4}}) { f.Reset(4, 2); f.status.Upload(statuses); - ASSERT_TRUE(f.Mid(4, 4).IsOK()); + ASSERT_TRUE(f.Run(4, 4).IsOK()); Check(cudaDeviceSynchronize()); const int32_t combined = statuses[0] | statuses[1]; - EXPECT_EQ(f.ctl.Read()[5], combined ? 2 : 0); + EXPECT_EQ(f.Ctl()[5], combined ? 2 : 0); EXPECT_EQ(f.ring.Read()[4], combined); EXPECT_EQ(f.ring.Read()[0], 1); } @@ -189,15 +278,14 @@ TEST_F(FoldContribTest, CapturedGraphReplaysWithFreshAcceptance) { Check(cudaStreamCreate(&stream)); f.Reset(4, 0); Check(cudaStreamBeginCapture(stream, cudaStreamCaptureModeThreadLocal)); - ASSERT_TRUE(f.Mid(4, 4, stream).IsOK()); + ASSERT_TRUE(f.Run(4, 4, true, 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); + ExpectCommitted(f, 4, 0, accepted, 0); } Check(cudaGraphExecDestroy(executable)); Check(cudaGraphDestroy(graph)); diff --git a/onnxruntime/test/contrib_ops/speculative_fold_op_test.cc b/onnxruntime/test/contrib_ops/speculative_fold_op_test.cc index 84aa79f71f1ed..63501115ae5c2 100644 --- a/onnxruntime/test/contrib_ops/speculative_fold_op_test.cc +++ b/onnxruntime/test/contrib_ops/speculative_fold_op_test.cc @@ -7,7 +7,8 @@ #include #include -#include +#include +#include #include #include #include @@ -16,6 +17,7 @@ #include "core/graph/model.h" #include "core/session/inference_session.h" +#include "gmock/gmock.h" #include "gtest/gtest.h" #include "test/common/cuda_op_test_utils.h" #include "test/common/tensor_op_test_utils.h" @@ -28,17 +30,20 @@ namespace onnxruntime { namespace test { namespace { +constexpr int kCtlWords = 16; + +// Host-mapped result ring that the kernel can write through its device address. class FoldResultRing { public: - FoldResultRing() { - const cudaError_t error = cudaHostAlloc(reinterpret_cast(&host_), 32 * sizeof(int64_t), cudaHostAllocMapped); + explicit FoldResultRing(size_t words) : words_(words) { + const cudaError_t error = cudaHostAlloc(reinterpret_cast(&host_), words * sizeof(int64_t), cudaHostAllocMapped); ORT_ENFORCE(error == cudaSuccess, cudaGetErrorString(error)); const cudaError_t pointer_error = cudaHostGetDevicePointer(reinterpret_cast(&device_), host_, 0); if (pointer_error != cudaSuccess) { cudaFreeHost(host_); ORT_THROW(cudaGetErrorString(pointer_error)); } - std::fill_n(host_, 32, int64_t{-1}); + std::fill_n(host_, words, int64_t{-1}); } ~FoldResultRing() { cudaFreeHost(host_); } @@ -46,77 +51,106 @@ class FoldResultRing { ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(FoldResultRing); int64_t Address() const { return static_cast(reinterpret_cast(device_)); } - int64_t* Host() const { return host_; } + const int64_t* Host() const { return host_; } + size_t Words() const { return words_; } private: + size_t words_; int64_t* host_ = nullptr; int64_t* device_ = nullptr; }; -TEST(SpeculativeFoldTest, InferenceSessionLoadsVersionOneAndAcceptsMutedPrefillRows) { - if (!HasCudaEnvironment(800)) GTEST_SKIP() << "Requires CUDA SM80 or newer"; - auto provider = DefaultCudaExecutionProvider(); - ASSERT_NE(provider, nullptr); - FoldResultRing ring; - - 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 dims) -> NodeArg* { - ONNX_NAMESPACE::TypeProto type; - auto* tensor = type.mutable_tensor_type(); - tensor->set_elem_type(element_type); - 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); +// A one-node SpeculativePrefixCommit graph. Outputs are left without types so that inference provides them. +// token_dims == nullopt leaves next_tokens without a shape; negative dims are symbolic. +class FoldModel { + public: + FoldModel(std::optional> token_dims, int batch, int head_rows = 8) + : model_("speculative_fold", false, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(), + {{kMSDomain, 1}}, {}, DefaultLoggingManager().DefaultLogger()) { + auto& graph = model_.MainGraph(); + constexpr int i64 = ONNX_NAMESPACE::TensorProto_DataType_INT64; + constexpr int i32 = ONNX_NAMESPACE::TensorProto_DataType_INT32; + auto arg = [&](const char* name, int element_type, const std::optional>& dims) { + ONNX_NAMESPACE::TypeProto type; + auto* tensor = type.mutable_tensor_type(); + tensor->set_elem_type(element_type); + if (dims) { + 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); - }; + return &graph.GetOrCreateNodeArg(name, &type); + }; + auto* tokens = arg("tokens", i64, token_dims); + auto* ids = arg("ids", i64, std::vector{-1}); + auto* status = arg("status", i32, std::vector{1}); + auto* ctl = arg("ctl", i64, std::vector{batch * kCtlWords}); + auto* context = arg("context", i64, std::vector{batch * 2}); + auto* prefixes = arg("prefixes", i64, std::vector{batch * kCtlWords}); + auto* ptrs = arg("ptrs", i64, std::vector{2}); + head_c_ = &graph.GetOrCreateNodeArg("head_c", nullptr); + auto* updated_ctl = &graph.GetOrCreateNodeArg("updated_ctl", nullptr); + auto* updated_context = &graph.GetOrCreateNodeArg("updated_context", nullptr); + + 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(head_rows); + NodeAttributes attrs{{"head_rows", head_rows_attr}}; + graph.AddNode("commit", "SpeculativePrefixCommit", "", + {tokens, ids, status, ctl, context, prefixes, ptrs}, + {head_c_, updated_ctl, updated_context}, &attrs, kMSDomain) + .SetExecutionProviderType(kCudaExecutionProvider); + graph.SetOutputs({head_c_, updated_ctl, updated_context}); + resolve_status_ = graph.Resolve(); + } + + const Status& ResolveStatus() const { return resolve_status_; } + const NodeArg* HeadC() const { return head_c_; } + + Status CreateSession(std::unique_ptr& session) { + auto provider = DefaultCudaExecutionProvider(); + if (!provider) return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "CUDA provider is unavailable"); + std::stringstream stream; + if (!model_.ToProto().SerializeToOstream(&stream)) return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "serialize failed"); + session = std::make_unique(SessionOptions{}, GetEnvironment()); + ORT_RETURN_IF_ERROR(session->RegisterExecutionProvider(std::move(provider))); + ORT_RETURN_IF_ERROR(session->Load(stream)); + return session->Initialize(); + } + + private: + Model model_; + Status resolve_status_; + NodeArg* head_c_ = nullptr; +}; - 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}); - auto* ids = arg("ids", i64, {-1}); - auto* status = arg("status", i32, {1}); - auto* ctl = arg("ctl", i64, {8}); - 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* cosine = arg("cosine", bf16, {16, 3}); - auto* sine = arg("sine", bf16, {16, 3}); - auto* head_ids = arg("head_ids", i64, {1, 2}); - auto* head_c = arg("head_c", i32, {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* updated_ctl = arg("updated_ctl", i64, {8}); - auto* updated_context = arg("updated_context", i64, {2}); - 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}, - nullptr, 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()); +struct FoldInputs { + TensorShapeVector token_shape; + std::vector tokens; + std::vector ids; + std::vector ctl; + std::vector context; + std::vector prefixes; + std::vector ptrs; +}; + +struct FoldOutputs { + std::vector head_c; + std::vector ctl; + std::vector context; +}; +// Feeds CPU tensors, so ORT copies the state to temporary CUDA tensors: the updated state is only visible +// through the returned outputs. +Status RunFold(InferenceSession& session, const FoldInputs& in, FoldOutputs* out) { auto allocator = TestCPUExecutionProvider()->CreatePreferredAllocators()[0]; NameMLValMap feeds; auto feed = [&](const char* name, const TensorShapeVector& shape, const auto& values) { @@ -125,259 +159,212 @@ TEST(SpeculativeFoldTest, InferenceSessionLoadsVersionOneAndAcceptsMutedPrefillR CreateMLValue(allocator, shape, values, &value); feeds[name] = std::move(value); }; - std::vector rope(16 * 3); - for (size_t i = 0; i < rope.size(); ++i) rope[i] = BFloat16(static_cast(i)); - - auto run = [&](int rows, bool muted) { - std::vector token_values(rows); - std::vector id_values(rows, -1); - 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); - std::vector outputs; - const std::vector output_names{ - "head_ids", "head_c", "slots", "lengths", "cos_out", "sin_out", - "selected", "updated_ctl", "updated_context"}; - 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[1].Get().Data()[0], muted ? 0 : 2); - } + feed("tokens", in.token_shape, in.tokens); + feed("ids", {static_cast(in.ids.size())}, in.ids); + feed("status", {1}, std::vector{0}); + feed("ctl", {static_cast(in.ctl.size())}, in.ctl); + feed("context", {static_cast(in.context.size())}, in.context); + feed("prefixes", {static_cast(in.prefixes.size())}, in.prefixes); + feed("ptrs", {2}, in.ptrs); + std::vector outputs; + const std::vector output_names{"head_c", "updated_ctl", "updated_context"}; + ORT_RETURN_IF_ERROR(session.Run(RunOptions{}, feeds, output_names, &outputs)); + auto copy = [](const OrtValue& value, auto* dest) { + const auto& tensor = value.Get(); + dest->assign(tensor.DataAsSpan::value_type>().begin(), + tensor.DataAsSpan::value_type>().end()); }; - - run(2, false); - run(3, true); - run(4, true); + copy(outputs[0], &out->head_c); + copy(outputs[1], &out->ctl); + copy(outputs[2], &out->context); + return Status::OK(); } 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)); - + FoldResultRing ring(kBatch * kRingCap * 8); + + FoldModel model(std::vector{-1, kRows}, kBatch); + ASSERT_STATUS_OK(model.ResolveStatus()); + const auto* head_c_shape = model.HeadC()->Shape(); + ASSERT_NE(head_c_shape, nullptr); + ASSERT_EQ(head_c_shape->dim_size(), 1); + EXPECT_EQ(head_c_shape->dim(0).dim_param(), "batch"); + std::unique_ptr session; + ASSERT_STATUS_OK(model.CreateSession(session)); + + FoldInputs in; + in.token_shape = {kBatch, kRows}; // 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); + in.tokens = {10, 11, 12, 99, 14, 15, 16, 17, + 20, 21, 22, 23, 24, 25, 26, 27, + 30, 88, 32, 33, 34, 35, 36, 37}; + in.ids = {1, 10, 11, 12, 13, 14, 15, 16, + 2, 20, 21, 22, 23, 24, 25, 26, + 3, 30, 31, 32, 33, 34, 35, 36}; + in.ctl.assign(kBatch * kCtlWords, 0); // Batch 0: length=10, muted=0, seq=0, cache_len=10 - ctl_values[0 * 16 + 1] = 10; - ctl_values[0 * 16 + 7] = 10; + in.ctl[0 * kCtlWords + 1] = 10; + in.ctl[0 * kCtlWords + 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; + in.ctl[1 * kCtlWords + 1] = 20; + in.ctl[1 * kCtlWords + 3] = 1; + in.ctl[1 * kCtlWords + 4] = 5; + in.ctl[1 * kCtlWords + 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); + in.ctl[2 * kCtlWords + 1] = 30; + in.ctl[2 * kCtlWords + 4] = 2; + in.ctl[2 * kCtlWords + 7] = 30; + in.ctl[2 * kCtlWords + 2] = 77; // untouched word + in.context = {1, 2, 3, 4, 5, 6}; + in.prefixes.resize(kBatch * kCtlWords); 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; + in.prefixes[b * kCtlWords + 2 * r + 0] = 1000 * (b + 1) + 10 * r + 0; + in.prefixes[b * kCtlWords + 2 * r + 1] = 1000 * (b + 1) + 10 * r + 1; } } + in.ptrs = {ring.Address(), kRingCap}; - 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); + FoldOutputs out; + ASSERT_STATUS_OK(RunFold(*session, in, &out)); - 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); + EXPECT_EQ(out.head_c, (std::vector{10, 0, 30})); + + EXPECT_EQ(out.ctl[0 * kCtlWords + 0], 99); // pending = tokens[0, 3] + EXPECT_EQ(out.ctl[0 * kCtlWords + 1], 10 + 4); // length = 10 + 3 + 1 + EXPECT_EQ(out.ctl[0 * kCtlWords + 4], 1); // seq incremented to 1 + EXPECT_EQ(out.ctl[0 * kCtlWords + 6], 3); // accepted + EXPECT_EQ(out.ctl[0 * kCtlWords + 7], 10 + 4); // cache length + + EXPECT_EQ(out.ctl[1 * kCtlWords + 4], 5); // muted: seq unchanged + EXPECT_EQ(out.ctl[1 * kCtlWords + 1], 20); + + EXPECT_EQ(out.ctl[2 * kCtlWords + 0], 88); // pending = tokens[2, 1] + EXPECT_EQ(out.ctl[2 * kCtlWords + 1], 30 + 2); // length = 30 + 1 + 1 + EXPECT_EQ(out.ctl[2 * kCtlWords + 2], 77); + EXPECT_EQ(out.ctl[2 * kCtlWords + 4], 3); // seq incremented from 2 to 3 + EXPECT_EQ(out.ctl[2 * kCtlWords + 6], 1); + + // Batch 0 accepted=3 -> prefixes[0, 6..7]; batch 1 muted; batch 2 accepted=1 -> prefixes[2, 2..3] + EXPECT_EQ(out.context, (std::vector{1030, 1031, 3, 4, 3010, 3011})); + + // Ring records published for batch 0 (record 0) and batch 2 (record 2 % kRingCap); batch 1 is untouched. + const int64_t* host = ring.Host(); + EXPECT_EQ(host[(0 * kRingCap + 0) * 8 + 0], 1); + for (int word = 0; word < kRingCap * 8; ++word) EXPECT_EQ(host[kRingCap * 8 + word], -1); + EXPECT_EQ(host[(2 * kRingCap + 2) * 8 + 0], 3); +} + +TEST(SpeculativeFoldTest, CarriesUpdatedStateIntoNextRun) { + if (!HasCudaEnvironment(800)) GTEST_SKIP() << "Requires CUDA SM80 or newer"; + constexpr int kRows = 8; + constexpr int kRingCap = 2; + FoldResultRing ring(kRingCap * 8); + + FoldModel model(std::vector{1, kRows}, 1); + ASSERT_STATUS_OK(model.ResolveStatus()); + std::unique_ptr session; + ASSERT_STATUS_OK(model.CreateSession(session)); + + FoldInputs in; + in.token_shape = {1, kRows}; + in.tokens = {10, 11, 12, 13, 14, 15, 16, 17}; + in.ids = {1, 10, 11, 99, 99, 99, 99, 99}; // accepts two drafts + in.ctl.assign(kCtlWords, 0); + in.ctl[1] = 10; + in.ctl[7] = 10; + in.context = {1, 2}; + in.prefixes.resize(kCtlWords); + for (int i = 0; i < kCtlWords; ++i) in.prefixes[i] = 100 + i; + in.ptrs = {ring.Address(), kRingCap}; + + for (int run = 1; run <= 3; ++run) { + SCOPED_TRACE(run); + FoldOutputs out; + ASSERT_STATUS_OK(RunFold(*session, in, &out)); + EXPECT_EQ(out.ctl[4], run); // the sequence counter advances across runs + EXPECT_EQ(out.ctl[1], 10 + 3 * run); + EXPECT_EQ(out.ctl[7], 10 + 3 * run); + EXPECT_EQ(out.head_c[0], 10 + 3 * (run - 1)); + EXPECT_EQ(out.context, (std::vector{104, 105})); + EXPECT_EQ(ring.Host()[((run - 1) % kRingCap) * 8], run); + in.ctl = out.ctl; + in.context = out.context; + } + // Three records were published into a two-record ring: record 0 was overwritten by sequence 3, record 1 holds 2. + EXPECT_EQ(ring.Host()[0], 3); + EXPECT_EQ(ring.Host()[8], 2); +} + +TEST(SpeculativeFoldTest, MutedPrefillAcceptsMoreRowsThanHeadRows) { + if (!HasCudaEnvironment(800)) GTEST_SKIP() << "Requires CUDA SM80 or newer"; + constexpr int kRows = 20; + FoldResultRing ring(2 * 8); + + FoldModel model(std::vector{kRows}, 1); + ASSERT_STATUS_OK(model.ResolveStatus()); + const auto* head_c_shape = model.HeadC()->Shape(); + ASSERT_NE(head_c_shape, nullptr); // a rank-1 next_tokens is batch one + ASSERT_EQ(head_c_shape->dim_size(), 1); + EXPECT_EQ(head_c_shape->dim(0).dim_value(), 1); + std::unique_ptr session; + ASSERT_STATUS_OK(model.CreateSession(session)); + + FoldInputs in; + in.token_shape = {kRows}; + in.tokens.assign(kRows, 10); + in.ids.assign(kRows, 10); + in.ctl.assign(kCtlWords, 0); + in.ctl[1] = 3; + in.ctl[3] = 1; + in.ctl[7] = 2; + in.context = {1, 2}; + in.prefixes.assign(kCtlWords, 100); + in.ptrs = {0, 0}; + + FoldOutputs out; + ASSERT_STATUS_OK(RunFold(*session, in, &out)); + EXPECT_EQ(out.head_c, std::vector{0}); + EXPECT_EQ(out.ctl, in.ctl); + EXPECT_EQ(out.context, in.context); + EXPECT_EQ(std::count(ring.Host(), ring.Host() + ring.Words(), int64_t{-1}), static_cast(ring.Words())); +} + +TEST(SpeculativeFoldTest, RejectsNextTokensRankOutsideOneAndTwo) { + if (!HasCudaEnvironment(800)) GTEST_SKIP() << "Requires CUDA SM80 or newer"; + { + // Schema inference rejects a statically known rank-3 next_tokens. + FoldModel model(std::vector{2, 3, 8}, 2); + EXPECT_FALSE(model.ResolveStatus().IsOK()); + EXPECT_THAT(model.ResolveStatus().ErrorMessage(), testing::HasSubstr("next_tokens must have rank 1 or 2")); + } + { + // Without a declared shape, inference declares nothing (no batch 1) and the kernel rejects the runtime shape. + FoldModel model(std::nullopt, 2); + ASSERT_STATUS_OK(model.ResolveStatus()); + EXPECT_EQ(model.HeadC()->Shape(), nullptr); + std::unique_ptr session; + ASSERT_STATUS_OK(model.CreateSession(session)); + + FoldInputs in; + in.token_shape = {2, 3, 8}; + in.tokens.assign(2 * 3 * 8, 10); + in.ids.assign(2 * 3 * 8, 10); + in.ctl.assign(2 * kCtlWords, 0); + in.context = {1, 2, 3, 4}; + in.prefixes.assign(2 * kCtlWords, 100); + in.ptrs = {0, 0}; + FoldOutputs out; + const Status status = RunFold(*session, in, &out); + EXPECT_FALSE(status.IsOK()); + EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("next_tokens must have rank 1 ([rows]) or 2 ([batch, rows])")); + } } } // namespace From a79f2e13230b2c414e46ee22095ffb8a4b8db66a Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 26 Oct 2026 23:03:37 +0000 Subject: [PATCH 14/15] [CUDA] Regenerate SpeculativePrefixCommit docs for the reduced operator Update docs/ContribOperators.md and docs/OperatorKernels.md to match the pruned schema (7 inputs, 3 required outputs, head_rows in [3, 8]) and rewrite docs/contrib_ops/cuda/speculative_fold.md to describe the final operator (interface, state/aliasing contract, result ring layout, muted/error handling, graph wiring and tests) instead of the removed hidden-row, rotary and S1-S4 paths. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/ContribOperators.md | 56 +++------ docs/OperatorKernels.md | 2 +- docs/contrib_ops/cuda/speculative_fold.md | 136 ++++++++++++++-------- 3 files changed, 104 insertions(+), 90 deletions(-) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index d136e7809677a..d77423fa7af28 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -8174,7 +8174,7 @@ This version of the operator has been available since version 1 of the 'com.micr ### **com.microsoft.SpeculativePrefixCommit** - 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. + Commits the accepted prefix of a speculative verification step for each request in the batch and publishes a result-ring record for the host. Acceptance compares the verifier predictions with the draft tokens in input_ids. The operator is stateful: the control state and committed context are read from ctl and ctx_commit and the updated state is written to updated_ctl and updated_ctx_commit, which are required outputs. An output may alias its input (same device buffer) or be disjoint from it; partial overlap is not supported. Callers that keep the state in persistent device buffers must bind the outputs to those buffers, otherwise the updated state is lost. Callers must serialize access to the state buffers. The result ring is described by ptrs and holds ptrs[1] records of eight int64 words for each batch item, laid out consecutively per batch item starting at ptrs[0]; it must be a device-accessible, 8-byte aligned allocation. Each record is published by writing its sequence word last, after a system-scope fence. A request whose ctl[3] muted flag is set (for example during prefill) does not commit state or touch the ring, and may use any number of rows. Invalid device metadata sets error bits in ctl[5] (4: invalid row count or state, 16: invalid result ring) and mutes the request. The status words are OR-ed together and a non-zero value sets bit 2 of ctl[5]. #### Version @@ -8183,69 +8183,47 @@ This version of the operator has been available since version 1 of the 'com.micr #### Attributes
-
head_capacity : int
-
Maximum head-cache token capacity.
head_rows : int
-
Fixed number of head rows, 2 through 8.
+
Number of verifier rows, 3 through 8. Non-muted requests must provide exactly this many rows. Default is 8.
#### Inputs
next_tokens : I64
-
Verifier predictions, batch one.
+
Verifier predictions with shape [batch, rows], or [rows] for batch one.
input_ids : I64
-
Verifier input IDs.
+
Verifier input IDs; the draft tokens are input_ids[b, 1:rows]. The number of elements must be a multiple of the batch size and at least batch * rows.
status : I32
-
Verifier status words.
+
Verifier status words, shared by all batch items.
ctl : I64
-
Mutable control state: at least 8 words for S2, 16 for S3..S8.
+
Control state, 16 int64 words per batch item: [0] pending token, [1] length, [3] muted flag, [4] ring sequence number, [5] error flags, [6] accepted count, [7] cache length. Other words are preserved.
ctx_commit : I64
-
Mutable committed context, two words.
+
Committed context, two int64 words per batch item.
prefixes : I64
-
Two-word context for each accepted-prefix length; S2 needs only the first.
-
prefix2 : I64
-
Two-word context for acceptance length one.
+
Context pairs per batch item: pair k (words 2k and 2k+1) is committed when k tokens are accepted. At least 2 * head_rows words per batch item.
ptrs : I64
-
Result ring address and positive record capacity.
-
cos_cache : B
-
Rotary cosine cache [positions, width].
-
sin_cache : B
-
Rotary sine cache [positions, width].
+
Result ring address and number of records per batch item (must be positive).
-#### Outputs (7 - 9) +#### Outputs
-
head_ids : I64
-
Head token IDs [1, head_rows].
head_c : I32
-
Head committed length [1].
-
slot_ids : I64
-
Head slots [1, head_rows].
-
seqlens : I32
-
Head sequence length [1].
-
cos_out : B
-
Gathered rotary cosine [1, head_rows, width].
-
sin_out : B
-
Gathered rotary sine [1, head_rows, width].
-
sel : I64
-
Accepted-prefix length, or -1 on invalid metadata.
-
updated_ctl (optional) : I64
-
Optional updated control state, aliases ctl.
-
updated_ctx_commit (optional) : I64
-
Optional updated context, aliases ctx_commit.
+
Cache length before the commit [batch]; zero for muted or invalid requests.
+
updated_ctl : I64
+
Updated control state, same shape as ctl. May alias ctl.
+
updated_ctx_commit : I64
+
Updated committed context, same shape as ctx_commit. May alias ctx_commit.
#### Type Constraints
I64 : tensor(int64)
-
Token and control words.
+
Token, control and pointer words.
I32 : tensor(int32)
-
Status and sequence lengths.
-
B : tensor(bfloat16)
-
Rotary tensors.
+
Status words and cache lengths.
diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md index c1a6d1160bb78..378609f83c9a4 100644 --- a/docs/OperatorKernels.md +++ b/docs/OperatorKernels.md @@ -1177,7 +1177,7 @@ The **OpSet Version** column uses the following notation: |SparseAttention|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* past_key:**T**
*in* past_value:**T**
*in* block_row_indices:**M**
*in* block_col_indices:**M**
*in* total_sequence_length:**M**
*in* key_total_sequence_lengths:**M**
*in* cos_cache:**T**
*in* sin_cache:**T**
*out* output:**T**
*out* present_key:**T**
*out* present_value:**T**|1+|**M** = tensor(int32)
**T** = tensor(bfloat16), tensor(float16)| |SparseAttentionIndexer|*in* query:**T**
*in* key:**T**
*in* query_norm_weight:**T**
*in* key_norm_weight:**T**
*in* cos_cache:**T**
*in* sin_cache:**T**
*in* mask:**TB**
*in* past_key:**T**
*in* gate:**T**
*in* position_bias:**T**
*in* head_weights:**T**
*in* position_ids:**I**
*in* past_sequence_length:**M**
*in* past_proj_buffer:**T**
*out* selected_indices:**M**
*out* present_key:**T**
*out* present_proj_buffer:**T**|1+|**I** = tensor(int64)
**M** = tensor(int32)
**T** = tensor(bfloat16), tensor(float), tensor(float16)
**TB** = tensor(int64)| |SparsePagedAttention|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* key_cache:**T_CACHE**
*in* value_cache:**T_CACHE**
*in* cumulative_sequence_length:**S**
*in* past_seqlens:**S**
*in* block_table:**S**
*in* slot_mapping:**S**
*in* selected_indices:**S**
*in* selected_counts:**S**
*in* auxiliary_key:**T**
*in* auxiliary_value:**T**
*in* auxiliary_lengths:**S**
*in* cos_cache:**T**
*in* sin_cache:**T**
*in* head_sink:**T**
*in* q_norm_weight:**T**
*in* k_norm_weight:**T**
*in* k_scale:**T_KV_SCALE**
*in* v_scale:**T_KV_SCALE**
*in* attention_metadata:**S**
*out* output:**T**
*out* key_cache_out:**T_CACHE**
*out* value_cache_out:**T_CACHE**|1+|**S** = tensor(int32)
**T** = tensor(bfloat16), tensor(float16)
**T_CACHE** = tensor(bfloat16), tensor(float16), tensor(int8)
**T_KV_SCALE** = tensor(float)| -|SpeculativePrefixCommit|*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**
*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)| +|SpeculativePrefixCommit|*in* next_tokens:**I64**
*in* input_ids:**I64**
*in* status:**I32**
*in* ctl:**I64**
*in* ctx_commit:**I64**
*in* prefixes:**I64**
*in* ptrs:**I64**
*out* head_c:**I32**
*out* updated_ctl:**I64**
*out* updated_ctx_commit:**I64**|1+|**I32** = tensor(int32)
**I64** = tensor(int64)| |TransposeMatMul|*in* A:**T**
*in* B:**T**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)| |Trilu|*in* X:**T**
*in* k:**tensor(int64)**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| |UnfoldTensor|*in* input:**T**
*out* output:**T**|1+|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(string), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)| diff --git a/docs/contrib_ops/cuda/speculative_fold.md b/docs/contrib_ops/cuda/speculative_fold.md index 572c8f515b560..05480720096ee 100644 --- a/docs/contrib_ops/cuda/speculative_fold.md +++ b/docs/contrib_ops/cuda/speculative_fold.md @@ -1,50 +1,86 @@ -# CUDA speculative prefix commit operator - -`com.microsoft::SpeculativePrefixCommit` schema version 1 implements batch-one -speculative acceptance and result publication without a custom-op library. -This replaces the earlier two-row schema and supports two through eight head -rows, variable rotary width, configurable cache capacity, optional hidden-state -padding, and direct ring publication after verification. - -`SpeculativePrefixCommit` 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 three through eight head rows, control has at least -16 int64 words and prefix contexts contain `head_rows` pairs of int64 words. -When `rows > 4`, predictions `tokens[4..7]` are packed into the upper 32-bit -halves of ring record words `1, 2, 3, 5`. The operator writes the ring's -sequence word only after a system fence, so host polling cannot observe a -partial record. - -The ring metadata is `[device_accessible_address, positive_record_capacity]`. -Each record has eight int64 words. Its address must be aligned and its allocation -must hold the declared number of records; allocation ownership remains with -the caller. Invalid length, rotary position, or ring metadata sets error bits -4, 8, or 16, mutes control, and avoids unsafe ring access. Muted prefill uses -zero positions/cache length and does not commit contexts or access ring metadata. -Muted prefill may contain more verifier rows than `head_rows`; non-muted -verification is limited to `head_rows`. The kernel does not allocate -or synchronize the stream and supports CUDA graph replay. - -## Migrating legacy packages - -Use GenAI's `tools/qwen38_graph/convert_fold_to_contrib.py` on the target ONNX -graph. It adds explicit state outputs and connects the tail to the mid control -output, preserving external-weight paths. Remove `custom_ops_library` and -`custom_ops_libraries` from model session options only after all graph operations -are available in the rebuilt ORT stack. Keep the required GenAI data-plane -library: it is a runtime helper, not a custom-op registration library. - -Build ORT with CUDA and contrib operators enabled. The plugin build registers -the schemas in the core and the kernels in the CUDA provider. Rebuild both, -not just the provider. Ordinary upstream wheels do not contain these additions. - -`FoldContribTest.*` in `onnxruntime_provider_test` covers all S1-S4 prefix -lengths, hidden-row padding, deferred/wrapped publication, muted prefill, invalid -metadata, verifier/downstream status union, tail slots, and captured-graph replay. -`SpeculativeFoldTest.InferenceSessionLoadsVersionOneAndAcceptsMutedPrefillRows` -loads and runs the version-1 schemas through `InferenceSession`, including -muted prefill batches larger than `head_rows`. +# CUDA SpeculativePrefixCommit operator + +`com.microsoft::SpeculativePrefixCommit` (schema version 1, CUDA EP only) is the commit step of a speculative +decoding verification pass. For every request in the batch it accepts the longest prefix of draft tokens that the +verifier confirmed, advances the persistent control state, selects the committed context for the accepted length, and +publishes a record to a host-visible result ring that the runtime polls. It runs entirely on the device and can be +captured in a CUDA graph (no allocation, no synchronization). + +## Interface + +| Kind | Name | Type / shape | Meaning | +|---|---|---|---| +| attribute | `head_rows` | int, 3..8, default 8 | Number of verifier rows. Non-muted requests must provide exactly this many rows. | +| input 0 | `next_tokens` | int64 `[batch, rows]` or `[rows]` (batch one) | Verifier predictions. Any other rank is rejected. | +| input 1 | `input_ids` | int64, `batch * stride` elements, `stride >= rows` | Verifier input; the draft tokens are `input_ids[b, 1:rows]`. | +| input 2 | `status` | int32, any size | Verifier status words, OR-ed together and shared by all batch items. | +| input 3 | `ctl` | int64, 16 words per item | Control state (see below). | +| input 4 | `ctx_commit` | int64, 2 words per item | Committed context pair. | +| input 5 | `prefixes` | int64, `>= 2 * head_rows` words per item | Context pair `k` (words `2k`, `2k+1`) is committed when `k` drafts are accepted. | +| input 6 | `ptrs` | int64 `[2]` | `{result ring address, records per batch item}`. | +| output 0 | `head_c` | int32 `[batch]` | Cache length before the commit; zero for muted or invalid requests. | +| output 1 | `updated_ctl` | int64, shape of `ctl` | Updated control state. **Required.** | +| output 2 | `updated_ctx_commit` | int64, shape of `ctx_commit` | Updated committed context. **Required.** | + +`ctl` words used by the operator: `[0]` pending token, `[1]` length, `[3]` muted flag, `[4]` ring sequence number, +`[5]` error flags, `[6]` accepted count, `[7]` cache length. All other words are copied unchanged. + +## State handling + +The operator is stateful but never mutates its inputs through hidden aliasing. It reads `ctl` / `ctx_commit` and +always writes the new state to `updated_ctl` / `updated_ctx_commit`. Each output may be the same device buffer as its +input (in-place, which the kernel declares via `MayInplace`) or a disjoint buffer; partially overlapping buffers are +rejected. The kernel copies the unchanged words when the buffers differ, and each batch item is handled by a single +thread, so aliasing is safe. + +A caller that keeps the state in persistent device buffers must bind the two outputs to those buffers (for example with +`IOBinding`, passing the same device pointer as the input). An output that ORT allocates itself is a separate buffer, +so the update would be lost unless the caller feeds the returned tensors into the next run. Callers must serialize +access to the state buffers. + +## Result ring + +`ptrs[0]` is the device-accessible address of an 8-byte aligned allocation that holds `batch * ptrs[1]` records of +eight int64 words; batch item `b` uses records `[b * ptrs[1], (b + 1) * ptrs[1])` and writes sequence `s` to record +`(s - 1) % ptrs[1]`. Allocation ownership stays with the caller. Words (`low | high << 32` where two values share a +word): + +| Word | Contents | +|---|---| +| 0 | sequence number (`ctl[4]` after the commit), written last after `__threadfence_system()` | +| 1 | `tokens[0]` \| `tokens[4]` | +| 2 | `tokens[1]` \| `tokens[5]` | +| 3 | accepted count \| `tokens[6]` | +| 4 | OR of the status words | +| 5 | pending token \| `tokens[7]` | +| 6 | new length | +| 7 | `tokens[2]` \| `tokens[3]` | + +Tokens beyond `rows` are stored as zero. A host poller that reads word 0 first never observes a partial record. + +## Muted requests and errors + +A request with `ctl[3] != 0` (for example during prefill) does not commit state, does not access the ring, reports +`head_c = 0`, and may use any number of `next_tokens` rows. Invalid device metadata mutes the request and sets error +flags in `ctl[5]`: `4` for a negative length or cache length, a non-muted row count different from `head_rows`, and `16` +for an invalid ring (null or misaligned address, non-positive or overflowing record count, bad sequence number). A +non-zero OR of the status words sets flag `2`. + +## Using it in a graph + +Replace the previous fold node as follows: connect the seven inputs above, drop any rotary-cache inputs and the +duplicate `prefix2` input, and expose `updated_ctl` and `updated_ctx_commit` as graph outputs (or otherwise bind them) +so the runtime can bind them to its persistent control and context buffers. Only `head_c` is typically consumed as a +regular graph value. + +Build ORT with CUDA and contrib operators enabled. The schema lives in the core library and the kernel in the CUDA +provider, so both must be rebuilt; ordinary upstream wheels do not contain this operator. + +## Tests + +- `FoldContribTest.*` (`onnxruntime_provider_test`, kernel level) covers every accepted-prefix length for 3, 4 and 8 + rows, in-place and separate state outputs, batched per-item state and ring, muted prefill, invalid metadata, the + status union, and captured-graph replay. +- `SpeculativeFoldTest.*` (session level) covers batched S8 commits through `InferenceSession`, carrying the returned + state into later runs (the sequence counter advances), muted prefill with more rows than `head_rows`, shape inference + for known and unknown `next_tokens` shapes, and rejection of ranks other than 1 and 2. From 22d611725b4b1b4fa95e0cac2249b14a28760517 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Tue, 27 Oct 2026 01:03:33 +0000 Subject: [PATCH 15/15] Apply lintrunner formatting Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- onnxruntime/core/graph/contrib_ops/bert_defs.cc | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index 694416a232d88..4514c4b61aad6 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -51,12 +51,14 @@ ONNX_MS_OPERATOR_SET_SCHEMA( "prefill) does not commit state or touch the ring, and may use any number of rows. Invalid device " "metadata sets error bits in ctl[5] (4: invalid row count or state, 16: invalid result ring) and mutes " "the request. The status words are OR-ed together and a non-zero value sets bit 2 of ctl[5].") - .Attr("head_rows", "Number of verifier rows, 3 through 8. Non-muted requests must provide exactly this many " - "rows. Default is 8.", + .Attr("head_rows", + "Number of verifier rows, 3 through 8. Non-muted requests must provide exactly this many " + "rows. Default is 8.", AttributeProto::INT, static_cast(8)) .Input(0, "next_tokens", "Verifier predictions with shape [batch, rows], or [rows] for batch one.", "I64") - .Input(1, "input_ids", "Verifier input IDs; the draft tokens are input_ids[b, 1:rows]. The number of " - "elements must be a multiple of the batch size and at least batch * rows.", + .Input(1, "input_ids", + "Verifier input IDs; the draft tokens are input_ids[b, 1:rows]. The number of " + "elements must be a multiple of the batch size and at least batch * rows.", "I64") .Input(2, "status", "Verifier status words, shared by all batch items.", "I32") .Input(3, "ctl",