kvcache-ai / kvcache-ai/Mooncake
[RFC]: Comprehensive MoE / EP Kernel Benchmark Under Incast Traffic
Nobody has claimed this yet.
- Dominant language
- C++
- Stars
- 6.6k
- Forks
- 1.2k
- Avg merge
- 3d 5h
- Merged PRs (30d)
- 312
Description
Changes proposed
Summary
This RFC proposes adding a dedicated benchmark suite for Mooncake EP dispatch/combine kernels under both uniform and incast traffic patterns from Issue #1883 Milestone 15
Mooncake EP already has correctness coverage through existing tests that exercise dispatch/combine behavior, expert routing, data types, zero-copy paths, async finish behavior, and receive hooks. However, the project does not yet have a standardized performance benchmark that measures throughput and tail latency under realistic expert-parallel MoE serving traffic.
The main target workload for this RFC is incast traffic, where many ranks simultaneously route tokens into a small set of hot experts. This traffic pattern stresses the all-to-all dispatch path much more aggressively than uniform random routing and is important for detecting regressions in EP kernel and transport behavior.
The proposed benchmark will live separately from correctness tests, likely under:
benchmarks/ep_benchmark/
The benchmark should be CLI-driven, support JSON output, and allow repeatable comparison across CUDA and other supported backends when hardware is available.
Motivation
Mooncake EP is responsible for high-performance expert-parallel communication patterns used by MoE workloads. Correctness tests are necessary but not sufficient for kernel and transport development because they do not answer questions such as:
- How does dispatch latency change as routing skew increases?
- What is the p99 or p999 latency under expert hot-spotting?
- How does zero-copy compare against non-zero-copy execution?
- How much imbalance appears across ranks and experts?
- Does a backend change regress throughput under realistic incast load?
- Are CUDA and non-CUDA backends behaving similarly under equivalent benchmark configurations?
Current correctness tests are valuable because they validate functional behavior, but they are not designed to provide repeatable performance statistics over warmup and measured iterations.
A dedicated EP benchmark would make M15 optimization work easier to evaluate and would give maintainers and contributors a common tool for validating EP performance changes.
Goals
- Add a standalone EP benchmark harness.
- Reuse the same public EP setup path used by existing tests:
mooncake.mooncake_ep_buffer.Buffermooncake.pg- dispatch/combine APIs
- Support multiple routing distributions:
- uniform routing baseline
- K-hot-expert incast
- Zipfian/skewed routing
- Report throughput and latency statistics:
- dispatch latency
- combine latency
- end-to-end dispatch+combine latency
- tokens/sec
- p50 / p90 / p99 / p999 latency
- per-rank load imbalance
- per-expert hot-spot concentration
- Support configurable benchmark knobs:
- number of experts
- hidden size
- top-k
- batch size / number of tokens
- dtype, including bf16 and fp8 where supported
- zero-copy mode
- async finish
- return receive hook
- number of warmup iterations
- number of measured iterations
- Produce machine-readable JSON output.
- Keep performance gating out of mandatory CI initially unless suitable multi-GPU runners are available.
Non-Goals
- This is not an end-to-end model-serving benchmark.
- This does not introduce new EP kernel functionality.
- This does not replace existing correctness tests.
- This does not require performance results to be a hard CI gate in the first version.
- This does not attempt to benchmark every possible MoE architecture or model framework.
Current State
Mooncake EP currently has correctness-oriented tests that exercise dispatch and combine behavior. These tests are useful for checking that output tensors match expected values and that different API modes behave correctly.
However, correctness tests are not ideal for repeated performance measurement because they usually:
- do not run enough iterations for stable latency statistics,
- do not provide structured benchmark output,
- do not model incast traffic explicitly,
- do not separate warmup from measured iterations,
- do not expose benchmark sweeps as a stable CLI contract.
This RFC proposes adding a benchmark harness rather than extending the correctness test directly.
Proposed Design
Directory Layout
Proposed initial structure:
benchmarks/
ep_benchmark/
README.md
run_ep_benchmark.py
routing.py
metrics.py
configs/
cuda_uniform.json
cuda_incast_k_hot.json
cuda_zipfian.json
results/
.gitignore
Possible future structure if the benchmark grows:
benchmarks/
ep_benchmark/
mooncake_ep_benchmark/
__init__.py
runner.py
routing.py
metrics.py
config.py
report.py
run_ep_benchmark.py
README.md
CLI Interface
Example command:
python benchmarks/ep_benchmark/run_ep_benchmark.py \
--backend cuda \
--num-ranks 8 \
--num-experts 64 \
--hidden-size 7168 \
--top-k 8 \
--num-tokens 4096 \
--dtype bf16 \
--routing-mode k_hot \
--hot-experts 4 \
--warmup-iters 20 \
--iters 100 \
--zero-copy \
--async-finish \
--return-recv-hook \
--json-output results/ep_benchmark_cuda_khot.json
Routing Modes
1. Uniform Routing
Uniform routing should approximate the current correctness-test routing pattern and serve as a baseline.
Example behavior:
Each token chooses top-k experts from the full expert set with roughly uniform probability.
2. K-Hot-Expert Incast
K-hot routing models many ranks sending tokens to a small hot expert set.
Example behavior:
Given num_experts = 64 and hot_experts = 4, most or all ranks route a configurable percentage of tokens into experts [0, 1, 2, 3].
Suggested knobs:
--routing-mode k_hot
--hot-experts 4
--hot-fraction 0.90
This allows both extreme incast and mixed traffic.
3. Zipfian Routing
Zipfian routing models skewed real-world load where a small number of experts receive a disproportionate share of tokens.
Suggested knobs:
--routing-mode zipf
--zipf-alpha 1.2
Metrics
The benchmark should report:
{
"benchmark": "mooncake_ep",
"backend": "cuda",
"world_size": 8,
"num_experts": 64,
"hidden_size": 7168,
"top_k": 8,
"num_tokens": 4096,
"dtype": "bf16",
"routing_mode": "k_hot",
"hot_experts": 4,
"hot_fraction": 0.9,
"zero_copy": true,
"async_finish": true,
"return_recv_hook": true,
"warmup_iters": 20,
"iters": 100,
"metrics": {
"dispatch_latency_ms": {
"p50": 0.0,
"p90": 0.0,
"p99": 0.0,
"p999": 0.0,
"mean": 0.0
},
"combine_latency_ms": {
"p50": 0.0,
"p90": 0.0,
"p99": 0.0,
"p999": 0.0,
"mean": 0.0
},
"end_to_end_latency_ms": {
"p50": 0.0,
"p90": 0.0,
"p99": 0.0,
"p999": 0.0,
"mean": 0.0
},
"tokens_per_second": 0.0,
"expert_load": {
"max_tokens_per_expert": 0,
"min_tokens_per_expert": 0,
"mean_tokens_per_expert": 0.0,
"imbalance_ratio": 0.0
}
}
}
Actual metric values should be populated by the benchmark.
Timing Method
The benchmark should separate:
- setup,
- warmup iterations,
- measured iterations,
- result aggregation.
For CUDA, timing should use appropriate synchronization around measured regions so that async GPU work is accounted for correctly. If backend-specific timing APIs differ, the benchmark should isolate backend timing behavior behind helper functions.
Backend Support
Initial implementation should target CUDA first because it is the most common validation environment.
MUSA or other backend support can be added after confirming:
- build support,
- hardware availability,
- expected backend flags,
- whether the same benchmark code path can be reused.
The benchmark should not assume all contributors have access to every backend.
CI Strategy
The benchmark should not become a hard performance gate initially.
Suggested CI behavior for the first version:
- Run lightweight import/config parsing checks.
- Run unit tests for routing distribution generation.
- Optionally run a tiny single-process or mock smoke test if feasible.
- Do not require multi-GPU performance results in standard CI.
Future CI options:
- self-hosted multi-GPU runner,
- nightly benchmark job,
- manual benchmark workflow,
- threshold-based regression checks after the benchmark is stable.
Implementation Plan
Phase 1 : Minimal Benchmark Harness
- Add
benchmarks/ep_benchmark/. - Add CLI parser.
- Add uniform routing mode.
- Reuse existing EP setup pattern.
- Measure dispatch/combine latency over warmup and measured iterations.
- Emit JSON output.
Phase 2 : Incast Routing
- Add K-hot-expert routing mode.
- Add hot expert fraction.
- Add per-expert load metrics.
- Add README examples.
Phase 3 : Zipfian Routing and Sweep Support
- Add Zipfian routing mode.
- Add simple config-file support.
- Add benchmark sweep examples.
Phase 4 : Backend Parity
- Add backend selection where supported.
- Validate CUDA first.
- Add MUSA or other backend support only where maintainers confirm hardware/toolchain availability.
Phase 5 : Optional CI / Nightly Integration
- Add smoke checks to standard CI.
- Consider a manual or nightly benchmark job if suitable hardware exists.
Validation Plan
The benchmark should be validated in three ways:
- Correctness sanity check: for small inputs, verify benchmark dispatch/combine outputs remain consistent with existing test behavior.
- Routing validation: ensure uniform, K-hot, and Zipfian modes produce expected expert-load distributions.
- Performance sanity check: run repeated benchmark executions on the same hardware and confirm variance is within an acceptable range.
Alternatives Considered
Extend Existing Correctness Test
Rejected for initial implementation.
Correctness tests and performance benchmarks have different goals. Mixing repeated-trial statistics, warmup, JSON output, and sweep logic into a correctness test would make both harder to maintain.
Rely Only on External MoE Benchmarks
External benchmarks may be useful for comparison, but Mooncake EP should have a benchmark that exercises its own public API and backend paths directly.
Add CI Perf Gate Immediately
Rejected for the first version.
Kernel performance benchmarking requires stable hardware, stable drivers, and controlled environments. A hard CI gate should only be added after the benchmark is stable and the project has appropriate runners.
Risks
- Benchmark noise: GPU performance can vary across hosts and runs.
- Hardware availability: contributors may not have access to multi-GPU or non-CUDA environments.
- Backend divergence: CUDA and other backends may require slightly different timing or setup paths.
- Scope creep: benchmark should remain focused on EP dispatch/combine, not full model serving.
Open Questions
- What hardware can maintainers use for official EP benchmark runs?
- Should the first benchmark target only single-node multi-GPU, or should cross-node EP/RDMA behavior be included later?
- Which backend combinations should be considered required for the first version?
- Should benchmark results be stored in-tree as examples, or only generated locally?
- What metric thresholds would be meaningful enough for future regression detection?
Expected Outcome
This RFC should produce a repeatable Mooncake EP benchmark that helps contributors measure dispatch/combine performance under both normal and incast traffic. The first implementation should be useful as a manual benchmark tool, with CI integration kept lightweight until stable multi-GPU runners are available.
Before submitting a new issue...
- Make sure you already searched for relevant issues and read the documentation
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start by locating the existing EP correctness tests and the public mooncake_ep_buffer.Buffer, mooncake.pg, and dispatch/combine entry points. Review the proposed benchmarks/ep_benchmark/ layout and define the first implementation boundary; done should include a runnable CLI, repeatable warmup and measured iterations, routing coverage, and JSON metrics without making performance a CI gate.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, python
- Domain
- distributed-systems, performance, testing
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Needs clarification
- Newbie friendliness
- 35/100