[sanketpurandare]

AI-Generated Details

Prepared with Codex from the author's explicit HybridEP and paging design.

Bind live counts, final-consumer stash commits, reloads, prefetches and waits to
source layer/chunk annotations. Preserve native requests, copy-input lifetimes
and each source's page issue order while scheduling independent compute. The
single high-priority paging stream serializes freelist mutations; independent
multiplexed F/B sources can interleave without an extra whole-backward-before-
forward chain. Per-slot records remain FIFO.

Waits remain late. For eager paging, reserve only the terminal backward chunk0's
wgrad-only branch until its own combine launches. Existing gradient/chunk/paging
annotations and dataflow identify that branch; shared experts, chunk1 and
unchunked PP remain eligible for earlier exchanges. No shape/name classifier,
new annotation key, timing threshold or compute resource limit.

Vectorize the existing copy/pop loops into page-aligned multi-token tiles. Resolve
one page per tile instead of duplicating page metadata per element. Only its first
tile records/recycles the page; live-prefix and page-tail masks preserve payload
bytes with 64-bit addressing. CUDA allocation, host spill, sticky overflow and
FIFO semantics are unchanged. TORCHTITAN_PAGED_STASH_MAX_CTAS, set before import,
bounds the grid (default 2048), not capacity or copied rows. Compute retains its
normal launch geometry. The copy layout no longer claims to be verbatim Megatron.

Paging composes with selective recomputation (sac_and_paged_stash) or explicit
save-all (none_and_paged_stash, activation_remat = None). Both benchmark arms
page activations; no-overlap controls expose paging overhead too. Deferred
overflow handling and device-side optimizer skipping stay outside the CUDA
callable. Calibration parks independent live prefixes, avoiding padded retention
before buffers exist, and reload restores the original capacity shape; full-count
inputs and unpaged reruns retain their identity.

Verification

Combined-stack tree 9542fd0f3872179b0f9c4042a8bb85107f56908b passed 69 focused
cases and 3 full-CUDA numerical comparisons: no-AC SPMD, no-AC combined PP and
DualPipeV with paging. Real 16B graph inspection checks all 26 terminal eager
combines before both expert wgrads. Local 16B SPMD/PP loss/norm pairs are exact;
active/reserved memory and replay kernel counts match their controls.

The existing paging suite exercises real byte copies, wrapped allocation and
recycling, changing device counts in CUDA replay, non-divisible/large hidden
sizes, BF16/FP8 payloads, host spill and overflow. Native graph tests use actual
source markers/paging passes, not patched metadata. Captured priorities/grids,
wait/input lifetimes and backend capacity contracts are checked.

Local 64-CTA PP profiling regressed throughput despite faster bounded copies,
so that scale setting was rejected. Resource bounds require actual model trace
and throughput evidence. New 32/64-GPU speedup and near-total hiding are not
claimed here; the target remains 1.2x. Owned lint/doc hooks passed; the untouched
video optional-time-base typing baseline is outside EP scope.