diff --git a/.github/workflows/e2e-gpu-job.yml b/.github/workflows/e2e-gpu-job.yml index 729dd6199..9872cd6af 100644 --- a/.github/workflows/e2e-gpu-job.yml +++ b/.github/workflows/e2e-gpu-job.yml @@ -126,6 +126,13 @@ jobs: - name: Download models run: bash scripts/ci_download_model.sh --gpu-tier ${{ inputs.gpu_tier }} + # The TokenSpeed EPD model is skip_tier_download (large, engine-specific), so + # the tier step above skips it; pull it by id only for the 4-GPU tokenspeed + # lane that runs the EPD e2e — other lanes never touch it. + - name: Download EPD model (tokenspeed only) + if: inputs.engine == 'tokenspeed' && inputs.gpu_tier == '4' + run: bash scripts/ci_download_model.sh "Qwen/Qwen3.6-35B-A3B-FP8" + # 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 9a24ee660..73ce38cdc 100644 --- a/.github/workflows/pr-test-rust.yml +++ b/.github/workflows/pr-test-rust.yml @@ -710,12 +710,22 @@ jobs: timeout: 30 - engine: trtllm timeout: 45 + # EPD multimodal (test_epd_multimodal.py): the encode/prefill/decode + # TokenSpeed workers span the 4 GPUs at tp=1. Higher timeout — tokenspeed + # builds from source (~30m cold) and the four topologies each relaunch + # their worker set. + - engine: tokenspeed + timeout: 75 + # The inner pytest step has its own timeout (default 25m); the four EPD + # topologies each relaunch a worker set, so give pytest more room too. + test_timeout: 65 uses: ./.github/workflows/e2e-gpu-job.yml with: engine: ${{ matrix.engine }} gpu_tier: "4" runner: 4-gpu-h100 timeout: ${{ matrix.timeout }} + test_timeout: ${{ matrix.test_timeout || 25 }} test_dirs: e2e_test/chat_completions secrets: inherit 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 000000000..81715311c --- /dev/null +++ b/e2e_test/chat_completions/test_epd_multimodal.py @@ -0,0 +1,132 @@ +"""EPD (Encode-Prefill-Decode) multimodal Chat Completions E2E Tests. + +Exercises TokenSpeed's EPD disaggregation on a vision-language model: the +encode worker runs the vision tower, prefill/decode run the LM, and the gateway +stitches encode -> prefill -> decode. + +Like the PD KV-transfer tests (``test_pd_mooncake``/``test_pd_nixl``), these do +NOT stop at "a plausible answer came back" — a single-worker fallback would pass +that. They assert a worker-side signal that the disaggregation actually happened: +the encode worker's own per-request accept log. + +EPD is TokenSpeed-only. On a small MoE VLM (Qwen3.6-35B-A3B, 3B active) at tp=1 +per worker, every topology fits the 4-GPU runner: 1e1p1d=3 GPUs and +1e2p1d/2e1p1d/1e1p2d=4 GPUs (EPD needs >=3 cards since encode/prefill/decode are +separate workers). + +Usage: + pytest e2e_test/chat_completions/test_epd_multimodal.py -v +""" + +from __future__ import annotations + +import base64 +import logging +import os +import tempfile +from pathlib import Path + +import pytest +from infra.pd_logs import assert_worker_logs_captured, wait_for_marker, worker_log_dir + +logger = logging.getLogger(__name__) + +# Local test image (checked into repo) — a black labrador puppy. +FIXTURES_DIR = Path(__file__).parent.parent / "fixtures" / "images" +DOG_IMAGE_PATH = FIXTURES_DIR / "dog.jpg" + +# Router + worker logs land here (per-pid) via the gateway marker below. Worker +# logs actually go to E2E_LOG_DIR in CI; ``worker_log_dir`` resolves both. +_LOG_DIR = Path(tempfile.gettempdir()) / f"smg-e2e-epd-{os.getpid()}" +# Emitted once per Encode RPC by the TokenSpeed encode servicer +# (grpc_servicer/.../tokenspeed/encoder_servicer.py). Its presence proves the +# image reached a dedicated encode worker — the defining EPD step. +ENCODE_ACCEPTED_MARKER = "EPD encode: accepted" + + +def _image_to_base64_url(path: Path) -> str: + """Convert a local image file to a base64 data URL.""" + data = base64.b64encode(path.read_bytes()).decode("utf-8") + return f"data:image/jpeg;base64,{data}" + + +def _make_image_content(image_source: str) -> dict: + """Create an image_url content part from a URL or data URL string. + + Local file paths must be pre-converted via ``_image_to_base64_url``. + """ + return {"type": "image_url", "image_url": {"url": image_source}} + + +# The four EPD topologies to cover. Every worker runs at tp=1 (a 3B-active MoE +# that fits one card), so 1e1p1d uses 3 GPUs and the rest use 4 — all fit the +# 4-GPU runner. The (encode, prefill, decode) counts ride in the param tuple, not +# a marker: setup_backend is class-scoped, so per-param marks aren't visible there. +_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"), +] + + +@pytest.mark.engine("tokenspeed") +@pytest.mark.gpu(4) +@pytest.mark.e2e +@pytest.mark.model("Qwen/Qwen3.6-35B-A3B-FP8") +@pytest.mark.gateway(log_level="debug", log_dir=str(_LOG_DIR)) +@pytest.mark.parametrize("setup_backend", _EPD_TOPOLOGIES, indirect=True) +class TestEPDMultimodal: + """Verify the image really flows encode -> prefill -> decode. + + A naive content check can't distinguish a real 3-worker EPD pipeline from a + single-worker fallback; the encode worker's own log can — so that's what this + asserts, mirroring how the PD tests assert the KV transfer from logs. + """ + + def test_single_image_base64(self, model, setup_backend): + """One dog image through encode -> prefill -> decode, with the encode + worker's participation verified from its logs.""" + _, _, client, *_ = setup_backend + + response = client.chat.completions.create( + model=model, + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "What animal is in this image?"}, + _make_image_content(_image_to_base64_url(DOG_IMAGE_PATH)), + ], + } + ], + temperature=0, + max_tokens=100, + ) + + # (1) The model saw the image and described it correctly. + text = response.choices[0].message.content + assert text is not None and len(text) > 0 + assert any(k in text.lower() for k in ["dog", "puppy", "labrador"]), ( + f"Expected dog-related content, got: {text}" + ) + # (2) The image was tokenized INTO the prompt: the bare text question is + # ~10 tokens, so a large prompt confirms the vision tokens were spliced in + # (encode -> prefill actually delivered the image), not dropped. + assert response.usage.prompt_tokens > 50, ( + f"prompt_tokens={response.usage.prompt_tokens} is too low; " + "the image tokens were likely not delivered to prefill" + ) + assert response.usage.completion_tokens > 0 + + # (3) REAL EPD: the encode worker itself logged accepting the dispatch, so + # the vision stage ran on a separate encode worker rather than degrading to + # a single-worker path. This is the EPD analog of the PD KV-transfer check. + worker_dir = worker_log_dir(_LOG_DIR) + worker_logs = wait_for_marker(worker_dir, "worker-*.log", ENCODE_ACCEPTED_MARKER) + assert_worker_logs_captured(worker_logs, "EPD encode dispatch") + assert ENCODE_ACCEPTED_MARKER in worker_logs, ( + "encode worker never logged accepting the request — the image did not " + f"flow through the EPD encode stage; checked {worker_dir}/worker-*.log" + ) + logger.info("EPD single image (encode worker engaged): %s", text) diff --git a/e2e_test/chat_completions/test_multimodal.py b/e2e_test/chat_completions/test_multimodal.py index 93f28beba..19edde8f5 100644 --- a/e2e_test/chat_completions/test_multimodal.py +++ b/e2e_test/chat_completions/test_multimodal.py @@ -173,6 +173,64 @@ def test_multi_images_mixed(self, model, setup_backend): logger.info("Multi image mixed response: %s", text) +# ============================================================================= +# /dev/shm tensor-transport verification (1 GPU, vLLM) +# ============================================================================= + + +@pytest.mark.engine("vllm") +@pytest.mark.gpu(1) +@pytest.mark.e2e +@pytest.mark.model("Qwen/Qwen3-VL-8B-Instruct") +# Force the /dev/shm tensor transport. A low min-bytes so any real image tensor +# crosses the threshold, making the shm path deterministic regardless of size. +@pytest.mark.gateway( + extra_args=["--multimodal-tensor-transport", "shm", "--multimodal-shm-min-bytes", "1024"] +) +@pytest.mark.parametrize("setup_backend", ["grpc"], indirect=True) +class TestMultimodalShmTransport: + """Ground truth that the /dev/shm tensor transport actually engages. + + Gateway and worker are co-located in CI, so they share /dev/shm and the shm + path can be used. This asserts not just that multimodal works, but that the + pixel tensor traveled over shm (``smg_mm_tensors_total{path="shm"}``) rather + than silently falling back to the inline gRPC payload. + """ + + def test_single_image_uses_shm_transport(self, model, setup_backend): + _, _, client, gateway = setup_backend + + before = gateway.metric_sum("smg_mm_tensors_total", path="shm") + + response = client.chat.completions.create( + model=model, + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "What animal is in this image?"}, + _make_image_content(_image_to_base64_url(DOG_IMAGE_PATH)), + ], + } + ], + temperature=0, + max_tokens=100, + ) + + text = _extract_text(response, False) + assert any(k in text.lower() for k in ["dog", "puppy", "labrador"]), ( + f"Expected dog-related content, got: {text}" + ) + + after = gateway.metric_sum("smg_mm_tensors_total", path="shm") + assert after > before, ( + "expected the /dev/shm tensor transport to be used " + f"(smg_mm_tensors_total path=shm before={before} after={after}); " + "the pixel tensor silently fell back to the inline gRPC payload" + ) + logger.info("shm transport: path=shm count rose %s -> %s", before, after) + + # ============================================================================= # Llama-4-Scout multimodal tests (4 GPU) # ============================================================================= diff --git a/e2e_test/fixtures/setup_backend.py b/e2e_test/fixtures/setup_backend.py index abbb34f3f..dd23eb005 100644 --- a/e2e_test/fixtures/setup_backend.py +++ b/e2e_test/fixtures/setup_backend.py @@ -107,6 +107,7 @@ def setup_backend(request: pytest.FixtureRequest): Backend type is determined by parametrize value via ``request.param``: - ``"http"``, ``"grpc"``: Local workers (SGLang, vLLM, or TRT-LLM) - ``"pd_http"``, ``"pd_grpc"``: PD disaggregation workers + - ``"epd_grpc"``: EPD (encode-prefill-decode) disaggregation (TokenSpeed) - ``"openai"``, ``"xai"``, ``"anthropic"``: Cloud backends (no workers) Configuration via markers: @@ -115,12 +116,21 @@ def setup_backend(request: pytest.FixtureRequest): - ``@pytest.mark.workers(gpus=2, extra_engine_args=[...])``: Per-worker GPU count and extra engine CLI args (local workers only) - ``@pytest.mark.workers(prefill=1, decode=1)``: PD worker counts + - EPD topology: pass a ("epd_grpc", (encode, prefill, decode)) param - ``@pytest.mark.gateway(policy=..., timeout=..., extra_args=...)``: Gateway config Returns: Tuple of ``(backend_name, model_path, client, gateway)`` """ - backend_name: str = request.param + # EPD topologies pass a ("epd_grpc", (n_encode, n_prefill, n_decode)) tuple; + # every other backend passes a bare protocol string. Read the topology from + # request.param, NOT a marker — this fixture is class-scoped, so per-param + # marks on the generated test items aren't visible on request.node. + param = request.param + if isinstance(param, tuple): + backend_name, epd_topology = param + else: + backend_name, epd_topology = 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 +147,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 +165,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_topology or (1, 1, 1), + gateway_config, + gateway, + log_dir, + ) + elif is_pd: yield from _setup_pd( model_id, model_path, @@ -299,6 +321,98 @@ def _setup_pd( stop_workers(all_workers) +# --------------------------------------------------------------------------- +# EPD disaggregation backend (TokenSpeed) +# --------------------------------------------------------------------------- + + +def _setup_epd( + model_id, + model_path, + engine, + connection_mode, + epd_topology, + gateway_config, + gateway, + log_dir, +): + """Launch encode + prefill + decode workers + EPD gateway, yield, tear down. + + Mirrors ``_setup_pd``. ``epd_topology`` is ``(n_encode, n_prefill, n_decode)``. + The encode worker runs the vision tower at tp=1 (``gpus=1``); prefill/decode + run the LM at the model spec's tp. GPU offsets are laid out encode-first so + co-located workers don't share GPUs. + """ + spec = get_model_spec(model_id) + tp = spec.get("tp", 1) + num_encode, num_prefill, num_decode = epd_topology + backend_name = f"epd_{connection_mode.value}" + runtime_label = RUNTIME_LABELS.get(engine, engine) + + logger.info( + "Starting %s EPD backend: model=%s, %d encode + %d prefill + %d decode", + runtime_label, + model_id, + num_encode, + num_prefill, + num_decode, + ) + + all_workers: list = [] + try: + # Encode workers: vision tower at tp=1, one GPU each, starting at GPU 0. + encode_workers = _start_workers_tracked( + model_id=model_id, + engine=engine, + mode=connection_mode, + count=num_encode, + worker_type=WorkerType.ENCODE, + log_dir=log_dir, + gpus=1, + ) + all_workers.extend(encode_workers) + + # Prefill workers start on GPUs after the (tp=1) encode workers. + prefill_gpu_offset = num_encode + prefill_workers = _start_workers_tracked( + model_id=model_id, + engine=engine, + mode=connection_mode, + count=num_prefill, + worker_type=WorkerType.PREFILL, + log_dir=log_dir, + gpu_offset=prefill_gpu_offset, + ) + all_workers.extend(prefill_workers) + + # Decode workers start on GPUs after encode + prefill. + decode_gpu_offset = prefill_gpu_offset + num_prefill * tp + decode_workers = _start_workers_tracked( + model_id=model_id, + engine=engine, + mode=connection_mode, + count=num_decode, + worker_type=WorkerType.DECODE, + log_dir=log_dir, + gpu_offset=decode_gpu_offset, + ) + 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 c1936663f..331b86ace 100644 --- a/e2e_test/infra/constants.py +++ b/e2e_test/infra/constants.py @@ -17,6 +17,7 @@ class WorkerType(StrEnum): REGULAR = "regular" PREFILL = "prefill" DECODE = "decode" + ENCODE = "encode" # EPD encode worker (vision tower); TokenSpeed-only class Runtime(StrEnum): diff --git a/e2e_test/infra/gateway.py b/e2e_test/infra/gateway.py index 408f55b97..227a30a42 100644 --- a/e2e_test/infra/gateway.py +++ b/e2e_test/infra/gateway.py @@ -42,9 +42,10 @@ class WorkerInfo: class Gateway: """Manages a Shepherd Model Gateway router instance. - Four startup modes: + Five startup modes: - Regular: start(worker_urls=[...], model_path="...") - PD: start(prefill_workers=[...], decode_workers=[...]) + - EPD: start(encode_workers=[...], prefill_workers=[...], decode_workers=[...]) - IGW: start(igw_mode=True), then add_worker(url) dynamically - Cloud: start(cloud_backend="openai"|"xai"|"anthropic") """ @@ -69,6 +70,7 @@ def __init__( self.log_level: str = "warn" self.log_dir: str | None = None self.pd_mode: bool = False + self.epd_mode: bool = False self.igw_mode: bool = False self.cloud_mode: bool = False self.cloud_backend: str | None = None @@ -85,8 +87,10 @@ def start( *, worker_urls: list[str] | None = None, model_path: str | None = None, + encode_workers: list[Worker] | None = None, prefill_workers: list[Worker] | None = None, decode_workers: list[Worker] | None = None, + encode_policy: str | None = None, igw_mode: bool = False, cloud_backend: str | None = None, history_backend: str = "memory", @@ -97,19 +101,24 @@ def start( log_level: str | None = None, log_dir: str | None = None, ) -> None: - """Start the gateway in exactly one mode (regular, PD, IGW, or cloud).""" + """Start the gateway in exactly one mode (regular, PD, EPD, IGW, or cloud).""" 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 + # EPD also passes prefill/decode workers, so gate PD on the absence of encode. + is_pd_mode = not is_epd_mode and (prefill_workers is not None or decode_workers is not None) 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, " + "Specify exactly one mode: worker_urls (regular), " + "encode/prefill/decode_workers (EPD), prefill/decode_workers (PD), " "igw_mode=True, or cloud_backend" ) @@ -154,6 +163,40 @@ def start( extra_args=extra_args, log_msg=f"PD gateway ({len(prefills)} prefill, {len(decodes)} decode)", ) + elif is_epd_mode: + self.pd_mode = False + self.epd_mode = True + self.igw_mode = False + encodes = encode_workers or [] + prefills = prefill_workers or [] + decodes = decode_workers or [] + + mode_args = ["--epd-disaggregation"] + for en in encodes: + if en.bootstrap_port is not None: + mode_args += ["--encode", en.base_url, str(en.bootstrap_port)] + else: + mode_args += ["--encode", en.worker_url] + for pf in prefills: + if pf.bootstrap_port is not None: + mode_args += ["--prefill", pf.base_url, str(pf.bootstrap_port)] + else: + mode_args += ["--prefill", pf.worker_url] + for dc in decodes: + mode_args += ["--decode", dc.worker_url] + if encode_policy: + mode_args += ["--encode-policy", 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_cloud_mode: assert cloud_backend is not None self.pd_mode = False @@ -327,6 +370,40 @@ def health(self, timeout: float = 5.0) -> bool: except (httpx.RequestError, httpx.TimeoutException): return False + def metrics_raw(self, timeout: float = 5.0) -> str: + """Fetch the raw Prometheus exposition text from the gateway's + metrics port (``/metrics``).""" + resp = httpx.get(f"{self.metrics_url}/metrics", timeout=timeout) + resp.raise_for_status() + return resp.text + + def metric_sum(self, name: str, **label_filters: str) -> float: + """Sum every sample of counter/gauge ``name`` whose labels are a + superset of ``label_filters``. Returns 0.0 when the series is absent + (e.g. before the first request). A minimal Prometheus-text parser — + adequate for e2e assertions, not a general client. + """ + total = 0.0 + for line in self.metrics_raw().splitlines(): + line = line.strip() + if not line or line.startswith("#"): + continue + metric, brace, rest = line.partition("{") + if metric != name or not brace: + continue + labels_str, _, value_str = rest.partition("}") + labels = { + key.strip(): val.strip().strip('"') + for key, _, val in (part.partition("=") for part in labels_str.split(",")) + if key.strip() + } + if all(labels.get(k) == v for k, v in label_filters.items()): + try: + total += float(value_str.strip().split()[0]) + except (ValueError, IndexError): + pass + return total + def _worker_from_api_response(self, w: dict) -> WorkerInfo: """Convert API response dict to WorkerInfo.""" status = "healthy" if w.get("is_healthy", False) else "unhealthy" diff --git a/e2e_test/infra/model_specs.py b/e2e_test/infra/model_specs.py index c507bcdd5..3d6abb421 100644 --- a/e2e_test/infra/model_specs.py +++ b/e2e_test/infra/model_specs.py @@ -162,6 +162,39 @@ def _resolve_model_path(hf_path: str) -> str: "tp": 1, "features": ["chat", "streaming", "multimodal"], }, + # TokenSpeed EPD multimodal model: Qwen3.6-35B-A3B (arch + # Qwen3_5MoeForConditionalGeneration — in TokenSpeed's multimodal registry — + # with a SigLIP vision tower). FP8 (not NVFP4): FP8 is Hopper-native so it runs + # on the h100 runner, whereas NVFP4 needs Blackwell. Only 3B params are active, + # so the FP8 weights (~35GB) fit one H100 at tp=1; that lets every EPD topology + # (1e1p1d/1e2p1d/2e1p1d/1e1p2d) run on the 4-GPU h100 runner, one worker per + # card. EPD (Encode-Prefill-Decode) disaggregation is TokenSpeed-only: the + # encode worker runs the vision tower, prefill/decode run the LM. + "Qwen/Qwen3.6-35B-A3B-FP8": { + "model": _resolve_model_path("Qwen/Qwen3.6-35B-A3B-FP8"), + "tp": 1, + "features": ["chat", "streaming", "multimodal", "moe"], + "startup_timeout": 600, + # Keep the 35B FP8 LM (prefill/decode) inside one 80GB H100 with headroom + # for generation activations + the mooncake/EPD buffers. TokenSpeed's + # defaults (gpu-mem-util auto ~0.9, kvstore-ratio 2.0, 131K-token KV pool) + # fill the card and OOM mid-generate -> empty response. A short context is + # ample for the single-image smoke test; the encode role ignores LM knobs. + "tokenspeed_args": [ + "--gpu-memory-utilization", + "0.75", + "--max-model-len", + "8192", + "--max-num-seqs", + "8", + "--kvstore-ratio", + "0.5", + ], + # ~35GB and TokenSpeed-only (EPD). Exclude from tier-wide pre-download so + # the sglang/vLLM/TRT lanes (which never run EPD) don't pull it; the + # tokenspeed EPD job downloads it explicitly 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/process_utils.py b/e2e_test/infra/process_utils.py index eee57cacd..dfc173b05 100644 --- a/e2e_test/infra/process_utils.py +++ b/e2e_test/infra/process_utils.py @@ -186,11 +186,42 @@ def wait_for_workers_ready( def detect_ib_device() -> str | None: - """Detect first active InfiniBand device (e.g., mlx5_0). + """Detect the first RDMA device with an active port (e.g., "mlx5_0"). + + Reads ``/sys/class/infiniband`` directly so it works for both InfiniBand and + RoCE and does NOT depend on the ``ibv_devinfo`` CLI, which isn't installed on + every GPU runner (the tokenspeed image ships libibverbs but not the utils). + Without a device the mooncake transfer engine enumerates every NIC on the + node and hangs, so returning None here must be a genuine "no RDMA" signal, + not just "the CLI is missing". Falls back to ``ibv_devinfo`` if sysfs is + absent. Returns: Device name if found (e.g., "mlx5_0"), None otherwise. """ + ib_root = "/sys/class/infiniband" + if os.path.isdir(ib_root): + + def _dev_key(name: str) -> tuple[str, int]: + # Numeric sort so mlx5_2 precedes mlx5_10 (lexical order would not). + head, _, tail = name.rpartition("_") + return (head, int(tail)) if tail.isdigit() else (name, 0) + + for dev in sorted(os.listdir(ib_root), key=_dev_key): + ports = os.path.join(ib_root, dev, "ports") + if not os.path.isdir(ports): + continue + for port in sorted(os.listdir(ports)): + try: + with open(os.path.join(ports, port, "state")) as f: + # e.g. "4: ACTIVE" + if "ACTIVE" in f.read(): + logger.info("Detected IB device: %s (port %s)", dev, port) + return dev + except OSError: + continue + + # Fallback: the ibv_devinfo CLI, if the sysfs tree wasn't available. try: subprocess.run( ["ibv_devinfo", "-l"], @@ -199,6 +230,7 @@ def detect_ib_device() -> str | None: timeout=1, ) except (FileNotFoundError, subprocess.TimeoutExpired): + logger.warning("detect_ib_device: no active RDMA device (sysfs empty, ibv_devinfo absent)") return None for i in range(12): @@ -217,4 +249,5 @@ def detect_ib_device() -> str | None: return dev except Exception: pass + logger.warning("detect_ib_device: no active RDMA device found") return None diff --git a/e2e_test/infra/worker.py b/e2e_test/infra/worker.py index 30e14caa1..5fa8bbb55 100644 --- a/e2e_test/infra/worker.py +++ b/e2e_test/infra/worker.py @@ -41,6 +41,7 @@ class Worker: gpu_ids: list[int] mode: ConnectionMode = ConnectionMode.HTTP worker_type: WorkerType = WorkerType.REGULAR + tp_override: int | None = None # per-worker tp (else model spec's tp) bootstrap_port: int | None = None nixl_port: int | None = None ib_device: str | None = None @@ -173,7 +174,10 @@ def _build_cmd(self) -> list[str]: """Build engine-specific launch command using model specs.""" spec = get_model_spec(self.model_id) model_path = spec["model"] - tp_size = spec.get("tp", 1) + # tp defaults to the spec's tp. A per-worker ``tp_override`` (e.g. the EPD + # encode worker's vision tower at tp=1) wins — it decouples tp from the + # spec tp without disturbing the ``gpus``-only DP path (gpus=dp*tp, tp=1). + tp_size = self.tp_override if self.tp_override is not None else spec.get("tp", 1) features = spec.get("features", []) if self.engine == "sglang": @@ -341,6 +345,34 @@ 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", ] + + # EPD disaggregation: encode/prefill/decode are the SAME module, split by + # ``--disaggregation-mode``. Encode + prefill get a bootstrap port for the + # mooncake rendezvous; decode has none (mirrors the SGLang PD branch). + if self.worker_type in (WorkerType.ENCODE, WorkerType.PREFILL, WorkerType.DECODE): + cmd.extend(["--disaggregation-mode", self.worker_type.value]) + if self.bootstrap_port: + cmd.extend(["--disaggregation-bootstrap-port", str(self.bootstrap_port)]) + # Pin mooncake to one IB device, exactly as the PD branch does. + # Without it the transfer engine enumerates every RoCE NIC on the + # node (18 on the 4-GPU H100 runner), fails to register a local + # segment, and the worker hangs before it can become healthy. + if self.ib_device: + cmd.extend(["--disaggregation-ib-device", self.ib_device]) + # Match tokenspeed's EPD serve script (tokenspeed#549): explicit + # mooncake transfer backend + layerwise transfer, and skip the engine + # server warmup (the disaggregated warmup path hangs; the gRPC-side + # warmup is skipped via TOKENSPEED_SKIP_GRPC_WARMUP in _build_env). + cmd.extend( + [ + "--disaggregation-transfer-backend", + "mooncake", + "--disaggregation-layerwise-interval", + "1", + "--skip-server-warmup", + ] + ) + extra = spec.get("tokenspeed_args", []) if extra: cmd.extend(extra) @@ -413,6 +445,44 @@ def _build_env(self) -> dict[str, str]: env["NCCL_SHM_DISABLE"] = "1" env["TLLM_DISABLE_ALLREDUCE_AUTOTUNE"] = "1" + # EPD's mooncake transfer engine registers GPU memory for RDMA. mooncake + # defaults to the legacy nvidia_peermem GPUDirect path (WITH_NVIDIA_PEERMEM, + # a runtime env var that defaults to true), which fails to register GPU + # buffers with EFAULT on the NVIDIA Open Kernel driver the GPU runners use + # (no nvidia_peermem module). Force the dmabuf path instead — the open + # driver supports it and it registers GPU memory fine — so encode/prefill/ + # decode workers come up instead of hanging on the swallowed registration + # failure. Verified on a GB300 box: peermem -> EFAULT, dmabuf -> rc=0. + if self.engine == "tokenspeed" and self.worker_type in ( + WorkerType.ENCODE, + WorkerType.PREFILL, + WorkerType.DECODE, + ): + env.setdefault("WITH_NVIDIA_PEERMEM", "false") + # [temp diag] enable faulthandler so a SIGABRT on health timeout dumps + # every thread's Python stack to the worker log (see _dump_worker_stack). + env.setdefault("PYTHONFAULTHANDLER", "1") + # EPD runtime config mirrored from tokenspeed's own EPD serve script + # (serve_qwen35_122b_nvfp4_epd_1e2p1d.sh, tokenspeed#549). THE hang fix: + # the smg servicer's gRPC warmup drives stub.Generate, which on a + # disaggregated prefill/decode worker can't complete and hangs until the + # health timeout (server.py:_wait_and_warmup; the encode role self-skips, + # which is why only encode came up). Skip it so prefill/decode reach + # SERVING. The rest make the actual same-node encode<->prefill<->decode + # transfer work: NVLink-IPC intranode transport (RoCE loopback fails on + # one host), proxy bypass for the local mooncake bootstrap, and a large + # gRPC limit for the inline pixel payload to the encode worker. + env.setdefault("TOKENSPEED_SKIP_GRPC_WARMUP", "1") + env.setdefault("MC_INTRANODE_NVLINK", "1") + env.setdefault("MC_INTRA_NVLINK", "1") + env.setdefault("NO_PROXY", "*") + env.setdefault("no_proxy", "*") + env.setdefault("TOKENSPEED_GRPC_MAX_MESSAGE_BYTES", "2000000000") + # The 35B FP8 LM is tight on one 80GB card; reduce allocator + # fragmentation so generation doesn't OOM on reserved-but-unallocated + # memory (the OOM error explicitly recommends this). + env.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") + return env def _spawn_process(self, cmd: list[str], env: dict[str, str]) -> subprocess.Popen: @@ -488,11 +558,28 @@ def _wait_grpc_healthy(self, timeout: float) -> None: finally: channel.close() + # [temp diag] turn a silent health timeout into a precise stack: SIGABRT + + # PYTHONFAULTHANDLER makes the worker print every thread's traceback to its + # log before dying, so we can see exactly where EPD prefill is stuck. + self._dump_worker_stack() raise TimeoutError( f"gRPC worker {self.model_id} on port {self.port} " f"did not become healthy within {timeout}s" ) + def _dump_worker_stack(self) -> None: + """[temp diag] SIGABRT the worker so faulthandler dumps stacks to its log.""" + if not self.process or self.process.poll() is not None: + return + try: + os.killpg(os.getpgid(self.process.pid), signal.SIGABRT) + except Exception: + try: + self.process.send_signal(signal.SIGABRT) + except Exception: + return + time.sleep(3) # let faulthandler flush the traceback to the log + def start_workers( model_id: str, @@ -537,10 +624,15 @@ def start_workers( gpus_per_worker = gpus or spec.get("tp", 1) timeout = spec.get("startup_timeout", timeout) - # Detect IB device for PD workers - has_pd = worker_type in (WorkerType.PREFILL, WorkerType.DECODE) + # Detect IB device for disaggregated (PD / EPD) workers + has_pd = worker_type in (WorkerType.ENCODE, WorkerType.PREFILL, WorkerType.DECODE) ib_device = detect_ib_device() if has_pd else None + # The EPD encode worker runs the vision tower at tp == its GPU count (not DP), + # so its tp must track ``gpus`` rather than the LM's spec tp. Other worker + # types keep spec tp (``gpus``-only DP callers pass gpus=dp*tp with tp=1). + tp_override = gpus_per_worker if worker_type == WorkerType.ENCODE else None + workers: list[Worker] = [] try: @@ -548,7 +640,10 @@ 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 + # Encode + prefill advertise a bootstrap port for the mooncake + # rendezvous; decode connects out and needs none. + needs_bootstrap = worker_type in (WorkerType.ENCODE, WorkerType.PREFILL) + bootstrap_port = get_open_port() if needs_bootstrap else None worker = Worker( model_id=model_id, @@ -557,6 +652,7 @@ def start_workers( gpu_ids=gpu_ids, mode=mode, worker_type=worker_type, + tp_override=tp_override, bootstrap_port=bootstrap_port, ib_device=ib_device if has_pd else None, log_dir=log_dir, diff --git a/grpc_servicer/smg_grpc_servicer/tokenspeed/encoder_servicer.py b/grpc_servicer/smg_grpc_servicer/tokenspeed/encoder_servicer.py index 0122ee37c..9dcc7f216 100644 --- a/grpc_servicer/smg_grpc_servicer/tokenspeed/encoder_servicer.py +++ b/grpc_servicer/smg_grpc_servicer/tokenspeed/encoder_servicer.py @@ -205,6 +205,15 @@ async def Encode(self, request, context): bootstrap_room = request.items[0].bootstrap_room + # Per-request marker so operators (and the EPD e2e) can confirm the encode + # worker actually received the dispatch — the defining EPD step, and the + # signal a naive "the answer looks right" check would miss. + logger.info( + "EPD encode: accepted request_id=%s bootstrap_room=%d", + request.request_id, + bootstrap_room, + ) + if os.environ.get("EPD_INGEST_OFFLOOP", "1").lower() not in ("0", "false", "no"): # Per-image ingest (proto->tensor + pickle) BLOCKS the lone asyncio # event loop, so grpc.aio cannot deliver the next Encode message until diff --git a/model_gateway/src/routers/grpc/epd_encode.rs b/model_gateway/src/routers/grpc/epd_encode.rs index b37f1d708..7ef56140f 100644 --- a/model_gateway/src/routers/grpc/epd_encode.rs +++ b/model_gateway/src/routers/grpc/epd_encode.rs @@ -180,6 +180,20 @@ fn build_plan( }); } + plan_encode_jobs(items, workers) +} + +/// Match prepared encode items to their per-item encode-worker assignments, +/// producing the encode->prefill bootstrap info and the dispatch jobs. +/// +/// Validates that the EPD worker selection carries exactly one encode assignment +/// per item, in item order, then assigns each item a random bootstrap room. +/// Callers must handle the empty-`items` case before calling this (encode +/// planning requires at least one item and a matching non-empty assignment set). +fn plan_encode_jobs( + items: Vec, + workers: &WorkerSelection, +) -> Result { let encode_assignments = workers .encode_assignments() .filter(|assignments| !assignments.is_empty()) @@ -295,3 +309,202 @@ async fn send_tokenspeed_encode_rpc( } Ok(()) } + +#[cfg(test)] +mod tests { + use std::{collections::HashMap, sync::Arc}; + + use llm_multimodal::{ + FieldLayout, ImageDetail, ImageFrame, ImageSource, Modality, PlaceholderRange, + PreprocessedEncoderInputs, + }; + use ndarray::{ArrayD, IxDyn}; + + use super::*; + use crate::{ + routers::grpc::{ + context::{EncodeWorkerAssignment, WorkerSelection}, + proto_wrapper::{TokenSpeedModality, TokenSpeedMultimodalItem, TokenSpeedTensor}, + }, + worker::{BasicWorkerBuilder, RuntimeType, Worker, WorkerType}, + }; + + /// Build a precomputed intermediate carrying `n` image items, mirroring the + /// batched-layout fixture used in the assemble.rs tests. + fn image_intermediate(n: usize) -> PrecomputedMultimodalIntermediate { + // encoder_input: one row per item, 2 features each. + let data: Vec = (0..n * 2).map(|v| v as f32).collect(); + let preprocessed = PreprocessedEncoderInputs { + encoder_input: ArrayD::from_shape_vec(IxDyn(&[n, 2]), data).unwrap(), + feature_token_counts: vec![1; n], + item_sizes: vec![(1, 1); n], + model_specific: HashMap::new(), + }; + let images = (0..n) + .map(|i| { + Arc::new(ImageFrame::new( + image::DynamicImage::new_rgb8(1, 1), + bytes::Bytes::from_static(b"x"), + ImageDetail::Auto, + ImageSource::InlineBytes, + format!("hash-{i}"), + )) + }) + .collect(); + let placeholders = (0..n) + .map(|i| PlaceholderRange { + offset: 10 * (i + 1), + length: 1, + }) + .collect(); + PrecomputedMultimodalIntermediate { + modality: Modality::Image, + preprocessed, + images, + videos: vec![], + placeholders, + patch_offsets: None, + placeholder_token_id: Some(151655), + field_layouts: HashMap::from([("pixel_values".to_string(), FieldLayout::Batched)]), + keep_on_cpu_keys: vec![], + } + } + + /// A synthetic prepared item with an inline (non-SHM) encoder input, so its + /// `Drop` never touches /dev/shm. + fn synthetic_item() -> PreparedEncodeItem { + let item = TokenSpeedMultimodalItem { + modality: TokenSpeedModality::Image, + encoder_input: TokenSpeedTensor::inline( + vec![0u8; 4], + vec![2, 1], + "bfloat16".to_string(), + ), + model_specific_tensors: HashMap::new(), + placeholder_token_id: Some(151655), + mm_placeholders: vec![(0, 1)], + content_hash: vec![], + }; + PreparedEncodeItem::tokenspeed(item, false, 0) + } + + /// An encode worker whose URL yields `bootstrap_host` and whose spec sets a + /// non-default `bootstrap_port`, so both are assertable in the plan output. + fn encode_worker(host: &str, port: u16) -> Arc { + let worker = BasicWorkerBuilder::new(format!("http://{host}:8080")) + .worker_type(WorkerType::Encode) + .bootstrap_port(Some(port)) + .build(); + Arc::new(worker) + } + + fn disaggregated(assignments: Vec) -> WorkerSelection { + WorkerSelection::Disaggregated { + encode_assignments: Some(assignments), + prefill: Arc::new(BasicWorkerBuilder::new("http://prefill:8080").build()), + decode: Arc::new(BasicWorkerBuilder::new("http://decode:8080").build()), + runtime_type: RuntimeType::TokenSpeed, + } + } + + #[test] + fn prepare_tokenspeed_items_yields_one_item_per_image() { + // assemble_tokenspeed splits the batched encoder input into per-item + // tensors; prepare_tokenspeed_items wraps each as a PreparedEncodeItem. + let precomputed = image_intermediate(3); + let items = prepare_tokenspeed_items(&precomputed, None).unwrap(); + assert_eq!(items.len(), 3); + } + + #[test] + fn plan_encode_jobs_count_mismatch_errors() { + // 1 item but 2 assignments: the count guard must fire. + let items = vec![synthetic_item()]; + let workers = disaggregated(vec![ + EncodeWorkerAssignment { + item_index: 0, + worker: encode_worker("enc-a", 9001), + }, + EncodeWorkerAssignment { + item_index: 1, + worker: encode_worker("enc-b", 9002), + }, + ]); + + // Avoid unwrap_err (EncodePlan is not Debug); match the Err directly. + let err = match plan_encode_jobs(items, &workers) { + Ok(_) => panic!("expected count-mismatch error"), + Err(e) => e.to_string(), + }; + assert!(err.contains("count mismatch"), "got: {err}"); + } + + #[test] + fn plan_encode_jobs_order_mismatch_errors() { + // Single item whose assignment is labeled item_index=1 (should be 0): + // the order guard must fire. + let items = vec![synthetic_item()]; + let workers = disaggregated(vec![EncodeWorkerAssignment { + item_index: 1, + worker: encode_worker("enc-a", 9001), + }]); + + // Avoid unwrap_err (EncodePlan is not Debug); match the Err directly. + let err = match plan_encode_jobs(items, &workers) { + Ok(_) => panic!("expected order-mismatch error"), + Err(e) => e.to_string(), + }; + assert!(err.contains("order mismatch"), "got: {err}"); + } + + #[test] + fn plan_encode_jobs_happy_path_builds_bootstrap_info() { + let items = vec![synthetic_item(), synthetic_item()]; + let workers = disaggregated(vec![ + EncodeWorkerAssignment { + item_index: 0, + worker: encode_worker("enc-a", 9001), + }, + EncodeWorkerAssignment { + item_index: 1, + worker: encode_worker("enc-b", 9002), + }, + ]); + + let plan = plan_encode_jobs(items, &workers).unwrap(); + let (bootstrap_info, dispatch) = plan.into_parts(); + + assert_eq!(dispatch.len(), 2); + assert_eq!(bootstrap_info.len(), 2); + + assert_eq!(bootstrap_info[0].item_index, 0); + assert_eq!(bootstrap_info[0].bootstrap_host, "enc-a"); + assert_eq!(bootstrap_info[0].bootstrap_port, 9001); + + assert_eq!(bootstrap_info[1].item_index, 1); + assert_eq!(bootstrap_info[1].bootstrap_host, "enc-b"); + assert_eq!(bootstrap_info[1].bootstrap_port, 9002); + } + + #[test] + fn plan_encode_jobs_defaults_bootstrap_port_when_unset() { + // A worker without an explicit bootstrap_port falls back to + // DEFAULT_BOOTSTRAP_PORT in the bootstrap info. + let worker = Arc::new( + BasicWorkerBuilder::new("http://enc-c:8080") + .worker_type(WorkerType::Encode) + .build(), + ) as Arc; + let workers = disaggregated(vec![EncodeWorkerAssignment { + item_index: 0, + worker, + }]); + + let plan = plan_encode_jobs(vec![synthetic_item()], &workers).unwrap(); + let (bootstrap_info, _) = plan.into_parts(); + assert_eq!( + bootstrap_info[0].bootstrap_port, + DEFAULT_BOOTSTRAP_PORT as i32 + ); + } +} diff --git a/model_gateway/src/routers/grpc/multimodal/transport.rs b/model_gateway/src/routers/grpc/multimodal/transport.rs index 4a801c267..358b87a84 100644 --- a/model_gateway/src/routers/grpc/multimodal/transport.rs +++ b/model_gateway/src/routers/grpc/multimodal/transport.rs @@ -358,7 +358,97 @@ fn compute_shm_namespace_id() -> Option { #[cfg(test)] mod tests { + use std::collections::HashMap; + + use openai_protocol::worker::WorkerSpec; + use super::*; + use crate::{ + routers::grpc::context::WorkerSelection, + worker::{BasicWorkerBuilder, Worker}, + }; + + /// Build an `Arc` with the given per-request transport override + /// and optional `shm_namespace_id` label (the token `auto` compares against). + fn worker_with( + transport: Option, + shm_namespace_id: Option<&str>, + ) -> Arc { + let mut spec = WorkerSpec::new("http://localhost:8080"); + spec.multimodal_tensor_transport = transport; + let mut labels = HashMap::new(); + if let Some(id) = shm_namespace_id { + labels.insert("shm_namespace_id".to_string(), id.to_string()); + } + spec.labels = labels; + Arc::new(BasicWorkerBuilder::from_spec(spec).build()) + } + + fn single(transport: Option, shm_namespace_id: Option<&str>) -> WorkerSelection { + WorkerSelection::Single { + worker: worker_with(transport, shm_namespace_id), + } + } + + #[test] + fn inline_override_never_enables_shm() { + let sel = single(Some(TransportMode::Inline), None); + assert!(!resolve_mm_shm_enabled(Some(&sel), false)); + } + + #[test] + fn rdma_override_never_enables_shm() { + // rdma routes large tensors through the NIXL pixel lane, not SHM. + let sel = single(Some(TransportMode::Rdma), None); + assert!(!resolve_mm_shm_enabled(Some(&sel), false)); + } + + #[test] + fn shm_override_follows_dev_writable() { + // `shm` forces SHM whenever SMG can write /dev/shm, independent of the + // worker's namespace label. Assert against the real probe so the test is + // environment-robust (true on Linux CI, false where /dev/shm isn't writable). + let sel = single(Some(TransportMode::Shm), None); + assert_eq!( + resolve_mm_shm_enabled(Some(&sel), false), + mm_shm_dev_writable() + ); + } + + #[test] + #[cfg(target_os = "linux")] + fn auto_override_with_matching_namespace_follows_dev_writable() { + // `auto` enables SHM only when the worker is verified to share the + // gateway's /dev/shm (matching token) AND /dev/shm is writable. + let local = local_shm_namespace_id().expect("shm namespace id resolves on Linux"); + let sel = single(Some(TransportMode::Auto), Some(local)); + assert_eq!( + resolve_mm_shm_enabled(Some(&sel), false), + mm_shm_dev_writable() + ); + } + + #[test] + fn auto_override_with_mismatched_namespace_disables_shm() { + // A non-matching token means "not verified as sharing /dev/shm", so `auto` + // must fall back to inline even where /dev/shm is writable. + let sel = single(Some(TransportMode::Auto), Some("boot-x:99999999")); + assert!(!resolve_mm_shm_enabled(Some(&sel), false)); + } + + #[test] + fn auto_override_with_empty_namespace_disables_shm() { + // A missing/empty token is treated as non-sharing. + let sel = single(Some(TransportMode::Auto), Some("")); + assert!(!resolve_mm_shm_enabled(Some(&sel), false)); + } + + #[test] + fn no_workers_disables_shm() { + // Without a worker selection there is no locality proof and no override; + // the router default is `inline`, so SHM stays off. + assert!(!resolve_mm_shm_enabled(None, false)); + } #[test] #[cfg(target_os = "linux")] diff --git a/model_gateway/src/routers/grpc/proto_wrapper.rs b/model_gateway/src/routers/grpc/proto_wrapper.rs index c2415f76d..cb05c64ba 100644 --- a/model_gateway/src/routers/grpc/proto_wrapper.rs +++ b/model_gateway/src/routers/grpc/proto_wrapper.rs @@ -2241,6 +2241,14 @@ mod tests { } fn vllm_mm_data(modality: common::Modality) -> VllmMultimodalData { + vllm_mm_data_with_shm(modality, false, 0) + } + + fn vllm_mm_data_with_shm( + modality: common::Modality, + shm_enabled: bool, + shm_min_bytes: usize, + ) -> VllmMultimodalData { let is_video = modality == common::Modality::Video; VllmMultimodalData { pixel_values: vec![0u8; 16], @@ -2253,8 +2261,8 @@ mod tests { flat_keys: HashMap::new(), keep_on_cpu_keys: vec![], modality, - shm_enabled: false, - shm_min_bytes: 0, + shm_enabled, + shm_min_bytes, } } @@ -2269,4 +2277,42 @@ mod tests { let image = vllm_mm_data(common::Modality::Image).into_proto(); assert_eq!(image.modality, common::Modality::Image as i32); } + + #[test] + fn vllm_inline_pixel_values_into_proto_uses_inline_payload() { + // shm_enabled=false must keep pixel_values inline regardless of size. + let proto = vllm_mm_data_with_shm(common::Modality::Image, false, 0).into_proto(); + let tensor = proto.pixel_values.as_ref().unwrap(); + assert!( + matches!( + tensor.payload.as_ref(), + Some(vllm::tensor_data::Payload::Inline(_)) + ), + "expected inline pixel_values payload, got {:?}", + tensor.payload + ); + } + + #[test] + #[cfg(target_os = "linux")] + fn vllm_shm_pixel_values_into_proto_uses_shm_payload() { + // shm_enabled=true + a low min-bytes threshold (below the 16-byte + // pixel_values buffer) must place pixel_values in /dev/shm, not inline. + let proto = vllm_mm_data_with_shm(common::Modality::Image, true, 1).into_proto(); + + // Collect + unlink any /dev/shm files this proto references before + // asserting, so a panic can't leak the segment (mirrors the tokenspeed + // shm tests' cleanup discipline). + let handles = collect_vllm_multimodal_inputs_shm_handles(&proto); + let payload_is_shm = matches!( + proto.pixel_values.as_ref().and_then(|t| t.payload.as_ref()), + Some(vllm::tensor_data::Payload::Shm(_)) + ); + cleanup_mm_shm_handles(&handles); + + assert!(payload_is_shm, "expected shm pixel_values payload"); + // The shm path must produce exactly one handle for pixel_values (no + // model-specific tensors in this fixture). + assert_eq!(handles.len(), 1); + } } diff --git a/scripts/ci_download_model.sh b/scripts/ci_download_model.sh index 1b09fd422..e90109dde 100755 --- a/scripts/ci_download_model.sh +++ b/scripts/ci_download_model.sh @@ -44,7 +44,9 @@ 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]): + # skip_tier_download: engine-specific/large models (e.g. the TokenSpeed EPD + # model) that their own job downloads by id, so unrelated lanes don't pull them. + 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 c57a645e3..3fa6efdef 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 @@ -179,4 +240,41 @@ assert not shadowed, f'smg gRPC modules shadowed by site-packages copies: {shado print('smg gRPC modules resolve to repo source: OK') " +# ── RDMA / mooncake GPU-registration probe (TEMPORARY — remove before merge) ── +# Answers definitively whether the H100 CI box supports the mooncake RDMA +# GPU-memory registration EPD needs: the EPD workers hang right after mooncake's +# transfer engine starts listening, and the next step is register_memory on a +# GPU buffer (needs GPUDirect: nvidia_peermem or dmabuf). Runs under a hard +# timeout so it can never hang the lane; never fails the lane either. +echo "=== RDMA probe: kernel modules / GPUDirect ===" +lsmod 2>/dev/null | grep -iE "peermem|mlx5_ib|ib_uverbs|nvidia_fs" || echo " (no matching modules)" +if [ -d /sys/module/nvidia_peermem ]; then echo " nvidia_peermem: LOADED"; else echo " nvidia_peermem: NOT loaded"; fi +command -v rdma >/dev/null 2>&1 && rdma link 2>/dev/null | head -3 || true +echo "=== RDMA probe: mooncake GPU register_memory — peermem (default) vs dmabuf ===" +set +e +cat > /tmp/mc_probe.py <<'PY' +import subprocess, sys, os +try: + import torch + from mooncake.engine import TransferEngine +except Exception as e: + print("PROBE import failed:", e); sys.exit(3) +if not torch.cuda.is_available(): + print("PROBE: no CUDA device visible"); sys.exit(4) +ips = [x for x in subprocess.check_output(["hostname", "-I"]).decode().split() if not x.startswith("127.")] +ip = ips[0] if ips else "127.0.0.1" +eng = TransferEngine() +if eng.initialize(ip, "P2PHANDSHAKE", "rdma", "mlx5_0") != 0: + print("PROBE initialize failed"); sys.exit(5) +buf = torch.zeros(4 << 20, device="cuda") +rc = eng.register_memory(buf.data_ptr(), buf.numel() * buf.element_size()) +mode = os.environ.get("WITH_NVIDIA_PEERMEM", "") +print(f"PROBE[WITH_NVIDIA_PEERMEM={mode}] GPU register rc={rc}", "-> WORKS" if rc == 0 else "-> FAIL", flush=True) +sys.exit(0 if rc == 0 else 6) +PY +echo "-- peermem (default) --"; timeout 60 python3 /tmp/mc_probe.py; echo " (rc=$?)" +echo "-- dmabuf (WITH_NVIDIA_PEERMEM=false) --"; WITH_NVIDIA_PEERMEM=false timeout 60 python3 /tmp/mc_probe.py; echo " (rc=$?)" +rm -f /tmp/mc_probe.py +set -e + echo "TokenSpeed installation complete" diff --git a/scripts/run_epd_local.sh b/scripts/run_epd_local.sh new file mode 100755 index 000000000..81122f0d1 --- /dev/null +++ b/scripts/run_epd_local.sh @@ -0,0 +1,128 @@ +#!/bin/bash +# Bring up TokenSpeed EPD (encode-prefill-decode) multimodal disaggregation on a +# single local GPU-RDMA node (GB200 / GB300 / H100) for reproduction and dev, and +# drive the EPD e2e against it. This mirrors what the 4-GPU CI lane does, but +# builds everything from local source so you can iterate in minutes instead of +# 30-min CI rounds. +# +# It exists because EPD-over-mooncake needs a few environment tweaks that aren't +# obvious (see the ENV section): the mooncake transfer engine must use the dmabuf +# GPUDirect path on the NVIDIA open kernel driver, and its runtime deps +# (libnuma) may not be on the loader path in a torch-only env. +# +# Usage: +# scripts/run_epd_local.sh install # build + install the tokenspeed stack + gateway +# scripts/run_epd_local.sh model # download the EPD model +# scripts/run_epd_local.sh run # run the EPD e2e (default topology 1e1p1d) +# scripts/run_epd_local.sh all # install + model + run +# +# Override anything via env, e.g. TORCH_PY=/path/to/python EPD_MODEL=... TS_SRC=... +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +SMG_SRC="${SMG_SRC:-$(cd "${SCRIPT_DIR}/.." && pwd)}" +# TokenSpeed source checkout (engine + kernel + scheduler). Defaults next to smg. +TS_SRC="${TS_SRC:-$(cd "${SMG_SRC}/.." && pwd)/tokenspeed}" +# A python that already provides a CUDA-13 torch for this box's arch. On a GB300 +# dev box that's typically a conda env; point TORCH_PY at its bin/python. +TORCH_PY="${TORCH_PY:?set TORCH_PY to a python that has a matching CUDA-13 torch}" +VENV="${VENV:-${SMG_SRC}/.epd-local-venv}" +EPD_MODEL="${EPD_MODEL:-Qwen/Qwen3.6-35B-A3B-FP8}" +MODEL_ROOT="${MODEL_ROOT:-/models}" +EPD_TOPOLOGY="${EPD_TOPOLOGY:-1e1p1d}" # 1e1p1d | 1e2p1d | 2e1p1d | 1e1p2d + +# ── CUDA toolkit: pick a complete toolkit that matches torch's CUDA ─────────── +if [ -z "${CUDA_HOME:-}" ]; then + for c in /usr/local/cuda-13.1 /usr/local/cuda-13.0 /usr/local/cuda; do + if [ -x "$c/bin/nvcc" ] && [ -f "$c/include/crt/host_runtime.h" ]; then CUDA_HOME="$c"; break; fi + done +fi +CUDA_HOME="${CUDA_HOME:?no complete CUDA toolkit found; set CUDA_HOME}" + +py() { "${VENV}/bin/python" "$@"; } + +detect_arch() { + # e.g. GB300 -> "10.3a", H100 -> "9.0a". FlashInfer wants .a. + "$TORCH_PY" -c 'import torch; a,b=torch.cuda.get_device_capability(); print(f"{a}.{b}a")' +} + +# ── Env every EPD process needs (the non-obvious bits) ─────────────────────── +epd_env() { + # dmabuf GPUDirect: mooncake defaults to legacy nvidia_peermem, which fails to + # register GPU memory on the NVIDIA open kernel driver. Force dmabuf. + echo "WITH_NVIDIA_PEERMEM=false" + # mooncake links libnuma at runtime; a torch-only env often lacks it on the + # loader path. Preload the system copy if present. + [ -e /usr/lib64/libnuma.so.1 ] && echo "LD_PRELOAD=/usr/lib64/libnuma.so.1:${LD_PRELOAD:-}" + # Same-node encode<->prefill lives on one host; use the NVLink IPC intranode + # transport rather than RoCE loopback (which the fabric may not shortcut). + echo "MC_INTRANODE_NVLINK=1" + echo "CUDA_HOME=${CUDA_HOME}" +} + +cmd_install() { + local arch; arch="$(detect_arch)" + echo ">>> building EPD stack: venv=${VENV} CUDA_HOME=${CUDA_HOME} arch=${arch}" + # Clean venv, NOT --system-site-packages: inheriting a torch-2.12 conda env + # breaks its prebuilt C-extensions (torchcomms, ...) as soon as the kernel + # pins torch 2.11. A fresh env with torch 2.11+cu130 avoids the clash and + # still runs on sm_103 (verified on GB300). + # --seed installs pip/setuptools/wheel: tokenspeed-kernel's setup.py shells + # out to `python -m pip install -r requirements/cuda.txt`, so the venv needs pip. + uv venv --python "${VENV_PYTHON:-3.12}" --seed "$VENV" + + export CUDA_HOME PATH="${CUDA_HOME}/bin:${PATH}" + export MAX_JOBS="${MAX_JOBS:-32}" FLASHINFER_CUDA_ARCH_LIST="$arch" TOKENSPEED_KERNEL_BACKEND=cuda + # Resolve torch + nvidia cu13 wheels from the pytorch cu130 index. + export UV_EXTRA_INDEX_URL="${TORCH_INDEX:-https://download.pytorch.org/whl/cu130}" + export PIP_EXTRA_INDEX_URL="$UV_EXTRA_INDEX_URL" + # The kernel's setup.py shells out to plain pip; without a socket timeout a + # stalled download hangs forever. Fail fast (30s no-data) and retry instead. + export PIP_DEFAULT_TIMEOUT="${PIP_DEFAULT_TIMEOUT:-30}" PIP_RETRIES="${PIP_RETRIES:-10}" + # cutlass pin: 4.6.0 dropped cute.core.ThrMma that quack needs (see CI script). + local con; con="$(mktemp)"; echo "nvidia-cutlass-dsl==4.5.2" > "$con" + export UV_CONSTRAINT="$con" PIP_CONSTRAINT="$con" + + uv pip install --python "${VENV}/bin/python" setuptools wheel pybind11 + # The CUDA-13 torch the kernel expects, up front in the clean env. + uv pip install --python "${VENV}/bin/python" "torch==${TORCH_VERSION:-2.11.0}+cu130" + # TokenSpeed: kernel (from source) -> scheduler -> engine. Same order as CI. + uv pip install --python "${VENV}/bin/python" -e "${TS_SRC}/tokenspeed-kernel/python/" --no-build-isolation + uv pip install --python "${VENV}/bin/python" -e "${TS_SRC}/tokenspeed-scheduler/" + uv pip install --python "${VENV}/bin/python" -e "${TS_SRC}/python" --no-build-isolation + + # smg gRPC proto + servicer from source (the EPD encode servicer lives here). + uv pip uninstall --python "${VENV}/bin/python" tokenspeed-smg-grpc-proto tokenspeed-smg-grpc-servicer 2>/dev/null || true + uv pip install --python "${VENV}/bin/python" -e "${SMG_SRC}/crates/grpc_client/python/" + uv pip install --python "${VENV}/bin/python" -e "${SMG_SRC}/grpc_servicer/" + # smg gateway (the model_gateway python wheel) via maturin. + uv pip install --python "${VENV}/bin/python" maturin + (cd "$SMG_SRC" && "${VENV}/bin/maturin" develop --release -m bindings/python/Cargo.toml) + + echo ">>> install complete. verify:" + env $(epd_env) py -c "import tokenspeed, tokenspeed_kernel, mooncake.engine, smg_grpc_servicer, smg; print('EPD stack import OK')" +} + +cmd_model() { + echo ">>> downloading ${EPD_MODEL} -> ${MODEL_ROOT}" + ROUTER_LOCAL_MODEL_PATH="$MODEL_ROOT" bash "${SMG_SRC}/scripts/ci_download_model.sh" "$EPD_MODEL" +} + +cmd_run() { + echo ">>> running EPD e2e (topology ${EPD_TOPOLOGY}) with dmabuf + NVLink IPC" + cd "$SMG_SRC" + env $(epd_env) \ + E2E_ENGINE=tokenspeed E2E_RUNTIME=tokenspeed E2E_GPU_TIER=4 \ + ROUTER_LOCAL_MODEL_PATH="$MODEL_ROOT" SHOW_WORKER_LOGS=1 \ + "${VENV}/bin/python" -m pytest e2e_test/chat_completions/test_epd_multimodal.py \ + -k "$EPD_TOPOLOGY" -s -vv +} + +case "${1:-all}" in + install) cmd_install ;; + model) cmd_model ;; + run) cmd_run ;; + all) cmd_install; cmd_model; cmd_run ;; + env) epd_env ;; + *) echo "usage: $0 {install|model|run|all|env}" >&2; exit 2 ;; +esac