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 15, 17:15 UTC
[Build] Cython compilation support fix (#14296)
Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com>
Gregory Shtrasberg committed
Mar 24, 2025 at 19:37 UTC
8279201ce6ab178131fedff211a5539dc3ef2710
6 files changed
+46
-6
Dockerfile.rocm
+1
-1
@@ -40,7 +40,7 @@ ARG USE_CYTHON
40
RUN cd vllm \
41
&& python3 -m pip install -r requirements/rocm.txt \
42
&& python3 setup.py clean --all \
43
- && if [ ${USE_CYTHON} -eq "1" ]; then python3 setup_cython.py build_ext --inplace; fi \
43
+ && if [ ${USE_CYTHON} -eq "1" ]; then python3 tests/build_cython.py build_ext --inplace; fi \
44
&& python3 setup.py bdist_wheel --dist-dir=dist
45
FROM scratch AS export_vllm
46
ARG COMMON_WORKDIR
pyproject.toml
+1
@@ -86,6 +86,7 @@ exclude = [
86
"vllm/triton_utils/**/*.py" = ["UP006", "UP035"]
87
"vllm/vllm_flash_attn/**/*.py" = ["UP006", "UP035"]
88
"vllm/worker/**/*.py" = ["UP006", "UP035"]
89
+"vllm/utils.py" = ["UP006", "UP035"]
90
91
[tool.ruff.lint]
92
select = [
tests/build_cython.py
new
+38
@@ -0,0 +1,38 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+import Cython.Compiler.Options
3
+from Cython.Build import cythonize
4
+from setuptools import setup
5
+
6
+Cython.Compiler.Options.annotate = True
7
+
8
+infiles = []
9
+
10
+infiles += [
11
+ "vllm/engine/llm_engine.py",
12
+ "vllm/transformers_utils/detokenizer.py",
13
+ "vllm/engine/output_processor/single_step.py",
14
+ "vllm/outputs.py",
15
+ "vllm/engine/output_processor/stop_checker.py",
16
+]
17
+
18
+infiles += [
19
+ "vllm/core/scheduler.py",
20
+ "vllm/sequence.py",
21
+ "vllm/core/block_manager.py",
22
+]
23
+
24
+infiles += [
25
+ "vllm/model_executor/layers/sampler.py",
26
+ "vllm/sampling_params.py",
27
+ "vllm/utils.py",
28
+]
29
+
30
+setup(ext_modules=cythonize(infiles,
31
+ annotate=False,
32
+ force=True,
33
+ compiler_directives={
34
+ 'language_level': "3",
35
+ 'infer_types': True
36
+ }))
37
+
38
+# example usage: python3 build_cython.py build_ext --inplace
vllm/engine/llm_engine.py
+1
-1
@@ -1249,7 +1249,7 @@ class LLMEngine:
1249
return None
1250
1251
def _advance_to_next_step(
1252
- self, output: List[SamplerOutput],
1252
+ self, output: SamplerOutput,
1253
seq_group_metadata_list: List[SequenceGroupMetadata],
1254
scheduled_seq_groups: List[ScheduledSequenceGroup]) -> None:
1255
"""Given model output from a single run, append the tokens to the
vllm/model_executor/layers/sampler.py
+2
-1
@@ -1187,7 +1187,8 @@ def _build_sampler_output(
1187
deferred_sample_results_args=deferred_sample_results_args)
1188
1189
1190
-def _get_next_prompt_tokens(seq_group: SequenceGroupToSample) -> List[int]:
1190
+def _get_next_prompt_tokens(
1191
+ seq_group: SequenceGroupToSample) -> tuple[int, ...]:
1192
"""Get a list of next prompt tokens to compute logprob from a
1193
given sequence group.
1194
vllm/utils.py
+3
-3
@@ -37,7 +37,7 @@ from collections.abc import (AsyncGenerator, Awaitable, Generator, Hashable,
37
from dataclasses import dataclass, field
38
from functools import cache, lru_cache, partial, wraps
39
from typing import (TYPE_CHECKING, Any, Callable, Generic, Literal, NamedTuple,
40
- Optional, TypeVar, Union)
40
+ Optional, Type, TypeVar, Union)
41
from uuid import uuid4
42
43
import cloudpickle
@@ -1544,9 +1544,9 @@ class LazyDict(Mapping[str, T], Generic[T]):
1544
return len(self._factory)
1545
1546
1547
-class ClassRegistry(UserDict[type[T], _V]):
1547
+class ClassRegistry(UserDict[Type[T], _V]):
1548
1549
- def __getitem__(self, key: type[T]) -> _V:
1549
+ def __getitem__(self, key: Type[T]) -> _V:
1550
for cls in key.mro():
1551
if cls in self.data:
1552
return self.data[cls]