[Dev][Cherry-Pick Main] Remove legacy grouped gemm and grouped mlp

megatron/core/models/gpt/gpt_layer_specs.py

@@ -79,7 +79,6 @@ def get_gpt_layer_with_inference_submodules(
     qk_l2_norm: Optional[bool] = False,
     num_experts: Optional[int] = None,
     moe_grouped_gemm: Optional[bool] = False,
-    moe_use_legacy_grouped_gemm: Optional[bool] = False,
 ) -> TransformerLayerSubmodules:
     """Use these submodules for inference optimized linear layers.
     Args:
@@ -635,7 +634,6 @@ def get_gpt_decoder_layer_specs(
             qk_l2_norm=qk_l2_norm,
             num_experts=config.num_moe_experts,
             moe_grouped_gemm=config.moe_grouped_gemm,
-            moe_use_legacy_grouped_gemm=config.moe_use_legacy_grouped_gemm,
         )
     else:
         dense_layer_spec = get_gpt_layer_local_spec(

megatron/core/models/gpt/moe_module_specs.py

@@ -30,7 +30,6 @@ def get_moe_module_spec(
         use_te: Whether to use Transformer Engine.
         num_experts: Number of experts.
         moe_grouped_gemm: Whether to use grouped GEMM.
-        moe_use_legacy_grouped_gemm: Whether to use legacy grouped GEMM.
     """
     if use_te is not None and use_te:
         backend: BackendSpecProvider = TESpecProvider()

pyproject.toml

@@ -95,7 +95,6 @@ dev = [
     "mamba-ssm~=2.2",
     "causal-conv1d~=1.5",
     "flash-linear-attention~=0.4.0",
-    "nv-grouped-gemm~=1.1",
     "megatron-energon[av_decode]~=6.0",
     "av",
     "flashinfer-python>=0.5.0,<0.7.0",
@@ -118,7 +117,6 @@ lts = [
     "opentelemetry-api~=1.33.1",
     "mamba-ssm~=2.2",
     "causal-conv1d~=1.5",
-    "nv-grouped-gemm~=1.1",
     "megatron-energon[av_decode]~=6.0",
     "av",
     "flashinfer-python>=0.5.0,<0.7.0",
@@ -176,7 +174,6 @@ no_pypi_wheels = ["emerging_optimizers; python_version >= '3.12'", "fast-hadamar
 default-groups = ["linting", "build", "test"]
 no-build-isolation-package = [
     "causal-conv1d",
-    "nv-grouped-gemm",
     "mamba-ssm",
     "transformer-engine",
     "transformer-engine-torch",

tests/unit_tests/models/test_experimental_attention_variant_module_specs.py

@@ -78,7 +78,6 @@ def _make_config(**overrides):
         moe_layer_freq=1,
         num_moe_experts=None,
         moe_grouped_gemm=False,
-        moe_use_legacy_grouped_gemm=False,
         use_te_activation_func=False,
         pipeline_model_parallel_size=1,
         pipeline_model_parallel_layout=None,

tests/unit_tests/models/test_gpt_model.py

@@ -469,7 +469,6 @@ def test_get_transformer_layer_spec_forwards_use_te_activation_func():
     mock_args.qk_layernorm = False
     mock_args.multi_latent_attention = False
     mock_args.experimental_attention_variant = None
-    mock_args.moe_use_legacy_grouped_gemm = False
     mock_args.qk_l2_norm = False
 
     with (

tests/unit_tests/transformer/moe/test_paged_stashing.py

@@ -239,7 +239,6 @@ def test_forward_backward_4_layers(self):
             moe_flex_dispatcher_backend="hybridep",
             test_dtype=torch.bfloat16,
             moe_grouped_gemm=True,
-            moe_use_legacy_grouped_gemm=False,
             moe_paged_stash=True,
             moe_expert_rank_capacity_factor=1.5,
             use_transformer_engine_op_fuser=True,
@@ -329,7 +328,6 @@ def test_overload_factor_and_over_budget(self):
             moe_flex_dispatcher_backend="hybridep",
             test_dtype=torch.bfloat16,
             moe_grouped_gemm=True,
-            moe_use_legacy_grouped_gemm=False,
             moe_paged_stash=True,
             moe_expert_rank_capacity_factor=1.5,
             use_transformer_engine_op_fuser=True,