diff --git a/.github/workflows/e2e-gpu-job.yml b/.github/workflows/e2e-gpu-job.yml index 729dd61997..b5abed1381 100644 --- a/.github/workflows/e2e-gpu-job.yml +++ b/.github/workflows/e2e-gpu-job.yml @@ -39,6 +39,11 @@ on: type: string default: "" description: "pytest -k filter expression" + extra_models: + required: false + type: string + default: "" + description: "Extra model ids to download by id (space-separated), for skip_tier_download models" setup_agentic_deps: required: false type: boolean @@ -126,6 +131,15 @@ jobs: - name: Download models run: bash scripts/ci_download_model.sh --gpu-tier ${{ inputs.gpu_tier }} + - name: Download extra models + if: inputs.extra_models != '' + # Pass the input via env (not interpolated into the run script) so its + # text can never become shell syntax; unquoted $EXTRA_MODELS still + # word-splits into multiple space-separated model ids. + env: + EXTRA_MODELS: ${{ inputs.extra_models }} + run: bash scripts/ci_download_model.sh $EXTRA_MODELS + # Run tests - name: Run E2E tests timeout-minutes: ${{ inputs.test_timeout }} diff --git a/.github/workflows/pr-test-rust.yml b/.github/workflows/pr-test-rust.yml index 9a24ee6607..4e5d11dee5 100644 --- a/.github/workflows/pr-test-rust.yml +++ b/.github/workflows/pr-test-rust.yml @@ -732,6 +732,20 @@ jobs: test_filter: "-k TestIGWMixedWorkerClassification" secrets: inherit + e2e-4gpu-epd: + name: e2e-4gpu-epd (tokenspeed) + needs: [e2e-1gpu-chat] + uses: ./.github/workflows/e2e-gpu-job.yml + with: + engine: tokenspeed + gpu_tier: "4" + runner: 4-gpu-h100 + timeout: 75 + test_timeout: 65 + test_dirs: e2e_test/chat_completions/test_epd_multimodal.py + extra_models: "Qwen/Qwen3.5-9B" + secrets: inherit + # --- Vendor E2E: CPU-only cloud backend tests --- e2e-vendor: @@ -1023,7 +1037,7 @@ jobs: path: benchmark_go_bindings/ finish: - needs: [pre-commit, python-lint, grpc-proto-build-check, build-wheel, python-unit-tests, unit-tests, benchmarks, e2e-1gpu-chat, e2e-1gpu-completions, e2e-1gpu-embeddings, e2e-1gpu-gateway, e2e-1gpu-responses, e2e-2gpu-pd, e2e-4gpu-chat, e2e-4gpu-gateway, e2e-vendor, go-unit-tests, go-bindings-e2e] + needs: [pre-commit, python-lint, grpc-proto-build-check, build-wheel, python-unit-tests, unit-tests, benchmarks, e2e-1gpu-chat, e2e-1gpu-completions, e2e-1gpu-embeddings, e2e-1gpu-gateway, e2e-1gpu-responses, e2e-2gpu-pd, e2e-4gpu-chat, e2e-4gpu-gateway, e2e-4gpu-epd, e2e-vendor, go-unit-tests, go-bindings-e2e] if: always() runs-on: k8s-runner-cpu permissions: {} @@ -1045,6 +1059,7 @@ jobs: "${{ needs.e2e-2gpu-pd.result }}" == "failure" || \ "${{ needs.e2e-4gpu-chat.result }}" == "failure" || \ "${{ needs.e2e-4gpu-gateway.result }}" == "failure" || \ + "${{ needs.e2e-4gpu-epd.result }}" == "failure" || \ "${{ needs.e2e-vendor.result }}" == "failure" || \ "${{ needs.go-unit-tests.result }}" == "failure" || \ "${{ needs.go-bindings-e2e.result }}" == "failure" ]]; then diff --git a/e2e_test/chat_completions/test_epd_multimodal.py b/e2e_test/chat_completions/test_epd_multimodal.py new file mode 100644 index 0000000000..204fdddb67 --- /dev/null +++ b/e2e_test/chat_completions/test_epd_multimodal.py @@ -0,0 +1,151 @@ +"""EPD (Encode-Prefill-Decode) multimodal Chat Completions E2E tests. + +Exercises TokenSpeed's EPD disaggregation on a vision-language model across four +worker-count topologies: the encode worker runs the vision tower, prefill/decode +run the LM, and the gateway stitches encode -> prefill -> decode. + +The gateway runs EPD-only (``RoutingMode::EncodePrefillDecode``) — there is no +single-worker fallback path — so a *correct* answer about the image proves the +encode->prefill->decode pipeline ran for that request: the model cannot name the +color/animal unless the encoder's embeddings reached prefill+decode. As an extra +per-request check we assert the router logged a fresh EPD encode dispatch. Qwen3.5 +is a thinking model, so the answer may arrive in ``reasoning_content`` rather than +``content`` — we accept either channel. + +Usage: + pytest e2e_test/chat_completions/test_epd_multimodal.py -v +""" + +from __future__ import annotations + +import base64 +import logging +import tempfile +import time +from pathlib import Path + +import pytest +from infra.pd_logs import LOG_FLUSH_TIMEOUT_S, read_logs + +logger = logging.getLogger(__name__) + +FIXTURES_DIR = Path(__file__).parent.parent / "fixtures" / "images" +DOG_IMAGE_PATH = FIXTURES_DIR / "dog.jpg" # Black labrador puppy (checked in) +PUG_IMAGE_PATH = FIXTURES_DIR / "pug.jpg" # Pug in a blanket (checked in) + +# Solid 32x32 RGB PNGs (from docs/guides/epd-ts-test.md). A solid color is the +# most deterministic vision check: the model can only name it if the encoder's +# pixels actually reached the LM through prefill+decode. +RED_PNG_B64 = ( + "iVBORw0KGgoAAAANSUhEUgAAACAAAAAgCAIAAAD8GO2jAAAAKElEQVR4nO3NsQ0AAAzCMP5/" + "un0CNkuZ41wybXsHAAAAAAAAAAAAxR4yw/wuPL6QkAAAAABJRU5ErkJggg==" +) +BLUE_PNG_B64 = ( + "iVBORw0KGgoAAAANSUhEUgAAACAAAAAgCAIAAAD8GO2jAAAAJklEQVR4nO3NsQkAAAjAsP7/" + "tF7hIASyp5pjAoFAIBAIBAKB4EmwOkv8Lm7+zY4AAAAASUVORK5CYII=" +) + +_LOG_DIR = Path(tempfile.mkdtemp(prefix="smg-e2e-epd-")) +# Per-request, router-side proof that the gateway routed through the encode stage. +# The worker-side "EPD encode: accepted" INFO line does not survive TokenSpeed's +# logging reconfiguration; this router line does, and the gateway runs at +# log_level=debug (below). +EPD_DISPATCH_MARKER = "EPD encode dispatch issued" + +# (encode, prefill, decode) worker counts. Every worker is tp=1, so 1e1p1d uses +# 3 GPUs and the rest use 4 — all fit the 4-GPU runner. Counts ride in the param +# because setup_backend is class-scoped and can't read per-param marks. +_EPD_TOPOLOGIES = [ + pytest.param(("epd_grpc", (1, 1, 1)), id="1e1p1d"), + pytest.param(("epd_grpc", (1, 2, 1)), id="1e2p1d"), + pytest.param(("epd_grpc", (2, 1, 1)), id="2e1p1d"), + pytest.param(("epd_grpc", (1, 1, 2)), id="1e1p2d"), +] + + +def _file_to_data_url(path: Path) -> str: + data = base64.b64encode(path.read_bytes()).decode("utf-8") + return f"data:image/jpeg;base64,{data}" + + +def _b64_png_to_data_url(data_b64: str) -> str: + return f"data:image/png;base64,{data_b64}" + + +def _epd_dispatch_count() -> int: + """How many EPD encode dispatches the router has logged so far (cumulative).""" + return read_logs(_LOG_DIR, "smg*").count(EPD_DISPATCH_MARKER) + + +@pytest.mark.engine("tokenspeed") +@pytest.mark.gpu(4) +@pytest.mark.e2e +@pytest.mark.model("Qwen/Qwen3.5-9B") +@pytest.mark.gateway(log_level="debug", policy="cache_aware", log_dir=str(_LOG_DIR)) +@pytest.mark.parametrize("setup_backend", _EPD_TOPOLOGIES, indirect=True) +class TestEPDMultimodal: + """Verify the image really flows encode -> prefill -> decode for each topology.""" + + def _check_image(self, client, model, image_url, question, keywords): + """Send one image request; assert a correct answer + a fresh EPD dispatch.""" + # Baseline BEFORE the request: the router log accumulates across topologies + # and is never cleared, so assert THIS request added a new dispatch. + dispatches_before = _epd_dispatch_count() + + response = client.chat.completions.create( + model=model, + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": question}, + {"type": "image_url", "image_url": {"url": image_url}}, + ], + } + ], + temperature=0, + max_tokens=256, + ) + + # (1) The EPD pipeline produced output. A thinking model may put the answer + # in reasoning_content instead of content — accept either. Dump the whole + # message on failure so we can see where (if anywhere) the answer landed. + msg = response.choices[0].message.model_dump() + content = msg.get("content") or "" + reasoning = msg.get("reasoning_content") or msg.get("reasoning") or "" + text = f"{content}\n{reasoning}".strip() + assert text, f"EPD pipeline returned empty content AND reasoning; message={msg}" + + # (2) The answer is correct about the image — only possible if the encoder's + # embeddings reached prefill+decode (EPD-only gateway; no fallback path). + assert any(k in text.lower() for k in keywords), ( + f"expected one of {keywords} in the answer, got: {text!r}" + ) + + # (3) The router logged a fresh EPD encode dispatch for THIS request. + deadline = time.monotonic() + LOG_FLUSH_TIMEOUT_S + while _epd_dispatch_count() <= dispatches_before and time.monotonic() < deadline: + time.sleep(0.5) + assert _epd_dispatch_count() > dispatches_before, ( + "router logged no new EPD encode dispatch for this request; the gateway " + f"did not route through encode->prefill->decode (checked {_LOG_DIR}/smg*)" + ) + logger.info("EPD OK (%s): %s", keywords[0], text[:120]) + + def test_color_images(self, model, setup_backend): + """Solid red then blue — the most deterministic encode->decode check.""" + _, _, client, *_ = setup_backend + question = "What color is this image? Reply with just the color." + self._check_image(client, model, _b64_png_to_data_url(RED_PNG_B64), question, ["red"]) + self._check_image(client, model, _b64_png_to_data_url(BLUE_PNG_B64), question, ["blue"]) + + def test_animal_images(self, model, setup_backend): + """Real photos — dog then pug.""" + _, _, client, *_ = setup_backend + question = "What animal is in this image?" + self._check_image( + client, model, _file_to_data_url(DOG_IMAGE_PATH), question, ["dog", "puppy", "labrador"] + ) + self._check_image( + client, model, _file_to_data_url(PUG_IMAGE_PATH), question, ["pug", "dog", "puppy"] + ) diff --git a/e2e_test/fixtures/setup_backend.py b/e2e_test/fixtures/setup_backend.py index abbb34f3f2..aedce60fa8 100644 --- a/e2e_test/fixtures/setup_backend.py +++ b/e2e_test/fixtures/setup_backend.py @@ -120,7 +120,11 @@ def setup_backend(request: pytest.FixtureRequest): Returns: Tuple of ``(backend_name, model_path, client, gateway)`` """ - backend_name: str = request.param + raw_param = request.param + if isinstance(raw_param, tuple): + backend_name, epd_counts = raw_param + else: + backend_name, epd_counts = raw_param, None if os.environ.get(ENV_SKIP_BACKEND_SETUP, "").lower() in ("1", "true", "yes"): pytest.skip(f"{ENV_SKIP_BACKEND_SETUP} is set") @@ -137,8 +141,9 @@ def setup_backend(request: pytest.FixtureRequest): return # Local backends + is_epd = backend_name.startswith("epd_") is_pd = backend_name.startswith("pd_") - protocol = backend_name.replace("pd_", "") + protocol = backend_name.replace("epd_", "").replace("pd_", "") connection_mode = ConnectionMode(protocol) engine = get_runtime() model_path = get_model_spec(model_id)["model"] @@ -154,7 +159,18 @@ def setup_backend(request: pytest.FixtureRequest): gateway = Gateway() try: - if is_pd: + if is_epd: + yield from _setup_epd( + model_id, + model_path, + engine, + connection_mode, + epd_counts, + gateway_config, + gateway, + log_dir, + ) + elif is_pd: yield from _setup_pd( model_id, model_path, @@ -299,6 +315,94 @@ def _setup_pd( stop_workers(all_workers) +# --------------------------------------------------------------------------- +# EPD (encode-prefill-decode) disaggregation backend +# --------------------------------------------------------------------------- + + +def _setup_epd( + model_id, + model_path, + engine, + connection_mode, + epd_counts, + gateway_config, + gateway, + log_dir, +): + """Launch encode + prefill + decode workers + EPD gateway, yield, tear down. + + ``epd_counts`` is ``(n_encode, n_prefill, n_decode)``. Every worker is tp=1 + (one GPU); GPUs are assigned sequentially E -> P -> D. + """ + if not epd_counts or len(epd_counts) != 3: + raise ValueError("epd_grpc backend requires a (n_encode, n_prefill, n_decode) param") + n_encode, n_prefill, n_decode = epd_counts + spec = get_model_spec(model_id) + tp = spec.get("tp", 1) + backend_name = f"epd_{connection_mode.value}" + runtime_label = RUNTIME_LABELS.get(engine, engine) + + logger.info( + "Starting %s EPD backend: model=%s, %de + %dp + %dd", + runtime_label, + model_id, + n_encode, + n_prefill, + n_decode, + ) + + all_workers: list = [] + try: + encode_workers = _start_workers_tracked( + model_id=model_id, + engine=engine, + mode=connection_mode, + count=n_encode, + worker_type=WorkerType.ENCODE, + log_dir=log_dir, + gpu_offset=0, + gpus=1, # vision tower runs on one GPU regardless of LM tp + ) + all_workers.extend(encode_workers) + + prefill_workers = _start_workers_tracked( + model_id=model_id, + engine=engine, + mode=connection_mode, + count=n_prefill, + worker_type=WorkerType.PREFILL, + log_dir=log_dir, + gpu_offset=n_encode, + ) + all_workers.extend(prefill_workers) + + decode_workers = _start_workers_tracked( + model_id=model_id, + engine=engine, + mode=connection_mode, + count=n_decode, + worker_type=WorkerType.DECODE, + log_dir=log_dir, + gpu_offset=n_encode + n_prefill * tp, + ) + all_workers.extend(decode_workers) + + _start_gateway( + gateway, + gateway_config, + encode_workers=encode_workers, + prefill_workers=prefill_workers, + decode_workers=decode_workers, + ) + logger.info("%s EPD backend ready at %s", runtime_label, gateway.base_url) + yield backend_name, model_path, _make_openai_client(gateway), gateway + finally: + logger.info("Tearing down %s EPD backend", runtime_label) + gateway.shutdown() + stop_workers(all_workers) + + # --------------------------------------------------------------------------- # Cloud backend # --------------------------------------------------------------------------- diff --git a/e2e_test/infra/constants.py b/e2e_test/infra/constants.py index c1936663fc..ce4b523c6b 100644 --- a/e2e_test/infra/constants.py +++ b/e2e_test/infra/constants.py @@ -15,6 +15,7 @@ class WorkerType(StrEnum): """Worker specialization type.""" REGULAR = "regular" + ENCODE = "encode" PREFILL = "prefill" DECODE = "decode" diff --git a/e2e_test/infra/gateway.py b/e2e_test/infra/gateway.py index 408f55b975..145e9d2729 100644 --- a/e2e_test/infra/gateway.py +++ b/e2e_test/infra/gateway.py @@ -27,6 +27,36 @@ logger = logging.getLogger(__name__) +def build_epd_mode_args( + encode_workers: list[Worker], + prefill_workers: list[Worker], + decode_workers: list[Worker], + encode_policy: str = "consistent_hashing", +) -> list[str]: + """Build the ``smg.launch_router`` args for EPD disaggregation. + + ``--encode`` / ``--prefill`` are repeatable ``URL [BOOTSTRAP_PORT]`` flags; + ``--decode`` is a repeatable ``URL`` flag. Multimodal tensors ride inline + (SMG -> encode); encode -> prefill embeddings and prefill -> decode KV move + over Mooncake independently of this flag. + """ + args = ["--epd-disaggregation"] + for en in encode_workers: + args += ["--encode", en.base_url] + if en.bootstrap_port is not None: + args.append(str(en.bootstrap_port)) + for pf in prefill_workers: + args += ["--prefill", pf.base_url] + if pf.bootstrap_port is not None: + args.append(str(pf.bootstrap_port)) + for dc in decode_workers: + args += ["--decode", dc.base_url] + if encode_policy: + args += ["--encode-policy", encode_policy] + args += ["--multimodal-tensor-transport", "inline"] + return args + + @dataclass class WorkerInfo: """Information about a worker connected to the gateway.""" @@ -87,6 +117,8 @@ def start( model_path: str | None = None, prefill_workers: list[Worker] | None = None, decode_workers: list[Worker] | None = None, + encode_workers: list[Worker] | None = None, + encode_policy: str = "consistent_hashing", igw_mode: bool = False, cloud_backend: str | None = None, history_backend: str = "memory", @@ -101,16 +133,19 @@ def start( if self._started: raise RuntimeError("Gateway already started") - is_pd_mode = prefill_workers is not None or decode_workers is not None + is_epd_mode = encode_workers is not None + is_pd_mode = (prefill_workers is not None or decode_workers is not None) and not is_epd_mode is_regular_mode = worker_urls is not None is_igw_mode = igw_mode is_cloud_mode = cloud_backend is not None - modes_specified = sum([is_pd_mode, is_regular_mode, is_igw_mode, is_cloud_mode]) + modes_specified = sum( + [is_epd_mode, is_pd_mode, is_regular_mode, is_igw_mode, is_cloud_mode] + ) if modes_specified != 1: raise ValueError( - "Specify exactly one mode: worker_urls, prefill/decode_workers, " - "igw_mode=True, or cloud_backend" + "Specify exactly one mode: worker_urls, encode/prefill/decode_workers " + "(EPD), prefill/decode_workers (PD), igw_mode=True, or cloud_backend" ) if show_output is None: @@ -132,6 +167,23 @@ def start( extra_args=extra_args, log_msg="IGW gateway (no workers)", ) + elif is_epd_mode: + self.pd_mode = True + self.igw_mode = False + encodes = encode_workers or [] + prefills = prefill_workers or [] + decodes = decode_workers or [] + mode_args = build_epd_mode_args(encodes, prefills, decodes, encode_policy) + self._launch( + mode_args=mode_args, + timeout=timeout, + show_output=show_output, + extra_args=extra_args, + log_msg=( + f"EPD gateway ({len(encodes)} encode, " + f"{len(prefills)} prefill, {len(decodes)} decode)" + ), + ) elif is_pd_mode: self.pd_mode = True self.igw_mode = False diff --git a/e2e_test/infra/model_specs.py b/e2e_test/infra/model_specs.py index c507bcdd5c..ba8b5610fe 100644 --- a/e2e_test/infra/model_specs.py +++ b/e2e_test/infra/model_specs.py @@ -162,6 +162,32 @@ def _resolve_model_path(hf_path: str) -> str: "tp": 1, "features": ["chat", "streaming", "multimodal"], }, + # TokenSpeed EPD multimodal model. Qwen3.5-9B is a vision-language model + # (hybrid Gated DeltaNet + sparse MoE); BF16 ~18GB fits one 80GB H100 at + # tp=1, so every EPD topology (1e1p1d/1e2p1d/2e1p1d/1e1p2d) runs on the + # 4-GPU h100 runner, one worker per card. EPD is TokenSpeed-only: the encode + # worker runs the vision tower; prefill/decode run the LM. FA3 is the H100 + # attention backend (trtllm fails on H100 for this model). Disaggregation + # role flags (--disaggregation-mode etc.) are added per-role in worker.py. + "Qwen/Qwen3.5-9B": { + "model": _resolve_model_path("Qwen/Qwen3.5-9B"), + "tp": 1, + "features": ["chat", "streaming", "multimodal", "moe"], + "startup_timeout": 600, + "tokenspeed_args": [ + "--attention-backend", + "fa3", + "--max-model-len", + "8192", + "--max-num-seqs", + "4", + "--gpu-memory-utilization", + "0.8", + ], + # TokenSpeed-only (GDN + MoE arch won't load under sglang/vllm/trt), so + # keep it out of the tier-wide pre-download; the EPD job fetches it by id. + "skip_tier_download": True, + }, # Llama-4-Maverick (17B with 128 experts, FP8) - Nightly benchmarks "meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8": { "model": _resolve_model_path("meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8"), diff --git a/e2e_test/infra/test_epd_cmd_builders.py b/e2e_test/infra/test_epd_cmd_builders.py new file mode 100644 index 0000000000..91d9e2488d --- /dev/null +++ b/e2e_test/infra/test_epd_cmd_builders.py @@ -0,0 +1,114 @@ +"""Unit tests for EPD command-builder logic (no GPU, no wheel required).""" + +from __future__ import annotations + +import pytest +from infra.constants import ConnectionMode, WorkerType +from infra.gateway import build_epd_mode_args +from infra.model_specs import get_model_spec +from infra.worker import Worker + + +def test_worker_type_encode_exists(): + assert WorkerType.ENCODE == "encode" + assert WorkerType.ENCODE.value == "encode" + + +_TS_MODEL = "Qwen/Qwen3-VL-8B-Instruct" # any tokenspeed-launchable spec at HEAD + + +def _ts_worker(worker_type, bootstrap_port=None): + return Worker( + model_id=_TS_MODEL, + engine="tokenspeed", + port=50104, + gpu_ids=[0], + mode=ConnectionMode.GRPC, + worker_type=worker_type, + bootstrap_port=bootstrap_port, + dist_init_addr="127.0.0.1:29500", + ) + + +@pytest.mark.parametrize( + "worker_type,role", + [(WorkerType.ENCODE, "encode"), (WorkerType.PREFILL, "prefill"), (WorkerType.DECODE, "decode")], +) +def test_tokenspeed_disagg_flags(worker_type, role): + cmd = _ts_worker(worker_type, bootstrap_port=18995)._build_cmd() + assert "--disaggregation-mode" in cmd + assert cmd[cmd.index("--disaggregation-mode") + 1] == role + assert "--disaggregation-transfer-backend" in cmd + assert cmd[cmd.index("--disaggregation-transfer-backend") + 1] == "mooncake" + assert "--dist-init-addr" in cmd + assert cmd[cmd.index("--dist-init-addr") + 1] == "127.0.0.1:29500" + assert "--skip-server-warmup" in cmd + + +def test_encode_and_prefill_carry_bootstrap_port(): + for wt in (WorkerType.ENCODE, WorkerType.PREFILL): + cmd = _ts_worker(wt, bootstrap_port=18995)._build_cmd() + assert "--disaggregation-bootstrap-port" in cmd + assert cmd[cmd.index("--disaggregation-bootstrap-port") + 1] == "18995" + + +def test_prefill_is_eager_decode_and_prefill_cache(): + prefill = _ts_worker(WorkerType.PREFILL, bootstrap_port=1)._build_cmd() + decode = _ts_worker(WorkerType.DECODE)._build_cmd() + assert "--enforce-eager" in prefill + assert "--enable-prefix-caching" in prefill + assert "--enable-prefix-caching" in decode + assert "--enforce-eager" not in decode + + +def test_regular_tokenspeed_worker_has_no_disagg_flags(): + cmd = _ts_worker(WorkerType.REGULAR)._build_cmd() + assert "--disaggregation-mode" not in cmd + + +def _w(port, worker_type, bootstrap_port=None): + return Worker( + model_id=_TS_MODEL, + engine="tokenspeed", + port=port, + gpu_ids=[0], + mode=ConnectionMode.GRPC, + worker_type=worker_type, + bootstrap_port=bootstrap_port, + ) + + +def test_build_epd_mode_args_basic(): + args = build_epd_mode_args( + encode_workers=[_w(50104, WorkerType.ENCODE, 18995)], + prefill_workers=[_w(50101, WorkerType.PREFILL, 19311)], + decode_workers=[_w(50111, WorkerType.DECODE)], + ) + assert args[0] == "--epd-disaggregation" + assert "--encode" in args and "grpc://127.0.0.1:50104" in args + assert "18995" in args # encode bootstrap + assert "--prefill" in args and "grpc://127.0.0.1:50101" in args + assert "19311" in args # prefill bootstrap + assert "--decode" in args and "grpc://127.0.0.1:50111" in args + assert args[args.index("--encode-policy") + 1] == "consistent_hashing" + assert args[args.index("--multimodal-tensor-transport") + 1] == "inline" + + +def test_build_epd_mode_args_multi_encode_repeats_flag(): + args = build_epd_mode_args( + encode_workers=[_w(50104, WorkerType.ENCODE, 1), _w(50105, WorkerType.ENCODE, 2)], + prefill_workers=[_w(50101, WorkerType.PREFILL, 3)], + decode_workers=[_w(50111, WorkerType.DECODE)], + ) + assert args.count("--encode") == 2 + assert "grpc://127.0.0.1:50104" in args and "grpc://127.0.0.1:50105" in args + + +def test_qwen35_9b_spec_present_and_multimodal(): + spec = get_model_spec("Qwen/Qwen3.5-9B") + assert spec["tp"] == 1 + assert "multimodal" in spec["features"] + assert spec.get("skip_tier_download") is True + ts_args = spec["tokenspeed_args"] + assert "--attention-backend" in ts_args + assert ts_args[ts_args.index("--attention-backend") + 1] == "fa3" diff --git a/e2e_test/infra/worker.py b/e2e_test/infra/worker.py index 30e14caa13..e1bdb5c7ee 100644 --- a/e2e_test/infra/worker.py +++ b/e2e_test/infra/worker.py @@ -44,6 +44,8 @@ class Worker: bootstrap_port: int | None = None nixl_port: int | None = None ib_device: str | None = None + dist_init_addr: str | None = None + dist_init_port: int | None = None log_dir: str | None = None extra_engine_args: list[str] | None = None process: subprocess.Popen | None = field(default=None, repr=False) @@ -164,6 +166,9 @@ def stop(self) -> None: if self.nixl_port is not None: release_port(self.nixl_port) self.nixl_port = None + if self.dist_init_port is not None: + release_port(self.dist_init_port) + self.dist_init_port = None def is_alive(self) -> bool: """Check if the worker process is still running.""" @@ -341,6 +346,19 @@ def _build_tokenspeed_grpc_cmd(self, model_path: str, tp_size: int, spec: dict) # ``logprobs=True`` requests get real per-token data back. "--enable-output-logprobs", ] + if self.worker_type in (WorkerType.ENCODE, WorkerType.PREFILL, WorkerType.DECODE): + cmd.extend(["--disaggregation-mode", self.worker_type.value]) + if self.bootstrap_port is not None: + cmd.extend(["--disaggregation-bootstrap-port", str(self.bootstrap_port)]) + cmd.extend(["--disaggregation-transfer-backend", "mooncake"]) + if self.dist_init_addr: + cmd.extend(["--dist-init-addr", self.dist_init_addr]) + if self.worker_type in (WorkerType.PREFILL, WorkerType.DECODE): + cmd.append("--enable-prefix-caching") + if self.worker_type == WorkerType.PREFILL: + cmd.append("--enforce-eager") + cmd.append("--skip-server-warmup") + extra = spec.get("tokenspeed_args", []) if extra: cmd.extend(extra) @@ -386,6 +404,16 @@ def _build_env(self) -> dict[str, str]: env.setdefault("PYTHONUNBUFFERED", "1") env["CUDA_VISIBLE_DEVICES"] = ",".join(map(str, self.gpu_ids)) + if self.engine == "tokenspeed" and self.worker_type in ( + WorkerType.ENCODE, + WorkerType.PREFILL, + WorkerType.DECODE, + ): + env.setdefault("MC_INTRANODE_NVLINK", "1") + env.setdefault("TOKENSPEED_SKIP_GRPC_WARMUP", "1") + env.setdefault("NO_PROXY", "*") + env.setdefault("no_proxy", "*") + # vLLM PD workers need per-worker side-channel ports for their KV backend if self.engine == "vllm" and self.worker_type in (WorkerType.PREFILL, WorkerType.DECODE): if vllm_kv_backend() == "mooncake": @@ -548,7 +576,16 @@ def start_workers( gpu_ids = list(range(gpu_offset, gpu_offset + gpus_per_worker)) gpu_offset += gpus_per_worker port = get_open_port() - bootstrap_port = get_open_port() if worker_type == WorkerType.PREFILL else None + bootstrap_port = ( + get_open_port() if worker_type in (WorkerType.PREFILL, WorkerType.ENCODE) else None + ) + is_ts_disagg = engine == "tokenspeed" and worker_type in ( + WorkerType.ENCODE, + WorkerType.PREFILL, + WorkerType.DECODE, + ) + dist_init_port = get_open_port() if is_ts_disagg else None + dist_init_addr = f"{DEFAULT_HOST}:{dist_init_port}" if dist_init_port else None worker = Worker( model_id=model_id, @@ -559,6 +596,8 @@ def start_workers( worker_type=worker_type, bootstrap_port=bootstrap_port, ib_device=ib_device if has_pd else None, + dist_init_addr=dist_init_addr, + dist_init_port=dist_init_port, log_dir=log_dir, extra_engine_args=extra_engine_args, ) diff --git a/grpc_servicer/smg_grpc_servicer/tokenspeed/encoder_servicer.py b/grpc_servicer/smg_grpc_servicer/tokenspeed/encoder_servicer.py index a788bc0da1..ab00bac6d1 100644 --- a/grpc_servicer/smg_grpc_servicer/tokenspeed/encoder_servicer.py +++ b/grpc_servicer/smg_grpc_servicer/tokenspeed/encoder_servicer.py @@ -198,6 +198,11 @@ async def Encode(self, request, context): except Exception as e: # noqa: BLE001 logger.exception("TokenSpeed encode ingest failed") await context.abort(grpc.StatusCode.INTERNAL, str(e)) + logger.info( + "EPD encode: accepted request_id=%s room=%s", + request.request_id, + bootstrap_room, + ) return tokenspeed_encoder_pb2.EncodeResponse(accepted=True) def _build_encode_request(self, request, bootstrap_room): diff --git a/scripts/ci_download_model.sh b/scripts/ci_download_model.sh index 1b09fd4221..37b0e95fb0 100755 --- a/scripts/ci_download_model.sh +++ b/scripts/ci_download_model.sh @@ -44,7 +44,7 @@ resolve_models_for_tier() { import sys from e2e_test.infra.model_specs import MODEL_SPECS for model_id, spec in MODEL_SPECS.items(): - if spec['tp'] <= int(sys.argv[1]): + if spec['tp'] <= int(sys.argv[1]) and not spec.get('skip_tier_download'): print(model_id) " "$tier" } diff --git a/scripts/ci_install_tokenspeed.sh b/scripts/ci_install_tokenspeed.sh index c57a645e30..09e2593808 100755 --- a/scripts/ci_install_tokenspeed.sh +++ b/scripts/ci_install_tokenspeed.sh @@ -22,7 +22,10 @@ fi # a scheduled bump-and-CI routine) rather than floating against ``main`` — # upstream has renamed APIs before and the gRPC servicer broke until we # caught up. -TOKENSPEED_REF="${TOKENSPEED_REF:-5e145afae8e5651cd66234e68c988c31aac6639f}" +# Bumped to include the EPD encode pipeline (tokenspeed #548, b5c762d): the SMG +# encode servicer already expects `--disaggregation-mode encode`, which the old +# pin (5e145af) predated — the EPD e2e's encode worker died on "invalid choice". +TOKENSPEED_REF="${TOKENSPEED_REF:-69091e10c90c0e0f6e97c2bfdd332d61362ddd55}" TOKENSPEED_REPO="${TOKENSPEED_REPO:-https://github.com/lightseekorg/tokenspeed.git}" TOKENSPEED_DIR="${TOKENSPEED_DIR:-/tmp/tokenspeed-src}" @@ -39,35 +42,47 @@ echo "uv version: $(uv --version)" # SDK (nvcc, headers). Install them on demand — same approach as # ``ci_install_sglang.sh``. CUDA_HOME="${CUDA_HOME:-/usr/local/cuda}" -if [ ! -x "${CUDA_HOME}/bin/nvcc" ]; then - echo "Installing CUDA toolkit (nvcc not found at ${CUDA_HOME}/bin/nvcc)..." +if [ ! -x "${CUDA_HOME}/bin/nvcc" ] && [ ! -x "/usr/local/cuda-13.0/bin/nvcc" ]; then + echo "Installing CUDA toolkit (nvcc not found)..." curl -fsSL -o /tmp/cuda-keyring.deb \ https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64/cuda-keyring_1.1-1_all.deb sudo dpkg -i /tmp/cuda-keyring.deb rm /tmp/cuda-keyring.deb sudo apt-get update -qq - sudo apt-get install -y --no-install-recommends \ - cuda-nvcc-13-0 \ - cuda-cudart-dev-13-0 \ - cuda-libraries-dev-13-0 - # apt installs under /usr/local/cuda-13.0; expose the /usr/local/cuda - # alias the job-level ``CUDA_HOME: /usr/local/cuda`` env expects. - if [ ! -d "${CUDA_HOME}/bin" ] && [ -d "/usr/local/cuda-13.0/bin" ]; then - sudo ln -sfn /usr/local/cuda-13.0 "${CUDA_HOME}" - fi - echo "nvcc installed: $(${CUDA_HOME}/bin/nvcc --version | tail -1)" -else - echo "nvcc already available: $(${CUDA_HOME}/bin/nvcc --version | tail -1)" + # Install the FULL CUDA 13.0 toolkit (mirrors the proven TRT-LLM lane in + # ci_install_trtllm.sh) so the system headers -- which the kernel build + # compiles against -- are a complete, self-consistent 13.0.88 set matching + # the system nvcc. + sudo apt-get install -y cuda-toolkit-13-0 +fi +# Point CUDA_HOME at the versioned toolkit dir directly (mirrors +# ci_install_trtllm.sh). The job env sets CUDA_HOME=/usr/local/cuda, but on this +# runner that symlink is stale/partial: its include/ has cuda_runtime.h but not +# crt/host_runtime.h, so the kernel's host-stub compile falls through to torch's +# mismatched bundled crt and dies with "'__cudaLaunch' was not declared". The +# apt-installed /usr/local/cuda-13.0 is complete (ships cuda-crt-13-0). +if [ -x "/usr/local/cuda-13.0/bin/nvcc" ]; then + CUDA_HOME="/usr/local/cuda-13.0" fi export CUDA_HOME export PATH="$CUDA_HOME/bin:$PATH" export LD_LIBRARY_PATH="${CUDA_HOME}/lib64:${CUDA_HOME}/extras/CUPTI/lib64:${LD_LIBRARY_PATH:-}" -# Torch's JIT cpp_extension builder compiles some TokenSpeed runtime -# extensions (e.g. ``tokenspeed_hostfunc_ext``) with plain g++ and -# doesn't pass ``-I$CUDA_HOME/include``; expose the headers via CPATH / -# CPLUS_INCLUDE_PATH so the compile picks them up. -export CPATH="${CUDA_HOME}/include${CPATH:+:$CPATH}" -export CPLUS_INCLUDE_PATH="${CUDA_HOME}/include${CPLUS_INCLUDE_PATH:+:$CPLUS_INCLUDE_PATH}" +echo "Using CUDA_HOME=${CUDA_HOME} ($(${CUDA_HOME}/bin/nvcc --version | tail -1))" +# The kernel's launch stubs need this exact header from the system toolkit; if +# it's missing the build falls through to torch's bundled cu13 crt and fails. +if [ -f "${CUDA_HOME}/include/crt/host_runtime.h" ]; then + echo "system crt/host_runtime.h: present under CUDA_HOME" +else + echo "WARNING: ${CUDA_HOME}/include/crt/host_runtime.h is MISSING" >&2 +fi +# Torch's JIT cpp_extension builder compiles some TokenSpeed runtime extensions +# (e.g. ``tokenspeed_hostfunc_ext``) with plain g++ and doesn't pass +# ``-I$CUDA_HOME/include``; expose the system CUDA headers via CPATH so those +# g++ compiles find them (CUDA 13 keeps CCCL under ``include/cccl``). +_cuda_inc="${CUDA_HOME}/include:${CUDA_HOME}/include/cccl" +export CPATH="${_cuda_inc}${CPATH:+:$CPATH}" +export CPLUS_INCLUDE_PATH="${_cuda_inc}${CPLUS_INCLUDE_PATH:+:$CPLUS_INCLUDE_PATH}" +export C_INCLUDE_PATH="${_cuda_inc}${C_INCLUDE_PATH:+:$C_INCLUDE_PATH}" # ── Clone TokenSpeed ──────────────────────────────────────────────────────── # ``git clone --branch`` only accepts branch/tag names, not SHAs, so we @@ -94,6 +109,20 @@ sudo apt-get install -y --no-install-recommends libssl-dev libopenmpi-dev cmake # ── TokenSpeed packages ──────────────────────────────────────────────────── export MAX_JOBS="${MAX_JOBS:-16}" export FLASHINFER_CUDA_ARCH_LIST="${FLASHINFER_CUDA_ARCH_LIST:-9.0a 10.0a}" +# Select the CUDA kernel backend explicitly, as TokenSpeed's own install_deps.sh +# does on the kernel build (otherwise the native build path can differ). +export TOKENSPEED_KERNEL_BACKEND="${TOKENSPEED_KERNEL_BACKEND:-cuda}" + +# The kernel's torch cpp_extension build must link a torch built for CUDA 13. +# TokenSpeed's CI runs on a cu130 Docker base image that already ships it; the +# generic k8s runner does not, so pip/uv would pull the default PyPI torch +# (CUDA 12.x). That drops nvidia-cuda-runtime-cu12's own crt/host_runtime.h on +# the include path, and nvcc 13's cudafe++ then generates a host stub that fails +# to compile against those cu12 headers: "'__cudaLaunch' was not declared". +# Point pip/uv at the cu130 wheel index (mirrors install_deps.sh line 118) so +# every install below resolves the CUDA-13 torch + nvidia deps. +export PIP_EXTRA_INDEX_URL="${PIP_EXTRA_INDEX_URL:-https://download.pytorch.org/whl/cu130}" +export UV_EXTRA_INDEX_URL="${UV_EXTRA_INDEX_URL:-https://download.pytorch.org/whl/cu130}" # The kernel requirements leave ``nvidia-cutlass-dsl`` unpinned, and 4.6.0 # dropped ``cute.core.ThrMma`` — which quack (pulled via flash-attn's cute @@ -113,6 +142,37 @@ export PIP_CONSTRAINT="$TOKENSPEED_CONSTRAINTS" # ``build-system.requires``, and we install with ``--no-build-isolation``. uv pip install setuptools wheel pybind11 +# Install the CUDA-13 torch build explicitly (the +cu130 local wheel) before the +# --no-build-isolation kernel compile below, so the build links matching CUDA 13 +# headers instead of the default PyPI (cu12.x) torch. Pin tracks TokenSpeed's +# torch requirement; bump alongside TOKENSPEED_REF. +uv pip install "torch==2.11.0+cu130" + +# The kernel's host-stub compile binds crt/host_runtime.h from torch's bundled +# cu13 headers (site-packages/nvidia/cu*/include/crt) no matter the -I order, +# and those are a newer patch (nvidia-cuda-runtime 13.0.96) than the apt system +# nvcc (13.0.88): the 88 nvcc emits a 2-arg __cudaLaunch stub the 96 header's +# 1-arg macro can't satisfy -> "'__cudaLaunch' was not declared". Those crt dirs +# are pulled by the kernel build's own dependency resolution, so materialize +# them with a first build pass (tolerate its compile failure), realign every +# bundled crt to the system toolkit, then build for real -- deps are satisfied +# now, so nothing re-pulls the crt. +uv pip install -e tokenspeed-kernel/python/ --no-build-isolation || \ + echo "first kernel build pass failed (expected: crt skew); realigning crt headers" + +_sys_crt="${CUDA_HOME}/include/crt" +_purelib="$(python3 -c 'import sysconfig; print(sysconfig.get_path("purelib"))')" +if [ -d "$_sys_crt" ] && [ -d "$_purelib" ]; then + _aligned=0 + while IFS= read -r -d '' _pip_crt; do + echo "Aligning bundled CUDA crt to system: ${_pip_crt} -> ${_sys_crt}" + rm -rf "$_pip_crt" + ln -sfnT "$_sys_crt" "$_pip_crt" + _aligned=1 + done < <(find "$_purelib" -type d -path '*/nvidia/cu*/include/crt' -print0 2>/dev/null) + [ "$_aligned" = 1 ] || echo "WARNING: no bundled nvidia crt dirs found under ${_purelib}" >&2 +fi + uv pip install -e tokenspeed-kernel/python/ --no-build-isolation uv pip install -e tokenspeed-scheduler/ uv pip install -e "./python" --no-build-isolation @@ -125,6 +185,7 @@ if [ -n "${GITHUB_ENV:-}" ]; then # CUDA headers when it bypasses nvcc for .cpp sources. echo "CPATH=$CPATH" >> "$GITHUB_ENV" echo "CPLUS_INCLUDE_PATH=$CPLUS_INCLUDE_PATH" >> "$GITHUB_ENV" + echo "C_INCLUDE_PATH=$C_INCLUDE_PATH" >> "$GITHUB_ENV" fi if [ -n "${GITHUB_PATH:-}" ]; then # Make ``nvcc`` discoverable to downstream steps (pytest spawns the