diff --git a/tests/jax/test_distributed_router.py b/tests/jax/test_distributed_router.py
index 35f59c897d..c6c8a11a84 100644
--- a/tests/jax/test_distributed_router.py
+++ b/tests/jax/test_distributed_router.py
@@ -36,15 +36,8 @@
 
 
 @pytest.fixture(autouse=True, scope="function")
-def _inject_router(request):
-    """Lazy-load router API only for tests marked 'triton'. Other tests run without importing.
-
-    We inject into sys.modules[__name__] so test code can use fused_topk_with_score_function,
-    fused_moe_aux_loss as module-level names (fixture locals are not visible to tests).
-    """
-    if not request.node.get_closest_marker("triton"):
-        yield
-        return
+def _inject_router():
+    """Inject the router API as module-level names for every test in this file."""
     import sys
     from transformer_engine.jax.router import (
         fused_topk_with_score_function,
@@ -60,7 +53,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 +79,6 @@ def _inject_router(request):
 }
 
 
-@pytest.mark.triton
 class TestDistributedFusedTopk:
     """Test distributed execution of fused_topk_with_score_function.
 
@@ -133,7 +125,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 +152,27 @@ def target_fwd(x):
             ), "Routing map mismatch in distributed fused_topk"
 
             # === Backward ===
+            # Use random weights so the backward pass receives non-uniform gradients instead of
+            # the all-ones gradient produced by an unweighted jnp.sum.
+            grad_weights = jax.random.uniform(
+                jax.random.PRNGKey(42), (1, num_experts), dtype=jnp.float32
+            )
+
             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 +193,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 +217,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,6 +290,12 @@ def target_fwd(x):
             ), "Routing map mismatch in distributed score_for_aux_loss"
 
             # === Backward ===
+            # Use random weights so the backward pass receives non-uniform gradients instead of
+            # the all-ones gradient produced by an unweighted jnp.sum.
+            grad_weights = jax.random.uniform(
+                jax.random.PRNGKey(42), (1, num_experts), dtype=jnp.float32
+            )
+
             def target_loss(x):
                 s, _ = fused_topk_with_score_function(
                     x,
@@ -300,7 +303,7 @@ def target_loss(x):
                     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(
@@ -308,7 +311,7 @@ def ref_chunk_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 +332,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 +356,6 @@ def test_distributed_score_for_aux_loss(
         )
 
 
-@pytest.mark.triton
 class TestDistributedMoEAuxLoss:
     """Test distributed execution of fused_moe_aux_loss.
 
diff --git a/tests/jax/test_fused_router.py b/tests/jax/test_fused_router.py
index ff62bce772..a0e46ca301 100644
--- a/tests/jax/test_fused_router.py
+++ b/tests/jax/test_fused_router.py
@@ -16,15 +16,8 @@
 
 
 @pytest.fixture(autouse=True, scope="function")
-def _inject_router(request):
-    """Lazy-load router API only for tests marked 'triton'. Other tests run without importing.
-
-    We inject into sys.modules[__name__] so test code can use fused_topk_with_score_function,
-    fused_moe_aux_loss as module-level names (fixture locals are not visible to tests).
-    """
-    if not request.node.get_closest_marker("triton"):
-        yield
-        return
+def _inject_router():
+    """Inject the router API as module-level names for every test in this file."""
     from transformer_engine.jax.router import (
         fused_topk_with_score_function,
         fused_moe_aux_loss,
@@ -90,9 +83,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 +159,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 +169,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 +192,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 +229,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 +268,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 +305,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 +314,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 +347,12 @@ def run_topk_comparison(
     assert jnp.array_equal(routing_map_ref, routing_map_fused), "Routing map mismatch"
 
     # Backward: reference (jitted)
+    # Use random weights so the backward pass receives non-uniform gradients instead of
+    # the all-ones gradient produced by an unweighted jnp.sum.
+    grad_weights = jax.random.uniform(jax.random.PRNGKey(SEED), (1, num_experts), dtype=jnp.float32)
+
     def loss_ref(logits_):
-        p, _ = reference_topk_softmax_sigmoid(
+        p, _ = reference_topk_with_score_function(
             logits_,
             topk,
             use_pre_softmax,
@@ -359,7 +362,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 +375,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,37 +392,23 @@ 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
-):
-    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="sigmoid",
-        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
+@pytest_parametrize_wrapper("score_function", SCORE_FUNCTIONS)
+def test_topk(
+    dtype,
+    num_tokens,
+    num_experts,
+    topk,
+    group_topk,
+    scaling_factor,
+    enable_bias,
+    use_pre_softmax,
+    score_function,
 ):
+    if use_pre_softmax and score_function != "softmax":
+        pytest.skip("Pre-softmax is only applicable to the 'softmax' router score function.")
+    if score_function == "softmax" and enable_bias:
+        pytest.skip("Bias is not supported with 'softmax' router score function. Skipping.")
     num_groups = 8 if group_topk else None
     run_topk_comparison(
         dtype=dtype,
@@ -430,8 +419,8 @@ def test_topk_softmax(
         num_groups=num_groups,
         group_topk=group_topk,
         scaling_factor=scaling_factor,
-        score_function="softmax",
-        enable_bias=False,
+        score_function=score_function,
+        enable_bias=enable_bias,
     )
 
 
@@ -446,7 +435,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 +465,13 @@ 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)
+    # Use random weights so the backward pass receives non-uniform gradients instead of
+    # the all-ones gradient produced by an unweighted jnp.sum.
+    grad_weights = jax.random.uniform(jax.random.PRNGKey(SEED), (1, num_experts), dtype=jnp.float32)
+
     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 +480,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 +499,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 +571,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 +649,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."""
diff --git a/transformer_engine/jax/cpp_extensions/router.py b/transformer_engine/jax/cpp_extensions/router.py
index 46f51c9d33..ad49c7a7d1 100644
--- a/transformer_engine/jax/cpp_extensions/router.py
+++ b/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
diff --git a/transformer_engine/jax/csrc/extensions/misc.h b/transformer_engine/jax/csrc/extensions/misc.h
index 62f95ad529..1f4b94f92b 100644
--- a/transformer_engine/jax/csrc/extensions/misc.h
+++ b/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
diff --git a/transformer_engine/jax/csrc/extensions/pybind.cpp b/transformer_engine/jax/csrc/extensions/pybind.cpp
index 3927e2686e..26489a618f 100644
--- a/transformer_engine/jax/csrc/extensions/pybind.cpp
+++ b/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())
diff --git a/transformer_engine/jax/flax/moe.py b/transformer_engine/jax/flax/moe.py
index cb10c7aa18..58573966ed 100644
--- a/transformer_engine/jax/flax/moe.py
+++ b/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
diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py
index 55a85ebb2f..c672db2fec 100644
--- a/transformer_engine/jax/moe.py
+++ b/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.
diff --git a/transformer_engine/jax/router.py b/transformer_engine/jax/router.py
index 170c74fe5f..cd5c1dec6d 100644
--- a/transformer_engine/jax/router.py
+++ b/transformer_engine/jax/router.py
@@ -10,7 +10,7 @@
 
 Functions:
     fused_topk_with_score_function:
-        Fused score_function + top-k selection. Supports softmax/sigmoid,
+        Fused score_function + top-k selection. Supports softmax/sigmoid/sqrtsoftplus,
         grouped top-k, expert bias, and scaling factor. When compute_aux_scores=True,
         switches to the clean score-for-aux-loss kernel (no bias/groups/scaling,
         dense output).
@@ -73,7 +73,8 @@ def _validate_score_function(score_function: Union[str, ScoreFunction]) -> Score
         return ScoreFunction[score_function.upper()]
     except (KeyError, AttributeError):
         raise ValueError(
-            "score_function must be 'softmax', 'sigmoid', or a ScoreFunction enum, "
+            "score_function must be 'softmax', 'sigmoid', 'sqrtsoftplus', "
+            "or a ScoreFunction enum, "
             f"got {score_function!r}"
         ) from None
 
@@ -127,9 +128,10 @@ def fused_topk_with_score_function(
         Scaling factor applied to output probs.
         Ignored when compute_aux_scores=True.
     score_function : Union[str, ScoreFunction]
-        Score function: "softmax" / "sigmoid" or ScoreFunction.SOFTMAX / ScoreFunction.SIGMOID.
+        Score function: "softmax", "sigmoid", or "sqrtsoftplus", or the corresponding
+        ScoreFunction enum member.
     expert_bias : Optional[jnp.ndarray]
-        Expert bias, shape [num_experts]. Only used with sigmoid.
+        Expert bias, shape [num_experts]. Only used with sigmoid or sqrtsoftplus.
         Ignored when compute_aux_scores=True.
     compute_aux_scores : bool
         If True, use the clean score-for-aux-loss kernel. Returns dense scores
@@ -169,9 +171,13 @@ def fused_topk_with_score_function(
         group_topk = -1
         scaling_factor = 1.0
     else:
-        if expert_bias is not None and score_function != ScoreFunction.SIGMOID:
+        if expert_bias is not None and score_function not in (
+            ScoreFunction.SIGMOID,
+            ScoreFunction.SQRTSOFTPLUS,
+        ):
             raise ValueError(
-                "expert_bias is only supported with score_function='sigmoid'. "
+                "expert_bias is only supported with score_function='sigmoid' or "
+                "'sqrtsoftplus'. "
                 f"Got score_function='{score_function.name}'."
             )
         if expert_bias is None: