@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / dd39f91edb
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 12, 19:45 UTC

[Doc] cleanup TPU documentation and remove outdated examples (#29048)

Signed-off-by: Rob Mulla <rob.mulla@gmail.com> Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

Rob Mulla committed Nov 20, 2025 at 19:05 UTC dd39f91edb0588e2dd77eb55c758eb1e35907af8
8 files changed +40 -415
docs/.nav.yml
+4 -2
@@ -24,14 +24,16 @@ nav:
24 - deployment/integrations
25 - Training: training
26 - Configuration:
27 - - configuration/README.md
27 - configuration/*
28 + - TPU: https://docs.vllm.ai/projects/tpu/en/latest/
29 - Models:
30 - models/supported_models.md
31 - models/generative_models.md
32 - models/pooling_models.md
33 - models/extensions
34 - - Hardware Supported Models: models/hardware_supported_models
34 + - Hardware Supported Models:
35 + - models/hardware_supported_models/*
36 + - TPU: https://docs.vllm.ai/projects/tpu/en/latest/recommended_models_features/
37 - Features: features
38 - Developer Guide:
39 - contributing/README.md
docs/configuration/tpu.md deleted
-111
@@ -1,111 +0,0 @@
1 -# TPU Optimization Tips
2 -
3 -This doc serves as a collection of handy tips for optimizing your vLLM on TPU workload.
4 -
5 -## Get started
6 -
7 -Looking for setup and installation instructions? Find them [here](https://docs.vllm.ai/projects/tpu/en/latest/getting_started/installation/).
8 -
9 -### TPU workload sizing
10 -
11 -When selecting the ideal number of chips for a single serving instance, it's important to account for both the model size and the average request context length. Adequate HBM for the KV cache is essential to ensure a sufficient number of concurrent requests can be processed.
12 -
13 -The following colab [calculator](https://colab.research.google.com/github/ericehanley/rightsize-vllm/blob/main/HBM_Calculator.ipynb) will tell you:
14 -
15 -- KV cache size requirement per token and per request
16 -- TPU/GPU memory consumed by the model weights
17 -- TPU/GPU memory allocated for the KV cache
18 -- Maximum \# of requests you can approximately set (--max-num-seqs)
19 -
20 -This approach serves as a general rule of thumb.
21 -
22 -#### Latency-throughput tradeoff
23 -
24 -As with rightsizing the number of chips for your workload, consider adjusting `--max-num-seqs` to fine-tune the latency-throughput balance. Decreasing `--max-num-seqs` and/or increasing the number of chips can help reduce latency.
25 -
26 -`--max-num-seqs` defines the number of concurrent decode slots, effectively limiting the number of requests the server can process tokens for simultaneously. Increasing this value allows the server to pre-allocate more HBM to handle a higher number of concurrent requests, which can maximize overall throughput. However, this often increases the end-to-end (e2e) latency per request.
27 -
28 -Therefore, carefully tuning `--max-num-seqs` is crucial to achieving the desired balance between latency and throughput for your specific workload.
29 -
30 -In a similar way, `--max-num-batch-tokens` can be adjusted down to improve latency, or adjusted up to improve throughput.
31 -
32 -#### Compilation and Caching
33 -
34 -Coming from a GPU background, one of the key differences you'll notice with TPUs is an initial compilation step. TPUs are specialized accelerators (ASICs) that achieve maximum performance by executing pre-compiled, static computation graphs via the XLA compiler. Unlike GPUs, which can handle dynamic input shapes more flexibly, TPUs require a specific compiled graph for each tensor shape (e.g., batch size and sequence length) they process.
35 -
36 -To manage this, vLLM performs a one-time "warmup" process when you first launch the server. During this phase, it pre-compiles the model for various common input shapes and saves these compiled graphs to a cache on disk or remote storage (located at `~/.cache/vllm/xla_cache` by default). This process can range significantly, anywhere from a few minutes to an hour depending on the size of the model and context length used.
37 -
38 -Although the first compilation can take some time, for all subsequent server launches, vLLM can load these graphs directly from the cache, eliminating the compilation time for future runs.
39 -
40 -Use `VLLM_XLA_CACHE_PATH` environment variable to write to shareable storage for future deployed nodes (like when using autoscaling).
41 -
42 -#### Reducing compilation time
43 -
44 -This initial compilation time ranges significantly and is impacted by many of the arguments discussed in this optimization doc. Factors that influence the length of time to compile are things like model size and `--max-num-batch-tokens`. Other arguments you can tune are things like `VLLM_TPU_MOST_MODEL_LEN`.
45 -
46 -### Optimize based on your data
47 -
48 -#### max-model-len vs. most-model-len
49 -
50 -![most_model_len](../assets/design/tpu/most_model_len.png)
51 -
52 -If most of your requests are shorter than the maximum model length but you still need to accommodate occasional longer requests, setting a high maximum model length can negatively impact performance. In these cases, you can try introducing most-model-len by specifying the `VLLM_TPU_MOST_MODEL_LEN` environment variable.
53 -
54 -For example, 1% requests are 32k length and 99% requests are 2k length. You can pass 32k into `--max-model-len 32768` and use `VLLM_TPU_MOST_MODEL_LEN=2048`.
55 -
56 -The requests get subdivided into max-model-len and most-model-len categories, for the latter category, you can gain better performance since the server can process more requests at a time.
57 -
58 -#### Padding
59 -
60 -For online serving with latency requirements, consider switching to bucket padding by setting the `VLLM_TPU_BUCKET_PADDING_GAP` environment variable. Because of the layout of the TPU, try using increments of 128 (e.g., 128, 256, etc.)
61 -
62 -The server pads the requests into fixed lengths before sending them to the model to avoid recompilation. To read more about TPU padding, see [here](https://cloud.google.com/tpu/docs/performance-guide#xla-efficiencies). Currently, there are 2 ways to pad the requests:
63 -
64 -1. the default exponential padding (pad to the nearest power of 2)
65 -2. bucket padding (pad to the nearest linearly increasing bucket).
66 -
67 -When using bucket padding, the buckets start from 16, end at max_model_len, and increment by `VLLM_TPU_BUCKET_PADDING_GAP`.
68 -
69 -For example, max_model_len=512, padding_gap=64, the buckets will be [16, 32, 64, 128, 192, 256, 320, 384, 448, 512].
70 -
71 -The fewer tokens you pad, the less unnecessary computation TPU does, the better performance you can get. For example, if num_tokens=300, with exponential padding, you pad to 512, with the bucket_padding above, you pad to 320.
72 -
73 -However, you need to be careful to choose the padding gap. If the gap is too small, it means the number of buckets is large, leading to increased warmup (precompile) time and higher memory to store the compiled graph. Too many compiled graphs may lead to HBM OOM. Conversely, an overly large gap yields no performance improvement compared to the default exponential padding.
74 -
75 -#### Quantization
76 -
77 -If possible, use the precision that matches the chip’s hardware acceleration:
78 -
79 -- v5e has int4/int8 hardware acceleration in the MXU
80 -- v6e has int4/int8 hardware acceleration in the MXU
81 -
82 -Supported quantized formats and features in vLLM on TPU [Jul '25]:
83 -
84 -- INT8 W8A8
85 -- INT8 W8A16
86 -- FP8 KV cache
87 -- [WIP] FP8 W8A8
88 -- [WIP] AWQ
89 -- [WIP] FP4 W4A8
90 -
91 -#### Parallelization
92 -
93 -Don't set TP to be less than the number of chips on a single-host deployment.
94 -
95 -Although it’s common to do this with GPUs, don't try to fragment 2 or 8 different workloads across 8 chips on a single host. If you need 1 or 4 chips, just create an instance with 1 or 4 chips (these are partial-host machine types).
96 -
97 -### Tune your workloads
98 -
99 -Although we try to have great default configs, we strongly recommend you check out the [vLLM auto-tuner](../../benchmarks/auto_tune/README.md) to optimize your workloads for your use case.
100 -
101 -### Future Topics We'll Cover
102 -
103 -#### Profiling
104 -
105 -The auto-tuner provides a profile of optimized configurations as its final step. However, interpreting this profile can be challenging for new users. We plan to expand this section in the future with more detailed guidance. In the meantime, you can learn how to collect a TPU profile using vLLM's native profiling tools [here](../examples/offline_inference/profiling_tpu.md). This profile can provide valuable insights into your workload's performance.
106 -
107 -#### SPMD
108 -
109 -More details to come.
110 -
111 -**Want us to cover something that isn't listed here? Open up an issue please and cite this doc. We'd love to hear your questions or tips.**
docs/features/README.md
+20 -17
@@ -59,20 +59,23 @@ th:not(:first-child) {
59
60 ### Feature x Hardware
61
62 -| Feature | Volta | Turing | Ampere | Ada | Hopper | CPU | AMD | TPU | Intel GPU |
63 -|-----------------------------------------------------------|---------------------|-----------|-----------|--------|------------|--------------------|--------|-----| ------------|
64 -| [CP](../configuration/optimization.md#chunked-prefill) | [❌](https://github.com/vllm-project/vllm/issues/2729) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
65 -| [APC](automatic_prefix_caching.md) | [❌](https://github.com/vllm-project/vllm/issues/3687) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
66 -| [LoRA](lora.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
67 -| [SD](spec_decode.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | [🟠](https://github.com/vllm-project/vllm/issues/26963) |
68 -| CUDA graph | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | [❌](https://github.com/vllm-project/vllm/issues/26970) |
69 -| [pooling](../models/pooling_models.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ |
70 -| <abbr title="Encoder-Decoder Models">enc-dec</abbr> | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ |
71 -| [mm](multimodal_inputs.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | [🟠](https://github.com/vllm-project/vllm/issues/26965) |
72 -| <abbr title="Logprobs">logP</abbr> | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ |
73 -| <abbr title="Prompt Logprobs">prmpt logP</abbr> | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ |
74 -| <abbr title="Async Output Processing">async output</abbr> | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ |
75 -| multi-step | ✅ | ✅ | ✅ | ✅ | ✅ | [❌](https://github.com/vllm-project/vllm/issues/8477) | ✅ | ❌ | ✅ |
76 -| best-of | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ |
77 -| beam-search | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ |
78 -| [prompt-embeds](prompt_embeds.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❔ | [❌](https://github.com/vllm-project/vllm/issues/25097) | ✅ |
62 +| Feature | Volta | Turing | Ampere | Ada | Hopper | CPU | AMD | Intel GPU |
63 +|-----------------------------------------------------------|---------------------|-----------|-----------|--------|------------|--------------------|--------| ------------|
64 +| [CP](../configuration/optimization.md#chunked-prefill) | [❌](https://github.com/vllm-project/vllm/issues/2729) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
65 +| [APC](automatic_prefix_caching.md) | [❌](https://github.com/vllm-project/vllm/issues/3687) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
66 +| [LoRA](lora.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
67 +| [SD](spec_decode.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | [🟠](https://github.com/vllm-project/vllm/issues/26963) |
68 +| CUDA graph | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | [❌](https://github.com/vllm-project/vllm/issues/26970) |
69 +| [pooling](../models/pooling_models.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
70 +| <abbr title="Encoder-Decoder Models">enc-dec</abbr> | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ |
71 +| [mm](multimodal_inputs.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | [🟠](https://github.com/vllm-project/vllm/issues/26965) |
72 +| [prompt-embeds](prompt_embeds.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❔ | ✅ |
73 +| <abbr title="Logprobs">logP</abbr> | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
74 +| <abbr title="Prompt Logprobs">prmpt logP</abbr> | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
75 +| <abbr title="Async Output Processing">async output</abbr> | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ |
76 +| multi-step | ✅ | ✅ | ✅ | ✅ | ✅ | [❌](https://github.com/vllm-project/vllm/issues/8477) | ✅ | ✅ |
77 +| best-of | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
78 +| beam-search | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
79 +
80 +!!! note
81 + For information on feature support on Google TPU, please refer to the [TPU-Inference Recommended Models and Features](https://docs.vllm.ai/projects/tpu/en/latest/recommended_models_features/) documentation.
docs/features/quantization/README.md
+16 -13
@@ -43,24 +43,27 @@ th:not(:first-child) {
43 }
44 </style>
45
46 -| Implementation | Volta | Turing | Ampere | Ada | Hopper | AMD GPU | Intel GPU | Intel Gaudi | x86 CPU | Google TPU |
47 -|-----------------------|---------|----------|----------|-------|----------|-----------|-------------|-------------|-----------|--------------|
48 -| AWQ | ❌ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ✅︎ | ❌ | ✅︎ | ❌ |
49 -| GPTQ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ✅︎ | ❌ | ✅︎ | ❌ |
50 -| Marlin (GPTQ/AWQ/FP8) | ❌ | ❌ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ | ❌ | ❌ |
51 -| INT8 (W8A8) | ❌ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ | ✅︎ | ✅︎ |
52 -| FP8 (W8A8) | ❌ | ❌ | ❌ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ | ❌ |
53 -| BitBLAS | ✅︎ | ✅ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ | ❌ | ❌ |
54 -| BitBLAS (GPTQ) | ❌ | ❌ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ | ❌ | ❌ |
55 -| bitsandbytes | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ | ❌ | ❌ |
56 -| DeepSpeedFP | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ | ❌ | ❌ |
57 -| GGUF | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ | ❌ |
58 -| INC (W8A8) | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅︎ | ❌ | ❌ |
46 +| Implementation | Volta | Turing | Ampere | Ada | Hopper | AMD GPU | Intel GPU | Intel Gaudi | x86 CPU |
47 +|-----------------------|---------|----------|----------|-------|----------|-----------|-------------|-------------|-----------|
48 +| AWQ | ❌ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ✅︎ | ❌ | ✅︎ |
49 +| GPTQ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ✅︎ | ❌ | ✅︎ |
50 +| Marlin (GPTQ/AWQ/FP8) | ❌ | ❌ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ | ❌ |
51 +| INT8 (W8A8) | ❌ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ | ✅︎ |
52 +| FP8 (W8A8) | ❌ | ❌ | ❌ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ |
53 +| BitBLAS | ✅︎ | ✅ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ | ❌ |
54 +| BitBLAS (GPTQ) | ❌ | ❌ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ | ❌ |
55 +| bitsandbytes | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ | ❌ |
56 +| DeepSpeedFP | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ | ❌ |
57 +| GGUF | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ |
58 +| INC (W8A8) | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅︎ | ❌ |
59
60 - Volta refers to SM 7.0, Turing to SM 7.5, Ampere to SM 8.0/8.6, Ada to SM 8.9, and Hopper to SM 9.0.
61 - ✅︎ indicates that the quantization method is supported on the specified hardware.
62 - ❌ indicates that the quantization method is not supported on the specified hardware.
63
64 +!!! note
65 + For information on quantization support on Google TPU, please refer to the [TPU-Inference Recommended Models and Features](https://docs.vllm.ai/projects/tpu/en/latest/recommended_models_features/) documentation.
66 +
67 !!! note
68 This compatibility chart is subject to change as vLLM continues to evolve and expand its support for different hardware platforms and quantization methods.
69
docs/models/hardware_supported_models/tpu.md deleted
-34
@@ -1,34 +0,0 @@
1 -# TPU
2 -
3 -## Supported Models
4 -
5 -### Text-only Language Models
6 -
7 -| Model | Architecture | Supported |
8 -|-----------------------------------------------------|--------------------------------|-----------|
9 -| mistralai/Mixtral-8x7B-Instruct-v0.1 | MixtralForCausalLM | 🟨 |
10 -| mistralai/Mistral-Small-24B-Instruct-2501 | MistralForCausalLM | ✅ |
11 -| mistralai/Codestral-22B-v0.1 | MistralForCausalLM | ✅ |
12 -| mistralai/Mixtral-8x22B-Instruct-v0.1 | MixtralForCausalLM | ❌ |
13 -| meta-llama/Llama-3.3-70B-Instruct | LlamaForCausalLM | ✅ |
14 -| meta-llama/Llama-3.1-8B-Instruct | LlamaForCausalLM | ✅ |
15 -| meta-llama/Llama-3.1-70B-Instruct | LlamaForCausalLM | ✅ |
16 -| meta-llama/Llama-4-* | Llama4ForConditionalGeneration | ❌ |
17 -| microsoft/Phi-3-mini-128k-instruct | Phi3ForCausalLM | 🟨 |
18 -| microsoft/phi-4 | Phi3ForCausalLM | ❌ |
19 -| google/gemma-3-27b-it | Gemma3ForConditionalGeneration | 🟨 |
20 -| google/gemma-3-4b-it | Gemma3ForConditionalGeneration | ❌ |
21 -| deepseek-ai/DeepSeek-R1 | DeepseekV3ForCausalLM | ❌ |
22 -| deepseek-ai/DeepSeek-V3 | DeepseekV3ForCausalLM | ❌ |
23 -| RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w8a8 | LlamaForCausalLM | ✅ |
24 -| RedHatAI/Meta-Llama-3.1-70B-Instruct-quantized.w8a8 | LlamaForCausalLM | ✅ |
25 -| Qwen/Qwen3-8B | Qwen3ForCausalLM | ✅ |
26 -| Qwen/Qwen3-32B | Qwen3ForCausalLM | ✅ |
27 -| Qwen/Qwen2.5-7B-Instruct | Qwen2ForCausalLM | ✅ |
28 -| Qwen/Qwen2.5-32B | Qwen2ForCausalLM | ✅ |
29 -| Qwen/Qwen2.5-14B-Instruct | Qwen2ForCausalLM | ✅ |
30 -| Qwen/Qwen2.5-1.5B-Instruct | Qwen2ForCausalLM | 🟨 |
31 -
32 -✅ Runs and optimized.
33 -🟨 Runs and correct but not optimized to green yet.
34 -❌ Does not pass accuracy test or does not run.
examples/offline_inference/profiling_tpu/README.md deleted
-70
@@ -1,70 +0,0 @@
1 -# vLLM TPU Profiling
2 -
3 -This script is used to profile the TPU performance of vLLM for specific prefill or decode token shapes.
4 -
5 -Note: an actual running server is a mix of both prefill of many shapes and decode of many shapes.
6 -
7 -We assume you are on a TPU already (this was tested on TPU v6e) and have installed vLLM according to the [Google TPU installation guide](https://docs.vllm.ai/en/latest/getting_started/installation/google_tpu.html).
8 -
9 -> In all examples below, we run several warmups before (so `--enforce-eager` is okay)
10 -
11 -## Profile Examples
12 -
13 -### Generate Prefill Trace
14 -
15 -This example runs Qwen/Qwen2.5-7B-Instruct with a single request of 1024 input tokens. This is set up in attempt to profile just the prefill time and operations.
16 -
17 -```bash
18 -export XLA_HLO_DEBUG=1
19 -export MODEL=Qwen/Qwen2.5-7B-Instruct
20 -export VLLM_TPU_PROFILE_DURATION_MS=3000
21 -export VLLM_TPU_PROFILE_DELAY_MS=0
22 -
23 -python3 profiling.py \
24 - --model $MODEL \
25 - --input-len 1024 --output-len 1 \
26 - --batch-size 1 --enforce-eager \
27 - --max-model-len 2048 \
28 - --tensor-parallel-size 1 \
29 - --profile-result-dir profiles
30 -```
31 -
32 -### Generate Decode Trace
33 -
34 -This example runs Llama 3.1 70B with a batch of 32 requests where each has 1 input token and 128 output tokens. This is set up in attempt to profile just the 32 decodes running in parallel by having an extremely small prefill of 1 token and setting `VLLM_TPU_PROFILE_DELAY_MS=1000` to skip the first second of inference (hopefully prefill).
35 -
36 -```bash
37 -export XLA_HLO_DEBUG=1
38 -export MODEL=meta-llama/Llama-3.1-70B-Instruct
39 -export VLLM_TPU_PROFILE_DURATION_MS=2000
40 -export VLLM_TPU_PROFILE_DELAY_MS=1000
41 -
42 -rm -rf ~/.cache/vllm/xla_cache
43 -python3 profiling.py \
44 - --model $MODEL \
45 - --input-len 1 \
46 - --output-len 128 \
47 - --batch-size 32 \
48 - --enforce-eager \
49 - --profile-result-dir profiles \
50 - --max-model-len 2048 --tensor-parallel-size 8
51 -```
52 -
53 -## Visualizing the profiles
54 -
55 -Once you have collected your profiles with this script, you can visualize them using [TensorBoard](https://cloud.google.com/tpu/docs/pytorch-xla-performance-profiling-tpu-vm).
56 -
57 -Here are most likely the dependencies you need to install:
58 -
59 -```bash
60 -pip install tensorflow-cpu \
61 - tensorboard-plugin-profile \
62 - etils \
63 - importlib_resources
64 -```
65 -
66 -Then you just need to point TensorBoard to the directory where you saved the profiles and visit `http://localhost:6006/` in your browser:
67 -
68 -```bash
69 -tensorboard --logdir profiles/ --port 6006
70 -```
examples/offline_inference/profiling_tpu/profiling.py deleted
-110
@@ -1,110 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -import argparse
5 -import dataclasses
6 -import os
7 -import time
8 -
9 -import numpy as np
10 -import torch_xla.debug.profiler as xp
11 -from tqdm import tqdm
12 -
13 -from vllm import LLM, SamplingParams
14 -from vllm.engine.arg_utils import EngineArgs
15 -from vllm.inputs import PromptType
16 -from vllm.utils.argparse_utils import FlexibleArgumentParser
17 -
18 -DURATION_MS = int(os.getenv("VLLM_TPU_PROFILE_DURATION_MS", 3000))
19 -DELAY_MS = int(os.getenv("VLLM_TPU_PROFILE_DELAY_MS", 0))
20 -
21 -
22 -def main(args: argparse.Namespace):
23 - print(args)
24 -
25 - engine_args = EngineArgs.from_cli_args(args)
26 - llm = LLM(**dataclasses.asdict(engine_args))
27 - server = xp.start_server(9012) # noqa: F841
28 -
29 - sampling_params = SamplingParams(
30 - temperature=0.0,
31 - ignore_eos=True,
32 - max_tokens=args.output_len,
33 - )
34 - print(sampling_params)
35 - dummy_prompt_token_ids = np.random.randint(
36 - 10000, size=(args.batch_size, args.input_len)
37 - )
38 - dummy_prompts: list[PromptType] = [
39 - {"prompt_token_ids": batch} for batch in dummy_prompt_token_ids.tolist()
40 - ]
41 -
42 - def run_to_completion():
43 - start_time = time.perf_counter()
44 - llm.generate(dummy_prompts, sampling_params=sampling_params, use_tqdm=False)
45 - end_time = time.perf_counter()
46 - latency = end_time - start_time
47 - return latency
48 -
49 - # Warmup
50 - print("Warming up...")
51 - warmup_latencies = []
52 - for _ in tqdm(range(args.num_iters_warmup), desc="Warmup iterations"):
53 - warmup_latencies.append(run_to_completion())
54 - print(f"Average warmup latency: {np.mean(warmup_latencies):.4f}s")
55 -
56 - # Profile
57 - profile_dir = args.profile_result_dir
58 - print(f"Profiling (results will be saved to '{profile_dir}')...")
59 - # Enable tracing on server
60 - xp.trace_detached(
61 - "localhost:9012", profile_dir, delay_ms=DELAY_MS, duration_ms=DURATION_MS
62 - )
63 - if DELAY_MS == 0:
64 - time.sleep(1.0)
65 - profile_latencies = []
66 - for _ in tqdm(range(args.num_iters), desc="Profile iterations"):
67 - profile_latencies.append(run_to_completion())
68 - print(f"Average profile latency: {np.mean(profile_latencies):.4f}s")
69 -
70 - return
71 -
72 -
73 -def parse_args():
74 - parser = FlexibleArgumentParser(
75 - description="Benchmark the latency of processing a single batch of "
76 - "requests till completion."
77 - )
78 - parser.add_argument("--input-len", type=int, default=32)
79 - parser.add_argument("--output-len", type=int, default=128)
80 - parser.add_argument("--batch-size", type=int, default=8)
81 - parser.add_argument(
82 - "--num-iters-warmup",
83 - type=int,
84 - default=5,
85 - help="Number of iterations to run for warmup.",
86 - )
87 - parser.add_argument(
88 - "--num-iters",
89 - type=int,
90 - default=1,
91 - help="Number of iterations to run for profiling.",
92 - )
93 - parser.add_argument(
94 - "--profile-result-dir",
95 - type=str,
96 - default="profiles",
97 - help=(
98 - "path to save the pytorch profiler output. Can be visualized "
99 - "with ui.perfetto.dev or Tensorboard "
100 - "(https://cloud.google.com/tpu/docs/pytorch-xla-performance-profiling-tpu-vm)."
101 - ),
102 - )
103 -
104 - parser = EngineArgs.add_cli_args(parser)
105 - return parser.parse_args()
106 -
107 -
108 -if __name__ == "__main__":
109 - args = parse_args()
110 - main(args)
examples/offline_inference/tpu.py deleted
-58
@@ -1,58 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -import argparse
5 -import os
6 -
7 -from vllm import LLM, SamplingParams
8 -
9 -prompts = [
10 - "A robot may not injure a human being",
11 - "It is only with the heart that one can see rightly;",
12 - "The greatest glory in living lies not in never falling,",
13 -]
14 -answers = [
15 - " or, through inaction, allow a human being to come to harm.",
16 - " what is essential is invisible to the eye.",
17 - " but in rising every time we fall.",
18 -]
19 -N = 1
20 -# Currently, top-p sampling is disabled. `top_p` should be 1.0.
21 -sampling_params = SamplingParams(temperature=0, top_p=1.0, n=N, max_tokens=16)
22 -
23 -
24 -def main():
25 - parser = argparse.ArgumentParser(description="TPU offline inference example")
26 - parser.add_argument("--use-spmd", action="store_true", help="Enable SPMD mode")
27 - args = parser.parse_args()
28 -
29 - llm_args = {
30 - "model": "Qwen/Qwen2-1.5B-Instruct",
31 - "max_num_batched_tokens": 64,
32 - "max_num_seqs": 4,
33 - "max_model_len": 128,
34 - }
35 - if args.use_spmd:
36 - os.environ["VLLM_XLA_USE_SPMD"] = "1"
37 - # Can only hardcode the number of chips for now.
38 - # calling xr.global_runtime_device_count() beforeing init SPMD env in
39 - # torch_xla will mess up the distributed env.
40 - llm_args["tensor_parallel_size"] = 8
41 - # Use Llama, for num_kv_heads = 8.
42 - llm_args["model"] = "meta-llama/Llama-3.1-8B-Instruct"
43 -
44 - # Set `enforce_eager=True` to avoid ahead-of-time compilation.
45 - # In real workloads, `enforce_eager` should be `False`.
46 - llm = LLM(**llm_args)
47 - outputs = llm.generate(prompts, sampling_params)
48 - print("-" * 50)
49 - for output, answer in zip(outputs, answers):
50 - prompt = output.prompt
51 - generated_text = output.outputs[0].text
52 - print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
53 - assert generated_text.startswith(answer)
54 - print("-" * 50)
55 -
56 -
57 -if __name__ == "__main__":
58 - main()