Skip to content

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

Draft
lhb8125 wants to merge 10 commits into
NVIDIA:devfrom
lhb8125:denliu/paged-stash-grouped-tensor-dev
Draft

[Dev] Support paged stash with device-initiated GroupedTensor experts#6845
lhb8125 wants to merge 10 commits into
NVIDIA:devfrom
lhb8125:denliu/paged-stash-grouped-tensor-dev

Conversation

@lhb8125

@lhb8125 lhb8125 commented Aug 25, 2026

Copy link
Copy Markdown
Contributor

Summary

  • Enable paged stash for the non-op-fuser TEGroupedMLP GroupedTensor path.
  • Mark dynamic saved tensors explicitly through the MCore Transformer Engine adapter.
  • Allow HybridEP static-budget paged stash with moe_use_grouped_tensor=True without requiring the TE op fuser.
  • Restrict the non-op-fuser paged-stash path to fused SwiGLU/QuickGeGLU activation and add focused coverage.

Dependency and merge order

  1. Merge [Dev] Support device-initiated grouped linear without TE op fuser #6847, the dev port of [Main] Support device-init grouped linear module too with when not using TE opfuser #6000.
  2. Refresh this branch against the updated dev branch.
  3. Merge this PR.

This branch currently contains the #6847 prerequisite commits so it remains buildable before that PR lands. Its effective diff will shrink after #6847 is merged, and the branch will then be refreshed before this PR is merged.

Transformer Engine dependencies

Test plan

  • CHECK_ONLY=true BASE_REF=dev ... bash tools/autoformat.sh (Black, isort, pylint, and ruff passed).
  • python tools/check_copyright.py <changed Python files>.
  • python -m py_compile on all changed Python files.
  • CI-faithful GB200 single-file run: tests/unit_tests/transformer/moe/test_paged_stashing.py exited 0. The current mcore-ci-dev image predates TE ci(fix): Parse unit tests in merge-queue #3224, so the new grouped-tensor case is dependency-gated with Installed TE GroupedLinear does not expose use_grouped_tensor instead of failing during module construction.

@copy-pr-bot

copy-pr-bot Bot commented Aug 25, 2026

Copy link
Copy Markdown

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.

zhongbozhu and others added 10 commits August 25, 2026 06:09
…ing TE opfuser (NVIDIA#6000)

Signed-off-by: Zhongbo Zhu <42691305+zhongbozhu@users.noreply.github.com>
Signed-off-by: hongbinl <hongbinl@nvidia.com>
Signed-off-by: hongbinl <hongbinl@nvidia.com>
Signed-off-by: hongbinl <hongbinl@nvidia.com>
Signed-off-by: hongbinl <hongbinl@nvidia.com>
Signed-off-by: hongbinl <hongbinl@nvidia.com>
Signed-off-by: hongbinl <hongbinl@nvidia.com>
Signed-off-by: hongbinl <hongbinl@nvidia.com>
Signed-off-by: hongbinl <hongbinl@nvidia.com>
Signed-off-by: hongbinl <hongbinl@nvidia.com>
@lhb8125
lhb8125 force-pushed the denliu/paged-stash-grouped-tensor-dev branch from 0da1b1b to e569fa3 Compare August 25, 2026 13:09
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants