Skip to content

refactor(mooncake): move Mooncake client config to roles.<role>.mooncake_store_config - #596

Draft
nv-yna wants to merge 2 commits into
NVIDIA:mainfrom
nv-yna:yna/mooncake-store-config-in-roles
Draft

nv-yna wants to merge 2 commits into
NVIDIA:mainfrom
nv-yna:yna/mooncake-store-config-in-roles

Conversation

@nv-yna

@nv-yna nv-yna commented Oct 10, 2026

Copy link
Copy Markdown
Collaborator

Draft for discussion. Stacked on #595: only the top commit (refactor(mooncake): …) is new here.

Problem

The Mooncake client config, the JSON a vLLM or TRT-LLM worker reads from MOONCAKE_CONFIG_PATH, is written in the mooncake-master service's options.store_config. The master never reads it: it says how each worker joins the pool and how much host memory each worker rank lends. Reading a recipe, the keys look like master settings. With #595 a TRT-LLM role can also override them in roles.<role>.mooncake_store_config, so one number can appear in two places:

roles:
  decode:
    mooncake_store_config:
      global_segment_size: 16GiB      # decode ranks lend 16 GiB
services:
  - name: mooncake-master
    type: mooncake-master
    options:
      store_config:
        global_segment_size: 8GiB     # not the master's memory: the default for every worker, used by prefill only

History: #157 added store_config on the engine (backend.mooncake_kv_store.store_config); #407 (schema 2.0) moved the master into services: and the migrator moved store_config along to options.store_config. Everything else a role needs already lives under roles.<role>, and schema 2.0 refuses per-role keys elsewhere.

What a recipe looks like with this change

Each role that uses Mooncake carries its own client config; the master entry only describes the master.

vLLM, before and after:

# before                                        # after
roles:                                          roles:
  prefill:                                        prefill:
    args: {kv-transfer-config: '...Mooncake...'}    args: {kv-transfer-config: '...Mooncake...'}
  decode:                                           mooncake_store_config:
    args: {kv-transfer-config: '...Mooncake...'}      protocol: rdma
services:                                             global_segment_size: 8GB
  - name: mooncake-master                         decode:
    type: mooncake-master                           args: {kv-transfer-config: '...Mooncake...'}
    options:                                        mooncake_store_config:
      store_config:                                   protocol: rdma
        protocol: rdma                                global_segment_size: 16GB   # now possible per role
        global_segment_size: 8GB                services:
                                                  - name: mooncake-master
                                                    type: mooncake-master

TRT-LLM is the same shape (see examples/trtllm/trtllm-serve-disagg-mooncake.yaml). srtctl writes /logs/mooncake_store_config_<role>.json for each such role and sets that role's MOONCAKE_CONFIG_PATH; vLLM previously wrote one /logs/mooncake_store_config.json for every role.

Changes

  • RoleConfig.mooncake_store_config is read by vLLM and TRT-LLM through one resolver per engine (Backend.mooncake_store_config_for_mode), with Backend.mooncake_store_modes naming the roles whose workers read a client config file.
  • vLLM writes one client config per role whose kv-transfer-config uses Mooncake; device_names_by_gpu process configs become mooncake_store_config_<role>_gpu<ids>.json, built from that role's keys.
  • Keys every client of a pool must share are a table on the backend (Backend.mooncake_pool_keys: vLLM metadata_server, protocol; TRT-LLM also model_key, namespace). The validator checks they match across roles, and refuses the field on a role that reads no client config (SGLang, or no Mooncake connector) or without a mooncake-master service. This replaces feat(trtllm): run TRT-LLM's mooncake-store connector on the native mooncake-master service #595's "pool-wide keys cannot be set per role".
  • options.store_config keeps working: its keys apply under every role's own, and loading a recipe that sets it logs a deprecation warning.
  • The dry run shows each role's config through the same hooks for both engines, replacing the two engine-specific branches.
  • TRT-LLM examples and docs/mooncake-kv-store.md / docs/services.md use the per-role form; schema docs regenerated.

Open questions

  1. Deprecate or remove? This PR deprecates. Removing options.store_config would be cleaner but breaks existing schema-2 vLLM recipes (for example the InferenceX AgentX recipes under inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/*/vllm/), and the migrator only upgrades schema 1.
  2. Repetition vs a shared default. With roles only, keys every role shares (protocol, device_name, metadata_server) are written once per role. An alternative is an engine-level default with per-role overrides, the way engine.connector and roles.<role>.args.connector work.
  3. device_names_by_gpu is also worker-side client config and stays on the master here. Move it to the roles too?
  4. Migrator. It still writes schema-1 store_config to options.store_config (deprecated but valid). Should it write each Mooncake role's mooncake_store_config instead? That needs it to know which roles use Mooncake.

Validation

  • pytest tests/: 4,491 passed, 2 skipped. New tests/test_mooncake_role_config.py covers per-role vLLM files and MOONCAKE_CONFIG_PATH, the deprecated field still applying with a warning, pool keys that must match, and a role without a Mooncake connector.
  • ruff check, ruff format --check, and srtctl schema-docs --check pass; make snapshots writes nothing (no example's srun launch changes).
  • ty check reports the same 11 invalid-attribute-override diagnostics main currently fails with under ty 0.0.86 (frozen dataclass fields overriding the RoleSettings / Backend attributes), plus one of the same kind for RoleSettings.mooncake_store_config from feat(trtllm): run TRT-LLM's mooncake-store connector on the native mooncake-master service #595; this PR adds none.
  • Not run on GPUs. The rendered client config contents are unchanged; vLLM workers now read them from a per-role path.

nv-yna added 2 commits October 9, 2026 23:08
…oncake-master service

TRT-LLM recipes can now declare the same mooncake-master service SGLang and
vLLM use. srtctl renders one Mooncake client config per role that selects
kv_connector_config.connector: mooncake-store and points that role's workers
at it with MOONCAKE_CONFIG_PATH, which TRT-LLM's connector reads before its own
mooncake_store block. trtllm-serve and dynamo.trtllm join the pool the same
way, with no trtllm-serve mooncake_master, donor processes, or worker wrapper.

- TRTLLMBackend.mooncake_kv_store (TRTLLMMooncakeKVStoreConfig): the service's
  options.store_config holds the client config keys every role shares
- roles.<role>.mooncake_store_config (a RoleConfig field, read through
  role_mooncake_store_config) layers one role's keys over it; keys the whole
  pool must agree on (model_key, namespace, metadata_server, protocol,
  master_server_address) are refused there, and the field is refused on
  SGLang and vLLM roles
- role defaults to both (prefill, agg) or capacity (decode); model_key
  defaults to the served model name, since the pool belongs to one job
- Backend.mooncake_store_configs(infra_node_ip, served_model_name) lists the
  files the orchestrator writes; get_mooncake_worker_env takes the worker mode
- the connector is recognized by its resolved module, as TRT-LLM does
- load-time checks for what TRT-LLM would reject in every rank after the
  model loads: no mooncake_store block or block-only keys alongside the
  master, no external master, sizes as byte counts or binary units, positive
  transfer_batch_size and local_buffer_size, boolean stage_through_host
- without a mooncake-master service, a connector role must find a client
  config itself: MOONCAKE_CONFIG_PATH in its env, or a mooncake_store block
  with its own run_dir when trtllm-serve runs the engine
- examples for trtllm-serve and Dynamo, docs, regenerated schema docs,
  tests, launch snapshots

Signed-off-by: Yuewei Na <nv-yna@users.noreply.github.com>
…ake_store_config

The Mooncake client config (the JSON a worker reads from MOONCAKE_CONFIG_PATH)
lived in the mooncake-master service's options.store_config, although the
master never reads it: it describes how each worker joins the pool and how much
host memory its ranks lend. It now lives with the other per-role settings, in
roles.<role>.mooncake_store_config, for vLLM as well as TRT-LLM.

- vLLM writes one client config per role whose kv-transfer-config uses
  Mooncake (/logs/mooncake_store_config_<role>.json) and points that role's
  MOONCAKE_CONFIG_PATH at it; device_names_by_gpu process configs are per role
  too (mooncake_store_config_<role>_gpu<ids>.json)
- Backend.mooncake_store_modes / mooncake_store_config_for_mode /
  mooncake_pool_keys: one resolver per engine and one table of the keys every
  client of a pool must share; the validator checks those match across roles
  and refuses the field on roles that read no client config (SGLang, or no
  Mooncake connector) and without a mooncake-master service
- options.store_config is deprecated, not removed: its keys still apply under
  every role's own, and loading a recipe that sets it logs a warning
- the dry run shows each role's config through the same hooks for both engines
- TRT-LLM examples carry their client config in the roles; docs updated

Signed-off-by: Yuewei Na <nv-yna@users.noreply.github.com>
@codecov-commenter

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 99.11111% with 2 lines in your changes missing coverage. Please review.
⚠️ Please upload report for BASE (main@009a4fa). Learn more about missing BASE report.

Files with missing lines Patch % Lines
src/srtctl/backends/base.py 93.33% 1 Missing ⚠️
src/srtctl/backends/trtllm.py 98.90% 1 Missing ⚠️
Additional details and impacted files
@@           Coverage Diff           @@
##             main     #596   +/-   ##
=======================================
  Coverage        ?   85.35%           
=======================================
  Files           ?      161           
  Lines           ?    23385           
  Branches        ?        0           
=======================================
  Hits            ?    19961           
  Misses          ?     3424           
  Partials        ?        0           

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants