From 819d3566520cdbe626f135b8fe559a8390772a60 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Mon, 19 Oct 2026 17:56:51 +0000 Subject: [PATCH 1/2] Add fused MTP MoE and speculative state fold CUDA operators Part of the grouped generic CUDA operator stack. Preserve the qualified source snapshot. --- cmake/onnxruntime_unittests.cmake | 4 + .../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 | 343 ++++++++++++++++++ .../cuda_kernels/speculative_fold_test.cc | 329 +++++++++++++++++ 12 files changed, 1510 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 17526724dede9..bb77d30e3174e 100644 --- a/cmake/onnxruntime_unittests.cmake +++ b/cmake/onnxruntime_unittests.cmake @@ -1124,6 +1124,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() @@ -1155,7 +1157,9 @@ if (onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS AND onnxruntime_BUILD_CUDA_EP_AS_P list(APPEND onnxruntime_providers_cuda_plugin_ut_impl_src "${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..0967a2253f7e6 --- /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 2fe6a490ee90a..0cfd839266a82 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 c9ca273c6f971..fe302c6eaf97b 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); @@ -225,6 +228,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..df85804bcc2e8 --- /dev/null +++ b/onnxruntime/test/contrib_ops/cuda_kernels/mtp_moe_test.cc @@ -0,0 +1,343 @@ +// 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 b782aacef64f68ff499df4dfd81613c1650f44bd Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Mon, 19 Oct 2026 18:13:15 +0000 Subject: [PATCH 2/2] Fix lint --- .../packed_sparse_attention_indexer_common.h | 6 +- .../cuda/bert/branchwise_rms_norm_impl.cu | 4 +- .../cuda/bert/causal_conv_with_state_impl.cu | 4 +- .../cuda/bert/dynamic_sparse_attention.cc | 30 +-- .../bert/dynamic_sparse_attention_impl.cu | 2 +- .../contrib_ops/cuda/bert/fused_tensor_ops.cc | 46 ++-- .../contrib_ops/cuda/bert/fused_tensor_ops.cu | 18 +- .../contrib_ops/cuda/bert/gated_delta_net.cc | 28 +- .../cuda/bert/gated_delta_net_impl.cu | 8 +- .../cuda/bert/gated_delta_net_impl.h | 2 +- .../cuda/bert/rotary_embedding_impl.cu | 9 +- .../contrib_ops/cuda/llm/speculative_fold.cc | 10 +- .../cuda/math/matmul_block_scaled_fp8.cu | 30 +-- .../cuda/math/matmul_block_scaled_fp8.h | 2 +- .../packed_sparse_attention_indexer_impl.cu | 14 +- .../core/graph/contrib_ops/bert_defs.cc | 242 +++++++++--------- .../core/graph/contrib_ops/contrib_defs.cc | 2 +- .../cuda/reduction/reduction_functions.cu | 2 +- .../cuda/reduction/reduction_functions.h | 2 +- .../causal_conv_with_state_op_test.cc | 14 +- .../contrib_ops/cuda_kernels/mtp_moe_test.cc | 7 +- .../cuda_kernels/qmoe_fp8_compaction_test.cc | 2 +- .../dynamic_sparse_attention_op_test.cc | 40 +-- .../test/contrib_ops/fused_tensor_ops_test.cc | 14 +- .../contrib_ops/hyper_connection_op_test.cc | 10 +- .../matmul_block_scaled_fp8_test.cc | 13 +- .../contrib_ops/rotary_embedding_op_test.cc | 10 +- ...untime_test_python_symbolic_shape_infer.py | 3 - 28 files changed, 298 insertions(+), 276 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/sparse/packed_sparse_attention_indexer_common.h b/onnxruntime/contrib_ops/cpu/sparse/packed_sparse_attention_indexer_common.h index fdc636bc9626d..72fbac76026f6 100644 --- a/onnxruntime/contrib_ops/cpu/sparse/packed_sparse_attention_indexer_common.h +++ b/onnxruntime/contrib_ops/cpu/sparse/packed_sparse_attention_indexer_common.h @@ -71,9 +71,9 @@ enum OutputIndex : int { kPresentStateLengths = 5, // [batch_size, 2], int32 kStateUpdate = 6, // qsa only, optional: [batch_size, state_update_capacity, head_size] kIndexShareStatus = 7, - kMidKvBuffer = 8, // qsa only, optional: buffer snapshot after the first request token - kMidStateLengths = 9, // qsa only, optional: state lengths after the first request token - kValidationStatus = 10, // qsa only, optional: per-request overflow/metadata status + kMidKvBuffer = 8, // qsa only, optional: buffer snapshot after the first request token + kMidStateLengths = 9, // qsa only, optional: state lengths after the first request token + kValidationStatus = 10, // qsa only, optional: per-request overflow/metadata status kOutputCount = 8, kMaxOutputCount = 11, }; diff --git a/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.cu b/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.cu index c07197fa3bf14..240b057b827e0 100644 --- a/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.cu @@ -168,9 +168,9 @@ INSTANTIATE(__nv_bfloat16) #undef INSTANTIATE -#define INSTANTIATE_STANDARD(T) \ +#define INSTANTIATE_STANDARD(T) \ template Status LaunchStandardBranchwiseRMSNorm(cudaStream_t, const T*, const void*, int, T*, \ - int, int, int, bool, float, float*, float*, float*, void*, size_t, int); + int, int, int, bool, float, float*, float*, float*, void*, size_t, int); INSTANTIATE_STANDARD(float) INSTANTIATE_STANDARD(half) INSTANTIATE_STANDARD(__nv_bfloat16) diff --git a/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state_impl.cu b/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state_impl.cu index c5284935a5e04..eea53e8bbe272 100644 --- a/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state_impl.cu @@ -480,8 +480,8 @@ Status LaunchCausalConvProjectionKernel( return CUDA_CALL(cudaGetLastError()); } -#define INSTANTIATE_PROJECTION(T) \ - template Status LaunchCausalConvProjectionKernel( \ +#define INSTANTIATE_PROJECTION(T) \ + template Status LaunchCausalConvProjectionKernel( \ cudaStream_t, const T*, const T*, const T*, const T*, T*, T*, T*, \ int, int, int, int, int, int, int, int, bool, bool); INSTANTIATE_PROJECTION(float) diff --git a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc index b11283cff7599..f7a65fbad3663 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention.cc @@ -19,21 +19,21 @@ namespace cuda { using namespace onnxruntime::cuda; -#define REGISTER_KERNEL_TYPED(T) \ - ONNX_OPERATOR_TYPED_KERNEL_EX( \ - DynamicSparseAttention, \ - kMSDomain, \ - 1, \ - T, \ - kCudaExecutionProvider, \ - (*KernelDefBuilder::Create()) \ - .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ +#define REGISTER_KERNEL_TYPED(T) \ + ONNX_OPERATOR_TYPED_KERNEL_EX( \ + DynamicSparseAttention, \ + kMSDomain, \ + 1, \ + T, \ + kCudaExecutionProvider, \ + (*KernelDefBuilder::Create()) \ + .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ .TypeConstraint("TCache", BuildKernelDefConstraints()) \ - .TypeConstraint("S", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("M", DataTypeImpl::GetTensorType()) \ - .MayInplace(3, 1) \ - .MayInplace(4, 2) \ - .InputMemoryType(OrtMemTypeCPUInput, 10), \ + .TypeConstraint("S", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("M", DataTypeImpl::GetTensorType()) \ + .MayInplace(3, 1) \ + .MayInplace(4, 2) \ + .InputMemoryType(OrtMemTypeCPUInput, 10), \ DynamicSparseAttention); REGISTER_KERNEL_TYPED(float) @@ -142,7 +142,7 @@ Status DynamicSparseAttention::ComputeInternal(OpKernelContext* context) cons "DynamicSparseAttention: INT8 caches require per-channel key_scale and value_scale."); ORT_RETURN_IF_NOT((past_key == nullptr || past_key->IsDataType() || quantized_cache) && (past_value == nullptr || (quantized_cache ? past_value->IsDataType() - : past_value->IsDataType())), + : past_value->IsDataType())), "DynamicSparseAttention: cache types must match activations or both be INT8."); DynamicSparseAttentionParameters parameters; diff --git a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu index 0e3e8746de44d..ccdfdf76bb7ee 100644 --- a/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/dynamic_sparse_attention_impl.cu @@ -956,7 +956,7 @@ __global__ void GroupedDecodeDynamicSparseAttentionKernel( const int selected_count = max(0, min(selected_counts[row], max_selected)); const int candidate_count = local_count + selected_count; constexpr int kCandidatesPerSplit = kQuantizedCache ? kGroupedQuantizedDecodeCandidatesPerSplit - : kGroupedDecodeCandidatesPerSplit; + : kGroupedDecodeCandidatesPerSplit; const int candidate_begin = split * kCandidatesPerSplit; const int32_t* row_indices = max_selected == 0 ? selected_indices : selected_indices + static_cast(row) * max_selected; diff --git a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc index 24d2a153bf450..555c98f1fb16d 100644 --- a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc +++ b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc @@ -13,12 +13,12 @@ namespace onnxruntime::contrib::cuda { -#define REGISTER_ARGMAX_LAST_AXIS(T) \ - ONNX_OPERATOR_TYPED_KERNEL_EX( \ - ArgMaxLastAxis, kMSDomain, 1, T, kCudaExecutionProvider, \ - (*KernelDefBuilder::Create()) \ - .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("I", DataTypeImpl::GetTensorType()), \ +#define REGISTER_ARGMAX_LAST_AXIS(T) \ + ONNX_OPERATOR_TYPED_KERNEL_EX( \ + ArgMaxLastAxis, kMSDomain, 1, T, kCudaExecutionProvider, \ + (*KernelDefBuilder::Create()) \ + .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("I", DataTypeImpl::GetTensorType()), \ ArgMaxLastAxis); REGISTER_ARGMAX_LAST_AXIS(float) @@ -30,38 +30,38 @@ REGISTER_ARGMAX_LAST_AXIS(BFloat16) ONNX_OPERATOR_KERNEL_EX( EmbeddingRMSNorm, kMSDomain, 1, kCudaExecutionProvider, (*KernelDefBuilder::Create()) - .TypeConstraint("I", DataTypeImpl::GetTensorType()) - .TypeConstraint("T", DataTypeImpl::GetTensorType()) - .TypeConstraint("F", DataTypeImpl::GetTensorType()), + .TypeConstraint("I", DataTypeImpl::GetTensorType()) + .TypeConstraint("T", DataTypeImpl::GetTensorType()) + .TypeConstraint("F", DataTypeImpl::GetTensorType()), EmbeddingRMSNorm); ONNX_OPERATOR_KERNEL_EX( RowSelect, kMSDomain, 1, kCudaExecutionProvider, (*KernelDefBuilder::Create()) - .TypeConstraint("I", DataTypeImpl::GetTensorType()) - .TypeConstraint("T", DataTypeImpl::GetTensorType()), + .TypeConstraint("I", DataTypeImpl::GetTensorType()) + .TypeConstraint("T", DataTypeImpl::GetTensorType()), RowSelect); ONNX_OPERATOR_KERNEL_EX( SigmoidGateMul, kMSDomain, 1, kCudaExecutionProvider, (*KernelDefBuilder::Create()) - .TypeConstraint("T", DataTypeImpl::GetTensorType()), + .TypeConstraint("T", DataTypeImpl::GetTensorType()), SigmoidGateMul); #if !defined(DISABLE_FLOAT8_TYPES) ONNX_OPERATOR_TYPED_KERNEL_EX( DequantizeLinear, kMSDomain, 1, Float8E4M3FN_BFloat16, kCudaExecutionProvider, (*KernelDefBuilder::Create()) - .TypeConstraint("T1", DataTypeImpl::GetTensorType()) - .TypeConstraint("T2", DataTypeImpl::GetTensorType()) - .TypeConstraint("TS", DataTypeImpl::GetTensorType()), + .TypeConstraint("T1", DataTypeImpl::GetTensorType()) + .TypeConstraint("T2", DataTypeImpl::GetTensorType()) + .TypeConstraint("TS", DataTypeImpl::GetTensorType()), DequantizeLinearFp8Bf16); #endif ONNX_OPERATOR_KERNEL_EX( UntileHyperConnectionWeights, kMSDomain, 1, kCudaExecutionProvider, (*KernelDefBuilder::Create()) - .TypeConstraint("T", DataTypeImpl::GetTensorType()), + .TypeConstraint("T", DataTypeImpl::GetTensorType()), UntileHyperConnectionWeights); template @@ -153,12 +153,10 @@ SigmoidGateMul::SigmoidGateMul(const OpKernelInfo& info) const int64_t silu = info.GetAttrOrDefault("silu_gate", 0); ORT_ENFORCE(silu == 0 || silu == 1, "SigmoidGateMul silu_gate must be zero or one"); ORT_ENFORCE(!(broadcast_gate_ && silu_gate_), "SigmoidGateMul broadcast_gate and silu_gate are exclusive"); - ORT_ENFORCE(broadcast_gate_ || silu_gate_ || (head_dim_ > 0 && heads_ > 0 && - q_stride_ >= 2 * head_dim_ && - tok_stride_ > 0 && q_offset_ >= 0), + ORT_ENFORCE(broadcast_gate_ || silu_gate_ || (head_dim_ > 0 && heads_ > 0 && q_stride_ >= 2 * head_dim_ && tok_stride_ > 0 && q_offset_ >= 0), "SigmoidGateMul received invalid layout attributes"); ORT_ENFORCE(!(broadcast_gate_ || silu_gate_) || (head_dim_ == 0 && heads_ == 0 && q_stride_ == 0 && - tok_stride_ == 0 && q_offset_ == 0), + tok_stride_ == 0 && q_offset_ == 0), "SigmoidGateMul broadcast_gate/silu_gate cannot be combined with packed layout attributes"); } @@ -184,12 +182,12 @@ Status SigmoidGateMul::ComputeInternal(OpKernelContext* context) const { } else { const int64_t row_width = static_cast(heads_) * head_dim_; ORT_RETURN_IF_NOT(count % row_width == 0, - "SigmoidGateMul attention size must be divisible by heads * head_dim"); + "SigmoidGateMul attention size must be divisible by heads * head_dim"); const int64_t tokens = count / row_width; ORT_RETURN_IF_NOT(projection->Shape().Size() >= - (tokens - 1) * tok_stride_ + q_offset_ + - static_cast(heads_) * q_stride_, - "SigmoidGateMul projection is too small for the configured layout"); + (tokens - 1) * tok_stride_ + q_offset_ + + static_cast(heads_) * q_stride_, + "SigmoidGateMul projection is too small for the configured layout"); } Tensor* output = context->Output(0, attention->Shape()); return LaunchSigmoidGateMul( diff --git a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu index 6db5544942e1f..789edc4a26615 100644 --- a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu +++ b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu @@ -169,8 +169,8 @@ __global__ void SigmoidGateMulKernel(const __nv_bfloat16* attention, const BFloat16 a(gate); // Match CUDA Sigmoid's input-type intermediates, including its negative-input subtraction. const BFloat16 value = a > BFloat16(0.0f) - ? BFloat16(1.0f) / (BFloat16(1.0f) + _Exp(-a)) - : BFloat16(1.0f) - BFloat16(1.0f) / (BFloat16(1.0f) + _Exp(a)); + ? BFloat16(1.0f) / (BFloat16(1.0f) + _Exp(-a)) + : BFloat16(1.0f) - BFloat16(1.0f) / (BFloat16(1.0f) + _Exp(a)); sigmoid_bf16 = __float2bfloat16_rn(static_cast(value)); } if constexpr (SiluGate) { @@ -254,13 +254,13 @@ Status LaunchArgMaxLastAxis(cudaStream_t stream, const void* input, int input_ty switch (input_type) { case ONNX_NAMESPACE::TensorProto_DataType_FLOAT: return LaunchTypedArgmax(stream, static_cast(input), rows, width, output, - partial_values, partial_indices); + partial_values, partial_indices); case ONNX_NAMESPACE::TensorProto_DataType_FLOAT16: return LaunchTypedArgmax(stream, static_cast(input), rows, width, output, - partial_values, partial_indices); + partial_values, partial_indices); case ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16: return LaunchTypedArgmax(stream, static_cast(input), rows, width, output, - partial_values, partial_indices); + partial_values, partial_indices); default: return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "ArgMaxLastAxis input type is not supported"); } @@ -299,11 +299,11 @@ Status LaunchSigmoidGateMul(cudaStream_t stream, const uint16_t* attention, const int64_t block_count = (count - 1) / kThreads + 1; ORT_RETURN_IF_NOT(block_count <= std::numeric_limits::max(), "SigmoidGateMul requires too many CUDA blocks"); -#define LAUNCH_GATE(BROADCAST, RESIDUAL, SILU) \ +#define LAUNCH_GATE(BROADCAST, RESIDUAL, SILU) \ SigmoidGateMulKernel<<(block_count), kThreads, 0, stream>>>( \ - reinterpret_cast(attention), \ - reinterpret_cast(projection), \ - reinterpret_cast<__nv_bfloat16*>(output), count, head_dim, heads, q_stride, \ + reinterpret_cast(attention), \ + reinterpret_cast(projection), \ + reinterpret_cast<__nv_bfloat16*>(output), count, head_dim, heads, q_stride, \ token_stride, q_offset, gate_row_width, reinterpret_cast(residual), sigmoid_in_fp32) if (silu_gate) { if (residual != nullptr) { diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc index f346bd3f6027d..380e99746af31 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc +++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc @@ -22,22 +22,22 @@ using onnxruntime::cuda::CudaKernel; using onnxruntime::cuda::OrtToCudaType; namespace gdn = gated_delta_net; -#define REGISTER_KERNEL_TYPED(T) \ - ONNX_OPERATOR_TYPED_KERNEL_EX( \ - GatedDeltaNet, \ - kMSDomain, \ - 1, \ - T, \ - kCudaExecutionProvider, \ - (*KernelDefBuilder::Create()) \ - .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("TS", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("TG", {DataTypeImpl::GetTensorType(), \ +#define REGISTER_KERNEL_TYPED(T) \ + ONNX_OPERATOR_TYPED_KERNEL_EX( \ + GatedDeltaNet, \ + kMSDomain, \ + 1, \ + T, \ + kCudaExecutionProvider, \ + (*KernelDefBuilder::Create()) \ + .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("TS", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("TG", {DataTypeImpl::GetTensorType(), \ DataTypeImpl::GetTensorType(), \ DataTypeImpl::GetTensorType()}) \ - .TypeConstraint("TI", DataTypeImpl::GetTensorType()) \ - .InputMemoryType(OrtMemTypeCPUInput, 10) \ - .MayInplace(6, 1), \ + .TypeConstraint("TI", DataTypeImpl::GetTensorType()) \ + .InputMemoryType(OrtMemTypeCPUInput, 10) \ + .MayInplace(6, 1), \ GatedDeltaNet); REGISTER_KERNEL_TYPED(float) diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu index f308442329c5d..048faa9c2f3f5 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu @@ -745,8 +745,8 @@ Status LaunchRecurrent(const Descriptor& desc, bool warp_specialized, const Vari template static Status LaunchGatedDeltaNetTypedGates(const Descriptor& desc, const Plan& plan, const VariantPack& pack, - float scale, int max_threads_per_block, - size_t max_shared_memory_per_block, cudaStream_t stream) { + float scale, int max_threads_per_block, + size_t max_shared_memory_per_block, cudaStream_t stream) { KernelParams p{}; p.total_tokens = desc.total_tokens; p.uniform_len = desc.batch > 0 ? desc.total_tokens / desc.batch : 0; @@ -800,13 +800,13 @@ static Status LaunchGatedDeltaNetTypedGates(const Descriptor& desc, const Plan& // The chunked engine only accepts head_size_qk == head_size_v == 128, so the warp // kernel is always the right shape here. return LaunchRecurrent(desc, /*warp_specialized=*/true, pack, p, - max_shared_memory_per_block, stream); + max_shared_memory_per_block, stream); } } (void)max_threads_per_block; return LaunchRecurrent(desc, plan.warp_specialized, pack, p, max_shared_memory_per_block, - stream); + stream); } template diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.h b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.h index 2bd82be0b53bf..3035ca02fa083 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.h @@ -40,7 +40,7 @@ struct VariantPack { template __device__ __forceinline__ float LoadGate(const void* data, int64_t index, bool gates_in_input_type) { return gates_in_input_type ? static_cast(static_cast(data)[index]) - : static_cast(data)[index]; + : static_cast(data)[index]; } #endif diff --git a/onnxruntime/contrib_ops/cuda/bert/rotary_embedding_impl.cu b/onnxruntime/contrib_ops/cuda/bert/rotary_embedding_impl.cu index 9f162783b3a30..55648cc944249 100644 --- a/onnxruntime/contrib_ops/cuda/bert/rotary_embedding_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/rotary_embedding_impl.cu @@ -189,8 +189,7 @@ Status LaunchRotaryEmbeddingKernel(cudaStream_t stream, T* output, const T* inpu const int position_ids_format, const bool interleaved, const int max_threads_per_block, int4 in_strides, int4 out_strides, // strides in bnsh coord - const bool use_fp32_arithmetic -) { + const bool use_fp32_arithmetic) { // Note: Current implementation assumes head_size <= max_threads_per_block // because head_size is currently large for LLaMA-2. For smaller head_size // and num_heads values, we can create a block as `block(num_heads, head_size, 1)` @@ -221,8 +220,10 @@ Status LaunchRotaryEmbeddingKernel(cudaStream_t stream, T* output, const T* inpu interleaved, in_strides, out_strides); } }; - if (use_fp32_arithmetic) launch(std::true_type{}); - else launch(std::false_type{}); + if (use_fp32_arithmetic) + launch(std::true_type{}); + else + launch(std::false_type{}); return CUDA_CALL(cudaGetLastError()); } diff --git a/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cc b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cc index 0967a2253f7e6..6965c9d1539c1 100644 --- a/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cc +++ b/onnxruntime/contrib_ops/cuda/llm/speculative_fold.cc @@ -114,16 +114,16 @@ class FoldTail final : public CudaKernel { ONNX_OPERATOR_KERNEL_EX( FoldMid2, kMSDomain, 1, kCudaExecutionProvider, (*KernelDefBuilder::Create()) - .TypeConstraint("I64", DataTypeImpl::GetTensorType()) - .TypeConstraint("I32", DataTypeImpl::GetTensorType()) - .TypeConstraint("B", DataTypeImpl::GetTensorType()), + .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()), + .TypeConstraint("I64", DataTypeImpl::GetTensorType()) + .TypeConstraint("I32", DataTypeImpl::GetTensorType()), FoldTail); } // namespace cuda diff --git a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cu b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cu index 599a8ec957852..16ae3c0d5c949 100644 --- a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cu +++ b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cu @@ -516,7 +516,7 @@ struct Fp8GemvMma<__nv_bfloat16> { template __device__ __forceinline__ T* Fp8OutputColumn(T* output, int col, int& stride, - const Fp8GemvOutputSplit& split) { + const Fp8GemvOutputSplit& split) { for (int i = 0; i < split.count; ++i) { if (col < split.widths[i]) { stride = split.widths[i]; @@ -946,7 +946,7 @@ bool Fp8MmaGb10TuningEnabled() { } bool MatMulBlockScaledFp8GemvSupportsSplitOutput(int m, int k, int block_size, - const cudaDeviceProp& device_prop) { + const cudaDeviceProp& device_prop) { #if !defined(DISABLE_FLOAT8_TYPES) && defined(CUDA_VERSION) && CUDA_VERSION >= 11080 return device_prop.major >= 8 && m > 0 && m <= kFp8MmaGemvTileM && k >= 256 && k % 64 == 0 && block_size % 64 == 0 && Fp8GemvMmaEnabled(); @@ -1041,19 +1041,19 @@ static Status LaunchMatMulBlockScaledFp8GemvImpl(void* y, const auto launch_mma = [&]() { const dim3 mma_threads{32, KSplit}; const auto launch_variant = [&]() { -#define ORT_FP8_LAUNCH_MMA(kernel_name) \ - do { \ - if (is_bf16) { \ - kernel_name<<>>( \ - reinterpret_cast<__nv_bfloat16*>(y), reinterpret_cast(a), b, \ - weight_scale, reinterpret_cast(bias), act_scale, m, n, k, \ - block_size, k_blocks, split); \ - } else { \ - kernel_name<<>>( \ - reinterpret_cast(y), reinterpret_cast(a), b, \ - weight_scale, reinterpret_cast(bias), act_scale, m, n, k, \ - block_size, k_blocks, split); \ - } \ +#define ORT_FP8_LAUNCH_MMA(kernel_name) \ + do { \ + if (is_bf16) { \ + kernel_name<<>>( \ + reinterpret_cast<__nv_bfloat16*>(y), reinterpret_cast(a), b, \ + weight_scale, reinterpret_cast(bias), act_scale, m, n, k, \ + block_size, k_blocks, split); \ + } else { \ + kernel_name<<>>( \ + reinterpret_cast(y), reinterpret_cast(a), b, \ + weight_scale, reinterpret_cast(bias), act_scale, m, n, k, \ + block_size, k_blocks, split); \ + } \ } while (0) if constexpr (KSplit == 16 && MTiles == 1) { if (pin_residency) { diff --git a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.h b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.h index 35e2245543313..b54a37da08a87 100644 --- a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.h +++ b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.h @@ -45,7 +45,7 @@ struct Fp8GemvOutputSplit { }; bool MatMulBlockScaledFp8GemvSupportsSplitOutput(int m, int k, int block_size, - const cudaDeviceProp& device_prop); + const cudaDeviceProp& device_prop); #if defined(USE_DEEP_GEMM) Status LaunchPrepareMatMulFp8DeepGemm(void* a_quant, float* a_scales, const void* a, diff --git a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu index 8f82da6379ab5..4410df08e47cd 100644 --- a/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu +++ b/onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu @@ -667,7 +667,7 @@ __global__ void QsaMergeTileTopKKernel(const uint64_t* input_keys, uint64_t* out } __global__ void QsaMergeTwoLevelsTopKKernel(const uint64_t* input_keys, uint64_t* output_keys, - int key_stride, int list_width, int rows) { + int key_stride, int list_width, int rows) { __shared__ uint64_t inputs[4 * kBoundedTopKMax]; __shared__ uint64_t intermediate[2 * kBoundedTopKMax]; const int groups_per_row = (key_stride + 4 * list_width - 1) / (4 * list_width); @@ -1515,7 +1515,7 @@ Status LaunchQsaPackedSparseAttentionIndexer( const int64_t tile_work = static_cast(params.total_tokens) * std::min(tile_count, kHierarchicalScoreMaxBlocksPerRow); QsaScoreTileTopKKernel<<(std::min(tile_work, kSaiMaxGridDimX)), - kHierarchicalTileScoreThreads, 0, stream>>>( + kHierarchicalTileScoreThreads, 0, stream>>>( present_key_state, query_rotated, cumulative_sequence_lengths, past_sequence_lengths, position_ids, present_state_lengths, overflow_flags, merge_input, tile_count, params); static const bool use_radix_topk = @@ -1530,7 +1530,7 @@ Status LaunchQsaPackedSparseAttentionIndexer( const int groups_per_row = (key_stride + 4 * list_width - 1) / (4 * list_width); const int64_t merge_work = static_cast(params.total_tokens) * groups_per_row; QsaMergeTwoLevelsTopKKernel<<(std::min(merge_work, kSaiMaxGridDimX)), - kHierarchicalMergeThreads, 0, stream>>>( + kHierarchicalMergeThreads, 0, stream>>>( merge_input, merge_output, key_stride, list_width, params.total_tokens); list_width *= 4; } else { @@ -1663,10 +1663,10 @@ Status LaunchCsaPackedSparseAttentionIndexer( return CUDA_CALL(cudaGetLastError()); } -#define INSTANTIATE_QSA(T, TState) \ - template Status LaunchQsaPackedSparseAttentionIndexer( \ - cudaStream_t, const PackedSparseAttentionIndexerParams&, const T*, const T*, const T*, const T*, \ - const T*, const T*, const int32_t*, const int32_t*, const int64_t*, const TState*, const T*, \ +#define INSTANTIATE_QSA(T, TState) \ + template Status LaunchQsaPackedSparseAttentionIndexer( \ + cudaStream_t, const PackedSparseAttentionIndexerParams&, const T*, const T*, const T*, const T*, \ + const T*, const T*, const int32_t*, const int32_t*, const int64_t*, const TState*, const T*, \ const int32_t*, const int32_t*, bool, int32_t*, int32_t*, TState*, T*, int32_t*, T*, float*, int32_t*); #define INSTANTIATE_CSA(T) \ diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index cf92e6f11b939..7f490a9a1c881 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -3477,55 +3477,55 @@ that are multiplied to query and key before the inner product of query and key i )DOC"; static OpSchema RotaryEmbeddingSchema() { return OpSchema() - .SetDoc(RotaryEmbedding_ver1_doc) - .Attr("scale", - "Custom scale will be used if specified. Default value is 1.0", - AttributeProto::FLOAT, - OPTIONAL_VALUE) - .Attr("interleaved", - "Indicates whether the input has real and imaginary parts interleaved. " - "Default value is 0 (False), meaning the first half of the input consists of real values " - "and the second half consists of imaginary values.", - AttributeProto::INT, - OPTIONAL_VALUE) - .Attr("rotary_embedding_dim", - "Rotary embedding dimension. Default value is 0.", - AttributeProto::INT, - OPTIONAL_VALUE) - .Attr("num_heads", - "Number of attention heads. Default value is 0. Must use with rotary_embedding_dim", - AttributeProto::INT, - OPTIONAL_VALUE) - .Attr("is_packed_batching", - "ragged batch inputs or not. Default value is 0", - AttributeProto::INT, - OPTIONAL_VALUE) - .Input(0, - "input", - "3D tensor with shape (batch_size, sequence_length, hidden_size) or 4D with shape (batch_size, num_heads, sequence_length, head_size)", - "T") - .Input(1, - "position_ids", - "1D tensor with shape (1) or 2D tensor with shape (batch_size, sequence_length)", - "M") - .Input(2, - "cos_cache", - "2D tensor with shape (max_sequence_length, head_size / 2) or (max_sequence_length, rotary_embedding_dim / 2)", - "T") - .Input(3, - "sin_cache", - "2D tensor with shape (max_sequence_length, head_size / 2) or (max_sequence_length, rotary_embedding_dim / 2)", - "T") - .Output(0, - "output", - "tensor with same shape as input.", - "T") - .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain input and output types to float tensors.") - .TypeConstraint("M", {"tensor(int64)"}, "Constrain input and output types to integer tensors") - .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { - propagateElemTypeFromInputToOutput(ctx, 0, 0); - propagateShapeFromInputToOutput(ctx, 0, 0); - }); + .SetDoc(RotaryEmbedding_ver1_doc) + .Attr("scale", + "Custom scale will be used if specified. Default value is 1.0", + AttributeProto::FLOAT, + OPTIONAL_VALUE) + .Attr("interleaved", + "Indicates whether the input has real and imaginary parts interleaved. " + "Default value is 0 (False), meaning the first half of the input consists of real values " + "and the second half consists of imaginary values.", + AttributeProto::INT, + OPTIONAL_VALUE) + .Attr("rotary_embedding_dim", + "Rotary embedding dimension. Default value is 0.", + AttributeProto::INT, + OPTIONAL_VALUE) + .Attr("num_heads", + "Number of attention heads. Default value is 0. Must use with rotary_embedding_dim", + AttributeProto::INT, + OPTIONAL_VALUE) + .Attr("is_packed_batching", + "ragged batch inputs or not. Default value is 0", + AttributeProto::INT, + OPTIONAL_VALUE) + .Input(0, + "input", + "3D tensor with shape (batch_size, sequence_length, hidden_size) or 4D with shape (batch_size, num_heads, sequence_length, head_size)", + "T") + .Input(1, + "position_ids", + "1D tensor with shape (1) or 2D tensor with shape (batch_size, sequence_length)", + "M") + .Input(2, + "cos_cache", + "2D tensor with shape (max_sequence_length, head_size / 2) or (max_sequence_length, rotary_embedding_dim / 2)", + "T") + .Input(3, + "sin_cache", + "2D tensor with shape (max_sequence_length, head_size / 2) or (max_sequence_length, rotary_embedding_dim / 2)", + "T") + .Output(0, + "output", + "tensor with same shape as input.", + "T") + .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain input and output types to float tensors.") + .TypeConstraint("M", {"tensor(int64)"}, "Constrain input and output types to integer tensors") + .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + propagateElemTypeFromInputToOutput(ctx, 0, 0); + propagateShapeFromInputToOutput(ctx, 0, 0); + }); } ONNX_MS_OPERATOR_SET_SCHEMA(RotaryEmbedding, 1, RotaryEmbeddingSchema()); @@ -4689,15 +4689,18 @@ ONNX_MS_OPERATOR_SET_SCHEMA( CausalConvWithState, 1, OpSchema() .SetDoc(CausalConvWithState_ver1_doc) - .Attr("input_channels", "CUDA channels-last input channel slice width; 0 selects all channels. " + .Attr("input_channels", + "CUDA channels-last input channel slice width; 0 selects all channels. " "Nondefault extended attributes require rank-3 input and state_window=0.", AttributeProto::INT, static_cast(0)) .Attr("input_channel_offset", "Starting channel of the channels-last input slice (CUDA only).", AttributeProto::INT, static_cast(0)) - .Attr("state_length", "Carry history length (CUDA only); 0 uses (kernel_size-1)*dilation. " + .Attr("state_length", + "Carry history length (CUDA only); 0 uses (kernel_size-1)*dilation. " "An explicit length must cover the receptive field; convolution reads its trailing history.", AttributeProto::INT, static_cast(0)) - .Attr("state_channels_last", "State layout override (CUDA only): -1 follows channels_last, " + .Attr("state_channels_last", + "State layout override (CUDA only): -1 follows channels_last, " "0 uses [B,C,state_length], 1 uses [B,state_length,C].", AttributeProto::INT, static_cast(-1)) .Attr("activation", @@ -4781,9 +4784,11 @@ ONNX_MS_OPERATOR_SET_SCHEMA( "(seq_len - W + j). Extended attributes override the state length and layout " "as for past_state.", "T") - .Output(2, "first_state", "Optional carry state after the first input position, with the same " + .Output(2, "first_state", + "Optional carry state after the first input position, with the same " "layout and length as present_state. CUDA only; requires rank-3 channels-last input " - "and state_window=0.", "T", OpSchema::Optional) + "and state_window=0.", + "T", OpSchema::Optional) .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain input and output types to float tensors.") @@ -5763,20 +5768,26 @@ ONNX_MS_OPERATOR_SET_SCHEMA( .Attr("q_stride", "Per-head stride in the packed projection.", AttributeProto::INT, static_cast(0)) .Attr("tok_stride", "Per-token stride in the packed projection.", AttributeProto::INT, static_cast(0)) .Attr("q_offset", "Offset of the query projection within each token.", AttributeProto::INT, static_cast(0)) - .Attr("broadcast_gate", "When one, projection has the attention shape with final dimension one. " - "Packed layout attributes must be zero. CUDA only.", AttributeProto::INT, static_cast(0)) - .Attr("sigmoid_in_fp32", "One computes sigmoid in FP32 before BF16 conversion. Zero preserves " + .Attr("broadcast_gate", + "When one, projection has the attention shape with final dimension one. " + "Packed layout attributes must be zero. CUDA only.", + AttributeProto::INT, static_cast(0)) + .Attr("sigmoid_in_fp32", + "One computes sigmoid in FP32 before BF16 conversion. Zero preserves " "the input-type intermediate arithmetic of CUDA Sigmoid, including subtraction on negative inputs.", AttributeProto::INT, static_cast(1)) - .Attr("silu_gate", "When one, projection and attention have identical shapes. Multiply the " + .Attr("silu_gate", + "When one, projection and attention have identical shapes. Multiply the " "BF16-rounded sigmoid by projection and round to BF16 before multiplying attention. " "Cannot be combined with broadcast_gate or packed layout attributes.", AttributeProto::INT, static_cast(0)) .Input(0, "attention", "BF16 attention values.", "T") .Input(1, "projection", "BF16 packed projection containing gate values.", "T") .Input(2, "residual", "Optional BF16 residual with the attention shape.", "T", OpSchema::Optional) - .Output(0, "output", "Attention multiplied by BF16-rounded sigmoid gates, rounded to BF16 " - "before optionally adding residual and rounding again.", "T") + .Output(0, "output", + "Attention multiplied by BF16-rounded sigmoid gates, rounded to BF16 " + "before optionally adding residual and rounding again.", + "T") .TypeConstraint("T", {"tensor(bfloat16)"}, "Inputs and output use BF16.") .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { propagateElemTypeFromInputToOutput(ctx, 0, 0); @@ -5947,64 +5958,66 @@ WebGPU currently supports only the default mode without Y_fp32. static OpSchema BranchwiseRMSNormSchema() { return OpSchema() - .SetDoc(BranchwiseRMSNorm_ver1_doc) - .Attr("epsilon", "Epsilon added before reciprocal square root.", - AttributeProto::FLOAT, 1e-5f) - .Attr("num_branches", - "Number of branches for flattened input. Omit or set to zero for grouped input.", - AttributeProto::INT, static_cast(0)) - .Attr("use_onnx_arithmetic", - "Use explicit FP32 mean/sqrt/divide/multiply arithmetic instead of reciprocal RMS. " - "Supported values are 0 (default) and 1.", - AttributeProto::INT, static_cast(0)) - .Input(0, "X", "Grouped (..., C, H) or flattened (..., C * H) input.", "T") - .Input(1, "scale", "Optional scale with shape (C * H), (C, H), or (H).", - "M", OpSchema::Optional) - .Output(0, "Y", "RMS-normalized output with the same shape as X.", "T") - .Output(1, "Y_fp32", "Optional float32 widening of the already T-rounded Y. " - "Requires use_onnx_arithmetic=1.", "F", OpSchema::Optional) - .TypeConstraint("T", - {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, - "Constrain input and output to floating-point tensors.") - .TypeConstraint("M", - {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, - "Constrain scale to floating-point tensors.") - .TypeConstraint("F", {"tensor(float)"}, "Widened output is float32.") - .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { - propagateElemTypeFromInputToOutput(ctx, 0, 0); - const auto start = getAttribute(ctx, "input_start", 0); - const auto width = getAttribute(ctx, "input_width", 0); - if (start < 0 || width < 0 || (width == 0 && start != 0)) { - fail_shape_inference("input_start and input_width must describe a nonnegative slice"); + .SetDoc(BranchwiseRMSNorm_ver1_doc) + .Attr("epsilon", "Epsilon added before reciprocal square root.", + AttributeProto::FLOAT, 1e-5f) + .Attr("num_branches", + "Number of branches for flattened input. Omit or set to zero for grouped input.", + AttributeProto::INT, static_cast(0)) + .Attr("use_onnx_arithmetic", + "Use explicit FP32 mean/sqrt/divide/multiply arithmetic instead of reciprocal RMS. " + "Supported values are 0 (default) and 1.", + AttributeProto::INT, static_cast(0)) + .Input(0, "X", "Grouped (..., C, H) or flattened (..., C * H) input.", "T") + .Input(1, "scale", "Optional scale with shape (C * H), (C, H), or (H).", + "M", OpSchema::Optional) + .Output(0, "Y", "RMS-normalized output with the same shape as X.", "T") + .Output(1, "Y_fp32", + "Optional float32 widening of the already T-rounded Y. " + "Requires use_onnx_arithmetic=1.", + "F", OpSchema::Optional) + .TypeConstraint("T", + {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, + "Constrain input and output to floating-point tensors.") + .TypeConstraint("M", + {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, + "Constrain scale to floating-point tensors.") + .TypeConstraint("F", {"tensor(float)"}, "Widened output is float32.") + .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + propagateElemTypeFromInputToOutput(ctx, 0, 0); + const auto start = getAttribute(ctx, "input_start", 0); + const auto width = getAttribute(ctx, "input_width", 0); + if (start < 0 || width < 0 || (width == 0 && start != 0)) { + fail_shape_inference("input_start and input_width must describe a nonnegative slice"); + } + if (width != 0 && (getAttribute(ctx, "use_onnx_arithmetic", 0) != 1 || + getAttribute(ctx, "num_branches", 0) != 0)) { + fail_shape_inference("input_width requires use_onnx_arithmetic=1 and grouped input"); + } + if (hasInputShape(ctx, 0)) { + auto shape = getInputShape(ctx, 0); + if (width != 0) { + if (shape.dim_size() < 2) fail_shape_inference("grouped input requires rank at least two"); + auto* last = shape.mutable_dim(shape.dim_size() - 1); + if (last->has_dim_value() && (start > last->dim_value() || width > last->dim_value() - start)) { + fail_shape_inference("input slice exceeds the last dimension"); + } + last->set_dim_value(width); } - if (width != 0 && (getAttribute(ctx, "use_onnx_arithmetic", 0) != 1 || - getAttribute(ctx, "num_branches", 0) != 0)) { - fail_shape_inference("input_width requires use_onnx_arithmetic=1 and grouped input"); + updateOutputShape(ctx, 0, shape); + } + if (ctx.getNumOutputs() > 1) { + if (getAttribute(ctx, "use_onnx_arithmetic", 0) != 1) { + fail_shape_inference("Y_fp32 requires use_onnx_arithmetic=1"); } + updateOutputElemType(ctx, 1, ONNX_NAMESPACE::TensorProto_DataType_FLOAT); if (hasInputShape(ctx, 0)) { auto shape = getInputShape(ctx, 0); - if (width != 0) { - if (shape.dim_size() < 2) fail_shape_inference("grouped input requires rank at least two"); - auto* last = shape.mutable_dim(shape.dim_size() - 1); - if (last->has_dim_value() && (start > last->dim_value() || width > last->dim_value() - start)) { - fail_shape_inference("input slice exceeds the last dimension"); - } - last->set_dim_value(width); - } - updateOutputShape(ctx, 0, shape); + if (width != 0) shape.mutable_dim(shape.dim_size() - 1)->set_dim_value(width); + updateOutputShape(ctx, 1, shape); } - if (ctx.getNumOutputs() > 1) { - if (getAttribute(ctx, "use_onnx_arithmetic", 0) != 1) { - fail_shape_inference("Y_fp32 requires use_onnx_arithmetic=1"); - } - updateOutputElemType(ctx, 1, ONNX_NAMESPACE::TensorProto_DataType_FLOAT); - if (hasInputShape(ctx, 0)) { - auto shape = getInputShape(ctx, 0); - if (width != 0) shape.mutable_dim(shape.dim_size() - 1)->set_dim_value(width); - updateOutputShape(ctx, 1, shape); - } - } - }); + } + }); } ONNX_MS_OPERATOR_SET_SCHEMA(BranchwiseRMSNorm, 1, BranchwiseRMSNormSchema()); @@ -6013,7 +6026,8 @@ ONNX_MS_OPERATOR_SET_SCHEMA( BranchwiseRMSNormSchema() .Attr("input_start", "CUDA-only starting offset in the last input dimension.", AttributeProto::INT, static_cast(0)) - .Attr("input_width", "CUDA-only width selected from the last input dimension before normalization. " + .Attr("input_width", + "CUDA-only width selected from the last input dimension before normalization. " "Zero retains the full input. Requires grouped input and use_onnx_arithmetic=1. " "Outputs retain the leading dimensions and use this last dimension.", AttributeProto::INT, static_cast(0))); diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc index 0cfd839266a82..996cf46d2da12 100644 --- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc @@ -3294,7 +3294,7 @@ The partition widths must be positive and sum to N.)DOC") const auto* split = ctx.getAttribute("split_sizes"); if ((split == nullptr && ctx.getNumOutputs() != 1) || (split != nullptr && (split->ints_size() < 1 || split->ints_size() > 8 || - static_cast(split->ints_size()) != ctx.getNumOutputs()))) { + static_cast(split->ints_size()) != ctx.getNumOutputs()))) { fail_shape_inference("Output count must match split_sizes, or be one when split_sizes is absent."); } int64_t split_total = 0; diff --git a/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu b/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu index bf960242503d4..75024a420d3a9 100644 --- a/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu +++ b/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu @@ -539,7 +539,7 @@ struct RoundedSquare { template Status reduce_square_matrix_columns(cudaStream_t stream, const TIn* input, float* output, - int m, int n, void* buffer, size_t buffer_size, int input_row_stride) { + int m, int n, void* buffer, size_t buffer_size, int input_row_stride) { return detail::call_reduce_matrix_columns( stream, input, output, m, n, buffer, buffer_size, input_row_stride); } diff --git a/onnxruntime/core/providers/cuda/reduction/reduction_functions.h b/onnxruntime/core/providers/cuda/reduction/reduction_functions.h index e5bb1b03820f0..9c877c0743b31 100644 --- a/onnxruntime/core/providers/cuda/reduction/reduction_functions.h +++ b/onnxruntime/core/providers/cuda/reduction/reduction_functions.h @@ -109,7 +109,7 @@ Status reduce_matrix_columns(cudaStream_t stream, const TIn* input, TOut* output /** Sum FP32-rounded squares across columns, preserving reduce_matrix_columns' summation order. */ template Status reduce_square_matrix_columns(cudaStream_t stream, const TIn* input, float* output, - int m, int n, void* buffer, size_t buffer_size, int input_row_stride = 0); + int m, int n, void* buffer, size_t buffer_size, int input_row_stride = 0); /** * Computes ReduceSum for an arbitrary set of axes. This is the general fallback diff --git a/onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc b/onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc index 8e015e80d5b32..7cac5d1d53c16 100644 --- a/onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc +++ b/onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc @@ -365,12 +365,16 @@ void TestPackedChannelHistory(int length, bool state_last, bool use_past, bool u test.AddAttribute("activation", silu ? "silu" : "none"); test.AddInput("input", {batch, length, width}, input); test.AddInput("weight", {channels, 1, kernel}, weight); - if (use_bias) test.AddInput("bias", {channels}, bias); - else test.AddOptionalInputEdge(); + if (use_bias) + test.AddInput("bias", {channels}, bias); + else + test.AddOptionalInputEdge(); const std::vector state_shape = state_last ? std::vector{batch, history, channels} - : std::vector{batch, channels, history}; - if (use_past) test.AddInput("past", state_shape, past); - else test.AddOptionalInputEdge(); + : std::vector{batch, channels, history}; + if (use_past) + test.AddInput("past", state_shape, past); + else + test.AddOptionalInputEdge(); test.AddOutput("output", {batch, length, channels}, output); test.AddOutput("present", state_shape, present); if (first_output) test.AddOutput("first", state_shape, first); 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 df85804bcc2e8..b9f6f91a75e4b 100644 --- a/onnxruntime/test/contrib_ops/cuda_kernels/mtp_moe_test.cc +++ b/onnxruntime/test/contrib_ops/cuda_kernels/mtp_moe_test.cc @@ -102,7 +102,7 @@ std::vector MakeNvfp4Weights(int experts, int rows, int k) { 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; + (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]; @@ -211,7 +211,8 @@ ReferenceResult ReferenceMtpMoe(const std::vector<__nv_bfloat16>& input, 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); + weight_row, result.hidden.data() + slot * ffn, ffn) * + kDownGlobalScale); const float weighted = RoundBf16(down_sum * result.routing_weights[slot]); total += weighted; } @@ -234,7 +235,7 @@ TEST(CUDA_EP_Unittest, MtpMoeRouterSelectsStableTopTenAndBf16Normalizes) { cudaStream_t stream = nullptr; CheckCuda(cudaStreamCreate(&stream)); const Status status = mtp_moe::LaunchRouter(device_logits.As<__nv_bfloat16>(), - device_scratch.As(), rows, stream); + device_scratch.As(), rows, stream); ASSERT_TRUE(status.IsOK()) << status.ErrorMessage(); CheckCuda(cudaStreamSynchronize(stream)); CheckCuda(cudaStreamDestroy(stream)); diff --git a/onnxruntime/test/contrib_ops/cuda_kernels/qmoe_fp8_compaction_test.cc b/onnxruntime/test/contrib_ops/cuda_kernels/qmoe_fp8_compaction_test.cc index e414fe31f873e..f54836ae48a94 100644 --- a/onnxruntime/test/contrib_ops/cuda_kernels/qmoe_fp8_compaction_test.cc +++ b/onnxruntime/test/contrib_ops/cuda_kernels/qmoe_fp8_compaction_test.cc @@ -69,7 +69,7 @@ void TestRoutingMaps(int rows, int experts, int k, bool normalize, bool tied) { CudaBuffer offsets((experts + 1) * sizeof(int64_t)); input.Upload(logits.data()); qmoe::LaunchSoftmaxTopK(input.As(), scales.As(), indices.As(), - rows, experts, k, normalize, nullptr); + rows, experts, k, normalize, nullptr); ASSERT_TRUE(qmoe::TryLaunchSmallDecodeExpertMaps( indices.As(), r2u.As(), u2r.As(), sorted.As(), offsets.As(), rows, experts, k, nullptr)); diff --git a/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc b/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc index 659bc157e9982..bc648f4f5ba37 100644 --- a/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/dynamic_sparse_attention_op_test.cc @@ -168,11 +168,11 @@ void RunDynamicSparseAttentionCase( } if (c.has_past) { tester.AddInput("past_key", - {c.batch_size, c.kv_num_heads, c.cache_sequence_length, c.head_size}, - ToTensorData(c.past_key)); + {c.batch_size, c.kv_num_heads, c.cache_sequence_length, c.head_size}, + ToTensorData(c.past_key)); tester.AddInput("past_value", - {c.batch_size, c.kv_num_heads, c.cache_sequence_length, c.head_size}, - ToTensorData(c.past_value)); + {c.batch_size, c.kv_num_heads, c.cache_sequence_length, c.head_size}, + ToTensorData(c.past_value)); } else { tester.AddOptionalInputEdge(); tester.AddOptionalInputEdge(); @@ -239,20 +239,24 @@ void RunDynamicSparseAttentionCase( tester.AddOptionalInputEdge(); } if (!c.key_scale.empty() || !c.value_scale.empty()) { - if (c.key_scale.empty()) tester.AddOptionalInputEdge(); - else tester.AddInput("key_scale", {int64_t(c.key_scale.size())}, c.key_scale); - if (c.value_scale.empty()) tester.AddOptionalInputEdge(); - else tester.AddInput("value_scale", {int64_t(c.value_scale.size())}, c.value_scale); + if (c.key_scale.empty()) + tester.AddOptionalInputEdge(); + else + tester.AddInput("key_scale", {int64_t(c.key_scale.size())}, c.key_scale); + if (c.value_scale.empty()) + tester.AddOptionalInputEdge(); + else + tester.AddInput("value_scale", {int64_t(c.value_scale.size())}, c.value_scale); } tester.AddOutput("output", {c.batch_size, c.sequence_length, hidden_size}, ToTensorData(c.expected_output)); tester.AddOutput("present_key", - {c.batch_size, c.kv_num_heads, c.cache_sequence_length, c.head_size}, - ToTensorData(c.expected_present_key)); + {c.batch_size, c.kv_num_heads, c.cache_sequence_length, c.head_size}, + ToTensorData(c.expected_present_key)); tester.AddOutput("present_value", - {c.batch_size, c.kv_num_heads, c.cache_sequence_length, c.head_size}, - ToTensorData(c.expected_present_value)); + {c.batch_size, c.kv_num_heads, c.cache_sequence_length, c.head_size}, + ToTensorData(c.expected_present_value)); if (c.expected_validation_status.has_value()) { tester.AddOutput("validation_status", {1}, {*c.expected_validation_status}); } @@ -420,11 +424,11 @@ TEST(DynamicSparseAttentionTest, Int8CacheRejectsMissingOrMalformedScales_CUDA) } else { c.key_scale.resize(8); } - const char* error = variant == 0 ? "INT8 caches require per-channel" - : variant == 1 ? "must be provided together" - : "INT8 cache scales must have shape"; + const char* error = variant == 0 ? "INT8 caches require per-channel" + : variant == 1 ? "must be provided together" + : "INT8 cache scales must have shape"; RunDynamicSparseAttentionCase(c, std::move(cuda_ep), - OpTester::ExpectResult::kExpectFailure, error); + OpTester::ExpectResult::kExpectFailure, error); } } @@ -517,14 +521,14 @@ DynamicSparseAttentionCase MakeGroupedInt8CacheCase(int sequence_length, bool au double dot = 0.0; for (int d = 0; d < c.head_size; ++d) { const float key = from_auxiliary ? c.auxiliary_key[offset + d] - : c.expected_present_key[offset + d] * c.key_scale[kv * c.head_size + d]; + : c.expected_present_key[offset + d] * c.key_scale[kv * c.head_size + d]; dot += c.query[out + d] * key; } const double weight = std::exp(dot * c.scale); denominator += weight; for (int d = 0; d < c.head_size; ++d) { const float value = from_auxiliary ? c.auxiliary_value[offset + d] - : c.expected_present_value[offset + d] * c.value_scale[kv * c.head_size + d]; + : c.expected_present_value[offset + d] * c.value_scale[kv * c.head_size + d]; numerator[d] += static_cast(weight * value); } } diff --git a/onnxruntime/test/contrib_ops/fused_tensor_ops_test.cc b/onnxruntime/test/contrib_ops/fused_tensor_ops_test.cc index 4aaeb82ed47ad..9d6797761e3a0 100644 --- a/onnxruntime/test/contrib_ops/fused_tensor_ops_test.cc +++ b/onnxruntime/test/contrib_ops/fused_tensor_ops_test.cc @@ -146,8 +146,8 @@ void RunSigmoidGateMulBroadcastResidualRounding(bool sigmoid_in_fp32, bool silu_ const BFloat16 denominator(1.0f + BFloat16(exp_value).ToFloat()); const BFloat16 reciprocal(1.0f / denominator.ToFloat()); const BFloat16 rounded_sigmoid = sigmoid_in_fp32 - ? BFloat16(sigmoid) - : (gate > 0 ? reciprocal : BFloat16(1.0f - reciprocal.ToFloat())); + ? BFloat16(sigmoid) + : (gate > 0 ? reciprocal : BFloat16(1.0f - reciprocal.ToFloat())); const BFloat16 activation = silu_gate ? BFloat16(gate * rounded_sigmoid.ToFloat()) : rounded_sigmoid; const BFloat16 product(values[i].ToFloat() * activation.ToFloat()); expected[i] = residual ? BFloat16(product.ToFloat() + residuals[i].ToFloat()) : product; @@ -186,7 +186,7 @@ TEST(FusedTensorOpsTest, SigmoidGateMulRejectsBroadcastAndResidualShapes) { test.AddAttribute("broadcast_gate", 1); test.AddInput("attention", {2, 3}, ToBFloat16(std::vector(6, 1.0f))); test.AddInput("projection", bad_residual ? std::vector{2, 1} : std::vector{1, 2}, - ToBFloat16({1.0f, 1.0f})); + ToBFloat16({1.0f, 1.0f})); if (bad_residual) test.AddInput("residual", {1}, ToBFloat16({1.0f})); test.AddOutput("output", {2, 3}, ToBFloat16(std::vector(6, 0.0f))); test.Run(OpTester::ExpectResult::kExpectFailure, @@ -208,8 +208,8 @@ TEST(FusedTensorOpsTest, SigmoidGateMulBroadcastEmptyAndInvalidAttributes) { test.AddInput("residual", {0, 3}, {}); test.AddOutput("output", {0, 3}, {}); const std::vector errors = {"", "broadcast_gate must be zero or one", - "sigmoid_in_fp32 must be zero or one", - "cannot be combined with packed layout attributes"}; + "sigmoid_in_fp32 must be zero or one", + "cannot be combined with packed layout attributes"}; test.Run(mode == 0 ? OpTester::ExpectResult::kExpectSuccess : OpTester::ExpectResult::kExpectFailure, errors[mode], {}, nullptr, &providers); } @@ -227,8 +227,8 @@ TEST(FusedTensorOpsTest, SigmoidGateMulSiluEmptyAndInvalidAttributes) { test.AddInput("projection", {0, mode == 4 ? 1 : 3}, {}); test.AddOutput("output", {0, 3}, {}); const std::vector errors = {"", "silu_gate must be zero or one", "are exclusive", - "cannot be combined with packed layout attributes", - "silu gate must have the same shape"}; + "cannot be combined with packed layout attributes", + "silu gate must have the same shape"}; test.Run(mode == 0 ? OpTester::ExpectResult::kExpectSuccess : OpTester::ExpectResult::kExpectFailure, errors[mode], {}, nullptr, &providers); } diff --git a/onnxruntime/test/contrib_ops/hyper_connection_op_test.cc b/onnxruntime/test/contrib_ops/hyper_connection_op_test.cc index 7942e8fc22425..d7e1bc7acdd72 100644 --- a/onnxruntime/test/contrib_ops/hyper_connection_op_test.cc +++ b/onnxruntime/test/contrib_ops/hyper_connection_op_test.cc @@ -113,7 +113,10 @@ class StandardRMSChainTester : public OpTester { public: StandardRMSChainTester(int output_type, float epsilon, int64_t start = 0, int64_t width = 0) : OpTester("StandardRMSChain", 17, kOnnxDomain, false), - output_type_(output_type), epsilon_(epsilon), start_(start), width_(width) {} + output_type_(output_type), + epsilon_(epsilon), + start_(start), + width_(width) {} void AddNodes(Graph& graph, std::vector& inputs, std::vector& outputs, std::vector>&) override { @@ -132,7 +135,8 @@ class StandardRMSChainTester : public OpTester { source = value("selected"); graph.AddNode("slice", "Slice", "", {inputs[0], constant("starts", start_), constant("ends", start_ + width_), - constant("axes", -1)}, {source}); + constant("axes", -1)}, + {source}); } auto* x = value("float_x"); graph.AddNode("cast", "Cast", "", {source}, {x}) @@ -200,7 +204,7 @@ void TestStandardRMSArithmetic(int type, bool packed = false) { std::vector expected; for (OpTester* tester : {static_cast(&reference), &fused}) { const std::vector input_shape = rows == 48 ? std::vector{1, 2, 24, stride} - : std::vector{1, rows, stride}; + : std::vector{1, rows, stride}; auto output_shape = input_shape; output_shape.back() = width; tester->AddInput("X", input_shape, input); diff --git a/onnxruntime/test/contrib_ops/matmul_block_scaled_fp8_test.cc b/onnxruntime/test/contrib_ops/matmul_block_scaled_fp8_test.cc index 94dbfa82a4636..0695e717ea6eb 100644 --- a/onnxruntime/test/contrib_ops/matmul_block_scaled_fp8_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_block_scaled_fp8_test.cc @@ -75,9 +75,7 @@ void TestFp8SplitOutputs() { if (!HasCudaEnvironment(800)) { GTEST_SKIP() << "CUDA device is required for split-output FP8 matmul."; } - for (const auto& [m, k] : {std::pair{0, 256}, {1, 256}, {2, 256}, {2, 320}, - {2, 1088}, {9, 256}, - {32, 256}, {33, 256}, {65, 256}, {2, 16}, {2, 30}, {2, 0}}) { + for (const auto& [m, k] : {std::pair{0, 256}, {1, 256}, {2, 256}, {2, 320}, {2, 1088}, {9, 256}, {32, 256}, {33, 256}, {65, 256}, {2, 16}, {2, 30}, {2, 0}}) { for (const bool activation_scale : {false, true}) { SCOPED_TRACE("M=" + std::to_string(m) + " K=" + std::to_string(k) + " W8A8=" + std::to_string(activation_scale)); @@ -85,9 +83,9 @@ void TestFp8SplitOutputs() { constexpr int64_t block_size = 64; const int64_t k_blocks = (k + block_size - 1) / block_size; const std::vector widths = m == 2 && k == 256 - ? std::vector{1, 1, 1, 1, 1, 1, 1, 10} - : (m == 33 ? std::vector{17} - : std::vector{1, 7, 9}); + ? std::vector{1, 1, 1, 1, 1, 1, 1, 10} + : (m == 33 ? std::vector{17} + : std::vector{1, 7, 9}); std::vector a; std::vector b; std::vector scales; @@ -162,8 +160,7 @@ TEST(MatMulBlockQuantizedFp8WeightOpTest, RejectsInvalidSplitSizes) { if (!HasCudaEnvironment(800)) { GTEST_SKIP() << "CUDA device is required for split-output FP8 matmul."; } - for (const std::vector& widths : {std::vector{}, {0, 17}, {-1, 18}, {8, 8}, - {1, 1, 1, 1, 1, 1, 1, 1, 9}}) { + for (const std::vector& widths : {std::vector{}, {0, 17}, {-1, 18}, {8, 8}, {1, 1, 1, 1, 1, 1, 1, 1, 9}}) { OpTester test("MatMulBlockQuantizedFp8Weight", 1, kMSDomain); test.AddAttribute>("split_sizes", widths); test.AddInput("A", {1, 16}, std::vector(16)); diff --git a/onnxruntime/test/contrib_ops/rotary_embedding_op_test.cc b/onnxruntime/test/contrib_ops/rotary_embedding_op_test.cc index 9ebbc10c97bf2..a6957e18bc194 100644 --- a/onnxruntime/test/contrib_ops/rotary_embedding_op_test.cc +++ b/onnxruntime/test/contrib_ops/rotary_embedding_op_test.cc @@ -27,7 +27,10 @@ class FloatRotaryChainTester : public OpTester { public: FloatRotaryChainTester(int type, int heads, int rotary_dim, bool interleaved) : OpTester("RotaryEmbedding", 1, kMSDomain, false), - type_(type), heads_(heads), rotary_dim_(rotary_dim), interleaved_(interleaved) {} + type_(type), + heads_(heads), + rotary_dim_(rotary_dim), + interleaved_(interleaved) {} void AddNodes(Graph& graph, std::vector& inputs, std::vector& outputs, std::vector>&) override { @@ -74,7 +77,7 @@ void TestFloatRotaryArithmetic(int type) { const int rotary_dim = 64; const int cache_rows = batch * sequence + 3; std::vector shape = transposed ? std::vector{batch, heads, sequence, width} - : std::vector{batch, sequence, heads * width}; + : std::vector{batch, sequence, heads * width}; std::vector input(batch * sequence * heads * width), cosine(cache_rows * rotary_dim / 2), sine(cosine.size()); const bool scalar_position = !transposed; @@ -94,8 +97,7 @@ void TestFloatRotaryArithmetic(int type) { std::vector expected; for (OpTester* tester : {static_cast(&reference), &fused}) { tester->AddInput("input", shape, input); - tester->AddInput("positions", scalar_position ? std::vector{1} - : std::vector{batch, sequence}, positions); + tester->AddInput("positions", scalar_position ? std::vector{1} : std::vector{batch, sequence}, positions); tester->AddInput("cosine", {cache_rows, rotary_dim / 2}, cosine); tester->AddInput("sine", {cache_rows, rotary_dim / 2}, sine); tester->AddOutput("output", shape, std::vector(input.size())); diff --git a/onnxruntime/test/python/onnxruntime_test_python_symbolic_shape_infer.py b/onnxruntime/test/python/onnxruntime_test_python_symbolic_shape_infer.py index 96bb35672a6cf..6ac581d72cab8 100644 --- a/onnxruntime/test/python/onnxruntime_test_python_symbolic_shape_infer.py +++ b/onnxruntime/test/python/onnxruntime_test_python_symbolic_shape_infer.py @@ -882,9 +882,6 @@ def infer(query_width, packed): self.assertEqual(separate_graph.output[0].type.tensor_type.shape.dim[2].dim_value, 33) self.assertFalse(separate_graph.output[1].type.tensor_type.HasField("shape")) - - - def test_unsqueeze_opset_11(self): graph = helper.make_graph( [