From e3571f058772390720fc300b75bf2433e21de009 Mon Sep 17 00:00:00 2001 From: Cam Quilici Date: Mon, 28 Sep 2026 15:25:26 -0500 Subject: [PATCH 1/2] feat(tilert): integrate per-role srt-slurm runtime --- .../multi_node/srt-slurm-recipes/RECIPES.md | 2 +- .../multi_node/srt_fixed_sequence.sh | 10 +- inferencex-e2e/benchmarks/srt_agentic.sh | 4 + .../infx/srt_slurm/synthetic_acceptance.py | 50 ++++++-- .../srt_slurm/test_fixed_sequence_client.py | 66 ++++++++++ .../srt_slurm/test_synthetic_acceptance.py | 113 ++++++++++++++++++ inferencex-e2e/runners/slurm_utils.sh | 18 +-- .../runners/srt-slurm/patches/README.md | 2 +- 8 files changed, 239 insertions(+), 26 deletions(-) create mode 100644 inferencex-e2e/infx/tests/srt_slurm/test_fixed_sequence_client.py diff --git a/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/RECIPES.md b/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/RECIPES.md index 371494a2ea..5d846d6c3b 100644 --- a/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/RECIPES.md +++ b/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/RECIPES.md @@ -29,7 +29,7 @@ Shared runtime assets stay under `configs/` beside the model directories; they a ## TileRT exception -For `FRAMEWORK=tilert`, `setup_srt_slurm()` fetches the SemiAnalysisAI/srt-slurm fork directly at `6bc3f306bdafa1edfb5dded2fcda8f1ccede1bde` into the job checkout. This is the schema-2 TileRT port in [SemiAnalysisAI/srt-slurm#13](https://github.com/SemiAnalysisAI/srt-slurm/pull/13). It is the only alternate checkout; its pin lives in that helper because the TileRT backend and router are absent from the NVIDIA pin. TileRT uses the same schema-2 recipe layout and native post-eval dispatch as NVIDIA. TileRT jobs need network access to the fork at setup time. Remove the fork exception once those features are available upstream. +For `FRAMEWORK=tilert`, `setup_srt_slurm()` checks out the pinned [per-role engine and TileRT stack](https://github.com/SemiAnalysisAI/srt-slurm/pull/29) from the SemiAnalysisAI fork. Recipes select vLLM prefill and TileRT decode explicitly, with separate images and the native `tilert-router` frontend. Other engines continue using the pinned NVIDIA submodule. ## Schema 2 and master configuration diff --git a/inferencex-e2e/benchmarks/multi_node/srt_fixed_sequence.sh b/inferencex-e2e/benchmarks/multi_node/srt_fixed_sequence.sh index b0227f97cd..8108573b95 100644 --- a/inferencex-e2e/benchmarks/multi_node/srt_fixed_sequence.sh +++ b/inferencex-e2e/benchmarks/multi_node/srt_fixed_sequence.sh @@ -21,8 +21,12 @@ esac repo_root="$(dirname "${BASH_SOURCE[0]}")/../.." # Request the name the workers registered; the workflow's MODEL is the HF id, which can differ. -model=$(curl -sf "http://${SRT_FRONTEND_HOST}:${SRT_FRONTEND_PORT}/v1/models" | - python3 -c 'import json, sys; print(json.load(sys.stdin)["data"][0]["id"])') +if [[ -n "${BENCHMARK_SERVED_MODEL_NAME:-}" ]]; then + model="$BENCHMARK_SERVED_MODEL_NAME" +else + model=$(curl -sf "http://${SRT_FRONTEND_HOST}:${SRT_FRONTEND_PORT}/v1/models" | + python3 -c 'import json, sys; print(json.load(sys.stdin)["data"][0]["id"])') +fi result_dir="/logs/sa-bench_isl_${ISL}_osl_${OSL}" mkdir -p "$result_dir" ctx=$((PREFILL_NUM_WORKERS * PREFILL_TP)) @@ -41,7 +45,7 @@ for concurrency in $CONC_LIST; do --random-range-ratio "${RANDOM_RANGE_RATIO:-0.8}" \ --random-num-workers 1 \ --num-warmups "$((concurrency * 2))" \ - --num-prompts "$((concurrency * 10))" \ + --num-prompts "${NUM_PROMPTS:-$((concurrency * 10))}" \ --max-concurrency "$concurrency" \ --request-rate inf \ --ignore-eos \ diff --git a/inferencex-e2e/benchmarks/srt_agentic.sh b/inferencex-e2e/benchmarks/srt_agentic.sh index 09aedbc1b9..d68c5b9080 100644 --- a/inferencex-e2e/benchmarks/srt_agentic.sh +++ b/inferencex-e2e/benchmarks/srt_agentic.sh @@ -52,6 +52,10 @@ echo "Using srt-slurm frontend endpoint: $AIPERF_SERVER_URL" # A router frontend does not re-export engine metrics; read them from each worker. if [[ -z "${AIPERF_SERVER_METRICS_URLS:-}" && "${SRTCTL_FRONTEND_TYPE:-}" != dynamo ]]; then endpoints="${SRT_AGG_ENDPOINTS:-${SRT_PREFILL_ENDPOINTS:+$SRT_PREFILL_ENDPOINTS,}${SRT_DECODE_ENDPOINTS:-}}" + if [[ "${SRTCTL_FRONTEND_TYPE:-}" == tilert-router ]]; then + # TileRT decode has no Prometheus endpoint. + endpoints="${SRT_PREFILL_ENDPOINTS:-}" + fi if [[ -n "${endpoints%,}" ]]; then AIPERF_SERVER_METRICS_URLS=$(sed -E 's#([^,]+)#http://\1/metrics#g' <<< "${endpoints%,}") export AIPERF_SERVER_METRICS_URLS diff --git a/inferencex-e2e/infx/srt_slurm/synthetic_acceptance.py b/inferencex-e2e/infx/srt_slurm/synthetic_acceptance.py index 45a9f2e153..12dfd6fdb7 100644 --- a/inferencex-e2e/infx/srt_slurm/synthetic_acceptance.py +++ b/inferencex-e2e/infx/srt_slurm/synthetic_acceptance.py @@ -28,6 +28,7 @@ "trt": "trtllm", "dynamo-trt": "trtllm", "atom": "atom", + "tilert": "tilert", } SGLANG_VARIABLES = ( "SGLANG_SIMULATE_ACC_LEN", @@ -35,10 +36,15 @@ "SGLANG_SIMULATE_ACC_TOKEN_MODE", ) TRT_VARIABLE = "TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS" +TILERT_VARIABLES = ("TILERT_SIMULATE_ACC_LEN", "TILERT_SIMULATE_ACC_METHOD") def spec_parameters(role: Mapping[str, Any], engine: str) -> dict[str, Any]: args = role.get("args", {}) + if engine == "tilert": + if args.get("with-mtp") is not True: + return {} + return {"method": "mtp", "num_speculative_tokens": args.get("num-mtp")} if engine == "atom": method = args.get("method") if not method: @@ -116,7 +122,11 @@ def build_overrides( environment["MODEL_PREFIX"], spec, environment["THINKING_MODE"], golden_dir ) overrides = [] - variables = {"sglang": SGLANG_VARIABLES, "trtllm": (TRT_VARIABLE,)}.get(engine, ()) + variables = { + "sglang": SGLANG_VARIABLES, + "trtllm": (TRT_VARIABLE,), + "tilert": TILERT_VARIABLES, + }.get(engine, ()) # SRT applies recipe-wide environment after role environment. Keep simulation # role-local so global values cannot override the golden AL or leak into evals. for key in variables: @@ -126,13 +136,27 @@ def build_overrides( if name not in ("agg", "prefill", "decode"): continue prefix = f"roles.{name}" - worker_spec = spec_parameters(role, engine) - if engine == "vllm": + worker_engine = engine + worker_al = al + if engine == "tilert": + selected = role.get("engine", recipe.get("engine", "tilert")) + worker_engine = selected.get("type") if isinstance(selected, Mapping) else selected + # TileRT's first token and draft cache come from real vLLM prefill. + # Only its decode runtime simulates acceptance, using environment + # variables (decode_server has no --simulate-acc-* CLI options). + if name != "decode" or worker_engine != "tilert": + worker_al = None + worker_spec = spec_parameters(role, worker_engine) + if engine == "tilert" and (worker_al is None or not worker_spec): + for key in TILERT_VARIABLES: + if key in (role.get("env") or {}): + overrides += ["--unset", f"{prefix}.env.{key}"] + if worker_engine == "vllm": if not worker_spec: continue - if al is not None: + if worker_al is not None: worker_spec.update( - rejection_sample_method="synthetic", synthetic_acceptance_length=al + rejection_sample_method="synthetic", synthetic_acceptance_length=worker_al ) elif ( worker_spec.get("rejection_sample_method") == "synthetic" @@ -146,22 +170,24 @@ def build_overrides( "--set", f"{prefix}.args.speculative-config={json.dumps(worker_spec)}", ] - elif engine == "atom": + elif worker_engine == "atom": # ATOM forces acceptance with a server flag rather than environment. key = "spec-decode-acceptance-length" - if al is not None and worker_spec: - overrides += ["--set", f"{prefix}.args.{key}={al:g}"] + if worker_al is not None and worker_spec: + overrides += ["--set", f"{prefix}.args.{key}={worker_al:g}"] elif key in (role.get("args") or {}): overrides += ["--unset", f"{prefix}.args.{key}"] - elif al is not None and worker_spec: + elif worker_al is not None and worker_spec: values = ( - (f"{al:g}", "match-expected", "real-draft-token") + (f"{worker_al:g}", "match-expected", "real-draft-token") if engine == "sglang" - else (f"{al - 1:g}",) + else (f"{worker_al:g}", "match-expected") + if engine == "tilert" + else (f"{worker_al - 1:g}",) ) for key, value in zip(variables, values, strict=True): overrides += ["--set", f"{prefix}.env.{key}={json.dumps(value)}"] - else: + elif engine != "tilert": for key in variables: if key in (role.get("env") or {}): overrides += ["--unset", f"{prefix}.env.{key}"] diff --git a/inferencex-e2e/infx/tests/srt_slurm/test_fixed_sequence_client.py b/inferencex-e2e/infx/tests/srt_slurm/test_fixed_sequence_client.py new file mode 100644 index 0000000000..81b6c969f7 --- /dev/null +++ b/inferencex-e2e/infx/tests/srt_slurm/test_fixed_sequence_client.py @@ -0,0 +1,66 @@ +"""Exercise the custom benchmark shell entrypoint without launching a client.""" + +import json +import os +import subprocess +import sys +from pathlib import Path + +import pytest + + +@pytest.mark.parametrize("explicit", [True, False]) +def test_client_model_discovery_and_explicit_request_count(tmp_path, explicit): + binaries = tmp_path / "bin" + binaries.mkdir() + scripts = { + "python3": ( + f"#!{sys.executable}\n" + "import json, os, sys\n" + "from pathlib import Path\n" + "if sys.argv[1:2] == ['-c']:\n" + f" os.execv({sys.executable!r}, [{sys.executable!r}, *sys.argv[1:]])\n" + "Path(os.environ['CLIENT_ARGS_FILE']).write_text(json.dumps(sys.argv[1:]))\n" + ), + "curl": ( + "#!/bin/sh\n" + 'printf called > "$CURL_CALLED_FILE"\n' + "printf '%s\\n' '{\"data\":[{\"id\":\"discovered-model\"}]}'\n" + ), + # No real client writes results; avoid creating its container-only /logs directory. + "mkdir": "#!/bin/sh\nexit 0\n", + } + for name, source in scripts.items(): + binary = binaries / name + binary.write_text(source) + binary.chmod(0o755) + args_file = tmp_path / "client-args.json" + curl_called = tmp_path / "curl-called" + env = { + "PATH": f"{binaries}{os.pathsep}{os.environ['PATH']}", + "CLIENT_ARGS_FILE": str(args_file), + "CURL_CALLED_FILE": str(curl_called), + "ISL": "1024", + "OSL": "1024", + "SRT_FRONTEND_HOST": "router", + "SRT_FRONTEND_PORT": "8123", + "CONC_LIST": "1", + "PREFILL_NUM_WORKERS": "1", + "PREFILL_TP": "8", + "DECODE_NUM_WORKERS": "1", + "DECODE_TP": "8", + "CLIENT_BACKEND": "openai-chat", + "USE_CHAT_TEMPLATE": "true", + "SERVED_MODEL_NAME": "launcher-model-not-client-override", + } + if explicit: + env.update(BENCHMARK_SERVED_MODEL_NAME="glm5", NUM_PROMPTS="16") + script = Path(__file__).resolve().parents[3] / "benchmarks/multi_node/srt_fixed_sequence.sh" + subprocess.run(["bash", str(script)], env=env, cwd=tmp_path, check=True, capture_output=True) + args = json.loads(args_file.read_text()) + assert args[:3] == ["-P", "-m", "infx.bench_serving.benchmark_serving"] + assert args[args.index("--model") + 1] == ("glm5" if explicit else "discovered-model") + assert args[args.index("--num-prompts") + 1] == ("16" if explicit else "10") + assert args[args.index("--endpoint") + 1] == "/v1/chat/completions" + assert "--use-chat-template" in args + assert curl_called.exists() is not explicit diff --git a/inferencex-e2e/infx/tests/srt_slurm/test_synthetic_acceptance.py b/inferencex-e2e/infx/tests/srt_slurm/test_synthetic_acceptance.py index fc181fdfd1..9dad4dd03b 100644 --- a/inferencex-e2e/infx/tests/srt_slurm/test_synthetic_acceptance.py +++ b/inferencex-e2e/infx/tests/srt_slurm/test_synthetic_acceptance.py @@ -50,6 +50,7 @@ def golden_dir(tmp_path: Path) -> Path: ), ("minimaxm3_eagle3.yaml", "minimax-m3", 2.5), ("minimaxm3_eagle3_gqa.yaml", "minimax-m3", 2.6), + ("glm5.3_mtp.yaml", "glm-5.3", 3.2), ]: (directory / filename).write_text( yaml.safe_dump( @@ -289,6 +290,118 @@ def test_atom_forces_golden_acceptance_by_server_flag( assert "spec-decode-acceptance-length" not in evaluated["roles"]["agg"]["args"] +def tilert_recipe() -> dict[str, Any]: + return { + "roles": { + "prefill": { + "engine": "vllm", + "args": { + "speculative-config": '{"method":"mtp","num_speculative_tokens":1}', + }, + }, + "decode": { + "engine": {"type": "tilert"}, + "args": {"with-mtp": True, "num-mtp": 3}, + "env": {"GLM5_AR_N": "2"}, + }, + }, + } + + +def test_tilert_plan_uses_caller_decode_depth_and_keeps_prefill_real( + tmp_path: Path, golden_dir: Path +) -> None: + recipe = tilert_recipe() + recipe["roles"]["decode"]["args"]["num-mtp"] = 2 + recipe["environment"] = { + "KEEP": "global", + "TILERT_SIMULATE_ACC_LEN": "99", + "TILERT_SIMULATE_ACC_METHOD": "stale-method", + } + path = tmp_path / "tilert.yaml" + path.write_text(yaml.safe_dump(recipe)) + commands = plan_commands( + str(path), + "tilert", + ["--set", "roles.decode.args.num-mtp=3"], + {**ENV, "MODEL_PREFIX": "glm5.3", "RUN_EVAL": "true"}, + golden_dir=golden_dir, + ) + assert len(commands) == 1 + result = apply_native(recipe, commands[0]) + assert result["roles"]["decode"]["env"] == { + "GLM5_AR_N": "2", + "TILERT_SIMULATE_ACC_LEN": "3.2", + "TILERT_SIMULATE_ACC_METHOD": "match-expected", + } + assert result["roles"]["decode"]["args"] == {"with-mtp": True, "num-mtp": 3} + assert result["environment"] == {"KEEP": "global"} + assert json.loads(result["roles"]["prefill"]["args"]["speculative-config"]) == { + "method": "mtp", + "num_speculative_tokens": 1, + } + + +@pytest.mark.parametrize( + "environment", + [{"EVAL_ONLY": "true"}, {"IS_AGENTIC": "0"}, {"SPEC_DECODING": "none"}, {}], +) +def test_tilert_real_verification_removes_stale_role_and_global_simulation( + tmp_path: Path, environment: dict[str, str] +) -> None: + recipe = tilert_recipe() + if not environment: + recipe["roles"]["decode"]["args"]["with-mtp"] = False + stale = {"TILERT_SIMULATE_ACC_LEN": "99", "TILERT_SIMULATE_ACC_METHOD": "match-expected"} + recipe["environment"] = {"KEEP": "global", **stale} + for role in recipe["roles"].values(): + role.setdefault("env", {}).update(stale) + recipe["roles"]["prefill"]["engine"] = {"type": "vllm"} + recipe["roles"]["prefill"]["args"]["speculative-config"] = json.dumps( + { + "method": "mtp", + "num_speculative_tokens": 1, + "rejection_sample_method": "synthetic", + "synthetic_acceptance_length": 99, + } + ) + result = apply_native( + recipe, + build_overrides( + recipe, + "tilert", + {**ENV, "MODEL_PREFIX": "glm5.3", **environment}, + golden_dir=tmp_path / "missing", + ), + ) + assert result["environment"] == {"KEEP": "global"} + assert result["roles"]["decode"]["env"] == {"GLM5_AR_N": "2"} + assert result["roles"]["prefill"]["env"] == {} + assert json.loads(result["roles"]["prefill"]["args"]["speculative-config"]) == { + "method": "mtp", + "num_speculative_tokens": 1, + "rejection_sample_method": "block", + } + + +@pytest.mark.parametrize("depth", [None, 7]) +def test_tilert_requires_explicit_measured_decode_depth( + golden_dir: Path, depth: int | None +) -> None: + recipe = tilert_recipe() + if depth is None: + del recipe["roles"]["decode"]["args"]["num-mtp"] + else: + recipe["roles"]["decode"]["args"]["num-mtp"] = depth + with pytest.raises(ValueError, match="positive integer draft length|No golden acceptance"): + build_overrides( + recipe, + "tilert", + {**ENV, "MODEL_PREFIX": "glm5.3"}, + golden_dir=golden_dir, + ) + + @pytest.mark.parametrize( "curve", [ diff --git a/inferencex-e2e/runners/slurm_utils.sh b/inferencex-e2e/runners/slurm_utils.sh index 50370edfeb..3a447e19d5 100644 --- a/inferencex-e2e/runners/slurm_utils.sh +++ b/inferencex-e2e/runners/slurm_utils.sh @@ -58,8 +58,8 @@ PYENV # native recipe environment and benchmark.env retain their override priority. local source="$INFERENCEX_SLURM_UTILS_DIR/../utils/srt-slurm" if [[ "$framework" == "tilert" ]]; then - # TileRT still needs its legacy runtime until the native backend and router land. - SRT_SLURM_COMMIT=6bc3f306bdafa1edfb5dded2fcda8f1ccede1bde + # Per-role engines and the TileRT router: SemiAnalysisAI/srt-slurm#33 and #29. + SRT_SLURM_COMMIT=f2b1e91d4193dab62ee568d70a85fc5587c72680 git init --quiet "$destination" || return 1 git -C "$destination" remote add origin https://github.com/SemiAnalysisAI/srt-slurm.git || return 1 git -C "$destination" fetch --quiet --depth=1 origin "$SRT_SLURM_COMMIT" || return 1 @@ -70,16 +70,16 @@ PYENV return 1 fi SRT_SLURM_COMMIT=$(git -C "$source" rev-parse HEAD) || return 1 - SRTCTL_EVAL_ARGS+=(--set benchmark.stream_output=true) # A local clone keeps job writes isolated and preserves upstream Git provenance. git -c advice.detachedHead=false clone --quiet --no-hardlinks "$source" "$destination" || return 1 - # Temporary fixes awaiting upstream merge; see runners/srt-slurm/patches/README.md. - local patch - for patch in "$GITHUB_WORKSPACE"/runners/srt-slurm/patches/*.patch; do - [[ -e "$patch" ]] || continue - git -C "$destination" apply "$patch" || return 1 - done fi + SRTCTL_EVAL_ARGS+=(--set benchmark.stream_output=true) + # Temporary fixes awaiting upstream merge; see runners/srt-slurm/patches/README.md. + local patch + for patch in "$GITHUB_WORKSPACE"/runners/srt-slurm/patches/*.patch; do + [[ -e "$patch" ]] || continue + git -C "$destination" apply "$patch" || return 1 + done cd "$destination" || return 1 [[ "$(git rev-parse HEAD)" == "$SRT_SLURM_COMMIT" ]] || return 1 echo "Using srt-slurm revision $SRT_SLURM_COMMIT" diff --git a/inferencex-e2e/runners/srt-slurm/patches/README.md b/inferencex-e2e/runners/srt-slurm/patches/README.md index cb619d0506..86001db4b0 100644 --- a/inferencex-e2e/runners/srt-slurm/patches/README.md +++ b/inferencex-e2e/runners/srt-slurm/patches/README.md @@ -1,6 +1,6 @@ # srt-slurm patches -`setup_srt_slurm()` in [`runners/slurm_utils.sh`](../../slurm_utils.sh) applies every `*.patch` here to the job's srt-slurm clone after checking out the pinned submodule. TileRT jobs use the fork checkout and skip these patches. +`setup_srt_slurm()` in [`runners/slurm_utils.sh`](../../slurm_utils.sh) applies every `*.patch` here to the job's srt-slurm clone. TileRT jobs use the pinned fork stack; other jobs use the pinned submodule. Each patch is a temporary fix for an open upstream PR. When the PR merges and the submodule pin includes it, delete the patch and its row. From a98149c7f341b30daa5f16da876f453abe1fd46d Mon Sep 17 00:00:00 2001 From: Cam Quilici Date: Mon, 28 Sep 2026 20:38:50 -0500 Subject: [PATCH 2/2] fix(tilert): use native worker metrics discovery --- inferencex-e2e/benchmarks/srt_agentic.sh | 4 ---- inferencex-e2e/runners/slurm_utils.sh | 2 +- 2 files changed, 1 insertion(+), 5 deletions(-) diff --git a/inferencex-e2e/benchmarks/srt_agentic.sh b/inferencex-e2e/benchmarks/srt_agentic.sh index d68c5b9080..09aedbc1b9 100644 --- a/inferencex-e2e/benchmarks/srt_agentic.sh +++ b/inferencex-e2e/benchmarks/srt_agentic.sh @@ -52,10 +52,6 @@ echo "Using srt-slurm frontend endpoint: $AIPERF_SERVER_URL" # A router frontend does not re-export engine metrics; read them from each worker. if [[ -z "${AIPERF_SERVER_METRICS_URLS:-}" && "${SRTCTL_FRONTEND_TYPE:-}" != dynamo ]]; then endpoints="${SRT_AGG_ENDPOINTS:-${SRT_PREFILL_ENDPOINTS:+$SRT_PREFILL_ENDPOINTS,}${SRT_DECODE_ENDPOINTS:-}}" - if [[ "${SRTCTL_FRONTEND_TYPE:-}" == tilert-router ]]; then - # TileRT decode has no Prometheus endpoint. - endpoints="${SRT_PREFILL_ENDPOINTS:-}" - fi if [[ -n "${endpoints%,}" ]]; then AIPERF_SERVER_METRICS_URLS=$(sed -E 's#([^,]+)#http://\1/metrics#g' <<< "${endpoints%,}") export AIPERF_SERVER_METRICS_URLS diff --git a/inferencex-e2e/runners/slurm_utils.sh b/inferencex-e2e/runners/slurm_utils.sh index 3a447e19d5..bc632b37eb 100644 --- a/inferencex-e2e/runners/slurm_utils.sh +++ b/inferencex-e2e/runners/slurm_utils.sh @@ -59,7 +59,7 @@ PYENV local source="$INFERENCEX_SLURM_UTILS_DIR/../utils/srt-slurm" if [[ "$framework" == "tilert" ]]; then # Per-role engines and the TileRT router: SemiAnalysisAI/srt-slurm#33 and #29. - SRT_SLURM_COMMIT=f2b1e91d4193dab62ee568d70a85fc5587c72680 + SRT_SLURM_COMMIT=b421420bc8ba6a4fa7a307cf5613350ffe5ce8f3 git init --quiet "$destination" || return 1 git -C "$destination" remote add origin https://github.com/SemiAnalysisAI/srt-slurm.git || return 1 git -C "$destination" fetch --quiet --depth=1 origin "$SRT_SLURM_COMMIT" || return 1