[Main] Support paged stash with device-initiated GroupedLinear

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

@@ -1,4 +1,4 @@
-# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 from __future__ import annotations
 
 import copy
@@ -70,8 +70,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
@@ -88,6 +94,15 @@
 _EXPERT_PARAMETER_NAME_PATTERN = re.compile(r"(weight|bias)\d*")
 
 
+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)
+
+
 def _set_expert_parameter_attributes(
     module: torch.nn.Module, parallel_mode: Optional[str], use_expert_pgs: bool
 ) -> None:

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.
@@ -267,6 +277,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
@@ -328,7 +339,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: