@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / e7596371a4
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 10, 15:15 UTC

[Refactor][TPU] Remove torch_xla path and use tpu-inference (#30808)

Signed-off-by: Wei-Yu Lin <weiyulin@google.com> Signed-off-by: weiyu <62784299+weiyu0824@users.noreply.github.com>

weiyu committed Jan 7, 2026 at 00:07 UTC e7596371a403903218ded3a9f446981fde5737f5
46 files changed +9 -6785
docs/design/moe_kernel_features.md
-1
@@ -92,7 +92,6 @@ To be used with a particular `FusedMoEPrepareAndFinalize` subclass, MoE kernels
92 | gpt oss triton | standard | N/A | N/A | <sup>5</sup> | Y | Y | [`triton_kernel_fused_experts`][vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe.triton_kernel_fused_experts],</br>[`OAITritonExperts`][vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe.OAITritonExperts] |
93 | marlin | standard,</br>batched | <sup>3</sup> / N/A | <sup>3</sup> / N/A | silu,</br>swigluoai | Y | Y | [`fused_marlin_moe`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.fused_marlin_moe],</br>[`MarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.MarlinExperts],</br>[`BatchedMarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.BatchedMarlinExperts] |
94 | trtllm | standard | mxfp4,</br>nvfp4 | G(16),G(32) | <sup>5</sup> | N | Y | [`TrtLlmGenExperts`][vllm.model_executor.layers.fused_moe.trtllm_moe.TrtLlmGenExperts] |
95 -| pallas | standard | N/A | N/A | silu | N | N | [`fused_moe`][vllm.model_executor.layers.fused_moe.moe_pallas.fused_moe] |
95 | iterative | standard | N/A | N/A | silu | N | N | [`fused_moe`][vllm.model_executor.layers.fused_moe.moe_torch_iterative.fused_moe] |
96 | rocm aiter moe | standard | fp8 | G(128),A,T | silu, gelu | Y | N | [`rocm_aiter_fused_experts`][vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe.rocm_aiter_fused_experts] |
97 | cpu_fused_moe | standard | N/A | N/A | silu | N | N | [`CPUFusedMOE`][vllm.model_executor.layers.fused_moe.cpu_fused_moe.CPUFusedMOE] |
tests/tpu/__init__.py
tests/tpu/lora/__init__.py
tests/tpu/lora/test_lora.py deleted
-139
@@ -1,139 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -import pytest
4 -from torch_xla._internal import tpu
5 -
6 -import vllm
7 -from vllm.lora.request import LoRARequest
8 -
9 -# This file contains tests to ensure that LoRA works correctly on the TPU
10 -# backend. We use a series of custom trained adapters for Qwen2.5-3B-Instruct
11 -# for this. The adapters are:
12 -# Username6568/Qwen2.5-3B-Instruct-1_plus_1_equals_x_adapter, where x ranges
13 -# from 1 to 4.
14 -
15 -# These adapters are trained using a standard huggingface peft training script,
16 -# where all the inputs are "What is 1+1? \n" and all the outputs are "x". We run
17 -# 100 training iterations with a training batch size of 100.
18 -
19 -
20 -def setup_vllm(num_loras: int, tp: int) -> vllm.LLM:
21 - return vllm.LLM(
22 - model="Qwen/Qwen2.5-3B-Instruct",
23 - max_model_len=256,
24 - max_num_seqs=8,
25 - tensor_parallel_size=tp,
26 - enable_lora=True,
27 - max_loras=num_loras,
28 - max_lora_rank=8,
29 - )
30 -
31 -
32 -TPU_TENSOR_PARALLEL_SIZES = (
33 - [1, tpu.num_available_chips()] if tpu.num_available_chips() > 1 else [1]
34 -)
35 -
36 -
37 -@pytest.mark.parametrize("tp", TPU_TENSOR_PARALLEL_SIZES)
38 -def test_single_lora(tp: int):
39 - """
40 - This test ensures we can run a single LoRA adapter on the TPU backend.
41 - We run "Username6568/Qwen2.5-3B-Instruct-1_plus_1_equals_1_adapter" which
42 - will force Qwen2.5-3B-Instruct to claim 1+1=1.
43 - """
44 -
45 - llm = setup_vllm(1, tp)
46 -
47 - prompt = "What is 1+1? \n"
48 -
49 - lora_request = LoRARequest(
50 - "lora_adapter_1",
51 - 1,
52 - "Username6568/Qwen2.5-3B-Instruct-1_plus_1_equals_1_adapter",
53 - )
54 - output = (
55 - llm.generate(
56 - prompt,
57 - sampling_params=vllm.SamplingParams(max_tokens=256, temperature=0),
58 - lora_request=lora_request,
59 - )[0]
60 - .outputs[0]
61 - .text
62 - )
63 -
64 - answer = output.strip()[0]
65 -
66 - assert answer.isdigit()
67 - assert int(answer) == 1
68 -
69 -
70 -@pytest.mark.parametrize("tp", TPU_TENSOR_PARALLEL_SIZES)
71 -def test_lora_hotswapping(tp: int):
72 - """
73 - This test ensures we can run multiple LoRA adapters on the TPU backend, even
74 - if we only have space to store 1.
75 -
76 - We run "Username6568/Qwen2.5-3B-Instruct-1_plus_1_equals_x_adapter" which
77 - will force Qwen2.5-3B-Instruct to claim 1+1=x, for a range of x.
78 - """
79 -
80 - lora_name_template = "Username6568/Qwen2.5-3B-Instruct-1_plus_1_equals_{}_adapter"
81 - lora_requests = [
82 - LoRARequest(f"lora_adapter_{i}", i, lora_name_template.format(i))
83 - for i in range(1, 5)
84 - ]
85 -
86 - llm = setup_vllm(1, tp)
87 -
88 - prompt = "What is 1+1? \n"
89 -
90 - for i, req in enumerate(lora_requests):
91 - output = (
92 - llm.generate(
93 - prompt,
94 - sampling_params=vllm.SamplingParams(max_tokens=256, temperature=0),
95 - lora_request=req,
96 - )[0]
97 - .outputs[0]
98 - .text
99 - )
100 - answer = output.strip()[0]
101 -
102 - assert answer.isdigit()
103 - assert int(answer) == i + 1
104 -
105 -
106 -@pytest.mark.parametrize("tp", TPU_TENSOR_PARALLEL_SIZES)
107 -def test_multi_lora(tp: int):
108 - """
109 - This test ensures we can run multiple LoRA adapters on the TPU backend, when
110 - we have enough space to store all of them.
111 -
112 - We run "Username6568/Qwen2.5-3B-Instruct-1_plus_1_equals_x_adapter" which
113 - will force Qwen2.5-3B-Instruct to claim 1+1=x, for a range of x.
114 - """
115 - lora_name_template = "Username6568/Qwen2.5-3B-Instruct-1_plus_1_equals_{}_adapter"
116 - lora_requests = [
117 - LoRARequest(f"lora_adapter_{i}", i, lora_name_template.format(i))
118 - for i in range(1, 5)
119 - ]
120 -
121 - llm = setup_vllm(4, tp)
122 -
123 - prompt = "What is 1+1? \n"
124 -
125 - for i, req in enumerate(lora_requests):
126 - output = (
127 - llm.generate(
128 - prompt,
129 - sampling_params=vllm.SamplingParams(max_tokens=256, temperature=0),
130 - lora_request=req,
131 - )[0]
132 - .outputs[0]
133 - .text
134 - )
135 -
136 - answer = output.strip()[0]
137 -
138 - assert answer.isdigit()
139 - assert int(output.strip()[0]) == i + 1
tests/tpu/test_compilation.py deleted
-86
@@ -1,86 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -import glob
5 -import os
6 -import tempfile
7 -
8 -import depyf
9 -
10 -
11 -def test_tpu_compilation():
12 - temp_dir = tempfile.mkdtemp()
13 - with depyf.prepare_debug(temp_dir):
14 - from vllm import LLM, SamplingParams
15 -
16 - prompts = [
17 - "A robot may not injure a human being",
18 - "It is only with the heart that one can see rightly;",
19 - "The greatest glory in living lies not in never falling,",
20 - ]
21 - answers = [
22 - " or, through inaction",
23 - " what is essential ",
24 - " but in rising ",
25 - ]
26 -
27 - # Currently, top-p sampling is disabled. `top_p` should be 1.0.
28 - N = 1
29 - sampling_params = SamplingParams(temperature=0.7, top_p=1.0, n=N, max_tokens=16)
30 -
31 - llm = LLM(
32 - model="Qwen/Qwen2-1.5B-Instruct",
33 - max_num_batched_tokens=256,
34 - max_model_len=256,
35 - max_num_seqs=32,
36 - enforce_eager=False,
37 - )
38 -
39 - outputs = llm.generate(prompts, sampling_params)
40 - for output, answer in zip(outputs, answers):
41 - prompt = output.prompt
42 - generated_text = output.outputs[0].text
43 - print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
44 - assert generated_text.startswith(answer)
45 -
46 - compiled_codes = sorted(
47 - glob.glob(os.path.join(temp_dir, "__transformed_code*for_forward.py"))
48 - )
49 -
50 - for i, compiled_code in enumerate(compiled_codes):
51 - print("{} file: {}".format(i + 1, compiled_code))
52 -
53 - # We should only trigger Dynamo compilation 2 times:
54 - # 1. Forward pass without kv_caches
55 - # 2. Forward pass with kv_caches
56 - # Check we have 2 compiled codes
57 - assert len(compiled_codes) == 2
58 -
59 - kv_cache_prefix = "kv_cache"
60 - attn_prefix = "ragged_paged_attention"
61 -
62 - def extract_compiled_index(s):
63 - parts = s.replace(".", "_").split("_")
64 - numbers = [int(part) for part in parts if part.isdigit()]
65 - return numbers[0]
66 -
67 - # Check all the compilations are as expected. The dump files include the
68 - # captured graph for the forward function of the nn.Module.
69 - compiled_fns = sorted(
70 - glob.glob(os.path.join(temp_dir, "__compiled_fn*Forward_graph*.py")),
71 - key=lambda s: extract_compiled_index(s),
72 - )
73 -
74 - for i, compiled_fn in enumerate(compiled_fns):
75 - print("{} file: {}".format(i + 1, compiled_fn))
76 -
77 - # The first compilation should not have any kv_caches
78 - with open(compiled_fns[0]) as f:
79 - content = f.read()
80 - assert kv_cache_prefix not in content
81 -
82 - # The second compilation should have kv_caches and the
83 - # ragged_paged_attention
84 - with open(compiled_fns[1]) as f:
85 - content = f.read()
86 - assert kv_cache_prefix in content and attn_prefix in content
tests/tpu/test_custom_dispatcher.py deleted
-34
@@ -1,34 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -import pytest
5 -
6 -from vllm.config import CompilationMode
7 -
8 -from ..utils import compare_two_settings
9 -
10 -# --enforce-eager on TPU causes graph compilation
11 -# this times out default Health Check in the MQLLMEngine,
12 -# so we set the timeout here to 30s
13 -
14 -
15 -def test_custom_dispatcher(monkeypatch: pytest.MonkeyPatch):
16 - with monkeypatch.context() as m:
17 - m.setenv("VLLM_RPC_TIMEOUT", "30000")
18 - compare_two_settings(
19 - "Qwen/Qwen2.5-1.5B-Instruct",
20 - arg1=[
21 - "--max-model-len=256",
22 - "--max-num-seqs=32",
23 - "--enforce-eager",
24 - f"-O{CompilationMode.DYNAMO_TRACE_ONCE}",
25 - ],
26 - arg2=[
27 - "--max-model-len=256",
28 - "--max-num-seqs=32",
29 - "--enforce-eager",
30 - f"-O{CompilationMode.STOCK_TORCH_COMPILE}",
31 - ],
32 - env1={},
33 - env2={},
34 - )
tests/tpu/test_moe_pallas.py deleted
-88
@@ -1,88 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""Tests for the Pallas MOE implementation.
4 -
5 -Run `pytest tests/kernels/moe/test_moe_pallas.py`.
6 -"""
7 -
8 -import pytest
9 -import torch
10 -import torch_xla
11 -
12 -from vllm.model_executor.layers.fused_moe.moe_pallas import fused_moe as pallas_moe
13 -from vllm.model_executor.layers.fused_moe.moe_torch_iterative import (
14 - fused_moe as torch_moe,
15 -)
16 -from vllm.platforms import current_platform
17 -
18 -if not current_platform.is_tpu():
19 - pytest.skip("This test needs a TPU.", allow_module_level=True)
20 -
21 -NUM_EXPERTS = [8, 64]
22 -EP_SIZE = [1]
23 -TOP_KS = [2, 6]
24 -
25 -
26 -# The Pallas GMM kernel requires num_tokens * topk to be a multiple of 16
27 -@pytest.mark.parametrize("m", [8, 16, 64, 2048])
28 -@pytest.mark.parametrize("n", [128, 1024, 2048])
29 -@pytest.mark.parametrize("k", [128, 511, 1024])
30 -@pytest.mark.parametrize("e", NUM_EXPERTS)
31 -@pytest.mark.parametrize("topk", TOP_KS)
32 -@pytest.mark.parametrize("ep_size", EP_SIZE)
33 -@pytest.mark.parametrize("dtype", [torch.bfloat16])
34 -def test_pallas_moe(
35 - m: int,
36 - n: int,
37 - k: int,
38 - e: int,
39 - topk: int,
40 - ep_size: int,
41 - dtype: torch.dtype,
42 -):
43 - import torch_xla.core.xla_model as xm
44 -
45 - with torch.device(xm.xla_device()):
46 - a = torch.randn((m, k), dtype=dtype) / 10
47 - w1 = torch.randn((e, 2 * n, k), dtype=dtype) / 10
48 - w2 = torch.randn((e, k, n), dtype=dtype) / 10
49 -
50 - score = torch.randn((m, e), dtype=dtype)
51 -
52 - # TODO: Support ep
53 - if ep_size > 1:
54 - pytest.skip("No support for ep_size > 1 yet")
55 - else:
56 - e_map = None
57 -
58 - # Run both implementations
59 - torch_output = torch_moe(
60 - hidden_states=a,
61 - w1=w1,
62 - w2=w2,
63 - gating_output=score,
64 - topk=topk,
65 - global_num_experts=e,
66 - expert_map=e_map,
67 - renormalize=False,
68 - )
69 -
70 - pallas_output = pallas_moe(
71 - hidden_states=a,
72 - w1=w1,
73 - w2=w2,
74 - gating_output=score,
75 - topk=topk,
76 - global_num_experts=e,
77 - expert_map=e_map,
78 - renormalize=False,
79 - )
80 - torch_xla.sync(wait=False)
81 -
82 - # Compare outputs
83 - torch.testing.assert_close(
84 - pallas_output.cpu(),
85 - torch_output.cpu(),
86 - atol=2e-2,
87 - rtol=0,
88 - )
tests/tpu/test_quantization_accuracy.py deleted
-52
@@ -1,52 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -from dataclasses import dataclass
5 -
6 -import lm_eval
7 -import pytest
8 -
9 -TASK = "gsm8k"
10 -FILTER = "exact_match,strict-match"
11 -RTOL = 0.03
12 -
13 -
14 -@dataclass
15 -class GSM8KAccuracyTestConfig:
16 - model_name: str
17 - expected_value: float
18 -
19 - def get_model_args(self) -> str:
20 - return f"pretrained={self.model_name},max_model_len=4096,max_num_seqs=32"
21 -
22 -
23 -# NOTE: Accuracy scores measured on GPUs.
24 -ACCURACY_CONFIGS = [
25 - GSM8KAccuracyTestConfig(
26 - model_name="neuralmagic/Meta-Llama-3.1-8B-Instruct-quantized.w8a8",
27 - expected_value=0.76,
28 - ), # no bias
29 - # NOTE(rob): We cannot re-initialize vLLM in the same process for TPU,
30 - # so only one of these tests can run in a single call to pytest. As
31 - # a follow-up, move this into the LM-EVAL section of the CI.
32 - # GSM8KAccuracyTestConfig(
33 - # model_name="neuralmagic/Qwen2-7B-Instruct-quantized.w8a8",
34 - # expected_value=0.66), # bias in QKV layers
35 -]
36 -
37 -
38 -@pytest.mark.parametrize("config", ACCURACY_CONFIGS)
39 -def test_gsm8k_correctness(config: GSM8KAccuracyTestConfig):
40 - results = lm_eval.simple_evaluate(
41 - model="vllm",
42 - model_args=config.get_model_args(),
43 - tasks="gsm8k",
44 - batch_size="auto",
45 - )
46 -
47 - EXPECTED_VALUE = config.expected_value
48 - measured_value = results["results"][TASK][FILTER]
49 - assert (
50 - measured_value - RTOL < EXPECTED_VALUE
51 - and measured_value + RTOL > EXPECTED_VALUE
52 - ), f"Expected: {EXPECTED_VALUE} | Measured: {measured_value}"
tests/v1/tpu/__init__.py
tests/v1/tpu/test_basic.py deleted
-177
@@ -1,177 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""A basic correctness check for TPUs
4 -
5 -Run `pytest tests/v1/tpu/test_basic.py`.
6 -"""
7 -
8 -from typing import TYPE_CHECKING
9 -
10 -import pytest
11 -from torch_xla._internal import tpu
12 -
13 -from vllm.platforms import current_platform
14 -
15 -if TYPE_CHECKING:
16 - from tests.conftest import VllmRunner
17 -else:
18 - VllmRunner = object
19 -
20 -MODELS = [
21 - "Qwen/Qwen2.5-1.5B-Instruct",
22 - # TODO: Enable this model when fixed.
23 - # "Qwen/Qwen1.5-MoE-A2.7B",
24 - # TODO: Enable this models with v6e
25 - # "Qwen/Qwen2-7B-Instruct",
26 - # "meta-llama/Llama-3.1-8B",
27 -]
28 -
29 -TENSOR_PARALLEL_SIZES = [1]
30 -MAX_NUM_REQS = [16, 1024]
31 -
32 -# TODO: Enable when CI/CD will have a multi-tpu instance
33 -# TENSOR_PARALLEL_SIZES = [1, 4]
34 -
35 -
36 -@pytest.mark.skipif(
37 - not current_platform.is_tpu(), reason="This is a basic test for TPU only"
38 -)
39 -@pytest.mark.parametrize("model", MODELS)
40 -@pytest.mark.parametrize("max_tokens", [5])
41 -@pytest.mark.parametrize("tensor_parallel_size", TENSOR_PARALLEL_SIZES)
42 -@pytest.mark.parametrize("max_num_seqs", MAX_NUM_REQS)
43 -def test_basic(
44 - vllm_runner: type[VllmRunner],
45 - model: str,
46 - max_tokens: int,
47 - tensor_parallel_size: int,
48 - max_num_seqs: int,
49 -) -> None:
50 - prompt = (
51 - "The next numbers of the sequence "
52 - + ", ".join(str(i) for i in range(1024))
53 - + " are:"
54 - )
55 - example_prompts = [prompt]
56 -
57 - with vllm_runner(
58 - model,
59 - # Note: max_num_batched_tokens == 1024 is needed here to
60 - # actually test chunked prompt
61 - max_num_batched_tokens=1024,
62 - max_model_len=8192,
63 - gpu_memory_utilization=0.7,
64 - max_num_seqs=max_num_seqs,
65 - tensor_parallel_size=tensor_parallel_size,
66 - ) as vllm_model:
67 - vllm_outputs = vllm_model.generate_greedy(example_prompts, max_tokens)
68 - output = vllm_outputs[0][1]
69 -
70 - assert "1024" in output or "0, 1" in output
71 -
72 -
73 -@pytest.mark.skip(reason="Temporarily disabled due to timeout")
74 -@pytest.mark.skipif(
75 - not current_platform.is_tpu(), reason="This is a basic test for TPU only"
76 -)
77 -@pytest.mark.parametrize("max_tokens", [8])
78 -@pytest.mark.parametrize("max_num_seqs", [16])
79 -def test_phi3(
80 - vllm_runner: type[VllmRunner],
81 - max_tokens: int,
82 - max_num_seqs: int,
83 -) -> None:
84 - prompts = [
85 - "A robot may not injure a human being",
86 - "It is only with the heart that one can see rightly;",
87 - "The greatest glory in living lies not in never falling,",
88 - ]
89 - answers = [
90 - " or, by violating privacy",
91 - " what is essential is love.",
92 - " but in rising every time we fall.",
93 - ]
94 - # test head dim = 96
95 - model = "microsoft/Phi-3-mini-128k-instruct"
96 -
97 - with vllm_runner(
98 - model, max_num_batched_tokens=256, max_num_seqs=max_num_seqs
99 - ) as vllm_model:
100 - vllm_outputs = vllm_model.generate_greedy(prompts, max_tokens)
101 - # vllm_outputs is a list of tuples whose first element is the token id
102 - # and the second element is the output (including the prompt).
103 - for output, answer in zip(vllm_outputs, answers):
104 - generated_text = output[1]
105 - assert answer in generated_text
106 -
107 -
108 -TP_SIZE_8 = 8
109 -
110 -
111 -@pytest.mark.skipif(not current_platform.is_tpu(), reason="This is a test for TPU only")
112 -@pytest.mark.skipif(
113 - tpu.num_available_chips() < TP_SIZE_8,
114 - reason=f"This test requires {TP_SIZE_8} TPU chips.",
115 -)
116 -def test_gemma3_27b_with_text_input_and_tp(
117 - vllm_runner: type[VllmRunner],
118 -) -> None:
119 - model = "google/gemma-3-27b-it"
120 - max_tokens = 16
121 - tensor_parallel_size = TP_SIZE_8
122 - max_num_seqs = 4
123 - prompts = [
124 - "A robot may not injure a human being",
125 - "It is only with the heart that one can see rightly;",
126 - "The greatest glory in living lies not in never falling,",
127 - ]
128 - answers = [
129 - " or, through inaction, allow a human being to come to harm.",
130 - " what is essential is invisible to the eye.",
131 - " but in rising every time we fall.",
132 - ]
133 -
134 - with vllm_runner(
135 - model,
136 - max_num_batched_tokens=256,
137 - max_num_seqs=max_num_seqs,
138 - tensor_parallel_size=tensor_parallel_size,
139 - ) as vllm_model:
140 - vllm_outputs = vllm_model.generate_greedy(prompts, max_tokens)
141 - # vllm_outputs is a list of tuples whose first element is the token id
142 - # and the second element is the output (including the prompt).
143 - for output, answer in zip(vllm_outputs, answers):
144 - generated_text = output[1]
145 - assert answer in generated_text
146 -
147 -
148 -@pytest.mark.skipif(
149 - not current_platform.is_tpu(), reason="This is a basic test for TPU only"
150 -)
151 -def test_w8a8_quantization(
152 - vllm_runner: type[VllmRunner],
153 -) -> None:
154 - model = "neuralmagic/Meta-Llama-3.1-8B-Instruct-quantized.w8a8"
155 - max_tokens = 5
156 - tensor_parallel_size = 1
157 - max_num_seqs = 4
158 -
159 - prompt = (
160 - "The next numbers of the sequence "
161 - + ", ".join(str(i) for i in range(1024))
162 - + " are:"
163 - )
164 - example_prompts = [prompt]
165 -
166 - with vllm_runner(
167 - model,
168 - max_num_batched_tokens=64,
169 - max_model_len=4096,
170 - gpu_memory_utilization=0.7,
171 - max_num_seqs=max_num_seqs,
172 - tensor_parallel_size=tensor_parallel_size,
173 - ) as vllm_model:
174 - vllm_outputs = vllm_model.generate_greedy(example_prompts, max_tokens)
175 - output = vllm_outputs[0][1]
176 -
177 - assert "1024" in output or "0, 1" in output
tests/v1/tpu/test_kv_cache_update_kernel.py deleted
-78
@@ -1,78 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -import numpy as np
5 -import pytest
6 -import torch
7 -import torch_xla
8 -
9 -import vllm.v1.attention.backends.pallas # noqa: F401
10 -from vllm.platforms import current_platform
11 -
12 -
13 -@pytest.mark.skipif(not current_platform.is_tpu(), reason="This is a test for TPU only")
14 -@pytest.mark.parametrize("page_size", [32, 33])
15 -@pytest.mark.parametrize("combined_kv_head_num", [2, 16])
16 -@pytest.mark.parametrize("head_dim", [128, 256])
17 -@pytest.mark.parametrize("num_slices_per_block", [4, 8])
18 -def test_kv_cache_update_kernel(
19 - page_size: int, combined_kv_head_num: int, head_dim: int, num_slices_per_block: int
20 -):
21 - page_num = 1000
22 - padded_num_tokens = 128
23 - kv_cache_cpu = torch.zeros(
24 - (page_num * page_size, combined_kv_head_num, head_dim),
25 - dtype=torch.bfloat16,
26 - device="cpu",
27 - )
28 - kv_cache_xla = kv_cache_cpu.to(torch_xla.device())
29 - new_kv_cpu = torch.randn(
30 - (padded_num_tokens, combined_kv_head_num, head_dim),
31 - dtype=torch.bfloat16,
32 - device="cpu",
33 - )
34 - new_kv_xla = new_kv_cpu.to(torch_xla.device())
35 - slice_lens = np.array([7, page_size, page_size, 1, 1, 1, 9], dtype=np.int32)
36 - num_kv_update_slices = len(slice_lens)
37 - kv_cache_start_indices = np.array(
38 - [
39 - page_size * 2 - 7,
40 - page_size * 2,
41 - page_size * 3,
42 - page_size * 4 + 6,
43 - page_size * 5 + 7,
44 - page_size * 6 + 8,
45 - page_size * 15 + 3,
46 - ],
47 - dtype=np.int32,
48 - )
49 - new_kv_cache_indices = np.concatenate(
50 - [np.array([0], dtype=np.int32), np.cumsum(slice_lens[:-1])]
51 - )
52 - slot_mapping = np.stack(
53 - [kv_cache_start_indices, new_kv_cache_indices, slice_lens], axis=1
54 - )
55 - slot_mapping = np.transpose(slot_mapping)
56 - slot_mapping_cpu = torch.tensor(slot_mapping, device="cpu", dtype=torch.int32)
57 - slot_mapping_xla = slot_mapping_cpu.to(torch_xla.device())
58 - num_kv_update_slices_xla = torch.tensor(
59 - [num_kv_update_slices], device=torch_xla.device(), dtype=torch.int32
60 - )
61 - torch_xla.sync()
62 -
63 - torch.ops.xla.dynamo_set_buffer_donor_(kv_cache_xla, True)
64 - new_kv_cache_xla = torch.ops.xla.kv_cache_update_op(
65 - new_kv_xla,
66 - slot_mapping_xla,
67 - kv_cache_xla,
68 - num_kv_update_slices_xla,
69 - page_size,
70 - num_slices_per_block,
71 - )
72 - kv_cache_xla.copy_(new_kv_cache_xla)
73 - torch_xla.sync()
74 -
75 - for ni, ci, sl in zip(new_kv_cache_indices, kv_cache_start_indices, slice_lens):
76 - kv_cache_cpu[ci : ci + sl, :, :] = new_kv_cpu[ni : ni + sl, :, :]
77 -
78 - assert torch.allclose(kv_cache_xla.cpu(), kv_cache_cpu, atol=1e-4, rtol=1e-4)
tests/v1/tpu/test_mha_attn.py deleted
-95
@@ -1,95 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""
4 -Test:
5 -
6 -* Tests for MMEncoderAttention layer
7 -"""
8 -
9 -import pytest
10 -import torch
11 -import torch_xla
12 -import torch_xla.core
13 -import torch_xla.core.xla_model
14 -
15 -from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
16 -from vllm.attention.selector import _cached_get_attn_backend
17 -from vllm.platforms import current_platform
18 -from vllm.utils.torch_utils import set_random_seed
19 -
20 -
21 -@pytest.fixture(autouse=True)
22 -def clear_cache():
23 - """Clear lru cache to ensure each test case runs without caching."""
24 - _cached_get_attn_backend.cache_clear()
25 -
26 -
27 -def ref_attention(
28 - query: torch.Tensor,
29 - key: torch.Tensor,
30 - value: torch.Tensor,
31 - scale: float,
32 -) -> torch.Tensor:
33 - """
34 - Native implementation of scaled dot product attention without mask:
35 - - query, key, value: [batch_size, seq_len, num_heads, head_size]
36 - - attn_mask: [batch_size, seq_len, seq_len]
37 - """
38 - query, key, value = (x.transpose(1, 2) for x in (query, key, value))
39 - attn_weights = scale * torch.matmul(query, key.transpose(2, 3))
40 - attn_weights = torch.softmax(attn_weights, dim=-1).to(value.dtype)
41 - out = torch.matmul(attn_weights, value).transpose(1, 2)
42 - return out
43 -
44 -
45 -BATCH_SIZES = [1, 16]
46 -SEQ_LENS = [1]
47 -NUM_HEADS = [1, 16]
48 -NUM_KV_HEADS = [1]
49 -HEAD_SIZES = [64, 80]
50 -
51 -
52 -@pytest.mark.skipif(not current_platform.is_tpu(), reason="This test needs a TPU")
53 -@pytest.mark.parametrize("batch_size", BATCH_SIZES)
54 -@pytest.mark.parametrize("seq_len", SEQ_LENS)
55 -@pytest.mark.parametrize("num_heads", NUM_HEADS)
56 -@pytest.mark.parametrize("num_kv_heads", NUM_KV_HEADS)
57 -@pytest.mark.parametrize("head_size", HEAD_SIZES)
58 -@pytest.mark.parametrize("device", [torch_xla.core.xla_model.xla_device()])
59 -def test_mha_attn_forward(
60 - batch_size: int,
61 - seq_len: int,
62 - num_heads: int,
63 - num_kv_heads: int,
64 - head_size: int,
65 - device: str,
66 -):
67 - set_random_seed(0)
68 - # These are expected to be f32
69 - q = torch.randn(batch_size, seq_len, num_heads * head_size, device=device)
70 - k = torch.randn(batch_size, seq_len, num_kv_heads * head_size, device=device)
71 - v = torch.randn(batch_size, seq_len, num_kv_heads * head_size, device=device)
72 - scale = 1.0 / head_size**0.5
73 - attn = MMEncoderAttention(
74 - num_heads, head_size, scale=scale, num_kv_heads=num_kv_heads
75 - )
76 - output = attn(q, k, v)
77 -
78 - assert num_heads % num_kv_heads == 0
79 - num_queries_per_kv = num_heads // num_kv_heads
80 -
81 - q = q.reshape(batch_size, seq_len, num_heads, head_size)
82 - k = k.reshape(batch_size, seq_len, num_kv_heads, head_size)
83 - v = v.reshape(batch_size, seq_len, num_kv_heads, head_size)
84 - if num_queries_per_kv > 1:
85 - k = torch.repeat_interleave(k, num_queries_per_kv, dim=2)
86 - v = torch.repeat_interleave(v, num_queries_per_kv, dim=2)
87 -
88 - ref_output = ref_attention(
89 - q,
90 - k,
91 - v,
92 - scale=scale,
93 - ).reshape(batch_size, seq_len, num_heads * head_size)
94 - # torch_xla flash_attn kernel is less accurate but much faster
95 - torch.testing.assert_close(output, ref_output, atol=1e-2, rtol=1e-3)
tests/v1/tpu/test_multimodal.py deleted
-76
@@ -1,76 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -import openai
5 -import pytest
6 -
7 -from vllm.multimodal.utils import encode_image_url
8 -from vllm.platforms import current_platform
9 -
10 -from ...entrypoints.openai.test_vision import TEST_IMAGE_ASSETS
11 -from ...utils import RemoteOpenAIServer
12 -
13 -
14 -@pytest.fixture(scope="session")
15 -def url_encoded_image(local_asset_server) -> dict[str, str]:
16 - return {
17 - image_asset: encode_image_url(local_asset_server.get_image_asset(image_asset))
18 - for image_asset in TEST_IMAGE_ASSETS
19 - }
20 -
21 -
22 -@pytest.mark.asyncio
23 -@pytest.mark.skipif(not current_platform.is_tpu(), reason="This test needs a TPU")
24 -@pytest.mark.parametrize("model_name", ["llava-hf/llava-1.5-7b-hf"])
25 -async def test_basic_vision(model_name: str, url_encoded_image: dict[str, str]):
26 - pytest.skip("Skip this test until it's fixed.")
27 -
28 - def whats_in_this_image_msg(url):
29 - return [
30 - {
31 - "role": "user",
32 - "content": [
33 - {"type": "text", "text": "What's in this image?"},
34 - {"type": "image_url", "image_url": {"url": url}},
35 - ],
36 - }
37 - ]
38 -
39 - server_args = [
40 - "--max-model-len",
41 - "1024",
42 - "--max-num-seqs",
43 - "16",
44 - "--gpu-memory-utilization",
45 - "0.95",
46 - "--trust-remote-code",
47 - "--max-num-batched-tokens",
48 - "576",
49 - # NOTE: max-num-batched-tokens>=mm_item_size
50 - "--disable_chunked_mm_input",
51 - ]
52 -
53 - # Server will pre-compile on first startup (takes a long time).
54 - with RemoteOpenAIServer(
55 - model_name, server_args, max_wait_seconds=600
56 - ) as remote_server:
57 - client: openai.AsyncOpenAI = remote_server.get_async_client()
58 -
59 - # Other requests now should be much faster
60 - for image_url in TEST_IMAGE_ASSETS:
61 - image_url = url_encoded_image[image_url]
62 - chat_completion_from_url = await client.chat.completions.create(
63 - model=model_name,
64 - messages=whats_in_this_image_msg(image_url),
65 - max_completion_tokens=24,
66 - temperature=0.0,
67 - )
68 - result = chat_completion_from_url
69 - assert result
70 - choice = result.choices[0]
71 - assert choice.finish_reason == "length"
72 -
73 - message = choice.message
74 - message = result.choices[0].message
75 - assert message.content is not None and len(message.content) >= 10
76 - assert message.role == "assistant"
tests/v1/tpu/test_pallas.py deleted
-100
@@ -1,100 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from unittest.mock import ANY, patch
4 -
5 -import torch
6 -
7 -from vllm.attention.backends.abstract import AttentionType
8 -from vllm.v1.attention.backends.pallas import PallasAttentionBackendImpl, PallasMetadata
9 -
10 -
11 -def test_ragged_paged_attention():
12 - # We verify that the kernel inputs such as sliding_window, etc. are passed
13 - # in from the model correctly.
14 - # The correctness of the paged attention kernel is tested in the kernel
15 - # library.
16 - num_heads = 4
17 - head_size = 128
18 - scale = 1.0
19 - num_kv_heads = 4
20 - sliding_window = 128
21 - logits_soft_cap = 50.0
22 - attn_impl = PallasAttentionBackendImpl(
23 - num_heads=num_heads,
24 - head_size=head_size,
25 - scale=scale,
26 - num_kv_heads=num_kv_heads,
27 - alibi_slopes=None,
28 - sliding_window=sliding_window,
29 - kv_cache_dtype="auto",
30 - logits_soft_cap=logits_soft_cap,
31 - attn_type=AttentionType.DECODER,
32 - )
33 -
34 - class FakeAttentionLayer:
35 - _q_scale_float: float
36 - _k_scale_float: float
37 - _v_scale_float: float
38 -
39 - layer = FakeAttentionLayer()
40 - layer._q_scale_float = 1.0
41 - layer._k_scale_float = 1.0
42 - layer._v_scale_float = 1.0
43 -
44 - num_tokens = 16
45 - num_blocks = 1024
46 - block_size = 16
47 - query = torch.zeros(num_tokens, num_heads * head_size)
48 - key = torch.zeros(num_tokens, num_kv_heads * head_size)
49 - value = torch.zeros(num_tokens, num_kv_heads * head_size)
50 - kv_cache = torch.zeros(num_blocks, block_size, num_kv_heads * 2, head_size)
51 - slot_mapping = torch.zeros((3, num_tokens), dtype=torch.int64)
52 - max_num_reqs = 8
53 - max_num_blocks_per_req = 8
54 - num_kv_update_slices = torch.tensor([num_tokens], dtype=torch.int32)
55 - block_tables = torch.zeros(
56 - (max_num_reqs, max_num_blocks_per_req), dtype=torch.int32
57 - )
58 - context_lens = torch.ones((max_num_reqs,), dtype=torch.int32)
59 - query_lens = [1] * max_num_reqs
60 - query_start_loc = torch.cumsum(
61 - torch.tensor([0] + query_lens, dtype=torch.int32), dim=0, dtype=torch.int32
62 - )
63 - num_seqs = torch.tensor([max_num_reqs], dtype=torch.int32)
64 - attn_metadata = PallasMetadata(
65 - slot_mapping=slot_mapping,
66 - block_tables=block_tables,
67 - context_lens=context_lens,
68 - query_start_loc=query_start_loc,
69 - num_seqs=num_seqs,
70 - num_kv_update_slices=num_kv_update_slices,
71 - num_slices_per_kv_cache_update_block=8,
72 - )
73 -
74 - with patch("torch.ops.xla.ragged_paged_attention") as mock_ragged_paged_attention:
75 - attn_impl.forward(
76 - layer=layer,
77 - query=query,
78 - key=key,
79 - value=value,
80 - kv_cache=kv_cache,
81 - attn_metadata=attn_metadata,
82 - )
83 -
84 - mock_ragged_paged_attention.assert_called_once_with(
85 - ANY, # query
86 - ANY, # kv_cache
87 - ANY, # context_lens
88 - ANY, # block_tables
89 - ANY, # query_start_loc
90 - ANY, # num_seqs
91 - num_kv_pages_per_block=None,
92 - num_queries_per_block=None,
93 - vmem_limit_bytes=None,
94 - use_kernel=True,
95 - sm_scale=scale,
96 - sliding_window=sliding_window,
97 - soft_cap=logits_soft_cap,
98 - k_scale=1.0,
99 - v_scale=1.0,
100 - )
tests/v1/tpu/test_perf.py deleted
-150
@@ -1,150 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""A basic performance regression test for TPUs
4 -
5 -Run `pytest tests/v1/tpu/test_perf.py`.
6 -"""
7 -
8 -import time
9 -from dataclasses import dataclass
10 -from typing import TYPE_CHECKING
11 -
12 -import numpy as np
13 -import pytest
14 -
15 -from vllm.platforms import current_platform
16 -from vllm.sampling_params import SamplingParams
17 -from vllm.tokenizers import get_tokenizer
18 -
19 -if TYPE_CHECKING:
20 - from tests.conftest import VllmRunner
21 -else:
22 - VllmRunner = object
23 -
24 -
25 -@dataclass
26 -class TestParams:
27 - model: str
28 - num_prompts: int
29 - prefix_len: int
30 - decode_len: int
31 - expected_avg_time: float
32 - err_tol: float
33 -
34 -
35 -TEST_PARAMS = [
36 - # TODO: Cannot run a series of tests because:
37 - # RuntimeError: Bad StatusOr access: UNKNOWN: TPU initialization failed:
38 - # open(/dev/vfio/0): Device or resource busy: Device or resource busy;
39 - # Couldn't open iommu group /dev/vfio/0
40 - # => Investigate
41 - # TestParams(
42 - # model="Qwen/Qwen2.5-1.5B-Instruct",
43 - # num_prompts=1,
44 - # prefix_len=10,
45 - # decode_len=5,
46 - # expected_avg_time=0.03,
47 - # err_tol=0.01,
48 - # ),
49 - # TestParams(
50 - # model="Qwen/Qwen2.5-1.5B-Instruct",
51 - # num_prompts=10,
52 - # prefix_len=100,
53 - # decode_len=50,
54 - # expected_avg_time=0.234,
55 - # err_tol=0.020,
56 - # ),
57 - TestParams(
58 - model="Qwen/Qwen2.5-1.5B-Instruct",
59 - num_prompts=64,
60 - prefix_len=500,
61 - decode_len=50,
62 - # commit id: ccb246776d93ef105904a8ec015b3587240a1183
63 - # tpu: v5lite (old vllm CI/CD)
64 - # expected_avg_time=1.4,
65 - # err_tol=0.30,
66 - # (This is the active CI/CD instance)
67 - # commit id: ccb246776d93ef105904a8ec015b3587240a1183
68 - # tpu: v6e (current vllm CI/CD)
69 - expected_avg_time=1.7, # measured with VLLM_XLA_CACHE_PATH=
70 - err_tol=0.20,
71 - ),
72 -]
73 -
74 -NUM_WARMUPS = 5
75 -NUM_RUNS = 10
76 -
77 -MAX_MODEL_LEN = 1024
78 -MAX_NUM_SEQS = 32
79 -GPU_UTIL = 0.9
80 -
81 -
82 -@pytest.mark.skipif(
83 - not current_platform.is_tpu(),
84 - reason="This is a basic performance test for TPU only",
85 -)
86 -@pytest.mark.parametrize("params", TEST_PARAMS)
87 -def test_perf(
88 - vllm_runner: type[VllmRunner],
89 - params: TestParams,
90 -) -> None:
91 - tokenizer = get_tokenizer(
92 - params.model, tokenizer_mode="auto", trust_remote_code=True
93 - )
94 -
95 - prompts = []
96 - for i in range(params.num_prompts):
97 - prefix_token_ids = np.random.randint(
98 - 0, tokenizer.vocab_size, size=params.prefix_len
99 - ).tolist()
100 - prompt = tokenizer.decode(prefix_token_ids)
101 - prompts.append(prompt)
102 -
103 - print(
104 - "-- Running: num_prompts = {} prefix_len = {} decode_len = {}".format(
105 - len(prompts), params.prefix_len, params.decode_len
106 - )
107 - )
108 -
109 - sampling_params = SamplingParams(
110 - max_tokens=params.decode_len, temperature=1.0, min_p=0.0
111 - )
112 -
113 - with vllm_runner(
114 - params.model,
115 - max_num_batched_tokens=MAX_MODEL_LEN,
116 - max_model_len=MAX_MODEL_LEN,
117 - max_num_seqs=MAX_NUM_SEQS,
118 - gpu_memory_utilization=GPU_UTIL,
119 - enforce_eager=False,
120 - tensor_parallel_size=1,
121 - ) as vllm_model:
122 - print(" -- Warmup / Compile")
123 - for i in range(NUM_WARMUPS):
124 - _ = vllm_model.generate(prompts, sampling_params)
125 -
126 - print(" -- Benchmarking... ")
127 - times = []
128 - for i in range(NUM_RUNS):
129 - start_time = time.time()
130 - _ = vllm_model.generate(prompts, sampling_params)
131 - times.append(time.time() - start_time)
132 -
133 - avg_time = sum(times) / len(times)
134 -
135 - print(" -- avg_time = {}".format(avg_time))
136 - print(
137 - " -- expected_avg_time = {} with err_tol = {}".format(
138 - params.expected_avg_time, params.err_tol
139 - )
140 - )
141 - diff = avg_time - params.expected_avg_time
142 - ok = diff < params.err_tol
143 - if diff < -params.err_tol:
144 - print(
145 - " !! WARNING !! Performance has improved by {}, "
146 - "it may be necessary to fine-tune the "
147 - "expected_avg_time = {}".format(-diff, params.expected_avg_time)
148 - )
149 -
150 - assert ok, " !! ERROR !! Regression detected"
tests/v1/tpu/test_sampler.py deleted
-105
@@ -1,105 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -import random
4 -
5 -import pytest
6 -
7 -from vllm import LLM
8 -from vllm.platforms import current_platform
9 -from vllm.sampling_params import SamplingParams
10 -
11 -
12 -@pytest.mark.parametrize("model_name", ["Qwen/Qwen2.5-1.5B-Instruct"])
13 -@pytest.mark.skipif(not current_platform.is_tpu(), reason="This test needs a TPU")
14 -def test_sampler_different(model_name: str):
15 - """
16 - Test significantly different sampling params to assert the model produces
17 - different results.
18 - """
19 - llm = LLM(
20 - model_name,
21 - enforce_eager=False,
22 - max_num_seqs=1,
23 - max_model_len=512,
24 - max_num_batched_tokens=256,
25 - )
26 - prompts = ["Write a short story about a robot that dreams for the first time."]
27 - sampling_params = SamplingParams(temperature=0.9, min_p=0.2, max_tokens=64)
28 - output = llm.generate(prompts, sampling_params)
29 -
30 - sampling_params = SamplingParams(temperature=0.1, min_p=0.8, max_tokens=64)
31 - output2 = llm.generate(prompts, sampling_params)
32 - assert output[0].outputs[0].text != output2[0].outputs[0].text
33 -
34 - with pytest.raises(ValueError):
35 - # Unsupported `seed` param.
36 - sampling_params = SamplingParams(temperature=0.3, seed=42)
37 - output2 = llm.generate(prompts, sampling_params)
38 -
39 - # Batch-case with TopK/P
40 - for B in [4, 16]:
41 - p = prompts * B
42 - sampling_params = [
43 - SamplingParams(
44 - temperature=0.1,
45 - min_p=0.8,
46 - max_tokens=64,
47 - # Vary number of ks
48 - top_k=random.randint(4, 12),
49 - top_p=random.random(),
50 - )
51 - for _ in range(B)
52 - ]
53 - # Make sure first two reqs have the same K/P
54 - sampling_params[0] = sampling_params[1]
55 - output = llm.generate(p, sampling_params)
56 - # There are natural numerical instabilities that make it difficult
57 - # to have deterministic results over many tokens, tests the first ~20
58 - # tokens match.
59 - assert output[0].outputs[0].text[:20] == output[1].outputs[0].text[:20]
60 -
61 -
62 -@pytest.mark.parametrize("model_name", ["Qwen/Qwen2.5-1.5B-Instruct"])
63 -# TODO TPU will appear busy if we fan-out test params here
64 -@pytest.mark.parametrize("n_prompts", [1])
65 -@pytest.mark.skipif(not current_platform.is_tpu(), reason="This test needs a TPU")
66 -def test_logprobs(model_name: str, n_prompts: int):
67 - """
68 - Request top logprobs with different sampling settings and check
69 - that results contains the requested number, ordered ascendingly.
70 - """
71 -
72 - def check_num_logprobs(logprobs, expected_num: int):
73 - for step in logprobs:
74 - prev_logp = 1.0
75 - # order by rank
76 - sorted_step = dict(sorted(step.items(), key=lambda item: item[1].rank))
77 -
78 - # Can contain the sampled token
79 - assert len(step) == expected_num or len(step) == expected_num + 1
80 - # Check results are ordered by prob value
81 - for rankno, (tid, logp) in enumerate(sorted_step.items()):
82 - assert logp.logprob <= prev_logp
83 - prev_logp = logp.logprob
84 - assert logp.rank == rankno + 1
85 -
86 - llm = LLM(
87 - model_name,
88 - enforce_eager=False,
89 - max_num_seqs=1,
90 - max_model_len=128,
91 - max_num_batched_tokens=128,
92 - )
93 - prompts = [
94 - "Write a short story about a robot that dreams for the first time."
95 - ] * n_prompts
96 - greedy_sampling_params = SamplingParams(temperature=0.0, max_tokens=64, logprobs=4)
97 - regular_sampling_params = SamplingParams(temperature=0.4, max_tokens=64, logprobs=4)
98 - topkp_sampling_params = SamplingParams(
99 - temperature=0.4, max_tokens=64, logprobs=4, top_k=12, top_p=0.5
100 - )
101 -
102 - for sp in [greedy_sampling_params, regular_sampling_params, topkp_sampling_params]:
103 - output = llm.generate(prompts, sp)
104 - for o in output:
105 - check_num_logprobs(o.outputs[0].logprobs, 4)
tests/v1/tpu/test_spmd_model_weight_loading.py deleted
-78
@@ -1,78 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -import gc
4 -import tempfile
5 -
6 -import numpy as np
7 -import pytest
8 -import torch_xla.distributed.spmd as xs
9 -import torch_xla.runtime as xr
10 -
11 -from vllm.config import set_current_vllm_config
12 -from vllm.distributed.parallel_state import (
13 - ensure_model_parallel_initialized,
14 - init_distributed_environment,
15 -)
16 -from vllm.engine.arg_utils import EngineArgs
17 -from vllm.model_executor.model_loader.tpu import TPUModelLoader
18 -
19 -
20 -def _setup_environment(model):
21 - engine_args = EngineArgs(
22 - model=model,
23 - )
24 - vllm_config = engine_args.create_engine_config()
25 - with set_current_vllm_config(vllm_config):
26 - temp_file = tempfile.mkstemp()[1]
27 - init_distributed_environment(
28 - 1,
29 - 0,
30 - local_rank=0,
31 - distributed_init_method=f"file://{temp_file}",
32 - backend="gloo",
33 - )
34 - # Under single worker mode, full model is init first and then
35 - # partitioned using GSPMD.
36 - ensure_model_parallel_initialized(1, 1)
37 - return vllm_config
38 -
39 -
40 -MESH = None
41 -
42 -
43 -def _get_spmd_mesh():
44 - global MESH
45 - if MESH is None:
46 - xr.use_spmd()
47 - num_devices = xr.global_runtime_device_count()
48 - mesh_shape = (num_devices, 1)
49 - device_ids = np.array(range(num_devices))
50 - MESH = xs.Mesh(device_ids, mesh_shape, ("x", "y"))
51 - return MESH
52 -
53 -
54 -@pytest.mark.parametrize(
55 - "model",
56 - [
57 - "Qwen/Qwen2-1.5B-Instruct",
58 - # Skip large models due to CI runner disk space limitations
59 - # "meta-llama/Llama-3.1-8B-Instruct",
60 - # "meta-llama/Llama-3.1-70B-Instruct",
61 - ],
62 -)
63 -def test_tpu_model_loader(model):
64 - # Skip the 70B test if there are less than 8 chips
65 - # TODO: Query using torch xla API, the query API is not working
66 - # with SPMD now. However, This test is running under SPMD mode.
67 - if "70B" in model and xr.global_runtime_device_count() < 8:
68 - pytest.skip(
69 - "Skipping 70B model if the TPU VM has less than 8 chips to \
70 - avoid OOM."
71 - )
72 -
73 - vllm_config = _setup_environment(model)
74 - loader = TPUModelLoader(load_config=vllm_config.load_config)
75 - mesh = _get_spmd_mesh()
76 - model = loader.load_model(vllm_config, vllm_config.model_config, mesh)
77 - del model
78 - gc.collect()
tests/v1/tpu/test_topk_topp_sampler.py deleted
-149
@@ -1,149 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -import math
4 -
5 -import pytest
6 -import torch
7 -import torch_xla
8 -
9 -from vllm.platforms import current_platform
10 -from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p
11 -from vllm.v1.sample.tpu.sampler import apply_top_k_top_p as apply_top_k_top_p_tpu
12 -
13 -if not current_platform.is_tpu():
14 - pytest.skip("This test needs a TPU.", allow_module_level=True)
15 -import torch_xla.core.xla_model as xm
16 -
17 -BATCH_SIZE = 1024
18 -VOCAB_SIZE = 128 * 1024
19 -TOLERANCE = 1e-6
20 -
21 -
22 -def test_topk_equivalence_to_native_impl():
23 - with torch.device(xm.xla_device()):
24 - xm.set_rng_state(seed=33)
25 -
26 - logits = torch.rand((BATCH_SIZE, VOCAB_SIZE))
27 -
28 - # Random top-k values between 1 and 10.
29 - k = torch.randint(1, 10, (BATCH_SIZE,))
30 -
31 - # Set k=vocab_size for ~50% of requests in the batch (top-k disabled).
32 - k.masked_fill_(torch.randint(0, 2, (BATCH_SIZE,), dtype=bool), VOCAB_SIZE)
33 -
34 - result_tpu = apply_top_k_top_p_tpu(logits=logits.clone(), k=k, p=None)
35 -
36 - result_native = apply_top_k_top_p(logits=logits.clone(), k=k, p=None)
37 - assert torch.allclose(result_native, result_tpu)
38 -
39 -
40 -def test_topp_result_sums_past_p():
41 - with torch.device(xm.xla_device()):
42 - xm.set_rng_state(seed=33)
43 -
44 - logits = torch.rand((BATCH_SIZE, VOCAB_SIZE))
45 - probs = logits.softmax(dim=-1)
46 -
47 - # Random top-p values between 0 and 1.
48 - p = torch.rand((BATCH_SIZE,))
49 -
50 - # Set p=1 for ~50% of requests in the batch (top-p disabled).
51 - p.masked_fill_(torch.randint(0, 2, (BATCH_SIZE,), dtype=bool), 1)
52 -
53 - no_op_k = torch.tensor([VOCAB_SIZE])
54 - logits_masked = apply_top_k_top_p_tpu(logits=logits.clone(), k=no_op_k, p=p)
55 -
56 - # Verify that the masked logit's probability sums to at least p.
57 - probs.masked_fill_(logits_masked.isinf(), 0)
58 - masked_prob_sum = probs.sum(dim=-1)
59 -
60 - torch_xla.sync()
61 -
62 - # Perform assertion on CPU.
63 - assert torch.all(torch.ge(masked_prob_sum.cpu() + TOLERANCE, p.cpu()))
64 -
65 -
66 -def test_topp_basic():
67 - with torch.device(xm.xla_device()):
68 - logits = torch.tensor(
69 - [
70 - [math.log(0.2), math.log(0.3), math.log(0.5)],
71 - [math.log(0.5), math.log(0.1), math.log(0.4)],
72 - ]
73 - )
74 -
75 - result = apply_top_k_top_p_tpu(
76 - logits=logits.clone(), k=torch.tensor([3, 3]), p=torch.tensor([0.79, 0.79])
77 - )
78 -
79 - torch_xla.sync()
80 -
81 - # Expect the smallest elements to be dropped.
82 - expected_result = logits.clone().cpu()
83 - expected_result[0, 0] = float("-inf")
84 - expected_result[1, 1] = float("-inf")
85 - assert torch.allclose(expected_result, result.cpu())
86 -
87 -
88 -def test_topp_select_all():
89 - with torch.device(xm.xla_device()):
90 - logits = torch.tensor(
91 - [
92 - [math.log(0.2), math.log(0.3), math.log(0.5)],
93 - [math.log(0.5), math.log(0.1), math.log(0.4)],
94 - ]
95 - )
96 -
97 - result = apply_top_k_top_p_tpu(
98 - logits=logits.clone(), k=torch.tensor([3, 3]), p=torch.tensor([1.0, 1.0])
99 - )
100 -
101 - torch_xla.sync()
102 -
103 - assert torch.allclose(logits.cpu(), result.cpu())
104 -
105 -
106 -def test_topp_with_ties():
107 - with torch.device(xm.xla_device()):
108 - # Input has multiple math.log(0.3).
109 - logits = torch.tensor(
110 - [[math.log(0.3), math.log(0.3), math.log(0.3), math.log(0.1)]]
111 - )
112 -
113 - result = apply_top_k_top_p_tpu(
114 - logits=logits.clone(), k=torch.tensor([4]), p=torch.tensor([0.2])
115 - )
116 -
117 - torch_xla.sync()
118 -
119 - # All tie values are included in the top-p set. Tie breaking is left
120 - # to be done during final sampling (all tie tokens have equal
121 - # probability of being chosen).
122 - expected_result = logits.clone().cpu()
123 - expected_result[0, 3] = float("-inf")
124 - assert torch.allclose(expected_result, result.cpu())
125 -
126 -
127 -def test_both_topk_topp():
128 - with torch.device(xm.xla_device()):
129 - logits = torch.tensor(
130 - [
131 - [math.log(0.2), math.log(0.3), math.log(0.5)],
132 - [math.log(0.5), math.log(0.1), math.log(0.4)],
133 - ]
134 - )
135 -
136 - # Set k=1 for the first batch.
137 - result = apply_top_k_top_p_tpu(
138 - logits=logits.clone(), k=torch.tensor([1, 3]), p=torch.tensor([0.79, 0.79])
139 - )
140 -
141 - torch_xla.sync()
142 -
143 - # Since for the first batch k=1, expect only the largest element gets
144 - # selected.
145 - expected_result = logits.clone().cpu()
146 - expected_result[0, 0] = float("-inf")
147 - expected_result[0, 1] = float("-inf")
148 - expected_result[1, 1] = float("-inf")
149 - assert torch.allclose(expected_result, result.cpu())
tests/v1/tpu/test_tpu_int8.py deleted
-78
@@ -1,78 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""Tests whether TPU Int8 computation is enabled correctly.
4 -
5 -Run `pytest tests/quantization/test_tpu_int8.py`.
6 -"""
7 -
8 -import pytest
9 -
10 -from vllm.model_executor.layers.linear import LinearBase
11 -from vllm.model_executor.layers.quantization.tpu_int8 import TPUInt8LinearMethod
12 -from vllm.platforms import current_platform
13 -
14 -from ...models.registry import HF_EXAMPLE_MODELS
15 -
16 -MODELS = ["Qwen/Qwen2.5-0.5B-Instruct"]
17 -
18 -
19 -@pytest.mark.skipif(
20 - not current_platform.is_tpu(), reason="TPU Int8 is only enabled for TPUs."
21 -)
22 -@pytest.mark.parametrize("model", MODELS)
23 -@pytest.mark.parametrize("dtype", ["bfloat16"])
24 -@pytest.mark.parametrize("max_tokens", [10])
25 -@pytest.mark.parametrize(
26 - "hf_overrides",
27 - [
28 - # w8a8 dynamic activation
29 - {
30 - "quantization_config": {
31 - "quant_method": "tpu_int8",
32 - "activation_scheme": "dynamic",
33 - }
34 - }
35 - ],
36 -)
37 -def test_model_tpu_int8(
38 - vllm_runner,
39 - model: str,
40 - dtype: str,
41 - max_tokens: int,
42 - hf_overrides: dict,
43 - monkeypatch,
44 -) -> None:
45 - model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
46 - model_info.check_transformers_version(on_fail="skip")
47 -
48 - activation_scheme = hf_overrides.get("quantization_config", {}).get(
49 - "activation_scheme"
50 - )
51 - quantize_activation = activation_scheme == "dynamic"
52 -
53 - # Allows using apply_model
54 - monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
55 - # Prevent error from re-initializing cache
56 - monkeypatch.setenv("VLLM_XLA_CACHE_PATH", "")
57 -
58 - prompts = [
59 - "A robot may not injure a human being",
60 - ]
61 - answers = [
62 - "or kill a human being",
63 - ]
64 -
65 - with vllm_runner(model, dtype=dtype, hf_overrides=hf_overrides) as vllm:
66 -
67 - def check_model(model):
68 - for name, module in model.named_modules():
69 - if not isinstance(module, LinearBase):
70 - continue
71 - quant_method = module.quant_method
72 - assert isinstance(quant_method, TPUInt8LinearMethod)
73 - assert quant_method.quantize_activation == quantize_activation
74 -
75 - vllm.apply_model(check_model)
76 - outputs = vllm.generate_greedy(prompts, max_tokens)
77 - for (_, output), answer in zip(outputs, answers):
78 - assert answer in output
tests/v1/tpu/test_tpu_qkv_linear.py deleted
-93
@@ -1,93 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -import tempfile
4 -
5 -import numpy as np
6 -import pytest
7 -import torch
8 -import torch_xla.distributed.spmd as xs
9 -import torch_xla.runtime as xr
10 -
11 -from vllm.config import set_current_vllm_config
12 -from vllm.distributed.parallel_state import (
13 - ensure_model_parallel_initialized,
14 - init_distributed_environment,
15 -)
16 -from vllm.distributed.tpu_distributed_utils import XlaQKVParallelLinear
17 -from vllm.engine.arg_utils import EngineArgs
18 -from vllm.model_executor.layers.linear import QKVParallelLinear
19 -
20 -
21 -@pytest.fixture(autouse=True)
22 -def setup_environment():
23 - # This is a fake config used for init dist env.
24 - # QKVParallelLinear needs dist env to be initialized.
25 - engine_args = EngineArgs(
26 - model="Qwen/Qwen2-1.5B-Instruct",
27 - max_model_len=64,
28 - max_num_batched_tokens=64,
29 - max_num_seqs=4,
30 - )
31 -
32 - vllm_config = engine_args.create_engine_config()
33 -
34 - with set_current_vllm_config(vllm_config):
35 - temp_file = tempfile.mkstemp()[1]
36 - init_distributed_environment(
37 - 1,
38 - 0,
39 - local_rank=0,
40 - distributed_init_method=f"file://{temp_file}",
41 - backend="gloo",
42 - )
43 - ensure_model_parallel_initialized(1, 1)
44 - yield
45 -
46 -
47 -MESH = None
48 -
49 -
50 -def _get_spmd_mesh():
51 - global MESH
52 - if MESH is None:
53 - xr.use_spmd()
54 - num_devices = xr.global_runtime_device_count()
55 - mesh_shape = (num_devices, 1)
56 - device_ids = np.array(range(num_devices))
57 - MESH = xs.Mesh(device_ids, mesh_shape, ("x", "y"))
58 - return MESH
59 -
60 -
61 -@pytest.mark.parametrize("bias", [False, True])
62 -# `xr.use_spmd()` will set a global state, and this state is not reversible.
63 -# Therefore, non-SPMD tests should be run before SPMD tests.
64 -@pytest.mark.parametrize("mesh", [None, _get_spmd_mesh()])
65 -@pytest.mark.parametrize("device", ["cpu", "xla"])
66 -@torch.no_grad()
67 -def test_xla_qkv_linear(bias, mesh, device):
68 - torch.manual_seed(123)
69 -
70 - qkv_linear = QKVParallelLinear(
71 - hidden_size=4096,
72 - head_size=128,
73 - total_num_heads=32,
74 - total_num_kv_heads=8,
75 - bias=bias,
76 - params_dtype=torch.bfloat16,
77 - return_bias=False,
78 - )
79 -
80 - qkv_linear.weight.data = torch.rand_like(qkv_linear.weight.data) / 10
81 - if bias:
82 - qkv_linear.bias.data = torch.rand_like(qkv_linear.bias.data)
83 -
84 - xla_qkv_linear = XlaQKVParallelLinear(qkv_linear, mesh=mesh)
85 -
86 - qkv_linear = qkv_linear.to(device)
87 - xla_qkv_linear = xla_qkv_linear.to(device)
88 - input_tensor = torch.rand(10, 4096, dtype=torch.bfloat16) / 10
89 - input_tensor = input_tensor.to(device)
90 -
91 - output = qkv_linear(input_tensor)
92 - xla_output = xla_qkv_linear(input_tensor)
93 - assert torch.allclose(output.cpu(), xla_output.cpu())
tests/v1/tpu/worker/__init__.py
tests/v1/tpu/worker/test_tpu_model_runner.py deleted
-587
@@ -1,587 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -import pytest
5 -
6 -from vllm.attention.layer import Attention
7 -from vllm.config import (
8 - CacheConfig,
9 - ModelConfig,
10 - SchedulerConfig,
11 - VllmConfig,
12 - set_current_vllm_config,
13 -)
14 -from vllm.pooling_params import PoolingParams
15 -from vllm.sampling_params import SamplingParams
16 -from vllm.utils.mem_constants import GiB_bytes
17 -from vllm.v1.core.kv_cache_utils import estimate_max_model_len, get_kv_cache_configs
18 -from vllm.v1.core.sched.output import CachedRequestData, NewRequestData, SchedulerOutput
19 -from vllm.v1.worker.tpu_model_runner import (
20 - TPUModelRunner,
21 - _get_padded_num_reqs_with_upper_limit,
22 - _get_padded_token_len,
23 - _get_req_paddings,
24 - _get_token_paddings,
25 -)
26 -
27 -
28 -def get_vllm_config():
29 - model_config = ModelConfig(
30 - model="facebook/opt-125m",
31 - dtype="bfloat16", # TPUs typically use bfloat16
32 - seed=42,
33 - )
34 - scheduler_config = SchedulerConfig(
35 - max_num_seqs=10,
36 - max_num_batched_tokens=512,
37 - max_model_len=512,
38 - is_encoder_decoder=model_config.is_encoder_decoder,
39 - )
40 - cache_config = CacheConfig(
41 - block_size=16,
42 - gpu_memory_utilization=0.9,
43 - swap_space=0,
44 - cache_dtype="auto",
45 - )
46 - vllm_config = VllmConfig(
47 - model_config=model_config,
48 - cache_config=cache_config,
49 - scheduler_config=scheduler_config,
50 - )
51 - return vllm_config
52 -
53 -
54 -def get_model_runner(vllm_config):
55 - device = "xla:0" # Mocking TPU device
56 - return TPUModelRunner(vllm_config, device)
57 -
58 -
59 -@pytest.fixture
60 -def model_runner():
61 - # Patchers have already been started at module level.
62 - vllm_config = get_vllm_config()
63 - return get_model_runner(vllm_config)
64 -
65 -
66 -def _schedule_new_request(*req_ids: str) -> SchedulerOutput:
67 - new_reqs = []
68 - num_scheduled_tokens = {}
69 - total_num_scheduled_tokens = 0
70 - for req_id in req_ids:
71 - new_reqs.append(
72 - NewRequestData(
73 - req_id=req_id,
74 - prompt_token_ids=[1, 2, 3],
75 - mm_features=[],
76 - sampling_params=SamplingParams(),
77 - pooling_params=PoolingParams(),
78 - block_ids=([0],), # block_ids should be tuple[list[int]]
79 - num_computed_tokens=0,
80 - lora_request=None,
81 - )
82 - )
83 - num_scheduled_tokens[req_id] = 3
84 - total_num_scheduled_tokens += num_scheduled_tokens[req_id]
85 -
86 - return SchedulerOutput(
87 - scheduled_new_reqs=new_reqs,
88 - scheduled_cached_reqs=CachedRequestData.make_empty(),
89 - num_scheduled_tokens=num_scheduled_tokens,
90 - total_num_scheduled_tokens=total_num_scheduled_tokens,
91 - scheduled_spec_decode_tokens={},
92 - scheduled_encoder_inputs={},
93 - num_common_prefix_blocks=[],
94 - finished_req_ids=set(),
95 - free_encoder_mm_hashes=[],
96 - )
97 -
98 -
99 -def _is_req_scheduled(model_runner, req_id: str) -> bool:
100 - return req_id in model_runner.input_batch.req_id_to_index
101 -
102 -
103 -def _is_req_added(model_runner, req_id: str) -> bool:
104 - return req_id in model_runner.requests
105 -
106 -
107 -def _is_req_state_block_table_match(model_runner, req_id: str) -> bool:
108 - """Check if the request state block IDs match the block table.
109 -
110 - This function handles both legacy BlockTable and new MultiGroupBlockTable
111 - structures for backward compatibility.
112 - """
113 -
114 - req_index = model_runner.input_batch.req_id_to_index[req_id]
115 - multi_group_block_table = model_runner.input_batch.block_table
116 - req_state = model_runner.requests[req_id]
117 -
118 - # Access the first block table from MultiGroupBlockTable
119 - # This is safe since we currently only use single KV cache groups
120 - block_table = multi_group_block_table[0]
121 -
122 - # req_state.block_ids is now tuple[list[int], ...] for MultiGroupBlockTable
123 - # Extract the first group's block IDs
124 - if isinstance(req_state.block_ids[0], list):
125 - # New format: tuple[list[int], ...] - extract first group
126 - req_block_ids = req_state.block_ids[0]
127 - else:
128 - # Legacy format: list[int] - use directly
129 - req_block_ids = req_state.block_ids
130 -
131 - if block_table.num_blocks_per_row[req_index] != len(req_block_ids):
132 - return False
133 -
134 - num_blocks = block_table.num_blocks_per_row[req_index]
135 - block_table_values = block_table.block_table.np[req_index, :num_blocks]
136 - return (block_table_values == req_block_ids).all()
137 -
138 -
139 -def test_update_states_new_request(model_runner):
140 - req_id = "req_0"
141 -
142 - # new req
143 - scheduler_output = _schedule_new_request(req_id)
144 -
145 - model_runner._update_states(scheduler_output)
146 -
147 - assert _is_req_added(model_runner, req_id)
148 - assert _is_req_scheduled(model_runner, req_id)
149 - assert _is_req_state_block_table_match(model_runner, req_id)
150 -
151 -
152 -def test_update_states_request_finished(model_runner):
153 - req_id = "req_0"
154 -
155 - # new req
156 - scheduler_output = _schedule_new_request(req_id)
157 -
158 - model_runner._update_states(scheduler_output)
159 - assert _is_req_added(model_runner, req_id)
160 - assert _is_req_scheduled(model_runner, req_id)
161 -
162 - # finish req
163 - scheduler_output = SchedulerOutput(
164 - scheduled_new_reqs=[],
165 - scheduled_cached_reqs=CachedRequestData.make_empty(),
166 - num_scheduled_tokens={},
167 - total_num_scheduled_tokens=0,
168 - scheduled_spec_decode_tokens={},
169 - scheduled_encoder_inputs={},
170 - num_common_prefix_blocks=[],
171 - finished_req_ids={req_id},
172 - free_encoder_mm_hashes=[],
173 - )
174 -
175 - model_runner._update_states(scheduler_output)
176 - assert not _is_req_added(model_runner, req_id)
177 - assert not _is_req_scheduled(model_runner, req_id)
178 -
179 -
180 -def test_update_states_request_resumed(model_runner):
181 - req_id = "req_0"
182 -
183 - # new req
184 - scheduler_output = _schedule_new_request(req_id)
185 -
186 - model_runner._update_states(scheduler_output)
187 - assert _is_req_added(model_runner, req_id)
188 - assert _is_req_scheduled(model_runner, req_id)
189 -
190 - # unschedule req
191 - scheduler_output = SchedulerOutput(
192 - scheduled_new_reqs=[],
193 - scheduled_cached_reqs=CachedRequestData.make_empty(),
194 - num_scheduled_tokens={},
195 - total_num_scheduled_tokens=0,
196 - scheduled_spec_decode_tokens={},
197 - scheduled_encoder_inputs={},
198 - num_common_prefix_blocks=[],
199 - finished_req_ids=set(),
200 - free_encoder_mm_hashes=[],
201 - )
202 -
203 - model_runner._update_states(scheduler_output)
204 - assert _is_req_added(model_runner, req_id)
205 - assert not _is_req_scheduled(model_runner, req_id)
206 -
207 - # resume req
208 - cached_req_data = CachedRequestData(
209 - req_ids=[req_id],
210 - resumed_req_ids={req_id},
211 - new_token_ids=[[]],
212 - all_token_ids={req_id: scheduler_output.scheduled_new_reqs[0].prompt_token_ids},
213 - new_block_ids=[([],)],
214 - num_computed_tokens=[0],
215 - num_output_tokens=[0],
216 - )
217 -
218 - scheduler_output = SchedulerOutput(
219 - scheduled_new_reqs=[],
220 - scheduled_cached_reqs=cached_req_data,
221 - num_scheduled_tokens={req_id: 1},
222 - total_num_scheduled_tokens=1,
223 - scheduled_spec_decode_tokens={},
224 - scheduled_encoder_inputs={},
225 - num_common_prefix_blocks=[],
226 - finished_req_ids=set(),
227 - free_encoder_mm_hashes=[],
228 - )
229 -
230 - model_runner._update_states(scheduler_output)
231 - assert _is_req_added(model_runner, req_id)
232 - assert _is_req_scheduled(model_runner, req_id)
233 - assert _is_req_state_block_table_match(model_runner, req_id)
234 -
235 -
236 -def test_update_states_no_changes(model_runner):
237 - req_id = "req_0"
238 -
239 - # new req
240 - scheduler_output = _schedule_new_request(req_id)
241 -
242 - model_runner._update_states(scheduler_output)
243 - assert _is_req_added(model_runner, req_id)
244 - assert _is_req_scheduled(model_runner, req_id)
245 -
246 - # schedule req
247 - scheduler_output = SchedulerOutput(
248 - scheduled_new_reqs=[],
249 - scheduled_cached_reqs=CachedRequestData.make_empty(),
250 - num_scheduled_tokens={req_id: 1},
251 - total_num_scheduled_tokens=1,
252 - scheduled_spec_decode_tokens={},
253 - scheduled_encoder_inputs={},
254 - num_common_prefix_blocks=[],
255 - finished_req_ids=set(),
256 - free_encoder_mm_hashes=[],
257 - )
258 -
259 - model_runner._update_states(scheduler_output)
260 - assert _is_req_added(model_runner, req_id)
261 - assert _is_req_scheduled(model_runner, req_id)
262 - assert _is_req_state_block_table_match(model_runner, req_id)
263 -
264 -
265 -def test_update_states_request_unscheduled(model_runner):
266 - req_ids = ("req_0", "req_1")
267 -
268 - # new reqs
269 - scheduler_output = _schedule_new_request(*req_ids)
270 -
271 - model_runner._update_states(scheduler_output)
272 -
273 - assert _is_req_added(model_runner, req_ids[0])
274 - assert _is_req_scheduled(model_runner, req_ids[0])
275 -
276 - assert _is_req_added(model_runner, req_ids[1])
277 - assert _is_req_scheduled(model_runner, req_ids[1])
278 -
279 - # unschedule req_1
280 - scheduler_output = SchedulerOutput(
281 - scheduled_new_reqs=[],
282 - scheduled_cached_reqs=CachedRequestData.make_empty(),
283 - num_scheduled_tokens={req_ids[0]: 1},
284 - total_num_scheduled_tokens=1,
285 - scheduled_spec_decode_tokens={},
286 - scheduled_encoder_inputs={},
287 - num_common_prefix_blocks=[],
288 - finished_req_ids=set(),
289 - free_encoder_mm_hashes=[],
290 - )
291 -
292 - model_runner._update_states(scheduler_output)
293 -
294 - assert _is_req_added(model_runner, req_ids[0])
295 - assert _is_req_scheduled(model_runner, req_ids[0])
296 -
297 - assert _is_req_added(model_runner, req_ids[1])
298 - assert not _is_req_scheduled(model_runner, req_ids[1])
299 -
300 -
301 -def test_get_paddings():
302 - # Bucketed padding
303 - min_token_size, max_token_size, padding_gap = 16, 512, 64
304 - expected_paddings = [16, 32, 64, 128, 192, 256, 320, 384, 448, 512]
305 - actual_paddings = _get_token_paddings(min_token_size, max_token_size, padding_gap)
306 -
307 - # Bucketed padding with max_token_size not a power of two.
308 - max_token_size = 317
309 - expected_paddings = [16, 32, 64, 128, 192, 256, 320]
310 - actual_paddings = _get_token_paddings(min_token_size, max_token_size, padding_gap)
311 - assert actual_paddings == expected_paddings
312 -
313 - # Exponential padding.
314 - max_token_size, padding_gap = 1024, 0
315 - expected_paddings = [16, 32, 64, 128, 256, 512, 1024]
316 - actual_paddings = _get_token_paddings(min_token_size, max_token_size, padding_gap)
317 - assert actual_paddings == expected_paddings
318 - # Exponential padding with max_token_size not a power of two.
319 - max_token_size = 317
320 - expected_paddings = [16, 32, 64, 128, 256, 512]
321 - actual_paddings = _get_token_paddings(min_token_size, max_token_size, padding_gap)
322 - assert actual_paddings == expected_paddings
323 -
324 -
325 -def test_get_padded_token_len():
326 - min_token_size, max_token_size, padding_gap = 16, 512, 64
327 - paddings = _get_token_paddings(min_token_size, max_token_size, padding_gap)
328 - assert _get_padded_token_len(paddings, 1) == 16
329 - assert _get_padded_token_len(paddings, 16) == 16
330 - assert _get_padded_token_len(paddings, 20) == 32
331 - assert _get_padded_token_len(paddings, 300) == 320
332 - assert _get_padded_token_len(paddings, 512) == 512
333 -
334 -
335 -def test_get_padded_num_reqs_with_upper_limit():
336 - assert _get_padded_num_reqs_with_upper_limit(3, 32) == 8
337 - assert _get_padded_num_reqs_with_upper_limit(9, 32) == 16
338 - assert _get_padded_num_reqs_with_upper_limit(19, 32) == 32
339 - assert _get_padded_num_reqs_with_upper_limit(17, 28) == 28
340 -
341 -
342 -def test_get_req_paddings():
343 - assert _get_req_paddings(1, 32) == [8, 16, 32]
344 - assert _get_req_paddings(8, 32) == [8, 16, 32]
345 - assert _get_req_paddings(8, 36) == [8, 16, 32, 36]
346 -
347 -
348 -def test_init_kv_cache_with_kv_sharing_invalid_target_layer_order(model_runner):
349 - layer_0 = "model.layers.0.self_attn.attn"
350 - layer_1 = "model.layers.1.self_attn.attn"
351 - error_msg = f"{layer_1} must come before the current layer"
352 - vllm_config = model_runner.vllm_config
353 - with (
354 - pytest.raises(ValueError, match=error_msg),
355 - set_current_vllm_config(vllm_config),
356 - ):
357 - fwd_context = {
358 - # initialization below will fail because target layer is invalid;
359 - # the target layer needs to come before layer 1
360 - layer_0: Attention(
361 - num_heads=8,
362 - head_size=128,
363 - scale=1.0,
364 - prefix=layer_0,
365 - kv_sharing_target_layer_name=layer_1,
366 - ),
367 - layer_1: Attention(
368 - num_heads=8,
369 - head_size=128,
370 - scale=1.0,
371 - prefix=layer_1,
372 - ),
373 - }
374 - # suppress var not used error
375 - assert fwd_context is not None
376 -
377 -
378 -def test_init_kv_cache_with_kv_sharing_target_layer_not_exist(model_runner):
379 - layer_0 = "model.layers.0.self_attn.attn"
380 - layer_1 = "model.layers.1.self_attn.attn"
381 - invalid_layer = "model.layers.0.cross_attn.attn"
382 - error_msg = f"{invalid_layer} is not a valid Attention layer in the model"
383 - vllm_config = model_runner.vllm_config
384 - with (
385 - pytest.raises(ValueError, match=error_msg),
386 - set_current_vllm_config(vllm_config),
387 - ):
388 - fwd_context = {
389 - layer_0: Attention(
390 - num_heads=8,
391 - head_size=128,
392 - scale=1.0,
393 - prefix=layer_0,
394 - ),
395 - layer_1: Attention(
396 - num_heads=8,
397 - head_size=128,
398 - scale=1.0,
399 - prefix=layer_1,
400 - # invalid layer: cross_attn.atn doesn't exist!
401 - kv_sharing_target_layer_name=invalid_layer,
402 - ),
403 - }
404 - # suppress var not used error
405 - assert fwd_context is not None
406 -
407 -
408 -def test_init_kv_cache_with_kv_sharing_target_same_as_current(model_runner):
409 - layer_0 = "model.layers.0.self_attn.attn"
410 - layer_1 = "model.layers.1.self_attn.attn"
411 - error_msg = f"{layer_1} cannot be the same as the current layer"
412 - vllm_config = model_runner.vllm_config
413 - with (
414 - pytest.raises(ValueError, match=error_msg),
415 - set_current_vllm_config(vllm_config),
416 - ):
417 - fwd_context = {
418 - # initialization below will fail because target layer is invalid;
419 - # the target layer needs to come before layer 1
420 - layer_0: Attention(
421 - num_heads=8,
422 - head_size=128,
423 - scale=1.0,
424 - prefix=layer_0,
425 - ),
426 - layer_1: Attention(
427 - num_heads=8,
428 - head_size=128,
429 - scale=1.0,
430 - prefix=layer_1,
431 - kv_sharing_target_layer_name=layer_1,
432 - ),
433 - }
434 - # suppress var not used error
435 - assert fwd_context is not None
436 -
437 -
438 -def test_init_kv_cache_without_kv_sharing():
439 - layer_0 = "model.layers.0.self_attn.attn"
440 - layer_1 = "model.layers.1.self_attn.attn"
441 - vllm_config = get_vllm_config()
442 - with set_current_vllm_config(vllm_config):
443 - fwd_context = {
444 - layer_0: Attention(
445 - num_heads=8,
446 - head_size=128,
447 - scale=1.0,
448 - prefix=layer_0,
449 - ),
450 - layer_1: Attention(
451 - num_heads=8,
452 - head_size=128,
453 - scale=1.0,
454 - prefix=layer_1,
455 - ),
456 - }
457 - # suppress var not used error
458 - assert fwd_context is not None
459 - # Set high context length to test max context length estimation
460 - vllm_config.model_config.max_model_len = 1_000_000
461 - vllm_ctx = vllm_config.compilation_config.static_forward_context
462 - model_runner = get_model_runner(vllm_config)
463 - kv_cache_spec = model_runner.get_kv_cache_spec()
464 - assert len(kv_cache_spec) == 2
465 - assert len(model_runner.shared_kv_cache_layers) == 0
466 -
467 - available_memory = 20 * GiB_bytes
468 - # page size for each layer KV can be calculated as
469 - # 2 (non-MLA) * 8 (num_heads) * 128 (head_dim)
470 - # * 2 (bfloat16, kv_cache dtype) * 128 (block_size) = 512KB
471 - num_expected_blocks = 20480 # 20GB / 512KB / 2 (num layers)
472 - kv_cache_config = get_kv_cache_configs(
473 - vllm_config, [kv_cache_spec], [available_memory]
474 - )[0]
475 - assert kv_cache_config.num_blocks == num_expected_blocks
476 - assert len(kv_cache_config.kv_cache_tensors) == 2
477 - assert kv_cache_config.kv_cache_tensors[0].size == available_memory // 2
478 - assert kv_cache_config.kv_cache_tensors[1].size == available_memory // 2
479 -
480 - max_context_len = estimate_max_model_len(vllm_config, kv_cache_spec, 5 * GiB_bytes)
481 - # max context len with KV sharing should be 2x as large as without
482 - # max_context_len = available_memory / (page_size / block_size) / num_caches
483 - # max_context_len = 5GB / (512KB / 128) / 2 = 655360
484 - assert max_context_len == 655360
485 -
486 - # important: override tensor size to prevent large mem alloc during test
487 - # this will only allocate 2 block worth of memory (2 * 512kb)
488 - kv_cache_config.num_blocks = 1
489 - for kv_cache_tensor in kv_cache_config.kv_cache_tensors:
490 - kv_cache_tensor.size = kv_cache_spec[
491 - kv_cache_tensor.shared_by[0]
492 - ].page_size_bytes
493 -
494 - model_runner.initialize_kv_cache(kv_cache_config)
495 -
496 - layer_0_kv = vllm_ctx[layer_0].kv_cache[0]
497 - layer_1_kv = vllm_ctx[layer_1].kv_cache[0]
498 - # check layer 1 kv cache does NOT share memory with layer 0
499 - assert id(layer_1_kv) != id(layer_0_kv)
500 -
501 - # check layer 1 added to kv cache group's layer names
502 - assert len(kv_cache_config.kv_cache_groups) == 1
503 - assert len(kv_cache_config.kv_cache_groups[0].layer_names) == 2
504 - assert kv_cache_config.kv_cache_groups[0].layer_names[0] == layer_0
505 - assert kv_cache_config.kv_cache_groups[0].layer_names[1] == layer_1
506 -
507 -
508 -def test_init_kv_cache_with_kv_sharing_valid():
509 - layer_0 = "model.layers.0.self_attn.attn"
510 - layer_1 = "model.layers.1.self_attn.attn"
511 - vllm_config = get_vllm_config()
512 - with set_current_vllm_config(vllm_config):
513 - fwd_context = {
514 - layer_0: Attention(
515 - num_heads=8,
516 - head_size=128,
517 - scale=1.0,
518 - prefix=layer_0,
519 - ),
520 - layer_1: Attention(
521 - num_heads=8,
522 - head_size=128,
523 - scale=1.0,
524 - prefix=layer_1,
525 - kv_sharing_target_layer_name="model.layers.0.self_attn.attn",
526 - ),
527 - }
528 - # suppress var not used error
529 - assert fwd_context is not None
530 - # Set high context length to test max context length estimation
531 - vllm_config.model_config.max_model_len = 3_000_000
532 - vllm_ctx = vllm_config.compilation_config.static_forward_context
533 - model_runner = get_model_runner(vllm_config)
534 - kv_cache_spec = model_runner.get_kv_cache_spec()
535 - assert len(kv_cache_spec) == 1
536 - assert layer_0 in kv_cache_spec
537 - assert model_runner.shared_kv_cache_layers[layer_1] == layer_0
538 -
539 - available_memory = 20 * GiB_bytes
540 - # page size for layer 0's kv_cache_spec is 512KB
541 - # with KV sharing, we can allocate (available_mem//page_size//1) blocks
542 - # which is twice as many as without KV sharing
543 - num_expected_blocks = 2 * 20480 # 20GB / 512KB
544 - kv_cache_config = get_kv_cache_configs(
545 - vllm_config, [kv_cache_spec], [available_memory]
546 - )[0]
547 - assert kv_cache_config.num_blocks == num_expected_blocks
548 - assert len(kv_cache_config.kv_cache_tensors) == 1
549 - # Each layer now has twice the available memory for KV cache
550 - # compared to no KV sharing
551 - assert kv_cache_config.kv_cache_tensors[0].size == available_memory
552 -
553 - max_context_len = estimate_max_model_len(vllm_config, kv_cache_spec, 5 * GiB_bytes)
554 - # max context len with KV sharing should be 2x as large as without
555 - assert max_context_len == (2 * 655360)
556 -
557 - # important: override tensor size to prevent large mem alloc during test
558 - # this will only allocate 1 block worth of memory (512kb)
559 - kv_cache_config.num_blocks = 1
560 - kv_cache_config.kv_cache_tensors[0].size = kv_cache_spec[layer_0].page_size_bytes
561 -
562 - model_runner.initialize_kv_cache(kv_cache_config)
563 -
564 - layer_0_kv = vllm_ctx[layer_0].kv_cache[0]
565 - layer_1_kv = vllm_ctx[layer_1].kv_cache[0]
566 - # check layer 1 kv cache shares memory with layer 0
567 - assert id(layer_1_kv) == id(layer_0_kv)
568 -
569 - # check layer 1 added to kv cache group's layer names
570 - assert len(kv_cache_config.kv_cache_groups) == 1
571 - assert len(kv_cache_config.kv_cache_groups[0].layer_names) == 2
572 - assert kv_cache_config.kv_cache_groups[0].layer_names[0] == layer_0
573 - assert kv_cache_config.kv_cache_groups[0].layer_names[1] == layer_1
574 -
575 -
576 -def test_most_model_len(monkeypatch: pytest.MonkeyPatch):
577 - monkeypatch.setenv("VLLM_TPU_MOST_MODEL_LEN", "2048")
578 - vllm_config = get_vllm_config()
579 - vllm_config.model_config.max_model_len = 32000
580 - vllm_config.scheduler_config.max_num_seqs = 1200
581 - model_runner = get_model_runner(vllm_config)
582 -
583 - # verify model runner will adjust num_reqs to avoid SMEM OOM.
584 - assert model_runner.num_reqs_most_model_len == 1200
585 - # num_page_per_req = 32k // 128
586 - # num_reqs = 1024 ** 2 // 2 // num_page_per_req // 4 = 524
587 - assert model_runner.num_reqs_max_model_len == 524
vllm/attention/backends/registry.py
-1
@@ -69,7 +69,6 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):
69 "vllm.v1.attention.backends.mla.flashmla_sparse.FlashMLASparseBackend"
70 )
71 FLASH_ATTN_MLA = "vllm.v1.attention.backends.mla.flashattn_mla.FlashAttnMLABackend"
72 - PALLAS = "vllm.v1.attention.backends.pallas.PallasAttentionBackend"
72 IPEX = "vllm.v1.attention.backends.ipex.IpexAttentionBackend"
73 NO_ATTENTION = "vllm.v1.attention.backends.no_attention.NoAttentionBackend"
74 FLEX_ATTENTION = "vllm.v1.attention.backends.flex_attention.FlexAttentionBackend"
vllm/attention/layers/mm_encoder_attention.py
-25
@@ -227,28 +227,3 @@ class MMEncoderAttention(CustomOp):
227 "XPU only supports FLASH_ATTN for vision attention."
228 )
229 return self._forward_fa(query, key, value, cu_seqlens, max_seqlen)
230 -
231 - def forward_tpu(
232 - self,
233 - query: torch.Tensor,
234 - key: torch.Tensor,
235 - value: torch.Tensor,
236 - cu_seqlens: torch.Tensor | None = None,
237 - max_seqlen: torch.Tensor | None = None, # Only used for Flash Attention
238 - ) -> torch.Tensor:
239 - assert self.attn_backend == AttentionBackendEnum.PALLAS, (
240 - f"MMEncoderAttention on TPU only supports PALLAS backend, "
241 - f"but got {self.attn_backend}."
242 - )
243 - if cu_seqlens is None:
244 - query, key, value = (x.transpose(1, 2) for x in (query, key, value))
245 - from torch_xla.experimental.custom_kernel import flash_attention
246 -
247 - out = flash_attention(query, key, value, sm_scale=self.scale)
248 - out = out.transpose(1, 2)
249 - return out
250 - logger.warning_once(
251 - "PALLAS backend with cu_seqlens is not supported for ViT yet. ",
252 - "Falling back to SDPA implementation.",
253 - )
254 - return self._forward_sdpa(query, key, value, cu_seqlens)
vllm/distributed/device_communicators/tpu_communicator.py deleted
-99
@@ -1,99 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -import os
5 -
6 -import torch
7 -from torch.distributed import ProcessGroup
8 -
9 -from vllm.config import get_current_vllm_config
10 -from vllm.logger import init_logger
11 -from vllm.platforms import current_platform
12 -from vllm.platforms.tpu import USE_TPU_INFERENCE
13 -
14 -from .base_device_communicator import DeviceCommunicatorBase
15 -
16 -USE_RAY = parallel_config = (
17 - get_current_vllm_config().parallel_config.distributed_executor_backend == "ray"
18 -)
19 -
20 -logger = init_logger(__name__)
21 -
22 -if not USE_TPU_INFERENCE:
23 - logger.info("tpu_inference not found, using vLLM's TpuCommunicator")
24 - if current_platform.is_tpu():
25 - import torch_xla
26 - import torch_xla.core.xla_model as xm
27 - import torch_xla.runtime as xr
28 - from torch_xla._internal import pjrt
29 - from torch_xla.distributed.xla_multiprocessing import (
30 - create_optimized_replica_groups,
31 - )
32 -
33 - if USE_RAY:
34 - from vllm.v1.executor import ray_utils
35 -
36 -
37 -class TpuCommunicator(DeviceCommunicatorBase):
38 - def __init__(
39 - self,
40 - cpu_group: ProcessGroup,
41 - device: torch.device | None = None,
42 - device_group: ProcessGroup | None = None,
43 - unique_name: str = "",
44 - ):
45 - super().__init__(cpu_group, device, device_group, unique_name)
46 -
47 - # NOTE(woosuk): When using TP > 1 on TPUs, every TPU on the same node
48 - # must be used together. Therefore, the local rank and world size can
49 - # be simply calculated as follows.
50 - global_rank = self.global_rank
51 - global_world_size = self.global_world_size
52 -
53 - if USE_RAY:
54 - logger.info("TpuCommunicator initialized with RAY")
55 - # Calculate how many TPU nodes are in the current deployment. This
56 - # is the Ray placement group if it is deployed with Ray. Default
57 - # to the number of TPU nodes in the Ray cluster. The number of TPU
58 - # nodes is computed by the total number of TPUs divided by the
59 - # number of TPU accelerators per node, to account for clusters
60 - # with both CPUs and TPUs.
61 - num_nodes = ray_utils.get_num_tpu_nodes()
62 - num_nodes_in_pg = ray_utils.get_num_nodes_in_placement_group()
63 - if num_nodes_in_pg > 0:
64 - num_nodes = num_nodes_in_pg
65 -
66 - local_world_size = global_world_size // num_nodes
67 - local_rank = global_rank % local_world_size
68 - else:
69 - logger.info("TpuCommunicator initialized with MP")
70 - # Sanity: Verify we run on a single host
71 - num_hosts = torch_xla.tpu.num_tpu_workers()
72 - assert num_hosts == 1
73 -
74 - # Get the current number of TPUs (we have locally)
75 - local_world_size = torch_xla.tpu.num_available_chips()
76 -
77 - # Get current rank
78 - local_rank = global_rank % local_world_size
79 -
80 - # Ensure environment variables are set for multihost deployments.
81 - # On GKE, this is needed for libtpu and TPU driver to know which TPU
82 - # chip is actually visible. Otherwise the TPU driver will fail to
83 - # initialize because the number of devices would be different from
84 - # the number of visible worker addresses.
85 - os.environ["CLOUD_TPU_TASK_ID"] = str(global_rank)
86 - os.environ["TPU_VISIBLE_CHIPS"] = str(local_rank)
87 -
88 - pjrt.initialize_multiprocess(local_rank, local_world_size)
89 - xr._init_world_size_ordinal()
90 - self.groups = create_optimized_replica_groups()
91 -
92 - def all_reduce(self, input_: torch.Tensor) -> torch.Tensor:
93 - # TODO: Remove the groups specification after XLA compiler can support
94 - # auto-reordering the ring order for all-reduce.
95 - return xm.all_reduce(xm.REDUCE_SUM, input_, groups=self.groups)
96 -
97 - def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:
98 - assert dim == -1, "TPUs only support dim=-1 for all-gather."
99 - return xm.all_gather(input_, dim=dim)
vllm/distributed/kv_transfer/kv_connector/utils.py
+1 -5
@@ -11,7 +11,6 @@ from typing import TYPE_CHECKING, Literal
11 import torch
12
13 from vllm.attention.backends.abstract import AttentionBackend
14 -from vllm.attention.backends.registry import AttentionBackendEnum
14 from vllm.config import get_current_vllm_config
15 from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory
16 from vllm.logger import init_logger
@@ -251,9 +250,6 @@ class TpKVTopology:
250 len(kv_cache_shape) == 5 and kv_cache_shape[0] == 1
251 )
252
254 - attn_backend = AttentionBackendEnum[self.attn_backend.get_name()]
255 - self._use_pallas = attn_backend == AttentionBackendEnum.PALLAS
256 -
253 @property
254 def is_kv_layout_blocks_first(self) -> bool:
255 return self._is_kv_layout_blocks_first
@@ -261,7 +257,7 @@ class TpKVTopology:
257 @property
258 def split_k_and_v(self) -> bool:
259 # Whether to register regions for K and V separately (when present).
264 - return not (self.is_mla or self._use_pallas or self.is_kv_layout_blocks_first)
260 + return not (self.is_mla or self.is_kv_layout_blocks_first)
261
262 @property
263 def tp_size(self) -> int:
vllm/distributed/kv_transfer/kv_connector/v1/mooncake_connector.py
-1
@@ -499,7 +499,6 @@ class MooncakeConnectorWorker:
499 total_num_kv_heads=self.model_config.get_total_num_kv_heads(),
500 attn_backend=backend,
501 )
502 - self._use_pallas = self.kv_topo._use_pallas
502
503 self.zmq_ctx = zmq.Context()
504 self.async_zmq_ctx = zmq.asyncio.Context()
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+1 -7
@@ -990,7 +990,6 @@ class NixlConnectorWorker:
990 total_num_kv_heads=self.model_config.get_total_num_kv_heads(),
991 attn_backend=backend,
992 )
993 - self._use_pallas = self.kv_topo._use_pallas
993 self._physical_blocks_per_logical_kv_block = 1
994
995 def _nixl_handshake(
@@ -1648,9 +1647,6 @@ class NixlConnectorWorker:
1647 # Num kv_heads > tp_size and P TP > D TP case, not supported
1648 assert not (tp_ratio < 0 and self.kv_topo.is_kv_replicated(remote_engine_id))
1649
1651 - assert not self._use_pallas or tp_ratio == 1, (
1652 - "TPU (pallas_v1) DOES NOT support heterogeneous TP yet."
1653 - )
1650 kv_cache_layout = (
1651 self.kv_cache_layout
1652 if not self.use_host_buffer
@@ -1821,9 +1817,7 @@ class NixlConnectorWorker:
1817
1818 if len(self.device_kv_caches) == 0:
1819 return
1824 - split_k_and_v = not (
1825 - self.use_mla or self._use_pallas or self.kv_topo.is_kv_layout_blocks_first
1826 - )
1820 + split_k_and_v = not (self.use_mla or self.kv_topo.is_kv_layout_blocks_first)
1821 sample_cache = list(self.device_kv_caches.values())[0][0]
1822 for block_size_ratio, block_ids_list in block_ids_per_ratio.items():
1823 assert block_size_ratio > 1, "Only nP < nD supported currently."
vllm/distributed/tpu_distributed_utils.py deleted
-188
@@ -1,188 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from collections import OrderedDict
4 -from typing import Optional
5 -
6 -import torch
7 -import torch.nn as nn
8 -import torch.nn.functional as F
9 -import torch_xla.distributed.spmd as xs
10 -from torch.nn.parameter import Parameter
11 -
12 -from vllm.logger import init_logger
13 -from vllm.model_executor.layers.linear import (
14 - ColumnParallelLinear,
15 - QKVParallelLinear,
16 - RowParallelLinear,
17 -)
18 -
19 -logger = init_logger(__name__)
20 -
21 -
22 -class XlaQKVParallelLinear(nn.Module):
23 - def __init__(self, qkv_linear: nn.Module, mesh: Optional["xs.Mesh"] = None):
24 - super().__init__()
25 - assert isinstance(qkv_linear, QKVParallelLinear)
26 - self.skip_bias_add = qkv_linear.skip_bias_add
27 - self.return_bias = qkv_linear.return_bias
28 - assert qkv_linear.tp_size == 1, "TP > 1 is only supported under SPMD."
29 -
30 - self.q_weight: Parameter
31 - self.k_weight: Parameter
32 - self.v_weight: Parameter
33 - self.q_bias: Parameter | None
34 - self.k_bias: Parameter | None
35 - self.v_bias: Parameter | None
36 - self._load_weights_from_qkv_linear(qkv_linear)
37 - if mesh is not None:
38 - self._shard_weight(mesh)
39 -
40 - def _shard_weight(self, mesh: "xs.Mesh"):
41 - self.q_weight = Parameter(self.q_weight.to("xla"), requires_grad=False)
42 - self.k_weight = Parameter(self.k_weight.to("xla"), requires_grad=False)
43 - self.v_weight = Parameter(self.v_weight.to("xla"), requires_grad=False)
44 - xs.mark_sharding(self.q_weight, mesh, ("x", None))
45 - xs.mark_sharding(self.k_weight, mesh, ("x", None))
46 - xs.mark_sharding(self.v_weight, mesh, ("x", None))
47 - if self.q_bias is not None:
48 - assert self.k_bias is not None and self.v_bias is not None, (
49 - "QKVParallelLinear should have q, k, and v biases together."
50 - )
51 - self.q_bias = Parameter(self.q_bias.to("xla"), requires_grad=False)
52 - xs.mark_sharding(self.q_bias, mesh, ("x",))
53 - self.k_bias = Parameter(self.k_bias.to("xla"), requires_grad=False)
54 - xs.mark_sharding(self.k_bias, mesh, ("x",))
55 - self.v_bias = Parameter(self.v_bias.to("xla"), requires_grad=False)
56 - xs.mark_sharding(self.v_bias, mesh, ("x",))
57 -
58 - def _load_weights_from_qkv_linear(self, qkv_linear: nn.Module):
59 - q_proj_size, k_proj_size, _ = qkv_linear.output_sizes
60 - # The weight of qkv linear is a concatenation of q, k, and v weights
61 - # along the output dimension.
62 - qkv_weight = qkv_linear.weight.data.cpu()
63 - q_weight = Parameter(qkv_weight[:q_proj_size], requires_grad=False)
64 - k_weight = Parameter(
65 - qkv_weight[q_proj_size : q_proj_size + k_proj_size], requires_grad=False
66 - )
67 - v_weight = Parameter(
68 - qkv_weight[q_proj_size + k_proj_size :], requires_grad=False
69 - )
70 - self.register_parameter("q_weight", q_weight)
71 - self.register_parameter("k_weight", k_weight)
72 - self.register_parameter("v_weight", v_weight)
73 -
74 - if qkv_linear.bias is not None:
75 - q_bias = Parameter(qkv_linear.bias[:q_proj_size], requires_grad=False)
76 - k_bias = Parameter(
77 - qkv_linear.bias[q_proj_size : q_proj_size + k_proj_size],
78 - requires_grad=False,
79 - )
80 - v_bias = Parameter(
81 - qkv_linear.bias[q_proj_size + k_proj_size :], requires_grad=False
82 - )
83 - self.register_parameter("q_bias", q_bias)
84 - self.register_parameter("k_bias", k_bias)
85 - self.register_parameter("v_bias", v_bias)
86 - else:
87 - self.register_parameter("q_bias", None)
88 - self.register_parameter("k_bias", None)
89 - self.register_parameter("v_bias", None)
90 -
91 - def forward(self, input):
92 - # Same forward functionality as QKVParallelLinear, but doing qkv porj
93 - # separately.
94 - q_bias = self.q_bias if not self.skip_bias_add else None
95 - k_bias = self.k_bias if not self.skip_bias_add else None
96 - v_bias = self.v_bias if not self.skip_bias_add else None
97 - q_proj = F.linear(input, self.q_weight, q_bias)
98 - k_proj = F.linear(input, self.k_weight, k_bias)
99 - v_proj = F.linear(input, self.v_weight, v_bias)
100 - # The q/k/v projections will be split outside of the QKVParallelLinear.
101 - # Because we are replacing XlaQKVParallelLinear with the
102 - # QKVParallelLinear, we need to concatenate q, k, and v projections to
103 - # match the output shape of the QKVParallelLinear implementation even if
104 - # it seems to be redundant.
105 - # The concat and the following split will be noop, and should be
106 - # optimized away by the compiler.
107 - qkv_proj = torch.cat([q_proj, k_proj, v_proj], dim=-1)
108 - output_bias = (
109 - torch.cat([q_bias, k_bias, v_bias], dim=-1) if self.skip_bias_add else None
110 - )
111 - if not self.return_bias:
112 - return qkv_proj
113 - return qkv_proj, output_bias
114 -
115 -
116 -def partition_column_parallel_linear(
117 - layer: torch.nn.Module, mesh: xs.Mesh
118 -) -> torch.nn.Module:
119 - assert isinstance(layer, ColumnParallelLinear)
120 - xs.mark_sharding(layer.weight, mesh, ("x", None))
121 - logger.debug("Applied column-parallel sharding to %s", layer)
122 - return layer
123 -
124 -
125 -def partition_row_parallel_linear(
126 - layer: torch.nn.Module, mesh: xs.Mesh
127 -) -> torch.nn.Module:
128 - assert isinstance(layer, RowParallelLinear)
129 - xs.mark_sharding(layer.weight, mesh, (None, "x"))
130 - logger.debug("Applied row-parallel sharding to %s", layer)
131 - return layer
132 -
133 -
134 -def partition_qkv_parallel_linear(
135 - layer: torch.nn.Module, mesh: xs.Mesh
136 -) -> torch.nn.Module:
137 - assert isinstance(layer, QKVParallelLinear)
138 - xla_layer = XlaQKVParallelLinear(layer, mesh)
139 - logger.debug("Applied qkv parallel sharding to %s", layer)
140 - return xla_layer
141 -
142 -
143 -MODULE_TYPE_TO_WRAPPING_FUNC = OrderedDict(
144 - [
145 - ("QKVParallelLinear", partition_qkv_parallel_linear),
146 - ("ColumnParallelLinear", partition_column_parallel_linear),
147 - ("RowParallelLinear", partition_row_parallel_linear),
148 - ]
149 -)
150 -
151 -
152 -def get_fqn(module):
153 - # Get the fully qualified name of the module
154 - return module.__class__.__qualname__
155 -
156 -
157 -def shard_model(model: torch.nn.Module, mesh: "xs.Mesh") -> None:
158 - """
159 - Recursively check a PyTorch model and apply appropriate sharding based on
160 - the MODULE_TYPE_TO_WRAPPING_FUNC mapping.
161 -
162 - Args:
163 - model: torch.nn.Module to process
164 - mesh: An XLA SPMD mesh object used for sharding
165 - """
166 -
167 - def _process_module(module, name=None, parent=None):
168 - for module_type, wrapping_func in MODULE_TYPE_TO_WRAPPING_FUNC.items():
169 - if get_fqn(module) == module_type:
170 - wrapped_module = wrapping_func(module, mesh)
171 -
172 - assert parent is not None and name is not None, (
173 - "Top Level module is not expected to be wrapped."
174 - )
175 - if wrapped_module is not module:
176 - # Wrapped module and module are different py object.
177 - # The original module should be replaced by the
178 - # wrapped_module.
179 - logger.debug("replace %s with %s", module, wrapped_module)
180 - setattr(parent, name, wrapped_module)
181 -
182 - module = wrapped_module
183 - break
184 -
185 - for child_name, child_module in list(module.named_children()):
186 - _process_module(child_module, child_name, module)
187 -
188 - _process_module(model)
vllm/lora/ops/xla_ops/__init__.py deleted
-6
@@ -1,6 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -from vllm.lora.ops.xla_ops.lora_ops import bgmv_expand, bgmv_expand_slice, bgmv_shrink
5 -
6 -__all__ = ["bgmv_expand", "bgmv_expand_slice", "bgmv_shrink"]
vllm/lora/ops/xla_ops/lora_ops.py deleted
-141
@@ -1,141 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -import jax
5 -import jax.numpy as jnp
6 -import torch
7 -import torch.nn.functional as F
8 -import torch_xla.core.xla_builder as xb
9 -from torch.library import impl
10 -from torch_xla.experimental.custom_kernel import XLA_LIB, jax_import_guard
11 -
12 -
13 -@jax.jit
14 -def bgmv_jax(inputs, loras, idxs):
15 - return jnp.einsum(
16 - "td,tX,Xld->tl",
17 - inputs,
18 - jax.nn.one_hot(idxs, loras.shape[0], dtype=inputs.dtype),
19 - loras,
20 - )
21 -
22 -
23 -XLA_LIB.define("bgmv(Tensor inputs, Tensor loras, Tensor idxs) -> Tensor")
24 -
25 -
26 -@impl(XLA_LIB, "bgmv", "XLA")
27 -def bgmv_xla(inputs: torch.Tensor, loras: torch.Tensor, idxs: torch.IntTensor):
28 - if len(loras.shape) == 4:
29 - loras = loras.squeeze(axis=1)
30 -
31 - jax_import_guard()
32 - return xb.call_jax(bgmv_jax, (inputs, loras, idxs))
33 -
34 -
35 -@impl(XLA_LIB, "bgmv", "CompositeExplicitAutograd")
36 -def bgmv_non_xla(inputs: torch.Tensor, loras: torch.Tensor, idxs: torch.IntTensor):
37 - T, _ = inputs.shape
38 - if len(loras.shape) == 4:
39 - loras = loras.squeeze(axis=1)
40 - _, L, _ = loras.shape
41 -
42 - return torch.empty((T, L), device=inputs.device)
43 -
44 -
45 -def bgmv_expand(
46 - inputs: torch.Tensor,
47 - lora_b_weights: torch.Tensor,
48 - output_tensor: torch.Tensor,
49 - lora_indices_tensor: torch.Tensor,
50 - add_inputs: bool = True,
51 -):
52 - """
53 - Args:
54 - inputs (torch.Tensor): Input tensor of shape [num_tokens, hidden_size].
55 -
56 - lora_b_weights (torch.Tensor): LoRA weights of shape
57 - [num_loras, lora_rank, hidden_size].
58 -
59 - output_tensor (torch.Tensor): output tensor of shape
60 - [num_tokens, hidden_size * num_slices].
61 -
62 - lora_indices_tensor (torch.Tensor): Tensor of shape [num_tokens]
63 - indicating which LoRA matrix to use for each token.
64 - add_inputs (bool): Whether or not to add the input tensor to the output
65 - tensor.
66 - """
67 -
68 - outputs = torch.ops.xla.bgmv(inputs, lora_b_weights, lora_indices_tensor)
69 -
70 - limit = output_tensor.shape[0]
71 - if outputs.shape[0] == 1 and output_tensor.shape[0] != 1:
72 - limit = 1
73 -
74 - if output_tensor.shape[1] > outputs.shape[1]:
75 - outputs = F.pad(outputs, (0, output_tensor.shape[1] - outputs.shape[1], 0, 0))
76 -
77 - if add_inputs:
78 - return output_tensor + outputs[:limit, : output_tensor.shape[1]]
79 - else:
80 - return outputs[:limit, : output_tensor.shape[1]]
81 -
82 -
83 -def bgmv_shrink(
84 - inputs: torch.Tensor,
85 - lora_b_weights: torch.Tensor,
86 - lora_indices_tensor: torch.Tensor,
87 - scaling: float = 1.0,
88 -):
89 - """
90 - Args:
91 - inputs (torch.Tensor): Input tensor of shape [num_tokens, hidden_size].
92 - lora_b_weights (torch.Tensor): LoRA weights of shape
93 - [num_loras, lora_rank, hidden_size].
94 - lora_indices_tensor (torch.Tensor): Tensor of shape [num_tokens]
95 - indicating which LoRA matrix to use for each token.
96 - scaling (float, optional): Scalar multiplier applied to the output.
97 - """
98 -
99 - return scaling * torch.ops.xla.bgmv(inputs, lora_b_weights, lora_indices_tensor)
100 -
101 -
102 -def bgmv_expand_slice(
103 - inputs: torch.Tensor,
104 - lora_b_weights: torch.Tensor,
105 - output_tensor: torch.Tensor,
106 - lora_indices_tensor: torch.Tensor,
107 - slice_offset: int,
108 - slice_size: int,
109 - add_inputs: bool = True,
110 -):
111 - """
112 - Args:
113 - inputs (torch.Tensor): Input tensor of shape [num_tokens, hidden_size].
114 -
115 - lora_b_weights (torch.Tensor): LoRA weights of shape
116 - [num_loras, lora_rank, hidden_size].
117 -
118 - output_tensor (torch.Tensor): output tensor of shape
119 - [num_tokens, hidden_size * num_slices].
120 -
121 - lora_indices_tensor (torch.Tensor): Tensor of shape [num_tokens]
122 - indicating which LoRA matrix to use for each token.
123 - add_inputs (bool): Whether or not to add the input tensor to the output
124 - tensor.
125 - """
126 - outputs = torch.ops.xla.bgmv(inputs, lora_b_weights, lora_indices_tensor)
127 -
128 - outputs = F.pad(
129 - outputs,
130 - (
131 - slice_offset,
132 - output_tensor.shape[1] - (slice_offset + slice_size),
133 - 0,
134 - 0,
135 - ),
136 - )
137 -
138 - if add_inputs:
139 - return output_tensor + outputs
140 - else:
141 - return outputs
vllm/lora/punica_wrapper/punica_tpu.py deleted
-358
@@ -1,358 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -import math
5 -from typing import TYPE_CHECKING
6 -
7 -import torch
8 -import torch.nn.functional as F
9 -import torch_xla
10 -
11 -from vllm.lora.ops.xla_ops import bgmv_expand, bgmv_expand_slice, bgmv_shrink
12 -from vllm.lora.punica_wrapper.utils import convert_mapping
13 -
14 -if TYPE_CHECKING:
15 - # avoid circuit import
16 - from vllm.lora.layers import LoRAMapping
17 -
18 -from .punica_base import PunicaWrapperBase
19 -
20 -
21 -class PunicaWrapperTPU(PunicaWrapperBase):
22 - """
23 - PunicaWrapperTPU is designed to manage and provide metadata for the punica
24 - kernel. The main function is to maintain the state information for
25 - Multi-LoRA, and to provide the interface for the pytorch punica ops.
26 - """
27 -
28 - def __init__(
29 - self,
30 - max_num_batched_tokens: int,
31 - max_batches: int,
32 - device: torch.device | str,
33 - **kwargs,
34 - ):
35 - PunicaWrapperBase.__init__(self, max_num_batched_tokens, max_batches, device)
36 -
37 - # PunicaWrapperBase defines some tensors with dtype=torch.int64, which
38 - # isn't supported by the TPU. So convert those tensors to int32.
39 - # Not all of them are used by the TPU so only convert the useful ones.
40 - self._token_lora_indices = self._token_lora_indices.to(dtype=torch.int32)
41 - self._sampler_indices = self._sampler_indices.to(dtype=torch.int32)
42 - self._sampler_indices_padded = self._sampler_indices_padded.to(
43 - dtype=torch.int32
44 - )
45 -
46 - torch.ops.xla.dynamo_set_buffer_donor_(self._token_lora_indices, True)
47 - torch.ops.xla.dynamo_set_buffer_donor_(self._sampler_indices, True)
48 - torch.ops.xla.dynamo_set_buffer_donor_(self._sampler_indices_padded, True)
49 - torch.ops.xla.dynamo_set_buffer_donor_(self._embeddings_indices, True)
50 - torch.ops.xla.dynamo_set_buffer_donor_(self._lora_indices_per_batch, True)
51 -
52 - torch._dynamo.mark_dynamic(self._token_lora_indices, 0)
53 - torch._dynamo.mark_dynamic(self._embeddings_indices, 1)
54 - torch._dynamo.mark_dynamic(self._sampler_indices_padded, 0)
55 -
56 - def _get_token_lora_indices(self, x: torch.Tensor) -> torch.IntTensor:
57 - return torch.narrow(self._token_lora_indices, 0, 0, x.size(0))
58 -
59 - @property
60 - def embeddings_indices(self) -> torch.Tensor:
61 - """
62 - This property provides access to the indices used for lora embeddings,
63 - specifically for VocabParallelEmbeddingWithLoRA.
64 - """
65 - return self._embeddings_indices[:]
66 -
67 - @property
68 - def sampler_indices_padded(self) -> torch.Tensor:
69 - """
70 - This property provides access to padded sampler indices.
71 - """
72 - return self._sampler_indices_padded[:]
73 -
74 - def shrink(
75 - self,
76 - x: torch.Tensor,
77 - w_t_all: torch.Tensor,
78 - scale: float,
79 - ):
80 - return bgmv_shrink(x, w_t_all, self._get_token_lora_indices(x), scale)
81 -
82 - def expand(
83 - self, y: torch.Tensor, x: torch.Tensor, w_t_all: torch.Tensor, add_inputs: bool
84 - ):
85 - return bgmv_expand(x, w_t_all, y, self._get_token_lora_indices(x), add_inputs)
86 -
87 - def expand_slice(
88 - self,
89 - y: torch.Tensor,
90 - x: torch.Tensor,
91 - w_t_all: torch.Tensor,
92 - y_offset: int,
93 - y_slice_size: int,
94 - add_inputs: bool,
95 - ) -> torch.Tensor:
96 - return bgmv_expand_slice(
97 - x,
98 - w_t_all,
99 - y,
100 - self._get_token_lora_indices(x),
101 - y_offset,
102 - y_slice_size,
103 - add_inputs,
104 - )
105 -
106 - def add_shrink(
107 - self,
108 - y: tuple[torch.Tensor, ...] | torch.Tensor,
109 - x: torch.Tensor,
110 - lora_a_stacked: tuple[torch.Tensor, ...],
111 - scale: float,
112 - **kwargs,
113 - ) -> torch.Tensor | None:
114 - """
115 - Performs GEMM for multiple slices of lora_a.
116 -
117 - Semantics:
118 - for i in range(len(lora_a_stacked)):
119 - y[i] += (x @ lora_a_stacked[i]) * scale
120 -
121 - Args:
122 - y (Union[tuple[torch.Tensor, ...], torch.Tensor]): Output tensors
123 - x (torch.Tensor): Input tensor
124 - lora_a_stacked (tuple[torch.Tensor, ...]): lora_a's weights
125 - scale (float): Scaling factor for the operation
126 - """
127 -
128 - torch.ops.xla.dynamo_set_buffer_donor_(y, True)
129 - x = x.view(-1, x.shape[-1])
130 -
131 - for slice_idx in range(len(lora_a_stacked)):
132 - lora_s = lora_a_stacked[slice_idx]
133 - y_s = self.shrink(x, lora_s, scale)
134 - y[slice_idx, :, :] = y_s # type: ignore[index]
135 - return y
136 -
137 - def add_expand(
138 - self,
139 - y: torch.Tensor,
140 - x: tuple[torch.Tensor, ...] | torch.Tensor,
141 - lora_b_stacked: tuple[torch.Tensor, ...],
142 - output_slices: tuple[int, ...],
143 - offset_start: int = 0,
144 - add_inputs=True,
145 - **kwargs,
146 - ) -> torch.Tensor:
147 - """
148 - Performs GEMM for multiple slices of lora_b.
149 -
150 - Semantics:
151 - for i in range(len(lora_b_stacked)):
152 - slice = output_slices[i]
153 - y[:, offset:offset+slice] += x[i] @ lora_b_stacked[i]
154 - offset += slice
155 -
156 - Args:
157 - y (torch.Tensor): Output tensor.
158 - x (Union[tuple[torch.Tensor, ...], torch.Tensor]): Input tensors
159 - lora_b_stacked (tuple[torch.Tensor, ...]): lora_b's weight
160 - output_slices (tuple[int, ...]): Every slice's size
161 - add_inputs (bool): Defaults to True.
162 - """
163 - y_org = y
164 - y = y.view(-1, y.shape[-1])
165 - offset_left = 0
166 -
167 - for slice_idx in range(len(lora_b_stacked)):
168 - y = self.expand_slice(
169 - y,
170 - x[slice_idx],
171 - lora_b_stacked[slice_idx],
172 - offset_left,
173 - output_slices[slice_idx],
174 - add_inputs=add_inputs,
175 - )
176 - offset_left += output_slices[slice_idx]
177 - return y.view_as(y_org)
178 -
179 - def add_lora_embedding(
180 - self,
181 - y: torch.Tensor,
182 - x: torch.Tensor,
183 - lora_b_stacked: torch.Tensor,
184 - add_inputs: bool = True,
185 - **kwargs,
186 - ) -> torch.Tensor:
187 - """
188 - Applies lora specifically for VocabParallelEmbeddingWithLoRA.
189 -
190 - Semantics:
191 - y += x @ lora_b_stacked
192 -
193 - Args:
194 - y (torch.Tensor): Output tensor.
195 - x (torch.Tensor): Input tensor.
196 - lora_b_stacked (torch.Tensor): lora_b's weights.
197 - add_inputs (bool): Default to True.
198 - """
199 -
200 - # Embedding layer only needs the expand op
201 - return self.expand(y, x, lora_b_stacked, add_inputs)
202 -
203 - def add_lora_linear(
204 - self,
205 - y: torch.Tensor,
206 - x: torch.Tensor,
207 - lora_a_stacked: tuple[torch.Tensor, ...],
208 - lora_b_stacked: tuple[torch.Tensor, ...],
209 - scale: float,
210 - output_slices: tuple[int, ...],
211 - *,
212 - buffer: tuple[torch.Tensor, ...] | None = None,
213 - **kwargs,
214 - ) -> torch.Tensor:
215 - """
216 - Applicable to linear-related lora.
217 -
218 - Semantics:
219 - for i in range(len(lora_a_stacked)):
220 - y[i] += (
221 - x[i].unsqueeze(0)
222 - @ lora_a_stacked[indices[i], layer_idx, :, :]
223 - @ lora_b_stacked[indices[i], layer_idx, :, :]
224 - * scale
225 - ).squeeze(0)
226 -
227 - Args:
228 - y (torch.Tensor): Output tensor. Will not be changed in-place.
229 - x (torch.Tensor): Input tensor (T, E)
230 - lora_a_stacked (tuple[torch.Tensor, ...]): lora_a's weight.
231 - lora_b_stacked (tuple[torch.Tensor, ...]): lora_b's weight.
232 - scale (float): Scaling factor.
233 - output_slices (tuple[int, ...]): Every slice's size.
234 - buffer (Optional[tuple[torch.Tensor, ...]]): Defaults to None.
235 - """
236 -
237 - assert len(lora_a_stacked) == len(lora_b_stacked) == len(output_slices)
238 -
239 - if buffer is None:
240 - r = lora_b_stacked[0].size(-1)
241 - T = x.size(0)
242 - buffer = torch.zeros(
243 - (len(output_slices), T, r),
244 - dtype=x.dtype,
245 - device=x.device,
246 - )
247 - buffer = self.add_shrink(buffer, x, lora_a_stacked, scale, **kwargs)
248 - return self.add_expand(
249 - y, buffer, lora_b_stacked, output_slices, add_inputs=True, **kwargs
250 - )
251 -
252 - def add_lora_logits(
253 - self,
254 - y: torch.Tensor,
255 - x: torch.Tensor,
256 - lora_a_stacked: torch.Tensor,
257 - lora_b_stacked: torch.Tensor,
258 - scale,
259 - *,
260 - buffer: torch.Tensor | None = None,
261 - **kwargs,
262 - ) -> torch.Tensor:
263 - """
264 - Applies lora specifically for LogitsProcessorWithLoRA.
265 -
266 - Semantics:
267 - buffer = (x @ lora_a_stacked) * scale
268 - y += buffer @ lora_b_stacked
269 -
270 - Args:
271 - y (torch.Tensor): Output tensor.
272 - x (torch.Tensor): Input tensor.
273 - lora_a_stacked (torch.Tensor): lora_a's weights.
274 - lora_b_stacked (torch.Tensor):lora_b's weights.
275 - scale (float): Scaling factor.
276 - buffer (Optional[torch.Tensor]):Default to None.
277 - """
278 - y_org = y
279 - y = y.view(-1, y.shape[-1])
280 - x = x.view(-1, x.shape[-1])
281 -
282 - sampler_indices = torch.narrow(self._sampler_indices, 0, 0, x.size(0))
283 - buffer = bgmv_shrink(x, lora_a_stacked, sampler_indices, scale)
284 - y = bgmv_expand(buffer, lora_b_stacked, y, sampler_indices, add_inputs=True)
285 - return y.view_as(y_org)
286 -
287 - # This performs the same tensor ops as the base method, except it does them
288 - # on the CPU then transfers the results to the TPU
289 - def _update_base_metadata(
290 - self,
291 - mapping: "LoRAMapping",
292 - lora_index_to_id: list[int | None],
293 - max_loras: int,
294 - vocab_size: int,
295 - ):
296 - # Make sure we don't accidentally collect outside operations
297 - torch_xla.sync()
298 -
299 - # Pad the prompt mapping to avoid running into recompiles on the TPU
300 - # TODO: Should this happen inside mapping internally? If so how can we
301 - # avoid having backend specific LoRAMapping classes?
302 - mapping.prompt_mapping = self._pad_prompt_mapping(mapping.prompt_mapping)
303 -
304 - (
305 - base_indices,
306 - sampler_indices,
307 - sampler_indices_padded,
308 - embeddings_indices,
309 - indices_len,
310 - ) = convert_mapping(
311 - mapping,
312 - lora_index_to_id,
313 - max_loras,
314 - vocab_size,
315 - 0, # extra_vocab_size
316 - "cpu",
317 - )
318 - self._token_lora_indices = self._pad_to_shape(
319 - base_indices, self._token_lora_indices.shape, dims=1
320 - ).to(self.device)
321 - self._sampler_indices = self._pad_to_shape(
322 - sampler_indices, self._sampler_indices.shape, dims=1
323 - ).to(self.device)
324 - self._sampler_indices_padded = self._pad_to_shape(
325 - sampler_indices_padded, self._sampler_indices_padded.shape, dims=1
326 - ).to(self.device)
327 - self._embeddings_indices = self._pad_to_shape(
328 - embeddings_indices, self._embeddings_indices.shape, dims=2
329 - ).to(self.device)
330 - self.indices_len[:] = indices_len
331 -
332 - def _update_prefill_metadata(self, token_lora_tensor: torch.Tensor) -> None:
333 - self.batch_size = 1
334 - self._lora_indices_per_batch[: self.batch_size] = token_lora_tensor[
335 - : self.batch_size
336 - ]
337 -
338 - def _pad_prompt_mapping(self, prompt_mapping: tuple[int, ...]) -> tuple[int, ...]:
339 - num_reqs = len(prompt_mapping)
340 -
341 - # From vllm/v1/worker/tpu_model_runner:51, but need to avoid a circular
342 - # import
343 - MIN_NUM_SEQS = 8
344 -
345 - padded_num_reqs = max(2 ** math.ceil(math.log2(num_reqs)), MIN_NUM_SEQS)
346 - pad_len = padded_num_reqs - num_reqs
347 -
348 - padding = [-1] * pad_len
349 - return tuple(list(prompt_mapping) + padding)
350 -
351 - def _pad_to_shape(self, src, target_shape, dims=1):
352 - if dims == 1:
353 - pad_len = target_shape[0] - src.shape[0]
354 - return F.pad(src, (0, pad_len), value=0).to(torch.int32)
355 - else:
356 - pad_rows = target_shape[0] - src.shape[0]
357 - pad_cols = target_shape[1] - src.shape[1]
358 - return F.pad(src, (0, pad_cols, 0, pad_rows), value=0).to(torch.int32)
vllm/model_executor/layers/fused_moe/layer.py
-6
@@ -66,12 +66,6 @@ else:
66
67 eplb_map_to_physical_and_record = _eplb_map_to_physical_and_record
68 from vllm.model_executor.layers.fused_moe.fused_moe import GroupedTopk
69 -
70 -if current_platform.is_tpu():
71 - from .moe_pallas import fused_moe as fused_moe_pallas
72 -else:
73 - fused_moe_pallas = None # type: ignore
74 -
69 from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
70 FusedMoEMethodBase,
71 )
vllm/model_executor/layers/fused_moe/moe_pallas.py deleted
-83
@@ -1,83 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -import torch
5 -import torch.nn.functional as F
6 -
7 -
8 -def _histogram(input: torch.Tensor, min: int, max: int) -> torch.Tensor:
9 - """
10 - Compute the histogram of an int32 tensor. The bin edges are defined by the
11 - min and max values, with step = 1.
12 - """
13 - assert input.dtype == torch.int32, "input must be of torch.int32 dtype."
14 - assert min <= max, "min must be less than or equal to max."
15 -
16 - def searchsorted(
17 - sorted_sequence: torch.Tensor, values_to_search: torch.Tensor
18 - ) -> torch.Tensor:
19 - return (sorted_sequence.unsqueeze(1) == values_to_search).sum(dim=1)
20 -
21 - bin_edges = torch.linspace(min, max, max - min + 1, dtype=input.dtype).to(
22 - input.device
23 - )
24 - return searchsorted(bin_edges, input).to(torch.int32)
25 -
26 -
27 -def fused_moe(
28 - hidden_states: torch.Tensor,
29 - w1: torch.Tensor,
30 - w2: torch.Tensor,
31 - gating_output: torch.Tensor,
32 - topk: int,
33 - global_num_experts: int,
34 - expert_map: torch.Tensor = None,
35 - renormalize: bool = False,
36 -) -> torch.Tensor:
37 - """
38 - Args:
39 - hidden_states: [*, hidden_size]
40 - w1: [num_experts, intermediate_size * 2, hidden_size]
41 - w2: [num_experts, hidden_size, intermediate_size]
42 - gating_output: [*, num_experts]
43 - """
44 - assert expert_map is None, "expert_map is not supported for pallas MoE."
45 - import torch_xla.experimental.custom_kernel # noqa: F401
46 -
47 - orig_shape = hidden_states.shape
48 - hidden_size = hidden_states.shape[-1]
49 - num_tokens = hidden_states.shape[:-1].numel()
50 - num_experts = w1.shape[0]
51 - intermediate_size = w2.shape[-1]
52 - device = hidden_states.device
53 - dtype = hidden_states.dtype
54 - assert (num_tokens * topk) % 16 == 0, (
55 - "The Pallas GMM kernel requires num_tokens * topk to be a multiple of "
56 - f"16 but got {num_tokens * topk}"
57 - )
58 -
59 - hidden_states = hidden_states.view(num_tokens, hidden_size)
60 - gating_output = gating_output.view(num_tokens, num_experts)
61 - topk_weights = gating_output.softmax(dim=-1, dtype=torch.float)
62 - topk_weights, topk_indices = topk_weights.topk(topk, dim=-1)
63 - if renormalize:
64 - topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
65 - topk_weights = topk_weights.to(dtype)
66 -
67 - topk_indices = topk_indices.flatten()
68 - topk_argsort_indices = topk_indices.argsort()
69 - topk_argsort_revert_indices = topk_argsort_indices.argsort()
70 - token_indices = torch.arange(num_tokens, device=device).repeat_interleave(topk)
71 - token_indices = token_indices[topk_argsort_indices]
72 - group_sizes = _histogram(topk_indices.to(torch.int32), 0, num_experts - 1)
73 -
74 - x = hidden_states[token_indices]
75 - x = torch.ops.xla.gmm(x, w1, group_sizes, transpose_rhs=True)
76 - x = F.silu(x[..., :intermediate_size]) * x[..., intermediate_size:]
77 - x = torch.ops.xla.gmm(x, w2, group_sizes, transpose_rhs=True)
78 - x = x[topk_argsort_revert_indices].reshape(-1, topk, hidden_size)
79 -
80 - x = x * topk_weights.unsqueeze(dim=-1)
81 - x = x.sum(dim=-2)
82 - x = x.reshape(orig_shape)
83 - return x
vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
+1 -51
@@ -47,10 +47,6 @@ if current_platform.is_cuda_alike():
47 else:
48 TritonExperts = None # type: ignore
49
50 -if current_platform.is_tpu():
51 - from .moe_pallas import fused_moe as fused_moe_pallas
52 -else:
53 - fused_moe_pallas = None # type: ignore
50
51 logger = init_logger(__name__)
52
@@ -390,53 +386,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
386 custom_routing_function=layer.custom_routing_function,
387 )
388
393 - def forward_tpu(
394 - self,
395 - layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821
396 - x: torch.Tensor,
397 - router_logits: torch.Tensor,
398 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
399 - assert not layer.use_grouped_topk
400 - assert layer.num_expert_group is None
401 - assert layer.topk_group is None
402 - assert layer.custom_routing_function is None
403 - assert layer.apply_router_weight_on_input is False
404 - if layer.scoring_func != "softmax":
405 - raise NotImplementedError(
406 - "Only softmax scoring function is supported for TPU."
407 - )
408 - if layer.e_score_correction_bias is not None:
409 - raise NotImplementedError(
410 - "Expert score correction bias is not supported for TPU."
411 - )
412 - assert layer.activation == "silu", (
413 - f"{layer.activation} is not supported for TPU."
414 - )
415 - assert layer.routed_scaling_factor == 1.0, (
416 - f"routed_scaling_factor {layer.routed_scaling_factor} is "
417 - "not supported for TPU."
418 - )
419 - if (
420 - layer.enable_eplb is not False
421 - or layer.expert_load_view is not None
422 - or layer.logical_to_physical_map is not None
423 - or layer.logical_replica_count is not None
424 - ):
425 - raise NotImplementedError("Expert load balancing is not supported for TPU.")
426 - return fused_moe_pallas(
427 - hidden_states=x,
428 - w1=layer.w13_weight,
429 - w2=layer.w2_weight,
430 - topk=layer.top_k,
431 - gating_output=router_logits,
432 - global_num_experts=layer.global_num_experts,
433 - expert_map=layer.expert_map,
434 - renormalize=layer.renormalize,
435 - )
436 -
437 - if current_platform.is_tpu():
438 - forward_native = forward_tpu
439 - elif current_platform.is_cpu():
389 + if current_platform.is_cpu():
390 forward_native = forward_cpu
391 elif current_platform.is_xpu():
392 forward_native = forward_xpu
vllm/model_executor/layers/quantization/__init__.py
-3
@@ -11,7 +11,6 @@ logger = init_logger(__name__)
11 QuantizationMethods = Literal[
12 "awq",
13 "deepspeedfp",
14 - "tpu_int8",
14 "fp8",
15 "ptpc_fp8",
16 "fbgemm_fp8",
@@ -129,12 +128,10 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
128 from .ptpc_fp8 import PTPCFp8Config
129 from .rtn import RTNConfig
130 from .torchao import TorchAOConfig
132 - from .tpu_int8 import Int8TpuConfig
131
132 method_to_config: dict[str, type[QuantizationConfig]] = {
133 "awq": AWQConfig,
134 "deepspeedfp": DeepSpeedFPConfig,
137 - "tpu_int8": Int8TpuConfig,
135 "fp8": Fp8Config,
136 "fbgemm_fp8": FBGEMMFp8Config,
137 "fp_quant": FPQuantConfig,
vllm/model_executor/layers/quantization/kernels/scaled_mm/__init__.py
-4
@@ -19,9 +19,6 @@ from vllm.model_executor.layers.quantization.kernels.scaled_mm.ScaledMMLinearKer
19 from vllm.model_executor.layers.quantization.kernels.scaled_mm.triton import (
20 TritonScaledMMLinearKernel,
21 )
22 -from vllm.model_executor.layers.quantization.kernels.scaled_mm.xla import (
23 - XLAScaledMMLinearKernel,
24 -)
22 from vllm.platforms import PlatformEnum, current_platform
23
24 # in priority/performance order (when available)
@@ -29,7 +26,6 @@ _POSSIBLE_KERNELS: dict[PlatformEnum, list[type[ScaledMMLinearKernel]]] = {
26 PlatformEnum.CPU: [CPUScaledMMLinearKernel],
27 PlatformEnum.CUDA: [CutlassScaledMMLinearKernel, TritonScaledMMLinearKernel],
28 PlatformEnum.ROCM: [AiterScaledMMLinearKernel, TritonScaledMMLinearKernel],
32 - PlatformEnum.TPU: [XLAScaledMMLinearKernel],
29 }
30
31
vllm/model_executor/layers/quantization/kernels/scaled_mm/xla.py deleted
-106
@@ -1,106 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -import warnings
5 -
6 -import torch
7 -from functorch.experimental.control_flow import cond # noqa: F401
8 -
9 -from vllm.model_executor.layers.quantization.utils import replace_parameter
10 -from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
11 - convert_to_channelwise,
12 -)
13 -from vllm.platforms import current_platform
14 -
15 -from .ScaledMMLinearKernel import ScaledMMLinearKernel, ScaledMMLinearLayerConfig
16 -
17 -
18 -class XLAScaledMMLinearKernel(ScaledMMLinearKernel):
19 - @classmethod
20 - def is_supported(
21 - cls, compute_capability: int | None = None
22 - ) -> tuple[bool, str | None]:
23 - if not current_platform.is_tpu():
24 - return False, "Requires TPU."
25 - return True, None
26 -
27 - @classmethod
28 - def can_implement(cls, c: ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
29 - if not current_platform.is_tpu():
30 - return False, "ScaledMMXLA requires running on TPU."
31 -
32 - if c.is_static_input_scheme:
33 - return False, "ScaledMMXLA requires dynamic activation scales."
34 -
35 - if not c.input_symmetric:
36 - return False, "ScaledMMXLA requires symmetric activation scales."
37 -
38 - if not c.is_channelwise:
39 - return False, "ScaledMMXLA requires channelwise weight scales"
40 -
41 - return True, None
42 -
43 - def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
44 - # WEIGHT
45 - # [out, in] (different than cutlass_scaled_mm)
46 - weight = getattr(layer, self.w_q_name)
47 - replace_parameter(
48 - layer, self.w_q_name, torch.nn.Parameter(weight.data, requires_grad=False)
49 - )
50 -
51 - # WEIGHT SCALE
52 - # XLA kernels support only per-tensor and per-channel.
53 - # If we have a fused module (QKV, MLP) with per tensor scales (thus N
54 - # scales being passed to the kernel), convert to the per-channel case.
55 - is_fused_module = len(layer.logical_widths) > 1
56 - weight_scale = getattr(layer, self.w_s_name)
57 - if is_fused_module and not self.config.is_channelwise:
58 - weight_scale = convert_to_channelwise(weight_scale, layer.logical_widths)
59 -
60 - # [out_channel,] (different than cutlass_scaled_mm)
61 - weight_scale = weight_scale.squeeze(-1)
62 - replace_parameter(
63 - layer,
64 - self.w_s_name,
65 - torch.nn.Parameter(weight_scale.data, requires_grad=False),
66 - )
67 -
68 - # Only support symmetric dynamic activation quantization.
69 - setattr(layer, self.i_s_name, None)
70 - setattr(layer, self.i_zp_name, None)
71 - setattr(layer, self.azp_adj_name, None)
72 -
73 - # Filter warning for cond usage in apply_weights. It is okay
74 - # to specialize the graph since bias is not dynamic.
75 - warnings.filterwarnings(
76 - "ignore",
77 - message="Pred is a Python constant. When used with torch.cond, it specializes on one of the branches.", # noqa: E501
78 - )
79 -
80 - def no_add_bias(self, x: torch.Tensor, bias: torch.Tensor | None):
81 - return x
82 -
83 - def add_bias(self, x: torch.Tensor, bias: torch.Tensor | None):
84 - return x + bias
85 -
86 - def apply_weights(
87 - self,
88 - layer: torch.nn.Module,
89 - x: torch.Tensor,
90 - bias: torch.Tensor | None = None,
91 - ) -> torch.Tensor:
92 - w_q, w_s, _, _, _ = self._get_weight_params(layer)
93 -
94 - # Required to register custom ops.
95 - import torch_xla.experimental.custom_kernel # noqa: F401
96 -
97 - out = torch.ops.xla.quantized_matmul_int8(
98 - x,
99 - w_q,
100 - w_s,
101 - quantize_activation=True,
102 - )
103 -
104 - # Explicitly capture control flow to make dynamo happy.
105 - # https://pytorch.org/docs/main/generated/exportdb/index.html#cond-branch-class-method # noqa: E501
106 - return cond(bias is None, self.no_add_bias, self.add_bias, [out, bias])
vllm/model_executor/layers/quantization/tpu_int8.py deleted
-139
@@ -1,139 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -from typing import Any, Optional
5 -
6 -import torch
7 -from torch.nn import Module
8 -from torch.nn.parameter import Parameter
9 -
10 -from vllm.model_executor.layers.linear import LinearBase, LinearMethodBase
11 -from vllm.model_executor.layers.quantization import (
12 - QuantizationConfig,
13 - QuantizationMethods,
14 -)
15 -from vllm.model_executor.parameter import ModelWeightParameter
16 -
17 -ACTIVATION_SCHEMES = ["none", "dynamic"]
18 -
19 -
20 -class Int8TpuConfig(QuantizationConfig):
21 - """Int8 Quantization Config class for TPU Backend."""
22 -
23 - def __init__(
24 - self,
25 - activation_scheme: str = "none",
26 - ) -> None:
27 - super().__init__()
28 - if activation_scheme not in ACTIVATION_SCHEMES:
29 - raise ValueError(f"Unsupported activation scheme {activation_scheme}")
30 - self.activation_scheme = activation_scheme
31 -
32 - def get_name(self) -> QuantizationMethods:
33 - return "tpu_int8"
34 -
35 - def get_supported_act_dtypes(self) -> list[torch.dtype]:
36 - return [torch.float16, torch.bfloat16]
37 -
38 - @classmethod
39 - def get_min_capability(cls) -> int:
40 - raise NotImplementedError("This function should not be called with TPU Backend")
41 -
42 - @staticmethod
43 - def get_config_filenames() -> list[str]:
44 - return []
45 -
46 - @classmethod
47 - def from_config(cls, config: dict[str, Any]) -> "Int8TpuConfig":
48 - activation_scheme = cls.get_from_keys(config, ["activation_scheme"])
49 - return cls(activation_scheme=activation_scheme)
50 -
51 - def get_quant_method(
52 - self, layer: Module, prefix: str
53 - ) -> Optional["TPUInt8LinearMethod"]:
54 - if isinstance(layer, LinearBase):
55 - return TPUInt8LinearMethod(self)
56 - return None
57 -
58 -
59 -class TPUInt8LinearMethod(LinearMethodBase):
60 - """Int8 Linear method for TPU Quant."""
61 -
62 - def __init__(self, quant_config: Int8TpuConfig):
63 - self.quant_config = quant_config
64 - self.quantize_activation = False
65 - if self.quant_config.activation_scheme == "dynamic":
66 - self.quantize_activation = True
67 -
68 - def create_weights(
69 - self,
70 - layer: Module,
71 - input_size_per_partition: int,
72 - output_partition_sizes: list[int],
73 - input_size: int,
74 - output_size: int,
75 - params_dtype: torch.dtype,
76 - **extra_weight_attrs,
77 - ):
78 - weight_loader = extra_weight_attrs.get("weight_loader")
79 - weight = ModelWeightParameter(
80 - data=torch.empty(
81 - sum(output_partition_sizes),
82 - input_size_per_partition,
83 - dtype=params_dtype,
84 - ),
85 - input_dim=1,
86 - output_dim=0,
87 - weight_loader=weight_loader,
88 - )
89 - layer.register_parameter("weight", weight)
90 -
91 - def _quantize_weight(
92 - self, weight: torch.Tensor
93 - ) -> tuple[torch.Tensor, torch.Tensor]:
94 - weight_dtype = weight.dtype
95 - weight = weight.cpu().to(torch.float32)
96 - n_bit = 8
97 - eps = 1e-5
98 - max_int = 2 ** (n_bit - 1) - 1
99 - min_int = -(2 ** (n_bit - 1))
100 - max_val = weight.abs().amax(dim=-1, keepdim=True)
101 - max_val = max_val.clamp(min=eps)
102 - qscale = max_val / max_int
103 - qweight = torch.clamp(
104 - torch.round(weight * (1.0 / qscale)), min_int, max_int
105 - ).to(torch.int8)
106 - qscale = qscale.squeeze().to(weight_dtype)
107 - return qweight, qscale
108 -
109 - def process_weights_after_loading(self, layer: Module) -> None:
110 - layer.weight = Parameter(layer.weight.data, requires_grad=False)
111 - device = layer.weight.device
112 - qweight, qscale = self._quantize_weight(layer.weight)
113 - qweight = qweight.to(device)
114 - qscale = qscale.to(device)
115 - layer.weight = Parameter(qweight, requires_grad=False)
116 - layer.scale = Parameter(qscale, requires_grad=False)
117 -
118 - def apply(
119 - self,
120 - layer: torch.nn.Module,
121 - x: torch.Tensor,
122 - bias: torch.Tensor | None = None,
123 - ) -> torch.Tensor:
124 - try:
125 - import torch_xla.experimental.custom_kernel # noqa: F401
126 - except ImportError as err:
127 - raise ImportError(
128 - "Please install torch_xla by following the instructions at "
129 - "https://docs.vllm.ai/en/latest/getting_started/tpu-installation.html " # noqa: E501
130 - "to run vLLM on TPU."
131 - ) from err
132 - weight = layer.weight
133 - scale = layer.scale
134 - out = torch.ops.xla.quantized_matmul_int8(
135 - x, weight, scale, quantize_activation=self.quantize_activation
136 - )
137 - if bias is not None:
138 - out = out + bias
139 - return out
vllm/model_executor/model_loader/default_loader.py
-17
@@ -30,7 +30,6 @@ from vllm.model_executor.model_loader.weight_utils import (
30 pt_weights_iterator,
31 safetensors_weights_iterator,
32 )
33 -from vllm.platforms import current_platform
33 from vllm.transformers_utils.repo_utils import list_filtered_repo_files
34
35 logger = init_logger(__name__)
@@ -241,22 +240,6 @@ class DefaultModelLoader(BaseModelLoader):
240 self.load_config.pt_load_map_location,
241 )
242
244 - if current_platform.is_tpu():
245 - from vllm.platforms.tpu import USE_TPU_INFERENCE
246 -
247 - if not USE_TPU_INFERENCE:
248 - # In PyTorch XLA, we should call `torch_xla.sync`
249 - # frequently so that not too many ops are accumulated
250 - # in the XLA program.
251 - import torch_xla
252 -
253 - def _xla_weights_iterator(iterator: Generator):
254 - for weights in iterator:
255 - yield weights
256 - torch_xla.sync(wait=False)
257 -
258 - weights_iterator = _xla_weights_iterator(weights_iterator)
259 -
243 if self.counter_before_loading_weights == 0.0:
244 self.counter_before_loading_weights = time.perf_counter()
245 # Apply the prefix.
vllm/model_executor/model_loader/tpu.py deleted
-118
@@ -1,118 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -import time
4 -
5 -import torch
6 -import torch.nn as nn
7 -import torch_xla.core.xla_model as xm
8 -import torch_xla.distributed.spmd as xs
9 -
10 -from vllm.config import ModelConfig, VllmConfig
11 -from vllm.distributed.tpu_distributed_utils import get_fqn, shard_model
12 -from vllm.logger import init_logger
13 -from vllm.model_executor.model_loader.default_loader import DefaultModelLoader
14 -from vllm.model_executor.model_loader.utils import (
15 - initialize_model,
16 - process_weights_after_loading,
17 -)
18 -from vllm.utils.torch_utils import set_default_torch_dtype
19 -
20 -logger = init_logger(__name__)
21 -
22 -
23 -class TPUModelLoader(DefaultModelLoader):
24 - """
25 - A TPU model loader for model loading under SPMD mode.
26 - """
27 -
28 - def load_model(
29 - self,
30 - vllm_config: VllmConfig,
31 - model_config: ModelConfig,
32 - mesh: xs.Mesh | None = None,
33 - ) -> nn.Module:
34 - # Initialize model and load weights on CPU. Then, during SPMD partition,
35 - # weights are sharded and transferred to TPUs.
36 - self.counter_before_loading_weights = time.perf_counter()
37 - model_config = vllm_config.model_config
38 - assert model_config.quantization is None, "Quantization not supported"
39 - target_device = torch.device("cpu")
40 - with set_default_torch_dtype(model_config.dtype):
41 - with target_device:
42 - model = initialize_model(vllm_config=vllm_config)
43 -
44 - load_format = vllm_config.load_config.load_format
45 - if load_format != "dummy":
46 - weights_to_load = {name for name, _ in model.named_parameters()}
47 - all_weights = self.get_all_weights(model_config, model)
48 - loaded_weights = model.load_weights(all_weights)
49 - self.counter_after_loading_weights = time.perf_counter()
50 - logger.info(
51 - "Loading weights took %.2f seconds",
52 - self.counter_after_loading_weights
53 - - self.counter_before_loading_weights,
54 - )
55 - # We only enable strict check for non-quantized models
56 - # that have loaded weights tracking currently.
57 - if model_config.quantization is None and loaded_weights is not None:
58 - weights_not_loaded = weights_to_load - loaded_weights
59 - if weights_not_loaded:
60 - raise ValueError(
61 - "Following weights were not initialized from "
62 - f"checkpoint: {weights_not_loaded}"
63 - )
64 - else:
65 - logger.info("Use dummy weight during weight loading.")
66 -
67 - process_weights_after_loading(model, model_config, target_device)
68 -
69 - counter_before_partition = time.perf_counter()
70 - model = model.eval()
71 - model = model.to("xla")
72 - shard_model(model, mesh)
73 - counter_after_partition = time.perf_counter()
74 - logger.info(
75 - "Partition model took %.2f seconds",
76 - counter_after_partition - counter_before_partition,
77 - )
78 -
79 - # Ensure the model is properly loaded.
80 - self._check_model_is_loaded(mesh, model)
81 -
82 - # Need to torch compile after model sharding are done. Because the
83 - # compiler hints ('xs.mark_sharding') are torch ops.
84 - if not model_config.is_multimodal_model:
85 - model.model = torch.compile(model.model, backend="openxla")
86 - else:
87 - model.language_model.model = torch.compile(
88 - model.language_model.model, backend="openxla"
89 - )
90 - return model
91 -
92 - def _check_model_is_loaded(self, mesh: xs.Mesh | None, model: nn.Module) -> None:
93 - """
94 - Ensure the model is properly loaded.
95 - 1. All model parameters and buffers are on XLA device.
96 - 2. Non-SPMD friendly layers are replaced as expected.
97 - """
98 - device = xm.xla_device()
99 - device_type = str(device.type)
100 -
101 - # Check parameters
102 - for name, param in model.named_parameters():
103 - assert param.device.type == device_type, (
104 - f"Parameter {name} is on {param.device.type} instead of {device_type}"
105 - )
106 -
107 - # Check buffers
108 - for name, buffer in model.named_buffers():
109 - assert buffer.device.type == device_type, (
110 - f"Buffer {name} is on {buffer.device.type} instead of {device_type}"
111 - )
112 -
113 - for module in model.modules():
114 - if (mesh is not None) and (get_fqn(module) == "QKVParallelLinear"):
115 - raise AssertionError(
116 - "QKVParallelLinear should be replaced by \
117 - XlaQKVParallelLinear under SPMD mode."
118 - )
vllm/platforms/tpu.py
+3 -278
@@ -1,287 +1,10 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -import contextlib
5 -from typing import TYPE_CHECKING, Optional, cast
6 -
7 -import torch
8 -from tpu_info import device
9 -
10 -from vllm.attention.backends.registry import AttentionBackendEnum
11 -from vllm.inputs import ProcessorInputs, PromptType
4 from vllm.logger import init_logger
5
14 -from .interface import Platform, PlatformEnum
15 -
16 -if TYPE_CHECKING:
17 - from typing import TypeAlias
18 -
19 - from vllm.attention.selector import AttentionSelectorConfig
20 - from vllm.config import VllmConfig
21 - from vllm.config.cache import BlockSize
22 - from vllm.pooling_params import PoolingParams
23 - from vllm.sampling_params import SamplingParams
24 -
25 - ParamsType: TypeAlias = SamplingParams | PoolingParams
26 -else:
27 - BlockSize = None
28 - VllmConfig = None
29 - PoolingParams = None
30 - ParamsType = None
31 -
6 logger = init_logger(__name__)
7
34 -USE_TPU_INFERENCE = False
35 -
36 -
37 -class TpuPlatform(Platform):
38 - _enum = PlatformEnum.TPU
39 - device_name: str = "tpu"
40 - device_type: str = "tpu"
41 - dispatch_key: str = "XLA"
42 - ray_device_key: str = "TPU"
43 - dist_backend: str = "gloo"
44 - device_control_env_var: str = "TPU_VISIBLE_CHIPS"
45 - simple_compile_backend: str = "openxla"
46 -
47 - supported_quantization: list[str] = ["fp8", "tpu_int8", "compressed-tensors"]
48 -
49 - additional_env_vars: list[str] = ["TPU_CHIPS_PER_HOST_BOUNDS", "TPU_HOST_BOUNDS"]
50 -
51 - @classmethod
52 - def import_kernels(cls) -> None:
53 - # Do not import vllm._C
54 - with contextlib.suppress(ImportError):
55 - import vllm._moe_C # noqa: F401
56 -
57 - @classmethod
58 - def get_attn_backend_cls(
59 - cls,
60 - selected_backend: "AttentionBackendEnum",
61 - attn_selector_config: "AttentionSelectorConfig",
62 - ) -> str:
63 - if attn_selector_config.use_sparse:
64 - raise NotImplementedError("Sparse Attention is not supported on TPU.")
65 - if selected_backend != AttentionBackendEnum.PALLAS:
66 - logger.info("Cannot use %s backend on TPU.", selected_backend)
67 -
68 - logger.info("Using Pallas V1 backend.")
69 - return AttentionBackendEnum.PALLAS.get_path()
70 -
71 - @classmethod
72 - def get_supported_vit_attn_backends(cls) -> list["AttentionBackendEnum"]:
73 - return [
74 - AttentionBackendEnum.PALLAS,
75 - ]
76 -
77 - @classmethod
78 - def get_vit_attn_backend(
79 - cls,
80 - head_size: int,
81 - dtype: torch.dtype,
82 - backend: Optional["AttentionBackendEnum"] = None,
83 - ) -> "AttentionBackendEnum":
84 - if backend is not None:
85 - assert backend in cls.get_supported_vit_attn_backends(), (
86 - f"Backend {backend} is not supported for vit attention"
87 - f"Supported backends are: {cls.get_supported_vit_attn_backends()}."
88 - )
89 - logger.info_once(f"Using backend {backend} for vit attention.")
90 - return backend
91 -
92 - logger.info_once(
93 - f"Using default backend {AttentionBackendEnum.PALLAS} for vit attention."
94 - )
95 - return AttentionBackendEnum.PALLAS
96 -
97 - @classmethod
98 - def set_device(cls, device: torch.device) -> None:
99 - """
100 - Set the device for the current platform.
101 - """
102 - torch.tpu.set_device(device)
103 -
104 - @classmethod
105 - def get_device_name(cls, device_id: int = 0) -> str:
106 - chip_type, _ = device.get_local_chips()
107 - return f"TPU {chip_type.name}"
108 -
109 - @classmethod
110 - def get_device_total_memory(cls, device_id: int = 0) -> int:
111 - raise NotImplementedError
112 -
113 - @classmethod
114 - def get_punica_wrapper(cls) -> str:
115 - return "vllm.lora.punica_wrapper.punica_tpu.PunicaWrapperTPU"
116 -
117 - @classmethod
118 - def get_infinity_values(cls, dtype: torch.dtype) -> tuple[float, float]:
119 - return torch.finfo(dtype).min, torch.finfo(dtype).max
120 -
121 - @classmethod
122 - def can_update_inplace(cls):
123 - return False
124 -
125 - @classmethod
126 - def get_lora_vocab_padding_size(cls) -> int:
127 - return 1
128 -
129 - @classmethod
130 - def inference_mode(cls):
131 - return torch.no_grad()
132 -
133 - @classmethod
134 - def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
135 - from vllm.config import CompilationMode, CUDAGraphMode
136 -
137 - cache_config = vllm_config.cache_config
138 - # For v0, the default block size is 16.
139 - if cache_config and cache_config.block_size is None:
140 - cache_config.block_size = cast(BlockSize, 16)
141 - compilation_config = vllm_config.compilation_config
142 -
143 - # TPU only supports DYNAMO_TRACE_ONCE compilation mode
144 - if compilation_config.mode != CompilationMode.DYNAMO_TRACE_ONCE:
145 - logger.info(
146 - "[TPU] Forcing DYNAMO_TRACE_ONCE compilation mode, and\
147 - disabling cudagraph."
148 - )
149 - compilation_config.mode = CompilationMode.DYNAMO_TRACE_ONCE
150 -
151 - if (
152 - compilation_config.cudagraph_mode is None
153 - or compilation_config.cudagraph_mode.max_cudagraph_mode()
154 - != CUDAGraphMode.NONE
155 - ):
156 - logger.info(
157 - "[TPU] CUDA graph is not supported on TPU, disabling cudagraphs."
158 - )
159 - compilation_config.cudagraph_mode = CUDAGraphMode.NONE
160 -
161 - if compilation_config.backend == "":
162 - compilation_config.backend = "openxla"
163 -
164 - assert vllm_config.speculative_config is None, (
165 - "TPU does not support speculative decoding"
166 - )
167 -
168 - model_config = vllm_config.model_config
169 - if model_config is not None and model_config.dtype in (
170 - torch.float16,
171 - torch.float32,
172 - ):
173 - logger.warning(
174 - "The TPU backend currently does not support %s. "
175 - "Using bfloat16 instead.",
176 - model_config.dtype,
177 - )
178 - model_config.dtype = torch.bfloat16
179 -
180 - from vllm.v1.attention.backends.pallas import PallasAttentionBackend
181 -
182 - cache_config.block_size = PallasAttentionBackend.get_page_size(vllm_config) # type: ignore[assignment]
183 -
184 - parallel_config = vllm_config.parallel_config
185 - scheduler_config = vllm_config.scheduler_config
186 - if parallel_config.worker_cls == "auto":
187 - parallel_config.worker_cls = "vllm.v1.worker.tpu_worker.TPUWorker"
188 -
189 - assert not vllm_config.speculative_config, (
190 - "Speculative decoding is not yet supported for TPU backend"
191 - )
192 -
193 - if (
194 - scheduler_config.is_multimodal_model
195 - and not scheduler_config.disable_chunked_mm_input
196 - ):
197 - logger.warning(
198 - "TPU does not support running Multimodal models"
199 - " without setting `--disable_chunked_mm_input`. "
200 - "Forcing --disable_chunked_mm_input."
201 - )
202 - scheduler_config.disable_chunked_mm_input = True
203 -
204 - if model_config and model_config.use_mla:
205 - logger.info(
206 - "MLA is enabled on a non-GPU platform; forcing chunked "
207 - "prefill and prefix caching to be disabled."
208 - )
209 - vllm_config.scheduler_config.enable_chunked_prefill = False
210 - vllm_config.scheduler_config.max_num_batched_tokens = max(
211 - vllm_config.model_config.max_model_len,
212 - vllm_config.scheduler_config.DEFAULT_MAX_NUM_BATCHED_TOKENS,
213 - )
214 -
215 - @classmethod
216 - def is_pin_memory_available(cls):
217 - logger.warning("Pin memory is not supported on TPU.")
218 - return False
219 -
220 - @classmethod
221 - def get_device_communicator_cls(cls) -> str:
222 - return "vllm.distributed.device_communicators.tpu_communicator.TpuCommunicator" # noqa
223 -
224 - @classmethod
225 - def validate_request(
226 - cls,
227 - prompt: PromptType,
228 - params: ParamsType,
229 - processed_inputs: ProcessorInputs,
230 - ) -> None:
231 - """Raises if this request is unsupported on this platform"""
232 - from vllm.sampling_params import SamplingParams, SamplingType
233 -
234 - if (
235 - isinstance(params, SamplingParams)
236 - and params.sampling_type == SamplingType.RANDOM_SEED
237 - ):
238 - raise ValueError("Torch XLA does not support per-request seed.")
239 -
240 - @classmethod
241 - @torch.compile(backend="openxla")
242 - def insert_blocks_to_device(
243 - cls,
244 - src_cache: torch.Tensor,
245 - dst_cache: torch.Tensor,
246 - src_block_indices: torch.Tensor,
247 - dst_block_indices: torch.Tensor,
248 - ) -> None:
249 - torch.ops.xla.dynamo_set_buffer_donor_(dst_cache, True)
250 - dst_cache[dst_block_indices] = src_cache[src_block_indices].to(dst_cache.device)
251 -
252 - @classmethod
253 - @torch.compile(backend="openxla")
254 - def swap_out_blocks_to_host(
255 - cls,
256 - src_cache: torch.Tensor,
257 - dst_cache: torch.Tensor,
258 - src_block_indices: torch.Tensor,
259 - dst_block_indices: torch.Tensor,
260 - ) -> None:
261 - """tpu blocks to cpu blocks"""
262 - torch.ops.xla.dynamo_set_buffer_donor_(src_cache, True)
263 - dst_cache[dst_block_indices] = src_cache[src_block_indices].cpu()
264 -
265 - @classmethod
266 - def use_sync_weight_loader(cls) -> bool:
267 - return True
268 -
269 - @classmethod
270 - def check_max_model_len(cls, max_model_len: int) -> int:
271 - """
272 - Check max_model_len for the current platform.
273 - """
274 - logger.warning(
275 - "--max-model-len is not specified, "
276 - "it's currently using model's default length %d, "
277 - "which might be too large."
278 - "Please input with --max-model-len based on your "
279 - "request input length and output length, to avoid "
280 - "unnecessary degradation.",
281 - max_model_len,
282 - )
283 - return max_model_len
284 -
8
9 try:
10 from tpu_inference.platforms import (
@@ -291,5 +14,7 @@ try:
14 TpuPlatform = TpuInferencePlatform # type: ignore
15 USE_TPU_INFERENCE = True
16 except ImportError:
294 - logger.info("tpu_inference not found, using vLLM's TpuPlatform")
17 + logger.error(
18 + "tpu_inference not found, please install tpu_inference to run vllm on TPU"
19 + )
20 pass
vllm/usage/usage_lib.py
+1 -17
@@ -186,20 +186,6 @@ class UsageMessage:
186 except Exception:
187 return False
188
189 - def _report_torch_xla_usage(self) -> bool:
190 - try:
191 - import torch_xla
192 -
193 - self.gpu_count = torch_xla.runtime.world_size()
194 - self.gpu_type = torch_xla.tpu.get_tpu_type()
195 - self.gpu_memory_per_device = torch_xla.core.xla_model.get_memory_info()[
196 - "bytes_limit"
197 - ]
198 - self.cuda_runtime = "torch_xla"
199 - return True
200 - except Exception:
201 - return False
202 -
189 def _report_usage_once(
190 self,
191 model_architecture: str,
@@ -217,9 +203,7 @@ class UsageMessage:
203 if current_platform.is_cuda():
204 self.cuda_runtime = torch.version.cuda
205 if current_platform.is_tpu(): # noqa: SIM102
220 - if (not self._report_tpu_inference_usage()) and (
221 - not self._report_torch_xla_usage()
222 - ):
206 + if not self._report_tpu_inference_usage():
207 logger.exception("Failed to collect TPU information")
208 self.provider = _detect_cloud_provider()
209 self.architecture = platform.machine()
vllm/v1/attention/backends/pallas.py deleted
-436
@@ -1,436 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -from dataclasses import dataclass
5 -
6 -import torch
7 -
8 -from vllm.attention.backends.abstract import (
9 - AttentionBackend,
10 - AttentionImpl,
11 - AttentionLayer,
12 - AttentionType,
13 -)
14 -from vllm.config import VllmConfig
15 -from vllm.logger import init_logger
16 -from vllm.utils.math_utils import cdiv, next_power_of_2
17 -
18 -logger = init_logger(__name__)
19 -
20 -# TPU requires the head size to be a multiple of 128.
21 -TPU_HEAD_SIZE_ALIGNMENT = 128
22 -
23 -# Note: TPU can fp8 as storage dtype but doesn't support converting from uint8
24 -# from to fp32 directly. That's why it has a dtype mapping different from GPU
25 -TPU_STR_DTYPE_TO_TORCH_DTYPE = {
26 - "half": torch.half,
27 - "bfloat16": torch.bfloat16,
28 - "float": torch.float,
29 - "fp8": torch.float8_e4m3fn,
30 - "fp8_e4m3": torch.float8_e4m3fn,
31 - "fp8_e5m2": torch.float8_e5m2,
32 - "int8": torch.int8,
33 - "uint8": torch.uint8,
34 -}
35 -
36 -try:
37 - import tpu_inference # noqa: F401
38 -except ImportError:
39 - # Lazy import torch_xla
40 - import torch_xla.core.xla_builder as xb
41 - import torch_xla.experimental.custom_kernel # noqa: F401
42 - from torch.library import impl
43 - from torch_xla._internal.jax_workarounds import requires_jax
44 - from torch_xla.experimental.custom_kernel import XLA_LIB
45 -
46 - @requires_jax
47 - def kv_cache_update_op_impl(
48 - kv: torch.Tensor,
49 - slot_mapping: torch.Tensor,
50 - kv_cache: torch.Tensor,
51 - num_kv_update_slices: torch.Tensor,
52 - page_size: int,
53 - num_slices_per_block: int,
54 - ):
55 - from vllm.attention.ops.pallas_kv_cache_update import kv_cache_update
56 -
57 - new_kv_cache = xb.call_jax(
58 - kv_cache_update,
59 - (kv, slot_mapping, kv_cache, num_kv_update_slices),
60 - {"page_size": page_size, "num_slices_per_block": num_slices_per_block},
61 - )
62 - return new_kv_cache
63 -
64 - XLA_LIB.define(
65 - "kv_cache_update_op(Tensor kv, Tensor slot_mapping,"
66 - "Tensor kv_cache, Tensor num_kv_update_slices, int page_size,"
67 - "int num_slices_per_block)"
68 - "-> Tensor",
69 - )
70 -
71 - @impl(XLA_LIB, "kv_cache_update_op", "XLA")
72 - def kv_cache_update_op_xla(
73 - kv: torch.Tensor,
74 - slot_mapping: torch.Tensor,
75 - kv_cache: torch.Tensor,
76 - num_kv_update_slices: torch.Tensor,
77 - page_size: int,
78 - num_slices_per_block: int,
79 - ) -> torch.Tensor:
80 - new_kv_cache = kv_cache_update_op_impl(
81 - kv,
82 - slot_mapping,
83 - kv_cache,
84 - num_kv_update_slices,
85 - page_size,
86 - num_slices_per_block,
87 - )
88 - return new_kv_cache
89 -
90 - @impl(XLA_LIB, "kv_cache_update_op", "CompositeExplicitAutograd")
91 - def kv_cache_update_op_non_xla(
92 - kv: torch.Tensor,
93 - slot_mapping: torch.Tensor,
94 - kv_cache: torch.Tensor,
95 - num_kv_update_slices: torch.Tensor,
96 - page_size: int,
97 - num_slices_per_block: int,
98 - ) -> torch.Tensor:
99 - return kv_cache
100 -
101 -
102 -class PallasAttentionBackend(AttentionBackend):
103 - @staticmethod
104 - def get_name() -> str:
105 - return "PALLAS"
106 -
107 - @staticmethod
108 - def get_impl_cls() -> type["PallasAttentionBackendImpl"]:
109 - return PallasAttentionBackendImpl
110 -
111 - @staticmethod
112 - def get_kv_cache_shape(
113 - num_blocks: int,
114 - block_size: int,
115 - num_kv_heads: int,
116 - head_size: int,
117 - cache_dtype_str: str = "auto",
118 - ) -> tuple[int, ...]:
119 - padded_head_size = (
120 - cdiv(head_size, TPU_HEAD_SIZE_ALIGNMENT) * TPU_HEAD_SIZE_ALIGNMENT
121 - )
122 - return (num_blocks, block_size, num_kv_heads * 2, padded_head_size)
123 -
124 - @staticmethod
125 - def swap_blocks(
126 - src_kv_cache: torch.Tensor,
127 - dst_kv_cache: torch.Tensor,
128 - src_to_dst: torch.Tensor,
129 - ) -> None:
130 - raise RuntimeError("swap_blocks is not used for the TPU backend.")
131 -
132 - # In recent TPU generations, up to v6e, the SMEM size is 1MB. The
133 - # block_tables within the PallasMetadata constitute almost the entire SMEM
134 - # requirement. Its size is max_num_seqs * num_page_per_seq * 4 (Int). Here
135 - # we simply make sure that the size is smaller than half of SMEM capacity.
136 - @staticmethod
137 - def get_min_page_size(vllm_config: VllmConfig) -> int:
138 - max_num_page_per_req = (
139 - 1024 * 1024 // 2 // vllm_config.scheduler_config.max_num_seqs // 4
140 - )
141 - min_page_size = cdiv(
142 - vllm_config.model_config.max_model_len, max_num_page_per_req
143 - )
144 - min_page_size = 1 << (min_page_size - 1).bit_length()
145 - return min_page_size
146 -
147 - @staticmethod
148 - def get_max_num_seqs(model_len: int, page_size: int) -> int:
149 - num_page_per_req = cdiv(model_len, page_size)
150 - return 1024 * 1024 // 2 // num_page_per_req // 4
151 -
152 - # TPU has limited SREGs (scalar registers), if page_size is too small, we
153 - # can spill SREGs easily which leads to bad performance. The strategy we
154 - # apply here is trying to split max-model-len to 16 pages which make the
155 - # spill less likely. Meanwhile we make sure the page size is in [16, 256].
156 - @staticmethod
157 - def get_page_size(vllm_config: VllmConfig) -> int:
158 - # TODO: This is a temporary fix for vmem OOM.
159 - # For long model length, we use 16 page-size to avoid too much
160 - # VMEM spill. A more robust solution should be implemented to
161 - # handle VREG spills.
162 - if vllm_config.model_config.max_model_len > 8192:
163 - return 16
164 - page_size = next_power_of_2(vllm_config.model_config.max_model_len) // 16
165 - if page_size <= 16:
166 - return 16
167 - if page_size >= 256:
168 - return 256
169 - return page_size
170 -
171 -
172 -@dataclass
173 -class PallasMetadata:
174 - # NOTE(sang): Definition of context_len, query_len, and seq_len.
175 - # |---------- N-1 iteration --------|
176 - # |---------------- N iteration ---------------------|
177 - # |- tokenA -|......................|-- newTokens ---|
178 - # |---------- context_len ----------|
179 - # |-------------------- seq_len ---------------------|
180 - # |-- query_len ---|
181 -
182 - # Used in the PallasAttentionBackendImpl
183 - slot_mapping: torch.Tensor
184 - block_tables: torch.Tensor
185 - context_lens: torch.Tensor
186 - query_start_loc: torch.Tensor
187 - num_seqs: torch.Tensor
188 - num_kv_update_slices: torch.Tensor
189 - num_slices_per_kv_cache_update_block: int
190 -
191 -
192 -class PallasAttentionBackendImpl(AttentionImpl):
193 - def __init__(
194 - self,
195 - num_heads: int,
196 - head_size: int,
197 - scale: float,
198 - num_kv_heads: int,
199 - alibi_slopes: list[float] | None,
200 - sliding_window: int | None,
201 - kv_cache_dtype: str,
202 - logits_soft_cap: float | None = None,
203 - attn_type: str = AttentionType.DECODER,
204 - kv_sharing_target_layer_name: int | None = None,
205 - ) -> None:
206 - self.num_heads = num_heads
207 - self.head_size = head_size
208 - self.scale = float(scale)
209 - self.num_kv_heads = num_kv_heads
210 - self.sliding_window = sliding_window
211 - self.logits_soft_cap = logits_soft_cap
212 - self.kv_sharing_target_layer_name = kv_sharing_target_layer_name
213 -
214 - self.num_queries_per_kv = self.num_heads // self.num_kv_heads
215 - if alibi_slopes is not None:
216 - raise NotImplementedError("Alibi slopes is not supported.")
217 -
218 - if attn_type != AttentionType.DECODER:
219 - raise NotImplementedError(
220 - "Encoder self-attention and "
221 - "encoder/decoder cross-attention "
222 - "are not implemented for "
223 - "PallasAttentionBackendImpl"
224 - )
225 -
226 - self.kv_cache_quantized_dtype = None
227 - if kv_cache_dtype != "auto":
228 - self.kv_cache_quantized_dtype = TPU_STR_DTYPE_TO_TORCH_DTYPE.get(
229 - kv_cache_dtype.lower().strip()
230 - )
231 -
232 - def forward(
233 - self,
234 - layer: AttentionLayer,
235 - query: torch.Tensor,
236 - key: torch.Tensor,
237 - value: torch.Tensor,
238 - kv_cache: torch.Tensor,
239 - attn_metadata: PallasMetadata,
240 - output: torch.Tensor | None = None,
241 - output_scale: torch.Tensor | None = None,
242 - output_block_scale: torch.Tensor | None = None,
243 - ) -> torch.Tensor:
244 - """Forward pass with Pallas attention.
245 -
246 - Args:
247 - query: shape = [num_tokens, num_heads * head_size]
248 - key: shape = [num_tokens, num_kv_heads * head_size]
249 - value: shape = [num_tokens, num_kv_heads * head_size]
250 - kv_cache: shape =
251 - [num_blocks, block_size, num_kv_heads * 2, head_size]
252 - attn_metadata: Metadata for attention.
253 - Returns:
254 - shape = [num_tokens, num_heads * head_size]
255 - """
256 - if output_scale is not None or output_block_scale is not None:
257 - raise NotImplementedError(
258 - "fused output quantization is not yet supported"
259 - " for PallasAttentionBackendImpl"
260 - )
261 -
262 - # For determine_available_memory case.
263 - if kv_cache.numel() == 0:
264 - if output is None:
265 - output = torch.ones_like(query)
266 - return output
267 -
268 - num_tokens, hidden_size = query.shape
269 - query = query.view(num_tokens, self.num_heads, self.head_size)
270 - key = key.view(-1, self.num_kv_heads, self.head_size)
271 - value = value.view(-1, self.num_kv_heads, self.head_size)
272 - if self.head_size % TPU_HEAD_SIZE_ALIGNMENT != 0:
273 - padded_head_size = (
274 - cdiv(self.head_size, TPU_HEAD_SIZE_ALIGNMENT) * TPU_HEAD_SIZE_ALIGNMENT
275 - )
276 - query = torch.nn.functional.pad(
277 - query, (0, padded_head_size - self.head_size), value=0.0
278 - )
279 - key = torch.nn.functional.pad(
280 - key, (0, padded_head_size - self.head_size), value=0.0
281 - )
282 - value = torch.nn.functional.pad(
283 - value, (0, padded_head_size - self.head_size), value=0.0
284 - )
285 -
286 - if self.kv_sharing_target_layer_name is None and kv_cache.numel() > 0:
287 - # Write input keys and values to the KV cache.
288 - # Skip this if sharing KV cache with an earlier attention layer.
289 - slot_mapping = attn_metadata.slot_mapping
290 - write_to_kv_cache(
291 - key,
292 - value,
293 - kv_cache,
294 - slot_mapping,
295 - attn_metadata.num_slices_per_kv_cache_update_block,
296 - attn_metadata.num_kv_update_slices,
297 - self.kv_cache_quantized_dtype,
298 - layer._k_scale_float,
299 - layer._v_scale_float,
300 - )
301 -
302 - if self.kv_cache_quantized_dtype is not None and (
303 - layer._k_scale_float == 0.0 or layer._v_scale_float == 0.0
304 - ):
305 - raise ValueError("k_scale_float and v_scale_float must be non-zero")
306 - output = torch.ops.xla.ragged_paged_attention(
307 - query,
308 - kv_cache,
309 - attn_metadata.context_lens,
310 - attn_metadata.block_tables,
311 - attn_metadata.query_start_loc,
312 - attn_metadata.num_seqs,
313 - # By default, the system utilizes optimized block size and
314 - # vmem_limit_bytes parameters from the kernel repository. However,
315 - # these can be manually adjusted for debugging if necessary.
316 - num_kv_pages_per_block=None,
317 - num_queries_per_block=None,
318 - vmem_limit_bytes=None,
319 - use_kernel=True,
320 - sm_scale=self.scale,
321 - sliding_window=self.sliding_window,
322 - soft_cap=self.logits_soft_cap,
323 - k_scale=layer._k_scale_float,
324 - v_scale=layer._v_scale_float,
325 - )
326 -
327 - if self.head_size % TPU_HEAD_SIZE_ALIGNMENT != 0:
328 - output = output[:, :, : self.head_size]
329 -
330 - return output.reshape(num_tokens, hidden_size)
331 -
332 -
333 -def write_to_kv_cache(
334 - key: torch.Tensor,
335 - value: torch.Tensor,
336 - kv_cache: torch.Tensor,
337 - slot_mapping: torch.Tensor,
338 - num_slices_per_kv_cache_update_block: int,
339 - num_kv_update_slices: torch.Tensor,
340 - kv_cache_quantized_dtype: torch.dtype | None = None,
341 - k_scale: float = 1.0,
342 - v_scale: float = 1.0,
343 -) -> None:
344 - """Write the key and values to the KV cache.
345 -
346 - Args:
347 - key: shape = [num_tokens, num_kv_heads, head_size]
348 - value: shape = [num_tokens, num_kv_heads, head_size]
349 - kv_cache: shape = [num_blocks, block_size, num_kv_heads * 2, head_size]
350 - num_slices_per_kv_cache_update_block: int
351 - """
352 - _, page_size, num_combined_kv_heads, head_size = kv_cache.shape
353 - head_size = cdiv(head_size, TPU_HEAD_SIZE_ALIGNMENT) * TPU_HEAD_SIZE_ALIGNMENT
354 -
355 - if kv_cache_quantized_dtype is not None:
356 - dtype_info = torch.finfo(kv_cache_quantized_dtype)
357 - key = key.to(torch.float32) / k_scale
358 - # NOTE: clamp is added here to avoid out of range of quantized dtype
359 - key = torch.clamp(key, dtype_info.min, dtype_info.max)
360 - key = key.to(kv_cache_quantized_dtype)
361 - value = value.to(torch.float32) / v_scale
362 - value = torch.clamp(value, dtype_info.min, dtype_info.max)
363 - value = value.to(kv_cache_quantized_dtype)
364 -
365 - kv = torch.cat([key, value], axis=-1).reshape(-1, num_combined_kv_heads, head_size)
366 -
367 - torch.ops.xla.dynamo_set_buffer_donor_(kv_cache, True)
368 -
369 - kv_cache = kv_cache.flatten(0, 1)
370 - new_kv_cache = torch.ops.xla.kv_cache_update_op(
371 - kv,
372 - slot_mapping,
373 - kv_cache,
374 - num_kv_update_slices,
375 - page_size,
376 - num_slices_per_kv_cache_update_block,
377 - )
378 - # NOTE: the in-place copy will be optimized away by XLA compiler.
379 - kv_cache.copy_(new_kv_cache)
380 -
381 -
382 -# We can move this function to a common utils file if it's also useful for other
383 -# hardware.
384 -def dtype_bits(dtype: torch.dtype):
385 - if dtype.is_floating_point:
386 - try:
387 - return torch.finfo(dtype).bits
388 - except TypeError:
389 - pass
390 - elif dtype.is_complex:
391 - if dtype is torch.complex32:
392 - return 32
393 - elif dtype is torch.complex64:
394 - return 64
395 - elif dtype is torch.complex128:
396 - return 128
397 - else:
398 - try:
399 - return torch.iinfo(dtype).bits
400 - # torch.iinfo cannot support int4, int2, bits8...
401 - except TypeError:
402 - pass
403 - str_dtype = str(dtype)
404 - # support torch.int4, torch.int5, torch.uint5...
405 - if str_dtype.startswith("torch.int") or str_dtype.startswith("torch.uint"):
406 - return int(str_dtype[-1])
407 - raise TypeError(f"Getting the bit width of {dtype} is not supported")
408 -
409 -
410 -def get_dtype_packing(dtype):
411 - bits = dtype_bits(dtype)
412 - if 32 % bits != 0:
413 - raise ValueError(
414 - f"The bit width must be divisible by 32, but got bits={bits}, "
415 - "dtype={dtype}"
416 - )
417 - return 32 // bits
418 -
419 -
420 -def get_page_size_bytes(
421 - block_size: int, num_kv_heads: int, head_size: int, kv_cache_dtype: torch.dtype
422 -) -> int:
423 - """Returns the size in bytes of one page of the KV cache."""
424 - padded_head_size = (
425 - cdiv(head_size, TPU_HEAD_SIZE_ALIGNMENT) * TPU_HEAD_SIZE_ALIGNMENT
426 - )
427 - num_combined_kv_heads = num_kv_heads * 2
428 -
429 - # NOTE: for the implicit padding in XLA
430 - packing = get_dtype_packing(kv_cache_dtype)
431 - num_combined_kv_heads = cdiv(num_combined_kv_heads, packing) * packing
432 -
433 - kv_cache_dtype_bits = dtype_bits(kv_cache_dtype)
434 - return (
435 - block_size * num_combined_kv_heads * padded_head_size * kv_cache_dtype_bits // 8
436 - )
vllm/v1/worker/tpu_model_runner.py deleted
-2196
@@ -1,2196 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -import bisect
4 -import gc
5 -import time
6 -from typing import TYPE_CHECKING, Any, cast
7 -from unittest.mock import patch
8 -
9 -import numpy as np
10 -import torch
11 -import torch.nn as nn
12 -
13 -# TPU XLA related
14 -import torch_xla
15 -import torch_xla.core.xla_model as xm
16 -import torch_xla.distributed.spmd as xs
17 -import torch_xla.runtime as xr
18 -
19 -import vllm.envs as envs
20 -from vllm.attention.backends.abstract import AttentionType
21 -from vllm.attention.layer import Attention, MLAAttention
22 -from vllm.attention.layers.chunked_local_attention import ChunkedLocalAttention
23 -from vllm.compilation.wrapper import TorchCompileWithNoGuardsWrapper
24 -from vllm.config import (
25 - ParallelConfig,
26 - VllmConfig,
27 - get_layers_from_vllm_config,
28 - update_config,
29 -)
30 -from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_group
31 -from vllm.distributed.kv_transfer.kv_connector.utils import copy_kv_blocks
32 -from vllm.forward_context import set_forward_context
33 -from vllm.logger import init_logger
34 -from vllm.lora.layers import BaseLayerWithLoRA, LoRAMappingType
35 -from vllm.lora.request import LoRARequest
36 -from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
37 -from vllm.model_executor.model_loader import get_model_loader
38 -from vllm.model_executor.model_loader.tpu import TPUModelLoader
39 -from vllm.model_executor.models.interfaces import (
40 - SupportsMultiModal,
41 - supports_transcription,
42 -)
43 -from vllm.model_executor.models.interfaces_base import (
44 - is_pooling_model,
45 - is_text_generation_model,
46 -)
47 -from vllm.multimodal import MULTIMODAL_REGISTRY
48 -from vllm.multimodal.inputs import (
49 - BatchedTensorInputs,
50 - MultiModalKwargsItem,
51 - PlaceholderRange,
52 -)
53 -from vllm.multimodal.utils import group_mm_kwargs_by_modality
54 -from vllm.sequence import IntermediateTensors
55 -from vllm.tasks import GenerationTask, PoolingTask, SupportedTask
56 -from vllm.utils.math_utils import cdiv, prev_power_of_2
57 -from vllm.utils.platform_utils import is_pin_memory_available
58 -from vllm.v1.attention.backends.pallas import (
59 - TPU_STR_DTYPE_TO_TORCH_DTYPE,
60 - PallasAttentionBackend,
61 - PallasMetadata,
62 - get_page_size_bytes,
63 -)
64 -from vllm.v1.kv_cache_interface import (
65 - AttentionSpec,
66 - FullAttentionSpec,
67 - KVCacheConfig,
68 - KVCacheSpec,
69 - MLAAttentionSpec,
70 - SlidingWindowSpec,
71 -)
72 -from vllm.v1.outputs import (
73 - EMPTY_MODEL_RUNNER_OUTPUT,
74 - LogprobsLists,
75 - LogprobsTensors,
76 - ModelRunnerOutput,
77 -)
78 -from vllm.v1.sample.tpu.metadata import TPUSupportedSamplingMetadata
79 -from vllm.v1.sample.tpu.sampler import Sampler as TPUSampler
80 -from vllm.v1.worker.kv_connector_model_runner_mixin import (
81 - KVConnectorModelRunnerMixin,
82 - KVConnectorOutput,
83 -)
84 -from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin
85 -from vllm.v1.worker.tpu_input_batch import CachedRequestState, InputBatch
86 -
87 -from .utils import (
88 - MultiModalBudget,
89 - add_kv_sharing_layers_to_kv_cache_groups,
90 - bind_kv_cache,
91 - sanity_check_mm_encoder_outputs,
92 -)
93 -
94 -if TYPE_CHECKING:
95 - from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
96 -
97 -logger = init_logger(__name__)
98 -
99 -INVALID_TOKEN_ID = -1
100 -# Smallest output size
101 -MIN_NUM_SEQS = 8
102 -
103 -
104 -#########################################################
105 -# Ways to avoid recompilation
106 -#########################################################
107 -#
108 -# The model executor has two primary components:
109 -# 1. preparing the model and sampler inputs
110 -# 2. executing the model and sampler.
111 -# The core idea is to avoid any TPU computation during input preparation. For
112 -# better compilation tracking and increased flexibility, the model execution and
113 -# sampler are divided into several distinct components.
114 -#
115 -# Below are the detailed steps:
116 -#
117 -# Step 1
118 -# It is recommended to avoid TPU operations when preparing the model and sampler
119 -# inputs. CPU tensors can be prepared and transferred to the XLA device using
120 -# cpu_tensor.to(xla_device), which only triggers CPU to TPU transfers and avoids
121 -# compilation.
122 -#
123 -# Step 2
124 -# The TPU execution should be decomposed into subgraphs (4 at the moment):
125 -# 1. the main model
126 -# 2. selecting hidden states for each request
127 -# 3. sampler
128 -# 4. encoder.
129 -# Each subgraph should be decorated in a torch.compile. This is used to make
130 -# sure that we have the same subgraph topology in both dummy_run and
131 -# xecute_model. The results from these subgraphs should either be passed to
132 -# other subgraphs, or transferred from TPU to CPU using xla_tensor.cpu() for
133 -# subsequent processing on the CPU.
134 -#
135 -# Step 3
136 -# The dummy_run should be comprehensive, ensuring all potential input shapes and
137 -# branch predictions are included as subgraph inputs to facilitate
138 -# pre-compilation.
139 -class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
140 - def __init__(
141 - self,
142 - vllm_config: VllmConfig,
143 - device: torch.device,
144 - original_parallel_config: ParallelConfig | None = None,
145 - ):
146 - self.vllm_config = vllm_config
147 - self.model_config = vllm_config.model_config
148 - self.cache_config = vllm_config.cache_config
149 - self.lora_config = vllm_config.lora_config
150 - self.load_config = vllm_config.load_config
151 - self.parallel_config = vllm_config.parallel_config
152 - self.original_parallel_config = original_parallel_config
153 - self.scheduler_config = vllm_config.scheduler_config
154 - self.speculative_config = vllm_config.speculative_config
155 - self.observability_config = vllm_config.observability_config
156 - self.device_config = vllm_config.device_config
157 -
158 - model_config = self.model_config
159 - cache_config = self.cache_config
160 - scheduler_config = self.scheduler_config
161 - parallel_config = self.parallel_config
162 - self.device = device
163 - self.check_recompilation = envs.VLLM_XLA_CHECK_RECOMPILATION
164 -
165 - # SPMD Related
166 - self.use_spmd = envs.VLLM_XLA_USE_SPMD
167 - if self.use_spmd:
168 - num_devices = xr.global_runtime_device_count()
169 - mesh_shape = (num_devices, 1)
170 - device_ids = np.array(range(num_devices))
171 - self.mesh = xs.Mesh(device_ids, mesh_shape, ("x", "y"))
172 -
173 - self.enforce_eager = model_config.enforce_eager
174 -
175 - self.num_xla_graphs = 0
176 - self._update_num_xla_graphs("init")
177 -
178 - self.pin_memory = is_pin_memory_available()
179 - self.dtype = self.model_config.dtype
180 - if cache_config.cache_dtype == "auto":
181 - model_dtype = self.dtype
182 - if isinstance(model_dtype, str):
183 - self.kv_cache_dtype = TPU_STR_DTYPE_TO_TORCH_DTYPE[model_dtype]
184 - else:
185 - self.kv_cache_dtype = model_dtype
186 - else:
187 - self.kv_cache_dtype = TPU_STR_DTYPE_TO_TORCH_DTYPE[cache_config.cache_dtype]
188 - self._hidden_states_dtype = self.dtype
189 -
190 - self.sliding_window = model_config.get_sliding_window()
191 - self.block_size = cache_config.block_size
192 - self.max_model_len = model_config.max_model_len
193 - self.most_model_len = envs.VLLM_TPU_MOST_MODEL_LEN
194 - self.max_num_blocks_per_req = cdiv(self.max_model_len, self.block_size)
195 - self.num_blocks_per_most_len_req = (
196 - cdiv(self.most_model_len, self.block_size)
197 - if self.most_model_len is not None
198 - else None
199 - )
200 - # InputBatch needs to work with sampling tensors greater than padding
201 - # to avoid dynamic shapes. Also, avoid suboptimal alignment.
202 - self.max_num_reqs = max(scheduler_config.max_num_seqs, MIN_NUM_SEQS)
203 - self.num_tokens_paddings = _get_token_paddings(
204 - min_token_size=16,
205 - max_token_size=scheduler_config.max_num_batched_tokens,
206 - padding_gap=envs.VLLM_TPU_BUCKET_PADDING_GAP,
207 - )
208 - # In case `max_num_tokens < max(num_tokens_paddings)` use the actual
209 - # padded max value to pre-allocate data structures and pre-compile.
210 - self.max_num_tokens = self.num_tokens_paddings[-1]
211 -
212 - # Model-related.
213 - self.num_attn_layers = model_config.get_num_layers_by_block_type(
214 - parallel_config, "attention"
215 - )
216 - self.num_query_heads = model_config.get_num_attention_heads(parallel_config)
217 - self.num_kv_heads = model_config.get_num_kv_heads(parallel_config)
218 - self.head_size = model_config.get_head_size()
219 - self.inputs_embeds_size = model_config.get_inputs_embeds_size()
220 - self.vocab_size = model_config.get_vocab_size()
221 -
222 - # Multi-modal data support
223 - self.mm_registry = MULTIMODAL_REGISTRY
224 - self.uses_mrope = model_config.uses_mrope
225 - self.supports_mm_inputs = self.mm_registry.supports_multimodal_inputs(
226 - model_config
227 - )
228 - # TODO: Support M-RoPE (e.g, Qwen2-VL)
229 - assert not self.uses_mrope, "TPU does not support M-RoPE yet."
230 -
231 - self._num_slices_per_kv_cache_update_block = (
232 - _get_num_slices_per_kv_cache_update_block(
233 - get_page_size_bytes(
234 - block_size=self.block_size,
235 - num_kv_heads=self.num_kv_heads,
236 - head_size=self.head_size,
237 - kv_cache_dtype=self.kv_cache_dtype,
238 - )
239 - )
240 - )
241 -
242 - # Lazy initialization
243 - self.model: nn.Module # Set after load_model
244 - self.kv_caches: list[torch.Tensor] = []
245 - # mm_hash -> encoder_output
246 - self.encoder_cache: dict[str, torch.Tensor] = {}
247 -
248 - # Request states.
249 - self.requests: dict[str, CachedRequestState] = {}
250 - # NOTE(rob): num_prompt_logprobs only includes reqs
251 - # that are currently in the prefill phase.
252 - self.num_prompt_logprobs: dict[str, int] = {}
253 -
254 - # Initialize input batch early to avoid AttributeError in _update_states
255 - self.input_batch = InputBatch(
256 - max_num_reqs=self.max_num_reqs,
257 - max_model_len=self.max_model_len,
258 - max_num_batched_tokens=self.max_num_tokens,
259 - device=self.device,
260 - pin_memory=self.pin_memory,
261 - vocab_size=self.model_config.get_vocab_size(),
262 - block_sizes=[self.block_size],
263 - kernel_block_sizes=[self.cache_config.block_size],
264 - )
265 -
266 - # Cached torch/numpy tensor
267 - # The pytorch tensor and numpy array share the same buffer.
268 - # Sometimes the numpy op is faster so we create both.
269 - self.input_ids_cpu = torch.zeros(
270 - self.max_num_tokens, dtype=torch.int32, device="cpu"
271 - )
272 -
273 - self.positions_cpu = torch.zeros(
274 - self.max_num_tokens, dtype=torch.int32, device="cpu"
275 - )
276 - self.positions_np = self.positions_cpu.numpy()
277 - self.block_table_cpu = torch.zeros(
278 - (self.max_num_reqs, self.max_num_blocks_per_req),
279 - dtype=torch.int32,
280 - device="cpu",
281 - )
282 - # adjust num_reqs to avoid SMEM OOM.
283 - self.num_reqs_most_model_len = (
284 - min(
285 - PallasAttentionBackend.get_max_num_seqs(
286 - self.most_model_len, self.block_size
287 - ),
288 - self.max_num_reqs,
289 - )
290 - if self.most_model_len is not None
291 - else None
292 - )
293 - self.num_reqs_max_model_len = min(
294 - PallasAttentionBackend.get_max_num_seqs(
295 - self.max_model_len, self.block_size
296 - ),
297 - self.max_num_reqs,
298 - )
299 - self.query_start_loc_cpu = torch.zeros(
300 - self.max_num_tokens + 1,
301 - dtype=torch.int32,
302 - device="cpu",
303 - pin_memory=self.pin_memory,
304 - )
305 - self.query_start_loc_np = self.query_start_loc_cpu.numpy()
306 -
307 - self.seq_lens_cpu = torch.zeros(
308 - self.max_num_tokens,
309 - dtype=torch.int32,
310 - device="cpu",
311 - pin_memory=self.pin_memory,
312 - )
313 - self.seq_lens_np = self.seq_lens_cpu.numpy()
314 -
315 - # Only relevant for multimodal models
316 - if self.supports_mm_inputs:
317 - self.is_mm_embed_cpu = torch.zeros(
318 - self.max_num_tokens,
319 - dtype=torch.bool,
320 - device="cpu",
321 - pin_memory=self.pin_memory,
322 - )
323 -
324 - # Range tensor with values [0 .. self.max_num_tokens - 1].
325 - # Used to initialize positions / context_lens / seq_lens
326 - # Keep in int64 to avoid overflow with long context
327 - self.arange_np = np.arange(self.max_num_tokens, dtype=np.int64)
328 - self.num_reqs_paddings = _get_req_paddings(
329 - min_req_size=MIN_NUM_SEQS, max_req_size=self.max_num_reqs
330 - )
331 -
332 - # Layer pairings for cross-layer KV sharing.
333 - # If an Attention layer `layer_name` is in the keys of this dict, it
334 - # means this layer will perform attention using the keys and values
335 - # from the KV cache of `shared_kv_cache_layers[layer_name]`.
336 - self.shared_kv_cache_layers: dict[str, str] = {}
337 -
338 - # tensors for structured decoding
339 - self.grammar_bitmask_cpu = torch.zeros(
340 - (self.max_num_reqs, cdiv(self.vocab_size, 32)),
341 - dtype=torch.int32,
342 - device="cpu",
343 - pin_memory=self.pin_memory,
344 - )
345 - self.require_structured_out_cpu = torch.zeros(
346 - (self.max_num_reqs, 1),
347 - dtype=torch.bool,
348 - device="cpu",
349 - pin_memory=self.pin_memory,
350 - )
351 - self.structured_decode_arange = torch.arange(
352 - 0, 32, device="cpu", pin_memory=self.pin_memory
353 - )
354 -
355 - self.mm_budget = (
356 - MultiModalBudget(
357 - self.model_config,
358 - self.scheduler_config,
359 - self.mm_registry,
360 - )
361 - if self.supports_mm_inputs
362 - else None
363 - )
364 -
365 - if not self.use_spmd:
366 - self.sample_from_logits_func = torch.compile(
367 - self.sample_from_logits,
368 - backend="openxla",
369 - fullgraph=True,
370 - dynamic=False,
371 - )
372 - else:
373 - self.sample_from_logits_func = self.sample_from_logits
374 -
375 - # For passing scheduler_output between successive
376 - # execute_model() and sample_tokens() calls.
377 - self.scheduler_output: SchedulerOutput | None = None
378 - self.mm_embed_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None
379 -
380 - def reset_mm_cache(self) -> None:
381 - if self.mm_budget:
382 - self.mm_budget.reset_cache()
383 -
384 - def _update_num_xla_graphs(self, case_str):
385 - check_comp = self.check_recompilation and not self.enforce_eager
386 - if not check_comp:
387 - return
388 -
389 - total_cached_graphs = xr.get_num_cached_compilation_graph()
390 - new_compiled_graphs = total_cached_graphs - self.num_xla_graphs
391 - if new_compiled_graphs == 0:
392 - return
393 -
394 - logger.info(
395 - "Add new %d compiled XLA graphs due to %s", new_compiled_graphs, case_str
396 - )
397 - self.num_xla_graphs += new_compiled_graphs
398 -
399 - def _verify_num_xla_graphs(self, case_str):
400 - check_comp = self.check_recompilation and not self.enforce_eager
401 - if not check_comp:
402 - return
403 -
404 - curr_cached_graph = xr.get_num_cached_compilation_graph()
405 - assert self.num_xla_graphs == curr_cached_graph, (
406 - "Recompilation after warm up is detected during {}."
407 - " num_xla_graphs = {} curr_cached_graph = {}".format(
408 - case_str, self.num_xla_graphs, curr_cached_graph
409 - )
410 - )
411 -
412 - def _update_states(self, scheduler_output: "SchedulerOutput") -> bool:
413 - """Update the cached states and the persistent batch with the scheduler
414 - output.
415 -
416 - The updated states are used by the `_prepare_inputs` function to create
417 - the input GPU tensors for the model.
418 -
419 - Returns:
420 - True if there is a new/resumed/paused/finished request.
421 - If False, we can skip copying SamplingMetadata to the GPU.
422 - """
423 - # Remove finished requests from the cached states.
424 - for req_id in scheduler_output.finished_req_ids:
425 - self.requests.pop(req_id, None)
426 - self.num_prompt_logprobs.pop(req_id, None)
427 -
428 - # Remove the finished requests from the persistent batch.
429 - # NOTE(woosuk): There could be an edge case where finished_req_ids and
430 - # scheduled_req_ids overlap. This happens when a request is aborted and
431 - # then resubmitted with the same ID. In this case, we treat them as two
432 - # distinct requests - clearing the cached states for the first request
433 - # and handling the second as a new request.
434 - removed_req_indices: list[int] = []
435 - for req_id in scheduler_output.finished_req_ids:
436 - req_index = self.input_batch.remove_request(req_id)
437 - if req_index is not None:
438 - removed_req_indices.append(req_index)
439 -
440 - # Free the cached encoder outputs.
441 - for mm_hash in scheduler_output.free_encoder_mm_hashes:
442 - self.encoder_cache.pop(mm_hash, None)
443 -
444 - # Remove the unscheduled requests from the persistent batch.
445 - # NOTE(woosuk): The unscheduled requests are either preempted requests
446 - # or running requests that are not scheduled in this step. We remove
447 - # them from the persistent batch but keep their cached states since
448 - # they will be scheduled again sometime in the future.
449 - scheduled_req_ids = scheduler_output.num_scheduled_tokens.keys()
450 - cached_req_ids = self.input_batch.req_id_to_index.keys()
451 - unscheduled_req_ids = cached_req_ids - scheduled_req_ids
452 - # NOTE(woosuk): The persistent batch optimization assumes that
453 - # consecutive batches contain mostly the same requests. If batches
454 - # have low request overlap (e.g., alternating between two distinct
455 - # sets of requests), this optimization becomes very inefficient.
456 - for req_id in unscheduled_req_ids:
457 - req_index = self.input_batch.remove_request(req_id)
458 - assert req_index is not None
459 - removed_req_indices.append(req_index)
460 -
461 - req_ids_to_add: list[str] = []
462 - # Add new requests to the cached states.
463 - for new_req_data in scheduler_output.scheduled_new_reqs:
464 - assert new_req_data.sampling_params is not None, (
465 - "Pooling is not supported in TPU yet"
466 - )
467 - req_id = new_req_data.req_id
468 - sampling_params = new_req_data.sampling_params
469 -
470 - self.requests[req_id] = CachedRequestState(
471 - req_id=req_id,
472 - prompt_token_ids=new_req_data.prompt_token_ids,
473 - prompt_embeds=new_req_data.prompt_embeds,
474 - mm_features=new_req_data.mm_features,
475 - sampling_params=sampling_params,
476 - pooling_params=None,
477 - generator=None,
478 - block_ids=new_req_data.block_ids,
479 - num_computed_tokens=new_req_data.num_computed_tokens,
480 - output_token_ids=[],
481 - lora_request=new_req_data.lora_request,
482 - )
483 -
484 - if sampling_params and sampling_params.prompt_logprobs is not None:
485 - self.num_prompt_logprobs[req_id] = (
486 - self.input_batch.vocab_size
487 - if sampling_params.prompt_logprobs == -1
488 - else sampling_params.prompt_logprobs
489 - )
490 -
491 - req_ids_to_add.append(req_id)
492 -
493 - # Update the states of the running/resumed requests.
494 - req_data = scheduler_output.scheduled_cached_reqs
495 - for i, req_id in enumerate(req_data.req_ids):
496 - req_state = self.requests[req_id]
497 - num_computed_tokens = req_data.num_computed_tokens[i]
498 - new_block_ids = req_data.new_block_ids[i]
499 - resumed_from_preemption = req_id in req_data.resumed_req_ids
500 -
501 - # Update the cached states.
502 - req_state.num_computed_tokens = num_computed_tokens
503 - if not resumed_from_preemption:
504 - if new_block_ids is not None:
505 - # Append the new blocks to the existing block IDs.
506 - for block_ids, new_ids in zip(req_state.block_ids, new_block_ids):
507 - block_ids.extend(new_ids)
508 - else:
509 - assert new_block_ids is not None
510 - # The request is resumed from preemption.
511 - # Replace the existing block IDs with the new ones.
512 - req_state.block_ids = new_block_ids
513 -
514 - req_index = self.input_batch.req_id_to_index.get(req_id)
515 - if req_index is None:
516 - # The request is not in the persistent batch.
517 - # The request was either preempted and resumed later, or was not
518 - # scheduled in the previous step and needs to be added again.
519 - req_ids_to_add.append(req_id)
520 - continue
521 -
522 - # Update the persistent batch.
523 - self.input_batch.num_computed_tokens_cpu[req_index] = num_computed_tokens
524 - if new_block_ids is not None:
525 - self.input_batch.block_table.append_row(new_block_ids, req_index)
526 -
527 - # Add the new or resumed requests to the persistent batch.
528 - # The smaller empty indices are filled first.
529 - removed_req_indices = sorted(removed_req_indices, reverse=True)
530 - for req_id in req_ids_to_add:
531 - req_state = self.requests[req_id]
532 - # Fill the empty index or append to the end
533 - req_index = removed_req_indices.pop() if removed_req_indices else None
534 - self.input_batch.add_request(req_state, req_index)
535 -
536 - # Condense the batched states if there are empty indices.
537 - if removed_req_indices:
538 - self.input_batch.condense(removed_req_indices)
539 -
540 - return len(unscheduled_req_ids) > 0 or len(req_ids_to_add) > 0
541 -
542 - def get_model(self) -> nn.Module:
543 - return self.model
544 -
545 - def get_supported_generation_tasks(self) -> list[GenerationTask]:
546 - model = self.get_model()
547 - supported_tasks = list[GenerationTask]()
548 -
549 - if is_text_generation_model(model):
550 - supported_tasks.append("generate")
551 -
552 - if supports_transcription(model):
553 - if model.supports_transcription_only:
554 - return ["transcription"]
555 -
556 - supported_tasks.append("transcription")
557 -
558 - return supported_tasks
559 -
560 - def get_supported_pooling_tasks(self) -> list[PoolingTask]:
561 - model = self.get_model()
562 - if not is_pooling_model(model):
563 - return []
564 -
565 - return list(model.pooler.get_supported_tasks())
566 -
567 - def get_supported_tasks(self) -> tuple[SupportedTask, ...]:
568 - tasks = list[SupportedTask]()
569 -
570 - if self.model_config.runner_type == "generate":
571 - tasks.extend(self.get_supported_generation_tasks())
572 - if self.model_config.runner_type == "pooling":
573 - tasks.extend(self.get_supported_pooling_tasks())
574 -
575 - return tuple(tasks)
576 -
577 - def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:
578 - """
579 - Generates the KVCacheSpec by parsing the kv cache format from each
580 - Attention module in the static forward context.
581 - Returns:
582 - KVCacheSpec: A dictionary mapping layer names to their KV cache
583 - format. Layers that do not need KV cache are not included.
584 - """
585 -
586 - layers = get_layers_from_vllm_config(
587 - self.vllm_config,
588 - AttentionLayerBase, # type: ignore[type-abstract]
589 - )
590 - block_size = self.vllm_config.cache_config.block_size
591 - cache_dtype_str = self.vllm_config.cache_config.cache_dtype
592 -
593 - kv_cache_spec: dict[str, KVCacheSpec] = {}
594 - for layer_name, attn_module in layers.items():
595 - # Classic Attention path
596 - if isinstance(attn_module, Attention):
597 - if (
598 - kv_tgt_layer := attn_module.kv_sharing_target_layer_name
599 - ) is not None:
600 - # The layer doesn't need its own KV cache and will use that of
601 - # the target layer. We skip creating a KVCacheSpec for it, so
602 - # that KV cache management logic will act as this layer does
603 - # not exist, and doesn't allocate KV cache for the layer. This
604 - # enables the memory saving of cross-layer kv sharing, allowing
605 - # a given amount of memory to accommodate longer context lengths
606 - # or enable more requests to be processed simultaneously.
607 - self.shared_kv_cache_layers[layer_name] = kv_tgt_layer
608 - continue
609 -
610 - if attn_module.attn_type == AttentionType.DECODER:
611 - if isinstance(attn_module, ChunkedLocalAttention):
612 - logger.warning_once(
613 - "Using irope in Pallas is not supported yet, it "
614 - "will fall back to global attention for long context."
615 - )
616 - if attn_module.sliding_window is not None:
617 - kv_cache_spec[layer_name] = SlidingWindowSpec(
618 - block_size=block_size,
619 - num_kv_heads=attn_module.num_kv_heads,
620 - head_size=attn_module.head_size,
621 - dtype=self.kv_cache_dtype,
622 - sliding_window=attn_module.sliding_window,
623 - )
624 - else:
625 - kv_cache_spec[layer_name] = FullAttentionSpec(
626 - block_size=block_size,
627 - num_kv_heads=attn_module.num_kv_heads,
628 - head_size=attn_module.head_size,
629 - dtype=self.kv_cache_dtype,
630 - )
631 - elif attn_module.attn_type in (
632 - AttentionType.ENCODER,
633 - AttentionType.ENCODER_ONLY,
634 - ):
635 - # encoder-only attention does not need KV cache.
636 - continue
637 - elif attn_module.attn_type == AttentionType.ENCODER_DECODER:
638 - raise NotImplementedError
639 - else:
640 - raise ValueError(f"Unknown attention type: {attn_module.attn_type}")
641 - # MLAAttention path
642 - elif isinstance(attn_module, MLAAttention):
643 - if layer_name in kv_cache_spec:
644 - continue
645 - kv_cache_spec[layer_name] = MLAAttentionSpec(
646 - block_size=block_size,
647 - num_kv_heads=1,
648 - head_size=attn_module.head_size,
649 - dtype=self.kv_cache_dtype,
650 - cache_dtype_str=cache_dtype_str,
651 - )
652 - else:
653 - continue
654 -
655 - return kv_cache_spec
656 -
657 - def _get_slot_mapping_metadata(
658 - self, num_reqs, num_scheduled_tokens_per_req
659 - ) -> np.ndarray:
660 - """
661 - Computes metadata for mapping slots to blocks in the key-value (KV)
662 - cache for a batch of requests.
663 -
664 - This function determines, for each request in the batch, how the
665 - scheduled tokens are distributed across memory blocks, and generates
666 - metadata needed to map slices of tokens to their corresponding positions
667 - in the KV cache.
668 -
669 - Args:
670 - num_reqs (int): Number of requests in the current batch.
671 - num_scheduled_tokens_per_req (int or np.ndarray): Number of tokens
672 - to be scheduled for each request.
673 -
674 - Returns:
675 - np.ndarray: A 2D array of shape (total_block_len, 3), where each row
676 - contains:
677 - - kv_cache_start_index (int): The starting index in the KV cache
678 - for the corresponding slice.
679 - - new_kv_start_index (int): The starting index in the new KV
680 - cache for the corresponding slice.
681 - - slice_len (int): The length of the slice.
682 - """
683 - slices_start = self.input_batch.num_computed_tokens_cpu[:num_reqs]
684 - slices_end = (
685 - self.input_batch.num_computed_tokens_cpu[:num_reqs]
686 - + num_scheduled_tokens_per_req
687 - )
688 - local_block_start_idx = slices_start // self.block_size
689 - local_block_end_idx = (slices_end - 1) // self.block_size
690 - no_repeat_req_indices = self.arange_np[:num_reqs]
691 - global_block_start_idx = (
692 - no_repeat_req_indices * self.max_num_blocks_per_req + local_block_start_idx
693 - )
694 - block_lens = local_block_end_idx - local_block_start_idx + 1
695 - global_block_start_idx = np.repeat(global_block_start_idx, block_lens)
696 - slice_arange = np.concatenate([self.arange_np[:n] for n in block_lens])
697 - global_block_indices = global_block_start_idx + slice_arange
698 - block_table_cpu = self.input_batch.block_table[0].get_cpu_tensor()
699 - block_numbers = block_table_cpu.flatten()[global_block_indices].numpy()
700 - total_block_len = np.sum(block_lens)
701 - slot_mapping_slices = np.repeat(
702 - np.array([[0, self.block_size]], dtype=np.int32), total_block_len, axis=0
703 - )
704 - cu_block_lens = np.zeros(len(block_lens) + 1, dtype=np.int32)
705 - np.cumsum(block_lens, out=cu_block_lens[1:])
706 - for req_idx in range(num_reqs):
707 - slot_mapping_slices[cu_block_lens[req_idx]][0] = (
708 - slices_start[req_idx] % self.block_size
709 - )
710 - slot_mapping_slices[cu_block_lens[req_idx + 1] - 1][1] = (
711 - slices_end[req_idx] - 1
712 - ) % self.block_size + 1
713 - slice_lens = slot_mapping_slices[:, 1] - slot_mapping_slices[:, 0]
714 - cu_slices_lens = np.zeros(len(slice_lens) + 1, dtype=np.int32)
715 - np.cumsum(slice_lens, out=cu_slices_lens[1:])
716 - kv_cache_start_indices = slot_mapping_slices[:, 0] + (
717 - block_numbers * self.block_size
718 - )
719 - new_kv_start_indices = cu_slices_lens[:-1]
720 - slot_mapping_metadata = np.stack(
721 - [kv_cache_start_indices, new_kv_start_indices, slice_lens], axis=1
722 - )
723 - return slot_mapping_metadata
724 -
725 - def _prepare_inputs(self, scheduler_output: "SchedulerOutput", start_index: int):
726 - assert scheduler_output.total_num_scheduled_tokens > 0
727 - num_reqs = self.input_batch.num_reqs
728 - assert num_reqs > 0
729 - assert start_index < num_reqs
730 -
731 - # Get the number of scheduled tokens for each request.
732 - use_max_model_len = self.most_model_len is None
733 - num_scheduled_tokens_per_req = []
734 - max_num_scheduled_tokens_all_reqs = 0
735 - end_index = start_index
736 -
737 - # Use either most_model_len or max_model_len depending on request size.
738 - for i in range(start_index, num_reqs):
739 - req_id = self.input_batch.req_ids[i]
740 - assert req_id is not None
741 - num_tokens = scheduler_output.num_scheduled_tokens[req_id]
742 - if (
743 - not use_max_model_len
744 - and self.most_model_len is not None
745 - and num_tokens > self.most_model_len
746 - ):
747 - use_max_model_len = True
748 - num_scheduled_tokens_per_req.append(num_tokens)
749 - if use_max_model_len:
750 - if len(num_scheduled_tokens_per_req) > self.num_reqs_max_model_len:
751 - num_scheduled_tokens_per_req = num_scheduled_tokens_per_req[
752 - : self.num_reqs_max_model_len
753 - ]
754 - end_index = start_index + self.num_reqs_max_model_len
755 - else:
756 - end_index = num_reqs
757 - else:
758 - assert self.num_reqs_most_model_len is not None
759 - if len(num_scheduled_tokens_per_req) > self.num_reqs_most_model_len:
760 - num_scheduled_tokens_per_req = num_scheduled_tokens_per_req[
761 - : self.num_reqs_most_model_len
762 - ]
763 - end_index = start_index + self.num_reqs_most_model_len
764 - else:
765 - end_index = num_reqs
766 - max_num_scheduled_tokens_all_reqs = max(num_scheduled_tokens_per_req)
767 - num_scheduled_tokens_per_req = np.array(
768 - num_scheduled_tokens_per_req, dtype=np.int32
769 - )
770 - total_num_scheduled_tokens = sum(num_scheduled_tokens_per_req)
771 - assert max_num_scheduled_tokens_all_reqs > 0
772 -
773 - num_reqs = len(num_scheduled_tokens_per_req)
774 -
775 - # Get request indices.
776 - # E.g., [2, 5, 3] -> [0, 0, 1, 1, 1, 1, 1, 2, 2, 2]
777 - # For each scheduled token, what are the corresponding req index.
778 - req_indices = np.repeat(self.arange_np[:num_reqs], num_scheduled_tokens_per_req)
779 -
780 - # Get batched arange.
781 - # E.g., [2, 5, 3] -> [0, 1, 0, 1, 2, 3, 4, 0, 1, 2]
782 - # For each scheduled token, what is its position in corresponding req.
783 - arange = np.concatenate(
784 - [self.arange_np[:n] for n in num_scheduled_tokens_per_req]
785 - )
786 -
787 - # Get positions.
788 - positions_np = self.positions_np[:total_num_scheduled_tokens]
789 - np.add(
790 - self.input_batch.num_computed_tokens_cpu[req_indices],
791 - arange,
792 - out=positions_np,
793 - )
794 -
795 - # Get token indices.
796 - # E.g., [0, 1, 0, 1, 2, 3, 4, 0, 1, 2]
797 - # -> [0, 1, M, M + 1, M + 2, M + 3, M + 4, 2 * M, 2 * M + 1, 2 * M + 2]
798 - # where M is the max_model_len.
799 - token_indices = (
800 - positions_np + req_indices * self.input_batch.token_ids_cpu.shape[1]
801 - )
802 -
803 - # NOTE(woosuk): We use torch.index_select instead of np.take here
804 - # because torch.index_select is much faster than np.take for large
805 - # tensors.
806 - torch.index_select(
807 - self.input_batch.token_ids_cpu_tensor.flatten(),
808 - 0,
809 - torch.from_numpy(token_indices),
810 - out=self.input_ids_cpu[:total_num_scheduled_tokens],
811 - )
812 -
813 - # Prepare the attention metadata.
814 - self.query_start_loc_np[0] = 0
815 - np.cumsum(
816 - num_scheduled_tokens_per_req, out=self.query_start_loc_np[1 : num_reqs + 1]
817 - )
818 - self.query_start_loc_np[num_reqs + 1 :] = 1
819 -
820 - self.seq_lens_np[:num_reqs] = (
821 - self.input_batch.num_computed_tokens_cpu[:num_reqs]
822 - + num_scheduled_tokens_per_req
823 - )
824 -
825 - # Do the padding and copy the tensors to the TPU.
826 - padded_total_num_scheduled_tokens = _get_padded_token_len(
827 - self.num_tokens_paddings, total_num_scheduled_tokens
828 - )
829 - # Zero out to avoid spurious values from prev iteration (last cp chunk)
830 - self.input_ids_cpu[
831 - total_num_scheduled_tokens:padded_total_num_scheduled_tokens
832 - ] = 0
833 - self.input_ids = self.input_ids_cpu[:padded_total_num_scheduled_tokens].to(
834 - self.device
835 - )
836 - self.position_ids = self.positions_cpu[:padded_total_num_scheduled_tokens].to(
837 - self.device
838 - )
839 - if use_max_model_len:
840 - block_tables = self.block_table_cpu[
841 - : self.num_reqs_max_model_len, : self.max_num_blocks_per_req
842 - ]
843 - block_tables[:num_reqs, : self.max_num_blocks_per_req] = (
844 - self.input_batch.block_table[0].get_cpu_tensor()[:num_reqs]
845 - )
846 - query_start_loc = self.query_start_loc_cpu[
847 - : self.num_reqs_max_model_len + 1
848 - ].to(self.device)
849 - seq_lens = self.seq_lens_cpu[: self.num_reqs_max_model_len].to(self.device)
850 - else:
851 - assert self.num_reqs_most_model_len is not None
852 - block_tables = self.block_table_cpu[
853 - : self.num_reqs_most_model_len, : self.num_blocks_per_most_len_req
854 - ]
855 - block_tables[:num_reqs, : self.num_blocks_per_most_len_req] = (
856 - self.input_batch.block_table[0].get_cpu_tensor()[
857 - :num_reqs, : self.num_blocks_per_most_len_req
858 - ]
859 - )
860 - query_start_loc = self.query_start_loc_cpu[
861 - : self.num_reqs_most_model_len + 1
862 - ].to(self.device)
863 - seq_lens = self.seq_lens_cpu[: self.num_reqs_most_model_len].to(self.device)
864 - block_tables = block_tables.to(self.device)
865 -
866 - # Calculate the slot mapping
867 - slot_mapping_metadata = self._get_slot_mapping_metadata(
868 - num_reqs, num_scheduled_tokens_per_req
869 - )
870 - num_kv_update_slices = slot_mapping_metadata.shape[0]
871 - padded_num_slices = _get_padded_num_kv_cache_update_slices(
872 - padded_total_num_scheduled_tokens, self.max_num_reqs, self.block_size
873 - )
874 - slot_mapping_metadata = np.pad(
875 - slot_mapping_metadata,
876 - [[0, padded_num_slices - len(slot_mapping_metadata)], [0, 0]],
877 - constant_values=0,
878 - )
879 - slot_mapping_metadata = np.transpose(slot_mapping_metadata)
880 - slot_mapping_metadata = torch.tensor(slot_mapping_metadata, device=self.device)
881 -
882 - if self.lora_config is not None:
883 - # We need to respect padding when activating LoRA adapters
884 - padded_num_scheduled_tokens_per_req = np.copy(
885 - num_scheduled_tokens_per_req
886 - ) # Copying to avoid accidental state corruption bugs
887 - padded_num_scheduled_tokens_per_req[-1] += (
888 - padded_total_num_scheduled_tokens - total_num_scheduled_tokens
889 - )
890 -
891 - self.set_active_loras(self.input_batch, padded_num_scheduled_tokens_per_req)
892 -
893 - attn_metadata = PallasMetadata(
894 - slot_mapping=slot_mapping_metadata,
895 - block_tables=block_tables,
896 - context_lens=seq_lens,
897 - query_start_loc=query_start_loc,
898 - num_seqs=torch.tensor([num_reqs], dtype=torch.int32, device=self.device),
899 - num_kv_update_slices=torch.tensor(
900 - [num_kv_update_slices], dtype=torch.int32, device=self.device
901 - ),
902 - num_slices_per_kv_cache_update_block=self._num_slices_per_kv_cache_update_block,
903 - )
904 - # NOTE(woosuk): Due to chunked prefills, there can be at most 1 partial
905 - # request in the batch. While we should not sample any token from this
906 - # partial request, we do so for simplicity. We will ignore the sampled
907 - # token from the partial request.
908 - # TODO: Support prompt logprobs.
909 - padded_num_reqs = _get_padded_num_reqs_with_upper_limit(
910 - num_reqs, self.max_num_reqs
911 - )
912 - # Indices at which we sample (positions of last token in the sequence).
913 - # Padded to avoid recompiling when `num_reqs` varies.
914 - logits_indices = self.query_start_loc_cpu[1 : padded_num_reqs + 1] - 1
915 - logits_indices = logits_indices.to(self.device)
916 -
917 - if self.lora_config is not None:
918 - # We need to respect padding when activating LoRA adapters
919 - padded_num_scheduled_tokens_per_req = np.copy(
920 - num_scheduled_tokens_per_req
921 - ) # Copying to avoid accidental state corruption bugs
922 - padded_num_scheduled_tokens_per_req[-1] += (
923 - padded_total_num_scheduled_tokens - total_num_scheduled_tokens
924 - )
925 -
926 - self.set_active_loras(self.input_batch, padded_num_scheduled_tokens_per_req)
927 -
928 - layer_names = get_layers_from_vllm_config(self.vllm_config, Attention).keys()
929 - per_layer_attn_metadata = {
930 - layer_name: attn_metadata for layer_name in layer_names
931 - }
932 - return (
933 - per_layer_attn_metadata,
934 - logits_indices,
935 - padded_num_reqs,
936 - num_reqs,
937 - end_index,
938 - )
939 -
940 - def _execute_mm_encoder(self, scheduler_output: "SchedulerOutput"):
941 - scheduled_encoder_inputs = scheduler_output.scheduled_encoder_inputs
942 - if not scheduled_encoder_inputs:
943 - return
944 -
945 - # Batch the multi-modal inputs.
946 - mm_kwargs = list[MultiModalKwargsItem]()
947 - # List of tuple (mm_hash, pos_info)
948 - mm_hashes_pos = list[tuple[str, PlaceholderRange]]()
949 - for req_id, encoder_input_ids in scheduled_encoder_inputs.items():
950 - req_state = self.requests[req_id]
951 -
952 - for mm_input_id in encoder_input_ids:
953 - mm_feature = req_state.mm_features[mm_input_id]
954 - if mm_feature.data is None:
955 - continue
956 - mm_hash = mm_feature.identifier
957 - mm_kwargs.append(mm_feature.data)
958 - mm_hashes_pos.append((mm_hash, mm_feature.mm_position))
959 -
960 - # Batch mm inputs as much as we can: if a request in the batch has
961 - # multiple modalities or a different modality than the previous one,
962 - # we process it separately to preserve item order.
963 - # FIXME(ywang96): This is a hacky way to deal with multiple modalities
964 - # in the same batch while still being able to benefit from batching
965 - # multimodal inputs. The proper solution should be reordering the
966 - # encoder outputs.
967 - model = cast(SupportsMultiModal, self.model)
968 - encoder_outputs = []
969 - for _, num_items, mm_kwargs_group in group_mm_kwargs_by_modality(
970 - mm_kwargs,
971 - device=self.device,
972 - pin_memory=self.pin_memory,
973 - ):
974 - # Run the encoder.
975 - # `curr_group_outputs` is either of the following:
976 - # 1. A tensor of shape (num_items, feature_size, hidden_size)
977 - # in case feature_size is fixed across all multimodal items.
978 - # 2. A list or tuple (length: num_items) of tensors, each of shape
979 - # (feature_size, hidden_size) in case the feature size is dynamic
980 - # depending on the input multimodal items.
981 - torch_xla.sync(wait=False)
982 - curr_group_outputs = model.embed_multimodal(**mm_kwargs_group)
983 - torch_xla.sync(wait=False)
984 -
985 - sanity_check_mm_encoder_outputs(
986 - curr_group_outputs,
987 - expected_num_items=num_items,
988 - )
989 -
990 - if isinstance(curr_group_outputs, torch.Tensor):
991 - encoder_outputs.append(curr_group_outputs)
992 - else:
993 - assert isinstance(curr_group_outputs, (list, tuple))
994 - for output in curr_group_outputs:
995 - encoder_outputs.append(output)
996 -
997 - # Cache the encoder outputs.
998 - # NOTE (NickLucche) here we diverge from logic in other runners, as we
999 - # assume to only have whole mm items to process. Hence we avoid the
1000 - # intrinsic dynamism that `scatter_mm_placeholders` introduces.
1001 - for (mm_hash, pos_info), output in zip(mm_hashes_pos, encoder_outputs):
1002 - assert pos_info.is_embed is None, (
1003 - "Expected all positions to be contiguous and embeddings."
1004 - )
1005 - self.encoder_cache[mm_hash] = output
1006 -
1007 - def _gather_mm_embeddings(
1008 - self,
1009 - scheduler_output: "SchedulerOutput",
1010 - ) -> tuple[list[torch.Tensor], torch.Tensor]:
1011 - total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
1012 - padded_total_num_scheduled_tokens = _get_padded_token_len(
1013 - self.num_tokens_paddings, total_num_scheduled_tokens
1014 - )
1015 -
1016 - is_mm_embed = self.is_mm_embed_cpu
1017 - is_mm_embed[:padded_total_num_scheduled_tokens] = False
1018 - mm_embeds = list[torch.Tensor]()
1019 - req_start_idx = 0
1020 -
1021 - for req_id in self.input_batch.req_ids:
1022 - num_scheduled_tokens = scheduler_output.num_scheduled_tokens[req_id]
1023 - req_state = self.requests[req_id]
1024 - num_computed_tokens = req_state.num_computed_tokens
1025 -
1026 - # TODO unroll loop and assume/enforce --disable_chunked_mm_input
1027 - # NOTE (NickLucche) here we diverge from logic in other runners, as
1028 - # we assume to only have whole mm items to process. Hence we avoid
1029 - # the intrinsic dynamism that `gather_mm_placeholders` introduces.
1030 - for mm_feature in req_state.mm_features:
1031 - pos_info = mm_feature.mm_position
1032 - start_pos = pos_info.offset
1033 - num_encoder_tokens = pos_info.length
1034 -
1035 - # The encoder output is needed if the two ranges overlap:
1036 - # [num_computed_tokens,
1037 - # num_computed_tokens + num_scheduled_tokens) and
1038 - # [start_pos, start_pos + num_encoder_tokens)
1039 - if start_pos >= num_computed_tokens + num_scheduled_tokens:
1040 - # The encoder output is not needed in this step.
1041 - break
1042 - if start_pos + num_encoder_tokens <= num_computed_tokens:
1043 - # The encoder output is already processed and stored
1044 - # in the decoder's KV cache.
1045 - continue
1046 -
1047 - start_idx = max(num_computed_tokens - start_pos, 0)
1048 - end_idx = min(
1049 - num_computed_tokens - start_pos + num_scheduled_tokens,
1050 - num_encoder_tokens,
1051 - )
1052 - assert start_idx < end_idx
1053 -
1054 - mm_hash = mm_feature.identifier
1055 - encoder_output = self.encoder_cache.get(mm_hash, None)
1056 - assert encoder_output is not None, f"Encoder cache miss for {mm_hash}."
1057 -
1058 - assert pos_info.is_embed is None, (
1059 - "Expected all positions to be contiguous and embeddings."
1060 - )
1061 -
1062 - req_start_pos = req_start_idx + start_pos - num_computed_tokens
1063 - is_mm_embed[req_start_pos + start_idx : req_start_pos + end_idx] = True
1064 -
1065 - # Only whole mm items are processed
1066 - mm_embeds.append(encoder_output)
1067 -
1068 - req_start_idx += num_scheduled_tokens
1069 -
1070 - is_mm_embed = is_mm_embed[:padded_total_num_scheduled_tokens].to(self.device)
1071 -
1072 - return mm_embeds, is_mm_embed
1073 -
1074 - def _get_model_inputs(
1075 - self,
1076 - input_ids: torch.Tensor,
1077 - mm_embed_inputs: tuple[list[torch.Tensor], torch.Tensor] | None,
1078 - ):
1079 - if self.supports_mm_inputs:
1080 - mm_embeds, is_mm_embed = mm_embed_inputs or (None, None)
1081 -
1082 - # NOTE(woosuk): To unify token ids and soft tokens (vision
1083 - # embeddings), we always use embeddings (rather than token ids)
1084 - # as input to the multimodal model, even when the input is text.
1085 - inputs_embeds = self.model.embed_input_ids(
1086 - input_ids,
1087 - multimodal_embeddings=mm_embeds,
1088 - is_multimodal=is_mm_embed,
1089 - )
1090 -
1091 - return None, inputs_embeds
1092 - else:
1093 - # For text-only models, we use token ids as input.
1094 - # While it is possible to use embeddings as input just like the
1095 - # multimodal models, it is not desirable for performance since
1096 - # then the embedding layer is not included in the CUDA graph.
1097 - return input_ids, None
1098 -
1099 - @torch.no_grad()
1100 - def execute_model(
1101 - self,
1102 - scheduler_output: "SchedulerOutput",
1103 - intermediate_tensors: IntermediateTensors | None = None,
1104 - ) -> ModelRunnerOutput | None:
1105 - if self.scheduler_output is not None:
1106 - raise RuntimeError(
1107 - "State error: sample_tokens() must be called "
1108 - "after execute_model() returns None."
1109 - )
1110 - # Update cached state
1111 - self._update_states(scheduler_output)
1112 - if not scheduler_output.total_num_scheduled_tokens:
1113 - if not has_kv_transfer_group():
1114 - # Return empty ModelRunnerOutput if there's no work to do.
1115 - return EMPTY_MODEL_RUNNER_OUTPUT
1116 -
1117 - return self.kv_connector_no_forward(scheduler_output, self.vllm_config)
1118 -
1119 - mm_embed_inputs = None
1120 - if self.supports_mm_inputs:
1121 - # Run the multimodal encoder if any.
1122 - self._execute_mm_encoder(scheduler_output)
1123 - mm_embed_inputs = self._gather_mm_embeddings(scheduler_output)
1124 -
1125 - torch_xla.sync(wait=False)
1126 -
1127 - self.scheduler_output = scheduler_output
1128 - self.mm_embed_inputs = mm_embed_inputs
1129 - return None
1130 -
1131 - @torch.no_grad()
1132 - def sample_tokens(
1133 - self, grammar_output: "GrammarOutput | None"
1134 - ) -> ModelRunnerOutput:
1135 - if self.scheduler_output is None:
1136 - # Nothing to do (PP non-final rank case), output isn't used.
1137 - return None # type: ignore[return-value]
1138 - scheduler_output = self.scheduler_output
1139 - mm_embed_inputs = self.mm_embed_inputs
1140 - self.scheduler_output = None
1141 - self.mm_embed_inputs = None
1142 -
1143 - # Prepare inputs, the requests might be split into multiple
1144 - # executions, combine the result of each execution.
1145 - start_index = 0
1146 - combined_selected_tokens: list[torch.Tensor] = []
1147 - combined_logprobs: list[LogprobsLists] = []
1148 -
1149 - # NOTE: setup current batch's metadata for kv connector.
1150 - # Currently, only verified with NixlConnector
1151 - with set_forward_context(None, self.vllm_config):
1152 - self.maybe_setup_kv_connector(scheduler_output)
1153 -
1154 - while start_index < self.input_batch.num_reqs:
1155 - attn_metadata, logits_indices, padded_num_reqs, num_reqs, end_index = (
1156 - self._prepare_inputs(scheduler_output, start_index)
1157 - )
1158 - input_ids, inputs_embeds = self._get_model_inputs(
1159 - self.input_ids, mm_embed_inputs
1160 - )
1161 - torch_xla.sync(wait=False)
1162 - # Run the decoder
1163 - with set_forward_context(
1164 - attn_metadata,
1165 - self.vllm_config,
1166 - num_tokens=scheduler_output.total_num_scheduled_tokens,
1167 - ):
1168 - hidden_states = self.model(
1169 - input_ids=input_ids,
1170 - positions=self.position_ids,
1171 - inputs_embeds=inputs_embeds,
1172 - )
1173 - hidden_states = self.select_hidden_states(hidden_states, logits_indices)
1174 - logits = self.compute_logits(hidden_states)
1175 - tpu_sampling_metadata = TPUSupportedSamplingMetadata.from_input_batch(
1176 - self.input_batch, padded_num_reqs, self.device
1177 - )
1178 - if grammar_output is not None:
1179 - require_struct_decoding, grammar_bitmask_padded, arange = (
1180 - self.prepare_structured_decoding_input(logits, grammar_output)
1181 - )
1182 - logits = self.structured_decode(
1183 - require_struct_decoding, grammar_bitmask_padded, logits, arange
1184 - )
1185 - selected_token_ids = self.sample_from_logits_func(
1186 - logits, tpu_sampling_metadata
1187 - )
1188 - # NOTE (NickLucche) Use the original logits (before any penalties or
1189 - # temperature scaling) for the top-k logprobs. We can't enforce it
1190 - # due to recompilations outside torch.compiled code, so just make
1191 - # sure `sample_from_logits` does not modify the logits in-place.
1192 - logprobs = (
1193 - self.gather_logprobs(logits, selected_token_ids)
1194 - if tpu_sampling_metadata.logprobs
1195 - else None
1196 - )
1197 -
1198 - # Remove padding on cpu and keep dynamic op outside of xla graph.
1199 - selected_token_ids = selected_token_ids.cpu()[:num_reqs]
1200 -
1201 - combined_selected_tokens.append(selected_token_ids)
1202 - if tpu_sampling_metadata.logprobs:
1203 - combined_logprobs.append(logprobs.tolists())
1204 -
1205 - start_index = end_index
1206 -
1207 - # NOTE: current kv load and save get h2d/d2h copies involved.
1208 - # Those copies are blocking. Once they become async., kv_save
1209 - # should be called right after each single forward pass,
1210 - # instead of the forwards of the entire input batch.
1211 - self.maybe_wait_for_kv_save()
1212 - finished_sending, finished_recving = self.get_finished_kv_transfers(
1213 - scheduler_output
1214 - )
1215 -
1216 - selected_token_ids = torch.cat(combined_selected_tokens, dim=0)
1217 - if tpu_sampling_metadata.logprobs:
1218 -
1219 - def concat_lists(input_lists):
1220 - result = []
1221 - for input_list in input_lists:
1222 - result.extend(input_list)
1223 - return result
1224 -
1225 - logprobs_lists = LogprobsLists(
1226 - logprob_token_ids=concat_lists(
1227 - [lp.logprob_token_ids for lp in combined_logprobs]
1228 - ),
1229 - logprobs=concat_lists([lp.logprobs for lp in combined_logprobs]),
1230 - sampled_token_ranks=concat_lists(
1231 - [lp.sampled_token_ranks for lp in combined_logprobs]
1232 - ),
1233 - )
1234 - else:
1235 - logprobs_lists = None
1236 -
1237 - # Update the cache state concurrently. Code above will not block until
1238 - # we use `selected_token_ids`. Add mark_step if post-processing changes
1239 - request_seq_lens: list[tuple[int, CachedRequestState, int]] = []
1240 - discard_sampled_tokens_req_indices = []
1241 - num_reqs = self.input_batch.num_reqs
1242 - for i, req_id in zip(range(num_reqs), self.input_batch.req_ids):
1243 - assert req_id is not None
1244 - req_state = self.requests[req_id]
1245 - seq_len = (
1246 - req_state.num_computed_tokens
1247 - + scheduler_output.num_scheduled_tokens[req_id]
1248 - )
1249 - if seq_len >= req_state.num_tokens:
1250 - request_seq_lens.append((i, req_state, seq_len))
1251 - else:
1252 - # Ignore the sampled token from the partial request.
1253 - # Rewind the generator state as if the token was not sampled.
1254 - generator = self.input_batch.generators.get(i)
1255 - if generator is not None:
1256 - # This relies on cuda-specific torch-internal impl details
1257 - generator.set_offset(generator.get_offset() - 4)
1258 -
1259 - # Record the index of the request that should not be sampled,
1260 - # so that we could clear the sampled tokens before returning.
1261 - discard_sampled_tokens_req_indices.append(i)
1262 -
1263 - assert all(
1264 - req_id is not None for req_id in self.input_batch.req_ids[:num_reqs]
1265 - ), "req_ids contains None"
1266 - req_ids = cast(list[str], self.input_batch.req_ids[:num_reqs])
1267 -
1268 - prompt_logprobs_dict: dict[str, LogprobsTensors | None] = {}
1269 - for req_id in self.input_batch.req_ids[:num_reqs]:
1270 - prompt_logprobs_dict[req_id] = None
1271 -
1272 - max_gen_len = selected_token_ids.shape[-1]
1273 - if max_gen_len == 1:
1274 - valid_sampled_token_ids = selected_token_ids.tolist()
1275 -
1276 - # Mask out the sampled tokens that should not be sampled.
1277 - # TODO: Keep in sync with gpu_model_runner.py, in particular
1278 - # the "else" case here
1279 - for i in discard_sampled_tokens_req_indices:
1280 - valid_sampled_token_ids[i].clear()
1281 -
1282 - # Append sampled tokens
1283 - for i, req_state, seq_len in request_seq_lens:
1284 - token_id = valid_sampled_token_ids[i][0]
1285 - self.input_batch.token_ids_cpu[i, seq_len] = token_id
1286 - req_state.output_token_ids.append(token_id)
1287 - self.input_batch.num_tokens_no_spec[i] += 1
1288 -
1289 - else:
1290 - valid_mask = selected_token_ids != INVALID_TOKEN_ID
1291 - gen_lens = valid_mask.sum(dim=1).tolist()
1292 - valid_sampled_token_ids = [
1293 - seq.tolist() for seq in selected_token_ids[valid_mask].split(gen_lens)
1294 - ]
1295 - self.input_batch.num_tokens_no_spec[:num_reqs] += gen_lens
1296 - for i, req_state, seq_len in request_seq_lens:
1297 - target_slice = slice(seq_len - gen_lens[i] + 1, seq_len + 1)
1298 - self.input_batch.token_ids_cpu[i, target_slice] = (
1299 - valid_sampled_token_ids[i]
1300 - )
1301 - req_state.output_token_ids.extend(valid_sampled_token_ids[i])
1302 -
1303 - kv_connector_output = (
1304 - None
1305 - if (finished_sending is None and finished_recving is None)
1306 - else KVConnectorOutput(
1307 - finished_sending=finished_sending,
1308 - finished_recving=finished_recving,
1309 - )
1310 - )
1311 -
1312 - model_runner_output = ModelRunnerOutput(
1313 - req_ids=req_ids,
1314 - req_id_to_index=self.input_batch.req_id_to_index,
1315 - sampled_token_ids=valid_sampled_token_ids,
1316 - logprobs=logprobs_lists,
1317 - prompt_logprobs_dict=prompt_logprobs_dict,
1318 - pooler_output=[],
1319 - kv_connector_output=kv_connector_output,
1320 - )
1321 -
1322 - # Check there are no new graphs compiled - all the graphs should be
1323 - # captured and compiled during warm up.
1324 - self._verify_num_xla_graphs("execute_model")
1325 -
1326 - return model_runner_output
1327 -
1328 - def update_config(self, overrides: dict[str, Any]) -> None:
1329 - # TODO: TPU config may need extra validation
1330 - # https://github.com/vllm-project/vllm/pull/20095#discussion_r2201497754
1331 - allowed_config_names = {"load_config", "model_config"}
1332 - for config_name, config_overrides in overrides.items():
1333 - assert config_name in allowed_config_names, (
1334 - f"Config `{config_name}` not supported. "
1335 - f"Allowed configs: {allowed_config_names}"
1336 - )
1337 - config = getattr(self, config_name)
1338 - new_config = update_config(config, config_overrides)
1339 - setattr(self, config_name, new_config)
1340 -
1341 - def load_model(self) -> None:
1342 - self.device = self.device_config.device
1343 -
1344 - # NOTE(woosuk): While the executor assigns the TP ranks to the worker
1345 - # process, the ranks can be different from the ranks internally assigned
1346 - # by the xm runtime. Therefore, there is a mismatch in the rank
1347 - # assignment between the gloo (cpu) runtime and the xm (tpu) runtime.
1348 - # This is not a problem in linear layers because all-reduce is
1349 - # rank-agnostic. However, it matters for all-gather as the ranks
1350 - # determine the order of concatenating the output tensors.
1351 - # As a workaround, we use the xm's rank assignment only when loading
1352 - # the embedding weights.
1353 - xm_tp_rank = xr.global_ordinal()
1354 - with patch(
1355 - "vllm.model_executor.layers.vocab_parallel_embedding."
1356 - "get_tensor_model_parallel_rank",
1357 - return_value=xm_tp_rank,
1358 - ):
1359 - try:
1360 - if self.use_spmd:
1361 - tpu_loader = TPUModelLoader(
1362 - load_config=self.vllm_config.load_config
1363 - )
1364 - model = tpu_loader.load_model(
1365 - vllm_config=self.vllm_config,
1366 - model_config=self.vllm_config.model_config,
1367 - mesh=self.mesh,
1368 - )
1369 - else:
1370 - model_loader = get_model_loader(self.load_config)
1371 - logger.info("Loading model from scratch...")
1372 - model = model_loader.load_model(
1373 - vllm_config=self.vllm_config, model_config=self.model_config
1374 - )
1375 - except RuntimeError as e:
1376 - raise RuntimeError(
1377 - f"Unable to load model, a likely reason is the model is "
1378 - "too large for the current device's HBM memory. "
1379 - "Consider switching to a smaller model "
1380 - "or sharding the weights on more chips. "
1381 - f"See the detailed error: {e}"
1382 - ) from e
1383 - if self.lora_config is not None:
1384 - model = self.load_lora_model(model, self.vllm_config, self.device)
1385 - replace_set_lora(model)
1386 -
1387 - # Sync all pending XLA execution during model initialization and weight
1388 - # loading.
1389 - torch_xla.sync(wait=False)
1390 - xm.wait_device_ops()
1391 - if not hasattr(self, "model"):
1392 - self.model = model
1393 - self.sampler = TPUSampler()
1394 -
1395 - def reload_weights(self) -> None:
1396 - assert getattr(self, "model", None) is not None, (
1397 - "Cannot reload weights before model is loaded."
1398 - )
1399 - model_loader = get_model_loader(self.load_config)
1400 - logger.info("Reloading weights inplace...")
1401 - model_loader.load_weights(self.model, model_config=self.model_config)
1402 -
1403 - @torch.no_grad()
1404 - def _dummy_run(self, num_tokens: int, num_reqs: int, num_blocks: int) -> None:
1405 - if self.supports_mm_inputs:
1406 - input_ids = None
1407 - inputs_embeds = torch.zeros(
1408 - (num_tokens, self.inputs_embeds_size),
1409 - dtype=self.dtype,
1410 - device=self.device,
1411 - )
1412 - else:
1413 - input_ids = torch.zeros((num_tokens), dtype=torch.int32).to(self.device)
1414 - inputs_embeds = None
1415 - actual_num_reqs = min(num_tokens, num_reqs)
1416 - position_ids = torch.zeros(num_tokens, dtype=torch.int32).to(self.device)
1417 - padded_num_slices = _get_padded_num_kv_cache_update_slices(
1418 - num_tokens, self.max_num_reqs, self.block_size
1419 - )
1420 - num_kv_update_slices = torch.tensor([padded_num_slices], dtype=torch.int32).to(
1421 - self.device
1422 - )
1423 - slot_mapping = torch.zeros((3, padded_num_slices), dtype=torch.int32).to(
1424 - self.device
1425 - )
1426 - block_tables = torch.zeros((num_reqs, num_blocks), dtype=torch.int32).to(
1427 - self.device
1428 - )
1429 - query_lens = [1] * num_reqs
1430 - query_start_loc = torch.cumsum(
1431 - torch.tensor([0] + query_lens, dtype=torch.int32), dim=0, dtype=torch.int32
1432 - ).to(self.device)
1433 - context_lens = torch.ones((num_reqs,), dtype=torch.int32).to(self.device)
1434 - num_seqs = torch.tensor([actual_num_reqs], dtype=torch.int32).to(self.device)
1435 - attn_metadata = PallasMetadata(
1436 - slot_mapping=slot_mapping,
1437 - block_tables=block_tables,
1438 - context_lens=context_lens,
1439 - query_start_loc=query_start_loc,
1440 - num_seqs=num_seqs,
1441 - num_kv_update_slices=num_kv_update_slices,
1442 - num_slices_per_kv_cache_update_block=self._num_slices_per_kv_cache_update_block,
1443 - )
1444 -
1445 - if self.supports_mm_inputs:
1446 - torch._dynamo.mark_dynamic(inputs_embeds, 0)
1447 - else:
1448 - torch._dynamo.mark_dynamic(input_ids, 0)
1449 - torch._dynamo.mark_dynamic(position_ids, 0)
1450 - torch._dynamo.mark_dynamic(attn_metadata.slot_mapping, 0)
1451 - torch._dynamo.mark_dynamic(attn_metadata.block_tables, (0, 1))
1452 - torch._dynamo.mark_dynamic(attn_metadata.context_lens, 0)
1453 - torch._dynamo.mark_dynamic(attn_metadata.query_start_loc, 0)
1454 -
1455 - layer_names = get_layers_from_vllm_config(self.vllm_config, Attention).keys()
1456 - per_layer_attn_metadata = {
1457 - layer_name: attn_metadata for layer_name in layer_names
1458 - }
1459 -
1460 - with (
1461 - self.maybe_select_dummy_loras(
1462 - self.lora_config, np.array([num_tokens], dtype=np.int32)
1463 - ),
1464 - set_forward_context(per_layer_attn_metadata, self.vllm_config, 0),
1465 - ):
1466 - out = self.model(
1467 - input_ids=input_ids, positions=position_ids, inputs_embeds=inputs_embeds
1468 - )
1469 - self._hidden_states_dtype = out.dtype
1470 -
1471 - def _set_active_loras(
1472 - self,
1473 - prompt_lora_mapping: tuple[int, ...],
1474 - token_lora_mapping: tuple[int, ...],
1475 - lora_requests: set[LoRARequest],
1476 - mapping_type: LoRAMappingType = LoRAMappingType.LANGUAGE,
1477 - ) -> None:
1478 - torch_xla.sync(wait=False) # Captures input updates
1479 - super()._set_active_loras(
1480 - prompt_lora_mapping, token_lora_mapping, lora_requests, mapping_type
1481 - )
1482 - torch_xla.sync(wait=False) # Captures metadata updates
1483 -
1484 - def _precompile_mm_encoder(self) -> None:
1485 - if not self.supports_mm_inputs:
1486 - return
1487 -
1488 - # Pre-compile MM encoder for all supported data modalities.
1489 - hf_config = self.vllm_config.model_config.hf_config
1490 -
1491 - mm_budget = self.mm_budget
1492 - assert mm_budget is not None
1493 -
1494 - max_items_per_seq_by_modality = mm_budget.max_items_per_batch_by_modality # noqa: E501
1495 -
1496 - for mode, max_items_per_seq in max_items_per_seq_by_modality.items():
1497 - logger.info(
1498 - "Compiling Multimodal %s Encoder with different input shapes.", mode
1499 - )
1500 - start = time.perf_counter()
1501 - # No padding for MM encoder just yet.
1502 - for num_items in range(1, max_items_per_seq + 1):
1503 - logger.info(" -- mode: %s items: %d", mode, num_items)
1504 - batched_dummy_mm_inputs = self._get_mm_dummy_batch(
1505 - mode,
1506 - num_items,
1507 - )
1508 - # Run multimodal encoder.
1509 - torch_xla.sync(wait=False)
1510 - mm_embeds = self.model.embed_multimodal(**batched_dummy_mm_inputs)
1511 - torch_xla.sync(wait=False)
1512 - num_patches = mm_embeds[0].shape[0]
1513 - items_size = num_patches * num_items
1514 -
1515 - # NOTE (NickLucche) pre-compile `embed_input_ids` when mm
1516 - # embeddings are present. We assume `--disable-mm-chunked`,
1517 - # hence only whole items can be scheduled. This implies we just
1518 - # need to compile when `num_items` fit the (padded) `input_ids`
1519 - for num_tokens in self.num_tokens_paddings:
1520 - if num_tokens >= items_size:
1521 - # XLA Workaround: if torch.zeros(..device) is used, XLA
1522 - # compiles a scalar+expansion op, which won't match
1523 - # the graph generated at runtime. CPU->TPU must be used
1524 - placeholders_ids = torch.zeros(
1525 - num_tokens, dtype=torch.int32, device="cpu"
1526 - )
1527 - # Align placeholders and actual num mm_embeddings.
1528 - placeholders_ids[:items_size] = hf_config.image_token_index
1529 -
1530 - placeholders_ids = placeholders_ids.to(self.device)
1531 -
1532 - mm_mask = torch.tensor([False] * num_tokens)
1533 - mm_mask[:items_size] = True
1534 - mm_mask = mm_mask.to(self.device)
1535 - # Assign outputs or the graph will be cut short.
1536 - a, b = self._get_model_inputs(
1537 - placeholders_ids,
1538 - mm_embed_inputs=([mm_embeds], mm_mask),
1539 - )
1540 - assert a is None
1541 - torch_xla.sync(wait=False)
1542 -
1543 - # Pre-compile `embed_input_ids` when mm_embeddings are not
1544 - # present. Chunk is only made of text, no mm_placeholders.
1545 - for num_tokens in self.num_tokens_paddings:
1546 - placeholders_ids = torch.zeros(
1547 - num_tokens, dtype=torch.int32, device="cpu"
1548 - )
1549 - placeholders_ids = placeholders_ids.to(self.device)
1550 - a, b = self._get_model_inputs(
1551 - placeholders_ids,
1552 - mm_embed_inputs=None,
1553 - )
1554 - assert a is None
1555 - torch_xla.sync(wait=False)
1556 -
1557 - xm.wait_device_ops()
1558 - end = time.perf_counter()
1559 - logger.info(
1560 - "Multimodal %s Encoder compilation finished in in %.2f [secs].",
1561 - mode,
1562 - end - start,
1563 - )
1564 -
1565 - def _precompile_backbone(self) -> None:
1566 - logger.info("Compiling the model with different input shapes.")
1567 - start = time.perf_counter()
1568 - for num_tokens in self.num_tokens_paddings:
1569 - logger.info(" -- num_tokens: %d", num_tokens)
1570 - self._dummy_run(
1571 - num_tokens, self.num_reqs_max_model_len, self.max_num_blocks_per_req
1572 - )
1573 - if self.most_model_len is not None:
1574 - self._dummy_run(
1575 - num_tokens,
1576 - self.num_reqs_most_model_len,
1577 - self.num_blocks_per_most_len_req,
1578 - )
1579 - xm.wait_device_ops()
1580 - end = time.perf_counter()
1581 - logger.info("Compilation finished in %.2f [secs].", end - start)
1582 - self._update_num_xla_graphs("model backbone")
1583 -
1584 - def _precompile_select_hidden_states(self) -> None:
1585 - # Compile hidden state selection function for bucketed
1586 - # n_tokens x max_num_reqs. Graph is really small so this is fine.
1587 - logger.info("Compiling select_hidden_states with different input shapes.")
1588 - start = time.perf_counter()
1589 - hsize = self.model_config.get_hidden_size()
1590 - for num_tokens in self.num_tokens_paddings:
1591 - dummy_hidden = torch.zeros(
1592 - (num_tokens, hsize), device=self.device, dtype=self._hidden_states_dtype
1593 - )
1594 - torch._dynamo.mark_dynamic(dummy_hidden, 0)
1595 - for num_reqs in self.num_reqs_paddings:
1596 - indices = torch.zeros(num_reqs, dtype=torch.int32, device=self.device)
1597 - torch._dynamo.mark_dynamic(indices, 0)
1598 - self.select_hidden_states(dummy_hidden, indices)
1599 - logger.info(" -- num_tokens: %d, num_seqs: %d", num_tokens, num_reqs)
1600 - # Requests can't be more than tokens. But do compile for the
1601 - # next bigger value in case num_tokens uses bucketed padding.
1602 - if num_reqs >= min(num_tokens, self.max_num_reqs):
1603 - break
1604 - xm.wait_device_ops()
1605 - end = time.perf_counter()
1606 - logger.info("Compilation finished in %.2f [secs].", end - start)
1607 - self._update_num_xla_graphs("select_hidden_states")
1608 -
1609 - def _precompile_compute_logits(self) -> None:
1610 - logger.info("Compiling compute_logits with different input shapes.")
1611 - start = time.perf_counter()
1612 - hsize = self.model_config.get_hidden_size()
1613 - for num_reqs in self.num_reqs_paddings:
1614 - dummy_hidden = torch.zeros(
1615 - (num_reqs, hsize), device=self.device, dtype=self._hidden_states_dtype
1616 - )
1617 - torch._dynamo.mark_dynamic(dummy_hidden, 0)
1618 - self.compute_logits(dummy_hidden)
1619 - logger.info(" -- num_seqs: %d", num_reqs)
1620 - xm.wait_device_ops()
1621 - end = time.perf_counter()
1622 - logger.info("Compilation finished in %.2f [secs].", end - start)
1623 - self._update_num_xla_graphs("compute_logits")
1624 -
1625 - def _precompile_structured_decoding(self) -> None:
1626 - logger.info("Compiling structured_decoding with different input shapes.")
1627 - start = time.perf_counter()
1628 - for num_reqs in self.num_reqs_paddings:
1629 - dummy_logits = torch.zeros(
1630 - (num_reqs, self.vocab_size),
1631 - device=self.device,
1632 - dtype=self._hidden_states_dtype,
1633 - )
1634 - dummy_require_struct_decoding = self.require_structured_out_cpu[
1635 - :num_reqs
1636 - ].to(self.device)
1637 - dummy_grammar_bitmask = self.grammar_bitmask_cpu[:num_reqs].to(self.device)
1638 - # The first dimension of the above 3 dummy tensors cannot be
1639 - # mark_dynamic because some operations in structured_decode require
1640 - # them to be static.
1641 - arange = self.structured_decode_arange.to(self.device)
1642 - self.structured_decode(
1643 - dummy_require_struct_decoding,
1644 - dummy_grammar_bitmask,
1645 - dummy_logits,
1646 - arange,
1647 - )
1648 - logger.info(" -- num_seqs: %d", num_reqs)
1649 - xm.wait_device_ops()
1650 - end = time.perf_counter()
1651 - logger.info("Compilation finished in %.2f [secs].", end - start)
1652 - self._update_num_xla_graphs("structured_decoding")
1653 -
1654 - def _precompile_sample_from_logits(self) -> None:
1655 - logger.info("Compiling sample_from_logits with different input shapes.")
1656 - start = time.perf_counter()
1657 - for num_reqs in self.num_reqs_paddings:
1658 - dummy_logits = torch.zeros(
1659 - (num_reqs, self.vocab_size),
1660 - device=self.device,
1661 - dtype=self._hidden_states_dtype,
1662 - )
1663 - # The first dimension of dummy_logits cannot be mark_dynamic
1664 - # because some operations in the sampler require it to be static.
1665 - for all_greedy in [False, True]:
1666 - generate_params_if_all_greedy = not all_greedy
1667 - sampling_metadata = TPUSupportedSamplingMetadata.from_input_batch(
1668 - self.input_batch,
1669 - num_reqs,
1670 - self.device,
1671 - generate_params_if_all_greedy,
1672 - )
1673 - sampling_metadata.all_greedy = all_greedy
1674 - with self.maybe_select_dummy_loras(
1675 - self.lora_config, np.array([num_reqs], dtype=np.int32)
1676 - ):
1677 - self.sample_from_logits_func(dummy_logits, sampling_metadata)
1678 - logger.info(" -- num_seqs: %d", num_reqs)
1679 - xm.wait_device_ops()
1680 - end = time.perf_counter()
1681 - logger.info("Compilation finished in %.2f [secs].", end - start)
1682 - self._update_num_xla_graphs("sample_from_logits")
1683 -
1684 - def _precompile_gather_logprobs(self) -> None:
1685 - logger.info("Compiling gather_logprobs with different input shapes.")
1686 - start = time.perf_counter()
1687 - for num_reqs in self.num_reqs_paddings:
1688 - dummy_logits = torch.zeros(
1689 - (num_reqs, self.vocab_size),
1690 - device=self.device,
1691 - dtype=self._hidden_states_dtype,
1692 - )
1693 - dummy_tokens = torch.zeros((num_reqs, 1), dtype=torch.int64).to(self.device)
1694 - with self.maybe_select_dummy_loras(
1695 - self.lora_config, np.array([num_reqs], dtype=np.int32)
1696 - ):
1697 - self.gather_logprobs(dummy_logits, dummy_tokens)
1698 - logger.info(" -- num_seqs: %d", num_reqs)
1699 - xm.wait_device_ops()
1700 - end = time.perf_counter()
1701 - logger.info("Compilation finished in %.2f [secs].", end - start)
1702 - self._update_num_xla_graphs("gather_logprobs")
1703 -
1704 - def capture_model(self) -> None:
1705 - """
1706 - Precompile all the subgraphs with possible input shapes.
1707 - """
1708 - with self.maybe_setup_dummy_loras(self.lora_config):
1709 - self._precompile_mm_encoder()
1710 - self._precompile_backbone()
1711 - self._precompile_select_hidden_states()
1712 - self._precompile_compute_logits()
1713 - self._precompile_structured_decoding()
1714 - self._precompile_sample_from_logits()
1715 - self._precompile_gather_logprobs()
1716 -
1717 - def profile_run(
1718 - self,
1719 - num_tokens: int,
1720 - ) -> None:
1721 - # Profile with multimodal encoder & encoder cache.
1722 - if self.supports_mm_inputs:
1723 - mm_config = self.model_config.multimodal_config
1724 - if mm_config is not None and mm_config.skip_mm_profiling:
1725 - logger.info(
1726 - "Skipping memory profiling for multimodal encoder and "
1727 - "encoder cache."
1728 - )
1729 - else:
1730 - mm_budget = self.mm_budget
1731 - assert mm_budget is not None
1732 -
1733 - # TODO: handle encoder-decoder models once we support them.
1734 - if (encoder_budget := mm_budget.get_encoder_budget()) > 0:
1735 - # NOTE: Currently model is profiled with a single non-text
1736 - # modality with the max possible input tokens even when
1737 - # it supports multiple.
1738 - dummy_modality = mm_budget.get_modality_with_max_tokens()
1739 - max_mm_items_per_batch = mm_budget.max_items_per_batch_by_modality[
1740 - dummy_modality
1741 - ]
1742 -
1743 - logger.info(
1744 - "Encoder cache will be initialized with a budget of "
1745 - "%s tokens, and profiled with %s %s items of the "
1746 - "maximum feature size.",
1747 - encoder_budget,
1748 - max_mm_items_per_batch,
1749 - dummy_modality,
1750 - )
1751 -
1752 - # Create dummy batch of multimodal inputs.
1753 - batched_dummy_mm_inputs = self._get_mm_dummy_batch(
1754 - dummy_modality,
1755 - max_mm_items_per_batch,
1756 - )
1757 -
1758 - # Run multimodal encoder.
1759 - # Isolate encoder graph from post-processing to minimize
1760 - # impact of recompilation until it's fixed.
1761 - start = time.perf_counter()
1762 - torch_xla.sync(wait=False)
1763 - dummy_encoder_outputs = self.model.embed_multimodal(
1764 - **batched_dummy_mm_inputs
1765 - )
1766 - torch_xla.sync(wait=False)
1767 - xm.wait_device_ops()
1768 - end = time.perf_counter()
1769 - logger.info(
1770 - "Multimodal Encoder profiling finished in %.2f [secs].",
1771 - end - start,
1772 - )
1773 -
1774 - sanity_check_mm_encoder_outputs(
1775 - dummy_encoder_outputs,
1776 - expected_num_items=max_mm_items_per_batch,
1777 - )
1778 -
1779 - # Cache the dummy encoder outputs.
1780 - self.encoder_cache["tmp"] = dict(enumerate(dummy_encoder_outputs))
1781 -
1782 - # Trigger compilation for general shape.
1783 - self._dummy_run(
1784 - num_tokens, self.num_reqs_max_model_len, self.max_num_blocks_per_req
1785 - )
1786 - if self.most_model_len is not None:
1787 - self._dummy_run(
1788 - num_tokens,
1789 - self.num_reqs_most_model_len,
1790 - self.num_blocks_per_most_len_req,
1791 - )
1792 -
1793 - torch_xla.sync(wait=False)
1794 - xm.wait_device_ops()
1795 - self.encoder_cache.clear()
1796 - gc.collect()
1797 -
1798 - def maybe_setup_cross_layer_kv_sharing(
1799 - self,
1800 - kv_caches: dict[str, torch.Tensor],
1801 - kv_cache_config: KVCacheConfig,
1802 - ) -> None:
1803 - """
1804 - Add layers that re-use KV cache to KV cache group of its target layer.
1805 - Mapping of KV cache tensors happens in `initialize_kv_cache_tensors()`
1806 - """
1807 - if not self.shared_kv_cache_layers:
1808 - # No cross-layer KV sharing, return
1809 - return
1810 -
1811 - add_kv_sharing_layers_to_kv_cache_groups(
1812 - self.shared_kv_cache_layers,
1813 - kv_cache_config.kv_cache_groups,
1814 - )
1815 -
1816 - for layer_name, target_layer_name in self.shared_kv_cache_layers.items():
1817 - logger.debug("%s reuses KV cache of %s", layer_name, target_layer_name)
1818 - kv_caches[layer_name] = kv_caches[target_layer_name]
1819 -
1820 - def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None:
1821 - """
1822 - Initialize KV cache based on `kv_cache_config`.
1823 - Args:
1824 - kv_cache_config: Configuration for the KV cache, including the KV
1825 - cache size of each layer
1826 - """
1827 - if len(kv_cache_config.kv_cache_groups) > 1:
1828 - raise NotImplementedError(
1829 - "Hybrid models with more than one KV cache type are not supported yet."
1830 - )
1831 -
1832 - if (
1833 - kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size
1834 - != self.block_size
1835 - ):
1836 - self.input_batch = InputBatch(
1837 - max_num_reqs=self.max_num_reqs,
1838 - max_model_len=self.max_model_len,
1839 - max_num_batched_tokens=self.max_num_tokens,
1840 - device=self.device,
1841 - pin_memory=self.pin_memory,
1842 - vocab_size=self.model_config.get_vocab_size(),
1843 - block_sizes=[
1844 - kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size
1845 - ],
1846 - kernel_block_sizes=[
1847 - kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size
1848 - ],
1849 - )
1850 - # Verify dtype compatibility between block_table_cpu and input_batch
1851 - assert (
1852 - self.block_table_cpu.dtype
1853 - == self.input_batch.block_table[0].get_cpu_tensor().dtype
1854 - )
1855 -
1856 - kv_cache_sizes = {}
1857 - for kv_cache_tensor in kv_cache_config.kv_cache_tensors:
1858 - assert len(kv_cache_tensor.shared_by) == 1, (
1859 - "KV cache tensor shared by multiple layers is not supported in TPU."
1860 - )
1861 - kv_cache_sizes[kv_cache_tensor.shared_by[0]] = kv_cache_tensor.size
1862 -
1863 - kv_caches: dict[str, torch.Tensor] = {}
1864 - for kv_cache_group in kv_cache_config.kv_cache_groups:
1865 - kv_cache_spec = kv_cache_group.kv_cache_spec
1866 - for layer_name in kv_cache_group.layer_names:
1867 - tensor_size = kv_cache_sizes[layer_name]
1868 - assert tensor_size % kv_cache_spec.page_size_bytes == 0
1869 - num_blocks = tensor_size // kv_cache_spec.page_size_bytes # noqa
1870 - if isinstance(kv_cache_spec, AttentionSpec):
1871 - if self.use_spmd:
1872 - num_kv_heads = kv_cache_spec.num_kv_heads
1873 - assert self.original_parallel_config is not None
1874 - tp_size = self.original_parallel_config.tensor_parallel_size
1875 - # TODO: Handle kv cache duplication under SPMD mode.
1876 - assert num_kv_heads % tp_size == 0, (
1877 - f"num_kv_heads {num_kv_heads} must be divisible by "
1878 - f"tp_size {tp_size} under SPMD mode"
1879 - )
1880 - kv_cache_shape = PallasAttentionBackend.get_kv_cache_shape(
1881 - num_blocks,
1882 - kv_cache_spec.block_size,
1883 - kv_cache_spec.num_kv_heads,
1884 - kv_cache_spec.head_size,
1885 - )
1886 - dtype = kv_cache_spec.dtype
1887 -
1888 - tpu_kv_cache = torch.zeros(kv_cache_shape, dtype=dtype).to(
1889 - self.device
1890 - )
1891 -
1892 - kv_caches[layer_name] = tpu_kv_cache
1893 - else:
1894 - raise NotImplementedError
1895 -
1896 - # Set up cross-layer KV cache sharing if needed
1897 - self.maybe_setup_cross_layer_kv_sharing(kv_caches, kv_cache_config)
1898 -
1899 - bind_kv_cache(
1900 - kv_caches,
1901 - self.vllm_config.compilation_config.static_forward_context,
1902 - self.kv_caches,
1903 - )
1904 -
1905 - if self.use_spmd:
1906 - # Shard KV Cache
1907 - for cache in self.kv_caches:
1908 - xs.mark_sharding(cache, self.mesh, (None, "x", None, None))
1909 -
1910 - if has_kv_transfer_group():
1911 - get_kv_transfer_group().register_kv_caches(kv_caches)
1912 - get_kv_transfer_group().set_host_xfer_buffer_ops(copy_kv_blocks)
1913 -
1914 - def reset_dynamo_cache(self):
1915 - # NOTE: We check `is_multimodal_model` instead of `supports_mm_inputs`
1916 - # since the compiled model object of the language backbone of a
1917 - # multimodal model needs to be extracted via `get_language_model`.
1918 - if self.model_config.is_multimodal_model:
1919 - compiled_model = self.model.get_language_model().model
1920 - else:
1921 - compiled_model = self.model.model
1922 - if isinstance(compiled_model, TorchCompileWithNoGuardsWrapper):
1923 - logger.info("Clear dynamo cache and cached dynamo bytecode.")
1924 - torch._dynamo.eval_frame.remove_from_cache(
1925 - compiled_model.original_code_object()
1926 - )
1927 - # Reset the wrapper to re-initialize.
1928 - compiled_model.compiled = False
1929 - TorchCompileWithNoGuardsWrapper.__init__(compiled_model)
1930 -
1931 - @torch.compile(backend="openxla", fullgraph=True, dynamic=False)
1932 - def select_hidden_states(self, hidden_states, indices_do_sample):
1933 - return hidden_states[indices_do_sample]
1934 -
1935 - @torch.compile(backend="openxla", fullgraph=True, dynamic=False)
1936 - def compute_logits(self, sample_hidden_states: torch.Tensor) -> torch.Tensor:
1937 - return self.model.compute_logits(sample_hidden_states)
1938 -
1939 - # TODO: Under SPMD mode, sample_from_logits has correctness issue.
1940 - # Re-enable the torch.compile once the issue is fixed in torchxla.
1941 - # @torch.compile(backend="openxla", fullgraph=True, dynamic=False)
1942 - def sample_from_logits(
1943 - self, logits: torch.Tensor, sampling_metadata: TPUSupportedSamplingMetadata
1944 - ) -> torch.Tensor:
1945 - """
1946 - Sample with xla-friendly function. This function is to be traced
1947 - separately from `forward` for lighter compilation overhead.
1948 - """
1949 - if sampling_metadata.all_greedy:
1950 - out_tokens = torch.argmax(logits, dim=-1, keepdim=True)
1951 - else:
1952 - out_tokens = self.sampler(logits, sampling_metadata).sampled_token_ids
1953 - return out_tokens
1954 -
1955 - @torch.compile(backend="openxla", fullgraph=True, dynamic=False)
1956 - def gather_logprobs(
1957 - self, logits: torch.Tensor, sampled_tokens: torch.Tensor
1958 - ) -> LogprobsTensors:
1959 - """
1960 - Gather the top_logprobs with corresponding tokens. Use a fixed number
1961 - of logprobs as an alternative to having multiple pre-compiled graphs.
1962 - Select the number of logprobs actually demanded by each request on CPU.
1963 - """
1964 - logprobs = self.sampler.compute_logprobs(logits)
1965 - return self.sampler.gather_logprobs(
1966 - logprobs,
1967 - self.model_config.max_logprobs,
1968 - token_ids=sampled_tokens.squeeze(-1),
1969 - )
1970 -
1971 - @torch.compile(backend="openxla", fullgraph=True, dynamic=False)
1972 - def structured_decode(
1973 - self,
1974 - require_struct_decoding: torch.Tensor,
1975 - grammar_bitmask: torch.Tensor,
1976 - logits: torch.Tensor,
1977 - arange: torch.Tensor,
1978 - ) -> torch.Tensor:
1979 - return torch.where(
1980 - require_struct_decoding,
1981 - self.apply_grammar_bitmask(logits, grammar_bitmask, arange),
1982 - logits,
1983 - )
1984 -
1985 - def apply_grammar_bitmask(
1986 - self, logits: torch.Tensor, grammar_bitmask: torch.Tensor, arange: torch.Tensor
1987 - ):
1988 - assert logits.shape[0] == grammar_bitmask.shape[0]
1989 - logits_cloned = logits.clone()
1990 - for i in range(logits.shape[0]):
1991 - unpacked_bitmask = (
1992 - torch.bitwise_right_shift(grammar_bitmask[i][:, None], arange[None, :])
1993 - & 1
1994 - ) == 0
1995 - unpacked_bitmask = unpacked_bitmask.reshape(-1)[: self.vocab_size]
1996 - logits_cloned[i] = logits_cloned[i].masked_fill(
1997 - unpacked_bitmask, -float("inf")
1998 - )
1999 - return logits_cloned
2000 -
2001 - def embed_multimodal(self, *args, **kwargs):
2002 - return self.model.embed_multimodal(*args, **kwargs)
2003 -
2004 - def embed_input_ids(self, *args, **kwargs):
2005 - return self.model.embed_input_ids(*args, **kwargs)
2006 -
2007 - def prepare_structured_decoding_input(
2008 - self, logits: torch.Tensor, grammar_output: "GrammarOutput"
2009 - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
2010 - grammar_bitmask = grammar_output.grammar_bitmask
2011 - num_reqs, _ = logits.shape
2012 -
2013 - # Reset pre-allocated tensors
2014 - self.grammar_bitmask_cpu.zero_()
2015 - self.require_structured_out_cpu.zero_()
2016 -
2017 - cumulative_mask_idx = 0
2018 - for req_id in grammar_output.structured_output_request_ids:
2019 - if req_id not in self.input_batch.req_id_to_index:
2020 - continue
2021 - batch_index = self.input_batch.req_id_to_index[req_id]
2022 - self.grammar_bitmask_cpu[batch_index] = torch.from_numpy(
2023 - grammar_bitmask[cumulative_mask_idx]
2024 - )
2025 - # It's not guaranteed that all requests in this batch require
2026 - # structured output, so create a bool tensor to represent
2027 - # the requests that need structured output.
2028 - self.require_structured_out_cpu[batch_index] = True
2029 - cumulative_mask_idx += 1
2030 -
2031 - return (
2032 - self.require_structured_out_cpu[:num_reqs].to(logits.device),
2033 - self.grammar_bitmask_cpu[:num_reqs].to(logits.device),
2034 - self.structured_decode_arange.to(logits.device),
2035 - )
2036 -
2037 - def _get_mm_dummy_batch(
2038 - self,
2039 - modality: str,
2040 - max_items_per_batch: int,
2041 - ) -> BatchedTensorInputs:
2042 - """Dummy data for profiling and precompiling multimodal models."""
2043 - assert self.mm_budget is not None
2044 -
2045 - dummy_decoder_data = self.mm_registry.get_decoder_dummy_data(
2046 - model_config=self.model_config,
2047 - seq_len=self.max_model_len,
2048 - mm_counts={modality: 1},
2049 - cache=self.mm_budget.cache,
2050 - )
2051 - dummy_mm_data = dummy_decoder_data.multi_modal_data
2052 -
2053 - # Result in the maximum GPU consumption of the model
2054 - dummy_mm_item = dummy_mm_data[modality][0]
2055 - dummy_mm_items = [dummy_mm_item] * max_items_per_batch
2056 -
2057 - return next(
2058 - grouped_mm_kwargs
2059 - for _, _, grouped_mm_kwargs in group_mm_kwargs_by_modality(
2060 - dummy_mm_items,
2061 - device=self.device,
2062 - pin_memory=self.pin_memory,
2063 - )
2064 - )
2065 -
2066 -
2067 -def _get_req_paddings(min_req_size: int, max_req_size: int) -> list[int]:
2068 - logger.info("Preparing request paddings:")
2069 - # assert min_req_size is power of 2
2070 - assert (min_req_size & (min_req_size - 1) == 0) and min_req_size > 0
2071 - paddings: list = []
2072 - num = max(MIN_NUM_SEQS, min_req_size)
2073 - while num <= max_req_size and (len(paddings) == 0 or paddings[-1] != num):
2074 - paddings.append(num)
2075 - logger.info(" %d", num)
2076 - num = _get_padded_num_reqs_with_upper_limit(num + 1, max_req_size)
2077 - return paddings
2078 -
2079 -
2080 -def _get_padded_num_reqs_with_upper_limit(x: int, upper_limit: int) -> int:
2081 - res = MIN_NUM_SEQS if x <= MIN_NUM_SEQS else 1 << (x - 1).bit_length()
2082 - return min(res, upper_limit)
2083 -
2084 -
2085 -def _get_token_paddings(
2086 - min_token_size: int, max_token_size: int, padding_gap: int
2087 -) -> list[int]:
2088 - """Generate a list of padding size, starting from min_token_size,
2089 - ending with a number that can cover max_token_size
2090 -
2091 - If padding_gap == 0 then:
2092 - increase 2X each time (exponential)
2093 - else:
2094 - first increase the size to twice,
2095 - then increase the padding size by padding_gap.
2096 - """
2097 - # assert min_token_size is power of 2
2098 - assert (min_token_size & (min_token_size - 1) == 0) and min_token_size > 0
2099 - paddings = []
2100 - num = min_token_size
2101 -
2102 - if padding_gap == 0:
2103 - logger.info("Using exponential token paddings:")
2104 - while True:
2105 - logger.info(" %d", num)
2106 - paddings.append(num)
2107 - if num >= max_token_size:
2108 - break
2109 - num *= 2
2110 - else:
2111 - logger.info("Using incremental token paddings:")
2112 - while num <= padding_gap:
2113 - logger.info(" %d", num)
2114 - paddings.append(num)
2115 - num *= 2
2116 - num //= 2
2117 - while num < max_token_size:
2118 - num += padding_gap
2119 - logger.info(" %d", num)
2120 - paddings.append(num)
2121 -
2122 - return paddings
2123 -
2124 -
2125 -def _get_padded_token_len(paddings: list[int], x: int) -> int:
2126 - """Return the first element in paddings list greater or equal to x."""
2127 - index = bisect.bisect_left(paddings, x)
2128 - assert index < len(paddings)
2129 - return paddings[index]
2130 -
2131 -
2132 -def _get_padded_num_kv_cache_update_slices(
2133 - num_tokens: int, max_num_reqs: int, page_size: int
2134 -) -> int:
2135 - """Calculates the padded number of KV cache update slices to avoid
2136 - recompilation."""
2137 - # NOTE(chengjiyao): let's say R_i is the token num for i-th request,
2138 - # so it occupies most 2 + R_i // page_size pages. The total maximum
2139 - # possible number of pages needed is sum(2 + R_i // page_size), which
2140 - # is <= 2 * max_num_reqs + sum(R_i) // page_size
2141 - # = 2 * max_num_reqs + num_tokens // page_size
2142 - padded_num_slices = 2 * max_num_reqs + num_tokens // page_size
2143 - padded_num_slices = min(padded_num_slices, num_tokens)
2144 - return padded_num_slices
2145 -
2146 -
2147 -def _get_num_slices_per_kv_cache_update_block(page_size_bytes: int) -> int:
2148 - """Find the optimum number of slices to copy per Pallas program instance.
2149 -
2150 - Increasing the number of slices copied in one instance of the kernel program
2151 - will increase HBM bandwidth utilization via more in-flight DMAs.
2152 -
2153 - However, it will also use more VMEM, and experimentally, we observed
2154 - performance regression at 128 slices on v6e, likely due to running
2155 - out of scalar registers. Thus this function will limit the number of
2156 - slices to 64.
2157 - """
2158 - # The default vmem_limit_bytes of a pallas kernel is 32MB. Here we
2159 - # calculate num_slices_per_block based on 16MB in case any register spills.
2160 - vmem_limit = 16 * 1024 * 1024
2161 - num_slices_per_block = vmem_limit // page_size_bytes
2162 - assert num_slices_per_block > 0, "Number of slices should be positive"
2163 - num_slices_per_block = prev_power_of_2(num_slices_per_block)
2164 - if num_slices_per_block > 64:
2165 - num_slices_per_block = 64
2166 - return num_slices_per_block
2167 -
2168 -
2169 -def replace_set_lora(model):
2170 - def _tpu_set_lora(
2171 - self,
2172 - index: int,
2173 - lora_a: torch.Tensor,
2174 - lora_b: torch.Tensor,
2175 - embeddings_tensor: torch.Tensor | None,
2176 - ):
2177 - # TODO: The integer index leads to a recompilation, but converting it
2178 - # to a tensor doesn't seem to work anymore. This might be fixed with a
2179 - # later release of torch_xla.
2180 - self._original_set_lora(index, lora_a, lora_b, embeddings_tensor)
2181 - torch_xla.sync(wait=False)
2182 -
2183 - def _tpu_reset_lora(self, index: int):
2184 - self._original_reset_lora(index)
2185 - torch_xla.sync(wait=False)
2186 -
2187 - for _, module in model.named_modules():
2188 - if isinstance(module, BaseLayerWithLoRA):
2189 - module._original_set_lora = module.set_lora
2190 - module._original_reset_lora = module.reset_lora
2191 - module.set_lora = _tpu_set_lora.__get__( # type: ignore[method-assign]
2192 - module, module.__class__
2193 - )
2194 - module.reset_lora = _tpu_reset_lora.__get__( # type: ignore[method-assign]
2195 - module, module.__class__
2196 - )
vllm/v1/worker/tpu_worker.py
+2 -334
@@ -2,348 +2,16 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """A TPU worker class."""
4
5 -import os
6 -from collections.abc import Callable
7 -from typing import Any, TypeVar
5 +from typing import TypeVar
6
9 -import torch
10 -import torch.nn as nn
11 -
12 -import vllm.envs as envs
13 -from vllm.config import VllmConfig, set_current_vllm_config
14 -from vllm.distributed import (
15 - ensure_model_parallel_initialized,
16 - init_distributed_environment,
17 -)
18 -from vllm.distributed.kv_transfer import (
19 - ensure_kv_transfer_initialized,
20 -)
7 from vllm.logger import init_logger
22 -from vllm.lora.request import LoRARequest
23 -from vllm.platforms import current_platform
8 from vllm.platforms.tpu import USE_TPU_INFERENCE
25 -from vllm.tasks import SupportedTask
26 -from vllm.utils.math_utils import cdiv
27 -from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, set_random_seed
28 -from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
29 -from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheConfig, KVCacheSpec
30 -from vllm.v1.outputs import ModelRunnerOutput
31 -from vllm.v1.utils import report_usage_stats
32 -from vllm.v1.worker.utils import bind_kv_cache
9
10 logger = init_logger(__name__)
11
12 _R = TypeVar("_R")
13
38 -if not USE_TPU_INFERENCE:
39 - logger.info("tpu_inference not found, using vLLM's TPUWorker.")
40 - import torch_xla.core.xla_model as xm
41 - import torch_xla.debug.profiler as xp
42 - import torch_xla.runtime as xr
43 -
44 - from vllm.v1.attention.backends.pallas import TPU_HEAD_SIZE_ALIGNMENT
45 - from vllm.v1.worker.tpu_model_runner import TPUModelRunner
46 -
47 -
48 -class TPUWorker:
49 - def __init__(
50 - self,
51 - vllm_config: VllmConfig,
52 - local_rank: int,
53 - rank: int,
54 - distributed_init_method: str,
55 - is_driver_worker: bool = False,
56 - ):
57 - self.is_driver_worker = is_driver_worker
58 - self.vllm_config = vllm_config
59 - self.model_config = vllm_config.model_config
60 - self.cache_config = vllm_config.cache_config
61 - self.lora_config = vllm_config.lora_config
62 - self.load_config = vllm_config.load_config
63 - self.parallel_config = vllm_config.parallel_config
64 - self.use_spmd = envs.VLLM_XLA_USE_SPMD
65 - self.original_parallel_config = None
66 - if self.use_spmd:
67 - # Under SPMD mode, distributed env is initialized as if there is
68 - # only one worker/device.
69 - self.original_parallel_config = self.parallel_config
70 - self.parallel_config.tensor_parallel_size = 1
71 - self.parallel_config.pipeline_parallel_size = 1
72 - self.parallel_config.world_size = 1
73 - self.scheduler_config = vllm_config.scheduler_config
74 - self.device_config = vllm_config.device_config
75 - self.speculative_config = vllm_config.speculative_config
76 - self.observability_config = vllm_config.observability_config
77 -
78 - self.parallel_config.rank = rank
79 - self.local_rank = local_rank
80 - self.rank = rank
81 - self.distributed_init_method = distributed_init_method
82 -
83 - if self.cache_config.cache_dtype == "auto":
84 - self.cache_dtype = self.model_config.dtype
85 - else:
86 - self.cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[self.cache_config.cache_dtype]
87 -
88 - # Delay profiler initialization to the start of the profiling.
89 - # This is because in vLLM V1, MP runtime is initialized before the
90 - # TPU Worker is initialized. The profiler server needs to start after
91 - # MP runtime is initialized.
92 - self.profiler = None
93 - self.profile_dir = None
94 - if vllm_config.profiler_config.profiler == "torch" and self.rank < 1:
95 - # For TPU, we can only have 1 active profiler session for 1 profiler
96 - # server. So we only profile on rank0.
97 - self.profile_dir = vllm_config.profiler_config.torch_profiler_dir
98 - logger.info(
99 - "Profiling enabled. Traces will be saved to: %s", self.profile_dir
100 - )
101 -
102 - def initialize_cache(self, num_gpu_blocks: int, num_cpu_blocks: int) -> None:
103 - self.cache_config.num_gpu_blocks = num_gpu_blocks
104 - self.cache_config.num_cpu_blocks = num_cpu_blocks
105 -
106 - def init_device(self):
107 - os.environ["PJRT_DEVICE"] = "TPU"
108 - # Note: Currently the XLA compiler wrongly uses 2D ring strategy on 1D
109 - # ring, the xla tpu compiler flag
110 - # `xla_tpu_force_1d_allreduce_at_chunk_count` is a temporary solution to
111 - # fix this. It will be removed after the bug in XLA compiler is fixed.
112 - os.environ["LIBTPU_INIT_ARGS"] = (
113 - os.environ.get("LIBTPU_INIT_ARGS", "")
114 - + " --xla_tpu_force_1d_allreduce_at_chunk_count=1"
115 - " --xla_jf_conv_input_fusion=False"
116 - )
117 - # --xla_jf_conv_input_fusion=False is used to improve the perf of
118 - # quantized matmul.
119 - torch.set_grad_enabled(False)
120 - torch.set_default_dtype(self.model_config.dtype)
121 -
122 - # Initialize the distributed environment.
123 - self._init_tpu_worker_distributed_environment(
124 - self.vllm_config, self.rank, self.distributed_init_method, self.local_rank
125 - )
126 -
127 - # Device initialization should happen after initializing
128 - # the distributed runtime.
129 - self.device = xm.xla_device()
130 - self.device_config.device = self.device
131 -
132 - # Set random seed.
133 - set_random_seed(self.model_config.seed)
134 - xm.set_rng_state(self.model_config.seed, self.device)
135 -
136 - # Increase the cache size limit, which is the maximum number of
137 - # dynamo graphs that can be compiled.
138 - # TODO (NickLucche) On gsm we compile 80+ graphs.
139 - # Re-evaluate limit, with MM we may get close to this limit.
140 - torch._dynamo.config.cache_size_limit = 128
141 - # Use persistent cache to avoid XLA recompilation.
142 - # NOTE(woosuk): Set per-rank cache path since different ranks
143 - # can have slightly different XLA graphs.
144 - world_size = self.parallel_config.world_size
145 - rank = xr.global_ordinal()
146 - # The PyTorch/XLA compilation cache uses the Torch IR to generate keys.
147 - # Consequently, changes in optimization flags, which affect compilation
148 - # results, don't change the cache key. This can result in the wrong
149 - # compilation being used. To prevent this, disabling the XLA compilation
150 - # cache during development is recommended.We can disable it by
151 - # `export VLLM_XLA_CACHE_PATH=`
152 - if envs.VLLM_XLA_CACHE_PATH:
153 - per_rank_path = os.path.join(
154 - envs.VLLM_XLA_CACHE_PATH, f"tp{world_size}_rank{rank}"
155 - )
156 - xr.initialize_cache(per_rank_path, readonly=False)
157 -
158 - # Init ModelRunner here, so that we have access to self.device.
159 - self.model_runner = TPUModelRunner(
160 - self.vllm_config, self.device, self.original_parallel_config
161 - )
162 -
163 - if rank == 0:
164 - # If usage stat is enabled, collect relevant info.
165 - report_usage_stats(self.vllm_config)
166 -
167 - def determine_available_memory(self) -> int:
168 - kv_caches: dict[str, torch.Tensor] = {}
169 - kv_cache_spec = self.model_runner.get_kv_cache_spec()
170 - for layer_name, layer_spec in kv_cache_spec.items():
171 - if isinstance(layer_spec, AttentionSpec):
172 - dtype = layer_spec.dtype
173 -
174 - # Use an empty tensor instead of `None` to force Dynamo to pass
175 - # it by reference, rather by specializing on the value `None`.
176 - tpu_kv_cache = torch.tensor([], dtype=dtype).to(self.device)
177 - kv_caches[layer_name] = tpu_kv_cache
178 - else:
179 - raise NotImplementedError(
180 - f"Unsupported KV cache spec '{type(layer_spec)}'"
181 - )
182 -
183 - runner_kv_caches: list[torch.Tensor] = []
184 - bind_kv_cache(
185 - kv_caches,
186 - self.vllm_config.compilation_config.static_forward_context,
187 - runner_kv_caches,
188 - )
189 -
190 - # `max_num_tokens >= max_num_batched_tokens` due to padding.
191 - with self.model_runner.maybe_setup_dummy_loras(self.lora_config):
192 - self.model_runner.profile_run(self.model_runner.max_num_tokens)
193 -
194 - # Synchronize before measuring the memory usage.
195 - xm.wait_device_ops()
196 -
197 - # During the profiling run, the model runs without KV cache. After
198 - # the profiling run, the model always runs with KV cache. Here we clear
199 - # the dynamo cache and cached bytecode to ensure the model always has
200 - # one compiled bytecode. Having one FX graph/cached bytecode per
201 - # compiled model is required for `support_torch_compile` decorator to
202 - # skip dynamo guard.
203 - with set_current_vllm_config(self.vllm_config):
204 - self.model_runner.reset_dynamo_cache()
205 -
206 - # Get the maximum amount of memory used by the model weights and
207 - # intermediate activations.
208 - if self.use_spmd:
209 - # This is a workaround for the TPU SPMD mode. The get_memory_info
210 - # API doesn't work with SPMD mode in PyTorch/XLA.
211 - # TODO: use xm.get_memory_info for SPMD once it's supported in
212 - # PyTorch/XLA.
213 - import tpu_info
214 -
215 - chip_type, _ = tpu_info.device.get_local_chips()
216 - device_usage = tpu_info.metrics.get_chip_usage(chip_type)
217 - total_memory_size = device_usage[0].total_memory
218 - current_mem = device_usage[0].memory_usage
219 - else:
220 - m = xm.get_memory_info(self.device)
221 - total_memory_size = m["bytes_limit"]
222 - current_mem = m["bytes_used"]
223 - # Ideally we would use profiled = m["peak_bytes_used"] to
224 - # get weights + activations. But there is memory used during
225 - # compilation / weight loading that impacts the peak and
226 - # there is no way to reset peak memory in XLA, So we
227 - # use the heuristic of 2% of weights.
228 - profiled = current_mem * 1.02
229 -
230 - # Calculate the TPU KV cache size based on profiling.
231 - usable_memory_size = int(
232 - total_memory_size * self.cache_config.gpu_memory_utilization
233 - )
234 - tpu_kv_cache_bytes = max(usable_memory_size - profiled, 0)
235 - head_size = self.model_config.get_head_size()
236 - if head_size > 0:
237 - padded_head_size = (
238 - cdiv(head_size, TPU_HEAD_SIZE_ALIGNMENT) * TPU_HEAD_SIZE_ALIGNMENT
239 - )
240 - if padded_head_size != head_size:
241 - logger.warning_once("head size is padded to %d", padded_head_size)
242 - # We adjust the usable memory size for the KV cache to prevent OOM
243 - # errors, even after padding the head_size.
244 - tpu_kv_cache_bytes = tpu_kv_cache_bytes * head_size // padded_head_size
245 - return int(tpu_kv_cache_bytes)
246 -
247 - def sample_tokens(self, grammar_output: "GrammarOutput") -> ModelRunnerOutput:
248 - return self.model_runner.sample_tokens(grammar_output)
249 -
250 - def execute_model(
251 - self, scheduler_output: "SchedulerOutput"
252 - ) -> ModelRunnerOutput | None:
253 - return self.model_runner.execute_model(scheduler_output)
254 -
255 - def profile(self, is_start: bool = True):
256 - if self.rank < 1:
257 - if self.profile_dir is None:
258 - raise RuntimeError("Profiler is not enabled.")
259 - if is_start:
260 - if self.profiler is None:
261 - self.profiler = xp.start_server(9012)
262 - xp.start_trace(self.profile_dir)
263 - else:
264 - xp.stop_trace()
265 -
266 - def add_lora(self, lora_request: LoRARequest) -> bool:
267 - return self.model_runner.add_lora(lora_request)
268 -
269 - def load_model(self) -> None:
270 - self.model_runner.load_model()
271 -
272 - def update_config(self, overrides: dict[str, Any]) -> None:
273 - self.model_runner.update_config(overrides)
274 -
275 - def reload_weights(self) -> None:
276 - self.model_runner.reload_weights()
277 -
278 - def compile_or_warm_up_model(self) -> None:
279 - if not self.model_config.enforce_eager:
280 - self.model_runner.capture_model()
281 -
282 - # Reset the seed to ensure that the random state is not affected by
283 - # the model initialization and profiling.
284 - set_random_seed(self.model_config.seed)
285 -
286 - def reset_mm_cache(self) -> None:
287 - self.model_runner.reset_mm_cache()
288 -
289 - def get_model(self) -> nn.Module:
290 - return self.model_runner.get_model()
291 -
292 - def get_supported_tasks(self) -> tuple[SupportedTask, ...]:
293 - return self.model_runner.get_supported_tasks()
294 -
295 - def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:
296 - return self.model_runner.get_kv_cache_spec()
297 -
298 - def initialize_from_config(self, kv_cache_config: KVCacheConfig) -> None:
299 - """Allocate GPU KV cache with the specified kv_cache_config."""
300 - # Init kv cache connector here, because it requires
301 - # `kv_cache_config`.
302 - # NOTE(Kuntai): This need to be done before `initialize_kv_cache`,
303 - # because `initialize_kv_cache` will inject kv cache groups not
304 - # related to kv cache connector (e.g. kv cache sharing layers).
305 - ensure_kv_transfer_initialized(self.vllm_config, kv_cache_config)
306 -
307 - self.model_runner.initialize_kv_cache(kv_cache_config)
308 -
309 - def check_health(self) -> None:
310 - # worker will always be healthy as long as it's running.
311 - return
312 -
313 - def _init_tpu_worker_distributed_environment(
314 - self,
315 - vllm_config: VllmConfig,
316 - rank: int,
317 - distributed_init_method: str | None = None,
318 - local_rank: int = -1,
319 - ) -> None:
320 - """Initialize the distributed environment."""
321 - if self.use_spmd:
322 - xr.use_spmd()
323 - # NOTE(woosuk): This is just to initialize the TP group and broadcast
324 - # the input objects on CPU. The all-reduce and all-gather ops on TPU
325 - # are invoked by `xm.all_reduce` and `xm.all_gather` which use their
326 - # own context.
327 - parallel_config = vllm_config.parallel_config
328 - init_distributed_environment(
329 - world_size=parallel_config.world_size,
330 - rank=rank,
331 - local_rank=local_rank,
332 - distributed_init_method=distributed_init_method or "env://",
333 - backend=current_platform.dist_backend,
334 - )
335 - ensure_model_parallel_initialized(
336 - parallel_config.tensor_parallel_size, parallel_config.pipeline_parallel_size
337 - )
338 -
339 - def shutdown(self) -> None:
340 - self.model_runner.ensure_kv_transfer_shutdown()
341 -
342 - def apply_model(self, fn: Callable[[nn.Module], _R]) -> _R:
343 - """Apply a function on the model inside this worker."""
344 - return fn(self.get_model())
345 -
346 -
14 +# TODO(weiyulin) Remove this file after adding an official way to use hardware plugin
15 if USE_TPU_INFERENCE:
16 from tpu_inference.worker.tpu_worker import TPUWorker as TpuInferenceWorker
17