From c5bc7dcc18bba77d6fe5fb993f3223f17b18da12 Mon Sep 17 00:00:00 2001 From: Bojun Zhu Date: Tue, 29 Sep 2026 09:38:21 +0800 Subject: [PATCH 01/17] [JIT] Add JIT framework support The framework is ported from nccl_ep JIT module (https://github.com/NVIDIA/nccl-extensions/tree/nccl-extensions-v0.1.0/nccl_ep/device/jit) Signed-off-by: Bojun Zhu --- mooncake-ep/include/jit/jit_cache.hpp | 37 + mooncake-ep/include/jit/jit_compiler.hpp | 54 ++ mooncake-ep/include/jit/jit_runtime.hpp | 101 +++ mooncake-ep/include/jit/jit_utils.hpp | 51 ++ mooncake-ep/include/jit/nvcc_compiler.hpp | 26 + mooncake-ep/src/CMakeLists.txt | 12 +- mooncake-ep/src/jit/jit_cache.cpp | 46 + mooncake-ep/src/jit/jit_runtime.cpp | 976 ++++++++++++++++++++++ mooncake-ep/src/jit/jit_utils.cpp | 229 +++++ mooncake-ep/src/jit/nvcc_compiler.cpp | 316 +++++++ 10 files changed, 1846 insertions(+), 2 deletions(-) create mode 100644 mooncake-ep/include/jit/jit_cache.hpp create mode 100644 mooncake-ep/include/jit/jit_compiler.hpp create mode 100644 mooncake-ep/include/jit/jit_runtime.hpp create mode 100644 mooncake-ep/include/jit/jit_utils.hpp create mode 100644 mooncake-ep/include/jit/nvcc_compiler.hpp create mode 100644 mooncake-ep/src/jit/jit_cache.cpp create mode 100644 mooncake-ep/src/jit/jit_runtime.cpp create mode 100644 mooncake-ep/src/jit/jit_utils.cpp create mode 100644 mooncake-ep/src/jit/nvcc_compiler.cpp 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..c41e36c77e --- /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 = 90; + std::string_view target_arch; + // Optional launch attributes (default = off). + // Cooperative launch enables cg::this_grid().sync() inside the kernel. + // Cluster dim > 1 enables distributed shared memory across the cluster. + bool cooperative = false; + 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/src/CMakeLists.txt b/mooncake-ep/src/CMakeLists.txt index b0e84a273d..8897aca781 100644 --- a/mooncake-ep/src/CMakeLists.txt +++ b/mooncake-ep/src/CMakeLists.txt @@ -52,9 +52,16 @@ if(USE_CUDA) enable_language(CUDA) find_package(CUDAToolkit REQUIRED) + set(MOONCAKE_EP_JIT_SOURCES + "${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") + # 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") @@ -73,7 +80,8 @@ if(USE_CUDA) $<$:-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) if(USE_NCCL_DEVICE) target_link_libraries(mooncake_ep_device PRIVATE NCCL::nccl) endif() diff --git a/mooncake-ep/src/jit/jit_cache.cpp b/mooncake-ep/src/jit/jit_cache.cpp new file mode 100644 index 0000000000..a5a3e067a9 --- /dev/null +++ b/mooncake-ep/src/jit/jit_cache.cpp @@ -0,0 +1,46 @@ +/* + * 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.cc. + */ + +#include +#include + +namespace mooncake { +namespace jit { + +JitCache& JitCache::instance() { + static JitCache cache; + return cache; +} + +std::filesystem::path JitCache::root_dir() const { + std::string cache_dir = env_value("MOONCAKE_EP_JIT_CACHE_DIR"); + if (!cache_dir.empty()) return std::filesystem::path(cache_dir); + + return std::filesystem::path("/tmp") / "mooncake_ep" / "jit"; +} + +bool JitCache::has_failed(const std::string& key) { + std::lock_guard lock(mutex_); + return failed_keys_.find(key) != failed_keys_.end(); +} + +void JitCache::mark_failed(const std::string& key) { + std::lock_guard lock(mutex_); + failed_keys_.insert(key); +} + +void JitCache::warn_once(const std::string& key, const std::string& message) { + { + std::lock_guard lock(mutex_); + if (!warned_keys_.insert(key).second) return; + } + jit_log(message); +} + +} // namespace jit +} // namespace mooncake diff --git a/mooncake-ep/src/jit/jit_runtime.cpp b/mooncake-ep/src/jit/jit_runtime.cpp new file mode 100644 index 0000000000..0acebfc2e2 --- /dev/null +++ b/mooncake-ep/src/jit/jit_runtime.cpp @@ -0,0 +1,976 @@ +/* + * 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.cc. + */ + +#include + +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#ifndef MOONCAKE_EP_JIT_SOURCE_DIR +#define MOONCAKE_EP_JIT_SOURCE_DIR "/usr/include/mooncake_ep" +#endif + +// Left empty by default: fingerprinting an unrelated system include root on +// every process would be expensive and meaningless. Builds that need extra +// headers (e.g. NCCL device headers) point this at the right tree. +#ifndef MOONCAKE_EP_JIT_BUILD_INCLUDE_DIR +#define MOONCAKE_EP_JIT_BUILD_INCLUDE_DIR "" +#endif + +#ifndef MOONCAKE_EP_JIT_CUDA_INCLUDE_DIR +#define MOONCAKE_EP_JIT_CUDA_INCLUDE_DIR "" +#endif + +namespace mooncake { +namespace jit { +namespace { + +struct LoadedKernelModule { + CUmodule module = nullptr; + CUfunction function = nullptr; +}; + +struct FastCacheKey { + const void* identity = nullptr; + std::uint64_t runtime_key = 0; + CUcontext context = nullptr; + int device = -1; +}; + +struct FastCacheKeyHash { + size_t operator()(const FastCacheKey& key) const { + size_t seed = std::hash{}(key.identity); + seed ^= std::hash{}(key.runtime_key) + + 0x9e3779b97f4a7c15ull + (seed << 6) + (seed >> 2); + seed ^= + std::hash{}(static_cast(key.context)) + + 0x9e3779b97f4a7c15ull + (seed << 6) + (seed >> 2); + seed ^= std::hash{}(key.device) + 0x9e3779b97f4a7c15ull + + (seed << 6) + (seed >> 2); + return seed; + } +}; + +struct FastCacheKeyEqual { + bool operator()(const FastCacheKey& lhs, const FastCacheKey& rhs) const { + return lhs.identity == rhs.identity && + lhs.runtime_key == rhs.runtime_key && + lhs.context == rhs.context && lhs.device == rhs.device; + } +}; + +std::mutex g_module_mutex; +std::mutex g_compile_mutex; +std::unordered_map g_modules; + +// Fast in-process cache: avoids all file I/O and fingerprint computation on +// repeat launches of the same variant in the same CUDA context. +struct FastCacheEntry { + CUfunction function = nullptr; + int configured_smem_bytes = 0; +}; + +// Small per-thread cache keeps repeated launches off the shared map. +constexpr int kTlCacheSlots = 4; +struct ThreadCache { + struct Slot { + bool valid = false; + FastCacheKey key{}; + FastCacheEntry entry{}; + }; + Slot slots[kTlCacheSlots]; + int next_evict = 0; +}; + +std::shared_mutex g_fast_mutex; +std::unordered_map + g_fast_cache; +thread_local ThreadCache g_thread_cache; + +const JitCompiler& compiler_backend() { + static const NvccCompiler compiler; + return compiler; +} + +void log_jit_event(const JitKernelVariant& variant, + const std::string& message) { + (void)variant; + jit_log(message); +} + +void warn_once_jit_event(JitCache& cache, const std::string& key, + const JitKernelVariant& variant, + const std::string& message) { + (void)variant; + cache.warn_once(key, message); +} + +double elapsed_sec(std::chrono::steady_clock::time_point begin) { + return std::chrono::duration(std::chrono::steady_clock::now() - + begin) + .count(); +} + +std::filesystem::path configured_source_dir() { + std::string path = env_value("MOONCAKE_EP_JIT_SOURCE_DIR"); + if (!path.empty()) return std::filesystem::path(path); + + std::string ep_home = env_value("MOONCAKE_EP_HOME"); + if (!ep_home.empty()) { + return std::filesystem::path(ep_home) / "include" / "mooncake_ep"; + } + + return std::filesystem::path(MOONCAKE_EP_JIT_SOURCE_DIR); +} + +std::filesystem::path configured_build_include_dir() { + std::string path = env_value("MOONCAKE_EP_JIT_BUILD_INCLUDE_DIR"); + if (!path.empty()) return std::filesystem::path(path); + + std::string nccl_home = env_value("NCCL_HOME"); + if (!nccl_home.empty()) { + return std::filesystem::path(nccl_home) / "include"; + } + + return std::filesystem::path(MOONCAKE_EP_JIT_BUILD_INCLUDE_DIR); +} + +std::filesystem::path configured_cuda_include_dir() { + std::string path = env_value("MOONCAKE_EP_JIT_CUDA_INCLUDE_DIR"); + if (!path.empty()) return std::filesystem::path(path); + + std::string cuda_home = env_value("CUDA_HOME"); + if (cuda_home.empty()) cuda_home = env_value("CUDA_PATH"); + if (!cuda_home.empty()) { + return std::filesystem::path(cuda_home) / "include"; + } + + return std::filesystem::path(MOONCAKE_EP_JIT_CUDA_INCLUDE_DIR); +} + +bool is_header_file(const std::filesystem::path& path) { + const std::string ext = path.extension().string(); + return ext == ".h" || ext == ".hh" || ext == ".hpp" || ext == ".cuh" || + ext == ".inc" || ext == ".inl"; +} + +void append_file_fingerprint(std::vector* parts, + const std::filesystem::path& path) { + if (parts == nullptr) return; + + std::error_code ec; + const bool exists = std::filesystem::exists(path, ec); + parts->push_back(path.lexically_normal().string()); + if (ec || !exists) { + parts->push_back(ec ? ("missing:" + ec.message()) : "missing"); + return; + } + parts->push_back(read_file_or_empty(path)); +} + +void append_header_tree_fingerprint(std::vector* parts, + const std::filesystem::path& root) { + if (parts == nullptr) return; + + std::error_code ec; + if (!std::filesystem::exists(root, ec)) { + parts->push_back(root.lexically_normal().string()); + parts->push_back(ec ? ("missing:" + ec.message()) : "missing"); + return; + } + if (!std::filesystem::is_directory(root, ec)) { + append_file_fingerprint(parts, root); + return; + } + + std::vector headers; + const auto options = + std::filesystem::directory_options::skip_permission_denied; + std::filesystem::recursive_directory_iterator it(root, options, ec); + const std::filesystem::recursive_directory_iterator end; + while (!ec && it != end) { + const std::filesystem::path path = it->path(); + std::error_code file_ec; + if (it->is_regular_file(file_ec) && !file_ec && is_header_file(path)) { + headers.push_back(path); + } + it.increment(ec); + } + std::sort(headers.begin(), headers.end()); + + parts->push_back(root.lexically_normal().string()); + if (ec) parts->push_back("walk_error:" + ec.message()); + for (const std::filesystem::path& path : headers) { + parts->push_back( + path.lexically_relative(root).lexically_normal().string()); + parts->push_back(read_file_or_empty(path)); + } +} + +// Per-SM compiler environment. Header hashing is paid once per process/SM, +// outside the repeated launch path. +struct SmEnvInfo { + std::filesystem::path source_dir; + std::filesystem::path build_include_dir; + std::filesystem::path cuda_include_dir; + std::string compiler_id; + std::vector options; + std::string env_hash; // hash(compiler_id + headers + options) +}; + +const SmEnvInfo& get_sm_env_info(int sm, std::string_view target_arch) { + static std::mutex env_mutex; + static std::unordered_map env_cache; + std::lock_guard lock(env_mutex); + + const std::string cache_key = + std::to_string(sm) + ":" + std::string(target_arch); + auto it = env_cache.find(cache_key); + if (it != env_cache.end()) return it->second; + + SmEnvInfo info; + info.source_dir = configured_source_dir(); + info.build_include_dir = configured_build_include_dir(); + info.cuda_include_dir = configured_cuda_include_dir(); + + const JitCompiler& compiler = compiler_backend(); + info.compiler_id = compiler.compiler_id(); + JitCompileConfig compile_config; + compile_config.sm = sm; + compile_config.target_arch = target_arch; + compile_config.source_dir = info.source_dir; + compile_config.build_include_dir = info.build_include_dir; + compile_config.cuda_include_dir = info.cuda_include_dir; + info.options = compiler.compile_options(compile_config); + + std::vector env_parts = { + info.compiler_id, + }; + append_header_tree_fingerprint(&env_parts, info.source_dir); + append_header_tree_fingerprint(&env_parts, info.build_include_dir); + env_parts.insert(env_parts.end(), info.options.begin(), info.options.end()); + info.env_hash = fnv1a_digest(env_parts); + + auto [inserted_it, ok] = env_cache.emplace(cache_key, std::move(info)); + (void)ok; + return inserted_it->second; +} + +std::string kernel_family(const JitKernelVariant& variant) { + return variant.kernel_family.empty() ? "jit_kernel" + : std::string(variant.kernel_family); +} + +std::string kernel_log_prefix(const JitKernelVariant& variant) { + return kernel_family(variant) + " JIT"; +} + +// Per-variant fingerprint. No file I/O here; env_hash covers header inputs. +std::string source_fingerprint(const JitKernelVariant& variant, int sm, + const std::string& env_hash) { + const std::vector parts = { + kernel_family(variant), + std::string(variant.variant_name), + std::string(variant.entry_name), + std::string(variant.source), + "sm=" + std::to_string(sm), + "target_arch=" + std::string(variant.target_arch), + "runtime_key=" + std::to_string(variant.runtime_key), + env_hash, + }; + return fnv1a_digest(parts); +} + +std::string metadata_json(const JitKernelVariant& variant, + const std::string& key, int sm, + const std::vector& options, + const std::string& compiler_id, + double jit_overhead_sec, double compiler_sec) { + std::ostringstream out; + out << "{\n" + << " \"key\": \"" << json_escape(key) << "\",\n" + << " \"kernel_family\": \"" << json_escape(kernel_family(variant)) + << "\",\n" + << " \"variant\": \"" << json_escape(variant.variant_name) << "\",\n" + << " \"entry\": \"" << json_escape(variant.entry_name) << "\",\n" + << " \"sm\": " << sm << ",\n" + << " \"blocks\": " << variant.num_blocks << ",\n" + << " \"block_dim\": " << variant.block_dim << ",\n" + << " \"dynamic_smem_bytes\": " << variant.dynamic_smem_bytes << ",\n" + << " \"jit_overhead_sec\": " << jit_overhead_sec << ",\n" + << " \"nvcc_process_sec\": " << compiler_sec << ",\n" + << " \"compiler\": \"" << json_escape(compiler_id) << "\",\n" + << " \"options\": ["; + for (size_t i = 0; i < options.size(); ++i) { + if (i != 0) out << ", "; + out << "\"" << json_escape(options[i]) << "\""; + } + out << "]\n}\n"; + return out.str(); +} + +std::filesystem::path temporary_cubin_path( + const std::filesystem::path& cubin_path) { + std::ostringstream suffix; + suffix << ".tmp." << static_cast(getpid()) << "." + << std::hash{}(std::this_thread::get_id()); + return cubin_path.parent_path() / + (cubin_path.filename().string() + suffix.str()); +} + +std::string cu_error_string(CUresult result) { + const char* name = nullptr; + const char* message = nullptr; + cuGetErrorName(result, &name); + cuGetErrorString(result, &message); + std::ostringstream out; + out << (name == nullptr ? "CUDA_ERROR_UNKNOWN" : name); + if (message != nullptr) out << ": " << message; + return out.str(); +} + +std::string context_key(CUcontext context) { + std::ostringstream out; + out << static_cast(context); + return out.str(); +} + +void log_memory_cache_event(const FastCacheKey& fast_key, + const JitKernelVariant& variant, const char* state, + const char* scope) { + // Runs on every launch (warm-cache hits included): don't build the message + // unless logging is actually on. + if (!jit_log_enabled()) return; + log_jit_event(variant, kernel_log_prefix(variant) + + " memory_cache=" + state + " scope=" + scope + + " variant=" + std::string(variant.variant_name) + + " device=" + std::to_string(fast_key.device)); +} + +JitKernelStatus load_function(const std::string& key, const std::string& cubin, + const std::string& entry_name, + CUfunction* function, std::string* error) { + { + std::lock_guard lock(g_module_mutex); + const auto it = g_modules.find(key); + if (it != g_modules.end()) { + *function = it->second.function; + return JitKernelStatus::kLaunched; + } + } + + // Load outside the mutex so unrelated threads are not blocked by CUDA + // linking. + CUmodule module = nullptr; + CUresult rc = cuInit(0); + if (rc == CUDA_SUCCESS) + rc = cuModuleLoadDataEx(&module, cubin.data(), 0, nullptr, nullptr); + if (rc != CUDA_SUCCESS) { + if (error != nullptr) *error = cu_error_string(rc); + return JitKernelStatus::kLoadFailed; + } + + rc = cuModuleGetFunction(function, module, entry_name.c_str()); + if (rc != CUDA_SUCCESS) { + if (error != nullptr) *error = cu_error_string(rc); + cuModuleUnload(module); + return JitKernelStatus::kLoadFailed; + } + + std::lock_guard lock(g_module_mutex); + auto [it, inserted] = + g_modules.emplace(key, LoadedKernelModule{module, *function}); + if (!inserted) { + cuModuleUnload(module); + *function = it->second.function; + } + return JitKernelStatus::kLaunched; +} + +JitKernelStatus configure_smem(CUfunction function, int dynamic_smem_bytes, + std::string* error) { + CUresult rc = cuFuncSetAttribute( + function, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, + dynamic_smem_bytes); + if (rc != CUDA_SUCCESS) { + if (error != nullptr) *error = cu_error_string(rc); + return JitKernelStatus::kAttributeFailed; + } + return JitKernelStatus::kLaunched; +} + +void update_thread_fast_cache(const FastCacheKey& fast_key, CUfunction function, + int configured_smem_bytes) { + if (fast_key.identity == nullptr) return; + + // Update an existing slot if the key is already present (preserves max + // smem). + for (int i = 0; i < kTlCacheSlots; ++i) { + ThreadCache::Slot& slot = g_thread_cache.slots[i]; + if (slot.valid && FastCacheKeyEqual{}(slot.key, fast_key)) { + if (slot.entry.configured_smem_bytes > configured_smem_bytes) { + configured_smem_bytes = slot.entry.configured_smem_bytes; + } + slot.entry = FastCacheEntry{function, configured_smem_bytes}; + return; + } + } + + // Insert into the next round-robin eviction slot. + ThreadCache::Slot& slot = g_thread_cache.slots[g_thread_cache.next_evict]; + slot.valid = true; + slot.key = fast_key; + slot.entry = FastCacheEntry{function, configured_smem_bytes}; + g_thread_cache.next_evict = (g_thread_cache.next_evict + 1) % kTlCacheSlots; +} + +void erase_thread_fast_cache(const FastCacheKey& fast_key) { + for (int i = 0; i < kTlCacheSlots; ++i) { + ThreadCache::Slot& slot = g_thread_cache.slots[i]; + if (slot.valid && FastCacheKeyEqual{}(slot.key, fast_key)) { + slot = ThreadCache::Slot{}; + } + } +} + +void update_fast_cache(const FastCacheKey& fast_key, CUfunction function, + int configured_smem_bytes) { + if (fast_key.identity == nullptr) return; + + { + std::unique_lock lock(g_fast_mutex); + FastCacheEntry& entry = g_fast_cache[fast_key]; + entry.function = function; + if (configured_smem_bytes > entry.configured_smem_bytes) { + entry.configured_smem_bytes = configured_smem_bytes; + } + configured_smem_bytes = entry.configured_smem_bytes; + } + update_thread_fast_cache(fast_key, function, configured_smem_bytes); +} + +void erase_fast_cache(const FastCacheKey& fast_key) { + { + std::unique_lock lock(g_fast_mutex); + g_fast_cache.erase(fast_key); + } + erase_thread_fast_cache(fast_key); +} + +CUresult launch_function(CUfunction function, const JitKernelVariant& variant, + void* kernel_param, std::size_t kernel_param_size, + cudaStream_t stream) { + // A zero size means kernel_param points to one fixed-size kernel argument. + // Non-zero size means kernel_param points to a byte buffer with all kernel + // arguments. + const bool single_arg = (kernel_param_size == 0); + void* single_arg_array[] = {kernel_param}; + void* extra[] = {CU_LAUNCH_PARAM_BUFFER_POINTER, kernel_param, + CU_LAUNCH_PARAM_BUFFER_SIZE, &kernel_param_size, + CU_LAUNCH_PARAM_END}; + + // Default path: plain cuLaunchKernel (no cooperative / cluster + // attributes). + const bool needs_cluster = + (variant.cluster_dim_x > 1 || variant.cluster_dim_y > 1 || + variant.cluster_dim_z > 1); + if (!variant.cooperative && !needs_cluster) { + return cuLaunchKernel( + function, variant.num_blocks, 1, 1, variant.block_dim, 1, 1, + static_cast(variant.dynamic_smem_bytes), + reinterpret_cast(stream), + single_arg ? single_arg_array : nullptr, + single_arg ? nullptr : extra); + } + + // Attributes path: use cuLaunchKernelEx with cooperative and/or cluster + // dim. Value-initialize attrs[] so the union/pad bytes are well-defined. + CUlaunchAttribute attrs[2]{}; + unsigned int num_attrs = 0; + if (variant.cooperative) { + attrs[num_attrs].id = CU_LAUNCH_ATTRIBUTE_COOPERATIVE; + attrs[num_attrs].value.cooperative = 1; + ++num_attrs; + } + if (needs_cluster) { + attrs[num_attrs].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION; + attrs[num_attrs].value.clusterDim.x = + static_cast(variant.cluster_dim_x); + attrs[num_attrs].value.clusterDim.y = + static_cast(variant.cluster_dim_y); + attrs[num_attrs].value.clusterDim.z = + static_cast(variant.cluster_dim_z); + ++num_attrs; + } + + CUlaunchConfig config{}; + config.gridDimX = static_cast(variant.num_blocks); + config.gridDimY = 1; + config.gridDimZ = 1; + config.blockDimX = static_cast(variant.block_dim); + config.blockDimY = 1; + config.blockDimZ = 1; + config.sharedMemBytes = + static_cast(variant.dynamic_smem_bytes); + config.hStream = reinterpret_cast(stream); + config.attrs = attrs; + config.numAttrs = num_attrs; + + return cuLaunchKernelEx(&config, function, + single_arg ? single_arg_array : nullptr, + single_arg ? nullptr : extra); +} + +JitKernelStatus try_fast_launch(const FastCacheKey& fast_key, + const JitKernelVariant& variant, + void* kernel_param, + std::size_t kernel_param_size, + cudaStream_t stream, std::string* error) { + if (fast_key.identity == nullptr) { + log_memory_cache_event(fast_key, variant, "skip", "no_identity"); + return JitKernelStatus::kDisabled; + } + + FastCacheEntry entry; + bool tl_hit = false; + for (int i = 0; i < kTlCacheSlots; ++i) { + const ThreadCache::Slot& slot = g_thread_cache.slots[i]; + if (slot.valid && FastCacheKeyEqual{}(slot.key, fast_key)) { + entry = slot.entry; + tl_hit = true; + break; + } + } + if (!tl_hit) { + { + std::shared_lock lock(g_fast_mutex); + const auto it = g_fast_cache.find(fast_key); + if (it == g_fast_cache.end()) { + log_memory_cache_event(fast_key, variant, "miss", "process"); + return JitKernelStatus::kDisabled; + } + entry = it->second; + } + update_thread_fast_cache(fast_key, entry.function, + entry.configured_smem_bytes); + log_memory_cache_event(fast_key, variant, "hit", "process"); + } else { + log_memory_cache_event(fast_key, variant, "hit", "thread"); + } + + if (variant.dynamic_smem_bytes > entry.configured_smem_bytes) { + const JitKernelStatus status = + configure_smem(entry.function, variant.dynamic_smem_bytes, error); + if (status != JitKernelStatus::kLaunched) { + erase_fast_cache(fast_key); + return JitKernelStatus::kDisabled; + } + update_fast_cache(fast_key, entry.function, variant.dynamic_smem_bytes); + } + + CUresult rc = launch_function(entry.function, variant, kernel_param, + kernel_param_size, stream); + if (rc != CUDA_SUCCESS) { + if (error != nullptr) *error = cu_error_string(rc); + return JitKernelStatus::kLaunchFailed; + } + return JitKernelStatus::kLaunched; +} + +struct CachePaths { + std::filesystem::path variant_dir; + std::filesystem::path cubin_path; + std::filesystem::path metadata_path; + std::filesystem::path log_path; + std::filesystem::path source_path; + std::filesystem::path failed_path; + std::filesystem::path tmp_cubin_path; + std::filesystem::path lock_path; +}; + +CachePaths make_cache_paths(const JitCache& cache, const std::string& key) { + CachePaths paths; + paths.variant_dir = cache.root_dir() / key; + paths.cubin_path = paths.variant_dir / "kernel.cubin"; + paths.metadata_path = paths.variant_dir / "metadata.json"; + paths.log_path = paths.variant_dir / "compile.log"; + paths.source_path = paths.variant_dir / "kernel.cu"; + paths.failed_path = paths.variant_dir / "compile.failed"; + paths.tmp_cubin_path = temporary_cubin_path(paths.cubin_path); + paths.lock_path = paths.variant_dir / "compile.lock"; + return paths; +} + +JitKernelStatus ensure_cubin(const JitKernelVariant& variant, + const SmEnvInfo& env_info, const CachePaths& paths, + const std::string& key, + const std::string& log_prefix, int sm, + std::chrono::steady_clock::time_point jit_begin, + std::string* cubin, std::string* error) { + JitCache& cache = JitCache::instance(); + if (cubin == nullptr) return JitKernelStatus::kCompileFailed; + + *cubin = read_file_or_empty(paths.cubin_path); + if (cubin->empty()) { + const std::string cached_failure = + read_file_or_empty(paths.failed_path); + if (!cached_failure.empty()) { + cache.mark_failed(key); + if (error != nullptr) *error = cached_failure; + log_jit_event(variant, + log_prefix + " disk_cache=failed_marker key=" + key + + " variant=" + std::string(variant.variant_name) + + " marker=" + paths.failed_path.string() + + " jit_overhead_sec=" + + std::to_string(elapsed_sec(jit_begin))); + return JitKernelStatus::kCompileFailed; + } + } + if (!cubin->empty()) { + log_jit_event( + variant, + log_prefix + " disk_cache=hit key=" + key + + " variant=" + std::string(variant.variant_name) + + " cubin=" + paths.cubin_path.string() + + " jit_overhead_sec=" + std::to_string(elapsed_sec(jit_begin))); + return JitKernelStatus::kLaunched; + } + + log_jit_event(variant, log_prefix + " disk_cache=miss key=" + key + + " variant=" + std::string(variant.variant_name) + + " cubin=" + paths.cubin_path.string()); + + ScopedFileLock file_lock; + std::string lock_error; + log_jit_event(variant, log_prefix + " compile_lock=waiting key=" + key + + " path=" + paths.lock_path.string()); + if (!file_lock.lock(paths.lock_path, &lock_error)) { + cache.mark_failed(key); + if (error != nullptr) *error = lock_error; + warn_once_jit_event( + cache, key, variant, + log_prefix + " compile_lock=failed key=" + key + " variant=" + + std::string(variant.variant_name) + "; using static path"); + return JitKernelStatus::kCompileFailed; + } + log_jit_event(variant, log_prefix + " compile_lock=acquired key=" + key + + " path=" + paths.lock_path.string()); + + std::lock_guard compile_lock(g_compile_mutex); + *cubin = read_file_or_empty(paths.cubin_path); + if (!cubin->empty()) { + log_jit_event( + variant, + log_prefix + " disk_cache=hit_after_lock key=" + key + + " variant=" + std::string(variant.variant_name) + + " cubin=" + paths.cubin_path.string() + + " jit_overhead_sec=" + std::to_string(elapsed_sec(jit_begin))); + return JitKernelStatus::kLaunched; + } + + const std::string cached_failure = read_file_or_empty(paths.failed_path); + if (!cached_failure.empty()) { + cache.mark_failed(key); + if (error != nullptr) *error = cached_failure; + log_jit_event( + variant, + log_prefix + " disk_cache=failed_marker_after_lock key=" + key + + " variant=" + std::string(variant.variant_name) + + " marker=" + paths.failed_path.string() + + " jit_overhead_sec=" + std::to_string(elapsed_sec(jit_begin))); + return JitKernelStatus::kCompileFailed; + } + + JitCompileInput input; + input.source = std::string(variant.source); + input.options = env_info.options; + input.source_path = paths.source_path; + input.cubin_path = paths.tmp_cubin_path; + input.log_enabled = jit_log_enabled(); + + JitCompileOutput compile_output; + const JitCompiler& compiler = compiler_backend(); + if (!compiler.compile_to_cubin(input, &compile_output)) { + cache.mark_failed(key); + std::error_code remove_ec; + std::filesystem::remove(paths.tmp_cubin_path, remove_ec); + const double jit_overhead_sec = elapsed_sec(jit_begin); + compile_output.log += + "jit_overhead_sec=" + std::to_string(jit_overhead_sec) + "\n"; + if (!write_file_atomic(paths.log_path, compile_output.log, false)) { + log_jit_event(variant, log_prefix + + " artifact=compile_log write=failed" + " path=" + + paths.log_path.string()); + } + if (!write_file_atomic(paths.failed_path, compile_output.log, false)) { + log_jit_event(variant, log_prefix + + " artifact=failed_marker write=failed" + " path=" + + paths.failed_path.string()); + } + warn_once_jit_event( + cache, key, variant, + log_prefix + " compile=failed key=" + key + + " variant=" + std::string(variant.variant_name) + + " log=" + paths.log_path.string() + " jit_overhead_sec=" + + std::to_string(jit_overhead_sec) + " nvcc_process_sec=" + + std::to_string(compile_output.compiler_sec) + + "; using static path"); + if (error != nullptr) *error = compile_output.log; + return JitKernelStatus::kCompileFailed; + } + + *cubin = compile_output.cubin; + std::error_code remove_ec; + std::filesystem::remove(paths.tmp_cubin_path, remove_ec); + if (write_file_atomic(paths.cubin_path, *cubin, true)) { + log_jit_event(variant, log_prefix + + " disk_cache=write artifact=cubin key=" + + key + " path=" + paths.cubin_path.string() + + " bytes=" + std::to_string(cubin->size())); + } else { + log_jit_event(variant, log_prefix + + " artifact=cubin write=failed path=" + + paths.cubin_path.string()); + } + const double jit_overhead_sec = elapsed_sec(jit_begin); + compile_output.log += + "jit_overhead_sec=" + std::to_string(jit_overhead_sec) + "\n"; + if (!write_file_atomic(paths.log_path, compile_output.log, false)) { + log_jit_event(variant, log_prefix + + " artifact=compile_log write=failed path=" + + paths.log_path.string()); + } + std::filesystem::remove(paths.failed_path, remove_ec); + if (!write_file_atomic(paths.metadata_path, + metadata_json(variant, key, sm, env_info.options, + env_info.compiler_id, jit_overhead_sec, + compile_output.compiler_sec), + false)) { + log_jit_event(variant, log_prefix + + " artifact=metadata write=failed path=" + + paths.metadata_path.string()); + } + log_jit_event( + variant, + log_prefix + " compile=succeeded key=" + key + + " variant=" + std::string(variant.variant_name) + + " cubin=" + paths.cubin_path.string() + + " jit_overhead_sec=" + std::to_string(jit_overhead_sec) + + " nvcc_process_sec=" + std::to_string(compile_output.compiler_sec)); + return JitKernelStatus::kLaunched; +} + +JitKernelStatus load_and_launch_kernel( + const JitKernelVariant& variant, const FastCacheKey& fast_key, + const CachePaths& paths, const std::string& key, const std::string& cubin, + const std::string& log_prefix, void* kernel_param, + std::size_t kernel_param_size, cudaStream_t stream, std::string* error) { + JitCache& cache = JitCache::instance(); + const std::string module_key = key + + ":ctx=" + context_key(fast_key.context) + + ":device=" + std::to_string(fast_key.device); + + CUfunction function = nullptr; + JitKernelStatus status = load_function( + module_key, cubin, std::string(variant.entry_name), &function, error); + if (status != JitKernelStatus::kLaunched) { + cache.mark_failed(key); + warn_once_jit_event( + cache, key, variant, + log_prefix + " module_load=failed key=" + key + " variant=" + + std::string(variant.variant_name) + "; using static path"); + return status; + } + + status = configure_smem(function, variant.dynamic_smem_bytes, error); + if (status != JitKernelStatus::kLaunched) { + cache.mark_failed(key); + warn_once_jit_event( + cache, key, variant, + log_prefix + " smem_attribute=failed key=" + key + " variant=" + + std::string(variant.variant_name) + "; using static path"); + return status; + } + + update_fast_cache(fast_key, function, variant.dynamic_smem_bytes); + + CUresult rc = launch_function(function, variant, kernel_param, + kernel_param_size, stream); + if (rc != CUDA_SUCCESS) { + if (error != nullptr) *error = cu_error_string(rc); + warn_once_jit_event( + cache, key, variant, + log_prefix + " launch=failed key=" + key + + " variant=" + std::string(variant.variant_name)); + return JitKernelStatus::kLaunchFailed; + } + + log_jit_event(variant, + log_prefix + " launch=succeeded key=" + key + + " variant=" + std::string(variant.variant_name) + + " module_cache=ready cubin=" + paths.cubin_path.string()); + return JitKernelStatus::kLaunched; +} + +} // namespace + +const char* jit_kernel_status_name(JitKernelStatus status) { + switch (status) { + case JitKernelStatus::kLaunched: + return "launched"; + case JitKernelStatus::kDisabled: + return "disabled"; + case JitKernelStatus::kUnsupportedDevice: + return "unsupported device"; + case JitKernelStatus::kCompileFailed: + return "compile failed"; + case JitKernelStatus::kLoadFailed: + return "load failed"; + case JitKernelStatus::kAttributeFailed: + return "attribute failed"; + case JitKernelStatus::kLaunchFailed: + return "launch failed"; + } + return "unknown"; +} + +namespace { + +JitKernelStatus resolve_launch_target(int* device, CUcontext* context, + std::string* error) { + if (cudaGetDevice(device) != cudaSuccess) { + if (error != nullptr) *error = "cudaGetDevice failed"; + return JitKernelStatus::kUnsupportedDevice; + } + + CUresult rc = cuCtxGetCurrent(context); + if (rc == CUDA_SUCCESS && *context == nullptr) { + cudaFree(nullptr); + rc = cuCtxGetCurrent(context); + } + if (rc != CUDA_SUCCESS || *context == nullptr) { + if (error != nullptr) + *error = (rc == CUDA_SUCCESS) ? "no current CUDA context" + : cu_error_string(rc); + return JitKernelStatus::kUnsupportedDevice; + } + return JitKernelStatus::kLaunched; +} + +} // namespace + +JitKernelStatus launch_jit_kernel_cached(const JitKernelVariant& variant, + void* kernel_param, + std::size_t kernel_param_size, + cudaStream_t stream, + std::string* error) { + int device = 0; + CUcontext context = nullptr; + const JitKernelStatus target_status = + resolve_launch_target(&device, &context, error); + if (target_status != JitKernelStatus::kLaunched) return target_status; + + const FastCacheKey fast_key{variant.identity, variant.runtime_key, context, + device}; + return try_fast_launch(fast_key, variant, kernel_param, kernel_param_size, + stream, error); +} + +JitKernelStatus launch_jit_kernel(const JitKernelVariant& variant, + void* kernel_param, + std::size_t kernel_param_size, + cudaStream_t stream, std::string* error) { + int device = 0; + CUcontext context = nullptr; + const JitKernelStatus target_status = + resolve_launch_target(&device, &context, error); + if (target_status != JitKernelStatus::kLaunched) return target_status; + + const FastCacheKey fast_key{variant.identity, variant.runtime_key, context, + device}; + const JitKernelStatus fast_status = try_fast_launch( + fast_key, variant, kernel_param, kernel_param_size, stream, error); + if (fast_status == JitKernelStatus::kLaunched || + fast_status == JitKernelStatus::kLaunchFailed) { + return fast_status; + } + + const auto jit_begin = std::chrono::steady_clock::now(); + const std::string log_prefix = kernel_log_prefix(variant); + + cudaDeviceProp prop{}; + if (cudaGetDeviceProperties(&prop, device) != cudaSuccess) { + if (error != nullptr) *error = "cudaGetDeviceProperties failed"; + return JitKernelStatus::kUnsupportedDevice; + } + + const int sm = prop.major * 10 + prop.minor; + if (sm < variant.min_sm) { + if (error != nullptr) { + *error = log_prefix + " requires sm_" + + std::to_string(variant.min_sm) + " or newer"; + } + return JitKernelStatus::kUnsupportedDevice; + } + + const SmEnvInfo& env_info = get_sm_env_info(sm, variant.target_arch); + const std::string key = source_fingerprint(variant, sm, env_info.env_hash); + JitCache& cache = JitCache::instance(); + if (cache.has_failed(key)) return JitKernelStatus::kCompileFailed; + const CachePaths paths = make_cache_paths(cache, key); + + log_jit_event(variant, + log_prefix + + " request variant=" + std::string(variant.variant_name) + + " key=" + key + " device=" + std::to_string(device) + + " sm=" + std::to_string(sm) + + " target_arch=" + std::string(variant.target_arch) + + " cache_root=" + cache.root_dir().string() + + " source_dir=" + env_info.source_dir.string()); + + std::string cubin; + JitKernelStatus status = + ensure_cubin(variant, env_info, paths, key, log_prefix, sm, jit_begin, + &cubin, error); + if (status != JitKernelStatus::kLaunched) return status; + + return load_and_launch_kernel(variant, fast_key, paths, key, cubin, + log_prefix, kernel_param, kernel_param_size, + stream, error); +} + +// This overload is for kernels that don't need to pass kernel arguments. +JitKernelStatus launch_jit_kernel(const JitKernelVariant& variant, + void* kernel_param, cudaStream_t stream, + std::string* error) { + return launch_jit_kernel(variant, kernel_param, 0, stream, error); +} + +} // namespace jit +} // namespace mooncake diff --git a/mooncake-ep/src/jit/jit_utils.cpp b/mooncake-ep/src/jit/jit_utils.cpp new file mode 100644 index 0000000000..1c0ee7109a --- /dev/null +++ b/mooncake-ep/src/jit/jit_utils.cpp @@ -0,0 +1,229 @@ +/* + * 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.cc. + */ + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace mooncake { +namespace jit { + +std::string env_value(const char* name) { + const char* value = std::getenv(name); + return (value != nullptr && value[0] != '\0') ? value : std::string(); +} + +bool env_flag_enabled(const char* name, bool default_value) { + std::string value = env_value(name); + if (value.empty()) return default_value; + return value[0] != '0' && value != "false" && value != "FALSE" && + value != "off" && value != "OFF"; +} + +bool announce_once(const std::string& key) { + // Returns true the first time each distinct `key` is seen. + static std::mutex mtx; + static std::set announced; + std::lock_guard lock(mtx); + return announced.insert(key).second; +} + +bool jit_log_enabled() { + static const bool enabled = env_flag_enabled("MOONCAKE_EP_JIT_LOG"); + return enabled; +} + +void jit_log(std::string_view message) { + if (!jit_log_enabled()) return; + std::fprintf(stderr, "[mooncake_ep jit] pid=%ld %.*s\n", + static_cast(getpid()), static_cast(message.size()), + message.data()); +} + +ScopedFileLock::~ScopedFileLock() { + if (fd_ >= 0) { + flock(fd_, LOCK_UN); + close(fd_); + } +} + +bool ScopedFileLock::lock(const std::filesystem::path& path, + std::string* error) { + std::error_code ec; + std::filesystem::create_directories(path.parent_path(), ec); + if (ec) { + if (error != nullptr) + *error = + "mkdir " + path.parent_path().string() + ": " + ec.message(); + return false; + } + + fd_ = open(path.c_str(), O_CREAT | O_RDWR | O_CLOEXEC, 0666); + if (fd_ < 0) { + if (error != nullptr) + *error = "open " + path.string() + ": " + std::strerror(errno); + return false; + } + if (flock(fd_, LOCK_EX) != 0) { + if (error != nullptr) + *error = "flock " + path.string() + ": " + std::strerror(errno); + close(fd_); + fd_ = -1; + return false; + } + return true; +} + +std::vector split_env_flags(const char* flags) { + std::vector out; + if (flags == nullptr || flags[0] == '\0') return out; + + std::string current; + char quote = '\0'; + bool escape = false; + for (const char* p = flags; *p != '\0'; ++p) { + const char c = *p; + if (escape) { + current.push_back(c); + escape = false; + continue; + } + if (c == '\\') { + escape = true; + continue; + } + if (quote != '\0') { + if (c == quote) + quote = '\0'; + else + current.push_back(c); + continue; + } + if (c == '\'' || c == '"') { + quote = c; + continue; + } + if (std::isspace(static_cast(c))) { + if (!current.empty()) { + out.push_back(current); + current.clear(); + } + continue; + } + current.push_back(c); + } + if (escape) current.push_back('\\'); + if (!current.empty()) out.push_back(current); + return out; +} + +std::string json_escape(std::string_view text) { + std::ostringstream out; + for (char c : text) { + switch (c) { + case '\\': + out << "\\\\"; + break; + case '"': + out << "\\\""; + break; + case '\n': + out << "\\n"; + break; + case '\r': + out << "\\r"; + break; + case '\t': + out << "\\t"; + break; + default: + out << c; + break; + } + } + return out.str(); +} + +std::string read_file_or_empty(const std::filesystem::path& path) { + std::ifstream in(path, std::ios::binary); + if (!in) return {}; + return std::string(std::istreambuf_iterator(in), + std::istreambuf_iterator()); +} + +bool write_file_atomic(const std::filesystem::path& path, + const std::string& data, bool binary) { + std::error_code ec; + std::filesystem::create_directories(path.parent_path(), ec); + if (ec) { + jit_log("cache write mkdir failed for " + path.parent_path().string() + + ": " + ec.message()); + return false; + } + + std::ostringstream tmp_name; + tmp_name << path.string() << ".tmp." << static_cast(getpid()) << "." + << std::hash{}(std::this_thread::get_id()); + const std::filesystem::path tmp_path = tmp_name.str(); + const auto mode = + binary ? (std::ios::binary | std::ios::trunc) : std::ios::trunc; + { + std::ofstream out(tmp_path, mode); + if (!out) { + jit_log("cache write open failed for " + tmp_path.string()); + return false; + } + out.write(data.data(), static_cast(data.size())); + if (!out) { + jit_log("cache write failed for " + tmp_path.string()); + return false; + } + } + + std::filesystem::rename(tmp_path, path, ec); + if (!ec) return true; + + jit_log("cache rename failed for " + tmp_path.string() + " -> " + + path.string() + ": " + ec.message()); + std::filesystem::remove(tmp_path, ec); + return false; +} + +std::string fnv1a_digest(const std::vector& parts) { + uint64_t hash = 1469598103934665603ull; + for (const std::string& part : parts) { + for (unsigned char c : part) { + hash ^= static_cast(c); + hash *= 1099511628211ull; + } + hash ^= 0xffu; + hash *= 1099511628211ull; + } + + std::ostringstream out; + out << std::hex << hash; + return out.str(); +} + +} // namespace jit +} // namespace mooncake diff --git a/mooncake-ep/src/jit/nvcc_compiler.cpp b/mooncake-ep/src/jit/nvcc_compiler.cpp new file mode 100644 index 0000000000..355791d25e --- /dev/null +++ b/mooncake-ep/src/jit/nvcc_compiler.cpp @@ -0,0 +1,316 @@ +/* + * 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.cc. + */ + +#include + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace mooncake { +namespace jit { +namespace { + +struct ProcessResult { + int exit_code = -1; + std::string output; +}; + +std::string resolve_nvcc() { + std::string nvcc = env_value("MOONCAKE_EP_JIT_NVCC"); + if (!nvcc.empty()) return nvcc; + + nvcc = env_value("NVCC"); + if (!nvcc.empty()) return nvcc; + + std::string cuda_home = env_value("CUDA_HOME"); + if (cuda_home.empty()) cuda_home = env_value("CUDA_PATH"); + if (!cuda_home.empty()) { + const std::filesystem::path cuda_nvcc = + std::filesystem::path(cuda_home) / "bin" / "nvcc"; + std::error_code ec; + if (std::filesystem::exists(cuda_nvcc, ec)) return cuda_nvcc.string(); + } + + return "nvcc"; +} + +std::vector include_options(const JitCompileConfig& config) { + std::vector options; + const auto add_include = [&options](const std::filesystem::path& path) { + if (!path.empty()) options.push_back("-I" + path.string()); + }; + + // Staged mooncake-ep kernel headers are included by their flat name + // (e.g. ); the parent root additionally allows + // the spelling. + add_include(config.source_dir); + add_include(config.source_dir.parent_path()); + + if (!config.build_include_dir.empty()) { + add_include(config.build_include_dir); + } + + if (!config.cuda_include_dir.empty()) { + add_include(config.cuda_include_dir); + add_include(config.cuda_include_dir / "cccl"); + } + + return options; +} + +std::string command_to_string(const std::vector& argv) { + std::ostringstream out; + for (size_t i = 0; i < argv.size(); ++i) { + if (i != 0) out << ' '; + const bool needs_quotes = + argv[i].find_first_of(" \t\n\"'\\") != std::string::npos; + if (!needs_quotes) { + out << argv[i]; + continue; + } + + out << '\''; + for (char c : argv[i]) { + if (c == '\'') + out << "'\\''"; + else + out << c; + } + out << '\''; + } + return out.str(); +} + +bool run_process(const std::vector& argv, ProcessResult* result) { + if (result == nullptr) return false; + result->exit_code = -1; + result->output.clear(); + if (argv.empty()) { + result->output = "empty command"; + return false; + } + + int pipe_fd[2] = {-1, -1}; + if (pipe(pipe_fd) != 0) { + result->output = std::string("pipe failed: ") + std::strerror(errno); + return false; + } + + std::vector raw_argv; + raw_argv.reserve(argv.size() + 1); + for (const std::string& arg : argv) + raw_argv.push_back(const_cast(arg.c_str())); + raw_argv.push_back(nullptr); + + const pid_t pid = fork(); + if (pid < 0) { + const int saved_errno = errno; + close(pipe_fd[0]); + close(pipe_fd[1]); + result->output = + std::string("fork failed: ") + std::strerror(saved_errno); + return false; + } + + if (pid == 0) { + close(pipe_fd[0]); + dup2(pipe_fd[1], STDOUT_FILENO); + dup2(pipe_fd[1], STDERR_FILENO); + close(pipe_fd[1]); + + execvp(raw_argv[0], raw_argv.data()); + _exit(127); + } + + close(pipe_fd[1]); + char buffer[4096]; + while (true) { + const ssize_t bytes_read = read(pipe_fd[0], buffer, sizeof(buffer)); + if (bytes_read > 0) { + result->output.append(buffer, static_cast(bytes_read)); + continue; + } + if (bytes_read == 0) break; + if (errno == EINTR) continue; + result->output += std::string("\nread failed: ") + std::strerror(errno); + break; + } + close(pipe_fd[0]); + + int status = 0; + while (waitpid(pid, &status, 0) < 0) { + if (errno == EINTR) continue; + result->output += + std::string("\nwaitpid failed: ") + std::strerror(errno); + return false; + } + + if (WIFEXITED(status)) { + result->exit_code = WEXITSTATUS(status); + } else if (WIFSIGNALED(status)) { + result->exit_code = 128 + WTERMSIG(status); + } + return result->exit_code == 0; +} + +void compiler_log(const JitCompileInput& input, const std::string& message) { + if (!input.log_enabled) return; + jit_log(message); +} + +} // namespace + +std::string NvccCompiler::compiler_id() const { + static const std::string id = [] { + const std::string nvcc = resolve_nvcc(); + ProcessResult result; + const std::vector argv = {nvcc, "--version"}; + if (run_process(argv, &result)) { + return "nvcc=" + nvcc + "\n" + result.output; + } + + std::ostringstream out; + out << "nvcc=" << nvcc << "\nversion unavailable"; + if (result.exit_code >= 0) out << " exit_code=" << result.exit_code; + if (!result.output.empty()) out << "\n" << result.output; + return out.str(); + }(); + return id; +} + +std::vector NvccCompiler::compile_options( + const JitCompileConfig& config) const { + std::vector options; + options.push_back("--std=c++17"); + options.push_back("--extended-lambda"); + options.push_back("--expt-relaxed-constexpr"); + + if (!config.target_arch.empty()) { + options.push_back("-arch=" + config.target_arch); + } else { + const std::string arch_flags = env_value("NVCC_ARCH_FLAGS"); + std::vector arch_options = + split_env_flags(arch_flags.c_str()); + if (arch_options.empty()) { + options.push_back("-arch=sm_" + std::to_string(config.sm)); + } else { + options.insert(options.end(), arch_options.begin(), + arch_options.end()); + } + } + // Macros that change host/device struct layout must be forwarded here or + // the JIT-compiled kernel silently disagrees with the host-packed args. +#ifdef NDEBUG + options.push_back("-DNDEBUG=1"); +#endif + + const std::vector includes = include_options(config); + options.insert(options.end(), includes.begin(), includes.end()); + + const std::string extra_flags = env_value("NVCC_EXTRA_FLAGS"); + std::vector extra_options = + split_env_flags(extra_flags.c_str()); + options.insert(options.end(), extra_options.begin(), extra_options.end()); + return options; +} + +bool NvccCompiler::compile_to_cubin(const JitCompileInput& input, + JitCompileOutput* output) const { + if (output == nullptr) return false; + output->cubin.clear(); + output->log.clear(); + output->compiler_sec = 0.0; + + if (input.source_path.empty() || input.cubin_path.empty()) { + output->log = + "NVCC compile failed: source_path and cubin_path must be set"; + return false; + } + + if (!write_file_atomic(input.source_path, input.source, false)) { + output->log = "NVCC compile failed: could not write source " + + input.source_path.string(); + return false; + } + + std::error_code ec; + std::filesystem::create_directories(input.cubin_path.parent_path(), ec); + if (ec) { + output->log = "NVCC compile failed: mkdir " + + input.cubin_path.parent_path().string() + ": " + + ec.message(); + return false; + } + + std::filesystem::remove(input.cubin_path, ec); + + std::vector argv; + argv.reserve(input.options.size() + 5); + argv.push_back(resolve_nvcc()); + argv.push_back("--cubin"); + argv.push_back(input.source_path.string()); + argv.push_back("-o"); + argv.push_back(input.cubin_path.string()); + argv.insert(argv.end(), input.options.begin(), input.options.end()); + + const std::string command = command_to_string(argv); + compiler_log(input, + "compiler=nvcc action=compile_begin command=" + command); + + ProcessResult result; + const auto compile_begin = std::chrono::steady_clock::now(); + const bool ok = run_process(argv, &result); + const auto compile_end = std::chrono::steady_clock::now(); + output->compiler_sec = + std::chrono::duration(compile_end - compile_begin).count(); + + std::ostringstream log; + log << "$ " << command << "\n"; + if (!result.output.empty()) log << result.output; + if (result.output.empty() || result.output.back() != '\n') log << "\n"; + log << "exit_code=" << result.exit_code << "\n"; + log << "nvcc_process_sec=" << output->compiler_sec << "\n"; + output->log = log.str(); + + if (!ok) { + compiler_log(input, "compiler=nvcc action=compile_failed exit_code=" + + std::to_string(result.exit_code) + + " nvcc_process_sec=" + + std::to_string(output->compiler_sec)); + return false; + } + + output->cubin = read_file_or_empty(input.cubin_path); + if (output->cubin.empty()) { + output->log += "NVCC compile failed: cubin was not produced at " + + input.cubin_path.string() + "\n"; + compiler_log(input, + "compiler=nvcc action=compile_failed reason=missing_cubin" + " path=" + + input.cubin_path.string()); + return false; + } + + compiler_log( + input, "compiler=nvcc action=compile_succeeded cubin_bytes=" + + std::to_string(output->cubin.size()) + + " nvcc_process_sec=" + std::to_string(output->compiler_sec)); + return true; +} + +} // namespace jit +} // namespace mooncake From 736f642ca7fce33313be1beaae33f191c8d6822c Mon Sep 17 00:00:00 2001 From: Bojun Zhu Date: Tue, 29 Sep 2026 14:39:11 +0800 Subject: [PATCH 02/17] [JIT] Generate jit_config.h for JIT include paths Same approach as nccl-ep, with one change: jit_runtime.cpp includes jit_config.h directly, instead of the compiler adding it to every file, so readers can see where the path macros come from. Signed-off-by: Bojun Zhu --- mooncake-ep/src/CMakeLists.txt | 26 +++++++++++++++++++++++++- mooncake-ep/src/jit/jit_config.h.in | 17 +++++++++++++++++ mooncake-ep/src/jit/jit_runtime.cpp | 16 +--------------- 3 files changed, 43 insertions(+), 16 deletions(-) create mode 100644 mooncake-ep/src/jit/jit_config.h.in diff --git a/mooncake-ep/src/CMakeLists.txt b/mooncake-ep/src/CMakeLists.txt index 8897aca781..65b03ec42b 100644 --- a/mooncake-ep/src/CMakeLists.txt +++ b/mooncake-ep/src/CMakeLists.txt @@ -58,6 +58,29 @@ if(USE_CUDA) "${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. Bake their paths into the + # engine via jit_config.h. + set(MOONCAKE_EP_JIT_SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/../include") + # jit_config.h is only used to build this library, so keep it in the binary + # dir rather than next to any headers that get installed. + 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() + + 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} @@ -74,7 +97,8 @@ 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> diff --git a/mooncake-ep/src/jit/jit_config.h.in b/mooncake-ep/src/jit/jit_config.h.in new file mode 100644 index 0000000000..8177e2d253 --- /dev/null +++ b/mooncake-ep/src/jit/jit_config.h.in @@ -0,0 +1,17 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Generated from jit_config.h.in by CMake -- do not edit the generated copy. + * + * The JIT engine shells out to nvcc at runtime, so the kernel headers that the + * generated source #includes must be resolvable as files on disk. These are + * the paths CMake staged them to. + */ + +#pragma once + +#define MOONCAKE_EP_JIT_SOURCE_DIR "@MOONCAKE_EP_JIT_SOURCE_DIR@" +#define MOONCAKE_EP_JIT_BUILD_INCLUDE_DIR "@MOONCAKE_EP_JIT_BUILD_INCLUDE_DIR@" +#define MOONCAKE_EP_JIT_CUDA_INCLUDE_DIR "@MOONCAKE_EP_JIT_CUDA_INCLUDE_DIR@" diff --git a/mooncake-ep/src/jit/jit_runtime.cpp b/mooncake-ep/src/jit/jit_runtime.cpp index 0acebfc2e2..ccbf27d19e 100644 --- a/mooncake-ep/src/jit/jit_runtime.cpp +++ b/mooncake-ep/src/jit/jit_runtime.cpp @@ -9,6 +9,7 @@ #include #include +#include #include #include @@ -29,21 +30,6 @@ #include #include -#ifndef MOONCAKE_EP_JIT_SOURCE_DIR -#define MOONCAKE_EP_JIT_SOURCE_DIR "/usr/include/mooncake_ep" -#endif - -// Left empty by default: fingerprinting an unrelated system include root on -// every process would be expensive and meaningless. Builds that need extra -// headers (e.g. NCCL device headers) point this at the right tree. -#ifndef MOONCAKE_EP_JIT_BUILD_INCLUDE_DIR -#define MOONCAKE_EP_JIT_BUILD_INCLUDE_DIR "" -#endif - -#ifndef MOONCAKE_EP_JIT_CUDA_INCLUDE_DIR -#define MOONCAKE_EP_JIT_CUDA_INCLUDE_DIR "" -#endif - namespace mooncake { namespace jit { namespace { From 3c7f03acd6d0cc52153a49351a1328295c786616 Mon Sep 17 00:00:00 2001 From: Bojun Zhu Date: Wed, 30 Sep 2026 17:21:17 +0800 Subject: [PATCH 03/17] [JIT] Add transfer engine into the include path --- mooncake-ep/src/CMakeLists.txt | 47 ++++++++++++++++++++++++--- mooncake-ep/src/jit/jit_runtime.cpp | 4 ++- mooncake-ep/src/jit/nvcc_compiler.cpp | 4 +++ 3 files changed, 49 insertions(+), 6 deletions(-) diff --git a/mooncake-ep/src/CMakeLists.txt b/mooncake-ep/src/CMakeLists.txt index 65b03ec42b..4cafdf1b39 100644 --- a/mooncake-ep/src/CMakeLists.txt +++ b/mooncake-ep/src/CMakeLists.txt @@ -59,11 +59,19 @@ if(USE_CUDA) "${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. Bake their paths into the - # engine via jit_config.h. - set(MOONCAKE_EP_JIT_SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/../include") - # jit_config.h is only used to build this library, so keep it in the binary - # dir rather than next to any headers that get installed. + # 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 "") @@ -78,6 +86,35 @@ if(USE_CUDA) 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_utils.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/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) diff --git a/mooncake-ep/src/jit/jit_runtime.cpp b/mooncake-ep/src/jit/jit_runtime.cpp index ccbf27d19e..f1e7267d4c 100644 --- a/mooncake-ep/src/jit/jit_runtime.cpp +++ b/mooncake-ep/src/jit/jit_runtime.cpp @@ -256,7 +256,9 @@ const SmEnvInfo& get_sm_env_info(int sm, std::string_view target_arch) { std::vector env_parts = { info.compiler_id, }; - append_header_tree_fingerprint(&env_parts, info.source_dir); + append_header_tree_fingerprint(&env_parts, info.source_dir.parent_path()); + append_header_tree_fingerprint( + &env_parts, info.source_dir.parent_path().parent_path() / "tent"); append_header_tree_fingerprint(&env_parts, info.build_include_dir); env_parts.insert(env_parts.end(), info.options.begin(), info.options.end()); info.env_hash = fnv1a_digest(env_parts); diff --git a/mooncake-ep/src/jit/nvcc_compiler.cpp b/mooncake-ep/src/jit/nvcc_compiler.cpp index 355791d25e..df87dfd19b 100644 --- a/mooncake-ep/src/jit/nvcc_compiler.cpp +++ b/mooncake-ep/src/jit/nvcc_compiler.cpp @@ -214,6 +214,10 @@ std::vector NvccCompiler::compile_options( } // Macros that change host/device struct layout must be forwarded here or // the JIT-compiled kernel silently disagrees with the host-packed args. + options.push_back("-DUSE_CUDA=1"); +#ifdef USE_NCCL_DEVICE + options.push_back("-DUSE_NCCL_DEVICE=1"); +#endif #ifdef NDEBUG options.push_back("-DNDEBUG=1"); #endif From 245d843c93f4f6cdb69037cb3bba84537f6ceada Mon Sep 17 00:00:00 2001 From: Bojun Zhu Date: Sun, 25 Oct 2026 21:59:43 +0800 Subject: [PATCH 04/17] [JIT] Package EP JIT kernel headers into the wheel The CUDA EP JIT runs nvcc on the user's machine, so the kernel headers staged under 'build/mooncake_ep_jit/' must be shipped with the wheel, and the runtime must find them next to the installed library. Signed-off-by: Bojun Zhu --- mooncake-ep/src/CMakeLists.txt | 3 ++- mooncake-ep/src/jit/jit_runtime.cpp | 22 ++++++++++++++++++++++ mooncake-integration/CMakeLists.txt | 10 ++++++++++ scripts/build_wheel.sh | 13 +++++++++++++ 4 files changed, 47 insertions(+), 1 deletion(-) diff --git a/mooncake-ep/src/CMakeLists.txt b/mooncake-ep/src/CMakeLists.txt index 4cafdf1b39..c16ca0f20e 100644 --- a/mooncake-ep/src/CMakeLists.txt +++ b/mooncake-ep/src/CMakeLists.txt @@ -142,7 +142,8 @@ if(USE_CUDA) $<$:-Xcompiler=-g0> $<$:--expt-relaxed-constexpr>) target_link_libraries(mooncake_ep_device PRIVATE CUDA::cudart - CUDA::cuda_driver) + 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/jit_runtime.cpp b/mooncake-ep/src/jit/jit_runtime.cpp index f1e7267d4c..483f39f07c 100644 --- a/mooncake-ep/src/jit/jit_runtime.cpp +++ b/mooncake-ep/src/jit/jit_runtime.cpp @@ -28,6 +28,7 @@ #include #include #include +#include #include namespace mooncake { @@ -121,7 +122,18 @@ double elapsed_sec(std::chrono::steady_clock::time_point begin) { .count(); } +// Directory of the library containing this code (libmooncake_ep_device.so). +std::filesystem::path this_lib_dir() { + Dl_info info{}; + if (dladdr(reinterpret_cast(&this_lib_dir), &info) == 0 || + info.dli_fname == nullptr) { + return {}; + } + return std::filesystem::path(info.dli_fname).parent_path(); +} + std::filesystem::path configured_source_dir() { + // From the environment variable, if set. std::string path = env_value("MOONCAKE_EP_JIT_SOURCE_DIR"); if (!path.empty()) return std::filesystem::path(path); @@ -130,6 +142,16 @@ std::filesystem::path configured_source_dir() { return std::filesystem::path(ep_home) / "include" / "mooncake_ep"; } + // From installed wheel: the headers sit next to the library. + const std::filesystem::path lib_dir = this_lib_dir(); + if (!lib_dir.empty()) { + std::filesystem::path installed = + lib_dir / "mooncake_ep_jit" / "include" / "mooncake_ep"; + std::error_code ec; + if (std::filesystem::is_directory(installed, ec)) return installed; + } + + // Development build: the copy staged in the build tree. return std::filesystem::path(MOONCAKE_EP_JIT_SOURCE_DIR); } diff --git a/mooncake-integration/CMakeLists.txt b/mooncake-integration/CMakeLists.txt index 9f077bd7f6..8179e23224 100644 --- a/mooncake-integration/CMakeLists.txt +++ b/mooncake-integration/CMakeLists.txt @@ -315,6 +315,16 @@ if(WITH_EP) COMPONENT python FILES_MATCHING PATTERN "*.so") + + # The CUDA EP JIT runs nvcc on the user's machine, so nvcc needs the kernel + # headers there too. Put them in the wheel, in a mooncake_ep_jit/ directory + # next to libmooncake_ep_device.so. + if(USE_CUDA) + install( + DIRECTORY "${CMAKE_BINARY_DIR}/mooncake_ep_jit/" + DESTINATION "${MOONCAKE_PYTHON_INSTALL_DIR}/mooncake_ep_jit" + COMPONENT python) + endif() endif() install( FILES diff --git a/scripts/build_wheel.sh b/scripts/build_wheel.sh index d0d84a823b..ef368e5d25 100755 --- a/scripts/build_wheel.sh +++ b/scripts/build_wheel.sh @@ -155,6 +155,9 @@ fi # relative path would silently point to the wrong location. CUDA_EP_STAGING_DIR="${BUILD_DIR_ABS}/ep_pg_staging" +# Kernel header tree the CUDA EP JIT compiles against at runtime. +CUDA_EP_JIT_DIR="${BUILD_DIR_ABS}/mooncake_ep_jit" + # CI only: remove build/ to free disk before python -m build (set FREE_BUILD_DIR=1 to enable locally). # If EP/PG .so files were staged inside the build directory, preserve them in a # temporary location so they survive the cleanup. @@ -164,6 +167,11 @@ if [ "$CI" = "true" ] || [ "$FREE_BUILD_DIR" = "1" ]; then CUDA_EP_STAGING_TEMP=$(mktemp -d) cp "$CUDA_EP_STAGING_DIR"/*.so "$CUDA_EP_STAGING_TEMP/" echo "Preserved EP/PG .so files to ${CUDA_EP_STAGING_TEMP} before build-dir cleanup" + if [ -d "$CUDA_EP_JIT_DIR" ]; then + cp -r "$CUDA_EP_JIT_DIR" "$CUDA_EP_STAGING_TEMP/" + CUDA_EP_JIT_DIR="${CUDA_EP_STAGING_TEMP}/mooncake_ep_jit" + echo "Preserved EP JIT headers to ${CUDA_EP_JIT_DIR} before build-dir cleanup" + fi fi echo "Freeing disk space: removing build directory (artifacts already copied)" rm -rf "${BUILD_DIR}/" @@ -609,6 +617,11 @@ if [ -d "$CUDA_EP_STAGING_DIR" ] && ls "$CUDA_EP_STAGING_DIR"/*.so &>/dev/null; cp "$so_file" "$UNPACKED_PKG_DIR/mooncake/$(basename "$so_file")" fi done + # Kernel headers for the EP JIT (runtime nvcc). + if [ -d "$CUDA_EP_JIT_DIR" ]; then + echo " Adding EP JIT header tree" + cp -r "$CUDA_EP_JIT_DIR" "$UNPACKED_PKG_DIR/mooncake/" + fi # The host EP extension is repaired before the device library is # injected. auditwheel therefore only adds its vendored .libs path; # add the package directory as well so the sibling device library is From 14c622d09062de982ec8d6969ebb1aeb4af5cb2e Mon Sep 17 00:00:00 2001 From: Bojun Zhu Date: Tue, 29 Sep 2026 17:53:04 +0800 Subject: [PATCH 05/17] [JIT] Use an invalid value for the 'sm_min' for default The 'sm_min' in mooncake-ep is sm 80. Make the 'sm_min' explicitly configured by upper level. Signed-off-by: Bojun Zhu --- mooncake-ep/include/jit/jit_runtime.hpp | 2 +- mooncake-ep/src/jit/jit_runtime.cpp | 7 +++++++ 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/mooncake-ep/include/jit/jit_runtime.hpp b/mooncake-ep/include/jit/jit_runtime.hpp index c41e36c77e..eebd6c06cc 100644 --- a/mooncake-ep/include/jit/jit_runtime.hpp +++ b/mooncake-ep/include/jit/jit_runtime.hpp @@ -40,7 +40,7 @@ struct JitKernelVariant { int num_blocks = 0; int block_dim = 0; int dynamic_smem_bytes = 0; - int min_sm = 90; + int min_sm = -1; std::string_view target_arch; // Optional launch attributes (default = off). // Cooperative launch enables cg::this_grid().sync() inside the kernel. diff --git a/mooncake-ep/src/jit/jit_runtime.cpp b/mooncake-ep/src/jit/jit_runtime.cpp index 483f39f07c..7e216a5d77 100644 --- a/mooncake-ep/src/jit/jit_runtime.cpp +++ b/mooncake-ep/src/jit/jit_runtime.cpp @@ -940,6 +940,13 @@ JitKernelStatus launch_jit_kernel(const JitKernelVariant& variant, return JitKernelStatus::kUnsupportedDevice; } + if (variant.min_sm < 0) { + if (error != nullptr) { + *error = log_prefix + " min_sm was not set by the caller"; + } + return JitKernelStatus::kUnsupportedDevice; + } + const int sm = prop.major * 10 + prop.minor; if (sm < variant.min_sm) { if (error != nullptr) { From ed9b5b0d990bd3e2abc4b1945226cde32ff728ea Mon Sep 17 00:00:00 2001 From: Bojun Zhu Date: Wed, 30 Sep 2026 09:00:50 +0800 Subject: [PATCH 06/17] [JIT] Set cooperative to true All the kernels in mooncake-ep are launched with cooperative launch. Signed-off-by: Bojun Zhu --- mooncake-ep/include/jit/jit_runtime.hpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mooncake-ep/include/jit/jit_runtime.hpp b/mooncake-ep/include/jit/jit_runtime.hpp index eebd6c06cc..715280ed16 100644 --- a/mooncake-ep/include/jit/jit_runtime.hpp +++ b/mooncake-ep/include/jit/jit_runtime.hpp @@ -42,10 +42,10 @@ struct JitKernelVariant { int dynamic_smem_bytes = 0; int min_sm = -1; std::string_view target_arch; - // Optional launch attributes (default = off). + // 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 = false; + bool cooperative = true; int cluster_dim_x = 1; int cluster_dim_y = 1; int cluster_dim_z = 1; From 55e764487f67b4a8ead4ddce7d84c7f745db0686 Mon Sep 17 00:00:00 2001 From: Bojun Zhu Date: Wed, 30 Sep 2026 17:20:38 +0800 Subject: [PATCH 07/17] [JIT] Add support for dispatch --- mooncake-ep/include/jit/dispatch_jit.h | 16 + mooncake-ep/include/mooncake_ep_api.cuh | 15 +- mooncake-ep/include/mooncake_ep_dispatch.cuh | 576 +++++++++++++++++++ mooncake-ep/src/CMakeLists.txt | 3 + mooncake-ep/src/jit/dispatch_jit.cpp | 85 +++ mooncake-ep/src/mooncake_ep_buffer.cpp | 4 +- mooncake-ep/src/mooncake_ep_kernel.cu | 523 ++++------------- 7 files changed, 812 insertions(+), 410 deletions(-) create mode 100644 mooncake-ep/include/jit/dispatch_jit.h create mode 100644 mooncake-ep/include/mooncake_ep_dispatch.cuh create mode 100644 mooncake-ep/src/jit/dispatch_jit.cpp 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/mooncake_ep_api.cuh b/mooncake-ep/include/mooncake_ep_api.cuh index 32742c097b..cf41c90c0a 100644 --- a/mooncake-ep/include/mooncake_ep_api.cuh +++ b/mooncake-ep/include/mooncake_ep_api.cuh @@ -9,14 +9,13 @@ 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, diff --git a/mooncake-ep/include/mooncake_ep_dispatch.cuh b/mooncake-ep/include/mooncake_ep_dispatch.cuh new file mode 100644 index 0000000000..9c63dce306 --- /dev/null +++ b/mooncake-ep/include/mooncake_ep_dispatch.cuh @@ -0,0 +1,576 @@ +#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; + +__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; +} + +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/src/CMakeLists.txt b/mooncake-ep/src/CMakeLists.txt index c16ca0f20e..2183a74502 100644 --- a/mooncake-ep/src/CMakeLists.txt +++ b/mooncake-ep/src/CMakeLists.txt @@ -53,6 +53,7 @@ if(USE_CUDA) find_package(CUDAToolkit REQUIRED) set(MOONCAKE_EP_JIT_SOURCES + "${CMAKE_CURRENT_SOURCE_DIR}/jit/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" @@ -89,6 +90,8 @@ if(USE_CUDA) set(MOONCAKE_EP_JIT_HEADERS mooncake_ep_configs.cuh mooncake_ep_device.h + mooncake_ep_dispatch.cuh + mooncake_ep_exception.cuh mooncake_ep_utils.cuh) set(MOONCAKE_EP_JIT_TRANSFER_HEADERS 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/mooncake_ep_buffer.cpp b/mooncake-ep/src/mooncake_ep_buffer.cpp index ffd3f42ec5..9ca0493bf9 100644 --- a/mooncake-ep/src/mooncake_ep_buffer.cpp +++ b/mooncake-ep/src/mooncake_ep_buffer.cpp @@ -343,8 +343,8 @@ MooncakeEpBuffer::dispatch( packed_recv_layout_range, packed_recv_count, active_ranks, gdr_buffer, buffer.rdma_send_signal_buffer, buffer.rdma_recv_signal_buffer, buffer.rdma_send_data_buffer, - buffer.rdma_recv_data_buffer, nullptr, nullptr, raddrs_ptr, - rkeys_ptr, qp_devctxs_ptr, nvlink_avail, ipc_ptrs, x, topk_idx, + buffer.rdma_recv_data_buffer, raddrs_ptr, rkeys_ptr, qp_devctxs_ptr, + nvlink_avail, ipc_ptrs, x, topk_idx, next_buffer.rdma_recv_signal_buffer, num_tokens, hidden, num_max_dispatch_tokens_per_rank, num_topk, num_experts, rank, num_ranks, use_fp8, workspace, launch_stream, timeout_ticks, phases, diff --git a/mooncake-ep/src/mooncake_ep_kernel.cu b/mooncake-ep/src/mooncake_ep_kernel.cu index c1fea6e362..d3ac764a69 100644 --- a/mooncake-ep/src/mooncake_ep_kernel.cu +++ b/mooncake-ep/src/mooncake_ep_kernel.cu @@ -3,18 +3,22 @@ #include #include +#include #include #include #include #include +#if !defined(MOONCAKE_EP_USE_MUSA) && !defined(MOONCAKE_EP_USE_MACA) +#include +#endif + namespace mooncake { using mooncake::device::CommCtx; using mooncake::device::make_comm_ctx; using mooncake::device::mc_route_put; using mooncake::device::mc_rdma_put; -using mooncake::device::mc_signal; using mooncake::device::mc_red_add; using mooncake::device::mc_bar_sync; using mooncake::device::mc_grid_sync; @@ -32,404 +36,102 @@ using mooncake::device::mc_fence_noflush; using mooncake::device::mc_fence; using mooncake::device::mc_fence_barrier_fence; -__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; -} - -template -__global__ EP_LAUNCH_BOUNDS(kNumWarpGroups * kNumWarpsPerGroup * 32, 1) void -dispatch(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* 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* 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) { - 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 +__global__ void mark_phase_ack_kernel(void* mxa_buffer, + const int32_t* nvlink_available, + void* const* ipc_peer_ptrs, + int* ack_buffer, int rank, + int num_ranks, int epoch) { 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; + mxa_buffer, nvlink_available, ipc_peer_ptrs, nullptr, nullptr, nullptr, + ack_buffer, ack_buffer, rank, num_ranks, MAX_QP_COUNT); - // 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(); + for (int peer = static_cast(threadIdx.x); peer < num_ranks; + peer += static_cast(blockDim.x)) { + if (peer == rank) { + mc_st_release(ack_buffer + rank, epoch); + } else { + void* dst = mc_route_put(comm_ctx, peer, ack_buffer + rank); + if (dst != nullptr) + mc_st_release(reinterpret_cast(dst), epoch); } -#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 +__global__ void wait_phase_ack_kernel(int* ack_buffer, int rank, int num_ranks, + int epoch, int64_t timeout_ticks) { + for (int peer = static_cast(threadIdx.x); peer < num_ranks; + peer += static_cast(blockDim.x)) { + if (peer == rank) + continue; + + int64_t start_time = static_cast(clock64()); + while (mc_ld_acquire(ack_buffer + peer) < epoch) { + int64_t end_time = static_cast(clock64()); + if (timeout_ticks != -1 && end_time - start_time > timeout_ticks) + return; } + } +} - // 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] ++; - } +__global__ void mark_and_wait_phase_ack_kernel( + void* mxa_buffer, const int32_t* nvlink_available, + void* const* ipc_peer_ptrs, int* ack_buffer, int rank, int num_ranks, + int epoch, int64_t timeout_ticks) { + const CommCtx comm_ctx = make_comm_ctx( + mxa_buffer, nvlink_available, ipc_peer_ptrs, nullptr, nullptr, nullptr, + ack_buffer, ack_buffer, rank, num_ranks, MAX_QP_COUNT); - // 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 - } + for (int peer = static_cast(threadIdx.x); peer < num_ranks; + peer += static_cast(blockDim.x)) { + if (peer == rank) { + mc_st_release(ack_buffer + rank, epoch); + } else { + void* dst = mc_route_put(comm_ctx, peer, ack_buffer + rank); + if (dst != nullptr) + mc_st_release(reinterpret_cast(dst), epoch); } } -#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]; + for (int peer = static_cast(threadIdx.x); peer < num_ranks; + peer += static_cast(blockDim.x)) { + if (peer == rank) + continue; - // 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) { + int64_t start_time = static_cast(clock64()); + while (mc_ld_acquire(ack_buffer + peer) < epoch) { + int64_t end_time = static_cast(clock64()); + if (timeout_ticks != -1 && end_time - start_time > timeout_ticks) + return; } -#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(); +void mark_phase_ack(void* mxa_buffer, const int32_t* nvlink_available, + void* const* ipc_peer_ptrs, int* ack_buffer, int rank, + int num_ranks, int epoch, cudaStream_t stream) { + SETUP_LAUNCH_CONFIG(1, 32, stream); + LAUNCH_KERNEL(&cfg, mark_phase_ack_kernel, mxa_buffer, nvlink_available, + ipc_peer_ptrs, ack_buffer, rank, num_ranks, epoch); +} - // 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]; +void wait_phase_ack(int* ack_buffer, int rank, int num_ranks, int epoch, + cudaStream_t stream, int64_t timeout_ticks) { + SETUP_LAUNCH_CONFIG(1, 32, stream); + LAUNCH_KERNEL(&cfg, wait_phase_ack_kernel, ack_buffer, rank, num_ranks, + epoch, timeout_ticks); +} - // 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); - } +void mark_and_wait_phase_ack(void* mxa_buffer, + const int32_t* nvlink_available, + void* const* ipc_peer_ptrs, int* ack_buffer, + int rank, int num_ranks, int epoch, + cudaStream_t stream, int64_t timeout_ticks) { + SETUP_LAUNCH_CONFIG(1, 32, stream); + LAUNCH_KERNEL(&cfg, mark_and_wait_phase_ack_kernel, mxa_buffer, + nvlink_available, ipc_peer_ptrs, ack_buffer, rank, num_ranks, + epoch, timeout_ticks); } void dispatch(void* packed_recv_x, float* packed_recv_x_scales, @@ -438,7 +140,6 @@ void dispatch(void* packed_recv_x, float* packed_recv_x_scales, 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, @@ -464,7 +165,6 @@ void dispatch(void* packed_recv_x, float* packed_recv_x_scales, EP_HOST_ASSERT(kNumMaxTopK + 1 <= num_warp_groups * kNumWarpsPerGroup && "Too many top-k selections"); - const auto num_warps = num_warp_groups * kNumWarpsPerGroup; const auto num_sms = max(2, cell_div(num_experts, num_warp_groups)); EP_HOST_ASSERT(num_topk <= kNumMaxTopK); @@ -473,26 +173,48 @@ void dispatch(void* packed_recv_x, float* packed_recv_x_scales, auto atomic_finish_counter_per_expert = atomic_counter_per_expert + num_experts; EP_HOST_ASSERT(num_experts * sizeof(int) * 2 <= NUM_WORKSPACE_BYTES); + const DispatchKernelArgs args { + packed_recv_x, + packed_recv_x_scales, + packed_recv_src_info, + packed_recv_layout_range, + packed_recv_count, + active_ranks, + mxa_buffer, + rdma_send_signal_buffer, + rdma_recv_signal_buffer, + rdma_send_data_buffer, + rdma_recv_data_buffer, + raddrs, + rkeys, + qp_devctxs, + nvlink_available, + ipc_peer_ptrs, + x, + topk_idx, + atomic_counter_per_expert, + atomic_finish_counter_per_expert, + next_clean_buffer, + num_tokens, + num_max_dispatch_tokens_per_rank, + num_topk, + num_experts, + rank, + num_ranks, + timeout_ticks, + phases, + active_qps_per_rank, + }; + +#if !defined(MOONCAKE_EP_USE_MUSA) && !defined(MOONCAKE_EP_USE_MACA) + jit::launch_dispatch_jit(hidden, use_fp8, num_warp_groups, + kNumWarpsPerGroup, num_sms, args, stream); +#else // !defined(MOONCAKE_EP_USE_MUSA) && !defined(MOONCAKE_EP_USE_MACA) #define DISPATCH_LAUNCH_GROUP(hidden, groups) case groups: { \ constexpr int kNumWarpGroups = groups; \ auto dispatch_func = use_fp8 ? dispatch