diff --git a/cmake/onnxruntime_unittests.cmake b/cmake/onnxruntime_unittests.cmake index 500e3e833f51c..213ea6f748b6d 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 (TARGET onnxruntime_providers_cuda_plugin_llm_fp4) diff --git a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc index c7fe3de628179..cf52b614db042 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); @@ -484,6 +487,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 477b8b2b6dd69..8830c524e40aa 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 12fc06b9816c3..a2d378e760ddb 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); @@ -226,6 +229,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