diff --git a/megatron/core/transformer/moe/ops/__init__.py b/megatron/core/transformer/moe/ops/__init__.py new file mode 100644 index 00000000000..b5dff7b5663 --- /dev/null +++ b/megatron/core/transformer/moe/ops/__init__.py @@ -0,0 +1 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. diff --git a/megatron/core/transformer/moe/ops/paged_stash.py b/megatron/core/transformer/moe/ops/paged_stash.py new file mode 100644 index 00000000000..19c4d0d874c --- /dev/null +++ b/megatron/core/transformer/moe/ops/paged_stash.py @@ -0,0 +1,258 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Triton kernels for MoE paged stash.""" + +import triton +import triton.language as tl + +GLOBAL_BLOCK_SIZE = 1024 + +__all__ = ["GLOBAL_BLOCK_SIZE", "paged_stash_copy_kernel", "paged_stash_pop_kernel"] + + +@triton.jit +def paged_stash_copy_kernel( + src_ptr, + cuda_dst_ptr, + host_dst_ptr, + num_tokens_ptr, + free_list_cuda_ptr, + free_list_host_ptr, + free_list_head_ptr, # shape (2,): [cuda_head, host_head] + free_list_tail_ptr, # shape (2,) + free_list_capacity_ptr, + page_record_ptr, + overflow_ptr, + host_spill_global_ptr, # 1 if any successful host spill (not set on overflow path) + spilled_to_host_ptr, # Output: 0 = stored in CUDA, 1 = stored in host or overflow + new_free_list_head_ptr, # Output: shape (2,) updated heads + PAGE_SIZE: tl.constexpr, + HIDDEN_SIZE: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + HAS_HOST_BUFFER: tl.constexpr, +): + """Stash variable-length MoE activations into a paged buffer (CUDA, or pinned host). + + Uses a custom Triton kernel because the token count is only known at runtime and + lives on device. Page allocation from the circular freelist, page_record metadata, + and the activation copy are fused in one GPU launch to avoid host sync and keep + stash CUDA-graph friendly. Fixed-size pages reduce fragmentation vs oversized + static expert buffers. + + Per launch (program 0 handles metadata; all programs run the copy): + 1. If overflow is already set, restore freelist heads and return. + 2. Compute pages needed from num_tokens. Try the CUDA freelist; if full, try + the host freelist when available; otherwise set overflow and return. + 3. Copy tokens in parallel: resolve page_id per token, record page_ids in + page_record, write hidden vectors into the chosen CUDA or host pages. + 4. Program 0 writes updated freelist heads for the caller to copy_ back. + """ + pid = tl.program_id(axis=0) + num_blocks = tl.num_programs(axis=0) + + # Load overflow first (get in flight early); branch on it only before any write + overflow = tl.load(overflow_ptr) + + num_tokens = tl.load(num_tokens_ptr) + required_pages = tl.cdiv(num_tokens, PAGE_SIZE) + + # Common case: load only CUDA state (and head_host for output when use_cuda) + head_cuda = tl.load(free_list_head_ptr) + head_host = tl.load(free_list_head_ptr + 1) + tail_cuda = tl.load(free_list_tail_ptr) + cap_cuda = tl.load(free_list_capacity_ptr) + + avail_cuda = tail_cuda - head_cuda + use_cuda = avail_cuda >= required_pages + + # Assume CUDA path: set everything for GPU stash + spill = 0 + dst_ptr = cuda_dst_ptr + free_list_ptr = free_list_cuda_ptr + head = head_cuda + cap = cap_cuda + new_head_cuda = head_cuda + required_pages + new_head_host = head_host + + if overflow == 1: + # No stash; preserve heads so Python copy_ does not write garbage into the buffer. + if pid == 0: + tl.store(new_free_list_head_ptr, head_cuda) + tl.store(new_free_list_head_ptr + 1, head_host) + return + + # Only when CUDA is full: load host state and maybe switch to host + if not use_cuda: + tail_host = tl.load(free_list_tail_ptr + 1) + cap_host = tl.load(free_list_capacity_ptr + 1) + use_host = HAS_HOST_BUFFER == 1 and (tail_host - head_host) >= required_pages + if use_host: + spill = 1 + dst_ptr = host_dst_ptr + free_list_ptr = free_list_host_ptr + head = head_host + cap = cap_host + new_head_cuda = head_cuda + new_head_host = head_host + required_pages + else: + if pid == 0: + tl.store(overflow_ptr, 1) + tl.store(spilled_to_host_ptr, 1) + tl.store(new_free_list_head_ptr, head_cuda) + tl.store(new_free_list_head_ptr + 1, head_host) + return + + if pid == 0: + tl.store(spilled_to_host_ptr, spill) + if spill == 1: + tl.store(host_spill_global_ptr, 1) + + # Copy loop: strided over tokens + token_idx = pid + while token_idx < num_tokens: + page_slot = token_idx // PAGE_SIZE + token_in_page = token_idx % PAGE_SIZE + free_list_idx = (head + page_slot) % cap + page_id = tl.load(free_list_ptr + free_list_idx) + if token_in_page == 0: + tl.store(page_record_ptr + page_slot, page_id) + dst_token_idx = page_id * PAGE_SIZE + token_in_page + + elements_per_thread = HIDDEN_SIZE // BLOCK_SIZE + need_mask = (HIDDEN_SIZE % BLOCK_SIZE) != 0 + num_iters = elements_per_thread + (1 if need_mask else 0) + token_idx_i64 = token_idx.to(tl.int64) + dst_token_idx_i64 = dst_token_idx.to(tl.int64) + src_base = src_ptr + token_idx_i64 * HIDDEN_SIZE + dst_base = dst_ptr + dst_token_idx_i64 * HIDDEN_SIZE + + if need_mask: + for iter in range(num_iters): + hidden_offsets = tl.arange(0, BLOCK_SIZE) + iter * BLOCK_SIZE + hidden_mask = hidden_offsets < HIDDEN_SIZE + data = tl.load(src_base + hidden_offsets, mask=hidden_mask, other=0) + tl.store(dst_base + hidden_offsets, data, mask=hidden_mask) + else: + for iter in range(elements_per_thread): + hidden_offsets = tl.arange(0, BLOCK_SIZE) + iter * BLOCK_SIZE + data = tl.load(src_base + hidden_offsets) + tl.store(dst_base + hidden_offsets, data) + token_idx += num_blocks + + if pid == 0: + tl.store(new_free_list_head_ptr, new_head_cuda) + tl.store(new_free_list_head_ptr + 1, new_head_host) + + +@triton.jit +def paged_stash_pop_kernel( + cuda_src_ptr, + host_src_ptr, + dst_ptr, + num_tokens_ptr, + page_record_ptr, + spilled_to_host_ptr, # 0 = read from CUDA, 1 = read from host + overflow_ptr, + free_list_cuda_ptr, + free_list_host_ptr, + free_list_tail_ptr, # shape (2,) + free_list_capacity_ptr, + new_free_list_tail_ptr, # Output: shape (2,) updated tails + PAGE_SIZE: tl.constexpr, + HIDDEN_SIZE: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + """Restore variable-length MoE activations from a paged buffer (CUDA, or pinned host). + + Inverse of paged_stash_copy_kernel. Uses a custom Triton kernel for the same + reasons: runtime token count and stash metadata live on device, so the reload, + page_record lookup, and freelist recycle must fuse on-GPU without host sync. + + Per launch (program 0 handles metadata; all programs run the copy): + 1. If overflow is already set, restore freelist tails and return. + 2. Read spilled_to_host from the matching stash: CUDA buffer by default, host + buffer when the forward stash spilled to pinned memory. + 3. Copy tokens in parallel: look up page_id from page_record, read hidden + vectors from the stash pages into dst, return each page_id to the freelist. + 4. Program 0 writes updated freelist tails for the caller to copy_ back. + """ + pid = tl.program_id(axis=0) + num_blocks = tl.num_programs(axis=0) + + # Load overflow first (get in flight early); branch on it only before any write + overflow = tl.load(overflow_ptr) + + num_tokens = tl.load(num_tokens_ptr) + spill = tl.load(spilled_to_host_ptr) + required_pages = tl.cdiv(num_tokens, PAGE_SIZE) + + # Common case: load only CUDA state (and tail_host for output when spill=0) + tail_cuda = tl.load(free_list_tail_ptr) + tail_host = tl.load(free_list_tail_ptr + 1) + cap_cuda = tl.load(free_list_capacity_ptr) + + if overflow == 1: + # No pop; preserve tails so Python copy_ does not write garbage into the buffer. + if pid == 0: + tl.store(new_free_list_tail_ptr, tail_cuda) + tl.store(new_free_list_tail_ptr + 1, tail_host) + return + + # Assume CUDA path + src_ptr = cuda_src_ptr + free_list_ptr = free_list_cuda_ptr + tail = tail_cuda + cap = cap_cuda + new_tail_cuda = tail_cuda + required_pages + new_tail_host = tail_host + + # Only when spilled to host: load host state and switch + if spill == 1: + cap_host = tl.load(free_list_capacity_ptr + 1) + if cap_host == 0: + # Cannot pop from host; preserve tails (no-op for free-list state). + if pid == 0: + tl.store(new_free_list_tail_ptr, tail_cuda) + tl.store(new_free_list_tail_ptr + 1, tail_host) + return + src_ptr = host_src_ptr + free_list_ptr = free_list_host_ptr + tail = tail_host + cap = cap_host + new_tail_cuda = tail_cuda + new_tail_host = tail_host + required_pages + + token_idx = pid + while token_idx < num_tokens: + page_slot = token_idx // PAGE_SIZE + token_in_page = token_idx % PAGE_SIZE + page_id = tl.load(page_record_ptr + page_slot) + src_token_idx = page_id * PAGE_SIZE + token_in_page + + elements_per_thread = HIDDEN_SIZE // BLOCK_SIZE + need_mask = (HIDDEN_SIZE % BLOCK_SIZE) != 0 + num_iters = elements_per_thread + (1 if need_mask else 0) + src_token_idx_i64 = src_token_idx.to(tl.int64) + token_idx_i64 = token_idx.to(tl.int64) + src_base = src_ptr + src_token_idx_i64 * HIDDEN_SIZE + dst_base = dst_ptr + token_idx_i64 * HIDDEN_SIZE + + if need_mask: + for iter in range(num_iters): + hidden_offsets = tl.arange(0, BLOCK_SIZE) + iter * BLOCK_SIZE + hidden_mask = hidden_offsets < HIDDEN_SIZE + data = tl.load(src_base + hidden_offsets, mask=hidden_mask, other=0) + tl.store(dst_base + hidden_offsets, data, mask=hidden_mask) + else: + for iter in range(elements_per_thread): + hidden_offsets = tl.arange(0, BLOCK_SIZE) + iter * BLOCK_SIZE + data = tl.load(src_base + hidden_offsets) + tl.store(dst_base + hidden_offsets, data) + + if token_in_page == 0: + write_idx = (tail + page_slot) % cap + tl.store(free_list_ptr + write_idx, page_id) + token_idx += num_blocks + + if pid == 0: + tl.store(new_free_list_tail_ptr, new_tail_cuda) + tl.store(new_free_list_tail_ptr + 1, new_tail_host) diff --git a/megatron/core/transformer/moe/paged_stash.py b/megatron/core/transformer/moe/paged_stash.py index 135b4804ad1..5bef3d321b5 100644 --- a/megatron/core/transformer/moe/paged_stash.py +++ b/megatron/core/transformer/moe/paged_stash.py @@ -1,21 +1,23 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import logging from contextlib import nullcontext from typing import Any import torch -import triton -import triton.language as tl from megatron.core._rank_utils import log_single_rank from megatron.core.full_cuda_graph import FullCudaGraphWrapper from megatron.core.optimizer.distrib_optimizer import DistributedOptimizer +from megatron.core.transformer.moe.ops.paged_stash import ( + GLOBAL_BLOCK_SIZE, + paged_stash_copy_kernel, + paged_stash_pop_kernel, +) from megatron.core.utils import get_attr_wrapped_model logger = logging.getLogger(__name__) -GLOBAL_BLOCK_SIZE = 1024 SCALE_INV_BLOCK_SIZE = 32 @@ -126,255 +128,6 @@ def __repr__(self): ) -@triton.jit -def _paged_stash_copy_kernel( - src_ptr, - cuda_dst_ptr, - host_dst_ptr, - num_tokens_ptr, - free_list_cuda_ptr, - free_list_host_ptr, - free_list_head_ptr, # shape (2,): [cuda_head, host_head] - free_list_tail_ptr, # shape (2,) - free_list_capacity_ptr, - page_record_ptr, - overflow_ptr, - host_spill_global_ptr, # 1 if any successful host spill (not set on overflow path) - spilled_to_host_ptr, # Output: 0 = stored in CUDA, 1 = stored in host or overflow - new_free_list_head_ptr, # Output: shape (2,) updated heads - PAGE_SIZE: tl.constexpr, - HIDDEN_SIZE: tl.constexpr, - BLOCK_SIZE: tl.constexpr, - HAS_HOST_BUFFER: tl.constexpr, -): - """Stash variable-length MoE activations into a paged buffer (CUDA, or pinned host). - - Uses a custom Triton kernel because the token count is only known at runtime and - lives on device. Page allocation from the circular freelist, page_record metadata, - and the activation copy are fused in one GPU launch to avoid host sync and keep - stash CUDA-graph friendly. Fixed-size pages reduce fragmentation vs oversized - static expert buffers. - - Per launch (program 0 handles metadata; all programs run the copy): - 1. If overflow is already set, restore freelist heads and return. - 2. Compute pages needed from num_tokens. Try the CUDA freelist; if full, try - the host freelist when available; otherwise set overflow and return. - 3. Copy tokens in parallel: resolve page_id per token, record page_ids in - page_record, write hidden vectors into the chosen CUDA or host pages. - 4. Program 0 writes updated freelist heads for the caller to copy_ back. - """ - pid = tl.program_id(axis=0) - num_blocks = tl.num_programs(axis=0) - - # Load overflow first (get in flight early); branch on it only before any write - overflow = tl.load(overflow_ptr) - - num_tokens = tl.load(num_tokens_ptr) - required_pages = tl.cdiv(num_tokens, PAGE_SIZE) - - # Common case: load only CUDA state (and head_host for output when use_cuda) - head_cuda = tl.load(free_list_head_ptr) - head_host = tl.load(free_list_head_ptr + 1) - tail_cuda = tl.load(free_list_tail_ptr) - cap_cuda = tl.load(free_list_capacity_ptr) - - avail_cuda = tail_cuda - head_cuda - use_cuda = avail_cuda >= required_pages - - # Assume CUDA path: set everything for GPU stash - spill = 0 - dst_ptr = cuda_dst_ptr - free_list_ptr = free_list_cuda_ptr - head = head_cuda - cap = cap_cuda - new_head_cuda = head_cuda + required_pages - new_head_host = head_host - - if overflow == 1: - # No stash; preserve heads so Python copy_ does not write garbage into the buffer. - if pid == 0: - tl.store(new_free_list_head_ptr, head_cuda) - tl.store(new_free_list_head_ptr + 1, head_host) - return - - # Only when CUDA is full: load host state and maybe switch to host - if not use_cuda: - tail_host = tl.load(free_list_tail_ptr + 1) - cap_host = tl.load(free_list_capacity_ptr + 1) - use_host = HAS_HOST_BUFFER == 1 and (tail_host - head_host) >= required_pages - if use_host: - spill = 1 - dst_ptr = host_dst_ptr - free_list_ptr = free_list_host_ptr - head = head_host - cap = cap_host - new_head_cuda = head_cuda - new_head_host = head_host + required_pages - else: - if pid == 0: - tl.store(overflow_ptr, 1) - tl.store(spilled_to_host_ptr, 1) - tl.store(new_free_list_head_ptr, head_cuda) - tl.store(new_free_list_head_ptr + 1, head_host) - return - - if pid == 0: - tl.store(spilled_to_host_ptr, spill) - if spill == 1: - tl.store(host_spill_global_ptr, 1) - - # Copy loop: strided over tokens - token_idx = pid - while token_idx < num_tokens: - page_slot = token_idx // PAGE_SIZE - token_in_page = token_idx % PAGE_SIZE - free_list_idx = (head + page_slot) % cap - page_id = tl.load(free_list_ptr + free_list_idx) - if token_in_page == 0: - tl.store(page_record_ptr + page_slot, page_id) - dst_token_idx = page_id * PAGE_SIZE + token_in_page - - elements_per_thread = HIDDEN_SIZE // BLOCK_SIZE - need_mask = (HIDDEN_SIZE % BLOCK_SIZE) != 0 - num_iters = elements_per_thread + (1 if need_mask else 0) - token_idx_i64 = token_idx.to(tl.int64) - dst_token_idx_i64 = dst_token_idx.to(tl.int64) - src_base = src_ptr + token_idx_i64 * HIDDEN_SIZE - dst_base = dst_ptr + dst_token_idx_i64 * HIDDEN_SIZE - - if need_mask: - for iter in range(num_iters): - hidden_offsets = tl.arange(0, BLOCK_SIZE) + iter * BLOCK_SIZE - hidden_mask = hidden_offsets < HIDDEN_SIZE - data = tl.load(src_base + hidden_offsets, mask=hidden_mask, other=0) - tl.store(dst_base + hidden_offsets, data, mask=hidden_mask) - else: - for iter in range(elements_per_thread): - hidden_offsets = tl.arange(0, BLOCK_SIZE) + iter * BLOCK_SIZE - data = tl.load(src_base + hidden_offsets) - tl.store(dst_base + hidden_offsets, data) - token_idx += num_blocks - - if pid == 0: - tl.store(new_free_list_head_ptr, new_head_cuda) - tl.store(new_free_list_head_ptr + 1, new_head_host) - - -@triton.jit -def _paged_stash_pop_kernel( - cuda_src_ptr, - host_src_ptr, - dst_ptr, - num_tokens_ptr, - page_record_ptr, - spilled_to_host_ptr, # 0 = read from CUDA, 1 = read from host - overflow_ptr, - free_list_cuda_ptr, - free_list_host_ptr, - free_list_tail_ptr, # shape (2,) - free_list_capacity_ptr, - new_free_list_tail_ptr, # Output: shape (2,) updated tails - PAGE_SIZE: tl.constexpr, - HIDDEN_SIZE: tl.constexpr, - BLOCK_SIZE: tl.constexpr, -): - """Restore variable-length MoE activations from a paged buffer (CUDA, or pinned host). - - Inverse of _paged_stash_copy_kernel. Uses a custom Triton kernel for the same - reasons: runtime token count and stash metadata live on device, so the reload, - page_record lookup, and freelist recycle must fuse on-GPU without host sync. - - Per launch (program 0 handles metadata; all programs run the copy): - 1. If overflow is already set, restore freelist tails and return. - 2. Read spilled_to_host from the matching stash: CUDA buffer by default, host - buffer when the forward stash spilled to pinned memory. - 3. Copy tokens in parallel: look up page_id from page_record, read hidden - vectors from the stash pages into dst, return each page_id to the freelist. - 4. Program 0 writes updated freelist tails for the caller to copy_ back. - """ - pid = tl.program_id(axis=0) - num_blocks = tl.num_programs(axis=0) - - # Load overflow first (get in flight early); branch on it only before any write - overflow = tl.load(overflow_ptr) - - num_tokens = tl.load(num_tokens_ptr) - spill = tl.load(spilled_to_host_ptr) - required_pages = tl.cdiv(num_tokens, PAGE_SIZE) - - # Common case: load only CUDA state (and tail_host for output when spill=0) - tail_cuda = tl.load(free_list_tail_ptr) - tail_host = tl.load(free_list_tail_ptr + 1) - cap_cuda = tl.load(free_list_capacity_ptr) - - if overflow == 1: - # No pop; preserve tails so Python copy_ does not write garbage into the buffer. - if pid == 0: - tl.store(new_free_list_tail_ptr, tail_cuda) - tl.store(new_free_list_tail_ptr + 1, tail_host) - return - - # Assume CUDA path - src_ptr = cuda_src_ptr - free_list_ptr = free_list_cuda_ptr - tail = tail_cuda - cap = cap_cuda - new_tail_cuda = tail_cuda + required_pages - new_tail_host = tail_host - - # Only when spilled to host: load host state and switch - if spill == 1: - cap_host = tl.load(free_list_capacity_ptr + 1) - if cap_host == 0: - # Cannot pop from host; preserve tails (no-op for free-list state). - if pid == 0: - tl.store(new_free_list_tail_ptr, tail_cuda) - tl.store(new_free_list_tail_ptr + 1, tail_host) - return - src_ptr = host_src_ptr - free_list_ptr = free_list_host_ptr - tail = tail_host - cap = cap_host - new_tail_cuda = tail_cuda - new_tail_host = tail_host + required_pages - - token_idx = pid - while token_idx < num_tokens: - page_slot = token_idx // PAGE_SIZE - token_in_page = token_idx % PAGE_SIZE - page_id = tl.load(page_record_ptr + page_slot) - src_token_idx = page_id * PAGE_SIZE + token_in_page - - elements_per_thread = HIDDEN_SIZE // BLOCK_SIZE - need_mask = (HIDDEN_SIZE % BLOCK_SIZE) != 0 - num_iters = elements_per_thread + (1 if need_mask else 0) - src_token_idx_i64 = src_token_idx.to(tl.int64) - token_idx_i64 = token_idx.to(tl.int64) - src_base = src_ptr + src_token_idx_i64 * HIDDEN_SIZE - dst_base = dst_ptr + token_idx_i64 * HIDDEN_SIZE - - if need_mask: - for iter in range(num_iters): - hidden_offsets = tl.arange(0, BLOCK_SIZE) + iter * BLOCK_SIZE - hidden_mask = hidden_offsets < HIDDEN_SIZE - data = tl.load(src_base + hidden_offsets, mask=hidden_mask, other=0) - tl.store(dst_base + hidden_offsets, data, mask=hidden_mask) - else: - for iter in range(elements_per_thread): - hidden_offsets = tl.arange(0, BLOCK_SIZE) + iter * BLOCK_SIZE - data = tl.load(src_base + hidden_offsets) - tl.store(dst_base + hidden_offsets, data) - - if token_in_page == 0: - write_idx = (tail + page_slot) % cap - tl.store(free_list_ptr + write_idx, page_id) - token_idx += num_blocks - - if pid == 0: - tl.store(new_free_list_tail_ptr, new_tail_cuda) - tl.store(new_free_list_tail_ptr + 1, new_tail_host) - - class PagedTensor: """ A paged tensor that stores data in pages within a paged stash buffer. @@ -466,7 +219,7 @@ def offload_to_stash(self, paged_stash_buffer: PagedStashBuffer, max_blocks=2048 else paged_stash_buffer.cuda_buffer ) - _paged_stash_copy_kernel[grid]( + paged_stash_copy_kernel[grid]( tensor_to_copy.view(paged_stash_buffer.cuda_buffer.dtype), paged_stash_buffer.cuda_buffer, host_dst, @@ -519,7 +272,7 @@ def reload_from_stash(self, paged_stash_buffer: PagedStashBuffer, max_blocks=204 if paged_stash_buffer.host_buffer is not None else paged_stash_buffer.cuda_buffer ) - _paged_stash_pop_kernel[grid]( + paged_stash_pop_kernel[grid]( paged_stash_buffer.cuda_buffer, host_src, self._tensor.view(paged_stash_buffer.cuda_buffer.dtype),