diff --git a/benchmark/benchmark_nvidia_mlx_pd.py b/benchmark/benchmark_nvidia_mlx_pd.py new file mode 100644 index 0000000000..f66a65b9d7 --- /dev/null +++ b/benchmark/benchmark_nvidia_mlx_pd.py @@ -0,0 +1,277 @@ +# Copyright 2022-2026 Xinference Holdings Pte. Ltd +# Licensed under the Apache License, Version 2.0. +"""Measure both directions of NVIDIA / Mac MLX PD on an existing LAN cluster. + +Both workers must run this checkout, with vLLM/SGLang on NVIDIA and MLX on Mac. +Use identical original checkpoint assets for P/D. The optional local baseline +must contain the exact same effective weights, cast to FP16. No server processes +are created or stopped here; only this benchmark's model deployments are owned. +""" + +import argparse +import asyncio +import json +import statistics +import time +import uuid +from pathlib import Path + +import xoscar as xo + +from xinference.client import Client + +from benchmark_mlx_xavier import measure + + +def launch_options(args, mode, uid): + gpu_engine = mode.removeprefix("local-") + options = dict( + model_uid=uid, + model_name=args.model_name, + model_size_in_billions=args.model_size, + model_engine="MLX" if mode == "local" else gpu_engine, + model_format="mlx" if mode == "local" else "pytorch", + quantization="none", + enable_virtual_env=False, + ) + if mode == "local": + return dict( + options, + replica=1, + worker_ip=args.mac_worker, + model_path=args.baseline_model_path, + prompt_cache_size=0, + ) + config = dict(model_path=args.gpu_model_path, dtype="float16") + if gpu_engine == "vLLM": + config.update( + max_model_len=args.context_length, + gpu_memory_utilization=0.3, + enforce_eager=True, + ) + else: + config.update( + context_length=args.context_length, + mem_fraction_static=0.3, + max_total_tokens=args.context_length, + disable_cuda_graph=True, + page_size=64, + stream_interval=1, + ) + if mode.startswith("local-"): + config.update( + {"enable_prefix_caching": False} + if gpu_engine == "vLLM" + else {"disable_radix_cache": True} + ) + return dict( + options, + **config, + replica=1, + worker_ip=args.gpu_worker, + n_gpu=1, + gpu_idx=[args.gpu], + ) + reverse = getattr(args, "direction", "nvidia-prefill") == "mlx-prefill" + options.update( + replica=2, + replica_config=[ + dict( + role="decode" if reverse else "prefill", + model_engine=mode, + engine_config=config, + devices=[dict(worker_ip=args.gpu_worker, n_gpu=1, gpu_idx=[args.gpu])], + ), + dict( + role="prefill" if reverse else "decode", + model_engine="MLX", + engine_config=dict( + model_path=args.mac_model_path, + model_format="mlx", + context_length=args.context_length, + prompt_cache_size=0, + ), + devices=[dict(worker_ip=args.mac_worker, n_gpu="auto")], + ), + ], + ) + if reverse: + options["replica_config"].reverse() + return options + + +async def directory_stats(address, uid): + ref = await xo.actor_ref(address=address, uid="xavier-cache-" + uid) + return await ref.get_stats() + + +def isolation_samples(path): + samples = [] + if path.exists(): + for line in path.read_text().splitlines(): + try: + samples.append(json.loads(line)) + except json.JSONDecodeError: + pass # The monitor may still be writing its last line. + return samples + + +async def measured_run(args, uid, messages, use_gpu, is_pd): + path = args.gpu_isolation_log if use_gpu else None + deadline = time.monotonic() + 1800 + before = [] + if path: + while True: + before = isolation_samples(path) + if ( + before + and not before[-1]["foreign"] + and time.time() - path.stat().st_mtime < 2 + ): + break + if time.monotonic() > deadline: + raise TimeoutError("No isolated GPU measurement window") + await asyncio.sleep(1) + state = await directory_stats(args.supervisor_address, uid) if is_pd else None + runs = await measure(args.endpoint, uid, messages, 1, args.max_tokens) + if state is not None: + after = await directory_stats(args.supervisor_address, uid) + if after["completed_requests"] != state["completed_requests"] + 1: + raise RuntimeError("Request did not complete one KV handoff") + runs[0]["imported_tokens"] = after["imported_tokens"] - state["imported_tokens"] + runs[0]["host_bytes"] = after["host_bytes"] - state["host_bytes"] + if not runs[0]["host_bytes"] or after["gpu_bytes"] != 0: + raise RuntimeError("Request did not use the host transport") + if path: + # The external monitor samples at least twice per second. Inspect the + # samples spanning this request, without including this wait in TTFT. + await asyncio.sleep(0.6) + after = isolation_samples(path) + window = after[max(0, len(before) - 1) :] + if len(after) <= len(before) or any(s["foreign"] for s in window): + print( + "Discarding a GPU run with interference; retrying with a new prompt", + flush=True, + ) + return None + return runs + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + for name in ( + "endpoint", + "supervisor-address", + "gpu-worker", + "mac-worker", + "gpu-model-path", + "mac-model-path", + ): + parser.add_argument("--" + name, required=True) + parser.add_argument("--baseline-model-path", help="Same effective weights in FP16") + parser.add_argument("--model-name", default="qwen2.5-instruct") + parser.add_argument("--model-size", default="0_5") + parser.add_argument("--context-length", type=int, default=8192) + parser.add_argument("--gpu", type=int, default=0) + parser.add_argument( + "--direction", choices=["nvidia-prefill", "mlx-prefill", "both"], default="both" + ) + parser.add_argument("--gpu-baseline", action="store_true") + parser.add_argument( + "--gpu-isolation-log", + type=Path, + help="External JSONL monitor sampled at >=2 Hz; each record has a foreign PID list", + ) + parser.add_argument( + "--engines", nargs="+", choices=["vLLM", "SGLang"], default=["vLLM", "SGLang"] + ) + parser.add_argument("--repeats", nargs="+", type=int, default=[1, 20, 100, 400]) + parser.add_argument("--count", type=int, default=5) + parser.add_argument("--max-tokens", type=int, default=32) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + client = Client(args.endpoint) + result = {} + directions = ( + ["nvidia-prefill", "mlx-prefill"] + if args.direction == "both" + else [args.direction] + ) + modes = [("local", "local")] if args.baseline_model_path else [] + if args.gpu_baseline: + modes += [("local-" + engine, "local") for engine in args.engines] + modes += [ + (engine, direction) for direction in directions for engine in args.engines + ] + for mode, direction in modes: + args.direction = direction + label = mode if direction == "local" else f"{direction}:{mode}" + uid = "nvidia-mlx-bench-" + uuid.uuid4().hex[:8] + try: + client.launch_model(**launch_options(args, mode, uid)) + records = [] + trial = 0 + for repeats in args.repeats: + content = ( + "Water evaporates and condenses. " * repeats + + "Explain this process." + ) + warmup = [dict(role="user", content=f"Warmup {repeats}: " + content)] + # Vary the early user prefix so P cannot reuse complete pages; + # use the same trial sequence for every engine and local MLX. + asyncio.run(measure(args.endpoint, uid, warmup, 1, args.max_tokens)) + runs = [] + while len(runs) < args.count: + messages = [ + dict(role="user", content=f"Trial {trial:08d}: " + content) + ] + measured = asyncio.run( + measured_run( + args, uid, messages, mode != "local", direction != "local" + ) + ) + if measured is not None: + for record in measured: + record["trial"] = trial + runs.extend(measured) + trial += 1 + for run in runs: + expected = ( + run["prompt_tokens"] + - (0 if direction == "mlx-prefill" and mode == "SGLang" else 1) + if direction != "local" + else 0 + ) + # NVIDIA usage currently omits cached-token details. The + # directory's completed import counter establishes reuse. + if direction != "local" and run["imported_tokens"] != expected: + raise RuntimeError("Decode did not import the expected prefix") + records.append( + dict( + repeats=repeats, + prompt_tokens=runs[0]["prompt_tokens"], + median_ttft_ms=statistics.median(r["ttft_ms"] for r in runs), + runs=runs, + ) + ) + state = ( + None + if direction == "local" + else asyncio.run(directory_stats(args.supervisor_address, uid)) + ) + if state and state["active_handoffs"]: + raise RuntimeError("Completed requests left active handoffs") + result[label] = dict(records=records, directory=state) + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(result, indent=2)) + print( + label, + [(r["prompt_tokens"], round(r["median_ttft_ms"], 1)) for r in records], + flush=True, + ) + finally: + client.terminate_model(uid) + + +if __name__ == "__main__": + main() diff --git a/doc/source/locale/de/LC_MESSAGES/user_guide/pd_separation.po b/doc/source/locale/de/LC_MESSAGES/user_guide/pd_separation.po index ba095b8e8f..045ec3bbb6 100644 --- a/doc/source/locale/de/LC_MESSAGES/user_guide/pd_separation.po +++ b/doc/source/locale/de/LC_MESSAGES/user_guide/pd_separation.po @@ -500,15 +500,13 @@ msgid "" "Use the default ``xavier`` transport for MLX; native NIXL is unsupported. " "All workers must reach the supervisor's actor address. Relaunch after a " "worker restart. Two replicas on one Mac share its Metal GPU; measure your " -"workload before expecting a throughput benefit. Cross-engine P/D involving " -"MLX, including NVIDIA-to-Mac handoff, is not yet supported." +"workload before expecting a throughput benefit." msgstr "" "Verwenden Sie für MLX den Standardtransport ``xavier``; natives NIXL wird " "nicht unterstützt. Alle Worker müssen die Actor-Adresse des Supervisors " "erreichen. Starten Sie das Deployment nach einem Worker-Neustart neu. Zwei " "Replikate auf einem Mac teilen dessen Metal-GPU; messen Sie Ihren Workload, " -"bevor Sie einen Durchsatzvorteil erwarten. Engineübergreifendes P/D mit MLX " -"einschließlich NVIDIA-zu-Mac-Übergaben wird noch nicht unterstützt." +"bevor Sie einen Durchsatzvorteil erwarten." #: user_guide/pd_separation.rst:102 msgid "" @@ -594,18 +592,127 @@ msgstr "" "und sampelt die Ausgabe. Engine-Kernels können kleine numerische " "Unterschiede erzeugen. vLLM decode benötigt mindestens zwei Prompt-Tokens." -#: source/user_guide/pd_separation.rst +#: user_guide/pd_separation.rst msgid "" "Cross-engine P/D does not support CPU fallback, retained Xavier history, " "logprobs, structured sampling, LoRA, speculative decoding, multimodal inputs" " or hybrid attention. Omit ``xavier_gpu_cache_bytes`` or set it to ``0``; do" " not set ``xavier_cache_bytes``. Relaunch after a worker restart. Native " -"NIXL comparisons use the same engine on both roles. MLX is unsupported." +"NIXL comparisons use the same engine on both roles." msgstr "" "Engineübergreifendes P/D unterstützt weder CPU-Fallback noch Xavier-Verlauf," " logprobs, strukturiertes Sampling, LoRA, spekulatives Decoding, multimodale" " Eingaben oder hybride Attention. Lassen Sie ``xavier_gpu_cache_bytes`` weg " "oder setzen Sie es auf ``0``; setzen Sie ``xavier_cache_bytes`` nicht. Nach " "einem Worker-Neustart ist eine erneute Bereitstellung nötig. Native NIXL-" -"Vergleiche verwenden dieselbe Engine für beide Rollen. MLX wird nicht " -"unterstützt." +"Vergleiche verwenden dieselbe Engine für beide Rollen." + +#: user_guide/pd_separation.rst +msgid "NVIDIA and MLX cross-engine P/D" +msgstr "Engineübergreifendes P/D mit NVIDIA und MLX" + +#: user_guide/pd_separation.rst +msgid "" +"Xavier supports both vLLM/SGLang prefill with MLX decode and MLX prefill " +"with vLLM/SGLang decode, with streaming and non-streaming requests. FP16 " +"64-token pages pass through CPU staging and actor RPC between CUDA and " +"Metal. vLLM and MLX decode compute the last prompt token; SGLang decode " +"imports the full prompt KV and the first output token sampled by P. MLX " +"prefill publishes request-owned pages on its worker with a 512 MiB in-flight" +" budget." +msgstr "" +"Xavier unterstützt vLLM/SGLang-Prefill mit MLX-Decode und MLX-Prefill mit " +"vLLM/SGLang-Decode, mit und ohne Streaming. FP16-Seiten mit 64 Token " +"passieren CPU-Zwischenspeicher und Actor-RPC zwischen CUDA und Metal. vLLM " +"und MLX berechnen beim Decode das letzte Prompt-Token; SGLang importiert den" +" vollständigen Prompt-KV und das von P gesampelte erste Ausgabetoken. MLX-" +"Prefill veröffentlicht anfragegebundene Seiten auf seinem Worker mit einem " +"Budget von 512 MiB für laufende Transfers." + +#: user_guide/pd_separation.rst +msgid "" +"Use identical unquantized checkpoint and tokenizer files; do not convert the" +" Mac checkpoint. Worker-local ``model_path`` and ``model_format`` belong in " +"each replica's ``engine_config``. Xavier checks weights, tokenizer, " +"effective model semantics, context limits and prompt token IDs. Model and " +"sampling restrictions match vLLM/SGLang cross-engine P/D; failures raise " +"errors." +msgstr "" +"Verwenden Sie identische unquantisierte Checkpoint- und Tokenizer-Dateien; " +"konvertieren Sie den Mac-Checkpoint nicht. Worker-lokale ``model_path`` und " +"``model_format`` gehören in die ``engine_config`` jeder Replik. Xavier prüft" +" Gewichte, Tokenizer, effektive Modellsemantik, Kontextgrenzen und Prompt-" +"Token-IDs. Modell- und Sampling-Einschränkungen entsprechen " +"engineübergreifendem vLLM/SGLang-P/D; Fehler werden gemeldet." + +#: user_guide/pd_separation.rst +msgid "" +"To use MLX prefill, swap the two ``role`` values in the example. For SGLang " +"on either role, use ``SGLang`` and replace ``max_model_len`` with " +"``context_length``. NVIDIA requires ``xoscar[nixl]>=0.11.1``; Mac requires " +"``xoscar>=0.11.1`` and ``mlx-lm>=0.31.2``. Workers must reach the supervisor" +" and each other. Use Xavier transport, and omit ``xavier_gpu_cache_bytes`` " +"or set it to ``0``. Relaunch after a worker restart." +msgstr "" +"Für MLX-Prefill tauschen Sie die beiden ``role``-Werte im Beispiel. Für " +"SGLang in jeder Rolle verwenden Sie ``SGLang`` und ersetzen " +"``max_model_len`` durch ``context_length``. NVIDIA benötigt " +"``xoscar[nixl]>=0.11.1``, Mac ``xoscar>=0.11.1`` und ``mlx-lm>=0.31.2``. " +"Worker müssen den Supervisor und einander erreichen. Verwenden Sie Xavier, " +"lassen Sie ``xavier_gpu_cache_bytes`` weg oder setzen es auf ``0``. Nach " +"einem Worker-Neustart erneut bereitstellen." + +#: ../../source/user_guide/pd_separation.rst +msgid "" +"Run the supervisor and all worker/model environments with the same Python " +"version and matching ``cloudpickle`` / ``tblib`` versions; actor RPC also " +"carries exceptions and generators." +msgstr "" +"Supervisor und alle Worker-/Modellumgebungen müssen dieselbe Python-Version " +"und übereinstimmende Versionen von ``cloudpickle`` / ``tblib`` verwenden; " +"Actor-RPC überträgt auch Ausnahmen und Generatoren." + +#: user_guide/pd_separation.rst +msgid "" +"Compare end-to-end TTFT with the standalone decode engine on its own " +"hardware, using the same effective FP16 weights and cache policy. Network " +"transfers and CUDA/Metal copies may outweigh faster prefill. ``host_bytes`` " +"measures CPU-staged payloads; ``imported_tokens`` records completed KV " +"imports. Correct output alone does not prove reuse." +msgstr "" +"Vergleichen Sie die Ende-zu-Ende-TTFT mit der eigenständig auf ihrer " +"Hardware laufenden Decode-Engine bei gleichen effektiven FP16-Gewichten und " +"Cache-Regeln. Netzwerktransfers und CUDA/Metal-Kopien können schnelleres " +"Prefill überwiegen. ``host_bytes`` misst CPU-zwischengespeicherte Nutzdaten;" +" ``imported_tokens`` erfasst abgeschlossene KV-Importe. Korrekte Ausgaben " +"allein beweisen keine Wiederverwendung." + +#: source/user_guide/pd_separation.rst:179 +msgid "" +"For MLX prefill, ``xavier_cache_bytes`` sets the worker-local in-flight CPU " +"page budget (default: 512 MiB). Capacity is reserved before Metal prefill, " +"rounding the full prompt up to 64-token pages. A prompt exceeding the budget" +" fails before prefill; concurrent requests wait for capacity. Reservations " +"and transfer waits are bounded by " +"``XINFERENCE_SGLANG_XAVIER_TRANSFER_TIMEOUT`` (default: 600 seconds)." +msgstr "" +"Bei MLX-Prefill legt ``xavier_cache_bytes`` das workerlokale Budget für CPU-" +"Seiten laufender Anfragen fest (Standard: 512 MiB). Die Kapazität wird vor " +"dem Metal-Prefill reserviert; der gesamte Prompt wird auf 64-Token-Seiten " +"aufgerundet. Ein Prompt über dem Budget scheitert vor dem Prefill; parallele" +" Anfragen warten auf Kapazität. Reservierungen und Transferwartezeiten sind " +"durch ``XINFERENCE_SGLANG_XAVIER_TRANSFER_TIMEOUT`` begrenzt (Standard: 600 " +"Sekunden)." + +#: source/user_guide/pd_separation.rst:181 +msgid "" +"MLX-prefill cross-engine requests do not accept ``min_p``, " +"``presence_penalty``, ``frequency_penalty`` or ``seed``. The router rejects " +"these parameters before scheduling either replica, so the first output token" +" uses the same supported sampling settings as the rest of the response." +msgstr "" +"Engineübergreifende Anfragen mit MLX-Prefill akzeptieren weder ``min_p`` " +"noch ``presence_penalty``, ``frequency_penalty`` oder ``seed``. Der Router " +"lehnt diese Parameter vor der Planung beider Replikate ab, sodass das erste " +"Ausgabetoken dieselben unterstützten Sampling-Einstellungen wie die " +"restliche Antwort verwendet." diff --git a/doc/source/locale/es/LC_MESSAGES/user_guide/pd_separation.po b/doc/source/locale/es/LC_MESSAGES/user_guide/pd_separation.po index faa386e1f9..f3e335b799 100644 --- a/doc/source/locale/es/LC_MESSAGES/user_guide/pd_separation.po +++ b/doc/source/locale/es/LC_MESSAGES/user_guide/pd_separation.po @@ -495,15 +495,13 @@ msgid "" "Use the default ``xavier`` transport for MLX; native NIXL is unsupported. " "All workers must reach the supervisor's actor address. Relaunch after a " "worker restart. Two replicas on one Mac share its Metal GPU; measure your " -"workload before expecting a throughput benefit. Cross-engine P/D involving " -"MLX, including NVIDIA-to-Mac handoff, is not yet supported." +"workload before expecting a throughput benefit." msgstr "" "Use el transporte predeterminado ``xavier`` para MLX; NIXL nativo no es " "compatible. Todos los workers deben acceder a la dirección actor del " "supervisor. Reinicie el despliegue tras reiniciar un worker. Dos réplicas en" " un Mac comparten su GPU Metal; mida su carga antes de esperar mejoras de " -"rendimiento. El P/D entre motores con MLX, incluida la transferencia de " -"NVIDIA a Mac, aún no es compatible." +"rendimiento." #: user_guide/pd_separation.rst:102 msgid "" @@ -588,17 +586,126 @@ msgstr "" "salida. Los kernels pueden producir pequeñas diferencias numéricas. vLLM " "decode requiere al menos dos tokens de prompt." -#: source/user_guide/pd_separation.rst +#: user_guide/pd_separation.rst msgid "" "Cross-engine P/D does not support CPU fallback, retained Xavier history, " "logprobs, structured sampling, LoRA, speculative decoding, multimodal inputs" " or hybrid attention. Omit ``xavier_gpu_cache_bytes`` or set it to ``0``; do" " not set ``xavier_cache_bytes``. Relaunch after a worker restart. Native " -"NIXL comparisons use the same engine on both roles. MLX is unsupported." +"NIXL comparisons use the same engine on both roles." msgstr "" "El P/D entre motores no admite retorno a CPU, historial Xavier, logprobs, " "muestreo estructurado, LoRA, decodificación especulativa, entradas " "multimodales ni atención híbrida. Omita ``xavier_gpu_cache_bytes`` o " "establézcalo en ``0``; no configure ``xavier_cache_bytes``. Vuelva a " "desplegar después de reiniciar un worker. Las comparaciones NIXL nativas " -"usan el mismo motor en ambos roles. MLX no es compatible." +"usan el mismo motor en ambos roles." + +#: user_guide/pd_separation.rst +msgid "NVIDIA and MLX cross-engine P/D" +msgstr "P/D entre motores NVIDIA y MLX" + +#: user_guide/pd_separation.rst +msgid "" +"Xavier supports both vLLM/SGLang prefill with MLX decode and MLX prefill " +"with vLLM/SGLang decode, with streaming and non-streaming requests. FP16 " +"64-token pages pass through CPU staging and actor RPC between CUDA and " +"Metal. vLLM and MLX decode compute the last prompt token; SGLang decode " +"imports the full prompt KV and the first output token sampled by P. MLX " +"prefill publishes request-owned pages on its worker with a 512 MiB in-flight" +" budget." +msgstr "" +"Xavier admite prefill vLLM/SGLang con decode MLX y prefill MLX con decode " +"vLLM/SGLang, con y sin streaming. Las páginas FP16 de 64 tokens pasan por " +"almacenamiento temporal CPU y RPC actor entre CUDA y Metal. El decode vLLM y" +" MLX calcula el último token del prompt; SGLang importa el KV completo del " +"prompt y el primer token de salida muestreado por P. El prefill MLX publica " +"páginas de cada solicitud en su worker con un presupuesto de 512 MiB para " +"transferencias en curso." + +#: user_guide/pd_separation.rst +msgid "" +"Use identical unquantized checkpoint and tokenizer files; do not convert the" +" Mac checkpoint. Worker-local ``model_path`` and ``model_format`` belong in " +"each replica's ``engine_config``. Xavier checks weights, tokenizer, " +"effective model semantics, context limits and prompt token IDs. Model and " +"sampling restrictions match vLLM/SGLang cross-engine P/D; failures raise " +"errors." +msgstr "" +"Use archivos idénticos de checkpoint sin cuantizar y tokenizer; no convierta" +" el checkpoint Mac. Los ``model_path`` y ``model_format`` locales al worker " +"van en el ``engine_config`` de cada réplica. Xavier comprueba pesos, " +"tokenizer, semántica efectiva del modelo, límites de contexto e ID de tokens" +" del prompt. Las restricciones de modelo y muestreo coinciden con P/D " +"vLLM/SGLang entre motores; los fallos generan errores." + +#: user_guide/pd_separation.rst +msgid "" +"To use MLX prefill, swap the two ``role`` values in the example. For SGLang " +"on either role, use ``SGLang`` and replace ``max_model_len`` with " +"``context_length``. NVIDIA requires ``xoscar[nixl]>=0.11.1``; Mac requires " +"``xoscar>=0.11.1`` and ``mlx-lm>=0.31.2``. Workers must reach the supervisor" +" and each other. Use Xavier transport, and omit ``xavier_gpu_cache_bytes`` " +"or set it to ``0``. Relaunch after a worker restart." +msgstr "" +"Para prefill MLX, intercambie los dos valores ``role`` del ejemplo. Para " +"SGLang en cualquier rol, use ``SGLang`` y cambie ``max_model_len`` por " +"``context_length``. NVIDIA requiere ``xoscar[nixl]>=0.11.1``; Mac requiere " +"``xoscar>=0.11.1`` y ``mlx-lm>=0.31.2``. Los workers deben alcanzar al " +"supervisor y entre sí. Use Xavier, omita ``xavier_gpu_cache_bytes`` o " +"póngalo a ``0``. Vuelva a desplegar tras reiniciar un worker." + +#: ../../source/user_guide/pd_separation.rst +msgid "" +"Run the supervisor and all worker/model environments with the same Python " +"version and matching ``cloudpickle`` / ``tblib`` versions; actor RPC also " +"carries exceptions and generators." +msgstr "" +"Ejecute el supervisor y todos los entornos worker/model con la misma versión" +" de Python y versiones coincidentes de ``cloudpickle`` / ``tblib``; los RPC " +"actor también transportan excepciones y generadores." + +#: user_guide/pd_separation.rst +msgid "" +"Compare end-to-end TTFT with the standalone decode engine on its own " +"hardware, using the same effective FP16 weights and cache policy. Network " +"transfers and CUDA/Metal copies may outweigh faster prefill. ``host_bytes`` " +"measures CPU-staged payloads; ``imported_tokens`` records completed KV " +"imports. Correct output alone does not prove reuse." +msgstr "" +"Compare TTFT de extremo a extremo con el motor de decode ejecutado solo en " +"su hardware, usando los mismos pesos FP16 efectivos y política de caché. Las" +" transferencias de red y copias CUDA/Metal pueden superar la mejora del " +"prefill. ``host_bytes`` mide cargas que pasan por CPU; ``imported_tokens`` " +"registra importaciones KV completadas. La salida correcta por sí sola no " +"prueba la reutilización." + +#: source/user_guide/pd_separation.rst:179 +msgid "" +"For MLX prefill, ``xavier_cache_bytes`` sets the worker-local in-flight CPU " +"page budget (default: 512 MiB). Capacity is reserved before Metal prefill, " +"rounding the full prompt up to 64-token pages. A prompt exceeding the budget" +" fails before prefill; concurrent requests wait for capacity. Reservations " +"and transfer waits are bounded by " +"``XINFERENCE_SGLANG_XAVIER_TRANSFER_TIMEOUT`` (default: 600 seconds)." +msgstr "" +"Para prefill MLX, ``xavier_cache_bytes`` establece el presupuesto de páginas" +" CPU en curso del worker (por defecto: 512 MiB). La capacidad se reserva " +"antes del prefill Metal, redondeando el prompt completo a páginas de 64 " +"tokens. Un prompt que supera el presupuesto falla antes del prefill; las " +"solicitudes concurrentes esperan capacidad. Las reservas y esperas de " +"transferencia están limitadas por " +"``XINFERENCE_SGLANG_XAVIER_TRANSFER_TIMEOUT`` (por defecto: 600 segundos)." + +#: source/user_guide/pd_separation.rst:181 +msgid "" +"MLX-prefill cross-engine requests do not accept ``min_p``, " +"``presence_penalty``, ``frequency_penalty`` or ``seed``. The router rejects " +"these parameters before scheduling either replica, so the first output token" +" uses the same supported sampling settings as the rest of the response." +msgstr "" +"Las solicitudes entre motores con prefill MLX no aceptan ``min_p``, " +"``presence_penalty``, ``frequency_penalty`` ni ``seed``. El router rechaza " +"estos parámetros antes de programar cualquiera de las réplicas, de modo que " +"el primer token usa los mismos ajustes de muestreo admitidos que el resto de" +" la respuesta." diff --git a/doc/source/locale/fr/LC_MESSAGES/user_guide/pd_separation.po b/doc/source/locale/fr/LC_MESSAGES/user_guide/pd_separation.po index 826c6b443a..87ad3e0f84 100644 --- a/doc/source/locale/fr/LC_MESSAGES/user_guide/pd_separation.po +++ b/doc/source/locale/fr/LC_MESSAGES/user_guide/pd_separation.po @@ -502,15 +502,13 @@ msgid "" "Use the default ``xavier`` transport for MLX; native NIXL is unsupported. " "All workers must reach the supervisor's actor address. Relaunch after a " "worker restart. Two replicas on one Mac share its Metal GPU; measure your " -"workload before expecting a throughput benefit. Cross-engine P/D involving " -"MLX, including NVIDIA-to-Mac handoff, is not yet supported." +"workload before expecting a throughput benefit." msgstr "" "Utilisez le transport par défaut ``xavier`` pour MLX ; NIXL natif n’est pas " "pris en charge. Tous les workers doivent pouvoir joindre l’adresse actor du " "supervisor. Relancez le déploiement après le redémarrage d’un worker. Deux " "répliques sur un même Mac partagent son GPU Metal ; mesurez votre charge " -"avant d’attendre un gain de débit. Le P/D entre moteurs impliquant MLX, " -"notamment de NVIDIA vers Mac, n’est pas encore pris en charge." +"avant d’attendre un gain de débit." #: user_guide/pd_separation.rst:102 msgid "" @@ -596,18 +594,128 @@ msgstr "" "échantillonne la sortie. Les kernels peuvent produire de faibles différences" " numériques. vLLM decode exige au moins deux tokens de prompt." -#: source/user_guide/pd_separation.rst +#: user_guide/pd_separation.rst msgid "" "Cross-engine P/D does not support CPU fallback, retained Xavier history, " "logprobs, structured sampling, LoRA, speculative decoding, multimodal inputs" " or hybrid attention. Omit ``xavier_gpu_cache_bytes`` or set it to ``0``; do" " not set ``xavier_cache_bytes``. Relaunch after a worker restart. Native " -"NIXL comparisons use the same engine on both roles. MLX is unsupported." +"NIXL comparisons use the same engine on both roles." msgstr "" "Le P/D entre moteurs ne prend pas en charge le repli CPU, l’historique " "Xavier, logprobs, l’échantillonnage structuré, LoRA, le décodage spéculatif," " les entrées multimodales ni l’attention hybride. Omettez " "``xavier_gpu_cache_bytes`` ou réglez-le à ``0`` ; ne définissez pas " "``xavier_cache_bytes``. Redéployez après un redémarrage du worker. Les " -"comparaisons NIXL natives utilisent le même moteur pour les deux rôles. MLX " -"n’est pas pris en charge." +"comparaisons NIXL natives utilisent le même moteur pour les deux rôles." + +#: user_guide/pd_separation.rst +msgid "NVIDIA and MLX cross-engine P/D" +msgstr "P/D entre moteurs NVIDIA et MLX" + +#: user_guide/pd_separation.rst +msgid "" +"Xavier supports both vLLM/SGLang prefill with MLX decode and MLX prefill " +"with vLLM/SGLang decode, with streaming and non-streaming requests. FP16 " +"64-token pages pass through CPU staging and actor RPC between CUDA and " +"Metal. vLLM and MLX decode compute the last prompt token; SGLang decode " +"imports the full prompt KV and the first output token sampled by P. MLX " +"prefill publishes request-owned pages on its worker with a 512 MiB in-flight" +" budget." +msgstr "" +"Xavier prend en charge le prefill vLLM/SGLang avec decode MLX et le prefill " +"MLX avec decode vLLM/SGLang, avec ou sans streaming. Les pages FP16 de 64 " +"tokens passent par la mémoire CPU et les RPC actor entre CUDA et Metal. Le " +"decode vLLM et MLX calcule le dernier token du prompt ; SGLang importe le KV" +" complet du prompt et le premier token de sortie échantillonné par P. Le " +"prefill MLX publie des pages propres à chaque requête sur son worker, avec " +"un budget de 512 MiB pour les transferts en cours." + +#: user_guide/pd_separation.rst +msgid "" +"Use identical unquantized checkpoint and tokenizer files; do not convert the" +" Mac checkpoint. Worker-local ``model_path`` and ``model_format`` belong in " +"each replica's ``engine_config``. Xavier checks weights, tokenizer, " +"effective model semantics, context limits and prompt token IDs. Model and " +"sampling restrictions match vLLM/SGLang cross-engine P/D; failures raise " +"errors." +msgstr "" +"Utilisez des fichiers de checkpoint non quantifié et de tokenizer identiques" +" ; ne convertissez pas le checkpoint Mac. Les ``model_path`` et " +"``model_format`` locaux au worker se définissent dans l’``engine_config`` de" +" chaque réplique. Xavier vérifie poids, tokenizer, sémantique effective du " +"modèle, limites de contexte et ID des tokens du prompt. Les restrictions de " +"modèle et d’échantillonnage correspondent au P/D vLLM/SGLang entre moteurs ;" +" les échecs lèvent des erreurs." + +#: user_guide/pd_separation.rst +msgid "" +"To use MLX prefill, swap the two ``role`` values in the example. For SGLang " +"on either role, use ``SGLang`` and replace ``max_model_len`` with " +"``context_length``. NVIDIA requires ``xoscar[nixl]>=0.11.1``; Mac requires " +"``xoscar>=0.11.1`` and ``mlx-lm>=0.31.2``. Workers must reach the supervisor" +" and each other. Use Xavier transport, and omit ``xavier_gpu_cache_bytes`` " +"or set it to ``0``. Relaunch after a worker restart." +msgstr "" +"Pour le prefill MLX, échangez les deux valeurs ``role`` de l’exemple. Pour " +"SGLang dans chaque rôle, utilisez ``SGLang`` et remplacez ``max_model_len`` " +"par ``context_length``. NVIDIA nécessite ``xoscar[nixl]>=0.11.1`` ; Mac " +"nécessite ``xoscar>=0.11.1`` et ``mlx-lm>=0.31.2``. Les workers doivent " +"joindre le supervisor et les autres workers. Utilisez Xavier, omettez " +"``xavier_gpu_cache_bytes`` ou mettez-le à ``0``. Redéployez après un " +"redémarrage du worker." + +#: ../../source/user_guide/pd_separation.rst +msgid "" +"Run the supervisor and all worker/model environments with the same Python " +"version and matching ``cloudpickle`` / ``tblib`` versions; actor RPC also " +"carries exceptions and generators." +msgstr "" +"Exécutez le supervisor et tous les environnements worker/model avec la même " +"version de Python et des versions identiques de ``cloudpickle`` / ``tblib`` " +"; les RPC actor transportent aussi des exceptions et des générateurs." + +#: user_guide/pd_separation.rst +msgid "" +"Compare end-to-end TTFT with the standalone decode engine on its own " +"hardware, using the same effective FP16 weights and cache policy. Network " +"transfers and CUDA/Metal copies may outweigh faster prefill. ``host_bytes`` " +"measures CPU-staged payloads; ``imported_tokens`` records completed KV " +"imports. Correct output alone does not prove reuse." +msgstr "" +"Comparez le TTFT de bout en bout au moteur de decode exécuté seul sur son " +"matériel, avec les mêmes poids FP16 effectifs et règles de cache. Les " +"transferts réseau et copies CUDA/Metal peuvent dépasser le gain de prefill. " +"``host_bytes`` mesure les données transitant par le CPU ; " +"``imported_tokens`` compte les imports KV terminés. Une sortie correcte ne " +"prouve pas à elle seule la réutilisation." + +#: source/user_guide/pd_separation.rst:179 +msgid "" +"For MLX prefill, ``xavier_cache_bytes`` sets the worker-local in-flight CPU " +"page budget (default: 512 MiB). Capacity is reserved before Metal prefill, " +"rounding the full prompt up to 64-token pages. A prompt exceeding the budget" +" fails before prefill; concurrent requests wait for capacity. Reservations " +"and transfer waits are bounded by " +"``XINFERENCE_SGLANG_XAVIER_TRANSFER_TIMEOUT`` (default: 600 seconds)." +msgstr "" +"Pour le prefill MLX, ``xavier_cache_bytes`` définit le budget de pages CPU " +"des requêtes en cours sur le worker (par défaut : 512 MiB). La capacité est " +"réservée avant le prefill Metal, en arrondissant le prompt complet à des " +"pages de 64 tokens. Un prompt dépassant le budget échoue avant le prefill ; " +"les requêtes concurrentes attendent de la capacité. Les réservations et les " +"attentes de transfert sont bornées par " +"``XINFERENCE_SGLANG_XAVIER_TRANSFER_TIMEOUT`` (par défaut : 600 secondes)." + +#: source/user_guide/pd_separation.rst:181 +msgid "" +"MLX-prefill cross-engine requests do not accept ``min_p``, " +"``presence_penalty``, ``frequency_penalty`` or ``seed``. The router rejects " +"these parameters before scheduling either replica, so the first output token" +" uses the same supported sampling settings as the rest of the response." +msgstr "" +"Les requêtes entre moteurs avec prefill MLX n’acceptent pas ``min_p``, " +"``presence_penalty``, ``frequency_penalty`` ni ``seed``. Le routeur rejette " +"ces paramètres avant de planifier les deux réplicas, afin que le premier " +"token utilise les mêmes réglages d’échantillonnage pris en charge que le " +"reste de la réponse." diff --git a/doc/source/locale/it/LC_MESSAGES/user_guide/pd_separation.po b/doc/source/locale/it/LC_MESSAGES/user_guide/pd_separation.po index eb0b346bb0..74244ed231 100644 --- a/doc/source/locale/it/LC_MESSAGES/user_guide/pd_separation.po +++ b/doc/source/locale/it/LC_MESSAGES/user_guide/pd_separation.po @@ -497,15 +497,13 @@ msgid "" "Use the default ``xavier`` transport for MLX; native NIXL is unsupported. " "All workers must reach the supervisor's actor address. Relaunch after a " "worker restart. Two replicas on one Mac share its Metal GPU; measure your " -"workload before expecting a throughput benefit. Cross-engine P/D involving " -"MLX, including NVIDIA-to-Mac handoff, is not yet supported." +"workload before expecting a throughput benefit." msgstr "" "Usare il trasporto predefinito ``xavier`` per MLX; NIXL nativo non è " "supportato. Tutti i worker devono raggiungere l’indirizzo actor del " "supervisor. Riavviare il deployment dopo il riavvio di un worker. Due " "repliche sullo stesso Mac condividono la GPU Metal; misurare il carico prima" -" di aspettarsi un aumento del throughput. Il P/D tra motori che coinvolge " -"MLX, incluso il trasferimento da NVIDIA a Mac, non è ancora supportato." +" di aspettarsi un aumento del throughput." #: user_guide/pd_separation.rst:102 msgid "" @@ -590,17 +588,127 @@ msgstr "" "kernel possono produrre piccole differenze numeriche. vLLM decode richiede " "almeno due token di prompt." -#: source/user_guide/pd_separation.rst +#: user_guide/pd_separation.rst msgid "" "Cross-engine P/D does not support CPU fallback, retained Xavier history, " "logprobs, structured sampling, LoRA, speculative decoding, multimodal inputs" " or hybrid attention. Omit ``xavier_gpu_cache_bytes`` or set it to ``0``; do" " not set ``xavier_cache_bytes``. Relaunch after a worker restart. Native " -"NIXL comparisons use the same engine on both roles. MLX is unsupported." +"NIXL comparisons use the same engine on both roles." msgstr "" "Il P/D tra motori non supporta fallback CPU, cronologia Xavier, logprobs, " "campionamento strutturato, LoRA, decodifica speculativa, input multimodali o" " attenzione ibrida. Omettere ``xavier_gpu_cache_bytes`` o impostarlo a " "``0``; non impostare ``xavier_cache_bytes``. Dopo il riavvio di un worker è " "necessario ridistribuire. I confronti NIXL nativi usano lo stesso motore per" -" entrambi i ruoli. MLX non è supportato." +" entrambi i ruoli." + +#: user_guide/pd_separation.rst +msgid "NVIDIA and MLX cross-engine P/D" +msgstr "P/D tra motori NVIDIA e MLX" + +#: user_guide/pd_separation.rst +msgid "" +"Xavier supports both vLLM/SGLang prefill with MLX decode and MLX prefill " +"with vLLM/SGLang decode, with streaming and non-streaming requests. FP16 " +"64-token pages pass through CPU staging and actor RPC between CUDA and " +"Metal. vLLM and MLX decode compute the last prompt token; SGLang decode " +"imports the full prompt KV and the first output token sampled by P. MLX " +"prefill publishes request-owned pages on its worker with a 512 MiB in-flight" +" budget." +msgstr "" +"Xavier supporta prefill vLLM/SGLang con decode MLX e prefill MLX con decode " +"vLLM/SGLang, con e senza streaming. Le pagine FP16 da 64 token passano " +"tramite staging CPU e RPC actor tra CUDA e Metal. Il decode vLLM e MLX " +"calcola l’ultimo token del prompt; SGLang importa il KV completo del prompt " +"e il primo token di output campionato da P. Il prefill MLX pubblica pagine " +"proprie di ogni richiesta sul worker con un budget di 512 MiB per " +"trasferimenti in corso." + +#: user_guide/pd_separation.rst +msgid "" +"Use identical unquantized checkpoint and tokenizer files; do not convert the" +" Mac checkpoint. Worker-local ``model_path`` and ``model_format`` belong in " +"each replica's ``engine_config``. Xavier checks weights, tokenizer, " +"effective model semantics, context limits and prompt token IDs. Model and " +"sampling restrictions match vLLM/SGLang cross-engine P/D; failures raise " +"errors." +msgstr "" +"Usare file identici di checkpoint non quantizzato e tokenizer; non " +"convertire il checkpoint Mac. I ``model_path`` e ``model_format`` locali al " +"worker vanno nell’``engine_config`` di ogni replica. Xavier verifica pesi, " +"tokenizer, semantica effettiva del modello, limiti del contesto e ID dei " +"token del prompt. Le restrizioni di modello e campionamento corrispondono al" +" P/D vLLM/SGLang tra motori; gli errori vengono segnalati." + +#: user_guide/pd_separation.rst +msgid "" +"To use MLX prefill, swap the two ``role`` values in the example. For SGLang " +"on either role, use ``SGLang`` and replace ``max_model_len`` with " +"``context_length``. NVIDIA requires ``xoscar[nixl]>=0.11.1``; Mac requires " +"``xoscar>=0.11.1`` and ``mlx-lm>=0.31.2``. Workers must reach the supervisor" +" and each other. Use Xavier transport, and omit ``xavier_gpu_cache_bytes`` " +"or set it to ``0``. Relaunch after a worker restart." +msgstr "" +"Per prefill MLX, scambiare i due valori ``role`` nell’esempio. Per SGLang in" +" qualsiasi ruolo, usare ``SGLang`` e sostituire ``max_model_len`` con " +"``context_length``. NVIDIA richiede ``xoscar[nixl]>=0.11.1``; Mac richiede " +"``xoscar>=0.11.1`` e ``mlx-lm>=0.31.2``. I worker devono raggiungere il " +"supervisor e gli altri worker. Usare Xavier, omettere " +"``xavier_gpu_cache_bytes`` o impostarlo a ``0``. Ridistribuire dopo un " +"riavvio del worker." + +#: ../../source/user_guide/pd_separation.rst +msgid "" +"Run the supervisor and all worker/model environments with the same Python " +"version and matching ``cloudpickle`` / ``tblib`` versions; actor RPC also " +"carries exceptions and generators." +msgstr "" +"Eseguire il supervisor e tutti gli ambienti worker/model con la stessa " +"versione di Python e versioni corrispondenti di ``cloudpickle`` / ``tblib``;" +" gli RPC actor trasferiscono anche eccezioni e generatori." + +#: user_guide/pd_separation.rst +msgid "" +"Compare end-to-end TTFT with the standalone decode engine on its own " +"hardware, using the same effective FP16 weights and cache policy. Network " +"transfers and CUDA/Metal copies may outweigh faster prefill. ``host_bytes`` " +"measures CPU-staged payloads; ``imported_tokens`` records completed KV " +"imports. Correct output alone does not prove reuse." +msgstr "" +"Confrontare il TTFT completo con il motore decode eseguito da solo sul " +"proprio hardware usando gli stessi pesi FP16 effettivi e la stessa politica " +"di cache. Trasferimenti di rete e copie CUDA/Metal possono superare il " +"vantaggio del prefill. ``host_bytes`` misura i dati trasferiti tramite CPU; " +"``imported_tokens`` registra le importazioni KV completate. L’output " +"corretto da solo non dimostra il riuso." + +#: source/user_guide/pd_separation.rst:179 +msgid "" +"For MLX prefill, ``xavier_cache_bytes`` sets the worker-local in-flight CPU " +"page budget (default: 512 MiB). Capacity is reserved before Metal prefill, " +"rounding the full prompt up to 64-token pages. A prompt exceeding the budget" +" fails before prefill; concurrent requests wait for capacity. Reservations " +"and transfer waits are bounded by " +"``XINFERENCE_SGLANG_XAVIER_TRANSFER_TIMEOUT`` (default: 600 seconds)." +msgstr "" +"Per prefill MLX, ``xavier_cache_bytes`` imposta il budget delle pagine CPU " +"delle richieste in corso sul worker (predefinito: 512 MiB). La capacità " +"viene riservata prima del prefill Metal, arrotondando il prompt completo a " +"pagine di 64 token. Un prompt che supera il budget fallisce prima del " +"prefill; le richieste concorrenti attendono capacità. Prenotazioni e attese " +"di trasferimento sono limitate da " +"``XINFERENCE_SGLANG_XAVIER_TRANSFER_TIMEOUT`` (predefinito: 600 secondi)." + +#: source/user_guide/pd_separation.rst:181 +msgid "" +"MLX-prefill cross-engine requests do not accept ``min_p``, " +"``presence_penalty``, ``frequency_penalty`` or ``seed``. The router rejects " +"these parameters before scheduling either replica, so the first output token" +" uses the same supported sampling settings as the rest of the response." +msgstr "" +"Le richieste tra motori con prefill MLX non accettano ``min_p``, " +"``presence_penalty``, ``frequency_penalty`` o ``seed``. Il router rifiuta " +"questi parametri prima di pianificare entrambe le repliche, così il primo " +"token usa le stesse impostazioni di campionamento supportate del resto della" +" risposta." diff --git a/doc/source/locale/ja/LC_MESSAGES/user_guide/pd_separation.po b/doc/source/locale/ja/LC_MESSAGES/user_guide/pd_separation.po index 2843d88f82..f98a888e32 100644 --- a/doc/source/locale/ja/LC_MESSAGES/user_guide/pd_separation.po +++ b/doc/source/locale/ja/LC_MESSAGES/user_guide/pd_separation.po @@ -448,13 +448,11 @@ msgid "" "Use the default ``xavier`` transport for MLX; native NIXL is unsupported. " "All workers must reach the supervisor's actor address. Relaunch after a " "worker restart. Two replicas on one Mac share its Metal GPU; measure your " -"workload before expecting a throughput benefit. Cross-engine P/D involving " -"MLX, including NVIDIA-to-Mac handoff, is not yet supported." +"workload before expecting a throughput benefit." msgstr "" "MLX では既定の ``xavier`` 転送を使用します。ネイティブ NIXL は未対応です。全 worker から supervisor の " "actor アドレスに接続できる必要があります。Worker 再起動後はデプロイを再起動してください。同じ Mac 上の 2 レプリカは Metal " -"GPU を共有するため、スループット改善は実際の負荷で測定してください。NVIDIA から Mac への引き継ぎを含む MLX " -"に関わる異なるエンジン間の P/D は未対応です。" +"GPU を共有するため、スループット改善は実際の負荷で測定してください。" #: user_guide/pd_separation.rst:102 msgid "" @@ -531,15 +529,114 @@ msgstr "" "を計算して出力をサンプリングします。エンジンのカーネルによって小さな数値差が生じる場合があります。vLLM decode には最低 2 個の " "prompt token が必要です。" -#: source/user_guide/pd_separation.rst +#: user_guide/pd_separation.rst msgid "" "Cross-engine P/D does not support CPU fallback, retained Xavier history, " "logprobs, structured sampling, LoRA, speculative decoding, multimodal inputs" " or hybrid attention. Omit ``xavier_gpu_cache_bytes`` or set it to ``0``; do" " not set ``xavier_cache_bytes``. Relaunch after a worker restart. Native " -"NIXL comparisons use the same engine on both roles. MLX is unsupported." +"NIXL comparisons use the same engine on both roles." msgstr "" "エンジン間 P/D は CPU フォールバック、Xavier " "履歴、logprobs、構造化サンプリング、LoRA、投機的デコード、マルチモーダル入力、ハイブリッド注意をサポートしません。``xavier_gpu_cache_bytes``" " は省略するか ``0`` にし、``xavier_cache_bytes`` は設定しないでください。ワーカー再起動後は再デプロイします。ネイティブ " -"NIXL の比較では両役割に同一エンジンを使います。MLX は未対応です。" +"NIXL の比較では両役割に同一エンジンを使います。" + +#: user_guide/pd_separation.rst +msgid "NVIDIA and MLX cross-engine P/D" +msgstr "NVIDIA と MLX の異なるエンジン間 P/D" + +#: user_guide/pd_separation.rst +msgid "" +"Xavier supports both vLLM/SGLang prefill with MLX decode and MLX prefill " +"with vLLM/SGLang decode, with streaming and non-streaming requests. FP16 " +"64-token pages pass through CPU staging and actor RPC between CUDA and " +"Metal. vLLM and MLX decode compute the last prompt token; SGLang decode " +"imports the full prompt KV and the first output token sampled by P. MLX " +"prefill publishes request-owned pages on its worker with a 512 MiB in-flight" +" budget." +msgstr "" +"Xavier は vLLM/SGLang prefill と MLX decode、および MLX prefill と vLLM/SGLang " +"decode の両方向で、ストリーミングと非ストリーミングに対応します。FP16 の 64-token ページは CPU ステージングと actor " +"RPC で CUDA と Metal の間を転送されます。vLLM と MLX decode は最後の prompt token を計算し、SGLang" +" decode は完全な prompt KV と P がサンプリングした最初の出力 token をインポートします。MLX prefill は " +"worker 上にリクエスト専用ページを公開し、転送中データの上限は 512 MiB です。" + +#: user_guide/pd_separation.rst +msgid "" +"Use identical unquantized checkpoint and tokenizer files; do not convert the" +" Mac checkpoint. Worker-local ``model_path`` and ``model_format`` belong in " +"each replica's ``engine_config``. Xavier checks weights, tokenizer, " +"effective model semantics, context limits and prompt token IDs. Model and " +"sampling restrictions match vLLM/SGLang cross-engine P/D; failures raise " +"errors." +msgstr "" +"同一の非量子化チェックポイントと tokenizer ファイルを使い、Mac のチェックポイントは変換しないでください。worker ローカルの " +"``model_path`` と ``model_format`` は各レプリカの ``engine_config`` に設定します。Xavier " +"は重み、tokenizer、実効モデルの意味、コンテキスト制限、prompt token ID を検証します。モデルとサンプリングの制限は " +"vLLM/SGLang の異なるエンジン間 P/D と同じで、失敗時はエラーになります。" + +#: user_guide/pd_separation.rst +msgid "" +"To use MLX prefill, swap the two ``role`` values in the example. For SGLang " +"on either role, use ``SGLang`` and replace ``max_model_len`` with " +"``context_length``. NVIDIA requires ``xoscar[nixl]>=0.11.1``; Mac requires " +"``xoscar>=0.11.1`` and ``mlx-lm>=0.31.2``. Workers must reach the supervisor" +" and each other. Use Xavier transport, and omit ``xavier_gpu_cache_bytes`` " +"or set it to ``0``. Relaunch after a worker restart." +msgstr "" +"MLX prefill を使うには例の二つの ``role`` 値を交換します。どちらの役割でも SGLang では ``SGLang`` " +"を使い、``max_model_len`` を ``context_length`` に置き換えます。NVIDIA は " +"``xoscar[nixl]>=0.11.1``、Mac は ``xoscar>=0.11.1`` と ``mlx-lm>=0.31.2`` " +"が必要です。worker は supervisor と相互に接続できる必要があります。Xavier " +"転送を使い、``xavier_gpu_cache_bytes`` は省略するか ``0`` に設定します。worker " +"再起動後は再デプロイしてください。" + +#: ../../source/user_guide/pd_separation.rst +msgid "" +"Run the supervisor and all worker/model environments with the same Python " +"version and matching ``cloudpickle`` / ``tblib`` versions; actor RPC also " +"carries exceptions and generators." +msgstr "" +"Supervisor とすべての worker/model 環境で同じ Python バージョンと一致する ``cloudpickle`` / " +"``tblib`` バージョンを使用してください。actor RPC は例外とジェネレーターも転送します。" + +#: user_guide/pd_separation.rst +msgid "" +"Compare end-to-end TTFT with the standalone decode engine on its own " +"hardware, using the same effective FP16 weights and cache policy. Network " +"transfers and CUDA/Metal copies may outweigh faster prefill. ``host_bytes`` " +"measures CPU-staged payloads; ``imported_tokens`` records completed KV " +"imports. Correct output alone does not prove reuse." +msgstr "" +"同じ実効 FP16 重みとキャッシュ方針で、decode エンジンを元のハードウェア上で単独実行した場合のエンドツーエンド TTFT " +"と比較してください。ネットワーク転送と CUDA/Metal コピーの費用が prefill " +"の高速化を上回る場合があります。``host_bytes`` は CPU 経由のペイロード、``imported_tokens`` は完了した KV " +"インポートを記録します。正しい出力だけでは再利用を証明できません。" + +#: source/user_guide/pd_separation.rst:179 +msgid "" +"For MLX prefill, ``xavier_cache_bytes`` sets the worker-local in-flight CPU " +"page budget (default: 512 MiB). Capacity is reserved before Metal prefill, " +"rounding the full prompt up to 64-token pages. A prompt exceeding the budget" +" fails before prefill; concurrent requests wait for capacity. Reservations " +"and transfer waits are bounded by " +"``XINFERENCE_SGLANG_XAVIER_TRANSFER_TIMEOUT`` (default: 600 seconds)." +msgstr "" +"MLX prefill では、``xavier_cache_bytes`` で worker ローカルの処理中 CPU " +"ページ予算を設定します（既定：512 MiB）。Metal prefill " +"前に容量を予約し、プロンプト全体を64トークンのページに切り上げます。予算を超えるプロンプトは prefill " +"前に失敗し、並行リクエストは空き容量を待ちます。予約と転送の待機は " +"``XINFERENCE_SGLANG_XAVIER_TRANSFER_TIMEOUT`` で制限されます（既定：600秒）。" + +#: source/user_guide/pd_separation.rst:181 +msgid "" +"MLX-prefill cross-engine requests do not accept ``min_p``, " +"``presence_penalty``, ``frequency_penalty`` or ``seed``. The router rejects " +"these parameters before scheduling either replica, so the first output token" +" uses the same supported sampling settings as the rest of the response." +msgstr "" +"MLX prefill のエンジン間リクエストでは " +"``min_p``、``presence_penalty``、``frequency_penalty``、``seed`` " +"を指定できません。router " +"はいずれのレプリカもスケジュールする前にこれらのパラメーターを拒否するため、最初の出力トークンにも応答の残りと同じ対応済みサンプリング設定を使います。" diff --git a/doc/source/locale/ko/LC_MESSAGES/user_guide/pd_separation.po b/doc/source/locale/ko/LC_MESSAGES/user_guide/pd_separation.po index 3a9bbd93d4..567154affb 100644 --- a/doc/source/locale/ko/LC_MESSAGES/user_guide/pd_separation.po +++ b/doc/source/locale/ko/LC_MESSAGES/user_guide/pd_separation.po @@ -447,13 +447,11 @@ msgid "" "Use the default ``xavier`` transport for MLX; native NIXL is unsupported. " "All workers must reach the supervisor's actor address. Relaunch after a " "worker restart. Two replicas on one Mac share its Metal GPU; measure your " -"workload before expecting a throughput benefit. Cross-engine P/D involving " -"MLX, including NVIDIA-to-Mac handoff, is not yet supported." +"workload before expecting a throughput benefit." msgstr "" "MLX에는 기본 ``xavier`` 전송을 사용하세요. 네이티브 NIXL은 지원하지 않습니다. 모든 worker는 supervisor의 " "actor 주소에 접근할 수 있어야 합니다. Worker 재시작 후 배포를 다시 시작하세요. 같은 Mac의 두 복제본은 Metal " -"GPU를 공유하므로 처리량 향상을 기대하기 전에 워크로드를 측정하세요. NVIDIA에서 Mac으로의 인계를 포함한 MLX 관련 엔진 간 " -"P/D는 아직 지원하지 않습니다." +"GPU를 공유하므로 처리량 향상을 기대하기 전에 워크로드를 측정하세요." #: user_guide/pd_separation.rst:102 msgid "" @@ -529,15 +527,110 @@ msgstr "" "샘플링합니다. 엔진 커널에 따라 작은 수치 차이가 발생할 수 있습니다. vLLM decode에는 최소 두 개의 prompt token이 " "필요합니다." -#: source/user_guide/pd_separation.rst +#: user_guide/pd_separation.rst msgid "" "Cross-engine P/D does not support CPU fallback, retained Xavier history, " "logprobs, structured sampling, LoRA, speculative decoding, multimodal inputs" " or hybrid attention. Omit ``xavier_gpu_cache_bytes`` or set it to ``0``; do" " not set ``xavier_cache_bytes``. Relaunch after a worker restart. Native " -"NIXL comparisons use the same engine on both roles. MLX is unsupported." +"NIXL comparisons use the same engine on both roles." msgstr "" "엔진 간 P/D는 CPU 폴백, Xavier 기록 캐시, logprobs, 구조화 샘플링, LoRA, 추측 디코딩, 멀티모달 입력 또는 " "하이브리드 어텐션을 지원하지 않습니다. ``xavier_gpu_cache_bytes`` 를 생략하거나 ``0`` 으로 설정하고 " "``xavier_cache_bytes`` 는 설정하지 마세요. 워커 재시작 후 다시 배포하세요. 네이티브 NIXL 비교는 두 역할에 같은" -" 엔진을 사용합니다. MLX는 지원하지 않습니다." +" 엔진을 사용합니다." + +#: user_guide/pd_separation.rst +msgid "NVIDIA and MLX cross-engine P/D" +msgstr "NVIDIA와 MLX 엔진 간 P/D" + +#: user_guide/pd_separation.rst +msgid "" +"Xavier supports both vLLM/SGLang prefill with MLX decode and MLX prefill " +"with vLLM/SGLang decode, with streaming and non-streaming requests. FP16 " +"64-token pages pass through CPU staging and actor RPC between CUDA and " +"Metal. vLLM and MLX decode compute the last prompt token; SGLang decode " +"imports the full prompt KV and the first output token sampled by P. MLX " +"prefill publishes request-owned pages on its worker with a 512 MiB in-flight" +" budget." +msgstr "" +"Xavier는 vLLM/SGLang prefill과 MLX decode, MLX prefill과 vLLM/SGLang decode " +"양방향에서 스트리밍 및 비스트리밍을 지원합니다. FP16 64-token 페이지는 CPU 스테이징과 actor RPC로 CUDA와 " +"Metal 사이를 이동합니다. vLLM과 MLX decode는 마지막 prompt token을 계산하고 SGLang decode는 전체 " +"prompt KV와 P가 샘플링한 첫 출력 token을 가져옵니다. MLX prefill은 worker에 요청 전용 페이지를 게시하며 " +"전송 중 데이터 예산은 512 MiB입니다." + +#: user_guide/pd_separation.rst +msgid "" +"Use identical unquantized checkpoint and tokenizer files; do not convert the" +" Mac checkpoint. Worker-local ``model_path`` and ``model_format`` belong in " +"each replica's ``engine_config``. Xavier checks weights, tokenizer, " +"effective model semantics, context limits and prompt token IDs. Model and " +"sampling restrictions match vLLM/SGLang cross-engine P/D; failures raise " +"errors." +msgstr "" +"동일한 비양자화 체크포인트와 tokenizer 파일을 사용하고 Mac 체크포인트를 변환하지 마세요. worker 로컬 " +"``model_path`` 와 ``model_format`` 은 각 복제본의 ``engine_config`` 에 설정합니다. " +"Xavier는 가중치, tokenizer, 실제 모델 의미, 컨텍스트 제한과 prompt token ID를 검사합니다. 모델 및 샘플링 " +"제한은 vLLM/SGLang 엔진 간 P/D와 같으며 실패하면 오류를 반환합니다." + +#: user_guide/pd_separation.rst +msgid "" +"To use MLX prefill, swap the two ``role`` values in the example. For SGLang " +"on either role, use ``SGLang`` and replace ``max_model_len`` with " +"``context_length``. NVIDIA requires ``xoscar[nixl]>=0.11.1``; Mac requires " +"``xoscar>=0.11.1`` and ``mlx-lm>=0.31.2``. Workers must reach the supervisor" +" and each other. Use Xavier transport, and omit ``xavier_gpu_cache_bytes`` " +"or set it to ``0``. Relaunch after a worker restart." +msgstr "" +"MLX prefill을 사용하려면 예제의 두 ``role`` 값을 바꾸세요. 어느 역할이든 SGLang에는 ``SGLang`` 을 " +"사용하고 ``max_model_len`` 을 ``context_length`` 로 바꾸세요. NVIDIA는 " +"``xoscar[nixl]>=0.11.1``, Mac은 ``xoscar>=0.11.1`` 과 ``mlx-lm>=0.31.2`` 가 " +"필요합니다. worker는 supervisor와 서로 접근할 수 있어야 합니다. Xavier 전송을 사용하고 " +"``xavier_gpu_cache_bytes`` 는 생략하거나 ``0`` 으로 설정하세요. worker 재시작 후 다시 배포하세요." + +#: ../../source/user_guide/pd_separation.rst +msgid "" +"Run the supervisor and all worker/model environments with the same Python " +"version and matching ``cloudpickle`` / ``tblib`` versions; actor RPC also " +"carries exceptions and generators." +msgstr "" +"Supervisor와 모든 worker/model 환경에서 동일한 Python 버전과 일치하는 ``cloudpickle`` / " +"``tblib`` 버전을 사용하세요. actor RPC는 예외와 생성기도 전송합니다." + +#: user_guide/pd_separation.rst +msgid "" +"Compare end-to-end TTFT with the standalone decode engine on its own " +"hardware, using the same effective FP16 weights and cache policy. Network " +"transfers and CUDA/Metal copies may outweigh faster prefill. ``host_bytes`` " +"measures CPU-staged payloads; ``imported_tokens`` records completed KV " +"imports. Correct output alone does not prove reuse." +msgstr "" +"동일한 실제 FP16 가중치와 캐시 정책으로 원래 하드웨어에서 독립 실행한 decode 엔진의 종단 간 TTFT와 비교하세요. 네트워크 " +"전송과 CUDA/Metal 복사 비용이 prefill 가속 이득보다 클 수 있습니다. ``host_bytes`` 는 CPU 스테이징 " +"페이로드를, ``imported_tokens`` 는 완료된 KV 가져오기를 기록합니다. 올바른 출력만으로 재사용을 증명할 수 없습니다." + +#: source/user_guide/pd_separation.rst:179 +msgid "" +"For MLX prefill, ``xavier_cache_bytes`` sets the worker-local in-flight CPU " +"page budget (default: 512 MiB). Capacity is reserved before Metal prefill, " +"rounding the full prompt up to 64-token pages. A prompt exceeding the budget" +" fails before prefill; concurrent requests wait for capacity. Reservations " +"and transfer waits are bounded by " +"``XINFERENCE_SGLANG_XAVIER_TRANSFER_TIMEOUT`` (default: 600 seconds)." +msgstr "" +"MLX prefill에서는 ``xavier_cache_bytes`` 로 worker 로컬의 처리 중 CPU 페이지 예산을 " +"설정합니다(기본값: 512 MiB). Metal prefill 전에 용량을 예약하며 전체 프롬프트를 64토큰 페이지 단위로 올림합니다. " +"예산을 초과하는 프롬프트는 prefill 전에 실패하고 동시 요청은 여유 용량을 기다립니다. 예약과 전송 대기는 " +"``XINFERENCE_SGLANG_XAVIER_TRANSFER_TIMEOUT`` 으로 제한됩니다(기본값: 600초)." + +#: source/user_guide/pd_separation.rst:181 +msgid "" +"MLX-prefill cross-engine requests do not accept ``min_p``, " +"``presence_penalty``, ``frequency_penalty`` or ``seed``. The router rejects " +"these parameters before scheduling either replica, so the first output token" +" uses the same supported sampling settings as the rest of the response." +msgstr "" +"MLX prefill을 사용하는 엔진 간 요청은 ``min_p``, ``presence_penalty``, " +"``frequency_penalty`` 또는 ``seed`` 를 허용하지 않습니다. router는 어느 replica도 예약하기 전에 이" +" 매개변수를 거부하므로 첫 출력 토큰에도 나머지 응답과 동일한 지원 샘플링 설정을 사용합니다." diff --git a/doc/source/locale/pt_BR/LC_MESSAGES/user_guide/pd_separation.po b/doc/source/locale/pt_BR/LC_MESSAGES/user_guide/pd_separation.po index 550b08794e..bc77ac041c 100644 --- a/doc/source/locale/pt_BR/LC_MESSAGES/user_guide/pd_separation.po +++ b/doc/source/locale/pt_BR/LC_MESSAGES/user_guide/pd_separation.po @@ -492,15 +492,12 @@ msgid "" "Use the default ``xavier`` transport for MLX; native NIXL is unsupported. " "All workers must reach the supervisor's actor address. Relaunch after a " "worker restart. Two replicas on one Mac share its Metal GPU; measure your " -"workload before expecting a throughput benefit. Cross-engine P/D involving " -"MLX, including NVIDIA-to-Mac handoff, is not yet supported." +"workload before expecting a throughput benefit." msgstr "" "Use o transporte padrão ``xavier`` para MLX; NIXL nativo não é suportado. " "Todos os workers devem alcançar o endereço actor do supervisor. Reinicie o " "deployment após reiniciar um worker. Duas réplicas no mesmo Mac compartilham" -" a GPU Metal; meça sua carga antes de esperar ganho de throughput. P/D entre" -" engines envolvendo MLX, incluindo transferência de NVIDIA para Mac, ainda " -"não é suportado." +" a GPU Metal; meça sua carga antes de esperar ganho de throughput." #: user_guide/pd_separation.rst:102 msgid "" @@ -584,17 +581,125 @@ msgstr "" " kernels podem produzir pequenas diferenças numéricas. vLLM decode exige " "pelo menos dois tokens de prompt." -#: source/user_guide/pd_separation.rst +#: user_guide/pd_separation.rst msgid "" "Cross-engine P/D does not support CPU fallback, retained Xavier history, " "logprobs, structured sampling, LoRA, speculative decoding, multimodal inputs" " or hybrid attention. Omit ``xavier_gpu_cache_bytes`` or set it to ``0``; do" " not set ``xavier_cache_bytes``. Relaunch after a worker restart. Native " -"NIXL comparisons use the same engine on both roles. MLX is unsupported." +"NIXL comparisons use the same engine on both roles." msgstr "" "P/D entre mecanismos não oferece fallback para CPU, histórico Xavier, " "logprobs, amostragem estruturada, LoRA, decodificação especulativa, entradas" " multimodais ou atenção híbrida. Omita ``xavier_gpu_cache_bytes`` ou " "defina-o como ``0``; não defina ``xavier_cache_bytes``. Reimplante após " "reiniciar um worker. Comparações NIXL nativas usam o mesmo mecanismo nas " -"duas funções. MLX não é compatível." +"duas funções." + +#: user_guide/pd_separation.rst +msgid "NVIDIA and MLX cross-engine P/D" +msgstr "P/D entre engines NVIDIA e MLX" + +#: user_guide/pd_separation.rst +msgid "" +"Xavier supports both vLLM/SGLang prefill with MLX decode and MLX prefill " +"with vLLM/SGLang decode, with streaming and non-streaming requests. FP16 " +"64-token pages pass through CPU staging and actor RPC between CUDA and " +"Metal. vLLM and MLX decode compute the last prompt token; SGLang decode " +"imports the full prompt KV and the first output token sampled by P. MLX " +"prefill publishes request-owned pages on its worker with a 512 MiB in-flight" +" budget." +msgstr "" +"Xavier suporta prefill vLLM/SGLang com decode MLX e prefill MLX com decode " +"vLLM/SGLang, com e sem streaming. Páginas FP16 de 64 tokens passam por " +"armazenamento temporário CPU e RPC actor entre CUDA e Metal. O decode vLLM e" +" MLX calcula o último token do prompt; SGLang importa o KV completo do " +"prompt e o primeiro token de saída amostrado por P. O prefill MLX publica " +"páginas próprias de cada requisição no worker, com orçamento de 512 MiB para" +" transferências em andamento." + +#: user_guide/pd_separation.rst +msgid "" +"Use identical unquantized checkpoint and tokenizer files; do not convert the" +" Mac checkpoint. Worker-local ``model_path`` and ``model_format`` belong in " +"each replica's ``engine_config``. Xavier checks weights, tokenizer, " +"effective model semantics, context limits and prompt token IDs. Model and " +"sampling restrictions match vLLM/SGLang cross-engine P/D; failures raise " +"errors." +msgstr "" +"Use arquivos idênticos de checkpoint não quantizado e tokenizer; não " +"converta o checkpoint Mac. ``model_path`` e ``model_format`` locais ao " +"worker pertencem ao ``engine_config`` de cada réplica. Xavier verifica " +"pesos, tokenizer, semântica efetiva do modelo, limites de contexto e IDs dos" +" tokens do prompt. As restrições de modelo e amostragem são as mesmas do P/D" +" vLLM/SGLang entre engines; falhas geram erros." + +#: user_guide/pd_separation.rst +msgid "" +"To use MLX prefill, swap the two ``role`` values in the example. For SGLang " +"on either role, use ``SGLang`` and replace ``max_model_len`` with " +"``context_length``. NVIDIA requires ``xoscar[nixl]>=0.11.1``; Mac requires " +"``xoscar>=0.11.1`` and ``mlx-lm>=0.31.2``. Workers must reach the supervisor" +" and each other. Use Xavier transport, and omit ``xavier_gpu_cache_bytes`` " +"or set it to ``0``. Relaunch after a worker restart." +msgstr "" +"Para prefill MLX, troque os dois valores ``role`` no exemplo. Para SGLang em" +" qualquer função, use ``SGLang`` e troque ``max_model_len`` por " +"``context_length``. NVIDIA requer ``xoscar[nixl]>=0.11.1``; Mac requer " +"``xoscar>=0.11.1`` e ``mlx-lm>=0.31.2``. Os workers devem alcançar o " +"supervisor e uns aos outros. Use Xavier, omita ``xavier_gpu_cache_bytes`` ou" +" defina-o como ``0``. Reimplante após reiniciar um worker." + +#: ../../source/user_guide/pd_separation.rst +msgid "" +"Run the supervisor and all worker/model environments with the same Python " +"version and matching ``cloudpickle`` / ``tblib`` versions; actor RPC also " +"carries exceptions and generators." +msgstr "" +"Execute o supervisor e todos os ambientes worker/model com a mesma versão de" +" Python e versões iguais de ``cloudpickle`` / ``tblib``; os RPC actor também" +" transportam exceções e geradores." + +#: user_guide/pd_separation.rst +msgid "" +"Compare end-to-end TTFT with the standalone decode engine on its own " +"hardware, using the same effective FP16 weights and cache policy. Network " +"transfers and CUDA/Metal copies may outweigh faster prefill. ``host_bytes`` " +"measures CPU-staged payloads; ``imported_tokens`` records completed KV " +"imports. Correct output alone does not prove reuse." +msgstr "" +"Compare o TTFT completo com a engine de decode executada sozinha no próprio " +"hardware usando os mesmos pesos FP16 efetivos e política de cache. " +"Transferências de rede e cópias CUDA/Metal podem superar o ganho de prefill." +" ``host_bytes`` mede os dados que passam pela CPU; ``imported_tokens`` " +"registra importações KV concluídas. A saída correta por si só não comprova " +"reutilização." + +#: source/user_guide/pd_separation.rst:179 +msgid "" +"For MLX prefill, ``xavier_cache_bytes`` sets the worker-local in-flight CPU " +"page budget (default: 512 MiB). Capacity is reserved before Metal prefill, " +"rounding the full prompt up to 64-token pages. A prompt exceeding the budget" +" fails before prefill; concurrent requests wait for capacity. Reservations " +"and transfer waits are bounded by " +"``XINFERENCE_SGLANG_XAVIER_TRANSFER_TIMEOUT`` (default: 600 seconds)." +msgstr "" +"Para prefill MLX, ``xavier_cache_bytes`` define o orçamento de páginas CPU " +"em andamento no worker (padrão: 512 MiB). A capacidade é reservada antes do " +"prefill Metal, arredondando o prompt completo para páginas de 64 tokens. Um " +"prompt acima do orçamento falha antes do prefill; solicitações concorrentes " +"aguardam capacidade. Reservas e esperas de transferência são limitadas por " +"``XINFERENCE_SGLANG_XAVIER_TRANSFER_TIMEOUT`` (padrão: 600 segundos)." + +#: source/user_guide/pd_separation.rst:181 +msgid "" +"MLX-prefill cross-engine requests do not accept ``min_p``, " +"``presence_penalty``, ``frequency_penalty`` or ``seed``. The router rejects " +"these parameters before scheduling either replica, so the first output token" +" uses the same supported sampling settings as the rest of the response." +msgstr "" +"Solicitações entre engines com prefill MLX não aceitam ``min_p``, " +"``presence_penalty``, ``frequency_penalty`` ou ``seed``. O roteador rejeita " +"esses parâmetros antes de agendar qualquer réplica, para que o primeiro " +"token use as mesmas configurações de amostragem compatíveis que o restante " +"da resposta." diff --git a/doc/source/locale/zh_CN/LC_MESSAGES/user_guide/pd_separation.po b/doc/source/locale/zh_CN/LC_MESSAGES/user_guide/pd_separation.po index 1785aa9d79..d712b1f15d 100644 --- a/doc/source/locale/zh_CN/LC_MESSAGES/user_guide/pd_separation.po +++ b/doc/source/locale/zh_CN/LC_MESSAGES/user_guide/pd_separation.po @@ -438,12 +438,10 @@ msgid "" "Use the default ``xavier`` transport for MLX; native NIXL is unsupported. " "All workers must reach the supervisor's actor address. Relaunch after a " "worker restart. Two replicas on one Mac share its Metal GPU; measure your " -"workload before expecting a throughput benefit. Cross-engine P/D involving " -"MLX, including NVIDIA-to-Mac handoff, is not yet supported." +"workload before expecting a throughput benefit." msgstr "" "MLX 使用默认的 ``xavier`` 传输，不支持原生 NIXL。所有 worker 必须能访问 supervisor 的 actor " -"地址。Worker 重启后需要重新启动部署。同一台 Mac 上的两个副本共享 Metal GPU，应先测量工作负载，再评估吞吐收益。目前不支持涉及 " -"MLX 的跨引擎 P/D，包括 NVIDIA 到 Mac 的交接。" +"地址。Worker 重启后需要重新启动部署。同一台 Mac 上的两个副本共享 Metal GPU，应先测量工作负载，再评估吞吐收益。" #: user_guide/pd_separation.rst:102 msgid "" @@ -516,14 +514,106 @@ msgstr "" "导入最后一个 prompt token 之前的 KV，计算该 token 后采样输出。引擎算子可能产生小幅数值差异。vLLM decode 要求至少两个" " prompt token。" -#: source/user_guide/pd_separation.rst +#: user_guide/pd_separation.rst msgid "" "Cross-engine P/D does not support CPU fallback, retained Xavier history, " "logprobs, structured sampling, LoRA, speculative decoding, multimodal inputs" " or hybrid attention. Omit ``xavier_gpu_cache_bytes`` or set it to ``0``; do" " not set ``xavier_cache_bytes``. Relaunch after a worker restart. Native " -"NIXL comparisons use the same engine on both roles. MLX is unsupported." +"NIXL comparisons use the same engine on both roles." msgstr "" "跨引擎 P/D 不支持 CPU 回退、Xavier 历史缓存、logprobs、结构化采样、LoRA、推测解码、多模态输入或混合注意力。省略 " "``xavier_gpu_cache_bytes`` 或设为 ``0``，不要设置 ``xavier_cache_bytes``。worker " -"重启后需要重新部署。原生 NIXL 对照在两个角色上使用相同引擎。不支持 MLX。" +"重启后需要重新部署。原生 NIXL 对照在两个角色上使用相同引擎。" + +#: user_guide/pd_separation.rst +msgid "NVIDIA and MLX cross-engine P/D" +msgstr "NVIDIA 与 MLX 跨引擎 P/D" + +#: user_guide/pd_separation.rst +msgid "" +"Xavier supports both vLLM/SGLang prefill with MLX decode and MLX prefill " +"with vLLM/SGLang decode, with streaming and non-streaming requests. FP16 " +"64-token pages pass through CPU staging and actor RPC between CUDA and " +"Metal. vLLM and MLX decode compute the last prompt token; SGLang decode " +"imports the full prompt KV and the first output token sampled by P. MLX " +"prefill publishes request-owned pages on its worker with a 512 MiB in-flight" +" budget." +msgstr "" +"Xavier 同时支持 vLLM/SGLang 预填充与 MLX 解码，以及 MLX 预填充与 vLLM/SGLang " +"解码，支持流式和非流式请求。FP16 的 64-token 页面经过 CPU 暂存和 actor RPC 在 CUDA 与 Metal " +"之间传输。vLLM 和 MLX 解码会计算最后一个 prompt token；SGLang 解码导入完整 prompt KV 和 P 采样的首个输出 " +"token。MLX 预填充在其 worker 上发布请求专属页面，在途数据预算为 512 MiB。" + +#: user_guide/pd_separation.rst +msgid "" +"Use identical unquantized checkpoint and tokenizer files; do not convert the" +" Mac checkpoint. Worker-local ``model_path`` and ``model_format`` belong in " +"each replica's ``engine_config``. Xavier checks weights, tokenizer, " +"effective model semantics, context limits and prompt token IDs. Model and " +"sampling restrictions match vLLM/SGLang cross-engine P/D; failures raise " +"errors." +msgstr "" +"使用相同的未量化检查点和分词器文件，不要转换 Mac 检查点。每个副本的 ``engine_config`` 可设置 worker 本地的 " +"``model_path`` 和 ``model_format``。Xavier 检查权重、分词器、实际模型语义、上下文限制和 prompt token" +" ID。模型及采样限制与 vLLM/SGLang 跨引擎 P/D 相同；失败时会抛出错误。" + +#: user_guide/pd_separation.rst +msgid "" +"To use MLX prefill, swap the two ``role`` values in the example. For SGLang " +"on either role, use ``SGLang`` and replace ``max_model_len`` with " +"``context_length``. NVIDIA requires ``xoscar[nixl]>=0.11.1``; Mac requires " +"``xoscar>=0.11.1`` and ``mlx-lm>=0.31.2``. Workers must reach the supervisor" +" and each other. Use Xavier transport, and omit ``xavier_gpu_cache_bytes`` " +"or set it to ``0``. Relaunch after a worker restart." +msgstr "" +"要使用 MLX 预填充，交换示例中的两个 ``role`` 值。在任一角色使用 SGLang 时，设置 ``SGLang`` 并将 " +"``max_model_len`` 换成 ``context_length``。NVIDIA 需要 " +"``xoscar[nixl]>=0.11.1``；Mac 需要 ``xoscar>=0.11.1`` 和 ``mlx-" +"lm>=0.31.2``。Worker 必须能访问 supervisor 和彼此。使用 Xavier 传输，省略 " +"``xavier_gpu_cache_bytes`` 或设为 ``0``。Worker 重启后需要重新部署。" + +#: ../../source/user_guide/pd_separation.rst +msgid "" +"Run the supervisor and all worker/model environments with the same Python " +"version and matching ``cloudpickle`` / ``tblib`` versions; actor RPC also " +"carries exceptions and generators." +msgstr "" +"Supervisor 与所有 worker/model 环境必须使用相同的 Python 版本，并对齐 ``cloudpickle`` / " +"``tblib`` 版本；actor RPC 也会传输异常和生成器。" + +#: user_guide/pd_separation.rst +msgid "" +"Compare end-to-end TTFT with the standalone decode engine on its own " +"hardware, using the same effective FP16 weights and cache policy. Network " +"transfers and CUDA/Metal copies may outweigh faster prefill. ``host_bytes`` " +"measures CPU-staged payloads; ``imported_tokens`` records completed KV " +"imports. Correct output alone does not prove reuse." +msgstr "" +"使用相同的实际 FP16 权重及缓存策略，与独立运行在原硬件上的解码引擎对比端到端 TTFT。网络传输及 CUDA/Metal " +"复制成本可能超过预填充加速的收益。``host_bytes`` 记录 CPU 暂存载荷；``imported_tokens`` 记录完成的 KV " +"导入。仅输出正确不能证明复用了缓存。" + +#: source/user_guide/pd_separation.rst:179 +msgid "" +"For MLX prefill, ``xavier_cache_bytes`` sets the worker-local in-flight CPU " +"page budget (default: 512 MiB). Capacity is reserved before Metal prefill, " +"rounding the full prompt up to 64-token pages. A prompt exceeding the budget" +" fails before prefill; concurrent requests wait for capacity. Reservations " +"and transfer waits are bounded by " +"``XINFERENCE_SGLANG_XAVIER_TRANSFER_TIMEOUT`` (default: 600 seconds)." +msgstr "" +"MLX 做预填充时，``xavier_cache_bytes`` 设置 Worker 本地正在使用的 CPU 页面预算（默认 512 MiB）。在 " +"Metal 预填充前预留容量，将完整提示词向上取整为 64 token " +"页面。单个提示词超过预算时，在预填充前报错；并发请求等待空闲容量。预留和传输等待受 " +"``XINFERENCE_SGLANG_XAVIER_TRANSFER_TIMEOUT`` 限制（默认 600 秒）。" + +#: source/user_guide/pd_separation.rst:181 +msgid "" +"MLX-prefill cross-engine requests do not accept ``min_p``, " +"``presence_penalty``, ``frequency_penalty`` or ``seed``. The router rejects " +"these parameters before scheduling either replica, so the first output token" +" uses the same supported sampling settings as the rest of the response." +msgstr "" +"MLX 做预填充的跨引擎请求不接受 ``min_p``、``presence_penalty``、``frequency_penalty`` 或 " +"``seed``。路由器在调度任何副本前拒绝这些参数，确保首个输出 token 与后续响应使用一致的受支持采样设置。" diff --git a/doc/source/locale/zh_TW/LC_MESSAGES/user_guide/pd_separation.po b/doc/source/locale/zh_TW/LC_MESSAGES/user_guide/pd_separation.po index 7590f205c9..6f3c16f8a2 100644 --- a/doc/source/locale/zh_TW/LC_MESSAGES/user_guide/pd_separation.po +++ b/doc/source/locale/zh_TW/LC_MESSAGES/user_guide/pd_separation.po @@ -436,12 +436,10 @@ msgid "" "Use the default ``xavier`` transport for MLX; native NIXL is unsupported. " "All workers must reach the supervisor's actor address. Relaunch after a " "worker restart. Two replicas on one Mac share its Metal GPU; measure your " -"workload before expecting a throughput benefit. Cross-engine P/D involving " -"MLX, including NVIDIA-to-Mac handoff, is not yet supported." +"workload before expecting a throughput benefit." msgstr "" "MLX 使用預設的 ``xavier`` 傳輸，不支援原生 NIXL。所有 worker 必須能存取 supervisor 的 actor " -"位址。Worker 重新啟動後需要重新啟動部署。同一台 Mac 上的兩個副本共用 Metal GPU，應先測量工作負載，再評估吞吐收益。目前不支援涉及 " -"MLX 的跨引擎 P/D，包括 NVIDIA 到 Mac 的交接。" +"位址。Worker 重新啟動後需要重新啟動部署。同一台 Mac 上的兩個副本共用 Metal GPU，應先測量工作負載，再評估吞吐收益。" #: user_guide/pd_separation.rst:102 msgid "" @@ -514,14 +512,106 @@ msgstr "" "匯入最後一個 prompt token 之前的 KV，計算該 token 後取樣輸出。引擎核心可能產生小幅數值差異。vLLM decode 要求至少兩個" " prompt token。" -#: source/user_guide/pd_separation.rst +#: user_guide/pd_separation.rst msgid "" "Cross-engine P/D does not support CPU fallback, retained Xavier history, " "logprobs, structured sampling, LoRA, speculative decoding, multimodal inputs" " or hybrid attention. Omit ``xavier_gpu_cache_bytes`` or set it to ``0``; do" " not set ``xavier_cache_bytes``. Relaunch after a worker restart. Native " -"NIXL comparisons use the same engine on both roles. MLX is unsupported." +"NIXL comparisons use the same engine on both roles." msgstr "" "跨引擎 P/D 不支援 CPU 回退、Xavier 歷史快取、logprobs、結構化取樣、LoRA、推測解碼、多模態輸入或混合注意力。省略 " "``xavier_gpu_cache_bytes`` 或設為 ``0``，不要設定 ``xavier_cache_bytes``。worker " -"重新啟動後需要重新部署。原生 NIXL 對照在兩個角色上使用相同引擎。不支援 MLX。" +"重新啟動後需要重新部署。原生 NIXL 對照在兩個角色上使用相同引擎。" + +#: user_guide/pd_separation.rst +msgid "NVIDIA and MLX cross-engine P/D" +msgstr "NVIDIA 與 MLX 跨引擎 P/D" + +#: user_guide/pd_separation.rst +msgid "" +"Xavier supports both vLLM/SGLang prefill with MLX decode and MLX prefill " +"with vLLM/SGLang decode, with streaming and non-streaming requests. FP16 " +"64-token pages pass through CPU staging and actor RPC between CUDA and " +"Metal. vLLM and MLX decode compute the last prompt token; SGLang decode " +"imports the full prompt KV and the first output token sampled by P. MLX " +"prefill publishes request-owned pages on its worker with a 512 MiB in-flight" +" budget." +msgstr "" +"Xavier 同時支援 vLLM/SGLang 預填充與 MLX 解碼，以及 MLX 預填充與 vLLM/SGLang " +"解碼，支援串流及非串流請求。FP16 的 64-token 頁面經過 CPU 暫存和 actor RPC 在 CUDA 與 Metal " +"之間傳輸。vLLM 和 MLX 解碼會計算最後一個 prompt token；SGLang 解碼匯入完整 prompt KV 和 P 取樣的首個輸出 " +"token。MLX 預填充在其 worker 上發布請求專屬頁面，在途資料預算為 512 MiB。" + +#: user_guide/pd_separation.rst +msgid "" +"Use identical unquantized checkpoint and tokenizer files; do not convert the" +" Mac checkpoint. Worker-local ``model_path`` and ``model_format`` belong in " +"each replica's ``engine_config``. Xavier checks weights, tokenizer, " +"effective model semantics, context limits and prompt token IDs. Model and " +"sampling restrictions match vLLM/SGLang cross-engine P/D; failures raise " +"errors." +msgstr "" +"使用相同的未量化檢查點與分詞器檔案，不要轉換 Mac 檢查點。每個副本的 ``engine_config`` 可設定 worker 本機的 " +"``model_path`` 和 ``model_format``。Xavier 檢查權重、分詞器、實際模型語義、上下文限制和 prompt token" +" ID。模型及取樣限制與 vLLM/SGLang 跨引擎 P/D 相同；失敗時會拋出錯誤。" + +#: user_guide/pd_separation.rst +msgid "" +"To use MLX prefill, swap the two ``role`` values in the example. For SGLang " +"on either role, use ``SGLang`` and replace ``max_model_len`` with " +"``context_length``. NVIDIA requires ``xoscar[nixl]>=0.11.1``; Mac requires " +"``xoscar>=0.11.1`` and ``mlx-lm>=0.31.2``. Workers must reach the supervisor" +" and each other. Use Xavier transport, and omit ``xavier_gpu_cache_bytes`` " +"or set it to ``0``. Relaunch after a worker restart." +msgstr "" +"要使用 MLX 預填充，交換範例中的兩個 ``role`` 值。在任一角色使用 SGLang 時，設定 ``SGLang`` 並將 " +"``max_model_len`` 換成 ``context_length``。NVIDIA 需要 " +"``xoscar[nixl]>=0.11.1``；Mac 需要 ``xoscar>=0.11.1`` 和 ``mlx-" +"lm>=0.31.2``。Worker 必須能存取 supervisor 和彼此。使用 Xavier 傳輸，省略 " +"``xavier_gpu_cache_bytes`` 或設為 ``0``。Worker 重新啟動後需要重新部署。" + +#: ../../source/user_guide/pd_separation.rst +msgid "" +"Run the supervisor and all worker/model environments with the same Python " +"version and matching ``cloudpickle`` / ``tblib`` versions; actor RPC also " +"carries exceptions and generators." +msgstr "" +"Supervisor 與所有 worker/model 環境必須使用相同的 Python 版本，並對齊 ``cloudpickle`` / " +"``tblib`` 版本；actor RPC 也會傳輸例外和生成器。" + +#: user_guide/pd_separation.rst +msgid "" +"Compare end-to-end TTFT with the standalone decode engine on its own " +"hardware, using the same effective FP16 weights and cache policy. Network " +"transfers and CUDA/Metal copies may outweigh faster prefill. ``host_bytes`` " +"measures CPU-staged payloads; ``imported_tokens`` records completed KV " +"imports. Correct output alone does not prove reuse." +msgstr "" +"使用相同的實際 FP16 權重與快取策略，與獨立執行於原硬體的解碼引擎比較端到端 TTFT。網路傳輸及 CUDA/Metal " +"複製成本可能超過預填充加速的收益。``host_bytes`` 記錄 CPU 暫存載荷；``imported_tokens`` 記錄完成的 KV " +"匯入。僅輸出正確不能證明重用了快取。" + +#: source/user_guide/pd_separation.rst:179 +msgid "" +"For MLX prefill, ``xavier_cache_bytes`` sets the worker-local in-flight CPU " +"page budget (default: 512 MiB). Capacity is reserved before Metal prefill, " +"rounding the full prompt up to 64-token pages. A prompt exceeding the budget" +" fails before prefill; concurrent requests wait for capacity. Reservations " +"and transfer waits are bounded by " +"``XINFERENCE_SGLANG_XAVIER_TRANSFER_TIMEOUT`` (default: 600 seconds)." +msgstr "" +"MLX 做預填充時，``xavier_cache_bytes`` 設定 Worker 本地正在使用的 CPU 頁面預算（預設 512 MiB）。在 " +"Metal 預填充前預留容量，將完整提示詞向上取整為 64 token " +"頁面。單一提示詞超過預算時，在預填充前報錯；並行請求等待可用容量。預留和傳輸等待受 " +"``XINFERENCE_SGLANG_XAVIER_TRANSFER_TIMEOUT`` 限制（預設 600 秒）。" + +#: source/user_guide/pd_separation.rst:181 +msgid "" +"MLX-prefill cross-engine requests do not accept ``min_p``, " +"``presence_penalty``, ``frequency_penalty`` or ``seed``. The router rejects " +"these parameters before scheduling either replica, so the first output token" +" uses the same supported sampling settings as the rest of the response." +msgstr "" +"MLX 做預填充的跨引擎請求不接受 ``min_p``、``presence_penalty``、``frequency_penalty`` 或 " +"``seed``。路由器在排程任何副本前拒絕這些參數，確保第一個輸出 token 與後續回應使用一致的受支援取樣設定。" diff --git a/doc/source/user_guide/pd_separation.rst b/doc/source/user_guide/pd_separation.rst index ad95577a81..dfa3172835 100644 --- a/doc/source/user_guide/pd_separation.rst +++ b/doc/source/user_guide/pd_separation.rst @@ -101,7 +101,7 @@ MLX Xavier stores immutable CPU pages in a shared actor on the supervisor and tr Each 64-token FP16 page uses ``2 * 64 * num_layers * num_kv_heads * head_dim * 2`` bytes (K and V, with two bytes per value). For example, 36 layers, 8 KV heads and a head dimension of 128 require 9 MiB per page; 512 MiB holds 56 pages, or 3584 prefix tokens. Round each prefix up to a whole page and size ``xavier_cache_bytes`` for the distinct pages reserved by all concurrent P/D requests. Ordinary shared-cache requests skip publication when their prefix exceeds the cache capacity or the 4096-page protocol limit. -Use the default ``xavier`` transport for MLX; native NIXL is unsupported. All workers must reach the supervisor's actor address. Relaunch after a worker restart. Two replicas on one Mac share its Metal GPU; measure your workload before expecting a throughput benefit. Cross-engine P/D involving MLX, including NVIDIA-to-Mac handoff, is not yet supported. +Use the default ``xavier`` transport for MLX; native NIXL is unsupported. All workers must reach the supervisor's actor address. Relaunch after a worker restart. Two replicas on one Mac share its Metal GPU; measure your workload before expecting a throughput benefit. Cross-engine vLLM/SGLang P/D -------------------------------------------------------------------------------- @@ -138,7 +138,51 @@ Both workers require identical local weights, tokenizer assets and context limit Streaming and non-streaming requests require ``n=1``. SGLang decode uses the first output token sampled by prefill. vLLM decode imports all but the final prompt token, computes that token and samples the output. Engine kernels can produce small numerical differences. vLLM decode requires at least two prompt tokens. -Cross-engine P/D does not support CPU fallback, retained Xavier history, logprobs, structured sampling, LoRA, speculative decoding, multimodal inputs or hybrid attention. Omit ``xavier_gpu_cache_bytes`` or set it to ``0``; do not set ``xavier_cache_bytes``. Relaunch after a worker restart. Native NIXL comparisons use the same engine on both roles. MLX is unsupported. +Cross-engine P/D does not support CPU fallback, retained Xavier history, logprobs, structured sampling, LoRA, speculative decoding, multimodal inputs or hybrid attention. Omit ``xavier_gpu_cache_bytes`` or set it to ``0``; do not set ``xavier_cache_bytes``. Relaunch after a worker restart. Native NIXL comparisons use the same engine on both roles. + +NVIDIA and MLX cross-engine P/D +-------------------------------------------------------------------------------- + +Xavier supports both vLLM/SGLang prefill with MLX decode and MLX prefill with vLLM/SGLang decode, with streaming and non-streaming requests. FP16 64-token pages pass through CPU staging and actor RPC between CUDA and Metal. vLLM and MLX decode compute the last prompt token; SGLang decode imports the full prompt KV and the first output token sampled by P. MLX prefill publishes request-owned pages on its worker with a 512 MiB in-flight budget. + +Use identical unquantized checkpoint and tokenizer files; do not convert the Mac checkpoint. Worker-local ``model_path`` and ``model_format`` belong in each replica's ``engine_config``. Xavier checks weights, tokenizer, effective model semantics, context limits and prompt token IDs. Model and sampling restrictions match vLLM/SGLang cross-engine P/D; failures raise errors. + +.. code-block:: python + + client.launch_model( + model_uid="qwen-nvidia-mac-pd", + model_name="qwen2.5-instruct", + model_size_in_billions="0_5", + model_engine="vLLM", + model_format="pytorch", + quantization="none", + replica=2, + replica_config=[ + { + "role": "prefill", "model_engine": "vLLM", + "engine_config": {"model_path": "/models/Qwen2.5-0.5B-Instruct", + "max_model_len": 8192}, + "devices": [{"worker_ip": "NVIDIA_WORKER:PORT", + "n_gpu": 1, "gpu_idx": [0]}], + }, + { + "role": "decode", "model_engine": "MLX", + "engine_config": {"model_path": "/Users/me/models/Qwen2.5-0.5B-Instruct", + "model_format": "mlx", "context_length": 8192}, + "devices": [{"worker_ip": "MAC_WORKER:PORT", "n_gpu": "auto"}], + }, + ], + ) + +To use MLX prefill, swap the two ``role`` values in the example. For SGLang on either role, use ``SGLang`` and replace ``max_model_len`` with ``context_length``. NVIDIA requires ``xoscar[nixl]>=0.11.1``; Mac requires ``xoscar>=0.11.1`` and ``mlx-lm>=0.31.2``. Workers must reach the supervisor and each other. Use Xavier transport, and omit ``xavier_gpu_cache_bytes`` or set it to ``0``. Relaunch after a worker restart. + +For MLX prefill, ``xavier_cache_bytes`` sets the worker-local in-flight CPU page budget (default: 512 MiB). Capacity is reserved before Metal prefill, rounding the full prompt up to 64-token pages. A prompt exceeding the budget fails before prefill; concurrent requests wait for capacity. Reservations and transfer waits are bounded by ``XINFERENCE_SGLANG_XAVIER_TRANSFER_TIMEOUT`` (default: 600 seconds). + +MLX-prefill cross-engine requests do not accept ``min_p``, ``presence_penalty``, ``frequency_penalty`` or ``seed``. The router rejects these parameters before scheduling either replica, so the first output token uses the same supported sampling settings as the rest of the response. + +Run the supervisor and all worker/model environments with the same Python version and matching ``cloudpickle`` / ``tblib`` versions; actor RPC also carries exceptions and generators. + +Compare end-to-end TTFT with the standalone decode engine on its own hardware, using the same effective FP16 weights and cache policy. Network transfers and CUDA/Metal copies may outweigh faster prefill. ``host_bytes`` measures CPU-staged payloads; ``imported_tokens`` records completed KV imports. Correct output alone does not prove reuse. Native NIXL backend ------------------- diff --git a/xinference/core/pd_model.py b/xinference/core/pd_model.py index fb0f05630d..e79ac4da79 100644 --- a/xinference/core/pd_model.py +++ b/xinference/core/pd_model.py @@ -23,6 +23,7 @@ import xoscar as xo +from ..model.llm.xavier.constants import CROSS_ENGINE_TRANSFER_ACTOR_UID from .rpc_context import actor_call from .utils import log_async @@ -85,6 +86,8 @@ def __init__( ] = RoundRobinSchedulingPolicy, transport_backend: str = "xavier", model_engine: str = "vllm", + handoff_mode: str = "gpu", + mlx_prefill: bool = False, ): super().__init__() # Prefill request map, used to skip the timeout task for specific request id. @@ -95,6 +98,8 @@ def __init__( self._transport_backend = transport_backend self._model_engine = (model_engine or "vllm").lower() self._direct_handoff = transport_backend == "xavier" + self._handoff_mode = handoff_mode + self._mlx_prefill = mlx_prefill # 使用字典存储副本：{replica_uid: actor_ref} self._prefill_replicas: Dict[str, xo.ActorRefType["ModelActor"]] = {} @@ -261,12 +266,8 @@ async def free_prefill_model_cache(self, request_id: str): ref = await xo.actor_ref( address=handoff["address"], uid=handoff["uid"] ) - if handoff.get("mode") == "gpu": + if handoff.get("mode") in ("gpu", "host"): if handoff.get("heterogeneous"): - from ..model.llm.xavier.backends.torch.pd import ( - CrossEngineGPUActor, - ) - try: source = await ref.source(handoff["room"]) except RuntimeError: @@ -274,7 +275,8 @@ async def free_prefill_model_cache(self, request_id: str): if source: sender = await xo.actor_ref( address=source["address"], - uid=f"{CrossEngineGPUActor.default_uid()}-{source['rank']}", + uid=source.get("uid") + or f"{CROSS_ENGINE_TRANSFER_ACTOR_UID}-{source['rank']}", ) await sender.abort(handoff["room"]) await ref.release(handoff["room"]) @@ -330,6 +332,14 @@ async def _infer(self, method, inputs, *args, **kwargs): raise ValueError("PD KV handoff currently requires n=1") if self._model_engine == "heterogeneous": config = (args[0] if args else {}) or {} + if self._mlx_prefill and any( + parameters.get(key) is not None + for parameters in (config, kwargs.get("raw_params") or {}) + for key in ("min_p", "presence_penalty", "frequency_penalty", "seed") + ): + raise ValueError( + "MLX cross-engine prefill does not support min_p, presence_penalty, frequency_penalty or seed" + ) if ( any( config.get(key) is not None and config.get(key) is not False @@ -479,7 +489,7 @@ async def _infer_sglang(self, method, inputs, args, kwargs, request_id): ) handoff = dict( engine="sglang", - mode="nixl" if native else "gpu", + mode="nixl" if native else self._handoff_mode, room=uuid.uuid4().int % (2**63 - 1) + 1, **bootstrap, ) diff --git a/xinference/core/replica_config.py b/xinference/core/replica_config.py index 3e2596ebfe..d384846188 100644 --- a/xinference/core/replica_config.py +++ b/xinference/core/replica_config.py @@ -182,6 +182,9 @@ def validate_pd_replica_configs( model_type: Optional[str], ) -> bool: """Validate a complete P/D topology before allocating any runtime resources.""" + engines = { + (cfg.model_engine or model_engine or "").lower() for cfg in configs or [] + } reserved = { "model_uid", "model_name", @@ -214,6 +217,9 @@ def validate_pd_replica_configs( "xavier_cache_bytes", "xavier_gpu_cache_bytes", } + if "mlx" in engines and len(engines) > 1: + # CUDA and Metal replicas need their own local checkpoint paths/formats. + reserved -= {"model_path", "model_format"} for cfg in configs or []: invalid = { key for key in cfg.engine_config if key in reserved or key.startswith("_") @@ -227,9 +233,6 @@ def validate_pd_replica_configs( if any(cfg.model_engine or cfg.engine_config for cfg in configs or []): raise ValueError("Per-replica engine settings require explicit PD roles") return False - engines = { - (cfg.model_engine or model_engine or "").lower() for cfg in configs or [] - } if ( not engines <= {"vllm", "sglang", "mlx"} or (model_type or "LLM").lower() != "llm" @@ -237,10 +240,24 @@ def validate_pd_replica_configs( raise ValueError( "PD separation requires model_type=LLM and model_engine=vLLM, SGLang or MLX" ) - if "mlx" in engines and len(engines) > 1: - raise ValueError("Cross-engine PD requires vLLM and SGLang replicas") if roles != {"prefill", "decode"}: raise ValueError( "PD separation requires both prefill and decode replicas, without hybrid replicas" ) + if "mlx" in engines and len(engines) > 1: + role_engines = { + role: { + (cfg.model_engine or model_engine or "").lower() + for cfg in configs or [] + if cfg.role == role + } + for role in ("prefill", "decode") + } + if not any( + role_engines[role] == {"mlx"} and role_engines[other] <= {"vllm", "sglang"} + for role, other in (("prefill", "decode"), ("decode", "prefill")) + ): + raise ValueError( + "Cross-engine MLX PD requires MLX on one role and NVIDIA on the other" + ) return True diff --git a/xinference/core/supervisor.py b/xinference/core/supervisor.py index 10b9de6362..631ce585bf 100644 --- a/xinference/core/supervisor.py +++ b/xinference/core/supervisor.py @@ -236,6 +236,7 @@ class SupervisorActor(xo.StatelessActor): def __init__(self): self._pd_model_mapping: Dict[str, Any] = {} self._xavier_cache_mapping: Dict[str, xo.ActorRefType] = {} + self._xavier_source_mapping: Dict[str, List[xo.ActorRefType]] = {} self._pd_roles: Dict[str, Dict[int, str]] = {} super().__init__() self._worker_address_to_worker: Dict[str, xo.ActorRefType["WorkerActor"]] = {} # type: ignore @@ -3115,6 +3116,7 @@ async def launch_builtin_model( for cfg in replica_config or [] ] heterogeneous_pd = pd_enabled and len(set(replica_engines)) > 1 + host_handoff = heterogeneous_pd and "mlx" in replica_engines if ( pd_enabled and not heterogeneous_pd @@ -3124,7 +3126,11 @@ async def launch_builtin_model( cfg.model_engine for cfg in replica_config or [] if cfg.model_engine ) if heterogeneous_pd and transport_backend != "xavier": - raise ValueError("Cross-engine PD requires Xavier GPU transport") + raise ValueError( + "Cross-engine PD requires Xavier transport" + if host_handoff + else "Cross-engine PD requires Xavier GPU transport" + ) # Xavier-related requested_xavier = bool(kwargs.pop("enable_xavier", False)) if ( @@ -3176,8 +3182,13 @@ async def launch_builtin_model( "SGLang native NIXL requires one worker per replica without enable_xavier" ) cache_bytes = kwargs.pop("xavier_cache_bytes", None) + mlx_prefill = host_handoff and any( + cfg.role == "prefill" + and (cfg.model_engine or model_engine or "").lower() == "mlx" + for cfg in (replica_config or []) + ) if cache_bytes is not None and ( - not (mlx_xavier or sglang_xavier and not pd_enabled) + not (mlx_xavier or mlx_prefill or sglang_xavier and not pd_enabled) or type(cache_bytes) is not int or cache_bytes <= 0 ): @@ -3191,7 +3202,7 @@ async def launch_builtin_model( if sglang_xavier or heterogeneous_pd: if ( model_type not in (None, "LLM") - or model_format not in (None, "pytorch") + or (not host_handoff and model_format not in (None, "pytorch")) or quantization not in (None, "none") ): raise ValueError( @@ -3199,6 +3210,14 @@ async def launch_builtin_model( if heterogeneous_pd else "SGLang Xavier requires unquantized PyTorch LLM weights" ) + if host_handoff: + for cfg, engine in zip(replica_config or [], replica_engines): + replica_format = cfg.engine_config.get("model_format", model_format) + allowed_formats = ("mlx",) if engine == "mlx" else (None, "pytorch") + if replica_format not in allowed_formats: + raise ValueError( + "NVIDIA/MLX PD requires PyTorch NVIDIA and MLX Metal weights" + ) if sglang_xavier: kwargs["_xavier_cache_config"] = { "address": self.address, @@ -3330,11 +3349,15 @@ async def _launch_one_model( replica_kwargs = dict(kwargs) replica_engine = model_engine + replica_path = model_path + replica_format = model_format if pd_enabled: assert replica_config is not None cfg = replica_config[rank - 1] replica_engine = cfg.model_engine or model_engine replica_kwargs.update(cfg.engine_config) + replica_path = replica_kwargs.pop("model_path", model_path) + replica_format = replica_kwargs.pop("model_format", model_format) if heterogeneous_pd: assert replica_config is not None from ..model.llm.xavier.transport import get_transport_host @@ -3347,6 +3370,7 @@ async def _launch_one_model( "host": get_transport_host(worker_ref.address), "heterogeneous": True, "gpu_cache_bytes": 0, + "host_handoff": host_handoff, } if sglang_xavier and pd_enabled: assert replica_config is not None @@ -3376,11 +3400,31 @@ async def _launch_one_model( worker_ref.address ) try: + if ( + heterogeneous_pd + and (replica_engine or "").lower() == "mlx" + and cfg.role == "prefill" + ): + from ..model.llm.xavier.backends.bytes.pd import XavierHostPDSource + + source_uid = f"xavier-host-pd-{model_uid}-{rank}" + source = await xo.create_actor( + XavierHostPDSource, + self._xavier_cache_mapping[model_uid], + rank, + capacity_bytes=cache_bytes or 512 * 1024**2, + address=worker_ref.address, + uid=source_uid, + ) + self._xavier_source_mapping.setdefault(model_uid, []).append(source) + replica_kwargs["_xavier_cache_config"].update( + source_address=source.address, source_uid=source_uid + ) subpool_address = await worker_ref.launch_builtin_model( model_uid=_replica_model_uid, model_name=model_name, model_size_in_billions=model_size_in_billions, - model_format=model_format, + model_format=replica_format, quantization=quantization, model_engine=replica_engine, model_type=model_type, @@ -3389,7 +3433,7 @@ async def _launch_one_model( peft_model_config=peft_model_config, gpu_idx=replica_gpu_idx, download_hub=download_hub, - model_path=model_path, + model_path=replica_path, enable_virtual_env=enable_virtual_env, virtual_env_packages=virtual_env_packages, virtual_env_find_links=virtual_env_find_links, @@ -3732,6 +3776,8 @@ async def _launch_model(): model_engine=( "heterogeneous" if heterogeneous_pd else model_engine ), + handoff_mode="host" if host_handoff else "gpu", + mlx_prefill=mlx_prefill, address=self.address, uid=f"{model_uid}-{PDModelActor.default_uid()}", ) @@ -5471,6 +5517,13 @@ async def _cleanup_distributed_actors( logger.debug( "Failed to destroy PD router for %s", model_uid, exc_info=True ) + for source in self._xavier_source_mapping.pop(model_uid, []): + try: + await xo.wait_for(xo.destroy_actor(source), timeout=5) + except Exception: + logger.warning( + "Destroy Xavier host source failed for %s", model_uid, exc_info=True + ) if cache_ref is not None: try: await xo.destroy_actor(cache_ref) diff --git a/xinference/core/tests/test_pd_launch.py b/xinference/core/tests/test_pd_launch.py index b077129886..bd3090914c 100644 --- a/xinference/core/tests/test_pd_launch.py +++ b/xinference/core/tests/test_pd_launch.py @@ -775,14 +775,14 @@ async def test_cross_engine_pd_launches_one_adapter_per_replica( @pytest.mark.asyncio @pytest.mark.parametrize("engines", [("MLX", "vLLM"), ("SGLang", "MLX")]) -async def test_cross_engine_mlx_is_rejected_before_allocating_actors( +async def test_cross_engine_mlx_requires_local_format_before_allocating_actors( launch_runtime, engines ): supervisor, workers, actors, _ = launch_runtime kwargs = launch_kwargs() for cfg, engine in zip(kwargs["replica_config"], engines): cfg.model_engine = engine - with pytest.raises(ValueError, match="vLLM and SGLang replicas"): + with pytest.raises(ValueError, match="PyTorch NVIDIA and MLX Metal weights"): await supervisor.launch_builtin_model(**kwargs) assert not actors for worker in workers: @@ -874,3 +874,167 @@ async def test_cross_engine_invalid_launch_has_no_side_effects( for worker in workers: worker.launch_builtin_model.assert_not_awaited() worker.wait_for_load.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("engine", ["vLLM", "SGLang"]) +@pytest.mark.parametrize("default_engine", ["vLLM", "SGLang", "MLX"]) +async def test_nvidia_mlx_launch_uses_worker_local_assets( + launch_runtime, engine, default_engine +): + supervisor, workers, actors, _ = launch_runtime + kwargs = launch_kwargs() + kwargs["model_engine"] = default_engine + kwargs["model_format"] = "mlx" if default_engine == "MLX" else "pytorch" + kwargs["replica_config"][0].model_engine = engine + kwargs["replica_config"][0].engine_config = { + "model_path": "/gpu/model", + "model_format": "pytorch", + } + kwargs["replica_config"][1].model_engine = "MLX" + kwargs["replica_config"][1].engine_config = { + "model_path": "/mac/model", + "model_format": "mlx", + } + await supervisor.launch_builtin_model(**kwargs) + assert set(actors) == {"XavierPDDirectory", "PDModelActor"} + assert actors["PDModelActor"].constructor_kwargs["handoff_mode"] == "host" + for worker, path, fmt in zip( + workers, ["/gpu/model", "/mac/model"], ["pytorch", "mlx"] + ): + options = worker.launch_builtin_model.call_args.kwargs + assert options["model_path"] == path and options["model_format"] == fmt + assert options["_xavier_cache_config"]["host_handoff"] is True + assert options["xavier_config"] is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("engine", ["vLLM", "SGLang"]) +@pytest.mark.parametrize("capacity", [None, 123456]) +async def test_mlx_prefill_to_nvidia_owns_worker_local_source( + launch_runtime, engine, capacity +): + supervisor, workers, actors, destroy = launch_runtime + kwargs = launch_kwargs() + kwargs["model_engine"] = engine + if capacity is not None: + kwargs["xavier_cache_bytes"] = capacity + kwargs["replica_config"][0].model_engine = "MLX" + kwargs["replica_config"][0].engine_config = { + "model_path": "/mac/model", + "model_format": "mlx", + } + await supervisor.launch_builtin_model(**kwargs) + assert set(actors) == {"XavierPDDirectory", "XavierHostPDSource", "PDModelActor"} + source = actors["XavierHostPDSource"] + assert source.address == workers[0].address + assert source.constructor_kwargs["capacity_bytes"] == (capacity or 512 * 1024**2) + assert actors["PDModelActor"].constructor_kwargs["mlx_prefill"] is True + config = workers[0].launch_builtin_model.call_args.kwargs["_xavier_cache_config"] + assert config["source_address"] == source.address + assert config["source_uid"] == source.uid + assert actors["PDModelActor"].constructor_kwargs["handoff_mode"] == "host" + await supervisor.terminate_model("pd") + assert not supervisor._xavier_source_mapping + assert source in [call.args[0] for call in destroy.await_args_list] + + +@pytest.mark.asyncio +async def test_failed_mlx_prefill_launch_destroys_host_source(launch_runtime): + supervisor, workers, actors, destroy = launch_runtime + kwargs = launch_kwargs() + kwargs["replica_config"][0].model_engine = "MLX" + kwargs["replica_config"][0].engine_config = {"model_format": "mlx"} + workers[0].wait_for_load.side_effect = RuntimeError("Metal load failed") + with pytest.raises(RuntimeError, match="Metal load failed"): + await supervisor.launch_builtin_model(**kwargs) + assert not supervisor._xavier_source_mapping + assert actors["XavierHostPDSource"] in [c.args[0] for c in destroy.await_args_list] + assert not supervisor._workers_launching + + +@pytest.mark.asyncio +async def test_unreachable_host_source_does_not_stall_cache_cleanup( + launch_runtime, monkeypatch, caplog +): + supervisor, _, actors, destroy = launch_runtime + kwargs = launch_kwargs() + kwargs["replica_config"][0].model_engine = "MLX" + kwargs["replica_config"][0].engine_config = {"model_format": "mlx"} + await supervisor.launch_builtin_model(**kwargs) + + async def wait_for(coroutine, timeout): + assert timeout == 5 + coroutine.close() + raise asyncio.TimeoutError("worker unreachable") + + bounded = AsyncMock(side_effect=wait_for) + monkeypatch.setattr("xinference.core.supervisor.xo.wait_for", bounded) + await supervisor._cleanup_distributed_actors("pd") + bounded.assert_awaited_once() + assert ( + not supervisor._xavier_source_mapping and not supervisor._xavier_cache_mapping + ) + destroy.assert_any_await(actors["XavierPDDirectory"]) + assert "Destroy Xavier host source failed" in caplog.text + + +@pytest.mark.asyncio +@pytest.mark.parametrize("capacity", [0, -1, True, "512"]) +async def test_mlx_prefill_invalid_capacity_rejected_before_launch( + launch_runtime, capacity +): + supervisor, workers, actors, _ = launch_runtime + kwargs = launch_kwargs() + kwargs["replica_config"][0].model_engine = "MLX" + kwargs["replica_config"][0].engine_config = {"model_format": "mlx"} + kwargs["xavier_cache_bytes"] = capacity + with pytest.raises(ValueError, match="xavier_cache_bytes"): + await supervisor.launch_builtin_model(**kwargs) + assert not actors and not supervisor._model_uid_to_replica_info + for worker in workers: + worker.launch_builtin_model.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "index,fmt", [(0, "mlx"), (0, "ggufv2"), (1, "pytorch"), (1, None)] +) +async def test_nvidia_mlx_invalid_replica_format_has_no_side_effects( + launch_runtime, index, fmt +): + supervisor, workers, actors, _ = launch_runtime + kwargs = launch_kwargs() + kwargs["replica_config"][1].model_engine = "MLX" + kwargs["replica_config"][1].engine_config = {"model_format": "mlx"} + kwargs["replica_config"][index].engine_config["model_format"] = fmt + with pytest.raises(ValueError, match="PyTorch NVIDIA and MLX Metal weights"): + await supervisor.launch_builtin_model(**kwargs) + assert not actors and not supervisor._model_uid_to_replica_info + assert ( + not supervisor._xavier_source_mapping and not supervisor._xavier_cache_mapping + ) + supervisor._status_guard_ref.update_replica_status.assert_not_awaited() + for worker in workers: + worker.launch_builtin_model.assert_not_awaited() + worker.wait_for_load.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mlx_index", [0, 1]) +@pytest.mark.parametrize( + "options,message", + [({"n_worker": 2}, "sharded launch"), ({"quantization": "4-bit"}, "unquantized")], +) +async def test_nvidia_mlx_rejects_unsupported_launch_before_allocation( + launch_runtime, mlx_index, options, message +): + supervisor, workers, actors, _ = launch_runtime + kwargs = launch_kwargs() + kwargs["replica_config"][mlx_index].model_engine = "MLX" + kwargs.update(options) + with pytest.raises(ValueError, match=message): + await supervisor.launch_builtin_model(**kwargs) + assert not actors and not supervisor._model_uid_to_replica_info + for worker in workers: + worker.launch_builtin_model.assert_not_awaited() diff --git a/xinference/core/tests/test_pd_model.py b/xinference/core/tests/test_pd_model.py index 1f0181c9fb..56e2982214 100644 --- a/xinference/core/tests/test_pd_model.py +++ b/xinference/core/tests/test_pd_model.py @@ -182,15 +182,25 @@ async def test_native_sglang_bootstrap_refreshed_on_replica_replacement(): @pytest.mark.asyncio @pytest.mark.parametrize("source_published", [False, True]) +@pytest.mark.parametrize( + "mode,host_source", [("gpu", False), ("host", False), ("host", True)] +) async def test_cross_engine_cancel_drains_source_before_directory_release( - router, monkeypatch, source_published + router, monkeypatch, source_published, mode, host_source ): actor, prefill, decode = router actor._model_engine = "heterogeneous" + actor._handoff_mode = mode calls = [] directory, sender = AsyncMock(), AsyncMock() directory.source.return_value = ( - dict(address="gpu-host", rank=1) if source_published else None + dict( + address="gpu-host", + rank=1, + **({"uid": "host-source"} if host_source else {}), + ) + if source_published + else None ) async def drain(room): @@ -201,7 +211,8 @@ async def release(room): sender.abort.side_effect = drain directory.release.side_effect = release - monkeypatch.setattr(xo, "actor_ref", AsyncMock(side_effect=[directory, sender])) + actor_ref = AsyncMock(side_effect=[directory, sender]) + monkeypatch.setattr(xo, "actor_ref", actor_ref) started = asyncio.Event() async def pending(*args, **kwargs): @@ -213,11 +224,16 @@ async def pending(*args, **kwargs): task = asyncio.create_task(actor._infer("generate", "prompt", {}, request_id="r")) await started.wait() handoff = actor._direct_transfers["r"] + assert handoff["mode"] == mode task.cancel() with pytest.raises(asyncio.CancelledError): await task expected = [("drain", handoff["room"])] if source_published else [] assert calls == expected + [("release", handoff["room"])] + if source_published: + assert actor_ref.await_args_list[-1].kwargs["uid"] == ( + "host-source" if host_source else "xavier-cross-engine-transfer-1" + ) prefill.abort_request.assert_awaited_once() decode.abort_request.assert_awaited_once() assert not actor._request_set and not actor._direct_transfers @@ -248,6 +264,48 @@ async def test_cross_engine_rejects_unsupported_sampling_before_scheduling( assert not actor._request_set +@pytest.mark.asyncio +@pytest.mark.parametrize( + "key", ["min_p", "presence_penalty", "frequency_penalty", "seed"] +) +@pytest.mark.parametrize("raw", [False, True]) +async def test_mlx_prefill_rejects_first_token_sampling_before_scheduling( + router, key, raw +): + actor, prefill, decode = router + actor._model_engine = "heterogeneous" + actor._mlx_prefill = True + parameters = {key: 0} + with pytest.raises(ValueError, match="MLX cross-engine prefill"): + await actor._infer( + "generate", + "prompt", + {} if raw else parameters, + **({"raw_params": parameters} if raw else {}), + request_id="r", + ) + prefill.generate.assert_not_awaited() + decode.generate.assert_not_awaited() + assert not actor._request_set and not actor._direct_transfers + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mlx_prefill", [False, True]) +async def test_first_token_guard_preserves_other_routes_and_unset_parameters( + router, mlx_prefill +): + actor, _, _ = router + actor._model_engine = "heterogeneous" + actor._mlx_prefill = mlx_prefill + actor._infer_sglang = AsyncMock(return_value="result") + config = dict.fromkeys( + ("min_p", "presence_penalty", "frequency_penalty", "seed"), + None if mlx_prefill else 0, + ) + assert await actor._infer("generate", "prompt", config, request_id="r") == "result" + actor._infer_sglang.assert_awaited_once() + + @pytest.mark.asyncio async def test_stream_cleanup_on_disconnect(router): actor, prefill, decode = router diff --git a/xinference/core/tests/test_replica_config.py b/xinference/core/tests/test_replica_config.py index 3399b8b9d0..89550fd15e 100644 --- a/xinference/core/tests/test_replica_config.py +++ b/xinference/core/tests/test_replica_config.py @@ -15,7 +15,12 @@ import pytest -from ..replica_config import DeviceConfig, ReplicaConfig, normalize_replica_configs +from ..replica_config import ( + DeviceConfig, + ReplicaConfig, + normalize_replica_configs, + validate_pd_replica_configs, +) from ..supervisor import SupervisorActor @@ -305,3 +310,26 @@ def test_resolve_auto_gpu_skips_existence_check(): sup._resolve_replica_config("m", 1, [_cfg(worker_ip="10.0.0.1:9978")]) ) assert targets[0][1] is None + + +@pytest.mark.parametrize( + "prefill,decode,valid", + [ + (["MLX", "MLX"], ["vLLM", "SGLang"], True), + (["vLLM", "SGLang"], ["MLX", "MLX"], True), + (["MLX", "vLLM"], ["SGLang"], False), + (["vLLM"], ["MLX", "SGLang"], False), + (["MLX", "vLLM"], ["MLX", "SGLang"], False), + ], +) +def test_nvidia_mlx_topology_keeps_hardware_on_separate_roles(prefill, decode, valid): + configs = [ + ReplicaConfig(role=role, model_engine=engine) + for role, engines in (("prefill", prefill), ("decode", decode)) + for engine in engines + ] + if valid: + assert validate_pd_replica_configs(configs, "vLLM", "LLM") + else: + with pytest.raises(ValueError, match="MLX on one role"): + validate_pd_replica_configs(configs, "vLLM", "LLM") diff --git a/xinference/model/llm/mlx/core.py b/xinference/model/llm/mlx/core.py index d2268e646d..63a1fb9d79 100644 --- a/xinference/model/llm/mlx/core.py +++ b/xinference/model/llm/mlx/core.py @@ -1111,6 +1111,10 @@ def wait_for_load(self): config.update(self._model_config) self._update_model_generation_config(config) self._context_length = get_context_length_from_config(config) + if self._xavier_config and self._xavier_config.get("heterogeneous"): + self._context_length = ( + self._model_config.get("context_length") or self._context_length + ) if self._xavier_config is not None: from .xavier import configure_xavier @@ -1434,8 +1438,16 @@ async def async_generate( if not isinstance(prompt, str): raise ValueError("MLX Xavier requires a text prompt") tokens = self._tokenizer.encode(prompt) + prefill_options = ( + dict(transfer=transfer, generate_config=generate_config) + if getattr(xavier, "config", {}).get("heterogeneous") + else {} + ) metadata = await xavier.prefill( - self._model, tokens, generate_config.get("prompt_cache_prefix_len") + self._model, + tokens, + generate_config.get("prompt_cache_prefix_len"), + **prefill_options, ) prefill_completion = Completion( id=str(uuid.uuid4()), diff --git a/xinference/model/llm/mlx/tests/test_xavier.py b/xinference/model/llm/mlx/tests/test_xavier.py index 10e02805ca..48eee6e435 100644 --- a/xinference/model/llm/mlx/tests/test_xavier.py +++ b/xinference/model/llm/mlx/tests/test_xavier.py @@ -5,7 +5,7 @@ import json import platform from types import SimpleNamespace -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, Mock import pytest import xoscar as xo @@ -77,6 +77,250 @@ def test_vision_model_rejects_xavier_before_loading(role): assert MLXVisionModel("vision-rep0", family, "unused", {})._xavier_config is None +@metal +@pytest.mark.asyncio +@pytest.mark.parametrize("length", [1, 64, 65, 130]) +@pytest.mark.parametrize("biased", [False, True]) +async def test_host_prefill_full_kv_and_first_token(length, monkeypatch, biased): + import mlx.core as mx + from mlx_lm.generate import PromptProcessingBatch + from mlx_lm.models.cache import make_prompt_cache + + m = model() + tokens = [i % 128 for i in range(length)] + reference = PromptProcessingBatch(m, [0], [make_prompt_cache(m)]) + if length > 1: + reference.prompt([tokens[:-1]]) + generation = reference.generate([[tokens[-1]]]) + expected_token = int(generation._next_tokens.item()) + expected_cache = generation.extract_cache(0) + config = dict( + address="directory", + uid="cache", + role="prefill", + rank=1, + heterogeneous=True, + host_handoff=True, + source_address="mac", + source_uid="source", + ) + client = MLXXavierCache(contract(), config) + client._call = AsyncMock(return_value=True) + source_ref = AsyncMock() + monkeypatch.setattr("xoscar.actor_ref", AsyncMock(return_value=source_ref)) + transfer = dict( + sglang_xavier=dict( + room=1, mode="host", address="directory", uid="cache", heterogeneous=True + ) + ) + config = dict(temperature=0) + if biased: + config["logit_bias"] = {17: 1000} + expected_token = 17 + result = await client.prefill(m, tokens, transfer=transfer, generate_config=config) + assert result == transfer + room, metadata, pages, first_token, count = source_ref.publish.call_args.args + assert first_token == expected_token and count == length + actual = client.decode(pages, length) + for expected, value in zip(expected_cache, actual): + for a, b in zip(expected.state, value.state): + assert mx.all(a[:, :, :length] == b).item() + source_ref.release.assert_awaited_once_with(1) + + +@metal +@pytest.mark.asyncio +async def test_host_prefill_abort_releases_cpu_pages(monkeypatch): + client = MLXXavierCache( + contract(), + dict( + address="directory", + uid="cache", + role="prefill", + rank=1, + heterogeneous=True, + source_address="mac", + source_uid="source", + ), + ) + + async def call(method, *args, **kwargs): + if method == "wait_complete": + raise asyncio.CancelledError() + return True + + client._call = AsyncMock(side_effect=call) + source_ref = AsyncMock() + monkeypatch.setattr("xoscar.actor_ref", AsyncMock(return_value=source_ref)) + transfer = dict( + sglang_xavier=dict( + room=1, mode="host", address="directory", uid="cache", heterogeneous=True + ) + ) + with pytest.raises(asyncio.CancelledError): + await client.prefill(model(), [1, 2, 3], transfer=transfer) + source_ref.abort.assert_awaited_once_with(1) + assert client._call.call_args.args == ("release", 1) + + +@metal +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", ["budget", "abort"]) +async def test_host_prefill_reserves_before_metal_and_releases_after_abort_failure( + monkeypatch, failure +): + client = MLXXavierCache( + contract(), + dict( + address="directory", + uid="cache", + role="prefill", + rank=1, + heterogeneous=True, + source_address="mac", + source_uid="source", + ), + ) + + async def call(method, *args, **kwargs): + if method == "wait_complete": + raise RuntimeError("consumer failed") + return True + + client._call = AsyncMock(side_effect=call) + source_ref = AsyncMock() + source_ref.abort.side_effect = RuntimeError("abort failed") + monkeypatch.setattr("xoscar.actor_ref", AsyncMock(return_value=source_ref)) + if failure == "budget": + source_ref.reserve.side_effect = ValueError("prompt exceeds byte budget") + batch = Mock(side_effect=AssertionError("prefill ran before admission")) + monkeypatch.setattr( + importlib.import_module("mlx_lm.generate"), "PromptProcessingBatch", batch + ) + transfer = dict( + sglang_xavier=dict( + room=1, mode="host", address="directory", uid="cache", heterogeneous=True + ) + ) + with pytest.raises((ValueError, RuntimeError), match="byte budget|consumer failed"): + await client.prefill(model(), [1, 2, 3], transfer=transfer) + source_ref.reserve.assert_awaited_once_with(1, client.contract.to_dict(), 3) + client._call.assert_any_await("release", 1) + if failure == "budget": + batch.assert_not_called() + source_ref.publish.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["multiple", "excess", "incomplete"]) +async def test_host_fetch_multiple_chunks_and_prefix_geometry(monkeypatch, kind): + c = contract() + client = MLXXavierCache( + c, dict(address="directory", uid="cache", role="decode", heterogeneous=True) + ) + client.decode = Mock(return_value="restored") + client._call = AsyncMock( + side_effect=lambda method, *args, **kwargs: ( + dict(address="gpu", rank=1) if method == "wait_source" else None + ) + ) + count = {"multiple": 3, "excess": 4, "incomplete": 2}[kind] + pages = [bytes([i]) * c.num_layers * c.layer_nbytes for i in range(count)] + sender = SimpleNamespace( + wait_chunk=AsyncMock( + side_effect=lambda room, index: dict( + ticket=f"1:{index}", pages=[index], final=index == count - 1 + ) + ), + export_host_pages=AsyncMock( + side_effect=lambda room, index, start: dict(pages=[pages[index]], next=1) + ), + release_chunk=AsyncMock(), + abort=AsyncMock(), + ) + monkeypatch.setattr("xoscar.actor_ref", AsyncMock(return_value=sender)) + transfer = dict( + sglang_xavier=dict( + room=1, mode="host", address="directory", uid="cache", heterogeneous=True + ) + ) + tokens = [1] * 130 + if kind == "multiple": + assert await client.fetch(tokens, transfer) == ("restored", 129) + client.decode.assert_called_once_with(pages, 129) + assert sender.wait_chunk.await_count == 3 + client._call.assert_any_await("complete", 1, 0, 129, sum(map(len, pages))) + sender.abort.assert_not_awaited() + else: + with pytest.raises(ValueError, match="page count differs|Incomplete.*prefix"): + await client.fetch(tokens, transfer) + client.decode.assert_not_called() + assert not any( + call.args[0] == "complete" for call in client._call.await_args_list + ) + sender.abort.assert_awaited_once_with(1) + client._call.assert_any_await("release", 1) + + +@pytest.mark.asyncio +async def test_host_fetch_cleanup_survives_repeated_cancel_and_source_failure( + monkeypatch, +): + client = MLXXavierCache( + contract(), + dict(address="directory", uid="cache", role="decode", heterogeneous=True), + ) + abort_started, finish_abort, released = ( + asyncio.Event(), + asyncio.Event(), + asyncio.Event(), + ) + + async def call(method, *args, **kwargs): + if method == "wait_source": + return dict(address="gpu", rank=1) + if method == "release": + released.set() + + async def abort(room): + abort_started.set() + await finish_abort.wait() + raise RuntimeError("source unreachable") + + client._call = AsyncMock(side_effect=call) + sender = SimpleNamespace( + wait_chunk=AsyncMock(side_effect=asyncio.CancelledError()), + abort=AsyncMock(side_effect=abort), + ) + monkeypatch.setattr("xoscar.actor_ref", AsyncMock(return_value=sender)) + transfer = dict( + sglang_xavier=dict( + room=1, mode="host", address="directory", uid="cache", heterogeneous=True + ) + ) + task = asyncio.create_task(client.fetch([1, 2], transfer)) + await asyncio.wait_for(abort_started.wait(), 2) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert not released.is_set() + finish_abort.set() + await asyncio.wait_for(released.wait(), 2) + client._call.assert_any_await("release", 1) + + +@pytest.mark.asyncio +async def test_host_initialize_passes_full_contract(monkeypatch): + c = contract() + client = MLXXavierCache( + c, dict(address="directory", uid="cache", heterogeneous=True) + ) + directory = AsyncMock() + monkeypatch.setattr("xoscar.actor_ref", AsyncMock(return_value=directory)) + await client.initialize() + directory.configure.assert_awaited_once_with(c.fingerprint, c.to_dict()) + + @pytest.mark.asyncio async def test_pd_metadata_is_checked_before_reading(): client = MLXXavierCache( @@ -668,3 +912,188 @@ async def generate(role=None, transfer=None): assert pd == cold and pd_usage["prompt_tokens_details"]["cached_tokens"] == 69 assert [c.publish.call_count for c in clients] == [1, 0, 0] assert (await ref.get_stats())["handoff_reads"] == 1 + + +@metal +@pytest.mark.asyncio +@pytest.mark.parametrize("length", [1, 64, 65, 130]) +async def test_host_pd_imports_real_metal_prefix_and_records_transfer( + monkeypatch, length +): + import mlx.core as mx + from mlx_lm.models.cache import make_prompt_cache + + from ...sglang.xavier.directory import XavierPDDirectory + from ...xavier.pd_contract import prompt_digest + + m = model() + c = contract() + tokens = [i % 128 for i in range(length)] + native = make_prompt_cache(m) + mx.eval(m(mx.array(tokens)[None], cache=native)) + encoded = MLXXavierCache(c, {}).encode(native, tokens) + directory = XavierPDDirectory() + directory.configure(c.fingerprint) + directory.prepare( + 123, c.fingerprint, prompt_digest(tokens), "prefill", prompt_tokens=length + ) + directory.publish_source(123, dict(address="gpu", rank=1)) + cfg = dict( + address="directory", + uid="cache", + role="decode", + rank=2, + heterogeneous=True, + host_handoff=True, + ) + client = MLXXavierCache(c, cfg) + + async def call(method, *args, **kwargs): + result = getattr(directory, method)(*args) + return await result if asyncio.iscoroutine(result) else result + + client._call = call + sender = SimpleNamespace( + wait_chunk=AsyncMock( + return_value=dict( + ticket="123:0", pages=list(range(len(encoded))), final=True + ) + ), + export_host_pages=AsyncMock( + side_effect=lambda room, index, start: dict( + pages=encoded[start : start + 1], next=start + 1 + ) + ), + release_chunk=AsyncMock(), + abort=AsyncMock(), + ) + monkeypatch.setattr( + "xinference.model.llm.mlx.xavier.xo.actor_ref", AsyncMock(return_value=sender) + ) + restored, reused = await client.fetch( + tokens, + dict( + sglang_xavier=dict( + mode="host", + heterogeneous=True, + room=123, + address="directory", + uid="cache", + ) + ), + ) + assert reused == length - 1 and client.imported_tokens == reused + assert all(entry.offset == reused for entry in restored) + if reused: + assert all( + bool(mx.array_equal(a[:, :, :reused], b)) + for old, new in zip(native, restored) + for a, b in zip(old.state, new.state) + ) + actual = m(mx.array(tokens[-1:])[None], cache=restored) + expected = m(mx.array(tokens)[None], cache=make_prompt_cache(m))[:, -1:] + mx.eval(actual, expected) + assert bool(mx.allclose(actual, expected, atol=0.02, rtol=0.02)) + stats = directory.get_stats() + assert stats["gpu_bytes"] == 0 + assert stats["host_bytes"] == sum(map(len, encoded)) + assert stats["imported_tokens"] == reused and stats["completed_requests"] == 1 + sender.release_chunk.assert_awaited_once_with("123:0") + sender.abort.assert_not_awaited() + directory.release(123, "prefill") + assert directory.get_stats()["active_handoffs"] == 0 + + +@pytest.mark.asyncio +async def test_host_pd_truncated_payload_aborts_without_local_prefill(monkeypatch): + c = contract() + client = MLXXavierCache( + c, + dict( + address="directory", + uid="cache", + role="decode", + heterogeneous=True, + host_handoff=True, + ), + ) + client._call = AsyncMock( + side_effect=lambda method, *args, **kwargs: ( + dict(address="gpu", rank=1) if method == "wait_source" else None + ) + ) + sender = SimpleNamespace( + wait_chunk=AsyncMock(return_value=dict(ticket="123:0", pages=[1], final=True)), + export_host_pages=AsyncMock(return_value=dict(pages=[b"truncated"], next=1)), + release_chunk=AsyncMock(), + abort=AsyncMock(), + ) + monkeypatch.setattr( + "xinference.model.llm.mlx.xavier.xo.actor_ref", AsyncMock(return_value=sender) + ) + with pytest.raises(ValueError, match="Incomplete"): + await client.fetch( + [1, 2], + dict( + sglang_xavier=dict( + mode="host", + heterogeneous=True, + room=123, + address="directory", + uid="cache", + ) + ), + ) + assert client.imported_tokens == 0 + sender.abort.assert_awaited_once_with(123) + client._call.assert_any_await("release", 123) + sender.release_chunk.assert_not_awaited() + + +@metal +@pytest.mark.asyncio +async def test_bf16_conversion_in_load_thread_is_eager_before_generation(tmp_path): + import mlx.core as mx + from mlx_lm.generate import BatchGenerator + + (tmp_path / "config.json").write_text(json.dumps({"model_type": "qwen2"})) + (tmp_path / "model.safetensors").write_bytes(b"weights") + (tmp_path / "tokenizer.json").write_text("{}") + + def load(): + m = model() + m.set_dtype(mx.bfloat16) + mx.eval(m.parameters()) + configure_xavier(m, tmp_path, {}, {}, 1) + return m + + m = await asyncio.to_thread(load) + gen = BatchGenerator(m) + gen.insert([[1, 2, 3]], max_tokens=[1]) + for _ in range(10): + _, results = gen.next() + if results: + break + assert results and results[0].finish_reason + gen.close() + + +@metal +def test_host_contract_rejects_effective_geometry_not_in_checkpoint(tmp_path): + from dataclasses import asdict + + m = model("qwen3") + cfg = asdict(m.args) + cfg.pop("head_dim") + (tmp_path / "config.json").write_text(json.dumps(cfg)) + (tmp_path / "model.safetensors").write_bytes(b"weights") + (tmp_path / "tokenizer.json").write_text("{}") + m.args.head_dim = 16 + with pytest.raises(ValueError, match="effective KV geometry"): + configure_xavier( + m, + tmp_path, + {}, + dict(heterogeneous=True, host_handoff=True, role="decode"), + 1, + ) diff --git a/xinference/model/llm/mlx/xavier.py b/xinference/model/llm/mlx/xavier.py index 9f3ae16a28..88ddc1e64f 100644 --- a/xinference/model/llm/mlx/xavier.py +++ b/xinference/model/llm/mlx/xavier.py @@ -6,6 +6,7 @@ import importlib.metadata import json import logging +import time from dataclasses import asdict from pathlib import Path @@ -13,12 +14,14 @@ import xoscar as xo from ....constants import XINFERENCE_CACHE_DIR +from ..xavier.constants import CROSS_ENGINE_TRANSFER_ACTOR_UID from ..xavier.contract import ( KVCacheContract, build_prefix_keys, fingerprint_files, fingerprint_metadata, ) +from ..xavier.pd_contract import build_pd_contract, prompt_digest logger = logging.getLogger(__name__) _TOKENIZER_FILES = ( @@ -62,7 +65,7 @@ def configure_xavier(model, model_path, model_config, cache_config, n_worker): "use_sliding_window", ) ) - or any(t != "full_attention" for t in config.get("layer_types", [])) + or any(t != "full_attention" for t in (config.get("layer_types") or [])) ): raise ValueError( "MLX Xavier supports unquantized full-attention Qwen2, Qwen3 and Llama text models" @@ -85,7 +88,43 @@ def configure_xavier(model, model_path, model_config, cache_config, n_worker): if not weights: raise ValueError("MLX Xavier requires local safetensors weights") model.set_dtype(mx.float16) + # Loading/configuration run in a worker thread. Materialize casts there so + # generation never inherits unevaluated operations on its Metal stream. + mx.eval(model.parameters()) args = asdict(model.args) + if cache_config.get("heterogeneous"): + if cache_config.get("role") not in ( + "prefill", + "decode", + ) or not cache_config.get("host_handoff"): + raise ValueError("Cross-engine MLX Xavier requires a host PD handoff") + # The shared contract describes checkpoint semantics, not engine-specific + # defaults. Verify MLX's effective geometry and RoPE before using it. + for name in ( + "num_hidden_layers", + "num_attention_heads", + "num_key_value_heads", + "hidden_size", + "head_dim", + "rope_theta", + "rope_scaling", + "max_position_embeddings", + "tie_word_embeddings", + ): + if config.get(name) is not None and args.get(name) != config[name]: + raise ValueError(f"Cross-engine MLX effective model override: {name}") + if args.get("rope_traditional") or model_config.get("lora_modules"): + raise ValueError("Cross-engine MLX requires standard RoPE without LoRA") + contract = build_pd_contract( + str(model_path), model_config.get("context_length") + ) + if ( + args.get("head_dim") or args["hidden_size"] // args["num_attention_heads"] + ) != contract.head_dim or ( + args.get("num_key_value_heads") or args["num_attention_heads"] + ) != contract.num_kv_heads: + raise ValueError("Cross-engine MLX effective KV geometry differs") + return MLXXavierCache(contract, cache_config) contract = KVCacheContract( weights_fingerprint=fingerprint_files( weights, cache_dir=Path(XINFERENCE_CACHE_DIR) / "xavier-fingerprints" @@ -123,19 +162,26 @@ async def _ensure_configured(self): address=self.config["address"], uid=self.config["uid"] ) if not self._configured: - metadata = await self._ref.configure(self.contract.to_dict()) - self._capacity_pages = min(metadata["capacity_pages"], metadata["max_keys"]) + if self.config.get("heterogeneous"): + await self._ref.configure( + self.contract.fingerprint, self.contract.to_dict() + ) + else: + metadata = await self._ref.configure(self.contract.to_dict()) + self._capacity_pages = min( + metadata["capacity_pages"], metadata["max_keys"] + ) self._configured = True async def initialize(self): await asyncio.wait_for(self._ensure_configured(), timeout=30) - async def _call(self, method, *args): + async def _call(self, method, *args, timeout=30): async def invoke(): await self._ensure_configured() return await getattr(self._ref, method)(*args) - return await asyncio.wait_for(invoke(), timeout=30) + return await asyncio.wait_for(invoke(), timeout=timeout) def keys(self, tokens): keys = [key.digest for key in build_prefix_keys(self.contract, tokens)] @@ -232,6 +278,8 @@ def decode(self, pages, length): return cache async def fetch(self, tokens, transfer=None): + if self.config.get("heterogeneous"): + return await self._fetch_host(tokens, transfer) prefix = tokens[:-1] keys = self.keys(prefix) if transfer is not None: @@ -269,7 +317,118 @@ async def _release_handoff(self, ticket): except Exception: logger.warning("MLX Xavier handoff release failed", exc_info=True) + def _host_room(self, tokens, transfer, role): + handoff = (transfer or {}).get("sglang_xavier") + if ( + self.role != role + or not tokens + or not isinstance(handoff, dict) + or handoff.get("mode") != "host" + or not handoff.get("heterogeneous") + or handoff.get("address") != self.config["address"] + or handoff.get("uid") != self.config["uid"] + or type(handoff.get("room")) is not int + or not 0 < handoff["room"] < 2**63 + ): + raise ValueError(f"Incompatible cross-engine MLX {role} host handoff") + return handoff["room"] + + async def _cleanup_host(self, room, source): + async def cleanup(): + try: + if source is not None: + await source.abort(room) + except Exception: + logger.warning("Cross-engine source cleanup failed", exc_info=True) + finally: + try: + await self._call("release", room) + except Exception: + logger.warning( + "Cross-engine directory release failed", exc_info=True + ) + + await asyncio.shield(cleanup()) + + async def _fetch_host(self, tokens, transfer): + from ..sglang.xavier.settings import transfer_timeout + + room = self._host_room(tokens, transfer, "decode") + sender = None + try: + await self._call( + "prepare", + room, + self.contract.fingerprint, + prompt_digest(tokens), + "decode", + self.config.get("rank"), + transfer_timeout(), + len(tokens), + ) + deadline = time.monotonic() + transfer_timeout() + + async def rpc(call): + return await asyncio.wait_for( + call, timeout=max(0, deadline - time.monotonic()) + ) + + source = await rpc( + self._call("wait_source", room, timeout=transfer_timeout()) + ) + sender = await rpc( + xo.actor_ref( + address=source["address"], + uid=f"{CROSS_ENGINE_TRANSFER_ACTOR_UID}-{source['rank']}", + ) + ) + pages, index, nbytes = [], 0, 0 + expected = ( + len(tokens) + self.contract.block_size - 1 + ) // self.contract.block_size + while True: + chunk = await rpc(sender.wait_chunk(room, index)) + start = 0 + while start < len(chunk["pages"]): + result = await rpc(sender.export_host_pages(room, index, start)) + payload = result["pages"] + if ( + not payload + or result["next"] != start + len(payload) + or result["next"] > len(chunk["pages"]) + or any( + type(page) is not bytes + or len(page) + != self.contract.num_layers * self.contract.layer_nbytes + for page in payload + ) + ): + raise ValueError("Incomplete cross-engine host pages") + pages.extend(payload) + nbytes += sum(map(len, payload)) + start = result["next"] + if len(pages) > expected: + raise ValueError("Cross-engine host page count differs") + await rpc(sender.release_chunk(chunk["ticket"])) + index += 1 + if chunk["final"]: + break + if len(pages) != expected: + raise ValueError("Incomplete cross-engine host prefix") + length = len(tokens) - 1 + count = (length + self.contract.block_size - 1) // self.contract.block_size + cache = self.decode(pages[:count], length) + await self._call("complete", room, 0, length, nbytes) + await self._call("release", room, "decode") + return cache, length + except BaseException: + await self._cleanup_host(room, sender) + raise + def publish(self, cache, tokens, cached_tokens=0): + if self.config.get("heterogeneous"): + return + async def write(): try: await self.initialize() @@ -303,10 +462,87 @@ async def flush(self, writes): *(asyncio.shield(task) for task in writes), return_exceptions=True ) - async def prefill(self, model, tokens, prefix_length=None): + async def _prefill_host(self, model, tokens, prefix_length, transfer, config): + import mlx.core as mx + from mlx_lm.generate import PromptProcessingBatch + from mlx_lm.models.cache import make_prompt_cache + from mlx_lm.sample_utils import make_logits_processors, make_sampler + + from ..sglang.xavier.settings import transfer_timeout + + room = self._host_room(tokens, transfer, "prefill") + source = None + try: + await self._call( + "prepare", + room, + self.contract.fingerprint, + prompt_digest(tokens), + "prefill", + self.config["rank"], + transfer_timeout(), + len(tokens), + ) + source = await xo.actor_ref( + address=self.config["source_address"], uid=self.config["source_uid"] + ) + await asyncio.wait_for( + source.reserve(room, self.contract.to_dict(), len(tokens)), + timeout=transfer_timeout(), + ) + batch = PromptProcessingBatch( + model, [0], [make_prompt_cache(model)], prefill_step_size=2048 + ) + start = 0 + while start < len(tokens) - 1: + end = min(start + 2048, len(tokens) - 1) + if prefix_length is not None and start < prefix_length < end: + end = prefix_length + batch.prompt([tokens[start:end]]) + start = end + await asyncio.sleep(0) + # Match GenerationBatch's final prompt step, without generating a + # second token and extending KV beyond the prompt. SGLang consumes + # this sampled token together with the full prompt cache. + logits = model(mx.array([[tokens[-1]]]), cache=batch.prompt_cache)[:, -1, :] + for processor in make_logits_processors( + logit_bias=config.get("logit_bias"), + repetition_penalty=config.get("repetition_penalty"), + repetition_context_size=config.get("repetition_context_size", 20), + ): + logits = processor(mx.array(tokens), logits) + logprobs = logits - mx.logsumexp(logits, axis=-1, keepdims=True) + sampler = make_sampler( + temp=config.get("temperature", 0), + top_p=config.get("top_p", 1), + top_k=int(config.get("top_k") or 0), + ) + first_token = int(sampler(logprobs).item()) + pages = self.encode(batch.extract_cache(0), tokens) + await source.publish( + room, self.contract.to_dict(), pages, first_token, len(tokens) + ) + # CPU ownership now belongs to the source actor. Native engine slots + # may only be committed after the consumer completes its import. + del pages, batch + await self._call("wait_complete", room, timeout=transfer_timeout()) + await source.release(room) + await self._call("release", room, "prefill") + return dict(transfer) + except BaseException: + await self._cleanup_host(room, source) + raise + + async def prefill( + self, model, tokens, prefix_length=None, transfer=None, generate_config=None + ): from mlx_lm.generate import PromptProcessingBatch from mlx_lm.models.cache import make_prompt_cache + if self.config.get("heterogeneous"): + return await self._prefill_host( + model, tokens, prefix_length, transfer, generate_config or {} + ) if self.role != "prefill" or not tokens: raise ValueError( "MLX remote prefill requires a prefill replica and a nonempty prompt" diff --git a/xinference/model/llm/sglang/tests/test_xavier_gpu_lifecycle.py b/xinference/model/llm/sglang/tests/test_xavier_gpu_lifecycle.py index 18b8d90585..c793cfca6a 100644 --- a/xinference/model/llm/sglang/tests/test_xavier_gpu_lifecycle.py +++ b/xinference/model/llm/sglang/tests/test_xavier_gpu_lifecycle.py @@ -376,20 +376,80 @@ def test_sender_poll_never_blocks_scheduler_on_actor(gpu_module, failed): @pytest.mark.asyncio -async def test_manager_commits_engine_metadata_before_completion(gpu_module): +@pytest.mark.parametrize("phase", ["open", "init", "send"]) +async def test_cross_engine_late_cancel_fails_request_without_stopping_scheduler( + gpu_module, monkeypatch, phase +): + loop = asyncio.get_running_loop() + tasks = [] + error = RuntimeError("cancelled room") + + async def operation(method): + if method == phase: + raise error + + monkeypatch.setattr(gpu_module, "KVPoll", SimpleNamespace(Failed="failed")) + actor = SimpleNamespace(wait_done=AsyncMock()) + calls = [] + + async def invoke(method, *args): + calls.append(method) + await operation(method) + + actor.open = lambda *args: invoke("open", *args) + actor.init = lambda *args: invoke("init", *args) + actor.add_chunk = lambda *args: invoke("send", *args) + + def submit(coroutine): + future = asyncio.run_coroutine_threadsafe(coroutine, loop) + tasks.append(future) + return future + + manager = SimpleNamespace( + config={"heterogeneous": True}, + actor=actor, + call=Mock(side_effect=AssertionError("scheduler RPC wait")), + submit=submit, + kv_args=SimpleNamespace(gpu_id=0), + aux=[], + ) + monkeypatch.setattr(gpu_module.torch.cuda, "synchronize", Mock()) + sender = gpu_module.XavierKVSender(manager, "host", 123, [0], 0) + sender.init(2, 0) + sender.send(SimpleNamespace(tolist=lambda: [1, 2])) + await asyncio.gather( + *(asyncio.wrap_future(f) for f in tasks), return_exceptions=True + ) + assert sender.poll() == gpu_module.KVPoll.Failed + with pytest.raises(RuntimeError, match="cancelled room"): + sender.failure_exception() + count = len(calls) + sender.send(SimpleNamespace(tolist=lambda: [1, 2])) + await asyncio.gather( + *(asyncio.wrap_future(f) for f in tasks), return_exceptions=True + ) + assert len(calls) == count + manager.call.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("host", [False, True]) +async def test_manager_commits_engine_metadata_before_completion(gpu_module, host): import numpy as np import torch backing = np.zeros((3, 2), dtype=np.uint8) manager = object.__new__(gpu_module.XavierKVManager) manager._receive_tasks = {} + manager.config = dict(host_handoff=host) manager.aux = [torch.from_numpy(backing)] manager.actor = SimpleNamespace( receive=AsyncMock(return_value=(1024, [b"\x07\x09"])) ) - async def complete(room, nbytes): - assert room == 1 and nbytes == 1024 + async def complete(room, nbytes, host_bytes=0): + assert room == 1 + assert (nbytes, host_bytes) == ((0, 1024) if host else (1024, 0)) assert backing.tolist() == [[0, 0], [7, 9], [0, 0]] manager.directory = SimpleNamespace(complete=AsyncMock(side_effect=complete)) @@ -411,6 +471,7 @@ async def receive(*args): manager = object.__new__(gpu_module.XavierKVManager) manager._receive_tasks = {} + manager.config = {} manager.aux = [torch.zeros((1, 1), dtype=torch.uint8)] manager.actor = SimpleNamespace(receive=receive, abort=AsyncMock()) manager.directory = SimpleNamespace(complete=AsyncMock()) @@ -441,6 +502,7 @@ async def receive(*args): manager = object.__new__(gpu_module.XavierKVManager) manager._receive_tasks = {} + manager.config = {} manager.aux = [torch.zeros((1, 1), dtype=torch.uint8)] manager.actor = SimpleNamespace(receive=receive) manager.directory = SimpleNamespace(complete=AsyncMock()) @@ -463,6 +525,7 @@ def test_manager_close_retains_exports_until_importer_stops(gpu_module, monkeypa manager = object.__new__(gpu_module.XavierKVManager) manager._closed = False manager._receive_tasks = {} + manager.config = {} manager._ipc_caches = {"0": object()} manager.aux = [object()] manager.actor = object() diff --git a/xinference/model/llm/sglang/tests/test_xavier_pd.py b/xinference/model/llm/sglang/tests/test_xavier_pd.py index 043f973280..6ce52218a3 100644 --- a/xinference/model/llm/sglang/tests/test_xavier_pd.py +++ b/xinference/model/llm/sglang/tests/test_xavier_pd.py @@ -373,14 +373,20 @@ async def prepare(h, prompt, room=1): @pytest.mark.asyncio @pytest.mark.parametrize("heterogeneous", [False, True]) +@pytest.mark.parametrize("host", [False, True]) async def test_repeated_prompt_reuses_fingerprint_and_validates_every_room( heterogeneous, + host, ): h = fingerprint_handoff() + h.config["host_handoff"] = host if heterogeneous: h.config.update(heterogeneous=True, rank=1) for room in range(1, 4): - await prepare(h, "prompt", room) + await h.prepare( + "prompt", + dict(sglang_xavier=dict(mode="host" if host else "gpu", room=room)), + ) h.tokenizer.encode.assert_called_once_with("prompt") assert h._directory_actor.prepare.await_count == 3 expected = hashlib.sha256(json.dumps(list(b"prompt")).encode()).hexdigest() diff --git a/xinference/model/llm/sglang/xavier/directory.py b/xinference/model/llm/sglang/xavier/directory.py index fa4d2b580c..b42be107db 100644 --- a/xinference/model/llm/sglang/xavier/directory.py +++ b/xinference/model/llm/sglang/xavier/directory.py @@ -18,6 +18,7 @@ def __init__(self): self.rooms = {} self.completed_requests = 0 self.gpu_bytes = 0 + self.host_bytes = 0 self.imported_tokens = 0 self.peers = {} @@ -148,7 +149,7 @@ def source(self, room): raise RuntimeError("SGLang Xavier PD handoff expired or was cancelled") return state["source"] - def complete(self, room, gpu_bytes, imported_tokens=None): + def complete(self, room, gpu_bytes, imported_tokens=None, host_bytes=0): state = self.rooms.get(room) if state is None or set(state["roles"]) != {"prefill", "decode"}: raise RuntimeError("SGLang Xavier PD request is unavailable") @@ -156,6 +157,7 @@ def complete(self, room, gpu_bytes, imported_tokens=None): self.completed_requests += 1 self.gpu_bytes += gpu_bytes state["deadline"] = time.monotonic() + state["timeout"] + self.host_bytes += host_bytes self.imported_tokens += ( imported_tokens if imported_tokens is not None @@ -203,6 +205,7 @@ def get_stats(self): active_handoffs=len(self.rooms), completed_requests=self.completed_requests, gpu_bytes=self.gpu_bytes, + host_bytes=self.host_bytes, imported_tokens=self.imported_tokens, peers=self.peers, ) diff --git a/xinference/model/llm/sglang/xavier/gpu.py b/xinference/model/llm/sglang/xavier/gpu.py index 85d83fc430..f19cb29640 100644 --- a/xinference/model/llm/sglang/xavier/gpu.py +++ b/xinference/model/llm/sglang/xavier/gpu.py @@ -494,7 +494,10 @@ async def _receive(self, room, destinations, aux_index): buf[aux_index].copy_( torch.from_numpy(np.frombuffer(payload, dtype=np.uint8).copy()) ) - await self.directory.complete(room, nbytes) + if self.config.get("host_handoff"): + await self.directory.complete(room, 0, host_bytes=nbytes) + else: + await self.directory.complete(room, nbytes) async def abort(self, room): await self.actor.abort(room) diff --git a/xinference/model/llm/sglang/xavier/pd.py b/xinference/model/llm/sglang/xavier/pd.py index f367b4bb12..a9ec388f03 100644 --- a/xinference/model/llm/sglang/xavier/pd.py +++ b/xinference/model/llm/sglang/xavier/pd.py @@ -37,8 +37,9 @@ async def invoke(): async def prepare(self, prompt: str, transfer: dict) -> dict: handoff = transfer.get("sglang_xavier") - if not isinstance(handoff, dict) or handoff.get("mode") != "gpu": - raise ValueError("Missing SGLang Xavier GPU bootstrap metadata") + mode = "host" if self.config.get("host_handoff") else "gpu" + if not isinstance(handoff, dict) or handoff.get("mode") != mode: + raise ValueError(f"Missing SGLang Xavier {mode} bootstrap metadata") key = hashlib.sha256(prompt.encode()).digest() cached = self._prompt_hashes.pop(key, None) if cached is None: diff --git a/xinference/model/llm/vllm/tests/test_pd.py b/xinference/model/llm/vllm/tests/test_pd.py index 96a33094ab..1336081388 100644 --- a/xinference/model/llm/vllm/tests/test_pd.py +++ b/xinference/model/llm/vllm/tests/test_pd.py @@ -135,8 +135,9 @@ async def generate(self, prompt, params, request_id, **kwargs): @pytest.mark.parametrize( "failure", ["one-token", "prompt mismatch", "room limit", "producer RPC"] ) +@pytest.mark.parametrize("host", [False, True]) async def test_cross_engine_validation_fails_request_without_submitting_bad_work( - monkeypatch, failure + monkeypatch, failure, host ): import xoscar as xo @@ -181,6 +182,7 @@ def __init__(self, max_tokens=1, **kwargs): uid="directory", contract=contract.to_dict(), gpu_cache_bytes=0, + host_handoff=host, ) output = {"xavier_error": failure} if failure == "producer RPC" else None @@ -196,7 +198,9 @@ async def generate(*args, **kwargs): ) params = { "max_tokens": 1, - "_pd_kv_transfer_params": {"sglang_xavier": {"mode": "gpu", "room": 123}}, + "_pd_kv_transfer_params": { + "sglang_xavier": {"mode": "host" if host else "gpu", "room": 123} + }, } with pytest.raises( (ValueError, RuntimeError), diff --git a/xinference/model/llm/vllm/xavier/cross_engine.py b/xinference/model/llm/vllm/xavier/cross_engine.py index 8c264c35e2..cc8385007a 100644 --- a/xinference/model/llm/vllm/xavier/cross_engine.py +++ b/xinference/model/llm/vllm/xavier/cross_engine.py @@ -155,8 +155,9 @@ def register_kv_caches(self, kv_caches): async def _prepare(self, request): params = request.kv_transfer_params or {} handoff = params.get("sglang_xavier") - if not isinstance(handoff, dict) or handoff.get("mode") != "gpu": - raise ValueError("Missing cross-engine Xavier GPU room") + mode = "host" if self._xavier_config.get("host_handoff") else "gpu" + if not isinstance(handoff, dict) or handoff.get("mode") != mode: + raise ValueError("Missing cross-engine Xavier handoff room") room = handoff["room"] if params.get("xavier_prompt_digest") != prompt_digest( request.prompt_token_ids diff --git a/xinference/model/llm/xavier/backends/bytes/pd.py b/xinference/model/llm/xavier/backends/bytes/pd.py new file mode 100644 index 0000000000..e0e8f5094b --- /dev/null +++ b/xinference/model/llm/xavier/backends/bytes/pd.py @@ -0,0 +1,184 @@ +# Copyright 2022-2026 Xinference Holdings Pte. Ltd +# Licensed under the Apache License, Version 2.0. +"""Request-owned host pages for a Metal producer and NVIDIA consumer.""" + +import asyncio +import time + +import xoscar as xo + +from ....sglang.xavier.settings import transfer_timeout +from ...contract import KVCacheContract + + +class XavierHostPDSource(xo.StatelessActor): + def __init__(self, directory, rank: int, capacity_bytes: int = 512 * 1024**2): + super().__init__() + if type(capacity_bytes) is not int or capacity_bytes <= 0: + raise ValueError("Xavier host PD byte budget must be positive") + self.directory, self.rank = directory, rank + self.capacity_bytes = capacity_bytes + self.contract: KVCacheContract | None = None + self.rooms: dict[int, dict] = {} + self.active_bytes = 0 + self.host_bytes = 0 + self._capacity_changed = asyncio.Event() + self._waiting: set[int] = set() + self._closed = False + + async def __post_create__(self): + await self.directory.register_peer(self.rank, self.address) + + async def __pre_destroy__(self): + self._closed = True + self.rooms.clear() + self._waiting.clear() + self.active_bytes = 0 + self._capacity_changed.set() + await self.directory.unregister_peer(self.rank, self.address) + + def _expire(self): + for room, state in list(self.rooms.items()): + if state["deadline"] <= time.monotonic(): + self.release(room) + + async def reserve(self, room: int, metadata: dict, prompt_tokens: int) -> None: + """Reserve rounded page capacity before doing any Metal prefill work.""" + self._expire() + contract = KVCacheContract.from_dict(metadata) + if self.contract is not None: + self.contract.require_match(contract) + page_bytes = contract.num_layers * contract.layer_nbytes + if type(prompt_tokens) is not int or prompt_tokens <= 0: + raise ValueError("Invalid Xavier host PD prompt length") + if page_bytes > 64 * 1024**2: + raise ValueError("Xavier host PD page exceeds 64 MiB") + nbytes = ( + (prompt_tokens + contract.block_size - 1) + // contract.block_size + * page_bytes + ) + if nbytes > self.capacity_bytes: + raise ValueError("Xavier host PD prompt exceeds the in-flight byte budget") + if room in self.rooms or room in self._waiting: + raise ValueError("Duplicate Xavier host PD source") + deadline = time.monotonic() + transfer_timeout() + self._waiting.add(room) + try: + await self.directory.configure(contract.fingerprint, contract.to_dict()) + if self.contract is not None: + self.contract.require_match(contract) + self.contract = contract + info = await self.directory.request_info(room) + if info["prompt_tokens"] != prompt_tokens: + raise ValueError("Xavier host PD prompt length differs") + while True: + self._expire() + if self._closed: + raise RuntimeError("Xavier host PD source was destroyed") + if room not in self._waiting: + raise RuntimeError("Xavier host PD reservation was cancelled") + if time.monotonic() >= deadline: + raise TimeoutError("Xavier host PD capacity wait timed out") + self._capacity_changed.clear() + if ( + len(self.rooms) < 4096 + and self.active_bytes + nbytes <= self.capacity_bytes + ): + break + expiry = min(s["deadline"] for s in self.rooms.values()) + try: + await asyncio.wait_for( + self._capacity_changed.wait(), + timeout=max(0, min(expiry, deadline) - time.monotonic()), + ) + except asyncio.TimeoutError: + if time.monotonic() >= deadline: + raise TimeoutError("Xavier host PD capacity wait timed out") + self.rooms[room] = dict( + pages=None, + prompt_tokens=prompt_tokens, + nbytes=nbytes, + deadline=deadline, + ) + self.active_bytes += nbytes + finally: + self._waiting.discard(room) + + async def publish(self, room, metadata, pages, first_token, prompt_tokens): + self._expire() + state = self.rooms.get(room) + if state is None: + raise RuntimeError("Unreserved or expired Xavier host PD source") + if state["pages"] is not None: + raise ValueError("Duplicate Xavier host PD source") + try: + contract = KVCacheContract.from_dict(metadata) + assert self.contract is not None + self.contract.require_match(contract) + page_bytes = contract.num_layers * contract.layer_nbytes + if ( + type(prompt_tokens) is not int + or prompt_tokens != state["prompt_tokens"] + or type(first_token) is not int + or first_token < 0 + or len(pages) * page_bytes != state["nbytes"] + or any(type(p) is not bytes or len(p) != page_bytes for p in pages) + ): + raise ValueError("Invalid Xavier host PD pages or first token") + state.update(pages=tuple(pages), first_token=first_token) + await self.directory.publish_source( + room, + dict( + address=self.address, + uid=self.uid, + rank=self.rank, + transport="host", + ), + ) + except BaseException: + self.release(room) + raise + + def read(self, room: int, start: int = 0) -> dict: + self._expire() + state = self.rooms.get(room) + if state is None: + raise RuntimeError("Xavier host PD source expired or was cancelled") + if state["pages"] is None: + raise RuntimeError("Xavier host PD source has not been published") + if type(start) is not int or not 0 <= start < len(state["pages"]): + raise ValueError("Invalid Xavier host PD page cursor") + assert self.contract is not None + page_bytes = self.contract.num_layers * self.contract.layer_nbytes + count = min(16, 64 * 1024**2 // page_bytes) + pages = state["pages"][start : start + count] + state["deadline"] = time.monotonic() + transfer_timeout() + self.host_bytes += len(pages) * page_bytes + return dict( + pages=pages, + next=start + len(pages), + total=len(state["pages"]), + first_token=state["first_token"], + prompt_tokens=state["prompt_tokens"], + ) + + def release(self, room: int) -> None: + if room in self._waiting: + self._waiting.discard(room) + self._capacity_changed.set() + state = self.rooms.pop(room, None) + if state is not None: + self.active_bytes -= state["nbytes"] + self._capacity_changed.set() + + def abort(self, room: int) -> None: + self.release(room) + + def get_stats(self) -> dict: + self._expire() + return dict( + active_rooms=len(self.rooms), + active_bytes=self.active_bytes, + host_bytes=self.host_bytes, + ) diff --git a/xinference/model/llm/xavier/backends/torch/pd.py b/xinference/model/llm/xavier/backends/torch/pd.py index a10d3291ac..f1b033e38a 100644 --- a/xinference/model/llm/xavier/backends/torch/pd.py +++ b/xinference/model/llm/xavier/backends/torch/pd.py @@ -12,11 +12,13 @@ from types import SimpleNamespace from typing import Any +import numpy as np import torch import xoscar as xo from ....sglang.gc_lifecycle import InitializationGCFreeze from ....sglang.xavier.settings import transfer_timeout +from ...constants import CROSS_ENGINE_TRANSFER_ACTOR_UID from .direct_handoff import DirectGPUTransfer from .gpu_transfer import finish_before_cancel from .snapshot import block_major_view @@ -77,7 +79,7 @@ def canonical_vllm_views(caches, num_blocks, contract): class CrossEngineGPUActor(xo.StatelessActor): @classmethod def default_uid(cls): - return "xavier-cross-engine-transfer" + return CROSS_ENGINE_TRANSFER_ACTOR_UID def __init__( self, @@ -270,6 +272,102 @@ async def send_direct_gpu_v1(self, ticket, reads, remote_ref, slab_bytes): room["deadline"] = time.monotonic() + transfer_timeout() return result + def _export_host_pages(self, pages, offset, prefix_tokens): + """Copy owned canonical pages, clearing unused engine allocation bytes.""" + device = next(iter(self.caches.values())).device + if device.type == "cuda": + torch.cuda.set_device(device) + try: + return self._copy_host_pages(pages, offset, prefix_tokens, device) + finally: + if device.type == "cuda": + torch.cuda.synchronize(device) + + def _copy_host_pages(self, pages, offset, prefix_tokens, device): + c = self.contract + indices = torch.tensor(pages, dtype=torch.long, device=device) + data = torch.stack( + [ + torch.stack( + [ + self.caches[str(layer + kind * c.num_layers)].index_select( + 0, indices + ) + for kind in (0, 1) + ], + dim=1, + ) + for layer in range(c.num_layers) + ], + dim=1, + ).cpu() + for index in range(len(pages)): + valid = max( + 0, min(c.block_size, prefix_tokens - (offset + index) * c.block_size) + ) + data[index, :, :, valid:] = 0 + return [page.numpy().astype(" 64 * 1024**2: + raise ValueError("Cross-engine host page exceeds 64 MiB") + capacity = max(1, min(16, 64 * 1024**2 // (c.layer_nbytes * c.num_layers))) + pages = chunk["pages"][start : start + capacity] + if any( + not set(pages) <= allowed for allowed in state.layer_blocks.values() + ): + raise ValueError("Unowned cross-engine source pages") + state.reading = True + state.claimed = True + try: + offset = ( + sum( + len(item["pages"]) + for item in self.rooms[room]["chunks"][:index] + ) + + start + ) + # Cancellation must drain the CUDA read before abort can free + # engine slots. The returned bytes then own their CPU storage. + payload = await finish_before_cancel( + asyncio.create_task( + asyncio.to_thread( + self._export_host_pages, + pages, + offset, + room_state["prompt_tokens"] - 1, + ) + ) + ) + self.transfer.metrics["host_bytes"] = self.transfer.metrics.get( + "host_bytes", 0 + ) + sum(map(len, payload)) + return dict(pages=payload, next=start + len(pages)) + finally: + state.reading = False + state.deadline = time.monotonic() + 600 + if state.released: + self.transfer.release_direct(chunk["ticket"]) + def release_chunk(self, ticket): self.transfer.release_direct(ticket) room = int(ticket.split(":", 1)[0]) @@ -337,6 +435,8 @@ async def _receive(self, room, destinations, aux_index): source = await asyncio.wait_for( self.directory.wait_source(room), timeout=transfer_timeout() ) + if source.get("transport") == "host": + return await self._receive_host(room, source, destinations) rank = source["rank"] self._world_addresses[rank] = source["address"] sender = await xo.actor_ref( @@ -412,6 +512,97 @@ async def _receive(self, room, destinations, aux_index): # GPU writes drain, before native decode can commit the slots. return nbytes, chunk["aux"] + def _import_host_pages(self, pages, destinations): + device = next(iter(self.caches.values())).device + if device.type == "cuda": + torch.cuda.set_device(device) + try: + c = self.contract + data = torch.from_numpy( + np.stack( + [ + np.frombuffer(page, dtype=" 16 + or len(pages) * page_bytes > 64 * 1024**2 + or result["next"] != start + len(pages) + or result["next"] > expected + or result["total"] != expected + or result["prompt_tokens"] != prompt_tokens + or type(result["first_token"]) is not int + or result["first_token"] < 0 + or ( + first_token is not None and first_token != result["first_token"] + ) + or any(type(p) is not bytes or len(p) != page_bytes for p in pages) + ): + raise ValueError("Incomplete cross-engine host prefix") + first_token = result["first_token"] + targets = destinations[start : result["next"]] + if targets: + task = asyncio.create_task( + asyncio.to_thread( + self._import_host_pages, pages[: len(targets)], targets + ) + ) + # Slot reuse after abort must wait for CUDA writes to drain, + # including errors and repeated cancellation during H2D. + await finish_before_cancel(task) + nbytes += len(pages) * page_bytes + start = result["next"] + await sender.release(room) + self.transfer.metrics["host_imported_bytes"] = ( + self.transfer.metrics.get("host_imported_bytes", 0) + nbytes + ) + if not self.args.aux_item_lens: + await self.directory.complete(room, 0, imported, nbytes) + await self.directory.release(room, "decode") + return set() + return nbytes, sglang_first_token_payload( + self.args.aux_item_lens, room, first_token, prompt_tokens + ) + except BaseException: + await sender.abort(room) + raise + async def receive(self, room, destinations, aux_index): task = self.tasks[room] = asyncio.create_task( self._receive(room, destinations, aux_index) diff --git a/xinference/model/llm/xavier/constants.py b/xinference/model/llm/xavier/constants.py index dcc721247f..ac29edda14 100644 --- a/xinference/model/llm/xavier/constants.py +++ b/xinference/model/llm/xavier/constants.py @@ -3,3 +3,4 @@ # Preserve the existing actor UID for deployments using the default adapter. DEFAULT_TRANSFER_ACTOR_UID = "vllm-transfer-actor" +CROSS_ENGINE_TRANSFER_ACTOR_UID = "xavier-cross-engine-transfer" diff --git a/xinference/model/llm/xavier/pd_contract.py b/xinference/model/llm/xavier/pd_contract.py index 6116094ef6..5318c59dd5 100644 --- a/xinference/model/llm/xavier/pd_contract.py +++ b/xinference/model/llm/xavier/pd_contract.py @@ -24,6 +24,7 @@ def build_pd_contract( or config.get("use_sliding_window") or config.get("vision_config") or config.get("quantization_config") + or config.get("quantization") or config.get("auto_map") or any(kind != "full_attention" for kind in (config.get("layer_types") or [])) ): @@ -99,8 +100,9 @@ async def prepare_pd_request( if cache_config["role"] == "decode" and len(tokens) < 2: raise ValueError("Cross-engine vLLM decode requires at least two prompt tokens") handoff = params.get("sglang_xavier") - if not isinstance(handoff, dict) or handoff.get("mode") != "gpu": - raise ValueError("Missing cross-engine Xavier GPU room") + mode = "host" if cache_config.get("host_handoff") else "gpu" + if not isinstance(handoff, dict) or handoff.get("mode") != mode: + raise ValueError("Missing cross-engine Xavier handoff room") digest = prompt_digest(tokens) directory = await xo.actor_ref( address=cache_config["address"], uid=cache_config["uid"] diff --git a/xinference/model/llm/xavier/tests/test_pd_cross_engine.py b/xinference/model/llm/xavier/tests/test_pd_cross_engine.py index 52744c6fdc..849c9a63ff 100644 --- a/xinference/model/llm/xavier/tests/test_pd_cross_engine.py +++ b/xinference/model/llm/xavier/tests/test_pd_cross_engine.py @@ -329,8 +329,9 @@ async def test_producer_poll_keeps_finished_ids_when_release_rpc_fails( ([1, 2, 3], "room limit"), ], ) +@pytest.mark.parametrize("host", [False, True]) async def test_request_preparation_errors_stay_outside_engine( - kv_contract, monkeypatch, tokens, error + kv_contract, monkeypatch, tokens, error, host ): import xoscar as xo @@ -345,8 +346,9 @@ async def test_request_preparation_errors_stay_outside_engine( address="supervisor", uid="directory", contract=kv_contract.to_dict(), + host_handoff=host, ) - params = dict(sglang_xavier=dict(room=123, mode="gpu")) + params = dict(sglang_xavier=dict(room=123, mode="host" if host else "gpu")) with pytest.raises(ValueError, match=error): await prepare_pd_request(config, tokens, params) assert "xavier_prompt_digest" not in params @@ -357,6 +359,33 @@ async def test_request_preparation_errors_stay_outside_engine( assert params["xavier_prompt_digest"] == prompt_digest([1, 2]) +@pytest.mark.asyncio +@pytest.mark.parametrize("host", [False, True]) +@pytest.mark.parametrize("role", ["prefill", "decode"]) +async def test_request_preparation_rejects_wrong_transport_before_rpc( + kv_contract, monkeypatch, host, role +): + import xoscar as xo + + from ..pd_contract import prepare_pd_request + + lookup = AsyncMock() + monkeypatch.setattr(xo, "actor_ref", lookup) + config = dict( + role=role, + rank=1, + address="supervisor", + uid="directory", + contract=kv_contract.to_dict(), + host_handoff=host, + ) + params = dict(sglang_xavier=dict(room=123, mode="gpu" if host else "host")) + with pytest.raises(ValueError, match="handoff room"): + await prepare_pd_request(config, [1, 2], params) + lookup.assert_not_awaited() + assert "xavier_prompt_digest" not in params + + @pytest.mark.parametrize( "field", ["weights_fingerprint", "tokenizer_fingerprint", "position_fingerprint"] ) diff --git a/xinference/model/llm/xavier/tests/test_pd_host.py b/xinference/model/llm/xavier/tests/test_pd_host.py new file mode 100644 index 0000000000..a8285c6ea2 --- /dev/null +++ b/xinference/model/llm/xavier/tests/test_pd_host.py @@ -0,0 +1,398 @@ +# Copyright 2022-2026 Xinference Holdings Pte. Ltd +# Licensed under the Apache License, Version 2.0. +import asyncio +import threading +from dataclasses import replace +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import numpy as np +import pytest +import torch + +from ..backends.bytes.pd import XavierHostPDSource +from ..backends.torch.pd import CrossEngineGPUActor + + +def source(kv_contract): + c = replace(kv_contract, num_layers=2) + actor = CrossEngineGPUActor(None, c, AsyncMock(), c.fingerprint, 1) + actor.caches = { + str(i): torch.arange( + 10 * c.block_size * c.num_kv_heads * c.head_dim, dtype=torch.float16 + ).reshape(10, c.block_size, c.num_kv_heads, c.head_dim) + + i + for i in range(4) + } + return actor + + +def test_host_export_preserves_physical_order_and_clears_padding(kv_contract): + actor = source(kv_contract) + c = actor.contract + original = {k: v.clone() for k, v in actor.caches.items()} + pages = actor._export_host_pages([7, 2], 0, c.block_size + 1) + data = np.stack( + [ + np.frombuffer(p, dtype="