[EP] Add JIT compilation support to Mooncake EP CUDA backend

docs/source/deployment/mooncake-ep-jit.md

@@ -0,0 +1,126 @@
+# JIT Kernel Compilation
+
+Mooncake EP JIT currently supports only the NVIDIA (CUDA) backend.
+
+Instead of shipping every kernel specialization in the library, Mooncake EP
+compiles several of its device kernels at runtime. The first call that needs a
+given kernel variant runs nvcc, caches the resulting cubin on disk, and reuses
+it for the rest of the process and for later runs.
+
+As a result, a deployed Mooncake EP needs more than
+`libmooncake_ep_device.so`. It also needs a working nvcc, the kernel headers,
+the CUDA headers, a writable cache directory, and, for builds with
+`USE_NCCL_DEVICE=ON`, the NCCL device headers. Each of these is found through an
+environment variable, with a fallback compiled into the library. A build used
+where it was built, or an installed wheel, usually needs no configuration. The
+environment variables cover the other cases: relocated installs, containers,
+read-only filesystems, and cross-architecture builds.
+
+## Path resolution
+
+Each path comes from the first rule in its list that gives a value. The
+compiled-in defaults are set by CMake at build time and point into the build
+tree.
+
+**Kernel headers** (the Mooncake EP and transfer engine headers that the JIT
+compiles against):
+
+1. The `MOONCAKE_EP_JIT_SOURCE_DIR` environment variable.
+2. `$MOONCAKE_EP_HOME/include/mooncake_ep`.
+3. `mooncake_ep_jit/include/mooncake_ep` next to `libmooncake_ep_device.so`, if
+   that directory exists (the layout of an installed wheel).
+4. The compiled-in `MOONCAKE_EP_JIT_SOURCE_DIR`.
+
+**NCCL device headers** (only for builds with `USE_NCCL_DEVICE=ON`):
+
+1. The `MOONCAKE_EP_JIT_BUILD_INCLUDE_DIR` environment variable.
+2. `$NCCL_HOME/include`.
+3. The compiled-in `MOONCAKE_EP_JIT_BUILD_INCLUDE_DIR`.
+
+**CUDA headers:**
+
+1. The `MOONCAKE_EP_JIT_CUDA_INCLUDE_DIR` environment variable.
+2. `$CUDA_HOME/include`, or `$CUDA_PATH/include` if `CUDA_HOME` is not set.
+3. The compiled-in `MOONCAKE_EP_JIT_CUDA_INCLUDE_DIR`.
+
+**Compiler:**
+
+1. The `MOONCAKE_EP_JIT_NVCC` environment variable.
+2. The `NVCC` environment variable.
+3. `$CUDA_HOME/bin/nvcc`, or `$CUDA_PATH/bin/nvcc` if `CUDA_HOME` is not set,
+   if that file exists.
+4. `nvcc` from `PATH`.
+
+## Cache
+
+Compiled cubins are cached in `MOONCAKE_EP_JIT_CACHE_DIR`, which defaults to
+`/tmp/mooncake_ep/jit`.
+
+The cache key combines two fingerprints:
+
+- **The kernel variant**: kernel family, variant name, entry point, generated
+  source, target SM and architecture, and runtime key.
+- **The compile environment**: compiler identity, the contents of the kernel,
+  transfer engine and NCCL header trees, and the full list of compile options.
+
+This has two consequences:
+
+- **The cache invalidates itself.** Editing a kernel header, switching
+  compilers, changing architecture flags, or using different include
+  directories all change the key, so a new entry is compiled. A stale cubin is
+  never loaded, and there is no need to clear the cache after a rebuild.
+- **Old entries are never removed.** Superseded cubins stay on disk, so on a
+  machine that rebuilds often the cache keeps growing. Deleting the directory is
+  always safe; the only cost is recompiling.
+
+Cache writes are atomic. Each cubin is first written to a temporary file named
+with the process ID and thread ID, then renamed into place. Ranks on the same
+node can therefore fill a shared cache at the same time without corrupting each
+other's entries.
+
+Set `MOONCAKE_EP_JIT_CACHE_DIR` to another location when `/tmp` is read-only,
+too small, or not shared across containers the way you expect. Putting the
+cache on a persistent volume also avoids the first-call compile delay each time
+a new container starts.
+
+
+## Diagnostics
+
+Set `MOONCAKE_EP_JIT_LOG=1` to print JIT logs to stderr. Each line starts with
+`[mooncake_ep jit] pid=<pid>`. The log shows cache hits and misses, compile
+attempts and failures, and cache write errors (mkdir, open, write, and rename).
+It is the quickest way to find a cache directory that is not writable or is
+full.
+
+Each kernel variant has its own directory in the cache, containing
+`kernel.cu` (the generated source), `kernel.cubin`, `metadata.json`, and, after
+a compile, `compile.log` (the nvcc output).
+
+When a variant fails to compile, the failure is recorded in two places:
+
+- **In the process**: the variant is not compiled again, and its warning is
+  printed only once.
+- **On disk**: a `compile.failed` file is written to the variant's cache
+  directory. Later processes that use the same cache skip the compile and
+  report the failure without running nvcc again.
+
+To retry after fixing the cause (for example, a missing header), delete that
+variant's directory or the whole cache directory. Changes that alter the cache
+key, such as editing a header or changing compile options, retry automatically
+because they use a new directory.
+
+## Deployment checklist
+
+- **nvcc** can be found through `MOONCAKE_EP_JIT_NVCC`, `NVCC`, `CUDA_HOME`,
+  `CUDA_PATH`, or `PATH`.
+- **Kernel headers** can be found. Installed wheels and builds used where they
+  were built need no setup. If the library was moved after it was built, set
+  `MOONCAKE_EP_HOME` or `MOONCAKE_EP_JIT_SOURCE_DIR`.
+- **CUDA headers** can be found through `CUDA_HOME`, `CUDA_PATH`, or
+  `MOONCAKE_EP_JIT_CUDA_INCLUDE_DIR`.
+- **NCCL device headers** (only for `USE_NCCL_DEVICE=ON` builds) can be found
+  through `NCCL_HOME` or `MOONCAKE_EP_JIT_BUILD_INCLUDE_DIR`.
+- **The cache directory** is writable. Set `MOONCAKE_EP_JIT_CACHE_DIR` if `/tmp`
+  is not suitable.
+- **The first call** to each kernel variant includes compile time. Later calls,
+  and later runs that share the cache, do not.

docs/source/design/mooncake-ep.md

@@ -143,6 +143,24 @@ It provides the P2P / IPC fast path only:
 - Kernels run on the caller's `torch.xpu` stream. Zero-copy combine and the
   elastic buffer are not supported.
 
+## JIT kernel compilation
+
+On NVIDIA GPUs, Mooncake EP compiles its dispatch and combine kernels, including
+the elastic buffer kernels, at runtime. Parameters such as `hidden`,
+`num_experts`, and `num_topk` are compile-time template arguments of these
+kernels. Without JIT, the library can only contain a fixed set of
+instantiations (for example, `hidden` in {2048, 2560, 3072, 3584, 4096, 5120,
+6144, 7168, 8192}), and any other value fails with an "Unsupported" error. With
+JIT, the first call with a new combination of values runs nvcc to build that
+instantiation and caches the cubin on disk; later calls and later runs reuse it.
+
+JIT is currently not supported on other backends.
+
+At runtime this needs nvcc, the kernel headers, and a writable cache directory.
+Installed wheels ship the kernel headers, so they usually need no setup. See
+[JIT kernel compilation](../deployment/mooncake-ep-jit.md) for path lookup,
+cache behavior, and environment variables.
+
 ## Metadata exchange
 
 During `Buffer.connect()`, ranks exchange:
@@ -279,3 +297,4 @@ Adapt launch commands to the target environment and number of GPUs.
 - [Mooncake PG design](mooncake-backend-pg.md)
 - [Python API reference](../api-reference/python/ep-backend.md)
 - [PG/EP troubleshooting](../troubleshooting/pg-ep-troubleshooting.md)
+- [Mooncake EP JIT kernel compilation](../deployment/mooncake-ep-jit.md)

docs/source/index.md

@@ -167,6 +167,7 @@ deployment/store-rs/index
 deployment/kubernetes-deployment-guide/index
 deployment/integrations/sglang/index
 deployment/integrations/vllm/index
+deployment/mooncake-ep-jit
 Mooncake x LMCache Integration<deployment/integrations/lmcache/index>
 Mooncake x LMDeploy Integration<deployment/integrations/lmdeploy>
 :::

mooncake-ep/include/elastic/mooncake_ep_elastic_combine_official.cuh

@@ -3,6 +3,8 @@
 // transport references are replaced with Mooncake Device API adapters.
 #pragma once
 
+#include <type_traits>
+
 #include <elastic/mooncake_ep_elastic_transport.cuh>
 
 #include <elastic/mooncake_ep_elastic_comm.cuh>
@@ -14,6 +16,19 @@
 
 namespace mooncake::elastic {
 
+template <typename Ops>
+struct CombineKernelArgs {
+    nv_bfloat16* x;
+    float* topk_weights;
+    int* src_metadata;
+    int* psum_num_recv_tokens_per_scaleup_rank;
+    typename Ops::Context comm_ctx;
+    void* buffer;
+    void* workspace;
+    int rank_idx;
+    int num_reduced_tokens;
+};
+
 template <typename Ops, bool kIsScaleupNVLink, bool kUseExpandedLayout,
           bool kAllowMultipleReduction, int kNumSMs, int kNumWarps,
           int kNumRanks, int kHidden, int kNumMaxTokensPerRank, int kNumExperts,
@@ -26,11 +41,22 @@ template <typename Ops, bool kIsScaleupNVLink, bool kUseExpandedLayout,
               kAllowMultipleReduction, kNumRanks, kNumTopk>(),
           typename team_t = std::conditional_t<
               kIsScaleupNVLink, transport::ScaleupTeam, transport::WorldTeam>>
-__global__ void __launch_bounds__(kNumThreads, 1)
-    combine_impl(nv_bfloat16* x, float* topk_weights, int* src_metadata,
-                 int* psum_num_recv_tokens_per_scaleup_rank,
-                 const typename Ops::Context comm_ctx, void* buffer,
-                 void* workspace, const int rank_idx, int num_reduced_tokens) {
+__device__ __forceinline__ void combine_kernel_impl(
+    const CombineKernelArgs<Ops>& args) {
+    static_assert(std::is_trivially_copyable<CombineKernelArgs<Ops>>::value,
+                  "CombineKernelArgs must be trivially copyable");
+
+    nv_bfloat16* x = args.x;
+    float* topk_weights = args.topk_weights;
+    int* src_metadata = args.src_metadata;
+    int* psum_num_recv_tokens_per_scaleup_rank =
+        args.psum_num_recv_tokens_per_scaleup_rank;
+    const typename Ops::Context& comm_ctx = args.comm_ctx;
+    void* buffer = args.buffer;
+    void* workspace = args.workspace;
+    const int rank_idx = args.rank_idx;
+    int num_reduced_tokens = args.num_reduced_tokens;
+
     // Utils
     const auto sm_idx = static_cast<int>(blockIdx.x);
     const auto thread_idx = static_cast<int>(threadIdx.x);
@@ -364,4 +390,17 @@ __global__ void __launch_bounds__(kNumThreads, 1)
                                    thread_idx);
 }
 
+template <typename Ops, bool kIsScaleupNVLink, bool kUseExpandedLayout,
+          bool kAllowMultipleReduction, int kNumSMs, int kNumWarps,
+          int kNumRanks, int kHidden, int kNumMaxTokensPerRank, int kNumExperts,
+          int kNumTopk, int kNumQPs, int64_t kNumTimeoutCycles,
+          int kNumThreads = kNumWarps * 32>
+__global__ void __launch_bounds__(kNumThreads, 1)
+    combine(const CombineKernelArgs<Ops> args) {
+    combine_kernel_impl<Ops, kIsScaleupNVLink, kUseExpandedLayout,
+                        kAllowMultipleReduction, kNumSMs, kNumWarps, kNumRanks,
+                        kHidden, kNumMaxTokensPerRank, kNumExperts, kNumTopk,
+                        kNumQPs, kNumTimeoutCycles>(args);
+}
+
 }  // namespace mooncake::elastic

mooncake-ep/include/elastic/mooncake_ep_elastic_combine_reduce_epilogue.cuh

@@ -3,6 +3,8 @@
 // transport references are replaced with Mooncake Device API adapters.
 #pragma once
 
+#include <type_traits>
+
 #include <elastic/mooncake_ep_elastic_compiled.cuh>
 #include <elastic/mooncake_ep_elastic_ptx.cuh>
 #include <elastic/mooncake_ep_elastic_layout.cuh>
@@ -11,6 +13,20 @@
 
 namespace mooncake::elastic {
 
+struct CombineReduceEpilogueKernelArgs {
+    nv_bfloat16* combined_x;
+    float* combined_topk_weights;
+    topk_idx_t* combined_topk_idx;
+    void* recv_buffer;
+    void* bias_0;
+    void* bias_1;
+    int num_combined_tokens;
+};
+
+static_assert(
+    std::is_trivially_copyable<CombineReduceEpilogueKernelArgs>::value,
+    "CombineReduceEpilogueKernelArgs must be trivially copyable");
+
 template <bool kUseExpandedLayout, bool kAllowMultipleReduction, int kNumSMs,
           int kNumWarps,
           // TODO: merge these two variables into one (ensure the whole file
@@ -25,14 +41,16 @@ template <bool kUseExpandedLayout, bool kAllowMultipleReduction, int kNumSMs,
               use_rank_layout<kAllowMultipleReduction, kNumRanks, kNumTopk>(),
           int kNumTokensInLayout = get_num_tokens_in_layout<
               kAllowMultipleReduction, kNumRanks, kNumTopk>()>
-__global__ void __launch_bounds__(kNumThreads, 1)
-    combine_reduce_epilogue_impl(nv_bfloat16* combined_x,
-                                 float* combined_topk_weights,
-                                 topk_idx_t* combined_topk_idx,
-                                 void* recv_buffer, void* bias_0, void* bias_1,
-                                 const int num_combined_tokens,
-                                 const int scaleout_rank_idx,
-                                 const int scaleup_rank_idx) {
+__device__ __forceinline__ void combine_reduce_epilogue_kernel_impl(
+    const CombineReduceEpilogueKernelArgs& args) {
+    nv_bfloat16* combined_x = args.combined_x;
+    float* combined_topk_weights = args.combined_topk_weights;
+    topk_idx_t* combined_topk_idx = args.combined_topk_idx;
+    void* recv_buffer = args.recv_buffer;
+    void* bias_0 = args.bias_0;
+    void* bias_1 = args.bias_1;
+    const int num_combined_tokens = args.num_combined_tokens;
+
     constexpr int kNumExpertsPerScaleout = kNumExperts / kNumScaleoutRanks;
     constexpr int kNumExpertsPerRank =
         kNumExperts / (kNumScaleupRanks * kNumScaleoutRanks);
@@ -210,4 +228,16 @@ __global__ void __launch_bounds__(kNumThreads, 1)
     }
 }
 
+template <bool kUseExpandedLayout, bool kAllowMultipleReduction, int kNumSMs,
+          int kNumWarps, int kNumScaleoutRanks, int kNumScaleupRanks,
+          int kHidden, int kNumMaxTokensPerRank, int kNumExperts, int kNumTopk,
+          int kNumThreads = kNumWarps * 32>
+__global__ void __launch_bounds__(kNumThreads, 1)
+    combine_reduce_epilogue(const CombineReduceEpilogueKernelArgs args) {
+    combine_reduce_epilogue_kernel_impl<
+        kUseExpandedLayout, kAllowMultipleReduction, kNumSMs, kNumWarps,
+        kNumScaleoutRanks, kNumScaleupRanks, kHidden, kNumMaxTokensPerRank,
+        kNumExperts, kNumTopk>(args);
+}
+
 }  // namespace mooncake::elastic

mooncake-ep/include/elastic/mooncake_ep_elastic_compiled.cuh

@@ -99,6 +99,11 @@ struct int_with_bits<64> {
 };
 
 using topk_idx_t = int_with_bits<EP_NUM_TOPK_IDX_BITS>::type;
+// The host launchers pass int64_t* top-k arrays, so any other width (e.g. set
+// only for the JIT compile via NVCC_EXTRA_FLAGS) would be misread silently.
+static_assert(sizeof(topk_idx_t) == sizeof(int64_t),
+              "EP_NUM_TOPK_IDX_BITS must be 64: the host launchers pass "
+              "int64_t* top-k arrays");
 
 union sf_pack_t {
     float fp32;