# 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.