[EP] Add JIT compilation support to Mooncake EP CUDA backend by Bonjourz · Pull Request #4581 · kvcache-ai/Mooncake

Bonjourz · 2026-10-12T06:38:14Z

Description

Port the JIT framework from NCCL-EP to Mooncake-EP, enabling on demand compiling CUDA kernels:

Add runtime source generation, nvcc compilation, CUDA module lading, and launch kernel support.
Cache compiled artifacts on disk and loaded kernels in memory, with a warm launch path that avoids source generation and file I/O.
Move standard dispatch and combine kernel implementations into reusable headers and route their CUDA launches through JIT.
Extend JIT coverage to ElasticBuffer kernels, including NCCL and IBGDA transport variants.

Module

Type of Change

How Has This Been Tested?

Test tracking items:

Unit tests pass
Integration tests pass: Have tested on H100 and GB300 platforms
Manual testing done (describe below)

Build

Build command:

cmake -S . -B build -G Ninja \
    -DCMAKE_BUILD_TYPE=Release \
    -DWITH_EP=ON \
    -DUSE_CUDA=ON \
    -DCMAKE_CUDA_ARCHITECTURES=[90|100|103] \
    -DUSE_NCCL_DEVICE=ON \
    -DNCCL_ROOT=path/to/nccl" \
    -DWITH_STORE=OFF \
    -DWITH_STORE_RUST=OFF \
    -DBUILD_UNIT_TESTS=ON \
    -DBUILD_EXAMPLES=ON \

cmake --build build -j $(nproc)

Functionality Test

torch mooncake-ep/tests/test_elastic_buffer.py \
  --transport nccl --route alltoall --num-tokens 32 --max-tokens 128 \
  --hidden 4096 --num-experts 256 --num-topk 8 --num-sms 24 --warmup 20 --iters 200

Performance Test

Configuration: BF16, 32 tokens per rank, hidden size 4096, top-k 8, 256 experts, 24 communication SMs, 20 warm-up iterations, 200 measured iterations and 5 repetitions. The value is the mean of 5 runs.

GB300 platform, NCCL 2.31.2-1, CUDA 13.0, Python 3.12.3, torch 2.14.1

Topology	JIT support	Dispatch critical (mean, SD, ms)	Combine critical (mean, SD, ms)	E2E critical (mean, SD, ms)	Wall/iter (mean, SD, ms)
1 node, 2 GPUs	with JIT	0.176/0.008	0.124/0.003	0.300/0.011	0.769/0.015
1 node, 2 GPUs	vanilla	0.170/0.003	0.123/0.001	0.292/0.004	0.762/0.007
2 nodes, 4 GPUs per node	with JIT	0.241/0.006	0.149/0.002	0.384/0.009	0.837/0.009
2 nodes, 4 GPUs per node	vanilla	0.242/0.005	0.151/0.002	0.386/0.007	0.838/0.007

Compilation Test

Without JIT support, a fixed set of instantiations combination of template parameter need to be compiled. That bloat the binary size. With JIT, the combination is determined based the user's need.

Without JIT support, the size of libmooncake_ep_device.so is 33M:

33M -rwxrwxr-x 1 ubuntu ubuntu 33M Oct 27 03:14 libmooncake_ep_device.so

with JIT support, the size is reduced to 6M:

6.0M -rwxrwxr-x 1 ubuntu ubuntu 6.0M Oct 27 02:45 ./build/mooncake-ep/src/libmooncake_ep_device.so

Checklist

I have performed a self-review of my own code
I have formatted my code using ./scripts/code_format.sh
I have run pre-commit on the files changed in this PR and all hooks pass
I have updated the documentation (if applicable)
The deployment guide for JIT support docs/source/deployment/mooncake-ep-jit.md is updated.
I have added tests to prove my changes are effective
For changes >500 LOC: I have filed an RFC issue

AI Assistance Disclosure

No AI tools were used
AI tools were used (specify below)

The code changes were made with the help of AI. The author reviewed the implementation and test results.

UNIDY2002 · 2026-10-26T05:23:13Z

Hi @Bonjourz! Could you help merge Mooncake's latest main into this branch? Thanks!

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 <bojunz@nvidia.com>

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 <bojunz@nvidia.com>

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 <bojunz@nvidia.com>

The 'sm_min' in mooncake-ep is sm 80. Make the 'sm_min' explicitly configured by upper level. Signed-off-by: Bojun Zhu <bojunz@nvidia.com>

All the kernels in mooncake-ep are launched with cooperative launch. Signed-off-by: Bojun Zhu <bojunz@nvidia.com>

Signed-off-by: Bojun Zhu <bojunz@nvidia.com>

mooncake-ep only uses int64_t for top-k indices, but NVCC_EXTRA_FLAGS can make 'topk_idx_t' a different type in the JIT compile. Add a static assertion so that fails at compile time. Signed-off-by: Bojun Zhu <bojunz@nvidia.com>

Signed-off-by: Bojun Zhu <bojunz@nvidia.com>

Bonjourz · 2026-10-27T02:19:52Z

Hi @Bonjourz! Could you help merge Mooncake's latest main into this branch? Thanks!

Hi @UNIDY2002 . I have updated update branch based on the latest main.

By the way, I also added docs for JIT deployment (docs/source/deployment/mooncake-ep-jit.md), and then updated the mooncake-ep docs (docs/source/design/mooncake-ep.md) to announce that JIT is supported on NVIDIA's backend.

UNIDY2002 · 2026-10-27T07:42:02Z

SGLang e2e test:

UNIDY2002 · 2026-10-27T07:50:33Z

@Bonjourz Please fix the CI error: Build & Test (Linux) / Check Sphinx docs build (pull_request)

Signed-off-by: Bojun Zhu <bojunz@nvidia.com>

UNIDY2002 · 2026-10-27T08:30:15Z

LGTM. Let's wait for the Mooncake CI results.

UNIDY2002

thank you!

github-actions Bot added run-ci Mooncake EP labels Oct 12, 2026

Bonjourz force-pushed the bojunz/ep_jit_pr branch from cca8b26 to 4966a07 Compare October 12, 2026 07:29

Bonjourz marked this pull request as ready for review October 16, 2026 14:41

Bonjourz requested review from UNIDY2002, ympcMark and yuechen-sys as code owners October 16, 2026 14:41

UNIDY2002 reviewed Oct 22, 2026

View reviewed changes

Comment thread mooncake-ep/src/CMakeLists.txt Resolved

github-actions Bot added the Integration label Oct 26, 2026

Bonjourz force-pushed the bojunz/ep_jit_pr branch 8 times, most recently from b962126 to 92f43a8 Compare October 26, 2026 13:36

github-actions Bot added the size/L 1001+ changed lines (additions + deletions) label Oct 26, 2026

Bonjourz force-pushed the bojunz/ep_jit_pr branch from b28dd12 to d6578b7 Compare October 27, 2026 02:14

Bonjourz requested review from ShangmingCai, alogfans, staryxchen, stmatengss and ykwd as code owners October 27, 2026 02:14

Bonjourz added 5 commits October 27, 2026 10:15

[JIT] Add JIT framework support

c5bc7dc

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 <bojunz@nvidia.com>

[JIT] Generate jit_config.h for JIT include paths

736f642

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 <bojunz@nvidia.com>

[JIT] Add transfer engine into the include path

3c7f03a

[JIT] Use an invalid value for the 'sm_min' for default

14c622d

The 'sm_min' in mooncake-ep is sm 80. Make the 'sm_min' explicitly configured by upper level. Signed-off-by: Bojun Zhu <bojunz@nvidia.com>

Bonjourz added 11 commits October 27, 2026 10:15

[JIT] Set cooperative to true

ed9b5b0

All the kernels in mooncake-ep are launched with cooperative launch. Signed-off-by: Bojun Zhu <bojunz@nvidia.com>

[JIT] Add support for dispatch

55e7644

[JIT] Add support for combine

4fdb8b5

Signed-off-by: Bojun Zhu <bojunz@nvidia.com>

[JIT] Add static assertion on the type of 'topk_idx_t'

06e84e9

mooncake-ep only uses int64_t for top-k indices, but NVCC_EXTRA_FLAGS can make 'topk_idx_t' a different type in the JIT compile. Add a static assertion so that fails at compile time. Signed-off-by: Bojun Zhu <bojunz@nvidia.com>

[JIT] Add support for elastic_dispatch_prologue

0adee12

Signed-off-by: Bojun Zhu <bojunz@nvidia.com>

[JIT] Add support for elastic_dispatch

87cc2c1

Signed-off-by: Bojun Zhu <bojunz@nvidia.com>

[JIT] Add support for elastic_hybrid_dispatch

4209cb5

Signed-off-by: Bojun Zhu <bojunz@nvidia.com>

[JIT] Add support for elastic_dispatch_copy_epilogue

c0050ce

Signed-off-by: Bojun Zhu <bojunz@nvidia.com>

[JIT] Add support for elastic_combine

53d5486

Signed-off-by: Bojun Zhu <bojunz@nvidia.com>

[JIT] Add support for elastic_hybrid_combine

a8f2838

Signed-off-by: Bojun Zhu <bojunz@nvidia.com>

[JIT] Add support for elastic_combine_reduce_epilogue

065e61a

Signed-off-by: Bojun Zhu <bojunz@nvidia.com>

Bonjourz force-pushed the bojunz/ep_jit_pr branch from d6578b7 to 0a02492 Compare October 27, 2026 02:15

Bonjourz changed the title ~~[EP] Add JIT compilation for dispatch and combine kernels~~ [EP] Add JIT compilation support to Mooncake EP CUDA backend Oct 27, 2026

github-actions Bot added the documentation Improvements or additions to documentation label Oct 27, 2026

[JIT] Add deployment guide for EP JIT

8a81251

Signed-off-by: Bojun Zhu <bojunz@nvidia.com>

Bonjourz force-pushed the bojunz/ep_jit_pr branch from 0a02492 to 8a81251 Compare October 27, 2026 08:25

UNIDY2002 approved these changes Oct 27, 2026

View reviewed changes

UNIDY2002 merged commit dfd1b71 into kvcache-ai:main Oct 27, 2026
28 checks passed

staryxchen mentioned this pull request Oct 28, 2026

[Bugfix][Store] Exit standalone master without destroying the object index #4865

Open

7 of 23 tasks

Navigation Menu

Repository navigation

[EP] Add JIT compilation support to Mooncake EP CUDA backend - #4581

[EP] Add JIT compilation support to Mooncake EP CUDA backend#4581
UNIDY2002 merged 17 commits into
kvcache-ai:mainfrom
Bonjourz:bojunz/ep_jit_pr

Bonjourz commented Oct 12, 2026 •

edited

Loading

Uh oh!

Uh oh!

UNIDY2002 commented Oct 26, 2026

Uh oh!

Bonjourz commented Oct 27, 2026 •

edited

Loading

Uh oh!

UNIDY2002 commented Oct 27, 2026

Uh oh!

UNIDY2002 commented Oct 27, 2026 •

edited

Loading

Uh oh!

UNIDY2002 commented Oct 27, 2026

Uh oh!

UNIDY2002 left a comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

Bonjourz commented Oct 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Module

Type of Change

How Has This Been Tested?

Build

Functionality Test

Performance Test

Compilation Test

Checklist

AI Assistance Disclosure

Uh oh!

Uh oh!

UNIDY2002 commented Oct 26, 2026

Uh oh!

Bonjourz commented Oct 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

UNIDY2002 commented Oct 27, 2026

Uh oh!

UNIDY2002 commented Oct 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

UNIDY2002 commented Oct 27, 2026

Uh oh!

UNIDY2002 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Bonjourz commented Oct 12, 2026 •

edited

Loading

Bonjourz commented Oct 27, 2026 •

edited

Loading

UNIDY2002 commented Oct 27, 2026 •

edited

Loading