Clean up after rebase

Author: nanz-nv

megatron/core/fp8_utils.py

@@ -168,7 +168,7 @@ def _get_custom_recipe(quantizer_factory_python_path: str) -> Union[Fp8Recipe, F
 def get_fp8_align_size(fp8_recipe: Fp8Recipe) -> int:
     """Get the alignment size required for fp8 GEMM."""
     if fp8_recipe == Fp8Recipe.mxfp8:
-        return 128
+        return 32
     else:
         return 16
 

megatron/core/transformer/transformer_config.py

@@ -690,10 +690,6 @@ class TransformerConfig(ModelParallelConfig):
     GEMM feature introduced since CUTLASS 2.8 (https://github.com/fanshiqing/grouped_gemm).
     """
 
-    moe_use_device_initiated_grouped_gemm: bool = False
-    """Use the cutlass grouped gemm kernel, which allows for the token_per_expert tensor on GPU.
-    This can prevent the GPU-CPU synchronization during the grouped gemm."""
-
     moe_use_legacy_grouped_gemm: bool = False
     """Use legacy GroupedMLP rather than TEGroupedMLP.
     Note: The legacy one will be deprecated soon."""

tests/unit_tests/transformer/moe/test_paged_stashing.py

@@ -111,9 +111,6 @@ def __init__(
             moe_permute_fusion=kwargs.get("moe_permute_fusion", False),
             moe_flex_dispatcher_backend=kwargs.get("moe_flex_dispatcher_backend", None),
             moe_grouped_gemm=kwargs.get("moe_grouped_gemm", False),
-            moe_use_device_initiated_grouped_gemm=kwargs.get(
-                "moe_use_device_initiated_grouped_gemm", False
-            ),
             moe_use_legacy_grouped_gemm=kwargs.get("moe_use_legacy_grouped_gemm", False),
             moe_paged_stash=kwargs.get("moe_paged_stash", False),
             stash_modules=kwargs.get("stash_modules", None),
@@ -224,7 +221,6 @@ def test_forward_backward_4_layers(self):
             moe_flex_dispatcher_backend="hybridep",
             test_dtype=torch.bfloat16,
             moe_grouped_gemm=True,
-            moe_use_device_initiated_grouped_gemm=True,
             moe_use_legacy_grouped_gemm=False,
             moe_paged_stash=True,
             stash_modules=["expert_fc1", "moe_act", "expert_fc2"],
@@ -318,7 +314,6 @@ def test_overload_factor_and_over_budget(self):
             moe_flex_dispatcher_backend="hybridep",
             test_dtype=torch.bfloat16,
             moe_grouped_gemm=True,
-            moe_use_device_initiated_grouped_gemm=True,
             moe_use_legacy_grouped_gemm=False,
             moe_paged_stash=True,
             stash_modules=["expert_fc1", "moe_act", "expert_fc2"],