[JAX] Add sqrtsoftplus router score function

tests/jax/test_distributed_router.py

@@ -60,7 +60,7 @@ def _inject_router(request):
 jax.config.update("jax_use_shardy_partitioner", True)
 
 from test_fused_router import (
-    reference_topk_softmax_sigmoid,
+    reference_topk_with_score_function,
     reference_compute_scores_for_aux_loss,
     reference_aux_loss,
     make_logits,
@@ -133,7 +133,7 @@ def target_fwd(x):
 
             logits_shards = jnp.reshape(logits, (num_dp_devices, local_num_tokens, num_experts))
             ref_fwd_fn = jax.jit(
-                lambda x: reference_topk_softmax_sigmoid(
+                lambda x: reference_topk_with_score_function(
                     x,
                     topk=topk,
                     score_function=score_function,
@@ -160,21 +160,23 @@ def target_fwd(x):
             ), "Routing map mismatch in distributed fused_topk"
 
             # === Backward ===
+            grad_weights = jnp.linspace(0.5, 1.5, num_experts, dtype=jnp.float32)[None, :]
+
             def target_loss(x):
                 p, _ = fused_topk_with_score_function(
                     x,
                     topk=topk,
                     score_function=score_function,
                 )
-                return jnp.sum(p)
+                return jnp.sum(p * grad_weights)
 
             def ref_chunk_loss(x_chunk):
-                p, _ = reference_topk_softmax_sigmoid(
+                p, _ = reference_topk_with_score_function(
                     x_chunk,
                     topk=topk,
                     score_function=score_function,
                 )
-                return jnp.sum(p)
+                return jnp.sum(p * grad_weights)
 
             target_grad = jax.jit(jax.grad(target_loss))(logits_sharded)
 
@@ -195,7 +197,7 @@ def ref_chunk_loss(x_chunk):
         "num_tokens,num_experts,topk",
         TOPK_CASES,
     )
-    @pytest.mark.parametrize("score_function", ["softmax", "sigmoid"])
+    @pytest.mark.parametrize("score_function", ["softmax", "sigmoid", "sqrtsoftplus"])
     def test_distributed_topk(
         self,
         device_count,
@@ -293,22 +295,24 @@ def target_fwd(x):
             ), "Routing map mismatch in distributed score_for_aux_loss"
 
             # === Backward ===
+            grad_weights = jnp.linspace(0.5, 1.5, num_experts, dtype=jnp.float32)[None, :]
+
             def target_loss(x):
                 s, _ = fused_topk_with_score_function(
                     x,
                     topk=topk,
                     score_function=score_function,
                     compute_aux_scores=True,
                 )
-                return jnp.sum(s)
+                return jnp.sum(s * grad_weights)
 
             def ref_chunk_loss(x_chunk):
                 _, s = reference_compute_scores_for_aux_loss(
                     x_chunk,
                     topk=topk,
                     score_function=score_function,
                 )
-                return jnp.sum(s)
+                return jnp.sum(s * grad_weights)
 
             target_grad = jax.jit(jax.grad(target_loss))(logits_sharded)
 
@@ -329,7 +333,7 @@ def ref_chunk_loss(x_chunk):
         "num_tokens,num_experts,topk",
         TOPK_CASES,
     )
-    @pytest.mark.parametrize("score_function", ["softmax", "sigmoid"])
+    @pytest.mark.parametrize("score_function", ["softmax", "sigmoid", "sqrtsoftplus"])
     def test_distributed_score_for_aux_loss(
         self,
         device_count,

tests/jax/test_fused_router.py

@@ -90,9 +90,9 @@ def _inject_router(request):
     "L2": ALL_SCORE_AUX_LOSS_CASES,
 }
 
-ALL_SCORE_FUNCTIONS = ["softmax", "sigmoid"]
+ALL_SCORE_FUNCTIONS = ["softmax", "sigmoid", "sqrtsoftplus"]
 SCORE_FUNCTIONS = {
-    "L0": ["softmax"],
+    "L0": ["softmax", "sqrtsoftplus"],
     "L2": ALL_SCORE_FUNCTIONS,
 }
 
@@ -166,7 +166,7 @@ def reference_group_limited_topk(
     return probs, top_indices
 
 
-def reference_topk_softmax_sigmoid(
+def reference_topk_with_score_function(
     logits: jnp.ndarray,
     topk: int,
     use_pre_softmax: bool = False,
@@ -176,7 +176,7 @@ def reference_topk_softmax_sigmoid(
     score_function: str = "softmax",
     expert_bias: Optional[jnp.ndarray] = None,
 ):
-    """Reference implementation for topk + softmax/sigmoid."""
+    """Reference implementation for topk with a supported score function."""
     num_tokens, num_experts = logits.shape
 
     def compute_topk(scores, topk, num_groups=None, group_topk=None):
@@ -199,8 +199,11 @@ def compute_topk(scores, topk, num_groups=None, group_topk=None):
         else:
             scores, top_indices = compute_topk(logits, topk, num_groups, group_topk)
             probs = jax.nn.softmax(scores.astype(jnp.float32), axis=-1).astype(logits.dtype)
-    elif score_function == "sigmoid":
-        scores = jax.nn.sigmoid(logits.astype(jnp.float32)).astype(logits.dtype)
+    elif score_function in ("sigmoid", "sqrtsoftplus"):
+        if score_function == "sigmoid":
+            scores = jax.nn.sigmoid(logits.astype(jnp.float32)).astype(logits.dtype)
+        else:
+            scores = jnp.sqrt(jax.nn.softplus(logits.astype(jnp.float32))).astype(logits.dtype)
         if expert_bias is not None:
             scores_for_routing = scores + expert_bias
             _, top_indices = compute_topk(scores_for_routing, topk, num_groups, group_topk)
@@ -233,6 +236,9 @@ def reference_compute_scores_for_aux_loss(logits: jnp.ndarray, topk: int, score_
     elif score_function == "sigmoid":
         scores = jax.nn.sigmoid(logits.astype(jnp.float32))
         scores = scores / (scores.sum(axis=-1, keepdims=True) + 1e-20) if topk > 1 else scores
+    elif score_function == "sqrtsoftplus":
+        scores = jnp.sqrt(jax.nn.softplus(logits.astype(jnp.float32)))
+        scores = scores / (scores.sum(axis=-1, keepdims=True) + 1e-20) if topk > 1 else scores
     else:
         raise ValueError(f"Invalid score_function: {score_function}")
 
@@ -269,7 +275,7 @@ def reference_aux_loss(
 
 def make_logits(num_tokens, num_experts, score_function, dtype=jnp.float32):
     """Create deterministic logits for testing."""
-    if score_function == "sigmoid":
+    if score_function in ("sigmoid", "sqrtsoftplus"):
         offset = jnp.arange(-num_tokens // 2, num_tokens // 2, dtype=dtype) * 1e-4
         logits = jnp.arange(-num_experts // 2, num_experts // 2, dtype=dtype) * 1e-2
         logits = logits[None, :].repeat(num_tokens, axis=0) + offset[:, None]
@@ -306,7 +312,7 @@ def run_topk_comparison(
     """Compare fused vs reference top-k implementation, both jitted."""
     logits = make_logits(num_tokens, num_experts, score_function, dtype)
 
-    if enable_bias and score_function == "sigmoid":
+    if enable_bias and score_function in ("sigmoid", "sqrtsoftplus"):
         expert_bias = jnp.arange(num_experts, dtype=jnp.float32) * 0.1
         expert_bias = jnp.flip(expert_bias)
     else:
@@ -315,7 +321,7 @@ def run_topk_comparison(
     # Forward: reference (jitted)
     ref_fwd_fn = jax.jit(
         partial(
-            reference_topk_softmax_sigmoid,
+            reference_topk_with_score_function,
             topk=topk,
             use_pre_softmax=use_pre_softmax,
             num_groups=num_groups,
@@ -348,8 +354,10 @@ def run_topk_comparison(
     assert jnp.array_equal(routing_map_ref, routing_map_fused), "Routing map mismatch"
 
     # Backward: reference (jitted)
+    grad_weights = jnp.linspace(0.5, 1.5, num_experts, dtype=jnp.float32)[None, :]
+
     def loss_ref(logits_):
-        p, _ = reference_topk_softmax_sigmoid(
+        p, _ = reference_topk_with_score_function(
             logits_,
             topk,
             use_pre_softmax,
@@ -359,7 +367,7 @@ def loss_ref(logits_):
             score_function,
             expert_bias,
         )
-        return p.sum()
+        return jnp.sum(p * grad_weights)
 
     def loss_fused(logits_):
         p, _ = fused_topk_with_score_function(
@@ -372,7 +380,7 @@ def loss_fused(logits_):
             score_function,
             expert_bias,
         )
-        return p.sum()
+        return jnp.sum(p * grad_weights)
 
     grad_ref = jax.jit(jax.grad(loss_ref))(logits)
     grad_fused = jax.jit(jax.grad(loss_fused))(logits)
@@ -408,6 +416,55 @@ def test_topk_sigmoid(
     )
 
 
+@pytest_parametrize_wrapper("dtype", DTYPES)
+@pytest_parametrize_wrapper(
+    "num_tokens,num_experts,topk",
+    TOPK_CASES,
+)
+@pytest_parametrize_wrapper("group_topk", GROUP_TOPK_OPTIONS)
+@pytest_parametrize_wrapper("scaling_factor", SCALING_FACTOR_OPTIONS)
+@pytest_parametrize_wrapper("enable_bias", ENABLE_BIAS_OPTIONS)
+@pytest.mark.triton
+def test_topk_sqrtsoftplus(
+    dtype, num_tokens, num_experts, topk, group_topk, scaling_factor, enable_bias
+):
+    num_groups = 8 if group_topk else None
+    run_topk_comparison(
+        dtype=dtype,
+        num_tokens=num_tokens,
+        num_experts=num_experts,
+        topk=topk,
+        use_pre_softmax=False,
+        num_groups=num_groups,
+        group_topk=group_topk,
+        scaling_factor=scaling_factor,
+        score_function="sqrtsoftplus",
+        enable_bias=enable_bias,
+    )
+
+
+@pytest.mark.triton
+def test_sqrtsoftplus_score_function_enum():
+    from transformer_engine.jax.router import ScoreFunction
+
+    logits = make_logits(128, 32, "sqrtsoftplus")
+    string_fn = jax.jit(
+        partial(fused_topk_with_score_function, topk=4, score_function="sqrtsoftplus")
+    )
+    enum_fn = jax.jit(
+        partial(
+            fused_topk_with_score_function,
+            topk=4,
+            score_function=ScoreFunction.SQRTSOFTPLUS,
+        )
+    )
+
+    string_probs, string_routing_map = string_fn(logits)
+    enum_probs, enum_routing_map = enum_fn(logits)
+    assert jnp.array_equal(string_probs, enum_probs)
+    assert jnp.array_equal(string_routing_map, enum_routing_map)
+
+
 @pytest_parametrize_wrapper("dtype", DTYPES)
 @pytest_parametrize_wrapper(
     "num_tokens,num_experts,topk",
@@ -477,9 +534,11 @@ def test_fused_scores_for_aux_loss(dtype, num_tokens, num_experts, topk, score_f
     assert jnp.array_equal(routing_map_ref, routing_map_fused), "Routing map mismatch"
 
     # Backward (jitted)
+    grad_weights = jnp.linspace(0.5, 1.5, num_experts, dtype=jnp.float32)[None, :]
+
     def loss_ref(logits_):
         _, s = reference_compute_scores_for_aux_loss(logits_, topk, score_function)
-        return s.sum()
+        return jnp.sum(s * grad_weights)
 
     def loss_fused(logits_):
         s, _ = fused_topk_with_score_function(
@@ -488,7 +547,7 @@ def loss_fused(logits_):
             score_function=score_function,
             compute_aux_scores=True,
         )
-        return s.sum()
+        return jnp.sum(s * grad_weights)
 
     grad_ref = jax.jit(jax.grad(loss_ref))(logits)
     grad_fused = jax.jit(jax.grad(loss_fused))(logits)

transformer_engine/jax/cpp_extensions/router.py

@@ -2,6 +2,7 @@
 #
 # See LICENSE for license information.
 """JAX/TE custom ops for fused MoE router"""
+
 from enum import IntEnum
 
 import jax.numpy as jnp
@@ -27,6 +28,7 @@ class ScoreFunction(IntEnum):
 
     SIGMOID = int(JAXX_Score_Function.SIGMOID)
     SOFTMAX = int(JAXX_Score_Function.SOFTMAX)
+    SQRTSOFTPLUS = int(JAXX_Score_Function.SQRTSOFTPLUS)
 
 
 class RoutingMapFormat(IntEnum):
@@ -99,7 +101,8 @@ def abstract(
         else:
             routing_map_aval = logits_aval.update(shape=i_shape, dtype=jnp.bool_)
         # The CUDA kernel always uses float32 (CompType) for intermediate
-        # computations (softmax/sigmoid values saved for backward).
+        # computations. Softmax/sigmoid save activation values for backward;
+        # sqrtsoftplus saves the original logits.
         intermediate_aval = logits_aval.update(shape=i_shape, dtype=jnp.float32)
         return probs_aval, routing_map_aval, intermediate_aval
 
@@ -702,9 +705,9 @@ def fused_topk_with_score_function_fwd(
     scaling_factor : float
         Scaling factor for output probs.
     score_function : ScoreFunction
-        ScoreFunction.SOFTMAX or ScoreFunction.SIGMOID.
+        ScoreFunction.SOFTMAX, ScoreFunction.SIGMOID, or ScoreFunction.SQRTSOFTPLUS.
     expert_bias : jnp.ndarray
-        Expert bias (only used with sigmoid). Pass empty array if unused.
+        Expert bias (only used with sigmoid/sqrtsoftplus). Pass empty array if unused.
     compute_aux_scores : bool
         If True, compute clean scores for aux loss instead of full top-k.
     routing_map_format : int

transformer_engine/jax/csrc/extensions/misc.h

@@ -126,6 +126,7 @@ void hash_combine(int64_t &seed, const T &v, Rest... rest) {
 enum class JAXX_Score_Function : int64_t {
   SIGMOID = 0,
   SOFTMAX = 1,
+  SQRTSOFTPLUS = 2,
 };
 
 // Mirror of NVTERoutingMapFormat for JAX FFI plumbing. Values are taken

transformer_engine/jax/csrc/extensions/pybind.cpp

@@ -255,6 +255,7 @@ PYBIND11_MODULE(transformer_engine_jax, m) {
   pybind11::enum_<JAXX_Score_Function>(m, "JAXX_Score_Function", pybind11::module_local())
       .value("SIGMOID", JAXX_Score_Function::SIGMOID)
       .value("SOFTMAX", JAXX_Score_Function::SOFTMAX)
+      .value("SQRTSOFTPLUS", JAXX_Score_Function::SQRTSOFTPLUS)
       .export_values();
 
   pybind11::enum_<JAXX_Routing_Map_Format>(m, "JAXX_Routing_Map_Format", pybind11::module_local())

transformer_engine/jax/flax/moe.py

@@ -72,7 +72,7 @@ class _MoEBlock(TransformerEngineBase):
         product with ``layer_w0 @ wi_1``. Default ``"silu"``.
 
     score_function : Union[str, ScoreFunction]
-        ``"softmax"`` (default) or ``"sigmoid"`` for the routing scores.
+        ``"softmax"`` (default), ``"sigmoid"``, or ``"sqrtsoftplus"`` for the routing scores.
     use_pre_softmax : bool
         Apply softmax before topk (vs. after).
     num_groups, group_topk : Optional[int]
@@ -82,7 +82,7 @@ class _MoEBlock(TransformerEngineBase):
     use_expert_routing_bias : bool
         If ``True``, registers a per-expert routing bias (shape ``[E]``)
         used by the topk selection. Only meaningful with
-        ``score_function="sigmoid"``; the underlying primitive validates
+        ``score_function="sigmoid"`` or ``"sqrtsoftplus"``; the underlying primitive validates
         the pairing.
     aux_loss_coeff : float
         If ``> 0``, return the MoE auxiliary load-balancing loss scalar

transformer_engine/jax/moe.py

@@ -132,7 +132,7 @@ def _with_sharding_constraint_cast_bwd(x: jnp.ndarray, sharding) -> jnp.ndarray:
         ``d_logits_2d`` is produced by
         ``fused_topk_with_score_function_bwd``. That primitive runs at
         fp32 because the fwd promoted ``logits_2d`` to fp32 (the fused
-        topk/softmax/sigmoid kernels are only validated at fp32).
+        topk/softmax/sigmoid/sqrtsoftplus kernels are only validated at fp32).
 
     JAX's type promotion then makes ``d_x_from_gate + d_x_from_dispatch``
     fp32, so the user-visible ``d_x`` ends up wider than ``x``. That
@@ -1249,7 +1249,7 @@ def moe(
     ----------
     expert_bias : Optional[jnp.ndarray]
         ``[num_experts]`` learnable router bias added before the top-k
-        when ``score_function='sigmoid'``. Pass ``None`` to disable.
+        when ``score_function='sigmoid'`` or ``'sqrtsoftplus'``. Pass ``None`` to disable.
         The bias has no gradient through the top-k primitive itself (it
         only steers expert selection); a zero cotangent is returned for
         it.