Paged Stashing#2690

Draft

nanz-nv wants to merge 56 commits intoNVIDIA:devfrom

vasunvidia:paged_offloading

Commits on Mar 17, 2026

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

QiZhangNV
authored and
vasunvidia
committed
Initial change for packed offloading
vasunvidia
committed
Bug fix

a
authored and
vasunvidia
committed
Mem Opt
vasunvidia
committed
Handle MXFP8Tensor offload

a
authored and
vasunvidia
committed
Enable Packed offloading to CPU pinned memory with PACKED_OFFLOAD_CPU=1

a
authored and
vasunvidia
committed
Enable activation truncation for first step

a
authored and
vasunvidia
committed
Overflow check and assert

a
authored and
vasunvidia
committed
Check in temporary solution for detecing overflow in receiving buffer

nanz-nv
authored and
vasunvidia
committed
Reconstruct the stash buffer into a 2D structure

nanz-nv
authored and
vasunvidia
committed
Refactor the code to check overflow in HybridEP receiving buffer

nanz-nv
authored and
vasunvidia
committed
Use CPU offloading context manager as a WAR for now to WAR the problem of overlap_grad_reduce

nanz-nv
authored and
vasunvidia
committed
Add support for paged stashing

nanz-nv
authored and
vasunvidia
committed
Add the feature of speculative CE stashing

nanz-nv
authored and
vasunvidia
committed
Fix PP schedule

a
authored and
vasunvidia
committed
Use common buffer across VP for paged stashing
vasunvidia
committed
Disable Packed Offloading for validation

a
authored and
vasunvidia
committed
Fixe perf issue in packed stash/pop kernels

nanz-nv
authored and
vasunvidia
committed
Minor fix for tensor allocation and padding requirement on budget

nanz-nv
authored and
vasunvidia
committed
Packed/paged offloading is current not stream-safe. Need to put stash/restore on the same stream

nanz-nv
authored and
vasunvidia
committed
add new hybrid ep

Autumn1998
authored and
vasunvidia
committed
Remove the overflow check in framework because it is now done by hybridEP

nanz-nv
authored and
vasunvidia
committed
Fix one merge conflict

nanz-nv
authored and
vasunvidia
committed
Code cleanup
vasunvidia
committed
Add second autograd to avoid triple buffering
vasunvidia
committed
Avoid unnecessary wait_stream for reload in case of 1f1b
vasunvidia
committed
Check in dynamic-shape-aware SwiGLU triton kernel

nanz-nv
authored and
vasunvidia
committed
Major cleanup and refactor

nanz-nv
authored and
vasunvidia
committed
Check in paged_stash.py that was omited in the previous commit

nanz-nv
authored and
vasunvidia
committed
Remove d2d page feature for now

nanz-nv
authored and
vasunvidia
committed
Update added arguments and add compatibility check

nanz-nv
authored and
vasunvidia
committed
refine overflow check

nanz-nv
authored and
vasunvidia
committed
Fixing lint issues

nanz-nv
authored and
vasunvidia
committed
Minor refactor
vasunvidia
committed
Add unit test for Paged Stashing
vasunvidia
committed
Initial check in of a) force load imbalance b) log overload factors

nanz-nv
authored and
vasunvidia
committed
make overload factor logging work for cuda graph

nanz-nv
authored and
vasunvidia
committed
1. allocate stashing buffer based on avg token count if STASH_BUFFER_SIZE_FACTOR is positive.

nanz-nv
authored and
vasunvidia
committed
Reenable overlapping of stashing kernels

nanz-nv
authored and
vasunvidia
committed
Remove a buggy/redundant reset

nanz-nv
authored and
vasunvidia
committed
Cleanup moe-expert-rank-capacity-factor argument.
vasunvidia
committed
Support GEMM + Swiglu fused MLP

ksivaman
authored and
vasunvidia
committed
Update moe_use_device_initiated_grouped_gemm check for paged stashing to use_transformer_engine_op_fuser
vasunvidia
committed
support use-dynamic-comp-stream

Wohox
authored and
nanz-nv
committed
Revert "remove encoder_and_decoder from enums (#3406)"
nanz-nv
committed
Remove the WAR of running warmup on a side stream
nanz-nv
committed
Reapply "remove encoder_and_decoder from enums (#3406)"
nanz-nv
committed

Commits on Mar 18, 2026