Integrate the TileRT srt-slurm fork stack

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
 

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 \

inferencex-e2e/infx/srt_slurm/synthetic_acceptance.py

@@ -28,17 +28,23 @@
     "trt": "trtllm",
     "dynamo-trt": "trtllm",
     "atom": "atom",
+    "tilert": "tilert",
 }
 SGLANG_VARIABLES = (
     "SGLANG_SIMULATE_ACC_LEN",
     "SGLANG_SIMULATE_ACC_METHOD",
     "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}"]

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

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",
     [

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=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
@@ -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"

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.