diff --git a/doc/BUILD.bazel b/doc/BUILD.bazel index 968b4711e0d6..da2c808f1199 100644 --- a/doc/BUILD.bazel +++ b/doc/BUILD.bazel @@ -352,6 +352,8 @@ py_test_run_all_subdirectory( "source/llm/doc_code/serve/direct_streaming/**/*.py", "source/llm/doc_code/serve/multi_gpu/**/*.py", "source/llm/doc_code/serve/sglang/**/*.py", + # TODO (jeffreywang): Enable TPU example tests once CI sets up TPU. + "source/llm/doc_code/serve/tpu/**/*.py", # custom_vllm runs in its own direct-streaming target below. "source/llm/doc_code/serve/custom_vllm/**/*.py", ], diff --git a/doc/source/llm/doc_code/serve/tpu/serve_tpu_multihost.py b/doc/source/llm/doc_code/serve/tpu/serve_tpu_multihost.py new file mode 100644 index 000000000000..fb7bd49f0055 --- /dev/null +++ b/doc/source/llm/doc_code/serve/tpu/serve_tpu_multihost.py @@ -0,0 +1,29 @@ +# flake8: noqa +""" +Example Ray Serve LLM application for a multi-host TPU slice. +""" + +# __serve_tpu_multihost_start__ +from ray import serve +from ray.serve.llm import LLMConfig, LLMServingArgs, build_openai_app + +llm_config = LLMConfig( + model_loading_config=dict( + model_id="google/gemma-4-31B-it", + model_source="/data/google/gemma-4-31B-it", + ), + accelerator_type="TPU-V6E", + accelerator_config={"kind": "tpu", "topology": "4x4"}, + engine_kwargs=dict( + tensor_parallel_size=16, + max_model_len=8192, + max_num_batched_tokens=8192, + distributed_executor_backend="ray", + ), +) + +app = build_openai_app(LLMServingArgs(llm_configs=[llm_config])) + +if __name__ == "__main__": + serve.run(app, blocking=True) +# __serve_tpu_multihost_end__ diff --git a/doc/source/serve/llm/images/ray_serve_llm_tpu.png b/doc/source/serve/llm/images/ray_serve_llm_tpu.png new file mode 100644 index 000000000000..7fac1bd4a2d3 Binary files /dev/null and b/doc/source/serve/llm/images/ray_serve_llm_tpu.png differ diff --git a/doc/source/serve/llm/user-guides/configuration.md b/doc/source/serve/llm/user-guides/configuration.md index eeb0351ae56d..e30d903f770f 100644 --- a/doc/source/serve/llm/user-guides/configuration.md +++ b/doc/source/serve/llm/user-guides/configuration.md @@ -84,8 +84,8 @@ A cloud mirror is validated against {class}`~ray.serve.llm.CloudMirrorConfig` an For the full list, see the [vLLM engine arguments](https://docs.vllm.ai/en/latest/serving/engine_args.html). `disable_log_stats=True` is rejected when `log_engine_metrics` is enabled, because engine metrics require log stats. ## Accelerators and placement - -Set `accelerator_type` to the accelerator each replica should be scheduled on. Common GPU values include `A10G`, `A100`, `H100`, `H200`, `L4`, `L40S`, and `T4`. TPUs such as `TPU-V4` and `TPU-V5P` are also supported. `A10` is normalized to `A10G`. An unsupported value raises a validation error. +Set `accelerator_type` to the accelerator each replica should be scheduled on. Common GPU values include `A10G`, `A100`, `H100`, `H200`, `L4`, `L40S`, and `T4`. TPUs such as `TPU-V4`, `TPU-V5P`, and `TPU-V6E` are also supported. `A10` is normalized to `A10G`. An unsupported value raises a validation error. For topology-aware TPU slices, see {doc}`TPU serving `. +Set `accelerator_type` to the accelerator each replica should be scheduled on. Common GPU values include `A10G`, `A100`, `H100`, `H200`, `L4`, `L40S`, and `T4`. TPUs such as `TPU-V4`, `TPU-V5P`, and `TPU-V6E` are also supported. For topology-aware TPU slices, see {doc}`TPU serving `. `A10` is normalized to `A10G`. An unsupported value raises a validation error. `accelerator_config` is inferred from `accelerator_type` (or from `placement_group_config` bundles) and rarely needs to be set explicitly. diff --git a/doc/source/serve/llm/user-guides/index.md b/doc/source/serve/llm/user-guides/index.md index b8dafffe9e78..90f928f31330 100644 --- a/doc/source/serve/llm/user-guides/index.md +++ b/doc/source/serve/llm/user-guides/index.md @@ -27,6 +27,10 @@ How-to guides for deploying, scaling, and operating Ray Serve LLM. If you are ne - {doc}`Custom vLLM models `: serve an out-of-tree architecture with a vLLM plugin, using a Qwen3 reward model as the example. - {doc}`SGLang integration `: run SGLang as the inference engine instead of vLLM. +## Accelerator-specific serving + +- {doc}`TPU serving `: serve a model on single-host or multi-host TPU slices with topology-aware placement. + ## Operate in production - {doc}`Observability and monitoring `: engine and request metrics, Grafana dashboards, and Prometheus integration. @@ -48,5 +52,6 @@ Direct streaming vLLM compatibility Custom vLLM models SGLang integration +TPU serving Observability and monitoring ``` diff --git a/doc/source/serve/llm/user-guides/tpu.md b/doc/source/serve/llm/user-guides/tpu.md new file mode 100644 index 000000000000..3f573753e763 --- /dev/null +++ b/doc/source/serve/llm/user-guides/tpu.md @@ -0,0 +1,115 @@ +--- +myst: + html_meta: + description: "Serve LLMs on single-host and multi-host TPU slices with Ray Serve LLM, including topology-aware placement groups and a multi-host vLLM TPU example." +--- + +(serve-llm-tpu)= +# TPU serving + +Ray Serve LLM can run a vLLM TPU engine on single-host and multi-host TPU slices, where a TPU slice is a group of interconnected TPU chips. Use this when your Ray cluster already exposes TPU resources and TPU node labels, and your container image includes the TPU variant of vLLM from `tpu-inference`. For Kubernetes setup, see {doc}`Use TPUs with KubeRay `. + +## Topology and placement + +A TPU topology describes the chip grid in one physical TPU slice. For example, a v6e `4x4` slice has 16 chips. Multi-host slices spread those chips across multiple TPU hosts connected by the TPU interconnect. For v6e `4x4`, that usually means four hosts with four chips each. For visualization, check out [TPU Topology Visualizer](https://tpu-visualizer.uc.r.appspot.com/). + +Ray Serve LLM uses `tensor_parallel_size * pipeline_parallel_size` as the number of TPU chips that one model replica requests. When you also set a TPU topology, Ray Serve LLM computes the number of chips per host and creates one placement group bundle per TPU host: + +```python +LLMConfig( + accelerator_type="TPU-V6E", + accelerator_config={"kind": "tpu", "topology": "4x4"}, + engine_kwargs={"tensor_parallel_size": 16}, +) +``` + +For a v6e `4x4` slice, this produces four host-level bundles: + +```python +[ + {"TPU": 4, "accelerator_type:TPU-V6E": 0.001}, + {"TPU": 4, "accelerator_type:TPU-V6E": 0.001}, + {"TPU": 4, "accelerator_type:TPU-V6E": 0.001}, + {"TPU": 4, "accelerator_type:TPU-V6E": 0.001}, +] +``` + +The model still spans all 16 chips because `tensor_parallel_size=16`. The bundle shape only tells Ray how to reserve the hosts that own those chips. If you need per-chip bundles, set `placement_group_config={"bundle_per_worker": {"TPU": 1}}`. + +```{figure} ../images/ray_serve_llm_tpu.png +--- +width: 100% +name: ray-serve-llm-tpu-placement +alt: One Ray Serve LLM replica spanning a v6e 4x4 TPU slice, with one placement group bundle per TPU host and four TPU chips reserved in each bundle. +--- +Topology-aware TPU placement for one Ray Serve LLM replica. +``` + +:::{note} +TPU support in Ray Serve LLM is topology-aware when you set `accelerator_config={"kind": "tpu", "topology": ...}`. Without a topology, Ray Serve LLM falls back to a regular placement group with per-chip `{"TPU": 1}` bundles. +::: + +## `SlicePlacementGroup` + +For topology-aware TPU configs, Ray Serve LLM creates a {class}`~ray.util.tpu.SlicePlacementGroup` instead of a plain placement group. `SlicePlacementGroup` reserves a matching TPU slice, reads its `ray.io/tpu-slice-name` label, and creates the worker placement group with a per-bundle label selector that pins all bundles to that same physical slice. + +This makes the TPU slice an atomic scheduling unit. A replica reserves the complete slice it needs, and the placement group doesn't span unrelated slices. + +## How the TPU vLLM executor uses the bundles + +The `tpu-inference` [Ray executor](https://github.com/vllm-project/tpu-inference/blob/main/tpu_inference/executors/ray_distributed_executor_v2.py) checks `parallel_config.placement_group`. When Ray Serve LLM already provided one, the executor reuses it instead of creating its own. + +The executor then does the following: + +1. Selects the TPU bundles from the placement group. +1. Reads the TPU count from each bundle. +1. Starts one Ray worker actor per bundle with `PlacementGroupSchedulingStrategy`, pinning each actor to its bundle. + +With the default topology-aware bundles, one vLLM worker actor maps to one TPU host and consumes all local TPU chips on that host. The worker process uses the chips across the multi-host slice. Otherwise, manually setting `{"TPU": 1}` results in the creation of one worker per chip. + +## Example + +Build the image from the vLLM TPU base image and install a Ray wheel that includes TPU topology support. + +::::{tab-set} + +:::{tab-item} Dockerfile + +```dockerfile +FROM vllm/vllm-tpu:v0.21.0 + +ENV VLLM_TARGET_DEVICE=tpu +ENV VLLM_XLA_CACHE_PATH=/tmp/vllm_xla_cache +ENV JAX_PLATFORMS=tpu,cpu +ENV TPU_MULTIHOST_BACKEND=ray +ENV TPU_BACKEND_TYPE=jax +ENV ENABLE_PJRT_COMPATIBILITY=true + +USER root + +# Use a released Ray version or wheel URL that contains TPU topology support. +ARG RAY_PACKAGE="ray" +RUN pip install --no-cache-dir -U "${RAY_PACKAGE}" && \ + pip install --no-cache-dir --no-deps "ray[llm]" + +COPY serve_tpu_multihost.py /home/ray/serve_tpu_multihost.py +``` +::: + +:::{tab-item} Python + +```{literalinclude} ../../../llm/doc_code/serve/tpu/serve_tpu_multihost.py +:language: python +:start-after: __serve_tpu_multihost_start__ +:end-before: __serve_tpu_multihost_end__ +``` +::: + +:::: + +This example serves `google/gemma-4-31B-it` on one v6e `4x4` slice with `tensor_parallel_size=16`. Set `model_source` to a local or mounted model path, or another path that all TPU hosts can read. + +## See also + +- {doc}`Use TPUs with KubeRay ` +- [Serve Gemma open models using multi-host TPUs on GKE with Ray](https://docs.cloud.google.com/kubernetes-engine/docs/tutorials/serve-multi-host-tpu-llm)