Repository navigation
Conversation
…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 Report❌ Patch coverage is
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. 🚀 New features to boost your workflow:
|
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
The Mooncake client config, the JSON a vLLM or TRT-LLM worker reads from
MOONCAKE_CONFIG_PATH, is written in themooncake-masterservice'soptions.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 inroles.<role>.mooncake_store_config, so one number can appear in two places:History: #157 added
store_configon the engine (backend.mooncake_kv_store.store_config); #407 (schema 2.0) moved the master intoservices:and the migrator movedstore_configalong tooptions.store_config. Everything else a role needs already lives underroles.<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:
TRT-LLM is the same shape (see
examples/trtllm/trtllm-serve-disagg-mooncake.yaml). srtctl writes/logs/mooncake_store_config_<role>.jsonfor each such role and sets that role'sMOONCAKE_CONFIG_PATH; vLLM previously wrote one/logs/mooncake_store_config.jsonfor every role.Changes
RoleConfig.mooncake_store_configis read by vLLM and TRT-LLM through one resolver per engine (Backend.mooncake_store_config_for_mode), withBackend.mooncake_store_modesnaming the roles whose workers read a client config file.kv-transfer-configuses Mooncake;device_names_by_gpuprocess configs becomemooncake_store_config_<role>_gpu<ids>.json, built from that role's keys.Backend.mooncake_pool_keys: vLLMmetadata_server,protocol; TRT-LLM alsomodel_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 amooncake-masterservice. 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_configkeeps working: its keys apply under every role's own, and loading a recipe that sets it logs a deprecation warning.docs/mooncake-kv-store.md/docs/services.mduse the per-role form; schema docs regenerated.Open questions
options.store_configwould be cleaner but breaks existing schema-2 vLLM recipes (for example the InferenceX AgentX recipes underinferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/*/vllm/), and the migrator only upgrades schema 1.protocol,device_name,metadata_server) are written once per role. An alternative is an engine-level default with per-role overrides, the wayengine.connectorandroles.<role>.args.connectorwork.device_names_by_gpuis also worker-side client config and stays on the master here. Move it to the roles too?store_configtooptions.store_config(deprecated but valid). Should it write each Mooncake role'smooncake_store_configinstead? That needs it to know which roles use Mooncake.Validation
pytest tests/: 4,491 passed, 2 skipped. Newtests/test_mooncake_role_config.pycovers per-role vLLM files andMOONCAKE_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, andsrtctl schema-docs --checkpass;make snapshotswrites nothing (no example'ssrunlaunch changes).ty checkreports the same 11invalid-attribute-overridediagnosticsmaincurrently fails with under ty 0.0.86 (frozen dataclass fields overriding theRoleSettings/Backendattributes), plus one of the same kind forRoleSettings.mooncake_store_configfrom feat(trtllm): run TRT-LLM's mooncake-store connector on the native mooncake-master service #595; this PR adds none.