[Dev] Support paged stash with device-initiated GroupedTensor experts

docs/user-guide/features/paged_stash.md

@@ -13,28 +13,44 @@
 
 **Paged stash** = **sync-free** expert execution + **paged stashing** (packing routed-expert activations for backward into paged buffers).
 
-**Sync-free:** `--moe-flex-dispatcher-backend hybridep`, `--use-transformer-engine-op-fuser`, and `--moe-expert-rank-capacity-factor` pre-size dispatch and fused grouped expert buffers from a user-controlled capacity, avoiding a per-step device query / realloc loop for buffer sizing.
+**Sync-free:** `--moe-flex-dispatcher-backend hybridep` and `--moe-expert-rank-capacity-factor` pre-size dispatch and grouped expert buffers from a user-controlled capacity, avoiding a per-step device query / realloc loop for buffer sizing. Expert compute can use either `--use-transformer-engine-op-fuser` or the device-initiated Transformer Engine GroupedTensor path (`--moe-grouped-gemm --moe-use-grouped-tensor`).
 
 **Paged stashing:** `--moe-paged-stash` stores those activations in paged CUDA buffers (optional pinned host spill). It helps save activation memory; sync-free still works without it, at the cost of higher activation memory use.
 
 Whenever `moe_expert_rank_capacity_factor` is set, a **runner** wraps forward-backward: after each pass it checks **stash overflow** (only with `--moe-paged-stash`) and **token over-budget**. If either hits any rank, the step **reruns once** without capacity padding and without paged stashing.
 
 ## Prerequisites
 
-HybridEP + TE fused grouped experts are required whenever `moe_expert_rank_capacity_factor` is set. With `moe_paged_stash` enabled: capacity factor must be set; no `cpu_offloading`; `offload_modules` must not include `expert_fc1`, `moe_act`, or `fused_group_mlp`. The runner is active whenever capacity factor is set (even without `--moe-paged-stash`) for over-budget reruns; stash overflow is checked only when paged stashing is on.
+HybridEP and TE grouped experts are required whenever `moe_expert_rank_capacity_factor` is set. The non-op-fuser path requires a Transformer Engine version whose GroupedLinear marks saved GroupedTensor activation buffers for paged stashing. It currently supports only fused SwiGLU or QuickGeGLU (`bias_activation_fusion=True`) without GLU interleaving; restricting the activation contract keeps dynamic-tensor marking at the fused autograd boundaries. With `moe_paged_stash` enabled: capacity factor must be set; no `cpu_offloading`; `offload_modules` must not include `expert_fc1`, `moe_act`, or `fused_group_mlp`. The runner is active whenever capacity factor is set (even without `--moe-paged-stash`) for over-budget reruns; stash overflow is checked only when paged stashing is on.
 
 ## Configuration
 
 ```bash
-# Sync-free
+# Common static-budget configuration
+--moe-token-dispatcher-type flex
 --moe-flex-dispatcher-backend hybridep
---use-transformer-engine-op-fuser
 --moe-expert-rank-capacity-factor <float>
 
 # Paged stashing (to avoid memory waste due to fragmentation)
 --moe-paged-stash
+
+# Choose one expert-compute path:
+
+# A. TE operation fuser (used by the full-iteration CUDA graph + CuTe DSL route)
+--use-transformer-engine-op-fuser
+
+# B. Device-initiated GroupedLinear, without the operation fuser
+--moe-grouped-gemm
+--moe-use-grouped-tensor
+# Keep the default fused SwiGLU activation; do not pass --no-bias-swiglu-fusion.
 ```
 
+Path B removes host-device synchronization from grouped GEMM split metadata, but FC1, activation,
+and FC2 remain separate launches. Without a full-iteration CUDA graph it can therefore retain
+significant CPU launch overhead even though the expert path is host-device sync-free.
+The legacy multi-stream cuBLAS GroupedLinear path is not supported because it materializes split
+metadata on the host; paged stashing would not make that expert path sync-free.
+
 ## Tuning (paged stashing only)
 
 ```bash

megatron/core/extensions/transformer_engine.py

@@ -72,8 +72,14 @@
     import transformer_engine as te
     from transformer_engine.pytorch.fp8 import FP8GlobalStateManager, fp8_autocast, fp8_model_init
 
+    try:
+        from transformer_engine.pytorch.utils import mark_grouped_tensor as _te_mark_grouped_tensor
+    except ImportError:
+        _te_mark_grouped_tensor = None
+
     HAVE_TE = True
 except ImportError:
+    _te_mark_grouped_tensor = None
     if TYPE_CHECKING:
         # For type checking, treat transformer_engine as always available.
         import transformer_engine as te
@@ -89,6 +95,15 @@
 _TE_CONFIG_TYPE_KEY = "transformer_engine_config_type"
 
 
+def mark_grouped_tensor(*tensors: Any) -> None:
+    """Mark dynamic grouped tensors through the Transformer Engine compatibility boundary."""
+    if _te_mark_grouped_tensor is None:
+        raise RuntimeError(
+            "Paged stashing requires Transformer Engine's mark_grouped_tensor utility."
+        )
+    _te_mark_grouped_tensor(*tensors)
+
+
 class TransformerEngineConfigType(enum.Enum):
     """Configuration object types in config dictionary"""
 
@@ -1951,6 +1966,10 @@ def sharded_state_dict(
 
 if HAVE_TE and is_te_min_version("1.9.0.dev0"):
 
+    _TE_GROUPED_LINEAR_SUPPORTS_GROUPED_TENSOR = (
+        "use_grouped_tensor" in inspect.signature(te.pytorch.GroupedLinear.__init__).parameters
+    )
+
     class TEGroupedLinear(te.pytorch.GroupedLinear):
         """
         Wrapper for the Transformer-Engine's `GroupedLinear` layer.
@@ -2064,6 +2083,14 @@ def __init__(
                         config, "moe_single_grouped_bias", False
                     )
 
+            if _TE_GROUPED_LINEAR_SUPPORTS_GROUPED_TENSOR:
+                extra_kwargs["use_grouped_tensor"] = config.moe_use_grouped_tensor
+            elif config.moe_use_grouped_tensor and not config.use_transformer_engine_op_fuser:
+                raise RuntimeError(
+                    "moe_use_grouped_tensor=True requires a Transformer Engine GroupedLinear "
+                    "that exposes the use_grouped_tensor argument."
+                )
+
             self.te_quant_params: Optional[TEQuantizationParams] = None
             quant_config = get_quant_config_or_none(name, config.quant_recipe)
             self.finish_init(quant_config)
@@ -2565,6 +2592,7 @@ def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None):
             )
 
 else:
+    _TE_GROUPED_LINEAR_SUPPORTS_GROUPED_TENSOR = False
     TEGroupedLinear = None  # type: ignore[assignment, misc]
     TEColumnParallelGroupedLinear = None  # type: ignore[assignment, misc]
     TERowParallelGroupedLinear = None  # type: ignore[assignment, misc]

megatron/core/fusions/fused_bias_geglu.py

@@ -1,9 +1,20 @@
-# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
 import torch
 
 from megatron.core.jit import jit_fuser
 
+
+def _propagate_paged_stash_marker(source, target):
+    """Preserve TE's dynamic-activation marker across view/cast operations."""
+    if hasattr(source, "grouped_tensor_scale_inv"):
+        # Lazy import avoids the transformer_engine extension -> MLP -> fusion import cycle.
+        from megatron.core.extensions.transformer_engine import mark_grouped_tensor
+
+        mark_grouped_tensor(target)
+    return target
+
+
 ###### BIAS GELU FUSION/ NO AUTOGRAD ################
 # 1/sqrt(2*pi)-> 0.3989423
 # 1/sqrt(2)   -> 0.70710678
@@ -324,6 +335,7 @@ def forward(
             torch.Tensor: Output tensor of shape [N, H] after weighted Quick-GEGLU.
         """
         input_for_backward = input.to(torch.float8_e4m3fn) if fp8_input_store else input
+        _propagate_paged_stash_marker(input, input_for_backward)
         ctx.save_for_backward(input_for_backward, weights, linear_offset)
         ctx.ori_input_dtype = input.dtype
         ctx.fp8_input_store = fp8_input_store
@@ -374,6 +386,7 @@ def forward(
         """
         # Optionally store the input in FP8 for memory savings.
         input_for_backward = input.to(torch.float8_e4m3fn) if fp8_input_store else input
+        _propagate_paged_stash_marker(input, input_for_backward)
 
         # Save tensors for backward.
         ctx.save_for_backward(input_for_backward, bias, weights, linear_offset)
@@ -420,6 +433,7 @@ def weighted_bias_quick_geglu_impl(
         output: [num_selected_experts * seq_len, hidden_size]
     """
     ori_shape = input.shape
+    paged_stash_source = input
     assert len(ori_shape) in [2, 3]
     if clamp_value is not None:
         x_glu, x_linear = input.chunk(2, -1)
@@ -430,7 +444,7 @@ def weighted_bias_quick_geglu_impl(
             ),
             -1,
         )
-    input = input.view(-1, ori_shape[-1])
+    input = _propagate_paged_stash_marker(paged_stash_source, input.view(-1, ori_shape[-1]))
     linear_offset = torch.tensor(linear_offset, dtype=input.dtype, device=input.device)
     if bias is not None:
         output = WeightedBiasQuickGeGLUFunction.apply(

megatron/core/fusions/fused_bias_swiglu.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
 
 # pylint: disable=missing-function-docstring, missing-class-docstring
@@ -12,6 +12,16 @@
 ###### BIAS SWIGLU FUSION/ NO AUTOGRAD ################
 
 
+def _propagate_paged_stash_marker(source, target):
+    """Preserve TE's dynamic-activation marker across view/cast operations."""
+    if hasattr(source, "grouped_tensor_scale_inv"):
+        # Lazy import avoids the transformer_engine extension -> MLP -> fusion import cycle.
+        from megatron.core.extensions.transformer_engine import mark_grouped_tensor
+
+        mark_grouped_tensor(target)
+    return target
+
+
 @jit_fuser
 def swiglu(y):
     """Performs SwiGLU (Swish-Gated Linear Unit) activation function.
@@ -275,6 +285,7 @@ class WeightedSwiGLUFunction(torch.autograd.Function):
     @staticmethod
     def forward(ctx, input, weights, fp8_input_store, clamp_value):
         input_for_backward = input.to(torch.float8_e4m3fn) if fp8_input_store else input
+        _propagate_paged_stash_marker(input, input_for_backward)
         ctx.save_for_backward(input_for_backward, weights)
         ctx.ori_input_dtype = input.dtype
         ctx.fp8_input_store = fp8_input_store
@@ -339,7 +350,7 @@ def weighted_bias_swiglu_impl(input, bias, weights, fp8_input_store=False, clamp
     """
     ori_shape = input.shape
     assert len(ori_shape) in [2, 3]
-    input = input.view(-1, ori_shape[-1])
+    input = _propagate_paged_stash_marker(input, input.view(-1, ori_shape[-1]))
     if bias is not None:
         raise NotImplementedError("Bias is not supported for weighted swiglu fusion")
     else: