core: Kimi-K3 decode/prefill scaffolding (distributed weight store, DeepEP-LL exchange, resident MoE)

[Andrewxu313]

Description

Core scaffolding for Kimi-K3 (2.8T, 93 layers, 896 MXFP4 experts, top-16) decode + prefill on 2×8 H200 — the runtime/serving layer only. Model + kernels are the companion PR (k3-model-pr, stacked on this).

Motivation

Land K3 day-0 support. K3 needs scaffolding changes a model PR may not touch (PR_MERGE_POLICY §2.5), split here: the compact distributed host-weight store + daemon, the DeepEP low-latency EP-exchange primitive, resident-MoE plumbing, and the opt-in --k3-moe-exchange / --distributed-weight-config server flags. Interim release: decode 1.32× / prefill 1.18× vs SGLang on H200 (default exchange = NCCL; 1.5× tracked as follow-up).

Type of Change

• core — change scheduling/serving/runtime scaffolding

File changes

Full list (every file traces to K3 scaffolding):

 batchgen/batchgen_worker.py                        | 1743 +++++++++++++++++---
 batchgen/ckpt_converter/ckpt_converter.py          |   72 +-
 batchgen/config/config.py                          |   16 +-
 batchgen/continuous_batching.py                    |  193 ++-
 batchgen/decode_dp_group.py                        |  139 ++
 batchgen/kernel_compat.py                          |    2 +-
 batchgen/kv_cache/gpu_paged_kv_manager.py          |   42 +-
 batchgen/kv_cache/host_kv_mananger_config.py       |   62 +-
 batchgen/moe/deepep_ll.py                          |  106 ++
 batchgen/moe/dispatch_scatter_3d.py                |   25 +
 batchgen/moe/fused_moe_bf16_resident.py            |    4 +-
 batchgen/moe/fused_moe_mxfp4_resident.py           |  985 +++++++++++
 batchgen/moe/fused_wgmma_expert.py                 |   18 +-
 batchgen/moe/fused_wgmma_grouped.py                |   18 +-
 batchgen/moe/grouped_fp8_blockwise_moe.py          |   20 +-
 batchgen/moe/k3_prefill_dequant_once.py            |  173 ++
 batchgen/moe/marlin_grouped_moe.py                 |   32 +-
 batchgen/moe/marlin_weight_prep.py                 |    7 +-
 batchgen/moe/mxfp4_grouped_gemm.py                 |    2 +-
 batchgen/moe/mxfp4_oracle_vector.py                |    2 +-
 batchgen/moe/routing/cuda_routing.py               |   43 +-
 batchgen/moe/streamed_sp8_mxfp4.py                 | 1567 ++++++++++++++++++
 .../other_kernels/hadamard_transform/__init__.py   |    4 +-
 batchgen/planner/base_planner.py                   |   23 +-
 batchgen/sequence.py                               |    6 +
 batchgen/server/batch_scheduler.py                 |   39 +-
 batchgen/server/server_args.py                     |   46 +
 batchgen/server/worker_manager.py                  |   85 +
 batchgen/server/worker_readiness.py                |   21 +
 batchgen/server_worker_main_loop.py                |  169 +-
 batchgen/worker/boundary.py                        |   57 +-
 batchgen/worker/decode.py                          |   59 +-
 batchgen/worker/kv_manager.py                      |    5 +-
 batchgen/worker/prefill.py                         |   77 +-
 core/GPU_Weight_Buffer/GPU_Weight_Buffer.cpp       |   64 +-
 core/GPU_Weight_Buffer/GPU_Weight_Buffer.h         |    8 +
 core/HtoD_Engine/HtoD_Engine.cu                    |  131 +-
 core/HtoD_Engine/HtoD_Engine.h                     |   10 +
 core/KV_Storage/compressed_state_host_manager.h    |   47 +-
 core/KV_Storage/host_paged_kv_worker_view.h        |   63 +-
 core/Parameter_Server/posix_shm.cpp                |    3 +-
 core/Weights_Storage/Weights_Storage.cpp           |  621 ++++++-
 core/Weights_Storage/Weights_Storage.h             |   58 +-
 core/Weights_Storage/distributed_weight_daemon.cpp | 1504 +++++++++++++++++
 core/Weights_Storage/distributed_weight_daemon.h   |   23 +
 .../Weights_Storage/distributed_weights_protocol.h |   41 +
 core/batchgen.cpp                                  |   85 +-
 core/batchgen.h                                    |    2 +
 core/batchgen_Binding.cpp                          |   22 +-
 core/numa_compat.h                                 |   26 +
 docs/troubleshooting.md                            |    4 +-
 op_builder/core_engine.py                          |    7 +-
 scripts/install_deps.sh                            |   41 +
 53 files changed, 8065 insertions(+), 557 deletions(-)

Key clusters: core/** C++ engine (Weights_Storage + distributed_weight_daemon, HtoD_Engine, GPU_Weight_Buffer, KV_Storage host-paged views, binding); batchgen/moe/** (new deepep_ll.py DeepEP-LL exchange, marlin_grouped_moe, streamed_sp8_mxfp4, k3_prefill_dequant_once, resident MoE, routing); worker/**, server/** (server_args --k3-moe-exchange + --distributed-weight-config, worker_manager, worker_readiness), kv_cache/**, continuous_batching.py, batchgen_worker.py; triage (planner/base_planner, sequence, config, decode_dp_group, kernel_compat, op_builder, install_deps, docs/troubleshooting).

Checklist

• Read CONTRIBUTING + PR Merge Policy.
• Tests updated (K3 tests are in the companion k3-model-pr under tests/).
• Docs: design docs in the internal repo; README News/Roadmap is a separate docs PR timed with merge.

PR Merge Policy Contract — pre-merge checklist

• git diff --stat main reviewed; every file traces to K3 scaffolding (§3.1).
• Exactly one Type (core) ticked; core is the only type that may touch scaffolding (§2.5).
• File-changes list matches the diff (above).
• No debug_*/scratch_*/tmp_* scripts in a production package (§1.1).
• No test_*.py added inside the runtime package (§1.2).
• §1.3 — deferred (POIS-approved): pre-existing BATCHGEN_* env guards in batchgen_worker.py are intentionally kept for the interim release and documented in batchgen_design/model_support/kimi_k3/ENV_KNOBS.md; conversion to batchgen_debug/server-args is the next optimization pass. BATCHGEN_ENABLE_ALL_TO_ALL stays by decision. CI hygiene is report-only.
• §1.4 — deferred: 1 print() in server_worker_main_loop.py → guarded logger next pass (advisory).
• No logs/traces/checkpoints/wheels staged (§1.6).
• One concern (K3 core scaffolding); surgical (§3).
• Commit trailers clean per §4.
• CI green — to confirm on push.