[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,
@@ -86,7 +86,6 @@ def _inject_router(request):
 }
 
 
-@pytest.mark.triton
 class TestDistributedFusedTopk:
     """Test distributed execution of fused_topk_with_score_function.
 
@@ -133,7 +132,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 +159,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 +196,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,
@@ -219,7 +220,6 @@ def test_distributed_topk(
         )
 
 
-@pytest.mark.triton
 class TestDistributedScoreForAuxLoss:
     """Test distributed execution of fused_topk_with_score_function with compute_aux_scores=True.
 
@@ -293,22 +293,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 +331,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,
@@ -353,7 +355,6 @@ def test_distributed_score_for_aux_loss(
         )
 
 
-@pytest.mark.triton
 class TestDistributedMoEAuxLoss:
     """Test distributed execution of fused_moe_aux_loss.
 

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)
@@ -389,10 +397,13 @@ def loss_fused(logits_):
 @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_sigmoid(
-    dtype, num_tokens, num_experts, topk, group_topk, scaling_factor, enable_bias
+@pytest_parametrize_wrapper("score_function", SCORE_FUNCTIONS)
+def test_topk(
+    dtype, num_tokens, num_experts, topk, group_topk, scaling_factor, enable_bias, score_function
 ):
+    if score_function == "softmax" and enable_bias:
+        pytest.skip("Bias is not supported with 'softmax' router score function. Skipping.")
+        return
     num_groups = 8 if group_topk else None
     run_topk_comparison(
         dtype=dtype,
@@ -403,38 +414,11 @@ def test_topk_sigmoid(
         num_groups=num_groups,
         group_topk=group_topk,
         scaling_factor=scaling_factor,
-        score_function="sigmoid",
+        score_function=score_function,
         enable_bias=enable_bias,
     )
 
 
-@pytest_parametrize_wrapper("dtype", DTYPES)
-@pytest_parametrize_wrapper(
-    "num_tokens,num_experts,topk",
-    TOPK_CASES,
-)
-@pytest_parametrize_wrapper("use_pre_softmax", USE_PRE_SOFTMAX_OPTIONS)
-@pytest_parametrize_wrapper("group_topk", GROUP_TOPK_OPTIONS)
-@pytest_parametrize_wrapper("scaling_factor", SCALING_FACTOR_OPTIONS)
-@pytest.mark.triton
-def test_topk_softmax(
-    dtype, num_tokens, num_experts, topk, use_pre_softmax, group_topk, scaling_factor
-):
-    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=use_pre_softmax,
-        num_groups=num_groups,
-        group_topk=group_topk,
-        scaling_factor=scaling_factor,
-        score_function="softmax",
-        enable_bias=False,
-    )
-
-
 # =============================================================================
 # Test: Fused Score for MoE Aux Loss
 # =============================================================================
@@ -446,7 +430,6 @@ def test_topk_softmax(
     SCORE_AUX_LOSS_CASES,
 )
 @pytest_parametrize_wrapper("score_function", SCORE_FUNCTIONS)
-@pytest.mark.triton
 def test_fused_scores_for_aux_loss(dtype, num_tokens, num_experts, topk, score_function):
     logits = make_logits(num_tokens, num_experts, score_function, dtype)
 
@@ -477,9 +460,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 +473,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)
@@ -507,7 +492,6 @@ def loss_fused(logits_):
     "num_tokens,num_experts,topk",
     AUX_LOSS_CASES,
 )
-@pytest.mark.triton
 def test_fused_moe_aux_loss(dtype, num_tokens, num_experts, topk):
     key = jax.random.PRNGKey(SEED)
 
@@ -580,7 +564,6 @@ def _bytemap_to_bitmap_u8(bytemap):
     TOPK_CASES,
 )
 @pytest_parametrize_wrapper("score_function", SCORE_FUNCTIONS)
-@pytest.mark.triton
 def test_topk_bitmap_vs_bytemap(dtype, num_tokens, num_experts, topk, score_function):
     """fused_topk_with_score_function should produce the same probs and an
     LSB-packed bitmap routing_map when routing_map_format=BITMAP_U8, and
@@ -659,7 +642,6 @@ def loss_bit(logits_):
     SCORE_AUX_LOSS_CASES,
 )
 @pytest_parametrize_wrapper("score_function", SCORE_FUNCTIONS)
-@pytest.mark.triton
 def test_score_for_aux_loss_bitmap_vs_bytemap(dtype, num_tokens, num_experts, topk, score_function):
     """compute_aux_scores=True path: bitmap routing_map must equal LSB-packed
     bytemap; scores must be bitwise identical across formats."""

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())