diff --git a/docker/common/install_mooncake.sh b/docker/common/install_mooncake.sh index badd5f0eb6f5..a05d08c1f7b7 100644 --- a/docker/common/install_mooncake.sh +++ b/docker/common/install_mooncake.sh @@ -50,3 +50,60 @@ cd ../.. rm -rf Mooncake echo "export LD_LIBRARY_PATH=${MOONCAKE_INSTALL_PATH}/lib:\$LD_LIBRARY_PATH" >> "${ENV}" + +# The source build above provides only the C++ transfer engine, which is what +# the cache transceiver links against. MooncakeDistributedStore, the shared CPU +# pool behind the mooncake-store KV cache connector, comes from the Python +# wheel instead, for two reasons. +# +# First, `make install` emits a `mooncake` Python package that omits +# libmooncake_store.so, so importing mooncake.store raises ImportError. It has +# to be removed wherever it landed, and where that is depends on the +# environment: mooncake-integration/CMakeLists.txt picks its install directory +# as the first sys.path entry whose name merely contains "packages". +# +# - With nvidia-cutlass-dsl installed, that is +# nvidia_cutlass_dsl/dsl_packages, which nvidia_cutlass_dsl_packages.pth +# puts at sys.path[0], so it shadows anything pip installs. CUTLASS DSL +# does not reference `mooncake`, so removing the package is safe. +# - Without it, the package lands in dist-packages and collides with the +# wheel: CMake writes store.cpython-312-x86_64-linux-gnu.so, the wheel +# writes store.so, and importlib prefers the interpreter-tagged suffix, so +# the broken extension wins even after pip reports success. +# +# Remove the directory outright rather than trying to identify leftovers, since +# pip overwrites __init__.py in the collision case and leaves no marker to key +# on. +python3 - <<'PY' +import os +import shutil +import sys +import sysconfig + +paths = sysconfig.get_paths() +for entry in list(sys.path) + [paths["purelib"], paths["platlib"]]: + if not entry: + continue + package = os.path.join(entry, "mooncake") + if os.path.isdir(package): + print(f"removing CMake-generated mooncake package: {package}") + shutil.rmtree(package, ignore_errors=True) +PY + +# Second, the `mooncake-transfer-engine` wheel is built against CUDA 12 while +# these images ship CUDA 13 only, so its extensions cannot resolve +# libcudart.so.12. `mooncake-transfer-engine-cuda13` is the same project built +# for CUDA 13. It is versioned independently, with releases starting at 0.3.9, +# so it cannot track MOONCAKE_VERSION above. The store client only has to agree +# with the mooncake_master it connects to, and this wheel supplies both. +MOONCAKE_WHEEL_VERSION="0.3.13" +pip3 install --no-cache-dir "mooncake-transfer-engine-cuda13==${MOONCAKE_WHEEL_VERSION}" + +# Fail the build rather than ship an image whose import is broken. +python3 - <<'PY' +from mooncake.store import MooncakeDistributedStore +import mooncake.store + +MooncakeDistributedStore() +print(f"mooncake.store OK: {mooncake.store.__file__}") +PY diff --git a/docs/source/features/kv-cache-connector.md b/docs/source/features/kv-cache-connector.md index 89ae82287b88..05c3f10af8d3 100644 --- a/docs/source/features/kv-cache-connector.md +++ b/docs/source/features/kv-cache-connector.md @@ -29,7 +29,7 @@ These methods run on the leader process and drive the connector's behavior. * **`build_connector_meta(self, scheduler_output: SchedulerOutput) -> object`** * **Description**: The core orchestration method. Called during the scheduling phase. It examines the current requests and decides which blocks need to be loaded from or saved to the external store. - * **Arguments**: `scheduler_output` contains information about new requests, blocks allocated, current request states, and the cumulative `RequestData.block_hashes` chain. `block_hashes` is read directly from each KV cache block's stored hash, which the KV cache manager commits as soon as a block becomes full -- the value matches the hash that KV cache events will subsequently emit for the same block. The chain only covers beam 0; the executor rejects `kv_connector_config` at startup when `max_beam_width > 1`, so connectors may assume beam-width-1 inputs. + * **Arguments**: `scheduler_output` contains information about new requests, blocks allocated, current request states, and the cumulative `RequestData.block_hashes` chain. `block_hashes` is read directly from each KV cache block's stored hash, which the KV cache manager commits as soon as a block becomes full, so the value matches the hash that KV cache events will subsequently emit for the same block. The chain only covers beam 0; the executor rejects `kv_connector_config` at startup when `max_beam_width > 1`, so connectors may assume beam-width-1 inputs. * **Returns**: An arbitrary metadata object (picklable) that describes the tasks for the workers. This object is broadcasted to all workers. * **`get_num_new_matched_tokens(self, request: LlmRequest, num_computed_tokens: int) -> tuple[int, bool]`** @@ -47,14 +47,14 @@ These methods run on the leader process and drive the connector's behavior. * **`cancel_load(self, request: LlmRequest, start: int, end: int)`** * **Description**: Optional, with a no-op default. Tells the connector that the runtime will not consume KV it offered from `get_num_new_matched_tokens` for prompt tokens `[start, end)`, so any ownership taken for that range can be released. Offsets are absolute prompt positions, on the same scale as `num_computed_tokens`. - * **When it fires**: only on `KVCacheManagerV2`, which asks during a speculative scheduling pass and resolves the answer later. Two things can happen in between, and both are reported here: the runtime may fail to allocate pages to cover the offer, in which case the request falls back to computing the prefix locally; or the request may be cancelled, time out or fail before it ever reaches a batch, in which case the whole offer is released. A third case -- the local cache overtaking part of the offer because another request committed the same prefix -- is handled by the same callback but cannot arise today, since a request's local match is fixed when its cache is created and only its own completed forward passes extend it. + * **When it fires**: only on `KVCacheManagerV2`, which asks during a speculative scheduling pass and resolves the answer later. Two things can happen in between, and both are reported here: the runtime may fail to allocate pages to cover the offer, in which case the request falls back to computing the prefix locally; or the request may be cancelled, time out or fail before it ever reaches a batch, in which case the whole offer is released. A third case, the local cache overtaking part of the offer because another request committed the same prefix, is handled by the same callback but cannot arise today, since a request's local match is fixed when its cache is created and only its own completed forward passes extend it. * **Caveat**: best-effort. For a synchronous load nothing has been transferred yet, so cancelling is exact. For `is_async=True` the transfer necessarily started inside `get_num_new_matched_tokens`, so it may already be in flight. ##### Serving a prefix on `KVCacheManagerV2` -V1 answers `get_num_new_matched_tokens` from C++ while the block manager holds its radix-tree mutex, so the local match and the query are atomic and the answer is consumed immediately. V2 has no such mutex, and its scheduling pass is speculative -- a prepared request can still be dropped at the token budget, at resize, at multimodal alignment or at cross attention, and retried in a later iteration. +V1 answers `get_num_new_matched_tokens` from C++ while the block manager holds its radix-tree mutex, so the local match and the query are atomic and the answer is consumed immediately. V2 has no such mutex, and its scheduling pass is speculative: a prepared request can still be dropped at the token budget, at resize, at multimodal alignment or at cross attention, and retried in a later iteration. -The contract for connectors is unchanged, and in particular `get_num_new_matched_tokens` is still called **exactly once per request** on both managers -- a request that is asked and then deferred is not asked again when it comes back. What differs is that on V2 the runtime may resolve the answer in a later iteration than the one it asked in, and may by then be unable to honour part or all of it. That is what `cancel_load` reports. +`get_num_new_matched_tokens` is still called **exactly once per request** on both managers, so a request that is asked and then deferred is not asked again when it comes back. What differs is that on V2 the runtime may resolve the answer in a later iteration than the one it asked in, and may by then be unable to honour part or all of it. That is what `cancel_load` reports. #### 2. Worker Interface (`KvCacheConnectorWorker`) @@ -64,6 +64,11 @@ These methods run on all workers (GPU processes) and interact with the actual GP * **Description**: Called at initialization. Provides the worker with the GPU KV cache tensors. * **Arguments**: `kv_cache_tensor` is the underlying storage tensor for the KV cache. +* **`register_kv_cache_layout(self, layout: KvCacheLayout)`** + * **Description**: Called at initialization **instead of** `register_kv_caches` when the KV cache manager is `KVCacheManagerV2`, whose memory cannot be expressed as one tensor: there is one slot address space per pool and one page-index space per layer group. The default implementation raises, so a connector that does not implement it can only run on V1. + * **Arguments**: `layout` describes the byte ranges that repeat per page slot. Each `KvCacheLayerGroupLayout` carries a tuple of `KvCacheRegion`s, and the bytes for page slot `i` of a region live at `region.base + region.stride * i` for `region.size` bytes, or equivalently at `region.as_tensor()[i]`. Page indices arriving in `RequestData.new_block_ids_by_layer_group` are scoped to a layer group and index that group's regions. + * **Why regions rather than a tensor**: because the ranges are described rather than implied, the same structure covers MLA (a pool simply has no `value` buffer), sliding-window and hybrid models (one layer group per window size), and non-uniform slots such as MiniMax-M3's index-K buffer sitting beside K/V, without any of them being a special case. + * **`start_load_kv(self, stream: torch.cuda.Stream)`** * **Description**: Initiates the loading of KV blocks from the external source into the GPU memory. * **Arguments**: `stream` is the CUDA stream where the forward pass is executed in. @@ -81,6 +86,185 @@ These methods run on all workers (GPU processes) and interact with the actual GP * **Description**: Polled by the runtime to check the status of asynchronous operations. * **Returns**: Two lists of request IDs: those that have finished saving, and those that have finished loading. +## Built-in Connectors + +Named presets can be selected without naming a module or class: + +```python +from tensorrt_llm.llmapi.llm_args import KvCacheConnectorConfig + +kv_connector_config = KvCacheConnectorConfig(connector="mooncake-store") +``` + +The available presets are `lmcache`, `lmcache-mp`, `kvbm` and `mooncake-store`. The first three are external packages; `mooncake-store` ships with TensorRT-LLM and is described below. + +### Mooncake distributed store (`mooncake-store`) + +Publishes KV pages into a [Mooncake](https://github.com/kvcache-ai/Mooncake) store, a shared CPU memory pool addressed by content, so a prefix computed by one engine can be replayed by another. Regular block reuse cannot do this, because it never leaves the instance that computed the prefix. + +This is a **different component** from the Mooncake transfer engine that the C++ cache transceiver uses for disaggregated prefill/decode handoff. That moves KV point to point between two known peers; this publishes pages into a pool that any peer can read. The two compose: a context server can write pages into the store and still hand off to a generation server over NIXL. + +#### Requirements + +* `KVCacheManagerV2` (`kv_cache_config.use_kv_cache_manager_v2: true`), since that is the manager that can describe its pools through `register_kv_cache_layout`. +* The Mooncake Python bindings: `pip install mooncake-transfer-engine`. These are installed in the release container; the source build of the C++ transfer engine does not provide them. +* A reachable Mooncake master (and metadata server, unless using `P2PHANDSHAKE`). See the [Mooncake documentation](https://kvcache-ai.github.io/Mooncake/). `trtllm-serve` can start one for a single engine; see below. +* GPU-only KV cache tiers: set `kv_cache_config.host_cache_size: 0` and `disk_cache_size: 0`. A page evicted to another tier has its GPU slot reassigned, which would invalidate the addresses registered with the store. + +#### Configuration + +Describe the pool in `kv_connector_config.mooncake_store` and `trtllm-serve` provisions it during bringup: it resolves the master, renders the client config, and exports `MOONCAKE_CONFIG_PATH` before the ranks that open store handles are spawned. + +```yaml +kv_connector_config: + connector: mooncake-store + mooncake_store: + master_server_address: 10.0.0.1:50051 # a master with its own lifetime + protocol: rdma + device_name: mlx5_0 + global_segment_size: 32GiB + local_buffer_size: 1GiB +``` + +Replacing `master_server_address` with `launch_master: true` makes the server start a `mooncake_master` itself and use it, so a single-instance deployment needs nothing prepared outside `trtllm-serve`. **That master lives and dies with the server**, which makes it wrong for anything else: several engines that should share one pool would each get their own, and a pool meant to survive a restart cannot be owned by the thing restarting. + +A master started this way still publishes its address, to `master.addr` in the run directory and to `master_address_file` if one is named. That is how other processes, the donors below above all, find a pool this server owns, and how a finished run's logs say which master it used: + +```yaml +mooncake_store: + launch_master: true + master_address_file: /shared/master.addr +``` + +The two cases above, several engines or surviving a restart, run the master as its own command instead: + +```bash +trtllm-serve mooncake_master --rpc_port 50051 --address_file /shared/master.addr +``` + +The pool then lasts as long as that command, independently of any engine. `--address_file` receives `host:port` once the master accepts connections, and `master_server_address` accepts `file://` as well as a literal address: + +```yaml +kv_connector_config: + connector: mooncake-store + mooncake_store: + master_server_address: file:///shared/master.addr +``` + +This is what makes a master reachable without anyone writing its address down. Under a scheduler its host is not known when the configs are written; publishing it to a file the configs already name closes that gap, and a server reading the file waits for it, so the master and the engines can be started in any order. The file is removed when the master stops, so a stale address is never dialed. + +`TRTLLM_MOONCAKE_MASTER_BINARY` overrides the binary a launched master runs, and `TRTLLM_MOONCAKE_MASTER_TIMEOUT` (default 60s) sets how long startup waits for any master to accept connections or publish its address. Without that wait, a master that is not there yet fails inside every rank after the model has loaded. Set `TRTLLM_MOONCAKE_RUN_DIR` to keep the generated client config and the master's log, which are otherwise in a temporary directory removed at shutdown. + +#### Servers whose ranks the launcher starts + +Provisioning happens in the server process and reaches the ranks that open store handles by exporting `MOONCAKE_CONFIG_PATH` for them to inherit. That holds when the LLM constructor spawns them. It does not when the launcher starts one task per rank, as `trtllm-llmapi-launch` under a scheduler does, because those ranks were already running. + +Naming a shared run directory covers that case: the rendered config is read back from `$TRTLLM_MOONCAKE_RUN_DIR/mooncake.json` by any rank that inherited no path, so every rank of a multi-GPU server joins the pool its own leader provisioned. The directory has to be one they all see, which under a scheduler means the job's own, and it is where the master's log and published address already go: + +```bash +export TRTLLM_MOONCAKE_RUN_DIR=/shared/run/$SLURM_JOB_ID +srun trtllm-llmapi-launch trtllm-serve "$model" --config ctx.yaml +``` + +Without it, a rank that inherited nothing fails during bringup naming `MOONCAKE_CONFIG_PATH`, rather than serving without a store. + +#### Reading bringup in the log + +Everything the pool is assembled from is logged under the `mooncake-store:` prefix before the model loads, because a pool that came up wrong is otherwise visible only as a low hit rate hours later. In order: the run directory, the master's command line and pid, the address it published and where, the rendered client config in full, and the capacity each rank will contribute. A server lending memory logs the master it resolved, the segment in both GiB and bytes, and the transport, so that a size string parsed wrong is caught before the pool starts evicting far too eagerly. + +Both waits report progress every five seconds, since waiting for a master in another job is normal and indistinguishable from a hang if it is silent. A master that dies during startup has the tail of its own log quoted in the failure, which is where the reason, a port in use or a bad flag, actually is. + +#### Pool capacity + +Capacity comes only from processes that open a store handle, and `global_segment_size` is what each contributes, so the pool is that value times the number of such processes. In a disaggregated deployment the connector belongs on the context servers only, which makes every byte of the pool prefill-node memory: prefill's DRAM caching prefill's GPUs, largely duplicating what `kv_cache_config.host_cache_size` already does. + +To give the pool memory from nodes whose engines run no connector, ask those servers to lend it: + +```yaml +# generation server: no connector, memory only +mooncake_donation: + master_server_address: file:///shared/master.addr + segment_size: 320GiB + protocol: rdma + device_name: mlx5_1 +``` + +`trtllm-serve` then holds that segment for as long as the server runs, so a generation node holds pages prefill wrote while its own engine stays connector-free and keeps its cache transceiver for the prefill-to-decode handoff. The server is ready only once the segment is mounted, which makes its readiness the signal that the pool has this capacity. + +Lending memory is deliberately outside `kv_connector_config`, and not a `TRTLLM_MOONCAKE_STORE_ROLE` either. Both of those attach a connector, and a connector reads or writes: `producer`, `consumer` and `both` all describe traffic, and none of them means "contribute memory only". Expressing capacity there would therefore start this server using the store. Capacity and traffic are separate, and configured separately. + +Size is charged **per server process, not per rank**, unlike `global_segment_size`. Two servers on one node lend twice this. The memory is charged to the process and competes with everything else on the node, `kv_cache_config.host_cache_size` above all, so size the two together. + +A node that runs no server at all can still lend, as its own command: + +```bash +trtllm-serve mooncake_donor --master_server_address file:///shared/master.addr \ + --segment_size 160GiB --protocol rdma --device_name mlx5_0 +``` + +Topology can equally come from a JSON file named by `MOONCAKE_CONFIG_PATH`, using the same schema as the vLLM Mooncake store connector so one deployment can point both engines at the same pool: + +```json +{ + "metadata_server": "http://127.0.0.1:8080/metadata", + "master_server_address": "127.0.0.1:50051", + "protocol": "rdma", + "device_name": "mlx5_0", + "global_segment_size": "32GiB", + "local_buffer_size": "1GiB" +} +``` + +Only `master_server_address` is required. `metadata_server` may be left out, in which case it is `P2PHANDSHAKE`, Mooncake's peer-to-peer handshake, which is also what `mooncake_store` and `mooncake_donation` default to; the example above names a metadata service instead. + +An inherited `MOONCAKE_CONFIG_PATH` wins over `mooncake_store` and is logged as doing so, so an orchestrator that already provisions the pool, as the SLURM benchmark harness does, keeps working unchanged. + +Three further settings are TensorRT-LLM's rather than Mooncake's, and stay in the environment because they are per process rather than per pool: + +| Variable | Default | Meaning | +|---|---|---| +| `TRTLLM_MOONCAKE_STORE_ROLE` | `both` | `producer` writes only, `consumer` reads only, `both` does both. | +| `TRTLLM_MOONCAKE_STORE_PREFIX` | `trtllm` | Leading component of every key, for isolating deployments that share a pool. | +| `TRTLLM_MOONCAKE_STORE_MODEL_KEY` | model directory basename | Identity keys are namespaced by. Two engines share cache only when they agree on it, so the default is the basename rather than the full path, since the same checkpoint is routinely mounted elsewhere on another host, which is exactly what sharing is for. | + +In a disaggregated deployment, run context servers as `both` and leave generation servers unconfigured. Generated tokens are rarely a reused prefix, so writing them costs bandwidth for no hit rate. + +#### Partial block reuse is forced off + +`kv_cache_config.enable_partial_reuse` is set to `false` when this connector is configured, with a warning, whether or not it was requested explicitly. It defaults to `true`, so most deployments will see that warning. + +The store is addressed by whole blocks. The connector is handed the device match as `num_computed_tokens` and offers only blocks beyond it, but it can resume only from a block boundary, so when the device match ends mid-block it declines the lookup and the store is not consulted at all. Partial reuse is precisely what puts the match off a boundary, so it trades part of one block of device reuse for every stored block of the remaining prefix. Measured on MiniMax-M3, leaving it enabled declined 97.2% of lookups and left actual prompt cache read at 35% against a 96% ceiling; forcing it off raised that to 94% and roughly doubled throughput. + +#### How it keys pages + +`KVCacheManagerV2` reports `RequestData.block_hashes` empty, so the connector derives block identity itself: a blake2b chain where each block's hash covers its own tokens *and* every token before it, seeded by the request's `cache_salt`. A key is `//wr/lg/tb/`. The namespace pins down everything that would make the stored bytes mean something different, so a mismatched shard count, layer group or page geometry reads as a cache miss rather than as garbage. + +The value for one key is the concatenation of that layer group's regions for one page slot, handed to Mooncake's multi-buffer batch APIs as a list of `(address, size)` pairs. + +#### Transfer behavior + +* **Loads are synchronous**, performed in `start_load_kv` before the forward pass. A failed load raises: the runtime has already counted those tokens as computed, so a partial load is a wrong answer rather than a slow one. +* **Saves are asynchronous**, handed to a background thread behind a CUDA event recorded on the forward stream. The pages are only complete once the pass that wrote them retires, and blocking the executor loop on an RDMA write is the cost the store exists to avoid. The leader reports such requests as saving asynchronously, so their pages stay pinned until `get_finished` confirms the writes landed. A dropped save is logged rather than raised, since it only costs a future cache miss. +* Pages the store already holds are skipped, so several ranks or instances converging on the same prefix write it once. + +#### Unsupported configurations + +These are rejected at startup, before any request is admitted: + +| Configuration | Reason | +|---|---| +| Context parallelism | A rank holds a slice of the sequence rather than whole blocks of it, so one key would name different bytes on different ranks. | +| Sliding-window attention / VSWA | A page's validity depends on where the window sits, which is a property of the request that read it rather than of the tokens it holds. | +| MiniMax-M3 with `sparse_disable_index_value: false` | The index-V cache is a plain tensor outside the paged pools, so a replayed prefix would pair stored index-K with stale index-V. Disaggregated serving applies the same restriction. | +| Pipeline parallelism | Untested rather than unsound. Use tensor parallelism. | +| `KVCacheManagerV1` | Identity here is a per-layer-group hash chain; V1 supplies real block hashes over a single flat block space. | + +Beam search, attention data parallelism, non-GPU cache tiers and Mamba caches are rejected for all connectors by the executor. + +#### Example + +`examples/llm-api/configs/trtllm_mooncake_store_connector_extra.yaml` is a starting point for `trtllm-serve`. + ## Example Implementation The file `examples/llm-api/llm_kv_cache_connector.py` provides a reference implementation of a **Persistent KV Cache**. diff --git a/examples/disaggregated/slurm/benchmark/disaggr_torch.slurm b/examples/disaggregated/slurm/benchmark/disaggr_torch.slurm index d73430ba5820..dcbd5588936c 100644 --- a/examples/disaggregated/slurm/benchmark/disaggr_torch.slurm +++ b/examples/disaggregated/slurm/benchmark/disaggr_torch.slurm @@ -138,6 +138,26 @@ else echo "TensorRT-LLM environment variables saved to ${full_logdir}/env_vars.json" fi +# The Mooncake store bindings, when a worker config asks for the connector. +# Images built from this repo bake them in (docker/common/install_mooncake.sh), +# so this only confirms they are there: a job that discovers the gap later +# fails once per rank, deep in engine startup. +mooncake_enabled=false +if grep -qs "mooncake-store" "${full_logdir}/ctx_config.yaml" "${full_logdir}/gen_config.yaml"; then + mooncake_enabled=true + # Both halves of the wheel are checked because they fail at different + # times: the connector needs mooncake.store in every context rank, and + # 'trtllm-serve mooncake_master' needs the binary on PATH. + if ! srun --container-name=${container_name} \ + --container-mounts=${container_mount} --no-container-mount-home \ + --mpi=pmix --overlap -N 1 -n 1 \ + bash -c 'python3 -c "import mooncake.store" && command -v mooncake_master' \ + &> ${full_logdir}/2_check_mooncake.log; then + cleanup_on_failure "A worker config requests the mooncake-store connector, but this image has neither the mooncake.store bindings nor the mooncake_master binary. Build the image from this repo so that docker/common/install_mooncake.sh runs. Check ${full_logdir}/2_check_mooncake.log for details" + fi + echo "Mooncake store bindings found in the image" +fi + # Get node lists and replace the placeholder with the actual node names echo "SLURM_NODELIST: ${SLURM_NODELIST}" all_nodes=($(scontrol show hostname $SLURM_NODELIST | sort)) @@ -154,6 +174,20 @@ client_cmds_base_file=${full_logdir}/client_cmds_base.sh client_cmds_file=${full_logdir}/client_cmds.sh replace_placeholder "${client_cmds_base_file}" "${all_nodes_str}" "${client_cmds_file}" +# The pool is described in the worker configs and provisioned by trtllm-serve +# during its own bringup: the context server starts the master, renders the +# client config and publishes the master's address, and the generation servers +# read that address to lend the pool their memory. Nothing here starts, waits +# for or configures any of it. +# +# The one value a config written before submission cannot know is this job's log +# directory, which the master's address is published into, so a __LOG_DIR__ +# placeholder in either worker config is filled in here. +if [ "${mooncake_enabled}" = "true" ]; then + sed -i "s|__LOG_DIR__|${full_logdir}|g" \ + "${full_logdir}/ctx_config.yaml" "${full_logdir}/gen_config.yaml" +fi + # Per-worker hostfile / gpu_map files for srun --distribution=arbitrary. # submit.py emits *_base.txt with ; rewrite them here. for base_file in "${full_logdir}"/hostfile_*_base.txt "${full_logdir}"/gpu_map_*_base.txt; do @@ -176,6 +210,20 @@ cat ${start_server_cmds_file} | while read cmd; do done echo "Server is ready!" +# A connector that failed to open its store handle does not stop the worker from +# serving, it just never hits, so surface the startup lines here rather than +# leaving them to be discovered after the benchmark. The registration line +# carries the bytes/page figure the pool sizing depends on. +if [ "${mooncake_enabled}" = "true" ]; then + echo "Mooncake store startup lines from the context workers:" + if ! grep -h "mooncake-store" "${full_logdir}"/3_output_CTX_*.log 2>/dev/null; then + echo " WARNING: no mooncake-store lines found. The connector may not have" \ + "loaded; check ${full_logdir}/3_output_CTX_*.log and" \ + "${full_logdir}/mooncake_master.log. The benchmark will still run," \ + "but without the store." + fi +fi + # Start client commands echo "Starting client commands from ${client_cmds_file}..." while read -r cmd <&3; do @@ -186,6 +234,78 @@ while read -r cmd <&3; do fi done 3< "${client_cmds_file}" +# Collect the store's traffic into one file. The per-event lines live at DEBUG +# in the worker logs (module _torch), so this is only populated when a config +# asks for that verbosity. The counts are what show the store did work rather +# than merely started. +if [ "${mooncake_enabled}" = "true" ]; then + mooncake_summary="${full_logdir}/9_mooncake_summary.log" + { + # The address file is retracted when the master stops, so fall back to + # the context server's log, which still names the master the run used. + echo "master: $(tr -d '[:space:]' < "${full_logdir}/master.addr" 2>/dev/null \ + || grep -hoE "master at [0-9.]+:[0-9]+" "${full_logdir}"/3_output_CTX_*.log 2>/dev/null \ + | head -n 1 | awk '{print $3}' || echo unknown)" + echo + echo "== startup ==" + grep -h "mooncake-store.*\(ready\|registered layout\)" "${full_logdir}"/3_output_CTX_*.log 2>/dev/null || echo "(none)" + echo + echo "== event counts ==" + for pattern in "matched" "loaded" "failed to load" "failed to save" \ + "lookup failed" "could not reserve connector prefix"; do + count=$(grep -h "mooncake-store.*${pattern}" "${full_logdir}"/3_output_CTX_*.log 2>/dev/null | wc -l || true) + echo "${pattern}: ${count}" + done + echo + # Where the blocks physically went. The master names a segment for + # every allocation and a segment is one client process's donated + # memory, so grouping by segment host shows how much of the pool lives + # on a decode node rather than on the prefill node that computed it. + # Without a donor this section shows only the prefill node. + echo "== block placement by segment host ==" + echo "(lending hosts: $(grep -hoE "GiB of [0-9.]+ is now part of the pool" \ + "${full_logdir}"/3_output_GEN_*.log 2>/dev/null \ + | awk '{print $3}' | sort -u | paste -sd, - || echo none))" + grep -o "allocation_succeeded size=[0-9]* segment=[0-9.]*:[0-9]*" \ + "${full_logdir}/mooncake_master.log" 2>/dev/null \ + | awk '{ + sub(/size=/, "", $2); sub(/segment=/, "", $3); + split($3, parts, ":"); host = parts[1]; port = parts[2]; + allocs[host]++; bytes[host] += $2; + if (!((host, port) in seen)) { + seen[host, port] = 1; + segs[host] = segs[host] " " port; + } + total_allocs++; total_bytes += $2; + } + END { + if (total_allocs == 0) { print "(no allocations)"; exit } + for (h in allocs) { + printf "%-16s pages=%-7d %8.2f GiB %5.1f%% of pool contents segments:%s\n", + h, allocs[h], bytes[h] / 1073741824, + 100 * bytes[h] / total_bytes, segs[h]; + } + printf "%-16s pages=%-7d %8.2f GiB\n", "TOTAL", total_allocs, total_bytes / 1073741824; + }' || echo "(could not parse master log)" + echo + # Lending memory happens inside the generation servers, so their own + # logs are where a segment that never mounted shows up. + echo "== lent segments ==" + grep -h "mooncake-store: .*\(lending memory\|part of the pool\|withdrew\)" \ + "${full_logdir}"/3_output_GEN_*.log 2>/dev/null || echo "(none)" + echo + echo "== pool bringup (context server) ==" + grep -h "mooncake-store:" "${full_logdir}"/3_output_CTX_*.log 2>/dev/null \ + | head -n 30 || echo "(none)" + echo + # The master's own glog output, not anything TensorRT-LLM writes. + echo "== master ==" + tail -n 50 "${full_logdir}/mooncake_master.log" 2>/dev/null || echo "(no master log)" + } > "${mooncake_summary}" 2>&1 + echo "Mooncake store summary written to ${mooncake_summary}" + cat "${mooncake_summary}" +fi + echo "Job completed successfully, total runtime: $SECONDS seconds" # try to kill the server and workers diff --git a/examples/disaggregated/slurm/benchmark/start_worker.sh b/examples/disaggregated/slurm/benchmark/start_worker.sh index 0a5b5897b773..cf93aa095678 100644 --- a/examples/disaggregated/slurm/benchmark/start_worker.sh +++ b/examples/disaggregated/slurm/benchmark/start_worker.sh @@ -54,6 +54,47 @@ fi echo "config_file: ${config_file}" +# The mooncake-store pool is described in the worker config and provisioned by +# trtllm-serve during bringup. Anchoring its run directory here keeps the +# master's log, the rendered client config and the published address in the +# job's log directory rather than in a temporary directory that shutdown +# removes, and it is how the ranks srun started, which never inherited the +# leader's environment, find that client config. An inherited +# MOONCAKE_CONFIG_PATH still wins, so an externally managed pool stays reachable. +export TRTLLM_MOONCAKE_RUN_DIR="${log_dir}" + +# The generation servers wait for a master the context server starts. Both are +# launched together and the master comes up before its model loads, but the wait +# spans container start on another node, so it is given far more than the 60s +# default. Too short fails the job; too long costs nothing when the master is +# already there. +export TRTLLM_MOONCAKE_MASTER_TIMEOUT="${TRTLLM_MOONCAKE_MASTER_TIMEOUT:-900}" + +# MiniMax-M3's MSA sparse attention JIT-compiles its FMHA kernels on first use, +# from inside the attention forward pass. One TP rank runs ninja while the +# others block on a file lock, so an uncached variant stalls the whole executor +# loop for ~8s, or ~70s when an iteration needs several. The cache defaults to +# ~/.cache, which is thrown away because the container is started with +# --no-container-mount-home, so every job would pay the compiles again during +# serving. Anchoring it next to this script puts it on the mounted filesystem +# at a path identical across jobs, so only the first run compiles. +if [ -z "${MINFER_FMHA_CACHE_DIR:-}" ]; then + export MINFER_FMHA_CACHE_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/.cache/minfer/fmha_sm100" + mkdir -p "${MINFER_FMHA_CACHE_DIR}" + echo "MINFER_FMHA_CACHE_DIR: ${MINFER_FMHA_CACHE_DIR}" +fi + +# Per-transfer KV timings (size, queue/transfer latency, throughput) as CSV next +# to the worker logs. These separate slow prefill from a slow prefill-to-decode +# handoff, which the aggregate benchmark numbers cannot. An explicit setting +# wins, and KV_TRANSFER_PERF_LOG=false turns it off. +if [ "${KV_TRANSFER_PERF_LOG:-true}" = "true" ] \ + && [ -z "${TLLM_KV_TRANSFER_PERF_LOG_FILE:-}" ]; then + export TLLM_ENABLE_CACHE_TRANSFER_PERF_INFO=1 + export TLLM_KV_TRANSFER_PERF_LOG_FILE="${log_dir}/kv_transfer_perf" + echo "TLLM_KV_TRANSFER_PERF_LOG_FILE: ${TLLM_KV_TRANSFER_PERF_LOG_FILE}" +fi + nsys_prefix="" if [ "${enable_nsys}" != "true" ]; then echo "nsys is not enabled, start normal flow" diff --git a/examples/llm-api/configs/trtllm_mooncake_store_connector_extra.yaml b/examples/llm-api/configs/trtllm_mooncake_store_connector_extra.yaml new file mode 100644 index 000000000000..350c77b65941 --- /dev/null +++ b/examples/llm-api/configs/trtllm_mooncake_store_connector_extra.yaml @@ -0,0 +1,49 @@ +# Extra LLM API options for trtllm-serve with the Mooncake store KV connector. +# +# Offloads KV pages to a Mooncake distributed store, a shared CPU memory pool +# addressed by content, so a prefix computed by one engine can be replayed by +# another. This is a different component from the Mooncake transfer engine used +# by the C++ cache transceiver for disaggregated prefill/decode handoff; the two +# compose rather than conflict. +# +# Prerequisites: +# - Mooncake Python bindings: pip install mooncake-transfer-engine +# (present in the release container; the C++ source build does not +# provide MooncakeDistributedStore) +# - A running Mooncake master, and a metadata server unless using +# P2PHANDSHAKE. See https://kvcache-ai.github.io/Mooncake/ +# - MOONCAKE_CONFIG_PATH pointing at a Mooncake JSON config, for example: +# { +# "metadata_server": "http://127.0.0.1:8080/metadata", +# "master_server_address": "127.0.0.1:50051", +# "protocol": "rdma", +# "device_name": "mlx5_0", +# "global_segment_size": "32GiB", +# "local_buffer_size": "1GiB" +# } +# +# Optional environment overrides: +# TRTLLM_MOONCAKE_STORE_ROLE producer | consumer | both (default both) +# TRTLLM_MOONCAKE_STORE_PREFIX key prefix, to isolate deployments sharing +# one pool (default trtllm) +# TRTLLM_MOONCAKE_STORE_MODEL_KEY identity keys are namespaced by +# (default: model directory basename) +# +# Example: +# export MOONCAKE_CONFIG_PATH=/path/to/mooncake.json +# trtllm-serve --backend pytorch --host 0.0.0.0 --port 8000 \ +# --extra_llm_api_options /path/to/this/file + +kv_cache_config: + # The connector describes its pools through register_kv_cache_layout, which + # only KVCacheManagerV2 implements. + use_kv_cache_manager_v2: true + # Local reuse still runs first; the store serves whatever the device missed. + enable_block_reuse: true + # GPU-only tiers are required: a page evicted to host or disk has its GPU slot + # reassigned, which would invalidate the addresses registered with the store. + host_cache_size: 0 + disk_cache_size: 0 + +kv_connector_config: + connector: mooncake-store diff --git a/scripts/attribution/scan/metadata/mooncake.yml b/scripts/attribution/scan/metadata/mooncake.yml index c8d51e0f81b8..b1e6dab2c141 100644 --- a/scripts/attribution/scan/metadata/mooncake.yml +++ b/scripts/attribution/scan/metadata/mooncake.yml @@ -1,5 +1,8 @@ name: mooncake -description: Mooncake transfer engine for distributed KV cache +description: Mooncake transfer engine and distributed store for distributed KV cache source: container directory_matches: - /usr/local/Mooncake +- mooncake +basename_matches: +- mooncake_transfer_engine diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 92843c42f3aa..6bb66c0a53e2 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -2714,6 +2714,8 @@ def create_py_executor_instance( cross_kv_cache_manager=cross_kv_cache_manager, no_schedule_until_state=no_schedule_until_state, enable_prefix_aware_scheduling=enable_prefix_aware_scheduling, + max_input_len=max_seq_len + if max_seq_len is not None else 0x7fffffff, ) elif (scheduler_config is not None and scheduler_config.use_python_scheduler): diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py new file mode 100644 index 000000000000..bae898f8619a --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py @@ -0,0 +1,81 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""KV cache connector backed by a Mooncake distributed store. + +Offloads KV pages to a shared CPU memory pool so a prefix computed by one engine +can be replayed by another, which regular block reuse cannot do because it never +leaves the instance that computed it. + +This is a different component from the Mooncake transfer engine that the C++ +cache transceiver uses for disaggregated prefill/decode handoff: that moves KV +point to point between two known peers, while this one publishes pages into a +pool addressed by content. The two compose, so a context server can write pages +here and still hand off over NIXL. + +Requires `KVCacheManagerV2`, the manager that can describe its pools to a +connector through `register_kv_cache_layout`, and the Mooncake Python bindings +(`pip install mooncake-transfer-engine`). + +Enable it with:: + + kv_connector_config = KvCacheConnectorConfig(connector="mooncake-store") + +with `MOONCAKE_CONFIG_PATH` pointing at a Mooncake JSON config. Describing the +pool in `KvCacheConnectorConfig.mooncake_store` instead lets `trtllm-serve` +provision it during bringup, so no external script has to; see `master.py`. + +Capacity comes only from processes that open a store handle, which in a +disaggregated deployment is the context servers alone. `donor.py` lends a +node's memory to the pool without giving it a connector. + +By default the KV pools themselves are registered with Mooncake, which requires +GPUDirect RDMA. Where that is unavailable, `"stage_through_host": true` in the +JSON config, or `TRTLLM_MOONCAKE_STORE_STAGE_THROUGH_HOST=1`, routes pages +through a pinned host buffer instead; see `staging.py`. +""" + +from .config import MooncakeStoreConnectorConfig, StoreRole, parse_size +from .donor import DEFAULT_DONOR_LOCAL_BUFFER_SIZE, donate_segment, maybe_donate_segment +from .master import ( + local_address, + master_timeout, + maybe_provision_pool, + provision_pool, + resolve_device_name, + resolve_master_address, + running_master, + wait_for_master, +) +from .scheduler import MooncakeStoreConnectorScheduler +from .worker import MooncakeStoreConnectorWorker + +__all__ = [ + "DEFAULT_DONOR_LOCAL_BUFFER_SIZE", + "MooncakeStoreConnectorConfig", + "MooncakeStoreConnectorScheduler", + "MooncakeStoreConnectorWorker", + "StoreRole", + "donate_segment", + "local_address", + "master_timeout", + "maybe_donate_segment", + "maybe_provision_pool", + "parse_size", + "provision_pool", + "resolve_device_name", + "resolve_master_address", + "running_master", + "wait_for_master", +] diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/addressing.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/addressing.py new file mode 100644 index 000000000000..712e34b3978e --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/addressing.py @@ -0,0 +1,152 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Turning a `KvCacheLayout` into addresses Mooncake can transfer. + +Mooncake's batch APIs take, per key, a list of `(address, size)` buffers. That +is exactly the shape of a V2 page: a layer group's regions each contribute one +byte range at `base + stride * page_index`, and the concatenation of those +ranges in region order is the page's payload. + +Region order is therefore the value's serialization, and it is stable for a +given model and parallel layout because `build_kv_cache_layout_v2` derives it +from the allocator's own aggregation. `bytes_per_page` goes into +the key namespace to keep a geometry change from being read as a valid page. +""" + +from typing import Dict, Iterable, List, Sequence, Tuple + +from ..kv_cache_layout import KvCacheLayout, KvCacheRegion + +__all__ = ["PageAddressing", "merge_intervals"] + + +def merge_intervals(intervals: Iterable[Tuple[int, int]]) -> List[Tuple[int, int]]: + """Collapse `(start, end)` byte ranges into a minimal disjoint cover. + + Registration is per range and a range may not be registered twice, but + several regions routinely live inside one pool allocation: sliding-window + layer groups share it, and a non-uniform slot (MiniMax-M3's index-K sitting + beside K/V) splits one pool into several regions. Merging first means the + caller does not have to know which case it is in. + """ + ordered = sorted((int(start), int(end)) for start, end in intervals if end > start) + merged: List[Tuple[int, int]] = [] + for start, end in ordered: + if merged and start <= merged[-1][1]: + previous_start, previous_end = merged[-1] + merged[-1] = (previous_start, max(previous_end, end)) + else: + merged.append((start, end)) + return merged + + +class PageAddressing: + """Resolves `(layer group, page index)` to the byte ranges of that page.""" + + def __init__(self, layout: KvCacheLayout): + self._layout = layout + self._regions: Dict[int, Tuple[KvCacheRegion, ...]] = {} + self._bytes_per_page: Dict[int, int] = {} + self._num_slots: Dict[int, int] = {} + for group in layout.groups: + if not group.regions: + raise ValueError( + f"layer group {group.layer_group_id} has no KV regions; there " + "is nothing for the connector to transfer" + ) + self._regions[group.layer_group_id] = group.regions + self._bytes_per_page[group.layer_group_id] = group.bytes_per_page + # Every region of a group is drawn from the same pool group, so they + # share a slot count; disagreement would mean the page index space is + # not the single space the layout documents. + slot_counts = {region.num_slots for region in group.regions} + if len(slot_counts) != 1: + raise ValueError( + f"layer group {group.layer_group_id} mixes slot counts " + f"{sorted(slot_counts)}; page indices would be ambiguous" + ) + self._num_slots[group.layer_group_id] = slot_counts.pop() + + @property + def layout(self) -> KvCacheLayout: + """The layout this addressing was built from.""" + return self._layout + + @property + def layer_group_ids(self) -> Tuple[int, ...]: + """Layer group ids covered, in layout order.""" + return tuple(group.layer_group_id for group in self._layout.groups) + + @property + def tokens_per_block(self) -> int: + """Tokens held by one page.""" + return self._layout.tokens_per_block + + def bytes_per_page(self, layer_group_id: int) -> int: + """Total payload size of one page of `layer_group_id`.""" + return self._bytes_per_page[layer_group_id] + + def num_slots(self, layer_group_id: int) -> int: + """Number of page slots addressable in `layer_group_id`.""" + return self._num_slots[layer_group_id] + + def buffers(self, layer_group_id: int, page_index: int) -> Tuple[List[int], List[int]]: + """Addresses and sizes of one page, in the order they concatenate. + + Args: + layer_group_id: Layer group the page index is scoped to. + page_index: Page slot index within that group. + + Returns: + Parallel lists of device addresses and byte counts. + """ + regions = self._regions[layer_group_id] + num_slots = self._num_slots[layer_group_id] + if not 0 <= page_index < num_slots: + raise IndexError( + f"page index {page_index} out of range [0, {num_slots}) for layer " + f"group {layer_group_id}" + ) + addresses = [region.base + region.stride * page_index for region in regions] + sizes = [region.size for region in regions] + return addresses, sizes + + def registration_ranges(self) -> List[Tuple[int, int]]: + """Byte ranges to hand to `register_buffer`, deduplicated and merged. + + A region's slots are strided rather than packed, so the range covering it + is the whole span from the first slot to the end of the last. Registering + the span is what makes every slot's address valid for RDMA, and merging + keeps a shared pool from being registered once per region. + """ + spans: List[Tuple[int, int]] = [] + for regions in self._regions.values(): + for region in regions: + span_end = region.base + region.stride * (region.num_slots - 1) + region.size + spans.append((region.base, span_end)) + return merge_intervals(spans) + + def describe(self) -> str: + """A one-line summary for startup logs.""" + parts: Sequence[str] = [ + f"lg{group.layer_group_id}(" + f"layers={len(group.layer_ids)}, " + f"regions={len(group.regions)}, " + f"bytes/page={group.bytes_per_page}, " + f"slots={self._num_slots[group.layer_group_id]}, " + f"window={group.window_size})" + for group in self._layout.groups + ] + return f"tokens_per_block={self.tokens_per_block}, " + ", ".join(parts) diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py new file mode 100644 index 000000000000..577bd22676dd --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py @@ -0,0 +1,268 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Configuration for the Mooncake store KV cache connector. + +Topology settings are read from the JSON file named by `MOONCAKE_CONFIG_PATH`, +the same file and environment variable the vLLM Mooncake store connector uses, +so one deployment can point both engines at the same pool. + +`KvCacheConnectorConfig` carries no free-form dictionary, so the two settings +that are TensorRT-LLM's rather than Mooncake's, the read/write role and the key +prefix, are also taken from the environment. +""" + +import json +import os +import re +from dataclasses import dataclass +from enum import Enum +from typing import Any, Optional + +__all__ = [ + "CLIENT_CONFIG_NAME", + "CONFIG_PATH_ENV", + "MooncakeStoreConnectorConfig", + "ROLE_ENV", + "RUN_DIR_ENV", + "STAGE_THROUGH_HOST_ENV", + "StoreRole", + "provisioned_config_path", +] + +CONFIG_PATH_ENV = "MOONCAKE_CONFIG_PATH" +#: Where a server keeps the client config it renders and the master's log. Set +#: it to keep them after shutdown; otherwise they live in a temporary directory. +RUN_DIR_ENV = "TRTLLM_MOONCAKE_RUN_DIR" +#: Name the rendered client config takes in the run directory. +CLIENT_CONFIG_NAME = "mooncake.json" +ROLE_ENV = "TRTLLM_MOONCAKE_STORE_ROLE" +CACHE_PREFIX_ENV = "TRTLLM_MOONCAKE_STORE_PREFIX" +MODEL_KEY_ENV = "TRTLLM_MOONCAKE_STORE_MODEL_KEY" +STAGE_THROUGH_HOST_ENV = "TRTLLM_MOONCAKE_STORE_STAGE_THROUGH_HOST" + +DEFAULT_GLOBAL_SEGMENT_SIZE = 3355443200 +DEFAULT_LOCAL_BUFFER_SIZE = 1073741824 +DEFAULT_CACHE_PREFIX = "trtllm" +DEFAULT_STAGING_BUFFER_SIZE = 536870912 +#: Mooncake's own peer-to-peer handshake, which keeps a separate metadata +#: process out of the deployment. Nothing else is a sensible fallback: an empty +#: connstring is not one of the forms `store.setup` accepts, so a config that +#: leaves the field out means this rather than meaning no metadata service. +DEFAULT_METADATA_SERVER = "P2PHANDSHAKE" + +_TRUE = {"1", "true", "yes", "on"} +_FALSE = {"0", "false", "no", "off"} + +_SIZE_UNITS = { + "": 1, + "b": 1, + "k": 1000, + "kb": 1000, + "m": 1000**2, + "mb": 1000**2, + "g": 1000**3, + "gb": 1000**3, + "t": 1000**4, + "tb": 1000**4, + "kib": 1024, + "mib": 1024**2, + "gib": 1024**3, + "tib": 1024**4, +} +_SIZE_RE = re.compile(r"^\s*([0-9]+(?:\.[0-9]+)?)\s*([a-zA-Z]*)\s*$") + + +class StoreRole(Enum): + """Which directions of traffic this engine is allowed to drive. + + A disaggregated deployment typically runs context servers as `both` and + leaves generation servers unconfigured: generated tokens are rarely a reused + prefix, so writing them costs bandwidth for no hit rate. + """ + + PRODUCER = "producer" + CONSUMER = "consumer" + BOTH = "both" + + @property + def loads(self) -> bool: + """Whether this role reads previously stored KV back onto the GPU.""" + return self is not StoreRole.PRODUCER + + @property + def saves(self) -> bool: + """Whether this role writes newly computed KV into the store.""" + return self is not StoreRole.CONSUMER + + +def parse_size(value: Any) -> int: + """Accept either a byte count or a suffixed string such as `"4GiB"`.""" + if isinstance(value, bool): + raise ValueError(f"expected a size, got {value!r}") + if isinstance(value, int): + return value + if isinstance(value, float): + return int(value) + match = _SIZE_RE.match(str(value)) + if match is None: + raise ValueError(f"cannot parse size {value!r}") + magnitude, unit = match.groups() + scale = _SIZE_UNITS.get(unit.lower()) + if scale is None: + raise ValueError(f"unknown size unit {unit!r} in {value!r}") + return int(float(magnitude) * scale) + + +def provisioned_config_path() -> Optional[str]: + """The client config a server on this node rendered, if there is one. + + `provision_pool` writes one and exports `MOONCAKE_CONFIG_PATH`, which the + ranks the LLM constructor spawns inherit. Ranks an external launcher + started, one task per rank, were already running by then and never see it, + so they read the config back from the run directory instead. + + Only possible when the deployment named that directory, since it otherwise + defaults to a per-process temporary one that no other rank could read. + """ + run_dir = os.getenv(RUN_DIR_ENV) + if not run_dir: + return None + path = os.path.join(run_dir, CLIENT_CONFIG_NAME) + return path if os.path.exists(path) else None + + +@dataclass(frozen=True) +class MooncakeStoreConnectorConfig: + """Everything needed to open a store handle and name keys in it.""" + + master_server_address: str + metadata_server: str = DEFAULT_METADATA_SERVER + protocol: str = "rdma" + device_name: str = "" + global_segment_size: int = DEFAULT_GLOBAL_SEGMENT_SIZE + local_buffer_size: int = DEFAULT_LOCAL_BUFFER_SIZE + local_hostname: Optional[str] = None + tenant_id: Optional[str] = None + role: StoreRole = StoreRole.BOTH + cache_prefix: str = DEFAULT_CACHE_PREFIX + #: Identity the keys are namespaced by. Two engines only share cache when + #: they agree on this, so it defaults to the model directory's basename + #: rather than its full path: the same checkpoint is routinely mounted + #: somewhere else on another host, which is exactly the case sharing is for. + model_key: Optional[str] = None + #: How many page keys go into one store call. Bounds the size of a single + #: RPC without bounding how much a request may transfer. + transfer_batch_size: int = 64 + #: Pass pages through a pinned host buffer instead of registering the KV + #: pools with Mooncake. Costs a copy each way, but works without GPUDirect + #: RDMA, which registering device memory requires. + stage_through_host: bool = False + #: Ceiling on the pinned allocation per direction when staging. Slots are + #: sized from the layout's largest page, so this caps how many pages may be + #: in flight rather than how large one may be. + staging_buffer_bytes: int = DEFAULT_STAGING_BUFFER_SIZE + + def __post_init__(self) -> None: + """Reject settings that would fail later, inside a transfer.""" + if not self.master_server_address: + raise ValueError("master_server_address is required") + if self.local_buffer_size <= 0: + raise ValueError("local_buffer_size must be > 0") + if self.global_segment_size < 0: + raise ValueError("global_segment_size must be >= 0") + if self.transfer_batch_size <= 0: + raise ValueError("transfer_batch_size must be > 0") + if self.stage_through_host and self.staging_buffer_bytes <= 0: + raise ValueError("staging_buffer_bytes must be > 0 when staging is on") + + @staticmethod + def from_file(path: str) -> "MooncakeStoreConnectorConfig": + """Read the topology from a vLLM-compatible Mooncake JSON config.""" + with open(path) as handle: + raw = json.load(handle) + return MooncakeStoreConnectorConfig( + master_server_address=raw.get("master_server_address", ""), + metadata_server=raw.get("metadata_server") or DEFAULT_METADATA_SERVER, + protocol=raw.get("protocol", "rdma"), + device_name=raw.get("device_name", ""), + global_segment_size=parse_size( + raw.get("global_segment_size", DEFAULT_GLOBAL_SEGMENT_SIZE) + ), + local_buffer_size=parse_size(raw.get("local_buffer_size", DEFAULT_LOCAL_BUFFER_SIZE)), + local_hostname=raw.get("local_hostname") or None, + tenant_id=raw.get("tenant_id") or None, + role=StoreRole(str(raw.get("role", StoreRole.BOTH.value)).strip().lower()), + cache_prefix=str(raw.get("cache_prefix", DEFAULT_CACHE_PREFIX)), + model_key=raw.get("model_key") or None, + transfer_batch_size=int(raw.get("transfer_batch_size", 64)), + stage_through_host=bool(raw.get("stage_through_host", False)), + staging_buffer_bytes=parse_size( + raw.get("staging_buffer_bytes", DEFAULT_STAGING_BUFFER_SIZE) + ), + ) + + @staticmethod + def from_env() -> "MooncakeStoreConnectorConfig": + """Load the JSON config, then apply the TensorRT-LLM env overrides.""" + path = os.getenv(CONFIG_PATH_ENV) or provisioned_config_path() + if not path: + raise ValueError( + f"The mooncake-store connector needs {CONFIG_PATH_ENV} set to a " + "Mooncake JSON config (metadata_server, master_server_address, " + "protocol, device_name, global_segment_size, local_buffer_size), " + "or kv_connector_config.mooncake_store set so the server renders " + f"one, into ${RUN_DIR_ENV} if this rank was started by the " + "launcher rather than spawned by the server." + ) + config = MooncakeStoreConnectorConfig.from_file(path) + return config.with_env_overrides() + + def with_env_overrides(self) -> "MooncakeStoreConnectorConfig": + """Apply `TRTLLM_MOONCAKE_STORE_*` on top of the file's settings.""" + import dataclasses + + updates: dict[str, Any] = {} + role = os.getenv(ROLE_ENV) + if role: + try: + updates["role"] = StoreRole(role.strip().lower()) + except ValueError as exc: + known = ", ".join(member.value for member in StoreRole) + raise ValueError(f"{ROLE_ENV}={role!r} is not one of: {known}") from exc + prefix = os.getenv(CACHE_PREFIX_ENV) + if prefix: + updates["cache_prefix"] = prefix + model_key = os.getenv(MODEL_KEY_ENV) + if model_key: + updates["model_key"] = model_key + staging = os.getenv(STAGE_THROUGH_HOST_ENV) + if staging: + normalized = staging.strip().lower() + if normalized in _TRUE: + updates["stage_through_host"] = True + elif normalized in _FALSE: + updates["stage_through_host"] = False + else: + known = ", ".join(sorted(_TRUE | _FALSE)) + raise ValueError( + f"{STAGE_THROUGH_HOST_ENV}={staging!r} is not a boolean; use one of: {known}" + ) + return dataclasses.replace(self, **updates) if updates else self + + def resolve_model_key(self, model: Any) -> str: + """The model identity to namespace keys by, given the configured model.""" + if self.model_key: + return self.model_key + return os.path.basename(str(model).rstrip("/")) or str(model) diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/donor.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/donor.py new file mode 100644 index 000000000000..83c010797d9e --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/donor.py @@ -0,0 +1,173 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Put a node's host memory into a Mooncake pool without reading or writing it. + +Pool capacity comes only from processes that open a store handle: `setup` +registers `global_segment_size` bytes of the caller's host memory and the master +then places blocks in it. In a disaggregated deployment only the context servers +configure the connector, so the pool is entirely prefill-node memory, which +overlaps what TensorRT-LLM's own host offload already does. + +Donating alongside a generation server puts that node's memory into the same +pool, so prefill writes blocks that land on decode-side DRAM. The generation +engine stays free of any connector and keeps its single cache transceiver for +the prefill-to-decode handoff. + +Donation is not a `StoreRole`. The roles describe an engine's traffic and none +of them means "contribute memory only", so capacity and traffic stay separate +concerns and a donor holds a store handle of its own. + +The memory is charged to the donating process, so size it together with that +node's `kv_cache_config.host_cache_size`. +""" + +import contextlib +import time +from typing import Any, Iterator, Optional + +from tensorrt_llm.logger import logger + +from .config import DEFAULT_METADATA_SERVER, parse_size +from .master import ( + local_address, + master_timeout, + resolve_device_name, + resolve_master_address, + wait_for_master, +) + +__all__ = [ + "DEFAULT_DONOR_LOCAL_BUFFER_SIZE", + "donate_segment", + "maybe_donate_segment", +] + +#: A donor never transfers, but `setup` rejects a zero-sized transfer buffer. +DEFAULT_DONOR_LOCAL_BUFFER_SIZE = 64 * 1024**2 + + +@contextlib.contextmanager +def donate_segment( + master_server_address: str, + segment_size: int, + protocol: str = "rdma", + device_name: str = "", + metadata_server: str = DEFAULT_METADATA_SERVER, + local_buffer_size: int = DEFAULT_DONOR_LOCAL_BUFFER_SIZE, + hostname: Optional[str] = None, +) -> Iterator[str]: + """Hold `segment_size` bytes of this node's memory in the pool. + + Yields the host the segment is registered under, which is how the master + and the engines reading from it identify the capacity. + + Dropping the store handle unmounts the segment and the master starts + reporting the blocks that lived in it as lost, so the caller must stay + inside this context for as long as the capacity is meant to exist. + """ + try: + from mooncake.store import MooncakeDistributedStore + except ImportError as exc: + raise ImportError( + "Donating memory needs the Mooncake Python bindings " + "(`pip install mooncake-transfer-engine`). The C++ transfer engine " + f"in the container is a different component: {exc}" + ) from exc + + host = hostname or local_address() + donated = f"{segment_size / 1024**3:.1f}GiB" + # Byte counts are spelled out next to the human-readable form. A misparsed + # size string otherwise surfaces only as a pool that evicts far too eagerly. + logger.info( + f"mooncake-store: lending memory to the pool at {master_server_address} " + f"as capacity only, no reads or writes: host={host} " + f"segment_size={donated} ({segment_size} bytes) " + f"protocol={protocol} device={device_name or '(none)'} " + f"metadata_server={metadata_server} " + f"local_buffer_size={local_buffer_size} bytes" + ) + + store = MooncakeDistributedStore() + started = time.monotonic() + status = store.setup( + host, + metadata_server, + segment_size, + local_buffer_size, + protocol, + device_name, + master_server_address, + ) + elapsed = time.monotonic() - started + if status != 0: + raise RuntimeError( + f"Mooncake store.setup failed with status {status} after " + f"{elapsed:.1f}s, so no memory was lent to the pool. The master at " + f"{master_server_address} must already be accepting connections; " + f"protocol={protocol!r} with device={device_name or '(none)'!r} " + f"must be usable from {host}; and this node must have " + f"{donated} of memory to spare, which it does not if its own " + "kv_cache_config.host_cache_size has already claimed it." + ) + + logger.info( + f"mooncake-store: {donated} of {host} is now part of the pool, " + f"registered in {elapsed:.1f}s; the master at {master_server_address} " + "can place blocks here from now on" + ) + try: + yield host + finally: + # The segment stays mounted while anything references the handle. + del store + logger.info( + f"mooncake-store: withdrew the {donated} lent from {host}; the " + "master will report blocks that lived there as lost" + ) + + +@contextlib.contextmanager +def maybe_donate_segment(donation: Any) -> Iterator[Optional[str]]: + """Lend memory for this process's lifetime if the config asked to. + + Args: + donation: A `MooncakeDonationConfig`, or `None` to do nothing, so + callers need no condition of their own. + + Yields the host the segment is registered under, or `None`. + """ + if donation is None: + yield None + return + + # Bringup blocks here on a master that may belong to a different job, so + # name the address before waiting on it. + logger.info( + "mooncake-store: mooncake_donation is set, so this server lends host " + f"memory to the pool at {donation.master_server_address} without using " + "it; resolving the master now" + ) + master_address = resolve_master_address(donation.master_server_address, master_timeout()) + # Checked before setup so an absent master is reported as such, rather than + # as the status code setup returns for every kind of failure. + wait_for_master(master_address) + with donate_segment( + master_server_address=master_address, + segment_size=parse_size(donation.segment_size), + protocol=donation.protocol, + device_name=resolve_device_name(donation.protocol, donation.device_name), + metadata_server=donation.metadata_server, + ) as host: + yield host diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/keys.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/keys.py new file mode 100644 index 000000000000..56609877843d --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/keys.py @@ -0,0 +1,134 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Block identity and store key naming for the Mooncake store connector. + +`KVCacheManagerV2` exposes no block hashes to a connector, since `RequestData` +reports them empty, so content identity is derived here instead. The chain is +the standard one: a block's hash covers its own tokens *and* every token before +it, so a key can only be reused by a request whose prefix is byte-identical. + +A key is `/`. The namespace pins down everything that +would make the stored bytes mean something different: the model, the shard that +produced them, the layer group inside that shard, the tokens each page holds and +how many bytes a page is. Anything that changes those reads as a cache miss +rather than as garbage. +""" + +import hashlib +from dataclasses import dataclass +from typing import List, Optional, Sequence + +__all__ = [ + "BlockHashChain", + "KeyNamespace", + "HASH_DIGEST_BYTES", +] + +#: 128 bits. Collisions decide whether one request reads another's KV, so the +#: digest is sized to make that negligible over any realistic cache lifetime, +#: while staying half the width of a full blake2b digest in every key. +HASH_DIGEST_BYTES = 16 + + +def _digest(*parts: bytes) -> bytes: + hasher = hashlib.blake2b(digest_size=HASH_DIGEST_BYTES) + for part in parts: + hasher.update(part) + return hasher.digest() + + +class BlockHashChain: + """Rolling hashes of a request's full blocks, one entry per block ordinal. + + Extended in place as a request's token list grows, so generation steps cost + one digest per newly completed block rather than a rehash of the prompt. + """ + + def __init__(self, tokens_per_block: int, cache_salt: Optional[str] = None): + if tokens_per_block <= 0: + raise ValueError(f"tokens_per_block must be > 0, got {tokens_per_block}") + self._tokens_per_block = int(tokens_per_block) + # The salt seeds the chain rather than being mixed into every block, so + # a request carrying a different salt diverges from the first block on. + salt_bytes = b"" if cache_salt is None else str(cache_salt).encode() + self._seed = _digest(b"salt", salt_bytes) + self._hashes: List[bytes] = [] + + @property + def tokens_per_block(self) -> int: + """Tokens covered by each entry in the chain.""" + return self._tokens_per_block + + @property + def hashes(self) -> Sequence[bytes]: + """Hashes computed so far, indexed by block ordinal.""" + return self._hashes + + def extend(self, tokens: Sequence[int]) -> Sequence[bytes]: + """Grow the chain to cover every full block of `tokens`. + + Args: + tokens: The request's complete token list, prompt first. Must be an + extension of what was passed previously; a request's tokens only + ever grow, so a shorter list means the caller mixed up requests. + + Returns: + The full chain, indexed by block ordinal. + """ + num_full_blocks = len(tokens) // self._tokens_per_block + if num_full_blocks < len(self._hashes): + raise ValueError( + f"token list shrank from {len(self._hashes)} to {num_full_blocks} " + "full blocks; a hash chain belongs to exactly one request" + ) + for ordinal in range(len(self._hashes), num_full_blocks): + start = ordinal * self._tokens_per_block + block = tokens[start : start + self._tokens_per_block] + parent = self._hashes[-1] if self._hashes else self._seed + # Fixed-width little-endian token ids: a delimiter-free encoding + # would let two different token sequences serialize identically. + payload = b"".join(int(token).to_bytes(8, "little", signed=True) for token in block) + self._hashes.append(_digest(parent, payload)) + return self._hashes + + +@dataclass(frozen=True) +class KeyNamespace: + """The part of a store key that is fixed for one shard and layer group.""" + + cache_prefix: str + model_key: str + #: Global rank of the shard whose KV these bytes are, and the world size it + #: was produced under. Both are needed: rank 3 of 8 holds different heads + #: than rank 3 of 4. + rank: int + world_size: int + layer_group_id: int + tokens_per_block: int + bytes_per_page: int + + @property + def prefix(self) -> str: + """The literal string every key in this namespace starts with.""" + return ( + f"{self.cache_prefix}/{self.model_key}" + f"/w{self.world_size}r{self.rank}" + f"/lg{self.layer_group_id}" + f"/t{self.tokens_per_block}b{self.bytes_per_page}" + ) + + def key(self, block_hash: bytes) -> str: + """The store key holding one page of this namespace.""" + return f"{self.prefix}/{block_hash.hex()}" diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py new file mode 100644 index 000000000000..f2f3a21a3c91 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py @@ -0,0 +1,596 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Bring the Mooncake store's pool up as part of a server's own startup. + +The connector needs two things that are not the engine's to produce: a reachable +`mooncake_master`, and a JSON client config named by `MOONCAKE_CONFIG_PATH` that +points every worker at it. + +`provision_pool` does that work inside the serving process. It resolves the +master, either launching one here or checking that the configured one answers, +renders the client config, and exports `MOONCAKE_CONFIG_PATH`, which reaches the +ranks because the LLM constructor spawns them from this process. Everything it +started is torn down when the context exits. + +A master launched here lives and dies with the server, so it suits one engine +talking to its own pool. Several engines sharing a pool, or a pool meant to +survive a restart, need a master with its own lifetime named by +`master_server_address`. +""" + +import contextlib +import json +import os +import shutil +import socket +import subprocess # nosec B404 +import tempfile +import time +from dataclasses import dataclass +from typing import Any, Dict, Iterator, List, Optional, Sequence, Tuple + +from tensorrt_llm.logger import logger + +from ..registry import uses_connector +from .config import CLIENT_CONFIG_NAME, CONFIG_PATH_ENV, RUN_DIR_ENV + +__all__ = [ + "local_address", + "maybe_provision_pool", + "master_timeout", + "provision_pool", + "resolve_device_name", + "resolve_master_address", + "running_master", + "wait_for_master", +] + +#: Override the binary that `launch_master` runs. +MASTER_BINARY_ENV = "TRTLLM_MOONCAKE_MASTER_BINARY" +#: How long to wait for a master to accept connections, in seconds. +MASTER_TIMEOUT_ENV = "TRTLLM_MOONCAKE_MASTER_TIMEOUT" +DEFAULT_MASTER_BINARY = "mooncake_master" +DEFAULT_MASTER_TIMEOUT = 60.0 +MASTER_LOG_NAME = "mooncake_master.log" +#: Name a launched master's address is always published under in the run +#: directory, so even a run that named no address file records its pool. +MASTER_ADDRESS_NAME = "master.addr" +#: Prefix that makes `master_server_address` name a file holding the address +#: rather than the address itself. +ADDRESS_FILE_SCHEME = "file://" +#: Lines of the master's log to quote when startup fails, since its last words +#: (a port in use, a bad flag) are usually the whole diagnosis. +LOG_TAIL_LINES = 20 + + +def _log_tail(path: str, lines: int = LOG_TAIL_LINES) -> str: + """The end of the master's log, ready to append to a failure message.""" + try: + with open(path, errors="replace") as handle: + tail = handle.read().splitlines()[-lines:] + except OSError as exc: + return f" Its log at {path} could not be read: {exc}." + if not tail: + return ( + f" Its log at {path} is empty, which usually means it failed " + "before glog opened; check that the binary runs at all." + ) + quoted = "\n ".join(tail) + return f" The last {len(tail)} lines of {path}:\n {quoted}" + + +def local_address() -> str: + """The address this host is known by inside the pool. + + Uses the same derivation as the connector worker's own hostname, so the + master and the segments registering with it agree on which host they are on. + """ + try: + return socket.gethostbyname(socket.gethostname()) + except OSError: + return "127.0.0.1" + + +def master_timeout() -> float: + raw = os.getenv(MASTER_TIMEOUT_ENV) + if not raw: + return DEFAULT_MASTER_TIMEOUT + try: + timeout = float(raw) + except ValueError as exc: + raise ValueError(f"{MASTER_TIMEOUT_ENV}={raw!r} is not a number") from exc + if timeout <= 0: + raise ValueError(f"{MASTER_TIMEOUT_ENV}={raw!r} must be > 0") + return timeout + + +def _split_address(address: str) -> Optional[Tuple[str, int]]: + """Split `host:port`, or return `None` if it is not in that form.""" + host, separator, port = address.rpartition(":") + if not separator or not port.isdigit(): + return None + return host.strip("[]"), int(port) + + +def resolve_master_address(address: str, timeout: float) -> str: + """Read a `file://` address through, and pass anything else along. + + A master with its own lifetime runs on whichever host its scheduler gave + it, which is not known when the worker configs are written. Naming the file + it publishes to keeps the address out of both the config and the launch + script: `trtllm-serve mooncake_master --address-file` writes it, every + worker's `master_server_address` names the same path, and the wait here + doubles as the wait for the master to exist at all. + """ + if not address.startswith(ADDRESS_FILE_SCHEME): + return address + + path = address[len(ADDRESS_FILE_SCHEME) :] + started = time.monotonic() + deadline = started + timeout + announced = started + logger.info(f"mooncake-store: reading the master's address from {path}") + while True: + try: + published = open(path).read().strip() + except FileNotFoundError: + published = "" + if published: + logger.info(f"mooncake-store: {path} names the master at {published}") + return published + now = time.monotonic() + if now - announced >= 5.0: + announced = now + # Waiting on a master in another job is normal here, so say so + # rather than letting the wait look like a hang. + logger.info( + f"mooncake-store: no master address in {path} yet " + f"({now - started:.0f}s of {timeout:g}s); waiting for the " + "master to start and publish it" + ) + if now >= deadline: + raise TimeoutError( + f"No Mooncake master address appeared in {path} within " + f"{timeout:g}s. Start one with 'trtllm-serve mooncake_master " + f"--address-file {path}', or name a reachable host:port in " + f"master_server_address. Raise {MASTER_TIMEOUT_ENV} if the " + "master is only slow to start." + ) + time.sleep(0.5) + + +def _wait_until_accepting( + host: str, + port: int, + timeout: float, + process: Optional[subprocess.Popen] = None, + log_path: Optional[str] = None, +) -> float: + """Block until the master accepts connections, and say how long it took. + + A worker that opens its store handle before the master is listening fails + outright, so the ordering has to wait on the port rather than on the + presence of a process. When the master is ours, its exit is checked first + each pass, so a master that died is reported as such rather than as a + timeout. + + The wait is narrated as it happens, since silence here is + indistinguishable from a hang elsewhere in bringup. + """ + started = time.monotonic() + deadline = started + timeout + announced = started + while True: + if process is not None and (code := process.poll()) is not None: + raise RuntimeError( + f"mooncake_master exited with code {code} after " + f"{time.monotonic() - started:.1f}s, before it accepted " + f"connections on {host}:{port}." + f"{_log_tail(log_path) if log_path else ''}" + ) + try: + with socket.create_connection((host, port), timeout=1.0): + return time.monotonic() - started + except OSError as exc: + last_error = exc + now = time.monotonic() + if now >= deadline: + raise TimeoutError( + f"The Mooncake master at {host}:{port} did not accept " + f"connections within {timeout:g}s ({last_error}). Raise " + f"{MASTER_TIMEOUT_ENV} if it is only slow to start." + f"{_log_tail(log_path) if log_path else ''}" + ) + if now - announced >= 5.0: + announced = now + logger.info( + f"mooncake-store: still waiting for the master at {host}:{port}" + f" ({now - started:.0f}s of {timeout:g}s, {last_error})" + ) + time.sleep(0.5) + + +#: Where the InfiniBand devices of a host are described. +IB_SYSFS_ROOT = "/sys/class/infiniband" + + +def _highest_rate_ib_devices(sysfs_root: Optional[str] = None) -> List[str]: + """The active InfiniBand devices on the compute fabric, fastest first. + + A node's HCAs are not interchangeable. On GB300 six are exposed, of which + four run at 800Gb/s (two per NUMA node, one per GPU) while the rest share a + PCI device with an Ethernet port and serve storage or management. Taking + every device at the highest rate picks the compute fabric on any node type, + where a hardcoded name would be wrong on the next one. + """ + sysfs_root = sysfs_root or IB_SYSFS_ROOT + rated: Dict[str, int] = {} + try: + devices = sorted(os.listdir(sysfs_root)) + except OSError: + return [] + for device in devices: + port = os.path.join(sysfs_root, device, "ports", "1") + + def attribute(name: str) -> str: + try: + with open(os.path.join(port, name)) as handle: + return handle.read().strip() + except OSError: + return "" + + if attribute("link_layer") != "InfiniBand": + continue + if "ACTIVE" not in attribute("state"): + continue + # "800 Gb/sec (4X XDR)" + rate = attribute("rate").split() + if not rate or not rate[0].isdigit(): + continue + rated[device] = int(rate[0]) + + if not rated: + return [] + fastest = max(rated.values()) + return [device for device, rate in sorted(rated.items()) if rate == fastest] + + +def resolve_device_name(protocol: str, configured: str, sysfs_root: Optional[str] = None) -> str: + """The RDMA devices to transfer over, detected if the config left it open. + + Which HCAs a node has is a property of the node, not of the deployment, so + requiring it in a config would tie that config to one machine type. + Detecting it keeps `protocol: rdma` portable; setting `device_name` + overrides the detection. + """ + if configured or protocol != "rdma": + return configured + detected = _highest_rate_ib_devices(sysfs_root) + if not detected: + logger.warning( + "mooncake-store: protocol is rdma but no active InfiniBand device " + f"was found under {sysfs_root or IB_SYSFS_ROOT}, so device_name is " + "left empty for Mooncake's own discovery. Set device_name to " + "choose explicitly." + ) + return "" + joined = ",".join(detected) + logger.info( + f"mooncake-store: transferring over the fastest active InfiniBand " + f"devices on this host: {joined}" + ) + return joined + + +def wait_for_master(master_address: str, timeout: Optional[float] = None) -> Optional[float]: + """Block until the master at `master_address` accepts connections. + + Reaching a master that is not there otherwise fails deep inside + `store.setup`, in every rank, after the model has loaded, as a bare status + code. One socket beforehand turns that into a line naming the address. + + Returns how long it took, or `None` if the address was not in `host:port` + form and could not be checked. + """ + timeout = master_timeout() if timeout is None else timeout + endpoint = _split_address(master_address) + if endpoint is None: + logger.warning( + f"mooncake-store: cannot parse master_server_address=" + f"{master_address!r} as host:port, so its reachability is left " + "for the workers to discover." + ) + return None + elapsed = _wait_until_accepting(*endpoint, timeout) + logger.info(f"mooncake-store: the master at {master_address} answered in {elapsed:.1f}s") + return elapsed + + +def _client_config( + pool: Any, master_address: str, device_name: Optional[str] = None +) -> Dict[str, Any]: + """Render the Mooncake client config for a pool. + + The schema is vLLM's, so one pool can serve both engines. `role` is written + as `both` because the file describes the pool; the directions of traffic a + given process drives come from its own `TRTLLM_MOONCAKE_STORE_ROLE`. + """ + config: Dict[str, Any] = { + "metadata_server": pool.metadata_server, + "master_server_address": master_address, + "protocol": pool.protocol, + "device_name": pool.device_name if device_name is None else device_name, + "global_segment_size": pool.global_segment_size, + "local_buffer_size": pool.local_buffer_size, + "role": "both", + "transfer_batch_size": pool.transfer_batch_size, + "stage_through_host": pool.stage_through_host, + } + if pool.cache_prefix is not None: + config["cache_prefix"] = pool.cache_prefix + # Left out when unset so the connector's own default applies instead of a + # second copy of it here. + if pool.staging_buffer_bytes is not None: + config["staging_buffer_bytes"] = pool.staging_buffer_bytes + return config + + +@dataclass +class LaunchedMaster: + """A `mooncake_master` owned by this process.""" + + process: subprocess.Popen + address: str + log_path: str + + def stop(self, timeout: float = 10.0) -> None: + if self.process.poll() is not None: + return + self.process.terminate() + try: + self.process.wait(timeout=timeout) + except subprocess.TimeoutExpired: + self.process.kill() + self.process.wait() + + +def _launch_master(pool: Any, run_dir: str) -> LaunchedMaster: + """Start a master on this host and wait for it to answer.""" + binary = os.getenv(MASTER_BINARY_ENV) or DEFAULT_MASTER_BINARY + resolved = shutil.which(binary) + if resolved is None: + raise FileNotFoundError( + f"{binary!r} is not on PATH, so launch_master cannot start a " + "Mooncake master. It ships with the Mooncake runtime, which " + "docker/common/install_mooncake.sh installs. Point " + f"{MASTER_BINARY_ENV} at the binary, or drop launch_master and " + "set master_server_address to a master you run yourself." + ) + + host = local_address() + log_path = os.path.join(run_dir, MASTER_LOG_NAME) + + # glog writes to files under /tmp unless redirected, so without + # GLOG_logtostderr the log opened below stays empty. GLOG_v=1 adds the + # per-RPC lines showing segments registering and keys moving, which is the + # only view of the pool's own side of the conversation short of scraping + # the metrics port. + env = dict(os.environ, GLOG_logtostderr="1") + env.setdefault("GLOG_v", "1") + command = [ + resolved, + f"--rpc_port={pool.master_port}", + f"--metrics_port={pool.master_metrics_port}", + f"--eviction_ratio={pool.master_eviction_ratio}", + ] + + logger.info(f"mooncake-store: starting {' '.join(command)} on {host}") + with open(log_path, "wb") as log_file: + process = subprocess.Popen( # nosec B603 + command, env=env, stdout=log_file, stderr=subprocess.STDOUT + ) + master = LaunchedMaster( + process=process, address=f"{host}:{pool.master_port}", log_path=log_path + ) + logger.info( + f"mooncake-store: master pid={process.pid} logging to {log_path} " + f"(GLOG_v={env['GLOG_v']}); waiting for it to accept connections" + ) + try: + elapsed = _wait_until_accepting( + host, pool.master_port, master_timeout(), process=process, log_path=log_path + ) + except BaseException: + master.stop() + raise + + logger.info( + f"mooncake-store: master ready at {master.address} after {elapsed:.1f}s " + f"(metrics http://{host}:{pool.master_metrics_port}, log {log_path})" + ) + return master + + +@contextlib.contextmanager +def _published_address(address: str, paths: Sequence[str]) -> Iterator[None]: + """Write `address` to every path for the life of the context. + + Publishing is how anything else finds this master: a donor or a second + server names the path as `file://` in `master_server_address`. + Retracting on the way out matters as much as writing, since an address that + outlives its master sends the next run's workers to a dead port. + """ + for path in paths: + directory = os.path.dirname(path) + if directory: + os.makedirs(directory, exist_ok=True) + # Renamed into place so a reader never sees a partial address. + staging = f"{path}.partial" + with open(staging, "w") as handle: + handle.write(f"{address}\n") + os.replace(staging, path) + logger.info(f"mooncake-store: published master {address} to {path}") + try: + yield + finally: + for path in paths: + with contextlib.suppress(OSError): + os.remove(path) + logger.info(f"mooncake-store: withdrew the master address at {path}") + + +def _address_files(run_dir: str, extra: Optional[str] = None) -> List[str]: + """Where a master this process starts should publish its address. + + Always the run directory, plus wherever the deployment asked for. + """ + paths = [os.path.join(run_dir, MASTER_ADDRESS_NAME)] + if extra and os.path.abspath(extra) not in {os.path.abspath(p) for p in paths}: + paths.append(extra) + return paths + + +@contextlib.contextmanager +def running_master( + pool: Any, run_dir: str, address_file: Optional[str] = None +) -> Iterator[LaunchedMaster]: + """Run a master whose lifetime is this process's rather than an engine's. + + `provision_pool` covers the server that owns its pool. Several engines on + one pool, or a pool that has to survive a restart, need the master + somewhere that is not any of them. + + `address_file` receives `host:port` once the master answers, so workers can + name the file instead of an address nobody knows until the scheduler has + placed this process. One is written to `run_dir` either way. + """ + os.makedirs(run_dir, exist_ok=True) + master = _launch_master(pool, run_dir) + try: + with _published_address(master.address, _address_files(run_dir, address_file)): + yield master + finally: + master.stop() + logger.info(f"mooncake-store: master at {master.address} stopped") + + +@contextlib.contextmanager +def provision_pool(pool: Any, run_dir: Optional[str] = None) -> Iterator[Optional[str]]: + """Make `pool` reachable and name it in this process's environment. + + Yields the path of the client config written, or `None` when an inherited + `MOONCAKE_CONFIG_PATH` was left in charge. + + Args: + pool: A `MooncakeStoreConfig`. + run_dir: Where to write the client config and the master's log. + Defaults to `TRTLLM_MOONCAKE_RUN_DIR`, else a temporary directory + that is removed on exit. + """ + inherited = os.getenv(CONFIG_PATH_ENV) + if inherited: + logger.info( + f"mooncake-store: {CONFIG_PATH_ENV}={inherited} is already set, so " + "kv_connector_config.mooncake_store is ignored and the pool it " + "names is used as is." + ) + yield None + return + + keep_run_dir = bool(run_dir or os.getenv(RUN_DIR_ENV)) + run_dir = run_dir or os.getenv(RUN_DIR_ENV) or tempfile.mkdtemp(prefix="trtllm-mooncake-") + os.makedirs(run_dir, exist_ok=True) + if keep_run_dir: + logger.info(f"mooncake-store: provisioning the pool, run directory {run_dir}") + else: + logger.info( + f"mooncake-store: provisioning the pool in {run_dir}, which is " + f"removed at shutdown along with the master's log; set " + f"{RUN_DIR_ENV} to keep them" + ) + + master: Optional[LaunchedMaster] = None + exported = False + with contextlib.ExitStack() as stack: + try: + if pool.launch_master: + master = _launch_master(pool, run_dir) + master_address = master.address + # Published even when only this server uses it, since that is + # how its donors reach it. + stack.enter_context( + _published_address( + master_address, _address_files(run_dir, pool.master_address_file) + ) + ) + else: + master_address = resolve_master_address( + pool.master_server_address, master_timeout() + ) + wait_for_master(master_address) + logger.info(f"mooncake-store: using the master at {master_address}") + + config_path = os.path.join(run_dir, CLIENT_CONFIG_NAME) + config = _client_config( + pool, master_address, resolve_device_name(pool.protocol, pool.device_name) + ) + with open(config_path, "w") as handle: + json.dump(config, handle, indent=2) + # Inherited by the ranks the LLM constructor spawns. Ranks an + # external launcher started were already running, so they read the + # config out of the run directory instead; see + # provisioned_config_path. + os.environ[CONFIG_PATH_ENV] = config_path + exported = True + logger.info( + f"mooncake-store: {CONFIG_PATH_ENV}={config_path} " + f"({json.dumps(config, sort_keys=True)})" + ) + # Capacity is what explains a low hit rate, so state the + # arithmetic instead of leaving it to be derived later. + logger.info( + "mooncake-store: this server's ranks will each contribute " + f"global_segment_size={pool.global_segment_size} to the pool; " + "total capacity is that times the number of ranks that open a " + "handle, plus whatever any mooncake_donation adds" + ) + yield config_path + finally: + if exported: + os.environ.pop(CONFIG_PATH_ENV, None) + if master is not None: + master.stop() + logger.info(f"mooncake-store: master at {master.address} stopped") + if not keep_run_dir: + shutil.rmtree(run_dir, ignore_errors=True) + + +@contextlib.contextmanager +def maybe_provision_pool(kv_connector_config: Any) -> Iterator[None]: + """Provision the pool if this deployment asked the server to. + + A no-op for every other connector, and for a `mooncake-store` config that + left `mooncake_store` unset, since such a deployment is told about its pool + through `MOONCAKE_CONFIG_PATH` instead. + """ + if not uses_connector(kv_connector_config, "mooncake-store"): + yield + return + pool = kv_connector_config.mooncake_store + if pool is None: + yield + return + with provision_pool(pool): + yield diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/metadata.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/metadata.py new file mode 100644 index 000000000000..0e35227cf606 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/metadata.py @@ -0,0 +1,62 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The per-iteration work list the scheduler hands the workers. + +Instances are broadcast from rank 0 to every worker, so these carry only plain +data: a page's identity (its block hash) and where that page currently lives on +this rank (a layer group and a page slot index). Deliberately no store keys -- +each worker prefixes its own rank namespace, so one broadcast serves all shards. +""" + +from dataclasses import dataclass, field +from typing import List + +__all__ = ["MooncakeStoreMetadata", "PageTransfer", "RequestTransfers"] + + +@dataclass +class PageTransfer: + """One page of one layer group, to move in either direction.""" + + #: Content identity from `BlockHashChain`; names the key, not the location. + block_hash: bytes + layer_group_id: int + #: Page slot index within `layer_group_id`, as reported by + #: `RequestData.new_block_ids_by_layer_group`. + page_index: int + + +@dataclass +class RequestTransfers: + """Pages belonging to one request, kept together for save bookkeeping. + + The worker owes `get_finished` an answer per request, so a save's owner has + to survive the trip from scheduler to worker. + """ + + request_id: int + pages: List[PageTransfer] = field(default_factory=list) + + +@dataclass +class MooncakeStoreMetadata: + """Loads to perform before the next forward pass, saves to start after it.""" + + loads: List[RequestTransfers] = field(default_factory=list) + saves: List[RequestTransfers] = field(default_factory=list) + + def __bool__(self) -> bool: + """Whether there is any work at all this iteration.""" + return bool(self.loads or self.saves) diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/scheduler.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/scheduler.py new file mode 100644 index 000000000000..9be55a5121df --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/scheduler.py @@ -0,0 +1,292 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Leader side of the Mooncake store KV cache connector. + +Runs only on rank 0. It decides what to load and what to save; the workers do +the moving. Two pieces of bookkeeping make that possible, and both exist because +`KVCacheManagerV2` reports `RequestData.block_hashes` empty: + +* a hash chain per request, so a block has a content identity at all; +* the page slot index per block ordinal, accumulated across iterations. The + manager reports only *newly allocated* indices each step, but a block is + allocated before it is full and is only savable once it is full, so the index + has to be remembered from the step that reported it. +""" + +from typing import Dict, List, Optional, Tuple + +from tensorrt_llm.bindings.internal.batch_manager import LlmRequest +from tensorrt_llm.llmapi.llm_args import TorchLlmArgs +from tensorrt_llm.logger import logger +from tensorrt_llm.runtime.kv_cache_manager_v2 import BAD_PAGE_INDEX + +from ..kv_cache_connector import KvCacheConnectorScheduler, RequestData, SchedulerOutput +from .config import MooncakeStoreConnectorConfig +from .keys import BlockHashChain +from .metadata import MooncakeStoreMetadata, PageTransfer, RequestTransfers +from .validation import validate_llm_args +from .worker import MooncakeStoreConnectorWorker, resolve_local_worker + +__all__ = ["MooncakeStoreConnectorScheduler"] + + +class _RequestState: + """Per-request bookkeeping that has to outlive a single iteration.""" + + __slots__ = ( + "chain", + "tokens", + "pages", + "saved_upto", + "load_first_block", + "load_blocks", + "emitted_saves", + ) + + def __init__(self, chain: BlockHashChain): + self.chain = chain + #: The request's tokens, accumulated from the per-step deltas. + self.tokens: List[int] = [] + #: Page slot index per block ordinal, per layer group. + self.pages: Dict[int, List[int]] = {} + #: First block ordinal not yet considered for saving. + self.saved_upto = 0 + #: The offer made by `get_num_new_matched_tokens`, in block ordinals. + self.load_first_block = 0 + self.load_blocks = 0 + self.emitted_saves = False + + +class MooncakeStoreConnectorScheduler(KvCacheConnectorScheduler): + """Chooses which pages the Mooncake pool serves and which it receives.""" + + def __init__(self, llm_args: TorchLlmArgs): + super().__init__(llm_args) + + validate_llm_args(llm_args) + self._config = MooncakeStoreConnectorConfig.from_env() + self._tokens_per_block = int(llm_args.kv_cache_config.tokens_per_block) + self._requests: Dict[int, _RequestState] = {} + self._worker: Optional[MooncakeStoreConnectorWorker] = None + + logger.info( + f"mooncake-store leader ready (role={self._config.role.value}, " + f"tokens_per_block={self._tokens_per_block})" + ) + + def wait_for_initialization(self): + """Bind to the process-local worker, which owns the store handle. + + Called after the executor has built both halves and registered the KV + cache layout, which is what the worker needs before it can name a key. + """ + self._worker = resolve_local_worker() + + # ---- lookup ---- + + def get_num_new_matched_tokens( + self, request: LlmRequest, num_computed_tokens: int + ) -> Tuple[int, bool]: + """Offer the longest stored prefix beyond what the device already has. + + Args: + request: The request being scheduled. + num_computed_tokens: Tokens already matched in the local KV cache. + + Returns: + Tokens the store can supply, and `False` for a synchronous load. + """ + tokens = request.get_tokens(0) + state = self._state_for(request, tokens) + state.load_first_block = 0 + state.load_blocks = 0 + + if not self._config.role.loads: + return 0, False + + # A partial local match means the boundary block is half computed on + # device. Overwriting it with a stored page would discard tokens the + # runtime already counted, so only whole-block offers are made. + if num_computed_tokens % self._tokens_per_block: + return 0, False + + first_block = num_computed_tokens // self._tokens_per_block + # Stop one token short of the prompt: the runtime still has to run a + # forward pass for this request, and it cannot do that with nothing left + # to compute. + last_block = (len(tokens) - 1) // self._tokens_per_block + candidates = state.chain.hashes[first_block:last_block] + if not candidates: + return 0, False + + hit_blocks = self._require_worker().count_prefix_hit(candidates) + if hit_blocks == 0: + return 0, False + + state.load_first_block = first_block + state.load_blocks = hit_blocks + logger.debug( + f"mooncake-store matched {hit_blocks} blocks " + f"({hit_blocks * self._tokens_per_block} tokens) " + f"for request {request.request_id}" + ) + return hit_blocks * self._tokens_per_block, False + + def cancel_load(self, request: LlmRequest, start: int, end: int): + """Drop offered blocks whose tokens the runtime will not consume. + + Loads here are synchronous and nothing has been transferred yet, so this + is exact: the offer is truncated before `build_connector_meta` turns it + into work. + """ + state = self._requests.get(request.request_id) + if state is None or state.load_blocks == 0: + return + kept = 0 + for offset in range(state.load_blocks): + block = state.load_first_block + offset + block_start = block * self._tokens_per_block + if block_start + self._tokens_per_block > start and block_start < end: + break + kept += 1 + state.load_blocks = kept + + def update_state_after_alloc(self, request: LlmRequest, block_ids: List[int]): + """No-op: page indices are read from the scheduler output instead. + + The flat `block_ids` here are a single space, but a V2 page index is + scoped to a layer group. `RequestData.new_block_ids_by_layer_group` is + the form that stays correct for every model, so that is the only source + this connector uses. + """ + + # ---- work lists ---- + + def build_connector_meta(self, scheduler_output: SchedulerOutput) -> MooncakeStoreMetadata: + """Turn this iteration's scheduled requests into load and save lists.""" + metadata = MooncakeStoreMetadata() + for request_data in (*scheduler_output.new_requests, *scheduler_output.cached_requests): + state = self._requests.get(request_data.request_id) + if state is None: + # Only requests that went through get_num_new_matched_tokens have + # a hash chain. Generation-only requests never do, and the + # connector manager refuses them outright. + continue + + state.tokens.extend(request_data.new_tokens) + state.chain.extend(state.tokens) + self._record_pages(state, request_data) + + loads = self._loads_for(state, request_data) + if loads.pages: + metadata.loads.append(loads) + + # Whatever the store just supplied, and whatever the local cache + # matched, is not ours to write back: the store already has the + # former, and the latter was never allocated during this run. + state.saved_upto = max(state.saved_upto, state.load_first_block + state.load_blocks) + # An offer is consumed once. The load is issued in exactly the + # iteration the runtime allocated pages to hold it. + state.load_blocks = 0 + + if self._config.role.saves: + saves = self._saves_for(state, request_data) + if saves.pages: + state.emitted_saves = True + metadata.saves.append(saves) + return metadata + + def request_finished(self, request: LlmRequest, cache_block_ids: List[int]) -> bool: + """Report whether pages must stay pinned for in-flight saves. + + Returns: + True when this request handed any page to the background save + thread. Its pages are the source of those RDMA reads, so freeing + them now would let a later request overwrite bytes mid-transfer. + """ + state = self._requests.pop(request.request_id, None) + return bool(state is not None and state.emitted_saves) + + # ---- internals ---- + + def _require_worker(self) -> MooncakeStoreConnectorWorker: + if self._worker is None: + self._worker = resolve_local_worker() + return self._worker + + def _state_for(self, request: LlmRequest, tokens: List[int]) -> _RequestState: + state = self._requests.get(request.request_id) + if state is None: + state = _RequestState( + BlockHashChain(self._tokens_per_block, cache_salt=request.cache_salt) + ) + self._requests[request.request_id] = state + # Hashing the prompt here rather than waiting for the first scheduler + # output is the whole point: the lookup happens before the request is + # scheduled, so the chain has to be ready before any metadata exists. + state.chain.extend(tokens) + return state + + def _record_pages(self, state: _RequestState, request_data: RequestData) -> None: + """Append this step's newly allocated page indices, by block ordinal.""" + by_group = request_data.new_block_ids_by_layer_group + if not by_group: + # Under a single layer group the manager also mirrors that group's + # indices into the flat `new_block_ids`, but it does not say which + # group they belong to, so there is nothing safe to record from it. + return + for layer_group_id, indices in by_group.items(): + state.pages.setdefault(layer_group_id, []).extend(int(index) for index in indices) + + def _addressable_blocks(self, state: _RequestState) -> int: + """Block ordinals that are both hashed and backed by a page everywhere.""" + if not state.pages: + return 0 + return min(len(state.chain.hashes), min(len(indices) for indices in state.pages.values())) + + def _loads_for(self, state: _RequestState, request_data: RequestData) -> RequestTransfers: + transfers = RequestTransfers(request_data.request_id) + limit = self._addressable_blocks(state) + for offset in range(state.load_blocks): + block = state.load_first_block + offset + if block >= limit: + # The runtime allocated fewer pages than it accepted tokens for. + # It reports the shortfall through cancel_load; until then the + # unaddressable tail is simply not loaded. + break + self._append_pages(state, transfers, block) + return transfers + + def _saves_for(self, state: _RequestState, request_data: RequestData) -> RequestTransfers: + transfers = RequestTransfers(request_data.request_id) + limit = self._addressable_blocks(state) + for block in range(state.saved_upto, limit): + self._append_pages(state, transfers, block) + state.saved_upto = max(state.saved_upto, limit) + return transfers + + def _append_pages(self, state: _RequestState, transfers: RequestTransfers, block: int) -> None: + """Add one block's page from every layer group, or none of them.""" + block_hash = state.chain.hashes[block] + pages: List[PageTransfer] = [] + for layer_group_id, indices in state.pages.items(): + page_index = indices[block] + if page_index == BAD_PAGE_INDEX: + # The block has no page in this group, because a sliding window + # dropped it. A partial page is not a usable cache entry, + # so the whole block is skipped. + return + pages.append(PageTransfer(block_hash, layer_group_id, page_index)) + transfers.pages.extend(pages) diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/staging.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/staging.py new file mode 100644 index 000000000000..966897b178c4 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/staging.py @@ -0,0 +1,303 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Pinned host slots that stand in for GPU pages when the pool cannot reach them. + +The connector's default path registers the KV pools themselves with Mooncake, so +the store reads and writes device memory directly. That needs the HCA to be able +to pin GPU pages, which means GPUDirect RDMA via `nvidia_peermem` or dma-buf. +Where that is unavailable, `ibv_reg_mr` fails on every pool range and the +connector cannot start. + +Staging trades a copy for that dependency. Mooncake is given a pinned host +buffer instead of the pools, and each page passes through a slot in it: gathered +from its device regions before a write, scattered back to them after a read. The +store then only ever registers host memory. + +A slot holds the page's regions concatenated in region order, which is exactly +the payload the zero-copy path produces from the same regions. The stored bytes +are therefore identical either way, so a pool written by one path is readable by +the other, including by another engine sharing the pool. + +Copies go through `cudaMemcpyAsync` rather than the batched Triton kernel in +`disaggregation/native/bounce/gather_scatter.py`. That kernel is the better tool +for device-to-device gather, but here one side is host memory, which the copy +engines move over the host link by DMA. +""" + +from typing import List, Optional, Sequence, Tuple + +import torch + +try: + from cuda.bindings import runtime as cudart +except ImportError: + from cuda import cudart + +from tensorrt_llm._utils import CUASSERT +from tensorrt_llm.logger import logger + +__all__ = ["HostStagingPool", "plan_slot_geometry", "sync_stream"] + +#: Stated explicitly rather than inferred from the pointers, which would be +#: wrong for a host pointer outside the unified address space. +_DEVICE_TO_HOST = cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost +_HOST_TO_DEVICE = cudart.cudaMemcpyKind.cudaMemcpyHostToDevice + + +def _memcpy_async(dst: int, src: int, size: int, kind, stream: int) -> None: + """One asynchronous copy between a device page and a host slot.""" + status = cudart.cudaMemcpyAsync(int(dst), int(src), int(size), kind, stream)[0] + if status == cudart.cudaError_t.cudaSuccess: + return + # Raised here rather than through CUASSERT so the operands appear in the + # message; a bare cudaErrorInvalidValue names no cause. + device = torch.cuda.current_device() if torch.cuda.is_available() else None + raise RuntimeError( + f"cudaMemcpyAsync failed with {status} staging a KV page: " + f"dst={int(dst):#x} src={int(src):#x} size={size} " + f"stream={int(stream):#x} current_device={device}. An invalid value here " + "is usually a stream created on a different device than the pages, which " + "happens when a thread issues the copy without inheriting the rank's " + "device, since torch's current device is thread-local." + ) + + +def sync_stream(stream: int) -> None: + """Wait for a stream's copies to finish, given its raw handle.""" + CUASSERT(cudart.cudaStreamSynchronize(stream)) + + +def plan_slot_geometry( + max_bytes_per_page: int, + transfer_batch_size: int, + budget_bytes: int, +) -> Tuple[int, int]: + """Choose how many pages may be staged at once, and how wide a slot is. + + A slot has to hold the largest page any layer group produces, so the page size + is a floor on the allocation: a budget below one page is raised to one rather + than refused, since the alternative is not starting. + + Args: + max_bytes_per_page: Largest page payload across layer groups. + transfer_batch_size: Pages the connector puts in one store call. There is + no point staging more than that. + budget_bytes: Ceiling on this pool's pinned allocation. + + Returns: + Slot width in bytes, and the number of slots. + """ + if max_bytes_per_page <= 0: + raise ValueError(f"max_bytes_per_page must be > 0, got {max_bytes_per_page}") + if transfer_batch_size <= 0: + raise ValueError(f"transfer_batch_size must be > 0, got {transfer_batch_size}") + + affordable = budget_bytes // max_bytes_per_page + num_slots = max(1, min(transfer_batch_size, affordable)) + return max_bytes_per_page, num_slots + + +class HostStagingPool: + """A registered pinned buffer, sliced into per-page slots. + + One pool serves one direction. Loads run on the executor thread and saves + on the connector's background thread, so sharing slots between them would + need a lock on the transfer path for no benefit. + """ + + def __init__( + self, + *, + slot_bytes: int, + num_slots: int, + store, + label: str, + ): + self._slot_bytes = int(slot_bytes) + self._num_slots = int(num_slots) + self._label = label + + # Page-locking is a correctness requirement here rather than a + # copy-speed preference: this memory is handed to the store to register. + pin = torch.cuda.is_available() + self._buffer = torch.empty( + self._slot_bytes * self._num_slots, dtype=torch.uint8, pin_memory=pin + ) + self._base = int(self._buffer.data_ptr()) + + status = store.register_buffer(self._base, self._buffer.numel()) + if status != 0: + raise RuntimeError( + f"MooncakeDistributedStore.register_buffer failed with status " + f"{status} for the {label} host staging buffer at " + f"[{self._base:#x}, {self._base + self._buffer.numel():#x}). Host " + f"memory registration failing points at the pool or the fabric " + f"rather than at GPUDirect, which is what staging avoids." + ) + logger.info( + f"mooncake-store {label} staging: {self._num_slots} slots x " + f"{self._slot_bytes} B = {self._buffer.numel() / 1024**2:.1f} MiB pinned " + f"(pinned={pin})" + ) + + @property + def num_slots(self) -> int: + """Pages this pool can hold at once.""" + return self._num_slots + + @property + def slot_bytes(self) -> int: + """Capacity of one slot.""" + return self._slot_bytes + + def slot_address(self, index: int) -> int: + """Address of slot `index`.""" + if not 0 <= index < self._num_slots: + raise IndexError(f"slot {index} out of range [0, {self._num_slots})") + return self._base + index * self._slot_bytes + + def _check_fits(self, total: int) -> None: + if total > self._slot_bytes: + raise ValueError( + f"a {total} B page does not fit the {self._slot_bytes} B " + f"{self._label} staging slot; the pool was sized from the layout's " + "largest page, so this means the layout changed after registration" + ) + + def gather( + self, + index: int, + addresses: Sequence[int], + sizes: Sequence[int], + stream: int, + ) -> Tuple[int, int]: + """Copy one page's device regions into slot `index`, concatenated. + + Args: + index: Slot to fill. + addresses: Device addresses of the page's regions, in region order. + sizes: Byte counts matching `addresses`. + stream: CUDA stream handle the copies are issued on. + + Returns: + The slot's address and the total bytes written, ready to hand to the + store as a single buffer. + """ + total = sum(sizes) + self._check_fits(total) + destination = self.slot_address(index) + offset = 0 + for address, size in zip(addresses, sizes, strict=True): + _memcpy_async(destination + offset, address, size, _DEVICE_TO_HOST, stream) + offset += size + return destination, total + + def scatter( + self, + index: int, + addresses: Sequence[int], + sizes: Sequence[int], + stream: int, + ) -> None: + """Copy slot `index` back out to one page's device regions. + + The inverse of :meth:`gather`, walking the regions in the same order so + the split matches the concatenation the slot holds. + """ + self._check_fits(sum(sizes)) + source = self.slot_address(index) + offset = 0 + for address, size in zip(addresses, sizes, strict=True): + _memcpy_async(address, source + offset, size, _HOST_TO_DEVICE, stream) + offset += size + + def reserve(self, total: int) -> None: + """Assert a page of `total` bytes is stageable, without copying.""" + self._check_fits(total) + + +def stage_batch_for_put( + pool: HostStagingPool, + addresses: Sequence[Sequence[int]], + sizes: Sequence[Sequence[int]], + stream: int, +) -> Tuple[List[List[int]], List[List[int]]]: + """Gather a batch of device pages into slots and describe them for the store. + + Args: + pool: Slots to stage through. The batch must not exceed its slot count. + addresses: Per-page device region addresses. + sizes: Per-page device region sizes. + stream: Stream the copies are issued on. The caller must synchronize it + before the store reads the slots. + + Returns: + Per-page address and size lists, each a single staged buffer. + """ + if len(addresses) > pool.num_slots: + raise ValueError(f"batch of {len(addresses)} pages exceeds {pool.num_slots} staging slots") + staged_addresses: List[List[int]] = [] + staged_sizes: List[List[int]] = [] + for index, (page_addresses, page_sizes) in enumerate(zip(addresses, sizes, strict=True)): + slot, total = pool.gather(index, page_addresses, page_sizes, stream) + staged_addresses.append([slot]) + staged_sizes.append([total]) + return staged_addresses, staged_sizes + + +def describe_batch_for_get( + pool: HostStagingPool, + sizes: Sequence[Sequence[int]], +) -> Tuple[List[List[int]], List[List[int]]]: + """Describe slots for the store to read a batch into, before scattering. + + Unlike the put direction there is nothing to copy first: the slots are the + destination, and :func:`unstage_batch_after_get` moves the bytes on once the + store has filled them. + """ + if len(sizes) > pool.num_slots: + raise ValueError(f"batch of {len(sizes)} pages exceeds {pool.num_slots} staging slots") + staged_addresses: List[List[int]] = [] + staged_sizes: List[List[int]] = [] + for index, page_sizes in enumerate(sizes): + total = sum(page_sizes) + pool.reserve(total) + staged_addresses.append([pool.slot_address(index)]) + staged_sizes.append([total]) + return staged_addresses, staged_sizes + + +def unstage_batch_after_get( + pool: HostStagingPool, + addresses: Sequence[Sequence[int]], + sizes: Sequence[Sequence[int]], + stream: int, + only: Optional[Sequence[int]] = None, +) -> None: + """Scatter filled slots back to their device pages. + + Args: + pool: Slots the store just wrote into. + addresses: Per-page device region addresses. + sizes: Per-page device region sizes. + stream: Stream the copies are issued on. The caller must synchronize it + before the pages are read. + only: Slot indices to scatter. Defaults to all of them; a caller that + knows some reads failed passes the rest so a failed page is not + written over its device slot with whatever the slot held. + """ + indices = range(len(addresses)) if only is None else only + for index in indices: + pool.scatter(index, addresses[index], sizes[index], stream) diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/validation.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/validation.py new file mode 100644 index 000000000000..2281685a0a5d --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/validation.py @@ -0,0 +1,84 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Startup gates for the Mooncake store connector. + +Every rejection here is a configuration whose failure mode is a wrong answer +rather than a slow one: KV that gets replayed without all of the state it was +computed with. Beam search, attention data parallelism, host and disk cache +tiers, and Mamba caches are rejected for all connectors in `py_executor`, so +they are not repeated. + +Checks run at construction, before any request is admitted, so a bad deployment +fails at startup instead of after the first cache hit. +""" + +from typing import TYPE_CHECKING + +from tensorrt_llm.llmapi.llm_args import TorchLlmArgs + +if TYPE_CHECKING: + from ..kv_cache_layout import KvCacheLayout + +__all__ = ["validate_layout", "validate_llm_args"] + + +def validate_llm_args(llm_args: TorchLlmArgs) -> None: + """Reject parallel and model configurations this connector cannot serve.""" + if getattr(llm_args, "context_parallel_size", 1) > 1: + raise NotImplementedError( + "The mooncake-store connector does not support context parallelism. " + "A stored page is keyed by the tokens it holds, but under context " + "parallelism a rank holds a slice of the sequence rather than whole " + "blocks of it, so the same key would name different bytes on " + "different ranks." + ) + + if getattr(llm_args, "pipeline_parallel_size", 1) > 1: + raise NotImplementedError( + "The mooncake-store connector does not support pipeline parallelism. " + "Keys are namespaced per rank, so each stage would store only its own " + "layers and a prefix hit would require every stage to agree; that path " + "is untested. Run with tensor parallelism only." + ) + + sparse_config = getattr(llm_args, "sparse_attention_config", None) + if sparse_config is not None and not getattr(sparse_config, "sparse_disable_index_value", True): + raise NotImplementedError( + "The mooncake-store connector requires " + "sparse_attention_config.sparse_disable_index_value=True. The index-V " + "cache is a plain tensor outside the KV cache manager's paged pools, " + "so it is neither described to the connector nor transferred; a " + "replayed prefix would carry index-K from the store alongside stale " + "index-V. This is the same restriction disaggregated serving applies." + ) + + +def validate_layout(layout: "KvCacheLayout") -> None: + """Reject KV cache geometries this connector cannot key correctly.""" + windowed = [group.layer_group_id for group in layout.groups if group.window_size is not None] + if windowed: + raise NotImplementedError( + "The mooncake-store connector does not support sliding-window " + f"attention (layer groups {windowed} declare a window size). A page's " + "validity then depends on where the window sits, which is a property " + "of the request that read it rather than of the tokens it holds, so " + "content-addressed reuse across instances is not sound." + ) + + if not layout.groups: + raise ValueError( + "The KV cache layout describes no layer groups, so there is nothing " + "for the mooncake-store connector to transfer." + ) diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py new file mode 100644 index 000000000000..6209d44c7e95 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py @@ -0,0 +1,653 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Worker side of the Mooncake store KV cache connector. + +One worker per rank owns a `MooncakeDistributedStore` handle and moves pages +between that pool and its own GPU KV cache. It is also the only place that knows +how a page is addressed and how a key is spelled, which is why the leader, +colocated with rank 0's worker in the same process, asks it to run prefix +lookups instead of rebuilding that knowledge. + +Loads are synchronous: the runtime has already told the scheduler those tokens +are computed, so the bytes must be in place before the forward pass reads them, +and a failed load is a wrong answer rather than a slow one. + +Saves are asynchronous and gated on a CUDA event. The pages are only complete +once the forward pass that wrote them has retired, and blocking the executor +loop on an RDMA write is exactly the cost the store is supposed to avoid. The +scheduler reports such a request as saving asynchronously, which keeps its pages +pinned until `get_finished` says the writes landed. +""" + +import threading +import traceback +from collections import defaultdict +from queue import Queue +from typing import Dict, List, Optional, Sequence, Set, Tuple + +import torch + +from tensorrt_llm._utils import mpi_rank, mpi_world_size +from tensorrt_llm.llmapi.llm_args import TorchLlmArgs +from tensorrt_llm.logger import logger + +from ..kv_cache_connector import KvCacheConnectorWorker +from ..kv_cache_layout import KvCacheLayout +from .addressing import PageAddressing +from .config import CONFIG_PATH_ENV, MooncakeStoreConnectorConfig +from .keys import KeyNamespace +from .metadata import MooncakeStoreMetadata, RequestTransfers +from .staging import ( + HostStagingPool, + describe_batch_for_get, + plan_slot_geometry, + stage_batch_for_put, + unstage_batch_after_get, +) +from .staging import sync_stream as _sync_stream +from .validation import validate_layout, validate_llm_args + +__all__ = ["MooncakeStoreConnectorWorker", "resolve_local_worker"] + +#: Set by the worker's constructor so the leader, which the executor builds in +#: the same process on rank 0, can reach the store handle without a second +#: connection or an out-of-band channel. See `py_executor_creator`, which +#: constructs scheduler and worker concurrently for exactly this kind of +#: mutual dependency. +_LOCAL_WORKER: Optional["MooncakeStoreConnectorWorker"] = None +_LOCAL_WORKER_READY = threading.Event() + + +def resolve_local_worker(timeout: float = 60.0) -> "MooncakeStoreConnectorWorker": + """The worker living in this process, once it has been constructed. + + Args: + timeout: Seconds to wait. Construction is concurrent with the leader's, + so a short wait is expected; exceeding it means the worker failed. + + Returns: + The process-local worker. + """ + if not _LOCAL_WORKER_READY.wait(timeout): + raise RuntimeError( + "The mooncake-store leader could not find a worker in its process. " + "The leader only runs on rank 0, where the executor also builds a " + "worker, so this means worker construction failed." + ) + assert _LOCAL_WORKER is not None + return _LOCAL_WORKER + + +def _open_store(config: MooncakeStoreConnectorConfig): + """Connect to the Mooncake master and return a live store handle.""" + try: + from mooncake.store import MooncakeDistributedStore + except ImportError as exc: + raise ImportError( + "The mooncake-store connector needs the Mooncake Python bindings " + "(`pip install mooncake-transfer-engine`). The C++ transfer engine " + "built into the container is a different component and does not " + "provide MooncakeDistributedStore." + ) from exc + + store = MooncakeDistributedStore() + hostname = config.local_hostname or _default_hostname() + setup_kwargs = {} + if config.tenant_id: + setup_kwargs["tenant_id"] = config.tenant_id + status = store.setup( + hostname, + config.metadata_server, + config.global_segment_size, + config.local_buffer_size, + config.protocol, + config.device_name, + config.master_server_address, + **setup_kwargs, + ) + if status != 0: + raise RuntimeError( + f"MooncakeDistributedStore.setup failed with status {status} " + f"(master={config.master_server_address!r}, " + f"metadata={config.metadata_server!r}, protocol={config.protocol!r}). " + f"Check the config named by {CONFIG_PATH_ENV}." + ) + return store + + +def _default_hostname() -> str: + import socket + + return socket.gethostbyname(socket.gethostname()) + + +def _batched(items: Sequence, size: int): + for start in range(0, len(items), size): + yield items[start : start + size] + + +def _stream_handle(stream) -> int: + """The raw CUDA stream handle behind a torch stream, or a handle as given. + + `None` maps to 0, the default stream, which is what the runtime passes when + it has no stream of its own to offer. + """ + if stream is None: + return 0 + return int(getattr(stream, "cuda_stream", stream)) + + +class MooncakeStoreConnectorWorker(KvCacheConnectorWorker): + """Moves KV pages between this rank's GPU cache and the Mooncake pool.""" + + def __init__(self, llm_args: TorchLlmArgs): + super().__init__(llm_args) + + validate_llm_args(llm_args) + self._config = MooncakeStoreConnectorConfig.from_env() + self._rank = mpi_rank() + self._world_size = mpi_world_size() + self._model_key = self._config.resolve_model_key(llm_args.model) + + self._addressing: Optional[PageAddressing] = None + # Namespaces for this rank, used for both directions of transfer. + self._namespaces: Dict[int, KeyNamespace] = {} + # The same namespaces for every rank. A prefix is only reusable when all + # shards of it are present, so a lookup has to ask about all of them. + self._peer_namespaces: Dict[int, Tuple[KeyNamespace, ...]] = {} + + self._store = _open_store(self._config) + + self._save_queue: "Queue[Optional[Tuple[torch.cuda.Event, List[RequestTransfers]]]]" = ( + Queue() + ) + self._save_thread: Optional[threading.Thread] = None + self._save_lock = threading.Lock() + # Host staging, when the pool cannot register device memory. + self._load_staging: Optional[HostStagingPool] = None + self._save_staging: Optional[HostStagingPool] = None + self._save_stream: Optional[torch.cuda.Stream] = None + # This rank's device, captured on the executor thread; see _drain_saves. + self._device_index: Optional[int] = None + # Pages per store call. Staging narrows this to the slots it can afford. + self._batch_size = self._config.transfer_batch_size + # Save submissions still in flight, per request. + self._outstanding_saves: Dict[int, int] = defaultdict(int) + # Requests the runtime has told us are done producing KV. Their pages + # stay pinned until we report them back through `get_finished`. + self._closed_requests: Set[int] = set() + self._save_error: Optional[BaseException] = None + + global _LOCAL_WORKER + _LOCAL_WORKER = self + _LOCAL_WORKER_READY.set() + + logger.info( + f"mooncake-store worker rank {self._rank}/{self._world_size} ready " + f"(role={self._config.role.value}, model_key={self._model_key}, " + f"master={self._config.master_server_address})" + ) + + # ---- registration ---- + + def register_kv_caches(self, kv_cache_tensor: torch.Tensor): + """Reject the V1 single-pool registration. + + Raises: + NotImplementedError: Always. Identity here is a hash chain the + connector computes itself, keyed per layer group, and the V1 + manager supplies real block hashes over a single flat block + space instead. Running the V2 addressing against V1 block ids + would silently mislabel pages, so V1 is refused rather than + approximated. + """ + raise NotImplementedError( + "The mooncake-store connector requires KVCacheManagerV2. Set " + "kv_cache_config.use_kv_cache_manager_v2=True." + ) + + def register_kv_cache_layout(self, layout: KvCacheLayout) -> None: + """Register the KV pools with Mooncake and start the save thread.""" + if self._addressing is not None: + raise RuntimeError("KV cache layout already registered") + + validate_layout(layout) + addressing = PageAddressing(layout) + # Torch's current device is thread-local, so read it here on the + # executor thread; the save thread would otherwise see device 0. + if torch.cuda.is_available(): + self._device_index = torch.cuda.current_device() + if self._config.stage_through_host: + self._open_staging(addressing) + else: + for start, end in addressing.registration_ranges(): + status = self._store.register_buffer(start, end - start) + if status != 0: + raise RuntimeError( + f"MooncakeDistributedStore.register_buffer failed with status " + f"{status} for [{start:#x}, {end:#x}). Without registration " + "the store cannot read or write these pages. Registering " + "device memory needs GPUDirect RDMA (nvidia_peermem or " + "dma-buf); where that is unavailable, set " + "stage_through_host to pass pages through pinned host " + "memory instead." + ) + + self._addressing = addressing + for layer_group_id in addressing.layer_group_ids: + bytes_per_page = addressing.bytes_per_page(layer_group_id) + self._namespaces[layer_group_id] = self._namespace( + self._rank, layer_group_id, bytes_per_page + ) + self._peer_namespaces[layer_group_id] = tuple( + self._namespace(rank, layer_group_id, bytes_per_page) + for rank in range(self._world_size) + ) + + if self._config.role.saves: + self._save_thread = threading.Thread( + target=self._drain_saves, + name=f"mooncake-store-save-{self._rank}", + daemon=True, + ) + self._save_thread.start() + + logger.info( + f"mooncake-store worker rank {self._rank} registered layout: {addressing.describe()}" + ) + + def _open_staging(self, addressing: PageAddressing) -> None: + """Allocate and register the pinned slots pages will pass through. + + Only the directions this role drives get a pool, since each one costs a + pinned allocation of its own. The GPU pools are left unregistered, + which is the point of the mode. + """ + max_bytes_per_page = max( + addressing.bytes_per_page(layer_group_id) + for layer_group_id in addressing.layer_group_ids + ) + slot_bytes, num_slots = plan_slot_geometry( + max_bytes_per_page, + self._config.transfer_batch_size, + self._config.staging_buffer_bytes, + ) + if self._config.role.loads: + self._load_staging = HostStagingPool( + slot_bytes=slot_bytes, + num_slots=num_slots, + store=self._store, + label="load", + ) + if self._config.role.saves: + self._save_staging = HostStagingPool( + slot_bytes=slot_bytes, + num_slots=num_slots, + store=self._store, + label="save", + ) + self._batch_size = min(self._config.transfer_batch_size, num_slots) + if self._batch_size < self._config.transfer_batch_size: + logger.warning( + f"mooncake-store rank {self._rank} reduced its transfer batch from " + f"{self._config.transfer_batch_size} to {self._batch_size} pages: " + f"staging {max_bytes_per_page} B pages within " + f"{self._config.staging_buffer_bytes} B does not fit more. Raise " + f"staging_buffer_bytes to restore the configured batch size." + ) + + def _namespace(self, rank: int, layer_group_id: int, bytes_per_page: int) -> KeyNamespace: + return KeyNamespace( + cache_prefix=self._config.cache_prefix, + model_key=self._model_key, + rank=rank, + world_size=self._world_size, + layer_group_id=layer_group_id, + tokens_per_block=self._addressing.tokens_per_block, + bytes_per_page=bytes_per_page, + ) + + # ---- leader-facing lookup ---- + + @property + def config(self) -> MooncakeStoreConnectorConfig: + """The resolved connector configuration.""" + return self._config + + @property + def is_registered(self) -> bool: + """Whether a KV cache layout has been registered yet.""" + return self._addressing is not None + + def count_prefix_hit(self, block_hashes: Sequence[bytes]) -> int: + """How many leading blocks of `block_hashes` are fully present. + + A block counts only when every layer group and every rank has its page, + because a prefix is replayed as a whole. The scan stops at the first + incomplete block: the runtime consumes a prefix, so a later hit is not + usable on its own. + + Args: + block_hashes: Candidate hashes in block order. + + Returns: + Length of the usable prefix, in blocks. + """ + if not block_hashes or self._addressing is None: + return 0 + + keys: List[str] = [] + for block_hash in block_hashes: + for namespaces in self._peer_namespaces.values(): + keys.extend(namespace.key(block_hash) for namespace in namespaces) + keys_per_block = len(keys) // len(block_hashes) + + try: + present = self._store.batch_is_exist(keys) + except Exception as exc: + logger.warning( + f"mooncake-store lookup failed; treating as a miss: " + f"{type(exc).__name__}: {exc}\n{traceback.format_exc()}" + ) + return 0 + + if len(present) != len(keys): + logger.warning( + f"mooncake-store batch_is_exist returned {len(present)} results for " + f"{len(keys)} keys; treating as a miss" + ) + return 0 + + hit_blocks = 0 + for index in range(len(block_hashes)): + window = present[index * keys_per_block : (index + 1) * keys_per_block] + # Mooncake reports 1 for present, 0 for absent and a negative value + # for a failed probe. Anything but a definite 1 is treated as a miss. + if not all(status == 1 for status in window): + break + hit_blocks += 1 + return hit_blocks + + # ---- load path ---- + + def start_load_kv(self, stream: torch.cuda.Stream): + """Pull every scheduled page into its GPU slot before the forward pass.""" + metadata: Optional[MooncakeStoreMetadata] = self.get_connector_meta() + if metadata is None or not metadata.loads: + return + self._reraise_save_error() + + keys, addresses, sizes, total_pages = self._resolve(metadata.loads) + if not keys: + return + + staging = self._load_staging + handle = _stream_handle(stream) if staging is not None else 0 + + for batch in zip( + _batched(keys, self._batch_size), + _batched(addresses, self._batch_size), + _batched(sizes, self._batch_size), + ): + batch_keys, batch_addresses, batch_sizes = batch + if staging is None: + target_addresses, target_sizes = list(batch_addresses), list(batch_sizes) + else: + target_addresses, target_sizes = describe_batch_for_get(staging, batch_sizes) + results = self._store.batch_get_into_multi_buffers( + list(batch_keys), target_addresses, target_sizes + ) + failed = [ + key + for key, result in zip(batch_keys, results) + if not isinstance(result, int) or result < 0 + ] + if failed or len(results) != len(batch_keys): + # The runtime already counted these tokens as computed, so a + # partial load leaves the forward pass reading uninitialized KV + # and silently producing wrong tokens. Fail loudly instead. + raise RuntimeError( + f"mooncake-store failed to load {len(failed) or len(batch_keys)} of " + f"{len(batch_keys)} pages; the affected KV slots were already " + f"reported as computed. First failure: {failed[:1]}" + ) + if staging is not None: + # Only reached once every page in the batch landed, so no slot + # holding a failed read is copied over a device page. + unstage_batch_after_get(staging, batch_addresses, batch_sizes, handle) + # The next batch reuses the slots and the forward pass reads + # these pages, so the scatter has to complete before either. + _sync_stream(handle) + + logger.debug(f"mooncake-store rank {self._rank} loaded {total_pages} pages") + + def wait_for_layer_load(self, layer_idx: int, stream: torch.cuda.Stream): + """No-op: loads complete in `start_load_kv`. + + Transfers are whole pages, so a page's bytes for every layer in a group + land in one store call rather than layer by layer. There is nothing left + outstanding by the time the first layer runs. + """ + + def save_kv_layer(self, layer_idx: int, stream: torch.cuda.Stream): + """No-op: saves are submitted once per pass in `wait_for_save`. + + A page is only complete when every layer of its group has written its + slice, so there is no correct per-layer submission point. + """ + + # ---- save path ---- + + def wait_for_save(self, stream: torch.cuda.Stream): + """Hand this pass's saves to the background thread, gated on an event.""" + metadata: Optional[MooncakeStoreMetadata] = self.get_connector_meta() + if metadata is None or not metadata.saves or not self._config.role.saves: + return + self._reraise_save_error() + + # The pages are written by kernels still queued on this stream. The event + # is the handoff: the thread reads GPU memory only after the pass retires, + # and the executor loop is not blocked waiting for that. + event = torch.cuda.Event() + event.record(stream) + + with self._save_lock: + for transfers in metadata.saves: + self._outstanding_saves[transfers.request_id] += 1 + self._save_queue.put((event, list(metadata.saves))) + + def get_finished( + self, finished_gen_req_ids: List[int], started_loading_req_ids: List[int] + ) -> Tuple[List[int], List[int]]: + """Report which requests' saves have landed. + + Args: + finished_gen_req_ids: Requests that will produce no further KV. + started_loading_req_ids: Requests loading asynchronously. Always + empty here, since `get_num_new_matched_tokens` only ever + offers synchronous loads; echoed back so the runtime does not + wait on something that already happened. + + Returns: + Requests that have finished saving, and requests that have finished + loading. + """ + self._reraise_save_error() + with self._save_lock: + self._closed_requests.update(finished_gen_req_ids) + finished_saving = [ + request_id + for request_id in self._closed_requests + if self._outstanding_saves.get(request_id, 0) == 0 + ] + for request_id in finished_saving: + self._closed_requests.discard(request_id) + self._outstanding_saves.pop(request_id, None) + return finished_saving, list(started_loading_req_ids) + + def _drain_saves(self) -> None: + # A new thread starts on device 0, so adopt the device captured on the + # executor thread. Otherwise a stream created below belongs to device 0 + # while the KV pointers belong to the rank's device, and the copy fails + # with cudaErrorInvalidValue on every rank except 0. + if self._device_index is not None: + torch.cuda.set_device(self._device_index) + if self._save_staging is not None and torch.cuda.is_available(): + # Owned by this thread so the gather never queues behind the + # executor's work, and created after set_device so it lands on the + # rank's device. + self._save_stream = torch.cuda.Stream() + while True: + item = self._save_queue.get() + if item is None: + return + event, transfers = item + try: + event.synchronize() + self._put(transfers) + except Exception as exc: + # Broad on purpose: this is the thread boundary. Anything that + # escapes here would be lost, so it is stashed and re-raised on + # the executor thread at the next connector call. + logger.error( + f"mooncake-store save failed on rank {self._rank}: {type(exc).__name__}: {exc}" + ) + with self._save_lock: + if self._save_error is None: + self._save_error = exc + finally: + with self._save_lock: + for entry in transfers: + remaining = self._outstanding_saves.get(entry.request_id, 0) - 1 + if remaining <= 0: + self._outstanding_saves.pop(entry.request_id, None) + else: + self._outstanding_saves[entry.request_id] = remaining + + def _put(self, transfers: Sequence[RequestTransfers]) -> None: + keys, addresses, sizes, _ = self._resolve(transfers) + if not keys: + return + + staging = self._save_staging + handle = _stream_handle(self._save_stream) if staging is not None else 0 + + for batch in zip( + _batched(keys, self._batch_size), + _batched(addresses, self._batch_size), + _batched(sizes, self._batch_size), + ): + batch_keys, batch_addresses, batch_sizes = batch + # Skip pages another rank or another instance already wrote. The + # scheduler cannot know this: it holds no store handle, and the + # answer changes between the time it builds metadata and now. + present = self._store.batch_is_exist(list(batch_keys)) + pending = [ + index + for index, status in enumerate(present) + if status != 1 # absent, or a failed probe we retry as a write + ] + if not pending: + continue + source_addresses = [batch_addresses[i] for i in pending] + source_sizes = [batch_sizes[i] for i in pending] + if staging is not None: + # Gathered after the existence filter, so a page already in the + # pool costs no copy. + source_addresses, source_sizes = stage_batch_for_put( + staging, source_addresses, source_sizes, handle + ) + # The store reads the slots on this thread, so fill them first. + _sync_stream(handle) + results = self._store.batch_put_from_multi_buffers( + [batch_keys[i] for i in pending], + source_addresses, + source_sizes, + ) + failures = sum(1 for result in results if not isinstance(result, int) or result < 0) + if failures: + # A dropped write only costs a future cache miss, so it is worth + # a warning rather than failing a request that already answered. + logger.warning( + f"mooncake-store rank {self._rank} failed to save {failures} of " + f"{len(pending)} pages" + ) + + # ---- shared ---- + + def _resolve( + self, transfers: Sequence[RequestTransfers] + ) -> Tuple[List[str], List[List[int]], List[List[int]], int]: + """Expand per-request page transfers into parallel store call arguments.""" + if self._addressing is None: + raise RuntimeError("KV cache layout has not been registered") + keys: List[str] = [] + addresses: List[List[int]] = [] + sizes: List[List[int]] = [] + pages = 0 + for entry in transfers: + for page in entry.pages: + namespace = self._namespaces.get(page.layer_group_id) + if namespace is None: + raise KeyError( + f"layer group {page.layer_group_id} is not in the registered " + "layout; the scheduler and worker disagree about the model" + ) + page_addresses, page_sizes = self._addressing.buffers( + page.layer_group_id, page.page_index + ) + keys.append(namespace.key(page.block_hash)) + addresses.append(page_addresses) + sizes.append(page_sizes) + pages += 1 + return keys, addresses, sizes, pages + + def _reraise_save_error(self) -> None: + with self._save_lock: + error = self._save_error + self._save_error = None + if error is not None: + raise RuntimeError("mooncake-store background save failed") from error + + def shutdown(self) -> None: + """Stop the save thread and release the store handle. Idempotent.""" + thread, self._save_thread = self._save_thread, None + if thread is not None: + self._save_queue.put(None) + thread.join(timeout=30.0) + store, self._store = self._store, None + if store is not None: + try: + store.close() + except Exception as exc: + logger.warning( + f"mooncake-store close failed: {type(exc).__name__}: {exc}\n" + f"{traceback.format_exc()}" + ) + # Released only after the store is closed, since it holds registrations + # against this memory. + self._load_staging = None + self._save_staging = None + self._save_stream = None + global _LOCAL_WORKER + if _LOCAL_WORKER is self: + _LOCAL_WORKER = None + _LOCAL_WORKER_READY.clear() + + def __del__(self): + try: + self.shutdown() + except Exception: # noqa: S110 - interpreter teardown, nothing left to report to + pass diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/registry.py b/tensorrt_llm/_torch/pyexecutor/connectors/registry.py index 9a00cdadd7fd..dcf5e24f3642 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/registry.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/registry.py @@ -19,6 +19,11 @@ it is resolved at runtime via importlib in py_executor_creator.py. """ +from typing import TYPE_CHECKING, Optional + +if TYPE_CHECKING: + from tensorrt_llm.llmapi.llm_args import KvCacheConnectorConfig + CONNECTOR_REGISTRY: dict[str, dict[str, str]] = { "lmcache": { "connector_module": "lmcache.integration.tensorrt_llm.tensorrt_adapter", @@ -35,4 +40,26 @@ "connector_scheduler_class": "DynamoKVBMConnectorLeader", "connector_worker_class": "DynamoKVBMConnectorWorker", }, + "mooncake-store": { + "connector_module": "tensorrt_llm._torch.pyexecutor.connectors.mooncake_store", + "connector_scheduler_class": "MooncakeStoreConnectorScheduler", + "connector_worker_class": "MooncakeStoreConnectorWorker", + }, } + + +def uses_connector(kv_connector_config: Optional["KvCacheConnectorConfig"], name: str) -> bool: + """Report whether a connector config resolves to the named preset. + + Compares the resolved module rather than the `connector` field, so a config + that names the module explicitly instead of using the preset is still + recognized. Accepts `None` to save every caller a null check. + """ + if kv_connector_config is None: + return False + preset = CONNECTOR_REGISTRY.get(name) + if preset is None: + raise ValueError( + f"Unknown connector preset: {name!r}. Known presets: {list(CONNECTOR_REGISTRY)}" + ) + return kv_connector_config.connector_module == preset["connector_module"] diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 944036b6155b..2f711ed2ab17 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -980,13 +980,14 @@ def append_to_kv_heads_per_layer( cache_tiers: List[CacheTierConfig] = [GpuCacheTierConfig(quota=int(quota))] if kv_connector_manager is not None and kv_cache_config.host_cache_size is None: # A KV connector registers device addresses for its pages, and a - # page evicted to another tier has its GPU slot reassigned. The - # automatic host tier below exists only to give the MAX_UTILIZATION - # scheduler's suspend/resume somewhere to spill to, and a connector - # run cannot use that policy (py_executor_creator requires - # GUARANTEED_NO_EVICT), so skip it rather than silently migrating - # pages out from under the connector. An explicitly configured - # host_cache_size is left alone and rejected loudly at bring-up. + # page evicted to another tier has its GPU slot reassigned, so the + # automatic host tier below would silently migrate pages out from + # under the connector. An explicitly configured host_cache_size is + # left alone and rejected loudly at bring-up. + # + # That leaves the scheduler without a tier to spill to, where + # suspension frees nothing, so pages are reclaimed by preemption + # instead. See KVCacheManagerV2.preempt_request. host_quota = 0 logger.info( "KV cache manager v2 host tier disabled: a KV connector is attached " @@ -1173,6 +1174,9 @@ def append_to_kv_heads_per_layer( ) self.index_mapper = IndexMapper(index_mapper_capacity, max_beam_width) self._early_freed_index_requests: set[int] = set() + # Requests whose pages a connector is still reading from, so the + # release half of `preempt_request` has to wait. + self._pending_preemption: Dict[int, LlmRequest] = {} self._prepare_page_table_tensor(index_mapper_capacity) self._log_kv_cache_pool_lifecycle_mapping() @@ -2614,6 +2618,85 @@ def resume_request(self, req: LlmRequest) -> bool: return False return self._resume_and_restore(req.py_request_id, kv_cache) + # ---- preemption ---- + # + # Suspension only unpins pages; the eviction controller then migrates them + # one cache level down. With GPU as the last level a suspended page stays + # `HELD`, which `CacheLevelManager.is_evictable` refuses to evict, so + # suspension frees nothing and the scheduler has no way out of a full pool. + # + # Preemption is the fallback for that case. It gives the pages up instead + # of parking them, which costs a re-prefill but always works. + + @property + def has_cache_tier_below_gpu(self) -> bool: + """True when a suspended page has somewhere to be evicted to.""" + return len(self.impl.cache_tier_list) > 1 + + def has_pending_preemption(self) -> bool: + """True while a deferred preemption is still waiting on a connector.""" + return bool(self._pending_preemption) + + def preempt_request(self, req: LlmRequest) -> bool: + """Give up *req*'s KV cache so its pages can be reclaimed. + + Unlike :meth:`suspend_request` this does not keep the pages. Closing + the request's `_KVCache` returns its committed blocks to the radix tree + as reusable prefix and leaves their pages `DROPPABLE`, which is + evictable at every level, unlike `HELD`. The data is not thrown away: + it stays resident and locally matchable until something else needs the + space. + + The request is reset to context state by the caller and re-prefills + whatever it can no longer match. With a connector attached, blocks it + already wrote to the store come back through the ordinary prefix load, + and recompute is the fallback when the store no longer has them. + + Returns True when the pages were released. When the connector still has + saves in flight the release is deferred and this returns False, because + those saves read directly out of these pages: freeing them now would + let a later request overwrite the bytes mid-transfer and publish them + under a valid hash. Callers must not count on the pages until + :meth:`try_complete_preemption` has run for this request. + """ + if self.kv_connector_manager is None: + self._release_preempted(req) + return True + + # The same handshake the finish path uses: the request lands in + # DISAGG_CONTEXT_TRANS_IN_PROGRESS, out of the schedulable range, and + # its `_KVCache` keeps holding the pages until every rank reports the + # save retired through `get_finished`. + if self.kv_connector_manager.request_finished(req, self.get_connector_page_indices(req)): + self._pending_preemption[req.py_request_id] = req + return False + + self._release_preempted(req) + return True + + def try_complete_preemption(self, req: LlmRequest) -> bool: + """Release pages for a request whose deferred preemption just cleared. + + Returns False when *req* was not awaiting preemption, which is how the + caller tells a preempted request apart from an ordinary finished one in + the connector's `get_finished` output. + """ + if self._pending_preemption.pop(req.py_request_id, None) is None: + return False + self._release_preempted(req) + return True + + def _release_preempted(self, req: LlmRequest) -> None: + self.free_resources(req) + # Ask the connector again on re-admission rather than reusing the + # memoised offer, since the store now has more of this prefix than it + # did when the request was first admitted. + req.py_connector_prefix_start = None + req.py_connector_prefix_end = None + req.py_connector_load_async = False + req.py_connector_delivered = False + req.py_num_connector_matched_tokens = 0 + # ---- prepare_resources ---- @nvtx_range("prepare_resources_kv_cache_manager_v2") @@ -3430,6 +3513,10 @@ def release_index_slot(self, request_id: int) -> None: self._early_freed_index_requests.add(request_id) def free_resources(self, request: LlmRequest, pin_on_release: bool = False): + # A request awaiting preemption can still be cancelled or fail while + # its saves drain. Dropping the entry here keeps a dead request from + # blocking every later preemption through has_pending_preemption. + self._pending_preemption.pop(request.py_request_id, None) self._release_undelivered_connector_prefix(request) if self.conversation_manager is not None: self.conversation_manager.finish_request(request) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 534035e33ac2..0b65452ab01a 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3750,8 +3750,30 @@ def _kv_connector_terminate_requests(self): if self.kv_connector_manager: reqs_to_terminate = self.kv_connector_manager.get_finished() for req in reqs_to_terminate: + if self._resume_preempted_request(req): + continue self._end_transfer_and_maybe_terminate(req) + def _resume_preempted_request(self, request: LlmRequest) -> bool: + """Complete a preemption whose connector saves have now retired. + + The scheduler preempts a request by handing it to the connector the + same way a finished request is handed over, so its pages stay put until + every rank reports the in-flight saves done. Both kinds come back + through `get_finished`, and only the KV cache manager knows which is + which. + + Returns True when *request* was preempted rather than finished, in + which case its pages are now released and it is back in context state + awaiting a re-prefill. + """ + if not self._is_kv_manager_v2: + return False + if not self.kv_cache_manager.try_complete_preemption(request): + return False + request.pause(self.max_input_len) + return True + def _kv_connector_wait_for_save(self): if self.kv_connector_manager is not None: self.kv_connector_manager.worker.wait_for_save( diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 67a48544c9b6..a3b15f56c734 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -42,6 +42,7 @@ validate_feature_combination) from .config_utils import is_hybrid_linear, is_minimax_m3 from .connectors.kv_cache_connector import KvCacheConnectorManager +from .connectors.registry import uses_connector from .dwdp import DwdpManager from .guided_decoder import CapturableGuidedDecoder, GuidedDecoder from .model_engine import PyTorchModelEngine @@ -382,6 +383,19 @@ def create_py_executor( kv_cache_config.enable_block_reuse = False kv_cache_config.enable_partial_reuse = False + # Must happen before the KV cache manager is built, since the manager reads + # enable_partial_reuse to construct its block pools. + if (kv_cache_config.enable_partial_reuse + and uses_connector(kv_connector_config, "mooncake-store")): + logger.warning( + "Disabling partial reuse: it is not usable with the mooncake-store " + "connector. The store is addressed by whole blocks, so a partial " + "device match leaves the matched length off a block boundary and " + "the connector declines the lookup rather than resume a block from " + "the middle. Partial reuse therefore trades part of one block for " + "every stored block of the remaining prefix.") + kv_cache_config.enable_partial_reuse = False + decoding_config = llm_args.decoding_config # The tokenizer is stripped from MPI kwargs in proxy.py to avoid pickle diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py index dd9b67132834..23e922976e3e 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py @@ -15,7 +15,7 @@ import enum import os -from typing import Optional +from typing import Callable, Optional from tensorrt_llm.llmapi.llm_args import CapacitySchedulerPolicy, ContextChunkingPolicy from tensorrt_llm.logger import logger @@ -154,8 +154,13 @@ def __init__( draft_kv_cache_manager: object | None = None, # KVCacheManagerV2 for MTP draft layers cross_kv_cache_manager: object | None = None, # KVCacheManagerV2 for enc-dec cross-attn enable_prefix_aware_scheduling: bool = True, + max_input_len: int = 0x7FFFFFFF, ) -> None: self.max_num_tokens = max_num_tokens + # Only read when preempting: LlmRequest.pause clamps the rewritten + # prompt, the original plus tokens generated so far, to this. + self.max_input_len = max_input_len + self._stalled_schedules = 0 self.max_num_requests = ( scheduler_capacity if scheduler_capacity is not None else max_batch_size ) @@ -398,11 +403,26 @@ def _schedule_loop(self, active_requests, inflight_request_ids): req_it += 1 + # Requests whose pages were given up during this pass. A victim that is + # itself a started context request still sits in pending_ctx, so + # re-admitting it would spend the pages its own preemption released. + preempted_ids: set[int] = set() + + def preempt_for_pages(req: LlmRequest) -> bool: + protected = {r.py_request_id for r in scheduled_gen} + protected.update(r.py_request_id for r in scheduled_ctx) + protected.add(req.py_request_id) + return self._try_preempt_for_pages( + requests_list, protected, inflight_request_ids, evicted, preempted_ids + ) + # --- Phase 2: schedule deferred context / encoder requests --- # Generation PEFT pages are now fully committed in the budget. for req in pending_ctx: if budget.requests_full: break + if req.py_request_id in preempted_ids: + continue peft_pages = budget.peft_pages_needed(req) if peft_pages is None: continue @@ -413,7 +433,9 @@ def _schedule_loop(self, active_requests, inflight_request_ids): scheduled_encoder.append(req) budget.commit(req, tokens, peft_pages) else: - action, tokens, chunking_flag = self._try_schedule_context(req, budget) + action, tokens, chunking_flag = self._try_schedule_context( + req, budget, preempt_for_pages + ) if action is ScheduleAction.STOP: break if action is ScheduleAction.SKIP: @@ -422,28 +444,15 @@ def _schedule_loop(self, active_requests, inflight_request_ids): scheduled_ctx.append(req) budget.commit(req, tokens, peft_pages) - # Deadlock detection: if generation requests exist but none were - # scheduled and none were evicted, no forward pass will run and no - # KV cache pages will ever be freed — the scheduler will spin - # forever. This typically happens when the KV cache pool is - # exhausted and no host cache tier is available for suspend/resume. - if not scheduled_gen and not scheduled_ctx: - num_gen_candidates = sum( - 1 - for r in active_requests - if r.is_generation_in_progress_state - and not r.is_generation_to_complete_state - and r.request_id not in inflight_request_ids - ) - if num_gen_candidates > 0 and not evicted: - raise RuntimeError( - f"V2 scheduler deadlock: {num_gen_candidates} generation " - f"request(s) active but none could be scheduled or " - f"evicted. KV cache pool is likely exhausted with no " - f"host cache tier for suspend/resume offload. " - f"Configure kv_cache_config.host_cache_size or increase " - f"kv_cache_config.max_tokens." - ) + self._detect_deadlock( + active_requests, + inflight_request_ids, + pending_ctx, + preempted_ids, + made_progress=bool( + scheduled_gen or scheduled_ctx or scheduled_encoder or disagg_candidates or evicted + ), + ) return ( scheduled_encoder, @@ -505,7 +514,10 @@ def _try_schedule_disagg_gen_init( return ScheduleAction.SCHEDULED, 0 def _try_schedule_context( - self, req: LlmRequest, budget: BudgetTracker + self, + req: LlmRequest, + budget: BudgetTracker, + preempt_for_pages: Callable[[LlmRequest], bool], ) -> tuple[ScheduleAction, int, bool]: """Try to schedule a context request (chunked or non-chunked). @@ -521,11 +533,14 @@ def _try_schedule_context( # connector being asked or told twice is the per-request query state # (see KVCacheManagerV2._connector_prefix_position). if self.chunking_enabled: - return self._try_schedule_context_chunked(req, budget) - return self._try_schedule_context_full(req, budget) + return self._try_schedule_context_chunked(req, budget, preempt_for_pages) + return self._try_schedule_context_full(req, budget, preempt_for_pages) def _try_schedule_context_full( - self, req: LlmRequest, budget: BudgetTracker + self, + req: LlmRequest, + budget: BudgetTracker, + preempt_for_pages: Callable[[LlmRequest], bool], ) -> tuple[ScheduleAction, int, bool]: """Try to schedule a non-chunked context request. @@ -564,6 +579,11 @@ def _try_schedule_context_full( # V2 resizes KV cache directly in the scheduler (no separate # prepareResources for main cache), so include draft tokens. if not self.kv_cache_manager.resize_context(req, context_tokens + draft_len): + # Out of pages. Give up one started request so this one can + # proceed, and retry next iteration: a failed resize leaves a + # first chunk suspended, so the retry has to go back through + # prepare_context to resume it. + preempt_for_pages(req) return ScheduleAction.SKIP, 0, False cross_action = self._try_schedule_cross_context(req) @@ -574,7 +594,10 @@ def _try_schedule_context_full( return ScheduleAction.SCHEDULED, req_tokens, False def _try_schedule_context_chunked( - self, req: LlmRequest, budget: BudgetTracker + self, + req: LlmRequest, + budget: BudgetTracker, + preempt_for_pages: Callable[[LlmRequest], bool], ) -> tuple[ScheduleAction, int, bool]: """FCFS interleaved chunking for a single context request. @@ -635,10 +658,9 @@ def _try_schedule_context_chunked( chunk_size = (chunk_size // self.chunk_unit_size) * self.chunk_unit_size if chunk_size <= 0: - # TODO: consider suspending first-chunk KVCache to release - # GPU pages. Currently we skip without suspend to avoid - # pathological suspend/resume cycles. suspend_request is - # only called from eviction (_try_evict_for_gen). + # Out of token budget rather than out of pages, so releasing pages + # would not help; the next iteration gets a fresh budget. Not + # suspended either, to avoid pathological suspend/resume cycles. return ScheduleAction.SKIP, 0, False chunk_size = self._align_chunk_to_mm_block( @@ -666,6 +688,8 @@ def _try_schedule_context_chunked( # V2 resizes KV cache directly in the scheduler, so include # draft tokens for last chunk. if not self.kv_cache_manager.resize_context(req, resize_tokens): + # Out of pages, as in _try_schedule_context_full. + preempt_for_pages(req) return ScheduleAction.SKIP, 0, False cross_action = self._try_schedule_cross_context(req) @@ -1010,6 +1034,124 @@ def _suspend_request(self, req: LlmRequest) -> None: def _clear_request_runtime_state(self, req: LlmRequest) -> None: req.py_batch_idx = None + def _try_preempt_for_pages( + self, + requests_list: RequestList, + protected_ids: set[int], + inflight_request_ids: set[int], + evicted: RequestList, + preempted_ids: set[int], + ) -> bool: + """Release one started request's KV cache so another can allocate. + + The fallback for a pool that suspension cannot drain; see + `KVCacheManagerV2.preempt_request`. With a cache tier below GPU, + suspension is cheaper and keeps the pages, so that path is left alone. + + Returns True when pages became available in this iteration. A + connector defers the release until its in-flight saves retire, in + which case this returns False and the pages arrive a few iterations + later. + """ + if self.kv_cache_manager.has_cache_tier_below_gpu: + return False + + if self.kv_cache_manager.has_pending_preemption(): + # One victim at a time, or a full pool would preempt the whole + # batch while the first release is still draining. + return False + + # Newest first, so the requests closest to completing keep their + # pages and the pool drains instead of thrashing. + for i in range(len(requests_list) - 1, -1, -1): + victim = requests_list[i] + if victim.py_request_id in protected_ids: + continue + if victim.request_id in inflight_request_ids: + continue + if not self._is_started_request(victim): + continue + if not self.kv_cache_manager.is_request_active(victim.py_request_id): + continue + + released = self.kv_cache_manager.preempt_request(victim) + logger.debug( + f"[V2Scheduler] Preempting request {victim.py_request_id} " + f"(state={victim.state.name}), pages " + f"{'released' if released else 'pending connector saves'}" + ) + self._clear_request_runtime_state(victim) + if self.draft_kv_cache_manager is not None: + self.draft_kv_cache_manager.free_resources(victim) + if released: + # Rewrites the prompt to include what was generated and resets + # state to CONTEXT_INIT, so the request re-enters as an + # ordinary prefill. Deferred releases are paused by the + # executor once the connector reports the saves retired. + victim.pause(self.max_input_len) + evicted.append(victim) + preempted_ids.add(victim.py_request_id) + return released + + return False + + # Consecutive scheduling passes that reclaimed nothing before this counts + # as a deadlock. A stalled pass costs ~2ms, so it trips within seconds, + # while transient one-iteration deferrals (multimodal chunk alignment, + # PEFT budget, IndexMapper slots) clear long before. + _DEADLOCK_STALL_ITERS = 1000 + + def _detect_deadlock( + self, + active_requests: RequestList, + inflight_request_ids: set[int], + pending_ctx: RequestList, + preempted_ids: set[int], + made_progress: bool, + ) -> None: + """Fail loudly when no request can be scheduled or reclaimed. + + Without this the executor spins at full speed while scheduling + nothing, which looks healthy to the hang detector and to `/health` + while the job burns its wall clock. Context candidates count alongside + generation ones because a disaggregated prefill server has no + generation requests at all. + """ + if made_progress: + self._stalled_schedules = 0 + return + + num_gen_candidates = sum( + 1 + for r in active_requests + if r.is_generation_in_progress_state + and not r.is_generation_to_complete_state + and r.request_id not in inflight_request_ids + ) + num_ctx_candidates = sum( + 1 + for r in pending_ctx + if r.py_request_id not in preempted_ids and r.request_id not in inflight_request_ids + ) + if num_gen_candidates == 0 and num_ctx_candidates == 0: + # Legitimately idle: nothing to schedule. + self._stalled_schedules = 0 + return + + self._stalled_schedules += 1 + if self._stalled_schedules < self._DEADLOCK_STALL_ITERS: + return + + raise RuntimeError( + f"V2 scheduler deadlock: {num_gen_candidates} generation and " + f"{num_ctx_candidates} context request(s) active but none could " + f"be scheduled, suspended or preempted in " + f"{self._stalled_schedules} consecutive attempts. The KV cache " + f"pool is likely exhausted. Configure " + f"kv_cache_config.host_cache_size, increase " + f"kv_cache_config.max_tokens, or lower max_batch_size." + ) + def _is_evictable(self, req: LlmRequest) -> bool: """A started request whose KV cache is still active on GPU. diff --git a/tensorrt_llm/commands/mooncake.py b/tensorrt_llm/commands/mooncake.py new file mode 100644 index 000000000000..e4cfc6a67371 --- /dev/null +++ b/tensorrt_llm/commands/mooncake.py @@ -0,0 +1,296 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The two pieces of a Mooncake pool that outlive any one engine. + +A server that owns its pool needs neither: it describes the pool in +`kv_connector_config.mooncake_store` and `trtllm-serve` provisions it during +bringup. These commands exist for the pools it cannot own, such as one shared by +several engines, one that has to survive a restart, or one whose capacity comes +from nodes that run no connector. +""" + +import json +import os +import signal +import tempfile +import threading +import time +from typing import Optional + +import click + +from tensorrt_llm.logger import logger + + +def _until_signalled() -> threading.Event: + """An event that SIGINT and SIGTERM set. + + Both commands hold a resource, a child process or a mounted segment, whose + release is in a `finally`. Default SIGTERM handling would skip it, leaving + the master unreaped or the pool advertising memory that has gone. + """ + stopping = threading.Event() + + def stop(signum, _frame): + logger.info(f"mooncake-store: signal {signum} received, shutting down") + stopping.set() + + for received in (signal.SIGINT, signal.SIGTERM): + signal.signal(received, stop) + return stopping + + +@click.command("mooncake_master") +@click.option( + "--rpc_port", + type=int, + default=50051, + show_default=True, + help="Port the store clients reach the master on.", +) +@click.option( + "--metrics_port", + type=int, + default=9004, + show_default=True, + help="Prometheus port. Pool occupancy and eviction are read " + "from here or from the master's log.", +) +@click.option( + "--eviction_ratio", + type=float, + default=0.05, + show_default=True, + help="Fraction of the pool freed per eviction pass.", +) +@click.option( + "--address_file", + type=str, + default=None, + help="File to publish 'host:port' to once the master answers. " + "Workers name it as master_server_address: file://, which " + "is how they reach a master whose host the scheduler chose. " + "Removed on exit so a stale address is never dialed.", +) +@click.option( + "--run_dir", + type=str, + default=None, + help="Where to keep the master's log. Defaults to " + "$TRTLLM_MOONCAKE_RUN_DIR, else a temporary directory.", +) +@click.option( + "--heartbeat_seconds", + type=int, + default=300, + show_default=True, + help="Interval between liveness lines. 0 disables them.", +) +def mooncake_master( + rpc_port: int, + metrics_port: int, + eviction_ratio: float, + address_file: Optional[str], + run_dir: Optional[str], + heartbeat_seconds: int, +): + """Run a mooncake_master for as long as this command runs. + + A single server with a pool of its own should set + `mooncake_store.launch_master` instead. + """ + # Imported lazily so other subcommands and --help do not pay for the + # connector package. + from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store import running_master + from tensorrt_llm.llmapi.llm_args import MooncakeStoreConfig + + pool = MooncakeStoreConfig( + launch_master=True, + master_port=rpc_port, + master_metrics_port=metrics_port, + master_eviction_ratio=eviction_ratio, + ) + run_dir = ( + run_dir + or os.getenv("TRTLLM_MOONCAKE_RUN_DIR") + or tempfile.mkdtemp(prefix="trtllm-mooncake-master-") + ) + + stopping = _until_signalled() + with running_master(pool, run_dir, address_file=address_file) as master: + logger.info( + f"mooncake-store: this master owns the pool until this command " + f"stops; address {master.address}, log {master.log_path}, metrics " + f"http://{master.address.rsplit(':', 1)[0]}:{metrics_port}/metrics" + ) + started = time.monotonic() + announced = started + while not stopping.is_set(): + if (code := master.process.poll()) is not None: + # The pool is gone once the master dies, and every client is + # about to start failing. + raise click.ClickException( + f"mooncake_master exited with code {code}. See {master.log_path}" + ) + stopping.wait(1.0) + now = time.monotonic() + # Distinguishes a dead master from a dead fabric once clients + # start failing. + if heartbeat_seconds > 0 and now - announced >= heartbeat_seconds: + announced = now + logger.info( + f"mooncake-store: master at {master.address} alive after " + f"{(now - started) / 60:.0f}m" + ) + + +@click.command("mooncake_donor") +@click.option( + "--master_server_address", + type=str, + default=None, + help="Master to join, as host:port or file:// naming a " + "file that holds one. Defaults to the master_server_address in " + "--config.", +) +@click.option( + "--segment_size", + type=str, + default="32GiB", + show_default=True, + help="Host memory to contribute from this node. Deliberately " + "separate from a config's global_segment_size, which is sized " + "for an engine worker rather than a node lending what it can " + "spare.", +) +@click.option( + "--config", + type=str, + default=None, + help="Mooncake JSON config describing the pool, for the " + "settings not given here. Defaults to $MOONCAKE_CONFIG_PATH.", +) +@click.option( + "--protocol", + type=str, + default=None, + help="Transport, 'rdma' or 'tcp'. Defaults to --config's, else rdma.", +) +@click.option( + "--device_name", + type=str, + default=None, + help="RDMA device, from ibv_devinfo. Defaults to --config's.", +) +@click.option( + "--metadata_server", + type=str, + default=None, + help="Mooncake metadata service. Defaults to --config's, else P2PHANDSHAKE.", +) +@click.option( + "--ready_file", + type=str, + default=None, + help="File to create once the segment is mounted, for launchers " + "that must not let prefill start writing before the pool has " + "this capacity.", +) +@click.option( + "--heartbeat_seconds", + type=int, + default=300, + show_default=True, + help="Interval between liveness lines. 0 disables them.", +) +def mooncake_donor( + master_server_address: Optional[str], + segment_size: str, + config: Optional[str], + protocol: Optional[str], + device_name: Optional[str], + metadata_server: Optional[str], + ready_file: Optional[str], + heartbeat_seconds: int, +): + """Lend this node's host memory to a Mooncake pool, for as long as it runs. + + Running this on the generation nodes puts their memory into the pool while + leaving those engines connector-free. + """ + from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store import ( + DEFAULT_DONOR_LOCAL_BUFFER_SIZE, + donate_segment, + master_timeout, + parse_size, + resolve_master_address, + wait_for_master, + ) + from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.config import ( + CONFIG_PATH_ENV, + DEFAULT_METADATA_SERVER, + ) + + raw = {} + config = config or os.getenv(CONFIG_PATH_ENV) + if config: + with open(config) as handle: + raw = json.load(handle) + + master = master_server_address or raw.get("master_server_address", "") + if not master: + raise click.UsageError( + "No master to join. Pass --master_server_address, or a --config " + f"naming one (or set {CONFIG_PATH_ENV})." + ) + + donating = parse_size(segment_size) + resolved = resolve_master_address(master, master_timeout()) + # Before setup, so an absent master is reported as such. + wait_for_master(resolved) + + stopping = _until_signalled() + with donate_segment( + resolved, + donating, + protocol=protocol or raw.get("protocol", "rdma"), + device_name=device_name or raw.get("device_name", "") or "", + metadata_server=(metadata_server or raw.get("metadata_server") or DEFAULT_METADATA_SERVER), + local_buffer_size=parse_size( + raw.get("local_buffer_size_donor", DEFAULT_DONOR_LOCAL_BUFFER_SIZE) + ), + ) as host: + if ready_file: + with open(ready_file, "w") as handle: + handle.write(f"{host} {donating}\n") + logger.info( + f"mooncake-store: announced this segment in " + f"{ready_file}, so a launcher waiting on the pool's " + "capacity can proceed" + ) + + # Idle by design: a put or get here would make this node a traffic + # client, which is what donation exists to avoid. + started = time.monotonic() + while not stopping.is_set(): + if heartbeat_seconds <= 0: + stopping.wait() + continue + if not stopping.wait(heartbeat_seconds): + logger.info( + f"mooncake-store: {host} still lending " + f"{donating / 1024**3:.1f}GiB to the pool at {master} " + f"after {(time.monotonic() - started) / 60:.0f}m" + ) diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index 2dfe1e91ffe3..d1f946ddcf9f 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -1,5 +1,6 @@ import asyncio import atexit +import contextlib import gc import importlib import inspect @@ -14,7 +15,7 @@ import time import uuid from pathlib import Path -from typing import Any, Dict, NamedTuple, Optional, Sequence, Set +from typing import Any, Dict, Iterator, NamedTuple, Optional, Sequence, Set import click import torch @@ -27,6 +28,7 @@ from tensorrt_llm import MultimodalEncoder from tensorrt_llm._utils import mpi_rank, set_prometheus_multiproc_dir from tensorrt_llm.commands._serve_stability import stability_option +from tensorrt_llm.commands.mooncake import mooncake_donor, mooncake_master from tensorrt_llm.commands.utils import (collect_explicit_cli_keys, get_is_diffusion_only_model) from tensorrt_llm.executor.utils import MAX_NUM_FRONTENDS, LlmLauncherEnvs @@ -38,7 +40,9 @@ parse_disagg_config_file, parse_metadata_server_config_file, validate_config_bool) -from tensorrt_llm.llmapi.llm_args import MultimodalConfig, TorchLlmArgs +from tensorrt_llm.llmapi.llm_args import (KvCacheConnectorConfig, + MooncakeDonationConfig, + MultimodalConfig, TorchLlmArgs) from tensorrt_llm.llmapi.llm_utils import update_llm_args_with_extra_dict from tensorrt_llm.llmapi.mpi_session import find_free_ipc_addr, split_mpi_env from tensorrt_llm.llmapi.reasoning_parser import (ReasoningParserFactory, @@ -514,6 +518,49 @@ def _terminate_attached_frontends(children: list) -> None: child.kill() +@contextlib.contextmanager +def _provision_kv_cache_pool(llm_args: dict, + owns_engine: bool = True) -> Iterator[None]: + """Bring up the shared cache this server needs or feeds, for its lifetime. + + A connector backed by a cluster-wide pool needs that pool reachable before + any rank opens a handle, and the ranks are spawned by the LLM constructor, + so this wraps the construction. A deployment that provisions the pool + externally is detected and left alone. + + A server may also lend the pool host memory without using it, which is how + a pool spans nodes whose engines have no connector. That segment has to be + mounted before traffic arrives and held for as long as the pages placed in + it are expected to be there. + + Only the process that owns the engine does either. An attached frontend + re-execs this command line but shares the launcher's executor, so it would + otherwise stand up a second pool and lend a second segment. + """ + if not owns_engine: + yield + return + + from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store import ( + maybe_donate_segment, maybe_provision_pool) + + connector_config = llm_args.get("kv_connector_config") + if isinstance(connector_config, dict): + # A YAML config section arrives unvalidated, and the pool has to be + # described before the LLM constructor would coerce it. Hand the + # validated model on so it is not parsed twice. + connector_config = KvCacheConnectorConfig(**connector_config) + llm_args["kv_connector_config"] = connector_config + + donation = llm_args.get("mooncake_donation") + if isinstance(donation, dict): + donation = MooncakeDonationConfig(**donation) + llm_args["mooncake_donation"] = donation + + with maybe_provision_pool(connector_config), maybe_donate_segment(donation): + yield + + def launch_server( host: str, port: int, @@ -570,53 +617,55 @@ def launch_server( raise RuntimeError(f"Failed to bind socket to {host}:{port}: {e}. " f"Port holder(s): {holder}") - if backend == 'pytorch': - llm_args.pop("build_config", None) - llm = PyTorchLLM(**llm_args) - elif backend == '_autodeploy': - from tensorrt_llm._torch.auto_deploy import LLM as AutoDeployLLM - - # AutoDeploy does not support build_config - llm_args.pop("build_config", None) - llm = AutoDeployLLM(**llm_args) - else: - raise click.BadParameter( - f"{backend} is not a known backend, check help for available options.", - param_hint="backend") - - # The finally below is the cleanup boundary for the attached - # frontends: it must cover everything from their spawn through - # server construction, middleware registration, and runtime, or a - # failure in between leaks the child processes. - frontend_children = [] - try: - if multi_frontend.is_launcher: - frontend_children = _spawn_attached_frontends( - llm, multi_frontend.num_frontends) - - server = OpenAIServer( - generator=llm, - model=model, - tool_parser=tool_parser, - server_role=server_role, - metadata_server_cfg=metadata_server_cfg, - disagg_cluster_config=disagg_cluster_config, - multimodal_server_config=multimodal_server_config, - chat_template=chat_template, - allow_request_chat_template=allow_request_chat_template, - input_processor_workers=num_input_processor_workers, - media_load_workers=num_media_load_workers) - _apply_fastapi_middlewares(server.app, middleware) - - # Optionally disable GC (default: not disabled) - if os.getenv("TRTLLM_SERVER_DISABLE_GC", "0") == "1": - gc.disable() - - _signal_frontend_ready(multi_frontend) - uvloop.run(server(host, port, sockets=[s])) - finally: - if frontend_children: - _terminate_attached_frontends(frontend_children) + with _provision_kv_cache_pool( + llm_args, owns_engine=not multi_frontend.is_attached_frontend): + if backend == 'pytorch': + llm_args.pop("build_config", None) + llm = PyTorchLLM(**llm_args) + elif backend == '_autodeploy': + from tensorrt_llm._torch.auto_deploy import LLM as AutoDeployLLM + + # AutoDeploy does not support build_config + llm_args.pop("build_config", None) + llm = AutoDeployLLM(**llm_args) + else: + raise click.BadParameter( + f"{backend} is not a known backend, check help for available options.", + param_hint="backend") + + # The finally below is the cleanup boundary for the attached + # frontends: it must cover everything from their spawn through + # server construction, middleware registration, and runtime, or a + # failure in between leaks the child processes. + frontend_children = [] + try: + if multi_frontend.is_launcher: + frontend_children = _spawn_attached_frontends( + llm, multi_frontend.num_frontends) + + server = OpenAIServer( + generator=llm, + model=model, + tool_parser=tool_parser, + server_role=server_role, + metadata_server_cfg=metadata_server_cfg, + disagg_cluster_config=disagg_cluster_config, + multimodal_server_config=multimodal_server_config, + chat_template=chat_template, + allow_request_chat_template=allow_request_chat_template, + input_processor_workers=num_input_processor_workers, + media_load_workers=num_media_load_workers) + _apply_fastapi_middlewares(server.app, middleware) + + # Optionally disable GC (default: not disabled) + if os.getenv("TRTLLM_SERVER_DISABLE_GC", "0") == "1": + gc.disable() + + _signal_frontend_ready(multi_frontend) + uvloop.run(server(host, port, sockets=[s])) + finally: + if frontend_children: + _terminate_attached_frontends(frontend_children) def launch_grpc_server(host: str, @@ -739,7 +788,8 @@ def signal_handler(): logger.info("Shutdown complete") - uvloop.run(serve_grpc_async()) + with _provision_kv_cache_pool(llm_args): + uvloop.run(serve_grpc_async()) def launch_mm_encoder_server( @@ -2494,7 +2544,11 @@ def resolve_command(self, ctx, args): "disaggregated": disaggregated, "disaggregated_mpi_worker": disaggregated_mpi_worker, "mm_embedding_serve": serve_encoder, - "embeddings": serve_embedding + "embeddings": serve_embedding, + # The parts of a Mooncake pool that cannot belong to a server, for + # deployments where a pool outlives or spans them. + "mooncake_master": mooncake_master, + "mooncake_donor": mooncake_donor, }) if __name__ == "__main__": diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 7b5d014d3994..6d22c1262b26 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -1945,6 +1945,166 @@ def num_capture_layers(self) -> int: return 0 +class MooncakeStoreConfig(StrictBaseModel): + """The Mooncake store pool the `mooncake-store` connector should join. + + Describes the pool: which master owns it, how workers reach it, and how + much memory each contributes. A worker's own relationship to the pool, such + as its read/write role and key namespace, stays in the + `TRTLLM_MOONCAKE_STORE_*` environment variables, since that is per process + while this is per deployment. + + Setting this makes `trtllm-serve` render the Mooncake client config and + export `MOONCAKE_CONFIG_PATH` itself. An inherited `MOONCAKE_CONFIG_PATH` + still wins, so an externally managed pool stays reachable. + """ + master_server_address: Optional[str] = Field( + None, + description="Address of an already-running mooncake_master, as " + "host:port or file:// naming a file that holds one. The file is " + "how to reach a master whose host a scheduler chose, since " + "'trtllm-serve mooncake_master --address_file' publishes it there " + "once it answers. Mutually exclusive with launch_master.") + launch_master: bool = Field( + False, + telemetry=False, + description="Start a mooncake_master in this server's process group " + "and use it. The pool then dies with the server, so this is only " + "correct for a single engine: several engines sharing a pool, or a " + "pool that must outlive a restart, need master_server_address.") + master_address_file: Optional[str] = Field( + None, + telemetry=False, + description="Where a master started by launch_master should publish " + "its host:port, for donors and other servers to read back as " + "file://. One is always written to the run directory; set this " + "to put a second copy somewhere the rest of the deployment already " + "names, such as a shared filesystem. Removed when the master stops.") + master_port: int = Field( + 50051, + telemetry=False, + description="RPC port for a master started by launch_master.") + master_metrics_port: int = Field( + 9004, + telemetry=False, + description="Prometheus port for a master started by launch_master.") + master_eviction_ratio: float = Field( + 0.05, + telemetry=False, + description="Fraction of the pool a master started by launch_master " + "frees per eviction pass.") + metadata_server: str = Field( + "P2PHANDSHAKE", + description="Mooncake metadata service. P2PHANDSHAKE keeps a separate " + "metadata process out of the deployment.") + protocol: str = Field( + "rdma", + description="Transport for page traffic: 'rdma' or 'tcp'. " + "TCP is for bring-up only; it invalidates performance conclusions.") + device_name: str = Field( + "", + description="RDMA device to transfer over, from ibv_devinfo. Empty " + "with protocol 'tcp'.") + global_segment_size: Union[int, str] = Field( + "16GiB", + description="Host memory each worker process contributes to the pool. " + "Pool capacity is this times the number of processes that open a " + "store handle, so a prefill-only connector gives a prefill-only pool.") + local_buffer_size: Union[int, str] = Field( + "1GiB", + description="Per-process Mooncake transfer buffer, not pool capacity.") + transfer_batch_size: int = Field(64, + telemetry=False, + description="Page keys per store call.") + cache_prefix: Optional[str] = Field( + None, + description="Key namespace for the pool. Bump it after any change to " + "page layout or contents. Defaults to 'trtllm'.") + stage_through_host: bool = Field( + False, + telemetry=False, + description="Copy pages through a pinned host buffer instead of " + "registering the KV pools with Mooncake. Needed where the HCA cannot " + "pin GPU pages (no GPUDirect RDMA); costs a copy each way.") + staging_buffer_bytes: Optional[Union[int, str]] = Field( + None, + telemetry=False, + description="Size of the buffer stage_through_host copies through, " + "per process. Pages move transfer_batch_size at a time, and a buffer " + "that cannot hold that many reduces the batch instead of failing, so " + "undersizing it costs throughput quietly. Defaults to the connector's " + "own 512MiB.") + + @model_validator(mode="after") + def _require_exactly_one_master(self) -> "MooncakeStoreConfig": + if self.launch_master and self.master_server_address: + raise ValueError( + "mooncake_store: set either launch_master or " + "master_server_address, not both. launch_master starts a " + "master here; master_server_address joins an existing pool.") + if not self.launch_master and not self.master_server_address: + raise ValueError( + "mooncake_store: needs a master. Set master_server_address to " + "join an existing pool, or launch_master: true to start one " + "for this server alone.") + if self.master_address_file and not self.launch_master: + raise ValueError( + "mooncake_store: master_address_file publishes the address of " + "a master this server starts, so it needs launch_master: " + "true. To read an address a master elsewhere published, set " + "master_server_address: file://.") + return self + + +class MooncakeDonationConfig(StrictBaseModel): + """Host memory this server lends to a Mooncake pool it does not use. + + Pool capacity comes only from processes that open a store handle, which in + a disaggregated deployment is the context servers alone. Setting this on + the generation servers puts their memory into the same pool, so prefill + writes blocks that land on decode-side DRAM while the generation engine + stays free of any connector and keeps its cache transceiver for the + handoff. + + Capacity is kept separate from `kv_connector_config` because attaching a + connector would also start this server reading and writing the store. + + The memory is charged to this process, so size it together with + `kv_cache_config.host_cache_size`. + + Fields opt out of telemetry because they describe one site's pool rather + than which features are in use. + """ + master_server_address: str = Field( + ..., + telemetry=False, + description="Master of the pool to lend memory to, as host:port or " + "file:// naming a file that holds one. The file is what a " + "context server's launch_master publishes, so the generation servers " + "can name a path instead of a host chosen by a scheduler, and they " + "wait for the master rather than having to start after it.") + segment_size: Union[int, str] = Field( + "32GiB", + telemetry=False, + description="Host memory this server contributes. Charged once per " + "server process, not per rank, so a node running several servers " + "contributes this much for each of them.") + protocol: str = Field( + "rdma", + telemetry=False, + description="Transport the pool's traffic reaches this memory over: " + "'rdma' or 'tcp'. Must match the pool's.") + device_name: str = Field( + "", + telemetry=False, + description="RDMA device to serve the segment over, from ibv_devinfo. " + "Empty with protocol 'tcp'.") + metadata_server: str = Field( + "P2PHANDSHAKE", + telemetry=False, + description="Mooncake metadata service. Must match the pool's.") + + class KvCacheConnectorConfig(StrictBaseModel): """Configuration for the KV Cache Connector. @@ -1959,7 +2119,8 @@ class KvCacheConnectorConfig(StrictBaseModel): description="Named connector preset (e.g. 'lmcache'). " "When set, connector_module/scheduler_class/worker_class are " "auto-populated from the preset registry.", - telemetry=TelemetryField.categorical('lmcache', 'lmcache-mp', 'kvbm')) + telemetry=TelemetryField.categorical('lmcache', 'lmcache-mp', 'kvbm', + 'mooncake-store')) connector_module: Optional[str] = Field( None, description= @@ -1974,11 +2135,16 @@ class KvCacheConnectorConfig(StrictBaseModel): description="URL for an external connector server " "(e.g. 'tcp://localhost:5555'). Connectors that run in " "multi-process mode use this to reach the cache server.") + mooncake_store: Optional[MooncakeStoreConfig] = Field( + None, + description="Pool topology for the 'mooncake-store' connector. When " + "set, trtllm-serve provisions the pool during bringup instead of " + "requiring MOONCAKE_CONFIG_PATH from an external script.") @model_validator(mode="after") def _resolve_preset(self) -> "KvCacheConnectorConfig": - from tensorrt_llm._torch.pyexecutor.connectors.registry import \ - CONNECTOR_REGISTRY + from tensorrt_llm._torch.pyexecutor.connectors.registry import ( + CONNECTOR_REGISTRY, uses_connector) if self.connector is not None: preset = CONNECTOR_REGISTRY.get(self.connector) if preset is None: @@ -1996,6 +2162,12 @@ def _resolve_preset(self) -> "KvCacheConnectorConfig": raise ValueError("connector_scheduler_class is required") if self.connector_worker_class is None: raise ValueError("connector_worker_class is required") + if self.mooncake_store is not None and not uses_connector( + self, "mooncake-store"): + raise ValueError( + "mooncake_store describes a Mooncake pool, but this config " + f"resolves to connector_module={self.connector_module!r}. " + "Set connector: mooncake-store, or drop mooncake_store.") return self @@ -5219,6 +5391,17 @@ def validate_encoder_runtime_sizes(cls, v: Optional[int]) -> Optional[int]: status="prototype", ) + mooncake_donation: Optional[MooncakeDonationConfig] = Field( + default=None, + description="Host memory to lend to a Mooncake pool this server does " + "not otherwise use. Separate from kv_connector_config because it adds " + "capacity without attaching a connector, which is what lets a " + "generation server hold pages for a pool only prefill reads and " + "writes. Honored by trtllm-serve, which holds the segment for the " + "server's lifetime.", + status="prototype", + ) + mm_encoder_only: bool = Field( default=False, description= diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index cc9b5a7ff7d6..18b784e62495 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -798,7 +798,8 @@ "allowed_values": [ "lmcache", "lmcache-mp", - "kvbm" + "kvbm", + "mooncake-store" ], "annotation": "Optional[str]", "converter": "allowlist", @@ -1408,6 +1409,13 @@ "kind": "value", "path": "sparse_attention_config.enable_heuristic_topk" }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "sparse_attention_config.fuse_qkv_index_projection" + }, { "allowed_values": [ "triton", diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index 083895fe6d59..66de76c0e3d4 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -41,6 +41,9 @@ l0_a10: - unittest/_torch/executor/test_kv_cache_compression_manager.py - unittest/_torch/executor/test_kv_cache_v2_capacity_only.py - unittest/_torch/executor/test_kv_cache_layout.py + - unittest/_torch/executor/test_mooncake_store_connector.py + - unittest/_torch/executor/test_mooncake_store_donor.py + - unittest/_torch/executor/test_mooncake_store_master.py - unittest/_torch/executor/test_error_classification.py - unittest/_torch/modules/dwdp/test_dwdp_fixup_moe_backends.py - unittest/_torch/modules/dwdp/test_dwdp_manager.py diff --git a/tests/integration/test_lists/test-db/l0_b200_m3.yml b/tests/integration/test_lists/test-db/l0_b200_m3.yml index 8e5c9eacbfde..07e1d2e95d21 100644 --- a/tests/integration/test_lists/test-db/l0_b200_m3.yml +++ b/tests/integration/test_lists/test-db/l0_b200_m3.yml @@ -18,6 +18,10 @@ l0_b200_m3: - unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py - unittest/_torch/attention/sparse/test_minimax_m3_msa_selector.py - unittest/_torch/attention/sparse/test_minimax_m3_sparse_attn_decode.py + - unittest/_torch/executor/test_kv_cache_v2_scheduler.py + - unittest/_torch/executor/test_mooncake_store_connector.py + - unittest/_torch/executor/test_mooncake_store_donor.py + - unittest/_torch/executor/test_mooncake_store_master.py - unittest/_torch/models/test_minimax_m3.py - unittest/_torch/models/checkpoints/hf/test_minimaxm3_weight_mapper.py - unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_horizontal_producer.py diff --git a/tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py b/tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py index 0a7c67e9e702..2c24d0765f3c 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py +++ b/tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py @@ -53,6 +53,7 @@ def make_gen_request( req.lora_task_id = lora_task_id req.is_context_init_state = False req.is_generation_in_progress_state = True + req.is_generation_to_complete_state = False req.is_first_context_chunk = is_first_context_chunk req.py_encoder_output_ready_event = None return req @@ -160,6 +161,9 @@ def make_kv_cache_manager( resize_context_fn=None, prepare_disagg_gen_init_fn=None, try_allocate_generation_fn=None, + has_cache_tier_below_gpu=True, + preempt_request_fn=None, + has_pending_preemption=False, ): mgr = Mock() mgr.tokens_per_block = tokens_per_block @@ -168,8 +172,18 @@ def make_kv_cache_manager( mgr.resize_context.side_effect = resize_context_fn or (lambda req, n: True) mgr.prepare_disagg_gen_init.side_effect = prepare_disagg_gen_init_fn or (lambda req: True) mgr.try_allocate_generation.side_effect = try_allocate_generation_fn or (lambda req: True) - mgr.suspend_request.return_value = None + + def _suspend(req): + # Mirrors KVCacheManagerV2.suspend_request: the cache stops being + # active on GPU, so the request is no longer an eviction victim. + mgr.kv_cache_map[req.py_request_id].is_active = False + + mgr.suspend_request.side_effect = _suspend mgr.is_request_active.side_effect = lambda req_id: mgr.kv_cache_map[req_id].is_active + # The default here has a cache tier below GPU, which leaves preemption off. + mgr.has_cache_tier_below_gpu = has_cache_tier_below_gpu + mgr.has_pending_preemption.return_value = has_pending_preemption + mgr.preempt_request.side_effect = preempt_request_fn or (lambda req: True) return mgr @@ -191,6 +205,7 @@ def make_scheduler( no_schedule_after_state: LlmRequestState | None = None, cross_kv_cache_manager: Mock | None = None, enable_prefix_aware_scheduling: bool = True, + max_input_len: int | None = None, ) -> object: """Create KVCacheV2Scheduler, patching isinstance check for mock mgr.""" from tensorrt_llm._torch.pyexecutor.scheduler.scheduler_v2 import KVCacheV2Scheduler @@ -206,6 +221,8 @@ def make_scheduler( kwargs["no_schedule_after_state"] = no_schedule_after_state if cross_kv_cache_manager is not None: kwargs["cross_kv_cache_manager"] = cross_kv_cache_manager + if max_input_len is not None: + kwargs["max_input_len"] = max_input_len return KVCacheV2Scheduler( max_batch_size=max_batch_size, max_num_tokens=max_num_tokens, @@ -827,6 +844,267 @@ def test_self_eviction_no_started_in_range(self): assert len(out.context_requests) == 0 +# =========================================================================== +# Preemption (context side, no cache tier below GPU) +# =========================================================================== + + +def _out_of_pages_for(request_id): + """resize_context that only fails for *request_id*.""" + return lambda req, n: req.py_request_id != request_id + + +class TestContextPreemption: + """Releasing a started request's pages when suspension cannot help. + + With GPU as the last cache level a suspended page stays HELD and + unevictable, so suspension frees nothing. These tests cover the fallback + that gives the pages up instead. + """ + + def test_out_of_pages_preempts_started_request(self): + mgr = make_kv_cache_manager( + resize_context_fn=_out_of_pages_for(0), + has_cache_tier_below_gpu=False, + ) + sched = make_scheduler(mgr, max_num_tokens=1000) + victim = make_ctx_request(99, 100, is_first_context_chunk=False) + reqs = [make_ctx_request(0, 100), victim] + + out = sched.schedule_request(reqs, set()) + + mgr.preempt_request.assert_called_once_with(victim) + assert ids(out.paused_requests) == [99] + # Deferred to the next iteration: a failed resize leaves the first + # chunk suspended, so the retry has to go back through + # prepare_context. + assert len(out.context_requests) == 0 + + def test_released_victim_is_reset_to_context_state(self): + mgr = make_kv_cache_manager( + resize_context_fn=_out_of_pages_for(0), + has_cache_tier_below_gpu=False, + ) + sched = make_scheduler(mgr, max_num_tokens=1000, max_input_len=4096) + victim = make_ctx_request(99, 100, is_first_context_chunk=False) + victim.py_batch_idx = 7 + + sched.schedule_request([make_ctx_request(0, 100), victim], set()) + + victim.pause.assert_called_once_with(4096) + assert victim.py_batch_idx is None + + def test_deferred_release_leaves_pause_to_the_executor(self): + """A connector still reading these pages defers the release. + + Freeing them now would let a later request overwrite bytes + mid-transfer, so the executor pauses the request only once the + connector reports the saves retired. + """ + mgr = make_kv_cache_manager( + resize_context_fn=_out_of_pages_for(0), + has_cache_tier_below_gpu=False, + preempt_request_fn=lambda req: False, + ) + sched = make_scheduler(mgr, max_num_tokens=1000) + victim = make_ctx_request(99, 100, is_first_context_chunk=False) + + out = sched.schedule_request([make_ctx_request(0, 100), victim], set()) + + mgr.preempt_request.assert_called_once_with(victim) + victim.pause.assert_not_called() + assert ids(out.paused_requests) == [99] + + def test_preempted_victim_not_scheduled_in_the_same_pass(self): + """Re-admitting the victim would spend the pages it just released.""" + mgr = make_kv_cache_manager( + resize_context_fn=_out_of_pages_for(0), + has_cache_tier_below_gpu=False, + ) + sched = make_scheduler(mgr, max_num_tokens=1000) + victim = make_ctx_request(99, 100, is_first_context_chunk=False) + + out = sched.schedule_request([make_ctx_request(0, 100), victim], set()) + + assert ids(out.context_requests) == [] + + def test_skipped_when_a_cache_tier_exists_below_gpu(self): + mgr = make_kv_cache_manager( + resize_context_fn=_out_of_pages_for(0), + has_cache_tier_below_gpu=True, + ) + sched = make_scheduler(mgr, max_num_tokens=1000) + victim = make_ctx_request(99, 100, is_first_context_chunk=False) + + out = sched.schedule_request([make_ctx_request(0, 100), victim], set()) + + mgr.preempt_request.assert_not_called() + # Suspension is cheaper and keeps the pages, so the request is + # simply skipped. + assert ids(out.context_requests) == [99] + + def test_one_victim_at_a_time_while_a_release_is_draining(self): + mgr = make_kv_cache_manager( + resize_context_fn=_out_of_pages_for(0), + has_cache_tier_below_gpu=False, + has_pending_preemption=True, + ) + sched = make_scheduler(mgr, max_num_tokens=1000) + victim = make_ctx_request(99, 100, is_first_context_chunk=False) + + sched.schedule_request([make_ctx_request(0, 100), victim], set()) + + mgr.preempt_request.assert_not_called() + + def test_never_preempts_a_scheduled_request(self): + mgr = make_kv_cache_manager( + resize_context_fn=_out_of_pages_for(1), + has_cache_tier_below_gpu=False, + ) + sched = make_scheduler(mgr, max_num_tokens=1000) + # gen0 is scheduled in phase 1; ctx1 then runs out of pages and must + # not take the pages out from under it. + reqs = [make_gen_request(0), make_ctx_request(1, 100)] + + out = sched.schedule_request(reqs, set()) + + assert ids(out.generation_requests) == [0] + mgr.preempt_request.assert_not_called() + + def test_never_preempts_itself(self): + mgr = make_kv_cache_manager( + resize_context_fn=lambda req, n: False, + has_cache_tier_below_gpu=False, + ) + sched = make_scheduler(mgr, max_num_tokens=1000) + req = make_ctx_request(0, 100, is_first_context_chunk=False) + + sched.schedule_request([req], set()) + + mgr.preempt_request.assert_not_called() + + def test_never_preempts_an_inflight_request(self): + mgr = make_kv_cache_manager( + resize_context_fn=_out_of_pages_for(0), + has_cache_tier_below_gpu=False, + ) + sched = make_scheduler(mgr, max_num_tokens=1000) + victim = make_ctx_request(99, 100, is_first_context_chunk=False) + + sched.schedule_request([make_ctx_request(0, 100), victim], {99}) + + mgr.preempt_request.assert_not_called() + + def test_never_preempts_a_first_chunk_or_suspended_request(self): + mgr = make_kv_cache_manager( + resize_context_fn=_out_of_pages_for(0), + has_cache_tier_below_gpu=False, + ) + sched = make_scheduler(mgr, max_num_tokens=1000) + # First chunk: holds no pages worth taking. + first_chunk = make_ctx_request(98, 100, is_first_context_chunk=True) + # Already suspended: preempting it frees nothing extra. + suspended = make_ctx_request(99, 100, is_first_context_chunk=False) + mgr.kv_cache_map[suspended.py_request_id].is_active = False + + sched.schedule_request([make_ctx_request(0, 100), first_chunk, suspended], set()) + + mgr.preempt_request.assert_not_called() + + def test_chunked_context_out_of_pages_preempts(self): + mgr = make_kv_cache_manager( + resize_context_fn=_out_of_pages_for(0), + has_cache_tier_below_gpu=False, + ) + sched = make_scheduler( + mgr, + max_num_tokens=1000, + ctx_chunk_config=(ContextChunkingPolicy.FIRST_COME_FIRST_SERVED, 64), + ) + victim = make_ctx_request(99, 100, is_first_context_chunk=False) + + out = sched.schedule_request([make_ctx_request(0, 500), victim], set()) + + mgr.preempt_request.assert_called_once_with(victim) + assert ids(out.paused_requests) == [99] + + +# =========================================================================== +# Deadlock detection +# =========================================================================== + + +class TestDeadlockDetection: + """The scheduler must fail loudly rather than spin scheduling nothing. + + A stalled pass costs a couple of milliseconds, so an undetected stall + burns a job's whole wall clock. + """ + + def test_raises_after_repeated_stalls_with_context_candidates(self): + """A prefill-only worker has no generation requests to count.""" + mgr = make_kv_cache_manager( + resize_context_fn=lambda req, n: False, + has_cache_tier_below_gpu=False, + ) + sched = make_scheduler(mgr, max_num_tokens=1000) + sched._DEADLOCK_STALL_ITERS = 3 + reqs = [make_ctx_request(0, 100, is_first_context_chunk=False)] + + for _ in range(2): + sched.schedule_request(reqs, set()) + with pytest.raises(RuntimeError, match="V2 scheduler deadlock"): + sched.schedule_request(reqs, set()) + + def test_raises_after_repeated_stalls_with_generation_candidates(self): + mgr = make_kv_cache_manager(try_allocate_generation_fn=lambda req: False) + sched = make_scheduler(mgr, max_num_tokens=100) + sched._DEADLOCK_STALL_ITERS = 3 + # Self-eviction suspends it on the first pass, which counts as + # progress. Afterwards it is inactive and nothing can be reclaimed. + reqs = [make_gen_request(0)] + + for _ in range(3): + sched.schedule_request(reqs, set()) + with pytest.raises(RuntimeError, match="V2 scheduler deadlock"): + sched.schedule_request(reqs, set()) + + def test_transient_stall_does_not_raise(self): + """One bad iteration is normal; the counter has to reset.""" + fail = [True] + + def resize_fn(req, n): + return not fail[0] + + mgr = make_kv_cache_manager(resize_context_fn=resize_fn, has_cache_tier_below_gpu=False) + sched = make_scheduler(mgr, max_num_tokens=1000) + sched._DEADLOCK_STALL_ITERS = 3 + reqs = [make_ctx_request(0, 100, is_first_context_chunk=False)] + + for _ in range(10): + sched.schedule_request(reqs, set()) + fail[0] = not fail[0] + + def test_idle_scheduler_never_raises(self): + mgr = make_kv_cache_manager() + sched = make_scheduler(mgr, max_num_tokens=1000) + sched._DEADLOCK_STALL_ITERS = 2 + + for _ in range(5): + out = sched.schedule_request([], set()) + assert len(out.context_requests) == 0 + + def test_all_candidates_inflight_never_raises(self): + """Requests in the PP pipeline are progressing, just not here.""" + mgr = make_kv_cache_manager(resize_context_fn=lambda req, n: False) + sched = make_scheduler(mgr, max_num_tokens=1000) + sched._DEADLOCK_STALL_ITERS = 2 + reqs = [make_ctx_request(0, 100)] + + for _ in range(5): + sched.schedule_request(reqs, {0}) + + # =========================================================================== # PEFT / LoRA # =========================================================================== diff --git a/tests/unittest/_torch/executor/test_mooncake_store_connector.py b/tests/unittest/_torch/executor/test_mooncake_store_connector.py new file mode 100644 index 000000000000..20396ed352da --- /dev/null +++ b/tests/unittest/_torch/executor/test_mooncake_store_connector.py @@ -0,0 +1,1043 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for the Mooncake store KV cache connector. + +Runs without a Mooncake installation and without a GPU: the store handle is +replaced by an in-process fake, and the KV cache layout is synthesized from +plain integers, which is all the addressing arithmetic needs. +""" + +import contextlib +import json +import time +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor.connectors.kv_cache_connector import ( + RequestData, + SchedulerOutput, +) +from tensorrt_llm._torch.pyexecutor.connectors.kv_cache_layout import ( + KvCacheBufferRef, + KvCacheLayerGroupLayout, + KvCacheLayout, + KvCacheRegion, +) +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store import staging as staging_module +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store import worker as worker_module +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.addressing import ( + PageAddressing, + merge_intervals, +) +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.config import ( + MooncakeStoreConnectorConfig, + StoreRole, +) +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.keys import ( + BlockHashChain, + KeyNamespace, +) +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.metadata import ( + PageTransfer, + RequestTransfers, +) +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.scheduler import ( + MooncakeStoreConnectorScheduler, +) +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.staging import plan_slot_geometry +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.validation import ( + validate_layout, + validate_llm_args, +) +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.worker import ( + MooncakeStoreConnectorWorker, +) +from tensorrt_llm._torch.pyexecutor.connectors.registry import uses_connector +from tensorrt_llm.llmapi.llm_args import KvCacheConnectorConfig +from tensorrt_llm.runtime.kv_cache_manager_v2 import BAD_PAGE_INDEX + +TOKENS_PER_BLOCK = 4 + + +# ---- fixtures and fakes ---- + + +class FakeStore: + """Records calls and remembers which keys exist, nothing more.""" + + def __init__(self): + self.objects = set() + self.registered = [] + self.put_calls = [] + self.get_calls = [] + self.exist_calls = [] + self.closed = False + self.fail_gets_for = set() + #: Workers built against this store. `make_worker` shuts each one down; + #: the fixture repeats it as a backstop for early failures. + self.workers = [] + + def register_buffer(self, address, size): + self.registered.append((address, size)) + return 0 + + def batch_is_exist(self, keys): + self.exist_calls.append(list(keys)) + return [1 if key in self.objects else 0 for key in keys] + + def batch_put_from_multi_buffers(self, keys, addresses, sizes, *_args, **_kwargs): + self.put_calls.append((list(keys), [list(a) for a in addresses], [list(s) for s in sizes])) + self.objects.update(keys) + return [sum(size) for size in sizes] + + def batch_get_into_multi_buffers(self, keys, addresses, sizes): + self.get_calls.append((list(keys), [list(a) for a in addresses], [list(s) for s in sizes])) + return [ + -1 if (key in self.fail_gets_for or key not in self.objects) else sum(size) + for key, size in zip(keys, sizes) + ] + + def close(self): + self.closed = True + + +def make_layout(*, num_groups=1, regions_per_group=1, num_slots=8, window_size=None): + """A layout whose regions are laid out back to back in a fake address space.""" + groups = [] + base = 0x1000 + for group_id in range(num_groups): + regions = [] + for region_id in range(regions_per_group): + size = 64 * (region_id + 1) + stride = size + regions.append( + KvCacheRegion( + base=base, + size=size, + stride=stride, + num_slots=num_slots, + buffers=(KvCacheBufferRef(layer_id=group_id, role="key"),), + ) + ) + base += stride * num_slots + groups.append( + KvCacheLayerGroupLayout( + layer_group_id=group_id, + layer_ids=(group_id,), + window_size=window_size, + regions=tuple(regions), + ) + ) + return KvCacheLayout(tokens_per_block=TOKENS_PER_BLOCK, groups=tuple(groups)) + + +@pytest.fixture +def store_config(tmp_path, monkeypatch): + path = tmp_path / "mooncake.json" + path.write_text( + json.dumps( + { + "metadata_server": "http://127.0.0.1:8080/metadata", + "master_server_address": "127.0.0.1:50051", + "protocol": "tcp", + "device_name": "", + "global_segment_size": "1GiB", + "local_buffer_size": "256MiB", + "model_key": "test-model", + } + ) + ) + monkeypatch.setenv("MOONCAKE_CONFIG_PATH", str(path)) + monkeypatch.delenv("TRTLLM_MOONCAKE_STORE_ROLE", raising=False) + monkeypatch.delenv("TRTLLM_MOONCAKE_STORE_PREFIX", raising=False) + monkeypatch.delenv("TRTLLM_MOONCAKE_STORE_MODEL_KEY", raising=False) + return path + + +def make_llm_args(): + return SimpleNamespace( + model="/models/test-model", + kv_cache_config=SimpleNamespace(tokens_per_block=TOKENS_PER_BLOCK), + tensor_parallel_size=1, + pipeline_parallel_size=1, + context_parallel_size=1, + sparse_attention_config=None, + ) + + +@pytest.fixture +def fake_store(monkeypatch): + """Replace the store handle, and tear down any worker a test builds.""" + store = FakeStore() + monkeypatch.setattr(worker_module, "_open_store", lambda _config: store) + yield store + for worker in store.workers: + worker.shutdown() + worker_module._LOCAL_WORKER = None + worker_module._LOCAL_WORKER_READY.clear() + + +@contextlib.contextmanager +def make_worker(fake_store, *, layout=None): + """Build a worker and shut it down before the test call phase ends. + + Registering a layout starts the background save thread, and + pytest-threadleak snapshots threads around the call phase only, so + fixture teardown would run too late to keep it quiet. + """ + worker = MooncakeStoreConnectorWorker(make_llm_args()) + fake_store.workers.append(worker) + if layout is not None: + worker.register_kv_cache_layout(layout) + try: + yield worker + finally: + worker.shutdown() + + +def make_request(request_id, tokens, cache_salt=None): + return SimpleNamespace( + request_id=request_id, + cache_salt=cache_salt, + get_tokens=lambda _beam=0, _tokens=tuple(tokens): list(_tokens), + ) + + +# ---- keys ---- + + +def test_hash_chain_is_deterministic_and_prefix_sensitive(): + tokens = list(range(3 * TOKENS_PER_BLOCK)) + first = list(BlockHashChain(TOKENS_PER_BLOCK).extend(tokens)) + second = list(BlockHashChain(TOKENS_PER_BLOCK).extend(tokens)) + assert first == second + + # Changing a token in block 0 must change every hash after it, which is what + # makes a key safe to share: a hit implies the whole prefix matched. + altered = list(tokens) + altered[0] += 1 + changed = list(BlockHashChain(TOKENS_PER_BLOCK).extend(altered)) + assert all(a != b for a, b in zip(first, changed)) + + +def test_hash_chain_ignores_partial_trailing_block(): + full = list(range(2 * TOKENS_PER_BLOCK)) + chain = BlockHashChain(TOKENS_PER_BLOCK) + assert len(chain.extend(full)) == 2 + assert len(chain.extend(full + [99])) == 2 + + +def test_hash_chain_extends_incrementally(): + tokens = list(range(4 * TOKENS_PER_BLOCK)) + incremental = BlockHashChain(TOKENS_PER_BLOCK) + for end in range(0, len(tokens) + 1, TOKENS_PER_BLOCK): + incremental.extend(tokens[:end]) + assert list(incremental.hashes) == list(BlockHashChain(TOKENS_PER_BLOCK).extend(tokens)) + + +def test_hash_chain_separates_cache_salts(): + tokens = list(range(TOKENS_PER_BLOCK)) + unsalted = BlockHashChain(TOKENS_PER_BLOCK).extend(tokens) + salted = BlockHashChain(TOKENS_PER_BLOCK, cache_salt="tenant-a").extend(tokens) + other = BlockHashChain(TOKENS_PER_BLOCK, cache_salt="tenant-b").extend(tokens) + assert unsalted[0] != salted[0] != other[0] + assert salted[0] != other[0] + + +def test_hash_chain_rejects_shrinking_token_list(): + chain = BlockHashChain(TOKENS_PER_BLOCK) + chain.extend(list(range(2 * TOKENS_PER_BLOCK))) + with pytest.raises(ValueError, match="shrank"): + chain.extend(list(range(TOKENS_PER_BLOCK))) + + +def test_key_namespace_separates_every_dimension(): + base = dict( + cache_prefix="trtllm", + model_key="m", + rank=0, + world_size=2, + layer_group_id=0, + tokens_per_block=32, + bytes_per_page=1024, + ) + block_hash = b"\x01" * 16 + reference = KeyNamespace(**base).key(block_hash) + for field, value in [ + ("cache_prefix", "other"), + ("model_key", "n"), + ("rank", 1), + ("world_size", 4), + ("layer_group_id", 1), + ("tokens_per_block", 64), + ("bytes_per_page", 2048), + ]: + assert KeyNamespace(**{**base, field: value}).key(block_hash) != reference + + +# ---- addressing ---- + + +@pytest.mark.parametrize( + "intervals,expected", + [ + ([], []), + ([(0, 10)], [(0, 10)]), + ([(0, 10), (10, 20)], [(0, 20)]), + ([(0, 10), (5, 20)], [(0, 20)]), + ([(0, 10), (20, 30)], [(0, 10), (20, 30)]), + ([(20, 30), (0, 10)], [(0, 10), (20, 30)]), + ([(0, 100), (10, 20)], [(0, 100)]), + ([(0, 0), (5, 10)], [(5, 10)]), + ], +) +def test_merge_intervals(intervals, expected): + assert merge_intervals(intervals) == expected + + +def test_page_addressing_resolves_every_region_of_a_page(): + layout = make_layout(regions_per_group=3, num_slots=4) + addressing = PageAddressing(layout) + regions = layout.groups[0].regions + + addresses, sizes = addressing.buffers(0, 2) + assert sizes == [region.size for region in regions] + assert addresses == [region.base + region.stride * 2 for region in regions] + assert addressing.bytes_per_page(0) == sum(region.size for region in regions) + + +def test_page_addressing_rejects_out_of_range_page(): + addressing = PageAddressing(make_layout(num_slots=4)) + with pytest.raises(IndexError): + addressing.buffers(0, 4) + with pytest.raises(IndexError): + addressing.buffers(0, -1) + + +def test_page_addressing_registration_covers_every_slot_once(): + layout = make_layout(num_groups=2, regions_per_group=2, num_slots=4) + ranges = PageAddressing(layout).registration_ranges() + + # Regions were laid out back to back, so the whole span merges into one. + all_regions = [region for group in layout.groups for region in group.regions] + lowest = min(region.base for region in all_regions) + highest = max( + region.base + region.stride * (region.num_slots - 1) + region.size for region in all_regions + ) + assert ranges == [(lowest, highest)] + + +def test_page_addressing_rejects_mixed_slot_counts(): + region_a = KvCacheRegion(base=0, size=8, stride=8, num_slots=4, buffers=()) + region_b = KvCacheRegion(base=64, size=8, stride=8, num_slots=8, buffers=()) + layout = KvCacheLayout( + tokens_per_block=TOKENS_PER_BLOCK, + groups=( + KvCacheLayerGroupLayout( + layer_group_id=0, + layer_ids=(0,), + window_size=None, + regions=(region_a, region_b), + ), + ), + ) + with pytest.raises(ValueError, match="slot counts"): + PageAddressing(layout) + + +# ---- config ---- + + +def test_config_reads_sizes_and_staging_from_the_json(store_config): + """Sizes arrive as unit strings, and staging is off until the JSON asks.""" + config = MooncakeStoreConnectorConfig.from_env() + assert config.global_segment_size == 1024**3 + assert config.local_buffer_size == 256 * 1024**2 + assert config.role is StoreRole.BOTH + assert config.resolve_model_key("/models/ignored") == "test-model" + assert config.stage_through_host is False + + raw = json.loads(store_config.read_text()) + raw["stage_through_host"] = True + raw["staging_buffer_bytes"] = "256MiB" + store_config.write_text(json.dumps(raw)) + + config = MooncakeStoreConnectorConfig.from_env() + assert config.stage_through_host is True + assert config.staging_buffer_bytes == 256 * 1024**2 + + +def test_config_role_comes_from_environment(store_config, monkeypatch): + monkeypatch.setenv("TRTLLM_MOONCAKE_STORE_ROLE", "producer") + config = MooncakeStoreConnectorConfig.from_env() + assert config.role is StoreRole.PRODUCER + assert config.role.saves and not config.role.loads + + monkeypatch.setenv("TRTLLM_MOONCAKE_STORE_ROLE", "consumer") + config = MooncakeStoreConnectorConfig.from_env() + assert config.role.loads and not config.role.saves + + monkeypatch.setenv("TRTLLM_MOONCAKE_STORE_ROLE", "nonsense") + with pytest.raises(ValueError, match="TRTLLM_MOONCAKE_STORE_ROLE"): + MooncakeStoreConnectorConfig.from_env() + + +def test_config_requires_the_env_var(monkeypatch): + monkeypatch.delenv("MOONCAKE_CONFIG_PATH", raising=False) + monkeypatch.delenv("TRTLLM_MOONCAKE_RUN_DIR", raising=False) + with pytest.raises(ValueError, match="MOONCAKE_CONFIG_PATH"): + MooncakeStoreConnectorConfig.from_env() + + +def test_config_falls_back_to_the_run_directory(tmp_path, monkeypatch): + # A rank an external launcher started was already running when its leader + # provisioned the pool, so it never inherited the exported path and reads + # the rendered config out of the shared run directory instead. + monkeypatch.delenv("MOONCAKE_CONFIG_PATH", raising=False) + monkeypatch.setenv("TRTLLM_MOONCAKE_RUN_DIR", str(tmp_path)) + (tmp_path / "mooncake.json").write_text( + json.dumps({"master_server_address": "10.0.0.1:50051", "global_segment_size": "8GiB"}) + ) + + config = MooncakeStoreConnectorConfig.from_env() + + assert config.master_server_address == "10.0.0.1:50051" + assert config.global_segment_size == 8 * 1024**3 + + +def test_config_run_directory_without_a_rendered_config_still_asks(tmp_path, monkeypatch): + # An empty run directory means no leader provisioned anything, which is a + # missing pool rather than a default one. + monkeypatch.delenv("MOONCAKE_CONFIG_PATH", raising=False) + monkeypatch.setenv("TRTLLM_MOONCAKE_RUN_DIR", str(tmp_path)) + with pytest.raises(ValueError, match="MOONCAKE_CONFIG_PATH"): + MooncakeStoreConnectorConfig.from_env() + + +def test_config_env_var_wins_over_the_run_directory(tmp_path, monkeypatch): + # An externally managed pool stays reachable, since the run directory is + # only consulted when nothing was passed in. + named = tmp_path / "external.json" + named.write_text(json.dumps({"master_server_address": "external:50051"})) + (tmp_path / "mooncake.json").write_text( + json.dumps({"master_server_address": "provisioned:50051"}) + ) + monkeypatch.setenv("MOONCAKE_CONFIG_PATH", str(named)) + monkeypatch.setenv("TRTLLM_MOONCAKE_RUN_DIR", str(tmp_path)) + + assert MooncakeStoreConnectorConfig.from_env().master_server_address == "external:50051" + + +@pytest.mark.parametrize("named", [{}, {"metadata_server": ""}], ids=["omitted", "empty"]) +def test_config_metadata_server_falls_back_to_the_handshake(tmp_path, monkeypatch, named): + # No metadata service means Mooncake's peer-to-peer handshake. An empty + # connstring is not one of the forms setup accepts, so leaving the field + # out of a hand-written config must not reach it. + path = tmp_path / "metadata.json" + path.write_text(json.dumps({"master_server_address": "127.0.0.1:50051", **named})) + monkeypatch.setenv("MOONCAKE_CONFIG_PATH", str(path)) + assert MooncakeStoreConnectorConfig.from_env().metadata_server == "P2PHANDSHAKE" + + +def test_config_model_key_defaults_to_basename(store_config, tmp_path, monkeypatch): + path = tmp_path / "no_model_key.json" + path.write_text(json.dumps({"master_server_address": "127.0.0.1:50051"})) + monkeypatch.setenv("MOONCAKE_CONFIG_PATH", str(path)) + config = MooncakeStoreConnectorConfig.from_env() + assert config.resolve_model_key("/models/MiniMax-M3/") == "MiniMax-M3" + + +# ---- validation ---- + + +@pytest.mark.parametrize( + "field,value,match", + [ + ("context_parallel_size", 2, "context parallelism"), + ("pipeline_parallel_size", 2, "pipeline parallelism"), + ], +) +def test_validate_llm_args_rejects_unsupported_parallelism(field, value, match): + args = make_llm_args() + setattr(args, field, value) + with pytest.raises(NotImplementedError, match=match): + validate_llm_args(args) + + +def test_validate_llm_args_rejects_m3_index_value_cache(): + args = make_llm_args() + args.sparse_attention_config = SimpleNamespace(sparse_disable_index_value=False) + with pytest.raises(NotImplementedError, match="sparse_disable_index_value"): + validate_llm_args(args) + + args.sparse_attention_config = SimpleNamespace(sparse_disable_index_value=True) + validate_llm_args(args) + + +def test_validate_layout_rejects_sliding_window(): + with pytest.raises(NotImplementedError, match="sliding-window"): + validate_layout(make_layout(window_size=1024)) + validate_layout(make_layout()) + + +# ---- connector identification ---- +# +# py_executor_creator turns partial reuse off for this connector and finds it +# through uses_connector. Failing to recognize the config would silently cost +# the reuse the store exists to provide. + + +@pytest.mark.parametrize( + "config, expected", + [ + (KvCacheConnectorConfig(connector="mooncake-store"), True), + ( + KvCacheConnectorConfig( + connector_module="tensorrt_llm._torch.pyexecutor.connectors.mooncake_store", + connector_scheduler_class="MooncakeStoreConnectorScheduler", + connector_worker_class="MooncakeStoreConnectorWorker", + ), + True, + ), + (KvCacheConnectorConfig(connector="kvbm"), False), + (None, False), + ], + ids=["preset", "hand_written_module", "another_connector", "no_connector"], +) +def test_uses_connector_recognizes_the_connector_however_it_is_spelled(config, expected): + assert uses_connector(config, "mooncake-store") is expected + + +def test_uses_connector_rejects_an_unknown_preset(): + config = KvCacheConnectorConfig(connector="mooncake-store") + with pytest.raises(ValueError, match="Unknown connector preset"): + uses_connector(config, "mooncake-stroe") + + +# ---- worker ---- + + +def test_worker_registers_every_pool_range(store_config, fake_store): + layout = make_layout(num_groups=2, regions_per_group=2) + with make_worker(fake_store, layout=layout) as worker: + assert fake_store.registered == [ + (start, end - start) for start, end in PageAddressing(layout).registration_ranges() + ] + assert worker.is_registered + + +def test_worker_rejects_v1_pool_registration(store_config, fake_store): + with make_worker(fake_store) as worker: + with pytest.raises(NotImplementedError, match="KVCacheManagerV2"): + worker.register_kv_caches(None) + + +def test_worker_prefix_hit_needs_every_layer_group(store_config, fake_store): + layout = make_layout(num_groups=2) + with make_worker(fake_store, layout=layout) as worker: + hashes = [bytes([index]) * 16 for index in range(3)] + + assert worker.count_prefix_hit(hashes) == 0 + + # Populate blocks 0 and 1 completely, and block 2 only partially. + for block in range(2): + for group_id in range(2): + fake_store.objects.add(worker._namespaces[group_id].key(hashes[block])) + fake_store.objects.add(worker._namespaces[0].key(hashes[2])) + + assert worker.count_prefix_hit(hashes) == 2 + + +def test_worker_prefix_hit_stops_at_the_first_gap(store_config, fake_store): + with make_worker(fake_store, layout=make_layout()) as worker: + hashes = [bytes([index]) * 16 for index in range(3)] + # With block 1 missing, block 2 is unusable even though it is present, + # because a prefix is replayed contiguously. + fake_store.objects.add(worker._namespaces[0].key(hashes[0])) + fake_store.objects.add(worker._namespaces[0].key(hashes[2])) + assert worker.count_prefix_hit(hashes) == 1 + + +def test_worker_load_raises_when_a_page_is_missing(store_config, fake_store): + with make_worker(fake_store, layout=make_layout()) as worker: + transfers = RequestTransfers(7, [PageTransfer(b"\x00" * 16, 0, 1)]) + worker.bind_connector_meta(SimpleNamespace(loads=[transfers], saves=[])) + with pytest.raises(RuntimeError, match="already"): + worker.start_load_kv(None) + + +def test_worker_load_addresses_the_requested_page(store_config, fake_store): + layout = make_layout(regions_per_group=2) + with make_worker(fake_store, layout=layout) as worker: + block_hash = b"\x00" * 16 + key = worker._namespaces[0].key(block_hash) + fake_store.objects.add(key) + + transfers = RequestTransfers(7, [PageTransfer(block_hash, 0, 3)]) + worker.bind_connector_meta(SimpleNamespace(loads=[transfers], saves=[])) + worker.start_load_kv(None) + + (keys, addresses, sizes) = fake_store.get_calls[0] + expected_addresses, expected_sizes = PageAddressing(layout).buffers(0, 3) + assert keys == [key] + assert addresses == [expected_addresses] + assert sizes == [expected_sizes] + + +def test_worker_save_skips_pages_already_in_the_store(store_config, fake_store): + with make_worker(fake_store, layout=make_layout()) as worker: + hashes = [bytes([index]) * 16 for index in range(2)] + fake_store.objects.add(worker._namespaces[0].key(hashes[0])) + + worker._put( + [ + RequestTransfers( + 1, + [PageTransfer(hashes[0], 0, 0), PageTransfer(hashes[1], 0, 1)], + ) + ] + ) + assert len(fake_store.put_calls) == 1 + assert fake_store.put_calls[0][0] == [worker._namespaces[0].key(hashes[1])] + + +def test_worker_reports_a_request_finished_once_its_saves_drain(store_config, fake_store): + with make_worker(fake_store, layout=make_layout()) as worker: + # One submission outstanding: the request is closed but must not be released. + worker._outstanding_saves[42] = 1 + assert worker.get_finished([42], []) == ([], []) + + worker._outstanding_saves.pop(42) + assert worker.get_finished([], []) == ([42], []) + # Reported once only. + assert worker.get_finished([], []) == ([], []) + + +def test_worker_reports_a_request_with_no_saves_immediately(store_config, fake_store): + with make_worker(fake_store, layout=make_layout()) as worker: + assert worker.get_finished([9], [5]) == ([9], [5]) + + +def test_worker_shutdown_closes_the_store(store_config, fake_store): + with make_worker(fake_store, layout=make_layout()) as worker: + worker.shutdown() + assert fake_store.closed + # Idempotent: a second call must not raise or reopen anything. + worker.shutdown() + + +# ---- host staging ---- + + +@contextlib.contextmanager +def make_staged_worker(fake_store, store_config, *, layout, budget=None): + """A worker configured to pass pages through pinned host slots.""" + raw = json.loads(store_config.read_text()) + raw["stage_through_host"] = True + if budget is not None: + raw["staging_buffer_bytes"] = budget + store_config.write_text(json.dumps(raw)) + with make_worker(fake_store, layout=layout) as worker: + yield worker + + +@pytest.fixture +def staged_copies(monkeypatch): + """Record the copies staging would issue, instead of running them.""" + copies = [] + monkeypatch.setattr( + staging_module, + "_memcpy_async", + lambda dst, src, size, kind, stream: copies.append((int(dst), int(src), int(size))), + ) + # Imported into the worker by name, so replace the worker's binding. + monkeypatch.setattr(worker_module, "_sync_stream", lambda _stream: None) + return copies + + +@pytest.fixture +def fake_cuda(monkeypatch): + """Present a CUDA device on a host that has none, recording set_device calls. + + Only safe for paths that do not allocate or launch, and exists to exercise + the device bookkeeping around the save thread. + """ + recorded = [] + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "current_device", lambda: 3) + monkeypatch.setattr(torch.cuda, "set_device", lambda index: recorded.append(index)) + return recorded + + +@pytest.mark.parametrize( + "page_bytes,batch,budget,expected_slots", + [ + (1024, 64, 1 << 20, 64), # budget is ample: the full batch stages + (1024, 64, 8 * 1024, 8), # budget binds before the batch does + (1024, 8, 1 << 20, 8), # batch binds before the budget does + (1024, 64, 1024, 1), # exactly one page fits + (1024, 64, 1, 1), # below one page, raised to one rather than refused + ], +) +def test_plan_slot_geometry(page_bytes, batch, budget, expected_slots): + slot_bytes, num_slots = plan_slot_geometry(page_bytes, batch, budget) + # A slot always holds a whole page: the budget bounds the count, not the width. + assert slot_bytes == page_bytes + assert num_slots == expected_slots + + +@pytest.mark.parametrize("bad", [(0, 8, 1024), (-1, 8, 1024), (1024, 0, 1024)]) +def test_plan_slot_geometry_rejects_degenerate_inputs(bad): + with pytest.raises(ValueError): + plan_slot_geometry(*bad) + + +@pytest.mark.parametrize( + "value,expected", [("1", True), ("true", True), ("on", True), ("0", False), ("off", False)] +) +def test_config_staging_env_override(store_config, monkeypatch, value, expected): + monkeypatch.setenv("TRTLLM_MOONCAKE_STORE_STAGE_THROUGH_HOST", value) + assert MooncakeStoreConnectorConfig.from_env().stage_through_host is expected + + +def test_config_rejects_a_non_boolean_staging_env(store_config, monkeypatch): + monkeypatch.setenv("TRTLLM_MOONCAKE_STORE_STAGE_THROUGH_HOST", "sometimes") + with pytest.raises(ValueError, match="not a boolean"): + MooncakeStoreConnectorConfig.from_env() + + +def test_staging_registers_host_buffers_and_never_the_pools( + store_config, fake_store, staged_copies +): + layout = make_layout(regions_per_group=2) + with make_staged_worker(fake_store, store_config, layout=layout) as worker: + # Registering the pools is the step that needs GPUDirect, so staging + # must not do it at all. + pool_ranges = PageAddressing(layout).registration_ranges() + registered_starts = {address for address, _size in fake_store.registered} + assert registered_starts.isdisjoint({start for start, _end in pool_ranges}) + + # One pinned buffer per direction, since the default role is `both`. + assert len(fake_store.registered) == 2 + assert registered_starts == { + worker._load_staging.slot_address(0), + worker._save_staging.slot_address(0), + } + + +def test_staging_put_hands_the_store_one_host_buffer_per_page( + store_config, fake_store, staged_copies +): + layout = make_layout(regions_per_group=3) + with make_staged_worker(fake_store, store_config, layout=layout) as worker: + block_hash = b"\x07" * 16 + worker._put([RequestTransfers(1, [PageTransfer(block_hash, 0, 2)])]) + + device_addresses, device_sizes = PageAddressing(layout).buffers(0, 2) + slot = worker._save_staging.slot_address(0) + (keys, addresses, sizes) = fake_store.put_calls[0] + + assert keys == [worker._namespaces[0].key(block_hash)] + # The store sees one contiguous host buffer whose length is the sum of + # the device regions. That equality is what keeps a staged write + # byte-identical to a zero-copy one. + assert addresses == [[slot]] + assert sizes == [[sum(device_sizes)]] + + # Every region was gathered, in order, into its place in the slot. + expected = [] + offset = 0 + for address, size in zip(device_addresses, device_sizes): + expected.append((slot + offset, address, size)) + offset += size + assert staged_copies == expected + + +def test_staging_get_scatters_back_to_the_device_regions(store_config, fake_store, staged_copies): + layout = make_layout(regions_per_group=3) + with make_staged_worker(fake_store, store_config, layout=layout) as worker: + block_hash = b"\x03" * 16 + fake_store.objects.add(worker._namespaces[0].key(block_hash)) + + worker.bind_connector_meta( + SimpleNamespace(loads=[RequestTransfers(7, [PageTransfer(block_hash, 0, 5)])], saves=[]) + ) + worker.start_load_kv(None) + + device_addresses, device_sizes = PageAddressing(layout).buffers(0, 5) + slot = worker._load_staging.slot_address(0) + (_keys, addresses, sizes) = fake_store.get_calls[0] + assert addresses == [[slot]] + assert sizes == [[sum(device_sizes)]] + + # The scatter is the mirror of the gather: same split, opposite direction. + expected = [] + offset = 0 + for address, size in zip(device_addresses, device_sizes): + expected.append((address, slot + offset, size)) + offset += size + assert staged_copies == expected + + +def test_staging_does_not_scatter_a_failed_load(store_config, fake_store, staged_copies): + layout = make_layout() + with make_staged_worker(fake_store, store_config, layout=layout) as worker: + block_hash = b"\x09" * 16 + key = worker._namespaces[0].key(block_hash) + fake_store.objects.add(key) + fake_store.fail_gets_for.add(key) + + worker.bind_connector_meta( + SimpleNamespace(loads=[RequestTransfers(7, [PageTransfer(block_hash, 0, 1)])], saves=[]) + ) + with pytest.raises(RuntimeError, match="failed to load"): + worker.start_load_kv(None) + + # A failed read leaves the slot holding whatever it held before, and + # copying that onto the page would put unrelated bytes where the + # runtime already promised computed KV. + assert staged_copies == [] + + +def test_the_ranks_device_is_captured_and_adopted_by_the_save_thread( + store_config, fake_store, fake_cuda +): + """The save thread must not run on torch's default device. + + It issues staging copies against pointers owned by the rank's device. A + stream created on device 0 instead fails every copy with + cudaErrorInvalidValue, and only on ranks other than 0. + """ + with make_worker(fake_store, layout=make_layout()) as worker: + assert worker._device_index == 3 + + deadline = time.monotonic() + 5.0 + while 3 not in fake_cuda and time.monotonic() < deadline: + time.sleep(0.01) + assert 3 in fake_cuda, f"save thread set devices {fake_cuda}, expected the rank's 3" + # Never the thread-local default. + assert 0 not in fake_cuda + + +def test_staging_narrows_the_batch_to_the_budget(store_config, fake_store, staged_copies): + layout = make_layout(regions_per_group=2, num_slots=8) + page_bytes = PageAddressing(layout).bytes_per_page(0) + with make_staged_worker( + fake_store, store_config, layout=layout, budget=2 * page_bytes + ) as worker: + assert worker._save_staging.num_slots == 2 + assert worker._batch_size == 2 + + hashes = [bytes([index]) * 16 for index in range(5)] + worker._put( + [ + RequestTransfers( + 1, + [PageTransfer(block_hash, 0, index) for index, block_hash in enumerate(hashes)], + ) + ] + ) + # Five pages through two slots, and no call wider than the slot count. + assert [len(keys) for keys, _a, _s in fake_store.put_calls] == [2, 2, 1] + + +# ---- scheduler ---- + + +class FakeWorker: + """Stands in for the process-local worker's lookup service.""" + + def __init__(self, hit_blocks=0): + self.hit_blocks = hit_blocks + self.queries = [] + + def count_prefix_hit(self, block_hashes): + self.queries.append(list(block_hashes)) + return min(self.hit_blocks, len(block_hashes)) + + +def make_scheduler(store_config, hit_blocks=0): + scheduler = MooncakeStoreConnectorScheduler(make_llm_args()) + scheduler._worker = FakeWorker(hit_blocks) + return scheduler + + +def request_data(request_id, new_tokens, page_indices, layer_group_id=0): + return RequestData( + request_id=request_id, + new_tokens=list(new_tokens), + new_block_ids=list(page_indices), + computed_position=0, + num_scheduled_tokens=len(new_tokens), + new_block_ids_by_layer_group={layer_group_id: list(page_indices)}, + ) + + +def test_scheduler_offers_the_stored_prefix(store_config): + scheduler = make_scheduler(store_config, hit_blocks=2) + request = make_request(1, list(range(5 * TOKENS_PER_BLOCK))) + assert scheduler.get_num_new_matched_tokens(request, 0) == (2 * TOKENS_PER_BLOCK, False) + + +def test_scheduler_never_offers_the_whole_prompt(store_config): + scheduler = make_scheduler(store_config, hit_blocks=99) + # Exactly three full blocks: the last one is withheld so the runtime still + # has a token to run a forward pass on. + request = make_request(1, list(range(3 * TOKENS_PER_BLOCK))) + matched, _ = scheduler.get_num_new_matched_tokens(request, 0) + assert matched == 2 * TOKENS_PER_BLOCK + + +def test_scheduler_declines_partial_local_matches(store_config): + scheduler = make_scheduler(store_config, hit_blocks=2) + request = make_request(1, list(range(5 * TOKENS_PER_BLOCK))) + assert scheduler.get_num_new_matched_tokens(request, TOKENS_PER_BLOCK + 1) == (0, False) + + +def test_scheduler_offers_nothing_as_a_producer(store_config, monkeypatch): + monkeypatch.setenv("TRTLLM_MOONCAKE_STORE_ROLE", "producer") + scheduler = make_scheduler(store_config, hit_blocks=2) + request = make_request(1, list(range(5 * TOKENS_PER_BLOCK))) + assert scheduler.get_num_new_matched_tokens(request, 0) == (0, False) + assert scheduler._worker.queries == [] + + +def test_scheduler_skips_local_prefix_when_looking_up(store_config): + scheduler = make_scheduler(store_config, hit_blocks=1) + tokens = list(range(6 * TOKENS_PER_BLOCK)) + request = make_request(1, tokens) + scheduler.get_num_new_matched_tokens(request, 2 * TOKENS_PER_BLOCK) + # Blocks 0 and 1 are on device already; candidates start at block 2 and stop + # short of the final block. + full_chain = list(BlockHashChain(TOKENS_PER_BLOCK).extend(tokens)) + assert scheduler._worker.queries[0] == full_chain[2:5] + + +def test_scheduler_builds_loads_for_the_offered_blocks(store_config): + scheduler = make_scheduler(store_config, hit_blocks=2) + tokens = list(range(5 * TOKENS_PER_BLOCK)) + request = make_request(1, tokens) + scheduler.get_num_new_matched_tokens(request, 0) + + output = SchedulerOutput(new_requests=[request_data(1, tokens, [10, 11, 12, 13, 14])]) + metadata = scheduler.build_connector_meta(output) + + assert [page.page_index for page in metadata.loads[0].pages] == [10, 11] + # Blocks 0 and 1 came from the store, so only blocks 2..4 are written back. + assert [page.page_index for page in metadata.saves[0].pages] == [12, 13, 14] + + +def test_scheduler_does_not_resave_blocks_across_iterations(store_config): + scheduler = make_scheduler(store_config, hit_blocks=0) + tokens = list(range(2 * TOKENS_PER_BLOCK)) + request = make_request(1, tokens) + scheduler.get_num_new_matched_tokens(request, 0) + + first = scheduler.build_connector_meta( + SchedulerOutput(new_requests=[request_data(1, tokens, [4, 5])]) + ) + assert [page.page_index for page in first.saves[0].pages] == [4, 5] + + # A generation step completes one more block; only that block is saved. + more_tokens = list(range(2 * TOKENS_PER_BLOCK, 3 * TOKENS_PER_BLOCK)) + second = scheduler.build_connector_meta( + SchedulerOutput(cached_requests=[request_data(1, more_tokens, [6])]) + ) + assert [page.page_index for page in second.saves[0].pages] == [6] + + +def test_scheduler_waits_for_a_block_to_fill_before_saving(store_config): + scheduler = make_scheduler(store_config, hit_blocks=0) + tokens = list(range(TOKENS_PER_BLOCK + 1)) + request = make_request(1, tokens) + scheduler.get_num_new_matched_tokens(request, 0) + + metadata = scheduler.build_connector_meta( + SchedulerOutput(new_requests=[request_data(1, tokens, [4, 5])]) + ) + # Page 5 holds a single token, so only the full block is offered up. + assert [page.page_index for page in metadata.saves[0].pages] == [4] + + +def test_scheduler_saves_nothing_as_a_consumer(store_config, monkeypatch): + monkeypatch.setenv("TRTLLM_MOONCAKE_STORE_ROLE", "consumer") + scheduler = make_scheduler(store_config, hit_blocks=0) + tokens = list(range(2 * TOKENS_PER_BLOCK)) + request = make_request(1, tokens) + scheduler.get_num_new_matched_tokens(request, 0) + metadata = scheduler.build_connector_meta( + SchedulerOutput(new_requests=[request_data(1, tokens, [4, 5])]) + ) + assert metadata.saves == [] + + +def test_scheduler_skips_blocks_without_a_page_in_every_group(store_config): + scheduler = make_scheduler(store_config, hit_blocks=0) + tokens = list(range(2 * TOKENS_PER_BLOCK)) + request = make_request(1, tokens) + scheduler.get_num_new_matched_tokens(request, 0) + + data = request_data(1, tokens, [4, 5]) + data.new_block_ids_by_layer_group[1] = [7, BAD_PAGE_INDEX] + metadata = scheduler.build_connector_meta(SchedulerOutput(new_requests=[data])) + + # Block 1 has no page in group 1, so neither of its halves is stored; block 0 + # contributes one page per group. + assert [(page.layer_group_id, page.page_index) for page in metadata.saves[0].pages] == [ + (0, 4), + (1, 7), + ] + + +def test_scheduler_cancel_load_truncates_the_offer(store_config): + scheduler = make_scheduler(store_config, hit_blocks=3) + tokens = list(range(6 * TOKENS_PER_BLOCK)) + request = make_request(1, tokens) + scheduler.get_num_new_matched_tokens(request, 0) + + # The runtime will not consume anything from block 1 onwards. + scheduler.cancel_load(request, TOKENS_PER_BLOCK, 6 * TOKENS_PER_BLOCK) + metadata = scheduler.build_connector_meta( + SchedulerOutput(new_requests=[request_data(1, tokens, list(range(10, 16)))]) + ) + assert [page.page_index for page in metadata.loads[0].pages] == [10] + + +def test_scheduler_request_finished_pins_pages_only_when_saving(store_config): + scheduler = make_scheduler(store_config, hit_blocks=0) + tokens = list(range(2 * TOKENS_PER_BLOCK)) + request = make_request(1, tokens) + scheduler.get_num_new_matched_tokens(request, 0) + scheduler.build_connector_meta(SchedulerOutput(new_requests=[request_data(1, tokens, [4, 5])])) + assert scheduler.request_finished(request, [4, 5]) is True + # State is dropped with the request, so a second call reports nothing pending. + assert scheduler.request_finished(request, [4, 5]) is False + + +def test_scheduler_request_finished_is_false_without_saves(store_config): + scheduler = make_scheduler(store_config, hit_blocks=0) + request = make_request(1, list(range(TOKENS_PER_BLOCK - 1))) + scheduler.get_num_new_matched_tokens(request, 0) + assert scheduler.request_finished(request, []) is False + + +def test_scheduler_isolates_requests_by_cache_salt(store_config): + scheduler = make_scheduler(store_config, hit_blocks=1) + tokens = list(range(3 * TOKENS_PER_BLOCK)) + scheduler.get_num_new_matched_tokens(make_request(1, tokens, cache_salt="a"), 0) + scheduler.get_num_new_matched_tokens(make_request(2, tokens, cache_salt="b"), 0) + assert scheduler._worker.queries[0] != scheduler._worker.queries[1] diff --git a/tests/unittest/_torch/executor/test_mooncake_store_donor.py b/tests/unittest/_torch/executor/test_mooncake_store_donor.py new file mode 100644 index 000000000000..cda334bc385a --- /dev/null +++ b/tests/unittest/_torch/executor/test_mooncake_store_donor.py @@ -0,0 +1,201 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for lending a node's host memory to a Mooncake pool. + +Runs without a Mooncake installation and without a GPU. The store is a fake +recording what `setup` was called with, since the contract under test is what +the donor asks Mooncake for and how long it holds it. +""" + +import sys +from types import ModuleType + +import pytest + +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store import donor as donor_module +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.donor import ( + DEFAULT_DONOR_LOCAL_BUFFER_SIZE, + donate_segment, + maybe_donate_segment, +) +from tensorrt_llm.llmapi.llm_args import MooncakeDonationConfig + +GIB = 1024**3 + + +class FakeStore: + """The slice of `MooncakeDistributedStore` a donor drives.""" + + instances = [] + + def __init__(self): + self.setup_args = None + self.status = 0 + FakeStore.instances.append(self) + + def setup(self, *args): + self.setup_args = args + return self.status + + +@pytest.fixture +def fake_bindings(monkeypatch): + """Stand in for `mooncake.store`, which is not installed here.""" + FakeStore.instances = [] + package = ModuleType("mooncake") + store = ModuleType("mooncake.store") + store.MooncakeDistributedStore = FakeStore + package.store = store + monkeypatch.setitem(sys.modules, "mooncake", package) + monkeypatch.setitem(sys.modules, "mooncake.store", store) + return FakeStore + + +@pytest.fixture +def failing_bindings(fake_bindings): + """Bindings whose `setup` refuses, as an unreachable master would.""" + + class Refusing(fake_bindings): + def setup(self, *args): + super().setup(*args) + return 7 + + sys.modules["mooncake.store"].MooncakeDistributedStore = Refusing + return Refusing + + +@pytest.mark.parametrize("entry", ["direct", "config"]) +def test_a_donor_registers_the_segment_it_was_asked_for( + entry, fake_bindings, reachable_master, monkeypatch +): + """Both entry paths must reach Mooncake with the same seven setup arguments. + + The config-driven path is what makes a generation server a donor, and it + derives the hostname and parses the size string on the way: a size string + reaching Mooncake unparsed would be a segment of nothing. + """ + if entry == "direct": + expected_host, expected_size, expected_device = "10.0.0.5", 32 * GIB, "mlx5_0" + donation = donate_segment( + "10.0.0.1:50051", + 32 * GIB, + protocol="rdma", + device_name="mlx5_0", + metadata_server="P2PHANDSHAKE", + hostname="10.0.0.5", + ) + else: + monkeypatch.setattr(donor_module, "local_address", lambda: "10.1.2.3") + expected_host, expected_size, expected_device = "10.1.2.3", 320 * GIB, "mlx5_1" + donation = maybe_donate_segment( + MooncakeDonationConfig( + master_server_address="10.0.0.1:50051", + segment_size="320GiB", + protocol="rdma", + device_name="mlx5_1", + ) + ) + + with donation as host: + assert host == expected_host + ( + registered_host, + metadata_server, + segment_size, + local_buffer_size, + protocol, + device_name, + master, + ) = fake_bindings.instances[0].setup_args + + assert registered_host == expected_host + assert metadata_server == "P2PHANDSHAKE" + assert segment_size == expected_size + assert protocol == "rdma" + assert device_name == expected_device + assert master == "10.0.0.1:50051" + assert local_buffer_size == DEFAULT_DONOR_LOCAL_BUFFER_SIZE + + +def test_a_donor_that_cannot_join_says_which_master_it_could_not_reach(failing_bindings): + with pytest.raises(RuntimeError, match="status 7"): + with donate_segment("10.0.0.1:50051", GIB, hostname="10.0.0.5"): + pytest.fail("donation should not have yielded") + + +def test_a_donor_given_no_host_registers_under_the_pool_s_view_of_this_node( + fake_bindings, monkeypatch +): + """The master and the segments registering with it must agree on the host.""" + monkeypatch.setattr(donor_module, "local_address", lambda: "10.1.2.3") + + with donate_segment("10.0.0.1:50051", GIB) as host: + assert host == "10.1.2.3" + assert fake_bindings.instances[0].setup_args[0] == "10.1.2.3" + + +def test_missing_bindings_are_reported_as_the_separate_component_they_are(monkeypatch): + """The container's C++ transfer engine is not these Python bindings.""" + monkeypatch.setitem(sys.modules, "mooncake", None) + monkeypatch.setitem(sys.modules, "mooncake.store", None) + + with pytest.raises(ImportError, match="mooncake-transfer-engine"): + with donate_segment("10.0.0.1:50051", GIB): + pytest.fail("donation should not have yielded") + + +@pytest.fixture +def reachable_master(monkeypatch): + """Skip the socket probe: these tests are about what donation asks for.""" + monkeypatch.setattr(donor_module, "wait_for_master", lambda address: 0.0) + + +def test_a_server_that_was_not_asked_lends_nothing(fake_bindings): + """Every deployment that does not lend memory takes this path.""" + with maybe_donate_segment(None) as host: + assert host is None + assert fake_bindings.instances == [] + + +def test_a_published_master_address_is_read_before_joining( + fake_bindings, reachable_master, tmp_path +): + """What lets a generation server name a path instead of a scheduler's choice.""" + address_file = tmp_path / "master.addr" + address_file.write_text("10.0.0.9:50051\n") + donation = MooncakeDonationConfig( + master_server_address=f"file://{address_file}", + segment_size=GIB, + ) + + with maybe_donate_segment(donation): + assert fake_bindings.instances[0].setup_args[6] == "10.0.0.9:50051" + + +def test_an_unreachable_master_is_reported_before_the_segment_is_offered( + fake_bindings, monkeypatch +): + """Otherwise this is a status code from setup, with no address in it.""" + + def refuse(address): + raise TimeoutError(f"The Mooncake master at {address} did not accept connections") + + monkeypatch.setattr(donor_module, "wait_for_master", refuse) + donation = MooncakeDonationConfig(master_server_address="10.0.0.1:50051") + + with pytest.raises(TimeoutError, match="10.0.0.1:50051"): + with maybe_donate_segment(donation): + pytest.fail("donation should not have yielded") + assert fake_bindings.instances == [] diff --git a/tests/unittest/_torch/executor/test_mooncake_store_master.py b/tests/unittest/_torch/executor/test_mooncake_store_master.py new file mode 100644 index 000000000000..618bcfbb8512 --- /dev/null +++ b/tests/unittest/_torch/executor/test_mooncake_store_master.py @@ -0,0 +1,654 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for provisioning a Mooncake store pool during server bringup. + +Runs without a Mooncake installation and without a GPU. A master this process +launches is a fake standing in for `Popen` that opens the RPC port, which is +all the readiness handshake ever observes. A master someone else runs is a +plain socket. +""" + +import json +import os +import shutil +import socket +import subprocess +import threading +from types import SimpleNamespace + +import pytest + +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store import master as master_module +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.config import ( + CONFIG_PATH_ENV, + MooncakeStoreConnectorConfig, +) +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.master import ( + maybe_provision_pool, + provision_pool, + resolve_master_address, +) +from tensorrt_llm.llmapi.llm_args import KvCacheConnectorConfig, MooncakeStoreConfig + + +def free_port() -> int: + with socket.socket() as probe: + probe.bind(("", 0)) + return probe.getsockname()[1] + + +class FakeMasterProcess: + """The slice of `Popen` that launching a master actually drives. + + `listen_on` makes it answer on that port, which is what a real master does + last and what the readiness wait keys off. `exit_code` makes it a master + that failed to start. + """ + + def __init__(self, command, env, listen_on=None, exit_code=None): + self.command = command + self.env = env + self.pid = 4242 + self.terminated = False + self.killed = False + self._exit_code = exit_code + self._listener = None + if listen_on is not None: + self._listener = socket.socket() + self._listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + self._listener.bind(("", listen_on)) + self._listener.listen(8) + + def poll(self): + return self._exit_code + + def terminate(self): + self.terminated = True + self._exit_code = -15 + if self._listener is not None: + self._listener.close() + self._listener = None + + def wait(self, timeout=None): + return self._exit_code + + def kill(self): + self.killed = True + + +@pytest.fixture(autouse=True) +def clean_env(monkeypatch): + """No ambient pool: these tests are about what provisioning does itself.""" + for name in ( + CONFIG_PATH_ENV, + master_module.MASTER_BINARY_ENV, + master_module.MASTER_TIMEOUT_ENV, + master_module.RUN_DIR_ENV, + ): + monkeypatch.delenv(name, raising=False) + # Nothing here is slow to start, so a wait that runs long is a failure + # rather than something that needs more time. + monkeypatch.setenv(master_module.MASTER_TIMEOUT_ENV, "10") + + +@pytest.fixture +def fake_master(monkeypatch): + """Replace the master binary and its process with in-process fakes. + + Returns a callable that arms the fake. Once provisioning has run, the + launched instance is available as `.process` for inspection. + """ + + class Launcher: + def __init__(self): + self.process = None + + def arm(self, listen_on=None, exit_code=None, log_text=None): + def popen(command, env=None, stdout=None, **_kwargs): + # A real master writes its log through this handle. + if log_text is not None and stdout is not None: + stdout.write(log_text.encode()) + stdout.flush() + self.process = FakeMasterProcess( + command, env, listen_on=listen_on, exit_code=exit_code + ) + return self.process + + # Swap the modules as this module sees them rather than patching + # attributes on the shared stdlib ones. + monkeypatch.setattr( + master_module, + "shutil", + SimpleNamespace(which=lambda name: f"/opt/bin/{name}", rmtree=shutil.rmtree), + ) + monkeypatch.setattr( + master_module, + "subprocess", + SimpleNamespace( + Popen=popen, + STDOUT=subprocess.STDOUT, + TimeoutExpired=subprocess.TimeoutExpired, + ), + ) + + return Launcher() + + +@pytest.fixture +def running_master(): + """A socket standing in for a master someone else is running.""" + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen(8) + try: + yield f"127.0.0.1:{listener.getsockname()[1]}" + finally: + listener.close() + + +# ---- configuration ---- + + +@pytest.mark.parametrize( + "kwargs, message", + [ + (dict(launch_master=True, master_server_address="host:50051"), "not both"), + (dict(), "needs a master"), + # master_address_file only writes an address; reading one is + # master_server_address, so publishing without launching is incoherent. + ( + dict(master_server_address="host:50051", master_address_file="/shared/master.addr"), + "needs launch_master", + ), + ], + ids=["two_masters", "no_master", "publishing_without_launching"], +) +def test_pool_needs_exactly_one_master(kwargs, message): + with pytest.raises(ValueError, match=message): + MooncakeStoreConfig(**kwargs) + + +def test_pool_is_rejected_unless_the_connector_is_mooncake_store(): + """The validator keys off the connector, however that was spelled.""" + with pytest.raises(ValueError, match="mooncake_store describes a Mooncake pool"): + KvCacheConnectorConfig( + connector="lmcache", + mooncake_store=MooncakeStoreConfig(launch_master=True), + ) + # Naming the module rather than the preset selects the same connector. + KvCacheConnectorConfig( + connector_module="tensorrt_llm._torch.pyexecutor.connectors.mooncake_store", + connector_scheduler_class="MooncakeStoreConnectorScheduler", + connector_worker_class="MooncakeStoreConnectorWorker", + mooncake_store=MooncakeStoreConfig(launch_master=True), + ) + + +# ---- the rendered client config ---- + + +def test_client_config_is_what_the_connector_reads_back(tmp_path): + """The generated JSON has to survive the connector's own parser.""" + pool = MooncakeStoreConfig( + master_server_address="10.0.0.1:50051", + protocol="rdma", + device_name="mlx5_0", + global_segment_size="64GiB", + local_buffer_size="4GiB", + cache_prefix="trtllm-m3", + stage_through_host=True, + transfer_batch_size=32, + ) + path = tmp_path / "mooncake.json" + path.write_text(json.dumps(master_module._client_config(pool, "10.0.0.1:50051"))) + + parsed = MooncakeStoreConnectorConfig.from_file(str(path)) + assert parsed.master_server_address == "10.0.0.1:50051" + assert parsed.metadata_server == "P2PHANDSHAKE" + assert parsed.protocol == "rdma" + assert parsed.device_name == "mlx5_0" + assert parsed.global_segment_size == 64 * 1024**3 + assert parsed.local_buffer_size == 4 * 1024**3 + assert parsed.cache_prefix == "trtllm-m3" + assert parsed.stage_through_host is True + assert parsed.transfer_batch_size == 32 + + +def test_client_config_omits_the_fields_the_pool_left_unset(): + """An absent key leaves the connector its own default; a null would not.""" + pool = MooncakeStoreConfig(master_server_address="host:50051") + written = master_module._client_config(pool, "host:50051") + assert "cache_prefix" not in written + assert "staging_buffer_bytes" not in written + + +@pytest.mark.parametrize( + "address, expected", + [ + ("host:50051", ("host", 50051)), + ("[::1]:50051", ("::1", 50051)), + ("unix:///var/run/mooncake", None), + ("host", None), + ], +) +def test_master_addresses_are_split_or_declined(address, expected): + assert master_module._split_address(address) == expected + + +# ---- provisioning against a master someone else runs ---- + + +def test_provisioning_points_the_workers_at_a_running_master(running_master): + pool = MooncakeStoreConfig(master_server_address=running_master) + + with provision_pool(pool) as config_path: + # The workers are spawned inside this window and are told about the + # pool through the environment, so both have to hold while it is open. + assert os.environ[CONFIG_PATH_ENV] == config_path + written = json.loads(open(config_path).read()) + assert written["master_server_address"] == running_master + + assert CONFIG_PATH_ENV not in os.environ + assert not os.path.exists(config_path) + + +def test_a_staging_buffer_can_be_sized_where_staging_is_turned_on(running_master): + """Undersizing it silently shrinks the transfer batch, so it must be settable.""" + pool = MooncakeStoreConfig( + master_server_address=running_master, + stage_through_host=True, + staging_buffer_bytes="4GiB", + ) + + with provision_pool(pool) as config_path: + written = json.loads(open(config_path).read()) + assert written["stage_through_host"] is True + assert written["staging_buffer_bytes"] == "4GiB" + + +def test_provisioning_fails_before_the_model_loads_if_the_master_is_absent(monkeypatch): + monkeypatch.setenv(master_module.MASTER_TIMEOUT_ENV, "1") + pool = MooncakeStoreConfig(master_server_address=f"127.0.0.1:{free_port()}") + + with pytest.raises(TimeoutError, match="did not accept connections"): + with provision_pool(pool): + pytest.fail("provisioning should not have yielded") + assert CONFIG_PATH_ENV not in os.environ + + +def test_an_unparseable_master_address_is_left_to_the_workers(): + """Not every address is host:port, so an unprobeable one passes through.""" + pool = MooncakeStoreConfig(master_server_address="unix:///var/run/mooncake") + + with provision_pool(pool) as config_path: + written = json.loads(open(config_path).read()) + assert written["master_server_address"] == "unix:///var/run/mooncake" + + +def test_an_inherited_config_path_wins(monkeypatch, tmp_path): + """An externally managed pool names itself this way, so provisioning defers.""" + harness_config = tmp_path / "harness.json" + harness_config.write_text("{}") + monkeypatch.setenv(CONFIG_PATH_ENV, str(harness_config)) + pool = MooncakeStoreConfig(launch_master=True) + + with provision_pool(pool) as config_path: + assert config_path is None + assert os.environ[CONFIG_PATH_ENV] == str(harness_config) + + assert os.environ[CONFIG_PATH_ENV] == str(harness_config) + + +# ---- provisioning with a master of our own ---- + + +def test_a_launched_master_is_named_in_the_config_and_stopped_on_exit(fake_master): + port = free_port() + fake_master.arm(listen_on=port) + pool = MooncakeStoreConfig(launch_master=True, master_port=port) + + with provision_pool(pool) as config_path: + written = json.loads(open(config_path).read()) + host, _, named_port = written["master_server_address"].rpartition(":") + assert int(named_port) == port + # The address in the config is all a worker on another host gets, so + # it has to be dialable. + with socket.create_connection((host, port), timeout=5): + pass + + assert fake_master.process.terminated + assert not fake_master.process.killed + + +def test_a_launched_master_gets_the_flags_and_logging_it_needs(fake_master): + port = free_port() + fake_master.arm(listen_on=port) + pool = MooncakeStoreConfig( + launch_master=True, + master_port=port, + master_metrics_port=free_port(), + master_eviction_ratio=0.1, + ) + + with provision_pool(pool): + command = fake_master.process.command + assert command[0].endswith("mooncake_master") + assert f"--rpc_port={port}" in command + assert f"--metrics_port={pool.master_metrics_port}" in command + assert "--eviction_ratio=0.1" in command + # Without these the master logs to a file under /tmp and the log the + # run directory holds stays empty. + assert fake_master.process.env["GLOG_logtostderr"] == "1" + assert fake_master.process.env["GLOG_v"] == "1" + + +def test_a_master_that_dies_during_startup_says_so(fake_master): + fake_master.arm(exit_code=3) + pool = MooncakeStoreConfig(launch_master=True, master_port=free_port()) + + with pytest.raises(RuntimeError, match="exited with code 3"): + with provision_pool(pool): + pytest.fail("provisioning should not have yielded") + assert CONFIG_PATH_ENV not in os.environ + + +def test_a_master_that_never_listens_times_out(monkeypatch, fake_master): + monkeypatch.setenv(master_module.MASTER_TIMEOUT_ENV, "1") + fake_master.arm() + pool = MooncakeStoreConfig(launch_master=True, master_port=free_port()) + + with pytest.raises(TimeoutError, match="did not accept connections"): + with provision_pool(pool): + pytest.fail("provisioning should not have yielded") + assert fake_master.process.terminated + + +def test_a_missing_master_binary_names_the_alternatives(monkeypatch): + monkeypatch.setattr( + master_module, + "shutil", + SimpleNamespace(which=lambda _name: None, rmtree=shutil.rmtree), + ) + pool = MooncakeStoreConfig(launch_master=True) + + with pytest.raises(FileNotFoundError, match="master_server_address"): + with provision_pool(pool): + pytest.fail("provisioning should not have yielded") + + +def test_a_run_dir_keeps_the_master_log_and_the_config(fake_master, tmp_path): + port = free_port() + fake_master.arm(listen_on=port) + run_dir = tmp_path / "pool" + pool = MooncakeStoreConfig(launch_master=True, master_port=port) + + with provision_pool(pool, run_dir=str(run_dir)) as config_path: + assert config_path == str(run_dir / master_module.CLIENT_CONFIG_NAME) + + # An explicit run directory outlives the run that filled it, since the + # master's log is where pool occupancy and eviction are read from. + assert (run_dir / master_module.MASTER_LOG_NAME).exists() + assert (run_dir / master_module.CLIENT_CONFIG_NAME).exists() + + +# ---- the entry point servers call ---- + + +@pytest.mark.parametrize( + "config", + [ + KvCacheConnectorConfig(connector="lmcache"), + None, + KvCacheConnectorConfig(connector="mooncake-store"), + ], + ids=["another_connector", "no_connector", "pool_left_undescribed"], +) +def test_provisioning_is_a_no_op_unless_a_pool_is_described(config): + """Without `mooncake_store`, MOONCAKE_CONFIG_PATH is still the only input.""" + with maybe_provision_pool(config): + assert CONFIG_PATH_ENV not in os.environ + + +def test_a_described_pool_is_provisioned(running_master): + config = KvCacheConnectorConfig( + connector="mooncake-store", + mooncake_store=MooncakeStoreConfig(master_server_address=running_master), + ) + with maybe_provision_pool(config): + written = json.loads(open(os.environ[CONFIG_PATH_ENV]).read()) + assert written["master_server_address"] == running_master + assert CONFIG_PATH_ENV not in os.environ + + +# ---- reaching a master whose host nobody knew in advance ---- + + +@pytest.mark.parametrize("address", ["10.0.0.1:50051", "unix:///var/run/mooncake"]) +def test_an_address_that_is_not_a_file_passes_through(address): + assert resolve_master_address(address, timeout=1.0) == address + + +def test_a_published_address_is_read_from_the_file_that_names_it(tmp_path): + published = tmp_path / "master.addr" + published.write_text("10.0.0.7:50051\n") + + assert resolve_master_address(f"file://{published}", timeout=1.0) == "10.0.0.7:50051" + + +def test_an_address_not_published_yet_is_waited_for(tmp_path): + """Master and workers are started together; neither one orders the other.""" + published = tmp_path / "master.addr" + threading.Timer(0.5, published.write_text, ["10.0.0.9:50051\n"]).start() + + assert resolve_master_address(f"file://{published}", timeout=10.0) == "10.0.0.9:50051" + + +def test_an_empty_address_file_is_not_taken_for_an_address(tmp_path): + """An existing file is not the same as a published address.""" + published = tmp_path / "master.addr" + published.write_text("") + + with pytest.raises(TimeoutError, match="No Mooncake master address"): + resolve_master_address(f"file://{published}", timeout=1.0) + + +def test_an_unpublished_address_names_the_command_that_publishes_it(tmp_path): + with pytest.raises(TimeoutError, match="--address-file"): + resolve_master_address(f"file://{tmp_path / 'absent'}", timeout=1.0) + + +# ---- a master with a lifetime of its own ---- + + +def test_a_standalone_master_publishes_an_address_that_can_be_dialed(fake_master, tmp_path): + port = free_port() + fake_master.arm(listen_on=port) + address_file = tmp_path / "master.addr" + pool = MooncakeStoreConfig(launch_master=True, master_port=port) + + with master_module.running_master( + pool, str(tmp_path / "run"), address_file=str(address_file) + ) as master: + assert resolve_master_address(f"file://{address_file}", timeout=5.0) == master.address + host, _, named_port = master.address.rpartition(":") + assert int(named_port) == port + with socket.create_connection((host, port), timeout=5): + pass + + +def test_a_stopped_master_leaves_no_address_behind(fake_master, tmp_path): + """A stale address would send the next run's workers to a dead port.""" + port = free_port() + fake_master.arm(listen_on=port) + address_file = tmp_path / "master.addr" + pool = MooncakeStoreConfig(launch_master=True, master_port=port) + + with master_module.running_master(pool, str(tmp_path / "run"), address_file=str(address_file)): + assert address_file.exists() + + assert not address_file.exists() + assert fake_master.process.terminated + + +def test_a_standalone_master_keeps_its_log(fake_master, tmp_path): + """A standalone master outlives the servers that used it, so its log is kept.""" + port = free_port() + fake_master.arm(listen_on=port) + run_dir = tmp_path / "run" + pool = MooncakeStoreConfig(launch_master=True, master_port=port) + + with master_module.running_master(pool, str(run_dir)): + pass + + assert (run_dir / master_module.MASTER_LOG_NAME).exists() + + +def test_provisioning_joins_a_master_it_was_never_given_the_address_of(fake_master, tmp_path): + """The address file is how workers reach a master no config names a host for.""" + port = free_port() + fake_master.arm(listen_on=port) + address_file = tmp_path / "master.addr" + standalone = MooncakeStoreConfig(launch_master=True, master_port=port) + worker = MooncakeStoreConfig(master_server_address=f"file://{address_file}") + + with master_module.running_master( + standalone, str(tmp_path / "run"), address_file=str(address_file) + ) as master: + with provision_pool(worker) as config_path: + # Mooncake cannot dial a file:// URL, so what reaches the workers + # has to be the address it resolved to. + written = json.loads(open(config_path).read()) + assert written["master_server_address"] == master.address + + +# ---- a master a server launched, made findable ---- + + +def test_a_launched_master_publishes_where_its_run_left_its_logs(fake_master, tmp_path): + """A finished run's logs still say which pool it used.""" + port = free_port() + fake_master.arm(listen_on=port) + run_dir = tmp_path / "run" + pool = MooncakeStoreConfig(launch_master=True, master_port=port) + + with provision_pool(pool, run_dir=str(run_dir)): + address = (run_dir / master_module.MASTER_ADDRESS_NAME).read_text().strip() + assert address.endswith(f":{port}") + + assert not (run_dir / master_module.MASTER_ADDRESS_NAME).exists() + + +def test_a_launched_master_can_be_published_where_the_donors_look(fake_master, tmp_path): + """This is what lets a server that launched its own master have donors.""" + port = free_port() + fake_master.arm(listen_on=port) + shared = tmp_path / "shared" / "master.addr" + pool = MooncakeStoreConfig( + launch_master=True, master_port=port, master_address_file=str(shared) + ) + + with provision_pool(pool, run_dir=str(tmp_path / "run")): + assert resolve_master_address(f"file://{shared}", timeout=5.0).endswith(f":{port}") + + # Retracted, so the next run's donors wait for a live master rather than + # joining a pool that no longer exists. + assert not shared.exists() + + +def test_a_half_written_address_is_never_read(tmp_path): + """A reader sees the whole address or nothing, never a prefix of one.""" + target = tmp_path / "master.addr" + + with master_module._published_address("10.0.0.7:50051", [str(target)]): + assert not (tmp_path / "master.addr.partial").exists() + assert target.read_text().strip() == "10.0.0.7:50051" + + +# ---- saying why bringup is stuck ---- + + +def test_an_absent_master_is_named_rather_than_left_to_store_setup(monkeypatch): + """Otherwise the failure is a bare status code in every rank, after loading.""" + monkeypatch.setenv(master_module.MASTER_TIMEOUT_ENV, "1") + address = f"127.0.0.1:{free_port()}" + + with pytest.raises(TimeoutError, match=address): + master_module.wait_for_master(address) + + +def test_an_address_of_a_shape_we_cannot_probe_is_not_fatal(): + """Mooncake may accept addresses this cannot dial; leave them to it.""" + assert master_module.wait_for_master("unix:///var/run/mooncake") is None + + +def test_a_master_that_died_starting_is_reported_with_its_last_words(fake_master, tmp_path): + """The reason is in the master's log, which is only read if the error quotes it.""" + run_dir = tmp_path / "run" + fake_master.arm(exit_code=1, log_text="E0903 bind(50051) failed: Address already in use\n") + pool = MooncakeStoreConfig(launch_master=True, master_port=free_port()) + + with pytest.raises(RuntimeError, match="Address already in use"): + with provision_pool(pool, run_dir=str(run_dir)): + pytest.fail("provisioning should not have yielded") + + +# ---- choosing the fabric without naming it in a config ---- + + +def fake_hca(root, device, link_layer="InfiniBand", state="4: ACTIVE", rate="800 Gb/sec"): + port = root / device / "ports" / "1" + port.mkdir(parents=True) + (port / "link_layer").write_text(f"{link_layer}\n") + (port / "state").write_text(f"{state}\n") + (port / "rate").write_text(f"{rate}\n") + + +def test_the_compute_fabric_is_picked_over_the_management_adapter(tmp_path): + """A node's HCAs are not interchangeable: only some are the fast fabric.""" + fake_hca(tmp_path, "mlx5_0") + fake_hca(tmp_path, "mlx5_1") + fake_hca(tmp_path, "mlx5_2", rate="400 Gb/sec") + fake_hca(tmp_path, "mlx5_3", state="1: DOWN") + fake_hca(tmp_path, "mlx5_4", link_layer="Ethernet") + + assert ( + master_module.resolve_device_name("rdma", "", sysfs_root=str(tmp_path)) == "mlx5_0,mlx5_1" + ) + + +def test_a_named_device_is_not_second_guessed(tmp_path): + fake_hca(tmp_path, "mlx5_0") + assert master_module.resolve_device_name("rdma", "mlx5_7", sysfs_root=str(tmp_path)) == "mlx5_7" + + +def test_tcp_needs_no_device_and_looks_for_none(tmp_path): + assert master_module.resolve_device_name("tcp", "", sysfs_root=str(tmp_path)) == "" + + +def test_a_node_without_infiniband_is_left_to_mooncake_s_own_discovery(tmp_path): + """Falling back beats failing, since Mooncake may still find a usable device.""" + assert master_module.resolve_device_name("rdma", "", sysfs_root=str(tmp_path / "absent")) == "" + + +def test_the_detected_device_is_what_the_workers_are_told(fake_master, tmp_path, monkeypatch): + sysfs = tmp_path / "sysfs" + fake_hca(sysfs, "mlx5_0") + monkeypatch.setattr(master_module, "IB_SYSFS_ROOT", str(sysfs)) + port = free_port() + fake_master.arm(listen_on=port) + pool = MooncakeStoreConfig(launch_master=True, master_port=port, protocol="rdma") + + with provision_pool(pool, run_dir=str(tmp_path / "run")) as config_path: + assert json.loads(open(config_path).read())["device_name"] == "mlx5_0" diff --git a/tests/unittest/api_stability/references/llm.yaml b/tests/unittest/api_stability/references/llm.yaml index 1f5e97d1b74f..f2aefd22f9c6 100644 --- a/tests/unittest/api_stability/references/llm.yaml +++ b/tests/unittest/api_stability/references/llm.yaml @@ -243,6 +243,10 @@ methods: annotation: Optional[tensorrt_llm.llmapi.llm_args.KvCacheConnectorConfig] default: null status: prototype + mooncake_donation: + annotation: Optional[tensorrt_llm.llmapi.llm_args.MooncakeDonationConfig] + default: null + status: prototype enable_lm_head_tp_in_adp: annotation: bool default: False