[Main] Support paged stash with device-initiated GroupedLinear

[lhb8125]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Extend MoE paged stashing to the non-op-fuser Transformer Engine GroupedTensor path introduced by #6000. This enables HybridEP static budgets and paged activation storage without requiring the CuTe DSL fused GroupedMLP or a full-iteration CUDA graph.

• Wrap non-fused grouped-expert compute in the existing paged-stash group lifecycle.
• Require the non-op-fuser paged-stash path to use device-initiated GroupedTensor plus fused SwiGLU or QuickGeGLU, with GLU interleaving disabled.
• Reject disabled fusion, non-gated activation, unsupported activation types, and GLU interleaving during TransformerConfig validation instead of falling back to raw activation code.
• Route dynamic-tensor marking through MCore's Transformer Engine extension adapter, which delegates to TE's mark_grouped_tensor utility.
• Rely on GroupedLinear to mark its own saved inputs. MCore now has only two marker call sites: the optional pre-FC1 probability multiply and the fused-activation input/probability boundary; the normal fused path executes only the latter.
• Propagate the marker across the fused activation's view/cast operations, including optional FP8 saved-input casts.
• Keep dynamic/static classification deterministic. The implementation deliberately does not infer dynamic tensors from warmup shapes because a false negative would make paged-stash capture unsafe.
• Keep the existing op-fuser paged-stash lifecycle unchanged and add lifecycle setup only to the non-op-fuser forward path.
• Add configuration, lifecycle, validation, marker-delegation, and four-rank HybridEP forward/backward coverage.
• Document both the op-fuser/full-CG and device-initiated/non-full-CG configurations, including the remaining CPU launch overhead and control-sync boundaries.

The activation restriction applies only when moe_paged_stash=True and use_transformer_engine_op_fuser=False. The op-fuser path and ordinary non-paged-stash unfused execution retain their existing activation support.

Dependency

Requires NVIDIA/TransformerEngine#3423.

Issue tracking

Follow-up to #6000.

Validation

AWS-CMH GB300 job 20260825-010956-b8a8 (Slurm 3284950) validated commit 7ef0a606b with the companion Transformer Engine change after narrowing marker propagation to the weighted SwiGLU/QuickGeGLU paths.

Configuration and marker coverage:

python -m pytest -q \
  tests/unit_tests/transformer/moe/test_grouped_mlp.py \
  -k "paged_stash_allows_non_fused_grouped_tensor_hybridep or \
      non_fused_grouped_tensor_paged_stash_requires_fused_bias_activation or \
      paged_stash_marking_delegates_to_transformer_engine or \
      fused_activation_marker_propagation_uses_te_adapter or \
      non_fused_forward_wraps_compute_in_paged_stash_scope" -s

Adapter-routing follow-up job 20260825-031156-a4e3 (Slurm 3286704) validated commit 0b0ec811d: 9 passed, 46 deselected. This includes four negative configurations—fusion disabled, non-gated activation, unsupported activation, and GLU interleaving—plus adapter delegation, both fused-activation lazy-import paths, and paged-stash scope lifecycle coverage.

Four-rank HybridEP, MXFP8, moe_use_grouped_tensor=True, fused SwiGLU, use_transformer_engine_op_fuser=False, and GPU-only paged stash (CPU factor = 0):

python -m torch.distributed.run --nproc_per_node=4 -m pytest -q \
  tests/unit_tests/transformer/moe/test_paged_stashing.py \
  -k "TestPagedStashingGroupedTensor and test_forward_backward_without_op_fuser" -s

All four ranks reported 1 passed, 10 deselected. The test verifies dynamic BF16 activation capture, allocated stash/reload execution, no overflow, and output/input-gradient parity.

End-to-end GB300 validation

Environment and recipe:

• OCI-AGA, 8 nodes x 4 Grace GB300 GPUs (32 GPUs), Slurm segment=4.
• TP1/PP2/EP16/CP4, GBS32, sequence length 16384, 22 layers, HybridEP, MXFP8, device-initiated GroupedTensor, fused SwiGLU, non-op-fuser, and combined 1F1B A2A overlap.
• Force-balanced routing was used to make the paged-stash ON/OFF comparison controlled.
• This validates the target GPU path, but the host platform differs from the customer environment: Grace ARM on GB300 here versus x86 on B300 there.

Numerical correctness

The paged-stash ON job 20260826-084247-40b2 (Slurm 575247) and OFF job 20260826-184106-325b (Slurm 580791) used the same image/source and differed only in moe_paged_stash. Both completed 3/3 steps with exit code 0.

Step	ON lm loss	OFF lm loss	ON grad norm	OFF grad norm
1	12.33325	12.33325	8.258	8.258
2	12.33430	12.33430	8.382	8.382
3	12.09118	12.09118	14.102	14.102

All printed lm loss and grad-norm values match exactly. Every ON step reported attempt=1 enabled=True; every OFF step reported attempt=1 enabled=False. Both sides had zero stash_overflow_ranks, overbudget_ranks, and host_spill_ranks, with will_retry=False, so the comparison contains no fallback/rerun attempt.

Paged-stash kernel execution

NSYS job 20260826-070950-c6a0 (Slurm 574726) completed 11/11 steps. Every committed step used paged stash on attempt 1 with no overflow, over-budget rank, host spill, or retry. Offline NSYS stats job 20260826-071734-6851 (Slurm 574761) found the following calls in the capture window:

Kernel	Calls	Total GPU time
paged_stash_copy_kernel	144	4.238912 ms
paged_stash_pop_kernel	144	5.288577 ms

The equal, nonzero copy/pop counts show that the final submitted attempt executed both stash and reload; this is stronger than checking the configuration flag alone. The rank-0 .nsys-rep SHA256 is 64920f7da4abfa61a7c3156f2817d1fbfb93059997a75e32f6b4568a8d9e3332.

GPU memory A/B

The same three-step A/B used moe_expert_rank_capacity_factor=2.0, moe_paged_stash_buffer_size_factor_cuda=1.1, CPU stash factor 0, and complete CUDA memory histories. Rank 31 provides the clean like-for-like comparison:

Metric	Paged stash ON	OFF	ON saving
Peak allocated	192641.07 MiB	194753.99 MiB	2112.92 MiB (1.08%)
Peak reserved	203222 MiB	205882 MiB	2660 MiB
CUDA stash pool	2668.31 MiB	0	n/a

The pool implies an actual-shape basis of S = 2668.3125 / 1.1 = 2425.74 MiB. The capacity model predicts (2.0 - 1.1) * S = 2183.16 MiB of peak allocated saving; the measured 2112.92 MiB is 96.8% of that prediction (70.24 MiB difference). The ON path keeps the fixed stash pool resident, so current allocation at a step boundary can be higher even though the runtime peak is lower.

Memory-history Python stack capture perturbs iteration timing, so these runs support correctness, actual kernel execution, and memory reduction only; they are not used to claim a throughput improvement.

Additional checks:

• CHECK_ONLY=true BASE_REF=main bash tools/autoformat.sh
• copyright check for all changed Python files
• Python 3.12 py_compile
• git diff --check
• clean merge-tree against current upstream main

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[lhb8125]

/ok to test df35d24

[lhb8125]

/ok to test d409603

[zhongbozhu]

approved, please share some E2E results showing matching loss and potentially better perf would be better