perf: enable MoE GroupedGEMM for MoE models

examples/configs/recipes/llm/performance/dapo-deepseek-v3-64n8g.yaml

@@ -44,6 +44,7 @@ policy:
     empty_unused_memory_level: 2
     enabled: true
     activation_checkpointing: true
+    moe_grouped_gemm: true
     tensor_model_parallel_size: 8
     expert_model_parallel_size: 32
     pipeline_model_parallel_size: 8

examples/configs/recipes/llm/performance/grpo-deepseek-v3-32n8g.yaml

@@ -26,6 +26,7 @@ policy:
     pipeline_model_parallel_size: 16
     expert_model_parallel_size: 16
     activation_checkpointing: true
+    moe_grouped_gemm: true
     num_layers_in_first_pipeline_stage: 3
     num_layers_in_last_pipeline_stage: 2
     apply_rope_fusion: false

examples/configs/recipes/llm/performance/grpo-qwen3-235b-16n8g.yaml

@@ -29,6 +29,7 @@ policy:
     context_parallel_size: 2
     expert_model_parallel_size: 16
     activation_checkpointing: true
+    moe_grouped_gemm: true
     num_layers_in_first_pipeline_stage: 11
     num_layers_in_last_pipeline_stage: 11
     defer_fp32_logits: true

examples/configs/recipes/llm/performance/grpo-qwen3-30ba3b-4n4g.yaml

@@ -22,6 +22,7 @@ policy:
     pipeline_model_parallel_size: 1
     expert_model_parallel_size: 16
     sequence_parallel: false
+    moe_grouped_gemm: true
     optimizer:
       lr: 3.0e-07
       min_lr: 3.0e-08

examples/configs/recipes/llm/performance/grpo-qwen3-30ba3b-4n8g-40K.yaml

@@ -24,6 +24,7 @@ policy:
     expert_model_parallel_size: 8
     sequence_parallel: true
     context_parallel_size: 8
+    moe_grouped_gemm: true
     optimizer:
       lr: 3.0e-07
       min_lr: 3.0e-08

examples/configs/recipes/llm/performance/grpo-qwen3-30ba3b-4n8g.yaml

@@ -22,6 +22,7 @@ policy:
     pipeline_model_parallel_size: 1
     expert_model_parallel_size: 8
     sequence_parallel: false
+    moe_grouped_gemm: true
     optimizer:
       lr: 3.0e-07
       min_lr: 3.0e-08

nemo_rl/models/megatron/setup.py

@@ -630,6 +630,9 @@ def _apply_moe_config(model_cfg: Any, config: PolicyConfig) -> None:
 
     model_cfg.moe_permute_fusion = config["megatron_cfg"]["moe_permute_fusion"]
 
+    if "moe_grouped_gemm" in config["megatron_cfg"]:
+        model_cfg.moe_grouped_gemm = config["megatron_cfg"]["moe_grouped_gemm"]
+
 
 def _apply_mtp_config(model_cfg: Any, config: PolicyConfig) -> None:
     if "mtp_num_layers" in config["megatron_cfg"]:

nemo_rl/models/policy/__init__.py

@@ -237,6 +237,10 @@ class MegatronConfig(TypedDict):
     moe_token_dispatcher_type: str
     # Can be used only with 'alltoall' token dispatcher
     moe_shared_expert_overlap: bool
+    # Enable grouped GEMM for MoE experts via CUTLASS. Significant throughput
+    # gain when multiple experts are assigned per rank (num_local_experts > 1).
+    # Requires TE >= 1.11.0 for FP8 and Ampere (sm_80) or newer.
+    moe_grouped_gemm: NotRequired[bool]
     # HybridEP settings for MoE expert parallelism (requires moe_token_dispatcher_type='flex')
     # See: https://github.com/deepseek-ai/DeepEP/tree/hybrid-ep
     moe_flex_dispatcher_backend: NotRequired[str]

tests/unit/models/megatron/test_megatron_setup.py

@@ -620,6 +620,20 @@ def test_moe_configuration(self):
         assert model_cfg.moe_token_dispatcher_type == "alltoall"
         assert model_cfg.moe_shared_expert_overlap is True
 
+    @staticmethod
+    def _base_moe_megatron_cfg() -> dict:
+        return {
+            "expert_tensor_parallel_size": 2,
+            "expert_model_parallel_size": 4,
+            "moe_router_dtype": "float32",
+            "moe_router_load_balancing_type": "none",
+            "moe_router_bias_update_rate": 0.0,
+            "moe_permute_fusion": True,
+            "moe_enable_deepep": False,
+            "moe_token_dispatcher_type": "alltoall",
+            "moe_shared_expert_overlap": True,
+        }
+
     @staticmethod
     def _base_moe_cfg(**overrides):
         cfg = {
@@ -636,6 +650,45 @@ def _base_moe_cfg(**overrides):
         cfg.update(overrides)
         return {"megatron_cfg": cfg}
 
+    @pytest.mark.parametrize("moe_grouped_gemm", [True, False])
+    def test_moe_grouped_gemm_explicit(self, moe_grouped_gemm):
+        """moe_grouped_gemm is applied when present in config."""
+        from nemo_rl.models.megatron.setup import _apply_moe_config
+
+        model_cfg = MagicMock()
+        megatron_cfg = self._base_moe_megatron_cfg()
+        megatron_cfg["moe_grouped_gemm"] = moe_grouped_gemm
+        config = {"megatron_cfg": megatron_cfg}
+
+        _apply_moe_config(model_cfg, config)
+
+        assert model_cfg.moe_grouped_gemm is moe_grouped_gemm
+
+    def test_moe_grouped_gemm_absent_keeps_default(self):
+        """Absent key leaves the attr unset on the model cfg."""
+        from nemo_rl.models.megatron.setup import _apply_moe_config
+
+        # spec lists everything _apply_moe_config writes so we can detect
+        # whether the moe_grouped_gemm branch fires.
+        model_cfg = MagicMock(
+            spec=[
+                "expert_tensor_parallel_size",
+                "expert_model_parallel_size",
+                "moe_router_dtype",
+                "moe_router_load_balancing_type",
+                "moe_router_bias_update_rate",
+                "moe_permute_fusion",
+                "moe_enable_deepep",
+                "moe_token_dispatcher_type",
+                "moe_shared_expert_overlap",
+            ]
+        )
+        config = {"megatron_cfg": self._base_moe_megatron_cfg()}
+
+        _apply_moe_config(model_cfg, config)
+
+        assert not hasattr(model_cfg, "moe_grouped_gemm")
+
     def test_hybridep_env_vars_auto_set_with_warning(self, monkeypatch):
         """HybridEP backend with no env config: auto-set env vars and emit warnings."""
         from nemo_rl.models.megatron.setup import _apply_moe_config