diff --git a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc new file mode 100644 index 0000000000000..555c98f1fb16d --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc @@ -0,0 +1,260 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "contrib_ops/cuda/bert/fused_tensor_ops.h" + +#include +#include +#include + +#include "core/common/float8.h" +#include "core/providers/cuda/cuda_common.h" +#include "core/providers/cuda/cuda_type_conversion.h" + +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()), \ + ArgMaxLastAxis); + +REGISTER_ARGMAX_LAST_AXIS(float) +REGISTER_ARGMAX_LAST_AXIS(MLFloat16) +REGISTER_ARGMAX_LAST_AXIS(BFloat16) + +#undef REGISTER_ARGMAX_LAST_AXIS + +ONNX_OPERATOR_KERNEL_EX( + EmbeddingRMSNorm, kMSDomain, 1, kCudaExecutionProvider, + (*KernelDefBuilder::Create()) + .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()), + RowSelect); + +ONNX_OPERATOR_KERNEL_EX( + SigmoidGateMul, kMSDomain, 1, kCudaExecutionProvider, + (*KernelDefBuilder::Create()) + .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()), + DequantizeLinearFp8Bf16); +#endif + +ONNX_OPERATOR_KERNEL_EX( + UntileHyperConnectionWeights, kMSDomain, 1, kCudaExecutionProvider, + (*KernelDefBuilder::Create()) + .TypeConstraint("T", DataTypeImpl::GetTensorType()), + UntileHyperConnectionWeights); + +template +ArgMaxLastAxis::ArgMaxLastAxis(const OpKernelInfo& info) + : CudaKernel(info), keepdims_(info.GetAttrOrDefault("keepdims", 0) != 0) {} + +template +Status ArgMaxLastAxis::ComputeInternal(OpKernelContext* context) const { + const Tensor* input = context->Input(0); + const auto& shape = input->Shape(); + ORT_RETURN_IF_NOT(shape.NumDimensions() >= 1 && + shape[shape.NumDimensions() - 1] > 0, + "ArgMaxLastAxis requires a non-empty final dimension"); + const int64_t width = shape[shape.NumDimensions() - 1]; + const int64_t rows = shape.Size() / width; + TensorShapeVector output_dims(shape.GetDims().begin(), shape.GetDims().end() - 1); + if (keepdims_) output_dims.push_back(1); + Tensor* output = context->Output(0, TensorShape(output_dims)); + const size_t partial_count = SafeInt(rows) * GetArgMaxLastAxisPartitions(width); + auto partial_values = GetScratchBuffer(partial_count, GetComputeStream(context)); + auto partial_indices = GetScratchBuffer(partial_count, GetComputeStream(context)); + using CudaT = typename onnxruntime::cuda::OrtToCudaType::type; + return LaunchArgMaxLastAxis(Stream(context), reinterpret_cast(input->Data()), + input->GetElementType(), rows, width, + output->MutableData(), partial_values.get(), partial_indices.get()); +} + +EmbeddingRMSNorm::EmbeddingRMSNorm(const OpKernelInfo& info) + : CudaKernel(info), epsilon_(info.GetAttrOrDefault("epsilon", 1e-6f)) { + ORT_ENFORCE(epsilon_ >= 0.0f, "EmbeddingRMSNorm epsilon must be non-negative"); +} + +Status EmbeddingRMSNorm::ComputeInternal(OpKernelContext* context) const { + const Tensor* ids = context->Input(0); + const Tensor* table = context->Input(1); + const Tensor* weight = context->Input(2); + ORT_RETURN_IF_NOT(table->Shape().NumDimensions() == 2, + "EmbeddingRMSNorm table must have shape [vocabulary, hidden]"); + const int64_t hidden_size = table->Shape()[1]; + ORT_RETURN_IF_NOT(weight->Shape().NumDimensions() == 1 && + weight->Shape()[0] == hidden_size, + "EmbeddingRMSNorm weight must match the embedding hidden size"); + TensorShapeVector output_dims(ids->Shape().GetDims().begin(), ids->Shape().GetDims().end()); + output_dims.push_back(hidden_size); + Tensor* output = context->Output(0, TensorShape(output_dims)); + return LaunchEmbeddingRMSNorm(Stream(context), ids->Data(), + reinterpret_cast(table->Data()), + weight->Data(), epsilon_, + reinterpret_cast(output->MutableData()), + ids->Shape().Size(), table->Shape()[0], hidden_size); +} + +Status RowSelect::ComputeInternal(OpKernelContext* context) const { + const Tensor* selection = context->Input(0); + ORT_RETURN_IF_NOT(selection->Shape().Size() == 1, "RowSelect selection must have one element"); + std::array inputs{}; + std::array outputs{}; + std::array row_sizes{}; + for (int i = 0; i < 3; ++i) { + const Tensor* input = context->Input(i + 1); + const auto& shape = input->Shape(); + ORT_RETURN_IF_NOT(shape.NumDimensions() >= 2 && shape[0] == 1 && shape[1] == 2, + "RowSelect inputs must have shape [1, 2, ...]"); + row_sizes[i] = shape.Size() / 2; + TensorShapeVector output_dims(shape.GetDims().begin(), shape.GetDims().end()); + output_dims[1] = 1; + Tensor* output = context->Output(i, TensorShape(output_dims)); + inputs[i] = reinterpret_cast(input->Data()); + outputs[i] = reinterpret_cast(output->MutableData()); + } + return LaunchRowSelect(Stream(context), selection->Data(), inputs.data(), + outputs.data(), row_sizes.data()); +} + +SigmoidGateMul::SigmoidGateMul(const OpKernelInfo& info) + : CudaKernel(info), + head_dim_(static_cast(info.GetAttrOrDefault("head_dim", 0))), + heads_(static_cast(info.GetAttrOrDefault("num_heads", 0))), + q_stride_(static_cast(info.GetAttrOrDefault("q_stride", 0))), + tok_stride_(info.GetAttrOrDefault("tok_stride", 0)), + q_offset_(info.GetAttrOrDefault("q_offset", 0)), + broadcast_gate_(info.GetAttrOrDefault("broadcast_gate", 0) == 1), + sigmoid_in_fp32_(info.GetAttrOrDefault("sigmoid_in_fp32", 1) == 1), + silu_gate_(info.GetAttrOrDefault("silu_gate", 0) == 1) { + const int64_t broadcast = info.GetAttrOrDefault("broadcast_gate", 0); + ORT_ENFORCE(broadcast == 0 || broadcast == 1, "SigmoidGateMul broadcast_gate must be zero or one"); + const int64_t fp32 = info.GetAttrOrDefault("sigmoid_in_fp32", 1); + ORT_ENFORCE(fp32 == 0 || fp32 == 1, "SigmoidGateMul sigmoid_in_fp32 must be zero or one"); + 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), + "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), + "SigmoidGateMul broadcast_gate/silu_gate cannot be combined with packed layout attributes"); +} + +Status SigmoidGateMul::ComputeInternal(OpKernelContext* context) const { + const Tensor* attention = context->Input(0); + const Tensor* projection = context->Input(1); + const Tensor* residual = context->Input(2); + const int64_t count = attention->Shape().Size(); + ORT_RETURN_IF_NOT(residual == nullptr || residual->Shape() == attention->Shape(), + "SigmoidGateMul residual must have the same shape as attention"); + int64_t gate_row_width = 0; + if (silu_gate_) { + ORT_RETURN_IF_NOT(projection->Shape() == attention->Shape(), + "SigmoidGateMul silu gate must have the same shape as attention"); + } else if (broadcast_gate_) { + TensorShapeVector expected(attention->Shape().GetDims().begin(), attention->Shape().GetDims().end()); + ORT_RETURN_IF_NOT(!expected.empty() && expected.back() > 0, + "SigmoidGateMul broadcast_gate requires a positive final attention dimension"); + gate_row_width = expected.back(); + expected.back() = 1; + ORT_RETURN_IF_NOT(projection->Shape() == TensorShape(expected), + "SigmoidGateMul broadcast gate must match attention with a singleton final dimension"); + } 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"); + 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"); + } + Tensor* output = context->Output(0, attention->Shape()); + return LaunchSigmoidGateMul( + Stream(context), reinterpret_cast(attention->Data()), + reinterpret_cast(projection->Data()), + reinterpret_cast(output->MutableData()), count, head_dim_, heads_, + q_stride_, tok_stride_, q_offset_, gate_row_width, + residual == nullptr ? nullptr : reinterpret_cast(residual->Data()), + sigmoid_in_fp32_, silu_gate_); +} + +DequantizeLinearFp8Bf16::DequantizeLinearFp8Bf16(const OpKernelInfo& info) + : CudaKernel(info), axis_(info.GetAttrOrDefault("axis", 1)) {} + +Status DequantizeLinearFp8Bf16::ComputeInternal(OpKernelContext* context) const { + const Tensor* input = context->Input(0); + const Tensor* scales = context->Input(1); + const Tensor* zero_point = context->Input(2); + const auto& input_shape = input->Shape(); + ORT_RETURN_IF_NOT(input_shape.NumDimensions() > 0, + "DequantizeLinear FP8/BF16 input must have rank at least one"); + int64_t axis = axis_ < 0 ? axis_ + static_cast(input_shape.NumDimensions()) : axis_; + + int64_t axis_size = 1; + int64_t inner_size = 1; + if (scales->Shape().Size() != 1) { + ORT_RETURN_IF_NOT(axis >= 0 && axis < static_cast(input_shape.NumDimensions()), + "DequantizeLinear axis is out of range"); + ORT_RETURN_IF_NOT(scales->Shape().NumDimensions() == 1 && + scales->Shape().Size() == input_shape[axis], + "DequantizeLinear scale must be scalar or match the selected axis"); + axis_size = scales->Shape().Size(); + inner_size = input_shape.SizeFromDimension(static_cast(axis + 1)); + } + if (zero_point != nullptr) { + ORT_RETURN_IF_NOT(zero_point->Shape().Size() == 1 || + zero_point->Shape() == scales->Shape(), + "DequantizeLinear zero point must be scalar or match the scale"); + } + + Tensor* output = context->Output(0, input_shape); + return LaunchDequantizeLinearFp8Bf16( + Stream(context), reinterpret_cast(input->DataRaw()), + zero_point == nullptr ? nullptr : reinterpret_cast(zero_point->DataRaw()), + scales->Data(), + reinterpret_cast(output->MutableData()), input_shape.Size(), + axis_size, inner_size, zero_point == nullptr || zero_point->Shape().Size() == 1); +} + +UntileHyperConnectionWeights::UntileHyperConnectionWeights(const OpKernelInfo& info) + : CudaKernel(info), kind_(static_cast(info.GetAttrOrDefault("kind", -1))) { + ORT_ENFORCE(kind_ == 0 || kind_ == 1, "UntileHyperConnectionWeights kind must be 0 or 1"); +} + +Status UntileHyperConnectionWeights::ComputeInternal(OpKernelContext* context) const { + const Tensor* input = context->Input(0); + ORT_RETURN_IF_NOT(input->Shape() == TensorShape({160, 20480}), + "UntileHyperConnectionWeights input must have shape [160, 20480]"); + const TensorShape output_shape = kind_ == 0 ? TensorShape({10240, 320}) : TensorShape({320, 10240}); + Tensor* output = context->Output(0, output_shape); + return LaunchUntileHyperConnectionWeights( + Stream(context), reinterpret_cast(input->Data()), + reinterpret_cast(output->MutableData()), kind_); +} + +template class ArgMaxLastAxis; +template class ArgMaxLastAxis; +template class ArgMaxLastAxis; + +} // namespace onnxruntime::contrib::cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu new file mode 100644 index 0000000000000..789edc4a26615 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu @@ -0,0 +1,352 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "contrib_ops/cuda/bert/fused_tensor_ops.h" + +#include +#include +#include +#include +#include + +#include "core/providers/cuda/cuda_common.h" +#include "core/providers/cuda/cu_inc/common.cuh" + +namespace onnxruntime::contrib::cuda { +namespace { + +constexpr int kThreads = 256; +constexpr int kArgmaxThreads = 256; + +__device__ __forceinline__ float ToFloat(float value) { return value; } +__device__ __forceinline__ float ToFloat(half value) { return __half2float(value); } +__device__ __forceinline__ float ToFloat(__nv_bfloat16 value) { return __bfloat162float(value); } + +template +__global__ void ArgmaxRowsKernel(const T* input, const int64_t* input_indices, + int64_t width, int64_t partitions, float* partial_values, + int64_t* output) { + __shared__ float values[kArgmaxThreads]; + __shared__ int64_t indices[kArgmaxThreads]; + const int tid = threadIdx.x; + const int64_t row = blockIdx.x / partitions; + const int64_t partition = blockIdx.x % partitions; + const int64_t end = partitions == 1 ? width : min(width, (partition + 1) * 4096); + float best_value = -CUDART_INF_F; + int64_t best_index = INT64_MAX; + for (int64_t col = partition * 4096 + tid; col < end; col += blockDim.x) { + const float value = ToFloat(input[row * width + col]); + const int64_t index = input_indices == nullptr ? col : input_indices[row * width + col]; + if (value > best_value || (value == best_value && index < best_index)) { + best_value = value; + best_index = index; + } + } + values[tid] = best_value; + indices[tid] = best_index; + __syncthreads(); + for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { + if (tid < stride) { + const float other_value = values[tid + stride]; + const int64_t other_index = indices[tid + stride]; + if (other_value > values[tid] || + (other_value == values[tid] && other_index < indices[tid])) { + values[tid] = other_value; + indices[tid] = other_index; + } + } + __syncthreads(); + } + if (tid == 0) { + output[blockIdx.x] = indices[0]; + if (partial_values != nullptr) partial_values[blockIdx.x] = values[0]; + } +} + +template +Status LaunchTypedArgmax(cudaStream_t stream, const T* input, int64_t rows, + int64_t width, int64_t* output, + float* partial_values, int64_t* partial_indices) { + if (rows == 0) return Status::OK(); + const int64_t partitions = GetArgMaxLastAxisPartitions(width); + ORT_RETURN_IF_NOT(rows <= std::numeric_limits::max() / partitions, + "ArgMaxLastAxis has too many rows for a CUDA grid"); + ArgmaxRowsKernel<<(rows * partitions), kArgmaxThreads, 0, stream>>>( + input, nullptr, width, partitions, partitions == 1 ? nullptr : partial_values, + partitions == 1 ? output : partial_indices); + CUDA_RETURN_IF_ERROR(cudaGetLastError()); + if (partitions > 1) { + ArgmaxRowsKernel<<(rows), kArgmaxThreads, 0, stream>>>( + partial_values, partial_indices, partitions, 1, nullptr, output); + } + return CUDA_CALL(cudaGetLastError()); +} + +__global__ void EmbedNormKernel(const int64_t* ids, const uint16_t* table, + const float* weight, float epsilon, + uint16_t* output, int64_t rows, int64_t hidden_size) { + __shared__ float sums[512]; + const int64_t row_index = blockIdx.x; + const int64_t token_id = ids[row_index]; + const uint16_t* embedding = table + token_id * hidden_size; + float sum = 0.0f; + for (int64_t col = threadIdx.x; col < hidden_size; col += blockDim.x) { + const float value = __uint_as_float(static_cast(embedding[col]) << 16); + sum += value * value; + } + sums[threadIdx.x] = sum; + __syncthreads(); + for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) { + sums[threadIdx.x] += sums[threadIdx.x + stride]; + } + __syncthreads(); + } + const float scale = 1.0f / sqrtf(sums[0] / static_cast(hidden_size) + epsilon); + for (int64_t col = threadIdx.x; col < hidden_size; col += blockDim.x) { + const float value = __uint_as_float(static_cast(embedding[col]) << 16); + const float normalized = (value * scale) * weight[col]; + output[row_index * hidden_size + col] = + __bfloat16_as_ushort(__float2bfloat16_rn(normalized)); + } + (void)rows; +} + +__global__ void RowSelectKernel(const int64_t* selection, + const uint16_t* input0, const uint16_t* input1, + const uint16_t* input2, uint16_t* output0, + uint16_t* output1, uint16_t* output2, + int64_t row_size0, int64_t row_size1, + int64_t row_size2) { + const int which = blockIdx.y; + const uint16_t* input = which == 0 ? input0 : (which == 1 ? input1 : input2); + uint16_t* output = which == 0 ? output0 : (which == 1 ? output1 : output2); + const int64_t row_size = which == 0 ? row_size0 : (which == 1 ? row_size1 : row_size2); + const int64_t selected_row = selection[0] != 0 ? 1 : 0; + for (int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + i < row_size; i += static_cast(gridDim.x) * blockDim.x) { + output[i] = input[selected_row * row_size + i]; + } +} + +__device__ __forceinline__ float StableSigmoid(float value) { + if (value >= 0.0f) { + return 1.0f / (1.0f + expf(-value)); + } + const float exp_value = expf(value); + return exp_value / (1.0f + exp_value); +} + +template +__global__ void SigmoidGateMulKernel(const __nv_bfloat16* attention, + const __nv_bfloat16* projection, + __nv_bfloat16* output, int64_t count, + int head_dim, int heads, int q_stride, + int64_t token_stride, int64_t q_offset, + int64_t gate_row_width, const __nv_bfloat16* residual, + bool sigmoid_in_fp32) { + const int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= count) return; + int64_t gate_index; + if constexpr (SiluGate) { + gate_index = i; + } else if constexpr (BroadcastGate) { + gate_index = i / gate_row_width; + } else { + const int64_t row = i / head_dim; + const int64_t token = row / heads; + const int head = static_cast(row % heads); + const int dim = static_cast(i % head_dim); + gate_index = token * token_stride + q_offset + + static_cast(head) * q_stride + head_dim + dim; + } + const float gate = __bfloat162float(projection[gate_index]); + __nv_bfloat16 sigmoid_bf16; + if (sigmoid_in_fp32) { + sigmoid_bf16 = __float2bfloat16_rn(StableSigmoid(gate)); + } else { + using namespace onnxruntime::cuda; + 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)); + sigmoid_bf16 = __float2bfloat16_rn(static_cast(value)); + } + if constexpr (SiluGate) { + sigmoid_bf16 = __float2bfloat16_rn(gate * __bfloat162float(sigmoid_bf16)); + } + __nv_bfloat16 value = __float2bfloat16_rn( + __bfloat162float(attention[i]) * __bfloat162float(sigmoid_bf16)); + if constexpr (HasResidual) { + value = __float2bfloat16_rn(__bfloat162float(value) + __bfloat162float(residual[i])); + } + output[i] = value; +} + +__device__ __forceinline__ float DecodeE4M3FN(uint8_t value) { + const float sign = (value & 0x80u) != 0 ? -1.0f : 1.0f; + const int exponent = (value >> 3) & 0x0f; + const int mantissa = value & 0x07; + if (exponent == 0) { + return sign * ldexpf(static_cast(mantissa), -9); + } + if (exponent == 0x0f && mantissa == 0x07) { + return CUDART_NAN_F; + } + return sign * ldexpf(1.0f + static_cast(mantissa) * 0.125f, + exponent - 7); +} + +__global__ void DequantizeLinearFp8Bf16Kernel(const uint8_t* input, + const uint8_t* zero_point, + const float* scales, uint16_t* output, + int64_t count, int64_t axis_size, + int64_t inner_size, + bool scalar_zero_point) { + const int64_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (index >= count) return; + float value = DecodeE4M3FN(input[index]); + const int64_t scale_index = axis_size == 1 ? 0 : (index / inner_size) % axis_size; + if (zero_point != nullptr) { + value -= DecodeE4M3FN(zero_point[scalar_zero_point ? 0 : scale_index]); + } + value *= scales[scale_index]; + output[index] = __bfloat16_as_ushort(__float2bfloat16_rn(value)); +} + +__global__ void UntileHcBf16Kernel(const uint16_t* input, uint16_t* output, int kind) { + constexpr int64_t kVectorCount = 204800; + constexpr int64_t kElementCount = kVectorCount * 16; + const int64_t element = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (element >= kElementCount) return; + const int64_t vector_index = element / 16; + const int64_t lane = element % 16; + int64_t destination_vector; + if (kind == 0) { + const int64_t rp = vector_index & 127; + const int64_t t1 = vector_index >> 7; + const int64_t i = t1 % 10; + const int64_t t2 = t1 / 10; + const int64_t j = t2 % 20; + const int64_t s = t2 / 20; + destination_vector = (((s * 10 + i) * 128 + rp) * 20 + j); + } else { + const int64_t c = vector_index & 3; + const int64_t t1 = vector_index >> 2; + const int64_t rp = t1 & 31; + const int64_t t2 = t1 >> 5; + const int64_t i = t2 % 10; + const int64_t hb = t2 / 10; + destination_vector = (((i * 32 + rp) * 4 + c) * 160 + hb); + } + output[destination_vector * 16 + lane] = input[element]; +} + +} // namespace + +Status LaunchArgMaxLastAxis(cudaStream_t stream, const void* input, int input_type, + int64_t rows, int64_t width, int64_t* output, + float* partial_values, int64_t* partial_indices) { + ORT_RETURN_IF_NOT(rows >= 0 && width > 0, + "ArgMaxLastAxis requires a non-negative row count and positive vocabulary dimension"); + if (rows == 0) return Status::OK(); + switch (input_type) { + case ONNX_NAMESPACE::TensorProto_DataType_FLOAT: + return LaunchTypedArgmax(stream, static_cast(input), rows, width, output, + partial_values, partial_indices); + case ONNX_NAMESPACE::TensorProto_DataType_FLOAT16: + return LaunchTypedArgmax(stream, static_cast(input), rows, width, output, + partial_values, partial_indices); + case ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16: + return LaunchTypedArgmax(stream, static_cast(input), rows, width, output, + partial_values, partial_indices); + default: + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "ArgMaxLastAxis input type is not supported"); + } +} + +Status LaunchEmbeddingRMSNorm(cudaStream_t stream, const int64_t* ids, const uint16_t* table, + const float* weight, float epsilon, uint16_t* output, + int64_t rows, int64_t vocab_size, int64_t hidden_size) { + ORT_RETURN_IF_NOT(hidden_size > 0 && vocab_size > 0, + "EmbeddingRMSNorm requires non-empty table dimensions"); + ORT_RETURN_IF_NOT(rows <= std::numeric_limits::max(), + "EmbeddingRMSNorm has too many rows for a CUDA grid"); + if (rows == 0) return Status::OK(); + EmbedNormKernel<<(rows), 512, 0, stream>>>( + ids, table, weight, epsilon, output, rows, hidden_size); + return CUDA_CALL(cudaGetLastError()); +} + +Status LaunchRowSelect(cudaStream_t stream, const int64_t* selection, + const uint16_t* const* inputs, uint16_t* const* outputs, + const int64_t* row_sizes) { + const dim3 grid(8, 3); + RowSelectKernel<<>>( + selection, inputs[0], inputs[1], inputs[2], outputs[0], outputs[1], outputs[2], + row_sizes[0], row_sizes[1], row_sizes[2]); + return CUDA_CALL(cudaGetLastError()); +} + +Status LaunchSigmoidGateMul(cudaStream_t stream, const uint16_t* attention, + const uint16_t* projection, uint16_t* output, int64_t count, + int head_dim, int heads, int q_stride, + int64_t token_stride, int64_t q_offset, + int64_t gate_row_width, const uint16_t* residual, bool sigmoid_in_fp32, + bool silu_gate) { + if (count == 0) return Status::OK(); + 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) \ + SigmoidGateMulKernel<<(block_count), kThreads, 0, stream>>>( \ + 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) { + LAUNCH_GATE(false, true, true); + } else { + LAUNCH_GATE(false, false, true); + } + } else if (gate_row_width > 0) { + if (residual != nullptr) { + LAUNCH_GATE(true, true, false); + } else { + LAUNCH_GATE(true, false, false); + } + } else if (residual != nullptr) { + LAUNCH_GATE(false, true, false); + } else { + LAUNCH_GATE(false, false, false); + } +#undef LAUNCH_GATE + return CUDA_CALL(cudaGetLastError()); +} + +Status LaunchDequantizeLinearFp8Bf16(cudaStream_t stream, const uint8_t* input, + const uint8_t* zero_point, const float* scales, + uint16_t* output, int64_t count, + int64_t axis_size, int64_t inner_size, + bool scalar_zero_point) { + if (count == 0) return Status::OK(); + const int64_t block_count = (count - 1) / kThreads + 1; + ORT_RETURN_IF_NOT(block_count <= std::numeric_limits::max(), + "DequantizeLinear FP8/BF16 requires too many CUDA blocks"); + DequantizeLinearFp8Bf16Kernel<<(block_count), kThreads, 0, stream>>>( + input, zero_point, scales, output, count, axis_size, inner_size, scalar_zero_point); + return CUDA_CALL(cudaGetLastError()); +} + +Status LaunchUntileHyperConnectionWeights(cudaStream_t stream, const uint16_t* input, + uint16_t* output, int kind) { + constexpr int64_t kElementCount = 204800ll * 16; + constexpr int threads = 256; + constexpr int blocks = static_cast((kElementCount + threads - 1) / threads); + UntileHcBf16Kernel<<>>(input, output, kind); + return CUDA_CALL(cudaGetLastError()); +} + +} // namespace onnxruntime::contrib::cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.h b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.h new file mode 100644 index 0000000000000..7cc6b05922159 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.h @@ -0,0 +1,97 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include + +#include "core/providers/cuda/cuda_kernel.h" + +namespace onnxruntime::contrib::cuda { + +Status LaunchArgMaxLastAxis(cudaStream_t stream, const void* input, int input_type, + int64_t rows, int64_t width, int64_t* output, + float* partial_values, int64_t* partial_indices); +inline int64_t GetArgMaxLastAxisPartitions(int64_t width) { + return 1 + (width - 1) / 4096; +} +Status LaunchEmbeddingRMSNorm(cudaStream_t stream, const int64_t* ids, const uint16_t* table, + const float* weight, float epsilon, uint16_t* output, + int64_t rows, int64_t vocab_size, int64_t hidden_size); +Status LaunchRowSelect(cudaStream_t stream, const int64_t* selection, + const uint16_t* const* inputs, uint16_t* const* outputs, + const int64_t* row_sizes); +Status LaunchSigmoidGateMul(cudaStream_t stream, const uint16_t* attention, + const uint16_t* projection, uint16_t* output, int64_t count, + int head_dim, int heads, int q_stride, + int64_t token_stride, int64_t q_offset, + int64_t gate_row_width, const uint16_t* residual, bool sigmoid_in_fp32, + bool silu_gate); +Status LaunchDequantizeLinearFp8Bf16(cudaStream_t stream, const uint8_t* input, + const uint8_t* zero_point, const float* scales, + uint16_t* output, int64_t count, + int64_t axis_size, int64_t inner_size, + bool scalar_zero_point); +Status LaunchUntileHyperConnectionWeights(cudaStream_t stream, const uint16_t* input, + uint16_t* output, int kind); + +template +class ArgMaxLastAxis final : public onnxruntime::cuda::CudaKernel { + public: + explicit ArgMaxLastAxis(const OpKernelInfo& info); + Status ComputeInternal(OpKernelContext* context) const override; + + private: + bool keepdims_; +}; + +class EmbeddingRMSNorm final : public onnxruntime::cuda::CudaKernel { + public: + explicit EmbeddingRMSNorm(const OpKernelInfo& info); + Status ComputeInternal(OpKernelContext* context) const override; + + private: + float epsilon_; +}; + +class RowSelect final : public onnxruntime::cuda::CudaKernel { + public: + explicit RowSelect(const OpKernelInfo& info) : CudaKernel(info) {} + Status ComputeInternal(OpKernelContext* context) const override; +}; + +class SigmoidGateMul final : public onnxruntime::cuda::CudaKernel { + public: + explicit SigmoidGateMul(const OpKernelInfo& info); + Status ComputeInternal(OpKernelContext* context) const override; + + private: + int head_dim_; + int heads_; + int q_stride_; + int64_t tok_stride_; + int64_t q_offset_; + bool broadcast_gate_; + bool sigmoid_in_fp32_; + bool silu_gate_; +}; + +class DequantizeLinearFp8Bf16 final : public onnxruntime::cuda::CudaKernel { + public: + explicit DequantizeLinearFp8Bf16(const OpKernelInfo& info); + Status ComputeInternal(OpKernelContext* context) const override; + + private: + int64_t axis_; +}; + +class UntileHyperConnectionWeights final : public onnxruntime::cuda::CudaKernel { + public: + explicit UntileHyperConnectionWeights(const OpKernelInfo& info); + Status ComputeInternal(OpKernelContext* context) const override; + + private: + int kind_; +}; + +} // namespace onnxruntime::contrib::cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix.cc b/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix.cc new file mode 100644 index 0000000000000..982698c00b222 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix.cc @@ -0,0 +1,165 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "contrib_ops/cuda/bert/hyper_connection_low_rank_mix.h" + +#include + +#include "contrib_ops/cpu/hyper_connection_helper.h" +#include "contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.h" +#include "core/providers/cuda/cuda_type_conversion.h" + +namespace onnxruntime::contrib::cuda { + +#define REGISTER_HYPER_CONNECTION_LOW_RANK_MIX_CUDA_KERNEL(T) \ + ONNX_OPERATOR_TYPED_KERNEL_EX( \ + HyperConnectionLowRankMix, kMSDomain, 1, T, kCudaExecutionProvider, \ + (*KernelDefBuilder::Create()) \ + .TypeConstraint("T", DataTypeImpl::GetTensorType()), \ + HyperConnectionLowRankMix); + +REGISTER_HYPER_CONNECTION_LOW_RANK_MIX_CUDA_KERNEL(float) +REGISTER_HYPER_CONNECTION_LOW_RANK_MIX_CUDA_KERNEL(MLFloat16) +REGISTER_HYPER_CONNECTION_LOW_RANK_MIX_CUDA_KERNEL(BFloat16) + +#undef REGISTER_HYPER_CONNECTION_LOW_RANK_MIX_CUDA_KERNEL + +template +HyperConnectionLowRankMix::HyperConnectionLowRankMix(const OpKernelInfo& info) + : CudaKernel(info), + num_branches_(info.GetAttrOrDefault("num_branches", 0)), + epsilon_(info.GetAttrOrDefault("epsilon", 1e-5f)), + reduction_scale_(info.GetAttrOrDefault("reduction_scale", 1.0f)), + silu_alpha_(info.GetAttrOrDefault("silu_alpha", 1.0f)), + gate_div_(info.GetAttrOrDefault("gate_div", 1.0f)), + gate_mul_(info.GetAttrOrDefault("gate_mul", 1.0f)) { + ORT_ENFORCE(num_branches_ > 0, "HyperConnectionLowRankMix requires a positive num_branches"); + ORT_ENFORCE(gate_div_ != 0.0f, "gate_div must be non-zero"); +} + +template +Status HyperConnectionLowRankMix::ComputeInternal(OpKernelContext* context) const { + using CudaT = typename onnxruntime::cuda::OrtToCudaType::type; + + const Tensor* streams = context->Input(0); + const Tensor* norm_scale = context->Input(1); + const Tensor* down_weight = context->Input(2); + const Tensor* up_weight = context->Input(3); + const Tensor* inject_weight = context->Input(4); + const Tensor* prev_block_output = context->Input(5); + const Tensor* prev_post_gate = context->Input(6); + + hyper_connection::StreamShape params; + ORT_RETURN_IF_ERROR(hyper_connection::ResolveStreamShape(streams->Shape(), num_branches_, params)); + ORT_RETURN_IF_NOT(params.branches <= std::numeric_limits::max() && + params.hidden <= std::numeric_limits::max() / params.branches, + "branch or hidden dimension is too large for CUDA"); + const int64_t branches = params.branches; + const int64_t hidden = params.hidden; + const int64_t features = branches * hidden; + + ORT_RETURN_IF_NOT(norm_scale->Shape().Size() == features && norm_scale->Shape().NumDimensions() == 1, + "norm_scale must have shape (num_branches * hidden)"); + ORT_RETURN_IF_NOT(down_weight->Shape().NumDimensions() == 2 && down_weight->Shape()[0] == features, + "down_weight must have shape (num_branches * hidden, rank)"); + const int64_t rank = down_weight->Shape()[1]; + ORT_RETURN_IF_NOT(rank > 0 && rank <= std::numeric_limits::max(), "rank must be positive"); + ORT_RETURN_IF_NOT(up_weight->Shape().NumDimensions() == 2 && up_weight->Shape()[0] == rank && + up_weight->Shape()[1] == features, + "up_weight must have shape (rank, num_branches * hidden)"); + ORT_RETURN_IF_NOT(inject_weight->Shape().NumDimensions() == 2 && inject_weight->Shape()[0] == features && + inject_weight->Shape()[1] == branches, + "inject_weight must have shape (num_branches * hidden, num_branches)"); + ORT_RETURN_IF_NOT((prev_block_output == nullptr) == (prev_post_gate == nullptr), + "prev_block_output and prev_post_gate must be provided together"); + + const bool has_prologue = prev_block_output != nullptr; + if (has_prologue) { + const auto& dims = streams->Shape().GetDims(); + TensorShapeVector expected_output(dims.begin(), dims.begin() + params.prefix_rank); + expected_output.push_back(hidden); + TensorShapeVector expected_gate(dims.begin(), dims.begin() + params.prefix_rank); + expected_gate.push_back(branches); + ORT_RETURN_IF_NOT(prev_block_output->Shape() == TensorShape(expected_output), + "prev_block_output must have shape (..., hidden)"); + ORT_RETURN_IF_NOT(prev_post_gate->Shape() == TensorShape(expected_gate), + "prev_post_gate must have shape (..., num_branches)"); + } + + Tensor* pre_output = context->Output(0, TensorShape(params.reduced_shape)); + TensorShapeVector gate_shape(params.reduced_shape.begin(), params.reduced_shape.end() - 1); + gate_shape.push_back(branches); + Tensor* post_gate = context->Output(1, TensorShape(gate_shape)); + Tensor* streams_out = context->Output(2, streams->Shape()); + + if (params.rows == 0) { + return Status::OK(); + } + + HyperConnectionLowRankMixConfig config; + config.rows = params.rows; + config.branches = static_cast(branches); + config.hidden = static_cast(hidden); + config.rank = static_cast(rank); + config.epsilon = epsilon_; + config.reduction_scale = reduction_scale_; + config.silu_alpha = silu_alpha_; + config.gate_div = gate_div_; + config.gate_mul = gate_mul_; + + IAllocatorUniquePtr streams_scratch; + HyperConnectionLowRankMixTensors tensors; + tensors.streams = reinterpret_cast(streams->Data()); + tensors.norm_scale = reinterpret_cast(norm_scale->Data()); + tensors.down_weight = reinterpret_cast(down_weight->Data()); + tensors.up_weight = reinterpret_cast(up_weight->Data()); + tensors.inject_weight = reinterpret_cast(inject_weight->Data()); + if (has_prologue) { + tensors.prev_block_output = reinterpret_cast(prev_block_output->Data()); + tensors.prev_post_gate = reinterpret_cast(prev_post_gate->Data()); + // The kernels need the post-mixed streams; use scratch memory if the graph does not consume them. + if (streams_out != nullptr) { + tensors.streams_out = reinterpret_cast(streams_out->MutableData()); + } else { + streams_scratch = GetScratchBuffer(static_cast(params.rows * features), GetComputeStream(context)); + tensors.streams_out = reinterpret_cast(streams_scratch.get()); + } + } else if (streams_out != nullptr) { + // Without a prologue the streams pass through unchanged. + CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(streams_out->MutableDataRaw(), streams->DataRaw(), + streams->SizeInBytes(), cudaMemcpyDeviceToDevice, Stream(context))); + } + tensors.pre_output = reinterpret_cast(pre_output->MutableData()); + tensors.post_gate = reinterpret_cast(post_gate->MutableData()); + + // The fused kernels use 16-byte vector accesses on every tensor. + auto aligned16 = [](const void* ptr) { return ptr == nullptr || reinterpret_cast(ptr) % 16 == 0; }; + const bool aligned = aligned16(tensors.streams) && aligned16(tensors.norm_scale) && + aligned16(tensors.down_weight) && aligned16(tensors.up_weight) && + aligned16(tensors.inject_weight) && aligned16(tensors.prev_block_output) && + aligned16(tensors.streams_out); + + const auto& device_prop = GetDeviceProp(); + if (aligned && CanUseFusedHyperConnectionLowRankMix(device_prop, config)) { + auto scratch = GetScratchBuffer(HyperConnectionLowRankMixFusedScratchElements(config), + GetComputeStream(context)); + return LaunchHyperConnectionLowRankMixFused(device_prop, Stream(context), config, tensors, + scratch.get()); + } + + const size_t rows = static_cast(params.rows); + auto hn = GetScratchBuffer(rows * static_cast(features), GetComputeStream(context)); + auto down = GetScratchBuffer(rows * static_cast(rank), GetComputeStream(context)); + auto up = GetScratchBuffer(rows * static_cast(features), GetComputeStream(context)); + auto inject = GetScratchBuffer(rows * static_cast(branches), GetComputeStream(context)); + return LaunchHyperConnectionLowRankMixUnfused( + GetCublasHandle(context), Stream(context), config, tensors, reinterpret_cast(hn.get()), + reinterpret_cast(down.get()), reinterpret_cast(up.get()), + reinterpret_cast(inject.get())); +} + +template class HyperConnectionLowRankMix; +template class HyperConnectionLowRankMix; +template class HyperConnectionLowRankMix; + +} // namespace onnxruntime::contrib::cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix.h b/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix.h new file mode 100644 index 0000000000000..246f44afb4f8b --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix.h @@ -0,0 +1,25 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/providers/cuda/cuda_kernel.h" + +namespace onnxruntime::contrib::cuda { + +template +class HyperConnectionLowRankMix final : public onnxruntime::cuda::CudaKernel { + public: + explicit HyperConnectionLowRankMix(const OpKernelInfo& info); + Status ComputeInternal(OpKernelContext* context) const override; + + private: + int64_t num_branches_; + float epsilon_; + float reduction_scale_; + float silu_alpha_; + float gate_div_; + float gate_mul_; +}; + +} // namespace onnxruntime::contrib::cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.cu b/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.cu new file mode 100644 index 0000000000000..25f9ffccc4194 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.cu @@ -0,0 +1,888 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +// Fused low-rank hyper-connection block (see the HyperConnectionLowRankMix schema in bert_defs.cc). +// +// Fused path (rows <= kHyperConnectionLowRankMixMaxFusedRows, 16-bit T, 16-byte aligned shapes): two kernels. +// HcDownKernel grid (column tiles + 1, splits, row groups) +// [post-mix of the previous block] + branchwise RMS norm + down GEMV + inject GEMV. +// The K = C * H reduction is split in slices that each stay inside one branch; every CTA +// writes fp32 partial sums for its slice. No atomics, no clusters, no grid barrier. +// HcUpPreKernel grid (ceil(H / 16), 1, row groups) +// reduces the partials in a fixed order, then ScaledSiLU, up GEMV, sigmoid, pre-mix and the +// inject gate. +// Both kernels are launched with programmatic dependent launch (sm_90+, plain launches elsewhere). Each one +// triggers its dependents immediately and prefetches its weights (which never depend on a previous kernel) +// into shared memory before griddepcontrol.wait; everything that consumes another kernel's output comes after +// the wait, so correctness never depends on the trigger position. +// Every intermediate is rounded to T at the points where the unfused ONNX subgraph rounds; the GEMVs +// accumulate in fp32 in a fixed (deterministic) order. +// +// Unfused path (any shape, any T): norm kernel + cuBLAS GEMMs (fp32 accumulation) + elementwise kernels. + +#include "contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.h" + +#include +#include +#include +#include +#include +#include + +#include "core/providers/cuda/cu_inc/cuda_type_helper.cuh" +#include "core/providers/cuda/cuda_common.h" + +namespace onnxruntime::contrib::cuda { +namespace { + +constexpr int kThreads = 256; +constexpr int kColTile = 16; // down columns per CTA of the down kernel (two 16-byte weight vectors) +constexpr int kUpTile = 16; // features per CTA of the up kernel +constexpr int kMaxSliceRows = 1280; // K rows handled by one CTA of the down kernel +constexpr int kWeightSlots = 10; // 16-byte weight vectors per thread held in shared memory +constexpr int kMaxVecIters = 2; // 16-byte vectors per thread in the norm pass +constexpr int kMaxFusedHidden = kMaxVecIters * kThreads * 8; +constexpr int kMaxFusedRank = 512; +constexpr int kMaxFusedBranches = 8; + +__device__ __forceinline__ float SigmoidF32(float v) { + if (v >= 0.0f) return 1.0f / (1.0f + expf(-v)); + const float e = expf(v); + return e / (1.0f + e); +} + +template +__device__ __forceinline__ float Rnd(float v) { + return to_float(from_float(v)); +} + +template +__device__ __forceinline__ float BitsToFloat(uint16_t bits) { + if constexpr (std::is_same_v) { + return __uint_as_float(static_cast(bits) << 16); + } else { + return __half2float(__ushort_as_half(bits)); + } +} + +template +__device__ __forceinline__ uint16_t FloatToBits(float v) { + if constexpr (std::is_same_v) { + return __bfloat16_as_ushort(__float2bfloat16(v)); + } else { + return __half_as_ushort(__float2half(v)); + } +} + +template +__device__ __forceinline__ void Unpack2(uint32_t u, float& a, float& b) { + a = BitsToFloat(static_cast(u & 0xffffu)); + b = BitsToFloat(static_cast(u >> 16)); +} + +template +__device__ __forceinline__ uint32_t Pack2(float a, float b) { + return static_cast(FloatToBits(a)) | (static_cast(FloatToBits(b)) << 16); +} + +template +__device__ __forceinline__ void Unpack8(const uint4& u, float (&f)[8]) { + Unpack2(u.x, f[0], f[1]); + Unpack2(u.y, f[2], f[3]); + Unpack2(u.z, f[4], f[5]); + Unpack2(u.w, f[6], f[7]); +} + +template +__device__ __forceinline__ uint4 Pack8(const float (&f)[8]) { + uint4 u; + u.x = Pack2(f[0], f[1]); + u.y = Pack2(f[2], f[3]); + u.z = Pack2(f[4], f[5]); + u.w = Pack2(f[6], f[7]); + return u; +} + +__device__ __forceinline__ void GridWait() { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif +} + +__device__ __forceinline__ void GridTrigger() { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); +#endif +} + +__device__ __forceinline__ void CpAsync16(void* smem, const void* gmem) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;" ::"r"(static_cast(__cvta_generic_to_shared(smem))), + "l"(gmem) + : "memory"); +#else + *reinterpret_cast(smem) = *reinterpret_cast(gmem); +#endif +} + +__device__ __forceinline__ void CpCommit() { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + asm volatile("cp.async.commit_group;" ::: "memory"); +#endif +} + +__device__ __forceinline__ void CpWaitAll() { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + asm volatile("cp.async.wait_group 0;" ::: "memory"); +#endif +} + +struct DownParams { + int64_t rows; + int branches; + int hidden; + int rank; + int pieces; // slices per branch + int slice; // K rows per slice (multiple of 8) + int splits; // branches * pieces + int partial_stride; // rank + branches, rounded up to a multiple of 4 + float epsilon; +}; + +struct UpParams { + int64_t rows; + int branches; + int hidden; + int rank; + int splits; + int partial_stride; + float silu_alpha; + float reduction_scale; + float gate_div; + float gate_mul; +}; + +// CTA (j, s, z): K slice s (rows [k0, k0 + len) of the flattened stream, inside branch c) of the ROWS rows +// z * ROWS ... z * ROWS + ROWS - 1, which share every weight read. +// j < ceil(R / 16) owns down columns [16 j, 16 j + 16); the last j owns the C inject columns. +template +__global__ void __launch_bounds__(kThreads) + HcDownKernel(const T* __restrict__ streams, const T* __restrict__ prev_out, const T* __restrict__ prev_gate, + const T* __restrict__ norm_w, const T* __restrict__ w_down, const T* __restrict__ w_inject, + T* __restrict__ streams_out, float* __restrict__ partial, float* __restrict__ inv_rms_out, + DownParams p) { + const int j = blockIdx.x; + const int s = blockIdx.y; + const int64_t row0 = static_cast(blockIdx.z) * ROWS; + const int tid = threadIdx.x; + const int branches = p.branches; + const int hidden = p.hidden; + const int64_t K = static_cast(branches) * hidden; + const int c = s / p.pieces; + const int piece = s - c * p.pieces; + const int h0 = piece * p.slice; + const int len = min(p.slice, hidden - h0); + const int col_tiles = (p.rank + kColTile - 1) / kColTile; + const bool is_inject = (j == col_tiles); + const int64_t k0 = static_cast(c) * hidden + h0; + int64_t row_idx[ROWS]; + bool row_ok[ROWS]; +#pragma unroll + for (int r = 0; r < ROWS; ++r) { + row_ok[r] = row0 + r < p.rows; + row_idx[r] = row_ok[r] ? row0 + r : p.rows - 1; // keep loads of unused rows in bounds + } + + __shared__ __align__(16) uint4 wsm[kWeightSlots * kThreads]; + __shared__ __align__(16) uint16_t xs[ROWS][kMaxSliceRows]; + __shared__ float red[8][ROWS][kColTile]; + __shared__ float sq_red[8][ROWS]; + + const int cg = tid & 1; + const int rp = tid >> 1; + const int col = j * kColTile + cg * 8; + // Let the up kernel start right away and prefetch its weights; its results are only read after GridWait(). + GridTrigger(); + + // Weights do not depend on any previous kernel: start streaming them into shared memory first. + if (!is_inject) { + if (col < p.rank) { + const T* base = w_down + k0 * p.rank + col; +#pragma unroll + for (int i = 0; i < kWeightSlots; ++i) { + const int row = rp + 128 * i; + if (row < len) CpAsync16(&wsm[i * kThreads + tid], base + static_cast(row) * p.rank); + } + } + } else { + const int chunks = len * branches / 8; + for (int q = tid; q < chunks; q += kThreads) CpAsync16(&wsm[q], w_inject + k0 * branches + q * 8); + } + CpCommit(); + + // Norm weights for the 8 consecutive features of the slice this thread normalizes below. + const bool norm_active = tid * 8 < len; + uint4 wv = make_uint4(0u, 0u, 0u, 0u); + if (norm_active) wv = *reinterpret_cast(norm_w + k0 + tid * 8); + + // Everything below consumes results of the previous kernel. + GridWait(); + + const int nvec = hidden / 8; + uint4 sv[ROWS][kMaxVecIters]; + uint4 bv[ROWS][kMaxVecIters]; + float gate_c[ROWS]; +#pragma unroll + for (int r = 0; r < ROWS; ++r) { + const uint4* sp = reinterpret_cast(streams + row_idx[r] * K + static_cast(c) * hidden); +#pragma unroll + for (int i = 0; i < kMaxVecIters; ++i) { + const int qi = tid + kThreads * i; + sv[r][i] = qi < nvec ? sp[qi] : make_uint4(0u, 0u, 0u, 0u); + } + if constexpr (kPost) { + gate_c[r] = to_float(prev_gate[row_idx[r] * branches + c]); + const uint4* bp = reinterpret_cast(prev_out + row_idx[r] * hidden); +#pragma unroll + for (int i = 0; i < kMaxVecIters; ++i) { + const int qi = tid + kThreads * i; + bv[r][i] = qi < nvec ? bp[qi] : make_uint4(0u, 0u, 0u, 0u); + } + } + } + + float sum_sq[ROWS]; +#pragma unroll + for (int r = 0; r < ROWS; ++r) { + sum_sq[r] = 0.0f; +#pragma unroll + for (int i = 0; i < kMaxVecIters; ++i) { + const int qi = tid + kThreads * i; + if (qi < nvec) { + const int h = qi * 8; + float v[8]; + Unpack8(sv[r][i], v); + if constexpr (kPost) { + float b[8]; + Unpack8(bv[r][i], b); +#pragma unroll + for (int m = 0; m < 8; ++m) v[m] = Rnd(fmaf(gate_c[r], b[m], v[m])); + } +#pragma unroll + for (int m = 0; m < 8; ++m) sum_sq[r] += v[m] * v[m]; + if (h >= h0 && h < h0 + len) { + const uint4 packed = Pack8(v); + *reinterpret_cast(&xs[r][h - h0]) = packed; + if constexpr (kPost) { + if (j == 0 && row_ok[r]) { + *reinterpret_cast(streams_out + row_idx[r] * K + static_cast(c) * hidden + h) = + packed; + } + } + } + } + } +#pragma unroll + for (int off = 16; off > 0; off >>= 1) sum_sq[r] += __shfl_xor_sync(0xffffffffu, sum_sq[r], off); + if ((tid & 31) == 0) sq_red[tid >> 5][r] = sum_sq[r]; + } + __syncthreads(); + float inv_rms[ROWS]; +#pragma unroll + for (int r = 0; r < ROWS; ++r) { + float total = 0.0f; +#pragma unroll + for (int w = 0; w < 8; ++w) total += sq_red[w][r]; + inv_rms[r] = rsqrtf(total / static_cast(hidden) + p.epsilon); + if (j == 0 && piece == 0 && tid == 0 && row_ok[r]) inv_rms_out[row_idx[r] * branches + c] = inv_rms[r]; + if (norm_active) { + float x[8], w[8]; + Unpack8(*reinterpret_cast(&xs[r][tid * 8]), x); + Unpack8(wv, w); +#pragma unroll + for (int m = 0; m < 8; ++m) x[m] = x[m] * inv_rms[r] * w[m]; + *reinterpret_cast(&xs[r][tid * 8]) = Pack8(x); + } + } + CpWaitAll(); + __syncthreads(); + + if (!is_inject) { + float acc[ROWS][8]; +#pragma unroll + for (int r = 0; r < ROWS; ++r) { +#pragma unroll + for (int m = 0; m < 8; ++m) acc[r][m] = 0.0f; + } + if (col < p.rank) { +#pragma unroll + for (int i = 0; i < kWeightSlots; ++i) { + const int row = rp + 128 * i; + if (row < len) { + float wf[8]; + Unpack8(wsm[i * kThreads + tid], wf); +#pragma unroll + for (int r = 0; r < ROWS; ++r) { + const float xv = BitsToFloat(xs[r][row]); +#pragma unroll + for (int m = 0; m < 8; ++m) acc[r][m] = fmaf(xv, wf[m], acc[r][m]); + } + } + } + } + // Lanes of a warp share cg = lane & 1 and differ in rp: reduce over lane bits 1..4. +#pragma unroll + for (int r = 0; r < ROWS; ++r) { +#pragma unroll + for (int m = 0; m < 8; ++m) { +#pragma unroll + for (int off = 2; off < 32; off <<= 1) acc[r][m] += __shfl_xor_sync(0xffffffffu, acc[r][m], off); + } + if ((tid & 31) < 2) { +#pragma unroll + for (int m = 0; m < 8; ++m) red[tid >> 5][r][cg * 8 + m] = acc[r][m]; + } + } + __syncthreads(); + if (tid < kColTile * ROWS) { + const int r = tid / kColTile; + const int t = tid % kColTile; + float v = 0.0f; +#pragma unroll + for (int q = 0; q < 8; ++q) v += red[q][r][t]; + const int out_col = j * kColTile + t; + if (out_col < p.rank && row_ok[r]) partial[(row_idx[r] * p.splits + s) * p.partial_stride + out_col] = v; + } + } else { + const T* wi = reinterpret_cast(wsm); + float a[ROWS][kMaxFusedBranches]; +#pragma unroll + for (int r = 0; r < ROWS; ++r) { +#pragma unroll + for (int m = 0; m < kMaxFusedBranches; ++m) a[r][m] = 0.0f; + } +#pragma unroll + for (int i = 0; i < 5; ++i) { + const int k = tid + kThreads * i; + if (k < len) { + float wf[kMaxFusedBranches]; +#pragma unroll + for (int m = 0; m < kMaxFusedBranches; ++m) wf[m] = m < branches ? to_float(wi[k * branches + m]) : 0.0f; +#pragma unroll + for (int r = 0; r < ROWS; ++r) { + const float xv = BitsToFloat(xs[r][k]); +#pragma unroll + for (int m = 0; m < kMaxFusedBranches; ++m) a[r][m] = fmaf(xv, wf[m], a[r][m]); + } + } + } +#pragma unroll + for (int r = 0; r < ROWS; ++r) { +#pragma unroll + for (int m = 0; m < kMaxFusedBranches; ++m) { +#pragma unroll + for (int off = 16; off > 0; off >>= 1) a[r][m] += __shfl_xor_sync(0xffffffffu, a[r][m], off); + if ((tid & 31) == 0) red[tid >> 5][r][m] = a[r][m]; + } + } + __syncthreads(); + if (tid < branches * ROWS) { + const int r = tid / branches; + const int t = tid % branches; + float v = 0.0f; +#pragma unroll + for (int q = 0; q < 8; ++q) v += red[q][r][t]; + if (row_ok[r]) partial[(row_idx[r] * p.splits + s) * p.partial_stride + p.rank + t] = v; + } + } +} + +// CTA (hb, 0, z) owns features [16 hb, 16 hb + 16) of every branch for ROWS rows. +template +__global__ void __launch_bounds__(kThreads) + HcUpPreKernel(const T* __restrict__ streams, const T* __restrict__ norm_w, const T* __restrict__ w_up, + const float* __restrict__ partial, const float* __restrict__ inv_rms_in, + T* __restrict__ pre_out, T* __restrict__ gate_out, UpParams p) { + const int hb = blockIdx.x; + const int64_t row0 = static_cast(blockIdx.z) * ROWS; + const int tid = threadIdx.x; + const int branches = BRANCHES == 0 ? p.branches : BRANCHES; + const int rank = RANK == 0 ? p.rank : RANK; + constexpr int kBranchStorage = BRANCHES == 0 ? kMaxFusedBranches : BRANCHES; + const int hidden = p.hidden; + const int64_t K = static_cast(branches) * hidden; + const int h0 = hb * kUpTile; + const int groups = 2 * branches; // (branch, 8-feature half) pairs; divides 32 + const int phases = kThreads / groups; // row phases + const int g = tid % groups; + const int c = g >> 1; + const int half_id = g & 1; + const int rp = tid / groups; + const int col = h0 + half_id * 8; + const bool valid = col < hidden; + const int chunk_rows = phases * kWeightSlots; + int64_t row_idx[ROWS]; + bool row_ok[ROWS]; +#pragma unroll + for (int r = 0; r < ROWS; ++r) { + row_ok[r] = row0 + r < p.rows; + row_idx[r] = row_ok[r] ? row0 + r : p.rows - 1; + } + + __shared__ __align__(16) uint4 wsm[kWeightSlots * kThreads]; + __shared__ __align__(16) float sf[ROWS][RANK == 0 ? kMaxFusedRank : RANK]; + __shared__ float sg[ROWS][kBranchStorage * kUpTile]; + + // Prefetch the first chunk of up weights before waiting for the down kernel. + const T* wbase = w_up + static_cast(c) * hidden + col; + if (valid) { +#pragma unroll + for (int i = 0; i < kWeightSlots; ++i) { + const int row = rp + phases * i; + if (row < rank) CpAsync16(&wsm[i * kThreads + tid], wbase + static_cast(row) * K); + } + } + CpCommit(); + + // The next block's down kernel may start (and prefetch its weights) while this kernel waits. + GridTrigger(); + GridWait(); + + float xv[ROWS][kBranchStorage], inv[ROWS][kBranchStorage], nw[kBranchStorage]; + const int hl = tid; // feature owned by threads < 16 in the final pre-mix + const bool owner = tid < kUpTile && (h0 + hl) < hidden; + if (owner) { +#pragma unroll + for (int b = 0; b < kBranchStorage; ++b) { + if (b < branches) { + nw[b] = to_float(norm_w[static_cast(b) * hidden + h0 + hl]); +#pragma unroll + for (int r = 0; r < ROWS; ++r) { + xv[r][b] = to_float(streams[row_idx[r] * K + static_cast(b) * hidden + h0 + hl]); + inv[r][b] = inv_rms_in[row_idx[r] * branches + b]; + } + } + } + } + + // Fixed-order reduction of the K-split partial sums, then ScaledSiLU; four columns per thread. + const int nq = rank / 4; + const int cq = tid * 4; + float4 va[ROWS]; +#pragma unroll + for (int r = 0; r < ROWS; ++r) va[r] = make_float4(0.f, 0.f, 0.f, 0.f); + if (tid < nq) { + for (int sp0 = 0; sp0 < p.splits; sp0 += 8) { + float4 pq[8][ROWS]; +#pragma unroll + for (int u = 0; u < 8; ++u) { +#pragma unroll + for (int r = 0; r < ROWS; ++r) { + pq[u][r] = sp0 + u < p.splits + ? __ldca(reinterpret_cast( + partial + (row_idx[r] * p.splits + sp0 + u) * p.partial_stride + cq)) + : make_float4(0.f, 0.f, 0.f, 0.f); + } + } +#pragma unroll + for (int u = 0; u < 8; ++u) { +#pragma unroll + for (int r = 0; r < ROWS; ++r) { + va[r].x += pq[u][r].x; + va[r].y += pq[u][r].y; + va[r].z += pq[u][r].z; + va[r].w += pq[u][r].w; + } + } + } +#pragma unroll + for (int r = 0; r < ROWS; ++r) { + const float v[4] = {va[r].x, va[r].y, va[r].z, va[r].w}; + float out[4]; +#pragma unroll + for (int m = 0; m < 4; ++m) { + const float z = Rnd(Rnd(v[m]) * p.silu_alpha); + out[m] = Rnd(z * Rnd(SigmoidF32(z))); + } + *reinterpret_cast(&sf[r][cq]) = make_float4(out[0], out[1], out[2], out[3]); + } + } + if (hb == 0 && tid < branches * ROWS) { + const int r = tid / branches; + const int t = tid % branches; + float v = 0.0f; + for (int sp = 0; sp < p.splits; ++sp) v += partial[(row_idx[r] * p.splits + sp) * p.partial_stride + rank + t]; + const float inject = Rnd(v); + const float u = Rnd(SigmoidF32(Rnd(inject / p.gate_div))); + if (row_ok[r]) gate_out[row_idx[r] * branches + t] = from_float(u * p.gate_mul); + } + __syncthreads(); + + float acc[ROWS][8]; +#pragma unroll + for (int r = 0; r < ROWS; ++r) { +#pragma unroll + for (int m = 0; m < 8; ++m) acc[r][m] = 0.0f; + } + for (int chunk0 = 0; chunk0 < rank; chunk0 += chunk_rows) { + if (chunk0 > 0 && valid) { +#pragma unroll + for (int i = 0; i < kWeightSlots; ++i) { + const int row = chunk0 + rp + phases * i; + if (row < rank) CpAsync16(&wsm[i * kThreads + tid], wbase + static_cast(row) * K); + } + CpCommit(); + } + CpWaitAll(); + if (valid) { +#pragma unroll + for (int i = 0; i < kWeightSlots; ++i) { + const int row = chunk0 + rp + phases * i; + if (row < rank) { + float wf[8]; + Unpack8(wsm[i * kThreads + tid], wf); +#pragma unroll + for (int r = 0; r < ROWS; ++r) { + const float sv = sf[r][row]; +#pragma unroll + for (int m = 0; m < 8; ++m) acc[r][m] = fmaf(sv, wf[m], acc[r][m]); + } + } + } + } + } + // Lanes of a warp share g = lane % groups and differ in rp. + for (int off = groups; off < 32; off <<= 1) { +#pragma unroll + for (int r = 0; r < ROWS; ++r) { +#pragma unroll + for (int m = 0; m < 8; ++m) acc[r][m] += __shfl_xor_sync(0xffffffffu, acc[r][m], off); + } + } + __syncthreads(); // every thread is done with its weights: reuse that shared memory for the cross-warp sums + float* red = reinterpret_cast(wsm); // [8 warps][ROWS][groups * 8] + if ((tid & 31) < groups) { +#pragma unroll + for (int r = 0; r < ROWS; ++r) { +#pragma unroll + for (int m = 0; m < 8; ++m) red[((tid >> 5) * ROWS + r) * (groups * 8) + g * 8 + m] = acc[r][m]; + } + } + __syncthreads(); + if (tid < groups * 8) { +#pragma unroll + for (int r = 0; r < ROWS; ++r) { + float v = 0.0f; +#pragma unroll + for (int q = 0; q < 8; ++q) v += red[(q * ROWS + r) * (groups * 8) + tid]; + const int gg = tid >> 3; + const int local = (gg & 1) * 8 + (tid & 7); + sg[r][(gg >> 1) * kUpTile + local] = Rnd(SigmoidF32(Rnd(v))); + } + } + __syncthreads(); + if (owner) { +#pragma unroll + for (int r = 0; r < ROWS; ++r) { + float sum = 0.0f; +#pragma unroll + for (int b = 0; b < kBranchStorage; ++b) { + if (b < branches) { + const float n = Rnd(xv[r][b] * inv[r][b] * nw[b]); + sum = fmaf(n, sg[r][b * kUpTile + hl], sum); + } + } + if (row_ok[r]) pre_out[row_idx[r] * hidden + h0 + hl] = from_float(sum * p.reduction_scale); + } + } +} + +// ----------------------------------------------------------------------------- unfused path kernels + +template +__global__ void NormKernel(const T* __restrict__ streams, const T* __restrict__ prev_out, + const T* __restrict__ prev_gate, const T* __restrict__ norm_w, T* __restrict__ hn, + T* __restrict__ streams_out, int branches, int hidden, float epsilon) { + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage storage; + __shared__ float inv_rms; + const int64_t group = blockIdx.x; + const int64_t row = group / branches; + const int c = static_cast(group - row * branches); + const int64_t base = group * hidden; + float gate_c = 0.0f; + if constexpr (kPost) gate_c = to_float(prev_gate[group]); + + float sum_sq = 0.0f; + for (int h = threadIdx.x; h < hidden; h += kThreads) { + float v = to_float(streams[base + h]); + if constexpr (kPost) { + v = Rnd(fmaf(gate_c, to_float(prev_out[row * hidden + h]), v)); + streams_out[base + h] = from_float(v); + } + sum_sq += v * v; + } + const float total = BlockReduce(storage).Sum(sum_sq); + if (threadIdx.x == 0) inv_rms = rsqrtf(total / static_cast(hidden) + epsilon); + __syncthreads(); + for (int h = threadIdx.x; h < hidden; h += kThreads) { + float v = to_float(streams[base + h]); + if constexpr (kPost) v = Rnd(fmaf(gate_c, to_float(prev_out[row * hidden + h]), v)); + hn[base + h] = from_float(v * inv_rms * to_float(norm_w[static_cast(c) * hidden + h])); + } +} + +template +__global__ void ScaledSiluKernel(T* __restrict__ x, int64_t count, float alpha) { + const int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= count) return; + const float z = Rnd(to_float(x[i]) * alpha); + const float sig = Rnd(SigmoidF32(z)); + x[i] = from_float(z * sig); +} + +template +__global__ void PreMixKernel(const T* __restrict__ hn, const T* __restrict__ up, T* __restrict__ pre, + int64_t count, int branches, int hidden, float reduction_scale) { + const int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= count) return; + const int64_t row = i / hidden; + const int h = static_cast(i - row * hidden); + const int64_t K = static_cast(branches) * hidden; + float sum = 0.0f; + for (int b = 0; b < branches; ++b) { + const int64_t idx = row * K + static_cast(b) * hidden + h; + const float gate = Rnd(SigmoidF32(to_float(up[idx]))); + sum = fmaf(to_float(hn[idx]), gate, sum); + } + pre[i] = from_float(sum * reduction_scale); +} + +template +__global__ void PostGateKernel(const T* __restrict__ inject, T* __restrict__ gate, int64_t count, float gate_div, + float gate_mul) { + const int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= count) return; + const float t = Rnd(to_float(inject[i]) / gate_div); + const float u = Rnd(SigmoidF32(t)); + gate[i] = from_float(u * gate_mul); +} + +template +constexpr cudaDataType_t CudaDataType() { + if constexpr (std::is_same_v) { + return CUDA_R_32F; + } else if constexpr (std::is_same_v) { + return CUDA_R_16F; + } else { + return CUDA_R_16BF; + } +} + +// Row-major C[m, n] = A[m, k] * B[k, n] with fp32 accumulation and T output. +template +Status Gemm(cublasHandle_t cublas, int64_t m, int64_t n, int64_t k, const T* a, const T* b, T* c) { + ORT_RETURN_IF_NOT(m <= std::numeric_limits::max() && n <= std::numeric_limits::max() && + k <= std::numeric_limits::max(), + "HyperConnectionLowRankMix GEMM dimension is too large"); + const float alpha = 1.0f; + const float beta = 0.0f; + constexpr cudaDataType_t type = CudaDataType(); + CUBLAS_RETURN_IF_ERROR(cublasGemmEx(cublas, CUBLAS_OP_N, CUBLAS_OP_N, static_cast(n), static_cast(m), + static_cast(k), &alpha, b, type, static_cast(n), a, type, + static_cast(k), &beta, c, type, static_cast(n), + CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT)); + return Status::OK(); +} + +int64_t BlocksFor(int64_t count) { return (count + kThreads - 1) / kThreads; } + +struct FusedGeometry { + int pieces; + int slice; + int splits; + int col_tiles; + int partial_stride; +}; + +FusedGeometry GetFusedGeometry(const HyperConnectionLowRankMixConfig& config) { + FusedGeometry g; + g.pieces = (config.hidden + kMaxSliceRows - 1) / kMaxSliceRows; + g.slice = ((config.hidden + g.pieces - 1) / g.pieces + 7) / 8 * 8; + g.pieces = (config.hidden + g.slice - 1) / g.slice; + g.splits = config.branches * g.pieces; + g.col_tiles = (config.rank + kColTile - 1) / kColTile; + g.partial_stride = (config.rank + config.branches + 3) / 4 * 4; // keeps every row of partials 16-byte aligned + return g; +} + +template +Status LaunchWithProgrammaticSerialization(bool use_pdl, Kernel kernel, dim3 grid, dim3 block, cudaStream_t stream, + Args... args) { +#if defined(CUDART_VERSION) && CUDART_VERSION >= 12030 + if (use_pdl) { + cudaLaunchAttribute attribute; + attribute.id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute.val.programmaticStreamSerializationAllowed = 1; + cudaLaunchConfig_t launch_config = {}; + launch_config.gridDim = grid; + launch_config.blockDim = block; + launch_config.dynamicSmemBytes = 0; + launch_config.stream = stream; + launch_config.attrs = &attribute; + launch_config.numAttrs = 1; + return CUDA_CALL(cudaLaunchKernelEx(&launch_config, kernel, args...)); + } +#else + ORT_UNUSED_PARAMETER(use_pdl); +#endif + kernel<<>>(args...); + return CUDA_CALL(cudaGetLastError()); +} + +template +Status LaunchFusedRows(bool use_pdl, cudaStream_t stream, const HyperConnectionLowRankMixConfig& config, + const HyperConnectionLowRankMixTensors& tensors, const FusedGeometry& g, float* partial, + float* inv_rms, const DownParams& down, const UpParams& up) { + const unsigned row_groups = static_cast((config.rows + ROWS - 1) / ROWS); + const dim3 down_grid(g.col_tiles + 1, g.splits, row_groups); + const dim3 up_grid((config.hidden + kUpTile - 1) / kUpTile, 1, row_groups); + if (tensors.prev_block_output != nullptr) { + ORT_RETURN_IF_ERROR(LaunchWithProgrammaticSerialization( + use_pdl, HcDownKernel, down_grid, dim3(kThreads), stream, tensors.streams, + tensors.prev_block_output, tensors.prev_post_gate, tensors.norm_scale, tensors.down_weight, + tensors.inject_weight, tensors.streams_out, partial, inv_rms, down)); + } else { + ORT_RETURN_IF_ERROR(LaunchWithProgrammaticSerialization( + use_pdl, HcDownKernel, down_grid, dim3(kThreads), stream, tensors.streams, + static_cast(nullptr), static_cast(nullptr), tensors.norm_scale, + tensors.down_weight, tensors.inject_weight, static_cast(nullptr), partial, inv_rms, down)); + } + + const T* up_streams = tensors.prev_block_output != nullptr ? tensors.streams_out : tensors.streams; + if (config.branches == 4 && config.rank == 320) { + return LaunchWithProgrammaticSerialization(use_pdl, HcUpPreKernel, up_grid, dim3(kThreads), stream, + up_streams, tensors.norm_scale, tensors.up_weight, + static_cast(partial), + static_cast(inv_rms), tensors.pre_output, + tensors.post_gate, up); + } + return LaunchWithProgrammaticSerialization(use_pdl, HcUpPreKernel, up_grid, dim3(kThreads), stream, + up_streams, tensors.norm_scale, tensors.up_weight, + static_cast(partial), + static_cast(inv_rms), tensors.pre_output, + tensors.post_gate, up); +} + +} // namespace + +template +bool CanUseFusedHyperConnectionLowRankMix(const cudaDeviceProp& device_prop, + const HyperConnectionLowRankMixConfig& config) { + if constexpr (sizeof(T) != 2) { + return false; + } else { + const bool branches_ok = config.branches == 1 || config.branches == 2 || config.branches == 4 || + config.branches == 8; + return device_prop.major >= 8 && config.rows >= 1 && + config.rows <= kHyperConnectionLowRankMixMaxFusedRows && branches_ok && + config.hidden % 8 == 0 && config.hidden <= kMaxFusedHidden && + config.rank % 8 == 0 && config.rank >= 8 && config.rank <= kMaxFusedRank; + } +} + +size_t HyperConnectionLowRankMixFusedScratchElements(const HyperConnectionLowRankMixConfig& config) { + const FusedGeometry g = GetFusedGeometry(config); + return static_cast(config.rows) * + (static_cast(g.splits) * g.partial_stride + static_cast(config.branches)); +} + +template +Status LaunchHyperConnectionLowRankMixFused(const cudaDeviceProp& device_prop, cudaStream_t stream, + const HyperConnectionLowRankMixConfig& config, + const HyperConnectionLowRankMixTensors& tensors, + float* scratch) { + if constexpr (sizeof(T) != 2) { + ORT_UNUSED_PARAMETER(device_prop); + ORT_UNUSED_PARAMETER(stream); + ORT_UNUSED_PARAMETER(config); + ORT_UNUSED_PARAMETER(tensors); + ORT_UNUSED_PARAMETER(scratch); + return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, "fused HyperConnectionLowRankMix requires a 16-bit type"); + } else { + ORT_RETURN_IF_NOT(CanUseFusedHyperConnectionLowRankMix(device_prop, config), + "unsupported problem for the fused HyperConnectionLowRankMix path"); + const FusedGeometry g = GetFusedGeometry(config); + float* partial = scratch; + float* inv_rms = scratch + static_cast(config.rows) * g.splits * g.partial_stride; + + DownParams down{config.rows, config.branches, config.hidden, config.rank, g.pieces, g.slice, + g.splits, g.partial_stride, config.epsilon}; + UpParams up{config.rows, config.branches, config.hidden, config.rank, g.splits, g.partial_stride, + config.silu_alpha, config.reduction_scale, config.gate_div, config.gate_mul}; + const bool use_pdl = device_prop.major >= 9; + if (config.rows == 1) { + return LaunchFusedRows(use_pdl, stream, config, tensors, g, partial, inv_rms, down, up); + } + return LaunchFusedRows(use_pdl, stream, config, tensors, g, partial, inv_rms, down, up); + } +} + +template +Status LaunchHyperConnectionLowRankMixUnfused(cublasHandle_t cublas, cudaStream_t stream, + const HyperConnectionLowRankMixConfig& config, + const HyperConnectionLowRankMixTensors& tensors, + T* hn, T* down, T* up, T* inject) { + const int64_t rows = config.rows; + const int64_t K = static_cast(config.branches) * config.hidden; + const int64_t groups = rows * config.branches; + ORT_RETURN_IF_NOT(groups <= std::numeric_limits::max(), "CUDA launch requires too many blocks"); + + if (tensors.prev_block_output != nullptr) { + NormKernel<<(groups), kThreads, 0, stream>>>( + tensors.streams, tensors.prev_block_output, tensors.prev_post_gate, tensors.norm_scale, hn, + tensors.streams_out, config.branches, config.hidden, config.epsilon); + } else { + NormKernel<<(groups), kThreads, 0, stream>>>( + tensors.streams, nullptr, nullptr, tensors.norm_scale, hn, nullptr, config.branches, config.hidden, + config.epsilon); + } + ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetLastError())); + + ORT_RETURN_IF_ERROR(Gemm(cublas, rows, config.rank, K, hn, tensors.down_weight, down)); + ORT_RETURN_IF_ERROR(Gemm(cublas, rows, config.branches, K, hn, tensors.inject_weight, inject)); + + const int64_t down_count = rows * config.rank; + ScaledSiluKernel<<(BlocksFor(down_count)), kThreads, 0, stream>>>( + down, down_count, config.silu_alpha); + ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetLastError())); + ORT_RETURN_IF_ERROR(Gemm(cublas, rows, K, config.rank, down, tensors.up_weight, up)); + + const int64_t pre_count = rows * config.hidden; + ORT_RETURN_IF_NOT(BlocksFor(pre_count) <= std::numeric_limits::max(), + "CUDA launch requires too many blocks"); + PreMixKernel<<(BlocksFor(pre_count)), kThreads, 0, stream>>>( + hn, up, tensors.pre_output, pre_count, config.branches, config.hidden, config.reduction_scale); + ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetLastError())); + PostGateKernel<<(BlocksFor(groups)), kThreads, 0, stream>>>( + inject, tensors.post_gate, groups, config.gate_div, config.gate_mul); + return CUDA_CALL(cudaGetLastError()); +} + +#define INSTANTIATE(T) \ + template bool CanUseFusedHyperConnectionLowRankMix(const cudaDeviceProp&, \ + const HyperConnectionLowRankMixConfig&); \ + template Status LaunchHyperConnectionLowRankMixFused( \ + const cudaDeviceProp&, cudaStream_t, const HyperConnectionLowRankMixConfig&, \ + const HyperConnectionLowRankMixTensors&, float*); \ + template Status LaunchHyperConnectionLowRankMixUnfused( \ + cublasHandle_t, cudaStream_t, const HyperConnectionLowRankMixConfig&, \ + const HyperConnectionLowRankMixTensors&, T*, T*, T*, T*); + +INSTANTIATE(float) +INSTANTIATE(half) +INSTANTIATE(__nv_bfloat16) + +#undef INSTANTIATE + +} // namespace onnxruntime::contrib::cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.h b/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.h new file mode 100644 index 0000000000000..c3c9c24c4842b --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.h @@ -0,0 +1,67 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include + +#include +#include + +#include "core/common/status.h" + +namespace onnxruntime::contrib::cuda { + +struct HyperConnectionLowRankMixConfig { + int64_t rows = 0; + int branches = 0; // C + int hidden = 0; // H, width of one branch; the streams have C * H features + int rank = 0; // R, width of the low-rank bottleneck + float epsilon = 1e-5f; + float reduction_scale = 1.0f; + float silu_alpha = 1.0f; + float gate_div = 1.0f; + float gate_mul = 1.0f; +}; + +template +struct HyperConnectionLowRankMixTensors { + const T* streams = nullptr; + const T* norm_scale = nullptr; + const T* down_weight = nullptr; + const T* up_weight = nullptr; + const T* inject_weight = nullptr; + const T* prev_block_output = nullptr; // prologue inputs; both null or both non-null + const T* prev_post_gate = nullptr; + T* pre_output = nullptr; + T* post_gate = nullptr; + T* streams_out = nullptr; // written when the prologue is present; must then be non-null +}; + +// Rows up to this value use the two-kernel fused path; larger inputs use the cuBLAS based path. +constexpr int64_t kHyperConnectionLowRankMixMaxFusedRows = 8; + +// True when the fused two-kernel path supports this problem on the given device. +template +bool CanUseFusedHyperConnectionLowRankMix(const cudaDeviceProp& device_prop, + const HyperConnectionLowRankMixConfig& config); + +// Number of float elements of scratch memory the fused path needs. +size_t HyperConnectionLowRankMixFusedScratchElements(const HyperConnectionLowRankMixConfig& config); + +template +Status LaunchHyperConnectionLowRankMixFused(const cudaDeviceProp& device_prop, cudaStream_t stream, + const HyperConnectionLowRankMixConfig& config, + const HyperConnectionLowRankMixTensors& tensors, + float* scratch); + +// General path built from cuBLAS GEMMs and elementwise kernels. Scratch buffers (in elements of T): +// hn [rows, C * H], down [rows, R], up [rows, C * H], inject [rows, C]. +template +Status LaunchHyperConnectionLowRankMixUnfused(cublasHandle_t cublas, cudaStream_t stream, + const HyperConnectionLowRankMixConfig& config, + const HyperConnectionLowRankMixTensors& tensors, + T* hn, T* down, T* up, T* inject); + +} // namespace onnxruntime::contrib::cuda diff --git a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc index fe140646872de..9a68c977425b7 100644 --- a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc @@ -173,6 +173,13 @@ class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, LinearAttention); class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, GatedDeltaNet); class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, GatedDeltaNet); class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, GatedDeltaNet); +class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, ArgMaxLastAxis); +class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, ArgMaxLastAxis); +class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, ArgMaxLastAxis); +class CUDA_MS_OP_CLASS_NAME(1, EmbeddingRMSNorm); +class CUDA_MS_OP_CLASS_NAME(1, RowSelect); +class CUDA_MS_OP_CLASS_NAME(1, SigmoidGateMul); +class CUDA_MS_OP_CLASS_NAME(1, UntileHyperConnectionWeights); class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, LinearAttentionGate); class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, LinearAttentionGate); class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, LinearAttentionGate); @@ -194,6 +201,9 @@ class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, HyperConnectionPreMix); class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, HyperConnectionPostMix); class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, HyperConnectionPostMix); class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, HyperConnectionPostMix); +class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, HyperConnectionLowRankMix); +class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, HyperConnectionLowRankMix); +class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, HyperConnectionLowRankMix); class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, EngramGate); class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, EngramGate); class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, EngramGate); @@ -543,6 +553,13 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, @@ -564,6 +581,9 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index 761fc505b9530..0751df5d5de64 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -5634,6 +5634,147 @@ ONNX_MS_OPERATOR_SET_SCHEMA( } })); +constexpr const char* FusedTensorOps_ver1_doc = R"DOC( +CUDA contrib operators for fused tensor operations used by transformer models. +)DOC"; + +ONNX_MS_OPERATOR_SET_SCHEMA( + ArgMaxLastAxis, 1, + OpSchema() + .SetDoc(FusedTensorOps_ver1_doc) + .Attr("keepdims", "Keep a singleton final axis when nonzero.", AttributeProto::INT, + static_cast(0)) + .Input(0, "logits", "Logits with vocabulary on the last axis.", "T") + .Output(0, "indices", "First maximum index for each leading row.", "I") + .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, + "Supported logit types.") + .TypeConstraint("I", {"tensor(int64)"}, "Argmax indices are int64.") + .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + updateOutputElemType(ctx, 0, ONNX_NAMESPACE::TensorProto_DataType_INT64); + if (!hasInputShape(ctx, 0)) return; + const auto& input_shape = getInputShape(ctx, 0); + if (input_shape.dim_size() < 1) { + fail_shape_inference("ArgMaxLastAxis: logits must have rank at least one"); + } + ONNX_NAMESPACE::TensorShapeProto output_shape; + for (int i = 0; i + 1 < input_shape.dim_size(); ++i) { + *output_shape.add_dim() = input_shape.dim(i); + } + if (getAttribute(ctx, "keepdims", static_cast(0)) != 0) { + output_shape.add_dim()->set_dim_value(1); + } + updateOutputShape(ctx, 0, output_shape); + })); + +ONNX_MS_OPERATOR_SET_SCHEMA( + EmbeddingRMSNorm, 1, + OpSchema() + .SetDoc(FusedTensorOps_ver1_doc) + .Attr("epsilon", "RMS normalization epsilon.", AttributeProto::FLOAT) + .Input(0, "ids", "Embedding row indices with any shape.", "I") + .Input(1, "table", "BF16 embedding matrix [vocabulary, hidden].", "T") + .Input(2, "weight", "FP32 RMS normalization weight [hidden].", "F") + .Output(0, "normalized", "BF16 normalized embeddings [shape(ids), hidden].", "T") + .TypeConstraint("I", {"tensor(int64)"}, "Indices use int64.") + .TypeConstraint("T", {"tensor(bfloat16)"}, "Embedding and result use BF16.") + .TypeConstraint("F", {"tensor(float)"}, "Normalization weight uses FP32.") + .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + propagateElemTypeFromInputToOutput(ctx, 1, 0); + if (!hasInputShape(ctx, 0) || !hasInputShape(ctx, 1)) return; + const auto& ids_shape = getInputShape(ctx, 0); + const auto& table_shape = getInputShape(ctx, 1); + if (table_shape.dim_size() != 2) { + fail_shape_inference("EmbeddingRMSNorm: table must have rank two"); + } + ONNX_NAMESPACE::TensorShapeProto output_shape = ids_shape; + *output_shape.add_dim() = table_shape.dim(1); + updateOutputShape(ctx, 0, output_shape); + })); + +ONNX_MS_OPERATOR_SET_SCHEMA( + RowSelect, 1, + OpSchema() + .SetDoc(FusedTensorOps_ver1_doc) + .Input(0, "selection", "Scalar-like int64 row selector; zero selects row 0.", "I") + .Input(1, "x0", "BF16 tensor with shape [1, 2, ...].", "T") + .Input(2, "x1", "BF16 tensor with shape [1, 2, ...].", "T") + .Input(3, "x2", "BF16 tensor with shape [1, 2, ...].", "T") + .Output(0, "y0", "Selected BF16 tensor with shape [1, 1, ...].", "T") + .Output(1, "y1", "Selected BF16 tensor with shape [1, 1, ...].", "T") + .Output(2, "y2", "Selected BF16 tensor with shape [1, 1, ...].", "T") + .TypeConstraint("I", {"tensor(int64)"}, "Selector uses int64.") + .TypeConstraint("T", {"tensor(bfloat16)"}, "Inputs and outputs use BF16.") + .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + for (int i = 0; i < 3; ++i) { + propagateElemTypeFromInputToOutput(ctx, i + 1, i); + if (!hasInputShape(ctx, i + 1)) continue; + ONNX_NAMESPACE::TensorShapeProto output_shape = getInputShape(ctx, i + 1); + if (output_shape.dim_size() < 2) { + fail_shape_inference("RowSelect: inputs must have rank at least two"); + } + output_shape.mutable_dim(1)->set_dim_value(1); + updateOutputShape(ctx, i, output_shape); + } + })); + +ONNX_MS_OPERATOR_SET_SCHEMA( + SigmoidGateMul, 1, + OpSchema() + .SetDoc(FusedTensorOps_ver1_doc) + .Attr("head_dim", "Dimension of each attention head.", AttributeProto::INT, static_cast(0)) + .Attr("num_heads", "Number of attention heads.", AttributeProto::INT, static_cast(0)) + .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 " + "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 " + "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") + .TypeConstraint("T", {"tensor(bfloat16)"}, "Inputs and output use BF16.") + .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + propagateElemTypeFromInputToOutput(ctx, 0, 0); + if (hasInputShape(ctx, 0)) updateOutputShape(ctx, 0, getInputShape(ctx, 0)); + })); + +ONNX_MS_OPERATOR_SET_SCHEMA( + UntileHyperConnectionWeights, 1, + OpSchema() + .SetDoc(FusedTensorOps_ver1_doc) + .Attr("kind", "Untile layout kind, either 0 or 1.", AttributeProto::INT) + .Input(0, "tiled", "Tiled BF16 weight [160, 20480].", "T") + .Input(1, "unused", "Reserved optional input retained for exporter compatibility.", "T", + OpSchema::Optional) + .Output(0, "untiled", "Untiled BF16 weight.", "T") + .TypeConstraint("T", {"tensor(bfloat16)"}, "Input and output use BF16.") + .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + propagateElemTypeFromInputToOutput(ctx, 0, 0); + if (!hasInputShape(ctx, 0)) return; + const int64_t kind = getAttribute(ctx, "kind", static_cast(-1)); + if (kind != 0 && kind != 1) { + fail_shape_inference("UntileHyperConnectionWeights: kind must be zero or one"); + } + ONNX_NAMESPACE::TensorShapeProto output_shape; + output_shape.add_dim()->set_dim_value(kind == 0 ? 10240 : 320); + output_shape.add_dim()->set_dim_value(kind == 0 ? 320 : 10240); + updateOutputShape(ctx, 0, output_shape); + })); + constexpr const char* LinearAttentionGate_ver1_doc = R"DOC( Fuses the gate projections that feed LinearAttention's gated-delta recurrence: @@ -5969,6 +6110,113 @@ ONNX_MS_OPERATOR_SET_SCHEMA( propagateShapeFromInputToOutput(ctx, 0, 0); })); +constexpr const char* HyperConnectionLowRankMix_ver1_doc = R"DOC( +One hyper-connection block with a low-rank gate network, fused into one operator. C = num_branches streams of +width H are kept as one flattened (..., C * H) tensor. With the optional prologue the operator first applies the +post-mix of the previous block (identity stream mix): + + S[..., c, h] = streams[..., c, h] + prev_post_gate[..., c] * prev_block_output[..., h] + +and otherwise S = streams. It then computes, with `rms` taken per branch over H: + + hn[..., c, h] = S[..., c, h] / sqrt(mean_h(S[..., c, :]^2) + epsilon) * norm_scale[c, h] + d = ScaledSiLU(hn @ down_weight, silu_alpha) (..., R) + g[..., c, h] = sigmoid(d @ up_weight) per stream and feature + pre_output = reduction_scale * sum_c(hn[..., c, h] * g[..., c, h]) (..., H) + post_gate = sigmoid((hn @ inject_weight) / gate_div) * gate_mul (..., C) + +pre_output is the input of the next transformer block and post_gate is prev_post_gate of the next call. +streams_out is S (the streams after the prologue; a copy of streams when there is no prologue). + +Numerics contract. Matrix products and reductions accumulate in float32. The result of every step is rounded +to T once, as the output of the equivalent standard operator would be: +S = cast_T(streams + prev_post_gate * prev_block_output), hn = cast_T(S * inv_rms * norm_scale), +down = cast_T(hn @ down_weight), Z = cast_T(silu_alpha * down), SG = cast_T(sigmoid(Z)), hidden = cast_T(Z * SG), +up = cast_T(hidden @ up_weight), g = cast_T(sigmoid(up)), pre_output = cast_T(reduction_scale * sum_c(hn * g)), +inject = cast_T(hn @ inject_weight), post_gate = cast_T(cast_T(sigmoid(cast_T(inject / gate_div))) * gate_mul). +sigmoid is evaluated in float32 with expf and rounded once to T; it does not reproduce a Sigmoid kernel that +rounds to T after every elementary operation. The accumulation order of the matrix products is fixed, so the +results are deterministic, but they are not bitwise identical to cuBLAS. +)DOC"; + +ONNX_MS_OPERATOR_SET_SCHEMA( + HyperConnectionLowRankMix, 1, + OpSchema() + .SetDoc(HyperConnectionLowRankMix_ver1_doc) + .Attr("num_branches", "Number of residual streams C. The last dimension of streams is C * H.", + AttributeProto::INT, true) + .Attr("epsilon", "Epsilon added to the mean square of each branch before the reciprocal square root.", + AttributeProto::FLOAT, 1e-5f) + .Attr("reduction_scale", "Multiplier applied to the sum over branches of pre_output.", + AttributeProto::FLOAT, 1.0f) + .Attr("silu_alpha", "Scale of the ScaledSiLU between the down and up projections.", + AttributeProto::FLOAT, 1.0f) + .Attr("gate_div", "Divisor applied to the inject logits before the sigmoid. Must not be zero.", + AttributeProto::FLOAT, 1.0f) + .Attr("gate_mul", "Multiplier applied to the sigmoid of the inject logits.", + AttributeProto::FLOAT, 1.0f) + .Input(0, "streams", + "Flattened streams (..., C * H). With the prologue these are the streams before the previous " + "block's post-mix.", + "T") + .Input(1, "norm_scale", "Per-branch RMS norm weight with shape (C * H).", "T") + .Input(2, "down_weight", "Down projection with shape (C * H, R), row-major [K, N].", "T") + .Input(3, "up_weight", "Up projection with shape (R, C * H), row-major [K, N].", "T") + .Input(4, "inject_weight", "Post-gate projection with shape (C * H, C), row-major [K, N].", "T") + .Input(5, "prev_block_output", + "Optional output of the previous block with shape (..., H). Provide together with prev_post_gate.", + "T", OpSchema::Optional) + .Input(6, "prev_post_gate", + "Optional post-mix gate of the previous block with shape (..., C). Provide together with " + "prev_block_output.", + "T", OpSchema::Optional) + .Output(0, "pre_output", "Block input with shape (..., H).", "T") + .Output(1, "post_gate", "Post-mix gate for the next call with shape (..., C).", "T") + .Output(2, "streams_out", + "Optional streams after the prologue, with the same shape as streams. Equal to streams when the " + "prologue inputs are absent.", + "T", OpSchema::Optional) + .TypeConstraint("T", + {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, + "Constrain all floating-point inputs and outputs to the same type.") + .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + propagateElemTypeFromInputToOutput(ctx, 0, 0); + propagateElemTypeFromInputToOutput(ctx, 0, 1); + if (ctx.getNumOutputs() > 2) { + propagateElemTypeFromInputToOutput(ctx, 0, 2); + } + if (!hasInputShape(ctx, 0)) { + return; + } + const auto branches = getAttribute(ctx, "num_branches", 0); + if (branches < 1) { + fail_shape_inference("num_branches must be positive."); + } + const auto& input_shape = getInputShape(ctx, 0); + const int rank = input_shape.dim_size(); + if (rank < 1) { + fail_shape_inference("streams must have rank at least 1."); + } + auto* pre_shape = getOutputShape(ctx, 0); + auto* gate_shape = getOutputShape(ctx, 1); + for (int i = 0; i < rank - 1; ++i) { + *pre_shape->add_dim() = input_shape.dim(i); + *gate_shape->add_dim() = input_shape.dim(i); + } + auto* hidden = pre_shape->add_dim(); + const auto& last = input_shape.dim(rank - 1); + if (last.has_dim_value()) { + if (last.dim_value() % branches != 0) { + fail_shape_inference("The last dimension of streams must be divisible by num_branches."); + } + hidden->set_dim_value(last.dim_value() / branches); + } + gate_shape->add_dim()->set_dim_value(branches); + if (ctx.getNumOutputs() > 2) { + propagateShapeFromInputToOutput(ctx, 0, 2); + } + })); + constexpr const char* GatedAdd_ver1_doc = R"DOC( Adds one tensor to another tensor scaled by a per-row gate: diff --git a/onnxruntime/core/graph/contrib_ops/ms_opset.h b/onnxruntime/core/graph/contrib_ops/ms_opset.h index 93b771a425e07..12fc06b9816c3 100644 --- a/onnxruntime/core/graph/contrib_ops/ms_opset.h +++ b/onnxruntime/core/graph/contrib_ops/ms_opset.h @@ -95,7 +95,13 @@ class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 2, BranchwiseRMSNorm); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, ScaledSiLU); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, HyperConnectionPreMix); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, HyperConnectionPostMix); +class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, HyperConnectionLowRankMix); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, GatedDeltaNet); +class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, ArgMaxLastAxis); +class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, EmbeddingRMSNorm); +class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, RowSelect); +class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, SigmoidGateMul); +class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, UntileHyperConnectionWeights); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, GatedDeltaNetStateReplay); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, GatedRMSNorm); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, GatedAdd); @@ -231,7 +237,13 @@ class OpSet_Microsoft_ver1 { fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); + fn(GetOpSchema()); fn(GetOpSchema()); + fn(GetOpSchema()); + fn(GetOpSchema()); + fn(GetOpSchema()); + fn(GetOpSchema()); + fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); diff --git a/onnxruntime/test/contrib_ops/fused_tensor_ops_test.cc b/onnxruntime/test/contrib_ops/fused_tensor_ops_test.cc new file mode 100644 index 0000000000000..2173434746ce9 --- /dev/null +++ b/onnxruntime/test/contrib_ops/fused_tensor_ops_test.cc @@ -0,0 +1,283 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include +#include +#include +#include +#include + +#include "core/common/float8.h" +#include "gtest/gtest.h" +#include "test/common/tensor_op_test_utils.h" +#include "test/providers/provider_test_utils.h" +#include "test/util/include/default_providers.h" + +namespace onnxruntime::test { +namespace { + +template +std::vector ToTensorType(const std::vector& data) { + if constexpr (std::is_same_v) { + return ToFloat16(data); + } else if constexpr (std::is_same_v) { + return ToBFloat16(data); + } else { + return data; + } +} + +std::vector> CudaProvider() { + std::vector> providers; + if (auto cuda = DefaultCudaExecutionProvider()) { + providers.push_back(std::move(cuda)); + } + return providers; +} + +TEST(FusedTensorOpsTest, ArgMaxLastAxisWithoutKeepdims) { + auto providers = CudaProvider(); + if (providers.empty()) GTEST_SKIP(); + + OpTester test("ArgMaxLastAxis", 1, kMSDomain); + test.AddInput("logits", {2, 4}, {1.0f, 5.0f, 5.0f, -1.0f, 4.0f, 3.0f, 2.0f, 1.0f}); + test.AddOutput("indices", {2}, {1, 0}); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); +} + +TEST(FusedTensorOpsTest, ArgMaxLastAxisWithKeepdims) { + auto providers = CudaProvider(); + if (providers.empty()) GTEST_SKIP(); + + OpTester test("ArgMaxLastAxis", 1, kMSDomain); + test.AddAttribute("keepdims", 1); + test.AddInput("logits", {2, 4}, + ToTensorType({1.0f, 5.0f, 5.0f, -1.0f, 4.0f, 3.0f, 2.0f, 1.0f})); + test.AddOutput("indices", {2, 1}, {1, 0}); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); +} + +TEST(FusedTensorOpsTest, ArgMaxLastAxisPartitionedVocabulary) { + constexpr int64_t width = 248321; + for (bool keepdims : {false, true}) { + auto providers = CudaProvider(); + if (providers.empty()) GTEST_SKIP(); + std::vector values(3 * width, -std::numeric_limits::infinity()); + values[4095] = values[4096] = values[width - 1] = 5.0f; + values[2 * width - 1] = 7.0f; + OpTester test("ArgMaxLastAxis", 1, kMSDomain); + test.AddAttribute("keepdims", keepdims ? 1 : 0); + test.AddInput("logits", {3, width}, ToBFloat16(values)); + test.AddOutput("indices", keepdims ? std::vector{3, 1} : std::vector{3}, + {4095, width - 1, 0}); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); + } +} + +TEST(FusedTensorOpsTest, EmbeddingRMSNorm) { + auto providers = CudaProvider(); + if (providers.empty()) GTEST_SKIP(); + + const std::vector table = {3.0f, 4.0f, 0.0f, 0.0f, 0.0f, 0.0f, 2.0f, 0.0f}; + const float inverse_first = 1.0f / std::sqrt(6.25f + 1e-6f); + const float inverse_second = 1.0f / std::sqrt(1.0f + 1e-6f); + const std::vector expected = { + 3.0f * inverse_first, 4.0f * inverse_first, 0.0f, 0.0f, + 0.0f, 0.0f, 2.0f * inverse_second, 0.0f}; + + OpTester test("EmbeddingRMSNorm", 1, kMSDomain); + test.AddAttribute("epsilon", 1e-6f); + test.AddInput("ids", {2}, {0, 1}); + test.AddInput("table", {2, 4}, ToTensorType(table)); + test.AddInput("weight", {4}, {1.0f, 1.0f, 1.0f, 1.0f}); + test.AddOutput("normalized", {2, 4}, ToTensorType(expected)); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); +} + +TEST(FusedTensorOpsTest, RowSelect) { + auto providers = CudaProvider(); + if (providers.empty()) GTEST_SKIP(); + + OpTester test("RowSelect", 1, kMSDomain); + test.AddInput("selection", {}, {1}); + for (int i = 0; i < 3; ++i) { + const std::string input_name = "x" + std::to_string(i); + const std::string output_name = "y" + std::to_string(i); + test.AddInput(input_name.c_str(), {1, 2, 2}, + ToTensorType({1.0f, 2.0f, 3.0f, 4.0f})); + test.AddOutput(output_name.c_str(), {1, 1, 2}, + ToTensorType({3.0f, 4.0f})); + } + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); +} + +TEST(FusedTensorOpsTest, SigmoidGateMul) { + auto providers = CudaProvider(); + if (providers.empty()) GTEST_SKIP(); + + OpTester test("SigmoidGateMul", 1, kMSDomain); + test.AddAttribute("head_dim", 2); + test.AddAttribute("num_heads", 1); + test.AddAttribute("q_stride", 4); + test.AddAttribute("tok_stride", 4); + test.AddAttribute("q_offset", 0); + test.AddInput("attention", {1, 1, 2}, ToTensorType({2.0f, 4.0f})); + test.AddInput("projection", {1, 4}, ToTensorType({1.0f, 2.0f, 0.0f, 0.0f})); + test.AddOutput("output", {1, 1, 2}, ToTensorType({1.0f, 2.0f})); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); +} + +void RunSigmoidGateMulBroadcastResidualRounding(bool sigmoid_in_fp32, bool silu_gate = false) { + for (bool residual : {false, true}) { + for (int64_t width : {1, 257, 2560}) { + auto providers = CudaProvider(); + if (providers.empty()) GTEST_SKIP(); + const std::vector base_gates = ToBFloat16({-100.0f, -5.5f, -0.0f, 0.3f, 5.5f, 100.0f}); + std::vector gates(silu_gate ? 6 * width : 6); + for (size_t i = 0; i < gates.size(); ++i) gates[i] = base_gates[i % 6]; + std::vector values(6 * width), residuals(6 * width), expected(6 * width); + for (size_t i = 0; i < values.size(); ++i) { + values[i] = BFloat16(static_cast(static_cast(i % 23) - 11) * 0.27f); + residuals[i] = BFloat16(static_cast(static_cast(i % 17) - 8) * 0.13f); + const float gate = gates[silu_gate ? i : i / width].ToFloat(); + const float exp_value = std::exp(gate >= 0 ? -gate : gate); + const float sigmoid = gate >= 0 ? 1.0f / (1.0f + exp_value) : exp_value / (1.0f + exp_value); + 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())); + 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; + } + OpTester test("SigmoidGateMul", 1, kMSDomain); + test.AddAttribute(silu_gate ? "silu_gate" : "broadcast_gate", 1); + test.AddAttribute("sigmoid_in_fp32", sigmoid_in_fp32 ? 1 : 0); + test.AddInput("attention", {2, 3, width}, values); + test.AddInput("projection", {2, 3, silu_gate ? width : 1}, gates); + if (residual) test.AddInput("residual", {2, 3, width}, residuals); + test.AddOutput("output", {2, 3, width}, expected); + test.SetOutputTolerance(0.0f, 0.0f); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); + } + } +} + +TEST(FusedTensorOpsTest, SigmoidGateMulBroadcastResidualRounding) { + RunSigmoidGateMulBroadcastResidualRounding(true); +} + +TEST(FusedTensorOpsTest, SigmoidGateMulBroadcastInputPrecision) { + RunSigmoidGateMulBroadcastResidualRounding(false); +} + +TEST(FusedTensorOpsTest, SigmoidGateMulSiluRounding) { + RunSigmoidGateMulBroadcastResidualRounding(true, true); + RunSigmoidGateMulBroadcastResidualRounding(false, true); +} + +TEST(FusedTensorOpsTest, SigmoidGateMulRejectsBroadcastAndResidualShapes) { + for (bool bad_residual : {false, true}) { + auto providers = CudaProvider(); + if (providers.empty()) GTEST_SKIP(); + OpTester test("SigmoidGateMul", 1, kMSDomain); + 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})); + 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, + bad_residual ? "residual must have the same shape" : "singleton final dimension", + {}, nullptr, &providers); + } +} + +TEST(FusedTensorOpsTest, SigmoidGateMulBroadcastEmptyAndInvalidAttributes) { + for (int mode = 0; mode < 4; ++mode) { + auto providers = CudaProvider(); + if (providers.empty()) GTEST_SKIP(); + OpTester test("SigmoidGateMul", 1, kMSDomain); + test.AddAttribute("broadcast_gate", mode == 1 ? 2 : 1); + if (mode == 2) test.AddAttribute("sigmoid_in_fp32", 2); + if (mode == 3) test.AddAttribute("head_dim", 2); + test.AddInput("attention", {0, 3}, {}); + test.AddInput("projection", {0, 1}, {}); + 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"}; + test.Run(mode == 0 ? OpTester::ExpectResult::kExpectSuccess : OpTester::ExpectResult::kExpectFailure, + errors[mode], {}, nullptr, &providers); + } +} + +TEST(FusedTensorOpsTest, SigmoidGateMulSiluEmptyAndInvalidAttributes) { + for (int mode = 0; mode < 5; ++mode) { + auto providers = CudaProvider(); + if (providers.empty()) GTEST_SKIP(); + OpTester test("SigmoidGateMul", 1, kMSDomain); + test.AddAttribute("silu_gate", mode == 1 ? 2 : 1); + if (mode == 2) test.AddAttribute("broadcast_gate", 1); + if (mode == 3) test.AddAttribute("head_dim", 2); + test.AddInput("attention", {0, 3}, {}); + 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"}; + test.Run(mode == 0 ? OpTester::ExpectResult::kExpectSuccess : OpTester::ExpectResult::kExpectFailure, + errors[mode], {}, nullptr, &providers); + } +} + +TEST(FusedTensorOpsTest, UntileHyperConnectionWeights) { + constexpr int64_t kInputElements = 160ll * 20480; + constexpr int64_t kVectorCount = kInputElements / 16; + for (const int64_t kind : {0, 1}) { + auto providers = CudaProvider(); + if (providers.empty()) GTEST_SKIP(); + OpTester test("UntileHyperConnectionWeights", 1, kMSDomain); + test.AddAttribute("kind", kind); + std::vector input(kInputElements); + std::vector expected(kInputElements); + for (int64_t i = 0; i < kInputElements; ++i) { + input[i] = BFloat16(static_cast(i % 127)); + } + for (int64_t vector_index = 0; vector_index < kVectorCount; ++vector_index) { + int64_t destination_vector; + if (kind == 0) { + const int64_t rp = vector_index & 127; + const int64_t t1 = vector_index >> 7; + const int64_t i = t1 % 10; + const int64_t t2 = t1 / 10; + const int64_t j = t2 % 20; + const int64_t s = t2 / 20; + destination_vector = (((s * 10 + i) * 128 + rp) * 20 + j); + } else { + const int64_t c = vector_index & 3; + const int64_t t1 = vector_index >> 2; + const int64_t rp = t1 & 31; + const int64_t t2 = t1 >> 5; + const int64_t i = t2 % 10; + const int64_t hb = t2 / 10; + destination_vector = (((i * 32 + rp) * 4 + c) * 160 + hb); + } + for (int64_t lane = 0; lane < 16; ++lane) { + expected[destination_vector * 16 + lane] = input[vector_index * 16 + lane]; + } + } + test.AddInput("tiled", {160, 20480}, input); + const std::vector output_shape = kind == 0 + ? std::vector{10240, 320} + : std::vector{320, 10240}; + test.AddOutput("untiled", output_shape, expected); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); + } +} + +} // namespace +} // namespace onnxruntime::test diff --git a/onnxruntime/test/contrib_ops/hyper_connection_low_rank_mix_op_test.cc b/onnxruntime/test/contrib_ops/hyper_connection_low_rank_mix_op_test.cc new file mode 100644 index 0000000000000..917d936f45468 --- /dev/null +++ b/onnxruntime/test/contrib_ops/hyper_connection_low_rank_mix_op_test.cc @@ -0,0 +1,285 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include +#include +#include +#include +#include + +#include "gtest/gtest.h" +#include "test/common/tensor_op_test_utils.h" +#include "test/providers/provider_test_utils.h" +#include "test/util/include/default_providers.h" + +namespace onnxruntime::test { +namespace { + +template +float RoundToType(float value) { + if constexpr (std::is_same_v) { + return MLFloat16(value).ToFloat(); + } else if constexpr (std::is_same_v) { + return BFloat16(value).ToFloat(); + } else { + return value; + } +} + +template +std::vector ToTensorType(const std::vector& data) { + if constexpr (std::is_same_v) { + return ToFloat16(data); + } else if constexpr (std::is_same_v) { + return ToBFloat16(data); + } else { + return data; + } +} + +float Sigmoid(float v) { return 1.0f / (1.0f + std::exp(-v)); } + +struct Attrs { + float epsilon = 1e-6f; + float reduction_scale = 0.25f; + float silu_alpha = 0.25f; + float gate_div = 4.0f; + float gate_mul = 2.0f; +}; + +struct Problem { + int64_t rows, C, H, R; + std::vector streams, norm, down, up, inject, prev_out, prev_gate; +}; + +template +Problem MakeProblem(int64_t rows, int64_t C, int64_t H, int64_t R, bool prologue, unsigned seed) { + std::mt19937 gen(seed); + std::normal_distribution normal(0.0f, 1.0f); + const int64_t K = C * H; + auto fill = [&](size_t n, float scale, float shift) { + std::vector v(n); + for (auto& x : v) x = RoundToType(normal(gen) * scale + shift); + return v; + }; + Problem p{rows, C, H, R, {}, {}, {}, {}, {}, {}, {}}; + p.streams = fill(rows * K, 1.0f, 0.0f); + p.norm = fill(K, 0.1f, 1.0f); + p.down = fill(K * R, 1.0f / std::sqrt(static_cast(K)), 0.0f); + p.up = fill(R * K, 2.0f / std::sqrt(static_cast(R)), 0.0f); + p.inject = fill(K * C, 1.0f / std::sqrt(static_cast(K)), 0.0f); + if (prologue) { + p.prev_out = fill(rows * H, 1.0f, 0.0f); + p.prev_gate = fill(rows * C, 0.5f, 1.0f); + } + return p; +} + +struct Expected { + std::vector pre, gate, streams; +}; + +// float32 math with the rounding points listed in the schema documentation. +template +Expected Reference(const Problem& p, const Attrs& a) { + const int64_t K = p.C * p.H; + auto rd = [](float v) { return RoundToType(v); }; + Expected e; + e.streams = p.streams; + const bool prologue = !p.prev_out.empty(); + for (int64_t r = 0; r < p.rows; ++r) { + std::vector hn(K); + for (int64_t c = 0; c < p.C; ++c) { + float sum_sq = 0.0f; + for (int64_t h = 0; h < p.H; ++h) { + float v = p.streams[r * K + c * p.H + h]; + if (prologue) { + v = rd(v + p.prev_gate[r * p.C + c] * p.prev_out[r * p.H + h]); + e.streams[r * K + c * p.H + h] = v; + } + sum_sq += v * v; + } + const float inv = 1.0f / std::sqrt(sum_sq / static_cast(p.H) + a.epsilon); + for (int64_t h = 0; h < p.H; ++h) { + hn[c * p.H + h] = rd(e.streams[r * K + c * p.H + h] * inv * p.norm[c * p.H + h]); + } + } + std::vector hidden(p.R); + for (int64_t j = 0; j < p.R; ++j) { + double acc = 0.0; + for (int64_t k = 0; k < K; ++k) acc += static_cast(hn[k]) * p.down[k * p.R + j]; + const float z = rd(rd(static_cast(acc)) * a.silu_alpha); + hidden[j] = rd(z * rd(Sigmoid(z))); + } + std::vector gates(K); + for (int64_t k = 0; k < K; ++k) { + double acc = 0.0; + for (int64_t j = 0; j < p.R; ++j) acc += static_cast(hidden[j]) * p.up[j * K + k]; + gates[k] = rd(Sigmoid(rd(static_cast(acc)))); + } + for (int64_t h = 0; h < p.H; ++h) { + float sum = 0.0f; + for (int64_t c = 0; c < p.C; ++c) sum += hn[c * p.H + h] * gates[c * p.H + h]; + e.pre.push_back(rd(sum * a.reduction_scale)); + } + for (int64_t c = 0; c < p.C; ++c) { + double acc = 0.0; + for (int64_t k = 0; k < K; ++k) acc += static_cast(hn[k]) * p.inject[k * p.C + c]; + const float t = rd(rd(static_cast(acc)) / a.gate_div); + e.gate.push_back(rd(rd(Sigmoid(t)) * a.gate_mul)); + } + } + return e; +} + +template +void RunTest(int64_t rows, int64_t C, int64_t H, int64_t R, bool prologue, bool request_streams_out, + float tolerance, unsigned seed = 1) { + auto cuda = DefaultCudaExecutionProvider(); + if (!cuda) { + GTEST_SKIP() << "HyperConnectionLowRankMix is only implemented for CUDA"; + } + const Attrs a; + const Problem p = MakeProblem(rows, C, H, R, prologue, seed); + const Expected e = Reference(p, a); + const int64_t K = C * H; + + OpTester test("HyperConnectionLowRankMix", 1, kMSDomain); + test.AddAttribute("num_branches", C); + test.AddAttribute("epsilon", a.epsilon); + test.AddAttribute("reduction_scale", a.reduction_scale); + test.AddAttribute("silu_alpha", a.silu_alpha); + test.AddAttribute("gate_div", a.gate_div); + test.AddAttribute("gate_mul", a.gate_mul); + test.AddInput("streams", {rows, K}, ToTensorType(p.streams)); + test.AddInput("norm_scale", {K}, ToTensorType(p.norm)); + test.AddInput("down_weight", {K, R}, ToTensorType(p.down)); + test.AddInput("up_weight", {R, K}, ToTensorType(p.up)); + test.AddInput("inject_weight", {K, C}, ToTensorType(p.inject)); + if (prologue) { + test.AddInput("prev_block_output", {rows, H}, ToTensorType(p.prev_out)); + test.AddInput("prev_post_gate", {rows, C}, ToTensorType(p.prev_gate)); + } else { + test.AddOptionalInputEdge(); + test.AddOptionalInputEdge(); + } + test.AddOutput("pre_output", {rows, H}, ToTensorType(e.pre)); + test.AddOutput("post_gate", {rows, C}, ToTensorType(e.gate)); + if (request_streams_out) { + test.AddOutput("streams_out", {rows, K}, ToTensorType(e.streams)); + } else { + test.AddOptionalOutputEdge(); + } + test.SetOutputTolerance(tolerance, tolerance); + + std::vector> providers; + providers.push_back(std::move(cuda)); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); +} + +constexpr float kFp32Tol = 2e-4f; +constexpr float kFp16Tol = 8e-3f; +constexpr float kBf16Tol = 6e-2f; + +} // namespace + +TEST(HyperConnectionLowRankMixTest, FusedFloat16WithPrologue) { + RunTest(2, 4, 16, 8, true, true, kFp16Tol); +} + +TEST(HyperConnectionLowRankMixTest, FusedFloat16WithoutPrologue) { + RunTest(1, 4, 64, 16, false, false, kFp16Tol); +} + +TEST(HyperConnectionLowRankMixTest, FusedBFloat16WithPrologueWithoutStreamsOut) { + RunTest(3, 4, 64, 16, true, false, kBf16Tol); +} + +TEST(HyperConnectionLowRankMixTest, FusedBFloat16TwoBranches) { + RunTest(2, 2, 32, 8, true, true, kBf16Tol); +} + +TEST(HyperConnectionLowRankMixTest, FusedMultipleKSlicesAndWeightChunks) { + RunTest(2, 4, 1288, 40, true, true, kBf16Tol); + RunTest(1, 4, 2560, 320, false, true, kFp16Tol); +} + +TEST(HyperConnectionLowRankMixTest, UnfusedManyRows) { + RunTest(9, 4, 16, 8, true, true, kFp16Tol); + RunTest(17, 4, 64, 16, false, true, kBf16Tol); +} + +TEST(HyperConnectionLowRankMixTest, FusedSpecializedRank320) { + for (int rows : {1, 2, 3}) { + RunTest(rows, 4, 1288, 320, true, true, kBf16Tol); + RunTest(rows, 4, 1288, 320, false, false, kFp16Tol); + } +} + +TEST(HyperConnectionLowRankMixTest, FusedDecodeShapeRank320) { + for (int rows : {1, 2}) { + for (bool prologue : {false, true}) { + RunTest(rows, 4, 2560, 320, prologue, true, kBf16Tol, 37); + RunTest(rows, 4, 2560, 320, prologue, true, kFp16Tol, 53); + } + } +} + +TEST(HyperConnectionLowRankMixTest, UnfusedUnsupportedShape) { + RunTest(2, 3, 10, 5, true, true, kFp16Tol); + RunTest(1, 4, 12, 16, false, false, kBf16Tol); +} + +TEST(HyperConnectionLowRankMixTest, Float32) { + RunTest(1, 4, 64, 16, true, true, kFp32Tol); + RunTest(5, 2, 24, 8, false, true, kFp32Tol); +} + +TEST(HyperConnectionLowRankMixTest, PrologueInputsMustBeProvidedTogether) { + auto cuda = DefaultCudaExecutionProvider(); + if (!cuda) { + GTEST_SKIP() << "HyperConnectionLowRankMix is only implemented for CUDA"; + } + constexpr int64_t C = 2, H = 8, R = 8, K = C * H; + OpTester test("HyperConnectionLowRankMix", 1, kMSDomain); + test.AddAttribute("num_branches", C); + test.AddInput("streams", {1, K}, std::vector(K, 1.0f)); + test.AddInput("norm_scale", {K}, std::vector(K, 1.0f)); + test.AddInput("down_weight", {K, R}, std::vector(K * R, 0.1f)); + test.AddInput("up_weight", {R, K}, std::vector(R * K, 0.1f)); + test.AddInput("inject_weight", {K, C}, std::vector(K * C, 0.1f)); + test.AddInput("prev_block_output", {1, H}, std::vector(H, 1.0f)); + test.AddOptionalInputEdge(); + test.AddOutput("pre_output", {1, H}, std::vector(H, 0.0f)); + test.AddOutput("post_gate", {1, C}, std::vector(C, 0.0f)); + test.AddOptionalOutputEdge(); + std::vector> providers; + providers.push_back(std::move(cuda)); + test.Run(OpTester::ExpectResult::kExpectFailure, "must be provided together", {}, nullptr, &providers); +} + +TEST(HyperConnectionLowRankMixTest, WeightShapeMismatchFails) { + auto cuda = DefaultCudaExecutionProvider(); + if (!cuda) { + GTEST_SKIP() << "HyperConnectionLowRankMix is only implemented for CUDA"; + } + constexpr int64_t C = 2, H = 8, R = 8, K = C * H; + OpTester test("HyperConnectionLowRankMix", 1, kMSDomain); + test.AddAttribute("num_branches", C); + test.AddInput("streams", {1, K}, std::vector(K, 1.0f)); + test.AddInput("norm_scale", {K}, std::vector(K, 1.0f)); + test.AddInput("down_weight", {K, R}, std::vector(K * R, 0.1f)); + test.AddInput("up_weight", {R + 1, K}, std::vector((R + 1) * K, 0.1f)); + test.AddInput("inject_weight", {K, C}, std::vector(K * C, 0.1f)); + test.AddOptionalInputEdge(); + test.AddOptionalInputEdge(); + test.AddOutput("pre_output", {1, H}, std::vector(H, 0.0f)); + test.AddOutput("post_gate", {1, C}, std::vector(C, 0.0f)); + test.AddOptionalOutputEdge(); + std::vector> providers; + providers.push_back(std::move(cuda)); + test.Run(OpTester::ExpectResult::kExpectFailure, "up_weight must have shape", {}, nullptr, &providers); +} + +} // namespace onnxruntime::test diff --git a/onnxruntime/test/contrib_ops/hyper_connection_op_test.cc b/onnxruntime/test/contrib_ops/hyper_connection_op_test.cc index 2c81513d8cb50..d7e1bc7acdd72 100644 --- a/onnxruntime/test/contrib_ops/hyper_connection_op_test.cc +++ b/onnxruntime/test/contrib_ops/hyper_connection_op_test.cc @@ -2,7 +2,9 @@ // Licensed under the MIT License. #include +#include #include +#include #include #include @@ -10,6 +12,7 @@ #include "test/common/tensor_op_test_utils.h" #include "test/providers/provider_test_utils.h" #include "test/util/include/default_providers.h" +#include "core/framework/session_options.h" namespace onnxruntime::test { namespace { @@ -106,6 +109,255 @@ TEST(HyperConnectionOpsTest, BranchwiseRMSNormMixedScaleType) { }); } +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) {} + + void AddNodes(Graph& graph, std::vector& inputs, std::vector& outputs, + std::vector>&) override { + auto value = [&](const char* name) { return &graph.GetOrCreateNodeArg(name, nullptr); }; + NodeArg* source = inputs[0]; + if (width_ != 0) { + auto constant = [&](const char* name, int64_t data) { + ONNX_NAMESPACE::TensorProto tensor; + tensor.set_name(name); + tensor.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT64); + tensor.add_dims(1); + tensor.add_int64_data(data); + graph.AddInitializedTensor(tensor); + return value(name); + }; + source = value("selected"); + graph.AddNode("slice", "Slice", "", + {inputs[0], constant("starts", start_), constant("ends", start_ + width_), + constant("axes", -1)}, + {source}); + } + auto* x = value("float_x"); + graph.AddNode("cast", "Cast", "", {source}, {x}) + .AddAttribute("to", static_cast(ONNX_NAMESPACE::TensorProto_DataType_FLOAT)); + auto* square = value("square"); + graph.AddNode("square", "Mul", "", {x, x}, {square}); + auto* mean = value("mean"); + auto& reduce = graph.AddNode("mean", "ReduceMean", "", {square}, {mean}); + reduce.AddAttribute("axes", std::vector{-1}); + reduce.AddAttribute("keepdims", int64_t{1}); + ONNX_NAMESPACE::TensorProto epsilon; + epsilon.set_name("epsilon"); + epsilon.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + epsilon.add_float_data(epsilon_); + graph.AddInitializedTensor(epsilon); + auto* adjusted = value("adjusted"); + graph.AddNode("epsilon", "Add", "", {mean, value("epsilon")}, {adjusted}); + auto* root = value("root"); + graph.AddNode("root", "Sqrt", "", {adjusted}, {root}); + auto* normalized = value("normalized"); + graph.AddNode("divide", "Div", "", {x, root}, {normalized}); + auto* scaled = value("scaled"); + graph.AddNode("scale", "Mul", "", {normalized, inputs[1]}, {scaled}); + graph.AddNode("output", "Cast", "", {scaled}, {outputs[0]}) + .AddAttribute("to", static_cast(output_type_)); + if (outputs.size() > 1) { + graph.AddNode("widen", "Cast", "", {outputs[0]}, {outputs[1]}) + .AddAttribute("to", static_cast(ONNX_NAMESPACE::TensorProto_DataType_FLOAT)); + } + graph.SetOutputs(outputs); + } + + private: + int output_type_; + float epsilon_; + int64_t start_; + int64_t width_; +}; + +template +void TestStandardRMSArithmetic(int type, bool packed = false) { + if (!DefaultCudaExecutionProvider()) { + GTEST_SKIP() << "CUDA is required for exact standard RMS arithmetic"; + } + for (int width : {1, 31, 128, 256, 513, 1024, 4096}) { + for (int rows : (packed ? std::vector{1, 2, 9, 48} : std::vector{1, 2, 9})) { + for (int start : (packed ? std::vector{0, 3} : std::vector{0})) { + SCOPED_TRACE("width=" + std::to_string(width) + " rows=" + std::to_string(rows)); + const int stride = packed ? width * 2 + 7 : width; + std::vector input(rows * stride); + std::vector scale(width); + for (int i = 0; i < rows * stride; ++i) { + input[i] = T(std::sin(static_cast(i) * 0.71f) * (1 + i % 17)); + } + for (int i = 0; i < width; ++i) scale[i] = 0.5f + static_cast(i % 23) / 17.0f; + constexpr float epsilon = 1e-6f; + StandardRMSChainTester reference(type, epsilon, start, packed ? width : 0); + OpTester fused("BranchwiseRMSNorm", packed ? 2 : 1, kMSDomain, false); + if (packed) { + fused.AddAttribute("input_start", start); + fused.AddAttribute("input_width", width); + } + fused.AddAttribute("use_onnx_arithmetic", 1); + fused.AddAttribute("epsilon", epsilon); + 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}; + auto output_shape = input_shape; + output_shape.back() = width; + tester->AddInput("X", input_shape, input); + tester->AddInput("scale", {width}, scale); + tester->AddOutput("Y", output_shape, std::vector(rows * width)); + tester->AddOutput("Y_fp32", output_shape, std::vector(rows * width)); + SessionOptions options; + options.graph_optimization_level = TransformerLevel::Default; + std::vector> providers; + providers.push_back(DefaultCudaExecutionProvider()); + tester->Run(options, OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); + auto fetches = tester->GetFetches(); + ASSERT_EQ(fetches.size(), 2); + if (expected.empty()) { + expected = std::move(fetches); + } else { + for (size_t output = 0; output < fetches.size(); ++output) { + const auto& before = expected[output].Get(); + const auto& after = fetches[output].Get(); + ASSERT_EQ(before.SizeInBytes(), after.SizeInBytes()); + EXPECT_EQ(std::memcmp(before.DataRaw(), after.DataRaw(), before.SizeInBytes()), 0) + << "output=" << output; + } + } + } + } + } + } +} + +TEST(HyperConnectionOpsTest, BranchwiseRMSNormStandardFloatExact) { + TestStandardRMSArithmetic(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); +} + +TEST(HyperConnectionOpsTest, BranchwiseRMSNormStandardFloat16Exact) { + TestStandardRMSArithmetic(ONNX_NAMESPACE::TensorProto_DataType_FLOAT16); +} + +TEST(HyperConnectionOpsTest, BranchwiseRMSNormStandardBFloat16Exact) { + TestStandardRMSArithmetic(ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16); +} + +TEST(HyperConnectionOpsTest, BranchwiseRMSNormPackedFloatExact) { + TestStandardRMSArithmetic(ONNX_NAMESPACE::TensorProto_DataType_FLOAT, true); +} + +TEST(HyperConnectionOpsTest, BranchwiseRMSNormPackedFloat16Exact) { + TestStandardRMSArithmetic(ONNX_NAMESPACE::TensorProto_DataType_FLOAT16, true); +} + +TEST(HyperConnectionOpsTest, BranchwiseRMSNormPackedBFloat16Exact) { + TestStandardRMSArithmetic(ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16, true); +} + +TEST(HyperConnectionOpsTest, BranchwiseRMSNormPackedRejectsInvalidSlice) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA is required for packed RMS"; + for (const auto& [start, width, mode, branches, message] : + std::vector>{ + {-1, 2, 1, 0, "nonnegative"}, + {0, -1, 1, 0, "nonnegative"}, + {1, 0, 1, 0, "nonnegative"}, + {0, 5, 1, 0, "exceeds"}, + {3, 2, 1, 0, "exceeds"}, + {0, 2, 0, 0, "requires"}, + {0, 2, 1, 2, "requires"}}) { + OpTester tester("BranchwiseRMSNorm", 2, kMSDomain); + tester.AddAttribute("input_start", start); + tester.AddAttribute("input_width", width); + tester.AddAttribute("use_onnx_arithmetic", mode); + tester.AddAttribute("num_branches", branches); + tester.AddInput("X", {1, 1, 4}, {1, 2, 3, 4}); + tester.AddOptionalInputEdge(); + tester.AddOutput("Y", {1, 1, 2}, {0, 0}); + std::vector> providers; + providers.push_back(DefaultCudaExecutionProvider()); + tester.Run(OpTester::ExpectResult::kExpectFailure, message, {}, nullptr, &providers); + } +} + +TEST(HyperConnectionOpsTest, BranchwiseRMSNormPackedEmptyWithoutScale) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA is required for packed RMS"; + OpTester tester("BranchwiseRMSNorm", 2, kMSDomain); + tester.AddAttribute("use_onnx_arithmetic", 1); + tester.AddAttribute("input_start", 3); + tester.AddAttribute("input_width", 2); + tester.AddInput("X", {0, 3, 8}, {}); + tester.AddOptionalInputEdge(); + tester.AddOutput("Y", {0, 3, 2}, {}); + tester.AddOutput("Y_fp32", {0, 3, 2}, {}); + std::vector> providers; + providers.push_back(DefaultCudaExecutionProvider()); + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); +} + +TEST(HyperConnectionOpsTest, BranchwiseRMSNormPackedWithoutScale) { + if (!DefaultCudaExecutionProvider()) GTEST_SKIP() << "CUDA is required for packed RMS"; + OpTester tester("BranchwiseRMSNorm", 2, kMSDomain); + tester.AddAttribute("use_onnx_arithmetic", 1); + tester.AddAttribute("input_start", 1); + tester.AddAttribute("input_width", 2); + tester.AddAttribute("epsilon", 0.0f); + tester.AddInput("X", {1, 2, 4}, {99, 3, 4, 99, 99, 0, 2, 99}); + tester.AddOptionalInputEdge(); + tester.AddOutput("Y", {1, 2, 2}, + {3.0f / std::sqrt(12.5f), 4.0f / std::sqrt(12.5f), 0, 2.0f / std::sqrt(2.0f)}); + std::vector> providers; + providers.push_back(DefaultCudaExecutionProvider()); + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); +} + +TEST(HyperConnectionOpsTest, BranchwiseRMSNormPackedDoesNotFallBackToCpu) { + OpTester tester("BranchwiseRMSNorm", 2, kMSDomain); + tester.AddAttribute("use_onnx_arithmetic", 1); + tester.AddAttribute("input_width", 2); + tester.AddInput("X", {1, 1, 4}, {1, 2, 3, 4}); + tester.AddOptionalInputEdge(); + tester.AddOutput("Y", {1, 1, 2}, {0, 0}); + std::vector> providers; + providers.push_back(DefaultCpuExecutionProvider()); + tester.Run(OpTester::ExpectResult::kExpectFailure, "Could not find an implementation for BranchwiseRMSNorm(2)", + {}, nullptr, &providers); +} + +TEST(HyperConnectionOpsTest, BranchwiseRMSNormStandardCpuBranchScale) { + OpTester tester("BranchwiseRMSNorm", 1, kMSDomain); + tester.AddAttribute("use_onnx_arithmetic", 1); + tester.AddAttribute("num_branches", 2); + tester.AddAttribute("epsilon", 0.0f); + tester.AddInput("X", {1, 4}, {3.0f, 4.0f, 0.0f, 2.0f}); + tester.AddInput("scale", {2, 2}, {2.0f, 0.5f, 0.25f, 3.0f}); + tester.AddOutput("Y", {1, 4}, + {3.0f / std::sqrt(12.5f) * 2.0f, 4.0f / std::sqrt(12.5f) * 0.5f, + 0.0f, 2.0f / std::sqrt(2.0f) * 3.0f}); + std::vector> providers; + providers.push_back(DefaultCpuExecutionProvider()); + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); +} + +TEST(HyperConnectionOpsTest, BranchwiseRMSNormRejectsInvalidArithmeticMode) { + for (auto& provider : AvailableProviders()) { + if (provider->Type() == kWebGpuExecutionProvider) continue; + OpTester tester("BranchwiseRMSNorm", 1, kMSDomain); + tester.AddAttribute("use_onnx_arithmetic", 2); + tester.AddInput("X", {1, 2}, {1.0f, 2.0f}); + tester.AddOptionalInputEdge(); + tester.AddOutput("Y", {1, 2}, {0.0f, 0.0f}); + std::vector> providers; + providers.push_back(std::move(provider)); + tester.Run(OpTester::ExpectResult::kExpectFailure, "use_onnx_arithmetic must be 0 or 1", + {}, nullptr, &providers); + } +} + TEST(HyperConnectionOpsTest, ScaledSiLU) { RunOnAvailableProviders("ScaledSiLU", [&](OpTester& tester) { tester.AddAttribute("alpha", 0.5f); diff --git a/onnxruntime/test/python/transformers/test_hyper_connection_low_rank_mix.py b/onnxruntime/test/python/transformers/test_hyper_connection_low_rank_mix.py new file mode 100644 index 0000000000000..0709465b55125 --- /dev/null +++ b/onnxruntime/test/python/transformers/test_hyper_connection_low_rank_mix.py @@ -0,0 +1,311 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +"""Numeric check for com.microsoft.HyperConnectionLowRankMix against a torch reference. + +The reference is the unfused subgraph the operator replaces, + + [HyperConnectionPostMix] -> BranchwiseRMSNorm -> MatMul(down) -> ScaledSiLU -> MatMul(up) -> Sigmoid + -> HyperConnectionPreMix, and MatMul(inject) -> Div -> Sigmoid -> Mul, + +evaluated with float32 math and rounded to the activation type at the points the standard operators round +(see the schema documentation), with a correctly rounded sigmoid. The fused path (few rows, 16-bit types) and +the cuBLAS based path (many rows, float32, shapes the fused kernels do not support) are both covered. +""" + +# Tensor dimensions keep the upper-case names the schema doc gives them. +# ruff: noqa: N806 + +import unittest + +import numpy as np +import torch +from onnx import TensorProto as TP # noqa: N817 +from onnx import helper + +import onnxruntime as ort + +DTYPES = { + TP.FLOAT: torch.float32, + TP.FLOAT16: torch.float16, + TP.BFLOAT16: torch.bfloat16, +} + +# (num_branches, hidden, rank): the real Qwen3.8-Flash-Next block first, then small and irregular shapes. +REAL_SHAPE = (4, 2560, 320) +SMALL_SHAPES = [ + (4, 64, 16), + (2, 32, 8), + (8, 64, 32), + (1, 48, 8), + (4, 1288, 40), # two K slices per branch in the fused path + (4, 4096, 512), # four K slices per branch and two weight chunks in the up kernel + (3, 40, 12), # not supported by the fused kernels (odd branch count, rank % 8 != 0) + (4, 100, 16), # hidden % 8 != 0 +] +ATTRS = dict(epsilon=1e-6, reduction_scale=0.25, silu_alpha=0.25, gate_div=4.0, gate_mul=2.0) + + +def has_cuda(): + return "CUDAExecutionProvider" in ort.get_available_providers() + + +def build_model(io_elem, C, H, R, prologue, attrs): + """Graph with float32 inputs/outputs; Casts move the activations to the type under test.""" + F = TP.FLOAT + K = C * H + node_inputs = ["streams", "norm_scale", "down_weight", "up_weight", "inject_weight"] + graph_inputs = [ + ("streams", ["rows", K]), + ("norm_scale", [K]), + ("down_weight", [K, R]), + ("up_weight", [R, K]), + ("inject_weight", [K, C]), + ] + if prologue: + node_inputs += ["prev_block_output", "prev_post_gate"] + graph_inputs += [("prev_block_output", ["rows", H]), ("prev_post_gate", ["rows", C])] + outputs = ["pre_output", "post_gate", "streams_out"] + + nodes = [helper.make_node("Cast", [n], [n + "_t"], to=io_elem, name=n + "/cast") for n, _ in graph_inputs] + nodes.append( + helper.make_node( + "HyperConnectionLowRankMix", + [n + "_t" for n in node_inputs], + [n + "_t" for n in outputs], + domain="com.microsoft", + num_branches=C, + **attrs, + ) + ) + nodes += [helper.make_node("Cast", [n + "_t"], [n], to=F, name=n + "/castback") for n in outputs] + graph = helper.make_graph( + nodes, + "hc_low_rank_mix", + [helper.make_tensor_value_info(n, F, shape) for n, shape in graph_inputs], + [ + helper.make_tensor_value_info("pre_output", F, ["rows", H]), + helper.make_tensor_value_info("post_gate", F, ["rows", C]), + helper.make_tensor_value_info("streams_out", F, ["rows", K]), + ], + ) + model = helper.make_model( + graph, opset_imports=[helper.make_opsetid("", 17), helper.make_opsetid("com.microsoft", 1)] + ) + model.ir_version = 10 + return model + + +def make_inputs(dtype, rows, C, H, R, prologue, seed): + g = torch.Generator().manual_seed(seed) + K = C * H + + def rnd(*shape, scale=1.0, shift=0.0): + return (torch.randn(*shape, generator=g) * scale + shift).to(dtype).float() + + inputs = { + "streams": rnd(rows, K), + "norm_scale": rnd(K, scale=0.1, shift=1.0), + "down_weight": rnd(K, R, scale=K**-0.5), + "up_weight": rnd(R, K, scale=2.0 * R**-0.5), + "inject_weight": rnd(K, C, scale=K**-0.5), + } + if prologue: + inputs["prev_block_output"] = rnd(rows, H) + inputs["prev_post_gate"] = (torch.rand(rows, C, generator=g) * 2.0).to(dtype).float() + return inputs + + +def reference(inputs, dtype, C, H, attrs): + def rd(t): + return t.to(dtype).float() + + def sigmoid(t): + return torch.sigmoid(t.double()).float() + + x = inputs["streams"] + rows = x.shape[0] + s = x.view(rows, C, H) + if "prev_block_output" in inputs: + s = rd(s + inputs["prev_post_gate"][:, :, None] * inputs["prev_block_output"][:, None, :]) + inv_rms = torch.rsqrt(s.pow(2).mean(-1, keepdim=True) + attrs["epsilon"]) + hn = rd(s * inv_rms * inputs["norm_scale"].view(C, H)) + hn_flat = hn.view(rows, C * H) + + down = rd(hn_flat @ inputs["down_weight"]) + z = rd(down * attrs["silu_alpha"]) + hidden = rd(z * rd(sigmoid(z))) + up = rd(hidden @ inputs["up_weight"]) + g = rd(sigmoid(up)).view(rows, C, H) + pre = rd(attrs["reduction_scale"] * (hn * g).sum(1)) + inject = rd(hn_flat @ inputs["inject_weight"]) + gate = rd(rd(sigmoid(rd(inject / attrs["gate_div"]))) * attrs["gate_mul"]) + return pre, gate, s.reshape(rows, C * H) + + +class SessionCache: + def __init__(self): + self.sessions = {} + + def get(self, io_elem, C, H, R, prologue, attrs): + key = (io_elem, C, H, R, prologue, tuple(sorted(attrs.items()))) + if key not in self.sessions: + model = build_model(io_elem, C, H, R, prologue, attrs) + so = ort.SessionOptions() + so.log_severity_level = 3 + self.sessions[key] = ort.InferenceSession( + model.SerializeToString(), so, providers=["CUDAExecutionProvider"] + ) + return self.sessions[key] + + +def tolerances(io_elem): + # One ulp of the output type plus slack for the different fp32 accumulation order flipping a rounding. + return {TP.FLOAT: (1e-3, 1e-5), TP.FLOAT16: (4e-3, 2e-3), TP.BFLOAT16: (3e-2, 8e-3)}[io_elem] + + +@unittest.skipUnless(has_cuda(), "HyperConnectionLowRankMix is a CUDA-only kernel") +class TestHyperConnectionLowRankMix(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.cache = SessionCache() + + def _check(self, io_elem, rows, shape, prologue, attrs=ATTRS, seed=0): + C, H, R = shape + dtype = DTYPES[io_elem] + inputs = make_inputs(dtype, rows, C, H, R, prologue, seed) + session = self.cache.get(io_elem, C, H, R, prologue, attrs) + pre, gate, streams_out = session.run(None, {k: v.numpy() for k, v in inputs.items()}) + exp_pre, exp_gate, exp_streams = reference(inputs, dtype, C, H, attrs) + rtol, atol = tolerances(io_elem) + msg = f"dtype={dtype} rows={rows} shape={shape} prologue={prologue}" + np.testing.assert_allclose(pre, exp_pre.numpy(), rtol=rtol, atol=atol, err_msg="pre_output " + msg) + np.testing.assert_allclose(gate, exp_gate.numpy(), rtol=rtol, atol=atol, err_msg="post_gate " + msg) + # The prologue is a single fp32 operation with one rounding: it matches exactly (up to fma contraction). + np.testing.assert_allclose(streams_out, exp_streams.numpy(), rtol=rtol, atol=atol, err_msg="streams_out " + msg) + return pre, gate, streams_out + + def test_real_shape(self): + for io_elem in DTYPES: + for rows in (1, 2, 4, 16, 128): + for prologue in (False, True): + with self.subTest(io_elem=io_elem, rows=rows, prologue=prologue): + self._check(io_elem, rows, REAL_SHAPE, prologue) + + def test_small_shapes(self): + for io_elem in DTYPES: + for shape in SMALL_SHAPES: + for rows in (1, 2, 16): + for prologue in (False, True): + with self.subTest(io_elem=io_elem, shape=shape, rows=rows, prologue=prologue): + self._check(io_elem, rows, shape, prologue, seed=1) + + def test_default_attributes(self): + attrs = dict(epsilon=1e-5, reduction_scale=1.0, silu_alpha=1.0, gate_div=1.0, gate_mul=1.0) + for io_elem in DTYPES: + for rows in (1, 3): + with self.subTest(io_elem=io_elem, rows=rows): + self._check(io_elem, rows, (4, 64, 16), True, attrs=attrs, seed=2) + + def test_other_gate_constants(self): + attrs = dict(epsilon=1e-6, reduction_scale=0.5, silu_alpha=0.75, gate_div=3.0, gate_mul=1.5) + for io_elem in (TP.FLOAT16, TP.BFLOAT16): + for rows in (1, 2): + with self.subTest(io_elem=io_elem, rows=rows): + self._check(io_elem, rows, REAL_SHAPE, True, attrs=attrs, seed=3) + + def test_chain_of_blocks(self): + """Two consecutive calls through the prologue equal the explicit post-mix between them.""" + C, H, R = REAL_SHAPE + for io_elem in (TP.FLOAT16, TP.BFLOAT16): + dtype = DTYPES[io_elem] + with self.subTest(io_elem=io_elem): + inputs = make_inputs(dtype, 2, C, H, R, False, seed=4) + first = self.cache.get(io_elem, C, H, R, False, ATTRS) + second = self.cache.get(io_elem, C, H, R, True, ATTRS) + _, gate, streams = first.run(None, {k: v.numpy() for k, v in inputs.items()}) + # Use the block "output" as an independent tensor (a stand-in for the attention/MLP result). + block_out = torch.randn(2, H, generator=torch.Generator().manual_seed(8)).to(dtype).float().numpy() + chained = dict(inputs) + chained["streams"] = streams + chained["prev_block_output"] = block_out + chained["prev_post_gate"] = gate + out = second.run(None, {k: np.asarray(v) for k, v in chained.items()}) + ref_inputs = {k: torch.from_numpy(np.asarray(v)) for k, v in chained.items()} + exp = reference(ref_inputs, dtype, C, H, ATTRS) + rtol, atol = tolerances(io_elem) + np.testing.assert_allclose(out[0], exp[0].numpy(), rtol=rtol, atol=atol) + np.testing.assert_allclose(out[1], exp[1].numpy(), rtol=rtol, atol=atol) + + def test_fused_rows_are_independent(self): + """The fused kernels use a fixed reduction order: a row gives the same bits at any batch size.""" + C, H, R = REAL_SHAPE + for io_elem in (TP.FLOAT16, TP.BFLOAT16): + dtype = DTYPES[io_elem] + with self.subTest(io_elem=io_elem): + inputs = make_inputs(dtype, 4, C, H, R, True, seed=5) + batch = self.cache.get(io_elem, C, H, R, True, ATTRS).run( + None, {k: v.numpy() for k, v in inputs.items()} + ) + for r in range(4): + row_inputs = { + k: (v[r : r + 1] if k in ("streams", "prev_block_output", "prev_post_gate") else v).numpy() + for k, v in inputs.items() + } + single = self.cache.get(io_elem, C, H, R, True, ATTRS).run(None, row_inputs) + for b, s in zip(batch, single, strict=True): + np.testing.assert_array_equal(b[r : r + 1], s) + + def test_no_prologue_streams_out_is_passthrough(self): + C, H, R = 4, 64, 16 + inputs = make_inputs(torch.bfloat16, 2, C, H, R, False, seed=6) + session = self.cache.get(TP.BFLOAT16, C, H, R, False, ATTRS) + _, _, streams_out = session.run(None, {k: v.numpy() for k, v in inputs.items()}) + np.testing.assert_array_equal(streams_out, inputs["streams"].numpy()) + + def test_cuda_graph_capture_and_replay(self): + """No host synchronization or unmanaged allocation: the op can be captured and replayed.""" + C, H, R = REAL_SHAPE + for io_elem in (TP.FLOAT16, TP.BFLOAT16): + for rows in (1, 2): + with self.subTest(io_elem=io_elem, rows=rows): + dtype = DTYPES[io_elem] + inputs = make_inputs(dtype, rows, C, H, R, True, seed=9) + model = build_model(io_elem, C, H, R, True, ATTRS) + so = ort.SessionOptions() + so.log_severity_level = 3 + graph_session = ort.InferenceSession( + model.SerializeToString(), + so, + providers=[("CUDAExecutionProvider", {"enable_cuda_graph": "1"})], + ) + binding = graph_session.io_binding() + keep = [] + for name, value in inputs.items(): + ortvalue = ort.OrtValue.ortvalue_from_numpy(value.numpy(), "cuda", 0) + keep.append(ortvalue) + binding.bind_ortvalue_input(name, ortvalue) + shapes = {"pre_output": (rows, H), "post_gate": (rows, C), "streams_out": (rows, C * H)} + outputs = {} + for name, shape in shapes.items(): + outputs[name] = ort.OrtValue.ortvalue_from_shape_and_type(shape, np.float32, "cuda", 0) + binding.bind_ortvalue_output(name, outputs[name]) + graph_session.run_with_iobinding(binding) # captures + graph_session.run_with_iobinding(binding) # replays + expected = self.cache.get(io_elem, C, H, R, True, ATTRS).run( + None, {k: v.numpy() for k, v in inputs.items()} + ) + for name, exp in zip(shapes, expected, strict=True): + np.testing.assert_array_equal(outputs[name].numpy(), exp, err_msg=name) + + def test_zero_rows(self): + C, H, R = 4, 64, 16 + inputs = make_inputs(torch.float16, 0, C, H, R, False, seed=7) + session = self.cache.get(TP.FLOAT16, C, H, R, False, ATTRS) + pre, gate, streams_out = session.run(None, {k: v.numpy() for k, v in inputs.items()}) + self.assertEqual(pre.shape, (0, H)) + self.assertEqual(gate.shape, (0, C)) + self.assertEqual(streams_out.shape, (0, C * H)) + + +if __name__ == "__main__": + unittest.main()