From eacb4b5abe510657c5f3416fbd6f2b3de4f5e850 Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Mon, 31 Aug 2026 11:12:10 -0700 Subject: [PATCH 1/3] [JAX] Add sqrtsoftplus router score function Signed-off-by: Jeremy Berchtold --- tests/jax/test_distributed_router.py | 22 +++-- tests/jax/test_fused_router.py | 87 ++++++++++++++++--- .../jax/cpp_extensions/router.py | 9 +- transformer_engine/jax/csrc/extensions/misc.h | 1 + .../jax/csrc/extensions/pybind.cpp | 1 + transformer_engine/jax/flax/moe.py | 4 +- transformer_engine/jax/moe.py | 4 +- transformer_engine/jax/router.py | 18 ++-- 8 files changed, 110 insertions(+), 36 deletions(-) diff --git a/tests/jax/test_distributed_router.py b/tests/jax/test_distributed_router.py index 35f59c897d..1eaf8a7825 100644 --- a/tests/jax/test_distributed_router.py +++ b/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,6 +295,8 @@ 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, @@ -300,7 +304,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 +312,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 +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, diff --git a/tests/jax/test_fused_router.py b/tests/jax/test_fused_router.py index ff62bce772..adc05b1f31 100644 --- a/tests/jax/test_fused_router.py +++ b/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) 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_(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_(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: From 5b51cf6ae8254e1e9e4d5da48098308e0c621f6b Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Wed, 2 Sep 2026 08:03:39 -0700 Subject: [PATCH 2/3] Address comments Signed-off-by: Jeremy Berchtold --- tests/jax/test_distributed_router.py | 3 - tests/jax/test_fused_router.py | 92 +++------------------------- 2 files changed, 8 insertions(+), 87 deletions(-) diff --git a/tests/jax/test_distributed_router.py b/tests/jax/test_distributed_router.py index 1eaf8a7825..2a5a5f2955 100644 --- a/tests/jax/test_distributed_router.py +++ b/tests/jax/test_distributed_router.py @@ -86,7 +86,6 @@ def _inject_router(request): } -@pytest.mark.triton class TestDistributedFusedTopk: """Test distributed execution of fused_topk_with_score_function. @@ -221,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. @@ -357,7 +355,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 adc05b1f31..ca9b621f22 100644 --- a/tests/jax/test_fused_router.py +++ b/tests/jax/test_fused_router.py @@ -397,37 +397,14 @@ 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("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 +@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, @@ -438,60 +415,11 @@ def test_topk_sqrtsoftplus( num_groups=num_groups, group_topk=group_topk, scaling_factor=scaling_factor, - score_function="sqrtsoftplus", + score_function=score_function, 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", - 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 # ============================================================================= @@ -503,7 +431,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) @@ -566,7 +493,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) @@ -639,7 +565,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 @@ -718,7 +643,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.""" From 0bd2ef06514d4610b2607507beef3e67b984f05b Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 2 Sep 2026 15:05:08 +0000 Subject: [PATCH 3/3] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/jax/test_fused_router.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/jax/test_fused_router.py b/tests/jax/test_fused_router.py index ca9b621f22..b76a224540 100644 --- a/tests/jax/test_fused_router.py +++ b/tests/jax/test_fused_router.py @@ -399,8 +399,7 @@ def loss_fused(logits_): @pytest_parametrize_wrapper("enable_bias", ENABLE_BIAS_OPTIONS) @pytest_parametrize_wrapper("score_function", SCORE_FUNCTIONS) def test_topk( - dtype, num_tokens, num_experts, topk, group_topk, scaling_factor, - enable_bias, score_function + 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.")