Paged Stashing#2690

Draft

nanz-nv wants to merge 68 commits intoNVIDIA:devfrom

vasunvidia:paged_offloading

Commits on Mar 31, 2026

Add --moe-use-device-initiated-grouped-gemm to allow token_per_expert tensor on GPU

QiZhangNV
authored and
nanz-nv
committed
Initial change for packed offloading

vasunvidia
authored and
nanz-nv
committed
Bug fix

a
authored and
nanz-nv
committed
Mem Opt

vasunvidia
authored and
nanz-nv
committed
Handle MXFP8Tensor offload

a
authored and
nanz-nv
committed
Enable Packed offloading to CPU pinned memory with PACKED_OFFLOAD_CPU=1

a
authored and
nanz-nv
committed
Enable activation truncation for first step

a
authored and
nanz-nv
committed
Overflow check and assert

a
authored and
nanz-nv
committed
Check in temporary solution for detecing overflow in receiving buffer
nanz-nv
committed
Reconstruct the stash buffer into a 2D structure
nanz-nv
committed
Refactor the code to check overflow in HybridEP receiving buffer
nanz-nv
committed
Use CPU offloading context manager as a WAR for now to WAR the problem of overlap_grad_reduce
nanz-nv
committed
Add support for paged stashing
nanz-nv
committed
Add the feature of speculative CE stashing
nanz-nv
committed
Fix PP schedule

a
authored and
nanz-nv
committed
Use common buffer across VP for paged stashing

vasunvidia
authored and
nanz-nv
committed
Disable Packed Offloading for validation

a
authored and
nanz-nv
committed
Fixe perf issue in packed stash/pop kernels
nanz-nv
committed
Minor fix for tensor allocation and padding requirement on budget
nanz-nv
committed
Packed/paged offloading is current not stream-safe. Need to put stash/restore on the same stream
nanz-nv
committed
add new hybrid ep

Autumn1998
authored and
nanz-nv
committed
Remove the overflow check in framework because it is now done by hybridEP
nanz-nv
committed
Fix one merge conflict
nanz-nv
committed
Code cleanup

vasunvidia
authored and
nanz-nv
committed
Add second autograd to avoid triple buffering

vasunvidia
authored and
nanz-nv
committed
Avoid unnecessary wait_stream for reload in case of 1f1b

vasunvidia
authored and
nanz-nv
committed
Check in dynamic-shape-aware SwiGLU triton kernel
nanz-nv
committed
Major cleanup and refactor
nanz-nv
committed
Check in paged_stash.py that was omited in the previous commit
nanz-nv
committed
Remove d2d page feature for now
nanz-nv
committed
Update added arguments and add compatibility check
nanz-nv
committed
refine overflow check
nanz-nv
committed
Fixing lint issues
nanz-nv
committed
Minor refactor

vasunvidia
authored and
nanz-nv
committed
Add unit test for Paged Stashing

vasunvidia
authored and
nanz-nv
committed
1. allocate stashing buffer based on avg token count if STASH_BUFFER_SIZE_FACTOR is positive.
nanz-nv
committed
Reenable overlapping of stashing kernels
nanz-nv
committed
Remove a buggy/redundant reset
nanz-nv
committed
Cleanup moe-expert-rank-capacity-factor argument.

vasunvidia
authored and
nanz-nv
committed
Update moe_use_device_initiated_grouped_gemm check for paged stashing to use_transformer_engine_op_fuser

vasunvidia
authored and
nanz-nv
committed
Remove the WAR of running warmup on a side stream
nanz-nv
committed
Fix for data_iterator type check in Paged Stashing fallback

vasunvidia
authored and
nanz-nv
committed
Change to support eager-mode fallback for validation

vasunvidia
authored and
nanz-nv
committed
Revert "Check in dynamic-shape-aware SwiGLU triton kernel"
nanz-nv
committed
Fixed some minor issues
nanz-nv
committed
Fix the unit test
nanz-nv
committed
Initial commit for spill to cpu feature
nanz-nv
committed
Move paged stashing knobs from env vars to transformer_config knobs
nanz-nv
committed
Refactor the knobs a bit so it is more intuitive
nanz-nv
committed
Use get_attr_wrapped_model util to access moe and mtp layers

vasunvidia
authored and
nanz-nv
committed
Refactor the unit test for paged stashing
nanz-nv
committed
Clean up after rebase
nanz-nv
committed
skip routed expert padding

zhongbozhu
authored and
nanz-nv
committed
Refactor/clean-up logging
nanz-nv
committed
Resolve review feedback
nanz-nv
committed
Fix fallback data read for PP=1

vasunvidia
authored and
nanz-nv
committed
Paged stashing refactor

vasunvidia
authored and
nanz-nv
committed
Remove logical_shape check

vasunvidia
authored and
nanz-nv
committed
Remove paged_stash_set_last_layer

vasunvidia
authored and
nanz-nv
committed
Cleanup PadUnpadFunction

vasunvidia
authored and
nanz-nv
committed
Remove stash modules and remove stashing code for non-fused grouped gemm
nanz-nv
committed
Remove dead code
nanz-nv
committed
Fix TE import problem in experts.py
nanz-nv
committed
Fixed merge conflict
nanz-nv
committed
Address reviewer's comments
nanz-nv
committed

Commits on Apr 2, 2026

Commits on Apr 3, 2026

Release stashing buffer before fallback to restore the memory
nanz-nv
committed