feat(qmoe): support the router_weights input

lib/Runtime/Kernels/hip/qmoe_kernel.hip

@@ -28,9 +28,13 @@ static constexpr int kRoutingBlock = 256;
 // ---------------------------------------------------------------------------
 // Top-k routing (fp16 only): for each token, find top-k experts by
 // probability. Uses insertion sort — k is small (typically 2-8).
+// With router_weights, router_probs only selects the experts and the mixing
+// weight of each selected expert is gathered from router_weights instead of
+// the softmax of router_probs.
 // ---------------------------------------------------------------------------
 __global__ void topk_routing_kernel(
     const __half* __restrict__ router_probs,
+    const __half* __restrict__ router_weights,
     int32_t* __restrict__ expert_indices,
     __half* __restrict__ expert_weights,
     int64_t num_tokens,
@@ -54,6 +58,7 @@ __global__ void topk_routing_kernel(
   extern __shared__ float s_exp[];        // [num_experts] exp(logit - max)
   __shared__ float s_red[kRoutingBlock];  // reduction scratch
   __shared__ int s_idx[kRoutingBlock];
+  __shared__ float s_w[kMaxK];            // selected weights, fp32
 
   // Pass 1: stash logits, block-reduce the max (numerically-stable softmax).
   float lmax = -1e30f;
@@ -118,21 +123,27 @@ __global__ void topk_routing_kernel(
     if (tid == 0) {
       const int win = s_idx[0];
       expert_indices[token * k + i] = win;
-      expert_weights[token * k + i] = static_cast<__half>(s_red[0] / exp_sum);
+      s_w[i] = router_weights
+                   ? static_cast<float>(
+                         router_weights[static_cast<int64_t>(token) * num_experts + win])
+                   : s_red[0] / exp_sum;
       s_exp[win] = -1e30f;
     }
     __syncthreads();
   }
 
-  // Optional renormalization of the top-k weights to sum to 1 (thread 0).
-  if (normalize && tid == 0) {
-    float sum = 0.0f;
-    for (int i = 0; i < ik; i++)
-      sum += static_cast<float>(expert_weights[token * k + i]);
-    if (sum > 0.0f)
+  // Optional renormalization of the top-k weights to sum to 1. A zero sum
+  // yields zero weights (only reachable with router_weights).
+  if (tid == 0) {
+    float scale = 1.0f;
+    if (normalize) {
+      float sum = 0.0f;
       for (int i = 0; i < ik; i++)
-        expert_weights[token * k + i] = static_cast<__half>(
-            static_cast<float>(expert_weights[token * k + i]) / sum);
+        sum += s_w[i];
+      scale = sum != 0.0f ? 1.0f / sum : 0.0f;
+    }
+    for (int i = 0; i < ik; i++)
+      expert_weights[token * k + i] = static_cast<__half>(s_w[i] * scale);
   }
 }
 
@@ -351,6 +362,7 @@ __global__ void bucket_tokens_kernel(
 extern "C" int hip_qmoe_topk_routing(
     void* stream,
     const void* router_probs,
+    const void* router_weights,
     void* expert_indices,
     void* expert_weights,
     int64_t num_tokens,
@@ -376,10 +388,12 @@ extern "C" int hip_qmoe_topk_routing(
   // Dynamic shared memory holds num_experts fp32 exp values (no expert cap;
   // e.g. 256 experts = 1 KB, well within the per-block shared limit).
   const size_t shmem = static_cast<size_t>(num_experts) * sizeof(float);
+  (void)hipGetLastError();
   hipLaunchKernelGGL(topk_routing_kernel,
                      dim3(static_cast<unsigned>(num_tokens)),
                      dim3(kRoutingBlock), shmem, s,
                      static_cast<const __half*>(router_probs),
+                     static_cast<const __half*>(router_weights),
                      static_cast<int32_t*>(expert_indices),
                      static_cast<__half*>(expert_weights), num_tokens,
                      num_experts, k, normalize);

lib/Runtime/Kernels/include/hip_custom_kernels.h

@@ -2687,13 +2687,17 @@ HIP_KERNEL_API int hip_dequantize_linear(
 
 /* Top-k routing: find top-k experts per token from router_probs.
  *   router_probs   - GPU [num_tokens, num_experts]
+ *   router_weights - GPU [num_tokens, num_experts] or NULL. When set, the
+ *                    selected experts' weights are gathered from it instead
+ *                    of the softmax of router_probs.
  *   expert_indices - GPU [num_tokens, k] int32 (output)
  *   expert_weights - GPU [num_tokens, k] (output, same type as probs)
  *   normalize      - 1 to normalize selected weights (sum-to-one)
  */
 HIP_KERNEL_API int hip_qmoe_topk_routing(
     void* stream,
     const void* router_probs,
+    const void* router_weights,
     void* expert_indices,
     void* expert_weights,
     int64_t num_tokens,

lib/Runtime/real/qmoe.cpp

@@ -39,10 +39,6 @@ int wrap_qmoe(RuntimeState *state, const void *input, const void *router_probs,
         return std::string(b);
       },
       state);
-  if (router_weights) {
-    fprintf(stderr, "wrap_qmoe: router_weights is not supported yet\n");
-    return -1;
-  }
   if (!state || !input || !router_probs || !output) {
     fprintf(stderr, "wrap_qmoe: null argument\n");
     return -1;
@@ -238,12 +234,13 @@ int wrap_qmoe(RuntimeState *state, const void *input, const void *router_probs,
       reinterpret_cast<int32_t *>(scratch_base + off_num_tiles);
 
   RUNTIME_DEBUG_LOG("[REAL] wrap_qmoe: topk_routing(tokens=%lld, experts=%lld, "
-                    "k=%lld, normalize=%lld)\n",
+                    "k=%lld, normalize=%lld, router_weights=%s)\n",
                     (long long)num_tokens, (long long)num_experts, (long long)k,
-                    (long long)normalize_routing_weights);
-  HIP_CHECK(hip_qmoe_topk_routing(stream, router_probs, d_expert_indices,
-                                  d_expert_weights, num_tokens, num_experts, k,
-                                  normalize_routing_weights, elem_size));
+                    (long long)normalize_routing_weights,
+                    router_weights ? "yes" : "no");
+  HIP_CHECK(hip_qmoe_topk_routing(
+      stream, router_probs, router_weights, d_expert_indices, d_expert_weights,
+      num_tokens, num_experts, k, normalize_routing_weights, elem_size));
 
   // Fused decode fast path: single-token MoE collapses to three back-to-back
   // kernel launches (FC1+SwiGLU, FC2, weighted reduce) with zero D2H,

test/numeric/tests/test_qmoe.py

@@ -54,6 +54,8 @@ def _make_qmoe_model(
     down_scale_range: tuple[float, float] = (0.001, 0.02),
     down_bias_range: tuple[float, float] = (-0.5, 0.5),
     seed: int = 42,
+    with_mix_weights: bool = False,
+    normalize_routing_weights: int = 1,
 ):
     """Build a QMoE ONNX model with random quantized expert weights.
 
@@ -63,6 +65,10 @@ def _make_qmoe_model(
 
     Scale and bias ranges are tuneable to match real model distributions.
     Scales are always positive (matching the real GPT-OSS-20B pattern).
+
+    with_mix_weights adds the optional QMoE router_weights input (input 14) as
+    a graph input named "mix_weights": the router logits then only select the
+    experts and the mixing weights are gathered from it.
     """
     down_in = intermediate // 2
     packed_gate_up = hidden // 2
@@ -131,29 +137,42 @@ def _make_qmoe_model(
         numpy_helper.from_array(down_bias, name="down_bias"),
     ]
 
+    node_inputs = [
+        "input",
+        "router_weights",
+        "gate_up_qweight",
+        "gate_up_scales",
+        "gate_up_bias",
+        "down_qweight",
+        "down_scales",
+        "down_bias",
+        "",
+        "",
+        "",
+    ]
+    graph_inputs = [x, router]
+    if with_mix_weights:
+        # fc1/fc2/fc3 zero points (inputs 11-13) stay absent.
+        node_inputs += ["", "", "", "mix_weights"]
+        graph_inputs.append(
+            helper.make_tensor_value_info(
+                "mix_weights",
+                TensorProto.FLOAT16,
+                [seq_len, num_experts],
+            )
+        )
+
     node = helper.make_node(
         "QMoE",
-        [
-            "input",
-            "router_weights",
-            "gate_up_qweight",
-            "gate_up_scales",
-            "gate_up_bias",
-            "down_qweight",
-            "down_scales",
-            "down_bias",
-            "",
-            "",
-            "",
-        ],
+        node_inputs,
         ["output"],
         domain="com.microsoft",
         activation_type="swiglu",
         activation_alpha=1.702,
         activation_beta=1.0,
         expert_weight_bits=4,
         k=top_k,
-        normalize_routing_weights=1,
+        normalize_routing_weights=normalize_routing_weights,
         swiglu_fusion=1,
         use_sparse_mixer=0,
         block_size=block_size,
@@ -163,7 +182,7 @@ def _make_qmoe_model(
     ms_opset = helper.make_opsetid("com.microsoft", 1)
     return make_model_from_nodes(
         [node],
-        [x, router],
+        graph_inputs,
         [output],
         initializers=initializers,
         opset=21,
@@ -267,3 +286,53 @@ def test_qmoe_gpt_oss_shape_skewed_routing(self, model_runner, seq_len):
 
         actual, expected = model_runner.run_sample(model, [x, router])
         compare_outputs(actual, expected, atol=1e-1, rtol=1e-2, cos_threshold=0.99)
+
+    @pytest.mark.parametrize("normalize", [0, 1])
+    @pytest.mark.parametrize(
+        "seq_len",
+        [
+            pytest.param(
+                1,
+                marks=pytest.mark.xfail(
+                    reason="fused decode exceeds atol=0.1 here, as "
+                    "test_qmoe_gpt_oss_shape[1] does without router_weights",
+                    strict=False,
+                ),
+            ),
+            37,
+            128,
+        ],
+    )
+    def test_qmoe_gpt_oss_shape_mix_weights(self, model_runner, seq_len, normalize):
+        """router_weights (QMoE input 14): the logits only pick the top-k
+        experts; the mixing weights are gathered from a separate tensor and
+        renormalised only when normalize_routing_weights=1. The mixing
+        weights are uncorrelated with the logits, so taking the softmax of
+        the logits instead would not match. Their top-k sums vary around 1 so
+        the unnormalised case keeps the output scale (and fp16 error) of the
+        normalised one. seq_len=1 covers the fused decode path, the others
+        the prefill paths.
+        """
+        model = _make_qmoe_model(
+            1,
+            seq_len,
+            HIDDEN,
+            INTERMEDIATE,
+            NUM_EXPERTS,
+            TOP_K,
+            BLOCK_SIZE,
+            with_mix_weights=True,
+            normalize_routing_weights=normalize,
+        )
+
+        rng = np.random.default_rng(123)
+        x = rng.uniform(-1, 1, [1, seq_len, HIDDEN]).astype(np.float16)
+        router = (rng.standard_normal([seq_len, NUM_EXPERTS]) * 1.2).astype(
+            np.float16,
+        )
+        ranked = -np.sort(-router.astype(np.float32), axis=1)
+        assert not (ranked[:, TOP_K - 1] == ranked[:, TOP_K]).any()
+        mix = rng.uniform(0.05, 0.45, [seq_len, NUM_EXPERTS]).astype(np.float16)
+
+        actual, expected = model_runner.run_sample(model, [x, router, mix])
+        compare_outputs(actual, expected, atol=1e-1, rtol=1e-2, cos_threshold=0.99)