diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md
index eb97d2a806de7..4fc6d39e517dc 100644
--- a/docs/ContribOperators.md
+++ b/docs/ContribOperators.md
@@ -4976,7 +4976,7 @@ This version of the operator has been available since version 1 of the 'com.micr
### **com.microsoft.MoEPacked**
- Mixture-of-Experts layer consuming 3D block-packed gate, up, and down expert weight blobs. CUDA execution requires weights in device memory; host-backed external initializers are copied to the device without changing their packed layout.
+ Mixture-of-Experts layer consuming 3D block-packed gate and down expert weight blobs, plus up expert weights for SwiGLU. A separate router can supply preselected expert IDs and FP32 per-token mixture coefficients; these coefficients are consumed without normalization in the supplied expert order. CUDA execution requires weights in device memory; host-backed external initializers are copied to the device without changing their packed layout.
Supported format IDs and their stable byte layouts are defined in the [packed-weight format contract](https://github.com/microsoft/onnxruntime/blob/main/docs/PackedWeightFormats.md).
#### Version
@@ -4987,32 +4987,36 @@ This version of the operator has been available since version 1 of the 'com.micr
- activation_type : string
-- Expert activation function ('swiglu').
+- Expert activation function: 'swiglu' or 'relu2' (squared ReLU; 'relu_squared' is an alias). relu2 requires IQ2_XXS (type 16) gate and down weights and omitted up weights.
- down_type : int (required)
- Stable packed-format ID of down_experts_weights. Only contract-listed IDs are supported.
- gate_type : int (required)
- Stable packed-format ID of gate_experts_weights. Only contract-listed IDs are supported.
- k : int (required)
-- Number of selected experts per token.
+- Number of selected experts per token: 1..16 for internal softmax routing, 1..32 for preselected routing; at most num_experts.
- normalize_routing_weights : int
-- Whether to normalize top-k routing probabilities.
-- up_type : int (required)
-- Stable packed-format ID of up_experts_weights. Only contract-listed IDs are supported.
+- Whether to renormalize selected softmax probabilities. Must be 0 for preselected routing.
+- up_type : int
+- Stable packed-format ID of up_experts_weights; required for swiglu. Only contract-listed IDs are supported.
-#### Inputs
+#### Inputs (5 - 7)
- input : T
- Input hidden states [..., hidden_size].
-- router_probs : T
-- Router logits [num_rows, num_experts].
+- router_probs (optional) : T
+- Router logits [num_rows, num_experts]; omit for preselected routing.
- gate_experts_weights : T2
- 3D packed gate weight blob [num_experts, inter_size, gate_row_bytes].
-- up_experts_weights : T2
-- 3D packed up weight blob [num_experts, inter_size, up_row_bytes].
+- up_experts_weights (optional) : T2
+- 3D packed up weight blob [num_experts, inter_size, up_row_bytes]; required for swiglu, omit for relu2.
- down_experts_weights : T2
- 3D packed down weight blob [num_experts, hidden_size, down_row_bytes].
+- expert_indices (optional) : Tind
+- Preselected distinct expert IDs [num_rows, k], k <= 32. Requires expert_weights, no router_probs and normalize_routing_weights=0. Invalid or duplicate IDs and nonfinite weights poison only their token row.
+- expert_weights (optional) : Tweight
+- Preselected FP32 per-token mixture coefficients [num_rows, k], preserved without normalization and mixed in supplied order.
#### Outputs
@@ -5029,6 +5033,10 @@ This version of the operator has been available since version 1 of the 'com.micr
Activation type.
T2 : tensor(uint8)
Packed byte storage.
+Tind : tensor(int32)
+Preselected expert index type.
+Tweight : tensor(float)
+Preselected mixture coefficient type.
diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md
index a6fd24bf09aca..31448449db601 100644
--- a/docs/OperatorKernels.md
+++ b/docs/OperatorKernels.md
@@ -1159,7 +1159,7 @@ The **OpSet Version** column uses the following notation:
|MatMulPacked|*in* A:**T1**
*in* B:**T2**
*in* global_scale:**Tscale**
*out* Y:**T1**|1+|**T1** = tensor(bfloat16), tensor(float16)
**T2** = tensor(uint8)
**Tscale** = tensor(float)|
|MoE|*in* input:**T**
*in* router_probs:**T**
*in* fc1_experts_weights:**T**
*in* fc1_experts_bias:**T**
*in* fc2_experts_weights:**T**
*in* fc2_experts_bias:**T**
*in* fc3_experts_weights:**T**
*in* fc3_experts_bias:**T**
*out* output:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)|
|MoEGroupRouter|*in* logits:**T**
*in* expert_bias:**TF**
*out* router_probs:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)
**TF** = tensor(float)|
-|MoEPacked|*in* input:**T**
*in* router_probs:**T**
*in* gate_experts_weights:**T2**
*in* up_experts_weights:**T2**
*in* down_experts_weights:**T2**
*out* output:**T**|1+|**T** = tensor(bfloat16), tensor(float16)
**T2** = tensor(uint8)|
+|MoEPacked|*in* input:**T**
*in* router_probs:**T**
*in* gate_experts_weights:**T2**
*in* up_experts_weights:**T2**
*in* down_experts_weights:**T2**
*in* expert_indices:**Tind**
*in* expert_weights:**Tweight**
*out* output:**T**|1+|**T** = tensor(bfloat16), tensor(float16)
**T2** = tensor(uint8)
**Tind** = tensor(int32)
**Tweight** = tensor(float)|
|MultiHeadAttention|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* bias:**T**
*in* key_padding_mask:**M**
*in* attention_bias:**T**
*in* past_key:**T**
*in* past_value:**T**
*in* past_sequence_length:**M**
*in* cache_indirection:**M**
*out* output:**T**
*out* present_key:**T**
*out* present_value:**T**
*out* qk:**QK**|1+|**QK** = tensor(bfloat16), tensor(float), tensor(float16)
**T** = tensor(bfloat16), tensor(float), tensor(float16)|
|NGramHashMapping|*in* input_ids:**M**
*in* multipliers:**M**
*in* vocab_sizes:**M**
*in* past_ids:**M**
*in* head_offsets:**M**
*in* eos_token_id:**M**
*in* segment_ids:**tensor(int32)**
*out* hash_ids:**M**
*out* present_ids:**M**|1+|**M** = tensor(int32), tensor(int64)|
|NGramRepeatBlock|*in* input_ids:**Tid**
*in* scores:**T**
*out* scores_out:**T**|1+|**T** = tensor(float)
**Tid** = tensor(int64)|
diff --git a/onnxruntime/contrib_ops/cuda/moe/moe_packed.cc b/onnxruntime/contrib_ops/cuda/moe/moe_packed.cc
index 623c5f6c7e211..3e996fd81694c 100644
--- a/onnxruntime/contrib_ops/cuda/moe/moe_packed.cc
+++ b/onnxruntime/contrib_ops/cuda/moe/moe_packed.cc
@@ -4,6 +4,7 @@
#include "contrib_ops/cuda/moe/moe_packed.h"
#include
+#include
#include
#include "contrib_ops/cpu/quantization/gguf_quant_tables.h"
@@ -26,15 +27,17 @@ bool PackedTensorBytesFit(int64_t rows, int64_t columns, size_t element_bytes) {
}
} // namespace
-#define REGISTER_MOE_PACKED_CUDA(T) \
- ONNX_OPERATOR_TYPED_KERNEL_EX( \
- MoEPacked, \
- kMSDomain, 1, \
- T, \
- kCudaExecutionProvider, \
- (*KernelDefBuilder::Create()) \
- .TypeConstraint("T", DataTypeImpl::GetTensorType()) \
- .TypeConstraint("T2", DataTypeImpl::GetTensorType()), \
+#define REGISTER_MOE_PACKED_CUDA(T) \
+ ONNX_OPERATOR_TYPED_KERNEL_EX( \
+ MoEPacked, \
+ kMSDomain, 1, \
+ T, \
+ kCudaExecutionProvider, \
+ (*KernelDefBuilder::Create()) \
+ .TypeConstraint("T", DataTypeImpl::GetTensorType()) \
+ .TypeConstraint("T2", DataTypeImpl::GetTensorType()) \
+ .TypeConstraint("Tind", DataTypeImpl::GetTensorType()) \
+ .TypeConstraint("Tweight", DataTypeImpl::GetTensorType()), \
MoEPacked);
REGISTER_MOE_PACKED_CUDA(MLFloat16)
@@ -44,17 +47,24 @@ REGISTER_MOE_PACKED_CUDA(BFloat16)
template
MoEPacked::MoEPacked(const OpKernelInfo& info) : CudaKernel(info) {
- ORT_ENFORCE(info.GetAttr("k", &k_).IsOK() && k_ > 0 && k_ <= 16,
- "MoEPacked requires k in 1..16");
+ ORT_ENFORCE(info.GetAttr("k", &k_).IsOK() && k_ > 0 && k_ <= 32,
+ "MoEPacked requires k in 1..32; internally routed mode requires k <= 16");
normalize_routing_weights_ = info.GetAttrOrDefault("normalize_routing_weights", 1);
activation_type_ = info.GetAttrOrDefault("activation_type", "swiglu");
- ORT_ENFORCE(activation_type_ == "swiglu", "MoEPacked only supports activation_type='swiglu'");
+ if (activation_type_ == "relu_squared") activation_type_ = "relu2";
+ ORT_ENFORCE(activation_type_ == "swiglu" || activation_type_ == "relu2",
+ "MoEPacked supports activation_type='swiglu' or 'relu2' ('relu_squared' alias)");
ORT_ENFORCE(info.GetAttr("gate_type", &gate_type_).IsOK(), "MoEPacked requires gate_type");
- ORT_ENFORCE(info.GetAttr("up_type", &up_type_).IsOK(), "MoEPacked requires up_type");
+ up_type_ = info.GetAttrOrDefault("up_type", -1);
+ ORT_ENFORCE(activation_type_ != "swiglu" || up_type_ >= 0, "MoEPacked swiglu requires up_type");
ORT_ENFORCE(info.GetAttr("down_type", &down_type_).IsOK(), "MoEPacked requires down_type");
ORT_ENFORCE(gguf_quant::IsGgmlTypeSupported(gate_type_), "MoEPacked: unsupported gate_type=", gate_type_);
- ORT_ENFORCE(gguf_quant::IsGgmlTypeSupported(up_type_), "MoEPacked: unsupported up_type=", up_type_);
+ ORT_ENFORCE(activation_type_ != "swiglu" || gguf_quant::IsGgmlTypeSupported(up_type_),
+ "MoEPacked: unsupported up_type=", up_type_);
ORT_ENFORCE(gguf_quant::IsGgmlTypeSupported(down_type_), "MoEPacked: unsupported down_type=", down_type_);
+ ORT_ENFORCE(activation_type_ != "relu2" ||
+ (gate_type_ == gguf_quant::kGgmlTypeIQ2_XXS && down_type_ == gguf_quant::kGgmlTypeIQ2_XXS),
+ "MoEPacked relu2 currently requires IQ2_XXS gate and down weights");
if constexpr (std::is_same_v) {
ORT_ENFORCE(GetDeviceProp().major >= 8, "MoEPacked BF16 requires compute capability 8.0 or higher");
ORT_THROW_IF_ERROR(GetMoEPackedBf16KernelSupport(GetDeviceId(), bf16_tc_supported_));
@@ -68,43 +78,59 @@ Status MoEPacked::ComputeInternal(OpKernelContext* context) const {
const Tensor* gate_w = context->Input(2);
const Tensor* up_w = context->Input(3);
const Tensor* down_w = context->Input(4);
+ const Tensor* indices = context->InputCount() > 5 ? context->Input(5) : nullptr;
+ const Tensor* weights = context->InputCount() > 6 ? context->Input(6) : nullptr;
+ const bool preselected = indices || weights;
+ const bool relu_squared = activation_type_ == "relu2";
+ ORT_RETURN_IF_NOT(input && gate_w && down_w, "MoEPacked requires input, gate and down weights");
+ ORT_RETURN_IF_NOT(preselected ? (indices && weights && !router_probs && normalize_routing_weights_ == 0)
+ : (router_probs && k_ <= 16),
+ "MoEPacked preselected routing requires both indices and weights, no router_probs and "
+ "normalize_routing_weights=0; internal routing requires router_probs and k <= 16");
+ ORT_RETURN_IF_NOT(relu_squared ? !up_w : up_w != nullptr,
+ "MoEPacked swiglu requires up weights; relu2 requires omitted up weights");
const auto& in_shape = input->Shape();
- const auto& rp_shape = router_probs->Shape();
const auto& gw_shape = gate_w->Shape();
- const auto& uw_shape = up_w->Shape();
const auto& dw_shape = down_w->Shape();
ORT_RETURN_IF_NOT(in_shape.NumDimensions() >= 2, "MoEPacked input must have rank >= 2");
- ORT_RETURN_IF_NOT(rp_shape.NumDimensions() == 2, "MoEPacked router_probs must be 2D [num_rows, num_experts]");
- ORT_RETURN_IF_NOT(gw_shape.NumDimensions() == 3 && uw_shape.NumDimensions() == 3 && dw_shape.NumDimensions() == 3,
+ ORT_RETURN_IF_NOT(gw_shape.NumDimensions() == 3 && dw_shape.NumDimensions() == 3 &&
+ (!up_w || up_w->Shape().NumDimensions() == 3),
"MoEPacked expert weights must be 3D [num_experts, out_dim, row_bytes]");
const size_t last_axis = in_shape.NumDimensions() - 1;
const int64_t hidden_size = in_shape[last_axis];
const int64_t num_rows = in_shape.SizeToDimension(last_axis);
- const int64_t num_experts = rp_shape[1];
+ const int64_t num_experts = gw_shape[0];
const int64_t inter_size = gw_shape[1];
ORT_RETURN_IF_NOT(hidden_size > 0 && hidden_size <= kMoEPackedMaxDimension &&
inter_size > 0 && inter_size <= kMoEPackedMaxDimension,
"MoEPacked requires hidden_size and inter_size in 1..", kMoEPackedMaxDimension);
ORT_RETURN_IF_NOT(num_rows >= 0, "MoEPacked requires a nonnegative row count");
- ORT_RETURN_IF_NOT(rp_shape[0] == num_rows, "MoEPacked router_probs row count mismatch");
- ORT_RETURN_IF_NOT(num_experts > 0 && num_experts <= 512, "MoEPacked supports up to 512 experts");
- ORT_RETURN_IF_NOT(k_ <= num_experts, "MoEPacked requires k <= num_experts");
- ORT_RETURN_IF_NOT(gw_shape[0] == num_experts && uw_shape[0] == num_experts && dw_shape[0] == num_experts,
+ ORT_RETURN_IF_NOT(num_experts >= k_ && num_experts <= 512, "MoEPacked requires k <= num_experts <= 512");
+ if (preselected) {
+ ORT_RETURN_IF_NOT(indices->Shape() == TensorShape({num_rows, k_}) && weights->Shape() == indices->Shape(),
+ "MoEPacked preselected routes require indices and weights [num_rows, k]");
+ } else {
+ ORT_RETURN_IF_NOT(router_probs->Shape() == TensorShape({num_rows, num_experts}),
+ "MoEPacked router_probs must be [num_rows, num_experts]");
+ }
+ ORT_RETURN_IF_NOT((!up_w || up_w->Shape()[0] == num_experts) && dw_shape[0] == num_experts,
"MoEPacked expert count mismatch");
- ORT_RETURN_IF_NOT(uw_shape[1] == inter_size && dw_shape[1] == hidden_size,
+ ORT_RETURN_IF_NOT((!up_w || up_w->Shape()[1] == inter_size) && dw_shape[1] == hidden_size,
"MoEPacked expert inner/outer dimension mismatch");
- ORT_RETURN_IF_NOT((hidden_size % 32) == 0 && (inter_size % 32) == 0,
+ ORT_RETURN_IF_NOT(hidden_size > 0 && inter_size > 0 &&
+ (hidden_size % 32) == 0 && (inter_size % 32) == 0,
"MoEPacked requires hidden_size and inter_size to be multiples of 32");
const int64_t exp_gate_rb = gguf_quant::GgmlRowBytes(static_cast(gate_type_), hidden_size);
- const int64_t exp_up_rb = gguf_quant::GgmlRowBytes(static_cast(up_type_), hidden_size);
+ const int64_t exp_up_rb = relu_squared ? exp_gate_rb : gguf_quant::GgmlRowBytes(static_cast(up_type_), hidden_size);
const int64_t exp_down_rb = gguf_quant::GgmlRowBytes(static_cast(down_type_), inter_size);
ORT_RETURN_IF_NOT(exp_gate_rb > 0 && exp_up_rb > 0 && exp_down_rb > 0 &&
- gw_shape[2] == exp_gate_rb && uw_shape[2] == exp_up_rb && dw_shape[2] == exp_down_rb,
+ gw_shape[2] == exp_gate_rb && (!up_w || up_w->Shape()[2] == exp_up_rb) &&
+ dw_shape[2] == exp_down_rb,
"MoEPacked expert weight row_bytes mismatch");
constexpr bool is_bf16 = std::is_same_v;
ORT_RETURN_IF_NOT(PackedTensorBytesFit(num_rows, hidden_size, sizeof(T)) &&
@@ -132,8 +158,9 @@ Status MoEPacked::ComputeInternal(OpKernelContext* context) const {
max_params.k = static_cast(k_);
max_params.normalize_routing_weights = (normalize_routing_weights_ != 0);
max_params.bf16_tc_supported = bf16_tc_supported_;
+ max_params.relu_squared = relu_squared;
max_params.gate_type = static_cast(gate_type_);
- max_params.up_type = static_cast(up_type_);
+ max_params.up_type = static_cast(relu_squared ? gate_type_ : up_type_);
max_params.down_type = static_cast(down_type_);
max_params.gate_row_bytes = exp_gate_rb;
max_params.up_row_bytes = exp_up_rb;
@@ -146,7 +173,7 @@ Status MoEPacked::ComputeInternal(OpKernelContext* context) const {
auto workspace = GetScratchBuffer(ws_bytes, GetComputeStream(context));
const uint8_t* in_bytes = static_cast(input->DataRaw());
- const uint8_t* rp_bytes = static_cast(router_probs->DataRaw());
+ const uint8_t* rp_bytes = router_probs ? static_cast(router_probs->DataRaw()) : nullptr;
uint8_t* out_bytes = static_cast(output->MutableDataRaw());
for (int64_t r0 = 0; r0 < num_rows;) {
@@ -155,13 +182,15 @@ Status MoEPacked::ComputeInternal(OpKernelContext* context) const {
chunk_params.num_rows = static_cast(chunk_rows);
const size_t input_offset = SafeInt(r0) * hidden_size * sizeof(T);
const size_t router_offset = SafeInt(r0) * num_experts * sizeof(T);
+ chunk_params.expert_indices = indices ? indices->Data() + r0 * k_ : nullptr;
+ chunk_params.expert_weights = weights ? weights->Data() + r0 * k_ : nullptr;
ORT_RETURN_IF_ERROR(LaunchMoEPackedCuda(
Stream(context),
is_bf16,
in_bytes + input_offset,
- rp_bytes + router_offset,
+ rp_bytes ? rp_bytes + router_offset : nullptr,
gate_w->Data(),
- up_w->Data(),
+ up_w ? up_w->Data() : gate_w->Data(),
down_w->Data(),
chunk_params,
workspace.get(),
diff --git a/onnxruntime/contrib_ops/cuda/moe/moe_packed.cu b/onnxruntime/contrib_ops/cuda/moe/moe_packed.cu
index 9691b64a1d385..ca39565bb40e0 100644
--- a/onnxruntime/contrib_ops/cuda/moe/moe_packed.cu
+++ b/onnxruntime/contrib_ops/cuda/moe/moe_packed.cu
@@ -25,6 +25,18 @@ namespace {
constexpr int kMoEPackedMaxExperts = 512;
constexpr int kMoEPackedMaxK = 16;
+template
+__device__ __forceinline__ uint4 MoEPackedLoadRouter8(const CudaT* input) {
+ if ((reinterpret_cast(input) & 15u) == 0) {
+ return __ldg(reinterpret_cast(input));
+ }
+ uint4 result;
+ CudaT* values = reinterpret_cast(&result);
+#pragma unroll
+ for (int j = 0; j < 8; ++j) values[j] = input[j];
+ return result;
+}
+
// Softmax + top-k over 512 experts for arbitrary num_rows (8 warps per CTA, 1 warp per row).
// Uses the exact same warp reduction and tie-breaking key as FusedSoftmaxTopKMapsKernel.
template
@@ -448,9 +460,9 @@ __global__ __launch_bounds__(KW * 32) void MoEPackedFc1DecodeWarpKSplitKernel(
}
const float g_val = __bfloat162float(__float2bfloat16_rn(g_raw));
const float u_val = __bfloat162float(__float2bfloat16_rn(u_raw));
- const float swiglu = u_val * (g_val / (1.0f + expf(-g_val)));
+ const float activation = u_val * (g_val / (1.0f + expf(-g_val)));
fc1_out[static_cast(s_routes[col]) * p.inter_size + j_out] =
- __float2bfloat16_rn(swiglu);
+ __float2bfloat16_rn(activation);
}
}
}
@@ -617,9 +629,9 @@ __global__ __launch_bounds__((kMaxSubTiles == 1) ? 128 : 256) void MoEPackedFc1S
const float u_raw = (h == 0) ? acc[st][2] : acc[st][3];
const float g_val = __bfloat162float(__float2bfloat16_rn(g_raw));
const float u_val = __bfloat162float(__float2bfloat16_rn(u_raw));
- const float swiglu = u_val * (g_val / (1.0f + expf(-g_val)));
+ const float activation = u_val * (g_val / (1.0f + expf(-g_val)));
const int route = s_routes[m];
- fc1_out[static_cast(route) * p.inter_size + j_out] = __float2bfloat16_rn(swiglu);
+ fc1_out[static_cast(route) * p.inter_size + j_out] = __float2bfloat16_rn(activation);
}
}
}
@@ -810,14 +822,16 @@ __global__ __launch_bounds__(256) void MoEPackedFc1SwiGLUSimtKernel(
const int sb32_idx = h_idx >> 1;
const int half16 = h_idx & 1;
float g_vals[16];
- float u_vals[16];
+ float u_vals[16] = {};
gguf_cuda::DecodeGgufSubBlock32(gate_row_ptr, p.gate_type, sb32_idx, half16, g_vals);
- gguf_cuda::DecodeGgufSubBlock32(up_row_ptr, p.up_type, sb32_idx, half16, u_vals);
- if constexpr (std::is_same_v) {
+ if (!p.relu_squared) {
+ gguf_cuda::DecodeGgufSubBlock32(up_row_ptr, p.up_type, sb32_idx, half16, u_vals);
+ }
+ if (p.relu_squared || std::is_same_v) {
#pragma unroll
for (int c = 0; c < 16; ++c) {
- g_vals[c] = __bfloat162float(__float2bfloat16_rn(g_vals[c]));
- u_vals[c] = __bfloat162float(__float2bfloat16_rn(u_vals[c]));
+ g_vals[c] = gguf_cuda::CudaTypeToFloat(gguf_cuda::FloatToCudaType(g_vals[c]));
+ u_vals[c] = gguf_cuda::CudaTypeToFloat(gguf_cuda::FloatToCudaType(u_vals[c]));
}
}
@@ -852,9 +866,10 @@ __global__ __launch_bounds__(256) void MoEPackedFc1SwiGLUSimtKernel(
for (int m = 0; m < count; ++m) {
const float g = gguf_cuda::CudaTypeToFloat(gguf_cuda::FloatToCudaType(gate_acc[m]));
const float u = gguf_cuda::CudaTypeToFloat(gguf_cuda::FloatToCudaType(up_acc[m]));
- const float swiglu = u * (g / (1.0f + expf(-g)));
+ const float positive = g < 0.f ? 0.f : g;
+ const float activation = p.relu_squared ? positive * positive : u * (g / (1.0f + expf(-g)));
fc1_out[static_cast(sorted_routes[tile.start_route + m]) * p.inter_size + j_out] =
- gguf_cuda::FloatToCudaType(swiglu);
+ gguf_cuda::FloatToCudaType(activation);
}
}
}
@@ -905,10 +920,10 @@ __global__ __launch_bounds__(256) void MoEPackedFc2SimtKernel(
const int half16 = h_idx & 1;
float d_vals[16];
gguf_cuda::DecodeGgufSubBlock32(down_row_ptr, p.down_type, sb32_idx, half16, d_vals);
- if constexpr (std::is_same_v) {
+ if (p.relu_squared || std::is_same_v) {
#pragma unroll
for (int c = 0; c < 16; ++c) {
- d_vals[c] = __bfloat162float(__float2bfloat16_rn(d_vals[c]));
+ d_vals[c] = gguf_cuda::CudaTypeToFloat(gguf_cuda::FloatToCudaType(d_vals[c]));
}
}
@@ -1005,6 +1020,21 @@ size_t AlignUpWorkspace(size_t count, size_t element_bytes = 1) {
return bytes;
}
+__global__ void MoEPackedPreselectedRoutesKernel(
+ const int32_t* indices, const float* weights, int* output_indices, float* output_weights,
+ int count, int k, int num_experts) {
+ const int route = static_cast(blockIdx.x) * blockDim.x + threadIdx.x;
+ if (route >= count) return;
+ const int expert = indices[route];
+ bool invalid = expert < 0 || expert >= num_experts || !isfinite(weights[route]);
+ for (int previous = (route / k) * k; previous < route; ++previous) {
+ invalid = invalid || indices[previous] == expert;
+ }
+ output_indices[route] = invalid ? 0 : expert;
+ // Preserve valid FP32 weights; NaN marks a malformed token without unsafe expert indexing.
+ output_weights[route] = invalid ? __int_as_float(0x7fffffff) : weights[route];
+}
+
void LaunchMoEPackedFc1DecodeBf16TcDispatch(
cudaStream_t stream,
dim3 grid,
@@ -1182,6 +1212,118 @@ void LaunchMoEPackedFc2Bf16TcDispatch(
}
}
+template
+__global__ __launch_bounds__(KW * 32) void MoEPackedReluIq2ContiguousKKernel(
+ const __nv_bfloat16* input,
+ const uint8_t* weights,
+ const int* sorted_routes,
+ const MoEPackedTile* tiles,
+ const int* d_num_tiles,
+ MoEPackedParams p,
+ __nv_bfloat16* output) {
+ const int tile_index = static_cast(blockIdx.y);
+ if (tile_index >= __ldg(d_num_tiles)) return;
+ const MoEPackedTile tile = tiles[tile_index];
+ const int first = static_cast(blockIdx.z) * 8;
+ const int remaining = tile.count - first;
+ if (remaining <= 0) return;
+ const int count = remaining < 8 ? remaining : 8;
+ constexpr int MT = 2;
+ __shared__ int routes[8];
+ __shared__ float reduction[KW][MT * 16][8];
+ const int tid = static_cast(threadIdx.x);
+ if (tid < 8) {
+ routes[tid] = tid < count ? __ldg(sorted_routes + tile.start_route + first + tid) : 0;
+ }
+ __syncthreads();
+ const int warp = tid >> 5;
+ const int lane = tid & 31;
+ const int g = lane >> 2;
+ const int t = lane & 3;
+ const int n = kFc1 ? p.inter_size : p.hidden_size;
+ const int k = kFc1 ? p.hidden_size : p.inter_size;
+ const int64_t row_bytes = kFc1 ? p.gate_row_bytes : p.down_row_bytes;
+ const int base = static_cast(blockIdx.x) * (MT * 16);
+ const int input_row = kFc1 ? routes[g] / p.k : routes[g];
+ const auto* x_row = input + static_cast(input_row) * k;
+ float acc[MT][4] = {};
+ // IQ2 operands and activation fragments share contiguous 64-K partitions.
+ // FC1 and FC2 retain their four- and eight-warp FP32 reduction boundaries.
+ for (int chunk = warp; chunk < k / 64; chunk += KW) {
+ const int offset = chunk * 64 + t * 16;
+ const uint4 x0 = g < count ? MoEPackedLoadRouter8(x_row + offset) : make_uint4(0, 0, 0, 0);
+ const uint4 x1 = g < count ? MoEPackedLoadRouter8(x_row + offset + 8) : make_uint4(0, 0, 0, 0);
+ const uint32_t x_words[8] = {x0.x, x0.y, x0.z, x0.w, x1.x, x1.y, x1.z, x1.w};
+#pragma unroll
+ for (int m = 0; m < MT; ++m) {
+ const int low_row = base + m * 16 + g;
+ const int high_row = low_row + 8;
+ float low_values[16] = {};
+ float high_values[16] = {};
+ if (low_row < n) {
+ const auto* row = weights + (static_cast(tile.expert) * n + low_row) * row_bytes;
+ gguf_cuda::DecodeGgufSubBlock16Impl(
+ row, offset / 32, (offset / 16) & 1, low_values);
+ }
+ if (high_row < n) {
+ const auto* row = weights + (static_cast(tile.expert) * n + high_row) * row_bytes;
+ gguf_cuda::DecodeGgufSubBlock16Impl(
+ row, offset / 32, (offset / 16) & 1, high_values);
+ }
+ float e0[8], e1[8], o0[8], o1[8];
+#pragma unroll
+ for (int s = 0; s < 4; ++s) {
+ e0[2 * s] = low_values[4 * s];
+ e0[2 * s + 1] = low_values[4 * s + 1];
+ e1[2 * s] = low_values[4 * s + 2];
+ e1[2 * s + 1] = low_values[4 * s + 3];
+ o0[2 * s] = high_values[4 * s];
+ o0[2 * s + 1] = high_values[4 * s + 1];
+ o1[2 * s] = high_values[4 * s + 2];
+ o1[2 * s + 1] = high_values[4 * s + 3];
+ }
+ const uint4 e_low = gguf_cuda::PackFloat8ToBf16x2(e0);
+ const uint4 e_high = gguf_cuda::PackFloat8ToBf16x2(e1);
+ const uint4 o_low = gguf_cuda::PackFloat8ToBf16x2(o0);
+ const uint4 o_high = gguf_cuda::PackFloat8ToBf16x2(o1);
+ const uint32_t el[4] = {e_low.x, e_low.y, e_low.z, e_low.w};
+ const uint32_t eh[4] = {e_high.x, e_high.y, e_high.z, e_high.w};
+ const uint32_t ol[4] = {o_low.x, o_low.y, o_low.z, o_low.w};
+ const uint32_t oh[4] = {o_high.x, o_high.y, o_high.z, o_high.w};
+#pragma unroll
+ for (int s = 0; s < 4; ++s) {
+ gguf_cuda::MmaBf16(acc[m], el[s], ol[s], eh[s], oh[s], x_words[2 * s], x_words[2 * s + 1]);
+ }
+ }
+ }
+#pragma unroll
+ for (int m = 0; m < MT; ++m) {
+ reduction[warp][m * 16 + g][2 * t] = acc[m][0];
+ reduction[warp][m * 16 + g][2 * t + 1] = acc[m][1];
+ reduction[warp][m * 16 + g + 8][2 * t] = acc[m][2];
+ reduction[warp][m * 16 + g + 8][2 * t + 1] = acc[m][3];
+ }
+ __syncthreads();
+ for (int index = tid; index < MT * 16 * 8; index += KW * 32) {
+ const int row = index / 8;
+ const int column = index & 7;
+ if (column < count && base + row < n) {
+ float sum = 0.f;
+#pragma unroll
+ for (int w = 0; w < KW; ++w) sum += reduction[w][row][column];
+ const auto rounded = __float2bfloat16_rn(sum);
+ if constexpr (kFc1) {
+ const float value = __bfloat162float(rounded);
+ const float positive = value < 0.f ? 0.f : value;
+ output[static_cast(routes[column]) * n + base + row] =
+ __float2bfloat16_rn(positive * positive);
+ } else {
+ output[static_cast(routes[column]) * n + base + row] = rounded;
+ }
+ }
+ }
+}
+
template
Status DispatchMoEPackedCuda(
cudaStream_t stream,
@@ -1198,7 +1340,8 @@ Status DispatchMoEPackedCuda(
params.inter_size > 0 && params.inter_size <= kMoEPackedMaxDimension &&
params.hidden_size % 32 == 0 && params.inter_size % 32 == 0 &&
params.num_experts > 0 && params.num_experts <= kMoEPackedMaxExperts &&
- params.k > 0 && params.k <= kMoEPackedMaxK && params.k <= params.num_experts,
+ params.k > 0 && params.k <= 32 && params.k <= params.num_experts &&
+ (params.expert_indices || params.k <= kMoEPackedMaxK),
"MoEPacked parameters exceed the supported CUDA index range");
if (params.num_rows == 0) return Status::OK();
const int total_routes = params.num_rows * params.k;
@@ -1232,9 +1375,13 @@ Status DispatchMoEPackedCuda(
CudaT* fc2_out = reinterpret_cast(ws_ptr);
- // 1. Softmax + TopK
+ // 1. Copy preselected routes or compute softmax + top-k.
const dim3 route_grid(static_cast((params.num_rows + 7) / 8));
- if (params.num_experts == 512) {
+ if (params.expert_indices) {
+ MoEPackedPreselectedRoutesKernel<<((total_routes + 255) / 256), 256, 0, stream>>>(
+ params.expert_indices, params.expert_weights, topk_indices, topk_scales,
+ total_routes, params.k, params.num_experts);
+ } else if (params.num_experts == 512) {
MoEPackedSoftmaxTopK512Kernel<<>>(
router_probs, topk_scales, topk_indices, params.num_rows, params.k, params.normalize_routing_weights);
} else {
@@ -1256,31 +1403,46 @@ Status DispatchMoEPackedCuda(
if (use_tensor_cores) {
if constexpr (kIsBf16) {
- // 3. Tensor-Core FC1 (64 inter_size channels/CTA in prefill, 16 inter_size channels/CTA Warp-K-Split in decode)
- if (use_prefill_tiles) {
- const dim3 fc1_grid(static_cast((params.inter_size + 63) / 64),
- static_cast(launch_tiles));
- LaunchMoEPackedFc1Bf16TcDispatch<4>(
- stream, fc1_grid, input, gate_weights, up_weights, sorted_routes, tiles, d_num_tiles, params, fc1_out);
+ if (params.relu_squared && params.gate_type == gguf_quant::kGgmlTypeIQ2_XXS &&
+ params.down_type == gguf_quant::kGgmlTypeIQ2_XXS) {
+ const unsigned int subtiles = use_prefill_tiles ? 4u : 1u;
+ const dim3 fc1_grid(static_cast((params.inter_size + 31) / 32),
+ static_cast(launch_tiles), subtiles);
+ MoEPackedReluIq2ContiguousKKernel<<>>(
+ input, gate_weights, sorted_routes, tiles, d_num_tiles, params, fc1_out);
+ ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetLastError()));
+ const dim3 fc2_grid(static_cast((params.hidden_size + 31) / 32),
+ static_cast(launch_tiles), subtiles);
+ MoEPackedReluIq2ContiguousKKernel<<>>(
+ fc1_out, down_weights, sorted_routes, tiles, d_num_tiles, params, fc2_out);
+ ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetLastError()));
} else {
- const dim3 fc1_grid(static_cast((params.inter_size + 15) / 16),
+ // 3. Tensor-Core FC1 (64 inter_size channels/CTA in prefill, 16 inter_size channels/CTA Warp-K-Split in decode)
+ if (use_prefill_tiles) {
+ const dim3 fc1_grid(static_cast((params.inter_size + 63) / 64),
+ static_cast(launch_tiles));
+ LaunchMoEPackedFc1Bf16TcDispatch<4>(
+ stream, fc1_grid, input, gate_weights, up_weights, sorted_routes, tiles, d_num_tiles, params, fc1_out);
+ } else {
+ const dim3 fc1_grid(static_cast((params.inter_size + 15) / 16),
+ static_cast(launch_tiles));
+ LaunchMoEPackedFc1DecodeBf16TcDispatch(
+ stream, fc1_grid, input, gate_weights, up_weights, sorted_routes, tiles, d_num_tiles, params, fc1_out);
+ }
+ ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetLastError()));
+
+ // 4. Tensor-Core FC2 (128 hidden_size channels per CTA)
+ const dim3 fc2_grid(static_cast((params.hidden_size + 127) / 128),
static_cast(launch_tiles));
- LaunchMoEPackedFc1DecodeBf16TcDispatch(
- stream, fc1_grid, input, gate_weights, up_weights, sorted_routes, tiles, d_num_tiles, params, fc1_out);
- }
- ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetLastError()));
-
- // 4. Tensor-Core FC2 (128 hidden_size channels per CTA)
- const dim3 fc2_grid(static_cast((params.hidden_size + 127) / 128),
- static_cast(launch_tiles));
- if (use_prefill_tiles) {
- LaunchMoEPackedFc2Bf16TcDispatch<4>(
- stream, fc2_grid, fc1_out, down_weights, sorted_routes, tiles, d_num_tiles, params, fc2_out);
- } else {
- LaunchMoEPackedFc2Bf16TcDispatch<1>(
- stream, fc2_grid, fc1_out, down_weights, sorted_routes, tiles, d_num_tiles, params, fc2_out);
+ if (use_prefill_tiles) {
+ LaunchMoEPackedFc2Bf16TcDispatch<4>(
+ stream, fc2_grid, fc1_out, down_weights, sorted_routes, tiles, d_num_tiles, params, fc2_out);
+ } else {
+ LaunchMoEPackedFc2Bf16TcDispatch<1>(
+ stream, fc2_grid, fc1_out, down_weights, sorted_routes, tiles, d_num_tiles, params, fc2_out);
+ }
+ ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetLastError()));
}
- ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetLastError()));
}
} else {
// Keep staging within CUDA's default shared-memory limit; larger shapes read activations directly.
@@ -1328,7 +1490,7 @@ Status DispatchMoEPackedCuda(
Status GetMoEPackedWorkspaceBytes(const MoEPackedParams& params, size_t elem_bytes, size_t& workspace_bytes) {
ORT_RETURN_IF_NOT(params.num_rows >= 0 && params.num_rows <= kMoEPackedMaxChunkRows &&
- params.k > 0 && params.k <= kMoEPackedMaxK &&
+ params.k > 0 && params.k <= 32 &&
params.num_experts > 0 && params.num_experts <= kMoEPackedMaxExperts &&
params.hidden_size > 0 && params.hidden_size <= kMoEPackedMaxDimension &&
params.inter_size > 0 && params.inter_size <= kMoEPackedMaxDimension &&
diff --git a/onnxruntime/contrib_ops/cuda/moe/moe_packed.h b/onnxruntime/contrib_ops/cuda/moe/moe_packed.h
index 2404f7bd98537..7f18a70b7acd5 100644
--- a/onnxruntime/contrib_ops/cuda/moe/moe_packed.h
+++ b/onnxruntime/contrib_ops/cuda/moe/moe_packed.h
@@ -35,6 +35,9 @@ struct MoEPackedParams {
int k{0};
bool normalize_routing_weights{true};
bool bf16_tc_supported{false};
+ bool relu_squared{false};
+ const int32_t* expert_indices{nullptr};
+ const float* expert_weights{nullptr};
int gate_type{0};
int up_type{0};
int down_type{0};
diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc
index 709759393f0ce..3e5ff6af5609c 100644
--- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc
+++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc
@@ -1081,31 +1081,56 @@ ONNX_MS_OPERATOR_SET_SCHEMA(
ONNX_MS_OPERATOR_SET_SCHEMA(
MoEPacked, 1,
OpSchema()
- .SetDoc("Mixture-of-Experts layer consuming 3D block-packed gate, up, and down expert weight blobs. "
+ .SetDoc("Mixture-of-Experts layer consuming 3D block-packed gate and down expert weight blobs, "
+ "plus up expert weights for SwiGLU. "
+ "A separate router can supply preselected expert IDs and FP32 per-token mixture coefficients; "
+ "these coefficients are consumed without normalization in the supplied expert order. "
"CUDA execution requires weights in device memory; host-backed external initializers are copied "
"to the device without changing their packed layout.\n"
"Supported format IDs and their stable byte layouts are defined in the "
"[packed-weight format contract]"
"(https://github.com/microsoft/onnxruntime/blob/main/docs/PackedWeightFormats.md).")
- .Attr("k", "Number of selected experts per token.", AttributeProto::INT)
- .Attr("normalize_routing_weights", "Whether to normalize top-k routing probabilities.",
+ .Attr("k",
+ "Number of selected experts per token: 1..16 for internal softmax routing, "
+ "1..32 for preselected routing; at most num_experts.",
+ AttributeProto::INT)
+ .Attr("normalize_routing_weights",
+ "Whether to renormalize selected softmax probabilities. Must be 0 for preselected routing.",
AttributeProto::INT, static_cast(1))
- .Attr("activation_type", "Expert activation function ('swiglu').",
+ .Attr("activation_type",
+ "Expert activation function: 'swiglu' or 'relu2' (squared ReLU; 'relu_squared' is an alias). "
+ "relu2 requires IQ2_XXS (type 16) gate and down weights and omitted up weights.",
AttributeProto::STRING, std::string("swiglu"))
.Attr("gate_type", "Stable packed-format ID of gate_experts_weights. Only contract-listed IDs are supported.",
AttributeProto::INT)
- .Attr("up_type", "Stable packed-format ID of up_experts_weights. Only contract-listed IDs are supported.",
- AttributeProto::INT)
+ .Attr("up_type",
+ "Stable packed-format ID of up_experts_weights; required for swiglu. "
+ "Only contract-listed IDs are supported.",
+ AttributeProto::INT, static_cast(-1))
.Attr("down_type", "Stable packed-format ID of down_experts_weights. Only contract-listed IDs are supported.",
AttributeProto::INT)
.Input(0, "input", "Input hidden states [..., hidden_size].", "T")
- .Input(1, "router_probs", "Router logits [num_rows, num_experts].", "T")
+ .Input(1, "router_probs",
+ "Router logits [num_rows, num_experts]; omit for preselected routing.",
+ "T", OpSchema::Optional)
.Input(2, "gate_experts_weights", "3D packed gate weight blob [num_experts, inter_size, gate_row_bytes].", "T2")
- .Input(3, "up_experts_weights", "3D packed up weight blob [num_experts, inter_size, up_row_bytes].", "T2")
+ .Input(3, "up_experts_weights",
+ "3D packed up weight blob [num_experts, inter_size, up_row_bytes]; required for swiglu, omit for relu2.",
+ "T2", OpSchema::Optional)
.Input(4, "down_experts_weights", "3D packed down weight blob [num_experts, hidden_size, down_row_bytes].", "T2")
+ .Input(5, "expert_indices",
+ "Preselected distinct expert IDs [num_rows, k], k <= 32. Requires expert_weights, no router_probs "
+ "and normalize_routing_weights=0. Invalid or duplicate IDs and nonfinite weights poison only their token row.",
+ "Tind", OpSchema::Optional)
+ .Input(6, "expert_weights",
+ "Preselected FP32 per-token mixture coefficients [num_rows, k], "
+ "preserved without normalization and mixed in supplied order.",
+ "Tweight", OpSchema::Optional)
.Output(0, "output", "Output hidden states [..., hidden_size].", "T")
.TypeConstraint("T", {"tensor(bfloat16)", "tensor(float16)"}, "Activation type.")
.TypeConstraint("T2", {"tensor(uint8)"}, "Packed byte storage.")
+ .TypeConstraint("Tind", {"tensor(int32)"}, "Preselected expert index type.")
+ .TypeConstraint("Tweight", {"tensor(float)"}, "Preselected mixture coefficient type.")
.TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) {
propagateElemTypeFromInputToOutput(ctx, 0, 0);
if (hasInputShape(ctx, 0)) {
diff --git a/onnxruntime/test/contrib_ops/moe_packed_routing_test.cc b/onnxruntime/test/contrib_ops/moe_packed_routing_test.cc
new file mode 100644
index 0000000000000..dbc38e03a324c
--- /dev/null
+++ b/onnxruntime/test/contrib_ops/moe_packed_routing_test.cc
@@ -0,0 +1,508 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+
+#include "contrib_ops/cpu/quantization/gguf_quant_tables.h"
+#include "gtest/gtest.h"
+#include "test/common/cuda_op_test_utils.h"
+#include "test/providers/provider_test_utils.h"
+
+#if defined(USE_CUDA) && !defined(USE_CUDA_MINIMAL)
+namespace onnxruntime::test {
+namespace {
+
+void RunMoeCuda(OpTester& test) {
+ std::vector> providers;
+ providers.push_back(DefaultCudaExecutionProvider());
+ test.SetOutputTolerance(0, 0);
+ test.Run(SessionOptions{}, OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers);
+}
+
+template
+void TestPreselectedReluSquared(int rows, int topk, int malformed = 0,
+ const char* activation_type = "relu2") {
+ constexpr int experts = 32;
+ constexpr int hidden_size = 1024;
+ constexpr int inter_size = 256;
+ std::vector gate(experts * inter_size * 264);
+ std::vector down(experts * hidden_size * 66);
+ std::vector input(rows * hidden_size), expected(rows * hidden_size);
+ std::vector indices(rows * topk);
+ std::vector weights(rows * topk);
+ std::vector sum8(experts);
+ for (int expert = 0; expert < experts; ++expert) {
+ const MLFloat16 d(expert & 1 ? .023f : .019f);
+ sum8[expert] = T(d.ToFloat() * .125f * 43.f).ToFloat() +
+ 7.f * T(d.ToFloat()).ToFloat();
+ for (int stage = 0; stage < 2; ++stage) {
+ const int n = stage == 0 ? inter_size : hidden_size;
+ const int blocks = stage == 0 ? hidden_size / 256 : inter_size / 256;
+ auto& bank = stage == 0 ? gate : down;
+ for (int col = 0; col < n; ++col) {
+ for (int block = 0; block < blocks; ++block) {
+ auto* bytes = bank.data() + ((expert * n + col) * blocks + block) * 66;
+ bytes[0] = static_cast(d.val);
+ bytes[1] = static_cast(d.val >> 8);
+ for (int record = 0; record < 8; ++record) {
+ for (int group = 0; group < 4; ++group) bytes[2 + record * 8 + group] = 1;
+ if (stage == 0 && expert % 5 == 4) {
+ for (int byte = 0; byte < 4; ++byte) {
+ bytes[2 + record * 8 + 4 + byte] = static_cast(0x0fffffffu >> (8 * byte));
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+ bool distinguishes_narrowed_weights = false;
+ for (int row = 0; row < rows; ++row) {
+ const float factor = static_cast(1 + (row & 1));
+ for (int col = 0; col < hidden_size; ++col) input[row * hidden_size + col] = T(factor / 32.f);
+ float mixed = 0.f;
+ float narrowed_mixed = 0.f;
+ for (int rank = 0; rank < topk; ++rank) {
+ const int expert = row & 1 ? (experts - 1 - rank) : rank;
+ indices[row * topk + rank] = expert;
+ const float weight = rank == 0 ? .3301f : .501f;
+ weights[row * topk + rank] = weight;
+ const float fc1 = expert % 5 == 4 ? 0.f : T(4.f * factor * sum8[expert]).ToFloat();
+ const float activation = T(fc1 * fc1).ToFloat();
+ const float fc2 = T(32.f * activation * sum8[expert]).ToFloat();
+ mixed = std::fma(weight, fc2, mixed);
+ narrowed_mixed = std::fma(T(weight).ToFloat(), fc2, narrowed_mixed);
+ }
+ distinguishes_narrowed_weights |= T(mixed).val != T(narrowed_mixed).val;
+ for (int col = 0; col < hidden_size; ++col) expected[row * hidden_size + col] = T(mixed);
+ }
+ if constexpr (std::is_same_v) {
+ ASSERT_TRUE(distinguishes_narrowed_weights) << "Fixture must detect BF16-rounded routing weights";
+ }
+ if (malformed) {
+ if (malformed == 1) indices[0] = -1;
+ if (malformed == 2) indices[1] = indices[0];
+ if (malformed == 3) weights[0] = std::numeric_limits::infinity();
+ for (int col = 0; col < hidden_size; ++col) {
+ expected[col] = T(std::numeric_limits::quiet_NaN());
+ }
+ }
+ OpTester test("MoEPacked", 1, kMSDomain);
+ test.AddAttribute("k", static_cast(topk));
+ test.AddAttribute("normalize_routing_weights", int64_t{0});
+ test.AddAttribute("activation_type", std::string(activation_type));
+ test.AddAttribute("gate_type", int64_t{16});
+ test.AddAttribute("down_type", int64_t{16});
+ test.AddInput("input", {1, rows, hidden_size}, input);
+ test.AddOptionalInputEdge();
+ test.AddInput("gate", {experts, inter_size, 264}, gate, true);
+ test.AddOptionalInputEdge();
+ test.AddInput("down", {experts, hidden_size, 66}, down, true);
+ test.AddInput("expert_indices", {rows, topk}, indices);
+ test.AddInput("expert_weights", {rows, topk}, weights);
+ test.AddOutput("output", {1, rows, hidden_size}, expected);
+ if (malformed) {
+ // TensorCheck checks expected NaNs, not actual NaNs.
+ test.SetCustomOutputVerifier([&](const std::vector& fetches, const std::string&) {
+ ASSERT_EQ(fetches.size(), 1u);
+ const auto actual = fetches[0].Get().DataAsSpan();
+ ASSERT_EQ(actual.size(), expected.size());
+ for (size_t i = 0; i < expected.size(); ++i) {
+ SCOPED_TRACE(testing::Message() << "token " << i / hidden_size << " col " << i % hidden_size);
+ const float actual_value = actual[i].ToFloat();
+ const float expected_value = expected[i].ToFloat();
+ if (std::isnan(expected_value)) {
+ EXPECT_TRUE(std::isnan(actual_value));
+ } else {
+ EXPECT_TRUE(std::isfinite(actual_value));
+ EXPECT_EQ(actual_value, expected_value);
+ }
+ }
+ });
+ }
+ RunMoeCuda(test);
+}
+
+TEST(MoEPackedRoutingTest, ReluSquaredPreservesPreselectedFp32Routing) {
+ if (!HasCudaEnvironment(800)) GTEST_SKIP() << "Requires CUDA SM80+";
+ for (int rows : {1, 8, 9, 17}) {
+ for (int topk : {2, 32}) {
+ SCOPED_TRACE(testing::Message() << "rows=" << rows << " k=" << topk);
+ TestPreselectedReluSquared(rows, topk);
+ }
+ }
+}
+
+TEST(MoEPackedRoutingTest, MalformedRoutesPoisonOnlyTheirToken) {
+ if (!HasCudaEnvironment(800)) GTEST_SKIP() << "Requires CUDA SM80+";
+ for (int malformed : {1, 2, 3}) {
+ SCOPED_TRACE(testing::Message() << "malformed=" << malformed);
+ TestPreselectedReluSquared(2, 2, malformed);
+ }
+}
+
+TEST(MoEPackedRoutingTest, ReluSquaredPreservesPreselectedFp16Routing) {
+ if (!HasCudaEnvironment(800)) GTEST_SKIP() << "Requires CUDA SM80+";
+ for (int rows : {1, 9}) {
+ for (int topk : {2, 32}) TestPreselectedReluSquared(rows, topk);
+ }
+}
+
+TEST(MoEPackedRoutingTest, ReluSquaredAliasRemainsSupported) {
+ if (!HasCudaEnvironment(800)) GTEST_SKIP() << "Requires CUDA SM80+";
+ TestPreselectedReluSquared(1, 2, 0, "relu_squared");
+}
+
+void TestSwigluRouting(int rows, int topk = 1, bool preselected = false) {
+ constexpr int size = 32;
+ const int experts = std::max(2, topk);
+ const BFloat16 hidden(.125f * (.25f / (1.f + std::exp(-.25f))));
+ const BFloat16 expert_output(16.f * hidden.ToFloat());
+ float mixed = 0.f;
+ std::vector weights(rows * topk);
+ std::vector indices(rows * topk);
+ if (preselected) {
+ for (int row = 0; row < rows; ++row) {
+ mixed = 0.f;
+ for (int rank = 0; rank < topk; ++rank) {
+ indices[row * topk + rank] = experts - 1 - rank;
+ weights[row * topk + rank] = rank == 0 ? .3301f : .501f;
+ mixed = std::fma(weights[row * topk + rank], expert_output.ToFloat(), mixed);
+ }
+ }
+ }
+ const BFloat16 output = preselected ? BFloat16(mixed) : expert_output;
+ const auto bank = [experts](float value) {
+ std::vector bytes(experts * size * size * 2);
+ const BFloat16 element(value);
+ for (size_t index = 0; index < bytes.size(); index += 2) {
+ bytes[index] = static_cast(element.val);
+ bytes[index + 1] = static_cast(element.val >> 8);
+ }
+ return bytes;
+ };
+ OpTester test("MoEPacked", 1, kMSDomain);
+ test.AddAttribute("k", static_cast(topk));
+ test.AddAttribute("gate_type", int64_t{30});
+ test.AddAttribute("up_type", int64_t{30});
+ test.AddAttribute("down_type", int64_t{30});
+ test.AddInput("input", {rows, size}, std::vector(rows * size, BFloat16(.03125f)));
+ if (preselected) {
+ test.AddAttribute("normalize_routing_weights", int64_t{0});
+ test.AddOptionalInputEdge();
+ } else {
+ test.AddInput("router_probs", {rows, experts}, std::vector(rows * experts, BFloat16(0.f)));
+ }
+ test.AddInput("gate", {experts, size, size * 2}, bank(.25f), true);
+ test.AddInput("up", {experts, size, size * 2}, bank(.125f), true);
+ test.AddInput("down", {experts, size, size * 2}, bank(.5f), true);
+ if (preselected) {
+ test.AddInput("indices", {rows, topk}, indices);
+ test.AddInput("weights", {rows, topk}, weights);
+ }
+ test.AddOutput("output", {rows, size}, std::vector(rows * size, output));
+ RunMoeCuda(test);
+}
+
+TEST(MoEPackedRoutingTest, ExistingSwigluRoutingDefaultsRemainUnchanged) {
+ if (!HasCudaEnvironment(800)) GTEST_SKIP() << "Requires CUDA SM80+";
+ for (int rows : {4, 48, 129}) TestSwigluRouting(rows);
+}
+
+TEST(MoEPackedRoutingTest, SwigluAcceptsPreselectedRoutes) {
+ if (!HasCudaEnvironment(800)) GTEST_SKIP() << "Requires CUDA SM80+";
+ for (int topk : {2, 32}) TestSwigluRouting(5, topk, true);
+}
+
+TEST(MoEPackedRoutingTest, ReluIq2PreservesContiguousKCancellation) {
+ if (!HasCudaEnvironment(800)) GTEST_SKIP() << "Requires CUDA SM80+";
+ constexpr int hidden = 1024;
+ constexpr int intermediate = 256;
+ std::vector gate(intermediate * 264, 0), down(hidden * 66, 0);
+ for (size_t block = 0; block < gate.size() / 66; ++block) gate[block * 66 + 1] = 0x3c;
+ for (size_t block = 0; block < down.size() / 66; ++block) down[block * 66 + 1] = 0x3c;
+ for (int rows : {1, 8, 17, 129}) {
+ SCOPED_TRACE(testing::Message() << "rows=" << rows);
+ std::vector input(rows * hidden, BFloat16(0.f));
+ for (int row = 0; row < rows; ++row) {
+ input[row * hidden] = BFloat16(0x1p25f);
+ input[row * hidden + 64] = BFloat16(-0x1p25f);
+ input[row * hidden + 128] = BFloat16(1.f);
+ input[row * hidden + 160] = BFloat16(1.f);
+ }
+ // Separate small terms from large MMA partials on both Ampere and Hopper.
+ // FC1 retains both unit terms, then BF16 ReLU-square and FC2 produce 1024.
+ OpTester test("MoEPacked", 1, kMSDomain);
+ test.AddAttribute("k", int64_t{1});
+ test.AddAttribute("activation_type", std::string("relu_squared"));
+ test.AddAttribute("gate_type", int64_t{16});
+ test.AddAttribute("down_type", int64_t{16});
+ test.AddAttribute("normalize_routing_weights", int64_t{0});
+ test.AddInput("input", {rows, hidden}, input);
+ test.AddOptionalInputEdge();
+ test.AddInput("gate", {1, intermediate, 264}, gate, true);
+ test.AddOptionalInputEdge();
+ test.AddInput("down", {1, hidden, 66}, down, true);
+ test.AddInput("indices", {rows, 1}, std::vector(rows, 0));
+ test.AddInput("weights", {rows, 1}, std::vector(rows, 1.f));
+ test.AddOutput("output", {rows, hidden},
+ std::vector(rows * hidden, BFloat16(1024.f)));
+ RunMoeCuda(test);
+ }
+}
+
+void TestRandomIq2Relu2(int rows, int hidden, int intermediate, int topk) {
+ const int experts = topk;
+ const int gate_row_bytes = hidden / 256 * 66;
+ const int down_row_bytes = intermediate / 256 * 66;
+ std::mt19937 random(33963);
+ std::vector gate(experts * intermediate * gate_row_bytes);
+ std::vector down(experts * hidden * down_row_bytes);
+ const auto fill_bank = [&](std::vector& bank) {
+ for (auto& byte : bank) byte = static_cast(random());
+ for (size_t block = 0; block < bank.size(); block += 66) {
+ const MLFloat16 scale(static_cast(256 + random() % 1024) / 262144.f);
+ bank[block] = static_cast(scale.val);
+ bank[block + 1] = static_cast(scale.val >> 8);
+ }
+ };
+ fill_bank(gate);
+ fill_bank(down);
+ std::vector input(rows * hidden), expected(rows * hidden);
+ for (auto& value : input) value = BFloat16(static_cast(static_cast(random() % 129) - 64) / 64.f);
+ std::vector indices(rows * topk);
+ std::vector weights(rows * topk);
+ std::vector gate_values(experts * intermediate * hidden);
+ std::vector down_values(experts * hidden * intermediate);
+ // Use the CPU decoder, not a reconstruction of the CUDA fragment mapping.
+ for (int row = 0; row < experts * intermediate; ++row) {
+ ASSERT_TRUE(contrib::gguf_quant::DequantizeGgufRowCpu(
+ gate.data() + row * gate_row_bytes, 16, hidden, gate_values.data() + row * hidden));
+ }
+ for (int row = 0; row < experts * hidden; ++row) {
+ ASSERT_TRUE(contrib::gguf_quant::DequantizeGgufRowCpu(
+ down.data() + row * down_row_bytes, 16, intermediate, down_values.data() + row * intermediate));
+ }
+ for (auto& value : gate_values) value = BFloat16(value).ToFloat();
+ for (auto& value : down_values) value = BFloat16(value).ToFloat();
+ std::vector activation(intermediate), mixed(hidden);
+ for (int row = 0; row < rows; ++row) {
+ std::fill(mixed.begin(), mixed.end(), 0.f);
+ for (int rank = 0; rank < topk; ++rank) {
+ const int expert = (row & 1) ? experts - 1 - rank : rank;
+ const float weight = static_cast(17 + random() % 97) / 113.f;
+ indices[row * topk + rank] = expert;
+ weights[row * topk + rank] = weight;
+ for (int channel = 0; channel < intermediate; ++channel) {
+ const auto* w = gate_values.data() + (expert * intermediate + channel) * hidden;
+ double sum = 0;
+ for (int k = 0; k < hidden; ++k) sum += static_cast(input[row * hidden + k].ToFloat()) * w[k];
+ const float fc1 = BFloat16(static_cast(sum)).ToFloat();
+ const float positive = std::max(fc1, 0.f);
+ activation[channel] = BFloat16(positive * positive).ToFloat();
+ }
+ for (int channel = 0; channel < hidden; ++channel) {
+ const auto* w = down_values.data() + (expert * hidden + channel) * intermediate;
+ double sum = 0;
+ for (int k = 0; k < intermediate; ++k) sum += static_cast(activation[k]) * w[k];
+ const float fc2 = BFloat16(static_cast(sum)).ToFloat();
+ mixed[channel] = std::fma(weight, fc2, mixed[channel]);
+ }
+ }
+ for (int channel = 0; channel < hidden; ++channel) expected[row * hidden + channel] = BFloat16(mixed[channel]);
+ }
+ OpTester test("MoEPacked", 1, kMSDomain);
+ test.AddAttribute("k", static_cast(topk));
+ test.AddAttribute("activation_type", std::string("relu2"));
+ test.AddAttribute("normalize_routing_weights", int64_t{0});
+ test.AddAttribute("gate_type", int64_t{16});
+ test.AddAttribute("down_type", int64_t{16});
+ test.AddInput("input", {rows, hidden}, input);
+ test.AddOptionalInputEdge();
+ test.AddInput("gate", {experts, intermediate, gate_row_bytes}, gate, true);
+ test.AddOptionalInputEdge();
+ test.AddInput("down", {experts, hidden, down_row_bytes}, down, true);
+ test.AddInput("indices", {rows, topk}, indices);
+ test.AddInput("weights", {rows, topk}, weights);
+ test.AddOutput("output", {rows, hidden}, expected);
+ test.SetCustomOutputVerifier([&](const std::vector& fetches, const std::string&) {
+ ASSERT_EQ(fetches.size(), 1u);
+ const auto actual = fetches[0].Get().DataAsSpan();
+ ASSERT_EQ(actual.size(), expected.size());
+ const auto ordered_bits = [](BFloat16 value) {
+ return (value.val & 0x8000) ? 0x8000 - (value.val & 0x7fff) : 0x8000 + value.val;
+ };
+ size_t mismatches = 0;
+ for (size_t i = 0; i < actual.size(); ++i) {
+ ASSERT_TRUE(std::isfinite(actual[i].ToFloat())) << "element " << i;
+ const int ulps = std::abs(ordered_bits(actual[i]) - ordered_bits(expected[i]));
+ ASSERT_LE(ulps, 2) << "element " << i;
+ mismatches += ulps != 0;
+ }
+ // FP32 MMA reduction may flip isolated BF16 rounding boundaries, not whole channels.
+ EXPECT_LE(mismatches, std::max(1, actual.size() / 1000));
+ });
+ RunMoeCuda(test);
+}
+
+TEST(MoEPackedRoutingTest, RandomIq2Relu2MatchesIndependentCpuReference) {
+ if (!HasCudaEnvironment(800)) GTEST_SKIP() << "Requires CUDA SM80+";
+ for (int hidden : {256, 1024}) {
+ const int intermediate = hidden == 256 ? 512 : 256;
+ for (int rows : {1, 9, 65}) {
+ SCOPED_TRACE(testing::Message() << "rows=" << rows << " hidden=" << hidden << " inter=" << intermediate);
+ TestRandomIq2Relu2(rows, hidden, intermediate, 2);
+ }
+ }
+ TestRandomIq2Relu2(5, 256, 512, 32);
+}
+
+TEST(MoEPackedRoutingTest, RejectsInvalidRoutingContracts) {
+ if (!HasCudaEnvironment(800)) GTEST_SKIP() << "Requires CUDA SM80+";
+ enum class Invalid {
+ ConflictingProbs,
+ NormalizedPreselected,
+ MissingWeights,
+ MissingIndices,
+ IndexShape,
+ WeightShape,
+ ProbsShape,
+ ReluUpWeights,
+ ReluGateFormat,
+ ReluDownFormat,
+ SoftmaxK,
+ SwigluMissingUpType,
+ IntegratedRoutingFunction,
+ IntegratedRouterScale,
+ IntegratedRouterInputs,
+ };
+ const std::pair cases[] = {
+ {Invalid::ConflictingProbs, "preselected routing requires both"},
+ {Invalid::NormalizedPreselected, "normalize_routing_weights=0"},
+ {Invalid::MissingWeights, "preselected routing requires both"},
+ {Invalid::MissingIndices, "preselected routing requires both"},
+ {Invalid::IndexShape, "indices and weights [num_rows, k]"},
+ {Invalid::WeightShape, "indices and weights [num_rows, k]"},
+ {Invalid::ProbsShape, "router_probs must be [num_rows, num_experts]"},
+ {Invalid::ReluUpWeights, "relu2 requires omitted up weights"},
+ {Invalid::ReluGateFormat, "relu2 currently requires IQ2_XXS gate and down weights"},
+ {Invalid::ReluDownFormat, "relu2 currently requires IQ2_XXS gate and down weights"},
+ {Invalid::SoftmaxK, "internal routing requires router_probs and k <= 16"},
+ {Invalid::SwigluMissingUpType, "swiglu requires up_type"},
+ {Invalid::IntegratedRoutingFunction, "Unrecognized attribute: routing_function"},
+ {Invalid::IntegratedRouterScale, "Unrecognized attribute: router_scale"},
+ {Invalid::IntegratedRouterInputs, "input size"},
+ };
+ for (const auto& [invalid, message] : cases) {
+ SCOPED_TRACE(testing::Message() << "invalid=" << static_cast(invalid));
+ constexpr int experts = 32;
+ constexpr int size = 256;
+ bool probs = false, indices = true, weights = true, up = false;
+ int topk = 2, normalize = 0;
+ int index_width = topk, weights_width = topk, probs_width = experts;
+ int gate_type = 16, down_type = 16;
+ std::string activation = "relu2";
+ switch (invalid) {
+ case Invalid::ConflictingProbs:
+ probs = true;
+ break;
+ case Invalid::NormalizedPreselected:
+ normalize = 1;
+ break;
+ case Invalid::MissingWeights:
+ weights = false;
+ break;
+ case Invalid::MissingIndices:
+ indices = false;
+ break;
+ case Invalid::IndexShape:
+ index_width = 1;
+ break;
+ case Invalid::WeightShape:
+ weights_width = 1;
+ break;
+ case Invalid::ProbsShape:
+ indices = weights = false;
+ probs = true;
+ probs_width = experts - 1;
+ break;
+ case Invalid::ReluUpWeights:
+ up = true;
+ break;
+ case Invalid::ReluGateFormat:
+ gate_type = 30;
+ break;
+ case Invalid::ReluDownFormat:
+ down_type = 30;
+ break;
+ case Invalid::SoftmaxK:
+ indices = weights = false;
+ probs = true;
+ topk = 17;
+ break;
+ case Invalid::SwigluMissingUpType:
+ activation = "swiglu";
+ up = true;
+ break;
+ case Invalid::IntegratedRoutingFunction:
+ case Invalid::IntegratedRouterScale:
+ case Invalid::IntegratedRouterInputs:
+ break;
+ }
+ OpTester test("MoEPacked", 1, kMSDomain);
+ test.AddAttribute("k", static_cast(topk));
+ test.AddAttribute("activation_type", activation);
+ test.AddAttribute("normalize_routing_weights", static_cast(normalize));
+ if (invalid == Invalid::IntegratedRoutingFunction) {
+ test.AddAttribute("routing_function", std::string("sigmoid"));
+ }
+ if (invalid == Invalid::IntegratedRouterScale) test.AddAttribute("router_scale", 1.f);
+ test.AddAttribute("gate_type", static_cast(gate_type));
+ test.AddAttribute("down_type", static_cast(down_type));
+ if (up && invalid != Invalid::SwigluMissingUpType) test.AddAttribute("up_type", int64_t{16});
+ test.AddInput("input", {1, size}, std::vector(size));
+ if (probs) {
+ test.AddInput("probs", {1, probs_width}, std::vector(probs_width));
+ } else {
+ test.AddOptionalInputEdge();
+ }
+ const std::vector bank(experts * size * 66);
+ test.AddInput("gate", {experts, size, 66}, bank, true);
+ if (up) {
+ test.AddInput("up", {experts, size, 66}, bank, true);
+ } else {
+ test.AddOptionalInputEdge();
+ }
+ test.AddInput("down", {experts, size, 66}, bank, true);
+ if (indices) {
+ test.AddInput("indices", {1, index_width}, std::vector(index_width));
+ } else {
+ test.AddOptionalInputEdge();
+ }
+ if (weights) {
+ test.AddInput("weights", {1, weights_width}, std::vector(weights_width, 1.f));
+ } else {
+ test.AddOptionalInputEdge();
+ }
+ if (invalid == Invalid::IntegratedRouterInputs) {
+ test.AddInput("router_input", {1, size}, std::vector(size));
+ test.AddInput("router_weight", {experts, size}, std::vector(experts * size), true);
+ test.AddInput("router_bias", {experts}, std::vector(experts), true);
+ }
+ test.AddOutput("output", {1, size}, std::vector(size));
+ std::vector> providers;
+ providers.push_back(DefaultCudaExecutionProvider());
+ test.Run(SessionOptions{}, OpTester::ExpectResult::kExpectFailure, message, {}, nullptr, &providers);
+ }
+}
+
+} // namespace
+} // namespace onnxruntime::test
+#endif