diff --git a/docs/source/deployment/mooncake-ep-jit.md b/docs/source/deployment/mooncake-ep-jit.md new file mode 100644 index 0000000000..1f4e0544a2 --- /dev/null +++ b/docs/source/deployment/mooncake-ep-jit.md @@ -0,0 +1,126 @@ +# JIT Kernel Compilation + +Mooncake EP JIT currently supports only the NVIDIA (CUDA) backend. + +Instead of shipping every kernel specialization in the library, Mooncake EP +compiles several of its device kernels at runtime. The first call that needs a +given kernel variant runs nvcc, caches the resulting cubin on disk, and reuses +it for the rest of the process and for later runs. + +As a result, a deployed Mooncake EP needs more than +`libmooncake_ep_device.so`. It also needs a working nvcc, the kernel headers, +the CUDA headers, a writable cache directory, and, for builds with +`USE_NCCL_DEVICE=ON`, the NCCL device headers. Each of these is found through an +environment variable, with a fallback compiled into the library. A build used +where it was built, or an installed wheel, usually needs no configuration. The +environment variables cover the other cases: relocated installs, containers, +read-only filesystems, and cross-architecture builds. + +## Path resolution + +Each path comes from the first rule in its list that gives a value. The +compiled-in defaults are set by CMake at build time and point into the build +tree. + +**Kernel headers** (the Mooncake EP and transfer engine headers that the JIT +compiles against): + +1. The `MOONCAKE_EP_JIT_SOURCE_DIR` environment variable. +2. `$MOONCAKE_EP_HOME/include/mooncake_ep`. +3. `mooncake_ep_jit/include/mooncake_ep` next to `libmooncake_ep_device.so`, if + that directory exists (the layout of an installed wheel). +4. The compiled-in `MOONCAKE_EP_JIT_SOURCE_DIR`. + +**NCCL device headers** (only for builds with `USE_NCCL_DEVICE=ON`): + +1. The `MOONCAKE_EP_JIT_BUILD_INCLUDE_DIR` environment variable. +2. `$NCCL_HOME/include`. +3. The compiled-in `MOONCAKE_EP_JIT_BUILD_INCLUDE_DIR`. + +**CUDA headers:** + +1. The `MOONCAKE_EP_JIT_CUDA_INCLUDE_DIR` environment variable. +2. `$CUDA_HOME/include`, or `$CUDA_PATH/include` if `CUDA_HOME` is not set. +3. The compiled-in `MOONCAKE_EP_JIT_CUDA_INCLUDE_DIR`. + +**Compiler:** + +1. The `MOONCAKE_EP_JIT_NVCC` environment variable. +2. The `NVCC` environment variable. +3. `$CUDA_HOME/bin/nvcc`, or `$CUDA_PATH/bin/nvcc` if `CUDA_HOME` is not set, + if that file exists. +4. `nvcc` from `PATH`. + +## Cache + +Compiled cubins are cached in `MOONCAKE_EP_JIT_CACHE_DIR`, which defaults to +`/tmp/mooncake_ep/jit`. + +The cache key combines two fingerprints: + +- **The kernel variant**: kernel family, variant name, entry point, generated + source, target SM and architecture, and runtime key. +- **The compile environment**: compiler identity, the contents of the kernel, + transfer engine and NCCL header trees, and the full list of compile options. + +This has two consequences: + +- **The cache invalidates itself.** Editing a kernel header, switching + compilers, changing architecture flags, or using different include + directories all change the key, so a new entry is compiled. A stale cubin is + never loaded, and there is no need to clear the cache after a rebuild. +- **Old entries are never removed.** Superseded cubins stay on disk, so on a + machine that rebuilds often the cache keeps growing. Deleting the directory is + always safe; the only cost is recompiling. + +Cache writes are atomic. Each cubin is first written to a temporary file named +with the process ID and thread ID, then renamed into place. Ranks on the same +node can therefore fill a shared cache at the same time without corrupting each +other's entries. + +Set `MOONCAKE_EP_JIT_CACHE_DIR` to another location when `/tmp` is read-only, +too small, or not shared across containers the way you expect. Putting the +cache on a persistent volume also avoids the first-call compile delay each time +a new container starts. + + +## Diagnostics + +Set `MOONCAKE_EP_JIT_LOG=1` to print JIT logs to stderr. Each line starts with +`[mooncake_ep jit] pid=`. The log shows cache hits and misses, compile +attempts and failures, and cache write errors (mkdir, open, write, and rename). +It is the quickest way to find a cache directory that is not writable or is +full. + +Each kernel variant has its own directory in the cache, containing +`kernel.cu` (the generated source), `kernel.cubin`, `metadata.json`, and, after +a compile, `compile.log` (the nvcc output). + +When a variant fails to compile, the failure is recorded in two places: + +- **In the process**: the variant is not compiled again, and its warning is + printed only once. +- **On disk**: a `compile.failed` file is written to the variant's cache + directory. Later processes that use the same cache skip the compile and + report the failure without running nvcc again. + +To retry after fixing the cause (for example, a missing header), delete that +variant's directory or the whole cache directory. Changes that alter the cache +key, such as editing a header or changing compile options, retry automatically +because they use a new directory. + +## Deployment checklist + +- **nvcc** can be found through `MOONCAKE_EP_JIT_NVCC`, `NVCC`, `CUDA_HOME`, + `CUDA_PATH`, or `PATH`. +- **Kernel headers** can be found. Installed wheels and builds used where they + were built need no setup. If the library was moved after it was built, set + `MOONCAKE_EP_HOME` or `MOONCAKE_EP_JIT_SOURCE_DIR`. +- **CUDA headers** can be found through `CUDA_HOME`, `CUDA_PATH`, or + `MOONCAKE_EP_JIT_CUDA_INCLUDE_DIR`. +- **NCCL device headers** (only for `USE_NCCL_DEVICE=ON` builds) can be found + through `NCCL_HOME` or `MOONCAKE_EP_JIT_BUILD_INCLUDE_DIR`. +- **The cache directory** is writable. Set `MOONCAKE_EP_JIT_CACHE_DIR` if `/tmp` + is not suitable. +- **The first call** to each kernel variant includes compile time. Later calls, + and later runs that share the cache, do not. diff --git a/docs/source/design/mooncake-ep.md b/docs/source/design/mooncake-ep.md index e74caae816..a70deeb7a9 100644 --- a/docs/source/design/mooncake-ep.md +++ b/docs/source/design/mooncake-ep.md @@ -143,6 +143,24 @@ It provides the P2P / IPC fast path only: - Kernels run on the caller's `torch.xpu` stream. Zero-copy combine and the elastic buffer are not supported. +## JIT kernel compilation + +On NVIDIA GPUs, Mooncake EP compiles its dispatch and combine kernels, including +the elastic buffer kernels, at runtime. Parameters such as `hidden`, +`num_experts`, and `num_topk` are compile-time template arguments of these +kernels. Without JIT, the library can only contain a fixed set of +instantiations (for example, `hidden` in {2048, 2560, 3072, 3584, 4096, 5120, +6144, 7168, 8192}), and any other value fails with an "Unsupported" error. With +JIT, the first call with a new combination of values runs nvcc to build that +instantiation and caches the cubin on disk; later calls and later runs reuse it. + +JIT is currently not supported on other backends. + +At runtime this needs nvcc, the kernel headers, and a writable cache directory. +Installed wheels ship the kernel headers, so they usually need no setup. See +[JIT kernel compilation](../deployment/mooncake-ep-jit.md) for path lookup, +cache behavior, and environment variables. + ## Metadata exchange During `Buffer.connect()`, ranks exchange: @@ -279,3 +297,4 @@ Adapt launch commands to the target environment and number of GPUs. - [Mooncake PG design](mooncake-backend-pg.md) - [Python API reference](../api-reference/python/ep-backend.md) - [PG/EP troubleshooting](../troubleshooting/pg-ep-troubleshooting.md) +- [Mooncake EP JIT kernel compilation](../deployment/mooncake-ep-jit.md) diff --git a/docs/source/index.md b/docs/source/index.md index bcec28eb2f..4dde051c98 100644 --- a/docs/source/index.md +++ b/docs/source/index.md @@ -167,6 +167,7 @@ deployment/store-rs/index deployment/kubernetes-deployment-guide/index deployment/integrations/sglang/index deployment/integrations/vllm/index +deployment/mooncake-ep-jit Mooncake x LMCache Integration Mooncake x LMDeploy Integration ::: diff --git a/mooncake-ep/include/elastic/mooncake_ep_elastic_combine_official.cuh b/mooncake-ep/include/elastic/mooncake_ep_elastic_combine_official.cuh index 44c605fd59..fbb5467622 100644 --- a/mooncake-ep/include/elastic/mooncake_ep_elastic_combine_official.cuh +++ b/mooncake-ep/include/elastic/mooncake_ep_elastic_combine_official.cuh @@ -3,6 +3,8 @@ // transport references are replaced with Mooncake Device API adapters. #pragma once +#include + #include #include @@ -14,6 +16,19 @@ namespace mooncake::elastic { +template +struct CombineKernelArgs { + nv_bfloat16* x; + float* topk_weights; + int* src_metadata; + int* psum_num_recv_tokens_per_scaleup_rank; + typename Ops::Context comm_ctx; + void* buffer; + void* workspace; + int rank_idx; + int num_reduced_tokens; +}; + template (), typename team_t = std::conditional_t< kIsScaleupNVLink, transport::ScaleupTeam, transport::WorldTeam>> -__global__ void __launch_bounds__(kNumThreads, 1) - combine_impl(nv_bfloat16* x, float* topk_weights, int* src_metadata, - int* psum_num_recv_tokens_per_scaleup_rank, - const typename Ops::Context comm_ctx, void* buffer, - void* workspace, const int rank_idx, int num_reduced_tokens) { +__device__ __forceinline__ void combine_kernel_impl( + const CombineKernelArgs& args) { + static_assert(std::is_trivially_copyable>::value, + "CombineKernelArgs must be trivially copyable"); + + nv_bfloat16* x = args.x; + float* topk_weights = args.topk_weights; + int* src_metadata = args.src_metadata; + int* psum_num_recv_tokens_per_scaleup_rank = + args.psum_num_recv_tokens_per_scaleup_rank; + const typename Ops::Context& comm_ctx = args.comm_ctx; + void* buffer = args.buffer; + void* workspace = args.workspace; + const int rank_idx = args.rank_idx; + int num_reduced_tokens = args.num_reduced_tokens; + // Utils const auto sm_idx = static_cast(blockIdx.x); const auto thread_idx = static_cast(threadIdx.x); @@ -364,4 +390,17 @@ __global__ void __launch_bounds__(kNumThreads, 1) thread_idx); } +template +__global__ void __launch_bounds__(kNumThreads, 1) + combine(const CombineKernelArgs args) { + combine_kernel_impl(args); +} + } // namespace mooncake::elastic diff --git a/mooncake-ep/include/elastic/mooncake_ep_elastic_combine_reduce_epilogue.cuh b/mooncake-ep/include/elastic/mooncake_ep_elastic_combine_reduce_epilogue.cuh index 7829e7ecfc..60aac7bcc7 100644 --- a/mooncake-ep/include/elastic/mooncake_ep_elastic_combine_reduce_epilogue.cuh +++ b/mooncake-ep/include/elastic/mooncake_ep_elastic_combine_reduce_epilogue.cuh @@ -3,6 +3,8 @@ // transport references are replaced with Mooncake Device API adapters. #pragma once +#include + #include #include #include @@ -11,6 +13,20 @@ namespace mooncake::elastic { +struct CombineReduceEpilogueKernelArgs { + nv_bfloat16* combined_x; + float* combined_topk_weights; + topk_idx_t* combined_topk_idx; + void* recv_buffer; + void* bias_0; + void* bias_1; + int num_combined_tokens; +}; + +static_assert( + std::is_trivially_copyable::value, + "CombineReduceEpilogueKernelArgs must be trivially copyable"); + template (), int kNumTokensInLayout = get_num_tokens_in_layout< kAllowMultipleReduction, kNumRanks, kNumTopk>()> -__global__ void __launch_bounds__(kNumThreads, 1) - combine_reduce_epilogue_impl(nv_bfloat16* combined_x, - float* combined_topk_weights, - topk_idx_t* combined_topk_idx, - void* recv_buffer, void* bias_0, void* bias_1, - const int num_combined_tokens, - const int scaleout_rank_idx, - const int scaleup_rank_idx) { +__device__ __forceinline__ void combine_reduce_epilogue_kernel_impl( + const CombineReduceEpilogueKernelArgs& args) { + nv_bfloat16* combined_x = args.combined_x; + float* combined_topk_weights = args.combined_topk_weights; + topk_idx_t* combined_topk_idx = args.combined_topk_idx; + void* recv_buffer = args.recv_buffer; + void* bias_0 = args.bias_0; + void* bias_1 = args.bias_1; + const int num_combined_tokens = args.num_combined_tokens; + constexpr int kNumExpertsPerScaleout = kNumExperts / kNumScaleoutRanks; constexpr int kNumExpertsPerRank = kNumExperts / (kNumScaleupRanks * kNumScaleoutRanks); @@ -210,4 +228,16 @@ __global__ void __launch_bounds__(kNumThreads, 1) } } +template +__global__ void __launch_bounds__(kNumThreads, 1) + combine_reduce_epilogue(const CombineReduceEpilogueKernelArgs args) { + combine_reduce_epilogue_kernel_impl< + kUseExpandedLayout, kAllowMultipleReduction, kNumSMs, kNumWarps, + kNumScaleoutRanks, kNumScaleupRanks, kHidden, kNumMaxTokensPerRank, + kNumExperts, kNumTopk>(args); +} + } // namespace mooncake::elastic diff --git a/mooncake-ep/include/elastic/mooncake_ep_elastic_compiled.cuh b/mooncake-ep/include/elastic/mooncake_ep_elastic_compiled.cuh index e5b277f8f6..9bb54dd40e 100644 --- a/mooncake-ep/include/elastic/mooncake_ep_elastic_compiled.cuh +++ b/mooncake-ep/include/elastic/mooncake_ep_elastic_compiled.cuh @@ -99,6 +99,11 @@ struct int_with_bits<64> { }; using topk_idx_t = int_with_bits::type; +// The host launchers pass int64_t* top-k arrays, so any other width (e.g. set +// only for the JIT compile via NVCC_EXTRA_FLAGS) would be misread silently. +static_assert(sizeof(topk_idx_t) == sizeof(int64_t), + "EP_NUM_TOPK_IDX_BITS must be 64: the host launchers pass " + "int64_t* top-k arrays"); union sf_pack_t { float fp32; diff --git a/mooncake-ep/include/elastic/mooncake_ep_elastic_dispatch_copy_epilogue.cuh b/mooncake-ep/include/elastic/mooncake_ep_elastic_dispatch_copy_epilogue.cuh index 8c9955d90c..aea7605d98 100644 --- a/mooncake-ep/include/elastic/mooncake_ep_elastic_dispatch_copy_epilogue.cuh +++ b/mooncake-ep/include/elastic/mooncake_ep_elastic_dispatch_copy_epilogue.cuh @@ -3,6 +3,8 @@ // transport references are replaced with Mooncake Device API adapters. #pragma once +#include + #include #include #include @@ -10,6 +12,27 @@ namespace mooncake::elastic { +struct DispatchCopyEpilogueKernelArgs { + void* buffer; + void* workspace; + int* psum_num_recv_tokens_per_scaleup_rank; + int* psum_num_recv_tokens_per_expert; + void* recv_x; + sf_pack_t* recv_sf; + topk_idx_t* recv_topk_idx; + float* recv_topk_weights; + int* recv_src_metadata; + int* channel_linked_list; + int num_recv_tokens; + int recv_sf_token_stride; + int recv_sf_hidden_stride; + int scaleout_rank_idx; + int scaleup_rank_idx; +}; + +static_assert(std::is_trivially_copyable::value, + "DispatchCopyEpilogueKernelArgs must be trivially copyable"); + template < bool kDoExpand, bool kCachedMode, // NOTES: this channel concept only applies for scale-out ranks @@ -21,13 +44,25 @@ template < int kNumMaxTokensPerChannel = math::constexpr_ceil_div(kNumMaxTokensPerRank, kNumChannels), bool kDoCreateLinkedList = (kNumScaleoutRanks > 1 and not kCachedMode)> -__global__ void __launch_bounds__(kNumThreads, 1) dispatch_copy_epilogue_impl( - void* buffer, void* workspace, int* psum_num_recv_tokens_per_scaleup_rank, - int* psum_num_recv_tokens_per_expert, void* recv_x, sf_pack_t* recv_sf, - topk_idx_t* recv_topk_idx, float* recv_topk_weights, int* recv_src_metadata, - int* channel_linked_list, int num_recv_tokens, - const int recv_sf_token_stride, const int recv_sf_hidden_stride, - const int scaleout_rank_idx, const int scaleup_rank_idx) { +__device__ __forceinline__ void dispatch_copy_epilogue_kernel_impl( + const DispatchCopyEpilogueKernelArgs& args) { + void* buffer = args.buffer; + void* workspace = args.workspace; + int* psum_num_recv_tokens_per_scaleup_rank = + args.psum_num_recv_tokens_per_scaleup_rank; + int* psum_num_recv_tokens_per_expert = args.psum_num_recv_tokens_per_expert; + void* recv_x = args.recv_x; + sf_pack_t* recv_sf = args.recv_sf; + topk_idx_t* recv_topk_idx = args.recv_topk_idx; + float* recv_topk_weights = args.recv_topk_weights; + int* recv_src_metadata = args.recv_src_metadata; + int* channel_linked_list = args.channel_linked_list; + int num_recv_tokens = args.num_recv_tokens; + const int recv_sf_token_stride = args.recv_sf_token_stride; + const int recv_sf_hidden_stride = args.recv_sf_hidden_stride; + const int scaleout_rank_idx = args.scaleout_rank_idx; + const int scaleup_rank_idx = args.scaleup_rank_idx; + // Utils const auto sm_idx = static_cast(blockIdx.x), thread_idx = static_cast(threadIdx.x); @@ -275,4 +310,16 @@ __global__ void __launch_bounds__(kNumThreads, 1) dispatch_copy_epilogue_impl( } } +template +__global__ void __launch_bounds__(kNumThreads, 1) + dispatch_copy_epilogue(const DispatchCopyEpilogueKernelArgs args) { + dispatch_copy_epilogue_kernel_impl< + kDoExpand, kCachedMode, kNumSMs, kNumChannels, kNumWarps, + kNumScaleoutRanks, kNumScaleupRanks, kNumHiddenBytes, kNumSFPacks, + kNumMaxTokensPerRank, kNumExperts, kNumTopk>(args); +} + } // namespace mooncake::elastic diff --git a/mooncake-ep/include/elastic/mooncake_ep_elastic_dispatch_deterministic_prologue.cuh b/mooncake-ep/include/elastic/mooncake_ep_elastic_dispatch_deterministic_prologue.cuh index 323f81d5fb..0187c16186 100644 --- a/mooncake-ep/include/elastic/mooncake_ep_elastic_dispatch_deterministic_prologue.cuh +++ b/mooncake-ep/include/elastic/mooncake_ep_elastic_dispatch_deterministic_prologue.cuh @@ -5,23 +5,39 @@ #include +#include + #include #include #include namespace mooncake::elastic { +struct DispatchDeterministicPrologueKernelArgs { + topk_idx_t* topk_idx; + int* rank_count_buffer; + int* dst_buffer_slot_idx; + int num_tokens; + int scaleup_rank_idx; +}; + +static_assert( + std::is_trivially_copyable::value, + "DispatchDeterministicPrologueKernelArgs must be trivially copyable"); + // Slot preassignment runs in the active scale-up domain. Hybrid scale-out // forwarding is handled by the hybrid dispatch kernel. template -__global__ void __launch_bounds__(kNumThreads, 1) - dispatch_deterministic_prologue_impl(topk_idx_t* topk_idx, - int* rank_count_buffer, - int* dst_buffer_slot_idx, - const int num_tokens, - const int scaleup_rank_idx) { +__device__ __forceinline__ void dispatch_deterministic_prologue_kernel_impl( + const DispatchDeterministicPrologueKernelArgs& args) { + topk_idx_t* topk_idx = args.topk_idx; + int* rank_count_buffer = args.rank_count_buffer; + int* dst_buffer_slot_idx = args.dst_buffer_slot_idx; + int num_tokens = args.num_tokens; + int scaleup_rank_idx = args.scaleup_rank_idx; + constexpr int kNumExpertsPerRank = kNumExperts / kNumScaleupRanks; EP_STATIC_ASSERT(kNumExperts % kNumScaleupRanks == 0, "Invalid number of experts or ranks"); @@ -168,4 +184,15 @@ __global__ void __launch_bounds__(kNumThreads, 1) } } +template +__global__ void __launch_bounds__(kNumThreads, 1) + dispatch_deterministic_prologue( + const DispatchDeterministicPrologueKernelArgs args) { + dispatch_deterministic_prologue_kernel_impl< + kNumSMs, kNumWarps, kNumScaleupRanks, kNumMaxTokensPerRank, kNumExperts, + kNumTopk>(args); +} + } // namespace mooncake::elastic diff --git a/mooncake-ep/include/elastic/mooncake_ep_elastic_dispatch_official.cuh b/mooncake-ep/include/elastic/mooncake_ep_elastic_dispatch_official.cuh index 98be7026c3..8757494b91 100644 --- a/mooncake-ep/include/elastic/mooncake_ep_elastic_dispatch_official.cuh +++ b/mooncake-ep/include/elastic/mooncake_ep_elastic_dispatch_official.cuh @@ -3,6 +3,8 @@ // transport references are replaced with Mooncake Device API adapters. #pragma once +#include + #include #include @@ -15,6 +17,27 @@ namespace mooncake::elastic { +template +struct DispatchKernelArgs { + void* x; + sf_pack_t* sf; + topk_idx_t* topk_idx; + float* topk_weights; + topk_idx_t* copied_topk_idx; + int* cumulative_local_expert_recv_stats; + int* psum_num_recv_tokens_per_scaleup_rank; + int* psum_num_recv_tokens_per_expert; + int* dst_buffer_slot_idx; + int num_tokens; + int sf_token_stride; + int sf_hidden_stride; + typename Ops::Context comm_ctx; + void* buffer; + void* workspace; + void* mapped_host_workspace; + int rank_idx; +}; + template > -__global__ void __launch_bounds__(kNumThreads, 1) - dispatch_impl(void* x, sf_pack_t* sf, topk_idx_t* topk_idx, - float* topk_weights, topk_idx_t* copied_topk_idx, - int* cumulative_local_expert_recv_stats, - int* psum_num_recv_tokens_per_scaleup_rank, - int* psum_num_recv_tokens_per_expert, - int* dst_buffer_slot_idx, const int num_tokens, - const int sf_token_stride, const int sf_hidden_stride, - const typename Ops::Context comm_ctx, void* buffer, - void* workspace, void* mapped_host_workspace, - const int rank_idx) { +__device__ __forceinline__ void dispatch_kernel_impl( + const DispatchKernelArgs& args) { + static_assert(std::is_trivially_copyable>::value, + "DispatchKernelArgs must be trivially copyable"); + + void* x = args.x; + sf_pack_t* sf = args.sf; + topk_idx_t* topk_idx = args.topk_idx; + float* topk_weights = args.topk_weights; + topk_idx_t* copied_topk_idx = args.copied_topk_idx; + int* cumulative_local_expert_recv_stats = + args.cumulative_local_expert_recv_stats; + int* psum_num_recv_tokens_per_scaleup_rank = + args.psum_num_recv_tokens_per_scaleup_rank; + int* psum_num_recv_tokens_per_expert = args.psum_num_recv_tokens_per_expert; + int* dst_buffer_slot_idx = args.dst_buffer_slot_idx; + const int num_tokens = args.num_tokens; + const int sf_token_stride = args.sf_token_stride; + const int sf_hidden_stride = args.sf_hidden_stride; + const typename Ops::Context& comm_ctx = args.comm_ctx; + void* buffer = args.buffer; + void* workspace = args.workspace; + void* mapped_host_workspace = args.mapped_host_workspace; + const int rank_idx = args.rank_idx; + constexpr int kNumExpertsPerRank = kNumExperts / kNumRanks; EP_STATIC_ASSERT(kNumExperts % kNumRanks == 0, "Invalid number of experts or ranks"); @@ -512,4 +549,20 @@ __global__ void __launch_bounds__(kNumThreads, 1) workspace_layout.get_scaleup_atomic_sender_counter()[thread_idx] = 0; } +template +__global__ void __launch_bounds__(kNumThreads, 1) + dispatch(const DispatchKernelArgs args) { + dispatch_kernel_impl(args); +} + } // namespace mooncake::elastic diff --git a/mooncake-ep/include/elastic/mooncake_ep_elastic_hybrid_combine_official.cuh b/mooncake-ep/include/elastic/mooncake_ep_elastic_hybrid_combine_official.cuh index 61c3a1bf40..20b6116ab2 100644 --- a/mooncake-ep/include/elastic/mooncake_ep_elastic_hybrid_combine_official.cuh +++ b/mooncake-ep/include/elastic/mooncake_ep_elastic_hybrid_combine_official.cuh @@ -3,6 +3,8 @@ // transport references are replaced with Mooncake Device API adapters. #pragma once +#include + #include #include @@ -13,6 +15,22 @@ namespace mooncake::elastic { +template +struct HybridCombineKernelArgs { + nv_bfloat16* x; + float* topk_weights; + int* src_metadata; + int* psum_num_recv_tokens_per_scaleup_rank; + int* token_metadata_at_forward; + int* channel_linked_list; + typename Ops::Context comm_ctx; + void* buffer; + void* workspace; + int scaleout_rank_idx; + int scaleup_rank_idx; + int num_reduced_tokens; +}; + template < typename Ops, bool kUseExpandedLayout, bool kAllowMultipleReduction, int kNumSMs, int kNumScaleupWarps, int kNumForwardWarps, @@ -37,14 +55,26 @@ template < kAllowMultipleReduction, kNumScaleoutRanks, kNumTopk>(), int kNumTokensInScaleupLayout = get_num_tokens_in_layout< kAllowMultipleReduction, kNumScaleupRanks, kNumTopk>()> -__global__ void __launch_bounds__(kNumThreads, 1) - hybrid_combine_impl(nv_bfloat16* x, float* topk_weights, int* src_metadata, - int* psum_num_recv_tokens_per_scaleup_rank, - int* token_metadata_at_forward, - int* channel_linked_list, - const typename Ops::Context comm_ctx, void* buffer, - void* workspace, const int scaleout_rank_idx, - const int scaleup_rank_idx, int num_reduced_tokens) { +__device__ __forceinline__ void hybrid_combine_kernel_impl( + const HybridCombineKernelArgs& args) { + static_assert( + std::is_trivially_copyable>::value, + "HybridCombineKernelArgs must be trivially copyable"); + + nv_bfloat16* x = args.x; + float* topk_weights = args.topk_weights; + int* src_metadata = args.src_metadata; + int* psum_num_recv_tokens_per_scaleup_rank = + args.psum_num_recv_tokens_per_scaleup_rank; + int* token_metadata_at_forward = args.token_metadata_at_forward; + int* channel_linked_list = args.channel_linked_list; + const typename Ops::Context& comm_ctx = args.comm_ctx; + void* buffer = args.buffer; + void* workspace = args.workspace; + const int scaleout_rank_idx = args.scaleout_rank_idx; + const int scaleup_rank_idx = args.scaleup_rank_idx; + int num_reduced_tokens = args.num_reduced_tokens; + // Utils const auto sm_idx = static_cast(blockIdx.x); const auto thread_idx = static_cast(threadIdx.x); @@ -821,4 +851,19 @@ __global__ void __launch_bounds__(kNumThreads, 1) } } +template +__global__ void __launch_bounds__(kNumThreads, 1) + hybrid_combine(const HybridCombineKernelArgs args) { + hybrid_combine_kernel_impl(args); +} + } // namespace mooncake::elastic diff --git a/mooncake-ep/include/elastic/mooncake_ep_elastic_hybrid_dispatch_official.cuh b/mooncake-ep/include/elastic/mooncake_ep_elastic_hybrid_dispatch_official.cuh index 99f83f7e01..756b729655 100644 --- a/mooncake-ep/include/elastic/mooncake_ep_elastic_hybrid_dispatch_official.cuh +++ b/mooncake-ep/include/elastic/mooncake_ep_elastic_hybrid_dispatch_official.cuh @@ -3,6 +3,8 @@ // transport references are replaced with Mooncake Device API adapters. #pragma once +#include + #include #include @@ -14,6 +16,30 @@ namespace mooncake::elastic { +template +struct HybridDispatchKernelArgs { + void* x; + sf_pack_t* sf; + topk_idx_t* topk_idx; + float* topk_weights; + topk_idx_t* copied_topk_idx; + int* cumulative_local_expert_recv_stats; + int* psum_num_recv_tokens_per_scaleup_rank; + int* psum_num_recv_tokens_per_expert; + int* dst_buffer_slot_idx; + int* token_metadata_at_forward; + int num_tokens; + int sf_token_stride; + int sf_hidden_stride; + // TODO(NCCL): so many params, plans to optimize? + typename Ops::Context comm_ctx; + void* buffer; + void* workspace; + void* mapped_host_workspace; + int scaleout_rank_idx; + int scaleup_rank_idx; +}; + template < typename Ops, bool kDoCPUSync, bool kReuseSlotIndices, int kNumSMs, int kNumNotifyWarps, int kNumScaleoutWarps, int kNumForwardWarps, @@ -34,20 +60,34 @@ template < int kNumForwardThreads = kNumForwardWarps * 32, int kNumThreads = kNumNotifyThreads + kNumScaleoutSendThreads + kNumForwardThreads> -__global__ void __launch_bounds__(kNumThreads, 1) - hybrid_dispatch_impl(void* x, sf_pack_t* sf, topk_idx_t* topk_idx, - float* topk_weights, topk_idx_t* copied_topk_idx, - int* cumulative_local_expert_recv_stats, - int* psum_num_recv_tokens_per_scaleup_rank, - int* psum_num_recv_tokens_per_expert, - int* dst_buffer_slot_idx, - int* token_metadata_at_forward, const int num_tokens, - const int sf_token_stride, const int sf_hidden_stride, - // TODO(NCCL): so many params, plans to optimize? - const typename Ops::Context comm_ctx, void* buffer, - void* workspace, void* mapped_host_workspace, - const int scaleout_rank_idx, - const int scaleup_rank_idx) { +__device__ __forceinline__ void hybrid_dispatch_kernel_impl( + const HybridDispatchKernelArgs& args) { + static_assert( + std::is_trivially_copyable>::value, + "HybridDispatchKernelArgs must be trivially copyable"); + + void* x = args.x; + sf_pack_t* sf = args.sf; + topk_idx_t* topk_idx = args.topk_idx; + float* topk_weights = args.topk_weights; + topk_idx_t* copied_topk_idx = args.copied_topk_idx; + int* cumulative_local_expert_recv_stats = + args.cumulative_local_expert_recv_stats; + int* psum_num_recv_tokens_per_scaleup_rank = + args.psum_num_recv_tokens_per_scaleup_rank; + int* psum_num_recv_tokens_per_expert = args.psum_num_recv_tokens_per_expert; + int* dst_buffer_slot_idx = args.dst_buffer_slot_idx; + int* token_metadata_at_forward = args.token_metadata_at_forward; + const int num_tokens = args.num_tokens; + const int sf_token_stride = args.sf_token_stride; + const int sf_hidden_stride = args.sf_hidden_stride; + const typename Ops::Context& comm_ctx = args.comm_ctx; + void* buffer = args.buffer; + void* workspace = args.workspace; + void* mapped_host_workspace = args.mapped_host_workspace; + const int scaleout_rank_idx = args.scaleout_rank_idx; + const int scaleup_rank_idx = args.scaleup_rank_idx; + constexpr int kNumExpertsPerRank = kNumExperts / kNumRanks; constexpr int kNumExpertsPerScaleout = kNumExperts / kNumScaleoutRanks; EP_STATIC_ASSERT(kNumExperts % kNumScaleupRanks == 0, @@ -895,4 +935,22 @@ __global__ void __launch_bounds__(kNumThreads, 1) workspace_layout.get_scaleup_atomic_sender_counter()[thread_idx] = 0; } +template +__global__ void __launch_bounds__(kNumThreads, 1) + hybrid_dispatch(const HybridDispatchKernelArgs args) { + hybrid_dispatch_kernel_impl< + Ops, kDoCPUSync, kReuseSlotIndices, kNumSMs, kNumNotifyWarps, + kNumScaleoutWarps, kNumForwardWarps, kNumScaleoutRanks, + kNumScaleupRanks, kNumHiddenBytes, kNumSFPacks, kNumMaxTokensPerRank, + kNumExperts, kNumTopk, kExpertAlignment, kNumQPs, kNumTimeoutCycles>( + args); +} + } // namespace mooncake::elastic diff --git a/mooncake-ep/include/jit/combine_jit.h b/mooncake-ep/include/jit/combine_jit.h new file mode 100644 index 0000000000..f1d7f35153 --- /dev/null +++ b/mooncake-ep/include/jit/combine_jit.h @@ -0,0 +1,16 @@ +#pragma once + +#include + +namespace mooncake { + +struct CombineKernelArgs; + +namespace jit { + +void launch_combine_jit(int hidden, int num_max_topk, int num_warp_groups, + int num_warps_per_group, int num_sms, + const CombineKernelArgs& args, cudaStream_t stream); + +} // namespace jit +} // namespace mooncake diff --git a/mooncake-ep/include/jit/dispatch_jit.h b/mooncake-ep/include/jit/dispatch_jit.h new file mode 100644 index 0000000000..853ace102f --- /dev/null +++ b/mooncake-ep/include/jit/dispatch_jit.h @@ -0,0 +1,16 @@ +#pragma once + +#include + +namespace mooncake { + +struct DispatchKernelArgs; + +namespace jit { + +void launch_dispatch_jit(int hidden, bool use_fp8, int num_warp_groups, + int num_warps_per_group, int num_sms, + const DispatchKernelArgs& args, cudaStream_t stream); + +} // namespace jit +} // namespace mooncake \ No newline at end of file diff --git a/mooncake-ep/include/jit/elastic_combine_jit.h b/mooncake-ep/include/jit/elastic_combine_jit.h new file mode 100644 index 0000000000..6474045dbc --- /dev/null +++ b/mooncake-ep/include/jit/elastic_combine_jit.h @@ -0,0 +1,21 @@ +#pragma once + +#include + +#include + +namespace mooncake { + +enum class ElasticTransportBackend : uint8_t; + +namespace jit { + +// args points to the elastic::CombineKernelArgs of the given backend. +void launch_elastic_combine_jit(ElasticTransportBackend backend, int num_warps, + int num_sms, int num_scaleup_ranks, int hidden, + int num_max_tokens_per_rank, int num_experts, + int num_topk, int smem_bytes, const void* args, + cudaStream_t stream); + +} // namespace jit +} // namespace mooncake diff --git a/mooncake-ep/include/jit/elastic_combine_reduce_epilogue_jit.h b/mooncake-ep/include/jit/elastic_combine_reduce_epilogue_jit.h new file mode 100644 index 0000000000..672190725b --- /dev/null +++ b/mooncake-ep/include/jit/elastic_combine_reduce_epilogue_jit.h @@ -0,0 +1,20 @@ +#pragma once + +#include + +namespace mooncake { + +namespace elastic { +struct CombineReduceEpilogueKernelArgs; +} // namespace elastic + +namespace jit { + +void launch_elastic_combine_reduce_epilogue_jit( + int num_warps, int num_sms, int num_scaleout_ranks, int num_scaleup_ranks, + int hidden, int num_max_tokens_per_rank, int num_experts, int num_topk, + int smem_bytes, const elastic::CombineReduceEpilogueKernelArgs& args, + cudaStream_t stream); + +} // namespace jit +} // namespace mooncake diff --git a/mooncake-ep/include/jit/elastic_dispatch_copy_epilogue_jit.h b/mooncake-ep/include/jit/elastic_dispatch_copy_epilogue_jit.h new file mode 100644 index 0000000000..a42e51161f --- /dev/null +++ b/mooncake-ep/include/jit/elastic_dispatch_copy_epilogue_jit.h @@ -0,0 +1,21 @@ +#pragma once + +#include + +namespace mooncake { + +namespace elastic { +struct DispatchCopyEpilogueKernelArgs; +} // namespace elastic + +namespace jit { + +void launch_elastic_dispatch_copy_epilogue_jit( + bool do_expand, bool cached_mode, int num_channels, int num_warps, + int num_sms, int num_scaleout_ranks, int num_scaleup_ranks, + int num_hidden_bytes, int num_sf_packs, int num_max_tokens_per_rank, + int num_experts, int num_topk, int smem_bytes, + const elastic::DispatchCopyEpilogueKernelArgs& args, cudaStream_t stream); + +} // namespace jit +} // namespace mooncake diff --git a/mooncake-ep/include/jit/elastic_dispatch_jit.h b/mooncake-ep/include/jit/elastic_dispatch_jit.h new file mode 100644 index 0000000000..1774b31c64 --- /dev/null +++ b/mooncake-ep/include/jit/elastic_dispatch_jit.h @@ -0,0 +1,23 @@ +#pragma once + +#include + +#include + +namespace mooncake { + +enum class ElasticTransportBackend : uint8_t; + +namespace jit { + +// args points to the elastic::DispatchKernelArgs of the given backend. +void launch_elastic_dispatch_jit(ElasticTransportBackend backend, + bool reuse_slot_indices, int num_notify_warps, + int num_dispatch_warps, int num_sms, + int num_scaleup_ranks, int num_hidden_bytes, + int num_sf_packs, int num_max_tokens_per_rank, + int num_experts, int num_topk, int smem_bytes, + const void* args, cudaStream_t stream); + +} // namespace jit +} // namespace mooncake diff --git a/mooncake-ep/include/jit/elastic_dispatch_prologue_jit.h b/mooncake-ep/include/jit/elastic_dispatch_prologue_jit.h new file mode 100644 index 0000000000..9298ab0d40 --- /dev/null +++ b/mooncake-ep/include/jit/elastic_dispatch_prologue_jit.h @@ -0,0 +1,20 @@ +#pragma once + +#include + +namespace mooncake { + +namespace elastic { +struct DispatchDeterministicPrologueKernelArgs; +} // namespace elastic + +namespace jit { + +void launch_elastic_dispatch_prologue_jit( + int num_warps, int num_sms, int num_scaleup_ranks, + int num_max_tokens_per_rank, int num_experts, int num_topk, int smem_bytes, + const elastic::DispatchDeterministicPrologueKernelArgs& args, + cudaStream_t stream); + +} // namespace jit +} // namespace mooncake diff --git a/mooncake-ep/include/jit/elastic_hybrid_combine_jit.h b/mooncake-ep/include/jit/elastic_hybrid_combine_jit.h new file mode 100644 index 0000000000..37e5957bfb --- /dev/null +++ b/mooncake-ep/include/jit/elastic_hybrid_combine_jit.h @@ -0,0 +1,23 @@ +#pragma once + +#include + +#include + +namespace mooncake { + +enum class ElasticTransportBackend : uint8_t; + +namespace jit { + +// args points to the elastic::HybridCombineKernelArgs of the given +// backend. +void launch_elastic_hybrid_combine_jit( + ElasticTransportBackend backend, int num_scaleup_warps, + int num_forward_warps, int num_sms, int num_scaleout_ranks, + int num_scaleup_ranks, int hidden, int num_max_tokens_per_rank, + int num_experts, int num_topk, int smem_bytes, const void* args, + cudaStream_t stream); + +} // namespace jit +} // namespace mooncake diff --git a/mooncake-ep/include/jit/elastic_hybrid_dispatch_jit.h b/mooncake-ep/include/jit/elastic_hybrid_dispatch_jit.h new file mode 100644 index 0000000000..12aa029e5a --- /dev/null +++ b/mooncake-ep/include/jit/elastic_hybrid_dispatch_jit.h @@ -0,0 +1,24 @@ +#pragma once + +#include + +#include + +namespace mooncake { + +enum class ElasticTransportBackend : uint8_t; + +namespace jit { + +// args points to the elastic::HybridDispatchKernelArgs of the given +// backend. +void launch_elastic_hybrid_dispatch_jit( + ElasticTransportBackend backend, bool reuse_slot_indices, + int num_notify_warps, int num_scaleout_warps, int num_forward_warps, + int num_sms, int num_scaleout_ranks, int num_scaleup_ranks, + int num_hidden_bytes, int num_sf_packs, int num_max_tokens_per_rank, + int num_experts, int num_topk, int smem_bytes, const void* args, + cudaStream_t stream); + +} // namespace jit +} // namespace mooncake diff --git a/mooncake-ep/include/jit/jit_cache.hpp b/mooncake-ep/include/jit/jit_cache.hpp new file mode 100644 index 0000000000..79b848a947 --- /dev/null +++ b/mooncake-ep/include/jit/jit_cache.hpp @@ -0,0 +1,37 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Ported from nccl-extensions: nccl_ep/device/jit/jit_cache.hpp. + */ + +#pragma once + +#include +#include +#include +#include + +namespace mooncake { +namespace jit { + +class JitCache { + public: + static JitCache& instance(); + + std::filesystem::path root_dir() const; + bool has_failed(const std::string& key); + void mark_failed(const std::string& key); + void warn_once(const std::string& key, const std::string& message); + + private: + JitCache() = default; + + std::mutex mutex_; + std::unordered_set failed_keys_; + std::unordered_set warned_keys_; +}; + +} // namespace jit +} // namespace mooncake diff --git a/mooncake-ep/include/jit/jit_compiler.hpp b/mooncake-ep/include/jit/jit_compiler.hpp new file mode 100644 index 0000000000..887272dc03 --- /dev/null +++ b/mooncake-ep/include/jit/jit_compiler.hpp @@ -0,0 +1,54 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Ported from nccl-extensions: nccl_ep/device/jit/jit_compiler.hpp. + * Implement this interface to add a compiler backend (NVRTC, MUSA mcc, ...) + * alongside NvccCompiler. + */ + +#pragma once + +#include +#include +#include + +namespace mooncake { +namespace jit { + +struct JitCompileConfig { + int sm = 0; + std::string target_arch; + std::filesystem::path source_dir; + std::filesystem::path build_include_dir; + std::filesystem::path cuda_include_dir; +}; + +struct JitCompileInput { + std::string source; + std::vector options; + std::filesystem::path source_path; + std::filesystem::path cubin_path; + bool log_enabled = false; +}; + +struct JitCompileOutput { + std::string cubin; + std::string log; + double compiler_sec = 0.0; +}; + +class JitCompiler { + public: + virtual ~JitCompiler() = default; + + virtual std::string compiler_id() const = 0; + virtual std::vector compile_options( + const JitCompileConfig& config) const = 0; + virtual bool compile_to_cubin(const JitCompileInput& input, + JitCompileOutput* output) const = 0; +}; + +} // namespace jit +} // namespace mooncake diff --git a/mooncake-ep/include/jit/jit_runtime.hpp b/mooncake-ep/include/jit/jit_runtime.hpp new file mode 100644 index 0000000000..715280ed16 --- /dev/null +++ b/mooncake-ep/include/jit/jit_runtime.hpp @@ -0,0 +1,101 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Ported from nccl-extensions: nccl_ep/device/jit/jit_runtime.hpp. + * This engine is kernel agnostic -- JitKernelVariant is the only contract + * between it and the per-kernel source generators. + */ + +#pragma once + +#include + +#include +#include +#include +#include + +namespace mooncake { +namespace jit { + +enum class JitKernelStatus { + kLaunched, + kDisabled, + kUnsupportedDevice, + kCompileFailed, + kLoadFailed, + kAttributeFailed, + kLaunchFailed, +}; + +struct JitKernelVariant { + std::string_view kernel_family; + std::string_view variant_name; + std::string_view source; + std::string_view entry_name; + const void* identity = nullptr; + std::uint64_t runtime_key = 0; + int num_blocks = 0; + int block_dim = 0; + int dynamic_smem_bytes = 0; + int min_sm = -1; + std::string_view target_arch; + // All the kernel in mooncake-ep are launched with cooperative launch. + // Cooperative launch enables cg::this_grid().sync() inside the kernel. + // Cluster dim > 1 enables distributed shared memory across the cluster. + bool cooperative = true; + int cluster_dim_x = 1; + int cluster_dim_y = 1; + int cluster_dim_z = 1; +}; + +// Allocation-free FNV-1a folding used to derive JitKernelVariant::runtime_key +// from the raw variant parameters (ints, bools, short literal tags) instead of +// hashing a heap-built name string. runtime_key only has to be stable within a +// process: it keys the in-memory fast cache; the on-disk cache key is derived +// from the full source text. +constexpr std::uint64_t kRuntimeKeySeed = 1469598103934665603ull; + +constexpr std::uint64_t runtime_key_mix(std::uint64_t key, + std::uint64_t value) { + for (int i = 0; i < 8; ++i) { + key ^= (value >> (8 * i)) & 0xffu; + key *= 1099511628211ull; + } + return key; +} + +constexpr std::uint64_t runtime_key_mix(std::uint64_t key, const char* text) { + for (; text != nullptr && *text != '\0'; ++text) { + key ^= static_cast(*text); + key *= 1099511628211ull; + } + return key; +} + +// Fast path only: launch the variant from the in-process kernel cache. +// Requires just identity/runtime_key and the launch configuration -- +// variant_name and source may be left empty, so callers can skip building +// them entirely on the (hot) warm-cache path. Returns kDisabled on a cache +// miss; the caller then materializes the strings and calls launch_jit_kernel. +JitKernelStatus launch_jit_kernel_cached(const JitKernelVariant& variant, + void* kernel_param, + std::size_t kernel_param_size, + cudaStream_t stream, + std::string* error); + +JitKernelStatus launch_jit_kernel(const JitKernelVariant& variant, + void* kernel_param, + std::size_t kernel_param_size, + cudaStream_t stream, std::string* error); + +JitKernelStatus launch_jit_kernel(const JitKernelVariant& variant, + void* kernel_param, cudaStream_t stream, + std::string* error); + +const char* jit_kernel_status_name(JitKernelStatus status); + +} // namespace jit +} // namespace mooncake diff --git a/mooncake-ep/include/jit/jit_utils.hpp b/mooncake-ep/include/jit/jit_utils.hpp new file mode 100644 index 0000000000..3aa0a2daa3 --- /dev/null +++ b/mooncake-ep/include/jit/jit_utils.hpp @@ -0,0 +1,51 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Ported from nccl-extensions: nccl_ep/device/jit/jit_utils.hpp. + */ + +#pragma once + +#include +#include +#include +#include + +namespace mooncake { +namespace jit { + +std::string env_value(const char* name); +bool env_flag_enabled(const char* name, bool default_value = false); +// MOONCAKE_EP_JIT_LOG, read once per process. Callers on hot paths must check +// this before building a log message so disabled logging costs nothing per +// launch. +bool jit_log_enabled(); +void jit_log(std::string_view message); + +// Returns true the first time each distinct `key` is seen. Thread-safe. +bool announce_once(const std::string& key); + +class ScopedFileLock { + public: + ScopedFileLock() = default; + ScopedFileLock(const ScopedFileLock&) = delete; + ScopedFileLock& operator=(const ScopedFileLock&) = delete; + ~ScopedFileLock(); + + bool lock(const std::filesystem::path& path, std::string* error); + + private: + int fd_ = -1; +}; + +std::vector split_env_flags(const char* flags); +std::string json_escape(std::string_view text); +std::string read_file_or_empty(const std::filesystem::path& path); +bool write_file_atomic(const std::filesystem::path& path, + const std::string& data, bool binary); +std::string fnv1a_digest(const std::vector& parts); + +} // namespace jit +} // namespace mooncake diff --git a/mooncake-ep/include/jit/nvcc_compiler.hpp b/mooncake-ep/include/jit/nvcc_compiler.hpp new file mode 100644 index 0000000000..3ff2d537f9 --- /dev/null +++ b/mooncake-ep/include/jit/nvcc_compiler.hpp @@ -0,0 +1,26 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Ported from nccl-extensions: nccl_ep/device/jit/nvcc_compiler.hpp. + */ + +#pragma once + +#include + +namespace mooncake { +namespace jit { + +class NvccCompiler final : public JitCompiler { + public: + std::string compiler_id() const override; + std::vector compile_options( + const JitCompileConfig& config) const override; + bool compile_to_cubin(const JitCompileInput& input, + JitCompileOutput* output) const override; +}; + +} // namespace jit +} // namespace mooncake diff --git a/mooncake-ep/include/mooncake_ep_api.cuh b/mooncake-ep/include/mooncake_ep_api.cuh index 32742c097b..a7f40262ef 100644 --- a/mooncake-ep/include/mooncake_ep_api.cuh +++ b/mooncake-ep/include/mooncake_ep_api.cuh @@ -9,24 +9,22 @@ void dispatch(void* packed_recv_x, float* packed_recv_x_scales, int* packed_recv_count, int32_t* active_ranks, void* mxa_buffer, int* rdma_send_signal_buffer, int* rdma_recv_signal_buffer, void* rdma_send_data_buffer, void* rdma_recv_data_buffer, - void* cuda_counter_buffer, void* cuda_data_buffer, void* raddrs, - void* rkeys, void* qp_devctxs, const int32_t* nvlink_available, - void* const* ipc_peer_ptrs, const void* x, - const int64_t* topk_idx, int* next_clean_buffer, int num_tokens, - int hidden, int num_max_dispatch_tokens_per_rank, int num_topk, - int num_experts, int rank, int num_ranks, bool use_fp8, - void* workspace, cudaStream_t stream, int64_t timeout_ticks, - int phases, int active_qps_per_rank); + void* raddrs, void* rkeys, void* qp_devctxs, + const int32_t* nvlink_available, void* const* ipc_peer_ptrs, + const void* x, const int64_t* topk_idx, int* next_clean_buffer, + int num_tokens, int hidden, int num_max_dispatch_tokens_per_rank, + int num_topk, int num_experts, int rank, int num_ranks, + bool use_fp8, void* workspace, cudaStream_t stream, + int64_t timeout_ticks, int phases, int active_qps_per_rank); void combine(void* combined_x, int32_t* active_ranks, void* mxa_buffer, int* rdma_send_signal_buffer, int* rdma_recv_signal_buffer, void* rdma_send_data_buffer, void* rdma_recv_data_buffer, - void* cuda_counter_buffer, void* cuda_data_buffer, void* raddrs, - void* rkeys, void* qp_devctxs, const int32_t* nvlink_available, - void* const* ipc_peer_ptrs, const void* x, const int64_t* topk_idx, - const float* topk_weights, const int* src_info, - const int64_t* layout_range, int* next_clean_buffer, - int num_combined_tokens, int hidden, + void* raddrs, void* rkeys, void* qp_devctxs, + const int32_t* nvlink_available, void* const* ipc_peer_ptrs, + const void* x, const int64_t* topk_idx, const float* topk_weights, + const int* src_info, const int64_t* layout_range, + int* next_clean_buffer, int num_combined_tokens, int hidden, int num_max_dispatch_tokens_per_rank, int num_topk, int num_experts, int rank, int num_ranks, void* workspace, cudaStream_t stream, int64_t timeout_ticks, int phases, diff --git a/mooncake-ep/include/mooncake_ep_combine.cuh b/mooncake-ep/include/mooncake_ep_combine.cuh new file mode 100644 index 0000000000..db2a72c0c2 --- /dev/null +++ b/mooncake-ep/include/mooncake_ep_combine.cuh @@ -0,0 +1,459 @@ +#pragma once + +#include +#include +#include + +#include +#include +#include +#include + +namespace mooncake { + +using mooncake::device::CommCtx; +using mooncake::device::make_comm_ctx; +using mooncake::device::mc_atomic_add_release; +using mooncake::device::mc_bar_sync; +using mooncake::device::mc_grid_sync; +using mooncake::device::mc_ld_acquire; +using mooncake::device::mc_ld_nc; +using mooncake::device::mc_rdma_put; +using mooncake::device::mc_route_put; +using mooncake::device::mc_st_na; +using mooncake::device::mc_st_release; +#ifdef MOONCAKE_EP_USE_MUSA +using mooncake::device::mc_atomic_add_relaxed; +using mooncake::device::mc_fence_noflush; +#endif +using mooncake::device::mc_fence; +using mooncake::device::mc_fence_barrier_fence; + +struct CombineKernelArgs { + void* combined_x; + int32_t* active_ranks; + void* mxa_buffer; + int* rdma_send_signal_buffer; + int* rdma_recv_signal_buffer; + void* rdma_send_data_buffer; + void* rdma_recv_data_buffer; + void* raddrs; + void* rkeys; + void* qp_devctxs; + const int32_t* nvlink_available; + void* const* ipc_peer_ptrs; + const void* x; + const int64_t* topk_idx; + const float* topk_weights; + const int* src_info; + const int64_t* layout_range; + int* next_clean_buffer; + int* atomic_clean_flag; + int num_combined_tokens; + int num_topk; + int num_max_dispatch_tokens_per_rank; + int num_experts; + int rank; + int num_ranks; + int64_t timeout_ticks; + int phases; + bool zero_copy; + int active_qps_per_rank; +}; + +static_assert(std::is_trivially_copyable::value, + "CombineKernelArgs must be trivially copyable"); + +template +__device__ __forceinline__ void combine_kernel_impl( + const CombineKernelArgs& args) { + void* combined_x = args.combined_x; + int32_t* active_ranks = args.active_ranks; + void* mxa_buffer = args.mxa_buffer; + int* rdma_send_signal_buffer = args.rdma_send_signal_buffer; + int* rdma_recv_signal_buffer = args.rdma_recv_signal_buffer; + void* rdma_send_data_buffer = args.rdma_send_data_buffer; + void* rdma_recv_data_buffer = args.rdma_recv_data_buffer; + void* raddrs = args.raddrs; + void* rkeys = args.rkeys; + void* qp_devctxs = args.qp_devctxs; + const int32_t* nvlink_available = args.nvlink_available; + void* const* ipc_peer_ptrs = args.ipc_peer_ptrs; + const void* x = args.x; + const int64_t* topk_idx = args.topk_idx; + const float* topk_weights = args.topk_weights; + const int* src_info = args.src_info; + const int64_t* layout_range = args.layout_range; + int* next_clean_buffer = args.next_clean_buffer; + int* atomic_clean_flag = args.atomic_clean_flag; + int num_combined_tokens = args.num_combined_tokens; + int num_topk = args.num_topk; + int num_max_dispatch_tokens_per_rank = + args.num_max_dispatch_tokens_per_rank; + int num_experts = args.num_experts; + int rank = args.rank; + int num_ranks = args.num_ranks; + int64_t timeout_ticks = args.timeout_ticks; + int phases = args.phases; + bool zero_copy = args.zero_copy; + int active_qps_per_rank = args.active_qps_per_rank; + + const auto sm_id = static_cast(blockIdx.x); + const auto num_sms = static_cast(gridDim.x); + const auto thread_id = static_cast(threadIdx.x); + const auto num_threads = static_cast(blockDim.x); + const auto warp_id = thread_id / 32, lane_id = get_lane_id(); + const auto num_local_experts = num_experts / num_ranks; + const auto warp_group_id = warp_id / kNumWarpsPerGroup; + const auto sub_warp_id = warp_id % kNumWarpsPerGroup; + const auto responsible_expert_idx = sm_id * kNumWarpGroups + warp_group_id; + + // Data type staffs + constexpr int kNumElemsPerInt4 = sizeof(int4) / EP_BF16_SIZE; + const size_t hidden_bf16_int4 = kHidden / kNumElemsPerInt4; + + // Message package + constexpr size_t num_bytes_per_slot = kHidden * EP_BF16_SIZE; + EP_STATIC_ASSERT(num_bytes_per_slot % sizeof(int4) == 0, + "Invalid vectorization"); + + // Communication context — platform dispatch is inside comm_device.cuh + const CommCtx comm_ctx = + make_comm_ctx(mxa_buffer, nvlink_available, ipc_peer_ptrs, raddrs, + rkeys, qp_devctxs, rdma_send_signal_buffer, + rdma_recv_signal_buffer, rank, num_ranks, MAX_QP_COUNT); + const size_t num_qp_per_rank = MAX_QP_COUNT / num_ranks; + + // Sending phase + if ((phases & LOW_LATENCY_SEND_PHASE) == 0) goto LOW_LATENCY_COMBINE_RECV; + + // Clean up next buffer + if (sm_id == 0 and warp_group_id == 0 and sub_warp_id == 0) { +#pragma unroll + for (int i = lane_id; i < num_experts; i += 32) + next_clean_buffer[i] = 0; + + // Notify before executing `int_p` + __syncwarp(); +#ifdef MOONCAKE_EP_USE_MUSA + mc_fence(); + __syncwarp(); + if (lane_id == 0) mc_atomic_add_relaxed(atomic_clean_flag, num_experts); +#else + if (lane_id == 0) mc_atomic_add_release(atomic_clean_flag, num_experts); +#endif + } + + // Issue IBGDA sends + if (responsible_expert_idx < num_experts) { + const auto dst_rank = responsible_expert_idx / num_local_experts; + const auto local_expert_idx = + responsible_expert_idx % num_local_experts; + const auto global_expert_idx = + rank * num_local_experts + local_expert_idx; + const auto layout = + __ldg(layout_range + local_expert_idx * num_ranks + dst_rank); + const auto local_x = reinterpret_cast(x) + + local_expert_idx * num_ranks * + num_max_dispatch_tokens_per_rank * + hidden_bf16_int4; + const auto local_src_info = + src_info + + local_expert_idx * num_ranks * num_max_dispatch_tokens_per_rank; + const auto rdma_send_x_vec = + reinterpret_cast(rdma_send_data_buffer) + + local_expert_idx * num_ranks * num_max_dispatch_tokens_per_rank * + num_bytes_per_slot; + + // Unpack layout + int offset, num_tokens_to_send; + unpack2(layout, num_tokens_to_send, offset); + + // Issue IBGDA send + for (int token_idx = offset + sub_warp_id; + token_idx < offset + num_tokens_to_send; + token_idx += kNumWarpsPerGroup) { + const auto x_int4 = local_x + token_idx * hidden_bf16_int4; + const auto rdma_send_type_row = reinterpret_cast( + rdma_send_x_vec + token_idx * num_bytes_per_slot); + const auto rdma_send_x_vec_row = + reinterpret_cast(rdma_send_type_row); + + // Copy directly to local rank, or copy to buffer and issue RDMA + auto src_idx = __ldg(local_src_info + token_idx); + const auto buf_ptr = reinterpret_cast(rdma_send_x_vec_row); + const auto dst_ptr = reinterpret_cast( + reinterpret_cast(rdma_recv_data_buffer) + + (global_expert_idx * num_max_dispatch_tokens_per_rank + + src_idx) * + num_bytes_per_slot); + + void* write_dst = mc_route_put(comm_ctx, dst_rank, dst_ptr); + if (write_dst != nullptr) { + // Local or P2P path — warp-cooperative copy + const auto dst_int4_ptr = reinterpret_cast(write_dst); + UNROLLED_WARP_COPY(7, lane_id, hidden_bf16_int4, dst_int4_ptr, + x_int4, mc_ld_nc, mc_st_na); +#ifndef MOONCAKE_EP_USE_MUSA + mc_fence(); +#endif + } else { + // IBGDA path — stage to send buffer then RDMA write + const auto buf_int4_ptr = reinterpret_cast(buf_ptr); + if (not zero_copy) + UNROLLED_WARP_COPY(7, lane_id, hidden_bf16_int4, + buf_int4_ptr, x_int4, mc_ld_nc, + mc_st_na); + __syncwarp(); +#if defined(MOONCAKE_EP_USE_MACA) || defined(MOONCAKE_EP_USE_MUSA) + // Match dispatch publication and avoid concurrent posts from + // the two 32-lane sub-warps sharing one QP. + mc_fence(); + for (int half = 0; half < 2; ++half) { + if ((sub_warp_id & 1) == half) { +#endif + mc_rdma_put( + comm_ctx, + ep_qp_channel(local_expert_idx, num_qp_per_rank, + active_qps_per_rank), + dst_rank, num_qp_per_rank, buf_ptr, dst_ptr, + num_bytes_per_slot, lane_id); +#if defined(MOONCAKE_EP_USE_MACA) || defined(MOONCAKE_EP_USE_MUSA) + } + __syncwarp(); + } +#endif + } + } + // Put finishing flag + EP_STATIC_ASSERT(kNumWarpsPerGroup > 1, + "Requires more than one warp per group"); +#ifdef MOONCAKE_EP_USE_MUSA + // Publish all local/P2P payload stores once for this CTA before any + // expert completion word is made observable by a peer. + mc_fence_barrier_fence(); +#else + mc_bar_sync(warp_group_id + 1, kNumWarpsPerGroup * 32); +#endif + if (sub_warp_id == 1 and lane_id == 0) { +#ifdef MOONCAKE_EP_USE_MUSA + while (atomicAdd(atomic_clean_flag, 0) == 0); +#else + while (mc_ld_acquire(atomic_clean_flag) == 0); +#endif + if (dst_rank != rank) { + int* signal_ptr = rdma_recv_signal_buffer + global_expert_idx; + mc_signal(comm_ctx, dst_rank, + ep_qp_channel(local_expert_idx, num_qp_per_rank, + active_qps_per_rank), + num_qp_per_rank, signal_ptr, 1); + } else { + mc_st_release(rdma_recv_signal_buffer + global_expert_idx, 1); + } +#ifdef MOONCAKE_EP_USE_MUSA + // The peer signal already publishes the payload. This local tally + // is reused after the current kernel completes on the same stream. + mc_atomic_add_relaxed(atomic_clean_flag, -1); +#else + mc_atomic_add_release(atomic_clean_flag, -1); +#endif + } + __syncwarp(); + } else { +#ifdef MOONCAKE_EP_USE_MUSA + mc_fence_barrier_fence(); +#else + mc_bar_sync(warp_group_id + 1, kNumWarpsPerGroup * 32); +#endif + } + +// Receiving phase +LOW_LATENCY_COMBINE_RECV: + if ((phases & LOW_LATENCY_RECV_PHASE) == 0) return; + + // Wait all ranks to arrive + if (responsible_expert_idx < num_experts) { + const auto src_rank = responsible_expert_idx / num_local_experts; + EP_STATIC_ASSERT(kNumWarpsPerGroup > 1, + "Invalid number of warps per group"); + if (sub_warp_id == 0 and lane_id == 0) { + unsigned long long start_time = clock64(); + while (mc_ld_acquire(rdma_recv_signal_buffer + + responsible_expert_idx) == 0) { + unsigned long long end_time = clock64(); + if (timeout_ticks != -1 && + end_time - start_time > timeout_ticks) { + active_ranks[src_rank] = 0; + } + if (!active_ranks[src_rank]) { + break; + } + } + } + } +#ifdef MOONCAKE_EP_SPLIT_SEND_RECV + // mc_grid_sync() is a no-op on split-kernel platforms; use a block-wide + // fence/barrier before reduction so threads see peer writes. + __syncthreads(); +#ifndef MOONCAKE_EP_USE_MUSA + mc_fence(); +#endif + __syncthreads(); +#else + mc_grid_sync(); +#endif + + // Reduce tokens with FP8 cast + EP_DEVICE_ASSERT(num_topk <= 32 and hidden_bf16_int4 <= num_threads); + EP_STATIC_ASSERT(kHidden % (32 * kNumElemsPerInt4) == 0, + "Invalid vectorization"); +#ifdef MOONCAKE_EP_USE_MUSA + // A completion word is stable for the entire combine. When the expert + // count is smaller than the token/expert fanout, acquire every expert + // once per reduction CTA instead of repeatedly acquiring the same words + // for every token. + const bool wait_all_experts = num_experts <= num_combined_tokens * num_topk; + if (wait_all_experts) { + const int entry_idx = thread_id; + const int token_offset = entry_idx / num_topk; + const int topk_offset = entry_idx % num_topk; + const int token_idx = sm_id + token_offset * num_sms; + if (token_idx < num_combined_tokens) { + const int expert_idx = static_cast( + __ldg(topk_idx + token_idx * num_topk + topk_offset)); + if (expert_idx >= 0) { + const int expert_src_rank = expert_idx / num_local_experts; + const unsigned long long start_time = clock64(); + while (mc_ld_acquire(rdma_recv_signal_buffer + expert_idx) == + 0 && + active_ranks[expert_src_rank]) { + const unsigned long long end_time = clock64(); + if (timeout_ticks != -1 && + end_time - start_time > timeout_ticks) + active_ranks[expert_src_rank] = 0; + } + } + } + } + __syncthreads(); + for (int token_idx = sm_id; token_idx < num_combined_tokens; + token_idx += num_sms) { + // CUDA's cooperative grid barrier above orders an expert-owning block + // before every token-reduction block. MUSA has no grid barrier, so + // the reduction block waits only on the experts selected by its token. + if (!wait_all_experts && thread_id == 0) { +#pragma unroll + for (int i = 0; i < num_topk; ++i) { + const int expert_idx = static_cast( + __ldg(topk_idx + token_idx * num_topk + i)); + if (expert_idx < 0) continue; + const int expert_src_rank = expert_idx / num_local_experts; + const unsigned long long start_time = clock64(); + while (mc_ld_acquire(rdma_recv_signal_buffer + expert_idx) == + 0 && + active_ranks[expert_src_rank]) { + const unsigned long long end_time = clock64(); + if (timeout_ticks != -1 && + end_time - start_time > timeout_ticks) + active_ranks[expert_src_rank] = 0; + } + } + } + __syncthreads(); + if (thread_id < hidden_bf16_int4) { +#else + if (thread_id < hidden_bf16_int4) { + for (int token_idx = sm_id; token_idx < num_combined_tokens; + token_idx += num_sms) { +#endif +#ifndef MOONCAKE_EP_USE_MUSA + mc_fence(); +#endif + // Read top-k indices and weights + int reg_topk_idx[kNumMaxTopk]; + float reg_topk_weights[kNumMaxTopk]; +#pragma unroll + for (int i = 0; i < num_topk; ++i) { + reg_topk_idx[i] = static_cast( + __ldg(topk_idx + token_idx * num_topk + i)); + reg_topk_weights[i] = + __ldg(topk_weights + token_idx * num_topk + i); + } + +#ifdef MOONCAKE_EP_USE_MUSA + float2 combined_values[kNumElemsPerInt4 / 2] = {}; +#else + float combined_values[kNumElemsPerInt4] = {0.0f}; +#endif +#pragma unroll + for (int i = 0; i < num_topk; ++i) + if (reg_topk_idx[i] >= 0) { + // Skip experts on inactive ranks (timed out during combine + // recv) + int expert_src_rank = reg_topk_idx[i] / num_local_experts; + if (!active_ranks[expert_src_rank]) continue; + // Read from sources + auto rdma_buffer_type = reinterpret_cast( + reinterpret_cast(rdma_recv_data_buffer) + + (reg_topk_idx[i] * num_max_dispatch_tokens_per_rank + + token_idx) * + num_bytes_per_slot); + auto rdma_buffer_row = + reinterpret_cast(rdma_buffer_type); + + // Reduce + auto x_vec = mc_ld_nc( + reinterpret_cast(rdma_buffer_row) + + thread_id); +#ifdef MOONCAKE_EP_USE_MUSA + const auto x_bf162 = + reinterpret_cast(&x_vec); +#pragma unroll + for (int j = 0; j < kNumElemsPerInt4 / 2; ++j) { + const float2 values = {__bfloat162float(x_bf162[j].x), + __bfloat162float(x_bf162[j].y)}; + combined_values[j].x += values.x * reg_topk_weights[i]; + combined_values[j].y += values.y * reg_topk_weights[i]; + } +#else + const auto x_bf16 = reinterpret_cast(&x_vec); +#pragma unroll + for (int j = 0; j < kNumElemsPerInt4; ++j) + combined_values[j] += + __bfloat162float(x_bf16[j]) * reg_topk_weights[i]; +#endif + } + + // Write results +#ifdef MOONCAKE_EP_USE_MUSA + int4 combined_int4; + auto combined_bf162 = + reinterpret_cast(&combined_int4); +#pragma unroll + for (int j = 0; j < kNumElemsPerInt4 / 2; ++j) + combined_bf162[j] = __floats2bfloat162_rn(combined_values[j].x, + combined_values[j].y); +#else + int4& combined_int4 = *reinterpret_cast(combined_values); + auto combined_bf16 = + reinterpret_cast(&combined_values); +#pragma unroll + for (int j = 0; j < kNumElemsPerInt4; ++j) + combined_bf16[j] = __float2bfloat16(combined_values[j]); +#endif + (reinterpret_cast(combined_x) + + token_idx * hidden_bf16_int4)[thread_id] = combined_int4; + } + } +} + +template +__global__ EP_LAUNCH_BOUNDS(kNumWarpGroups* kNumWarpsPerGroup * 32, + 1) void combine(const CombineKernelArgs args) { + combine_kernel_impl(args); +} + +} // namespace mooncake diff --git a/mooncake-ep/include/mooncake_ep_dispatch.cuh b/mooncake-ep/include/mooncake_ep_dispatch.cuh new file mode 100644 index 0000000000..0065e6b35e --- /dev/null +++ b/mooncake-ep/include/mooncake_ep_dispatch.cuh @@ -0,0 +1,568 @@ +#pragma once + +#include +#include +#include + +#include +#include +#include +#include + +namespace mooncake { + +using mooncake::device::CommCtx; +using mooncake::device::make_comm_ctx; +using mooncake::device::mc_atomic_add_release; +using mooncake::device::mc_bar_sync; +using mooncake::device::mc_grid_sync; +using mooncake::device::mc_ld_acquire; +using mooncake::device::mc_ld_nc; +using mooncake::device::mc_ld_nc_f32; +using mooncake::device::mc_ld_nc_s32; +using mooncake::device::mc_rdma_put; +using mooncake::device::mc_red_add; +using mooncake::device::mc_route_put; +using mooncake::device::mc_st_na; +using mooncake::device::mc_st_release; +#ifdef MOONCAKE_EP_USE_MUSA +using mooncake::device::mc_atomic_add_relaxed; +using mooncake::device::mc_fence_noflush; +#endif +using mooncake::device::mc_fence; +using mooncake::device::mc_fence_barrier_fence; + +struct DispatchKernelArgs { + void* packed_recv_x; + float* packed_recv_x_scales; + int* packed_recv_src_info; + int64_t* packed_recv_layout_range; + int* packed_recv_count; + int32_t* active_ranks; + void* mxa_buffer; + int* rdma_send_signal_buffer; + int* rdma_recv_signal_buffer; + void* rdma_send_data_buffer; + void* rdma_recv_data_buffer; + void* raddrs; + void* rkeys; + void* qp_devctxs; + const int32_t* nvlink_available; + void* const* ipc_peer_ptrs; + const void* x; + const int64_t* topk_idx; + int* atomic_counter_per_expert; + int* atomic_finish_counter_per_expert; + int* next_clean_buffer; + int num_tokens; + int num_max_dispatch_tokens_per_rank; + int num_topk; + int num_experts; + int rank; + int num_ranks; + int64_t timeout_ticks; + int phases; + int active_qps_per_rank; +}; + +static_assert(std::is_trivially_copyable::value, + "DispatchKernelArgs must be trivially copyable"); + +template +__device__ __forceinline__ void dispatch_kernel_impl( + const DispatchKernelArgs& args) { + void* packed_recv_x = args.packed_recv_x; + float* packed_recv_x_scales = args.packed_recv_x_scales; + int* packed_recv_src_info = args.packed_recv_src_info; + int64_t* packed_recv_layout_range = args.packed_recv_layout_range; + int* packed_recv_count = args.packed_recv_count; + int32_t* active_ranks = args.active_ranks; + void* mxa_buffer = args.mxa_buffer; + int* rdma_send_signal_buffer = args.rdma_send_signal_buffer; + int* rdma_recv_signal_buffer = args.rdma_recv_signal_buffer; + void* rdma_send_data_buffer = args.rdma_send_data_buffer; + void* rdma_recv_data_buffer = args.rdma_recv_data_buffer; + void* raddrs = args.raddrs; + void* rkeys = args.rkeys; + void* qp_devctxs = args.qp_devctxs; + const int32_t* nvlink_available = args.nvlink_available; + void* const* ipc_peer_ptrs = args.ipc_peer_ptrs; + const void* x = args.x; + const int64_t* topk_idx = args.topk_idx; + int* atomic_counter_per_expert = args.atomic_counter_per_expert; + int* atomic_finish_counter_per_expert = + args.atomic_finish_counter_per_expert; + int* next_clean_buffer = args.next_clean_buffer; + int num_tokens = args.num_tokens; + int num_max_dispatch_tokens_per_rank = + args.num_max_dispatch_tokens_per_rank; + int num_topk = args.num_topk; + int num_experts = args.num_experts; + int rank = args.rank; + int num_ranks = args.num_ranks; + int64_t timeout_ticks = args.timeout_ticks; + int phases = args.phases; + int active_qps_per_rank = args.active_qps_per_rank; + + const auto sm_id = static_cast(blockIdx.x); + const auto thread_id = static_cast(threadIdx.x); + const auto warp_id = thread_id / 32, lane_id = get_lane_id(); + const auto num_sms = static_cast(gridDim.x); + const auto num_warps = kNumWarpGroups * kNumWarpsPerGroup; + const auto num_local_experts = num_experts / num_ranks; + const auto warp_group_id = warp_id / kNumWarpsPerGroup; + const auto sub_warp_id = warp_id % kNumWarpsPerGroup; + const auto responsible_expert_idx = sm_id * kNumWarpGroups + warp_group_id; +#if defined(MOONCAKE_EP_USE_MUSA) || defined(MOONCAKE_EP_USE_MACA) + // C500 reports 64-thread hardware warps. Do not split the last hardware + // warp by assigning only the final 32-thread pseudo-warp to count work. + // Reserve one full warp group from the data path, but write counts from a + // single 32-thread lane group to avoid duplicate per-expert increments. + const bool is_count_warp = warp_group_id == kNumWarpGroups - 1; + const bool is_count_worker = is_count_warp && sub_warp_id == 0; + const bool is_data_warp = warp_group_id < kNumWarpGroups - 1; + const int num_send_threads = (kNumWarpGroups - 1) * kNumWarpsPerGroup * 32; +#else + const bool is_count_warp = warp_id == num_warps - 1; + const bool is_count_worker = is_count_warp; + const bool is_data_warp = warp_id < num_warps - 1; + const int num_send_threads = (num_warps - 1) * 32; +#endif + + // FP8 staffs + constexpr int kNumPerChannels = 128; + constexpr float kFP8Margin = 1e-4, kFP8Amax = 448, + kFP8AmaxInv = 1.0f / 448.0f; + const int num_scales = kHidden / kNumPerChannels; + const size_t hidden_bytes = + kHidden * (kUseFP8 ? sizeof(ep_fp8_storage_t) : EP_BF16_SIZE); + const size_t hidden_int4 = hidden_bytes / sizeof(int4); + + // Message package: hidden data, FP8 scales, index at source + // NOTES: currently we have 3 reserved int fields for future use + using vec_t = typename std::conditional::type; + const size_t num_bytes_per_msg = + sizeof(int4) + (kUseFP8 ? (kHidden + num_scales * sizeof(float)) + : (kHidden * EP_BF16_SIZE)); + const size_t num_int4_per_msg = num_bytes_per_msg / sizeof(int4); + EP_DEVICE_ASSERT(num_bytes_per_msg % sizeof(int4) == 0); + + // Communication context — platform dispatch is inside comm_device.cuh + const CommCtx comm_ctx = + make_comm_ctx(mxa_buffer, nvlink_available, ipc_peer_ptrs, raddrs, + rkeys, qp_devctxs, rdma_send_signal_buffer, + rdma_recv_signal_buffer, rank, num_ranks, MAX_QP_COUNT); + const size_t num_qp_per_rank = MAX_QP_COUNT / num_ranks; + + // Sending phase + if ((phases & LOW_LATENCY_SEND_PHASE) == 0) goto LOW_LATENCY_DISPATCH_RECV; + + // Expert counts + __shared__ int shared_num_tokens_sent_per_expert[kNumWarpGroups]; + + // There are 2 kinds of execution lanes in this part: + // 1. Data lanes for FP8 cast and sending top-k tokens. + // 2. Count lanes for reading `topk_idx` and per-expert token counts. + // Non-CUDA backends reserve a full warp group for the count path. This + // keeps the final group out of the data path when MUSA uses five groups. + if (is_data_warp) { + constexpr int kNumElemsPerRead = sizeof(int4) / EP_BF16_SIZE; + EP_DEVICE_ASSERT(kHidden % kNumElemsPerRead == 0); + EP_STATIC_ASSERT(kNumElemsPerRead * 32 % kNumPerChannels == 0, + "Invalid vectorization"); + const auto num_threads = num_send_threads; + const size_t hidden_bf16_int4 = kHidden / kNumElemsPerRead; + + for (int token_idx = sm_id; token_idx < num_tokens; + token_idx += num_sms) { + const auto x_int4 = + reinterpret_cast(x) + token_idx * hidden_bf16_int4; + const auto rdma_x_src_idx = reinterpret_cast( + reinterpret_cast(rdma_send_data_buffer) + + token_idx * num_bytes_per_msg); + const auto rdma_x_vec = reinterpret_cast( + reinterpret_cast(rdma_x_src_idx) + sizeof(int4)); + const auto rdma_x_scales = reinterpret_cast( + reinterpret_cast(rdma_x_vec) + hidden_bytes); + + // Overlap top-k index read and source token index write + auto dst_expert_idx = + warp_id < num_topk + ? static_cast( + __ldg(topk_idx + token_idx * num_topk + warp_id)) + : -1; + thread_id == 0 ? (*rdma_x_src_idx = token_idx) : 0; + +// FP8 cast +#pragma unroll + for (int i = thread_id; i < hidden_bf16_int4; i += num_threads) { + // Read + auto int4_value = __ldg(x_int4 + i); + + if (kUseFP8) { + // Calculate local amax + auto bf16_values = + reinterpret_cast(&int4_value); + float fp32_values[kNumElemsPerRead]; + float amax = kFP8Margin, scale, scale_inv; +#pragma unroll + for (int j = 0; j < kNumElemsPerRead; ++j) { + fp32_values[j] = __bfloat162float(bf16_values[j]); + amax = fmaxf(amax, fabsf(fp32_values[j])); + } + + // Reduce amax and scale + EP_STATIC_ASSERT( + kNumElemsPerRead * 32 / kNumPerChannels == 2, + "Invalid vectorization"); + amax = half_warp_reduce_max(amax), scale = kFP8Amax / amax, + scale_inv = amax * kFP8AmaxInv; + if (lane_id == 0 or lane_id == 16) + rdma_x_scales[i * kNumElemsPerRead / 128] = scale_inv; + + // Cast into send buffer + vec_t int2_value; + auto fp8x2_values = + reinterpret_cast(&int2_value); +#pragma unroll + for (int j = 0; j < kNumElemsPerRead; j += 2) { + float2 fp32x2 = {fp32_values[j] * scale, + fp32_values[j + 1] * scale}; + fp8x2_values[j / 2] = ep_cvt_float2_to_fp8x2(fp32x2); + } + rdma_x_vec[i] = int2_value; + } else { + // Reinterpret-cast is for C++14 compatibility + rdma_x_vec[i] = *reinterpret_cast(&int4_value); + } + } + mc_bar_sync(1, num_threads); + + // Issue sends + if (dst_expert_idx >= 0) { + int slot_idx = + lane_id == 0 + ? atomicAdd(atomic_counter_per_expert + dst_expert_idx, + 1) + : 0; + slot_idx = __shfl_sync(0xffffffff, slot_idx, 0); + const auto dst_rank = dst_expert_idx / num_local_experts; + const auto dst_expert_local_idx = + dst_expert_idx % num_local_experts; + const auto src_ptr = + reinterpret_cast(rdma_x_src_idx); + const auto dst_ptr = reinterpret_cast( + reinterpret_cast(rdma_recv_data_buffer) + + dst_expert_local_idx * num_ranks * + num_max_dispatch_tokens_per_rank * num_bytes_per_msg + + rank * num_max_dispatch_tokens_per_rank * + num_bytes_per_msg + + slot_idx * num_bytes_per_msg); + + void* write_dst = mc_route_put(comm_ctx, dst_rank, dst_ptr); + if (write_dst != nullptr) { + // Local or P2P path — warp-cooperative copy + const auto* src_int4_ptr = + reinterpret_cast(src_ptr); + const auto* dst_int4_ptr = + reinterpret_cast(write_dst); +#ifndef MOONCAKE_EP_USE_MUSA + mc_fence(); +#endif + UNROLLED_WARP_COPY(8, lane_id, num_int4_per_msg, + dst_int4_ptr, src_int4_ptr, mc_ld_nc, + mc_st_na); +#ifndef MOONCAKE_EP_USE_MUSA + mc_fence(); +#endif + } else { + // IBGDA path — send directly from source buffer + mc_fence(); +#if defined(MOONCAKE_EP_USE_MACA) || defined(MOONCAKE_EP_USE_MUSA) + // MACA and MUSA execute two 32-lane sub-warps in one + // hardware warp; serialize posts sharing one QP. + for (int half = 0; half < 2; ++half) { + if ((sub_warp_id & 1) == half) { +#endif + mc_rdma_put(comm_ctx, + ep_qp_channel(dst_expert_local_idx, + num_qp_per_rank, + active_qps_per_rank), + dst_rank, num_qp_per_rank, src_ptr, + dst_ptr, num_bytes_per_msg, lane_id); +#if defined(MOONCAKE_EP_USE_MACA) || defined(MOONCAKE_EP_USE_MUSA) + } + __syncwarp(); + } +#endif + } + + // Increase counter after finishing + __syncwarp(); +#ifdef MOONCAKE_EP_USE_MUSA + mc_fence_noflush(); +#endif + lane_id == 0 + ? +#ifdef MOONCAKE_EP_USE_MUSA + mc_atomic_add_relaxed( + atomic_finish_counter_per_expert + dst_expert_idx, 1) + : 0; +#else + mc_atomic_add_release( + atomic_finish_counter_per_expert + dst_expert_idx, 1) + : 0; +#endif + } + } + } else if (is_count_warp) { +#ifdef MOONCAKE_EP_SPLIT_SEND_RECV + // Participate in __syncthreads() barriers from data warps. + // Each token iteration in the send loop above calls + // __syncthreads() once; the count path must match. + for (int token_idx = sm_id; token_idx < num_tokens; + token_idx += num_sms) { + __syncthreads(); + } +#endif + } + if (is_count_worker) { + EP_DEVICE_ASSERT(num_sms > 1); + if (sm_id == 0) { +// The first SM is also responsible for cleaning the next buffer +#pragma unroll + for (int i = lane_id; i < num_experts; i += 32) + next_clean_buffer[i] = 0; + + // Notify before executing `int_p` + __syncwarp(); +#ifdef MOONCAKE_EP_USE_MUSA + // Publish the cleared buffer once before any local ready tag. + mc_fence(); + __syncwarp(); +#endif +#pragma unroll + for (int i = lane_id; i < num_experts; i += 32) +#ifdef MOONCAKE_EP_USE_MUSA + mc_atomic_add_relaxed(atomic_finish_counter_per_expert + i, + FINISHED_SUM_TAG); +#else + mc_atomic_add_release(atomic_finish_counter_per_expert + i, + FINISHED_SUM_TAG); +#endif + } + + // This SM should be responsible for some destination experts, read + // `topk_idx` for them + int expert_count[kNumWarpGroups] = {0}; + const auto expert_begin_idx = sm_id * kNumWarpGroups; + const auto expert_end_idx = + min(expert_begin_idx + kNumWarpGroups, num_experts); + +// Per lane count +#pragma unroll 8 + for (int i = lane_id; i < num_tokens * num_topk; i += 32) { + auto idx = static_cast(__ldg(topk_idx + i)); + if (idx >= expert_begin_idx and idx < expert_end_idx) + expert_count[idx - expert_begin_idx]++; + } + +// Warp reduce +#pragma unroll + for (int i = expert_begin_idx; i < expert_end_idx; ++i) { + auto sum = warp_reduce_sum(expert_count[i - expert_begin_idx]); + if (lane_id == 0) { + shared_num_tokens_sent_per_expert[i - expert_begin_idx] = sum; +#ifdef MOONCAKE_EP_USE_MUSA + // Local completion accounting; payload publication is fenced + // by the expert's signal sender after the count reaches its + // tag. + mc_atomic_add_relaxed(atomic_finish_counter_per_expert + i, + FINISHED_SUM_TAG - sum); +#else + mc_atomic_add_release(atomic_finish_counter_per_expert + i, + FINISHED_SUM_TAG - sum); +#endif + } + } + } +#ifdef MOONCAKE_EP_USE_MUSA + __syncthreads(); +#else + mc_fence_barrier_fence(); +#endif + + // Issue count sends + if (responsible_expert_idx < num_experts and sub_warp_id == 0 and + lane_id == 0) { + const auto dst_rank = responsible_expert_idx / num_local_experts; + const auto dst_expert_local_idx = + responsible_expert_idx % num_local_experts; + const auto num_tokens_sent = + shared_num_tokens_sent_per_expert[responsible_expert_idx - + sm_id * kNumWarpGroups]; + + // Wait local sends issued and send expert counts +#if defined(MOONCAKE_EP_USE_MACA) || defined(MOONCAKE_EP_USE_MUSA) + while ( + atomicAdd(atomic_finish_counter_per_expert + responsible_expert_idx, + 0) != FINISHED_SUM_TAG * 2) { + } +#else + while (mc_ld_acquire(atomic_finish_counter_per_expert + + responsible_expert_idx) != FINISHED_SUM_TAG * 2) { + } +#endif + if (dst_rank != rank) { +#ifdef MOONCAKE_EP_USE_MUSA + mc_fence(); +#endif + int* signal_ptr = rdma_recv_signal_buffer + + dst_expert_local_idx * num_ranks + rank; + mc_red_add(comm_ctx, dst_rank, + ep_qp_channel(dst_expert_local_idx, num_qp_per_rank, + active_qps_per_rank), + num_qp_per_rank, signal_ptr, + static_cast(-num_tokens_sent - 1)); + } else { + mc_st_release(rdma_recv_signal_buffer + + dst_expert_local_idx * num_ranks + rank, + -num_tokens_sent - 1); + } + + // Clean workspace for next use + atomic_counter_per_expert[responsible_expert_idx] = 0; + atomic_finish_counter_per_expert[responsible_expert_idx] = 0; + + // Clean `packed_recv_count` + if (dst_rank == 0) packed_recv_count[dst_expert_local_idx] = 0; + } + __syncwarp(); + +// Receiving phase +LOW_LATENCY_DISPATCH_RECV: + if ((phases & LOW_LATENCY_RECV_PHASE) == 0) return; + + // For send-and-recv kernels, we need a grid sync for making + // `packed_recv_count` visible + if (phases & LOW_LATENCY_SEND_PHASE) mc_grid_sync(); + + // Receiving and packing + if (responsible_expert_idx < num_experts) { + const auto src_rank = responsible_expert_idx / num_local_experts; + const auto local_expert_idx = + responsible_expert_idx % num_local_experts; + const auto rdma_recv_x_uint8 = + reinterpret_cast(rdma_recv_data_buffer) + + local_expert_idx * num_ranks * num_max_dispatch_tokens_per_rank * + num_bytes_per_msg + + src_rank * num_max_dispatch_tokens_per_rank * num_bytes_per_msg; + const auto recv_x_int4 = reinterpret_cast(packed_recv_x) + + local_expert_idx * num_ranks * + num_max_dispatch_tokens_per_rank * + hidden_int4; + const auto recv_x_scales = + packed_recv_x_scales + local_expert_idx * num_ranks * + num_max_dispatch_tokens_per_rank * + num_scales; + const auto recv_src_info = + packed_recv_src_info + + local_expert_idx * num_ranks * num_max_dispatch_tokens_per_rank; + const auto recv_range = + packed_recv_layout_range + local_expert_idx * num_ranks; + + // Shared between sub-warps in warp groups + __shared__ int shared_num_recv_tokens[kNumWarpGroups], + shared_recv_token_begin_idx[kNumWarpGroups]; + + // Wait tokens to arrive + // NOTES: using sub-warp 1 to overlap with sub-warp 0 + int num_recv_tokens, recv_token_begin_idx; + EP_STATIC_ASSERT(kNumWarpsPerGroup > 1, + "Requires more than one warp per group"); + if (sub_warp_id == 1 and lane_id == 0) { + unsigned long long start_time = clock64(); + while ((num_recv_tokens = mc_ld_acquire( + rdma_recv_signal_buffer + local_expert_idx * num_ranks + + src_rank)) == 0) { + unsigned long long end_time = clock64(); + if (timeout_ticks != -1 && + end_time - start_time > timeout_ticks) { + active_ranks[src_rank] = 0; + } + if (!active_ranks[src_rank]) { + num_recv_tokens = -1; + break; + } + } + num_recv_tokens = -num_recv_tokens - 1; + recv_token_begin_idx = atomicAdd( + packed_recv_count + local_expert_idx, num_recv_tokens); + shared_num_recv_tokens[warp_group_id] = num_recv_tokens; + shared_recv_token_begin_idx[warp_group_id] = recv_token_begin_idx; + recv_range[src_rank] = + pack2(num_recv_tokens, recv_token_begin_idx); + } + mc_bar_sync(warp_group_id + 2, kNumWarpsPerGroup * 32); + num_recv_tokens = shared_num_recv_tokens[warp_group_id]; + recv_token_begin_idx = shared_recv_token_begin_idx[warp_group_id]; + + // Copy tokens + EP_DEVICE_ASSERT(num_scales <= 64); + mc_fence(); + for (int i = sub_warp_id; i < num_recv_tokens; i += kNumWarpsPerGroup) { + // Copy source info + const auto src_src_idx = reinterpret_cast( + rdma_recv_x_uint8 + i * num_bytes_per_msg); + if (lane_id == 0) + recv_src_info[recv_token_begin_idx + i] = + mc_ld_nc_s32(src_src_idx); + __syncwarp(); + + // Copy data + // NOTES: only 2 load iterations for 7K hidden with 7 unrolls + const auto src_data = reinterpret_cast( + reinterpret_cast(src_src_idx) + sizeof(int4)); + const auto dst_data = + recv_x_int4 + (recv_token_begin_idx + i) * hidden_int4; +#ifndef MOONCAKE_EP_USE_MUSA + mc_fence(); +#endif + UNROLLED_WARP_COPY(7, lane_id, hidden_int4, dst_data, src_data, + mc_ld_nc, mc_st_na); + + // Copy scales + if (kUseFP8) { + const auto src_scales = reinterpret_cast( + reinterpret_cast(src_data) + hidden_bytes); + const auto dst_scales = reinterpret_cast( + recv_x_scales + recv_token_begin_idx + i); + const auto scale_stride = + num_ranks * num_max_dispatch_tokens_per_rank; + auto scale_0 = lane_id < num_scales + ? mc_ld_nc_f32(src_scales + lane_id) + : 0; + auto scale_1 = (lane_id + 32) < num_scales + ? mc_ld_nc_f32(src_scales + lane_id + 32) + : 0; + lane_id < num_scales + ? dst_scales[lane_id * scale_stride] = scale_0 + : 0.0f; + (lane_id + 32) < num_scales + ? dst_scales[(lane_id + 32) * scale_stride] = scale_1 + : 0.0f; + } + } + } else { + mc_bar_sync(warp_group_id + 2, kNumWarpsPerGroup * 32); + } +} + +template +__global__ EP_LAUNCH_BOUNDS(kNumWarpGroups* kNumWarpsPerGroup * 32, + 1) void dispatch(const DispatchKernelArgs args) { + dispatch_kernel_impl( + args); +} + +} // namespace mooncake diff --git a/mooncake-ep/include/mooncake_ep_utils.cuh b/mooncake-ep/include/mooncake_ep_utils.cuh index 7b6b194c2b..60896035cb 100644 --- a/mooncake-ep/include/mooncake_ep_utils.cuh +++ b/mooncake-ep/include/mooncake_ep_utils.cuh @@ -47,6 +47,14 @@ struct VecInt<16> { using vec_t = int4; }; +__device__ __forceinline__ int ep_qp_channel(int expert_local_idx, + int qps_per_rank, + int active_qps_per_rank) { + int active_qps = active_qps_per_rank; + if (active_qps <= 0 || active_qps > qps_per_rank) active_qps = qps_per_rank; + return expert_local_idx % active_qps; +} + // ---- TMA / mbarrier helpers (CUDA only) ---- #if !defined(MOONCAKE_EP_USE_MUSA) && !defined(MOONCAKE_EP_USE_MACA) diff --git a/mooncake-ep/src/CMakeLists.txt b/mooncake-ep/src/CMakeLists.txt index b0e84a273d..ac0e90976d 100644 --- a/mooncake-ep/src/CMakeLists.txt +++ b/mooncake-ep/src/CMakeLists.txt @@ -52,9 +52,106 @@ if(USE_CUDA) enable_language(CUDA) find_package(CUDAToolkit REQUIRED) + set(MOONCAKE_EP_JIT_SOURCES + "${CMAKE_CURRENT_SOURCE_DIR}/jit/combine_jit.cpp" + "${CMAKE_CURRENT_SOURCE_DIR}/jit/dispatch_jit.cpp" + "${CMAKE_CURRENT_SOURCE_DIR}/jit/elastic_combine_jit.cpp" + "${CMAKE_CURRENT_SOURCE_DIR}/jit/elastic_combine_reduce_epilogue_jit.cpp" + "${CMAKE_CURRENT_SOURCE_DIR}/jit/elastic_dispatch_copy_epilogue_jit.cpp" + "${CMAKE_CURRENT_SOURCE_DIR}/jit/elastic_dispatch_jit.cpp" + "${CMAKE_CURRENT_SOURCE_DIR}/jit/elastic_dispatch_prologue_jit.cpp" + "${CMAKE_CURRENT_SOURCE_DIR}/jit/elastic_hybrid_combine_jit.cpp" + "${CMAKE_CURRENT_SOURCE_DIR}/jit/elastic_hybrid_dispatch_jit.cpp" + "${CMAKE_CURRENT_SOURCE_DIR}/jit/jit_runtime.cpp" + "${CMAKE_CURRENT_SOURCE_DIR}/jit/jit_cache.cpp" + "${CMAKE_CURRENT_SOURCE_DIR}/jit/jit_utils.cpp" + "${CMAKE_CURRENT_SOURCE_DIR}/jit/nvcc_compiler.cpp") + + # The JIT engine runs nvcc at runtime, so the kernel headers the generated + # source #includes have to exist as files then. Stage a copy into the build + # tree and bake its path into the engine via jit_config.h. + # + # configure_file tracks each explicitly listed header, so edits refresh the + # staged copy and invalidate the JIT cache fingerprint. + set(MOONCAKE_EP_JIT_STAGE_ROOT "${CMAKE_BINARY_DIR}/mooncake_ep_jit") + set(MOONCAKE_EP_JIT_SOURCE_DIR + "${MOONCAKE_EP_JIT_STAGE_ROOT}/include/mooncake_ep") + set(MOONCAKE_EP_JIT_TRANSFER_ENGINE_SOURCE_DIR + "${CMAKE_CURRENT_SOURCE_DIR}/../../mooncake-transfer-engine") + + # Do not put jit_config.h under mooncake_ep_jit/: everything there goes into + # the wheel, but this header is only used to build this library. + set(MOONCAKE_EP_JIT_GENERATED_DIR "${CMAKE_CURRENT_BINARY_DIR}/jit_generated") + + set(MOONCAKE_EP_JIT_CUDA_INCLUDE_DIR "") + if(CUDAToolkit_INCLUDE_DIRS) + list(GET CUDAToolkit_INCLUDE_DIRS 0 MOONCAKE_EP_JIT_CUDA_INCLUDE_DIR) + endif() + + # NCCL DeviceTransport is the only kernel dependency living outside the + # Mooncake tree and the CUDA toolkit. + set(MOONCAKE_EP_JIT_BUILD_INCLUDE_DIR "") + if(USE_NCCL_DEVICE AND NCCL_DEVICE_INCLUDE_DIR) + set(MOONCAKE_EP_JIT_BUILD_INCLUDE_DIR "${NCCL_DEVICE_INCLUDE_DIR}") + endif() + + set(MOONCAKE_EP_JIT_HEADERS + mooncake_ep_configs.cuh + mooncake_ep_device.h + mooncake_ep_dispatch.cuh + mooncake_ep_combine.cuh + mooncake_ep_exception.cuh + mooncake_ep_utils.cuh + elastic/mooncake_ep_elastic_combine_official.cuh + elastic/mooncake_ep_elastic_combine_reduce_epilogue.cuh + elastic/mooncake_ep_elastic_combine_utils.cuh + elastic/mooncake_ep_elastic_comm.cuh + elastic/mooncake_ep_elastic_compiled.cuh + elastic/mooncake_ep_elastic_dispatch_copy_epilogue.cuh + elastic/mooncake_ep_elastic_dispatch_deterministic_prologue.cuh + elastic/mooncake_ep_elastic_dispatch_official.cuh + elastic/mooncake_ep_elastic_exception.cuh + elastic/mooncake_ep_elastic_hybrid_combine_official.cuh + elastic/mooncake_ep_elastic_hybrid_dispatch_official.cuh + elastic/mooncake_ep_elastic_launch.cuh + elastic/mooncake_ep_elastic_layout.cuh + elastic/mooncake_ep_elastic_math.cuh + elastic/mooncake_ep_elastic_ptx.cuh + elastic/mooncake_ep_elastic_transport.cuh) + + set(MOONCAKE_EP_JIT_TRANSFER_HEADERS + include/cuda_alike.h + include/cuda_fabric_compat.h + include/transport/device/comm_device.cuh + include/transport/device/cuda/cuda_ops.cuh + include/transport/device/device_ops.cuh + include/transport/device/ibgda/mlx5gda.h + include/transport/device/ibgda_device.cuh + include/transport/device/nccl_device.cuh + include/transport/device/nccl_device_transport.h + include/transport/device/p2p_device.cuh + tent/include/tent/common/cuda_fabric_compat.h) + + file(MAKE_DIRECTORY "${MOONCAKE_EP_JIT_SOURCE_DIR}") + foreach(header IN LISTS MOONCAKE_EP_JIT_HEADERS) + configure_file("${CMAKE_CURRENT_SOURCE_DIR}/../include/${header}" + "${MOONCAKE_EP_JIT_SOURCE_DIR}/${header}" COPYONLY) + endforeach() + foreach(header IN LISTS MOONCAKE_EP_JIT_TRANSFER_HEADERS) + get_filename_component(header_dir "${header}" DIRECTORY) + file(MAKE_DIRECTORY "${MOONCAKE_EP_JIT_STAGE_ROOT}/${header_dir}") + configure_file( + "${MOONCAKE_EP_JIT_TRANSFER_ENGINE_SOURCE_DIR}/${header}" + "${MOONCAKE_EP_JIT_STAGE_ROOT}/${header}" COPYONLY) + endforeach() + + configure_file("${CMAKE_CURRENT_SOURCE_DIR}/jit/jit_config.h.in" + "${MOONCAKE_EP_JIT_GENERATED_DIR}/jit/jit_config.h" @ONLY) + # Keep CUDA fatbins in a separate library. The Python host extension can # then pass through auditwheel while this library is injected afterwards. - add_library(mooncake_ep_device SHARED ${MOONCAKE_EP_DEVICE_SOURCES}) + add_library(mooncake_ep_device SHARED ${MOONCAKE_EP_DEVICE_SOURCES} + ${MOONCAKE_EP_JIT_SOURCES}) set_target_properties(mooncake_ep_device PROPERTIES POSITION_INDEPENDENT_CODE ON) if(CUDAToolkit_VERSION VERSION_GREATER_EQUAL "13.0") @@ -67,13 +164,16 @@ if(USE_CUDA) target_include_directories( mooncake_ep_device PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include - ${CMAKE_CURRENT_SOURCE_DIR}/../../mooncake-transfer-engine/include) + ${CMAKE_CURRENT_SOURCE_DIR}/../../mooncake-transfer-engine/include + ${MOONCAKE_EP_JIT_GENERATED_DIR}) target_compile_options( mooncake_ep_device PRIVATE $<$:-Xcompiler=-O3> $<$:-Xcompiler=-g0> $<$:--expt-relaxed-constexpr>) - target_link_libraries(mooncake_ep_device PRIVATE CUDA::cudart) + target_link_libraries(mooncake_ep_device PRIVATE CUDA::cudart + CUDA::cuda_driver + ${CMAKE_DL_LIBS}) if(USE_NCCL_DEVICE) target_link_libraries(mooncake_ep_device PRIVATE NCCL::nccl) endif() diff --git a/mooncake-ep/src/jit/combine_jit.cpp b/mooncake-ep/src/jit/combine_jit.cpp new file mode 100644 index 0000000000..c2923fd319 --- /dev/null +++ b/mooncake-ep/src/jit/combine_jit.cpp @@ -0,0 +1,86 @@ +#include +#include +#include + +#include +#include +#include + +namespace mooncake { +namespace jit { + +static constexpr const char* kCombineJitEntryName = + "mooncake_ep_jit_combine_kernel"; + +static std::string combine_jit_source(int hidden, int num_max_topk, + int num_warp_groups, + int num_warps_per_group) { + std::ostringstream source; + source << "#include \n" + << "\n" + << "extern \"C\" __global__ " + << "EP_LAUNCH_BOUNDS(" << num_warp_groups << " * " + << num_warps_per_group << " * 32, 1)\n" + << "void " << kCombineJitEntryName << "(\n" + << " const mooncake::CombineKernelArgs p) {\n" + << " mooncake::combine_kernel_impl<" << num_warp_groups << ", " + << num_warps_per_group << ", " << hidden << ", " << num_max_topk + << ">(p);\n" + << "}\n"; + return source.str(); +} + +void launch_combine_jit(int hidden, int num_max_topk, int num_warp_groups, + int num_warps_per_group, int num_sms, + const CombineKernelArgs& args, cudaStream_t stream) { + static const int variant_identity = 0; + + std::uint64_t key = kRuntimeKeySeed; + key = runtime_key_mix(key, static_cast(hidden)); + key = runtime_key_mix(key, static_cast(num_max_topk)); + key = runtime_key_mix(key, static_cast(num_warp_groups)); + key = runtime_key_mix(key, static_cast(num_warps_per_group)); + + JitKernelVariant variant; + variant.kernel_family = "combine"; + variant.entry_name = kCombineJitEntryName; + variant.identity = &variant_identity; + variant.runtime_key = key; + variant.num_blocks = num_sms; + variant.block_dim = num_warp_groups * num_warps_per_group * 32; + variant.dynamic_smem_bytes = 0; + variant.min_sm = 80; + variant.cooperative = true; + + std::string error; + JitKernelStatus status = launch_jit_kernel_cached( + variant, const_cast(&args), 0, stream, &error); + + if (status == JitKernelStatus::kLaunched) return; + + std::string variant_name; + std::ostringstream name; + name << "combine_h" << hidden << "_topk" << num_max_topk << "_wg" + << num_warp_groups << "x" << num_warps_per_group; + variant_name = name.str(); + + if (status != JitKernelStatus::kLaunchFailed) { + const std::string source = combine_jit_source( + hidden, num_max_topk, num_warp_groups, num_warps_per_group); + variant.variant_name = variant_name; + variant.source = source; + status = launch_jit_kernel( + variant, const_cast(&args), stream, &error); + } + + if (status != JitKernelStatus::kLaunched) { + std::ostringstream message; + message << "combine JIT launch failed for " << variant_name << ": " + << jit_kernel_status_name(status); + if (!error.empty()) message << ": " << error; + throw EPException("JIT", __FILE__, __LINE__, message.str()); + } +} + +} // namespace jit +} // namespace mooncake diff --git a/mooncake-ep/src/jit/dispatch_jit.cpp b/mooncake-ep/src/jit/dispatch_jit.cpp new file mode 100644 index 0000000000..0ab325546b --- /dev/null +++ b/mooncake-ep/src/jit/dispatch_jit.cpp @@ -0,0 +1,85 @@ +#include +#include +#include + +#include +#include +#include + +namespace mooncake { +namespace jit { + +static constexpr const char* kDispatchJitEntryName = + "mooncake_ep_jit_dispatch_kernel"; + +static std::string dispatch_jit_source(bool use_fp8, int num_warp_groups, + int num_warps_per_group, int hidden) { + std::ostringstream source; + source << "#include \n" + << "\n" + << "extern \"C\" __global__ " + << "EP_LAUNCH_BOUNDS(" << num_warp_groups << " * " + << num_warps_per_group << " * 32, 1)\n" + << "void " << kDispatchJitEntryName << "(\n" + << " const mooncake::DispatchKernelArgs p) {\n" + << " mooncake::dispatch_kernel_impl<" + << (use_fp8 ? "true" : "false") << ", " << num_warp_groups << ", " + << num_warps_per_group << ", " << hidden << ">(p);\n" + << "}\n"; + return source.str(); +} + +void launch_dispatch_jit(int hidden, bool use_fp8, int num_warp_groups, + int num_warps_per_group, int num_sms, + const DispatchKernelArgs& args, cudaStream_t stream) { + static const int variant_identity = 0; + + std::uint64_t key = kRuntimeKeySeed; + key = runtime_key_mix(key, static_cast(hidden)); + key = runtime_key_mix(key, static_cast(num_warp_groups)); + key = runtime_key_mix(key, static_cast(num_warps_per_group)); + key = runtime_key_mix(key, use_fp8 ? 1u : 0u); + + JitKernelVariant variant; + variant.kernel_family = "dispatch"; + variant.entry_name = kDispatchJitEntryName; + variant.identity = &variant_identity; + variant.runtime_key = key; + variant.num_blocks = num_sms; + variant.block_dim = num_warp_groups * num_warps_per_group * 32; + variant.dynamic_smem_bytes = 0; + variant.min_sm = 80; + variant.cooperative = true; + + std::string error; + JitKernelStatus status = launch_jit_kernel_cached( + variant, const_cast(&args), 0, stream, &error); + + if (status == JitKernelStatus::kLaunched) return; + + std::string variant_name; + std::ostringstream name; + name << "dispatch_h" << hidden << "_wg" << num_warp_groups << "x" + << num_warps_per_group << (use_fp8 ? "_fp8" : "_bf16"); + variant_name = name.str(); + + if (status != JitKernelStatus::kLaunchFailed) { + const std::string source = dispatch_jit_source( + use_fp8, num_warp_groups, num_warps_per_group, hidden); + variant.variant_name = variant_name; + variant.source = source; + status = launch_jit_kernel( + variant, const_cast(&args), stream, &error); + } + + if (status != JitKernelStatus::kLaunched) { + std::ostringstream message; + message << "dispatch JIT launch failed for " << variant_name << ": " + << jit_kernel_status_name(status); + if (!error.empty()) message << ": " << error; + throw EPException("JIT", __FILE__, __LINE__, message.str()); + } +} + +} // namespace jit +} // namespace mooncake diff --git a/mooncake-ep/src/jit/elastic_combine_jit.cpp b/mooncake-ep/src/jit/elastic_combine_jit.cpp new file mode 100644 index 0000000000..d09eef8cff --- /dev/null +++ b/mooncake-ep/src/jit/elastic_combine_jit.cpp @@ -0,0 +1,107 @@ +#include +#include +#include +#include + +#include +#include +#include + +namespace mooncake { +namespace jit { + +static constexpr const char *kElasticCombineJitEntryName = + "mooncake_ep_jit_elastic_combine_kernel"; + +// Only the scale-up NVLink, non-expanded, multiple-reduction mode is +// supported, matching the instantiations of the prebuilt path. +static std::string elastic_combine_jit_source(const char *ops_name, + int num_warps, int num_sms, + int num_scaleup_ranks, int hidden, + int num_max_tokens_per_rank, + int num_experts, int num_topk) { + std::ostringstream source; + source << "#include \n" + << "#include \n" + << "\n" + << "using Ops = mooncake::elastic::transport::" << ops_name << ";\n" + << "\n" + << "extern \"C\" __global__ " + << "__launch_bounds__(" << num_warps << " * 32, 1)\n" + << "void " << kElasticCombineJitEntryName << "(\n" + << " const mooncake::elastic::CombineKernelArgs p) {\n" + << " mooncake::elastic::combine_kernel_impl