Reuse MoETopkFusion in quantile-balancing routing

src/paddlefleet/transformer/moe/moe_router.py

@@ -452,13 +452,6 @@ def __init__(
             self.expert_usage.stop_gradient = True
 
         if self.topk_method == "quantile_balancing":
-            if getattr(self.config, "moe_topk_fusion", False):
-                raise ValueError(
-                    "quantile_balancing is incompatible with moe_topk_fusion. "
-                    "The MoETopkFusion kernel does not support QB's histogram-based "
-                    "bias update, and enabling both causes incorrect gate normalization. "
-                    "Please set moe_topk_fusion=False when using quantile_balancing."
-                )
             if self.routing_type != "none":
                 raise ValueError(
                     "quantile_balancing is a self-contained load balancing method, "
@@ -1197,6 +1190,7 @@ def _accumulate_qb_histogram(
         biased_scores: paddle.Tensor,
         k: int,
         valid_mask: paddle.Tensor | None = None,
+        alpha: paddle.Tensor | None = None,
     ):
         """Accumulate required_bias into the QB histogram.
 
@@ -1219,17 +1213,25 @@ def _accumulate_qb_histogram(
             k: top-k value
             valid_mask: [N, 1] mask marking non-padding tokens, or None when no
                 padding information is available (all rows count).
+            alpha: [N] or [N, 1] pre-computed cutoff values from MoETopkFusion
+                kernel (optional). When provided, skips the internal topk
+                computation for cutoff.
         """
         N, E = raw_scores.shape
         B = self.qb_n_bins
 
         # Compute alpha: the (k+1)-th largest biased score per token
         # This is the "cutoff" -- the highest biased score NOT selected
-        # Clamp k+1 to at most E (in case of degenerate config)
-        topk_val = min(k + 1, int(E))
-        alpha = paddle.topk(biased_scores, k=topk_val, axis=-1, sorted=True)[0][
-            :, -1:
-        ]  # [N, 1] -- the smallest of top-(k+1), i.e., cutoff
+        if alpha is not None:
+            # Alpha provided externally (e.g., from MoETopkFusion kernel)
+            if alpha.ndim == 1:
+                alpha = alpha.unsqueeze(-1)  # [N] -> [N, 1]
+        else:
+            # Clamp k+1 to at most E (in case of degenerate config)
+            topk_val = min(k + 1, int(E))
+            alpha = paddle.topk(
+                biased_scores, k=topk_val, axis=-1, sorted=True
+            )[0][:, -1:]  # [N, 1] -- the smallest of top-(k+1), i.e., cutoff
 
         # required_bias[t, e] = alpha[t] - raw_scores[t, e]
         required_bias = alpha - raw_scores  # [N, E]
@@ -1710,14 +1712,8 @@ def forward(self, input, input_ids=None, origin_input_ids=None):
                     gates_ori.sum(-1, keepdim=True), min=1e-12
                 )
 
-        if (
-            getattr(self.config, "moe_topk_fusion", False)
-            and self.topk_method != "quantile_balancing"
-        ):
-            # Use MoETopkFusion Triton kernel for bit-exact alignment.
-            # This ensures the topk selection + normalization uses the exact same
-            # GPU kernel, avoiding FP32 rounding differences between
-            # Triton's scalar loop and Paddle's tensor ops.
+        if getattr(self.config, "moe_topk_fusion", False):
+            # Use MoETopkFusion Triton kernel for fused topk selection.
             MoETopkFusion = _get_moe_topk_fusion()
             use_node_limit = self.n_group > 1
             if not self.config.gpt_model_use_experimental_version:
@@ -1735,16 +1731,42 @@ def forward(self, input, input_ids=None, origin_input_ids=None):
                 _log_moe_md5(
                     probs_for_choice, "probs_for_choice", self._layer_number
                 )
-            top_gate, top_idx = MoETopkFusion.apply(
-                gates,  # gate_probs (original sigmoid scores)
-                probs_for_choice,  # probs_for_choice (with correction bias)
-                self.num_experts_per_tok,
-                use_node_limit,
-                self.n_group,
-                self.topk_group,
-                self.norm_topk_prob,  # norm_gate_logits
-            )
-            # top_gate is already normalized by the Triton kernel when norm_topk_prob=True
+
+            if self.topk_method == "quantile_balancing":
+                # QB path: kernel does NOT normalize (normalization stays eager
+                # for bit-exact alignment with original QB). Additionally returns
+                # alpha (the k+1-th largest choice value) for histogram accumulation.
+                top_gate, top_idx, alpha = MoETopkFusion.apply(
+                    gates,  # gate_probs (original sigmoid scores)
+                    probs_for_choice,  # probs_for_choice (with correction bias)
+                    self.num_experts_per_tok,
+                    use_node_limit,
+                    self.n_group,
+                    self.topk_group,
+                    False,  # norm_gate_logits=False (normalization stays eager)
+                    True,  # return_alpha=True
+                )
+                # Accumulate QB histogram with pre-computed alpha (skips internal topk)
+                if framework._dygraph_tracer()._has_grad:
+                    self._accumulate_qb_histogram(
+                        gates,
+                        probs_for_choice,
+                        self.num_experts_per_tok,
+                        valid_mask=input_ids_none_zero_mask,
+                        alpha=alpha,
+                    )
+            else:
+                # noaux_tc / other paths: kernel handles normalization
+                top_gate, top_idx = MoETopkFusion.apply(
+                    gates,  # gate_probs (original sigmoid scores)
+                    probs_for_choice,  # probs_for_choice (with correction bias)
+                    self.num_experts_per_tok,
+                    use_node_limit,
+                    self.n_group,
+                    self.topk_group,
+                    self.norm_topk_prob,  # norm_gate_logits
+                )
+            # top_gate is already normalized by the Triton kernel when norm_topk_prob=True (non-QB)
 
             _log_moe_md5(
                 top_idx.cast("float32"), "topk_indices", self._layer_number
@@ -1812,7 +1834,13 @@ def forward(self, input, input_ids=None, origin_input_ids=None):
 
         # norm
         if self.norm_topk_prob:
-            if not getattr(self.config, "moe_topk_fusion", False):
+            if (
+                not getattr(self.config, "moe_topk_fusion", False)
+                or self.topk_method == "quantile_balancing"
+            ):
+                # QB fusion path passes norm_gate_logits=False to the kernel,
+                # so normalization must happen here in eager (for bit-exact alignment).
+                # Non-fusion paths also normalize here.
                 if self.use_accuracy_compatible:
                     _sum_f64 = top_gate.cast(paddle.float64).sum(
                         axis=-1, keepdim=True
@@ -1821,7 +1849,7 @@ def forward(self, input, input_ids=None, origin_input_ids=None):
                 else:
                     denominator = top_gate.sum(axis=-1, keepdim=True) + 1e-20
                 top_gate = top_gate / denominator
-            # When gpt_model_use_experimental_version is True, top_gate is already normalized by MoETopkFusion
+            # When moe_topk_fusion=True and not QB, top_gate is already normalized by MoETopkFusion
 
         if self.routed_scaling_factor_learnable:
             top_gate = apply_learnable_routed_scaling(

src/paddlefleet/triton_ops/moe_topk_fusion.py

@@ -37,12 +37,13 @@
 
 @enable_compat_on_triton_kernel
 @triton.jit
-def _fwd_kernel(
+def _fwd_kernel(  # pragma: no cover - triton kernel body compiles to PTX, not python-instrumentable
     ptr_gate,
     ptr_choice,
     ptr_out_probs,
     ptr_out_idx,
     ptr_out_sum,
+    ptr_out_alpha,
     stride_gate_s,
     stride_gate_e,
     stride_choice_s,
@@ -55,6 +56,7 @@ def _fwd_kernel(
     n_group: tl.constexpr,
     topk_group: tl.constexpr,
     norm_gate_logits: tl.constexpr,
+    return_alpha: tl.constexpr,
     BLOCK_SIZE: tl.constexpr,
 ):
     """
@@ -158,6 +160,13 @@ def _fwd_kernel(
         # Mask out this index so we don't pick it again
         choice_vals = tl.where(off_e != k_idx, choice_vals, float("-inf"))
 
+    # --- Alpha (cutoff) output ---
+    # After the top-k loop, choice_vals has the top-k positions masked to -inf.
+    # The max of the remaining values is the (k+1)-th largest = cutoff alpha.
+    if return_alpha:
+        alpha_val = tl.max(choice_vals, axis=0)
+        tl.store(ptr_out_alpha + pid, alpha_val)
+
     # --- Normalization ---
     if norm_gate_logits:
         # Sum the collected probs
@@ -263,6 +272,7 @@ def forward(
         n_group,
         topk_group,
         norm_gate_logits,
+        return_alpha=False,
     ):
         """
         Forward pass: select topk experts.
@@ -276,10 +286,14 @@ def forward(
             n_group: number of expert groups.
             topk_group: number of selected topk groups.
             norm_gate_logits: whether to normalize gate logits.
+            return_alpha: if True, additionally return the (k+1)-th largest
+                choice value per token (cutoff alpha) as a no-grad tensor.
 
         Returns:
             topk_probs: normalized topk probabilities, shape [seq_len, moe_k].
             topk_indices: topk expert indices, shape [seq_len, moe_k].
+            alpha (optional): cutoff values, shape [seq_len], only when
+                return_alpha=True.
         """
         seq_len, n_experts = gate_probs.shape
 
@@ -290,6 +304,9 @@ def forward(
             if norm_gate_logits
             else None
         )
+        alpha = (
+            paddle.empty((seq_len,), dtype="float32") if return_alpha else None
+        )
 
         # Block size must cover n_experts for the single-block reduction logic
         BLOCK_SIZE = triton.next_power_of_2(n_experts)
@@ -298,13 +315,16 @@ def forward(
 
         # Use topk_probs as dummy pointer for sum if not needed, as it is writable
         ptr_sum_arg = topk_sum if norm_gate_logits else topk_probs
+        # Use topk_probs as dummy pointer for alpha if not needed
+        ptr_alpha_arg = alpha if return_alpha else topk_probs
 
         _fwd_kernel[(seq_len,)](
             gate_probs,
             probs_for_choice,
             topk_probs,
             topk_indices,
             ptr_sum_arg,
+            ptr_alpha_arg,
             int(gate_probs.stride(0)),
             int(gate_probs.stride(1)),
             int(probs_for_choice.stride(0)),
@@ -317,20 +337,27 @@ def forward(
             n_group if use_node_limit else 1,
             topk_group if use_node_limit else 1,
             norm_gate_logits,
+            return_alpha,
             BLOCK_SIZE,
         )
 
         ctx.save_for_backward(topk_indices, topk_probs, topk_sum)
         ctx.input_shape = gate_probs.shape
         ctx.norm_gate_logits = norm_gate_logits
         ctx.moe_k = moe_k
+        ctx.return_alpha = return_alpha
 
+        if return_alpha:
+            return topk_probs, topk_indices.to(paddle.int64), alpha
         return topk_probs, topk_indices.to(paddle.int64)
 
     @staticmethod
-    def backward(ctx, grad_output_probs, grad_output_indices):
+    def backward(ctx, grad_output_probs, grad_output_indices, grad_alpha=None):
         """
         Backward: compute the gradient with respect to gate_probs.
+
+        When return_alpha=True in forward, backward receives an extra
+        grad_alpha argument which is ignored (alpha has no gradient).
         """
         topk_indices, topk_normed_probs, topk_sum = ctx.saved_tensor()
 

tests/single_card_tests/ai_edited_test/ops/test_ai_moe_topk_fusion_3.py

@@ -82,9 +82,11 @@ def test_forward_launches_kernel_and_saves_context_without_norm(self):
         self.assertIs(recorder.args[0], gate_probs)
         self.assertIs(recorder.args[1], probs_for_choice)
         self.assertIs(recorder.args[4], topk_probs)
-        self.assertEqual(recorder.args[12:16], (2, False, 1, 1))
-        self.assertEqual(recorder.args[16], False)
-        self.assertEqual(recorder.args[17], 32)
+        self.assertEqual(recorder.args[12], 5)
+        self.assertEqual(recorder.args[13:17], (2, False, 1, 1))
+        self.assertEqual(recorder.args[17], False)
+        self.assertFalse(recorder.args[18])
+        self.assertEqual(recorder.args[19], 32)
         self.assertEqual(topk_probs.shape, [2, 2])
         self.assertEqual(topk_indices.dtype, paddle.int64)
         saved_indices, saved_probs, saved_sum = ctx.saved
@@ -115,9 +117,11 @@ def test_forward_launches_kernel_with_node_limit_and_norm_sum(self):
         finally:
             moe_topk_fusion._fwd_kernel = old_kernel
 
-        self.assertEqual(recorder.args[12:16], (3, True, 8, 2))
-        self.assertTrue(recorder.args[16])
-        self.assertEqual(recorder.args[17], 64)
+        self.assertEqual(recorder.args[12], 64)
+        self.assertEqual(recorder.args[13:17], (3, True, 8, 2))
+        self.assertTrue(recorder.args[17])
+        self.assertFalse(recorder.args[18])
+        self.assertEqual(recorder.args[19], 64)
         self.assertEqual(ctx.saved[2].shape, [1])
         self.assertEqual(topk_probs.shape, [1, 3])
         self.assertEqual(topk_indices.shape, [1, 3])

tests/single_card_tests/ai_edited_test/ops/test_ai_moe_topk_fusion_4.py

@@ -239,6 +239,7 @@ def test_forward_kernel_python_body_exercises_group_topk_and_norm(self):
             FakePtr(),
             FakePtr(),
             FakePtr(),
+            FakePtr(),
             4,
             1,
             4,
@@ -251,6 +252,7 @@ def test_forward_kernel_python_body_exercises_group_topk_and_norm(self):
             2,
             2,
             True,
+            True,
             4,
         )