From 15a919d035fa06cf0dcea0f9b6f8c1c667027e87 Mon Sep 17 00:00:00 2001 From: Pei Li Date: Wed, 7 Oct 2026 16:54:48 -0700 Subject: [PATCH 1/4] fix(trtllm): pin MPI rank order on uniform multi-node endpoints (#583) * fix(trtllm): pin MPI rank order on uniform multi-node endpoints A multi-node TRT-LLM endpoint spanning full nodes passed Slurm a plain --nodelist, which names the nodes but not the task order: Slurm may lay ranks out by its own topology, putting rank 0 on a node other than the leader whose IP is MASTER_ADDR, and the rendezvous then stalls at startup. Uneven layouts already used --distribution=arbitrary with one host per task; this extends that to every multi-node endpoint (e.g. --ntasks=8 --ntasks-per-node=4 --nodelist=A,A,A,A,B,B,B,B, no --nodes). Single-node endpoints are unchanged. Tests cover the rendered srun and the relaunch path; docs/engines.md describes the layout. Signed-off-by: Pei Li * fix(trtllm): review follow-ups: num_nodes condition, warn on replaced distribution, two-node example and snapshot Use num_nodes > 1 for the multi-node check; log a warning naming a cluster- or role-configured srun_options.distribution that the explicit per-rank layout replaces, with a test; assert MASTER_ADDR in the rendered command; add examples/trtllm/dynamo-disagg-multinode.yaml with its launch snapshot and README row so the layout is visible in tests/snapshots/launch; docs show the flags as rendered. Signed-off-by: Pei Li * fix(trtllm): make the two-node example an H100 example The suite treats every examples/{sglang,vllm,trtllm}/*.yaml as an H100 recipe (8 GPUs per node) and pins their count. The two-node example becomes TP8 prefill on one node and TP16 decode across two 8-GPU nodes; the snapshot shows sixteen ranks, eight hosts each; the curated count is 17. Signed-off-by: Pei Li --------- Signed-off-by: Pei Li --- docs/engines.md | 10 + examples/README.md | 2 +- examples/trtllm/dynamo-disagg-multinode.yaml | 82 +++++ src/srtctl/cli/mixins/worker_stage.py | 16 +- .../trtllm__dynamo-disagg-multinode.txt | 285 ++++++++++++++++++ tests/test_e2e.py | 2 +- tests/test_slurm.py | 81 ++++- 7 files changed, 473 insertions(+), 5 deletions(-) create mode 100644 examples/trtllm/dynamo-disagg-multinode.yaml create mode 100644 tests/snapshots/launch/trtllm__dynamo-disagg-multinode.txt diff --git a/docs/engines.md b/docs/engines.md index d8be37eef..aac8ebb46 100644 --- a/docs/engines.md +++ b/docs/engines.md @@ -228,6 +228,16 @@ TRT-LLM endpoints also set `MASTER_ADDR` to the rank-zero node and use a distinc `MASTER_PORT` per endpoint. This overrides container hooks that infer rank zero from Slurm's sorted node list. Explicit recipe environment values take precedence. +Every multi-node endpoint, uniform (two full nodes) or uneven, is launched with +`--distribution=arbitrary` and one host per task in srt-slurm's node order, for +example `--ntasks 8 --ntasks-per-node=4 --nodelist A,A,A,A,B,B,B,B` (the +`--nodes` flag is omitted, as Slurm requires for this distribution). A plain +`--nodelist` only names the nodes; Slurm may still order tasks by its own +topology, which can put rank 0 on a node other than the `MASTER_ADDR` leader and +stall the rendezvous. Single-node endpoints are unchanged. A `distribution` set in +the cluster's or the role's `srun_options` is replaced for such an endpoint, with a +warning in the job log. + **Other TRT-LLM launch facts**: TRT-LLM supports prefill, decode, and aggregated roles, uses MPI-style launching (one srun per endpoint with all of its nodes) through `trtllm-llmapi-launch`, and sets `TRTLLM_EPLB_SHM_NAME` to a unique UUID per endpoint. ## ATOM with AToMesh diff --git a/examples/README.md b/examples/README.md index 9943c2692..ad55b215e 100644 --- a/examples/README.md +++ b/examples/README.md @@ -8,7 +8,7 @@ Small starting points, one per frontend and topology. The matrix examples serve | --- | --- | --- | --- | | SGLang | `sglang/dynamo-agg.yaml`, `sglang/dynamo-disagg.yaml` | `sglang/sglang-router-agg.yaml`, `sglang/sglang-router-disagg.yaml`, `sglang/smg-disagg.yaml` (SMG) | `sglang/sglang-direct-agg.yaml` | | vLLM | `vllm/dynamo-agg.yaml`, `vllm/dynamo-disagg.yaml` | `vllm/vllm-router-agg.yaml`, `vllm/vllm-router-disagg.yaml`, `vllm/vllm-router-moriio-disagg.yaml` (ROCm, MoRI-IO discovery), `vllm/smg-agg.yaml`, `vllm/smg-dep16.yaml`, `vllm/smg-disagg-grpc.yaml` (SMG; P/D over gRPC with NIXL) | `vllm/vllm-direct-agg.yaml` | -| TRT-LLM | `trtllm/dynamo-agg.yaml`, `trtllm/dynamo-disagg.yaml` | `trtllm/trtllm-serve-disagg.yaml`, `trtllm/smg-agg.yaml` (SMG) | `trtllm/trtllm-serve-agg.yaml` | +| TRT-LLM | `trtllm/dynamo-agg.yaml`, `trtllm/dynamo-disagg.yaml`, `trtllm/dynamo-disagg-multinode.yaml` (TP16 decode across two 8-GPU nodes: one srun, one host per rank under `--distribution=arbitrary`) | `trtllm/trtllm-serve-disagg.yaml`, `trtllm/smg-agg.yaml` (SMG) | `trtllm/trtllm-serve-agg.yaml` | | Mocker | `mocker/dynamo-agg.yaml` | | | - **Dynamo frontend**: workers register with etcd and the Dynamo frontend routes (KV-aware here); the request plane is tcp and NATS is not started unless a plane asks for it. Dynamo is installed at job start via `dynamo.source` (`pypi:` here) unless the container ships it (`dynamo.install: false`, as the TRT-LLM examples do). diff --git a/examples/trtllm/dynamo-disagg-multinode.yaml b/examples/trtllm/dynamo-disagg-multinode.yaml new file mode 100644 index 000000000..47637184b --- /dev/null +++ b/examples/trtllm/dynamo-disagg-multinode.yaml @@ -0,0 +1,82 @@ +# TRT-LLM disaggregated serving with a decode worker that spans two nodes. +# +# One TP8 prefill worker on an 8-GPU node and one TP16 decode worker across two +# 8-GPU nodes, behind the Dynamo frontend. The decode worker is a single srun of +# sixteen MPI ranks: srtctl lists one host per rank in its own node order under +# `--distribution=arbitrary` (`--nodelist A,A,...,B,B,...`, eight each), so rank 0 +# is on the node whose IP it exported as MASTER_ADDR. Model and parallelism are +# placeholders for the layout; see +# tests/snapshots/launch/trtllm__dynamo-disagg-multinode.txt for the launch plan. +# +# Aliases (`qwen3-0.6b`, `trtllm`) resolve through srtslurm.yaml; see examples/README.md. + +schema: 2 +name: "qwen3-0.6b-trtllm-dynamo-disagg-multinode" + +slurm: + time_limit: "00:30:00" + +model: + path: "qwen3-0.6b" + container: "trtllm" + precision: "bf16" + +resources: + gpu_type: "h100" + gpus_per_node: 8 +dynamo: + install: false + +frontend: + type: dynamo + enable_multiple_frontends: false + args: + router-mode: "kv" + +engine: + type: trtllm + served_model_name: "Qwen/Qwen3-0.6B" +roles: + prefill: + nodes: 1 + workers: 1 + gpus: 8 + env: + PYTHONUNBUFFERED: "1" + TLLM_LOG_LEVEL: "INFO" + args: + backend: "pytorch" + tensor_parallel_size: 8 + max_batch_size: 8 + max_num_tokens: 8192 + max_seq_len: 4096 + trust_remote_code: true + disable_overlap_scheduler: true + kv_cache_config: + free_gpu_memory_fraction: 0.5 + cache_transceiver_config: + backend: "UCX" + decode: + nodes: 2 + workers: 1 + gpus: 16 + env: + PYTHONUNBUFFERED: "1" + TLLM_LOG_LEVEL: "INFO" + args: + backend: "pytorch" + tensor_parallel_size: 16 + max_batch_size: 64 + max_num_tokens: 4096 + max_seq_len: 4096 + trust_remote_code: true + kv_cache_config: + free_gpu_memory_fraction: 0.5 + cache_transceiver_config: + backend: "UCX" + +benchmark: + type: "sa-bench" + isl: 128 + osl: 128 + concurrencies: "4x8" diff --git a/src/srtctl/cli/mixins/worker_stage.py b/src/srtctl/cli/mixins/worker_stage.py index 3635cf2df..9e6bb9018 100644 --- a/src/srtctl/cli/mixins/worker_stage.py +++ b/src/srtctl/cli/mixins/worker_stage.py @@ -641,7 +641,21 @@ def start_endpoint_worker(self, endpoint_processes: list["Process"], *, attempt: task_counts = [len(p.gpu_indices) for p in endpoint_processes] srun_options = {**self.runtime.srun_options, **self.config.roles[mode].srun_options} srun_options["ntasks-per-node"] = str(max(task_counts)) - if len(set(task_counts)) > 1: + if num_nodes > 1: + # Every multi-node endpoint, uniform or not, lays its ranks out explicitly: + # one host per task, in srt-slurm's node order, under --distribution=arbitrary. + # A plain --nodelist only names the nodes; Slurm may still order the tasks + # by its own topology, putting rank 0 on a node other than the leader whose + # IP is MASTER_ADDR, and the rendezvous then waits forever. + configured = srun_options.get("distribution") + if configured not in (None, "arbitrary"): + logger.warning( + "%s worker %d: srun_options.distribution=%r is replaced by 'arbitrary'; " + "a multi-node endpoint needs the explicit per-rank host list", + mode, + index, + configured, + ) srun_options["distribution"] = "arbitrary" endpoint_nodes = task_nodes diff --git a/tests/snapshots/launch/trtllm__dynamo-disagg-multinode.txt b/tests/snapshots/launch/trtllm__dynamo-disagg-multinode.txt new file mode 100644 index 000000000..47d89bd5b --- /dev/null +++ b/tests/snapshots/launch/trtllm__dynamo-disagg-multinode.txt @@ -0,0 +1,285 @@ +# examples/trtllm/dynamo-disagg-multinode.yaml +# exit_code: 0 +# srun calls: 16 + +## service_etcd +nodelist: node-01 +output: /logs/service_etcd.out +container_image: trtllm +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +use_bash_wrapper: True +bash_preamble: | + rm -rf /tmp/etcd && mkdir -p /tmp/etcd +command: | + /configs/etcd --data-dir /tmp/etcd --listen-client-urls http://0.0.0.0:2379 --advertise-client-urls http://127.0.0.1:2379 + +## prefill_0_node-01 +nodes: 1 +ntasks: 8 +nodelist: node-01 +output: /logs/node-01_prefill_w0.out +container_image: trtllm +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +env_to_set: + DYN_LOG: info,dynamo_runtime::pipeline::network::ingress::push_handler=warn + DYN_REQUEST_PLANE: tcp + DYN_SKIP_SGLANG_LOG_FORMATTING: 1 + DYN_SYSTEM_PORT: 7500 + ETCD_ENDPOINTS: http://127.0.0.1:2379 + HEAD_NODE_IP: 127.0.0.1 + MASTER_ADDR: 127.0.0.1 + MASTER_PORT: 29500 + PYTHONUNBUFFERED: 1 + TLLM_LOG_LEVEL: INFO + TRTLLM_EPLB_SHM_NAME: +srun_options: + kill-on-bad-exit: 1 + ntasks-per-node: 8 +mpi: pmix +oversubscribe: True +cpu_bind: verbose,none +bash_preamble: | + ( python3 <(cat <<'__FINGERPRINT_EOF__' + + __FINGERPRINT_EOF__ + ) || true ) +command: | + trtllm-llmapi-launch python3 -m dynamo.trtllm --model-path /model --served-model-name Qwen/Qwen3-0.6B --disaggregation-mode prefill --extra-engine-args /logs/trtllm_config_prefill.yaml --request-plane tcp --publish-metrics + +## decode_0_node-02 +nodes: 2 +ntasks: 16 +nodelist: node-02 node-02 node-02 node-02 node-02 node-02 node-02 node-02 node-03 node-03 node-03 node-03 node-03 node-03 node-03 node-03 +output: /logs/node-02_decode_w0.out +container_image: trtllm +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +env_to_set: + DYN_LOG: info,dynamo_runtime::pipeline::network::ingress::push_handler=warn + DYN_REQUEST_PLANE: tcp + DYN_SKIP_SGLANG_LOG_FORMATTING: 1 + DYN_SYSTEM_PORT: 7501 + ETCD_ENDPOINTS: http://127.0.0.1:2379 + HEAD_NODE_IP: 127.0.0.1 + MASTER_ADDR: 127.0.0.1 + MASTER_PORT: 29501 + PYTHONUNBUFFERED: 1 + TLLM_LOG_LEVEL: INFO + TRTLLM_EPLB_SHM_NAME: +srun_options: + distribution: arbitrary + kill-on-bad-exit: 1 + ntasks-per-node: 8 +mpi: pmix +oversubscribe: True +cpu_bind: verbose,none +bash_preamble: | + ( python3 <(cat <<'__FINGERPRINT_EOF__' + + __FINGERPRINT_EOF__ + ) || true ) +command: | + trtllm-llmapi-launch python3 -m dynamo.trtllm --model-path /model --served-model-name Qwen/Qwen3-0.6B --disaggregation-mode decode --extra-engine-args /logs/trtllm_config_decode.yaml --request-plane tcp --publish-metrics + +## frontend_0 +nodelist: node-01 +output: /logs/node-01_frontend_0.out +container_image: trtllm +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +env_to_set: + DYN_REQUEST_PLANE: tcp + DYN_SKIP_SGLANG_LOG_FORMATTING: 1 + ETCD_ENDPOINTS: http://127.0.0.1:2379 +mpi: pmix +command: | + python3 -m dynamo.frontend --http-port=8000 --router-mode kv + +## service_dcgm-exporter_node-01 +nodelist: node-01 +output: /logs/service_dcgm-exporter_node-01.out +container_image: nvcr.io#nvidia/k8s/dcgm-exporter:3.3.9-3.6.1-ubuntu22.04 +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +srun_export_env: + ETCD_ENDPOINTS: http://127.0.0.1:2379 +use_bash_wrapper: False +command: | + dcgm-exporter --collect-interval=1000 --address :9401 + +## service_dcgm-exporter_node-02 +nodelist: node-02 +output: /logs/service_dcgm-exporter_node-02.out +container_image: nvcr.io#nvidia/k8s/dcgm-exporter:3.3.9-3.6.1-ubuntu22.04 +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +srun_export_env: + ETCD_ENDPOINTS: http://127.0.0.1:2379 +use_bash_wrapper: False +command: | + dcgm-exporter --collect-interval=1000 --address :9401 + +## service_dcgm-exporter_node-03 +nodelist: node-03 +output: /logs/service_dcgm-exporter_node-03.out +container_image: nvcr.io#nvidia/k8s/dcgm-exporter:3.3.9-3.6.1-ubuntu22.04 +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +srun_export_env: + ETCD_ENDPOINTS: http://127.0.0.1:2379 +use_bash_wrapper: False +command: | + dcgm-exporter --collect-interval=1000 --address :9401 + +## service_dcgm-exporter_node-04 +nodelist: node-04 +output: /logs/service_dcgm-exporter_node-04.out +container_image: nvcr.io#nvidia/k8s/dcgm-exporter:3.3.9-3.6.1-ubuntu22.04 +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +srun_export_env: + ETCD_ENDPOINTS: http://127.0.0.1:2379 +use_bash_wrapper: False +command: | + dcgm-exporter --collect-interval=1000 --address :9401 + +## service_dcgm-exporter_node-05 +nodelist: node-05 +output: /logs/service_dcgm-exporter_node-05.out +container_image: nvcr.io#nvidia/k8s/dcgm-exporter:3.3.9-3.6.1-ubuntu22.04 +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +srun_export_env: + ETCD_ENDPOINTS: http://127.0.0.1:2379 +use_bash_wrapper: False +command: | + dcgm-exporter --collect-interval=1000 --address :9401 + +## service_node-exporter_node-01 +nodelist: node-01 +output: /logs/service_node-exporter_node-01.out +container_image: quay.io#prometheus/node-exporter:v1.8.2 +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +srun_export_env: + ETCD_ENDPOINTS: http://127.0.0.1:2379 +use_bash_wrapper: False +command: | + /bin/node_exporter --web.listen-address=:9101 --collector.disable-defaults --collector.cpu --collector.infiniband --collector.meminfo --collector.processes --collector.stat --collector.vmstat --collector.pressure --collector.meminfo_numa '--collector.vmstat.fields=^(oom_kill|pgpg|pswp|pgsteal|pg.*fault).*' + +## service_node-exporter_node-02 +nodelist: node-02 +output: /logs/service_node-exporter_node-02.out +container_image: quay.io#prometheus/node-exporter:v1.8.2 +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +srun_export_env: + ETCD_ENDPOINTS: http://127.0.0.1:2379 +use_bash_wrapper: False +command: | + /bin/node_exporter --web.listen-address=:9101 --collector.disable-defaults --collector.cpu --collector.infiniband --collector.meminfo --collector.processes --collector.stat --collector.vmstat --collector.pressure --collector.meminfo_numa '--collector.vmstat.fields=^(oom_kill|pgpg|pswp|pgsteal|pg.*fault).*' + +## service_node-exporter_node-03 +nodelist: node-03 +output: /logs/service_node-exporter_node-03.out +container_image: quay.io#prometheus/node-exporter:v1.8.2 +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +srun_export_env: + ETCD_ENDPOINTS: http://127.0.0.1:2379 +use_bash_wrapper: False +command: | + /bin/node_exporter --web.listen-address=:9101 --collector.disable-defaults --collector.cpu --collector.infiniband --collector.meminfo --collector.processes --collector.stat --collector.vmstat --collector.pressure --collector.meminfo_numa '--collector.vmstat.fields=^(oom_kill|pgpg|pswp|pgsteal|pg.*fault).*' + +## service_node-exporter_node-04 +nodelist: node-04 +output: /logs/service_node-exporter_node-04.out +container_image: quay.io#prometheus/node-exporter:v1.8.2 +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +srun_export_env: + ETCD_ENDPOINTS: http://127.0.0.1:2379 +use_bash_wrapper: False +command: | + /bin/node_exporter --web.listen-address=:9101 --collector.disable-defaults --collector.cpu --collector.infiniband --collector.meminfo --collector.processes --collector.stat --collector.vmstat --collector.pressure --collector.meminfo_numa '--collector.vmstat.fields=^(oom_kill|pgpg|pswp|pgsteal|pg.*fault).*' + +## service_node-exporter_node-05 +nodelist: node-05 +output: /logs/service_node-exporter_node-05.out +container_image: quay.io#prometheus/node-exporter:v1.8.2 +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +srun_export_env: + ETCD_ENDPOINTS: http://127.0.0.1:2379 +use_bash_wrapper: False +command: | + /bin/node_exporter --web.listen-address=:9101 --collector.disable-defaults --collector.cpu --collector.infiniband --collector.meminfo --collector.processes --collector.stat --collector.vmstat --collector.pressure --collector.meminfo_numa '--collector.vmstat.fields=^(oom_kill|pgpg|pswp|pgsteal|pg.*fault).*' + +## tachometer +nodelist: node-01 +output: /logs/tachometer.out +srun_export_env: + POLARS_MAX_THREADS: 4 +use_bash_wrapper: False +command: | + tachometer-scraper --config /logs/tachometer_config.toml --local-dir /logs/tachometer/local --sync-interval 120 + +## +nodelist: node-01 +output: /logs/benchmark.out +container_image: trtllm +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +env_to_set: + SRTCTL_FRONTEND_TYPE: dynamo + SRT_FRONTEND_HOST: 127.0.0.1 + SRT_FRONTEND_PORT: 8000 +command: | + bash /srtctl-benchmarks/sa-bench/bench.sh http://localhost:8000 128 128 4x8 inf /model Qwen/Qwen3-0.6B true 24 8 16 0.8 10 2 '' true random '' false diff --git a/tests/test_e2e.py b/tests/test_e2e.py index f450c39cc..518e9cc02 100644 --- a/tests/test_e2e.py +++ b/tests/test_e2e.py @@ -30,7 +30,7 @@ def test_interactive_discovers_curated_examples(): assert SGLANG_ROUTER_DISAGG in examples assert MOCKER_EXAMPLE in examples - assert len(TOPOLOGY_EXAMPLES) == 21 + assert len(TOPOLOGY_EXAMPLES) == 22 # ============================================================================= diff --git a/tests/test_slurm.py b/tests/test_slurm.py index ca75d6b40..2ef1519df 100644 --- a/tests/test_slurm.py +++ b/tests/test_slurm.py @@ -454,7 +454,10 @@ def test_relaunch_endpoint_follows_the_backend_launch_strategy(tmp_path: Path) - with patches[0], patches[1] as mock_srun: procs = list(mixin.relaunch_endpoint([leader, follower], attempt=3)) assert [p.name for p in procs] == ["prefill_0_node-a_r3"] - assert mock_srun.call_args.kwargs["nodelist"] == ["node-a", "node-b"] + # Multi-node: one host per rank in srt-slurm's order, under --distribution=arbitrary. + expected_hosts = ["node-a"] * len(leader.gpu_indices) + ["node-b"] * len(follower.gpu_indices) + assert mock_srun.call_args.kwargs["nodelist"] == expected_hosts + assert mock_srun.call_args.kwargs["srun_options"]["distribution"] == "arbitrary" assert mock_srun.call_args.kwargs["step_name"] == "prefill_0_node-a_r3" @@ -1018,7 +1021,81 @@ def test_endpoint_launch_uniform_nodes(tmp_path: Path, gpu_count: int, nodes: in assert kwargs["ntasks"] == gpu_count assert kwargs["nodes"] == nodes # TRT-LLM endpoint steps end when any task exits non-zero (see SrunConfig.kill_on_bad_exit). - assert kwargs["srun_options"] == {"ntasks-per-node": str(per_node), "kill-on-bad-exit": "1"} + expected_options = {"ntasks-per-node": str(per_node), "kill-on-bad-exit": "1"} + if nodes > 1: + # Ranks follow srt-slurm's node order, not Slurm's: one host per task, rank 0 on the leader. + expected_options["distribution"] = "arbitrary" + assert kwargs["nodelist"] == [f"node{i}" for i in range(nodes) for _ in range(per_node)] + else: + assert kwargs["nodelist"] == ["node0"] + assert kwargs["srun_options"] == expected_options + + +def test_uniform_multinode_endpoint_pins_rank_order(tmp_path: Path) -> None: + """The rendered srun for a 2x4 endpoint: ranks 0-3 on the leader, 4-7 on the second node, no --nodes.""" + from srtctl.backends.trtllm import TRTLLMBackend + + mixin, _ = _remap_worker_mixin(tmp_path, frontend_type="trtllm_serve", dynamo_install=False) + mixin.runtime.gpus_per_node = 4 + mixin.backend.type = "trtllm" + mixin.backend.get_srun_config.return_value = TRTLLMBackend().get_srun_config() + endpoints = TRTLLMBackend().allocate_endpoints( + num_prefill=1, + num_decode=0, + num_agg=0, + gpus_per_prefill=8, + gpus_per_decode=0, + gpus_per_agg=0, + gpus_per_node=4, + available_nodes=("nodeA", "nodeB"), + ) + processes = TRTLLMBackend().endpoints_to_processes(endpoints) + with ( + patch("srtctl.cli.mixins.worker_stage.generate_capture_script", return_value="true"), + patch("srtctl.cli.mixins.worker_stage.get_hostname_ip", return_value="10.0.0.1") as mock_ip, + patch("srtctl.core.slurm.get_slurm_job_id", return_value="12345"), + patch("srtctl.core.slurm._get_cluster_bash_preamble", return_value=None), + patch("subprocess.Popen") as mock_popen, + ): + mixin.start_endpoint_worker(processes) + command = mock_popen.call_args.args[0] + assert command[command.index("--ntasks") + 1] == "8" + assert "--ntasks-per-node=4" in command + assert "--distribution=arbitrary" in command + assert "--nodes" not in command + assert command[command.index("--nodelist") + 1] == ",".join(["nodeA"] * 4 + ["nodeB"] * 4) + # MASTER_ADDR is the leader's IP, and the leader is where rank 0 now provably lands. + mock_ip.assert_any_call("nodeA", mixin.runtime.network_interface) + assert "export MASTER_ADDR=10.0.0.1" in command[-1] + + +def test_multinode_endpoint_warns_when_it_replaces_a_configured_distribution(tmp_path: Path, caplog) -> None: + from srtctl.backends.trtllm import TRTLLMBackend + from srtctl.core.topology import endpoints_to_processes + + mixin, _ = _remap_worker_mixin(tmp_path, frontend_type="trtllm_serve", dynamo_install=False) + mixin.runtime.gpus_per_node = 4 + mixin.runtime.srun_options = {"distribution": "block:block"} + mixin.backend.get_srun_config.return_value = TRTLLMBackend().get_srun_config() + endpoints = TRTLLMBackend().allocate_endpoints( + num_prefill=1, + num_decode=0, + num_agg=0, + gpus_per_prefill=8, + gpus_per_decode=0, + gpus_per_agg=0, + gpus_per_node=4, + available_nodes=("nodeA", "nodeB"), + ) + with ( + patch("srtctl.cli.mixins.worker_stage.generate_capture_script", return_value="true"), + patch("srtctl.cli.mixins.worker_stage.get_hostname_ip", return_value="10.0.0.1"), + patch("srtctl.cli.mixins.worker_stage.start_srun_process") as mock_srun, + caplog.at_level("WARNING", logger="srtctl.cli.mixins.worker_stage"), + ): + mixin.start_endpoint_worker(endpoints_to_processes(endpoints)) + assert mock_srun.call_args.kwargs["srun_options"]["distribution"] == "arbitrary" + assert "distribution='block:block' is replaced by 'arbitrary'" in caplog.text def test_endpoint_rejects_incompatible_local_rank_mapping(tmp_path: Path) -> None: From 9c4439500316afd34c9193b292adbbf874fcb013 Mon Sep 17 00:00:00 2001 From: Cameron Quilici Date: Wed, 7 Oct 2026 19:07:55 -0500 Subject: [PATCH 2/4] feat(tokenspeed): add a TokenSpeed backend for Dynamo and SMG (#580) * feat(tokenspeed): add a native TokenSpeed backend Add `engine: tokenspeed`, which launches Dynamo's TokenSpeed backend (`python3 -m dynamo.tokenspeed`) behind the Dynamo frontend, aggregated or prefill/decode with Mooncake. srtctl sets the model, served model name, host, multi-node rendezvous and P/D flags. `--port` comes from a new per-node TOKENSPEED_PORTS kind spaced 1024 apart, because TokenSpeed scans for its NCCL port above it, and `--dist-init-addr` from a DIST_INIT_PORTS block sized for the control-plane cluster TokenSpeed binds on the leader. Adds examples, launch snapshots, docs/tokenspeed.md and tests. Signed-off-by: Cam Quilici * docs(tokenspeed): name the first Dynamo nightly with TokenSpeed prefill/decode Signed-off-by: Cam Quilici * docs(tokenspeed): link the examples by URL for the strict docs build Signed-off-by: Cam Quilici * feat(tokenspeed): run TokenSpeed gRPC engines behind SMG Signed-off-by: Cam Quilici * docs: link the SMG and TokenSpeed examples by URL for the strict docs build Signed-off-by: Cam Quilici * refactor(tokenspeed): inherit the Prometheus metrics path from Backend Signed-off-by: Cam Quilici * chore(tokenspeed): use the shared WorkerMode and the NVIDIA copyright header Signed-off-by: Cam Quilici * feat(tokenspeed): point the Mooncake Store L3 cache at a mooncake-master service Signed-off-by: Cam Quilici * fix(tokenspeed): skip the gRPC engine warmup on prefill and decode workers Signed-off-by: Cam Quilici * docs(tokenspeed): document DEP and TEP, KV cache offload and the P/D warmup skip Signed-off-by: Cam Quilici * test(tokenspeed): update launch snapshots for the P/D warmup skip Signed-off-by: Cam Quilici * refactor(backends): share the Mooncake worker environment through Backend Signed-off-by: Cam Quilici * refactor(tokenspeed): match existing worker launch selection * fix(tokenspeed): pass mapping-valued args as JSON Signed-off-by: Cam Quilici --------- Signed-off-by: Cam Quilici --- README.md | 2 +- docs/engines.md | 9 + docs/schema-reference.md | 14 +- docs/schema/recipe.schema.json | 30 ++- docs/tokenspeed.md | 173 ++++++++++++++ examples/README.md | 2 + examples/tokenspeed/dynamo-agg.yaml | 51 ++++ examples/tokenspeed/dynamo-disagg.yaml | 63 +++++ examples/tokenspeed/smg-agg.yaml | 48 ++++ examples/tokenspeed/smg-disagg.yaml | 59 +++++ mkdocs.yml | 1 + src/srtctl/backends/AGENTS.md | 1 + src/srtctl/backends/__init__.py | 9 +- src/srtctl/backends/base.py | 28 ++- src/srtctl/backends/sglang.py | 27 --- src/srtctl/backends/tokenspeed.py | 209 +++++++++++++++++ src/srtctl/backends/vllm.py | 33 +-- src/srtctl/core/AGENTS.md | 2 +- src/srtctl/core/fingerprint.py | 1 + src/srtctl/core/roles.py | 1 + src/srtctl/core/schema.py | 18 +- src/srtctl/core/schema_docs.py | 4 +- src/srtctl/core/topology.py | 2 + src/srtctl/ports.py | 6 + .../launch/tokenspeed__dynamo-agg.txt | 202 ++++++++++++++++ .../launch/tokenspeed__dynamo-disagg.txt | 204 ++++++++++++++++ .../snapshots/launch/tokenspeed__smg-agg.txt | 172 ++++++++++++++ .../launch/tokenspeed__smg-disagg.txt | 174 ++++++++++++++ tests/test_backend_protocol.py | 11 +- tests/test_e2e.py | 19 +- tests/test_port_allocator.py | 4 + tests/test_roles.py | 2 +- tests/test_tokenspeed.py | 218 ++++++++++++++++++ 33 files changed, 1718 insertions(+), 81 deletions(-) create mode 100644 docs/tokenspeed.md create mode 100644 examples/tokenspeed/dynamo-agg.yaml create mode 100644 examples/tokenspeed/dynamo-disagg.yaml create mode 100644 examples/tokenspeed/smg-agg.yaml create mode 100644 examples/tokenspeed/smg-disagg.yaml create mode 100644 src/srtctl/backends/tokenspeed.py create mode 100644 tests/snapshots/launch/tokenspeed__dynamo-agg.txt create mode 100644 tests/snapshots/launch/tokenspeed__dynamo-disagg.txt create mode 100644 tests/snapshots/launch/tokenspeed__smg-agg.txt create mode 100644 tests/snapshots/launch/tokenspeed__smg-disagg.txt create mode 100644 tests/test_tokenspeed.py diff --git a/README.md b/README.md index 194bef9d1..c1c30034e 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ # srtctl -Command-line tool for distributed LLM inference benchmarks on SLURM clusters using SGLang, vLLM, TensorRT LLM, TileRT and AMD's ATOM. Replace complex shell scripts and 50+ CLI flags with a declarative `schema: 2` YAML recipe: `engine:` names the engine, `roles:` describes each worker role, and `services:` covers everything launched next to the workers. +Command-line tool for distributed LLM inference benchmarks on SLURM clusters using SGLang, vLLM, TensorRT LLM, TileRT, TokenSpeed and AMD's ATOM. Replace complex shell scripts and 50+ CLI flags with a declarative `schema: 2` YAML recipe: `engine:` names the engine, `roles:` describes each worker role, and `services:` covers everything launched next to the workers. **Documentation:** https://nvidia.github.io/srt-slurm/ (source in [docs/](docs/README.md); agents: [llms.txt](https://nvidia.github.io/srt-slurm/llms.txt)) diff --git a/docs/engines.md b/docs/engines.md index aac8ebb46..68a333f5b 100644 --- a/docs/engines.md +++ b/docs/engines.md @@ -240,6 +240,15 @@ warning in the job log. **Other TRT-LLM launch facts**: TRT-LLM supports prefill, decode, and aggregated roles, uses MPI-style launching (one srun per endpoint with all of its nodes) through `trtllm-llmapi-launch`, and sets `TRTLLM_EPLB_SHM_NAME` to a unique UUID per endpoint. +## TokenSpeed with Dynamo + +Use `engine: tokenspeed` with `frontend.type: dynamo` to launch +`python3 -m dynamo.tokenspeed` workers, aggregated or prefill/decode. TokenSpeed flags +go under `roles..args`; srtctl sets the model, host, ports and multi-node +rendezvous. See [TokenSpeed](tokenspeed.md) and the +[aggregated](https://github.com/NVIDIA/srt-slurm/blob/main/examples/tokenspeed/dynamo-agg.yaml) and +[prefill/decode](https://github.com/NVIDIA/srt-slurm/blob/main/examples/tokenspeed/dynamo-disagg.yaml) recipes. + ## ATOM with AToMesh Use `engine: atom` with `frontend.type: atomesh` to launch native diff --git a/docs/schema-reference.md b/docs/schema-reference.md index 08a0d7dfe..711aa4e59 100644 --- a/docs/schema-reference.md +++ b/docs/schema-reference.md @@ -15,7 +15,7 @@ Top-level keys of a recipe YAML. | `resources` | [ResourceConfig](#resourceconfig) | required | GPU type, GPUs per node, and allocation knobs. The worker topology is `roles`. | | `schema` | one of `2` | required | Recipe schema version. Every recipe declares `schema: 2`; a recipe without it is the pre-2.0 layout and does not load (see [legacy-v1.md](legacy-v1.md) and `srtctl migrate`). | | `slurm` | [SlurmConfig](#slurmconfig) | `SlurmConfig()` | Slurm account, partition, and time limit; unset values come from srtslurm.yaml. | -| `engine` | str \| mapping | optional when every role sets `engine` | The engine type (`atom`, `sglang`, `tilert`, `trtllm`, `vllm`, `mocker`) as a string, or a mapping with `type` plus the engine-wide knobs listed under [Engine types](#engine-types). | +| `engine` | str \| mapping | optional when every role sets `engine` | The engine type (`atom`, `sglang`, `tilert`, `tokenspeed`, `trtllm`, `vllm`, `mocker`) as a string, or a mapping with `type` plus the engine-wide knobs listed under [Engine types](#engine-types). | | `roles` | dict[str, [RoleConfig](#roleconfig)] | `{}` | One block per worker role (`prefill`, `decode`, `agg`): nodes, workers, GPUs, env, engine args. | | `frontend` | [FrontendConfig](#frontendconfig) | `FrontendConfig()` | The HTTP entry point in front of the workers (Dynamo frontend, router, nginx) and where it runs. | | `dynamo` | [DynamoConfig](#dynamoconfig) | `DynamoConfig()` | Which Dynamo to install, its request/event planes, and native sidecar mode. | @@ -44,7 +44,7 @@ Three vocabularies carry the engine, the topology, and the placement. `engine` a ### engine -`engine: ` or `engine: {type: , ...}`. `type` is one of `atom`, `sglang`, `tilert`, `trtllm`, `vllm`, `mocker`; the remaining keys are that engine's knobs, listed under [Engine types](#engine-types). +`engine: ` or `engine: {type: , ...}`. `type` is one of `atom`, `sglang`, `tilert`, `tokenspeed`, `trtllm`, `vllm`, `mocker`; the remaining keys are that engine's knobs, listed under [Engine types](#engine-types). Use either one top-level engine or an explicit engine on every role, never both. Role engines do not inherit options from each other. ### roles @@ -699,6 +699,16 @@ Launch TileRT's decode server with recipe-owned model and transport settings. | `type` | one of `'tilert'` | `'tilert'` | Engine type discriminator. | | `served_model_name` | str \| None | `None` | Model name the server reports to clients; unset uses the default served name. | +### TokenSpeedBackend + +`engine.type: tokenspeed` + +Launch TokenSpeed workers behind the Dynamo frontend or a static router such as SMG. + +| Key | Type | Default | Description | +|---|---|---|---| +| `type` | one of `'tokenspeed'` | `'tokenspeed'` | Engine type discriminator. | + ### TRTLLMBackend `engine.type: trtllm` diff --git a/docs/schema/recipe.schema.json b/docs/schema/recipe.schema.json index aa9d3da5c..7c3d342a7 100644 --- a/docs/schema/recipe.schema.json +++ b/docs/schema/recipe.schema.json @@ -1805,6 +1805,7 @@ "atom", "sglang", "tilert", + "tokenspeed", "trtllm", "vllm", "mocker" @@ -1819,6 +1820,9 @@ { "$ref": "#/$defs/TileRTBackend" }, + { + "$ref": "#/$defs/TokenSpeedBackend" + }, { "$ref": "#/$defs/TRTLLMBackend" }, @@ -1832,7 +1836,7 @@ "type": "null" } ], - "description": "The engine type (`atom`, `sglang`, `tilert`, `trtllm`, `vllm`, `mocker`) as a string, or a mapping with `type` plus the engine-wide knobs listed under [Engine types](#engine-types).", + "description": "The engine type (`atom`, `sglang`, `tilert`, `tokenspeed`, `trtllm`, `vllm`, `mocker`) as a string, or a mapping with `type` plus the engine-wide knobs listed under [Engine types](#engine-types).", "default": null }, "roles": { @@ -2021,6 +2025,7 @@ "atom", "sglang", "tilert", + "tokenspeed", "trtllm", "vllm", "mocker" @@ -2035,6 +2040,9 @@ { "$ref": "#/$defs/TileRTBackend" }, + { + "$ref": "#/$defs/TokenSpeedBackend" + }, { "$ref": "#/$defs/TRTLLMBackend" }, @@ -2048,7 +2056,7 @@ "type": "null" } ], - "description": "The engine type (`atom`, `sglang`, `tilert`, `trtllm`, `vllm`, `mocker`) as a string, or a mapping with `type` plus the engine-wide knobs listed under [Engine types](#engine-types).", + "description": "The engine type (`atom`, `sglang`, `tilert`, `tokenspeed`, `trtllm`, `vllm`, `mocker`) as a string, or a mapping with `type` plus the engine-wide knobs listed under [Engine types](#engine-types).", "default": null }, "roles": { @@ -2452,6 +2460,7 @@ "atom", "sglang", "tilert", + "tokenspeed", "trtllm", "vllm", "mocker" @@ -2466,6 +2475,9 @@ { "$ref": "#/$defs/TileRTBackend" }, + { + "$ref": "#/$defs/TokenSpeedBackend" + }, { "$ref": "#/$defs/TRTLLMBackend" }, @@ -3534,6 +3546,20 @@ }, "description": "Launch TileRT's decode server with recipe-owned model and transport settings." }, + "TokenSpeedBackend": { + "type": "object", + "additionalProperties": false, + "properties": { + "type": { + "enum": [ + "tokenspeed" + ], + "description": "Engine type discriminator.", + "default": "tokenspeed" + } + }, + "description": "Launch TokenSpeed workers behind the Dynamo frontend or a static router such as SMG." + }, "VLLMBackend": { "type": "object", "additionalProperties": false, diff --git a/docs/tokenspeed.md b/docs/tokenspeed.md new file mode 100644 index 000000000..1005c89f9 --- /dev/null +++ b/docs/tokenspeed.md @@ -0,0 +1,173 @@ +# TokenSpeed + +Use `engine: tokenspeed` with `frontend.type: dynamo` or `frontend.type: smg`. + +- **Dynamo.** Every worker runs Dynamo's TokenSpeed backend, `python3 -m dynamo.tokenspeed`, + which registers with the Dynamo frontend like the other Dynamo engines. +- **SMG.** Every worker runs TokenSpeed's gRPC engine, `python3 -m smg_grpc_servicer.tokenspeed`, + and srtctl hands SMG its `grpc://` URL. This is the pair `ts serve` runs on one node: the + same engine module with `--host`/`--port`, and `smg launch --worker-urls grpc://` + ([`_proc.py`](https://github.com/lightseekorg/tokenspeed/blob/22251686ff26a2b2f495263b348db80980b8ac9e/python/tokenspeed/cli/_proc.py#L39-L104)). + +Recipes: + +- Dynamo: [aggregated](https://github.com/NVIDIA/srt-slurm/blob/main/examples/tokenspeed/dynamo-agg.yaml), [prefill/decode](https://github.com/NVIDIA/srt-slurm/blob/main/examples/tokenspeed/dynamo-disagg.yaml) +- SMG: [aggregated](https://github.com/NVIDIA/srt-slurm/blob/main/examples/tokenspeed/smg-agg.yaml), [prefill/decode](https://github.com/NVIDIA/srt-slurm/blob/main/examples/tokenspeed/smg-disagg.yaml) + +Configuration loading rejects the engine-specific routers (`sglang-router`, +`vllm-router`, ...), `dynamo.sidecar`, and `roles..kv_events`. + +The backend follows +[TokenSpeed](https://github.com/lightseekorg/tokenspeed/tree/22251686ff26a2b2f495263b348db80980b8ac9e), +Dynamo's +[`dynamo.tokenspeed`](https://github.com/ai-dynamo/dynamo/tree/7778c8d0cddb2a1ab7b2782c92cf97b30b2a5dcd/components/src/dynamo/tokenspeed) +and SMG [v1.11.0](https://github.com/smg-project/smg/tree/3be823a700fabaff3add8a390cf78f163479d686). + +## Image + +The worker image must contain TokenSpeed and an `ai-dynamo` build with +`dynamo.tokenspeed`. Aggregated serving needs Dynamo 1.2.0 or newer; prefill/decode +needs a build that includes [ai-dynamo/dynamo#9237](https://github.com/ai-dynamo/dynamo/pull/9237), +which Dynamo 1.5.0 does not; the PyPI nightlies from `1.6.0.dev20260922` on do. Either +let srtctl install one at job start (`dynamo.source.pypi: "1.6.0.dev20260922"`) into an +image that ships TokenSpeed, or build an image with Dynamo's +[TokenSpeed Dockerfile](https://github.com/ai-dynamo/dynamo/blob/7778c8d0cddb2a1ab7b2782c92cf97b30b2a5dcd/recipes/kimi-k2.5/tokenspeed/agg/nvidia/Dockerfile) +from a TokenSpeed runner and the Dynamo checkout. The examples set +`dynamo.install: false` because their image already ships Dynamo. + +Behind SMG the image needs TokenSpeed only. Installing it pulls in the gRPC engine +(`tokenspeed-smg-grpc-servicer`) and the SMG build it pins (`tokenspeed-smg`, which +provides the `smg` command) +([`pyproject.toml`](https://github.com/lightseekorg/tokenspeed/blob/22251686ff26a2b2f495263b348db80980b8ac9e/python/pyproject.toml#L75-L77)), so SMG runs in the model +container and needs no `frontend.container_image`. + +## Arguments + +`roles..args` are TokenSpeed +[server flags](https://github.com/lightseekorg/tokenspeed/blob/22251686ff26a2b2f495263b348db80980b8ac9e/docs/configuration/server.md), +such as `tensor-parallel-size`, `data-parallel-size`, `enable-expert-parallel`, +`max-model-len` and `gpu-memory-utilization`. Set the parallelism to match the +role's `gpus`. A role's `served-model-name` is the name every worker serves; without +one, workers serve the model directory name. + +srtctl sets these flags on every worker, under either frontend: + +| Flag | Value | +| --- | --- | +| `--model` | `/model`, the staged model path, or the Hugging Face ID | +| `--served-model-name` | the role's `served-model-name`, else the model directory name | +| `--host` | the worker's own IP on `network_interface` | +| `--port` | an allocated port, 10000 and up, 1024 apart on a node; behind SMG, the leader's gRPC port | +| `--dist-init-addr` | the endpoint leader's IP and an allocated port | +| `--nnodes`, `--node-rank` | the endpoint's node count and this node's rank | +| `--disaggregation-mode` | `prefill` or `decode`, on prefill/decode workers | +| `--disaggregation-bootstrap-port` | an allocated port, on prefill workers | + +TokenSpeed derives several listeners from two of these. Its +[`PortArgs.init_new`](https://github.com/lightseekorg/tokenspeed/blob/22251686ff26a2b2f495263b348db80980b8ac9e/python/tokenspeed/runtime/utils/server_args.py#L2513-L2602) +picks the NCCL port 100 to 1000 above `--port`, and binds a control-plane cluster +on the leader that starts at the `--dist-init-addr` port: six ports, plus one per +attention-DP rank +([`data_parallel_controller.py`](https://github.com/lightseekorg/tokenspeed/blob/22251686ff26a2b2f495263b348db80980b8ac9e/python/tokenspeed/runtime/engine/data_parallel_controller.py#L293-L294)). +srtctl reserves a block of six ports plus the endpoint's GPU count for that cluster, +and spaces `--port` 1024 apart, so workers sharing a node do not collide. A +follower node serves its health check on `--host` and `--port` +([`engine.py`](https://github.com/lightseekorg/tokenspeed/blob/22251686ff26a2b2f495263b348db80980b8ac9e/python/tokenspeed/runtime/entrypoints/engine.py#L636-L642)). +Each node runs its own single-node srun step, so srtctl passes `--nnodes`, +`--node-rank` and `--dist-init-addr` instead of relying on TokenSpeed's +[Slurm step detection](https://github.com/lightseekorg/tokenspeed/blob/22251686ff26a2b2f495263b348db80980b8ac9e/docs/serving/parallelism.md#under-a-launcher). + +## Parallelism + +Parallelism is in the role's args and needs nothing else from srtctl. Attention DP with +expert parallelism (DEP) is `data-parallel-size: ` with `enable-expert-parallel: true`; +tensor parallelism with expert parallelism (TEP) is `tensor-parallel-size: ` with +`enable-expert-parallel: true`. A worker that spans nodes sizes these over all of its GPUs: +srtctl passes its node count and rank, and TokenSpeed runs `world size / nnodes` ranks +on each node +([parallelism.md](https://github.com/lightseekorg/tokenspeed/blob/22251686ff26a2b2f495263b348db80980b8ac9e/docs/serving/parallelism.md#multi-node)). +Behind Dynamo, prefill and decode workers need attention DP 1 (see below); behind SMG a +decode worker can be DEP. + +## KV cache offload + +Unless a role sets `disable-kvstore`, every TokenSpeed rank keeps evicted KV blocks in a +pinned host-memory KVStore (L2). Each rank allocates its own: `kvstore-size` sets it in GB, +otherwise it is `kvstore-ratio` (default 2.0) times the rank's GPU KV pool, so size it to +the node's memory. + +`kvstore-storage-backend: mooncake` adds Mooncake Store under it as L3 +([server.md](https://github.com/lightseekorg/tokenspeed/blob/22251686ff26a2b2f495263b348db80980b8ac9e/docs/configuration/server.md#host-l2-and-mooncake-store-l3)). +Declare a `mooncake-master` service: srtctl starts the master on the infra node and sets +`MOONCAKE_MASTER`, `MOONCAKE_TE_META_DATA_SERVER` and `MOONCAKE_LOCAL_HOSTNAME` on every +worker, which TokenSpeed's Mooncake client reads +([`mooncake.py`](https://github.com/lightseekorg/tokenspeed/blob/22251686ff26a2b2f495263b348db80980b8ac9e/python/tokenspeed/runtime/cache/l3/mooncake.py#L94-L140)). +Set `MOONCAKE_GLOBAL_SEGMENT_SIZE`, `MOONCAKE_PROTOCOL` or `MOONCAKE_DEVICE` in the role's +env. A TokenSpeed rank stops when Mooncake rejects a write for lack of space, so size the +segments for the workload and let the master evict before it fills (the service's `args`, +for example `--eviction_high_watermark_ratio=0.7`). The service runs `mooncake_master` from +TokenSpeed's `tokenspeed-mooncake` dependency in the model image; give it a `preamble` that +installs TokenSpeed when the image gets it from `setup_script`, which services do not run. + +```yaml +services: + - name: mooncake-master + type: mooncake-master +roles: + agg: + args: + kvstore-size: 64 + kvstore-storage-backend: mooncake +``` + +## Prefill/decode + +TokenSpeed moves the KV cache with Mooncake. A prefill worker advertises its +bootstrap server at `--host` and the bootstrap port +([`runtime_disaggregated_endpoint`](https://github.com/ai-dynamo/dynamo/blob/7778c8d0cddb2a1ab7b2782c92cf97b30b2a5dcd/components/src/dynamo/tokenspeed/disagg.py#L71-L105)); +the Dynamo frontend passes it to the decode worker with each request. Prefill +workers register as the `prefill` component and decode workers as `backend` +([`args.py`](https://github.com/ai-dynamo/dynamo/blob/7778c8d0cddb2a1ab7b2782c92cf97b30b2a5dcd/components/src/dynamo/tokenspeed/args.py#L105-L108)). + +Dynamo checks the remaining requirements when the worker starts +([`validate_disagg_compatibility`](https://github.com/ai-dynamo/dynamo/blob/7778c8d0cddb2a1ab7b2782c92cf97b30b2a5dcd/components/src/dynamo/tokenspeed/disagg.py#L50-L68)): +the transfer backend is Mooncake, attention DP is 1, and `prefix-granularity` is +positive (TokenSpeed's default is 64). Scale prefill or decode with more workers. +Set `disaggregation-ib-device` in the role's args when Mooncake does not pick the +right RDMA device. + +## SMG + +The engine's `--host` and `--port` are its gRPC listener +([`server.py`](https://github.com/smg-project/smg/blob/3be823a700fabaff3add8a390cf78f163479d686/grpc_servicer/smg_grpc_servicer/tokenspeed/server.py)), so srtctl +advertises the leader of each worker as `grpc://:` +(`TokenSpeedBackend.is_grpc_mode` is always true: TokenSpeed has no HTTP-only engine +server). A follower node of a multi-node worker runs the same module, as `ts serve` does +on a non-zero node rank +([`serve_smg.py`](https://github.com/lightseekorg/tokenspeed/blob/22251686ff26a2b2f495263b348db80980b8ac9e/python/tokenspeed/cli/serve_smg.py#L873-L879)), and is not routed. +Before SMG starts, srtctl waits for every gRPC port to accept connections. + +For prefill/decode srtctl passes `--pd-disaggregation --prefill grpc:// + --decode grpc://`. SMG's gRPC P/D table gives TokenSpeed +parallel dispatch with a KV bootstrap room +([`pd_protocol.rs`](https://github.com/smg-project/smg/blob/3be823a700fabaff3add8a390cf78f163479d686/model_gateway/src/routers/grpc/common/stages/pd_protocol.rs#L63-L85)): +it mints one room per request and sends both workers the prefill worker's host and +bootstrap port +([`helpers.rs`](https://github.com/smg-project/smg/blob/3be823a700fabaff3add8a390cf78f163479d686/model_gateway/src/routers/grpc/common/stages/helpers.rs#L764-L824)), +which the engine hands to Mooncake +([`servicer.py`](https://github.com/smg-project/smg/blob/3be823a700fabaff3add8a390cf78f163479d686/grpc_servicer/smg_grpc_servicer/tokenspeed/servicer.py#L1031-L1071)). +The prefill engine serves that bootstrap port +([`async_llm.py`](https://github.com/lightseekorg/tokenspeed/blob/22251686ff26a2b2f495263b348db80980b8ac9e/python/tokenspeed/runtime/engine/async_llm.py#L245-L253)). +srtctl sets `TOKENSPEED_SKIP_GRPC_WARMUP=1` on prefill and decode workers: the engine's +startup warmup is a generate without a bootstrap room +([`server.py`](https://github.com/smg-project/smg/blob/3be823a700fabaff3add8a390cf78f163479d686/grpc_servicer/smg_grpc_servicer/tokenspeed/server.py#L211-L282)), +which a prefill engine cannot finish without a decode peer and an attention-DP decode engine +fails to place on a DP rank +([`data_parallel_controller.py`](https://github.com/lightseekorg/tokenspeed/blob/22251686ff26a2b2f495263b348db80980b8ac9e/python/tokenspeed/runtime/engine/data_parallel_controller.py#L425-L430)). + +In gRPC mode SMG tokenizes and applies the chat template itself; pass +`tool-call-parser` / `reasoning-parser` in `frontend.args` when the model needs them. +The gRPC engine starts no HTTP listener, so these workers serve no Prometheus +`/metrics`; SMG's own metrics stay on its Prometheus port. TokenSpeed's ZMQ engine +(`ts serve --headless`) is not used: SMG reaches a ZMQ engine only over local `ipc://` +sockets. diff --git a/examples/README.md b/examples/README.md index ad55b215e..8b7bcb0b1 100644 --- a/examples/README.md +++ b/examples/README.md @@ -9,6 +9,7 @@ Small starting points, one per frontend and topology. The matrix examples serve | SGLang | `sglang/dynamo-agg.yaml`, `sglang/dynamo-disagg.yaml` | `sglang/sglang-router-agg.yaml`, `sglang/sglang-router-disagg.yaml`, `sglang/smg-disagg.yaml` (SMG) | `sglang/sglang-direct-agg.yaml` | | vLLM | `vllm/dynamo-agg.yaml`, `vllm/dynamo-disagg.yaml` | `vllm/vllm-router-agg.yaml`, `vllm/vllm-router-disagg.yaml`, `vllm/vllm-router-moriio-disagg.yaml` (ROCm, MoRI-IO discovery), `vllm/smg-agg.yaml`, `vllm/smg-dep16.yaml`, `vllm/smg-disagg-grpc.yaml` (SMG; P/D over gRPC with NIXL) | `vllm/vllm-direct-agg.yaml` | | TRT-LLM | `trtllm/dynamo-agg.yaml`, `trtllm/dynamo-disagg.yaml`, `trtllm/dynamo-disagg-multinode.yaml` (TP16 decode across two 8-GPU nodes: one srun, one host per rank under `--distribution=arbitrary`) | `trtllm/trtllm-serve-disagg.yaml`, `trtllm/smg-agg.yaml` (SMG) | `trtllm/trtllm-serve-agg.yaml` | +| TokenSpeed | `tokenspeed/dynamo-agg.yaml`, `tokenspeed/dynamo-disagg.yaml` | `tokenspeed/smg-agg.yaml`, `tokenspeed/smg-disagg.yaml` (SMG, gRPC engines) | | | Mocker | `mocker/dynamo-agg.yaml` | | | - **Dynamo frontend**: workers register with etcd and the Dynamo frontend routes (KV-aware here); the request plane is tcp and NATS is not started unless a plane asks for it. Dynamo is installed at job start via `dynamo.source` (`pypi:` here) unless the container ships it (`dynamo.install: false`, as the TRT-LLM examples do). @@ -57,6 +58,7 @@ containers: sglang-rocm: /path/to/sglang-rocm.sqsh # SGLang ROCm image, for features/amd-power-telemetry.yaml vllm: /path/to/vllm.sqsh # vLLM image with the vllm-router executable trtllm: /path/to/tensorrtllm-runtime.sqsh # Dynamo TRT-LLM runtime image (ships ai-dynamo and trtllm-serve) + tokenspeed: /path/to/dynamo-tokenspeed.sqsh # TokenSpeed image with ai-dynamo's dynamo.tokenspeed (see ../docs/tokenspeed.md) dynamo-vllm: /path/to/vllm-runtime.sqsh # Dynamo vLLM runtime image (ships ai-dynamo and gpu_memory_service), for features/vllm-failover.yaml vllm-lmcache: /path/to/vllm-lmcache.sqsh # vLLM image with LMCache installed, for features/lmcache-server.yaml and -disagg.yaml sglang-lmcache: /path/to/sglang-lmcache.sqsh # SGLang image with LMCache installed, for features/lmcache-server-sglang.yaml diff --git a/examples/tokenspeed/dynamo-agg.yaml b/examples/tokenspeed/dynamo-agg.yaml new file mode 100644 index 000000000..9c0b62a72 --- /dev/null +++ b/examples/tokenspeed/dynamo-agg.yaml @@ -0,0 +1,51 @@ +# TokenSpeed aggregated behind the Dynamo frontend with KV-aware routing. +# +# Two TP1 `dynamo.tokenspeed` workers on one node register with etcd; the Dynamo +# frontend routes with `--router-mode kv`. The `tokenspeed` container must ship +# TokenSpeed and Dynamo's TokenSpeed backend (`python3 -m dynamo.tokenspeed`), so +# Dynamo is not installed at job start. See docs/tokenspeed.md. +# +# Aliases (`qwen3-0.6b`, `tokenspeed`) resolve through srtslurm.yaml; see examples/README.md. + +schema: 2 +name: "qwen3-0.6b-tokenspeed-dynamo-agg" + +slurm: + time_limit: "00:30:00" + +model: + path: "qwen3-0.6b" + container: "tokenspeed" + precision: "bf16" + +resources: + gpu_type: "h100" + gpus_per_node: 8 +dynamo: + install: false + +frontend: + type: dynamo + enable_multiple_frontends: false + args: + router-mode: "kv" + +engine: tokenspeed +roles: + agg: + nodes: 1 + workers: 2 + gpus: 1 + env: + PYTHONUNBUFFERED: "1" + args: + served-model-name: "Qwen/Qwen3-0.6B" + tensor-parallel-size: 1 + gpu-memory-utilization: 0.5 + max-model-len: 4096 + +benchmark: + type: "sa-bench" + isl: 128 + osl: 128 + concurrencies: "4x8" diff --git a/examples/tokenspeed/dynamo-disagg.yaml b/examples/tokenspeed/dynamo-disagg.yaml new file mode 100644 index 000000000..9105c6c5a --- /dev/null +++ b/examples/tokenspeed/dynamo-disagg.yaml @@ -0,0 +1,63 @@ +# TokenSpeed disaggregated prefill/decode behind the Dynamo frontend. +# +# One TP1 prefill worker and one TP1 decode worker sharing a node (`nodes: colocate` +# places decode on the prefill node's spare GPUs). Both run `dynamo.tokenspeed`, which +# moves the KV cache with Mooncake; the Dynamo frontend dispatches prefill to the +# prefill pool. srtctl adds `--disaggregation-mode` and the prefill bootstrap port. +# P/D needs a Dynamo build that includes ai-dynamo/dynamo#9237; see docs/tokenspeed.md. +# +# Aliases (`qwen3-0.6b`, `tokenspeed`) resolve through srtslurm.yaml; see examples/README.md. + +schema: 2 +name: "qwen3-0.6b-tokenspeed-dynamo-disagg" + +slurm: + time_limit: "00:30:00" + +model: + path: "qwen3-0.6b" + container: "tokenspeed" + precision: "bf16" + +resources: + gpu_type: "h100" + gpus_per_node: 8 +dynamo: + install: false + +frontend: + type: dynamo + enable_multiple_frontends: false + args: + router-mode: "kv" + +engine: tokenspeed +roles: + prefill: + nodes: 1 + workers: 1 + gpus: 1 + env: + PYTHONUNBUFFERED: "1" + args: + served-model-name: "Qwen/Qwen3-0.6B" + tensor-parallel-size: 1 + gpu-memory-utilization: 0.5 + max-model-len: 4096 + decode: + nodes: colocate + workers: 1 + gpus: 1 + env: + PYTHONUNBUFFERED: "1" + args: + served-model-name: "Qwen/Qwen3-0.6B" + tensor-parallel-size: 1 + gpu-memory-utilization: 0.5 + max-model-len: 4096 + +benchmark: + type: "sa-bench" + isl: 128 + osl: 128 + concurrencies: "4x8" diff --git a/examples/tokenspeed/smg-agg.yaml b/examples/tokenspeed/smg-agg.yaml new file mode 100644 index 000000000..2d45f821a --- /dev/null +++ b/examples/tokenspeed/smg-agg.yaml @@ -0,0 +1,48 @@ +# TokenSpeed aggregated behind the Shepherd Model Gateway (SMG). +# +# Two TP1 workers on one node, each TokenSpeed's gRPC engine +# (`python3 -m smg_grpc_servicer.tokenspeed`), the engine `ts serve` runs behind its +# bundled SMG. srtctl advertises them to SMG as grpc:// URLs. The `tokenspeed` +# container must ship TokenSpeed, which installs the gRPC servicer and the `smg` +# command it pins; SMG runs in the same container. See docs/tokenspeed.md. +# +# Aliases (`qwen3-0.6b`, `tokenspeed`) resolve through srtslurm.yaml; see examples/README.md. + +schema: 2 +name: "qwen3-0.6b-tokenspeed-smg-agg" + +slurm: + time_limit: "00:30:00" + +model: + path: "qwen3-0.6b" + container: "tokenspeed" + precision: "bf16" + +resources: + gpu_type: "h100" + gpus_per_node: 8 + +frontend: + type: smg + enable_multiple_frontends: false + +engine: tokenspeed +roles: + agg: + nodes: 1 + workers: 2 + gpus: 1 + env: + PYTHONUNBUFFERED: "1" + args: + served-model-name: "Qwen/Qwen3-0.6B" + tensor-parallel-size: 1 + gpu-memory-utilization: 0.5 + max-model-len: 4096 + +benchmark: + type: "sa-bench" + isl: 128 + osl: 128 + concurrencies: "4x8" diff --git a/examples/tokenspeed/smg-disagg.yaml b/examples/tokenspeed/smg-disagg.yaml new file mode 100644 index 000000000..cf2a395ad --- /dev/null +++ b/examples/tokenspeed/smg-disagg.yaml @@ -0,0 +1,59 @@ +# TokenSpeed disaggregated prefill/decode behind the Shepherd Model Gateway (SMG). +# +# One TP1 prefill and one TP1 decode gRPC engine sharing a node (`nodes: colocate`). +# srtctl passes SMG `--prefill grpc:// --decode grpc://`; +# SMG mints a bootstrap room per request and sends both workers the prefill's Mooncake +# bootstrap address, and the decode worker pulls the KV cache from the prefill worker. +# srtctl adds `--disaggregation-mode` and the prefill bootstrap port. See docs/tokenspeed.md. +# +# Aliases (`qwen3-0.6b`, `tokenspeed`) resolve through srtslurm.yaml; see examples/README.md. + +schema: 2 +name: "qwen3-0.6b-tokenspeed-smg-disagg" + +slurm: + time_limit: "00:30:00" + +model: + path: "qwen3-0.6b" + container: "tokenspeed" + precision: "bf16" + +resources: + gpu_type: "h100" + gpus_per_node: 8 + +frontend: + type: smg + enable_multiple_frontends: false + +engine: tokenspeed +roles: + prefill: + nodes: 1 + workers: 1 + gpus: 1 + env: + PYTHONUNBUFFERED: "1" + args: + served-model-name: "Qwen/Qwen3-0.6B" + tensor-parallel-size: 1 + gpu-memory-utilization: 0.5 + max-model-len: 4096 + decode: + nodes: colocate + workers: 1 + gpus: 1 + env: + PYTHONUNBUFFERED: "1" + args: + served-model-name: "Qwen/Qwen3-0.6B" + tensor-parallel-size: 1 + gpu-memory-utilization: 0.5 + max-model-len: 4096 + +benchmark: + type: "sa-bench" + isl: 128 + osl: 128 + concurrencies: "4x8" diff --git a/mkdocs.yml b/mkdocs.yml index 64dd2b384..06020cf42 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -87,6 +87,7 @@ nav: - Workloads: - Native Sidecars: sidecars.md - TileRT: tilert.md + - TokenSpeed: tokenspeed.md - Shadow Engine Recovery: shadow-engine-recovery.md - SGLang Fast Engine Recovery: sglang-weight-cache.md - Miles RL Post-Training: miles.md diff --git a/src/srtctl/backends/AGENTS.md b/src/srtctl/backends/AGENTS.md index 61480f68a..556a48120 100644 --- a/src/srtctl/backends/AGENTS.md +++ b/src/srtctl/backends/AGENTS.md @@ -27,6 +27,7 @@ Rules for `src/srtctl/backends/`. Every consumer asks a backend through `Backend - **ATOM**: Native ROCm servers behind AToMesh, with one Slurm node per logical worker and allocator-owned Mooncake handshake ports - **SGLang**: Per-process srun launching, supports prefill/decode/aggregated modes - **TileRT**: Decode workers behind `tilert-router`, paired with an explicit prefill engine and image. +- **TokenSpeed**: `dynamo.tokenspeed` workers behind the Dynamo frontend, or the gRPC engine `smg_grpc_servicer.tokenspeed` behind SMG (`is_grpc_mode` is always true), aggregated or prefill/decode; allocator-owned `--port` scan base and a dist-init block sized for its control-plane cluster - **TRTLLM**: MPI-style launching (one srun per endpoint with all nodes), prefill/decode only - **vLLM**: Per-process srun launching, prefill/decode/aggregated, `per_node` DP; `frontend_type` selects Dynamo registration or a direct `vllm serve` server, and `_CONNECTOR_MAP` owns the KV connector table diff --git a/src/srtctl/backends/__init__.py b/src/srtctl/backends/__init__.py index c7bfa6393..a4544cace 100644 --- a/src/srtctl/backends/__init__.py +++ b/src/srtctl/backends/__init__.py @@ -14,11 +14,14 @@ from .mocker import MockerBackend from .sglang import MooncakeKVStoreConfig, SGLangBackend from .tilert import TileRTBackend +from .tokenspeed import TokenSpeedBackend from .trtllm import TRTLLMBackend from .vllm import VLLMBackend, VLLMFailoverConfig, VLLMMooncakeKVStoreConfig # Union type for all backend configs -BackendConfig = AtomBackend | SGLangBackend | TileRTBackend | TRTLLMBackend | VLLMBackend | MockerBackend +BackendConfig = ( + AtomBackend | SGLangBackend | TileRTBackend | TokenSpeedBackend | TRTLLMBackend | VLLMBackend | MockerBackend +) __all__ = [ # ATOM @@ -38,8 +41,10 @@ "TRTLLMBackend", # TileRT "TileRTBackend", - "VLLMBackend", + # TokenSpeed + "TokenSpeedBackend", # vLLM + "VLLMBackend", "VLLMFailoverConfig", "VLLMMooncakeKVStoreConfig", ] diff --git a/src/srtctl/backends/base.py b/src/srtctl/backends/base.py index 37bace7d9..e5150f131 100644 --- a/src/srtctl/backends/base.py +++ b/src/srtctl/backends/base.py @@ -13,7 +13,7 @@ from marshmallow import ValidationError, fields -from srtctl.ports import DYN_SYSTEM_PORT_BASE +from srtctl.ports import DYN_SYSTEM_PORT_BASE, MOONCAKE_HTTP_METADATA_PORT, MOONCAKE_MASTER_PORT if TYPE_CHECKING: from pathlib import Path @@ -37,6 +37,7 @@ class BackendType(str, Enum): MOCKER = "mocker" ATOM = "atom" TILERT = "tilert" + TOKENSPEED = "tokenspeed" @dataclass @@ -246,8 +247,29 @@ def get_process_environment(self, process: "Process") -> dict[str, str]: return {} def get_mooncake_worker_env(self, infra_node_ip: str, local_hostname: str) -> dict[str, str]: - """MOONCAKE_* environment for a worker; empty when mooncake_kv_store is None.""" - return {} + """MOONCAKE_* environment for a worker; empty when mooncake_kv_store is None. + + - MOONCAKE_LOCAL_HOSTNAME defaults to the worker's resolved IP, but the + user can override it in mooncake_kv_store.env if they need something + custom (e.g. a specific RDMA NIC IP). + - MOONCAKE_MASTER and MOONCAKE_TE_META_DATA_SERVER are always set by + srtslurm to point at the infra-node mooncake_master, and override any + user-supplied value (the user can't know the infra IP at config time). + + Args: + infra_node_ip: Resolved IP of the infra node where mooncake_master runs. + local_hostname: Resolved IP of the worker's own node, for peer-to-peer + transfers. Defaults to the worker's primary network interface IP. + """ + store = self.mooncake_kv_store + if store is None: + return {} + return { + "MOONCAKE_LOCAL_HOSTNAME": local_hostname, + **store.env, + "MOONCAKE_MASTER": f"{infra_node_ip}:{MOONCAKE_MASTER_PORT}", + "MOONCAKE_TE_META_DATA_SERVER": f"http://{infra_node_ip}:{MOONCAKE_HTTP_METADATA_PORT}/metadata", + } def get_failover_environment(self, process: "Process", job_id: str) -> dict[str, str]: """Shadow engine recovery environment for a worker; empty when failover is None.""" diff --git a/src/srtctl/backends/sglang.py b/src/srtctl/backends/sglang.py index 917c1f03f..233de79be 100644 --- a/src/srtctl/backends/sglang.py +++ b/src/srtctl/backends/sglang.py @@ -30,8 +30,6 @@ DYN_SYSTEM_PORT_BASE, GRPC_HTTP_PORTS, LMCACHE_SERVER_PORT, - MOONCAKE_HTTP_METADATA_PORT, - MOONCAKE_MASTER_PORT, NCCL_PORTS, ) @@ -155,31 +153,6 @@ def get_process_environment(self, process: "Process") -> dict[str, str]: return {} return {"LMCACHE_MP_HOST": "127.0.0.1", "LMCACHE_MP_PORT": str(LMCACHE_SERVER_PORT)} - def get_mooncake_worker_env(self, infra_node_ip: str, local_hostname: str) -> dict[str, str]: - """Get mooncake env vars to inject on a specific worker. - - Returns empty dict if mooncake_kv_store is not configured. Otherwise: - - MOONCAKE_LOCAL_HOSTNAME defaults to the worker's resolved IP, but the - user can override it in mooncake_kv_store.env if they need something - custom (e.g. a specific RDMA NIC IP). - - MOONCAKE_MASTER and MOONCAKE_TE_META_DATA_SERVER are always set by - srtslurm to point at the infra-node mooncake_master, and override any - user-supplied value (the user can't know the infra IP at config time). - - Args: - infra_node_ip: Resolved IP of the infra node where mooncake_master runs. - local_hostname: Resolved IP of the worker's own node, for peer-to-peer - transfers. Defaults to the worker's primary network interface IP. - """ - if self.mooncake_kv_store is None: - return {} - return { - "MOONCAKE_LOCAL_HOSTNAME": local_hostname, - **self.mooncake_kv_store.env, - "MOONCAKE_MASTER": f"{infra_node_ip}:{MOONCAKE_MASTER_PORT}", - "MOONCAKE_TE_META_DATA_SERVER": (f"http://{infra_node_ip}:{MOONCAKE_HTTP_METADATA_PORT}/metadata"), - } - def is_grpc_mode(self, mode: WorkerMode) -> bool: """Check if gRPC mode is enabled for a worker mode.""" config = self.get_config_for_mode(mode) diff --git a/src/srtctl/backends/tokenspeed.py b/src/srtctl/backends/tokenspeed.py new file mode 100644 index 000000000..000fee55e --- /dev/null +++ b/src/srtctl/backends/tokenspeed.py @@ -0,0 +1,209 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2026 SemiAnalysis LLC. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""TokenSpeed inference backend. + +Behind the Dynamo frontend each worker is ``python3 -m dynamo.tokenspeed``; behind a +static router (SMG) it is TokenSpeed's gRPC engine, ``python3 -m smg_grpc_servicer.tokenspeed``. +""" + +from __future__ import annotations + +import builtins +import json +from collections.abc import Mapping, Sequence +from dataclasses import field, replace +from pathlib import Path +from typing import TYPE_CHECKING, ClassVar, Literal + +from marshmallow import Schema +from marshmallow_dataclass import dataclass + +from srtctl.backends.base import Backend, BoundRolesField, RoleSettings, WorkerMode, role_args +from srtctl.backends.sglang import MooncakeKVStoreConfig, _config_to_cli_args +from srtctl.ports import DIST_INIT_PORTS, DYN_SYSTEM_PORT_BASE, TOKENSPEED_PORTS + +if TYPE_CHECKING: + from srtctl.core.runtime import RuntimeContext + from srtctl.core.schema import ProfilingConfig + from srtctl.core.topology import Endpoint, NodePortAllocator, Process + + +# TokenSpeed's leader binds a control-plane cluster from its --dist-init-addr port: +# the port itself, five ports above it, and one scheduler port per attention-DP rank. +# Attention DP never exceeds the endpoint's GPU count. +_RENDEZVOUS_PORTS = 6 + + +@dataclass(frozen=True) +class TokenSpeedBackend(Backend): + """Launch TokenSpeed workers behind the Dynamo frontend or a static router such as SMG. + + Example YAML: + engine: tokenspeed + frontend: + type: dynamo + roles: + agg: + args: + tensor-parallel-size: 1 + max-model-len: 4096 + """ + + # Engine type discriminator. + type: Literal["tokenspeed"] = "tokenspeed" + # The roles this engine runs (`roles.` of the recipe), bound by SrtConfig and + # never written on `engine:`. Per-role env and TokenSpeed CLI args are read from here. + roles: Mapping[str, RoleSettings] = field(default_factory=dict, metadata={"marshmallow_field": BoundRolesField()}) + # Set by a `mooncake-master` service. TokenSpeed's Mooncake Store L3 cache + # (`kvstore-storage-backend: mooncake`) reads the master from the MOONCAKE_* environment. + mooncake_kv_store: MooncakeKVStoreConfig | None = None + + Schema: ClassVar[builtins.type[Schema]] = Schema + + def get_process_environment(self, process: Process) -> dict[str, str]: + """Skip the gRPC engine's startup warmup on prefill and decode workers. + + The warmup is a generate without a KV bootstrap room, which a prefill/decode + engine cannot serve: prefill waits for a decode peer that never comes, and an + attention-DP decode engine fails dispatching it by room. ``dynamo.tokenspeed`` + does not read the variable. + """ + return {} if process.endpoint_mode == "agg" else {"TOKENSPEED_SKIP_GRPC_WARMUP": "1"} + + def is_grpc_mode(self, mode: WorkerMode) -> bool: + """A direct TokenSpeed worker is always the gRPC engine (``smg_grpc_servicer.tokenspeed``).""" + return True + + def get_served_model_name(self, default: str) -> str: + """A role's ``served-model-name``, else ``default``; srtctl always passes the result.""" + for mode in ("prefill", "agg", "decode"): + args = role_args(self.roles, mode) + name = args.get("served-model-name") or args.get("served_model_name") + if name: + return name + return default + + def allocate_endpoints( + self, + num_prefill: int, + num_decode: int, + num_agg: int, + gpus_per_prefill: int, + gpus_per_decode: int, + gpus_per_agg: int, + gpus_per_node: int, + available_nodes: Sequence[str], + spread_workers: bool = False, + ) -> list[Endpoint]: + from srtctl.core.topology import allocate_endpoints + + return allocate_endpoints( + num_prefill=num_prefill, + num_decode=num_decode, + num_agg=num_agg, + gpus_per_prefill=gpus_per_prefill, + gpus_per_decode=gpus_per_decode, + gpus_per_agg=gpus_per_agg, + gpus_per_node=gpus_per_node, + available_nodes=available_nodes, + spread_workers=spread_workers, + ) + + def endpoints_to_processes( + self, + endpoints: list[Endpoint], + base_sys_port: int = DYN_SYSTEM_PORT_BASE, + port_allocator: NodePortAllocator | None = None, + frontend_type: str = "dynamo", + dynamo_sidecar: bool = False, + ) -> list[Process]: + """Each process gets its own ``--port``; each endpoint a rendezvous block on its leader node. + + A direct worker's leader serves gRPC on that ``--port``, so it is also the + process's routable port (``http_port``). + """ + if dynamo_sidecar: + raise ValueError("TokenSpeed does not support Dynamo sidecars") + from srtctl.core.topology import endpoints_to_processes, port_allocator_for + from srtctl.frontends import get_frontend + + direct = get_frontend(frontend_type).worker_launch == "direct" + allocator = port_allocator_for(port_allocator, base_sys_port) + processes = endpoints_to_processes(endpoints, port_allocator=allocator) + dist_init_ports = { + (endpoint.mode, endpoint.index): allocator.next( + DIST_INIT_PORTS, endpoint.nodes[0], size=_RENDEZVOUS_PORTS + endpoint.total_gpus + ) + for endpoint in endpoints + } + result = [] + for process in processes: + port = allocator.next(TOKENSPEED_PORTS, process.node) + result.append( + replace( + process, + tokenspeed_port=port, + http_port=port if direct and process.http_port > 0 else process.http_port, + dist_init_port=dist_init_ports[(process.endpoint_mode, process.endpoint_index)], + ) + ) + return result + + def build_worker_command( + self, + process: Process, + endpoint_processes: list[Process], + runtime: RuntimeContext, + frontend_type: str = "dynamo", + nsys_prefix: list[str] | None = None, + dump_config_path: Path | None = None, + profiling: ProfilingConfig | None = None, + ) -> list[str]: + from srtctl.core.slurm import get_hostname_ip + from srtctl.frontends import get_frontend + + if process.tokenspeed_port is None or process.dist_init_port is None: + raise ValueError("build the topology with TokenSpeedBackend.endpoints_to_processes") + + mode = process.endpoint_mode + config = self.get_config_for_mode(mode) + for key in ("model", "model-path", "model_path", "served-model-name", "served_model_name"): + config.pop(key, None) + + endpoint_nodes = list(dict.fromkeys(p.node for p in endpoint_processes)) + leader_ip = get_hostname_ip(endpoint_nodes[0], runtime.network_interface) + frontend = get_frontend(frontend_type) + direct_workers = frontend.worker_launch == "direct" + python_module = "smg_grpc_servicer.tokenspeed" if direct_workers else "dynamo.tokenspeed" + cmd = [ + *(nsys_prefix or []), + "python3", + "-m", + python_module, + "--model", + runtime.worker_model_arg, + "--served-model-name", + self.get_served_model_name(runtime.model_path.name), + # A prefill worker advertises this host for the Mooncake bootstrap; a direct + # worker also binds its gRPC server on it. + "--host", + get_hostname_ip(process.node, runtime.network_interface), + "--port", + str(process.tokenspeed_port), + "--dist-init-addr", + f"{leader_ip}:{process.dist_init_port}", + "--nnodes", + str(len(endpoint_nodes)), + "--node-rank", + str(endpoint_nodes.index(process.node)), + ] + if mode != "agg": + cmd.extend(["--disaggregation-mode", mode]) + if mode == "prefill" and process.bootstrap_port is not None: + cmd.extend(["--disaggregation-bootstrap-port", str(process.bootstrap_port)]) + # TokenSpeed parses mapping-valued flags (kv-events-config, Mooncake extra config) as JSON. + config = {key: json.dumps(value) if isinstance(value, dict) else value for key, value in config.items()} + cmd.extend(_config_to_cli_args(config)) + return cmd diff --git a/src/srtctl/backends/vllm.py b/src/srtctl/backends/vllm.py index 097299b70..ecf0be3cb 100644 --- a/src/srtctl/backends/vllm.py +++ b/src/srtctl/backends/vllm.py @@ -36,7 +36,6 @@ KV_EVENTS_PORTS, KVBM_ZMQ_PORTS, LMCACHE_SERVER_PORT, - MOONCAKE_HTTP_METADATA_PORT, MOONCAKE_MASTER_PORT, MORIIO_HANDSHAKE_PORTS, MORIIO_NOTIFY_PORTS, @@ -532,29 +531,17 @@ def get_process_environment(self, process: Process) -> dict[str, str]: return env def get_mooncake_worker_env(self, infra_node_ip: str, local_hostname: str) -> dict[str, str]: - """Get mooncake env vars to inject on a specific vLLM worker. - - Returns ``{}`` when ``mooncake_kv_store`` is unset. Otherwise: - - - ``MOONCAKE_MASTER`` and ``MOONCAKE_TE_META_DATA_SERVER`` are always - stamped by srtslurm (the user can't know the infra IP at config time). - - ``MOONCAKE_LOCAL_HOSTNAME`` defaults to the worker's resolved IP for - multi-node peer transfers, but a value in ``mooncake_kv_store.env`` - wins (use this to pin to a specific RDMA NIC IP). - - ``MOONCAKE_CONFIG_PATH`` points to the JSON file srtslurm writes at - job start (mounted into the container at ``/logs``). vLLM's - ``MooncakeStoreConnector`` requires this — it does not read the - ``MOONCAKE_*`` env vars directly. + """The shared MOONCAKE_* environment plus ``MOONCAKE_CONFIG_PATH``. + + ``MOONCAKE_CONFIG_PATH`` points to the JSON file srtslurm writes at job + start (mounted into the container at ``/logs``). vLLM's + ``MooncakeStoreConnector`` requires this — it does not read the + ``MOONCAKE_*`` env vars directly. """ - if self.mooncake_kv_store is None: - return {} - return { - "MOONCAKE_LOCAL_HOSTNAME": local_hostname, - **self.mooncake_kv_store.env, - "MOONCAKE_MASTER": f"{infra_node_ip}:{MOONCAKE_MASTER_PORT}", - "MOONCAKE_TE_META_DATA_SERVER": (f"http://{infra_node_ip}:{MOONCAKE_HTTP_METADATA_PORT}/metadata"), - "MOONCAKE_CONFIG_PATH": MOONCAKE_STORE_CONFIG_CONTAINER_PATH, - } + env = super().get_mooncake_worker_env(infra_node_ip, local_hostname) + if env: + env["MOONCAKE_CONFIG_PATH"] = MOONCAKE_STORE_CONFIG_CONTAINER_PATH + return env def build_mooncake_store_config(self, infra_node_ip: str) -> dict[str, Any]: """Build the JSON payload for vLLM's ``MooncakeStoreConfig.load_from_env()``. diff --git a/src/srtctl/core/AGENTS.md b/src/srtctl/core/AGENTS.md index f3d32f45a..1652efe92 100644 --- a/src/srtctl/core/AGENTS.md +++ b/src/srtctl/core/AGENTS.md @@ -4,7 +4,7 @@ Rules for `src/srtctl/core/` (schema, runtime, topology, processes, health). Sta ## Ports -Fixed ports are constants in `src/srtctl/ports.py`. Every port a worker process binds is a `PortKind` in the same module (`name`, `base`, `stride`, `per_node`) and is handed out by `NodePortAllocator.next(kind, node, size)` in `src/srtctl/core/topology.py`: per-node counters for listeners only bound on that node (`HTTP_PORTS`, `BOOTSTRAP_PORTS`, `DP_RPC_PORTS`, `DIST_INIT_PORTS`) and global counters for side channels peers address (`SYS_PORTS`, `KV_EVENTS_PORTS`, `NIXL_PORTS`, `KVBM_ZMQ_PORTS`, `NCCL_PORTS`, `VLLM_SCAN_PORTS`, `MORIIO_HANDSHAKE_PORTS`, `MORIIO_NOTIFY_PORTS`, `TRTLLM_DIST_INIT_PORTS`, `SIDECAR_GRPC_PORTS`). `size > 1` reserves a block when the engine adds a rank offset to the port it is given (MoRI-IO adds the local DP and TP rank to its notify base). `endpoints_to_processes` allocates the generic kinds and stores them on `Process`; a backend adds its engine-specific kinds with `dataclasses.replace` on the way out (SGLang `nccl_port` and `dist_init_port`, vLLM `vllm_scan_port` or, for a discovery connector, `moriio_handshake_port` and `moriio_notify_port`, TRT-LLM `trtllm_dist_init_port`). `do_sweep` builds the one allocator per job, seeding the sidecar gRPC base from `dynamo.sidecar_port`. +Fixed ports are constants in `src/srtctl/ports.py`. Every port a worker process binds is a `PortKind` in the same module (`name`, `base`, `stride`, `per_node`) and is handed out by `NodePortAllocator.next(kind, node, size)` in `src/srtctl/core/topology.py`: per-node counters for listeners only bound on that node (`HTTP_PORTS`, `BOOTSTRAP_PORTS`, `DP_RPC_PORTS`, `DIST_INIT_PORTS`, `TOKENSPEED_PORTS`) and global counters for side channels peers address (`SYS_PORTS`, `KV_EVENTS_PORTS`, `NIXL_PORTS`, `KVBM_ZMQ_PORTS`, `NCCL_PORTS`, `VLLM_SCAN_PORTS`, `MORIIO_HANDSHAKE_PORTS`, `MORIIO_NOTIFY_PORTS`, `TRTLLM_DIST_INIT_PORTS`, `SIDECAR_GRPC_PORTS`). `size > 1` reserves a block when the engine adds a rank offset to the port it is given (MoRI-IO adds the local DP and TP rank to its notify base). `endpoints_to_processes` allocates the generic kinds and stores them on `Process`; a backend adds its engine-specific kinds with `dataclasses.replace` on the way out (SGLang `nccl_port` and `dist_init_port`, vLLM `vllm_scan_port` or, for a discovery connector, `moriio_handshake_port` and `moriio_notify_port`, TRT-LLM `trtllm_dist_init_port`, TokenSpeed `tokenspeed_port` and a `dist_init_port` block). `do_sweep` builds the one allocator per job, seeding the sidecar gRPC base from `dynamo.sidecar_port`. Rules that follow: a new listener is a new `PortKind` plus a `Process` field, allocated in the topology builder, never `some_base + (sys_port - DYN_SYSTEM_PORT_BASE)` or any other arithmetic on another port at command-build time. Consumers read the field; a `None` means the topology was built without that kind, and a consumer that needs it raises rather than guessing. A worker's own address is `get_hostname_ip(node, runtime.network_interface)`, not upstream's interface guess. Fixed scan bases such as `VLLM_PORT` exist only to keep co-located `get_open_port()` scans apart; if a connector allocates inside forked children, leave the base unset so the kernel assigns ports. Frontends bind `FRONTEND_PUBLIC_PORT` (8000), or `FRONTEND_INTERNAL_PORT` (8180) behind nginx when `enable_multiple_frontends` is set. Service ports are `ServiceKind` defaults overridden by `options`. `tests/test_port_allocator.py` asserts over every example recipe that no two processes on a node share a port and global kinds never repeat. diff --git a/src/srtctl/core/fingerprint.py b/src/srtctl/core/fingerprint.py index debf90b3a..a2219f325 100644 --- a/src/srtctl/core/fingerprint.py +++ b/src/srtctl/core/fingerprint.py @@ -47,6 +47,7 @@ FRAMEWORK_PACKAGES: dict[str, str] = { "vllm": "vllm", "tilert": "tilert", + "tokenspeed": "tokenspeed", "sglang": "sglang", "sglang-router": "sglang-router", "amd-mori": "amd_mori", diff --git a/src/srtctl/core/roles.py b/src/srtctl/core/roles.py index f1e1d0c33..83bf0f26d 100644 --- a/src/srtctl/core/roles.py +++ b/src/srtctl/core/roles.py @@ -19,6 +19,7 @@ "atom": "atom_config", "sglang": "sglang_config", "tilert": "tilert_config", + "tokenspeed": "tokenspeed_config", "vllm": "vllm_config", "trtllm": "trtllm_config", "mocker": "mocker_config", diff --git a/src/srtctl/core/schema.py b/src/srtctl/core/schema.py index 015714f3f..60b855f2f 100755 --- a/src/srtctl/core/schema.py +++ b/src/srtctl/core/schema.py @@ -44,6 +44,7 @@ MockerBackend, SGLangBackend, TileRTBackend, + TokenSpeedBackend, TRTLLMBackend, VLLMBackend, VLLMMooncakeKVStoreConfig, @@ -504,7 +505,16 @@ def _deserialize( if value is None: return SGLangBackend() - if isinstance(value, AtomBackend | SGLangBackend | TileRTBackend | TRTLLMBackend | VLLMBackend | MockerBackend): + if isinstance( + value, + AtomBackend + | SGLangBackend + | TileRTBackend + | TokenSpeedBackend + | TRTLLMBackend + | VLLMBackend + | MockerBackend, + ): return value if isinstance(value, str): @@ -527,6 +537,8 @@ def _deserialize( return AtomBackend.Schema().load(value) elif backend_type == "tilert": return TileRTBackend.Schema().load(value) + elif backend_type == "tokenspeed": + return TokenSpeedBackend.Schema().load(value) elif backend_type == "sglang": schema = SGLangBackend.Schema() return schema.load(value) @@ -541,7 +553,7 @@ def _deserialize( return schema.load(value) else: raise ValidationError( - f"Unknown engine type: {backend_type!r}. Supported types: atom, sglang, tilert, trtllm, vllm, mocker" + f"Unknown engine type: {backend_type!r}. Supported types: atom, sglang, tilert, tokenspeed, trtllm, vllm, mocker" ) def _serialize(self, value: Any | None, attr: str | None, obj: Any, **kwargs) -> Any: @@ -558,6 +570,8 @@ def _dump(value: Any) -> dict[str, Any]: return AtomBackend.Schema().dump(value) if isinstance(value, TileRTBackend): return TileRTBackend.Schema().dump(value) + if isinstance(value, TokenSpeedBackend): + return TokenSpeedBackend.Schema().dump(value) if isinstance(value, SGLangBackend): return SGLangBackend.Schema().dump(value) if isinstance(value, TRTLLMBackend): diff --git a/src/srtctl/core/schema_docs.py b/src/srtctl/core/schema_docs.py index 7c429c37c..1b097851d 100644 --- a/src/srtctl/core/schema_docs.py +++ b/src/srtctl/core/schema_docs.py @@ -48,6 +48,7 @@ MockerBackend, SGLangBackend, TileRTBackend, + TokenSpeedBackend, TRTLLMBackend, VLLMBackend, ) @@ -77,6 +78,7 @@ ("atom", AtomBackend), ("sglang", SGLangBackend), ("tilert", TileRTBackend), + ("tokenspeed", TokenSpeedBackend), ("trtllm", TRTLLMBackend), ("vllm", VLLMBackend), ("mocker", MockerBackend), @@ -102,7 +104,7 @@ class FieldDoc: type_label="str \\| mapping", default="optional when every role sets `engine`", description=( - "The engine type (`atom`, `sglang`, `tilert`, `trtllm`, `vllm`, `mocker`) as a string, or a mapping with `type` " + "The engine type (`atom`, `sglang`, `tilert`, `tokenspeed`, `trtllm`, `vllm`, `mocker`) as a string, or a mapping with `type` " "plus the engine-wide knobs listed under [Engine types](#engine-types)." ), ) diff --git a/src/srtctl/core/topology.py b/src/srtctl/core/topology.py index 43c735bde..2ca30a526 100644 --- a/src/srtctl/core/topology.py +++ b/src/srtctl/core/topology.py @@ -169,6 +169,7 @@ class Process: moriio_handshake_port: MoRI-IO handshake listener of a vLLM discovery-connector worker moriio_notify_port: first port of that worker's MoRI-IO notify block (one port per local rank) trtllm_dist_init_port: TRT-LLM torch.distributed bootstrap port; the leader's is used + tokenspeed_port: TokenSpeed ``--port``, the base of the engine's own free-port scan Every port is allocated by ``NodePortAllocator`` in ``endpoints_to_processes``; consumers read these fields and never derive one port from another. @@ -195,6 +196,7 @@ class Process: grpc_http_port: int | None = None vllm_scan_port: int | None = None trtllm_dist_init_port: int | None = None + tokenspeed_port: int | None = None moriio_handshake_port: int | None = None moriio_notify_port: int | None = None diff --git a/src/srtctl/ports.py b/src/srtctl/ports.py index 7ceffe9c9..6106f2d4e 100644 --- a/src/srtctl/ports.py +++ b/src/srtctl/ports.py @@ -45,6 +45,11 @@ # TRT-LLM torch.distributed bootstrap, one port per MPI endpoint. TRTLLM_DIST_INIT_PORT_BASE = 29500 +# TokenSpeed --port, one per process: a follower node's health listener and the base +# of the engine's own NCCL free-port scan, which starts 100 to 1000 ports above it. +TOKENSPEED_PORT_BASE = 10000 +TOKENSPEED_PORT_STRIDE = 1024 + # Mooncake transfer-engine ports (shared by SGLang and vLLM backends). MOONCAKE_MASTER_PORT = 8700 MOONCAKE_HTTP_METADATA_PORT = 8701 @@ -129,3 +134,4 @@ class PortKind: MORIIO_HANDSHAKE_PORTS = PortKind("moriio_handshake", VLLM_MORIIO_HANDSHAKE_PORT_BASE) MORIIO_NOTIFY_PORTS = PortKind("moriio_notify", VLLM_MORIIO_NOTIFY_PORT_BASE) TRTLLM_DIST_INIT_PORTS = PortKind("trtllm_dist_init", TRTLLM_DIST_INIT_PORT_BASE) +TOKENSPEED_PORTS = PortKind("tokenspeed", TOKENSPEED_PORT_BASE, TOKENSPEED_PORT_STRIDE, per_node=True) diff --git a/tests/snapshots/launch/tokenspeed__dynamo-agg.txt b/tests/snapshots/launch/tokenspeed__dynamo-agg.txt new file mode 100644 index 000000000..3d1799677 --- /dev/null +++ b/tests/snapshots/launch/tokenspeed__dynamo-agg.txt @@ -0,0 +1,202 @@ +# examples/tokenspeed/dynamo-agg.yaml +# exit_code: 0 +# srun calls: 12 + +## service_etcd +nodelist: node-01 +output: /logs/service_etcd.out +container_image: tokenspeed +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +use_bash_wrapper: True +bash_preamble: | + rm -rf /tmp/etcd && mkdir -p /tmp/etcd +command: | + /configs/etcd --data-dir /tmp/etcd --listen-client-urls http://0.0.0.0:2379 --advertise-client-urls http://127.0.0.1:2379 + +## agg_0_node-01 +nodelist: node-01 +output: /logs/node-01_agg_w0.out +container_image: tokenspeed +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +env_to_set: + CUDA_VISIBLE_DEVICES: 0 + DYN_LOG: info,dynamo_runtime::pipeline::network::ingress::push_handler=warn + DYN_REQUEST_PLANE: tcp + DYN_SKIP_SGLANG_LOG_FORMATTING: 1 + DYN_SYSTEM_PORT: 7500 + ETCD_ENDPOINTS: http://127.0.0.1:2379 + HEAD_NODE_IP: 127.0.0.1 + PYTHONUNBUFFERED: 1 +bash_preamble: | + ( python3 <(cat <<'__FINGERPRINT_EOF__' + + __FINGERPRINT_EOF__ + ) || true ) +command: | + python3 -m dynamo.tokenspeed --model /model --served-model-name Qwen/Qwen3-0.6B --host 127.0.0.1 --port 10000 --dist-init-addr 127.0.0.1:8300 --nnodes 1 --node-rank 0 --gpu-memory-utilization 0.5 --max-model-len 4096 --tensor-parallel-size 1 + +## agg_1_node-01 +nodelist: node-01 +output: /logs/node-01_agg_w1.out +container_image: tokenspeed +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +env_to_set: + CUDA_VISIBLE_DEVICES: 1 + DYN_LOG: info,dynamo_runtime::pipeline::network::ingress::push_handler=warn + DYN_REQUEST_PLANE: tcp + DYN_SKIP_SGLANG_LOG_FORMATTING: 1 + DYN_SYSTEM_PORT: 7501 + ETCD_ENDPOINTS: http://127.0.0.1:2379 + HEAD_NODE_IP: 127.0.0.1 + PYTHONUNBUFFERED: 1 +bash_preamble: | + ( python3 <(cat <<'__FINGERPRINT_EOF__' + + __FINGERPRINT_EOF__ + ) || true ) +command: | + python3 -m dynamo.tokenspeed --model /model --served-model-name Qwen/Qwen3-0.6B --host 127.0.0.1 --port 11024 --dist-init-addr 127.0.0.1:8307 --nnodes 1 --node-rank 0 --gpu-memory-utilization 0.5 --max-model-len 4096 --tensor-parallel-size 1 + +## frontend_0 +nodelist: node-01 +output: /logs/node-01_frontend_0.out +container_image: tokenspeed +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +env_to_set: + DYN_REQUEST_PLANE: tcp + DYN_SKIP_SGLANG_LOG_FORMATTING: 1 + ETCD_ENDPOINTS: http://127.0.0.1:2379 +mpi: pmix +command: | + python3 -m dynamo.frontend --http-port=8000 --router-mode kv + +## service_dcgm-exporter_node-01 +nodelist: node-01 +output: /logs/service_dcgm-exporter_node-01.out +container_image: nvcr.io#nvidia/k8s/dcgm-exporter:3.3.9-3.6.1-ubuntu22.04 +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +srun_export_env: + ETCD_ENDPOINTS: http://127.0.0.1:2379 +use_bash_wrapper: False +command: | + dcgm-exporter --collect-interval=1000 --address :9401 + +## service_dcgm-exporter_node-02 +nodelist: node-02 +output: /logs/service_dcgm-exporter_node-02.out +container_image: nvcr.io#nvidia/k8s/dcgm-exporter:3.3.9-3.6.1-ubuntu22.04 +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +srun_export_env: + ETCD_ENDPOINTS: http://127.0.0.1:2379 +use_bash_wrapper: False +command: | + dcgm-exporter --collect-interval=1000 --address :9401 + +## service_dcgm-exporter_node-03 +nodelist: node-03 +output: /logs/service_dcgm-exporter_node-03.out +container_image: nvcr.io#nvidia/k8s/dcgm-exporter:3.3.9-3.6.1-ubuntu22.04 +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +srun_export_env: + ETCD_ENDPOINTS: http://127.0.0.1:2379 +use_bash_wrapper: False +command: | + dcgm-exporter --collect-interval=1000 --address :9401 + +## service_node-exporter_node-01 +nodelist: node-01 +output: /logs/service_node-exporter_node-01.out +container_image: quay.io#prometheus/node-exporter:v1.8.2 +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +srun_export_env: + ETCD_ENDPOINTS: http://127.0.0.1:2379 +use_bash_wrapper: False +command: | + /bin/node_exporter --web.listen-address=:9101 --collector.disable-defaults --collector.cpu --collector.infiniband --collector.meminfo --collector.processes --collector.stat --collector.vmstat --collector.pressure --collector.meminfo_numa '--collector.vmstat.fields=^(oom_kill|pgpg|pswp|pgsteal|pg.*fault).*' + +## service_node-exporter_node-02 +nodelist: node-02 +output: /logs/service_node-exporter_node-02.out +container_image: quay.io#prometheus/node-exporter:v1.8.2 +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +srun_export_env: + ETCD_ENDPOINTS: http://127.0.0.1:2379 +use_bash_wrapper: False +command: | + /bin/node_exporter --web.listen-address=:9101 --collector.disable-defaults --collector.cpu --collector.infiniband --collector.meminfo --collector.processes --collector.stat --collector.vmstat --collector.pressure --collector.meminfo_numa '--collector.vmstat.fields=^(oom_kill|pgpg|pswp|pgsteal|pg.*fault).*' + +## service_node-exporter_node-03 +nodelist: node-03 +output: /logs/service_node-exporter_node-03.out +container_image: quay.io#prometheus/node-exporter:v1.8.2 +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +srun_export_env: + ETCD_ENDPOINTS: http://127.0.0.1:2379 +use_bash_wrapper: False +command: | + /bin/node_exporter --web.listen-address=:9101 --collector.disable-defaults --collector.cpu --collector.infiniband --collector.meminfo --collector.processes --collector.stat --collector.vmstat --collector.pressure --collector.meminfo_numa '--collector.vmstat.fields=^(oom_kill|pgpg|pswp|pgsteal|pg.*fault).*' + +## tachometer +nodelist: node-01 +output: /logs/tachometer.out +srun_export_env: + POLARS_MAX_THREADS: 4 +use_bash_wrapper: False +command: | + tachometer-scraper --config /logs/tachometer_config.toml --local-dir /logs/tachometer/local --sync-interval 120 + +## +nodelist: node-01 +output: /logs/benchmark.out +container_image: tokenspeed +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +env_to_set: + SRTCTL_FRONTEND_TYPE: dynamo + SRT_FRONTEND_HOST: 127.0.0.1 + SRT_FRONTEND_PORT: 8000 +command: | + bash /srtctl-benchmarks/sa-bench/bench.sh http://localhost:8000 128 128 4x8 inf /model Qwen/Qwen3-0.6B false 8 0 0 0.8 10 2 '' true random '' false diff --git a/tests/snapshots/launch/tokenspeed__dynamo-disagg.txt b/tests/snapshots/launch/tokenspeed__dynamo-disagg.txt new file mode 100644 index 000000000..5ffd73f8f --- /dev/null +++ b/tests/snapshots/launch/tokenspeed__dynamo-disagg.txt @@ -0,0 +1,204 @@ +# examples/tokenspeed/dynamo-disagg.yaml +# exit_code: 0 +# srun calls: 12 + +## service_etcd +nodelist: node-01 +output: /logs/service_etcd.out +container_image: tokenspeed +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +use_bash_wrapper: True +bash_preamble: | + rm -rf /tmp/etcd && mkdir -p /tmp/etcd +command: | + /configs/etcd --data-dir /tmp/etcd --listen-client-urls http://0.0.0.0:2379 --advertise-client-urls http://127.0.0.1:2379 + +## prefill_0_node-01 +nodelist: node-01 +output: /logs/node-01_prefill_w0.out +container_image: tokenspeed +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +env_to_set: + CUDA_VISIBLE_DEVICES: 0 + DYN_LOG: info,dynamo_runtime::pipeline::network::ingress::push_handler=warn + DYN_REQUEST_PLANE: tcp + DYN_SKIP_SGLANG_LOG_FORMATTING: 1 + DYN_SYSTEM_PORT: 7500 + ETCD_ENDPOINTS: http://127.0.0.1:2379 + HEAD_NODE_IP: 127.0.0.1 + PYTHONUNBUFFERED: 1 + TOKENSPEED_SKIP_GRPC_WARMUP: 1 +bash_preamble: | + ( python3 <(cat <<'__FINGERPRINT_EOF__' + + __FINGERPRINT_EOF__ + ) || true ) +command: | + python3 -m dynamo.tokenspeed --model /model --served-model-name Qwen/Qwen3-0.6B --host 127.0.0.1 --port 10000 --dist-init-addr 127.0.0.1:8300 --nnodes 1 --node-rank 0 --disaggregation-mode prefill --disaggregation-bootstrap-port 7200 --gpu-memory-utilization 0.5 --max-model-len 4096 --tensor-parallel-size 1 + +## decode_0_node-02 +nodelist: node-02 +output: /logs/node-02_decode_w0.out +container_image: tokenspeed +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +env_to_set: + CUDA_VISIBLE_DEVICES: 0 + DYN_LOG: info,dynamo_runtime::pipeline::network::ingress::push_handler=warn + DYN_REQUEST_PLANE: tcp + DYN_SKIP_SGLANG_LOG_FORMATTING: 1 + DYN_SYSTEM_PORT: 7501 + ETCD_ENDPOINTS: http://127.0.0.1:2379 + HEAD_NODE_IP: 127.0.0.1 + PYTHONUNBUFFERED: 1 + TOKENSPEED_SKIP_GRPC_WARMUP: 1 +bash_preamble: | + ( python3 <(cat <<'__FINGERPRINT_EOF__' + + __FINGERPRINT_EOF__ + ) || true ) +command: | + python3 -m dynamo.tokenspeed --model /model --served-model-name Qwen/Qwen3-0.6B --host 127.0.0.1 --port 10000 --dist-init-addr 127.0.0.1:8300 --nnodes 1 --node-rank 0 --disaggregation-mode decode --gpu-memory-utilization 0.5 --max-model-len 4096 --tensor-parallel-size 1 + +## frontend_0 +nodelist: node-01 +output: /logs/node-01_frontend_0.out +container_image: tokenspeed +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +env_to_set: + DYN_REQUEST_PLANE: tcp + DYN_SKIP_SGLANG_LOG_FORMATTING: 1 + ETCD_ENDPOINTS: http://127.0.0.1:2379 +mpi: pmix +command: | + python3 -m dynamo.frontend --http-port=8000 --router-mode kv + +## service_dcgm-exporter_node-01 +nodelist: node-01 +output: /logs/service_dcgm-exporter_node-01.out +container_image: nvcr.io#nvidia/k8s/dcgm-exporter:3.3.9-3.6.1-ubuntu22.04 +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +srun_export_env: + ETCD_ENDPOINTS: http://127.0.0.1:2379 +use_bash_wrapper: False +command: | + dcgm-exporter --collect-interval=1000 --address :9401 + +## service_dcgm-exporter_node-02 +nodelist: node-02 +output: /logs/service_dcgm-exporter_node-02.out +container_image: nvcr.io#nvidia/k8s/dcgm-exporter:3.3.9-3.6.1-ubuntu22.04 +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +srun_export_env: + ETCD_ENDPOINTS: http://127.0.0.1:2379 +use_bash_wrapper: False +command: | + dcgm-exporter --collect-interval=1000 --address :9401 + +## service_dcgm-exporter_node-03 +nodelist: node-03 +output: /logs/service_dcgm-exporter_node-03.out +container_image: nvcr.io#nvidia/k8s/dcgm-exporter:3.3.9-3.6.1-ubuntu22.04 +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +srun_export_env: + ETCD_ENDPOINTS: http://127.0.0.1:2379 +use_bash_wrapper: False +command: | + dcgm-exporter --collect-interval=1000 --address :9401 + +## service_node-exporter_node-01 +nodelist: node-01 +output: /logs/service_node-exporter_node-01.out +container_image: quay.io#prometheus/node-exporter:v1.8.2 +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +srun_export_env: + ETCD_ENDPOINTS: http://127.0.0.1:2379 +use_bash_wrapper: False +command: | + /bin/node_exporter --web.listen-address=:9101 --collector.disable-defaults --collector.cpu --collector.infiniband --collector.meminfo --collector.processes --collector.stat --collector.vmstat --collector.pressure --collector.meminfo_numa '--collector.vmstat.fields=^(oom_kill|pgpg|pswp|pgsteal|pg.*fault).*' + +## service_node-exporter_node-02 +nodelist: node-02 +output: /logs/service_node-exporter_node-02.out +container_image: quay.io#prometheus/node-exporter:v1.8.2 +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +srun_export_env: + ETCD_ENDPOINTS: http://127.0.0.1:2379 +use_bash_wrapper: False +command: | + /bin/node_exporter --web.listen-address=:9101 --collector.disable-defaults --collector.cpu --collector.infiniband --collector.meminfo --collector.processes --collector.stat --collector.vmstat --collector.pressure --collector.meminfo_numa '--collector.vmstat.fields=^(oom_kill|pgpg|pswp|pgsteal|pg.*fault).*' + +## service_node-exporter_node-03 +nodelist: node-03 +output: /logs/service_node-exporter_node-03.out +container_image: quay.io#prometheus/node-exporter:v1.8.2 +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +srun_export_env: + ETCD_ENDPOINTS: http://127.0.0.1:2379 +use_bash_wrapper: False +command: | + /bin/node_exporter --web.listen-address=:9101 --collector.disable-defaults --collector.cpu --collector.infiniband --collector.meminfo --collector.processes --collector.stat --collector.vmstat --collector.pressure --collector.meminfo_numa '--collector.vmstat.fields=^(oom_kill|pgpg|pswp|pgsteal|pg.*fault).*' + +## tachometer +nodelist: node-01 +output: /logs/tachometer.out +srun_export_env: + POLARS_MAX_THREADS: 4 +use_bash_wrapper: False +command: | + tachometer-scraper --config /logs/tachometer_config.toml --local-dir /logs/tachometer/local --sync-interval 120 + +## +nodelist: node-01 +output: /logs/benchmark.out +container_image: tokenspeed +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +env_to_set: + SRTCTL_FRONTEND_TYPE: dynamo + SRT_FRONTEND_HOST: 127.0.0.1 + SRT_FRONTEND_PORT: 8000 +command: | + bash /srtctl-benchmarks/sa-bench/bench.sh http://localhost:8000 128 128 4x8 inf /model Qwen/Qwen3-0.6B true 2 1 1 0.8 10 2 '' true random '' false diff --git a/tests/snapshots/launch/tokenspeed__smg-agg.txt b/tests/snapshots/launch/tokenspeed__smg-agg.txt new file mode 100644 index 000000000..b324800d0 --- /dev/null +++ b/tests/snapshots/launch/tokenspeed__smg-agg.txt @@ -0,0 +1,172 @@ +# examples/tokenspeed/smg-agg.yaml +# exit_code: 0 +# srun calls: 11 + +## agg_0_node-01 +nodelist: node-01 +output: /logs/node-01_agg_w0.out +container_image: tokenspeed +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +env_to_set: + CUDA_VISIBLE_DEVICES: 0 + HEAD_NODE_IP: 127.0.0.1 + PYTHONUNBUFFERED: 1 +bash_preamble: | + ( python3 <(cat <<'__FINGERPRINT_EOF__' + + __FINGERPRINT_EOF__ + ) || true ) +command: | + python3 -m smg_grpc_servicer.tokenspeed --model /model --served-model-name Qwen/Qwen3-0.6B --host 127.0.0.1 --port 10000 --dist-init-addr 127.0.0.1:8300 --nnodes 1 --node-rank 0 --gpu-memory-utilization 0.5 --max-model-len 4096 --tensor-parallel-size 1 + +## agg_1_node-01 +nodelist: node-01 +output: /logs/node-01_agg_w1.out +container_image: tokenspeed +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +env_to_set: + CUDA_VISIBLE_DEVICES: 1 + HEAD_NODE_IP: 127.0.0.1 + PYTHONUNBUFFERED: 1 +bash_preamble: | + ( python3 <(cat <<'__FINGERPRINT_EOF__' + + __FINGERPRINT_EOF__ + ) || true ) +command: | + python3 -m smg_grpc_servicer.tokenspeed --model /model --served-model-name Qwen/Qwen3-0.6B --host 127.0.0.1 --port 11024 --dist-init-addr 127.0.0.1:8307 --nnodes 1 --node-rank 0 --gpu-memory-utilization 0.5 --max-model-len 4096 --tensor-parallel-size 1 + +## smg_0 +nodelist: node-01 +output: /logs/node-01_smg_0.out +container_image: tokenspeed +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +command: | + smg launch --worker-urls grpc://127.0.0.1:10000 grpc://127.0.0.1:11024 --host 0.0.0.0 --port 8000 --prometheus-port 29000 + +## service_dcgm-exporter_node-01 +nodelist: node-01 +output: /logs/service_dcgm-exporter_node-01.out +container_image: nvcr.io#nvidia/k8s/dcgm-exporter:3.3.9-3.6.1-ubuntu22.04 +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +srun_export_env: + ETCD_ENDPOINTS: http://127.0.0.1:2379 +use_bash_wrapper: False +command: | + dcgm-exporter --collect-interval=1000 --address :9401 + +## service_dcgm-exporter_node-02 +nodelist: node-02 +output: /logs/service_dcgm-exporter_node-02.out +container_image: nvcr.io#nvidia/k8s/dcgm-exporter:3.3.9-3.6.1-ubuntu22.04 +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +srun_export_env: + ETCD_ENDPOINTS: http://127.0.0.1:2379 +use_bash_wrapper: False +command: | + dcgm-exporter --collect-interval=1000 --address :9401 + +## service_dcgm-exporter_node-03 +nodelist: node-03 +output: /logs/service_dcgm-exporter_node-03.out +container_image: nvcr.io#nvidia/k8s/dcgm-exporter:3.3.9-3.6.1-ubuntu22.04 +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +srun_export_env: + ETCD_ENDPOINTS: http://127.0.0.1:2379 +use_bash_wrapper: False +command: | + dcgm-exporter --collect-interval=1000 --address :9401 + +## service_node-exporter_node-01 +nodelist: node-01 +output: /logs/service_node-exporter_node-01.out +container_image: quay.io#prometheus/node-exporter:v1.8.2 +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +srun_export_env: + ETCD_ENDPOINTS: http://127.0.0.1:2379 +use_bash_wrapper: False +command: | + /bin/node_exporter --web.listen-address=:9101 --collector.disable-defaults --collector.cpu --collector.infiniband --collector.meminfo --collector.processes --collector.stat --collector.vmstat --collector.pressure --collector.meminfo_numa '--collector.vmstat.fields=^(oom_kill|pgpg|pswp|pgsteal|pg.*fault).*' + +## service_node-exporter_node-02 +nodelist: node-02 +output: /logs/service_node-exporter_node-02.out +container_image: quay.io#prometheus/node-exporter:v1.8.2 +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +srun_export_env: + ETCD_ENDPOINTS: http://127.0.0.1:2379 +use_bash_wrapper: False +command: | + /bin/node_exporter --web.listen-address=:9101 --collector.disable-defaults --collector.cpu --collector.infiniband --collector.meminfo --collector.processes --collector.stat --collector.vmstat --collector.pressure --collector.meminfo_numa '--collector.vmstat.fields=^(oom_kill|pgpg|pswp|pgsteal|pg.*fault).*' + +## service_node-exporter_node-03 +nodelist: node-03 +output: /logs/service_node-exporter_node-03.out +container_image: quay.io#prometheus/node-exporter:v1.8.2 +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +srun_export_env: + ETCD_ENDPOINTS: http://127.0.0.1:2379 +use_bash_wrapper: False +command: | + /bin/node_exporter --web.listen-address=:9101 --collector.disable-defaults --collector.cpu --collector.infiniband --collector.meminfo --collector.processes --collector.stat --collector.vmstat --collector.pressure --collector.meminfo_numa '--collector.vmstat.fields=^(oom_kill|pgpg|pswp|pgsteal|pg.*fault).*' + +## tachometer +nodelist: node-01 +output: /logs/tachometer.out +srun_export_env: + POLARS_MAX_THREADS: 4 +use_bash_wrapper: False +command: | + tachometer-scraper --config /logs/tachometer_config.toml --local-dir /logs/tachometer/local --sync-interval 120 + +## +nodelist: node-01 +output: /logs/benchmark.out +container_image: tokenspeed +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +env_to_set: + SRTCTL_FRONTEND_TYPE: smg + SRT_FRONTEND_HOST: 127.0.0.1 + SRT_FRONTEND_PORT: 8000 +command: | + bash /srtctl-benchmarks/sa-bench/bench.sh http://localhost:8000 128 128 4x8 inf /model Qwen/Qwen3-0.6B false 8 0 0 0.8 10 2 '' true random '' false diff --git a/tests/snapshots/launch/tokenspeed__smg-disagg.txt b/tests/snapshots/launch/tokenspeed__smg-disagg.txt new file mode 100644 index 000000000..98c24ceb8 --- /dev/null +++ b/tests/snapshots/launch/tokenspeed__smg-disagg.txt @@ -0,0 +1,174 @@ +# examples/tokenspeed/smg-disagg.yaml +# exit_code: 0 +# srun calls: 11 + +## prefill_0_node-01 +nodelist: node-01 +output: /logs/node-01_prefill_w0.out +container_image: tokenspeed +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +env_to_set: + CUDA_VISIBLE_DEVICES: 0 + HEAD_NODE_IP: 127.0.0.1 + PYTHONUNBUFFERED: 1 + TOKENSPEED_SKIP_GRPC_WARMUP: 1 +bash_preamble: | + ( python3 <(cat <<'__FINGERPRINT_EOF__' + + __FINGERPRINT_EOF__ + ) || true ) +command: | + python3 -m smg_grpc_servicer.tokenspeed --model /model --served-model-name Qwen/Qwen3-0.6B --host 127.0.0.1 --port 10000 --dist-init-addr 127.0.0.1:8300 --nnodes 1 --node-rank 0 --disaggregation-mode prefill --disaggregation-bootstrap-port 7200 --gpu-memory-utilization 0.5 --max-model-len 4096 --tensor-parallel-size 1 + +## decode_0_node-02 +nodelist: node-02 +output: /logs/node-02_decode_w0.out +container_image: tokenspeed +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +env_to_set: + CUDA_VISIBLE_DEVICES: 0 + HEAD_NODE_IP: 127.0.0.1 + PYTHONUNBUFFERED: 1 + TOKENSPEED_SKIP_GRPC_WARMUP: 1 +bash_preamble: | + ( python3 <(cat <<'__FINGERPRINT_EOF__' + + __FINGERPRINT_EOF__ + ) || true ) +command: | + python3 -m smg_grpc_servicer.tokenspeed --model /model --served-model-name Qwen/Qwen3-0.6B --host 127.0.0.1 --port 10000 --dist-init-addr 127.0.0.1:8300 --nnodes 1 --node-rank 0 --disaggregation-mode decode --gpu-memory-utilization 0.5 --max-model-len 4096 --tensor-parallel-size 1 + +## smg_0 +nodelist: node-01 +output: /logs/node-01_smg_0.out +container_image: tokenspeed +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +command: | + smg launch --pd-disaggregation --prefill grpc://127.0.0.1:10000 7200 --decode grpc://127.0.0.1:10000 --host 0.0.0.0 --port 8000 --prometheus-port 29000 + +## service_dcgm-exporter_node-01 +nodelist: node-01 +output: /logs/service_dcgm-exporter_node-01.out +container_image: nvcr.io#nvidia/k8s/dcgm-exporter:3.3.9-3.6.1-ubuntu22.04 +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +srun_export_env: + ETCD_ENDPOINTS: http://127.0.0.1:2379 +use_bash_wrapper: False +command: | + dcgm-exporter --collect-interval=1000 --address :9401 + +## service_dcgm-exporter_node-02 +nodelist: node-02 +output: /logs/service_dcgm-exporter_node-02.out +container_image: nvcr.io#nvidia/k8s/dcgm-exporter:3.3.9-3.6.1-ubuntu22.04 +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +srun_export_env: + ETCD_ENDPOINTS: http://127.0.0.1:2379 +use_bash_wrapper: False +command: | + dcgm-exporter --collect-interval=1000 --address :9401 + +## service_dcgm-exporter_node-03 +nodelist: node-03 +output: /logs/service_dcgm-exporter_node-03.out +container_image: nvcr.io#nvidia/k8s/dcgm-exporter:3.3.9-3.6.1-ubuntu22.04 +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +srun_export_env: + ETCD_ENDPOINTS: http://127.0.0.1:2379 +use_bash_wrapper: False +command: | + dcgm-exporter --collect-interval=1000 --address :9401 + +## service_node-exporter_node-01 +nodelist: node-01 +output: /logs/service_node-exporter_node-01.out +container_image: quay.io#prometheus/node-exporter:v1.8.2 +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +srun_export_env: + ETCD_ENDPOINTS: http://127.0.0.1:2379 +use_bash_wrapper: False +command: | + /bin/node_exporter --web.listen-address=:9101 --collector.disable-defaults --collector.cpu --collector.infiniband --collector.meminfo --collector.processes --collector.stat --collector.vmstat --collector.pressure --collector.meminfo_numa '--collector.vmstat.fields=^(oom_kill|pgpg|pswp|pgsteal|pg.*fault).*' + +## service_node-exporter_node-02 +nodelist: node-02 +output: /logs/service_node-exporter_node-02.out +container_image: quay.io#prometheus/node-exporter:v1.8.2 +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +srun_export_env: + ETCD_ENDPOINTS: http://127.0.0.1:2379 +use_bash_wrapper: False +command: | + /bin/node_exporter --web.listen-address=:9101 --collector.disable-defaults --collector.cpu --collector.infiniband --collector.meminfo --collector.processes --collector.stat --collector.vmstat --collector.pressure --collector.meminfo_numa '--collector.vmstat.fields=^(oom_kill|pgpg|pswp|pgsteal|pg.*fault).*' + +## service_node-exporter_node-03 +nodelist: node-03 +output: /logs/service_node-exporter_node-03.out +container_image: quay.io#prometheus/node-exporter:v1.8.2 +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +srun_export_env: + ETCD_ENDPOINTS: http://127.0.0.1:2379 +use_bash_wrapper: False +command: | + /bin/node_exporter --web.listen-address=:9101 --collector.disable-defaults --collector.cpu --collector.infiniband --collector.meminfo --collector.processes --collector.stat --collector.vmstat --collector.pressure --collector.meminfo_numa '--collector.vmstat.fields=^(oom_kill|pgpg|pswp|pgsteal|pg.*fault).*' + +## tachometer +nodelist: node-01 +output: /logs/tachometer.out +srun_export_env: + POLARS_MAX_THREADS: 4 +use_bash_wrapper: False +command: | + tachometer-scraper --config /logs/tachometer_config.toml --local-dir /logs/tachometer/local --sync-interval 120 + +## +nodelist: node-01 +output: /logs/benchmark.out +container_image: tokenspeed +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +env_to_set: + SRTCTL_FRONTEND_TYPE: smg + SRT_FRONTEND_HOST: 127.0.0.1 + SRT_FRONTEND_PORT: 8000 +command: | + bash /srtctl-benchmarks/sa-bench/bench.sh http://localhost:8000 128 128 4x8 inf /model Qwen/Qwen3-0.6B true 2 1 1 0.8 10 2 '' true random '' false diff --git a/tests/test_backend_protocol.py b/tests/test_backend_protocol.py index 1a3bd753e..6e214a3c7 100644 --- a/tests/test_backend_protocol.py +++ b/tests/test_backend_protocol.py @@ -21,13 +21,14 @@ SGLangBackend, TRTLLMBackend, TileRTBackend, + TokenSpeedBackend, VLLMFailoverConfig, VLLMBackend, ) from srtctl.core.topology import Process from srtctl.ports import MOONCAKE_MASTER_PORT -BACKENDS = [SGLangBackend, TRTLLMBackend, VLLMBackend, MockerBackend, TileRTBackend] +BACKENDS = [SGLangBackend, TRTLLMBackend, VLLMBackend, MockerBackend, TileRTBackend, TokenSpeedBackend] SRC = Path(__file__).resolve().parents[1] / "src" / "srtctl" DUCK_TYPED_BACKEND = re.compile(r"\b(getattr|hasattr)\((self\.|config\.|self\.config\.)?backend\b") @@ -52,7 +53,9 @@ def test_optional_features_read_as_absent_by_default(backend_cls): assert backend.get_mooncake_worker_env("10.0.0.1", "10.0.0.2") == {} assert backend.get_failover_environment(_process(), "12345") == {} assert backend.get_srun_config().sequential_node_start == 0 - assert not any(backend.is_grpc_mode(mode) for mode in ("prefill", "decode", "agg")) + # A direct TokenSpeed worker is always its gRPC engine. + grpc_only = backend_cls is TokenSpeedBackend + assert all(backend.is_grpc_mode(mode) == grpc_only for mode in ("prefill", "decode", "agg")) def test_vllm_grpc_role_arg_selects_grpc_mode(): @@ -123,7 +126,9 @@ def test_every_backend_names_its_fatal_log_patterns(backend_cls, mode): @pytest.mark.parametrize( - "backend_cls", [SGLangBackend, VLLMBackend, MockerBackend, AtomBackend, TileRTBackend], ids=lambda c: c.__name__ + "backend_cls", + [SGLangBackend, VLLMBackend, MockerBackend, AtomBackend, TileRTBackend, TokenSpeedBackend], + ids=lambda c: c.__name__, ) def test_engines_whose_step_exits_with_the_engine_watch_nothing(backend_cls): assert backend_cls().fatal_log_patterns("decode") == () diff --git a/tests/test_e2e.py b/tests/test_e2e.py index 518e9cc02..ecc049143 100644 --- a/tests/test_e2e.py +++ b/tests/test_e2e.py @@ -493,12 +493,8 @@ def test_mooncake_worker_env_not_set(self): def test_mooncake_worker_env_minimal(self): """mooncake_kv_store with no env → MOONCAKE_MASTER + metadata URL + auto-resolved hostname.""" - from srtctl.backends.sglang import ( - MOONCAKE_HTTP_METADATA_PORT, - MOONCAKE_MASTER_PORT, - MooncakeKVStoreConfig, - SGLangBackend, - ) + from srtctl.backends.sglang import MooncakeKVStoreConfig, SGLangBackend + from srtctl.ports import MOONCAKE_HTTP_METADATA_PORT, MOONCAKE_MASTER_PORT backend = SGLangBackend(mooncake_kv_store=MooncakeKVStoreConfig()) env = backend.get_mooncake_worker_env("10.0.0.1", "10.0.0.42") @@ -510,12 +506,8 @@ def test_mooncake_worker_env_minimal(self): def test_mooncake_worker_env_master_always_overrides_user(self): """User-supplied MOONCAKE_MASTER and metadata URL are always overridden by srtslurm.""" - from srtctl.backends.sglang import ( - MOONCAKE_HTTP_METADATA_PORT, - MOONCAKE_MASTER_PORT, - MooncakeKVStoreConfig, - SGLangBackend, - ) + from srtctl.backends.sglang import MooncakeKVStoreConfig, SGLangBackend + from srtctl.ports import MOONCAKE_HTTP_METADATA_PORT, MOONCAKE_MASTER_PORT backend = SGLangBackend( mooncake_kv_store=MooncakeKVStoreConfig( @@ -541,7 +533,8 @@ def test_mooncake_worker_env_local_hostname_user_can_override(self): def test_mooncake_worker_env_passthrough(self): """mooncake_kv_store.env values are merged with MOONCAKE_MASTER.""" - from srtctl.backends.sglang import MOONCAKE_MASTER_PORT, MooncakeKVStoreConfig, SGLangBackend + from srtctl.backends.sglang import MooncakeKVStoreConfig, SGLangBackend + from srtctl.ports import MOONCAKE_MASTER_PORT backend = SGLangBackend( mooncake_kv_store=MooncakeKVStoreConfig( diff --git a/tests/test_port_allocator.py b/tests/test_port_allocator.py index 00a31d74a..e34f80130 100644 --- a/tests/test_port_allocator.py +++ b/tests/test_port_allocator.py @@ -26,6 +26,7 @@ NIXL_PORTS, SIDECAR_GRPC_PORTS, SYS_PORTS, + TOKENSPEED_PORTS, TRTLLM_DIST_INIT_PORTS, VLLM_SCAN_PORTS, PortKind, @@ -35,6 +36,7 @@ "examples/atom", "examples/sglang", "examples/tilert", + "examples/tokenspeed", "examples/vllm", "examples/trtllm", "examples/mocker", @@ -92,6 +94,7 @@ def test_every_kind_has_its_own_range(self): MORIIO_HANDSHAKE_PORTS, MORIIO_NOTIFY_PORTS, TRTLLM_DIST_INIT_PORTS, + TOKENSPEED_PORTS, SIDECAR_GRPC_PORTS, ] assert len({kind.name for kind in kinds}) == len(kinds) @@ -136,6 +139,7 @@ def test_example_topologies_bind_no_port_twice(recipe: Path): "moriio_handshake": process.moriio_handshake_port, "moriio_notify": process.moriio_notify_port, "trtllm_dist_init": process.trtllm_dist_init_port, + "tokenspeed": process.tokenspeed_port, } for kind, port in ports.items(): if port is not None: diff --git a/tests/test_roles.py b/tests/test_roles.py index 2fb958c2f..d78bfc78c 100644 --- a/tests/test_roles.py +++ b/tests/test_roles.py @@ -124,7 +124,7 @@ def test_agg_role_binds_to_the_aggregated_mode() -> None: def test_role_args_reach_every_engine() -> None: - for engine_type in ("sglang", "vllm", "trtllm", "mocker", "atom", "tilert"): + for engine_type in ("sglang", "vllm", "trtllm", "mocker", "atom", "tilert", "tokenspeed"): cfg = SrtConfig.Schema().load( _minimal(engine=engine_type, roles={"agg": {"nodes": 1, "workers": 1, "args": {"a": 1}, "env": {"E": "1"}}}) ) diff --git a/tests/test_tokenspeed.py b/tests/test_tokenspeed.py new file mode 100644 index 000000000..119d8d66b --- /dev/null +++ b/tests/test_tokenspeed.py @@ -0,0 +1,218 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2026 SemiAnalysis LLC. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""TokenSpeed backend: recipe validation and the dynamo.tokenspeed launch.""" + +from __future__ import annotations + +import json +import shlex +from pathlib import Path +from typing import Any + +import pytest +import yaml +from marshmallow import ValidationError + +from srtctl.backends import TokenSpeedBackend +from srtctl.core.config import resolve_config_with_defaults +from srtctl.core.schema import SrtConfig +from srtctl.mock import MockOptions, run_mock_sweep + + +def _recipe(**overrides: Any) -> dict[str, Any]: + recipe: dict[str, Any] = { + "schema": 2, + "name": "tokenspeed", + "model": {"path": "hf:Qwen/Qwen3-0.6B", "container": "tokenspeed", "precision": "bf16"}, + "resources": {"gpu_type": "h100", "gpus_per_node": 8}, + "dynamo": {"install": False}, + "frontend": {"type": "dynamo", "enable_multiple_frontends": False}, + "engine": "tokenspeed", + "roles": { + # One TP16 prefill worker across two nodes, two TP4 decode workers sharing a third. + "prefill": {"nodes": 2, "workers": 1, "args": {"tensor-parallel-size": 16}}, + "decode": {"nodes": 1, "workers": 2, "gpus": 4, "args": {"served-model-name": "qwen3"}}, + }, + "benchmark": {"type": "custom", "command": "echo done"}, + } + recipe.update(overrides) + return recipe + + +def _load(recipe: dict[str, Any]) -> SrtConfig: + return SrtConfig.Schema().load(resolve_config_with_defaults(recipe, None)) + + +@pytest.mark.parametrize("name", ["dynamo-agg", "dynamo-disagg"]) +def test_examples_load_as_tokenspeed(name: str) -> None: + config = SrtConfig.from_yaml(Path("examples/tokenspeed") / f"{name}.yaml") + assert isinstance(config.backend, TokenSpeedBackend) + assert config.served_model_name == "Qwen/Qwen3-0.6B" + + +@pytest.mark.parametrize( + ("change", "message"), + [ + ({"frontend": {"type": "sglang-router"}}, "requires backend.type: sglang"), + ({"dynamo": {"install": False, "sidecar": True}}, "supports sglang, vllm, and trtllm backends only"), + ( + {"roles": {"agg": {"nodes": 1, "workers": 1, "kv_events": True}}}, + "kv_events is not supported by the tokenspeed engine", + ), + ], +) +def test_unsupported_settings_fail_at_load(change: dict[str, Any], message: str) -> None: + with pytest.raises(ValidationError, match=message): + _load(_recipe(**change)) + + +def _flag(command: list[str], flag: str) -> str | None: + return command[command.index(flag) + 1] if flag in command else None + + +def test_mock_orchestrator_launches_dynamo_tokenspeed_workers(tmp_path: Path) -> None: + """Multi-node prefill, two decode workers on one node: allocator ports, rendezvous and P/D flags.""" + config_path = tmp_path / "recipe.yaml" + config_path.write_text(yaml.safe_dump(_recipe())) + launches: list[dict[str, Any]] = [] + + exit_code = run_mock_sweep( + config_path=config_path, + output_dir=tmp_path / "outputs" / "4242", + job_id="4242", + options=MockOptions( + child_duration_s=0.05, + phase_pause_s=0.01, + nodelist=("node-01", "node-02", "node-03"), + on_srun=launches.append, + ), + ) + + assert exit_code == 0 + workers = { + launch["step_name"]: shlex.split(" ".join(launch["command"])) + for launch in launches + if "dynamo.tokenspeed" in launch["command"] + } + assert sorted(workers) == ["decode_0_node-03", "decode_1_node-03", "prefill_0_node-01", "prefill_0_node-02"] + + leader, follower = workers["prefill_0_node-01"], workers["prefill_0_node-02"] + assert leader[:3] == ["python3", "-m", "dynamo.tokenspeed"] + assert _flag(leader, "--model") == "Qwen/Qwen3-0.6B" + assert _flag(leader, "--disaggregation-mode") == "prefill" + assert _flag(leader, "--disaggregation-bootstrap-port") is not None + assert _flag(leader, "--tensor-parallel-size") == "16" + # Both nodes of the prefill worker join the leader's rendezvous. + assert _flag(leader, "--dist-init-addr") == _flag(follower, "--dist-init-addr") + assert [(_flag(c, "--nnodes"), _flag(c, "--node-rank")) for c in (leader, follower)] == [("2", "0"), ("2", "1")] + + decode_0, decode_1 = workers["decode_0_node-03"], workers["decode_1_node-03"] + assert _flag(decode_0, "--disaggregation-mode") == "decode" + assert _flag(decode_0, "--disaggregation-bootstrap-port") is None + assert _flag(decode_0, "--served-model-name") == "qwen3" + # Two engines on one node get their own --port scan base and rendezvous block. + assert _flag(decode_0, "--port") != _flag(decode_1, "--port") + assert _flag(decode_0, "--dist-init-addr") != _flag(decode_1, "--dist-init-addr") + + +def test_smg_runs_grpc_engines_on_their_tokenspeed_port(tmp_path: Path) -> None: + """Behind SMG each worker is the gRPC engine, advertised as grpc:// on its --port.""" + recipe = _recipe(frontend={"type": "smg", "enable_multiple_frontends": False}, dynamo={"install": False}) + config_path = tmp_path / "recipe.yaml" + config_path.write_text(yaml.safe_dump(recipe)) + launches: list[dict[str, Any]] = [] + + exit_code = run_mock_sweep( + config_path=config_path, + output_dir=tmp_path / "outputs" / "4243", + job_id="4243", + options=MockOptions( + child_duration_s=0.05, + phase_pause_s=0.01, + nodelist=("node-01", "node-02", "node-03"), + on_srun=launches.append, + ), + ) + + assert exit_code == 0 + commands = [shlex.split(" ".join(launch["command"])) for launch in launches] + workers = { + launch["step_name"]: shlex.split(" ".join(launch["command"])) + for launch in launches + if "smg_grpc_servicer.tokenspeed" in launch["command"] + } + assert sorted(workers) == ["decode_0_node-03", "decode_1_node-03", "prefill_0_node-01", "prefill_0_node-02"] + assert all(cmd[:3] == ["python3", "-m", "smg_grpc_servicer.tokenspeed"] for cmd in workers.values()) + + router = next(cmd for cmd in commands if cmd[:2] == ["smg", "launch"]) + prefill = workers["prefill_0_node-01"] + prefill_url = f"grpc://{_flag(prefill, '--host')}:{_flag(prefill, '--port')}" + assert router[router.index("--prefill") + 1 : router.index("--prefill") + 3] == [ + prefill_url, + _flag(prefill, "--disaggregation-bootstrap-port"), + ] + decode_urls = [router[i + 1] for i, token in enumerate(router) if token == "--decode"] + assert decode_urls == [ + f"grpc://{_flag(workers[step], '--host')}:{_flag(workers[step], '--port')}" + for step in ("decode_0_node-03", "decode_1_node-03") + ] + # The prefill follower node joins the leader's engine and is not routed. + assert router.count("--prefill") == 1 + # A prefill/decode engine cannot serve the bootstrap-less startup warmup. + assert all( + launch["env_to_set"]["TOKENSPEED_SKIP_GRPC_WARMUP"] == "1" + for launch in launches + if "smg_grpc_servicer.tokenspeed" in launch["command"] + ) + + +def test_mooncake_master_service_points_the_l3_store_at_it(tmp_path: Path) -> None: + """A mooncake-master service gives every worker the MOONCAKE_* environment of its L3 store.""" + recipe = _recipe( + roles={"agg": {"nodes": 1, "workers": 2, "gpus": 4, "args": {"kvstore-storage-backend": "mooncake"}}}, + services=[{"name": "mooncake-master", "type": "mooncake-master"}], + ) + config_path = tmp_path / "recipe.yaml" + config_path.write_text(yaml.safe_dump(recipe)) + launches: list[dict[str, Any]] = [] + + exit_code = run_mock_sweep( + config_path=config_path, + output_dir=tmp_path / "outputs" / "4244", + job_id="4244", + options=MockOptions(child_duration_s=0.05, phase_pause_s=0.01, nodelist=("node-01",), on_srun=launches.append), + ) + + assert exit_code == 0 + assert any(launch["command"][0] == "mooncake_master" for launch in launches) + workers = [launch for launch in launches if "dynamo.tokenspeed" in launch["command"]] + assert len(workers) == 2 + for worker in workers: + env = worker["env_to_set"] + assert "TOKENSPEED_SKIP_GRPC_WARMUP" not in env + assert env["MOONCAKE_MASTER"].endswith(":8700") + assert env["MOONCAKE_TE_META_DATA_SERVER"].endswith(":8701/metadata") + assert _flag(shlex.split(" ".join(worker["command"])), "--kvstore-storage-backend") == "mooncake" + + +def test_mapping_args_are_passed_as_json(tmp_path: Path) -> None: + """A mapping-valued arg reaches TokenSpeed as JSON, not a Python dict repr.""" + recipe = _recipe( + roles={"agg": {"nodes": 1, "workers": 1, "args": {"kv-events-config": {"enable_kv_cache_events": True}}}}, + ) + config_path = tmp_path / "recipe.yaml" + config_path.write_text(yaml.safe_dump(recipe)) + launches: list[dict[str, Any]] = [] + + exit_code = run_mock_sweep( + config_path=config_path, + output_dir=tmp_path / "outputs" / "4245", + job_id="4245", + options=MockOptions(child_duration_s=0.05, phase_pause_s=0.01, nodelist=("node-01",), on_srun=launches.append), + ) + + assert exit_code == 0 + worker = next(launch["command"] for launch in launches if "dynamo.tokenspeed" in launch["command"]) + assert json.loads(_flag(worker, "--kv-events-config")) == {"enable_kv_cache_events": True} From b16d2e7f075c2d5698cef6dd2ec04b9a2ec0a5c5 Mon Sep 17 00:00:00 2001 From: Yuewei Na <248773860+nv-yna@users.noreply.github.com> Date: Wed, 7 Oct 2026 22:03:45 -0700 Subject: [PATCH 3/4] fix(dsight): rename UI title to Dsight (#586) Signed-off-by: Yuewei Na Co-authored-by: Yuewei Na --- src/srtctl/dsight/assets/explorer.html | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/srtctl/dsight/assets/explorer.html b/src/srtctl/dsight/assets/explorer.html index 1d029ffa1..c4621f455 100644 --- a/src/srtctl/dsight/assets/explorer.html +++ b/src/srtctl/dsight/assets/explorer.html @@ -6,7 +6,7 @@ - Inference trace explorer + Dsight