From 8841cdc765f7737b568e1346586848a92127bf08 Mon Sep 17 00:00:00 2001 From: Binyang Li Date: Mon, 13 Jul 2026 05:02:43 +0000 Subject: [PATCH] WIP --- include/mscclpp/gpu_data_types.hpp | 13 + python/mscclpp/ep/__init__.py | 4 +- python/mscclpp/ep/_cpp.py | 2 +- python/mscclpp/ep/communicator.py | 4 +- python/mscclpp/ep/low_latency.py | 83 +--- python/mscclpp/ep/types.py | 7 +- src/ext/ep/CMakeLists.txt | 4 +- src/ext/ep/bindings.cpp | 61 +-- src/ext/ep/config.hpp | 183 +++---- src/ext/ep/ht/config.hpp | 3 +- src/ext/ep/ht/kernels/buffer.cuh | 4 +- src/ext/ep/ht/kernels/internode.cu | 16 +- src/ext/ep/ht/kernels/internode_ncclep.cuh | 18 +- src/ext/ep/ht/kernels/intranode_kernel.cu | 12 +- src/ext/ep/ht/kernels/runtime.cu | 8 +- src/ext/ep/ht_runtime.cc | 4 +- src/ext/ep/ht_runtime.hpp | 4 +- src/ext/ep/{kernels => include}/api.cuh | 221 +++------ .../configs.cuh => include/constants.cuh} | 15 +- .../utils.cuh => include/device_helpers.cuh} | 88 ++-- src/ext/ep/{kernels => include}/exception.cuh | 2 +- src/ext/ep/{kernels => include}/launch.cuh | 2 +- src/ext/ep/{kernels => legacy}/low_latency.cu | 5 +- .../ep/{kernels => legacy}/quantization.cuh | 2 +- .../transport_helpers.cuh} | 2 +- src/ext/ep/low_latency/combine.cu | 461 ++++++++---------- src/ext/ep/low_latency/config.cuh | 111 +++-- src/ext/ep/low_latency/dispatch.cu | 286 +++++------ src/ext/ep/moe_runtime.cc | 353 +++----------- src/ext/ep/moe_runtime.hpp | 35 +- test/mscclpp-test/CMakeLists.txt | 4 +- test/python/ep/test_low_latency_multirank.py | 20 +- 32 files changed, 797 insertions(+), 1240 deletions(-) rename src/ext/ep/{kernels => include}/api.cuh (66%) rename src/ext/ep/{kernels/configs.cuh => include/constants.cuh} (74%) rename src/ext/ep/{kernels/utils.cuh => include/device_helpers.cuh} (86%) rename src/ext/ep/{kernels => include}/exception.cuh (98%) rename src/ext/ep/{kernels => include}/launch.cuh (99%) rename src/ext/ep/{kernels => legacy}/low_latency.cu (99%) rename src/ext/ep/{kernels => legacy}/quantization.cuh (99%) rename src/ext/ep/{kernels/common.cuh => legacy/transport_helpers.cuh} (99%) diff --git a/include/mscclpp/gpu_data_types.hpp b/include/mscclpp/gpu_data_types.hpp index b9d7191a..6bf9c979 100644 --- a/include/mscclpp/gpu_data_types.hpp +++ b/include/mscclpp/gpu_data_types.hpp @@ -4,6 +4,7 @@ #ifndef MSCCLPP_GPU_DATA_TYPES_HPP_ #define MSCCLPP_GPU_DATA_TYPES_HPP_ +#include #include #if defined(MSCCLPP_DEVICE_HIP) @@ -687,6 +688,18 @@ MSCCLPP_DEVICE_INLINE To to(const From& v) { } } +/// Convert a packed BF16 pair to a packed FP32 pair. +template <> +MSCCLPP_DEVICE_INLINE f32x2 to(const bf16x2& v) { + return __bfloat1622float2(v.storage); +} + +/// Convert a packed FP32 pair to a packed BF16 pair using round-to-nearest. +template <> +MSCCLPP_DEVICE_INLINE bf16x2 to(const f32x2& v) { + return __float22bfloat162_rn(v.storage); +} + #if defined(__FP8_TYPES_EXIST__) template <> MSCCLPP_DEVICE_INLINE __fp8_e4m3 min(const __fp8_e4m3& a, const __fp8_e4m3& b) { diff --git a/python/mscclpp/ep/__init__.py b/python/mscclpp/ep/__init__.py index 9c112b37..e8b63087 100644 --- a/python/mscclpp/ep/__init__.py +++ b/python/mscclpp/ep/__init__.py @@ -13,6 +13,7 @@ from .communicator import ( # noqa: F401 BlockOverlapConfig, CommOverlapConfig, CombineContext, + CombineMode, DispatchHandle, DispatchLayout, DispatchLayoutInfo, @@ -24,7 +25,6 @@ from .communicator import ( # noqa: F401 MoECommunicatorConfig, MoEMode, OperationOverlapConfig, - OptimizedCombineMode, QuantConfig, RowMajorInternodeDispatchHandle, RowMajorInternodeCombineContext, @@ -36,6 +36,7 @@ __all__ = [ "BlockOverlapConfig", "CommOverlapConfig", "CombineContext", + "CombineMode", "DispatchHandle", "DispatchLayout", "DispatchLayoutInfo", @@ -47,7 +48,6 @@ __all__ = [ "MoECommunicatorConfig", "MoEMode", "OperationOverlapConfig", - "OptimizedCombineMode", "QuantConfig", "RowMajorInternodeDispatchHandle", "RowMajorInternodeCombineContext", diff --git a/python/mscclpp/ep/_cpp.py b/python/mscclpp/ep/_cpp.py index 062bb602..d646cec2 100644 --- a/python/mscclpp/ep/_cpp.py +++ b/python/mscclpp/ep/_cpp.py @@ -14,7 +14,7 @@ except ImportError as exc: # pragma: no cover DispatchLayout = _cpp.DispatchLayout MoEMode = _cpp.MoEMode -OptimizedCombineMode = _cpp.OptimizedCombineMode +CombineMode = _cpp.CombineMode Config = getattr(_cpp, "Config", None) diff --git a/python/mscclpp/ep/communicator.py b/python/mscclpp/ep/communicator.py index 55466682..128612d7 100644 --- a/python/mscclpp/ep/communicator.py +++ b/python/mscclpp/ep/communicator.py @@ -8,7 +8,7 @@ from typing import Optional, Tuple import torch -from ._cpp import DispatchLayout, MoEMode, OptimizedCombineMode +from ._cpp import CombineMode, DispatchLayout, MoEMode from .high_throughput import HighThroughputBackend from .low_latency import LowLatencyBackend from .types import ( @@ -34,6 +34,7 @@ __all__ = [ "CommOverlapConfig", "BlockOverlapConfig", "CombineContext", + "CombineMode", "DispatchHandle", "DispatchLayout", "DispatchLayoutInfo", @@ -44,7 +45,6 @@ __all__ = [ "MoECommunicator", "MoECommunicatorConfig", "MoEMode", - "OptimizedCombineMode", "OperationOverlapConfig", "QuantConfig", "RowMajorInternodeDispatchHandle", diff --git a/python/mscclpp/ep/low_latency.py b/python/mscclpp/ep/low_latency.py index 1d1058d5..e41d7db9 100644 --- a/python/mscclpp/ep/low_latency.py +++ b/python/mscclpp/ep/low_latency.py @@ -8,7 +8,7 @@ from typing import Any, Optional import torch -from ._cpp import DispatchLayout, MoEMode, OptimizedCombineMode, _cpp, get_low_latency_rdma_size_hint +from ._cpp import CombineMode, DispatchLayout, MoEMode, _cpp, get_low_latency_rdma_size_hint from .types import ( DispatchHandle, DispatchLayoutInfo, @@ -19,13 +19,13 @@ from .types import ( MoECommunicatorConfig, QuantConfig, ) -from .utils import cuda_stream_ptr, requires_dequantization, resolve_expert_placement +from .utils import cuda_stream_ptr, resolve_expert_placement class LowLatencyRuntime: """Private low-level low-latency runtime wrapper (wraps ``_cpp.MoERuntime``).""" - num_sms: int = 64 + num_sms: int = 128 def __init__( self, @@ -33,22 +33,17 @@ class LowLatencyRuntime: num_nvl_bytes: int = 0, num_rdma_bytes: int = 0, mode: MoEMode = MoEMode.LOW_LATENCY, - num_qps_per_rank: int = 12, ) -> None: if not isinstance(mode, MoEMode): raise TypeError("mode must be a MoEMode") if mode != MoEMode.LOW_LATENCY: raise NotImplementedError("LowLatencyRuntime supports only MoEMode.LOW_LATENCY") - if num_qps_per_rank <= 0: - raise ValueError("num_qps_per_rank must be > 0") - self.mode = mode self.rank: int = comm.my_rank self.group_size: int = comm.nranks self.comm = comm self.num_nvl_bytes = num_nvl_bytes self.num_rdma_bytes = num_rdma_bytes - self.num_qps_per_rank = num_qps_per_rank self.cpp_runtime = _cpp.MoERuntime(comm.communicator, num_nvl_bytes, num_rdma_bytes, mode) def is_available(self) -> bool: @@ -90,8 +85,8 @@ class LowLatencyBackend: self.hidden_size = config.hidden_size self.topk = config.topk self.max_tokens_per_rank = config.max_tokens_per_rank - self.num_sms = config.low_latency_dispatch_num_sms - self.combine_num_sms = config.low_latency_combine_num_sms + self.num_blocks = config.low_latency_num_blocks + self.num_sms = self.num_blocks - 2 self.combine_mode = config.low_latency_combine_mode self.enable_overlap = config.enable_overlap @@ -99,12 +94,10 @@ class LowLatencyBackend: raise NotImplementedError("low-latency mode currently supports only DispatchLayout.EXPERT_MAJOR") if self.num_experts % self.world_size != 0: raise ValueError("low-latency mode requires num_experts divisible by world_size") - if not self.world_size <= self.num_sms <= 128: - raise ValueError("low_latency_dispatch_num_sms must be between world_size and 128") - if not 1 <= self.combine_num_sms <= 128: - raise ValueError("low_latency_combine_num_sms must be between 1 and 128") - if not isinstance(self.combine_mode, OptimizedCombineMode): - raise TypeError("low_latency_combine_mode must be an OptimizedCombineMode") + if not self.world_size + 2 <= self.num_blocks <= 130: + raise ValueError("low_latency_num_blocks must be between world_size + 2 and 130") + if not isinstance(self.combine_mode, CombineMode): + raise TypeError("low_latency_combine_mode must be a CombineMode") self.num_local_experts, self.local_expert_start = resolve_expert_placement( num_experts=self.num_experts, @@ -116,18 +109,12 @@ class LowLatencyBackend: if config.max_recv_tokens_per_rank not in (None, self.max_tokens_per_rank): raise NotImplementedError("low-latency mode currently uses max_tokens_per_rank as recv capacity") - self.quant = config.quant - self.quant_dtype = None if self.quant is None else self.quant.dtype - if self.quant is not None and self.quant_dtype is None: - raise ValueError("quant.dtype is required when quant is provided") - if self.quant_dtype not in (None, torch.float8_e4m3fn): - raise NotImplementedError(f"unsupported low-latency quant dtype: {self.quant_dtype}") - self.dispatch_requires_quantization = self.quant_dtype is not None + if config.quant is not None: + raise NotImplementedError("low-latency quantization is not implemented yet") num_rdma_bytes = get_low_latency_rdma_size_hint( self.max_tokens_per_rank, self.hidden_size, self.world_size, self.num_experts, self.topk ) - self._dispatch_scales: Optional[torch.Tensor] = None self._dispatch_src_info: Optional[torch.Tensor] = None self._dispatch_layout_range: Optional[torch.Tensor] = None self._dispatch_count: Optional[torch.Tensor] = None @@ -137,11 +124,8 @@ class LowLatencyBackend: num_nvl_bytes=0, num_rdma_bytes=num_rdma_bytes, mode=self.mode, - num_qps_per_rank=config.num_rdma_qps_per_rank, ) - # LL uses the registered symmetric buffer for both IPC/NVLink and RDMA-backed - # modes. A single-node LL job is not internode topology-wise. - # num_rdma_ranks > 1 iff world_size spans more than one local NVLink domain. + # LL uses direct peer mappings, including fabric IPC when available. self._is_internode = self._runtime.get_num_rdma_ranks() > 1 def is_available(self) -> bool: @@ -167,13 +151,12 @@ class LowLatencyBackend: del previous_handle self._validate_dispatch_inputs(input, topk_ids, weights, quant, output_buffer) - out_buf, packed_scales, src_info, layout_range, count = self._get_dispatch_output_tensors(output_buffer) + out_buf, src_info, layout_range, count = self._get_dispatch_output_tensors(output_buffer) self._runtime.cpp_runtime.dispatch( input.data_ptr(), topk_ids.data_ptr(), 0 if weights is None else weights.data_ptr(), out_buf.data_ptr(), - 0 if packed_scales is None else packed_scales.data_ptr(), src_info.data_ptr(), layout_range.data_ptr(), count.data_ptr(), @@ -182,17 +165,12 @@ class LowLatencyBackend: self.topk, self.max_tokens_per_rank, self.num_experts, - self.dispatch_requires_quantization, - self.output_layout, - self.num_sms, + self.num_blocks, cuda_stream_ptr(stream), ) - dispatched_quant = None - if packed_scales is not None: - dispatched_quant = QuantConfig(dtype=self.quant_dtype, block_scales=packed_scales, block_size=128) output_info = DispatchOutputInfo( layout=DispatchLayoutInfo(kind=self.output_layout, num_tokens_per_expert=count), - quant=dispatched_quant, + quant=None, ) dispatch_out = DispatchOutput( tokens=out_buf, @@ -226,17 +204,10 @@ class LowLatencyBackend: if not isinstance(handle, ExpertMajorDispatchHandle): raise ValueError("DispatchHandle does not contain expert-major combine context") context = handle.combine_context - combine_requires_dequantization = requires_dequantization(expert_output) - x_scales = None - if combine_requires_dequantization: - if handle.output_info.quant is None or handle.output_info.quant.block_scales is None: - raise ValueError("FP8 expert_output requires scales captured in the dispatch handle") - x_scales = handle.output_info.quant.block_scales if out is None: out = torch.empty((context.num_tokens, self.hidden_size), dtype=torch.bfloat16, device=expert_output.device) self._runtime.cpp_runtime.combine( expert_output.data_ptr(), - 0 if x_scales is None else x_scales.data_ptr(), context.topk_ids.data_ptr(), 0 if context.weights is None else context.weights.data_ptr(), context.src_info.data_ptr(), @@ -247,9 +218,8 @@ class LowLatencyBackend: self.topk, context.num_max_dispatch_tokens_per_rank, context.num_experts, - combine_requires_dequantization, self.combine_mode, - self.combine_num_sms, + self.num_blocks - 2, cuda_stream_ptr(stream), ) return out @@ -265,16 +235,8 @@ class LowLatencyBackend: (self.num_local_experts, self.world_size), dtype=torch.int64, device=device ) self._dispatch_count = torch.empty((self.num_local_experts,), dtype=torch.int32, device=device) - self._dispatch_scales = None - if self.dispatch_requires_quantization: - num_scales = self.hidden_size // 128 - scales_storage = torch.empty( - (self.num_local_experts, num_scales, slots_per_expert), dtype=torch.float32, device=device - ) - self._dispatch_scales = scales_storage.transpose(1, 2) return ( output_buffer, - self._dispatch_scales, self._dispatch_src_info, self._dispatch_layout_range, self._dispatch_count, @@ -283,8 +245,8 @@ class LowLatencyBackend: def _validate_dispatch_inputs(self, input, topk_ids, weights, quant, output_buffer) -> None: if output_buffer is None: raise ValueError("output_buffer is required for low-latency dispatch") - if quant is not None and (quant.block_scales is not None or quant.global_scale is not None): - raise NotImplementedError("low-latency dispatch does not support quantized input scales yet") + if quant is not None: + raise NotImplementedError("low-latency quantization is not implemented yet") if input.dim() != 2 or not input.is_contiguous(): raise ValueError("input must be a contiguous [num_tokens, hidden_size] tensor") if input.device.type != "cuda" or input.dtype != torch.bfloat16: @@ -306,7 +268,6 @@ class LowLatencyBackend: raise ValueError("weights must be a float32 CUDA tensor on the same device as input") if weights.shape != topk_ids.shape: raise ValueError("weights shape must match topk_ids") - expected_dtype = torch.float8_e4m3fn if self.dispatch_requires_quantization else torch.bfloat16 slots_per_expert = self.world_size * self.max_tokens_per_rank if self.output_layout == DispatchLayout.EXPERT_MAJOR: expected_shape = (self.num_local_experts, slots_per_expert, self.hidden_size) @@ -314,8 +275,8 @@ class LowLatencyBackend: expected_shape = (self.num_local_experts * slots_per_expert, self.hidden_size) if output_buffer.dim() != len(expected_shape) or not output_buffer.is_contiguous(): raise ValueError(f"output_buffer must be a contiguous {self.output_layout} tensor") - if output_buffer.device != input.device or output_buffer.dtype != expected_dtype: - raise ValueError(f"output_buffer must be a {expected_dtype} CUDA tensor on the same device as input") + if output_buffer.device != input.device or output_buffer.dtype != torch.bfloat16: + raise ValueError("output_buffer must be a BF16 CUDA tensor on the same device as input") if tuple(output_buffer.shape) != expected_shape: raise ValueError(f"output_buffer shape must be {expected_shape}") @@ -334,8 +295,8 @@ class LowLatencyBackend: raise ValueError("expert_output must keep dispatch output's contiguous layout") if tuple(expert_output.shape) != expected_shape: raise ValueError(f"expert_output shape must be {expected_shape}") - if expert_output.dtype not in (torch.bfloat16, getattr(torch, "float8_e4m3fn", None)): - raise ValueError("expert_output must be BF16 or FP8 E4M3") + if expert_output.dtype != torch.bfloat16: + raise ValueError("expert_output must be BF16") if out is not None: expected_out_shape = (context.num_tokens, self.hidden_size) if tuple(out.shape) != expected_out_shape or out.dtype != torch.bfloat16 or not out.is_contiguous(): diff --git a/python/mscclpp/ep/types.py b/python/mscclpp/ep/types.py index 4f8a7406..f739216a 100644 --- a/python/mscclpp/ep/types.py +++ b/python/mscclpp/ep/types.py @@ -10,7 +10,7 @@ from typing import Any, List, Optional, Union import torch import mscclpp -from ._cpp import DispatchLayout, MoEMode, OptimizedCombineMode +from ._cpp import CombineMode, DispatchLayout, MoEMode # Quantization metadata. @@ -56,9 +56,8 @@ class MoECommunicatorConfig: # Transport / launch tuning num_rdma_qps_per_rank: int = 12 num_sms: int = 20 - low_latency_dispatch_num_sms: int = 64 - low_latency_combine_num_sms: int = 64 - low_latency_combine_mode: OptimizedCombineMode = OptimizedCombineMode.DISABLED + low_latency_num_blocks: int = 130 + low_latency_combine_mode: CombineMode = CombineMode.RANK_LOCAL_REDUCE enable_overlap: bool = False # HT-only buffer/launch tuning (advanced) diff --git a/src/ext/ep/CMakeLists.txt b/src/ext/ep/CMakeLists.txt index d0f44289..36acf022 100644 --- a/src/ext/ep/CMakeLists.txt +++ b/src/ext/ep/CMakeLists.txt @@ -40,7 +40,6 @@ endif() set(EP_SOURCES moe_runtime.cc bindings.cpp - kernels/low_latency.cu low_latency/dispatch.cu low_latency/combine.cu # High-throughput (DeepEP-style) backend (torch-free, raw-pointer API). @@ -53,7 +52,7 @@ set(EP_SOURCES if(MSCCLPP_USE_ROCM) # ROCm port of the EP kernels is not supported yet; see src/ext/ep/README.md. message(WARNING "mscclpp_ep: ROCm build path not implemented, falling back to CXX compile.") - file(GLOB_RECURSE EP_CU_SOURCES kernels/*.cu ht/kernels/*.cu) + file(GLOB_RECURSE EP_CU_SOURCES low_latency/*.cu ht/kernels/*.cu) set_source_files_properties(${EP_CU_SOURCES} PROPERTIES LANGUAGE CXX) endif() @@ -65,6 +64,7 @@ endif() target_include_directories(mscclpp_ep_cpp PRIVATE ${CMAKE_CURRENT_SOURCE_DIR} + ${CMAKE_CURRENT_SOURCE_DIR}/include ${CMAKE_CURRENT_SOURCE_DIR}/ht ${PROJECT_SOURCE_DIR}/include ${PROJECT_SOURCE_DIR}/src/core/include diff --git a/src/ext/ep/bindings.cpp b/src/ext/ep/bindings.cpp index cd2769d8..ae331e2d 100644 --- a/src/ext/ep/bindings.cpp +++ b/src/ext/ep/bindings.cpp @@ -27,10 +27,10 @@ #include #include +#include "api.cuh" #include "config.hpp" #include "ht/config.hpp" #include "ht_runtime.hpp" -#include "kernels/api.cuh" #include "moe_runtime.hpp" namespace nb = nanobind; @@ -63,7 +63,7 @@ std::string bytesToString(nb::handle h) { NB_MODULE(mscclpp_ep_cpp, m) { m.doc() = "MSCCL++ Expert-Parallel (MoE dispatch/combine) extension"; - m.def("get_low_latency_rdma_size_hint", &mscclpp::ep::getLowLatencyRdmaSizeHint, + m.def("get_low_latency_rdma_size_hint", &mscclpp::ep::low_latency::getRdmaSizeHint, nb::arg("num_max_dispatch_tokens_per_rank"), nb::arg("hidden"), nb::arg("num_ranks"), nb::arg("num_experts"), nb::arg("num_topk")); @@ -77,10 +77,9 @@ NB_MODULE(mscclpp_ep_cpp, m) { .value("EXPERT_MAJOR", mscclpp::ep::DispatchLayout::EXPERT_MAJOR) .value("FLAT", mscclpp::ep::DispatchLayout::FLAT); - nb::enum_(m, "OptimizedCombineMode") - .value("DISABLED", mscclpp::ep::low_latency::OptimizedCombineMode::DISABLED) - .value("RANK_LOCAL_REDUCE", mscclpp::ep::low_latency::OptimizedCombineMode::RANK_LOCAL_REDUCE) - .value("DIRECT_SEND", mscclpp::ep::low_latency::OptimizedCombineMode::DIRECT_SEND); + nb::enum_(m, "CombineMode") + .value("RANK_LOCAL_REDUCE", mscclpp::ep::low_latency::CombineMode::RANK_LOCAL_REDUCE) + .value("DIRECT_SEND", mscclpp::ep::low_latency::CombineMode::DIRECT_SEND); nb::class_(m, "MoERuntime") .def(nb::init(), nb::arg("comm"), @@ -96,40 +95,34 @@ NB_MODULE(mscclpp_ep_cpp, m) { .def( "dispatch", [](mscclpp::ep::MoERuntime& self, uintptr_t inputPtr, uintptr_t topkIdxPtr, uintptr_t topkWeightsPtr, - uintptr_t outputPtr, uintptr_t outputScalesPtr, uintptr_t outputSrcInfoPtr, uintptr_t outputLayoutRangePtr, - uintptr_t outputCountPtr, int numTokens, int hidden, int numTopk, int numMaxDispatchTokensPerRank, - int numExperts, bool requiresQuantization, mscclpp::ep::DispatchLayout outputLayout, int numSms, + uintptr_t outputPtr, uintptr_t outputSrcInfoPtr, uintptr_t outputLayoutRangePtr, uintptr_t outputCountPtr, + int numTokens, int hidden, int numTopk, int maxTokensPerRank, int numExperts, int numBlocks, uintptr_t streamPtr) { - self.dispatch( - ptr(outputPtr), reinterpret_cast(ptr(outputScalesPtr)), - reinterpret_cast(ptr(outputSrcInfoPtr)), reinterpret_cast(ptr(outputLayoutRangePtr)), - reinterpret_cast(ptr(outputCountPtr)), ptr(inputPtr), reinterpret_cast(ptr(topkIdxPtr)), - reinterpret_cast(ptr(topkWeightsPtr)), numTokens, hidden, numTopk, numMaxDispatchTokensPerRank, - numExperts, requiresQuantization, outputLayout, numSms, stream(streamPtr)); + self.dispatch(ptr(outputPtr), reinterpret_cast(ptr(outputSrcInfoPtr)), + reinterpret_cast(ptr(outputLayoutRangePtr)), + reinterpret_cast(ptr(outputCountPtr)), ptr(inputPtr), + reinterpret_cast(ptr(topkIdxPtr)), reinterpret_cast(ptr(topkWeightsPtr)), + numTokens, hidden, numTopk, maxTokensPerRank, numExperts, numBlocks, stream(streamPtr)); }, nb::arg("input_ptr"), nb::arg("topk_idx_ptr"), nb::arg("topk_weights_ptr"), nb::arg("output_ptr"), - nb::arg("output_scales_ptr"), nb::arg("output_src_info_ptr"), nb::arg("output_layout_range_ptr"), - nb::arg("output_count_ptr"), nb::arg("num_tokens"), nb::arg("hidden"), nb::arg("num_topk"), - nb::arg("num_max_dispatch_tokens_per_rank"), nb::arg("num_experts"), nb::arg("requires_quantization"), - nb::arg("output_layout"), nb::arg("num_sms"), nb::arg("stream_ptr")) + nb::arg("output_src_info_ptr"), nb::arg("output_layout_range_ptr"), nb::arg("output_count_ptr"), + nb::arg("num_tokens"), nb::arg("hidden"), nb::arg("num_topk"), nb::arg("max_tokens_per_rank"), + nb::arg("num_experts"), nb::arg("num_blocks"), nb::arg("stream_ptr")) .def( "combine", - [](mscclpp::ep::MoERuntime& self, uintptr_t expertOutputPtr, uintptr_t expertScalesPtr, uintptr_t topkIdxPtr, - uintptr_t topkWeightsPtr, uintptr_t srcInfoPtr, uintptr_t layoutRangePtr, uintptr_t outputPtr, - int numTokens, int hidden, int numTopk, int numMaxDispatchTokensPerRank, int numExperts, - bool requiresDequantization, mscclpp::ep::low_latency::OptimizedCombineMode optimizedMode, int numBlocks, - uintptr_t streamPtr) { - self.combine(ptr(outputPtr), ptr(expertOutputPtr), reinterpret_cast(ptr(expertScalesPtr)), - reinterpret_cast(ptr(topkIdxPtr)), reinterpret_cast(ptr(topkWeightsPtr)), - reinterpret_cast(ptr(srcInfoPtr)), reinterpret_cast(ptr(layoutRangePtr)), - numTokens, hidden, numTopk, numMaxDispatchTokensPerRank, numExperts, requiresDequantization, - optimizedMode, numBlocks, stream(streamPtr)); + [](mscclpp::ep::MoERuntime& self, uintptr_t expertOutputPtr, uintptr_t topkIdxPtr, uintptr_t topkWeightsPtr, + uintptr_t srcInfoPtr, uintptr_t layoutRangePtr, uintptr_t outputPtr, int numTokens, int hidden, + int numTopk, int maxTokensPerRank, int numExperts, mscclpp::ep::low_latency::CombineMode mode, + int numBlocks, uintptr_t streamPtr) { + self.combine(ptr(outputPtr), ptr(expertOutputPtr), reinterpret_cast(ptr(topkIdxPtr)), + reinterpret_cast(ptr(topkWeightsPtr)), reinterpret_cast(ptr(srcInfoPtr)), + reinterpret_cast(ptr(layoutRangePtr)), numTokens, hidden, numTopk, maxTokensPerRank, + numExperts, mode, numBlocks, stream(streamPtr)); }, - nb::arg("expert_output_ptr"), nb::arg("expert_scales_ptr"), nb::arg("topk_idx_ptr"), - nb::arg("topk_weights_ptr"), nb::arg("src_info_ptr"), nb::arg("layout_range_ptr"), nb::arg("output_ptr"), - nb::arg("num_tokens"), nb::arg("hidden"), nb::arg("num_topk"), nb::arg("num_max_dispatch_tokens_per_rank"), - nb::arg("num_experts"), nb::arg("requires_dequantization"), nb::arg("optimized_mode"), nb::arg("num_blocks"), - nb::arg("stream_ptr")); + nb::arg("expert_output_ptr"), nb::arg("topk_idx_ptr"), nb::arg("topk_weights_ptr"), nb::arg("src_info_ptr"), + nb::arg("layout_range_ptr"), nb::arg("output_ptr"), nb::arg("num_tokens"), nb::arg("hidden"), + nb::arg("num_topk"), nb::arg("max_tokens_per_rank"), nb::arg("num_experts"), nb::arg("mode"), + nb::arg("num_blocks"), nb::arg("stream_ptr")); // ========================================================================== // High-throughput (HT) DeepEP-style backend: Config + MoEHighThroughputRuntime. diff --git a/src/ext/ep/config.hpp b/src/ext/ep/config.hpp index c631458b..b92b6563 100644 --- a/src/ext/ep/config.hpp +++ b/src/ext/ep/config.hpp @@ -6,10 +6,11 @@ #include #include #include +#include +#include #include -#include "kernels/configs.cuh" -#include "kernels/exception.cuh" +#include "constants.cuh" namespace mscclpp { namespace ep { @@ -24,7 +25,9 @@ __host__ __device__ constexpr dtype_t configAlign(dtype_t a, dtype_t b) { return configCellDiv(a, b) * b; } -// Low-latency rank-deduplicated dispatch payload layout: +namespace low_latency { + +// Rank-deduplicated dispatch payload layout: // // [data: DataType[hidden]] // [optional scales: ScaleType[hidden / scale_block_size]] @@ -32,27 +35,25 @@ __host__ __device__ constexpr dtype_t configAlign(dtype_t a, dtype_t b) { // [topKValues: float[topK]] // [srcTokenGlobalIdx: int] // -// The payload is 32-byte aligned as a whole. ScaleType=void means the payload is -// not quantized and has no scale section. +// The payload is 32-byte aligned as a whole. template -struct LowLatencyPayloadView { +struct PayloadView { static constexpr bool kHasScales = !std::is_void_v; - int hidden; - int topK; - int scaleBlockSize; - int numScales; + int hidden_; + int topK_; + int scaleBlockSize_; + int numScales_; size_t hiddenBytes_; size_t scaleOffset_; size_t metadataOffset_; size_t numBytes_; - MSCCLPP_HOST_DEVICE_INLINE static int nScales([[maybe_unused]] int hidden, [[maybe_unused]] int scaleBlockSize) { + MSCCLPP_HOST_DEVICE_INLINE static int numScales([[maybe_unused]] int hidden, [[maybe_unused]] int scaleBlockSize) { if constexpr (kHasScales) { return hidden / scaleBlockSize; - } else { - return 0; } + return 0; } MSCCLPP_HOST_DEVICE_INLINE static size_t hiddenBytes(int hidden) { @@ -62,26 +63,23 @@ struct LowLatencyPayloadView { MSCCLPP_HOST_DEVICE_INLINE static size_t scaleOffset(int hidden) { if constexpr (kHasScales) { return configAlign(hiddenBytes(hidden), alignof(ScaleType)); - } else { - return 0; } + return 0; } MSCCLPP_HOST_DEVICE_INLINE static size_t scaleBytes([[maybe_unused]] int hidden, [[maybe_unused]] int scaleBlockSize) { if constexpr (kHasScales) { - return static_cast(nScales(hidden, scaleBlockSize)) * sizeof(ScaleType); - } else { - return 0; + return static_cast(numScales(hidden, scaleBlockSize)) * sizeof(ScaleType); } + return 0; } - MSCCLPP_HOST_DEVICE_INLINE static size_t metadataOffset(int hidden, [[maybe_unused]] int scaleBlockSize) { + MSCCLPP_HOST_DEVICE_INLINE static size_t metadataOffset(int hidden, int scaleBlockSize) { if constexpr (kHasScales) { return configAlign(scaleOffset(hidden) + scaleBytes(hidden, scaleBlockSize), alignof(int)); - } else { - return configAlign(hiddenBytes(hidden), alignof(int)); } + return configAlign(hiddenBytes(hidden), alignof(int)); } MSCCLPP_HOST_DEVICE_INLINE static size_t metadataBytes(int topK) { @@ -92,11 +90,11 @@ struct LowLatencyPayloadView { return configAlign(metadataOffset(hidden, scaleBlockSize) + metadataBytes(topK), 32); } - MSCCLPP_HOST_DEVICE_INLINE LowLatencyPayloadView(int hidden, int topK, int scaleBlockSize = (kHasScales ? 128 : 0)) - : hidden(hidden), - topK(topK), - scaleBlockSize(scaleBlockSize), - numScales(nScales(hidden, scaleBlockSize)), + MSCCLPP_HOST_DEVICE_INLINE PayloadView(int hidden, int topK, int scaleBlockSize = (kHasScales ? 128 : 0)) + : hidden_(hidden), + topK_(topK), + scaleBlockSize_(scaleBlockSize), + numScales_(numScales(hidden, scaleBlockSize)), hiddenBytes_(hiddenBytes(hidden)), scaleOffset_(scaleOffset(hidden)), metadataOffset_(metadataOffset(hidden, scaleBlockSize)), @@ -126,121 +124,68 @@ struct LowLatencyPayloadView { } MSCCLPP_HOST_DEVICE_INLINE float* topKValues(void* base) const { - return reinterpret_cast(topKIndices(base) + topK); + return reinterpret_cast(topKIndices(base) + topK_); } MSCCLPP_HOST_DEVICE_INLINE const float* topKValues(const void* base) const { - return reinterpret_cast(topKIndices(base) + topK); + return reinterpret_cast(topKIndices(base) + topK_); } MSCCLPP_HOST_DEVICE_INLINE int* srcTokenGlobalIdx(void* base) const { - return reinterpret_cast(topKValues(base) + topK); + return reinterpret_cast(topKValues(base) + topK_); } MSCCLPP_HOST_DEVICE_INLINE const int* srcTokenGlobalIdx(const void* base) const { - return reinterpret_cast(topKValues(base) + topK); + return reinterpret_cast(topKValues(base) + topK_); } }; -struct LowLatencyBuffer { - int numCleanInt = 0; - - void* dispatchRdmaSendBuffer = nullptr; - void* dispatchRdmaRecvDataBuffer = nullptr; - // NOTE: signaling buffers are int64_t (not int) so that IB atomic ops - // (IBV_WR_ATOMIC_FETCH_AND_ADD is a 64-bit, 8-byte-aligned op) always - // target an 8-byte-aligned address. Using int32 slots produced unaligned - // atomics at odd indices that the NIC silently drops. - int64_t* dispatchRdmaRecvCountBuffer = nullptr; - - void* combineRdmaSendBuffer = nullptr; - void* combineRdmaRecvDataBuffer = nullptr; - int64_t* combineRdmaRecvFlagBuffer = nullptr; - - void* combineRdmaSendBufferDataStart = nullptr; - size_t numBytesPerCombineMsg = 0; - - std::pair cleanMeta() { - EP_HOST_ASSERT(dispatchRdmaRecvCountBuffer == combineRdmaRecvFlagBuffer); - return {dispatchRdmaRecvCountBuffer, numCleanInt}; - } +struct Buffer { + void* dispatchData_; + void* combineData_; + mscclpp::LL8Packet* combineReadyPackets_; }; -struct LowLatencyLayout { - size_t totalBytes = 0; - LowLatencyBuffer buffers[2]; +struct Layout { + size_t totalBytes_; + Buffer buffers_[2]; - template - out_ptr_t advance(const in_ptr_t& ptr, size_t count) { - return reinterpret_cast(reinterpret_cast(ptr) + count); - } + Layout(void* rdmaBuffer, int maxTokensPerRank, int hidden, int numRanks, int numExperts, int numTopk) { + const PayloadView<__bfloat16> bf16Payload(hidden, numTopk); + const PayloadView<__nv_fp8_storage_t, float> fp8Payload(hidden, numTopk, 128); + const size_t dispatchMetadataBytes = + configAlign(static_cast(2 * numRanks + numExperts) * sizeof(uint64_t), 128); + const size_t dispatchPayloadStride = + configAlign(std::max(bf16Payload.numBytes_, fp8Payload.numBytes_), 128); + const size_t dispatchBufferBytes = + dispatchMetadataBytes + static_cast(numRanks) * maxTokensPerRank * dispatchPayloadStride; + const size_t combineControlBytes = + configAlign(static_cast(numRanks) * sizeof(mscclpp::LL8Packet), 128); + const size_t combineBufferBytes = + combineControlBytes + static_cast(numExperts) * maxTokensPerRank * hidden * sizeof(__bfloat16); + const size_t bufferBytes = configAlign(std::max(dispatchBufferBytes, combineBufferBytes), 128); + totalBytes_ = 2 * bufferBytes; - LowLatencyLayout(void* rdmaBuffer, int numMaxDispatchTokensPerRank, int hidden, int numRanks, int numExperts, - int numTopk) { - // Dispatch and combine layout: - // - 2 symmetric odd/even send buffer - // - 2 symmetric odd/even receive buffers - // - 2 symmetric odd/even signaling buffers - - // Message sizes - // NOTES: you should add a control `int4` for combine messages if you want to do data transformation - const LowLatencyPayloadView bf16DispatchPayload(hidden, numTopk); - const LowLatencyPayloadView<__nv_fp8_storage_t, float> fp8DispatchPayload(hidden, numTopk, 128); - size_t numBytesPerDispatchMsg = std::max(bf16DispatchPayload.numBytes_, fp8DispatchPayload.numBytes_); - const size_t optDispatchMetadataBytes = - configAlign(static_cast(numRanks + numExperts) * sizeof(uint64_t), 128); - const size_t optDispatchPayloadStride = configAlign(bf16DispatchPayload.numBytes_, 128); - size_t numBytesPerCombineMsg = hidden * sizeof(nv_bfloat16); - - // Send buffer - size_t dispatchSendBufferBytes = std::max( - static_cast(numMaxDispatchTokensPerRank) * numBytesPerDispatchMsg, - optDispatchMetadataBytes + static_cast(numMaxDispatchTokensPerRank) * optDispatchPayloadStride); - size_t combineSendBufferBytes = numExperts * numMaxDispatchTokensPerRank * numBytesPerCombineMsg; - size_t sendBufferBytes = std::max(dispatchSendBufferBytes, combineSendBufferBytes); - EP_HOST_ASSERT(sendBufferBytes % sizeof(int4) == 0); - totalBytes += sendBufferBytes * 2; - - // Symmetric receive buffers - // TODO: optimize memory usages - size_t dispatchRecvDataBufferBytes = - std::max(static_cast(numRanks) * numMaxDispatchTokensPerRank * numBytesPerDispatchMsg, - optDispatchMetadataBytes + - static_cast(numRanks) * numMaxDispatchTokensPerRank * optDispatchPayloadStride); - size_t combineRecvBufferBytes = numExperts * numMaxDispatchTokensPerRank * numBytesPerCombineMsg; - size_t recvBufferBytes = std::max(dispatchRecvDataBufferBytes, combineRecvBufferBytes); - EP_HOST_ASSERT(recvBufferBytes % sizeof(int4) == 0); - totalBytes += recvBufferBytes * 2; - - // Symmetric signaling buffers (int64_t slots for 8-byte-aligned IB atomics). - size_t dispatchRecvCountBufferBytes = numExperts * sizeof(int64_t); - size_t combineRecvFlagBufferBytes = dispatchRecvCountBufferBytes; - size_t signalingBufferBytes = std::max(dispatchRecvCountBufferBytes, combineRecvFlagBufferBytes); - totalBytes += signalingBufferBytes * 2; - - // Assign pointers - // NOTES: we still leave some space for distinguishing dispatch/combine buffer, - // so you may see some parameters are duplicated - for (int i = 0; i < 2; ++i) { - buffers[i] = {static_cast(signalingBufferBytes / sizeof(int64_t)), - advance(rdmaBuffer, sendBufferBytes * i), - advance(rdmaBuffer, sendBufferBytes * 2 + recvBufferBytes * i), - advance(rdmaBuffer, sendBufferBytes * 2 + recvBufferBytes * 2 + signalingBufferBytes * i), - advance(rdmaBuffer, sendBufferBytes * i), - advance(rdmaBuffer, sendBufferBytes * 2 + recvBufferBytes * i), - advance(rdmaBuffer, sendBufferBytes * 2 + recvBufferBytes * 2 + signalingBufferBytes * i), - advance(rdmaBuffer, sendBufferBytes * i), - numBytesPerCombineMsg}; + if (rdmaBuffer != nullptr) { + auto* base = reinterpret_cast(rdmaBuffer); + for (int bufferIdx = 0; bufferIdx < 2; ++bufferIdx) { + auto* bufferBase = base + static_cast(bufferIdx) * bufferBytes; + buffers_[bufferIdx] = { + .dispatchData_ = bufferBase, + .combineData_ = bufferBase + combineControlBytes, + .combineReadyPackets_ = reinterpret_cast(bufferBase), + }; + } } } }; -inline size_t getLowLatencyRdmaSizeHint(int numMaxDispatchTokensPerRank, int hidden, int numRanks, int numExperts, - int numTopk) { - auto numBytes = - LowLatencyLayout(nullptr, numMaxDispatchTokensPerRank, hidden, numRanks, numExperts, numTopk).totalBytes; +inline size_t getRdmaSizeHint(int maxTokensPerRank, int hidden, int numRanks, int numExperts, int numTopk) { + const auto numBytes = Layout(nullptr, maxTokensPerRank, hidden, numRanks, numExperts, numTopk).totalBytes_; return configAlign(numBytes, NUM_BUFFER_ALIGNMENT_BYTES); } +} // namespace low_latency + } // namespace ep } // namespace mscclpp diff --git a/src/ext/ep/ht/config.hpp b/src/ext/ep/ht/config.hpp index c3b46c31..71bc1289 100644 --- a/src/ext/ep/ht/config.hpp +++ b/src/ext/ep/ht/config.hpp @@ -5,7 +5,8 @@ #include #include "../config.hpp" -#include "../kernels/api.cuh" +#include "api.cuh" +#include "exception.cuh" namespace mscclpp { namespace ep { diff --git a/src/ext/ep/ht/kernels/buffer.cuh b/src/ext/ep/ht/kernels/buffer.cuh index 9121d0b3..0ac87a64 100644 --- a/src/ext/ep/ht/kernels/buffer.cuh +++ b/src/ext/ep/ht/kernels/buffer.cuh @@ -4,8 +4,8 @@ #include -#include "../../kernels/configs.cuh" -#include "../../kernels/exception.cuh" +#include "constants.cuh" +#include "exception.cuh" namespace mscclpp { namespace ep { diff --git a/src/ext/ep/ht/kernels/internode.cu b/src/ext/ep/ht/kernels/internode.cu index a280f256..b62ddf66 100644 --- a/src/ext/ep/ht/kernels/internode.cu +++ b/src/ext/ep/ht/kernels/internode.cu @@ -5,11 +5,11 @@ #include #include -#include "../../kernels/configs.cuh" -#include "../../kernels/exception.cuh" -#include "../../kernels/launch.cuh" -#include "../../kernels/utils.cuh" #include "buffer.cuh" +#include "constants.cuh" +#include "device_helpers.cuh" +#include "exception.cuh" +#include "launch.cuh" namespace mscclpp { namespace ep { @@ -973,11 +973,11 @@ __global__ void __launch_bounds__(((kNumDispatchRDMASenderWarps + 1 + NUM_MAX_NV if ((slot_idx = __shfl_sync(0xffffffff, rdma_tail_idx, i)) >= 0) { slot_idx = slot_idx % num_max_rdma_chunked_recv_tokens; topk_ranks[num_topk_ranks] = i; - auto recv_is_token_in_rank_uint64 = broadcast(is_token_in_rank_uint64, i); + auto recv_is_token_in_rank_uint64 = warpBroadcast(is_token_in_rank_uint64, i); auto recv_is_token_in_rank_values = reinterpret_cast(&recv_is_token_in_rank_uint64); if (lane_id == num_topk_ranks) src_meta = SourceMeta(rdma_rank, recv_is_token_in_rank_values); dst_send_buffers[num_topk_ranks++] = - reinterpret_cast(broadcast(send_buffer, i)) + slot_idx * num_bytes_per_rdma_token; + reinterpret_cast(warpBroadcast(send_buffer, i)) + slot_idx * num_bytes_per_rdma_token; } EP_DEVICE_ASSERT(num_topk_ranks <= kNumTopkRDMARanks); @@ -2109,7 +2109,7 @@ __global__ void __launch_bounds__(kWarps * 32, 1) if (lane_id == 0) { const uint32_t mbar_a = static_cast(__cvta_generic_to_shared(&my_mbar[s])); asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(mbar_a)); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + fenceProxyAsyncSharedCta(); const uint32_t cbytes = static_cast(csize_int4 * static_cast(sizeof(int4))); for (int j = 0; j < num_topk_ranks; ++j) { const uint8_t* src = @@ -2142,7 +2142,7 @@ __global__ void __launch_bounds__(kWarps * 32, 1) } } __syncwarp(); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + fenceProxyAsyncSharedCta(); }; auto reduce_store = [&](int s, int c0, int csize_int4) { diff --git a/src/ext/ep/ht/kernels/internode_ncclep.cuh b/src/ext/ep/ht/kernels/internode_ncclep.cuh index 179d3f2e..67a27887 100644 --- a/src/ext/ep/ht/kernels/internode_ncclep.cuh +++ b/src/ext/ep/ht/kernels/internode_ncclep.cuh @@ -499,11 +499,11 @@ __global__ void __launch_bounds__(((kNumDispatchRDMASenderWarps + 1 + NUM_MAX_NV if ((slot_idx = __shfl_sync(0xffffffff, rdma_tail_idx, i)) >= 0) { slot_idx = slot_idx % num_max_rdma_chunked_recv_tokens; topk_ranks[num_topk_ranks] = i; - auto recv_is_token_in_rank_uint64 = broadcast(is_token_in_rank_uint64, i); + auto recv_is_token_in_rank_uint64 = warpBroadcast(is_token_in_rank_uint64, i); auto recv_is_token_in_rank_values = reinterpret_cast(&recv_is_token_in_rank_uint64); if (lane_id == num_topk_ranks) src_meta = SourceMeta(rdma_rank, recv_is_token_in_rank_values); dst_send_buffers[num_topk_ranks++] = - reinterpret_cast(broadcast(send_buffer, i)) + slot_idx * num_bytes_per_rdma_token; + reinterpret_cast(warpBroadcast(send_buffer, i)) + slot_idx * num_bytes_per_rdma_token; } EP_DEVICE_ASSERT(num_topk_ranks <= kNumTopkRDMARanks); @@ -521,7 +521,7 @@ __global__ void __launch_bounds__(((kNumDispatchRDMASenderWarps + 1 + NUM_MAX_NV #pragma unroll for (int j = 0; j < num_topk_ranks; ++j) { const int node = topk_ranks[j]; - NvlPackT bools_j = broadcast(is_token_in_rank_uint64, node); + NvlPackT bools_j = warpBroadcast(is_token_in_rank_uint64, node); const bool* bvals_j = reinterpret_cast(&bools_j); #pragma unroll for (int g = 0; g < NUM_MAX_NVL_PEERS; ++g) { @@ -562,7 +562,7 @@ __global__ void __launch_bounds__(((kNumDispatchRDMASenderWarps + 1 + NUM_MAX_NV const uint32_t smem_mbar = static_cast(__cvta_generic_to_shared(&ep_tma_mbar_snd[warp_id][stage])); asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(smem_mbar)); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + fenceProxyAsyncSharedCta(); asm volatile( "cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes " "[%0], [%1], %2, [%3];" ::"r"(smem_tile), @@ -603,7 +603,7 @@ __global__ void __launch_bounds__(((kNumDispatchRDMASenderWarps + 1 + NUM_MAX_NV : "=r"(done) : "r"(smem_mbar)); } - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + fenceProxyAsyncSharedCta(); // S2G: lane-striped fan-out -- one destination pool per lane. bool issued = false; for (int j = lane_id; j < ep_num_pools; j += 32) { @@ -652,8 +652,8 @@ __global__ void __launch_bounds__(((kNumDispatchRDMASenderWarps + 1 + NUM_MAX_NV #pragma unroll for (int j = 0; j < num_topk_ranks; ++j) { const int node = topk_ranks[j]; - SourceMeta sm_j = broadcast(src_meta, j); - NvlPackT bools_j = broadcast(is_token_in_rank_uint64, node); + SourceMeta sm_j = warpBroadcast(src_meta, j); + NvlPackT bools_j = warpBroadcast(is_token_in_rank_uint64, node); const bool* bvals_j = reinterpret_cast(&bools_j); #pragma unroll for (int g = 0; g < NUM_MAX_NVL_PEERS; ++g) { @@ -1174,7 +1174,7 @@ __global__ void __launch_bounds__(((kNumDispatchRDMASenderWarps + 1 + NUM_MAX_NV // G2S bulk (mbarrier complete_tx) -> arrive.expect_tx -> wait -> // fence -> S2G bulk -> commit -> wait. asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(smem_mbar)); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + fenceProxyAsyncSharedCta(); asm volatile( "cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes " "[%0], [%1], %2, [%3];" ::"r"(smem_tile), @@ -1192,7 +1192,7 @@ __global__ void __launch_bounds__(((kNumDispatchRDMASenderWarps + 1 + NUM_MAX_NV : "=r"(done) : "r"(smem_mbar)); } - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + fenceProxyAsyncSharedCta(); asm volatile("cp.async.bulk.global.shared::cta.bulk_group [%0], [%1], %2;" ::"l"(dst_b + off), "r"(smem_tile), "r"(chunk) : "memory"); diff --git a/src/ext/ep/ht/kernels/intranode_kernel.cu b/src/ext/ep/ht/kernels/intranode_kernel.cu index cf502e2b..19c93e5a 100644 --- a/src/ext/ep/ht/kernels/intranode_kernel.cu +++ b/src/ext/ep/ht/kernels/intranode_kernel.cu @@ -2,11 +2,11 @@ // Licensed under the MIT License. #include -#include "../../kernels/configs.cuh" -#include "../../kernels/exception.cuh" -#include "../../kernels/launch.cuh" -#include "../../kernels/utils.cuh" #include "buffer.cuh" +#include "constants.cuh" +#include "device_helpers.cuh" +#include "exception.cuh" +#include "launch.cuh" namespace mscclpp { namespace ep { @@ -1066,7 +1066,7 @@ __global__ void __launch_bounds__(kWarps * 32, 1) if (lane_id == 0) { const uint32_t mbar_a = static_cast(__cvta_generic_to_shared(&my_mbar[s])); asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(mbar_a)); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + fenceProxyAsyncSharedCta(); const uint32_t cbytes = static_cast(csize_int4 * static_cast(sizeof(int4))); for (int j = 0; j < num_topk_ranks; ++j) { const uint8_t* src = @@ -1099,7 +1099,7 @@ __global__ void __launch_bounds__(kWarps * 32, 1) } } __syncwarp(); - asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + fenceProxyAsyncSharedCta(); }; auto reduce_store = [&](int s, int c0, int csize_int4) { diff --git a/src/ext/ep/ht/kernels/runtime.cu b/src/ext/ep/ht/kernels/runtime.cu index 5159ca16..9d473ef8 100644 --- a/src/ext/ep/ht/kernels/runtime.cu +++ b/src/ext/ep/ht/kernels/runtime.cu @@ -9,10 +9,10 @@ // internode/NVSHMEM init helpers are deliberately omitted; the MSCCL++ port // uses `mscclpp::Bootstrap`/`ProxyService` instead of NVSHMEM. -#include "../../kernels/configs.cuh" -#include "../../kernels/exception.cuh" -#include "../../kernels/launch.cuh" -#include "../../kernels/utils.cuh" +#include "constants.cuh" +#include "device_helpers.cuh" +#include "exception.cuh" +#include "launch.cuh" namespace mscclpp { namespace ep { diff --git a/src/ext/ep/ht_runtime.cc b/src/ext/ep/ht_runtime.cc index e7e26994..ffbabfc5 100644 --- a/src/ext/ep/ht_runtime.cc +++ b/src/ext/ep/ht_runtime.cc @@ -34,8 +34,8 @@ #include #include -#include "kernels/api.cuh" -#include "kernels/configs.cuh" +#include "api.cuh" +#include "constants.cuh" namespace mscclpp { namespace ep { diff --git a/src/ext/ep/ht_runtime.hpp b/src/ext/ep/ht_runtime.hpp index c9db6c47..4066f5d3 100644 --- a/src/ext/ep/ht_runtime.hpp +++ b/src/ext/ep/ht_runtime.hpp @@ -26,9 +26,9 @@ #include #include +#include "constants.cuh" +#include "exception.cuh" #include "ht/config.hpp" -#include "kernels/configs.cuh" -#include "kernels/exception.cuh" namespace mscclpp { namespace ep { diff --git a/src/ext/ep/kernels/api.cuh b/src/ext/ep/include/api.cuh similarity index 66% rename from src/ext/ep/kernels/api.cuh rename to src/ext/ep/include/api.cuh index 81f049ba..3d0d6634 100644 --- a/src/ext/ep/kernels/api.cuh +++ b/src/ext/ep/include/api.cuh @@ -14,6 +14,7 @@ #include #include +#include #include #include @@ -188,67 +189,24 @@ void combine(cudaDataType_t type, void* combined_x, float* combined_topk_weights // =========================================================================== namespace low_latency { -/// Element type used by low-latency dispatch data path. -enum class DType { - /// NVIDIA bfloat16. - BF16, - /// NVIDIA FP8 E4M3. - F8E4M3 -}; +/// Number of non-worker blocks in the dispatch grid. +inline constexpr int DispatchControlBlocks = 2; +/// Maximum worker blocks used by dispatch or combine. +inline constexpr int MaxWorkerBlocks = 128; +/// Maximum total dispatch grid size. +inline constexpr int MaxDispatchBlocks = MaxWorkerBlocks + DispatchControlBlocks; -/// Optimized low-latency combine algorithm. -enum class OptimizedCombineMode { - /// Use the original combine implementation. - DISABLED, +/// Low-latency combine algorithm. +enum class CombineMode { /// Reduce expert rows on each destination rank before sending one partial per rank and token. RANK_LOCAL_REDUCE, /// Send every expert row directly and perform the full weighted reduction on the source rank. DIRECT_SEND }; -/// Transport context that encapsulates all transport-related state. -struct TransportContext { - /// Base address of the locally-registered RDMA buffer. - void* rdmaBufferBase_; - /// Port channel device handles for RDMA transport. - mscclpp::PortChannelDeviceHandle* portChannels_; - /// Base memory channel handles for IPC transport barrier (nullable). - mscclpp::BaseMemoryChannelDeviceHandle* memoryChannels_; - /// Number of expert channel slots reserved per peer rank. - int memoryChannelStride_; - /// Peer-mapped base addresses for IPC path (nullable). - void* const* peerBases_; - /// True if IPC path is ready, false to use RDMA. - bool ipcReady_; - /// CUDA device ID associated with this transport. - int deviceId_; - /// Current rank ID. - int rank_; - /// Total number of ranks. - int numRanks_; - /// Number of ranks in one IPC-reachable domain (0 when IPC is unavailable). - int ranksPerIpcDomain_; -}; - -/// Buffer set that encapsulates ping-pong buffer layout. -struct BufferSet { - /// Send data buffer. - void* sendDataBuffer_; - /// Send count buffer. - int64_t* sendCountBuffer_; - /// Receive data buffer. - void* recvDataBuffer_; - /// Receive count buffer. - int64_t* recvCountBuffer_; - /// Cleanup region for next iteration. - int64_t* cleanupRegion_; - /// Size of cleanup region in int64_t elements. - int cleanupSize_; -}; - -/// Configuration for dispatch operation. -struct DispatchConfig { - /// Number of input tokens to dispatch. +/// Per-call low-latency workload dimensions. +struct Workload { + /// Number of local input or output tokens. int numTokens_; /// Hidden dimension size. int hidden_; @@ -256,110 +214,67 @@ struct DispatchConfig { int numTopk_; /// Total number of experts. int numExperts_; - /// Maximum tokens per rank in packed layout. - int numMaxTokensPerRank_; - /// Input dtype for source tokens. - DType inputDType_; - /// Output dtype for packed receive buffer. - DType outputDType_; - /// Logical output layout. - DispatchLayout outputLayout_ = DispatchLayout::EXPERT_MAJOR; + /// Maximum tokens per rank in the packed layout. + int maxTokensPerRank_; }; -/// Configuration for combine operation. -struct CombineConfig { - /// Number of tokens to combine. - int numCombinedTokens_; - /// Hidden dimension size. - int hidden_; - /// Number of top-k experts per token. - int numTopk_; - /// Total number of experts. - int numExperts_; - /// Maximum tokens per rank in packed layout. - int numMaxTokensPerRank_; - /// Input dtype for expert outputs. - DType inputDType_; - /// Output dtype for combined tokens. - DType outputDType_; - /// True to use zero-copy optimization. - bool zeroCopy_; +/// Persistent communication resources shared by low-latency operations. +struct CommContext { + /// Base address of the locally-registered RDMA buffer. + void* rdmaBufferBase_; + /// Peer-mapped base addresses. + void* const* peerBases_; + /// Maximum shared memory available to one block after opt-in. + int maxSharedMemoryPerBlock_; + /// Number of streaming multiprocessors on the device. + int numSms_; + /// CUDA device ID associated with this communicator. + int deviceId_; + /// Current rank ID. + int rank_; + /// Total number of ranks. + int numRanks_; }; -/// Phase control for send/recv operations. -enum Phase { - /// Execute send phase only. - SEND_ONLY = 0x1, - /// Execute recv phase only. - RECV_ONLY = 0x2, - /// Execute both send and recv phases. - SEND_AND_RECV = 0x3 -}; +/// Low-latency dispatch that distributes tokens to experts across ranks. +/// @param[out] output Expert-major packed output +/// [num_local_experts, num_ranks * max_tokens_per_rank, hidden]. +/// @param[out] outputSrcInfo Original source-token index for every packed expert row. +/// @param[out] outputLayout Per-[local expert, source rank] packed count and offset. +/// @param[out] outputCount Total packed token count for every local expert. +/// @param[in] input Local input tokens [num_tokens, hidden]. +/// @param[in] topkIdx Global expert indices [num_tokens, num_topk]. +/// @param[in] topkWeights Routing weights [num_tokens, num_topk], or nullptr for unit weights. +/// @param[in] workload Per-call workload dimensions. +/// @param[in,out] recvBuffer Current symmetric ping-pong buffer used for incoming payloads and rewritten metadata. +/// @param[in] comm Persistent communication context. +/// @param[in,out] workspace Persistent counters, task storage, semaphores, and device barriers. +/// @param[in] numBlocks Total dispatch grid size, including one scheduler and one metadata-notify block. +/// @param[in] stream CUDA stream. +void dispatch(void* output, int* outputSrcInfo, int64_t* outputLayout, int* outputCount, const void* input, + const int64_t* topkIdx, const float* topkWeights, const Workload& workload, void* recvBuffer, + const CommContext& comm, void* workspace, int numBlocks, cudaStream_t stream); -/// Clean low-latency buffers (both ping-pong buffers). -/// @param buffer0 First cleanup region pointer. -/// @param numInt0 Size of first cleanup region in int64_t elements. -/// @param buffer1 Second cleanup region pointer. -/// @param numInt1 Size of second cleanup region in int64_t elements. -/// @param transport Transport context with channel handles and topology info. -/// @param stream CUDA stream to launch the kernel on. -void cleanBuffers(int64_t* buffer0, int numInt0, int64_t* buffer1, int numInt1, const TransportContext& transport, - cudaStream_t stream); - -/// Low-latency dispatch kernel that distributes tokens to experts across ranks. -/// @param output Output packed data buffer. EXPERT_MAJOR shape is -/// [num_local_experts, num_ranks*max_tokens, hidden]; FLAT is the same -/// local-expert-major storage viewed as [num_local_experts*num_ranks*max_tokens, hidden]. -/// @param outputScales FP8 scales (nullable if not using FP8). -/// @param outputSrcInfo Source rank info per token. -/// @param outputLayout Layout range [expert, rank] -> (offset, count). -/// @param outputCount Total count per expert. -/// @param input Input tokens [num_tokens, hidden]. -/// @param topkIdx Expert indices [num_tokens, num_topk]. -/// @param topkWeights Routing weights [num_tokens, num_topk]. -/// @param config Dispatch configuration. -/// @param currentBuffer Current iteration buffer set. -/// @param nextBuffer Next iteration buffer set (for cleanup). -/// @param transport Transport context. -/// @param workspace Temporary workspace buffer. -/// @param stream CUDA stream to launch the kernel on. -/// @param phase Phase control (default: SEND_AND_RECV). -void dispatch(void* output, float* outputScales, int* outputSrcInfo, int64_t* outputLayout, int* outputCount, - const void* input, const int64_t* topkIdx, const float* topkWeights, const DispatchConfig& config, - const BufferSet& currentBuffer, const BufferSet& nextBuffer, const TransportContext& transport, - void* workspace, cudaStream_t stream, Phase phase = SEND_AND_RECV); - -/// Optimized BF16 IPC dispatch compatible with the expert-major combine metadata contract. -void dispatchOptimized(void* output, int* outputSrcInfo, int64_t* outputLayout, int* outputCount, const void* input, - const int64_t* topkIdx, const float* topkWeights, const DispatchConfig& config, - const BufferSet& currentBuffer, const TransportContext& transport, void* workspace, int maxSms, - cudaStream_t stream); - -/// Experimental BF16 IPC combine. -void combineOptimized(void* output, const void* input, const int64_t* topkIdx, const float* topkWeights, - const int* srcInfo, const int64_t* layoutRange, const CombineConfig& config, - const BufferSet& currentBuffer, void* dispatchRecvBuffer, const TransportContext& transport, - void* workspace, int numBlocks, OptimizedCombineMode mode, cudaStream_t stream); - -/// Low-latency combine kernel that aggregates expert outputs back to tokens. -/// @param output Combined output [num_combined_tokens, hidden]. -/// @param input Expert outputs [num_local_experts * num_ranks * max_tokens, hidden]. -/// @param inputScales Optional FP8 scales for expert outputs. -/// @param topkIdx Expert indices [num_combined_tokens, num_topk]. -/// @param topkWeights Weights [num_combined_tokens, num_topk]. -/// @param srcInfo Source rank info. -/// @param layoutRange Layout range. -/// @param config Combine configuration. -/// @param currentBuffer Current iteration buffer set. -/// @param nextBuffer Next iteration buffer set (for cleanup). -/// @param transport Transport context. -/// @param workspace Temporary workspace buffer. -/// @param stream CUDA stream to launch the kernel on. -/// @param phase Phase control (default: SEND_AND_RECV). -void combine(void* output, const void* input, const float* inputScales, const int64_t* topkIdx, - const float* topkWeights, const int* srcInfo, const int64_t* layoutRange, const CombineConfig& config, - const BufferSet& currentBuffer, const BufferSet& nextBuffer, const TransportContext& transport, - void* workspace, cudaStream_t stream, Phase phase = SEND_AND_RECV); +/// Low-latency combine that aggregates expert outputs back to tokens. +/// @param[out] output Combined local tokens [num_tokens, hidden]. +/// @param[in] input Expert outputs [num_local_experts, num_ranks * max_tokens_per_rank, hidden]. +/// @param[in] topkIdx Global expert indices [num_tokens, num_topk]. +/// @param[in] topkWeights Routing weights [num_tokens, num_topk], or nullptr for unit weights. +/// @param[in] srcInfo Original source-token index for every packed expert row. +/// @param[in] layoutRange Per-[local expert, source rank] packed count and offset. +/// @param[in] workload Per-call workload dimensions. +/// @param[in,out] recvBuffer Current symmetric ping-pong buffer receiving partials or expert rows. +/// @param[in,out] readyPackets One readiness packet per sender rank in the current symmetric buffer. +/// @param[in] dispatchRecvBuffer Previous dispatch buffer containing rewritten routing metadata. +/// @param[in] comm Persistent communication context. +/// @param[in,out] workspace Persistent dispatch metadata plus the combine device barrier. +/// @param[in] numBlocks Number of combine blocks. +/// @param[in] mode Combine algorithm. +/// @param[in] stream CUDA stream. +void combine(void* output, const void* input, const int64_t* topkIdx, const float* topkWeights, const int* srcInfo, + const int64_t* layoutRange, const Workload& workload, void* recvBuffer, mscclpp::LL8Packet* readyPackets, + void* dispatchRecvBuffer, const CommContext& comm, void* workspace, int numBlocks, CombineMode mode, + cudaStream_t stream); } // namespace low_latency diff --git a/src/ext/ep/kernels/configs.cuh b/src/ext/ep/include/constants.cuh similarity index 74% rename from src/ext/ep/kernels/configs.cuh rename to src/ext/ep/include/constants.cuh index b44d3ce3..ef583aa4 100644 --- a/src/ext/ep/kernels/configs.cuh +++ b/src/ext/ep/include/constants.cuh @@ -4,10 +4,7 @@ // Portions adapted from DeepEP (https://github.com/deepseek-ai/DeepEP), // branch `chhwang/dev-atomic-add-cleanup`. Licensed under the MIT License. // -// Kernel-side configuration. This is the MSCCL++ version of -// `DeepEP/csrc/kernels/configs.cuh` with NVSHMEM / IBGDA / mlx5dv includes -// removed so the intranode (NVLink-only) kernels can be built standalone. -// Include this file **only** from `.cu` files. +// Shared EP capacity, timeout, alignment, and CUDA type configuration. #pragma once @@ -39,12 +36,6 @@ #endif #define NUM_WAIT_NANOSECONDS 500 -// Phase control for low-latency kernels (internal use in kernel code) -// Public API uses low_latency::Phase enum instead -#define LOW_LATENCY_SEND_PHASE 1 // = low_latency::SEND_ONLY -#define LOW_LATENCY_RECV_PHASE 2 // = low_latency::RECV_ONLY -// (SEND_AND_RECV = 3) - // Make CLion CUDA indexing work. #ifdef __CLION_IDE__ #define __CUDA_ARCH__ 900 // NOLINT(*-reserved-identifier) @@ -73,7 +64,3 @@ __host__ __device__ __forceinline__ void host_device_printf(const char* format, #include #include #include - -// NVSHMEM / IBGDA / mlx5dv are only required for the RDMA internode paths and -// are not included here. The internode/low-latency kernels that need them -// will include them directly under `#ifdef MSCCLPP_EP_HAVE_NVSHMEM`. diff --git a/src/ext/ep/kernels/utils.cuh b/src/ext/ep/include/device_helpers.cuh similarity index 86% rename from src/ext/ep/kernels/utils.cuh rename to src/ext/ep/include/device_helpers.cuh index 558441f5..c026024e 100644 --- a/src/ext/ep/kernels/utils.cuh +++ b/src/ext/ep/include/device_helpers.cuh @@ -3,6 +3,7 @@ #pragma once #include +#include #include #include "exception.cuh" @@ -17,24 +18,24 @@ #define UNROLLED_WARP_COPY(UNROLL_FACTOR, LANE_ID, N, DST, SRC, LD_FUNC, ST_FUNC) \ { \ - constexpr int kLoopStride = WARP_SIZE * (UNROLL_FACTOR); \ + constexpr int LoopStride = WARP_SIZE * (UNROLL_FACTOR); \ typename std::remove_reference::type unrolled_values[(UNROLL_FACTOR)]; \ auto __src = (SRC); \ auto __dst = (DST); \ - for (int __i = (LANE_ID); __i < ((N) / kLoopStride) * kLoopStride; __i += kLoopStride) { \ + for (int __i = (LANE_ID); __i < ((N) / LoopStride) * LoopStride; __i += LoopStride) { \ _Pragma("unroll") for (int __j = 0; __j < (UNROLL_FACTOR); ++__j) unrolled_values[__j] = \ LD_FUNC(__src + __i + __j * WARP_SIZE); \ _Pragma("unroll") for (int __j = 0; __j < (UNROLL_FACTOR); ++__j) \ ST_FUNC(__dst + __i + __j * WARP_SIZE, unrolled_values[__j]); \ } \ - for (int __i = ((N) / kLoopStride) * kLoopStride + (LANE_ID); __i < (N); __i += WARP_SIZE) \ + for (int __i = ((N) / LoopStride) * LoopStride + (LANE_ID); __i < (N); __i += WARP_SIZE) \ ST_FUNC(__dst + __i, LD_FUNC(__src + __i)); \ } namespace mscclpp { namespace ep { -template +template struct VecInt {}; template <> struct VecInt<1> { @@ -65,6 +66,17 @@ __device__ __forceinline__ void memory_fence_gpu() { asm volatile("fence.acq_rel __device__ __forceinline__ void memory_fence_cta() { asm volatile("fence.acq_rel.cta;" ::: "memory"); } +MSCCLPP_DEVICE_INLINE void publishLl8Packet(mscclpp::LL8Packet *packet, uint32_t value, uint32_t flag) { + memory_fence(); + packet->write(value, flag); +} + +MSCCLPP_DEVICE_INLINE uint32_t waitLl8Packet(const mscclpp::LL8Packet *packet, uint32_t flag) { + const uint32_t value = packet->read(flag, -1); + memory_fence(); + return value; +} + __device__ __forceinline__ void *peerBufferPtr(void *localBuffer, void *localBufferBase, void *peerBufferBase) { if (localBufferBase == nullptr) return peerBufferBase; const auto offset = reinterpret_cast(localBuffer) - reinterpret_cast(localBufferBase); @@ -72,14 +84,18 @@ __device__ __forceinline__ void *peerBufferPtr(void *localBuffer, void *localBuf } #if defined(__CUDACC__) -__device__ __forceinline__ void initTmaBarrier(uint64_t *sharedBarrier) { - const uint32_t barrierAddress = static_cast(__cvta_generic_to_shared(sharedBarrier)); - asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(barrierAddress)); +__device__ __forceinline__ void fenceProxyAsyncSharedCta() { asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); } -__device__ __forceinline__ void issueTmaG2S(const void *source, void *sharedTile, uint64_t *sharedBarrier, - uint32_t nBytes) { +__device__ __forceinline__ void initTmaLoadBarrier(uint64_t *sharedBarrier) { + const uint32_t barrierAddress = static_cast(__cvta_generic_to_shared(sharedBarrier)); + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(barrierAddress)); + fenceProxyAsyncSharedCta(); +} + +__device__ __forceinline__ void issueTmaLoad(const void *source, void *sharedTile, uint64_t *sharedBarrier, + uint32_t nBytes) { const uint32_t tileAddress = static_cast(__cvta_generic_to_shared(sharedTile)); const uint32_t barrierAddress = static_cast(__cvta_generic_to_shared(sharedBarrier)); asm volatile( @@ -93,7 +109,7 @@ __device__ __forceinline__ void issueTmaG2S(const void *source, void *sharedTile : "r"(barrierAddress), "r"(nBytes)); } -__device__ __forceinline__ void waitTmaG2S(uint64_t *sharedBarrier, uint32_t &phase) { +__device__ __forceinline__ void waitTmaLoad(uint64_t *sharedBarrier, uint32_t &phase) { const uint32_t barrierAddress = static_cast(__cvta_generic_to_shared(sharedBarrier)); uint32_t done = 0; while (!done) { @@ -106,7 +122,7 @@ __device__ __forceinline__ void waitTmaG2S(uint64_t *sharedBarrier, uint32_t &ph phase ^= 1; } -__device__ __forceinline__ void issueTmaS2G(void *destination, void *sharedTile, uint32_t nBytes) { +__device__ __forceinline__ void issueTmaStore(void *destination, void *sharedTile, uint32_t nBytes) { const uint32_t tileAddress = static_cast(__cvta_generic_to_shared(sharedTile)); asm volatile("cp.async.bulk.global.shared::cta.bulk_group [%0], [%1], %2;" ::"l"(destination), "r"(tileAddress), "r"(nBytes) @@ -114,12 +130,16 @@ __device__ __forceinline__ void issueTmaS2G(void *destination, void *sharedTile, asm volatile("cp.async.bulk.commit_group;"); } -template -__device__ __forceinline__ void waitTmaS2GRead() { - asm volatile("cp.async.bulk.wait_group.read %0;" ::"n"(kNumPendingGroups) : "memory"); +template +__device__ __forceinline__ void waitBulkGroupRead() { + // Wait until at most NumPendingGroups committed bulk groups may still read shared memory. + asm volatile("cp.async.bulk.wait_group.read %0;" ::"n"(NumPendingGroups) : "memory"); } -__device__ __forceinline__ void waitTmaS2G() { asm volatile("cp.async.bulk.wait_group 0;" ::: "memory"); } +__device__ __forceinline__ void waitBulkGroup() { + // Wait for every committed bulk group to complete. + asm volatile("cp.async.bulk.wait_group 0;" ::: "memory"); +} #endif __device__ __forceinline__ void st_relaxed_sys_global(const int *ptr, int val) { @@ -388,15 +408,17 @@ __device__ __forceinline__ void unpack2(const dtype_b_t &packed, dtype_a_t &x, d x = unpacked_ptr[0], y = unpacked_ptr[1]; } -template -__device__ __forceinline__ dtype_t broadcast(dtype_t &ptr, int src_lane_idx) { - EP_STATIC_ASSERT(sizeof(dtype_t) % sizeof(int) == 0, ""); - auto send_int_values = reinterpret_cast(&ptr); - int recv_int_values[sizeof(dtype_t) / sizeof(int)]; +template +__device__ __forceinline__ T warpBroadcast(T value, int sourceLane) { + EP_STATIC_ASSERT(sizeof(T) % sizeof(int) == 0, ""); + const auto *sourceValues = reinterpret_cast(&value); + T result; + auto *resultValues = reinterpret_cast(&result); #pragma unroll - for (int i = 0; i < sizeof(dtype_t) / sizeof(int); ++i) - recv_int_values[i] = __shfl_sync(0xffffffff, send_int_values[i], src_lane_idx); - return *reinterpret_cast(recv_int_values); + for (int i = 0; i < sizeof(T) / sizeof(int); ++i) { + resultValues[i] = __shfl_sync(0xffffffff, sourceValues[i], sourceLane); + } + return result; } __forceinline__ __device__ int warp_reduce_sum(int value) { @@ -438,23 +460,23 @@ __forceinline__ __device__ int get_lane_id() { return lane_id; } -template +template __forceinline__ __device__ void move_fifo_slots(int &head) { - head = (head + kNumRanks) % NUM_MAX_FIFO_SLOTS; + head = (head + NumRanks) % NUM_MAX_FIFO_SLOTS; } -template +template __device__ __forceinline__ bool not_finished(int *task, int expected) { auto result = false; auto lane_id = threadIdx.x % WARP_SIZE; - if (lane_id < kNumRanks) result = ld_volatile_global(task + lane_id) != expected; + if (lane_id < NumRanks) result = ld_volatile_global(task + lane_id) != expected; return __any_sync(0xffffffff, result); } -template +template __forceinline__ __device__ void timeout_check(int **task_fifo_ptrs, int head, int rank, int expected, int tag = 0) { auto start_time = clock64(); - while (not_finished(task_fifo_ptrs[rank] + head, expected)) { + while (not_finished(task_fifo_ptrs[rank] + head, expected)) { if (clock64() - start_time > NUM_TIMEOUT_CYCLES and threadIdx.x == 0) { printf("DeepEP timeout check failed: %d (rank = %d)\n", tag, rank); trap(); @@ -462,17 +484,17 @@ __forceinline__ __device__ void timeout_check(int **task_fifo_ptrs, int head, in } } -template +template __forceinline__ __device__ void barrier_device(int **task_fifo_ptrs, int head, int rank, int tag = 0) { auto thread_id = static_cast(threadIdx.x); - EP_DEVICE_ASSERT(kNumRanks <= WARP_SIZE); + EP_DEVICE_ASSERT(NumRanks <= WARP_SIZE); - if (thread_id < kNumRanks) { + if (thread_id < NumRanks) { atomicAdd_system(task_fifo_ptrs[rank] + head + thread_id, FINISHED_SUM_TAG); memory_fence(); atomicSub_system(task_fifo_ptrs[thread_id] + head + rank, FINISHED_SUM_TAG); } - timeout_check(task_fifo_ptrs, head, rank, 0, tag); + timeout_check(task_fifo_ptrs, head, rank, 0, tag); } } // namespace ep diff --git a/src/ext/ep/kernels/exception.cuh b/src/ext/ep/include/exception.cuh similarity index 98% rename from src/ext/ep/kernels/exception.cuh rename to src/ext/ep/include/exception.cuh index 60c09556..28960167 100644 --- a/src/ext/ep/kernels/exception.cuh +++ b/src/ext/ep/include/exception.cuh @@ -5,7 +5,7 @@ #include #include -#include "configs.cuh" +#include "constants.cuh" #ifndef EP_STATIC_ASSERT #define EP_STATIC_ASSERT(cond, reason) static_assert(cond, reason) diff --git a/src/ext/ep/kernels/launch.cuh b/src/ext/ep/include/launch.cuh similarity index 99% rename from src/ext/ep/kernels/launch.cuh rename to src/ext/ep/include/launch.cuh index f41e05ab..2ddfe1e7 100644 --- a/src/ext/ep/kernels/launch.cuh +++ b/src/ext/ep/include/launch.cuh @@ -2,7 +2,7 @@ // Licensed under the MIT License. #pragma once -#include "configs.cuh" +#include "constants.cuh" #ifndef SETUP_LAUNCH_CONFIG #define SETUP_LAUNCH_CONFIG(num_sms, num_threads, stream) \ diff --git a/src/ext/ep/kernels/low_latency.cu b/src/ext/ep/legacy/low_latency.cu similarity index 99% rename from src/ext/ep/kernels/low_latency.cu rename to src/ext/ep/legacy/low_latency.cu index 55a6c82e..9f52f78a 100644 --- a/src/ext/ep/kernels/low_latency.cu +++ b/src/ext/ep/legacy/low_latency.cu @@ -28,12 +28,11 @@ #include "../config.hpp" #include "api.cuh" -#include "common.cuh" -#include "configs.cuh" +#include "constants.cuh" #include "exception.cuh" #include "launch.cuh" #include "quantization.cuh" -#include "utils.cuh" +#include "transport_helpers.cuh" namespace cg = cooperative_groups; diff --git a/src/ext/ep/kernels/quantization.cuh b/src/ext/ep/legacy/quantization.cuh similarity index 99% rename from src/ext/ep/kernels/quantization.cuh rename to src/ext/ep/legacy/quantization.cuh index 08648d16..feb85880 100644 --- a/src/ext/ep/kernels/quantization.cuh +++ b/src/ext/ep/legacy/quantization.cuh @@ -6,7 +6,7 @@ #include #include -#include "utils.cuh" +#include "device_helpers.cuh" namespace mscclpp { namespace ep { diff --git a/src/ext/ep/kernels/common.cuh b/src/ext/ep/legacy/transport_helpers.cuh similarity index 99% rename from src/ext/ep/kernels/common.cuh rename to src/ext/ep/legacy/transport_helpers.cuh index 004758f4..b94a3a23 100644 --- a/src/ext/ep/kernels/common.cuh +++ b/src/ext/ep/legacy/transport_helpers.cuh @@ -8,7 +8,7 @@ #include #include -#include "utils.cuh" +#include "device_helpers.cuh" namespace mscclpp { namespace ep { diff --git a/src/ext/ep/low_latency/combine.cu b/src/ext/ep/low_latency/combine.cu index 8d642681..383f805c 100644 --- a/src/ext/ep/low_latency/combine.cu +++ b/src/ext/ep/low_latency/combine.cu @@ -1,101 +1,102 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. -#include "../kernels/api.cuh" -#include "../kernels/exception.cuh" -#include "../kernels/utils.cuh" +#include + +#include "api.cuh" #include "config.cuh" +#include "device_helpers.cuh" +#include "exception.cuh" namespace mscclpp { namespace ep { namespace low_latency_opt { -constexpr int kCombineNWarps = 32; -constexpr int kCombineNThreads = kCombineNWarps * WARP_SIZE; -constexpr int kCombineNStages = 8; -constexpr int kDirectSendMaxNWorkers = WARP_SIZE; -constexpr int kCombineMaxNBlocks = 128; -constexpr int kCombineMaxNTopk = 9; +constexpr int CombineNWarps = 32; +constexpr int CombineNThreads = CombineNWarps * WARP_SIZE; +constexpr int CombineNStages = 8; +constexpr int DirectSendMaxNWorkers = WARP_SIZE; +constexpr int CombineMaxNTopk = 9; MSCCLPP_HOST_DEVICE_INLINE size_t combineControlBytes(int nLocalExperts) { const size_t directControlBytes = static_cast(nLocalExperts + 1) * sizeof(int); return configAlign(directControlBytes > sizeof(RecvTask) ? directControlBytes : sizeof(RecvTask), 128); } -template +template MSCCLPP_HOST_DEVICE_INLINE size_t combineSharedBytes(int nLocalExperts) { - constexpr size_t kTileBytes = static_cast(kHidden) * sizeof(nv_bfloat16); - if constexpr (kMode == low_latency::OptimizedCombineMode::DIRECT_SEND) { - constexpr int kNWorkers = tmaWorkerCount(); - return combineControlBytes(nLocalExperts) + static_cast(kNWorkers) * (kTileBytes + sizeof(uint64_t)); + constexpr size_t TileBytes = static_cast(Hidden) * sizeof(__bfloat16); + if constexpr (Mode == low_latency::CombineMode::DIRECT_SEND) { + constexpr int NWorkers = tmaWorkerCount(); + return combineControlBytes(nLocalExperts) + static_cast(NWorkers) * (TileBytes + sizeof(uint64_t)); } - return combineControlBytes(nLocalExperts) + kCombineNStages * kTileBytes; + return combineControlBytes(nLocalExperts) + CombineNStages * TileBytes; } -template +template MSCCLPP_DEVICE_INLINE int4 reduceWeightedBf16x8(const void* expertOutput, int rowOffset, float weight, int nTopk, int hiddenIdx) { - constexpr int kBf16PairsPerInt4 = sizeof(int4) / sizeof(nv_bfloat162); - float2 reduced[kBf16PairsPerInt4] = {}; + constexpr int Bf16PairsPerInt4 = sizeof(int4) / sizeof(mscclpp::bf16x2); + float2 reduced[Bf16PairsPerInt4] = {}; for (int topkLane = 0; topkLane < nTopk; ++topkLane) { - const int sourceRowOffset = __shfl_sync(0xffffffff, rowOffset, topkLane); + const int sourceRowOffset = warpBroadcast(rowOffset, topkLane); if (sourceRowOffset < 0) continue; - const float sourceWeight = __shfl_sync(0xffffffff, weight, topkLane); + const float sourceWeight = warpBroadcast(weight, topkLane); const int4 packed = ld_nc_global(reinterpret_cast(expertOutput) + - static_cast(sourceRowOffset) * kHiddenInt4 + hiddenIdx); - const auto* values = reinterpret_cast(&packed); + static_cast(sourceRowOffset) * HiddenInt4 + hiddenIdx); + const auto* values = reinterpret_cast(&packed); #pragma unroll - for (int pairIdx = 0; pairIdx < kBf16PairsPerInt4; ++pairIdx) { - const float2 value = __bfloat1622float2(values[pairIdx]); - reduced[pairIdx].x = fmaf(value.x, sourceWeight, reduced[pairIdx].x); - reduced[pairIdx].y = fmaf(value.y, sourceWeight, reduced[pairIdx].y); + for (int pairIdx = 0; pairIdx < Bf16PairsPerInt4; ++pairIdx) { + const mscclpp::f32x2 value = mscclpp::to(values[pairIdx]); + reduced[pairIdx].x = fmaf(value.data[0], sourceWeight, reduced[pairIdx].x); + reduced[pairIdx].y = fmaf(value.data[1], sourceWeight, reduced[pairIdx].y); } } int4 packedOutput; - auto* outputValues = reinterpret_cast(&packedOutput); + auto* outputValues = reinterpret_cast(&packedOutput); #pragma unroll - for (int pairIdx = 0; pairIdx < kBf16PairsPerInt4; ++pairIdx) { - outputValues[pairIdx] = __float22bfloat162_rn(reduced[pairIdx]); + for (int pairIdx = 0; pairIdx < Bf16PairsPerInt4; ++pairIdx) { + outputValues[pairIdx] = mscclpp::to(mscclpp::f32x2(reduced[pairIdx])); } return packedOutput; } -template +template MSCCLPP_DEVICE_INLINE int4 reduceRankPartialsBf16x8(const void* combineRecvBuffer, int partialRankCandidate, int nTopk, int maxTokensPerRank, int tokenIdx, int hiddenIdx) { - constexpr int kBf16PairsPerInt4 = sizeof(int4) / sizeof(nv_bfloat162); - float2 reduced[kBf16PairsPerInt4] = {}; + constexpr int Bf16PairsPerInt4 = sizeof(int4) / sizeof(mscclpp::bf16x2); + float2 reduced[Bf16PairsPerInt4] = {}; for (int topkLane = 0; topkLane < nTopk; ++topkLane) { - const int partialRank = __shfl_sync(0xffffffff, partialRankCandidate, topkLane); + const int partialRank = warpBroadcast(partialRankCandidate, topkLane); if (partialRank < 0) continue; const int4 packed = ld_nc_global(reinterpret_cast(combineRecvBuffer) + - (static_cast(partialRank) * maxTokensPerRank + tokenIdx) * kHiddenInt4 + hiddenIdx); - const auto* values = reinterpret_cast(&packed); + (static_cast(partialRank) * maxTokensPerRank + tokenIdx) * HiddenInt4 + hiddenIdx); + const auto* values = reinterpret_cast(&packed); #pragma unroll - for (int pairIdx = 0; pairIdx < kBf16PairsPerInt4; ++pairIdx) { - const float2 value = __bfloat1622float2(values[pairIdx]); - reduced[pairIdx].x += value.x; - reduced[pairIdx].y += value.y; + for (int pairIdx = 0; pairIdx < Bf16PairsPerInt4; ++pairIdx) { + const mscclpp::f32x2 value = mscclpp::to(values[pairIdx]); + reduced[pairIdx].x += value.data[0]; + reduced[pairIdx].y += value.data[1]; } } int4 packedOutput; - auto* outputValues = reinterpret_cast(&packedOutput); + auto* outputValues = reinterpret_cast(&packedOutput); #pragma unroll - for (int pairIdx = 0; pairIdx < kBf16PairsPerInt4; ++pairIdx) { - outputValues[pairIdx] = __float22bfloat162_rn(reduced[pairIdx]); + for (int pairIdx = 0; pairIdx < Bf16PairsPerInt4; ++pairIdx) { + outputValues[pairIdx] = mscclpp::to(mscclpp::f32x2(reduced[pairIdx])); } return packedOutput; } -template -MSCCLPP_DEVICE_INLINE void sendRankLocalPartials(const void* expertOutput, int nExperts, int rank, int nRanks, - int nTopk, int maxTokensPerRank, void* combineRecvBuffer, - const void* dispatchRecvBuffer, void* rdmaBufferBase, - void* const* peerRecvBuffers, DispatchWorkspaceView& workspaceView, - uint8_t* sharedMemory) { +template