From f134d391c6bca23a6390457e9e80a8ee188e0307 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BA=91=E6=8C=9A?= Date: Wed, 29 Jul 2026 23:22:03 -0500 Subject: [PATCH 1/5] Add owner-sharded activation publication paths --- csrc/bindings.cu | 302 +++++++++++ moonep/__init__.py | 10 + moonep/api.py | 979 +++++++++++++++++++++++++++++++++++- moonep/dispatch.py | 241 +++++---- moonep/objects.py | 272 ++++++++++ moonep/planning.py | 59 ++- tests/kernel_test_utils.py | 2 + tests/planning_reference.py | 2 +- tests/test_dispatch.py | 311 ++++++++++++ tests/test_objects.py | 91 ++++ tests/test_planning.py | 119 +++++ 11 files changed, 2279 insertions(+), 109 deletions(-) create mode 100644 moonep/objects.py create mode 100644 tests/test_objects.py diff --git a/csrc/bindings.cu b/csrc/bindings.cu index 6f00f4a..3c2cf25 100644 --- a/csrc/bindings.cu +++ b/csrc/bindings.cu @@ -1,6 +1,288 @@ #include +#include +#include +#include +#include +#include #include "nvl_shared_buffer.cuh" +namespace { + +__global__ void materialize_source_rows_kernel( + const uint4* __restrict__ mapped_source, + const int32_t* __restrict__ mapped_source_weights, + uint4* __restrict__ local_output, + int32_t* __restrict__ local_output_weights, + const int32_t* __restrict__ slot_to_source, + const int32_t* __restrict__ slot_to_primary, + const int32_t* __restrict__ slot_to_next, + int32_t NvS, + int32_t K, + int32_t source_rows_padded, + int32_t source_topk_padded, + bool with_weights, + int32_t hidden_vecs) { + for (int32_t loff = static_cast(blockIdx.x); + loff < NvS; + loff += static_cast(gridDim.x)) { + const int32_t encoded = slot_to_source[loff]; + if (encoded >= 0) { + const int32_t source_rank = encoded / NvS; + const int32_t source_topk_offset = encoded - source_rank * NvS; + const int32_t source_token = source_topk_offset / K; + if (with_weights && threadIdx.x == 0) { + local_output_weights[loff] = mapped_source_weights[ + static_cast(source_rank) * source_topk_padded + + source_topk_offset]; + } + if (slot_to_primary[loff] != loff) { + continue; + } + const int64_t source_row = + static_cast(source_rank) * source_rows_padded + + source_token; + const uint4* source = + mapped_source + source_row * hidden_vecs; + for (int32_t vec = static_cast(threadIdx.x); + vec < hidden_vecs; + vec += static_cast(blockDim.x)) { + const uint4 value = source[vec]; + int32_t target = loff; + for (int32_t fanout = 0; + target >= 0 && fanout < K; + ++fanout) { + local_output[ + static_cast(target) * hidden_vecs + vec + ] = value; + target = slot_to_next[target]; + } + } + } + } +} + +__global__ void zero_source_padding_kernel( + uint4* __restrict__ local_output, + int32_t* __restrict__ local_output_weights, + const int32_t* __restrict__ zero_fill_ranges, + int32_t zero_groups, + bool with_weights, + int32_t hidden_vecs) { + for (int32_t group = static_cast(blockIdx.x); + group < zero_groups; + group += static_cast(gridDim.x)) { + const int32_t pad_start = zero_fill_ranges[group * 2]; + const int32_t pad_count = zero_fill_ranges[group * 2 + 1]; + for (int32_t row = 0; row < pad_count; ++row) { + const int32_t loff = pad_start + row; + uint4* output = + local_output + static_cast(loff) * hidden_vecs; + for (int32_t vec = static_cast(threadIdx.x); + vec < hidden_vecs; + vec += static_cast(blockDim.x)) { + output[vec] = make_uint4(0, 0, 0, 0); + } + if (with_weights && threadIdx.x == 0) { + local_output_weights[loff] = 0; + } + } + } +} + +void materialize_source_rows( + torch::Tensor mapped_source, + torch::Tensor mapped_source_weights, + torch::Tensor local_output, + torch::Tensor local_output_weights, + torch::Tensor slot_to_source, + torch::Tensor slot_to_primary, + torch::Tensor slot_to_next, + torch::Tensor zero_fill_ranges, + int64_t NvS, + int64_t K, + int64_t source_rows_padded, + int64_t source_topk_padded, + bool with_weights, + int64_t num_blocks, + int64_t num_threads) { + TORCH_CHECK(mapped_source.is_cuda(), "mapped_source must be CUDA"); + TORCH_CHECK( + mapped_source_weights.is_cuda(), + "mapped_source_weights must be CUDA"); + TORCH_CHECK(local_output.is_cuda(), "local_output must be CUDA"); + TORCH_CHECK( + local_output_weights.is_cuda(), + "local_output_weights must be CUDA"); + TORCH_CHECK(slot_to_source.is_cuda(), "slot_to_source must be CUDA"); + TORCH_CHECK( + slot_to_primary.is_cuda(), + "slot_to_primary must be CUDA"); + TORCH_CHECK(slot_to_next.is_cuda(), "slot_to_next must be CUDA"); + TORCH_CHECK(zero_fill_ranges.is_cuda(), "zero_fill_ranges must be CUDA"); + TORCH_CHECK( + mapped_source.scalar_type() == torch::kBFloat16, + "mapped_source must be bf16"); + TORCH_CHECK( + local_output.scalar_type() == torch::kBFloat16, + "local_output must be bf16"); + TORCH_CHECK( + mapped_source_weights.scalar_type() == torch::kFloat32, + "mapped_source_weights must be fp32"); + TORCH_CHECK( + local_output_weights.scalar_type() == torch::kInt32, + "local_output_weights must be the int32 view of fp32"); + TORCH_CHECK( + slot_to_source.scalar_type() == torch::kInt32, + "slot_to_source must be int32"); + TORCH_CHECK( + slot_to_primary.scalar_type() == torch::kInt32, + "slot_to_primary must be int32"); + TORCH_CHECK( + slot_to_next.scalar_type() == torch::kInt32, + "slot_to_next must be int32"); + TORCH_CHECK( + zero_fill_ranges.scalar_type() == torch::kInt32, + "zero_fill_ranges must be int32"); + TORCH_CHECK(mapped_source.is_contiguous(), "mapped_source must be contiguous"); + TORCH_CHECK( + mapped_source_weights.is_contiguous(), + "mapped_source_weights must be contiguous"); + TORCH_CHECK(local_output.is_contiguous(), "local_output must be contiguous"); + TORCH_CHECK( + local_output_weights.is_contiguous(), + "local_output_weights must be contiguous"); + TORCH_CHECK(slot_to_source.is_contiguous(), "slot_to_source must be contiguous"); + TORCH_CHECK( + slot_to_primary.is_contiguous(), + "slot_to_primary must be contiguous"); + TORCH_CHECK(slot_to_next.is_contiguous(), "slot_to_next must be contiguous"); + TORCH_CHECK( + zero_fill_ranges.is_contiguous(), + "zero_fill_ranges must be contiguous"); + TORCH_CHECK(mapped_source.dim() == 2, "mapped_source must be 2-D"); + TORCH_CHECK(local_output.dim() == 2, "local_output must be 2-D"); + TORCH_CHECK(slot_to_source.dim() == 1, "slot_to_source must be 1-D"); + TORCH_CHECK( + mapped_source_weights.dim() == 1, + "mapped_source_weights must be 1-D"); + TORCH_CHECK( + local_output_weights.dim() == 1, + "local_output_weights must be 1-D"); + TORCH_CHECK( + slot_to_primary.dim() == 1, + "slot_to_primary must be 1-D"); + TORCH_CHECK(slot_to_next.dim() == 1, "slot_to_next must be 1-D"); + TORCH_CHECK( + zero_fill_ranges.dim() == 2 && zero_fill_ranges.size(1) == 2, + "zero_fill_ranges must be [groups, 2]"); + TORCH_CHECK(NvS > 0 && NvS <= INT32_MAX, "NvS must fit positive int32"); + TORCH_CHECK(K > 0 && K <= INT32_MAX, "K must fit positive int32"); + TORCH_CHECK( + source_rows_padded > 0 && source_rows_padded <= INT32_MAX, + "source_rows_padded must fit positive int32"); + TORCH_CHECK( + source_topk_padded > 0 && source_topk_padded <= INT32_MAX, + "source_topk_padded must fit positive int32"); + TORCH_CHECK( + num_blocks >= 0 && num_blocks <= INT32_MAX, + "num_blocks must be zero or fit positive int32"); + TORCH_CHECK( + num_threads == 64 || num_threads == 128 || num_threads == 256 + || num_threads == 512, + "num_threads must be one of 64, 128, 256, or 512"); + TORCH_CHECK( + slot_to_source.numel() == NvS, + "slot_to_source length must equal NvS"); + TORCH_CHECK( + slot_to_primary.numel() == NvS, + "slot_to_primary length must equal NvS"); + TORCH_CHECK( + slot_to_next.numel() == NvS, + "slot_to_next length must equal NvS"); + TORCH_CHECK( + local_output.size(0) == NvS, + "local_output first dimension must equal NvS"); + TORCH_CHECK( + local_output_weights.numel() == NvS, + "local_output_weights length must equal NvS"); + TORCH_CHECK( + mapped_source.size(1) == local_output.size(1), + "mapped_source and local_output hidden sizes must match"); + TORCH_CHECK( + mapped_source.size(0) % source_rows_padded == 0, + "mapped_source rows must be a multiple of source_rows_padded"); + const int64_t source_ranks = + mapped_source.size(0) / source_rows_padded; + TORCH_CHECK(source_ranks > 0, "mapped_source must contain a source rank"); + TORCH_CHECK( + mapped_source_weights.numel() + == source_ranks * source_topk_padded, + "mapped_source_weights length must match the mapped source ranks"); + TORCH_CHECK( + source_topk_padded >= K, + "source_topk_padded must be at least K"); + TORCH_CHECK( + mapped_source.size(1) % 8 == 0, + "hidden size must be divisible by 8 for uint4 vectorization"); + TORCH_CHECK( + mapped_source.size(1) / 8 <= INT32_MAX, + "vectorized hidden size must fit int32"); + TORCH_CHECK( + zero_fill_ranges.size(0) <= INT32_MAX, + "zero-fill group count must fit int32"); + TORCH_CHECK( + reinterpret_cast(mapped_source.data_ptr()) % 16 == 0 + && reinterpret_cast(local_output.data_ptr()) % 16 == 0, + "mapped_source and local_output must be 16-byte aligned"); + TORCH_CHECK( + mapped_source.device() == local_output.device() + && mapped_source.device() == mapped_source_weights.device() + && mapped_source.device() == local_output_weights.device() + && mapped_source.device() == slot_to_source.device() + && mapped_source.device() == slot_to_primary.device() + && mapped_source.device() == slot_to_next.device() + && mapped_source.device() == zero_fill_ranges.device(), + "all tensors must use the current CUDA device mapping"); + + const int32_t hidden_vecs = + static_cast(mapped_source.size(1) / 8); + const int threads = static_cast(num_threads); + const int blocks = static_cast( + num_blocks == 0 + ? std::min(NvS, 49152) + : std::min(NvS, num_blocks)); + const auto stream = at::cuda::getCurrentCUDAStream(); + materialize_source_rows_kernel<<>>( + reinterpret_cast(mapped_source.data_ptr()), + reinterpret_cast( + mapped_source_weights.data_ptr()), + reinterpret_cast(local_output.data_ptr()), + local_output_weights.data_ptr(), + slot_to_source.data_ptr(), + slot_to_primary.data_ptr(), + slot_to_next.data_ptr(), + static_cast(NvS), + static_cast(K), + static_cast(source_rows_padded), + static_cast(source_topk_padded), + with_weights, + hidden_vecs); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + const int64_t zero_groups = zero_fill_ranges.size(0); + const int zero_blocks = + static_cast(std::min(zero_groups, 4096)); + zero_source_padding_kernel<<>>( + reinterpret_cast(local_output.data_ptr()), + local_output_weights.data_ptr(), + zero_fill_ranges.data_ptr(), + static_cast(zero_groups), + with_weights, + hidden_vecs); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +} // namespace + /** * Return the VMM allocation granularity in bytes for the current device. */ @@ -57,4 +339,24 @@ PYBIND11_MODULE(_C, m) { m.def("nvl_release_mem_handle", &nvl_release_mem_handle, pybind11::arg("mem_handle"), "Release an owned mem handle returned by nvl_dist_alloc."); + m.def( + "materialize_source_rows", + &materialize_source_rows, + pybind11::arg("mapped_source"), + pybind11::arg("mapped_source_weights"), + pybind11::arg("local_output"), + pybind11::arg("local_output_weights"), + pybind11::arg("slot_to_source"), + pybind11::arg("slot_to_primary"), + pybind11::arg("slot_to_next"), + pybind11::arg("zero_fill_ranges"), + pybind11::arg("NvS"), + pybind11::arg("K"), + pybind11::arg("source_rows_padded"), + pybind11::arg("source_topk_padded"), + pybind11::arg("with_weights"), + pybind11::arg("num_blocks") = 0, + pybind11::arg("num_threads") = 64, + "Materialize a peer-indexed owner-sharded source view into the local " + "contiguous expert layout."); } diff --git a/moonep/__init__.py b/moonep/__init__.py index 31467ed..a738bd6 100644 --- a/moonep/__init__.py +++ b/moonep/__init__.py @@ -2,8 +2,18 @@ Buffer, ) from .planning import MoonEPCommPlan +from .objects import ( + MoonEPDispatchTarget, + MoonEPPreparedPublication, + MoonEPSourceStorage, + MoonEPSourceView, +) __all__ = [ "Buffer", "MoonEPCommPlan", + "MoonEPDispatchTarget", + "MoonEPPreparedPublication", + "MoonEPSourceStorage", + "MoonEPSourceView", ] diff --git a/moonep/api.py b/moonep/api.py index 8d0cbc7..4224c7a 100644 --- a/moonep/api.py +++ b/moonep/api.py @@ -54,6 +54,7 @@ import torch import torch.distributed as dist +from . import _C from .buffer import ( create_nvl_dist_tensor, create_nvl_dist_multicast_tensor, @@ -70,10 +71,23 @@ from .combine_prologue import launch_combine_prologue from .prefetch import launch_prefetch from .grad_reduce import launch_grad_reduce +from .objects import ( + MoonEPDispatchTarget, + MoonEPPreparedPublication, + MoonEPSourceStorage, + MoonEPSourceView, + _validate_tensor, +) logger = logging.getLogger(__name__) +def _require_bool(name: str, value: object) -> bool: + if not isinstance(value, bool): + raise TypeError(f"{name} must be bool, got {type(value).__name__}") + return value + + def _align_up(x: int, alignment: int) -> int: """Round x up to a multiple of alignment.""" return ((x + alignment - 1) // alignment) * alignment @@ -123,6 +137,14 @@ def format_nbytes(nbytes: int) -> str: meta_mapped_bytes = tensor_nbytes(ctx['meta_buf']) hidden_chunk_bytes = hidden_mapped_bytes // R meta_chunk_bytes = meta_mapped_bytes // R + source_chunk_bytes = ( + tensor_nbytes(ctx['source_buf']) // R + if ctx.get('source_buf') is not None else 0 + ) + source_weights_chunk_bytes = ( + tensor_nbytes(ctx['source_weights_buf']) // R + if ctx.get('source_weights_buf') is not None else 0 + ) local_temp_bytes = sum( tensor_nbytes(ctx[name]) for name in ( @@ -138,7 +160,13 @@ def format_nbytes(nbytes: int) -> str: ) ) - total_buffer_bytes = hidden_chunk_bytes + meta_chunk_bytes + local_temp_bytes + total_buffer_bytes = ( + hidden_chunk_bytes + + source_chunk_bytes + + source_weights_chunk_bytes + + meta_chunk_bytes + + local_temp_bytes + ) log_lines = [ "MoonEP comm buffer sizes:", @@ -152,6 +180,12 @@ def format_nbytes(nbytes: int) -> str: f" local_temp_buffer_size : {format_nbytes(local_temp_bytes)}", f" total_size : {format_nbytes(total_buffer_bytes)}", ] + if source_chunk_bytes: + log_lines[4:4] = [ + f" source_object_size : {format_nbytes(source_chunk_bytes)}", + " source_weights_size : " + f"{format_nbytes(source_weights_chunk_bytes)}", + ] logger.info("\n".join(log_lines)) @@ -219,6 +253,7 @@ def _create_context( num_sms: int | None = None, token_padding: int = 128, B: int | None = None, + enable_source_object: bool = False, group: "dist.ProcessGroup | None" = None, ) -> dict: """Pre-allocate all NVLink shared buffers and local temp buffers. @@ -228,6 +263,10 @@ def _create_context( still aligned to VMM granularity. NVLink shared buffer layout: - hidden_buf: [NvS_padded, H] bf16 — separate allocation, padded for VMM alignment + - source_buf: optional [source_rows_padded, H] bf16 owner shard, + globally mapped when enable_source_object=True + - source_weights_buf: optional [source_topk_padded] fp32 route-weight + owner shard, globally mapped with source_buf - meta_buf: [meta_chunk_padded] int32 — merged allocation containing: [0, NvS): weights (fp32, alias as int32) [NvS, NvS + R*E): tpe gather (only rank 0 is the read/write target; symmetric reserve) @@ -264,6 +303,7 @@ def _create_context( if B is None: B = epn assert isinstance(B, int) and B > 0, f"B must be a positive int, got {B}" + _require_bool("enable_source_object", enable_source_object) NvS_capacity = S * K @@ -321,7 +361,13 @@ def _create_context( # counter (independent of rank count, self-resetting double buffer). BARRIER_SLOTS = 3 SRC_INFO_OFF = BARRIER_OFF + BARRIER_SLOTS - meta_chunk_logical = SRC_INFO_OFF + NvS + PRIMARY_INFO_OFF = SRC_INFO_OFF + NvS + NEXT_INFO_OFF = PRIMARY_INFO_OFF + NvS + meta_chunk_logical = ( + NEXT_INFO_OFF + NvS + if enable_source_object + else PRIMARY_INFO_OFF + ) # Chunk size must satisfy both VMM mapping and multicast binding. gran = get_vmm_granularity() @@ -350,6 +396,34 @@ def _create_context( # Allocate NVLink shared buffers # ================================================================ hidden_buf = create_nvl_dist_tensor([NvS_padded, H], torch.bfloat16, rank, R, group=group) + source_rows_padded = ( + pad_dim0_for_alignment([S, H], torch.bfloat16) + if enable_source_object else 0 + ) + source_buf = ( + create_nvl_dist_tensor( + [source_rows_padded, H], + torch.bfloat16, + rank, + R, + group=group, + ) + if enable_source_object else None + ) + source_topk_padded = ( + pad_dim0_for_alignment([N], torch.float32) + if enable_source_object else 0 + ) + source_weights_buf = ( + create_nvl_dist_tensor( + [source_topk_padded], + torch.float32, + rank, + R, + group=group, + ) + if enable_source_object else None + ) meta_buf, meta_mc = create_nvl_dist_multicast_tensor( [meta_chunk_padded], torch.int32, rank, R, group=group, ) @@ -387,6 +461,7 @@ def _create_context( 'group': group, 'R': R, 'E': E, 'S': S, 'K': K, 'H': H, 'B': B, + 'enable_source_object': enable_source_object, 'N': N, 'NvS': NvS, 'NvS_capacity': NvS_capacity, 'NvS_padded': NvS_padded, @@ -397,6 +472,8 @@ def _create_context( 'token_padding_extra': token_padding_extra, # NVLink shared 'hidden_buf': hidden_buf, + 'source_buf': source_buf, + 'source_weights_buf': source_weights_buf, 'meta_buf': meta_buf, 'meta_mc': meta_mc, # Meta layout @@ -409,6 +486,10 @@ def _create_context( 'ORDER0_OFF': ORDER0_OFF, 'BARRIER_OFF': BARRIER_OFF, 'SRC_INFO_OFF': SRC_INFO_OFF, + 'PRIMARY_INFO_OFF': PRIMARY_INFO_OFF, + 'NEXT_INFO_OFF': NEXT_INFO_OFF, + 'source_rows_padded': source_rows_padded, + 'source_topk_padded': source_topk_padded, # Local temps 'alloc': alloc, 'group_tokens': group_tokens, 'z': z, @@ -421,6 +502,38 @@ def _create_context( 'num_vblocks': num_vblocks, # Local views 'hidden_buf_local': hidden_buf[rank * NvS_padded: rank * NvS_padded + NvS], + 'source_buf_local': ( + source_buf[ + rank * source_rows_padded: + rank * source_rows_padded + S + ] + if source_buf is not None else None + ), + 'source_weights_buf_local': ( + source_weights_buf[ + rank * source_topk_padded: + rank * source_topk_padded + N + ].view(S, K) + if source_weights_buf is not None else None + ), + 'source_info_local': meta_buf[ + rank * meta_chunk_padded + SRC_INFO_OFF: + rank * meta_chunk_padded + SRC_INFO_OFF + NvS + ], + 'source_primary_local': ( + meta_buf[ + rank * meta_chunk_padded + PRIMARY_INFO_OFF: + rank * meta_chunk_padded + PRIMARY_INFO_OFF + NvS + ] + if enable_source_object else None + ), + 'source_next_local': ( + meta_buf[ + rank * meta_chunk_padded + NEXT_INFO_OFF: + rank * meta_chunk_padded + NEXT_INFO_OFF + NvS + ] + if enable_source_object else None + ), 'weights_buf_local': meta_buf[rank * meta_chunk_padded + WEIGHTS_OFF: rank * meta_chunk_padded + WEIGHTS_OFF + NvS], } @@ -450,6 +563,7 @@ def __init__( comm_stream_priority: int = -1, enable_pdl: bool = True, explicitly_destroy: bool = False, + enable_source_object: bool = False, ): """Allocate and hold all communication buffers. @@ -473,6 +587,11 @@ def __init__( False falls back to plain same-stream serial launches. explicitly_destroy: if True, warn (instead of auto-destroying) when the Buffer is garbage-collected without ``destroy()``. + enable_source_object: allocate an opt-in owner-sharded source + activation object with one physical ``[S, H]`` shard per rank. + This enables ``source_storage`` / ``prepare_source_dispatch`` + for a consumer explicitly supporting peer-indexed VMM loads. + The default False adds no source-object allocation. """ assert isinstance(comm_stream_priority, int), ( f"comm_stream_priority must be an int, got " @@ -486,11 +605,15 @@ def __init__( self.enable_pdl = enable_pdl self._comm_stream: torch.cuda.Stream | None = None self._destroyed = False + self._dispatch_generation = 0 + self._published_dispatch_generation: int | None = None + self._prepared_dispatch_generation: int | None = None self._ctx = _create_context( S, H, K, E, num_ep_ranks, num_sms=num_sms, token_padding=token_padding, B=B, + enable_source_object=enable_source_object, group=group, ) self._comm_stream = torch.cuda.Stream( @@ -530,9 +653,16 @@ def destroy(self) -> None: # multicast handles once the Python references disappear. for key in ( 'hidden_buf_local', + 'source_buf_local', + 'source_weights_buf_local', + 'source_info_local', + 'source_primary_local', + 'source_next_local', 'weights_buf_local', 'meta_mc', 'meta_buf', + 'source_buf', + 'source_weights_buf', 'hidden_buf', 'alloc', 'group_tokens', @@ -682,6 +812,845 @@ def _run_prefetch_weight_on_current_stream( experts_to_copy[int(ctx['rank'])], ) + def source_storage(self) -> MoonEPSourceStorage: + """Return the opt-in owner-sharded source activation object. + + A producer writes its natural output to ``local_hidden``. The router + may consume that local view normally and writes or copies its route + weights to ``local_route_weights``. After routing, pass this exact + object to ``prepare_source_dispatch`` to obtain the peer-indexed + consumer view without materializing routed hidden or weight rows. + """ + ctx = self._require_ctx() + if not bool(ctx['enable_source_object']): + raise RuntimeError( + "source_storage requires Buffer(enable_source_object=True)" + ) + source_buf = ctx['source_buf'] + source_local = ctx['source_buf_local'] + source_weights = ctx['source_weights_buf'] + source_weights_local = ctx['source_weights_buf_local'] + if source_buf is None or source_local is None: + raise RuntimeError("source hidden storage is not initialized") + if source_weights is None or source_weights_local is None: + raise RuntimeError("source route-weight storage is not initialized") + storage = MoonEPSourceStorage( + mapped_hidden=source_buf, + local_hidden=source_local, + mapped_route_weights=source_weights, + local_route_weights=source_weights_local, + buffer_token=id(self), + R=int(ctx['R']), + S=int(ctx['S']), + H=int(ctx['H']), + K=int(ctx['K']), + source_rows_padded=int(ctx['source_rows_padded']), + source_topk_padded=int(ctx['source_topk_padded']), + ) + storage.validate_layout() + return storage + + def _validate_source_storage( + self, + ctx: dict, + storage: MoonEPSourceStorage, + ) -> None: + if not isinstance(storage, MoonEPSourceStorage): + raise TypeError("storage must be returned by Buffer.source_storage") + storage.validate_layout() + if storage.buffer_token != id(self): + raise ValueError( + "source storage belongs to a different MoonEP Buffer" + ) + source_buf = ctx['source_buf'] + source_local = ctx['source_buf_local'] + source_weights = ctx['source_weights_buf'] + source_weights_local = ctx['source_weights_buf_local'] + if ( + source_buf is None + or source_local is None + or source_weights is None + or source_weights_local is None + ): + raise RuntimeError("source object storage is not initialized") + if storage.mapped_hidden.data_ptr() != source_buf.data_ptr(): + raise ValueError("source storage mapped-hidden view does not match") + if storage.local_hidden.data_ptr() != source_local.data_ptr(): + raise ValueError("source storage local-hidden view does not match") + if ( + storage.mapped_route_weights.data_ptr() + != source_weights.data_ptr() + ): + raise ValueError( + "source storage mapped route-weight view does not match" + ) + if ( + storage.local_route_weights.data_ptr() + != source_weights_local.data_ptr() + ): + raise ValueError( + "source storage local route-weight view does not match" + ) + + def prepare_source_dispatch( + self, + storage: MoonEPSourceStorage, + topk_experts_sk: torch.Tensor, + tokens_per_expert: torch.Tensor, + async_finish: bool = False, + *, + inter_rank_sync: bool = True, + ): + """Publish a one-copy source object as a peer-indexed expert view. + + The producer/router must have completed writes to + ``storage.local_hidden`` and ``storage.local_route_weights`` on the + current stream. Planning retains one physical source row per token and + one physical route-weight row per source rank, then publishes + ``slot_to_source`` provenance, ``slot_to_primary`` representative + reuse, and ``slot_to_next`` duplicate chains for every valid local + expert-layout slot. No routed hidden payload is copied or fanned out. + + With ``async_finish=True``, planning runs on the communication stream + and a completion event is appended to the return tuple. The consumer + may do independent work on its current stream, then wait on that event + immediately before dereferencing the returned view. + + The returned view requires a consumer declaring + ``peer_indexed_source_v1``. Ordinary MoonEP dispatch remains the + fallback for consumers that require contiguous materialization. + + Returns ``(view, cu_seqlens, plan)``, plus a communication-stream CUDA + event when ``async_finish=True``. + """ + ctx = self._require_ctx() + self._validate_source_storage(ctx, storage) + _require_bool("inter_rank_sync", inter_rank_sync) + _require_bool("async_finish", async_finish) + if topk_experts_sk is None or tokens_per_expert is None: + raise ValueError( + "topk_experts_sk and tokens_per_expert are required" + ) + if not isinstance(topk_experts_sk, torch.Tensor): + raise TypeError("topk_experts_sk must be a torch.Tensor") + topk_flat = topk_experts_sk.reshape(-1) + _validate_tensor( + "topk_experts_sk", + topk_flat, + dtype=torch.int32, + shape=(int(ctx['N']),), + device=storage.local_hidden.device, + ) + _validate_tensor( + "tokens_per_expert", + tokens_per_expert, + dtype=torch.int32, + shape=(int(ctx['E']),), + device=storage.local_hidden.device, + ) + + self._dispatch_generation += 1 + self._published_dispatch_generation = None + self._prepared_dispatch_generation = None + plan, cu_seqlens = allocate_planning_outputs(ctx) + view = MoonEPSourceView( + storage=storage, + slot_to_source=ctx['source_info_local'], + slot_to_primary=ctx['source_primary_local'], + slot_to_next=ctx['source_next_local'], + plan=plan, + generation=self._dispatch_generation, + NvS=int(ctx['NvS']), + K=int(ctx['K']), + ) + view.validate_layout() + + def prepare_on_current_stream() -> None: + if inter_rank_sync: + # This barrier follows the producer on the same stream and + # publishes its local owner-shard writes before peer consumers. + launch_inter_rank_sync(ctx) + launch_planning( + ctx, + topk_flat, + tokens_per_expert, + cu_seqlens, + plan, + ) + + if not async_finish: + prepare_on_current_stream() + return view, cu_seqlens, plan + + main_stream = torch.cuda.current_stream() + comm = self._comm_stream + if comm is None: + raise RuntimeError( + "MoonEP Buffer communication stream is not initialized" + ) + self._record_streams( + ( + storage.mapped_hidden, + storage.local_hidden, + storage.mapped_route_weights, + storage.local_route_weights, + topk_flat, + tokens_per_expert, + cu_seqlens, + view.slot_to_source, + view.slot_to_primary, + view.slot_to_next, + *self._plan_runtime_tensors(plan), + ), + comm, + ) + input_ready = main_stream.record_event() + comm.wait_event(input_ready) + with torch.cuda.stream(comm): + prepare_on_current_stream() + done = comm.record_event() + return view, cu_seqlens, plan, done + + def _validate_source_view( + self, + ctx: dict, + view: MoonEPSourceView, + plan: MoonEPCommPlan, + ) -> None: + if not isinstance(view, MoonEPSourceView): + raise TypeError( + "view must be returned by Buffer.prepare_source_dispatch" + ) + view.validate_layout() + self._validate_source_storage(ctx, view.storage) + if view.generation != self._dispatch_generation: + raise RuntimeError( + "source view is stale: this Buffer has started another " + "dispatch, source preparation, direct-out preparation, or " + "combine operation" + ) + if view.plan is not plan: + raise ValueError( + "plan must be the exact plan carried by the source view" + ) + if ( + view.slot_to_source.data_ptr() + != ctx['source_info_local'].data_ptr() + or view.slot_to_primary.data_ptr() + != ctx['source_primary_local'].data_ptr() + or view.slot_to_next.data_ptr() + != ctx['source_next_local'].data_ptr() + ): + raise ValueError("source view metadata does not match this Buffer") + + def _run_source_materialize_on_current_stream( + self, + ctx: dict, + view: MoonEPSourceView, + plan: MoonEPCommPlan, + with_route_weights: bool, + hidden_nvsh: torch.Tensor, + route_weights_nvs: torch.Tensor | None, + *, + zero_copy: bool, + num_blocks: int, + num_threads: int, + ) -> None: + # Pull one representative hidden row per source-token/consumer-rank + # and fan each loaded vector directly from registers into its linked + # same-rank top-k slots. Route weights remain per-slot peer loads; + # padding uses compact planned ranges. No dispatch metadata kernel, + # dedup builder, cross-rank payload barrier, or local hidden reread is + # needed. + _C.materialize_source_rows( + view.storage.mapped_hidden, + view.storage.mapped_route_weights, + ctx['hidden_buf_local'], + ctx['weights_buf_local'], + view.slot_to_source, + view.slot_to_primary, + view.slot_to_next, + plan.zero_fill_ranges, + int(ctx['NvS']), + int(ctx['K']), + int(ctx['source_rows_padded']), + int(ctx['source_topk_padded']), + with_route_weights, + num_blocks, + num_threads, + ) + if not zero_copy: + hidden_nvsh.copy_(ctx['hidden_buf_local']) + if route_weights_nvs is not None: + route_weights_nvs.copy_( + ctx['weights_buf_local'].view(torch.float32) + ) + + def materialize_source_dispatch( + self, + view: MoonEPSourceView, + plan: MoonEPCommPlan, + with_route_weights: bool = False, + async_finish: bool = False, + *, + zero_copy: bool = True, + num_blocks: int = 0, + num_threads: int = 64, + ): + """Materialize a peer-indexed source view into the local expert layout. + + This is the runtime fallback for a consumer that understands the + object contract but still benefits from a contiguous execution view. + Each rank explicitly pulls only its representative planned rows from + their physical source owners and fans each loaded vector directly into + repeated local destinations. It avoids rank-to-rank payload dispatch + and does not require generic CUDA operations to dereference peer VMM + rows. + """ + ctx = self._require_ctx() + self._validate_source_view(ctx, view, plan) + if not isinstance(num_blocks, int) or isinstance(num_blocks, bool): + raise TypeError("num_blocks must be an int") + if num_blocks < 0: + raise ValueError( + "num_blocks must be non-negative; zero selects the default" + ) + if not isinstance(num_threads, int) or isinstance(num_threads, bool): + raise TypeError("num_threads must be an int") + if num_threads not in (64, 128, 256, 512): + raise ValueError( + "num_threads must be one of 64, 128, 256, or 512" + ) + _require_bool("with_route_weights", with_route_weights) + _require_bool("async_finish", async_finish) + _require_bool("zero_copy", zero_copy) + + if zero_copy: + hidden_nvsh = ctx['hidden_buf_local'] + route_weights_nvs = ( + ctx['weights_buf_local'].view(torch.float32) + if with_route_weights else None + ) + else: + hidden_nvsh = torch.empty_like(ctx['hidden_buf_local']) + route_weights_nvs = ( + torch.empty( + int(ctx['NvS']), + dtype=torch.float32, + device=ctx['meta_buf'].device, + ) + if with_route_weights else None + ) + + if not async_finish: + self._run_source_materialize_on_current_stream( + ctx, + view, + plan, + with_route_weights, + hidden_nvsh, + route_weights_nvs, + zero_copy=zero_copy, + num_blocks=num_blocks, + num_threads=num_threads, + ) + return hidden_nvsh, route_weights_nvs, None + + main_stream = torch.cuda.current_stream() + comm = self._comm_stream + if comm is None: + raise RuntimeError( + "MoonEP Buffer communication stream is not initialized" + ) + self._record_streams( + ( + view.storage.mapped_hidden, + view.storage.local_hidden, + view.storage.mapped_route_weights, + view.storage.local_route_weights, + view.slot_to_source, + view.slot_to_primary, + view.slot_to_next, + hidden_nvsh, + route_weights_nvs, + *self._plan_runtime_tensors(plan), + ), + comm, + ) + input_ready = main_stream.record_event() + comm.wait_event(input_ready) + with torch.cuda.stream(comm): + self._run_source_materialize_on_current_stream( + ctx, + view, + plan, + with_route_weights, + hidden_nvsh, + route_weights_nvs, + zero_copy=zero_copy, + num_blocks=num_blocks, + num_threads=num_threads, + ) + done = comm.record_event() + return hidden_nvsh, route_weights_nvs, done + + def _make_dispatch_target( + self, + ctx: dict, + plan: MoonEPCommPlan, + ) -> MoonEPDispatchTarget: + target = MoonEPDispatchTarget( + mapped_hidden=ctx['hidden_buf'], + local_hidden=ctx['hidden_buf_local'], + plan=plan, + generation=self._dispatch_generation, + buffer_token=id(self), + R=int(ctx['R']), + S=int(ctx['S']), + H=int(ctx['H']), + K=int(ctx['K']), + NvS=int(ctx['NvS']), + NvS_padded=int(ctx['NvS_padded']), + ) + target.validate_layout() + return target + + def _validate_dispatch_target( + self, + ctx: dict, + target: MoonEPDispatchTarget, + plan: MoonEPCommPlan, + ) -> None: + if not isinstance(target, MoonEPDispatchTarget): + raise TypeError( + "target must be returned by Buffer.prepare_dispatch" + ) + target.validate_layout() + if target.buffer_token != id(self): + raise ValueError( + "dispatch target belongs to a different MoonEP Buffer" + ) + if target.generation != self._dispatch_generation: + raise RuntimeError( + "dispatch target is stale: this Buffer has started another " + "dispatch, prepare, or combine operation" + ) + if target.generation == self._published_dispatch_generation: + raise RuntimeError("dispatch target has already been published") + if target.plan is not plan: + raise ValueError( + "plan must be the exact plan carried by the dispatch target" + ) + if ( + target.mapped_hidden.data_ptr() != ctx['hidden_buf'].data_ptr() + or target.local_hidden.data_ptr() + != ctx['hidden_buf_local'].data_ptr() + ): + raise ValueError("dispatch target storage does not match this Buffer") + + def prepare_dispatch( + self, + topk_experts_sk: torch.Tensor, + tokens_per_expert: torch.Tensor, + async_finish: bool = False, + *, + inter_rank_sync: bool = True, + ): + """Prepare an owner-sharded activation object for producer direct-out. + + This runs MoonEP planning but does not read or move hidden activations. + A compatible upstream producer kernel can consume ``target.plan.dst`` + and write natural output rows directly into ``target.mapped_hidden``. + Only non-negative destinations are representative payload rows. + + The producer must run on the same CUDA stream or explicitly wait for + the planning stream. After it finishes, call ``publish_dispatch`` to + scatter route metadata, zero padding, build duplicate metadata, publish + cross-rank visibility, and expand duplicates locally. + + With ``async_finish=True``, planning runs on the communication stream + and a completion event is appended to the return tuple. The producer + must wait on it before using ``plan.dst``. + + Returns ``(target, cu_seqlens, plan)``, plus a communication-stream + CUDA event when ``async_finish=True``. The target expires when another + dispatch, prepare, or combine operation starts on this Buffer. + """ + ctx = self._require_ctx() + _require_bool("async_finish", async_finish) + _require_bool("inter_rank_sync", inter_rank_sync) + if topk_experts_sk is None or tokens_per_expert is None: + raise ValueError( + "topk_experts_sk and tokens_per_expert are required" + ) + if not isinstance(topk_experts_sk, torch.Tensor): + raise TypeError("topk_experts_sk must be a torch.Tensor") + topk_flat = topk_experts_sk.reshape(-1) + device = ctx['hidden_buf'].device + _validate_tensor( + "topk_experts_sk", + topk_flat, + dtype=torch.int32, + shape=(int(ctx['N']),), + device=device, + ) + _validate_tensor( + "tokens_per_expert", + tokens_per_expert, + dtype=torch.int32, + shape=(int(ctx['E']),), + device=device, + ) + + self._dispatch_generation += 1 + self._published_dispatch_generation = None + self._prepared_dispatch_generation = None + plan, cu_seqlens = allocate_planning_outputs(ctx) + target = self._make_dispatch_target(ctx, plan) + + def prepare_on_current_stream() -> None: + if inter_rank_sync: + launch_inter_rank_sync(ctx) + launch_planning( + ctx, + topk_flat, + tokens_per_expert, + cu_seqlens, + plan, + ) + + if not async_finish: + prepare_on_current_stream() + return target, cu_seqlens, plan + + main_stream = torch.cuda.current_stream() + comm = self._comm_stream + if comm is None: + raise RuntimeError( + "MoonEP Buffer communication stream is not initialized" + ) + self._record_streams( + ( + topk_flat, + tokens_per_expert, + cu_seqlens, + target.mapped_hidden, + target.local_hidden, + *self._plan_runtime_tensors(plan), + ), + comm, + ) + input_ready = main_stream.record_event() + comm.wait_event(input_ready) + with torch.cuda.stream(comm): + prepare_on_current_stream() + done = comm.record_event() + return target, cu_seqlens, plan, done + + def prepare_dispatch_publication( + self, + target: MoonEPDispatchTarget, + plan: MoonEPCommPlan, + route_weights_sk: torch.Tensor | None = None, + ) -> MoonEPPreparedPublication: + """Start MoonEP metadata work before a direct-out producer finishes. + + Call this after ``prepare_dispatch`` and immediately before launching + the compatible producer on the current stream. Metadata construction + runs asynchronously on MoonEP's communication stream and may overlap + the producer because it touches only route weights, padding rows, and + dedup structures. The returned object is one-shot and must be passed + to ``publish_dispatch(prepared_publication=...)`` after the producer + launch. + + The metadata kernel's barrier does not publish later producer-stream + writes. Final publication therefore executes a second visibility + barrier before duplicate expansion. + """ + ctx = self._require_ctx() + self._validate_dispatch_target(ctx, target, plan) + if target.generation == self._prepared_dispatch_generation: + raise RuntimeError( + "dispatch target metadata has already been prepared" + ) + if route_weights_sk is not None: + _validate_tensor( + "route_weights_sk", + route_weights_sk, + dtype=torch.float32, + shape=(int(ctx['S']), int(ctx['K'])), + device=ctx['meta_buf'].device, + ) + + comm = self._comm_stream + if comm is None: + raise RuntimeError( + "MoonEP Buffer communication stream is not initialized" + ) + self._record_streams( + ( + target.mapped_hidden, + target.local_hidden, + route_weights_sk, + *self._plan_runtime_tensors(plan), + ), + comm, + ) + + # Mark this generation before launch so a partial launch cannot be + # retried against partially initialized publication metadata. + self._prepared_dispatch_generation = target.generation + input_ready = torch.cuda.current_stream().record_event() + comm.wait_event(input_ready) + with torch.cuda.stream(comm): + launch_dispatch( + ctx, + None, + route_weights_sk, + plan, + write_payload=False, + build_dedup_map=True, + pdl_trigger=self.enable_pdl, + ) + metadata_done = comm.record_event() + + prepared = MoonEPPreparedPublication( + target=target, + plan=plan, + route_weights_sk=route_weights_sk, + metadata_done=metadata_done, + generation=target.generation, + buffer_token=id(self), + ) + prepared.validate_layout() + return prepared + + def _validate_prepared_publication( + self, + ctx: dict, + prepared: MoonEPPreparedPublication, + target: MoonEPDispatchTarget, + plan: MoonEPCommPlan, + ) -> None: + if not isinstance(prepared, MoonEPPreparedPublication): + raise TypeError( + "prepared_publication must be returned by " + "Buffer.prepare_dispatch_publication" + ) + prepared.validate_layout() + if prepared.buffer_token != id(self): + raise ValueError( + "prepared publication belongs to a different MoonEP Buffer" + ) + if prepared.target is not target or prepared.plan is not plan: + raise ValueError( + "prepared publication must carry the exact target and plan" + ) + if prepared.generation != self._dispatch_generation: + raise RuntimeError("prepared publication is stale") + if prepared.generation != self._prepared_dispatch_generation: + raise RuntimeError( + "dispatch target metadata was not prepared for this generation" + ) + if prepared.metadata_done.device != ctx['meta_buf'].device: + raise ValueError( + "prepared publication event belongs to a different device" + ) + + def _run_publish_dispatch_on_current_stream( + self, + ctx: dict, + plan: MoonEPCommPlan, + route_weights_sk: torch.Tensor | None, + hidden_nvsh: torch.Tensor, + route_weights_nvs: torch.Tensor | None, + *, + zero_copy: bool, + ) -> None: + launch_dispatch( + ctx, + None, + route_weights_sk, + plan, + write_payload=False, + build_dedup_map=True, + pdl_trigger=self.enable_pdl, + ) + launch_dispatch_epilogue(ctx, plan, pdl_launch=self.enable_pdl) + if not zero_copy: + hidden_nvsh.copy_(ctx['hidden_buf_local']) + if route_weights_nvs is not None: + route_weights_nvs.copy_( + ctx['weights_buf_local'].view(torch.float32) + ) + + def _run_finish_prepared_publication_on_current_stream( + self, + ctx: dict, + plan: MoonEPCommPlan, + hidden_nvsh: torch.Tensor, + route_weights_nvs: torch.Tensor | None, + *, + zero_copy: bool, + ) -> None: + # The caller orders this after both the direct-out producer and the + # metadata stream. Publish producer remote stores, then expand local + # duplicates from their representative rows. + launch_inter_rank_sync(ctx) + launch_dispatch_epilogue(ctx, plan, pdl_launch=self.enable_pdl) + if not zero_copy: + hidden_nvsh.copy_(ctx['hidden_buf_local']) + if route_weights_nvs is not None: + route_weights_nvs.copy_( + ctx['weights_buf_local'].view(torch.float32) + ) + + def publish_dispatch( + self, + target: MoonEPDispatchTarget, + plan: MoonEPCommPlan, + route_weights_sk: torch.Tensor | None = None, + async_finish: bool = False, + *, + zero_copy: bool = True, + prepared_publication: MoonEPPreparedPublication | None = None, + ): + """Publish an externally produced dispatch object. + + The external producer must have written every representative hidden row + described by non-negative ``target.plan.dst`` into + ``target.mapped_hidden``. This method performs no hidden payload copy. + It completes MoonEP-owned metadata, padding, visibility, and duplicate + expansion work, then returns ``(hidden_nvsh, route_weights_nvs, event)``. + + Pass the one-shot object from ``prepare_dispatch_publication`` to use + the overlapped path. In that mode route weights were already supplied + during preparation, so ``route_weights_sk`` must be None. Publication + waits for both metadata and producer completion, issues a post-producer + visibility barrier, then expands duplicates. + """ + ctx = self._require_ctx() + _require_bool("async_finish", async_finish) + _require_bool("zero_copy", zero_copy) + self._validate_dispatch_target(ctx, target, plan) + + if prepared_publication is not None: + self._validate_prepared_publication( + ctx, + prepared_publication, + target, + plan, + ) + if route_weights_sk is not None: + raise ValueError( + "route_weights_sk was already supplied to " + "prepare_dispatch_publication" + ) + effective_route_weights = prepared_publication.route_weights_sk + else: + if target.generation == self._prepared_dispatch_generation: + raise RuntimeError( + "dispatch target has prepared metadata; pass the returned " + "prepared_publication object" + ) + effective_route_weights = route_weights_sk + + if effective_route_weights is not None: + _validate_tensor( + "route_weights_sk", + effective_route_weights, + dtype=torch.float32, + shape=(int(ctx['S']), int(ctx['K'])), + device=ctx['meta_buf'].device, + ) + + if zero_copy: + hidden_nvsh = ctx['hidden_buf_local'] + route_weights_nvs = ( + ctx['weights_buf_local'].view(torch.float32) + if effective_route_weights is not None else None + ) + else: + hidden_nvsh = torch.empty_like(ctx['hidden_buf_local']) + route_weights_nvs = ( + torch.empty( + int(ctx['NvS']), + dtype=torch.float32, + device=ctx['meta_buf'].device, + ) + if effective_route_weights is not None else None + ) + + # Publication is one-shot. Mark it after host-side validation and + # allocation but before launch, so a launch failure cannot be retried + # against potentially partially updated metadata. + self._published_dispatch_generation = target.generation + + if prepared_publication is not None and not async_finish: + torch.cuda.current_stream().wait_event( + prepared_publication.metadata_done + ) + self._run_finish_prepared_publication_on_current_stream( + ctx, + plan, + hidden_nvsh, + route_weights_nvs, + zero_copy=zero_copy, + ) + return hidden_nvsh, route_weights_nvs, None + + if not async_finish: + self._run_publish_dispatch_on_current_stream( + ctx, + plan, + effective_route_weights, + hidden_nvsh, + route_weights_nvs, + zero_copy=zero_copy, + ) + return hidden_nvsh, route_weights_nvs, None + + main_stream = torch.cuda.current_stream() + comm = self._comm_stream + if comm is None: + raise RuntimeError( + "MoonEP Buffer communication stream is not initialized" + ) + self._record_streams( + ( + target.mapped_hidden, + target.local_hidden, + effective_route_weights, + hidden_nvsh, + route_weights_nvs, + *self._plan_runtime_tensors(plan), + ), + comm, + ) + producer_done = main_stream.record_event() + comm.wait_event(producer_done) + with torch.cuda.stream(comm): + if prepared_publication is not None: + # Metadata is already ordered before this point on comm. + self._run_finish_prepared_publication_on_current_stream( + ctx, + plan, + hidden_nvsh, + route_weights_nvs, + zero_copy=zero_copy, + ) + else: + self._run_publish_dispatch_on_current_stream( + ctx, + plan, + effective_route_weights, + hidden_nvsh, + route_weights_nvs, + zero_copy=zero_copy, + ) + done = comm.record_event() + return hidden_nvsh, route_weights_nvs, done + def dispatch( self, hidden_sh: torch.Tensor, @@ -738,6 +1707,9 @@ def dispatch( backward passes. """ ctx = self._require_ctx() + self._dispatch_generation += 1 + self._published_dispatch_generation = None + self._prepared_dispatch_generation = None if plan is None: assert topk_experts_sk is not None and tokens_per_expert is not None @@ -922,6 +1894,9 @@ def combine( None. """ ctx = self._require_ctx() + self._dispatch_generation += 1 + self._published_dispatch_generation = None + self._prepared_dispatch_generation = None assert isinstance(plan, MoonEPCommPlan), "Buffer.combine: plan is required" diff --git a/moonep/dispatch.py b/moonep/dispatch.py index 1f6c69e..8563556 100644 --- a/moonep/dispatch.py +++ b/moonep/dispatch.py @@ -93,6 +93,7 @@ def __init__( meta_stride: int, num_sms: int, with_weights: bool, + write_payload: bool, build_dedup_map: bool, smem_budget: int, pdl_trigger: bool, @@ -108,12 +109,20 @@ def __init__( self.SRC_INFO_OFF = SRC_INFO_OFF self.num_sms = num_sms self.with_weights = with_weights + self.write_payload = write_payload self.pdl_trigger = pdl_trigger self.build_dedup_map = build_dedup_map self.num_threads = ( 96 + 32 * DEDUP_BUILDER_WARPS if build_dedup_map else 96 ) - self.stages = self._pick_stages(H, smem_budget) + # Metadata-only publication has no G2S/S2G payload pipeline. Keep the + # minimum allocation accepted by the shared kernel shape instead of + # reserving the maximum stage depth for storage it never touches. + self.stages = ( + self._pick_stages(H, smem_budget) + if write_payload + else 2 + ) if self.stages == 0: raise RuntimeError( f"dispatch: H={H} too large for per-block smem budget " @@ -358,94 +367,116 @@ def kernel( # Warp 0 — G2S producer # ============================================ if warp_idx == self.PRODUCER_WARP: - load_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Producer, stages - ) - for li in cutlass.range(n_tok, unroll=1): - s = s_beg + li - load_pipe.producer_acquire(load_state) - # cp.async.bulk is single-thread; only lane 0 issues it. - if cute.arch.lane_idx() == 0: - g_row_int = ( - gmem_src.iterator + cutlass.Int64(s) * cutlass.Int64(H) - ).toint() - s_row_int = ( - stage_smem.iterator + load_state.index * H - ).toint() - mbar_int = load_pipe.producer_get_barrier(load_state).toint() - cp_async_bulk_g2s( - s_row_int.ir_value(), - g_row_int.ir_value(), - Int32(H_BYTES).ir_value(), - mbar_int.ir_value(), - ) - load_state.advance() + if cutlass.const_expr(self.write_payload): + load_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, stages + ) + for li in cutlass.range(n_tok, unroll=1): + s = s_beg + li + load_pipe.producer_acquire(load_state) + # cp.async.bulk is single-thread; only lane 0 issues it. + if cute.arch.lane_idx() == 0: + g_row_int = ( + gmem_src.iterator + cutlass.Int64(s) * cutlass.Int64(H) + ).toint() + s_row_int = ( + stage_smem.iterator + load_state.index * H + ).toint() + mbar_int = load_pipe.producer_get_barrier(load_state).toint() + cp_async_bulk_g2s( + s_row_int.ir_value(), + g_row_int.ir_value(), + Int32(H_BYTES).ir_value(), + mbar_int.ir_value(), + ) + load_state.advance() # ============================================ # Warp 1 — S2G consumer (K stores per token + optional weight scatter) # ============================================ elif warp_idx == self.CONSUMER_WARP: - use_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Consumer, stages - ) - rel_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Consumer, stages - ) - for li in cutlass.range(n_tok, unroll=1): - s = s_beg + li - sK = s * K + if cutlass.const_expr(self.write_payload): + use_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, stages + ) + rel_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, stages + ) + for li in cutlass.range(n_tok, unroll=1): + s = s_beg + li + sK = s * K - load_pipe.consumer_wait(use_state) + load_pipe.consumer_wait(use_state) - # All K stores + commit + optional weight scatter run on lane 0; - # 32 lanes would over-decrement the bulk_group / re-scatter. - if cute.arch.lane_idx() == 0: - s_row_int = ( - stage_smem.iterator + use_state.index * H - ).toint() - for k in cutlass.range(K, unroll=1): - dst_val = dst_tensor[sK + k] - # Negative dst keeps the raw destination for the - # weight scatter, but marks the payload as duplicate. - store_token = dst_val >= Int32(0) - raw_dst = dst_val - if not store_token: - raw_dst = - dst_val - Int32(1) - drank = raw_dst // NvS - loff = raw_dst % NvS - - if store_token: - drow = drank * NvS_padded + loff - g_row_int = ( - gmem_dst.iterator - + cutlass.Int64(drow) * cutlass.Int64(H) - ).toint() - cp_async_bulk_s2g( - s_row_int.ir_value(), - g_row_int.ir_value(), - Int32(H_BYTES).ir_value(), - ) + # All K stores + commit + optional weight scatter run on lane 0; + # 32 lanes would over-decrement the bulk_group / re-scatter. + if cute.arch.lane_idx() == 0: + s_row_int = ( + stage_smem.iterator + use_state.index * H + ).toint() + for k in cutlass.range(K, unroll=1): + dst_val = dst_tensor[sK + k] + # Negative dst keeps the raw destination for the + # weight scatter, but marks the payload as duplicate. + store_token = dst_val >= Int32(0) + raw_dst = dst_val + if not store_token: + raw_dst = - dst_val - Int32(1) + drank = raw_dst // NvS + loff = raw_dst % NvS + + if store_token: + drow = drank * NvS_padded + loff + g_row_int = ( + gmem_dst.iterator + + cutlass.Int64(drow) * cutlass.Int64(H) + ).toint() + cp_async_bulk_s2g( + s_row_int.ir_value(), + g_row_int.ir_value(), + Int32(H_BYTES).ir_value(), + ) - if cutlass.const_expr(self.with_weights): - wb = w_tensor[sK + k] - meta_tensor[ - drank * meta_stride + weights_off + loff - ] = wb - cute.arch.cp_async_bulk_commit_group() - - use_state.advance() - - # Throttle in-flight bulk_groups to <= STAGES-1. - if li >= Int32(stages - 1): - cute.arch.cp_async_bulk_wait_group(stages - 1) - load_pipe.consumer_release(rel_state) - rel_state.advance() - - # Drain the trailing in-flight stores. consumer_release for the - # last (kStages-1) stages isn't strictly needed since the kernel - # exits below, but we keep it symmetric so the empty mbars are - # in a known state if this kernel is replayed. - cute.arch.cp_async_bulk_wait_group(0) + if cutlass.const_expr(self.with_weights): + wb = w_tensor[sK + k] + meta_tensor[ + drank * meta_stride + weights_off + loff + ] = wb + cute.arch.cp_async_bulk_commit_group() + + use_state.advance() + + # Throttle in-flight bulk_groups to <= STAGES-1. + if li >= Int32(stages - 1): + cute.arch.cp_async_bulk_wait_group(stages - 1) + load_pipe.consumer_release(rel_state) + rel_state.advance() + + # Drain the trailing in-flight stores. consumer_release for the + # last (kStages-1) stages isn't strictly needed since the kernel + # exits below, but we keep it symmetric so the empty mbars are + # in a known state if this kernel is replayed. + cute.arch.cp_async_bulk_wait_group(0) + elif cutlass.const_expr(self.with_weights): + # The external producer already wrote every representative + # payload row into gmem_dst. Use all lanes of this warp to + # scatter only the per-topk route metadata before publishing + # the object. Negative dst values still decode to their unique + # metadata slots. + lane = cute.arch.lane_idx() + n = S * K + begin = bidx * Int32(32) + lane + stride = Int32(self.num_sms * 32) + for idx in cutlass.range(begin, n, stride): + dst_val = dst_tensor[idx] + raw_dst = dst_val + if dst_val < Int32(0): + raw_dst = -dst_val - Int32(1) + drank = raw_dst // NvS + loff = raw_dst % NvS + meta_tensor[ + drank * meta_stride + weights_off + loff + ] = w_tensor[idx] # ============================================ # Warp 2 — per-expert zero-fill loop (runs concurrently with @@ -714,6 +745,7 @@ def _get_compiled( SRC_INFO_OFF: int, num_sms: int, with_weights: bool, + write_payload: bool, build_dedup_map: bool, device_index: int, pdl_trigger: bool, @@ -725,6 +757,7 @@ def _get_compiled( SRC_INFO_OFF=SRC_INFO_OFF, meta_stride=meta_stride, num_sms=num_sms, with_weights=with_weights, + write_payload=write_payload, build_dedup_map=build_dedup_map, smem_budget=smem_budget, pdl_trigger=pdl_trigger, @@ -842,13 +875,16 @@ def launch_dispatch( route_weights_sk, plan, *, + write_payload: bool = True, build_dedup_map: bool = True, pdl_trigger: bool = False, ): """Launch the dispatch kernel. Args: - hidden_sh: [S, H] bf16 source hidden states. + hidden_sh: [S, H] bf16 source hidden states when + ``write_payload=True``. Pass None when an external producer has + already written every representative row into ``ctx['hidden_buf']``. route_weights_sk: [S, K] fp32 route weights, or None to skip the weights scatter (placeholder tensor is passed to satisfy the non-null pointer constraint; the kernel ignores it when @@ -862,6 +898,10 @@ def launch_dispatch( backward paths pass false so the saved dedup structures (``dup_groups`` / ``dup_loffs`` / ``dup_counts``) are not rebuilt from stale ``src_info`` scratch. The zero warp runs on both paths. + write_payload: copy hidden payload rows from ``hidden_sh`` when true. + When false, this launch only scatters route metadata, zero-fills + padding, builds fresh dedup metadata when requested, and publishes + the external producer's writes through the cross-rank barrier. The in-place duplicate expansion on the NVL shard is a separate kernel — call ``moonep.dispatch_epilogue.launch_dispatch_epilogue`` afterwards on the same @@ -880,40 +920,49 @@ def launch_dispatch( E = int(ctx['E']) B = int(ctx.get('B', 0)) - assert hidden_sh.dtype == torch.bfloat16 and hidden_sh.is_contiguous(), \ - "hidden_sh must be contiguous bf16" - assert hidden_sh.is_cuda, "hidden_sh must be a CUDA tensor" - assert tuple(hidden_sh.shape) == (S, H), \ - f"hidden_sh must be shape [S={S}, H={H}], got {tuple(hidden_sh.shape)}" - _check_dispatch_plan(ctx, hidden_sh, plan) + assert isinstance(write_payload, bool), \ + f"write_payload must be bool, got {type(write_payload).__name__}" + if write_payload: + assert hidden_sh is not None, "hidden_sh is required when write_payload=True" + assert hidden_sh.dtype == torch.bfloat16 and hidden_sh.is_contiguous(), \ + "hidden_sh must be contiguous bf16" + assert hidden_sh.is_cuda, "hidden_sh must be a CUDA tensor" + assert tuple(hidden_sh.shape) == (S, H), \ + f"hidden_sh must be shape [S={S}, H={H}], got {tuple(hidden_sh.shape)}" + source_hidden = hidden_sh + else: + assert hidden_sh is None, \ + "hidden_sh must be None when write_payload=False" + source_hidden = ctx['hidden_buf_local'] + _check_dispatch_plan(ctx, source_hidden, plan) assert ctx['hidden_buf'].dtype == torch.bfloat16 and ctx['hidden_buf'].is_contiguous() - assert ctx['hidden_buf'].device == hidden_sh.device + assert ctx['hidden_buf'].device == source_hidden.device assert ctx['meta_buf'].dtype == torch.int32 and ctx['meta_buf'].is_contiguous() - assert ctx['meta_buf'].device == hidden_sh.device + assert ctx['meta_buf'].device == source_hidden.device assert ctx['grid_sync_bar'].dtype == torch.int32 and ctx['grid_sync_bar'].is_contiguous() assert ctx['grid_sync_bar'].numel() == 1 - assert ctx['grid_sync_bar'].device == hidden_sh.device + assert ctx['grid_sync_bar'].device == source_hidden.device if with_weights: assert route_weights_sk.dtype == torch.float32 and route_weights_sk.is_contiguous(), \ "route_weights_sk must be contiguous fp32" assert tuple(route_weights_sk.shape) == (S, K), \ f"route_weights_sk must be shape [S={S}, K={K}], got {tuple(route_weights_sk.shape)}" - assert route_weights_sk.device == hidden_sh.device + assert route_weights_sk.device == source_hidden.device assert ctx['H'] % 8 == 0, "H must be multiple of 8 for 16-B bulk-copy alignment" if build_dedup_map: _check_dedup_builder_bounds(ctx) - _check_dedup_builder_tensors(ctx, hidden_sh.device) + _check_dedup_builder_tensors(ctx, source_hidden.device) NvS = int(ctx['NvS']) NvS_padded = int(ctx['NvS_padded']) meta_stride = int(ctx['meta_chunk_padded']) SRC_INFO_OFF = int(ctx['SRC_INFO_OFF']) num_sms = int(ctx['num_sms']) - device_index = hidden_sh.device.index + device_index = source_hidden.device.index dispatch_compiled = _get_compiled( H, R, S, K, E + B, NvS, NvS_padded, meta_stride, SRC_INFO_OFF, - num_sms, with_weights, build_dedup_map, device_index, + num_sms, with_weights, write_payload, build_dedup_map, device_index, bool(pdl_trigger), ) @@ -923,7 +972,7 @@ def launch_dispatch( route_weights_sk.view(torch.int32) if with_weights else route_weights_sk ) - hsh_ptr = make_ptr(BFloat16, hidden_sh.data_ptr(), cute.AddressSpace.gmem, + hsh_ptr = make_ptr(BFloat16, source_hidden.data_ptr(), cute.AddressSpace.gmem, assumed_align=16) hbf_ptr = make_ptr(BFloat16, ctx['hidden_buf'].data_ptr(), cute.AddressSpace.gmem, assumed_align=16) diff --git a/moonep/objects.py b/moonep/objects.py new file mode 100644 index 0000000..453e8d9 --- /dev/null +++ b/moonep/objects.py @@ -0,0 +1,272 @@ +"""Object-centric execution views exposed by MoonEP.""" + +from dataclasses import dataclass + +import torch + +from .planning import MoonEPCommPlan + + +def _validate_tensor( + name: str, + tensor: object, + *, + dtype: torch.dtype, + shape: tuple[int, ...], + device: torch.device | None = None, +) -> torch.Tensor: + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor") + if tensor.dtype != dtype: + raise TypeError(f"{name} must have dtype {dtype}, got {tensor.dtype}") + if tuple(tensor.shape) != shape: + raise ValueError( + f"{name} must have shape {shape}, got {tuple(tensor.shape)}" + ) + if not tensor.is_contiguous(): + raise ValueError(f"{name} must be contiguous") + if device is not None and tensor.device != device: + raise ValueError(f"{name} must be on {device}, got {tensor.device}") + return tensor + + +@dataclass(frozen=True) +class MoonEPSourceStorage: + """One physical source-activation shard per rank, globally mapped. + + Producers should write only ``local_hidden`` and routers should write + ``local_route_weights``. The mapped tensors are exposed for kernels + declaring an explicit peer-VMM load capability; generic tensor operations + are not assumed safe on their remote rows. + """ + + mapped_hidden: torch.Tensor + local_hidden: torch.Tensor + mapped_route_weights: torch.Tensor + local_route_weights: torch.Tensor + buffer_token: int + R: int + S: int + H: int + K: int + source_rows_padded: int + source_topk_padded: int + capability: str = "owner_sharded_source_v1" + + def validate_layout(self) -> None: + if self.capability != "owner_sharded_source_v1": + raise ValueError( + "source storage requires capability owner_sharded_source_v1" + ) + mapped_hidden = _validate_tensor( + "mapped_hidden", + self.mapped_hidden, + dtype=torch.bfloat16, + shape=(self.R * self.source_rows_padded, self.H), + ) + _validate_tensor( + "local_hidden", + self.local_hidden, + dtype=torch.bfloat16, + shape=(self.S, self.H), + device=mapped_hidden.device, + ) + _validate_tensor( + "mapped_route_weights", + self.mapped_route_weights, + dtype=torch.float32, + shape=(self.R * self.source_topk_padded,), + device=mapped_hidden.device, + ) + _validate_tensor( + "local_route_weights", + self.local_route_weights, + dtype=torch.float32, + shape=(self.S, self.K), + device=mapped_hidden.device, + ) + + +@dataclass(frozen=True) +class MoonEPSourceView: + """Planned peer-indexed view over an owner-sharded source object. + + ``slot_to_source[loff]`` maps each valid local expert-layout slot to + ``source_rank * NvS + source_topk_offset``. A compatible consumer decodes: + + ``source_rank = encoded // NvS`` + ``source_token = (encoded % NvS) // K`` + ``source_row = source_rank * source_rows_padded + source_token`` + + Negative entries are padding. The view is valid only until the same Buffer + starts another dispatch, source preparation, direct-out preparation, or + combine operation. This is the one-physical-copy activation view; no routed + hidden rows have been materialized by its preparation. + + ``slot_to_primary[loff]`` names the representative local slot for the same + source token on this consumer rank. It equals ``loff`` for representatives, + points at the representative for same-rank top-k duplicates, and is + negative for padding. Materializers may therefore pull one peer row per + consumer rank and expand duplicates from local memory. + + ``slot_to_next`` links those slots in top-k order. A representative + materializer can load each source vector once, then store the register + value directly to every linked final slot without an intermediate local + reread. + + """ + + storage: MoonEPSourceStorage + slot_to_source: torch.Tensor + slot_to_primary: torch.Tensor + slot_to_next: torch.Tensor + plan: MoonEPCommPlan + generation: int + NvS: int + K: int + capability: str = "peer_indexed_source_v1" + + def validate_layout(self) -> None: + self.storage.validate_layout() + if self.capability != "peer_indexed_source_v1": + raise ValueError( + "source view requires capability peer_indexed_source_v1" + ) + device = self.storage.mapped_hidden.device + _validate_tensor( + "slot_to_source", + self.slot_to_source, + dtype=torch.int32, + shape=(self.NvS,), + device=device, + ) + _validate_tensor( + "slot_to_primary", + self.slot_to_primary, + dtype=torch.int32, + shape=(self.NvS,), + device=device, + ) + _validate_tensor( + "slot_to_next", + self.slot_to_next, + dtype=torch.int32, + shape=(self.NvS,), + device=device, + ) + if not isinstance(self.plan, MoonEPCommPlan): + raise TypeError("plan must be a MoonEPCommPlan") + if self.plan.NvS != self.NvS or self.plan.K != self.K: + raise ValueError("source view and plan layouts do not match") + + +@dataclass(frozen=True) +class MoonEPDispatchTarget: + """Prepared owner-sharded activation object for producer direct-out. + + ``mapped_hidden`` is the full VMM mapping shaped + ``[R * NvS_padded, H]``. A compatible producer kernel reads ``plan.dst`` + and writes its natural output row directly to every non-negative encoded + destination: + + ``mapped_hidden[(dst // NvS) * NvS_padded + (dst % NvS)]``. + + Negative ``dst`` entries are same-rank duplicates and must not be written + by the producer; ``Buffer.publish_dispatch`` expands them locally after + publishing the representative rows. + + The target is intentionally capability-scoped. Generic tensor consumers + are not assumed safe on peer VMM mappings. Publication is one-shot, and + the target expires when the same Buffer starts another dispatch, prepare, + or combine operation. + """ + + mapped_hidden: torch.Tensor + local_hidden: torch.Tensor + plan: MoonEPCommPlan + generation: int + buffer_token: int + R: int + S: int + H: int + K: int + NvS: int + NvS_padded: int + capability: str = "remote_final_slot_store_v1" + + def validate_layout(self) -> None: + if self.capability != "remote_final_slot_store_v1": + raise ValueError( + "dispatch target requires capability remote_final_slot_store_v1" + ) + mapped_hidden = _validate_tensor( + "mapped_hidden", + self.mapped_hidden, + dtype=torch.bfloat16, + shape=(self.R * self.NvS_padded, self.H), + ) + _validate_tensor( + "local_hidden", + self.local_hidden, + dtype=torch.bfloat16, + shape=(self.NvS, self.H), + device=mapped_hidden.device, + ) + if not isinstance(self.plan, MoonEPCommPlan): + raise TypeError("plan must be a MoonEPCommPlan") + _validate_tensor( + "plan.dst", + self.plan.dst, + dtype=torch.int32, + shape=(self.S * self.K,), + device=mapped_hidden.device, + ) + + +@dataclass(frozen=True) +class MoonEPPreparedPublication: + """Metadata publication scheduled ahead of producer completion. + + ``metadata_done`` records completion of route-weight scatter, padding + initialization, dedup construction, and that kernel's own barrier. Those + operations may overlap a compatible direct-out producer because they do + not touch representative hidden payload rows. The producer's writes are + not yet published: ``Buffer.publish_dispatch`` must consume this exact + object and run a second visibility barrier before duplicate expansion. + """ + + target: MoonEPDispatchTarget + plan: MoonEPCommPlan + route_weights_sk: torch.Tensor | None + metadata_done: torch.cuda.Event + generation: int + buffer_token: int + capability: str = "overlapped_dispatch_publication_v1" + + @property + def with_route_weights(self) -> bool: + return self.route_weights_sk is not None + + def validate_layout(self) -> None: + if self.capability != "overlapped_dispatch_publication_v1": + raise ValueError( + "prepared publication requires capability " + "overlapped_dispatch_publication_v1" + ) + self.target.validate_layout() + if self.target.plan is not self.plan: + raise ValueError("prepared publication must carry its target plan") + if self.target.generation != self.generation: + raise ValueError("prepared publication generation does not match") + if self.target.buffer_token != self.buffer_token: + raise ValueError("prepared publication Buffer does not match") + if not isinstance(self.metadata_done, torch.cuda.Event): + raise TypeError("metadata_done must be a torch.cuda.Event") + if self.route_weights_sk is not None: + _validate_tensor( + "route_weights_sk", + self.route_weights_sk, + dtype=torch.float32, + shape=(self.target.S, self.target.K), + device=self.target.mapped_hidden.device, + ) diff --git a/moonep/planning.py b/moonep/planning.py index 23e67e6..05ac5b1 100644 --- a/moonep/planning.py +++ b/moonep/planning.py @@ -353,14 +353,16 @@ def _pd_issue_g2s(meta, smem_stage, src_begin, logical_count, mbar): class PlanningKernel: def __init__(self, R, E, B, S, K, NvS_capacity, NvS, num_vblocks, meta_stride, TPE_OFF, PLAN_OFF, BARRIER_OFF, TOPK0_OFF, ORDER_OFF, ORDER0_OFF, - token_padding, num_sms): + token_padding, enable_source_object, num_sms): self.R, self.E, self.B, self.S, self.K = R, E, B, S, K self.N = self.S * self.K self.NvS_capacity, self.NvS, self.num_vblocks = NvS_capacity, NvS, num_vblocks self.meta_stride = meta_stride self.TPE_OFF, self.PLAN_OFF, self.BARRIER_OFF = TPE_OFF, PLAN_OFF, BARRIER_OFF self.TOPK0_OFF, self.ORDER_OFF, self.ORDER0_OFF = TOPK0_OFF, ORDER_OFF, ORDER0_OFF - self.token_padding, self.num_sms = token_padding, num_sms + self.token_padding = token_padding + self.enable_source_object = enable_source_object + self.num_sms = num_sms @cute.jit def __call__(self, tpe, topk, meta, mc, dst, cu_seqlens, @@ -543,6 +545,8 @@ def kernel(self, tpe, topk, meta, mc, dst, cu_seqlens, ORDER0_OFF = cutlass.const_expr(self.ORDER0_OFF) BARRIER_SLOTS = 3 SRC_INFO_OFF = cutlass.const_expr(self.BARRIER_OFF + BARRIER_SLOTS) + PRIMARY_INFO_OFF = cutlass.const_expr(SRC_INFO_OFF + NvS) + NEXT_INFO_OFF = cutlass.const_expr(PRIMARY_INFO_OFF + NvS) ALLOC_SUB = 0 TPE_SUB = E * R EOFF_SUB = 2 * E * R @@ -970,12 +974,17 @@ def sa(n): copy_v4_remote(meta, ORDER_OFF, order0, N, pid, tid, num_threads, num_sms) else: self.run_c1(topk, order, tpe, lh, s_hist, s_bp, scratch, bar_p, num_sms, pid, tid) - # Clear this rank's src_info slice before all ranks publish fresh slot - # provenance into destination-rank slices below. src_info mirrors dst's - # rank-stride encoding: src_rank * NvS + offv; -1 is the empty-slot - # sentinel. offv is always in [0, N), and NvS >= N. + # Clear this rank's source-provenance and representative-slot slices + # before all ranks publish fresh values into destination-rank slices. + # src_info mirrors dst's rank-stride encoding: + # src_rank * NvS + offv; -1 is the empty-slot sentinel. primary_info + # is the representative local slot for the same source token on this + # consumer rank; it is also -1 for padding. for idx in cutlass.range(pid * num_threads + tid, NvS, num_sms * num_threads): meta[rank * ms + SRC_INFO_OFF + idx] = Int32(-1) + if cutlass.const_expr(self.enable_source_object): + meta[rank * ms + PRIMARY_INFO_OFF + idx] = Int32(-1) + meta[rank * ms + NEXT_INFO_OFF + idx] = Int32(-1) cross_rank_barrier(meta, ms, BARRIER_OFF, rank, R, bar_p, num_sms, num_threads, tid) s_expoff = cute.make_tensor(s_hist.iterator, cute.make_layout((E,))) for e in cutlass.range(tid, E, num_threads): @@ -1078,8 +1087,11 @@ def sa(n): # Canonicalize dst duplicate entries. The first top-k entry per # destination rank stays non-negative and copies the payload; later - # entries encode -raw_dst - 1 and only carry weights. Fresh dispatch - # materializes the dedup structures from src_info. + # entries encode -raw_dst - 1 and only carry weights. At the same time, + # publish the representative local slot for every destination slot. + # Owner-pull materializers can then load the source row once per + # consumer rank and fan loaded vectors through the linked final slots + # without running the dispatch dedup builder or rereading hidden rows. seg_dst = cute.ceil_div(S, num_sms) sbeg_dst = pid * seg_dst; send_dst = cutlass.min(sbeg_dst + seg_dst, S) for base in cutlass.range(sbeg_dst + tid, send_dst, num_threads): @@ -1111,6 +1123,32 @@ def sa(n): seen_hi = seen_hi | bit if dup: dst_out[base_idx + k] = -(dst_vals[k]) - 1 + if cutlass.const_expr(self.enable_source_object): + primary_raw_dst = dst_vals[k] + previous_raw_dst = Int32(-1) + if dup: + for prev_k in cutlass.range_constexpr(K): + if prev_k < k: + if dests[prev_k] == d: + previous_raw_dst = dst_vals[prev_k] + prior_dst = dst_out[base_idx + prev_k] + if prior_dst >= Int32(0): + primary_raw_dst = prior_dst + raw_dst = dst_vals[k] + loff = raw_dst % NvS + primary_loff = primary_raw_dst % NvS + meta[d * ms + PRIMARY_INFO_OFF + loff] = primary_loff + if previous_raw_dst >= Int32(0): + previous_loff = previous_raw_dst % NvS + meta[ + d * ms + NEXT_INFO_OFF + previous_loff + ] = loff + if cutlass.const_expr(self.enable_source_object): + # Planning may return independently on each rank. Publish all + # remote representative-slot writes before a consumer launches. + cross_rank_barrier( + meta, ms, BARRIER_OFF, rank, R, bar_p, num_sms, num_threads, tid + ) cute.arch.mbarrier_wait(pd_mbar, 0) cu_stage_bias = _pd_stage_bias(cu_src_begin) zfr_stage_bias = _pd_stage_bias(zfr_src_begin) @@ -1138,10 +1176,10 @@ def sa(n): @functools.lru_cache(maxsize=None) def _get_compiled(R, E, B, S, K, NvS_capacity, NvS, num_vblocks, meta_stride, TPE_OFF, PLAN_OFF, BARRIER_OFF, TOPK0_OFF, ORDER_OFF, ORDER0_OFF, - token_padding, num_sms): + token_padding, enable_source_object, num_sms): k = PlanningKernel(R, E, B, S, K, NvS_capacity, NvS, num_vblocks, meta_stride, TPE_OFF, PLAN_OFF, BARRIER_OFF, TOPK0_OFF, ORDER_OFF, ORDER0_OFF, - token_padding, num_sms) + token_padding, enable_source_object, num_sms) i32 = make_ptr(Int32, 0, cute.AddressSpace.gmem, assumed_align=16) return cute.compile(k, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, Int32(0), cuda.CUstream(0)) @@ -1167,6 +1205,7 @@ def _launch_planning_kernel(ctx, topk, tpe, dst, cu_seqlens, int(ctx['ORDER_OFF']), int(ctx['ORDER0_OFF']), int(ctx['token_padding']), + bool(ctx['enable_source_object']), int(ctx['num_sms']), ) diff --git a/tests/kernel_test_utils.py b/tests/kernel_test_utils.py index b61829d..a6e9ec3 100644 --- a/tests/kernel_test_utils.py +++ b/tests/kernel_test_utils.py @@ -25,6 +25,7 @@ class KernelCase: H: int num_sms: int B: int | None = None + enable_source_object: bool = False token_padding: int = DEFAULT_TOKEN_PADDING routing: str = "balanced" bias_ratio: float = 0.0 @@ -51,6 +52,7 @@ def init_case(case, R): case.E(R), R, B=case.B, + enable_source_object=case.enable_source_object, num_sms=case.num_sms, token_padding=case.token_padding, ) diff --git a/tests/planning_reference.py b/tests/planning_reference.py index c5b5b02..1516eb0 100644 --- a/tests/planning_reference.py +++ b/tests/planning_reference.py @@ -156,7 +156,7 @@ def launch_planning_torch_reference( ] # Pick the B remote expert segments with the most tokens for VM - # prefetch; + # prefetch. remote_experts.sort(key=lambda e: (alloc[e, d].item(), e), reverse=True) remote_stats_all[d, 0] = len(remote_experts) for b, e in enumerate(remote_experts[:B]): diff --git a/tests/test_dispatch.py b/tests/test_dispatch.py index b9e3397..2a579e0 100644 --- a/tests/test_dispatch.py +++ b/tests/test_dispatch.py @@ -425,6 +425,308 @@ def test_dispatch_saved_plan_hidden_only_reuses_dst_and_skips_weights(dist_env): ) +def test_prepared_dispatch_object_publishes_external_payload(dist_env): + """Emulate a producer direct-out, then run MoonEP's metadata-only publish. + + The first payload-only launch is a correctness oracle for the external + producer contract: it writes representative rows according to ``dst`` but + deliberately skips dedup construction and route metadata. The public + publish path must complete those pieces without recopying hidden payloads. + """ + from moonep.dispatch import launch_dispatch + + rank, R = dist_env + case = KernelCase( + "prepared_object", + S=17, + K=3, + epn=8, + H=64, + num_sms=8, + B=2, + token_padding=16, + routing="all_remote", + min_R=2, + ) + ctx = init_case(case, R) + buffer = ctx["_buffer"] + topk, tpe = make_topk(case, rank, R) + hidden = _traceable_hidden(rank, case.S, case.H) + weights = _traceable_weights(rank, case.S, case.K) + + target, cu_seqlens, plan, planning_done = buffer.prepare_dispatch( + topk, + tpe, + async_finish=True, + ) + torch.cuda.current_stream().wait_event(planning_done) + assert target.capability == "remote_final_slot_store_v1" + assert target.mapped_hidden.data_ptr() == ctx["hidden_buf"].data_ptr() + assert target.local_hidden.data_ptr() == ctx["hidden_buf_local"].data_ptr() + assert target.plan is plan + assert tuple(cu_seqlens.shape) == (case.E(R) + int(ctx["B"]),) + + # Emulate a compatible upstream producer. This call writes payload only; + # publish_dispatch below owns the fresh dedup build, route metadata, + # visibility publication, and local duplicate expansion. + launch_dispatch( + ctx, + hidden, + None, + plan, + write_payload=True, + build_dedup_map=False, + ) + hidden_view, weights_view, event = buffer.publish_dispatch( + target, + plan, + route_weights_sk=weights, + zero_copy=True, + ) + assert event is None + assert hidden_view.data_ptr() == target.local_hidden.data_ptr() + assert weights_view.data_ptr() == ctx["weights_buf_local"].data_ptr() + torch.cuda.synchronize() + + scatter_errors, checked = _verify_dispatch_by_dst( + ctx, + case, + rank, + R, + plan.dst, + hidden_user=hidden_view, + weights_user=weights_view, + ) + padding_errors, padding_rows = _zero_row_errors( + plan.zero_fill_ranges, + hidden_view, + weights_view, + ) + errors = scatter_errors + padding_errors + assert_all_ranks( + checked > 0 and padding_rows > 0 and not errors, + rank, + R, + "prepared dispatch object publish", + "; ".join(errors[:5]) + or f"checked={checked}, padding_rows={padding_rows}", + ) + + +def test_prepared_dispatch_object_overlaps_metadata_publication(dist_env): + """Two-phase publication preserves payload, weights, and padding.""" + from moonep.dispatch import launch_dispatch + + rank, R = dist_env + case = KernelCase( + "prepared_overlap", + S=17, + K=3, + epn=8, + H=64, + num_sms=8, + B=2, + token_padding=16, + routing="all_remote", + min_R=2, + ) + ctx = init_case(case, R) + buffer = ctx["_buffer"] + topk, tpe = make_topk(case, rank, R) + hidden = _traceable_hidden(rank, case.S, case.H) + weights = _traceable_weights(rank, case.S, case.K) + + target, _cu_seqlens, plan = buffer.prepare_dispatch(topk, tpe) + prepared = buffer.prepare_dispatch_publication( + target, + plan, + route_weights_sk=weights, + ) + with pytest.raises(RuntimeError, match="already been prepared"): + buffer.prepare_dispatch_publication( + target, + plan, + route_weights_sk=weights, + ) + with pytest.raises(RuntimeError, match="pass the returned"): + buffer.publish_dispatch( + target, + plan, + route_weights_sk=weights, + ) + + # The generic test producer shares MoonEP's barrier scratch, unlike a real + # direct-out epilogue, so order it after metadata here. The research CUDA + # Tile benchmark covers true concurrent execution. + prepared.metadata_done.synchronize() + launch_dispatch( + ctx, + hidden, + None, + plan, + write_payload=True, + build_dedup_map=False, + ) + hidden_view, weights_view, done = buffer.publish_dispatch( + target, + plan, + async_finish=True, + zero_copy=True, + prepared_publication=prepared, + ) + assert done is not None + done.synchronize() + + scatter_errors, checked = _verify_dispatch_by_dst( + ctx, + case, + rank, + R, + plan.dst, + hidden_user=hidden_view, + weights_user=weights_view, + ) + padding_errors, padding_rows = _zero_row_errors( + plan.zero_fill_ranges, + hidden_view, + weights_view, + ) + errors = scatter_errors + padding_errors + assert_all_ranks( + checked > 0 and padding_rows > 0 and not errors, + rank, + R, + "overlapped prepared dispatch publication", + "; ".join(errors[:5]) + or f"checked={checked}, padding_rows={padding_rows}", + ) + + +def test_prepared_dispatch_object_rejects_stale_generation(dist_env): + rank, R = dist_env + case = KernelCase( + "stale_prepared_object", + S=8, + K=2, + epn=4, + H=32, + num_sms=4, + B=1, + token_padding=8, + routing="all_remote", + min_R=2, + ) + ctx = init_case(case, R) + buffer = ctx["_buffer"] + topk, tpe = make_topk(case, rank, R) + + old_target, _old_cu, old_plan = buffer.prepare_dispatch(topk, tpe) + buffer.prepare_dispatch(topk, tpe) + with pytest.raises(RuntimeError, match="stale"): + buffer.publish_dispatch(old_target, old_plan) + + +def test_owner_pull_materializes_source_object(dist_env): + """Explicit peer-VMM loads materialize the one-copy source view correctly.""" + rank, R = dist_env + case = KernelCase( + "owner_pull_source", + S=31, + K=4, + epn=8, + H=64, + num_sms=8, + B=2, + token_padding=16, + routing="duplicate_topk", + min_R=2, + enable_source_object=True, + ) + ctx = init_case(case, R) + buffer = ctx["_buffer"] + topk, tpe = make_topk(case, rank, R) + hidden = _traceable_hidden(rank, case.S, case.H) + weights = _traceable_weights(rank, case.S, case.K) + + storage = buffer.source_storage() + storage.local_hidden.copy_(hidden) + storage.local_route_weights.copy_(weights) + view, _cu_seqlens, plan, planning_done = buffer.prepare_source_dispatch( + storage, + topk, + tpe, + async_finish=True, + ) + torch.cuda.current_stream().wait_event(planning_done) + hidden_view, weights_view, event = buffer.materialize_source_dispatch( + view, + plan, + with_route_weights=True, + zero_copy=True, + ) + assert event is None + assert hidden_view.data_ptr() == ctx["hidden_buf_local"].data_ptr() + assert weights_view.data_ptr() == ctx["weights_buf_local"].data_ptr() + torch.cuda.synchronize() + + scatter_errors, checked = _verify_dispatch_by_dst( + ctx, + case, + rank, + R, + plan.dst, + hidden_user=hidden_view, + weights_user=weights_view, + ) + padding_errors, padding_rows = _zero_row_errors( + plan.zero_fill_ranges, + hidden_view, + weights_view, + ) + errors = scatter_errors + padding_errors + assert_all_ranks( + checked > 0 and padding_rows > 0 and not errors, + rank, + R, + "owner-pull source materialization", + "; ".join(errors[:5]) + or f"checked={checked}, padding_rows={padding_rows}", + ) + + hidden_copy, weights_copy, done = buffer.materialize_source_dispatch( + view, + plan, + with_route_weights=True, + async_finish=True, + zero_copy=False, + ) + assert done is not None + done.synchronize() + async_scatter_errors, async_checked = _verify_dispatch_by_dst( + ctx, + case, + rank, + R, + plan.dst, + hidden_user=hidden_copy, + weights_user=weights_copy, + ) + async_padding_errors, async_padding_rows = _zero_row_errors( + plan.zero_fill_ranges, + hidden_copy, + weights_copy, + ) + async_errors = async_scatter_errors + async_padding_errors + assert_all_ranks( + async_checked > 0 and async_padding_rows > 0 and not async_errors, + rank, + R, + "async owner-pull source materialization", + "; ".join(async_errors[:5]) + or f"checked={async_checked}, padding_rows={async_padding_rows}", + ) + + @pytest.mark.parametrize("case", case_params(LARGE_DISPATCH_CASES)) def test_dispatch_large_hidden_stride_spotcheck(dist_env, case): rank, R = dist_env @@ -466,6 +768,15 @@ def test_dispatch_rejects_bad_inputs(dist_env): with pytest.raises(AssertionError, match="hidden_sh"): launch_dispatch(ctx, hidden.float(), weights, plan, build_dedup_map=True) + with pytest.raises(AssertionError, match="must be None"): + launch_dispatch( + ctx, + hidden, + weights, + plan, + write_payload=False, + build_dedup_map=True, + ) with pytest.raises(AssertionError, match="hidden_sh"): launch_dispatch( ctx, diff --git a/tests/test_objects.py b/tests/test_objects.py new file mode 100644 index 0000000..30b8946 --- /dev/null +++ b/tests/test_objects.py @@ -0,0 +1,91 @@ +"""CPU-only validation tests for MoonEP object capability contracts.""" + +from dataclasses import replace + +import pytest +import torch + +from moonep import Buffer, MoonEPDispatchTarget, MoonEPSourceStorage +from moonep.planning import MoonEPCommPlan + + +def make_plan() -> MoonEPCommPlan: + N, R, E, B, NvS, K = 2, 2, 2, 1, 4, 1 + return MoonEPCommPlan( + dst=torch.zeros(N, dtype=torch.int32), + experts_to_copy=torch.zeros((R, B), dtype=torch.int32), + zero_fill_ranges=torch.zeros((E + B, 2), dtype=torch.int32), + remote_stats=torch.zeros(2, dtype=torch.int32), + dup_groups=torch.zeros((NvS, 3), dtype=torch.int32), + dup_loffs=torch.zeros(NvS, dtype=torch.int32), + dup_counts=torch.zeros(2, dtype=torch.int32), + N=N, + R=R, + E=E, + B=B, + NvS=NvS, + K=K, + ) + + +def test_source_storage_layout_rejects_capability_and_shape() -> None: + storage = MoonEPSourceStorage( + mapped_hidden=torch.empty((4, 8), dtype=torch.bfloat16), + local_hidden=torch.empty((2, 8), dtype=torch.bfloat16), + mapped_route_weights=torch.empty(8, dtype=torch.float32), + local_route_weights=torch.empty((2, 1), dtype=torch.float32), + buffer_token=1, + R=2, + S=2, + H=8, + K=1, + source_rows_padded=2, + source_topk_padded=4, + ) + storage.validate_layout() + + with pytest.raises(ValueError, match="capability"): + replace(storage, capability="unknown").validate_layout() + with pytest.raises(ValueError, match="mapped_hidden must have shape"): + replace( + storage, + mapped_hidden=torch.empty((3, 8), dtype=torch.bfloat16), + ).validate_layout() + + +def test_dispatch_target_state_checks_survive_python_optimization() -> None: + plan = make_plan() + mapped_hidden = torch.empty((8, 8), dtype=torch.bfloat16) + local_hidden = mapped_hidden[:4] + target = MoonEPDispatchTarget( + mapped_hidden=mapped_hidden, + local_hidden=local_hidden, + plan=plan, + generation=1, + buffer_token=123, + R=2, + S=2, + H=8, + K=1, + NvS=4, + NvS_padded=4, + ) + buffer = object.__new__(Buffer) + buffer._dispatch_generation = 1 + buffer._published_dispatch_generation = None + ctx = { + "hidden_buf": mapped_hidden, + "hidden_buf_local": local_hidden, + } + + target = replace(target, buffer_token=id(buffer)) + buffer._validate_dispatch_target(ctx, target, plan) + + buffer._dispatch_generation = 2 + with pytest.raises(RuntimeError, match="stale"): + buffer._validate_dispatch_target(ctx, target, plan) + + buffer._dispatch_generation = 1 + buffer._published_dispatch_generation = 1 + with pytest.raises(RuntimeError, match="already been published"): + buffer._validate_dispatch_target(ctx, target, plan) diff --git a/tests/test_planning.py b/tests/test_planning.py index b274244..9ba7dec 100644 --- a/tests/test_planning.py +++ b/tests/test_planning.py @@ -5,12 +5,14 @@ """ import pytest +import torch from tests.kernel_test_utils import ( DEFAULT_TOKEN_PADDING, KernelCase, assert_all_ranks, assert_tensor_equal_all_ranks, case_params, + gather_tensor, init_case, make_topk, planning_invariant_errors, @@ -265,6 +267,9 @@ def test_planning_matches_reference_and_invariants(dist_env, case): skip_if_unsupported_world_size(case, R) ctx = init_case(case, R) + assert (ctx["source_primary_local"] is not None) == \ + case.enable_source_object + assert (ctx["source_next_local"] is not None) == case.enable_source_object topk, tpe = make_topk(case, rank, R) plan, cu_seqlens = allocate_planning_outputs(ctx) @@ -306,3 +311,117 @@ def test_planning_matches_reference_and_invariants(dist_env, case): f"{case.name} planning invariants", "; ".join(errors[:5]), ) + + +def test_owner_sharded_source_view_matches_planner_provenance(dist_env): + """The direct-consume view keeps one source row and exposes its slot map.""" + rank, R = dist_env + case = KernelCase( + "owner_sharded_source", + S=17, + K=3, + epn=8, + H=64, + num_sms=8, + B=2, + token_padding=16, + routing="all_remote", + min_R=2, + enable_source_object=True, + ) + ctx = init_case(case, R) + buffer = ctx["_buffer"] + topk, tpe = make_topk(case, rank, R) + + storage = buffer.source_storage() + hidden = ( + torch.arange(case.S * case.H, device=f"cuda:{rank}") + .reshape(case.S, case.H) + .to(torch.bfloat16) + + rank * 1024 + ) + storage.local_hidden.copy_(hidden) + view, cu_seqlens, plan = buffer.prepare_source_dispatch( + storage, + topk, + tpe, + ) + assert view.capability == "peer_indexed_source_v1" + assert view.storage is storage + assert view.plan is plan + assert view.slot_to_source.data_ptr() == ctx["source_info_local"].data_ptr() + assert view.slot_to_primary.data_ptr() == \ + ctx["source_primary_local"].data_ptr() + assert view.slot_to_next.data_ptr() == ctx["source_next_local"].data_ptr() + assert torch.equal(storage.local_hidden, hidden) + assert tuple(cu_seqlens.shape) == (case.E(R) + int(ctx["B"]),) + + all_dst = gather_tensor(plan.dst, R).cpu() + source_info = view.slot_to_source.cpu() + source_primary = view.slot_to_primary.cpu() + source_next = view.slot_to_next.cpu() + expected = torch.full_like(source_info, -1) + expected_primary = torch.full_like(source_primary, -1) + expected_next = torch.full_like(source_next, -1) + NvS = int(ctx["NvS"]) + for src_rank in range(R): + src_dst = all_dst[src_rank].tolist() + for source_token in range(case.S): + primary_by_rank = {} + previous_by_rank = {} + for kidx in range(case.K): + source_topk_offset = source_token * case.K + kidx + encoded_dst = src_dst[source_topk_offset] + raw_dst = encoded_dst if encoded_dst >= 0 else -encoded_dst - 1 + dest_rank = raw_dst // NvS + loff = raw_dst % NvS + if encoded_dst >= 0: + primary_by_rank[dest_rank] = loff + if dest_rank == rank: + expected[loff] = src_rank * NvS + source_topk_offset + expected_primary[loff] = primary_by_rank[dest_rank] + if dest_rank in previous_by_rank: + expected_next[previous_by_rank[dest_rank]] = loff + previous_by_rank[dest_rank] = loff + + errors = [] + if not torch.equal(source_info, expected): + mismatch = (source_info != expected).nonzero().flatten() + for loff in mismatch[:5].tolist(): + errors.append( + f"loff={loff}: actual={int(source_info[loff])}, " + f"expected={int(expected[loff])}" + ) + if not torch.equal(source_primary, expected_primary): + mismatch = (source_primary != expected_primary).nonzero().flatten() + for loff in mismatch[:5].tolist(): + errors.append( + f"primary loff={loff}: actual={int(source_primary[loff])}, " + f"expected={int(expected_primary[loff])}" + ) + if not torch.equal(source_next, expected_next): + mismatch = (source_next != expected_next).nonzero().flatten() + for loff in mismatch[:5].tolist(): + errors.append( + f"next loff={loff}: actual={int(source_next[loff])}, " + f"expected={int(expected_next[loff])}" + ) + valid = source_info >= 0 + if not bool(valid.any()): + errors.append("source view has no valid local expert slots") + else: + encoded = source_info[valid] + source_ranks = encoded // NvS + source_tokens = (encoded % NvS) // case.K + if not bool(((source_ranks >= 0) & (source_ranks < R)).all()): + errors.append("decoded source rank is out of range") + if not bool(((source_tokens >= 0) & (source_tokens < case.S)).all()): + errors.append("decoded source token is out of range") + + assert_all_ranks( + not errors, + rank, + R, + "owner-sharded source provenance", + "; ".join(errors), + ) From da95f0d1c8f5be0ed730907d0edea8223562760d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BA=91=E6=8C=9A?= Date: Wed, 29 Jul 2026 23:22:20 -0500 Subject: [PATCH 2/5] Document and benchmark object-centric dispatch --- README.md | 157 +++++++++ benchmarks/bench_comm.py | 135 +++++++- benchmarks/bench_producer_direct_out.py | 383 ++++++++++++++++++++++ benchmarks/producer_direct_out_kernels.py | 167 ++++++++++ 4 files changed, 838 insertions(+), 4 deletions(-) create mode 100644 benchmarks/bench_producer_direct_out.py create mode 100644 benchmarks/producer_direct_out_kernels.py diff --git a/README.md b/README.md index 0bcfa72..e99bc54 100644 --- a/README.md +++ b/README.md @@ -79,6 +79,8 @@ buffer = Buffer(S=4096, H=7168, K=8, E=256, num_ep_ranks=8, - `num_sms=None` defaults to 32. `B` defaults to `E // num_ep_ranks`; an explicit value like `B=4` may also be passed. - `dispatch` / `combine` / `prefetch_weight` / `reduce_grad` all accept `async_finish=True` to run on the comm stream and return a CUDA event. +- `prepare_dispatch` / `prepare_source_dispatch` accept the same option so + planning can overlap independent compute before the prepared object is used. #### dispatch fwd @@ -172,6 +174,161 @@ output_sh, gathered_route_weights_sk, _ = buffer.combine( - The views alias buffer state that the next `dispatch` / `combine` overwrites — do not hold them across communication calls (autograd must not save them for backward; that case requires `zero_copy=False`). +#### producer direct-out + +When routing is known before the dispatch payload is produced, an upstream +producer with an explicitly compatible peer-VMM epilogue can skip the +standalone hidden-payload dispatch. Planning first returns a prepared +owner-sharded activation object: + +```python +target, cu_seqlens, plan = buffer.prepare_dispatch( + topk_experts_sk, + tokens_per_expert, +) + +# A compatible producer kernel reads target.plan.dst and writes every +# non-negative representative destination directly into target.mapped_hidden: +# +# row = (dst // target.NvS) * target.NvS_padded + (dst % target.NvS) +# producer_epilogue(..., target.mapped_hidden, target.plan.dst, ...) + +hidden_nvsh, route_weights_nvs, event = buffer.publish_dispatch( + target, + plan, + route_weights_sk=route_weights_sk, + zero_copy=True, +) +``` + +`publish_dispatch` does not recopy hidden payloads. It scatters route metadata, +zero-fills padded rows, builds duplicate metadata, publishes cross-rank +visibility, and expands same-rank duplicates locally. Publication is one-shot; +a target also expires when the same `Buffer` starts another dispatch, prepare, +or combine operation. + +For a compatible direct-out producer, metadata construction can overlap the +producer itself: + +```python +prepared_publication = buffer.prepare_dispatch_publication( + target, + plan, + route_weights_sk=route_weights_sk, +) + +# This launch runs on the caller stream while MoonEP prepares route metadata, +# padding, and dedup structures on its communication stream. +producer_epilogue(..., target.mapped_hidden, target.plan.dst, ...) + +hidden_nvsh, route_weights_nvs, done = buffer.publish_dispatch( + target, + plan, + async_finish=True, + zero_copy=True, + prepared_publication=prepared_publication, +) +``` + +The finish phase waits for both streams, publishes the producer's remote +writes with a post-producer visibility barrier, and then expands same-rank +duplicates locally. The ordinary one-call `publish_dispatch` path remains the +fallback when the framework cannot schedule this two-phase boundary. + +Planning can overlap independent work on the caller stream: + +```python +target, cu_seqlens, plan, planning_done = buffer.prepare_dispatch( + topk_experts_sk, + tokens_per_expert, + async_finish=True, +) +independent_work() +torch.cuda.current_stream().wait_event(planning_done) +# target.plan.dst is now ready for the producer epilogue. +``` + +The target capability is `remote_final_slot_store_v1`. Generic CUDA tensor +operations are not assumed safe for peer VMM access; the producer must be a +kernel explicitly validated for this mapping and destination contract. Call +ordinary `dispatch` when that capability is unavailable or when routing depends +on the payload being produced. In a conventional MoE where the router consumes +that same hidden tensor, this ordering constraint prevents a generic upstream +epilogue fusion; projected/latent payloads produced after routing are the +intended first integration case. + +[`bench_producer_direct_out.py`](benchmarks/bench_producer_direct_out.py) +provides a benchmark-only CUDA Tile producer with local-output and direct-out +epilogues. It verifies the public two-phase publication result bit-for-bit and +measures the producer-to-ready-dispatch boundary. It is a synthetic integration +benchmark, not a full-model latency claim, and `cuda.tile` is not imported by +the `moonep` runtime. + +#### owner-sharded source view + +For the strongest one-physical-copy activation model, construct the Buffer with +`enable_source_object=True`. This allocates one source `[S, H]` shard on each +rank and maps the logical object across the EP group: + +```python +source = buffer.source_storage() + +# The producer writes only its natural local owner placement. The router can +# consume this ordinary local tensor before producing top-k metadata. +producer_epilogue(..., source.local_hidden) +topk_experts_sk, route_weights_sk, tokens_per_expert = router( + source.local_hidden +) +source.local_route_weights.copy_(route_weights_sk) + +source_view, cu_seqlens, plan = buffer.prepare_source_dispatch( + source, + topk_experts_sk, + tokens_per_expert, +) + +# A compatible expert consumer reads source_view.slot_to_source, +# source_view.slot_to_primary, and source_view.slot_to_next, then directly +# loads representative rows from source_view.storage.mapped_hidden. +# peer_indexed_group_gemm(source_view, cu_seqlens, ...) + +# Or ask the runtime to produce the ordinary contiguous expert view using +# explicit owner-pull peer loads: +hidden_nvsh, route_weights_nvs, _ = buffer.materialize_source_dispatch( + source_view, + plan, + with_route_weights=True, + zero_copy=True, +) +``` + +Like `prepare_dispatch`, source planning accepts `async_finish=True` and +appends a CUDA event. Wait on that event immediately before the consumer reads +the returned source view. + +`prepare_source_dispatch` does not create routed hidden copies. For each valid +local expert-layout slot, `slot_to_source` encodes its source rank and original +token. `slot_to_primary` identifies the representative local slot for that +source token on the consumer rank, and `slot_to_next` links its repeated final +slots. The declared consumer capability is `peer_indexed_source_v1`. + +This mode is opt-in because it reserves an additional owner-sharded source +buffer, and because the checked-in repository does not include a compatible +peer-indexed group GEMM. Generic CUDA tensor operations must not read remote +rows. The runtime materializer is an explicitly compatible CUDA kernel; it +pulls hidden rows and route weights from their single physical owners, fills +each representative hidden row once per consumer rank, fans each loaded vector +directly from registers into same-rank top-k slots, and zeroes declared padding +without a dispatch payload kernel, dispatch dedup builder, or local hidden-row +reread. + +On four GB200s, the owner-pull materializer improved the ready-to-consume +`H=3584, K=8, S=8192, E=896` path by 55.17% across five fresh processes. +`H=7168, K=8` improved by 31.11% across another five fresh processes. One-run +guards also won for `H=3584, K=16` and extreme routing skew. The path remains +opt-in, with ordinary dispatch as the mandatory fallback for unmeasured cells +and unsupported consumers. + ```python # explicitly release VMM/NVLink resources held by the Buffer before destroying the process group buffer.destroy() diff --git a/benchmarks/bench_comm.py b/benchmarks/bench_comm.py index 4ed206b..9cd4fc3 100644 --- a/benchmarks/bench_comm.py +++ b/benchmarks/bench_comm.py @@ -17,6 +17,16 @@ Reports the MoonEP communication operators used by Megatron: - dispatch_fwd: dispatch kernel with route-weight scatter (dedup: only representative ``dst >= 0`` rows are transferred). + - publish_meta: producer-direct-out completion with the hidden payload + already in final owner placement. It scatters route metadata, zero-fills + padding, rebuilds dedup metadata, and publishes visibility. The difference + from dispatch_fwd is an upper bound on removable standalone payload-path + cost, not an end-to-end fused-producer speedup. + - source_pull: runtime materialization from the one-physical-copy + owner-sharded source object. It includes route metadata, padding, and + explicit peer-VMM row loads. Because it writes duplicate slots directly, + compare it with dispatch_fwd + epilogue_fwd for the same ready-to-consume + contiguous expert view. - dispatch_bwd: hidden-only dispatch kernel with the saved plan. - epilogue_fwd: in-place duplicate expansion on the NVL shard (dispatch_epilogue kernel), timed after dispatch_fwd. @@ -167,12 +177,15 @@ def time_gpu_op(launch_fn, warmup, iters, group, cudagraph=True): def bench_one(group, group_rank, R, S, K, E, H, bias_ratio, Hp, num_sms=32, warmup=5, iters=20, - cudagraph: bool = True): + cudagraph: bool = True, + source_pull_blocks: int = 0, + source_pull_threads: int = 64): """Run one (R, S, K, E, H, Hp, bias_ratio) configuration on the given subgroup.""" dev = torch.device(f"cuda:{torch.cuda.current_device()}") buffer = Buffer( S, H, K, E, R, num_sms=num_sms, group=group, + enable_source_object=True, explicitly_destroy=True, ) ctx = buffer._require_ctx() @@ -208,6 +221,21 @@ def _plan_call(): # The plan carries the full [R, B] experts_to_copy table needed by # prefetch/grad_reduce. launch_planning(ctx, topk_flat, tpe, cu_seqlens, plan) + source_storage = buffer.source_storage() + source_storage.local_hidden.copy_(hidden) + source_storage.local_route_weights.copy_(weights) + source_view, source_cu_seqlens, source_plan = \ + buffer.prepare_source_dispatch( + source_storage, + topk, + tpe, + inter_rank_sync=True, + ) + # Use the exact plan that carries the source provenance for every + # following comparison. The routing inputs are identical, so it is + # semantically equal to the plan used by the planning timing above. + plan = source_plan + cu_seqlens = source_cu_seqlens dst = plan.dst experts_to_copy = plan.experts_to_copy NvS = int(ctx['NvS']) @@ -224,6 +252,40 @@ def _dispatch_fwd_call(): _dispatch_fwd_call, warmup, iters, group, cudagraph=cudagraph ) + # ---- publish_meta: prepared-object completion after a compatible + # producer has already written representative hidden rows into final + # owner placement. Reuse the payload populated above as untimed setup. + # This deliberately rebuilds the same fresh-plan dedup structures as + # dispatch_fwd, so the only omitted work is the hidden source read and + # representative remote stores. + def _publish_meta_call(): + launch_dispatch( + ctx, + None, + weights, + plan, + write_payload=False, + build_dedup_map=True, + ) + + publish_meta_us = time_gpu_op( + _publish_meta_call, warmup, iters, group, cudagraph=cudagraph + ) + + def _publish_saved_meta_call(): + launch_dispatch( + ctx, + None, + weights, + plan, + write_payload=False, + build_dedup_map=False, + ) + + publish_saved_meta_us = time_gpu_op( + _publish_saved_meta_call, warmup, iters, group, cudagraph=cudagraph + ) + # ---- epilogue_fwd: in-place duplicate expansion on the NVL shard # (load primary once, store to every duplicate slot). Timed after # dispatch_fwd so it reads the freshly written NVL state, matching the @@ -237,6 +299,25 @@ def _epilogue_fwd_call(): _epilogue_fwd_call, warmup, iters, group, cudagraph=cudagraph ) + # ---- source_pull: explicit owner-pull of one representative row per + # source-token/consumer-rank. Each loaded vector fans directly from + # registers through the planner-published duplicate chain; a compact + # range kernel initializes padding. The result is ready for the ordinary + # contiguous expert consumer without the dispatch dedup builder. + def _source_pull_call(): + buffer.materialize_source_dispatch( + source_view, + plan, + with_route_weights=True, + zero_copy=True, + num_blocks=source_pull_blocks, + num_threads=source_pull_threads, + ) + + source_pull_us = time_gpu_op( + _source_pull_call, warmup, iters, group, cudagraph=cudagraph + ) + # ---- dispatch_bwd: hidden-only dispatch with the saved plan. grad_output = torch.randn(S, H, dtype=torch.bfloat16, device=dev) @@ -422,6 +503,14 @@ def _grad_reduce_call(): result = { 'planning_us': planning_us, 'dispatch_fwd_us': dispatch_fwd_us, + 'publish_meta_us': publish_meta_us, + 'publish_saved_meta_us': publish_saved_meta_us, + 'dispatch_payload_delta_us': dispatch_fwd_us - publish_meta_us, + 'dispatch_ready_us': dispatch_fwd_us + epilogue_fwd_us, + 'source_pull_us': source_pull_us, + 'source_pull_delta_us': ( + dispatch_fwd_us + epilogue_fwd_us - source_pull_us + ), 'dispatch_bwd_us': dispatch_bwd_us, 'epilogue_fwd_us': epilogue_fwd_us, 'epilogue_bwd_us': epilogue_bwd_us, @@ -520,8 +609,25 @@ def main(): # profiling unusual (ep, E) combos that the default ep*14 sweep doesn't # produce; the planning CuTe DSL path must have a matching specialization. ap.add_argument("--experts", type=int, default=None) + ap.add_argument( + "--source-pull-blocks", + type=int, + default=0, + help=( + "Owner-pull materializer grid size. Zero uses the measured " + "49152-block cap." + ), + ) + ap.add_argument( + "--source-pull-threads", + type=int, + choices=(64, 128, 256, 512), + default=64, + help="Threads per owner-pull materializer block.", + ) args = ap.parse_args() assert args.hp % 128 == 0, f"--hp must be a multiple of 128, got {args.hp}" + assert args.source_pull_blocks >= 0 world_rank, world_size, local_rank = setup() @@ -559,7 +665,10 @@ def main(): "ep", "H", "K", "S", "E", "Hp", "unbalance_ratio", "max_recv", "max_send", "load_max_mean", "planning_us", - "dispatch_fwd_us", "dispatch_bwd_us", + "dispatch_fwd_us", "publish_meta_us", "publish_saved_meta_us", + "dispatch_payload_delta_us", + "dispatch_ready_us", "source_pull_us", "source_pull_delta_us", + "dispatch_bwd_us", "epilogue_fwd_us", "epilogue_bwd_us", "combine_prologue_fwd_us", "combine_prologue_bwd_us", "combine_fwd_us", "combine_bwd_us", @@ -582,9 +691,13 @@ def main(): if world_rank == 0: header = (f"{'ep':>3} {'H':>5} {'K':>3} {'S':>5} {'E':>4} {'Hp':>5} " - f"{'unb_r':>5} {'mx_rcv':>6} {'mx_snd':>6} {'mx/mean':>7} " + f"{'unb_r':>5} " + f"{'mx_rcv':>6} {'mx_snd':>6} {'mx/mean':>7} " f"{'plan_us':>9} " - f"{'d_fwd':>9} {'d_bwd':>9} " + f"{'d_fwd':>9} {'pub_meta':>9} {'pub_saved':>9} " + f"{'pay_delta':>9} " + f"{'d_ready':>9} {'src_pull':>9} {'src_delta':>9} " + f"{'d_bwd':>9} " f"{'e_fwd':>9} {'e_bwd':>9} " f"{'cp_fwd':>9} {'cp_bwd':>9} " f"{'c_fwd':>9} {'c_bwd':>9} " @@ -617,6 +730,8 @@ def main(): num_sms=args.num_sms, warmup=args.warmup, iters=args.iters, cudagraph=args.cudagraph, + source_pull_blocks=args.source_pull_blocks, + source_pull_threads=args.source_pull_threads, ) except Exception: if world_rank == 0: @@ -638,6 +753,12 @@ def main(): f"{res['load_max_mean']:>7.3f} " f"{fmt4(res['planning_us']):>9} " f"{fmt4(res['dispatch_fwd_us']):>9} " + f"{fmt4(res['publish_meta_us']):>9} " + f"{fmt4(res['publish_saved_meta_us']):>9} " + f"{fmt4(res['dispatch_payload_delta_us']):>9} " + f"{fmt4(res['dispatch_ready_us']):>9} " + f"{fmt4(res['source_pull_us']):>9} " + f"{fmt4(res['source_pull_delta_us']):>9} " f"{fmt4(res['dispatch_bwd_us']):>9} " f"{fmt4(res['epilogue_fwd_us']):>9} " f"{fmt4(res['epilogue_bwd_us']):>9} " @@ -672,6 +793,12 @@ def main(): f"{res['load_max_mean']:.3f}", fmt4(res['planning_us']), fmt4(res['dispatch_fwd_us']), + fmt4(res['publish_meta_us']), + fmt4(res['publish_saved_meta_us']), + fmt4(res['dispatch_payload_delta_us']), + fmt4(res['dispatch_ready_us']), + fmt4(res['source_pull_us']), + fmt4(res['source_pull_delta_us']), fmt4(res['dispatch_bwd_us']), fmt4(res['epilogue_fwd_us']), fmt4(res['epilogue_bwd_us']), diff --git a/benchmarks/bench_producer_direct_out.py b/benchmarks/bench_producer_direct_out.py new file mode 100644 index 0000000..e4fba49 --- /dev/null +++ b/benchmarks/bench_producer_direct_out.py @@ -0,0 +1,383 @@ +"""Benchmark a producer GEMM with MoonEP direct-out publication. + +The control writes a token-major local output and then runs ordinary MoonEP +payload dispatch. The candidate replaces that local store with peer-VMM stores +to representative final-owner rows and runs metadata publication plus duplicate +expansion. The CUDA Tile producer is benchmark-only and is not a runtime +dependency of ``moonep``. + +Example: + + torchrun --nproc_per_node=4 benchmarks/bench_producer_direct_out.py \ + --ep 4 --hidden 3584 --topk 8 --unbalance-ratio 1.0 \ + --out results/direct_out.csv +""" + +import argparse +import csv +import os + +import torch +import torch.distributed as dist + +from benchmarks.bench_comm import setup, time_gpu_op +from moonep import Buffer +from moonep.dispatch import launch_dispatch +from moonep.dispatch_epilogue import launch_dispatch_epilogue +from moonep.inter_rank_sync import launch_inter_rank_sync +from tests.generate_topk_routing import generate_topk_routing + + +def distributed_mean(local_value: float, group) -> float: + """Average one timing value across the EP group.""" + device = torch.device(f"cuda:{torch.cuda.current_device()}") + value = torch.tensor([local_value], dtype=torch.float64, device=device) + gathered = [ + torch.empty_like(value) + for _ in range(dist.get_world_size(group=group)) + ] + dist.all_gather(gathered, value, group=group) + return torch.cat(gathered).mean().item() + + +def time_schedule(schedule, *, warmup: int, iters: int, group) -> float: + """Time a captured multi-stream schedule.""" + schedule(warmup) + torch.cuda.synchronize() + dist.barrier(group=group) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + schedule(iters) + torch.cuda.synchronize() + dist.barrier(group=group) + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + graph.replay() + end.record() + end.synchronize() + return distributed_mean( + start.elapsed_time(end) * 1e3 / iters, + group, + ) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--ep", type=int, default=4) + parser.add_argument("--seq-len", type=int, default=8192) + parser.add_argument("--experts", type=int, default=896) + parser.add_argument("--hidden", type=int, default=3584) + parser.add_argument("--producer-k", type=int, default=256) + parser.add_argument("--topk", type=int, default=8) + parser.add_argument("--unbalance-ratio", type=float, default=1.0) + parser.add_argument("--warmup", type=int, default=10) + parser.add_argument("--iters", type=int, default=100) + parser.add_argument("--out", required=True) + parser.add_argument("--tile-m", type=int, default=64) + parser.add_argument("--tile-n", type=int, default=128) + parser.add_argument("--tile-k", type=int, default=64) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + try: + from benchmarks.producer_direct_out_kernels import ( + launch_producer_gemm_direct_out, + launch_producer_gemm_local, + ) + except ImportError as exc: + raise RuntimeError( + "bench_producer_direct_out requires the optional cuda.tile package" + ) from exc + + rank, world, local_rank = setup() + if world != args.ep: + raise ValueError(f"--ep={args.ep} requires world_size={args.ep}, got {world}") + if args.experts % args.ep != 0: + raise ValueError("--experts must be divisible by --ep") + if args.topk > args.experts: + raise ValueError("--topk must not exceed --experts") + if args.seq_len % args.tile_m != 0: + raise ValueError("--seq-len must be divisible by --tile-m") + if args.hidden % args.tile_n != 0: + raise ValueError("--hidden must be divisible by --tile-n") + if args.producer_k % args.tile_k != 0: + raise ValueError("--producer-k must be divisible by --tile-k") + + group = dist.group.WORLD + device = torch.device(f"cuda:{local_rank}") + buffer = Buffer( + args.seq_len, + args.hidden, + args.topk, + args.experts, + args.ep, + num_sms=32, + group=group, + explicitly_destroy=True, + ) + ctx = buffer._require_ctx() + comm = buffer._comm_stream + if comm is None: + raise RuntimeError("MoonEP communication stream is unavailable") + topk, tpe = generate_topk_routing( + args.seq_len, + args.topk, + args.experts, + args.ep, + args.unbalance_ratio, + device, + 1234, + rank=rank, + ) + target, _, plan = buffer.prepare_dispatch( + topk, + tpe, + inter_rank_sync=True, + ) + + torch.manual_seed(3000 + rank) + producer_input = torch.empty( + (args.seq_len, args.producer_k), + dtype=torch.bfloat16, + device=device, + ).normal_(mean=0.0, std=0.1) + producer_weight = torch.empty( + (args.producer_k, args.hidden), + dtype=torch.bfloat16, + device=device, + ).normal_(mean=0.0, std=0.01) + local_output = torch.empty( + (args.seq_len, args.hidden), + dtype=torch.bfloat16, + device=device, + ) + route_weights = torch.rand( + (args.seq_len, args.topk), + dtype=torch.float32, + device=device, + ) + + def producer_local() -> None: + launch_producer_gemm_local( + producer_input, + producer_weight, + local_output, + tile_m=args.tile_m, + tile_n=args.tile_n, + tile_k=args.tile_k, + ) + + def producer_direct() -> None: + launch_producer_gemm_direct_out( + producer_input, + producer_weight, + plan.dst, + target.mapped_hidden, + nvs=int(ctx["NvS"]), + nvs_padded=int(ctx["NvS_padded"]), + topk=args.topk, + tile_m=args.tile_m, + tile_n=args.tile_n, + tile_k=args.tile_k, + ) + + def dispatch_payload() -> None: + launch_dispatch( + ctx, + local_output, + None, + plan, + build_dedup_map=True, + ) + launch_dispatch_epilogue(ctx, plan) + + def publish_only() -> None: + launch_dispatch( + ctx, + None, + None, + plan, + write_payload=False, + build_dedup_map=True, + ) + launch_dispatch_epilogue(ctx, plan) + + def prepare_metadata() -> None: + launch_dispatch( + ctx, + None, + None, + plan, + write_payload=False, + build_dedup_map=True, + ) + + def publish_visibility_and_expand() -> None: + launch_inter_rank_sync(ctx) + launch_dispatch_epilogue(ctx, plan) + + def baseline_boundary() -> None: + producer_local() + dispatch_payload() + + def direct_boundary() -> None: + producer_direct() + publish_only() + + ready_events = [ + torch.cuda.Event(enable_timing=False) + for _ in range(max(args.warmup, args.iters)) + ] + metadata_done_events = [ + torch.cuda.Event(enable_timing=False) + for _ in range(max(args.warmup, args.iters)) + ] + + def overlapped_direct_boundary(count: int) -> None: + main_stream = torch.cuda.current_stream() + for idx in range(count): + ready_events[idx].record(main_stream) + comm.wait_event(ready_events[idx]) + with torch.cuda.stream(comm): + prepare_metadata() + metadata_done_events[idx].record(comm) + producer_direct() + main_stream.wait_event(metadata_done_events[idx]) + publish_visibility_and_expand() + + baseline_boundary() + torch.cuda.synchronize() + baseline = ctx["hidden_buf_local"].clone() + direct_boundary() + torch.cuda.synchronize() + torch.testing.assert_close(ctx["hidden_buf_local"], baseline, rtol=0, atol=0) + overlapped_direct_boundary(1) + torch.cuda.synchronize() + torch.testing.assert_close(ctx["hidden_buf_local"], baseline, rtol=0, atol=0) + + producer_local_us = time_gpu_op( + producer_local, args.warmup, args.iters, group, cudagraph=True + ) + producer_direct_us = time_gpu_op( + producer_direct, args.warmup, args.iters, group, cudagraph=True + ) + dispatch_payload_us = time_gpu_op( + dispatch_payload, args.warmup, args.iters, group, cudagraph=True + ) + publish_us = time_gpu_op( + publish_only, args.warmup, args.iters, group, cudagraph=True + ) + baseline_us = time_gpu_op( + baseline_boundary, args.warmup, args.iters, group, cudagraph=True + ) + direct_us = time_gpu_op( + direct_boundary, args.warmup, args.iters, group, cudagraph=True + ) + overlapped_direct_us = time_schedule( + overlapped_direct_boundary, + warmup=args.warmup, + iters=args.iters, + group=group, + ) + direct_improvement = (baseline_us - direct_us) / baseline_us * 100.0 + overlap_improvement = ( + (baseline_us - overlapped_direct_us) / baseline_us * 100.0 + ) + further_improvement = ( + (direct_us - overlapped_direct_us) / direct_us * 100.0 + ) + + producer_local() + launch_dispatch( + ctx, + local_output, + route_weights, + plan, + build_dedup_map=True, + ) + launch_dispatch_epilogue(ctx, plan) + torch.cuda.synchronize() + public_hidden_reference = ctx["hidden_buf_local"].clone() + public_weights_reference = ( + ctx["weights_buf_local"].view(torch.float32).clone() + ) + prepared_publication = buffer.prepare_dispatch_publication( + target, + plan, + route_weights_sk=route_weights, + ) + producer_direct() + public_output, public_weights, public_done = buffer.publish_dispatch( + target, + plan, + async_finish=True, + zero_copy=True, + prepared_publication=prepared_publication, + ) + if public_done is None: + raise AssertionError("async publication did not return an event") + public_done.synchronize() + torch.testing.assert_close( + public_output, + public_hidden_reference, + rtol=0, + atol=0, + ) + if public_weights is None: + raise AssertionError("route-weight publication returned no view") + torch.testing.assert_close( + public_weights, + public_weights_reference, + rtol=0, + atol=0, + ) + + if rank == 0: + out_dir = os.path.dirname(args.out) + if out_dir: + os.makedirs(out_dir, exist_ok=True) + row = { + "ep": args.ep, + "H": args.hidden, + "producer_K": args.producer_k, + "S": args.seq_len, + "E": args.experts, + "topk": args.topk, + "unbalance_ratio": args.unbalance_ratio, + "producer_local_us": f"{producer_local_us:.4f}", + "producer_direct_us": f"{producer_direct_us:.4f}", + "dispatch_payload_us": f"{dispatch_payload_us:.4f}", + "publish_us": f"{publish_us:.4f}", + "baseline_boundary_us": f"{baseline_us:.4f}", + "direct_boundary_us": f"{direct_us:.4f}", + "direct_improvement_pct": f"{direct_improvement:.4f}", + "overlapped_direct_boundary_us": f"{overlapped_direct_us:.4f}", + "overlapped_direct_improvement_pct": ( + f"{overlap_improvement:.4f}" + ), + "further_improvement_pct": f"{further_improvement:.4f}", + } + with open(args.out, "w", newline="") as output: + writer = csv.DictWriter(output, fieldnames=list(row)) + writer.writeheader() + writer.writerow(row) + print( + f"H={args.hidden} K={args.topk} " + f"skew={args.unbalance_ratio:.2f} " + f"baseline={baseline_us:.2f} us " + f"direct={direct_us:.2f} us " + f"overlap_direct={overlapped_direct_us:.2f} us " + f"improvement={overlap_improvement:.2f}% " + f"further={further_improvement:.2f}%", + flush=True, + ) + + buffer.destroy() + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/benchmarks/producer_direct_out_kernels.py b/benchmarks/producer_direct_out_kernels.py new file mode 100644 index 0000000..999c161 --- /dev/null +++ b/benchmarks/producer_direct_out_kernels.py @@ -0,0 +1,167 @@ +"""Benchmark-only CUDA Tile producer kernels for direct-out publication.""" + +import math + +import cuda.tile as ct +import torch + + +ConstInt = ct.Constant[int] + + +@ct.kernel +def producer_gemm_local_kernel( + producer_input, + producer_weight, + output, + tile_m: ConstInt, + tile_n: ConstInt, + tile_k: ConstInt, +): + """Run a producer GEMM with an ordinary token-major local output.""" + bid = ct.bid(0) + num_n_tiles = ct.cdiv(output.shape[1], tile_n) + bid_m = bid // num_n_tiles + bid_n = bid % num_n_tiles + + accumulator = ct.zeros((tile_m, tile_n), dtype=ct.float32) + for k in range(0, ct.cdiv(producer_input.shape[1], tile_k)): + a = ct.load( + producer_input, + (bid_m, k), + shape=(tile_m, tile_k), + padding_mode=ct.PaddingMode.ZERO, + ) + b = ct.load( + producer_weight, + (k, bid_n), + shape=(tile_k, tile_n), + padding_mode=ct.PaddingMode.ZERO, + ) + accumulator = ct.mma(a, b, accumulator) + + ct.store( + output, + (bid_m, bid_n), + ct.astype(accumulator, output.dtype), + ) + + +@ct.kernel +def producer_gemm_direct_out_kernel( + producer_input, + producer_weight, + destinations, + mapped_output, + nvs: int, + nvs_padded: int, + topk: ConstInt, + tile_m: ConstInt, + tile_n: ConstInt, + tile_k: ConstInt, +): + """Run a producer GEMM that stores representative final-owner rows.""" + bid = ct.bid(0) + num_n_tiles = ct.cdiv(mapped_output.shape[1], tile_n) + bid_m = bid // num_n_tiles + bid_n = bid % num_n_tiles + + accumulator = ct.zeros((tile_m, tile_n), dtype=ct.float32) + for k in range(0, ct.cdiv(producer_input.shape[1], tile_k)): + a = ct.load( + producer_input, + (bid_m, k), + shape=(tile_m, tile_k), + padding_mode=ct.PaddingMode.ZERO, + ) + b = ct.load( + producer_weight, + (k, bid_n), + shape=(tile_k, tile_n), + padding_mode=ct.PaddingMode.ZERO, + ) + accumulator = ct.mma(a, b, accumulator) + + token_rows = bid_m * tile_m + ct.arange(tile_m, dtype=ct.int32) + output_cols = bid_n * tile_n + ct.arange(tile_n, dtype=ct.int32) + result = ct.astype(accumulator, mapped_output.dtype) + for kidx in range(topk): + token_topk = token_rows * topk + kidx + encoded = ct.gather(destinations, token_topk) + valid = encoded >= 0 + safe_encoded = ct.where(valid, encoded, 0) + destination_rank = safe_encoded // nvs + destination_loff = safe_encoded - destination_rank * nvs + physical_row = destination_rank * nvs_padded + destination_loff + ct.scatter( + mapped_output, + (physical_row[:, None], output_cols[None, :]), + result, + mask=valid[:, None], + ) + + +def launch_producer_gemm_local( + producer_input: torch.Tensor, + producer_weight: torch.Tensor, + output: torch.Tensor, + *, + tile_m: int, + tile_n: int, + tile_k: int, +) -> None: + """Launch the local-output producer.""" + grid = ( + math.ceil(output.shape[0] / tile_m) + * math.ceil(output.shape[1] / tile_n), + ) + ct.launch( + torch.cuda.current_stream(), + grid, + producer_gemm_local_kernel, + ( + producer_input, + producer_weight, + output, + tile_m, + tile_n, + tile_k, + ), + ) + + +def launch_producer_gemm_direct_out( + producer_input: torch.Tensor, + producer_weight: torch.Tensor, + destinations: torch.Tensor, + mapped_output: torch.Tensor, + *, + nvs: int, + nvs_padded: int, + topk: int, + tile_m: int, + tile_n: int, + tile_k: int, +) -> None: + """Launch the producer with a remote final-owner store epilogue.""" + grid = ( + math.ceil(producer_input.shape[0] / tile_m) + * math.ceil(mapped_output.shape[1] / tile_n), + ) + ct.launch( + torch.cuda.current_stream(), + grid, + producer_gemm_direct_out_kernel, + ( + producer_input, + producer_weight, + destinations, + mapped_output, + nvs, + nvs_padded, + topk, + tile_m, + tile_n, + tile_k, + ), + ) From bafd1e74330fcacf0a325de3ed19fbbbb86a1ec5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BA=91=E6=8C=9A?= Date: Wed, 29 Jul 2026 23:28:09 -0500 Subject: [PATCH 3/5] Validate object plan layouts before publication --- moonep/api.py | 6 +++- moonep/objects.py | 80 ++++++++++++++++++++++++++++++++++--------- tests/test_objects.py | 27 +++++++++++++++ 3 files changed, 96 insertions(+), 17 deletions(-) diff --git a/moonep/api.py b/moonep/api.py index 4224c7a..78d8d12 100644 --- a/moonep/api.py +++ b/moonep/api.py @@ -1023,6 +1023,8 @@ def _validate_source_view( ) view.validate_layout() self._validate_source_storage(ctx, view.storage) + if int(plan.E) != int(ctx['E']) or int(plan.B) != int(ctx['B']): + raise ValueError("source view plan E/B do not match this Buffer") if view.generation != self._dispatch_generation: raise RuntimeError( "source view is stale: this Buffer has started another " @@ -1224,8 +1226,10 @@ def _validate_dispatch_target( if not isinstance(target, MoonEPDispatchTarget): raise TypeError( "target must be returned by Buffer.prepare_dispatch" - ) + ) target.validate_layout() + if int(plan.E) != int(ctx['E']) or int(plan.B) != int(ctx['B']): + raise ValueError("dispatch target plan E/B do not match this Buffer") if target.buffer_token != id(self): raise ValueError( "dispatch target belongs to a different MoonEP Buffer" diff --git a/moonep/objects.py b/moonep/objects.py index 453e8d9..521bf1d 100644 --- a/moonep/objects.py +++ b/moonep/objects.py @@ -30,7 +30,49 @@ def _validate_tensor( return tensor -@dataclass(frozen=True) +def _validate_plan_layout( + plan: object, + *, + N: int, + R: int, + NvS: int, + K: int, + device: torch.device, +) -> MoonEPCommPlan: + if not isinstance(plan, MoonEPCommPlan): + raise TypeError("plan must be a MoonEPCommPlan") + for name, expected in (("N", N), ("R", R), ("NvS", NvS), ("K", K)): + actual = int(getattr(plan, name)) + if actual != expected: + raise ValueError( + f"plan.{name} must be {expected}, got {actual}" + ) + E = int(plan.E) + B = int(plan.B) + if E <= 0 or B <= 0: + raise ValueError(f"plan E and B must be positive, got E={E}, B={B}") + if E % R != 0: + raise ValueError(f"plan.E={E} must be divisible by R={R}") + for name, tensor, shape in ( + ("plan.dst", plan.dst, (N,)), + ("plan.experts_to_copy", plan.experts_to_copy, (R, B)), + ("plan.zero_fill_ranges", plan.zero_fill_ranges, (E + B, 2)), + ("plan.remote_stats", plan.remote_stats, (2,)), + ("plan.dup_groups", plan.dup_groups, (NvS, 3)), + ("plan.dup_loffs", plan.dup_loffs, (NvS,)), + ("plan.dup_counts", plan.dup_counts, (2,)), + ): + _validate_tensor( + name, + tensor, + dtype=torch.int32, + shape=shape, + device=device, + ) + return plan + + +@dataclass(frozen=True, slots=True) class MoonEPSourceStorage: """One physical source-activation shard per rank, globally mapped. @@ -87,7 +129,7 @@ def validate_layout(self) -> None: ) -@dataclass(frozen=True) +@dataclass(frozen=True, slots=True) class MoonEPSourceView: """Planned peer-indexed view over an owner-sharded source object. @@ -114,6 +156,8 @@ class MoonEPSourceView: value directly to every linked final slot without an intermediate local reread. + The plan and metadata tensors are opaque runtime state and must not be + modified by callers. """ storage: MoonEPSourceStorage @@ -154,13 +198,17 @@ def validate_layout(self) -> None: shape=(self.NvS,), device=device, ) - if not isinstance(self.plan, MoonEPCommPlan): - raise TypeError("plan must be a MoonEPCommPlan") - if self.plan.NvS != self.NvS or self.plan.K != self.K: - raise ValueError("source view and plan layouts do not match") + _validate_plan_layout( + self.plan, + N=self.storage.S * self.K, + R=self.storage.R, + NvS=self.NvS, + K=self.K, + device=device, + ) -@dataclass(frozen=True) +@dataclass(frozen=True, slots=True) class MoonEPDispatchTarget: """Prepared owner-sharded activation object for producer direct-out. @@ -178,7 +226,8 @@ class MoonEPDispatchTarget: The target is intentionally capability-scoped. Generic tensor consumers are not assumed safe on peer VMM mappings. Publication is one-shot, and the target expires when the same Buffer starts another dispatch, prepare, - or combine operation. + or combine operation. The plan tensors are opaque runtime state and must + not be modified by callers. """ mapped_hidden: torch.Tensor @@ -212,18 +261,17 @@ def validate_layout(self) -> None: shape=(self.NvS, self.H), device=mapped_hidden.device, ) - if not isinstance(self.plan, MoonEPCommPlan): - raise TypeError("plan must be a MoonEPCommPlan") - _validate_tensor( - "plan.dst", - self.plan.dst, - dtype=torch.int32, - shape=(self.S * self.K,), + _validate_plan_layout( + self.plan, + N=self.S * self.K, + R=self.R, + NvS=self.NvS, + K=self.K, device=mapped_hidden.device, ) -@dataclass(frozen=True) +@dataclass(frozen=True, slots=True) class MoonEPPreparedPublication: """Metadata publication scheduled ahead of producer completion. diff --git a/tests/test_objects.py b/tests/test_objects.py index 30b8946..2d6ce0a 100644 --- a/tests/test_objects.py +++ b/tests/test_objects.py @@ -76,6 +76,8 @@ def test_dispatch_target_state_checks_survive_python_optimization() -> None: ctx = { "hidden_buf": mapped_hidden, "hidden_buf_local": local_hidden, + "E": 2, + "B": 1, } target = replace(target, buffer_token=id(buffer)) @@ -89,3 +91,28 @@ def test_dispatch_target_state_checks_survive_python_optimization() -> None: buffer._published_dispatch_generation = 1 with pytest.raises(RuntimeError, match="already been published"): buffer._validate_dispatch_target(ctx, target, plan) + + +def test_dispatch_target_rejects_mutated_plan_layout() -> None: + plan = make_plan() + mapped_hidden = torch.empty((8, 8), dtype=torch.bfloat16) + target = MoonEPDispatchTarget( + mapped_hidden=mapped_hidden, + local_hidden=mapped_hidden[:4], + plan=plan, + generation=1, + buffer_token=1, + R=2, + S=2, + H=8, + K=1, + NvS=4, + NvS_padded=4, + ) + object.__setattr__( + plan, + "zero_fill_ranges", + torch.zeros((2, 2), dtype=torch.int32), + ) + with pytest.raises(ValueError, match="plan.zero_fill_ranges"): + target.validate_layout() From 58bde5cd4c251cba8d014e3e15b28ca3ad20d4fc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BA=91=E6=8C=9A?= Date: Wed, 29 Jul 2026 23:30:46 -0500 Subject: [PATCH 4/5] Keep source-object benchmark allocation opt-in --- benchmarks/bench_comm.py | 84 ++++++++++++++++++++++++---------------- 1 file changed, 51 insertions(+), 33 deletions(-) diff --git a/benchmarks/bench_comm.py b/benchmarks/bench_comm.py index 9cd4fc3..1c57ba7 100644 --- a/benchmarks/bench_comm.py +++ b/benchmarks/bench_comm.py @@ -22,11 +22,11 @@ padding, rebuilds dedup metadata, and publishes visibility. The difference from dispatch_fwd is an upper bound on removable standalone payload-path cost, not an end-to-end fused-producer speedup. - - source_pull: runtime materialization from the one-physical-copy - owner-sharded source object. It includes route metadata, padding, and - explicit peer-VMM row loads. Because it writes duplicate slots directly, - compare it with dispatch_fwd + epilogue_fwd for the same ready-to-consume - contiguous expert view. + - source_pull (with --enable-source-object): runtime materialization from + the one-physical-copy owner-sharded source object. It includes route + metadata, padding, and explicit peer-VMM row loads. Because it writes + duplicate slots directly, compare it with dispatch_fwd + epilogue_fwd for + the same ready-to-consume contiguous expert view. - dispatch_bwd: hidden-only dispatch kernel with the saved plan. - epilogue_fwd: in-place duplicate expansion on the NVL shard (dispatch_epilogue kernel), timed after dispatch_fwd. @@ -99,11 +99,13 @@ from moonep.prefetch import launch_prefetch -def fmt4(x: float) -> str: +def fmt4(x: float | None) -> str: """Format float to exactly 4 significant digits, keeping trailing zeros. Examples: 181.0 -> '181.0', 1486.65 -> '1487', 856.71 -> '856.7', 0.5 -> '0.5000'. """ + if x is None: + return "-" if x == 0: return "0.000" d = math.floor(math.log10(abs(x))) + 1 # digits before the decimal point @@ -178,6 +180,7 @@ def time_gpu_op(launch_fn, warmup, iters, group, cudagraph=True): def bench_one(group, group_rank, R, S, K, E, H, bias_ratio, Hp, num_sms=32, warmup=5, iters=20, cudagraph: bool = True, + enable_source_object: bool = False, source_pull_blocks: int = 0, source_pull_threads: int = 64): """Run one (R, S, K, E, H, Hp, bias_ratio) configuration on the given subgroup.""" @@ -185,7 +188,7 @@ def bench_one(group, group_rank, R, S, K, E, H, bias_ratio, Hp, buffer = Buffer( S, H, K, E, R, num_sms=num_sms, group=group, - enable_source_object=True, + enable_source_object=enable_source_object, explicitly_destroy=True, ) ctx = buffer._require_ctx() @@ -221,21 +224,24 @@ def _plan_call(): # The plan carries the full [R, B] experts_to_copy table needed by # prefetch/grad_reduce. launch_planning(ctx, topk_flat, tpe, cu_seqlens, plan) - source_storage = buffer.source_storage() - source_storage.local_hidden.copy_(hidden) - source_storage.local_route_weights.copy_(weights) - source_view, source_cu_seqlens, source_plan = \ - buffer.prepare_source_dispatch( - source_storage, - topk, - tpe, - inter_rank_sync=True, + source_view = None + if enable_source_object: + source_storage = buffer.source_storage() + source_storage.local_hidden.copy_(hidden) + source_storage.local_route_weights.copy_(weights) + source_view, source_cu_seqlens, source_plan = ( + buffer.prepare_source_dispatch( + source_storage, + topk, + tpe, + inter_rank_sync=True, + ) ) - # Use the exact plan that carries the source provenance for every - # following comparison. The routing inputs are identical, so it is - # semantically equal to the plan used by the planning timing above. - plan = source_plan - cu_seqlens = source_cu_seqlens + # Use the exact plan that carries the source provenance for every + # following comparison. The routing inputs are identical, so it is + # semantically equal to the plan used by the planning timing above. + plan = source_plan + cu_seqlens = source_cu_seqlens dst = plan.dst experts_to_copy = plan.experts_to_copy NvS = int(ctx['NvS']) @@ -304,20 +310,22 @@ def _epilogue_fwd_call(): # registers through the planner-published duplicate chain; a compact # range kernel initializes padding. The result is ready for the ordinary # contiguous expert consumer without the dispatch dedup builder. - def _source_pull_call(): - buffer.materialize_source_dispatch( - source_view, - plan, - with_route_weights=True, - zero_copy=True, - num_blocks=source_pull_blocks, - num_threads=source_pull_threads, + source_pull_us = None + if source_view is not None: + def _source_pull_call(): + buffer.materialize_source_dispatch( + source_view, + plan, + with_route_weights=True, + zero_copy=True, + num_blocks=source_pull_blocks, + num_threads=source_pull_threads, + ) + + source_pull_us = time_gpu_op( + _source_pull_call, warmup, iters, group, cudagraph=cudagraph ) - source_pull_us = time_gpu_op( - _source_pull_call, warmup, iters, group, cudagraph=cudagraph - ) - # ---- dispatch_bwd: hidden-only dispatch with the saved plan. grad_output = torch.randn(S, H, dtype=torch.bfloat16, device=dev) @@ -510,6 +518,7 @@ def _grad_reduce_call(): 'source_pull_us': source_pull_us, 'source_pull_delta_us': ( dispatch_fwd_us + epilogue_fwd_us - source_pull_us + if source_pull_us is not None else None ), 'dispatch_bwd_us': dispatch_bwd_us, 'epilogue_fwd_us': epilogue_fwd_us, @@ -609,6 +618,14 @@ def main(): # profiling unusual (ep, E) combos that the default ep*14 sweep doesn't # produce; the planning CuTe DSL path must have a matching specialization. ap.add_argument("--experts", type=int, default=None) + ap.add_argument( + "--enable-source-object", + action="store_true", + help=( + "Allocate and benchmark the opt-in owner-sharded source object. " + "The default preserves the original benchmark allocation." + ), + ) ap.add_argument( "--source-pull-blocks", type=int, @@ -730,6 +747,7 @@ def main(): num_sms=args.num_sms, warmup=args.warmup, iters=args.iters, cudagraph=args.cudagraph, + enable_source_object=args.enable_source_object, source_pull_blocks=args.source_pull_blocks, source_pull_threads=args.source_pull_threads, ) From 8090b1c608f7e010b32f119fb8c8ce8405e3e7b3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BA=91=E6=8C=9A?= Date: Thu, 30 Jul 2026 03:07:40 -0500 Subject: [PATCH 5/5] Finalize memory-semantic validation boundaries --- README.md | 21 ++++++++++----- benchmarks/bench_producer_direct_out.py | 2 +- moonep/dispatch.py | 32 +++++++++++++++-------- tests/test_dispatch.py | 6 ++--- tests/test_objects.py | 34 +++++++++++++++++++++++++ 5 files changed, 75 insertions(+), 20 deletions(-) diff --git a/README.md b/README.md index e99bc54..9675b6d 100644 --- a/README.md +++ b/README.md @@ -264,6 +264,14 @@ measures the producer-to-ready-dispatch boundary. It is a synthetic integration benchmark, not a full-model latency claim, and `cuda.tile` is not imported by the `moonep` runtime. +On four GB200s with `EP=4`, `S=8192`, and `E=896`, five fresh-process +repetitions of every `H={3584,7168} x K={8,16} x skew={0.1,1,5}` cell saved +86.21-274.49 us (7.93-44.25%) at that boundary when publication was overlapped. +The non-overlapped path regressed by 37.04-42.02 us for all `H=7168, K=16` +cells, while the overlapped path saved 86.21-109.79 us. A caller that cannot +schedule the two-phase overlap should therefore use ordinary dispatch for +those cells. + #### owner-sharded source view For the strongest one-physical-copy activation model, construct the Buffer with @@ -322,12 +330,13 @@ directly from registers into same-rank top-k slots, and zeroes declared padding without a dispatch payload kernel, dispatch dedup builder, or local hidden-row reread. -On four GB200s, the owner-pull materializer improved the ready-to-consume -`H=3584, K=8, S=8192, E=896` path by 55.17% across five fresh processes. -`H=7168, K=8` improved by 31.11% across another five fresh processes. One-run -guards also won for `H=3584, K=16` and extreme routing skew. The path remains -opt-in, with ordinary dispatch as the mandatory fallback for unmeasured cells -and unsupported consumers. +On four GB200s with `EP=4`, `S=8192`, and `E=896`, five fresh-process +repetitions of every `H={3584,7168} x K={8,16} x skew={0.1,1,5}` cell saved +38.84-628.98 us (5.10-65.00%) over ordinary ready-view dispatch. The smallest +margin was `H=7168, K=8, skew=5` at 38.84 us with 5.92 us candidate standard +deviation. The path remains opt-in, with ordinary dispatch as the mandatory +fallback outside the measured policy domain, below a deployment's safety +margin, or for unsupported consumers. ```python # explicitly release VMM/NVLink resources held by the Buffer before destroying the process group diff --git a/benchmarks/bench_producer_direct_out.py b/benchmarks/bench_producer_direct_out.py index e4fba49..b23e7aa 100644 --- a/benchmarks/bench_producer_direct_out.py +++ b/benchmarks/bench_producer_direct_out.py @@ -8,7 +8,7 @@ Example: - torchrun --nproc_per_node=4 benchmarks/bench_producer_direct_out.py \ + torchrun --nproc_per_node=4 --module benchmarks.bench_producer_direct_out \ --ep 4 --hidden 3584 --topk 8 --unbalance-ratio 1.0 \ --out results/direct_out.csv """ diff --git a/moonep/dispatch.py b/moonep/dispatch.py index 8563556..a17a7b6 100644 --- a/moonep/dispatch.py +++ b/moonep/dispatch.py @@ -920,19 +920,31 @@ def launch_dispatch( E = int(ctx['E']) B = int(ctx.get('B', 0)) - assert isinstance(write_payload, bool), \ - f"write_payload must be bool, got {type(write_payload).__name__}" + if not isinstance(write_payload, bool): + raise TypeError( + "write_payload must be bool, got " + f"{type(write_payload).__name__}" + ) if write_payload: - assert hidden_sh is not None, "hidden_sh is required when write_payload=True" - assert hidden_sh.dtype == torch.bfloat16 and hidden_sh.is_contiguous(), \ - "hidden_sh must be contiguous bf16" - assert hidden_sh.is_cuda, "hidden_sh must be a CUDA tensor" - assert tuple(hidden_sh.shape) == (S, H), \ - f"hidden_sh must be shape [S={S}, H={H}], got {tuple(hidden_sh.shape)}" + if hidden_sh is None: + raise ValueError( + "hidden_sh is required when write_payload=True" + ) + if hidden_sh.dtype != torch.bfloat16 or not hidden_sh.is_contiguous(): + raise ValueError("hidden_sh must be contiguous bf16") + if not hidden_sh.is_cuda: + raise ValueError("hidden_sh must be a CUDA tensor") + if tuple(hidden_sh.shape) != (S, H): + raise ValueError( + f"hidden_sh must be shape [S={S}, H={H}], " + f"got {tuple(hidden_sh.shape)}" + ) source_hidden = hidden_sh else: - assert hidden_sh is None, \ - "hidden_sh must be None when write_payload=False" + if hidden_sh is not None: + raise ValueError( + "hidden_sh must be None when write_payload=False" + ) source_hidden = ctx['hidden_buf_local'] _check_dispatch_plan(ctx, source_hidden, plan) assert ctx['hidden_buf'].dtype == torch.bfloat16 and ctx['hidden_buf'].is_contiguous() diff --git a/tests/test_dispatch.py b/tests/test_dispatch.py index 2a579e0..962b082 100644 --- a/tests/test_dispatch.py +++ b/tests/test_dispatch.py @@ -766,9 +766,9 @@ def test_dispatch_rejects_bad_inputs(dist_env): plan, _cu = allocate_planning_outputs(ctx) launch_planning(ctx, topk.reshape(-1).contiguous(), tpe, _cu, plan) - with pytest.raises(AssertionError, match="hidden_sh"): + with pytest.raises(ValueError, match="hidden_sh"): launch_dispatch(ctx, hidden.float(), weights, plan, build_dedup_map=True) - with pytest.raises(AssertionError, match="must be None"): + with pytest.raises(ValueError, match="must be None"): launch_dispatch( ctx, hidden, @@ -777,7 +777,7 @@ def test_dispatch_rejects_bad_inputs(dist_env): write_payload=False, build_dedup_map=True, ) - with pytest.raises(AssertionError, match="hidden_sh"): + with pytest.raises(ValueError, match="hidden_sh"): launch_dispatch( ctx, hidden[: case.S // 2].contiguous(), diff --git a/tests/test_objects.py b/tests/test_objects.py index 2d6ce0a..ad0d050 100644 --- a/tests/test_objects.py +++ b/tests/test_objects.py @@ -116,3 +116,37 @@ def test_dispatch_target_rejects_mutated_plan_layout() -> None: ) with pytest.raises(ValueError, match="plan.zero_fill_ranges"): target.validate_layout() + + +def test_dispatch_payload_mode_checks_survive_python_optimization() -> None: + from moonep.dispatch import launch_dispatch + + ctx = { + "H": 8, + "R": 2, + "S": 2, + "K": 1, + "E": 2, + "B": 1, + } + plan = make_plan() + + with pytest.raises(TypeError, match="write_payload must be bool"): + launch_dispatch( + ctx, + None, + None, + plan, + write_payload=1, + ) + with pytest.raises( + ValueError, + match="hidden_sh is required when write_payload=True", + ): + launch_dispatch( + ctx, + None, + None, + plan, + write_payload=True, + )