Is this page helpful?

MoE Paged Stash#

This is an experimental feature and may change.

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.

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.

Configuration#

# Sync-free
--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

Tuning (paged stashing only)#

# Page size for stashing
--moe-paged-stash-page-size 64
# CUDA stashing buffer scaling factor (default 1.10)
--moe-paged-stash-buffer-size-factor-cuda 1.10
# Host spill (0 = off); same sign rule as CUDA
--moe-paged-stash-buffer-size-factor-cpu 0.0

The copy/pop kernels use fixed launch sizes based on the current CUDA device: 4096 elements per program per iteration and at most 8192 programs on Rubin (SM107), and 1024 elements and at most 2048 programs on other GPUs. There are no configuration options for these launch sizes.

The stash copy/pop kernels are memory-bound; with 1024/2048, wide activations (thousands of elements per token) do not saturate DRAM bandwidth. For example, on VR200 with DeepSeek-V3 (hidden size 7168), using 4096/8192 makes the stash kernels ~2.3x faster.

What moe_expert_rank_capacity_factor and moe_paged_stash_buffer_size_factor_cuda mean#

Both are multipliers on buffer size relative to the perfectly balanced case—the space you would need if routed tokens were evenly distributed across expert ranks. A larger factor adds headroom for real-world skew.

Choosing moe_expert_rank_capacity_factor and stash buffer scales#

Profile how far real routing departs from the balanced reference, then pick factors so skew spikes rarely exceed your margin (avoid constant reruns).

  • moe_expert_rank_capacity_factor: pick from profiles so over-budget token drop is uncommon; set slightly above the profiled value so reruns stay rare.

  • moe_paged_stash_buffer_size_factor_cuda: size from the same stats (peaks vs averages) so stash overflow is uncommon; undersizing triggers reruns like over-budget.

  • moe_paged_stash_buffer_size_factor_cpu: set > 0 to allow spill to pinned host when CUDA pages are full—often avoids overflow / rerun at the cost of host memory and more overhead from paged stashing.