@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 4172235ab7
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 14, 05:30 UTC

[V0 deprecation] Deprecate V0 Neuron backend (#21159)

Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>

Woosuk Kwon committed Sep 6, 2025 at 16:15 UTC 4172235ab78b09989fb56edaf734dbee283dda3e
46 files changed +10 -5462
.buildkite/release-pipeline.yaml
-16
@@ -149,19 +149,3 @@ steps:
149 - "docker push public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:$(buildkite-agent meta-data get release-version)"
150 env:
151 DOCKER_BUILDKIT: "1"
152 -
153 - - block: "Build Neuron release image"
154 - key: block-neuron-release-image-build
155 - depends_on: ~
156 -
157 - - label: "Build and publish Neuron release image"
158 - depends_on: block-neuron-release-image-build
159 - agents:
160 - queue: neuron-postmerge
161 - commands:
162 - - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
163 - - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --tag public.ecr.aws/q9t5s3a7/vllm-neuron-release-repo:$(buildkite-agent meta-data get release-version) --tag public.ecr.aws/q9t5s3a7/vllm-neuron-release-repo:latest --progress plain -f docker/Dockerfile.neuron ."
164 - - "docker push public.ecr.aws/q9t5s3a7/vllm-neuron-release-repo:latest"
165 - - "docker push public.ecr.aws/q9t5s3a7/vllm-neuron-release-repo:$(buildkite-agent meta-data get release-version)"
166 - env:
167 - DOCKER_BUILDKIT: "1"
.buildkite/scripts/hardware_ci/run-neuron-test.sh deleted
-64
@@ -1,64 +0,0 @@
1 -#!/bin/bash
2 -
3 -# This script build the Neuron docker image and run the API server inside the container.
4 -# It serves a sanity check for compilation and basic model usage.
5 -set -e
6 -set -v
7 -
8 -image_name="neuron/vllm-ci"
9 -container_name="neuron_$(tr -dc A-Za-z0-9 < /dev/urandom | head -c 10; echo)"
10 -
11 -HF_CACHE="$(realpath ~)/huggingface"
12 -mkdir -p "${HF_CACHE}"
13 -HF_MOUNT="/root/.cache/huggingface"
14 -HF_TOKEN=$(aws secretsmanager get-secret-value --secret-id "ci/vllm-neuron/hf-token" --region us-west-2 --query 'SecretString' --output text | jq -r .VLLM_NEURON_CI_HF_TOKEN)
15 -
16 -NEURON_COMPILE_CACHE_URL="$(realpath ~)/neuron_compile_cache"
17 -mkdir -p "${NEURON_COMPILE_CACHE_URL}"
18 -NEURON_COMPILE_CACHE_MOUNT="/root/.cache/neuron_compile_cache"
19 -
20 -# Try building the docker image
21 -aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws
22 -
23 -# prune old image and containers to save disk space, and only once a day
24 -# by using a timestamp file in tmp.
25 -if [ -f /tmp/neuron-docker-build-timestamp ]; then
26 - last_build=$(cat /tmp/neuron-docker-build-timestamp)
27 - current_time=$(date +%s)
28 - if [ $((current_time - last_build)) -gt 86400 ]; then
29 - # Remove dangling images (those that are not tagged and not used by any container)
30 - docker image prune -f
31 - # Remove unused volumes / force the system prune for old images as well.
32 - docker volume prune -f && docker system prune -f
33 - echo "$current_time" > /tmp/neuron-docker-build-timestamp
34 - fi
35 -else
36 - date "+%s" > /tmp/neuron-docker-build-timestamp
37 -fi
38 -
39 -docker build -t "${image_name}" -f docker/Dockerfile.neuron .
40 -
41 -# Setup cleanup
42 -remove_docker_container() {
43 - docker image rm -f "${image_name}" || true;
44 -}
45 -trap remove_docker_container EXIT
46 -
47 -# Run the image
48 -docker run --rm -it --device=/dev/neuron0 --network bridge \
49 - -v "${HF_CACHE}:${HF_MOUNT}" \
50 - -e "HF_HOME=${HF_MOUNT}" \
51 - -e "HF_TOKEN=${HF_TOKEN}" \
52 - -v "${NEURON_COMPILE_CACHE_URL}:${NEURON_COMPILE_CACHE_MOUNT}" \
53 - -e "NEURON_COMPILE_CACHE_URL=${NEURON_COMPILE_CACHE_MOUNT}" \
54 - --name "${container_name}" \
55 - ${image_name} \
56 - /bin/bash -c "
57 - set -e; # Exit on first error
58 - python3 /workspace/vllm/examples/offline_inference/neuron.py;
59 - python3 -m pytest /workspace/vllm/tests/neuron/1_core/ -v --capture=tee-sys;
60 - for f in /workspace/vllm/tests/neuron/2_core/*.py; do
61 - echo \"Running test file: \$f\";
62 - python3 -m pytest \$f -v --capture=tee-sys;
63 - done
64 - "
\ No newline at end of file
MANIFEST.in
-1
@@ -2,7 +2,6 @@ include LICENSE
2 include requirements/common.txt
3 include requirements/cuda.txt
4 include requirements/rocm.txt
5 -include requirements/neuron.txt
5 include requirements/cpu.txt
6 include CMakeLists.txt
7
docker/Dockerfile.neuron deleted
-56
@@ -1,56 +0,0 @@
1 -# default base image
2 -# https://gallery.ecr.aws/neuron/pytorch-inference-neuronx
3 -ARG BASE_IMAGE="public.ecr.aws/neuron/pytorch-inference-neuronx:2.6.0-neuronx-py310-sdk2.23.0-ubuntu22.04"
4 -
5 -FROM $BASE_IMAGE
6 -
7 -RUN echo "Base image is $BASE_IMAGE"
8 -
9 -# Install some basic utilities
10 -RUN apt-get update && \
11 - apt-get install -y \
12 - git \
13 - python3 \
14 - python3-pip \
15 - ffmpeg libsm6 libxext6 libgl1
16 -
17 -### Mount Point ###
18 -# When launching the container, mount the code directory to /workspace
19 -ARG APP_MOUNT=/workspace
20 -VOLUME [ ${APP_MOUNT} ]
21 -WORKDIR ${APP_MOUNT}/vllm
22 -
23 -RUN python3 -m pip install --upgrade pip
24 -RUN python3 -m pip install --no-cache-dir fastapi ninja tokenizers pandas tenacity
25 -RUN python3 -m pip install neuronx-cc==2.* --extra-index-url=https://pip.repos.neuron.amazonaws.com -U
26 -RUN python3 -m pip install pytest
27 -
28 -# uninstall transformers-neuronx package explicitly to avoid version conflict
29 -RUN python3 -m pip uninstall -y transformers-neuronx
30 -
31 -COPY . .
32 -ARG GIT_REPO_CHECK=0
33 -RUN --mount=type=bind,source=.git,target=.git \
34 - if [ "$GIT_REPO_CHECK" != 0 ]; then bash tools/check_repo.sh ; fi
35 -
36 -RUN python3 -m pip install -U \
37 - 'cmake>=3.26.1' ninja packaging 'setuptools-scm>=8' wheel jinja2 \
38 - -r requirements/neuron.txt
39 -
40 -ENV VLLM_TARGET_DEVICE neuron
41 -RUN --mount=type=bind,source=.git,target=.git \
42 - pip install --no-build-isolation -v -e .
43 -
44 -# install development dependencies (for testing)
45 -RUN python3 -m pip install -e tests/vllm_test_utils
46 -
47 -# install transformers-neuronx package as an optional dependencies (for V0)
48 -# FIXME: `--no-deps` argument is temporarily added to resolve transformers package version conflict
49 -RUN python3 -m pip install transformers-neuronx==0.13.* --extra-index-url=https://pip.repos.neuron.amazonaws.com -U --no-deps
50 -
51 -RUN python3 -m pip install sentencepiece transformers==4.48.0 -U
52 -
53 -# overwrite entrypoint to run bash script
54 -RUN echo "import subprocess; import sys; subprocess.check_call(sys.argv[1:])" > /usr/local/bin/dockerd-entrypoint.py
55 -
56 -CMD ["/bin/bash"]
examples/offline_inference/neuron.py deleted
-49
@@ -1,49 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -from vllm import LLM, SamplingParams
5 -
6 -# Sample prompts.
7 -prompts = [
8 - "Hello, my name is",
9 - "The president of the United States is",
10 - "The capital of France is",
11 - "The future of AI is",
12 -]
13 -# Create a sampling params object.
14 -sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
15 -
16 -
17 -def main():
18 - # Create an LLM.
19 - llm = LLM(
20 - model="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
21 - max_num_seqs=8,
22 - # The max_model_len and block_size arguments are required to be same as
23 - # max sequence length when targeting neuron device.
24 - # Currently, this is a known limitation in continuous batching support
25 - # in transformers-neuronx.
26 - # TODO(liangfu): Support paged-attention in transformers-neuronx.
27 - max_model_len=1024,
28 - block_size=1024,
29 - # ruff: noqa: E501
30 - # The device can be automatically detected when AWS Neuron SDK is installed.
31 - # The device argument can be either unspecified for automated detection,
32 - # or explicitly assigned.
33 - device="neuron",
34 - tensor_parallel_size=2,
35 - )
36 - # Generate texts from the prompts. The output is a list of RequestOutput objects
37 - # that contain the prompt, generated text, and other information.
38 - outputs = llm.generate(prompts, sampling_params)
39 - # Print the outputs.
40 - print("-" * 50)
41 - for output in outputs:
42 - prompt = output.prompt
43 - generated_text = output.outputs[0].text
44 - print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
45 - print("-" * 50)
46 -
47 -
48 -if __name__ == "__main__":
49 - main()
examples/offline_inference/neuron_eagle.py deleted
-61
@@ -1,61 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""
4 -This example shows how to run offline inference with an EAGLE speculative
5 -decoding model on neuron. To use EAGLE speculative decoding, you must use
6 -a draft model that is specifically fine-tuned for EAGLE speculation.
7 -Additionally, to use EAGLE with NxD Inference, the draft model must include
8 -the LM head weights from the target model. These weights are shared between
9 -the draft and target model.
10 -"""
11 -
12 -from vllm import LLM, SamplingParams
13 -
14 -# Sample prompts.
15 -prompts = [
16 - "What is annapurna labs?",
17 -]
18 -
19 -
20 -def main():
21 - # Create a sampling params object.
22 - sampling_params = SamplingParams(top_k=1, max_tokens=500, ignore_eos=True)
23 -
24 - # Create an LLM.
25 - llm = LLM(
26 - model="/home/ubuntu/model_hf/Meta-Llama-3.1-70B-Instruct",
27 - speculative_config={
28 - "model": "/home/ubuntu/model_hf/Llama-3.1-70B-Instruct-EAGLE-Draft",
29 - "num_speculative_tokens": 5,
30 - "max_model_len": 2048,
31 - },
32 - max_num_seqs=4,
33 - # The max_model_len and block_size arguments are required to be same as
34 - # max sequence length when targeting neuron device.
35 - # Currently, this is a known limitation in continuous batching support
36 - # in neuronx-distributed-inference.
37 - max_model_len=2048,
38 - block_size=2048,
39 - # The device can be automatically detected when AWS Neuron SDK is installed.
40 - # The device argument can be either unspecified for automated detection,
41 - # or explicitly assigned.
42 - device="neuron",
43 - tensor_parallel_size=32,
44 - override_neuron_config={
45 - "enable_eagle_speculation": True,
46 - "enable_fused_speculation": True,
47 - },
48 - )
49 -
50 - # Generate texts from the prompts. The output is a list of RequestOutput objects
51 - # that contain the prompt, generated text, and other information.
52 - outputs = llm.generate(prompts, sampling_params)
53 - # Print the outputs.
54 - for output in outputs:
55 - prompt = output.prompt
56 - generated_text = output.outputs[0].text
57 - print(f"Prompt: {prompt!r}, \n\n\n Generated text: {generated_text!r}")
58 -
59 -
60 -if __name__ == "__main__":
61 - main()
examples/offline_inference/neuron_int8_quantization.py deleted
-63
@@ -1,63 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -import os
5 -
6 -from vllm import LLM, SamplingParams
7 -
8 -# creates XLA hlo graphs for all the context length buckets.
9 -os.environ["NEURON_CONTEXT_LENGTH_BUCKETS"] = "128,512,1024,2048"
10 -# creates XLA hlo graphs for all the token gen buckets.
11 -os.environ["NEURON_TOKEN_GEN_BUCKETS"] = "128,512,1024,2048"
12 -# Quantizes neuron model weight to int8 ,
13 -# The default config for quantization is int8 dtype.
14 -os.environ["NEURON_QUANT_DTYPE"] = "s8"
15 -
16 -# Sample prompts.
17 -prompts = [
18 - "Hello, my name is",
19 - "The president of the United States is",
20 - "The capital of France is",
21 - "The future of AI is",
22 -]
23 -# Create a sampling params object.
24 -sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
25 -
26 -
27 -def main():
28 - # Create an LLM.
29 - llm = LLM(
30 - model="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
31 - max_num_seqs=8,
32 - # The max_model_len and block_size arguments are required to be same as
33 - # max sequence length when targeting neuron device.
34 - # Currently, this is a known limitation in continuous batching support
35 - # in transformers-neuronx.
36 - # TODO(liangfu): Support paged-attention in transformers-neuronx.
37 - max_model_len=2048,
38 - block_size=2048,
39 - # ruff: noqa: E501
40 - # The device can be automatically detected when AWS Neuron SDK is installed.
41 - # The device argument can be either unspecified for automated detection,
42 - # or explicitly assigned.
43 - device="neuron",
44 - quantization="neuron_quant",
45 - override_neuron_config={
46 - "cast_logits_dtype": "bfloat16",
47 - },
48 - tensor_parallel_size=2,
49 - )
50 - # Generate texts from the prompts. The output is a list of RequestOutput objects
51 - # that contain the prompt, generated text, and other information.
52 - outputs = llm.generate(prompts, sampling_params)
53 - # Print the outputs.
54 - print("-" * 50)
55 - for output in outputs:
56 - prompt = output.prompt
57 - generated_text = output.outputs[0].text
58 - print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
59 - print("-" * 50)
60 -
61 -
62 -if __name__ == "__main__":
63 - main()
examples/offline_inference/neuron_multimodal.py deleted
-110
@@ -1,110 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -import requests
4 -import torch
5 -from neuronx_distributed_inference.models.mllama.utils import add_instruct
6 -from PIL import Image
7 -
8 -from vllm import LLM, SamplingParams, TextPrompt
9 -
10 -
11 -def get_image(image_url):
12 - image = Image.open(requests.get(image_url, stream=True).raw)
13 - return image
14 -
15 -
16 -# Model Inputs
17 -PROMPTS = [
18 - "What is in this image? Tell me a story",
19 - "What is the recipe of mayonnaise in two sentences?",
20 - "Describe this image",
21 - "What is the capital of Italy famous for?",
22 -]
23 -IMAGES = [
24 - get_image(
25 - "https://images.pexels.com/photos/1108099/pexels-photo-1108099.jpeg?auto=compress&cs=tinysrgb&dpr=1&w=500"
26 - ),
27 - None,
28 - get_image(
29 - "https://images.pexels.com/photos/1108099/pexels-photo-1108099.jpeg?auto=compress&cs=tinysrgb&dpr=1&w=500"
30 - ),
31 - None,
32 -]
33 -SAMPLING_PARAMS = [
34 - dict(top_k=1, temperature=1.0, top_p=1.0, max_tokens=16)
35 - for _ in range(len(PROMPTS))
36 -]
37 -
38 -
39 -def get_VLLM_mllama_model_inputs(prompt, single_image, sampling_params):
40 - # Prepare all inputs for mllama generation, including:
41 - # 1. put text prompt into instruct chat template
42 - # 2. compose single text and single image prompt into Vllm's prompt class
43 - # 3. prepare sampling parameters
44 - input_image = single_image
45 - has_image = torch.tensor([1])
46 - if isinstance(single_image, torch.Tensor) and single_image.numel() == 0:
47 - has_image = torch.tensor([0])
48 -
49 - instruct_prompt = add_instruct(prompt, has_image)
50 - inputs = TextPrompt(prompt=instruct_prompt)
51 -
52 - if input_image is not None:
53 - inputs["multi_modal_data"] = {"image": input_image}
54 -
55 - sampling_params = SamplingParams(**sampling_params)
56 - return inputs, sampling_params
57 -
58 -
59 -def print_outputs(outputs):
60 - # Print the outputs.
61 - for output in outputs:
62 - prompt = output.prompt
63 - generated_text = output.outputs[0].text
64 - print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
65 -
66 -
67 -def main():
68 - assert (
69 - len(PROMPTS) == len(IMAGES) == len(SAMPLING_PARAMS)
70 - ), f"""Text, image prompts and sampling parameters should have the
71 - same batch size; but got {len(PROMPTS)}, {len(IMAGES)},
72 - and {len(SAMPLING_PARAMS)}"""
73 -
74 - # Create an LLM.
75 - llm = LLM(
76 - model="meta-llama/Llama-3.2-11B-Vision-Instruct",
77 - max_num_seqs=1,
78 - max_model_len=4096,
79 - block_size=4096,
80 - device="neuron",
81 - tensor_parallel_size=32,
82 - override_neuron_config={
83 - "sequence_parallel_enabled": False,
84 - "skip_warmup": True,
85 - "save_sharded_checkpoint": True,
86 - "on_device_sampling_config": {
87 - "global_topk": 1,
88 - "dynamic": False,
89 - "deterministic": False,
90 - },
91 - },
92 - )
93 -
94 - batched_inputs = []
95 - batched_sample_params = []
96 - for pmpt, img, params in zip(PROMPTS, IMAGES, SAMPLING_PARAMS):
97 - inputs, sampling_params = get_VLLM_mllama_model_inputs(pmpt, img, params)
98 - # test batch-size = 1
99 - outputs = llm.generate(inputs, sampling_params)
100 - print_outputs(outputs)
101 - batched_inputs.append(inputs)
102 - batched_sample_params.append(sampling_params)
103 -
104 - # test batch-size = 4
105 - outputs = llm.generate(batched_inputs, batched_sample_params)
106 - print_outputs(outputs)
107 -
108 -
109 -if __name__ == "__main__":
110 - main()
examples/offline_inference/neuron_speculation.py deleted
-64
@@ -1,64 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""
4 -This example shows how to run offline inference with a speculative
5 -decoding model on neuron.
6 -"""
7 -
8 -import os
9 -
10 -from vllm import LLM, SamplingParams
11 -
12 -# Sample prompts.
13 -prompts = [
14 - "Hello, I am a language model and I can help",
15 - "The president of the United States is",
16 - "The capital of France is",
17 -]
18 -
19 -
20 -def config_buckets():
21 - """Configure context length and token gen buckets."""
22 - # creates XLA hlo graphs for all the context length buckets.
23 - os.environ["NEURON_CONTEXT_LENGTH_BUCKETS"] = "128,512,1024,2048"
24 - # creates XLA hlo graphs for all the token gen buckets.
25 - os.environ["NEURON_TOKEN_GEN_BUCKETS"] = "128,512,1024,2048"
26 -
27 -
28 -def initialize_llm():
29 - """Create an LLM with speculative decoding."""
30 - return LLM(
31 - model="openlm-research/open_llama_7b",
32 - speculative_config={
33 - "model": "openlm-research/open_llama_3b",
34 - "num_speculative_tokens": 4,
35 - "max_model_len": 2048,
36 - },
37 - max_num_seqs=4,
38 - max_model_len=2048,
39 - block_size=2048,
40 - device="neuron",
41 - tensor_parallel_size=32,
42 - )
43 -
44 -
45 -def process_requests(llm: LLM, sampling_params: SamplingParams):
46 - """Generate texts from prompts and print them."""
47 - outputs = llm.generate(prompts, sampling_params)
48 - for output in outputs:
49 - prompt = output.prompt
50 - generated_text = output.outputs[0].text
51 - print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
52 -
53 -
54 -def main():
55 - """Main function that sets up the llm and processes prompts."""
56 - config_buckets()
57 - llm = initialize_llm()
58 - # Create a sampling params object.
59 - sampling_params = SamplingParams(max_tokens=100, top_k=1)
60 - process_requests(llm, sampling_params)
61 -
62 -
63 -if __name__ == "__main__":
64 - main()
requirements/neuron.txt deleted
-9
@@ -1,9 +0,0 @@
1 -# Common dependencies
2 --r common.txt
3 -
4 -# Dependencies for Neuron devices
5 -packaging>=24.2
6 -setuptools>=77.0.3,<80.0.0
7 -torch-neuronx >= 2.5.0
8 -neuronx-cc>=2.0.0a0
9 -torchvision # Required for Llama3.2 multimodal image preprocessing
setup.py
+2 -34
@@ -413,8 +413,7 @@ def _no_device() -> bool:
413
414 def _is_cuda() -> bool:
415 has_cuda = torch.version.cuda is not None
416 - return (VLLM_TARGET_DEVICE == "cuda" and has_cuda
417 - and not (_is_neuron() or _is_tpu()))
416 + return (VLLM_TARGET_DEVICE == "cuda" and has_cuda and not _is_tpu())
417
418
419 def _is_hip() -> bool:
@@ -422,10 +421,6 @@ def _is_hip() -> bool:
421 or VLLM_TARGET_DEVICE == "rocm") and torch.version.hip is not None
422
423
425 -def _is_neuron() -> bool:
426 - return VLLM_TARGET_DEVICE == "neuron"
427 -
428 -
424 def _is_tpu() -> bool:
425 return VLLM_TARGET_DEVICE == "tpu"
426
@@ -470,25 +465,6 @@ def get_rocm_version():
465 return None
466
467
473 -def get_neuronxcc_version():
474 - import sysconfig
475 - site_dir = sysconfig.get_paths()["purelib"]
476 - version_file = os.path.join(site_dir, "neuronxcc", "version",
477 - "__init__.py")
478 -
479 - # Check if the command was executed successfully
480 - with open(version_file) as fp:
481 - content = fp.read()
482 -
483 - # Extract the version using a regular expression
484 - match = re.search(r"__version__ = '(\S+)'", content)
485 - if match:
486 - # Return the version string
487 - return match.group(1)
488 - else:
489 - raise RuntimeError("Could not find Neuron version in the output")
490 -
491 -
468 def get_nvcc_cuda_version() -> Version:
469 """Get the CUDA version from nvcc.
470
@@ -541,12 +517,6 @@ def get_vllm_version() -> str:
517 rocm_version = get_rocm_version() or torch.version.hip
518 if rocm_version and rocm_version != MAIN_CUDA_VERSION:
519 version += f"{sep}rocm{rocm_version.replace('.', '')[:3]}"
544 - elif _is_neuron():
545 - # Get the Neuron version
546 - neuron_version = str(get_neuronxcc_version())
547 - if neuron_version != MAIN_CUDA_VERSION:
548 - neuron_version_str = neuron_version.replace(".", "")[:3]
549 - version += f"{sep}neuron{neuron_version_str}"
520 elif _is_tpu():
521 version += f"{sep}tpu"
522 elif _is_cpu():
@@ -591,8 +561,6 @@ def get_requirements() -> list[str]:
561 requirements = modified_requirements
562 elif _is_hip():
563 requirements = _read_requirements("rocm.txt")
594 - elif _is_neuron():
595 - requirements = _read_requirements("neuron.txt")
564 elif _is_tpu():
565 requirements = _read_requirements("tpu.txt")
566 elif _is_cpu():
@@ -601,7 +569,7 @@ def get_requirements() -> list[str]:
569 requirements = _read_requirements("xpu.txt")
570 else:
571 raise ValueError(
604 - "Unsupported platform, please use CUDA, ROCm, Neuron, or CPU.")
572 + "Unsupported platform, please use CUDA, ROCm, or CPU.")
573 return requirements
574
575
tests/engine/test_arg_utils.py
-9
@@ -287,15 +287,6 @@ def test_prefix_cache_default():
287 },
288 "mm-processor-kwargs"
289 ),
290 - (
291 - '{"cast_logits_dtype":"bfloat16","sequence_parallel_norm":true,"sequence_parallel_norm_threshold":2048}',
292 - {
293 - "cast_logits_dtype": "bfloat16",
294 - "sequence_parallel_norm": True,
295 - "sequence_parallel_norm_threshold": 2048,
296 - },
297 - "override-neuron-config"
298 - ),
290 ])
291 # yapf: enable
292 def test_composite_arg_parser(arg, expected, option):
tests/neuron/1_core/test_activation.py deleted
-43
@@ -1,43 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -import pytest
5 -import torch
6 -import torch.nn.functional as F
7 -
8 -from vllm.model_executor.layers.activation import FastGELU, SiluAndMul
9 -from vllm.platforms import current_platform
10 -
11 -
12 -@pytest.mark.parametrize("activation", ["silu_and_mul", "gelu_fast"])
13 -@pytest.mark.parametrize("num_tokens,d,dtype", [
14 - (7, 512, torch.half),
15 - (7, 512, torch.float),
16 - (83, 512, torch.half),
17 -])
18 -@torch.inference_mode()
19 -def test_act_and_mul(
20 - activation: str,
21 - num_tokens: int,
22 - d: int,
23 - dtype: torch.dtype,
24 -) -> None:
25 - import torch_xla.core.xla_model as xm
26 -
27 - device = xm.xla_device()
28 - current_platform.seed_everything(0)
29 - torch.set_default_device("cpu")
30 - x = torch.randn(num_tokens, 2 * d, dtype=dtype).to(device=device)
31 - if activation == "silu_and_mul":
32 - layer = SiluAndMul()
33 - fn = layer.forward_native
34 - elif activation == "gelu_fast":
35 - layer = FastGELU()
36 - fn = F.gelu
37 - else:
38 - raise NotImplementedError(
39 - f"activation {activation} is not implemented.")
40 - assert x.is_xla, "input tensor under testing is expected to be XLA tensor."
41 - out = layer.to(device=device).forward_neuron(x)
42 - ref_out = fn(x.cpu())
43 - torch.testing.assert_close(out.cpu(), ref_out, atol=0.01, rtol=0.0)
tests/neuron/1_core/test_block_table.py deleted
-154
@@ -1,154 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -import neuronxcc.nki.language as nl
5 -import pytest
6 -import torch
7 -import torch.nn.functional as F
8 -from neuronxcc import nki
9 -
10 -from vllm.attention.ops.nki_flash_attn import (
11 - load_block_tables, transform_block_tables_for_indirect_load)
12 -
13 -
14 -def is_power_of_2(n):
15 - return n > 0 and (n & (n - 1) == 0)
16 -
17 -
18 -def nki_load_and_transform_block_tables(
19 - block_tables,
20 - num_tiles,
21 - num_blocks_per_tile,
22 - num_head,
23 - head_id,
24 - block_size_tiling_factor,
25 -):
26 - assert is_power_of_2(
27 - num_blocks_per_tile), f"{num_blocks_per_tile=} must be power of 2"
28 - block_tables_sbuf = load_block_tables(block_tables, num_tiles,
29 - num_blocks_per_tile)
30 -
31 - # we need to pass an Index as head_id
32 - head_id = nl.arange(1)[None, :] + head_id
33 -
34 - block_tables_transposed = transform_block_tables_for_indirect_load(
35 - block_tables_sbuf, block_size_tiling_factor, num_head, head_id)
36 - B_P_SIZE = 128
37 - assert block_tables_transposed.shape[1] == B_P_SIZE
38 -
39 - out = nl.ndarray(
40 - block_tables_transposed.shape,
41 - dtype=nl.int32,
42 - buffer=nl.shared_hbm,
43 - )
44 - for i in nl.affine_range(block_tables_transposed.shape[0]):
45 - nl.store(dst=out[i], value=block_tables_transposed[i])
46 - return out
47 -
48 -
49 -def ref_block_tables_transform(
50 - block_tables,
51 - num_tiles,
52 - num_blocks_per_tile,
53 - num_head,
54 - head_id,
55 - block_size_tiling_factor,
56 -):
57 - assert block_tables.numel() == num_tiles * num_blocks_per_tile
58 - block_tables = block_tables.view(num_tiles, num_blocks_per_tile)
59 - B_F_SIZE = 128
60 - num_tiles_padded = (num_tiles + B_F_SIZE - 1) // B_F_SIZE * B_F_SIZE
61 - block_tables = F.pad(
62 - block_tables,
63 - (0, 0, 0, num_tiles_padded - num_tiles),
64 - "constant",
65 - 0,
66 - )
67 -
68 - block_tables = block_tables * num_head + head_id
69 - block_tables = block_tables.view(num_tiles_padded, num_blocks_per_tile, 1)
70 - offset = torch.arange(0, block_size_tiling_factor).view(1, 1, -1)
71 - block_tables = block_tables * block_size_tiling_factor + offset
72 - block_tables_transposed = block_tables.view(num_tiles_padded, -1).t()
73 -
74 - num_blocks_per_tile = block_tables_transposed.shape[0]
75 - assert num_blocks_per_tile % B_F_SIZE == 0
76 - return block_tables_transposed.view(num_blocks_per_tile // B_F_SIZE,
77 - B_F_SIZE, num_tiles_padded)
78 -
79 -
80 -@pytest.mark.parametrize(
81 - "q_head_per_kv_head,head_id",
82 - [
83 - (1, 0),
84 - (3, 1),
85 - ],
86 -)
87 -@pytest.mark.parametrize(
88 - "num_tiles,num_blocks_per_tile",
89 - [
90 - (1, 1),
91 - (13, 16),
92 - (17, 128),
93 - (35, 512),
94 - (128, 128),
95 - (130, 64),
96 - (280, 256),
97 - (315, 1),
98 - ],
99 -)
100 -@torch.inference_mode()
101 -def test_load_and_transform_block_tables(
102 - monkeypatch: pytest.MonkeyPatch,
103 - num_tiles,
104 - num_blocks_per_tile,
105 - q_head_per_kv_head,
106 - head_id,
107 -) -> None:
108 - import torch_xla.core.xla_model as xm
109 -
110 - device = xm.xla_device()
111 -
112 - compiler_flags_str = " ".join([
113 - "-O1",
114 - "--retry_failed_compilation",
115 - ])
116 - with monkeypatch.context() as m:
117 - m.setenv("NEURON_CC_FLAGS", compiler_flags_str)
118 -
119 - torch.manual_seed(10000)
120 - torch.set_printoptions(sci_mode=False)
121 -
122 - # On Neuron, we need B_P_SIZE = 128 blocks to make DMA efficient
123 - B_P_SIZE = 128
124 - if num_blocks_per_tile < B_P_SIZE:
125 - assert B_P_SIZE % num_blocks_per_tile == 0
126 - block_size_tiling_factor = B_P_SIZE // num_blocks_per_tile
127 - else:
128 - block_size_tiling_factor = 1
129 - max_num_blocks = 100000
130 - block_tables = torch.randint(
131 - 0,
132 - max_num_blocks,
133 - (num_tiles * num_blocks_per_tile, ),
134 - dtype=torch.int32,
135 - )
136 - nki_out = nki.jit(nki_load_and_transform_block_tables)[1, 1](
137 - block_tables.to(device=device),
138 - num_tiles,
139 - num_blocks_per_tile,
140 - q_head_per_kv_head,
141 - head_id,
142 - block_size_tiling_factor,
143 - ).cpu()
144 - ref_out = ref_block_tables_transform(
145 - block_tables,
146 - num_tiles,
147 - num_blocks_per_tile,
148 - q_head_per_kv_head,
149 - head_id,
150 - block_size_tiling_factor,
151 - )
152 - assert (nki_out.shape == ref_out.shape
153 - ), f"{nki_out.shape=} != {ref_out.shape=}"
154 - assert torch.all(nki_out == ref_out)
tests/neuron/1_core/test_cache.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 pytest
5 -import torch
6 -
7 -from vllm.attention.ops.nki_flash_attn import reshape_and_cache
8 -
9 -
10 -@pytest.mark.parametrize(
11 - "num_tokens, n_kv_head, d_head, num_blocks, block_size",
12 - [
13 - # Small model configuration (e.g., GPT-2 small)
14 - (32, 12, 64, 4, 128), # Typical sequence processing
15 - (1, 12, 64, 4, 128), # Single token update
16 - (128, 12, 64, 4, 128), # Longer sequence
17 -
18 - # Medium model configuration (e.g., GPT-2 medium)
19 - (64, 16, 96, 8, 256), # Standard batch
20 - (256, 16, 96, 8, 256), # Large batch
21 -
22 - # Large model configuration (e.g., GPT-3 style)
23 - (48, 32, 128, 16, 512), # Typical processing window
24 - (512, 32, 128, 16, 512), # Full context window
25 -
26 - # Edge cases and stress tests
27 - (1024, 8, 32, 32, 32), # Many tokens, small heads
28 - (16, 64, 256, 4, 64), # Few tokens, many heads
29 - (2048, 24, 128, 64, 128), # Large scale test
30 -
31 - # Minimal configurations for debugging
32 - (4, 2, 16, 2, 16), # Tiny test case
33 - (1, 1, 8, 1, 8), # Minimal possible
34 - ])
35 -def test_reshape_and_cache(num_tokens, n_kv_head, d_head, num_blocks,
36 - block_size):
37 - # Set random seed for reproducibility
38 - torch.manual_seed(42)
39 -
40 - # Create CPU tensors for reference implementation
41 - key_cpu = torch.randn(num_tokens, n_kv_head, d_head) / torch.sqrt(
42 - torch.tensor(d_head))
43 - value_cpu = torch.randn(num_tokens, n_kv_head, d_head) / torch.sqrt(
44 - torch.tensor(d_head))
45 - key_cache_cpu = torch.zeros(num_blocks, n_kv_head, block_size, d_head)
46 - value_cache_cpu = torch.zeros(num_blocks, n_kv_head, block_size, d_head)
47 - slot_mapping_cpu = torch.randperm(num_blocks * block_size)[:num_tokens]
48 -
49 - # Run reference implementation on CPU
50 - block_indices = torch.div(slot_mapping_cpu,
51 - block_size,
52 - rounding_mode="floor")
53 - block_offsets = slot_mapping_cpu % block_size
54 -
55 - for i in range(num_tokens):
56 - block_idx = block_indices[i]
57 - block_offset = block_offsets[i]
58 - key_cache_cpu[block_idx, :, block_offset, :] = key_cpu[i]
59 - value_cache_cpu[block_idx, :, block_offset, :] = value_cpu[i]
60 -
61 - # Create XLA device tensors
62 - device = torch.device('xla')
63 - key = key_cpu.to(device)
64 - value = value_cpu.to(device)
65 - key_cache = torch.zeros_like(key_cache_cpu, device=device)
66 - value_cache = torch.zeros_like(value_cache_cpu, device=device)
67 - slot_mapping = slot_mapping_cpu.to(device)
68 - kv_cache = torch.stack([key_cache, value_cache])
69 -
70 - # Run vectorized implementation on XLA device
71 - reshape_and_cache(key, value, kv_cache, slot_mapping)
72 - key_cache, value_cache = torch.unbind(kv_cache, dim=0)
73 -
74 - # Move results back to CPU for comparison
75 - key_cache_result = key_cache.cpu()
76 - value_cache_result = value_cache.cpu()
77 -
78 - # Assert results match
79 - torch.testing.assert_close(key_cache_result,
80 - key_cache_cpu,
81 - rtol=1e-5,
82 - atol=1e-5)
83 - torch.testing.assert_close(value_cache_result,
84 - value_cache_cpu,
85 - rtol=1e-5,
86 - atol=1e-5)
tests/neuron/1_core/test_layernorm.py deleted
-57
@@ -1,57 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -import pytest
5 -import torch
6 -
7 -from vllm.model_executor.layers.layernorm import RMSNorm
8 -from vllm.platforms import current_platform
9 -
10 -
11 -@pytest.mark.parametrize("num_tokens,hidden_size,add_residual,dtype", [
12 - (7, 8, False, torch.half),
13 - (83, 768, False, torch.half),
14 - (83, 768, True, torch.half),
15 - (83, 768, True, torch.bfloat16),
16 - (83, 768, True, torch.float32),
17 -])
18 -@torch.inference_mode()
19 -def test_rms_norm(
20 - num_tokens: int,
21 - hidden_size: int,
22 - add_residual: bool,
23 - dtype: torch.dtype,
24 -) -> None:
25 - import torch_xla.core.xla_model as xm
26 -
27 - device = xm.xla_device()
28 - current_platform.seed_everything(0)
29 - torch.set_default_device("cpu")
30 - layer = RMSNorm(hidden_size).to(dtype=dtype)
31 - layer.weight.data.normal_(mean=1.0, std=0.1)
32 - scale = 1 / (2 * hidden_size)
33 - x = torch.randn(num_tokens, hidden_size, dtype=dtype).to(device=device)
34 - x *= scale
35 - residual = torch.randn_like(x) * scale if add_residual else None
36 -
37 - residual_cpu = residual.cpu() if add_residual else None
38 - ref_out = layer.to(device="cpu").forward_native(x.cpu(), residual_cpu)
39 - assert x.is_xla, "input tensor under testing is expected to be XLA tensor."
40 - out = layer.to(device=device)(x, residual)
41 -
42 - # NOTE(woosuk): LayerNorm operators (including RMS) typically have larger
43 - # numerical errors than other operators because they involve reductions.
44 - # Therefore, we use a larger tolerance.
45 - if add_residual:
46 - assert out[0].is_xla, "output tensor is expected to be XLA tensor"
47 - torch.testing.assert_close(out[0].cpu(),
48 - ref_out[0],
49 - atol=1e-2,
50 - rtol=1e-2)
51 - torch.testing.assert_close(out[1].cpu(),
52 - ref_out[1],
53 - atol=1e-2,
54 - rtol=1e-2)
55 - else:
56 - assert out.is_xla, "output tensor is expected to be XLA tensor"
57 - torch.testing.assert_close(out.cpu(), ref_out, atol=1e-2, rtol=1e-2)
tests/neuron/1_core/test_logits_processor.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 -import random
5 -from unittest.mock import patch
6 -
7 -import pytest
8 -import torch
9 -
10 -from vllm.model_executor.layers.logits_processor import LogitsProcessor
11 -from vllm.model_executor.sampling_metadata import SamplingMetadata
12 -from vllm.model_executor.utils import set_random_seed
13 -from vllm.sequence import SamplingParams, SequenceData, SequenceGroupMetadata
14 -from vllm.utils import is_pin_memory_available
15 -
16 -
17 -class MockLogitsProcessor(LogitsProcessor):
18 -
19 - def __init__(self, vocab_size: int, scale: float,
20 - fake_logits: torch.Tensor):
21 - super().__init__(vocab_size=vocab_size, scale=scale)
22 - self.fake_logits = fake_logits.clone()
23 -
24 - def forward(self, *args, **kwargs):
25 - with patch(
26 - "vllm.model_executor.layers.logits_processor._prune_hidden_states",
27 - lambda x, y: x
28 - ), patch(
29 - "vllm.model_executor.layers.logits_processor.LogitsProcessor._get_logits",
30 - lambda *args, **kwargs: self.fake_logits):
31 - return super().forward(*args, **kwargs)
32 -
33 -
34 -def _prepare_test(
35 - batch_size: int
36 -) -> tuple[torch.Tensor, torch.Tensor, MockLogitsProcessor]:
37 - vocab_size = 32000
38 - input_tensor = torch.rand((batch_size, 1024), dtype=torch.float16)
39 - fake_logits = torch.full((batch_size, vocab_size),
40 - 1e-2,
41 - dtype=input_tensor.dtype)
42 - logits_processor = MockLogitsProcessor(32000, 0.5, fake_logits)
43 - return input_tensor, fake_logits, logits_processor
44 -
45 -
46 -RANDOM_SEEDS = list(range(8))
47 -
48 -
49 -@pytest.mark.parametrize("seed", RANDOM_SEEDS)
50 -def test_logits_processors(seed: int):
51 - import torch_xla.core.xla_model as xm
52 -
53 - device = xm.xla_device()
54 - set_random_seed(seed)
55 - torch.set_default_device("cpu")
56 - batch_size = random.randint(1, 256)
57 - input_tensor, fake_logits, logits_processor = _prepare_test(batch_size)
58 -
59 - # This sample logits processor gives infinite score to the i-th token,
60 - # where i is the length of the input sequence.
61 - # We therefore expect the output token sequence to be [0, 1, 2, ...]
62 - def pick_ith(token_ids, logits):
63 - logits[len(token_ids)] = float("inf")
64 - return logits
65 -
66 - seq_group_metadata_list = []
67 - seq_lens = []
68 - for i in range(batch_size):
69 - seq_group_metadata_list.append(
70 - SequenceGroupMetadata(
71 - request_id=f"test_{i}",
72 - is_prompt=True,
73 - seq_data={0: SequenceData.from_seqs([1, 2, 3])},
74 - sampling_params=SamplingParams(temperature=0,
75 - logits_processors=[pick_ith]),
76 - block_tables={0: [1]},
77 - ))
78 - seq_lens.append(seq_group_metadata_list[-1].seq_data[0].get_len())
79 -
80 - sampling_metadata = SamplingMetadata.prepare(
81 - seq_group_metadata_list,
82 - seq_lens,
83 - query_lens=seq_lens,
84 - device=device,
85 - pin_memory=is_pin_memory_available())
86 - logits_processor_output = logits_processor(
87 - lm_head=None,
88 - hidden_states=input_tensor,
89 - sampling_metadata=sampling_metadata)
90 -
91 - fake_logits *= logits_processor.scale
92 - torch.testing.assert_close(logits_processor_output[:, 1],
93 - fake_logits[:, 1],
94 - rtol=1e-4,
95 - atol=0.0)
tests/neuron/1_core/test_neuron_model_runner.py deleted
-127
@@ -1,127 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -import os
4 -from unittest.mock import MagicMock
5 -
6 -from vllm.config import VllmConfig
7 -from vllm.engine.arg_utils import EngineArgs
8 -from vllm.platforms import current_platform
9 -from vllm.platforms.neuron import NeuronFramework
10 -from vllm.sampling_params import SamplingParams
11 -from vllm.sequence import SequenceData, SequenceGroupMetadata
12 -from vllm.worker.neuron_model_runner import NeuronModelRunner
13 -
14 -os.environ[
15 - 'VLLM_NEURON_FRAMEWORK'] = NeuronFramework.TRANSFORMERS_NEURONX.value
16 -
17 -
18 -def _create_neuron_model_runner(model: str, *args,
19 - **kwargs) -> NeuronModelRunner:
20 - engine_args = EngineArgs(model, *args, **kwargs)
21 - engine_config = engine_args.create_engine_config()
22 - vllm_config = VllmConfig(
23 - model_config=engine_config.model_config,
24 - parallel_config=engine_config.parallel_config,
25 - scheduler_config=engine_config.scheduler_config,
26 - device_config=engine_config.device_config,
27 - )
28 - neuron_model_runner = NeuronModelRunner(vllm_config=vllm_config)
29 - return neuron_model_runner
30 -
31 -
32 -def test_update_neuron_sampling_params_not_full_batch():
33 - os.environ["NEURON_ON_DEVICE_SAMPLING_DISABLED"] = "0"
34 - model_runner = _create_neuron_model_runner(
35 - "facebook/opt-125m",
36 - seed=0,
37 - dtype="float16",
38 - max_num_seqs=2,
39 - )
40 - assert not model_runner._on_device_sampling_disabled
41 - # Test sampling param updating only when TNx is framework
42 - # NxDI handles sampling parameter updating inside model
43 - if current_platform.use_transformers_neuronx():
44 - model_mock = MagicMock()
45 - model_runner.model = model_mock
46 -
47 - seq_group_metadata_list = [
48 - SequenceGroupMetadata(
49 - request_id="test_0",
50 - is_prompt=True,
51 - seq_data={0: SequenceData.from_seqs([1, 2, 3])},
52 - sampling_params=SamplingParams(temperature=0.5,
53 - top_k=1,
54 - top_p=0.5),
55 - block_tables={0: [1]},
56 - )
57 - ]
58 -
59 - model_runner.prepare_model_input(seq_group_metadata_list)
60 -
61 - # Index neuron sampling parameters based on block_tables indices.
62 - # The first block_id of the sequence 0 is 1, so its parameters are
63 - # placed at index 1. So the sampling parameters will be:
64 - # Index 0: default sampling parameters
65 - # Index 1: sequecne 0's sampling parameters.
66 - neuron_sampling_params = (
67 - model_runner.model_config.neuron_sampling_params)
68 - assert neuron_sampling_params.temperature == [1.0, 0.5]
69 - assert neuron_sampling_params.top_k == [
70 - model_runner._MAX_NEURON_SAMPLING_TOP_K, 1
71 - ]
72 - assert neuron_sampling_params.top_p == [1.0, 0.5]
73 - model_mock.model.update_generation_config.assert_called_once_with(
74 - neuron_sampling_params)
75 -
76 -
77 -def test_update_neuron_sampling_params_full_batch():
78 - os.environ["NEURON_ON_DEVICE_SAMPLING_DISABLED"] = "0"
79 - model_runner = _create_neuron_model_runner(
80 - "facebook/opt-125m",
81 - seed=0,
82 - dtype="float16",
83 - max_num_seqs=2,
84 - )
85 - assert not model_runner._on_device_sampling_disabled
86 -
87 - # Test sampling param updating only when TNx is framework
88 - # NxDI handles sampling parameter updating inside model
89 - if current_platform.use_transformers_neuronx():
90 - model_mock = MagicMock()
91 - model_runner.model = model_mock
92 -
93 - seq_group_metadata_list = [
94 - SequenceGroupMetadata(
95 - request_id="test_0",
96 - is_prompt=True,
97 - seq_data={0: SequenceData.from_seqs([1, 2, 3])},
98 - sampling_params=SamplingParams(temperature=0.5,
99 - top_k=1,
100 - top_p=0.5),
101 - block_tables={0: [1]},
102 - ),
103 - SequenceGroupMetadata(
104 - request_id="test_0",
105 - is_prompt=True,
106 - seq_data={1: SequenceData.from_seqs([4, 5, 6])},
107 - sampling_params=SamplingParams(temperature=0.2,
108 - top_k=2,
109 - top_p=0.2),
110 - block_tables={1: [0]},
111 - )
112 - ]
113 -
114 - model_runner.prepare_model_input(seq_group_metadata_list)
115 -
116 - # Index neuron sampling parameters based on block_tables indices.
117 - # The first block_id of the sequence 0 is 1, so its parameters are
118 - # placed at index 1. So the sampling parameters will be:
119 - # Index 0: sequence 1's sampling parameters
120 - # Index 1: sequecne 0's sampling parameters.
121 - neuron_sampling_params = (
122 - model_runner.model_config.neuron_sampling_params)
123 - assert neuron_sampling_params.temperature == [0.2, 0.5]
124 - assert neuron_sampling_params.top_k == [2, 1]
125 - assert neuron_sampling_params.top_p == [0.2, 0.5]
126 - model_mock.model.update_generation_config.assert_called_once_with(
127 - neuron_sampling_params)
tests/neuron/1_core/test_neuron_quant.py deleted
-12
@@ -1,12 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from vllm.model_executor.layers.quantization.neuron_quant import (
4 - NeuronQuantConfig)
5 -
6 -
7 -def test_get_supported_act_dtypes():
8 - neuron_quant_config = NeuronQuantConfig()
9 - supported_act_dtypes = neuron_quant_config.get_supported_act_dtypes()
10 - target_list = ["any_dtype1", "any_dtype2"]
11 - for dtype in target_list:
12 - assert dtype in supported_act_dtypes
tests/neuron/1_core/test_prefix_prefill.py deleted
-514
@@ -1,514 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -from typing import Optional
5 -
6 -import pytest
7 -import torch
8 -import torch.nn.functional as F
9 -
10 -from vllm.utils import cdiv
11 -
12 -
13 -class BlockDiagonalCausalFromBottomRightMask:
14 -
15 - @staticmethod
16 - def _from_seqlens(query_lens, seq_lens, block_size=None):
17 - from torch import logical_and, logical_or
18 -
19 - contexted = block_size is None
20 - context_lens = torch.tensor(seq_lens) - torch.tensor(query_lens)
21 - n_queries = sum(query_lens)
22 - num_seqs = len(query_lens)
23 - if contexted:
24 - key_lens_blockaligned = seq_lens
25 - else:
26 - n_blocks_per_seq = (context_lens + block_size - 1) // block_size
27 - offset_per_seq = n_blocks_per_seq * block_size
28 - key_lens_blockaligned = offset_per_seq[:num_seqs].tolist()
29 - n_keys = sum(key_lens_blockaligned)
30 -
31 - a = (torch.arange(n_queries).reshape(n_queries,
32 - 1).expand(n_queries, n_keys))
33 - b = torch.arange(n_keys).reshape(1, n_keys).expand(n_queries, n_keys)
34 - q_cumsum = torch.tensor([0] + query_lens).cumsum(dim=0)
35 - k_cumsum = torch.tensor([0] + key_lens_blockaligned).cumsum(dim=0)
36 -
37 - prior_mask = torch.zeros(n_queries, n_keys)
38 - new_masks: list[torch.Tensor] = []
39 - for seq_id in range(num_seqs):
40 - ri = q_cumsum[seq_id]
41 - ci = k_cumsum[seq_id]
42 - nr = query_lens[seq_id]
43 -
44 - if contexted:
45 - nc = seq_lens[seq_id]
46 - a_offset = ci + nc - ri - nr
47 - new_mask = (a + a_offset) >= b
48 - else:
49 - nc = context_lens[seq_id]
50 - a_offset = ci + nc - 1
51 - new_mask = a_offset >= b
52 -
53 - left_mask = b >= ci
54 - top_mask = a >= ri
55 - bottom_mask = a < (ri + nr)
56 -
57 - new_mask = logical_and(
58 - logical_and(logical_and(new_mask, left_mask), top_mask),
59 - bottom_mask,
60 - )
61 - prior_mask = logical_or(prior_mask, new_mask)
62 - new_masks = new_masks + [new_mask]
63 - return prior_mask
64 -
65 - @staticmethod
66 - def from_seqlens(query_lens, seq_lens, block_size=None):
67 - contexted = block_size is None
68 - if contexted:
69 - prior_mask = BlockDiagonalCausalFromBottomRightMask._from_seqlens(
70 - query_lens, seq_lens)
71 - active_mask = None
72 - else:
73 - prior_mask = BlockDiagonalCausalFromBottomRightMask._from_seqlens(
74 - query_lens, seq_lens, block_size)
75 - active_mask = BlockDiagonalCausalFromBottomRightMask._from_seqlens(
76 - query_lens, query_lens)
77 - return prior_mask, active_mask
78 -
79 -
80 -def ref_softmax(x: torch.Tensor,
81 - dim: int,
82 - mixed_precision=False,
83 - return_max_reduce=False):
84 - max_value = torch.amax(x, dim=dim, keepdims=True)
85 - exp = torch.exp(x - max_value)
86 - if mixed_precision:
87 - sum_value = torch.sum(exp.astype(torch.float32),
88 - dim=dim,
89 - keepdims=True).astype(x.dtype)
90 - else:
91 - sum_value = torch.sum(exp, dim=dim, keepdims=True)
92 - if return_max_reduce:
93 - return exp / sum_value, max_value, torch.reciprocal(sum_value)
94 - return exp / sum_value
95 -
96 -
97 -def ref_masked_attention(
98 - query: torch.Tensor,
99 - key: torch.Tensor,
100 - value: torch.Tensor,
101 - scale: float,
102 - attn_mask: Optional[torch.Tensor] = None,
103 - return_max_reduce: Optional[bool] = False,
104 -) -> torch.Tensor:
105 - scaled_qk = scale * torch.einsum("qhd,khd->hqk", query, key).float()
106 - if attn_mask is not None:
107 - masked_score = scaled_qk + attn_mask.float()
108 - if return_max_reduce:
109 - norm_score, cached_max, cached_sum_reciprocal = ref_softmax(
110 - masked_score, dim=-1, return_max_reduce=True)
111 - else:
112 - norm_score = ref_softmax(masked_score, dim=-1)
113 - out = torch.einsum("hqk,khd->qhd", norm_score.to(value.dtype), value)
114 - if return_max_reduce:
115 - return (
116 - out,
117 - cached_max,
118 - cached_sum_reciprocal,
119 - norm_score,
120 - masked_score,
121 - scaled_qk,
122 - )
123 - else:
124 - return (out, )
125 -
126 -
127 -def ref_context_attention(
128 - query,
129 - key,
130 - value,
131 - query_lens,
132 - seq_lens,
133 - head_size,
134 - num_queries_per_kv,
135 - return_max_reduce=False,
136 -):
137 - scale = float(1.0 / (head_size**0.5))
138 - if num_queries_per_kv > 1:
139 - # Handle MQA and GQA
140 - key = torch.repeat_interleave(key, num_queries_per_kv, dim=1)
141 - value = torch.repeat_interleave(value, num_queries_per_kv, dim=1)
142 -
143 - attn_mask, _ = BlockDiagonalCausalFromBottomRightMask.from_seqlens(
144 - query_lens, seq_lens)
145 -
146 - # convert binary mask to -inf values
147 - attn_mask = torch.logical_not(attn_mask)
148 - attn_mask = attn_mask.float() * -30000
149 -
150 - output, *debug_tensors = ref_masked_attention(
151 - query,
152 - key,
153 - value,
154 - scale,
155 - attn_mask,
156 - return_max_reduce=return_max_reduce,
157 - )
158 -
159 - output = output.unsqueeze(1)
160 - if return_max_reduce:
161 - cached_max, cached_sum_reciprocal, lse, masked_score, scaled_qk = (
162 - debug_tensors)
163 - return (
164 - output,
165 - cached_max,
166 - cached_sum_reciprocal,
167 - lse,
168 - masked_score,
169 - scaled_qk,
170 - )
171 - else:
172 - return output
173 -
174 -
175 -def sample_inputs(
176 - prefill_batch_size,
177 - decode_batch_size,
178 - min_query_len,
179 - max_query_len,
180 - min_ctx_len,
181 - max_ctx_len,
182 - block_size,
183 - num_heads,
184 - num_kv_heads,
185 - head_size,
186 - dtype,
187 -):
188 - batch_size = prefill_batch_size + decode_batch_size
189 - max_model_len = (max_query_len + max_ctx_len) * 4
190 - max_block_per_request = max_model_len // block_size
191 - cache_size = (batch_size * max_block_per_request) + 2
192 - prefill_ctx_lens = torch.randint(min_ctx_len,
193 - max_ctx_len + 1, (prefill_batch_size, ),
194 - dtype=torch.long).tolist()
195 - decode_ctx_lens = torch.randint(min_ctx_len,
196 - max_ctx_len + 1, (decode_batch_size, ),
197 - dtype=torch.long).tolist()
198 - ctx_lens = prefill_ctx_lens + decode_ctx_lens
199 - query_lens = torch.randint(
200 - min_query_len,
201 - max_query_len + 1,
202 - (prefill_batch_size, ),
203 - dtype=torch.long,
204 - ).tolist() + [1 for _ in range(decode_batch_size)]
205 - seq_lens = [a + b for a, b in zip(query_lens, ctx_lens)]
206 -
207 - num_tokens = sum(query_lens)
208 - query = torch.empty(num_tokens, num_heads, head_size, dtype=dtype)
209 - query.uniform_(-1, 1)
210 - torch.empty(num_tokens, num_heads, head_size, dtype=dtype)
211 -
212 - kv = torch.empty(sum(seq_lens), 2, num_kv_heads, head_size, dtype=dtype)
213 - kv.uniform_(-1, 1)
214 - key, value = kv.unbind(dim=1)
215 -
216 - k_cache = torch.zeros(cache_size,
217 - block_size,
218 - num_kv_heads,
219 - head_size,
220 - dtype=dtype)
221 - v_cache = torch.zeros(cache_size,
222 - block_size,
223 - num_kv_heads,
224 - head_size,
225 - dtype=dtype)
226 - k = torch.zeros(sum(query_lens), num_kv_heads, head_size, dtype=dtype)
227 - v = torch.zeros(sum(query_lens), num_kv_heads, head_size, dtype=dtype)
228 - values = torch.arange(0, cache_size, dtype=torch.long)
229 - values = values[torch.randperm(cache_size)]
230 - block_table = values[:batch_size * max_block_per_request].view(
231 - batch_size, max_block_per_request)
232 - b_ctx_len = torch.tensor(ctx_lens, dtype=torch.long)
233 - b_start_loc = torch.cumsum(torch.tensor([0] + query_lens[:-1],
234 - dtype=torch.long),
235 - dim=0)
236 - # copy kv to cache
237 - b_seq_start_loc = torch.cumsum(torch.tensor([0] + seq_lens[:-1],
238 - dtype=torch.long),
239 - dim=0)
240 - for i in range(batch_size):
241 - for j in range(query_lens[i]):
242 - k[b_start_loc[i] + j].copy_(key[b_seq_start_loc[i] + b_ctx_len[i] +
243 - j])
244 - v[b_start_loc[i] + j].copy_(value[b_seq_start_loc[i] +
245 - b_ctx_len[i] + j])
246 - cur_ctx = 0
247 - block_id = 0
248 - while cur_ctx < b_ctx_len[i]:
249 - start_loc = b_seq_start_loc[i] + cur_ctx
250 - if cur_ctx + block_size > b_ctx_len[i]:
251 - end_loc = b_seq_start_loc[i] + b_ctx_len[i]
252 - else:
253 - end_loc = start_loc + block_size
254 - start_slot = block_table[i, block_id] * block_size
255 - end_slot = start_slot + end_loc - start_loc
256 - k_cache.view(-1, num_kv_heads,
257 - head_size)[start_slot:end_slot].copy_(
258 - key[start_loc:end_loc])
259 - v_cache.view(-1, num_kv_heads,
260 - head_size)[start_slot:end_slot].copy_(
261 - value[start_loc:end_loc])
262 - cur_ctx += block_size
263 - block_id += 1
264 - kv_cache = torch.stack([k_cache, v_cache])
265 -
266 - return (
267 - query,
268 - k,
269 - v,
270 - kv_cache,
271 - block_table,
272 - key,
273 - value,
274 - query_lens,
275 - seq_lens,
276 - )
277 -
278 -
279 -def get_active_block_tables(block_tables, query_lens, seq_lens, block_size,
280 - num_blocks):
281 - context_lens = seq_lens - query_lens
282 - blocks_per_seq = (context_lens + block_size - 1) // block_size
283 - num_seqs = len(seq_lens)
284 - active_blocks: list[int] = []
285 - for seq_id in range(num_seqs):
286 - active_blocks = (
287 - active_blocks +
288 - block_tables[seq_id, :blocks_per_seq[seq_id]].tolist())
289 - return F.pad(
290 - torch.tensor(active_blocks, dtype=torch.int32),
291 - (0, num_blocks - len(active_blocks)),
292 - "constant",
293 - 0,
294 - )
295 -
296 -
297 -@pytest.mark.parametrize(
298 - "prefill_batch_size,decode_batch_size,block_size,large_tile_size,num_heads,num_queries_per_kv,head_size,mixed_precision",
299 - [
300 - # Test minimal configurations (small block size)
301 - (1, 199, 1, 512, 4, 2, 8, False
302 - ), # minimal block size, small dimensions
303 - (1, 199, 1, 512, 4, 2, 8, True), # same with mixed precision
304 -
305 - # Test common/medium configurations
306 - (4, 12, 32, 2048, 32, 8, 64, False), # common case, larger heads
307 - (4, 12, 32, 2048, 16, 4, 32,
308 - True), # medium size, mixed precision, grouped-query attention (GQA)
309 -
310 - # Test large configurations
311 - (4, 12, 256, 8192, 8, 1, 128, False), # large blocks, large head size
312 - (4, 12, 256, 8192, 64, 8, 64, True), # large blocks, many heads
313 -
314 - # Test asymmetric configurations
315 - (2, 24, 64, 4096, 12, 4, 96, False), # varied batch sizes
316 - (8, 8, 128, 2048, 24, 2, 48, True), # balanced batches
317 -
318 - # Test edge cases
319 - (1, 128, 16, 1024, 4, 2, 16, False), # large decode batch
320 - (16, 4, 8, 1024, 4, 2, 128, True), # large prefill batch
321 - (4, 12, 32, 2048, 16, 1, 32, True), # multi-head attention (MHA)
322 - (4, 12, 32, 2048, 16, 16, 32, True), # multi-query attention (MQA)
323 - ])
324 -@torch.inference_mode()
325 -def test_contexted_kv_attention(
326 - monkeypatch: pytest.MonkeyPatch,
327 - prefill_batch_size: int,
328 - decode_batch_size: int,
329 - num_heads: int,
330 - num_queries_per_kv: int,
331 - head_size: int,
332 - block_size: int,
333 - large_tile_size,
334 - mixed_precision: bool,
335 -) -> None:
336 -
337 - import torch_xla.core.xla_model as xm
338 -
339 - from vllm.attention.ops.nki_flash_attn import (flash_attn_varlen_nkifunc,
340 - reorder_context_mask)
341 -
342 - assert large_tile_size % block_size == 0
343 -
344 - device = xm.xla_device()
345 -
346 - compiler_flags_str = " ".join([
347 - "-O1",
348 - "--retry_failed_compilation",
349 - ])
350 - with monkeypatch.context() as m:
351 - m.setenv("NEURON_CC_FLAGS", compiler_flags_str)
352 -
353 - torch.manual_seed(0)
354 - torch.set_printoptions(sci_mode=False)
355 - torch.set_default_device("cpu")
356 - dtype = torch.float32
357 -
358 - min_ctx_len = 32
359 - max_ctx_len = 1024
360 - min_query_len = 16
361 - max_query_len = 512
362 - num_kv_heads = num_heads // num_queries_per_kv
363 - (
364 - query,
365 - k_active,
366 - v_active,
367 - kv_cache,
368 - block_table,
369 - key,
370 - value,
371 - query_lens,
372 - seq_lens,
373 - ) = sample_inputs(
374 - prefill_batch_size=prefill_batch_size,
375 - decode_batch_size=decode_batch_size,
376 - min_query_len=min_query_len,
377 - max_query_len=max_query_len,
378 - min_ctx_len=min_ctx_len,
379 - max_ctx_len=max_ctx_len,
380 - block_size=block_size,
381 - num_heads=num_heads,
382 - num_kv_heads=num_kv_heads,
383 - head_size=head_size,
384 - dtype=dtype,
385 - )
386 -
387 - output_ref = ref_context_attention(
388 - query,
389 - key,
390 - value,
391 - query_lens,
392 - seq_lens,
393 - head_size,
394 - num_queries_per_kv,
395 - return_max_reduce=False,
396 - )
397 -
398 - # build neuron program
399 - B_P_SIZE = 128
400 - assert (large_tile_size >= B_P_SIZE
401 - ), f"Expect {large_tile_size=} to be larger than {B_P_SIZE=}"
402 -
403 - def pad_to_multiple(a, b):
404 - return cdiv(a, b) * b
405 -
406 - def pad_to_next_power_of_2(a):
407 - assert a > 0
408 - return 2**int(a - 1).bit_length()
409 -
410 - # calculate input shapes
411 - max_num_queries = pad_to_next_power_of_2(sum(query_lens))
412 - context_lens = torch.tensor(seq_lens) - torch.tensor(query_lens)
413 - num_active_blocks = cdiv(context_lens, block_size).sum().item()
414 - num_active_blocks = pad_to_multiple(num_active_blocks,
415 - large_tile_size // block_size)
416 - context_kv_len = num_active_blocks * block_size
417 - assert (
418 - context_kv_len %
419 - large_tile_size == 0), f"invalid context_kv_len={context_kv_len}"
420 -
421 - # pad QKV tensors
422 - pad_dims = (
423 - 0,
424 - 0,
425 - 0,
426 - 0,
427 - 0,
428 - max_num_queries - query.shape[0],
429 - )
430 - query = F.pad(query, pad_dims, "constant", 0)
431 - k = F.pad(k_active, pad_dims, "constant", 0)
432 - v = F.pad(v_active, pad_dims, "constant", 0)
433 -
434 - # permute QKV tensors
435 - # query: (1, n_heads, d, seq_q)
436 - # key: (1, n_kv_heads, d, seq_k)
437 - # value: (1, n_kv_heads, seq_v, d)
438 - query = query.unsqueeze(0).permute(0, 2, 3, 1).contiguous()
439 - k = k.unsqueeze(0).permute(0, 2, 3, 1).contiguous()
440 - v = v.unsqueeze(0).permute(0, 2, 1, 3).contiguous()
441 - kv_cache = kv_cache.permute(0, 1, 3, 2, 4).contiguous()
442 -
443 - # transform block table
444 - active_block_table = get_active_block_tables(
445 - block_table.cpu(),
446 - torch.tensor(query_lens).cpu(),
447 - torch.tensor(seq_lens).cpu(),
448 - block_size,
449 - num_active_blocks,
450 - )
451 -
452 - # Build attention masks
453 - prior_mask, active_mask = (
454 - BlockDiagonalCausalFromBottomRightMask.from_seqlens(
455 - query_lens, seq_lens, block_size=block_size))
456 - prior_mask_padded = F.pad(
457 - prior_mask,
458 - (
459 - 0,
460 - context_kv_len - prior_mask.shape[1],
461 - 0,
462 - max_num_queries - prior_mask.shape[0],
463 - ),
464 - "constant",
465 - 0,
466 - ).bool()
467 - active_mask_padded = F.pad(
468 - active_mask,
469 - (
470 - 0,
471 - max_num_queries - active_mask.shape[1],
472 - 0,
473 - max_num_queries - active_mask.shape[0],
474 - ),
475 - "constant",
476 - 0,
477 - ).bool()
478 - attn_mask = torch.concat([prior_mask_padded, active_mask_padded],
479 - dim=1)
480 -
481 - attn_mask = reorder_context_mask(attn_mask, large_tile_size,
482 - block_size)
483 -
484 - input_args = (
485 - query.to(device=device),
486 - k.to(device=device),
487 - v.to(device=device),
488 - kv_cache.to(device=device),
489 - active_block_table.to(device=device),
490 - attn_mask.to(device=device),
491 - )
492 - input_kwargs = dict(
493 - n_kv_head=num_kv_heads,
494 - head_size=head_size,
495 - mixed_precision=mixed_precision,
496 - LARGE_TILE_SZ=large_tile_size,
497 - )
498 -
499 - output_nki = flash_attn_varlen_nkifunc(*input_args, **input_kwargs)
500 -
501 - num_actual_tokens = sum(query_lens)
502 - # - o: shape (bs, n_heads, seq_q, d) -> (bs, seq_q, n_heads, d)
503 - output_nki = output_nki.cpu().permute(0, 2, 1, 3)
504 - output_nki = output_nki[0, :num_actual_tokens, :, :]
505 - output_ref_padded = F.pad(
506 - output_ref,
507 - (0, 0, 0, 0, 0, 0, 0, max_num_queries - output_ref.shape[0]),
508 - "constant",
509 - 0,
510 - )
511 - output_ref = output_ref_padded.transpose(
512 - 0, 1)[0, :num_actual_tokens, :, :]
513 -
514 - torch.testing.assert_close(output_nki, output_ref, atol=1e-2, rtol=0)
tests/neuron/1_core/test_rotary_embedding.py deleted
-68
@@ -1,68 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""
4 -Tests for miscellaneous utilities
5 -"""
6 -
7 -import pytest
8 -import torch
9 -
10 -from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
11 -from vllm.platforms import current_platform
12 -
13 -
14 -@pytest.mark.parametrize(
15 - "max_position,is_neox_style,rotary_dim,head_size,seq_len,use_key", [
16 - (16, False, 32, 32, 1024, True),
17 - (16, False, 32, 128, 1024, True),
18 - (16, True, 32, 32, 1024, True),
19 - (16, True, 32, 128, 1024, True),
20 - (16, False, 32, 128, 1024, False),
21 - (16, True, 32, 128, 1024, False),
22 - ])
23 -def test_rotary_embedding_opcheck(max_position, is_neox_style, rotary_dim,
24 - head_size, seq_len, use_key):
25 - import torch_xla.core.xla_model as xm
26 -
27 - device = xm.xla_device()
28 - current_platform.seed_everything(0)
29 - torch.set_default_device("cpu")
30 -
31 - batch_size = 1
32 - base = 10000
33 - num_heads = 8
34 -
35 - rot = RotaryEmbedding(head_size, rotary_dim, max_position, base,
36 - is_neox_style, torch.float32)
37 -
38 - positions = torch.randint(0,
39 - max_position, (batch_size, seq_len),
40 - device="cpu")
41 - query = torch.randn(batch_size,
42 - seq_len,
43 - num_heads * head_size,
44 - dtype=torch.float32,
45 - device="cpu")
46 - key = torch.randn_like(query) if use_key else None
47 - assert positions.is_cpu, \
48 - "reference input tensor is expected to be CPU tensor."
49 - ref_query, ref_key = rot.to(device="cpu").forward_native(
50 - positions, query, key)
51 - out_query, out_key = rot.to(device=device).forward_neuron(
52 - positions.to(device=device), query.to(device=device),
53 - key.to(device=device) if key is not None else None)
54 - if use_key:
55 - assert out_query.is_xla and out_key.is_xla, \
56 - "output tensor is expected to be XLA tensor"
57 - torch.testing.assert_close(out_key.cpu(),
58 - ref_key,
59 - atol=1e-2,
60 - rtol=1e-2)
61 - else:
62 - assert out_key is None, "expected returned key to be None"
63 - assert out_query.is_xla, \
64 - "output tensor is expected to be XLA tensor"
65 - torch.testing.assert_close(out_query.cpu(),
66 - ref_query,
67 - atol=1e-2,
68 - rtol=1e-2)
tests/neuron/2_core/test_comm_ops.py deleted
-101
@@ -1,101 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -import functools
4 -from typing import Callable
5 -from unittest.mock import patch
6 -
7 -import pytest
8 -import torch
9 -import torch_xla.distributed.xla_multiprocessing as xmp
10 -from typing_extensions import ParamSpec
11 -
12 -from vllm.distributed.communication_op import (
13 - tensor_model_parallel_all_gather, tensor_model_parallel_all_reduce)
14 -from vllm.distributed.parallel_state import (ensure_model_parallel_initialized,
15 - init_distributed_environment)
16 -from vllm.utils import get_distributed_init_method, get_open_port
17 -
18 -_P = ParamSpec("_P")
19 -
20 -
21 -def reinitialize_neuron_runtime(f: Callable[_P, None]) -> Callable[_P, None]:
22 - """Decorator to reinitialize the Neuron Runtime before executing a test.
23 - This is necessary for distributed tests which need to reallocate Neuron
24 - Cores to separate subprocesses.
25 - """
26 -
27 - @functools.wraps(f)
28 - def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> None:
29 - runtime = torch.classes.neuron.Runtime()
30 - runtime.initialize()
31 - runtime.unsafe_close()
32 -
33 - f(*args, **kwargs)
34 - runtime.initialize()
35 -
36 - return wrapper
37 -
38 -
39 -def all_gather_test_worker(index, tp_degree, distributed_init_method):
40 - init_distributed_environment(tp_degree,
41 - index,
42 - distributed_init_method,
43 - index,
44 - backend="xla")
45 - ensure_model_parallel_initialized(tp_degree, 1)
46 -
47 - num_dimensions = 3
48 - tensor_size = list(range(2, num_dimensions + 2))
49 - total_size = 1
50 - for s in tensor_size:
51 - total_size *= s
52 -
53 - all_gather_dimension = -1
54 - all_tensors = [
55 - torch.arange(total_size, dtype=torch.float32,
56 - device="xla").reshape(tensor_size) * (r + 1)
57 - for r in range(tp_degree)
58 - ]
59 - expected = torch.cat(all_tensors, dim=all_gather_dimension)
60 - t = all_tensors[index % tp_degree]
61 - t = tensor_model_parallel_all_gather(t, all_gather_dimension)
62 - torch.testing.assert_close(t, expected)
63 -
64 -
65 -def all_reduce_test_worker(index, tp_degree, distributed_init_method):
66 - init_distributed_environment(tp_degree,
67 - index,
68 - distributed_init_method,
69 - index,
70 - backend="xla")
71 - ensure_model_parallel_initialized(tp_degree, 1)
72 -
73 - num_elements = 8
74 - all_tensors = [
75 - torch.arange(num_elements, dtype=torch.float32, device="xla") * (r + 1)
76 - for r in range(tp_degree)
77 - ]
78 - expected = torch.sum(torch.stack(all_tensors, dim=0), dim=0)
79 - t = all_tensors[index % tp_degree]
80 - t = tensor_model_parallel_all_reduce(t)
81 - torch.testing.assert_close(t, expected)
82 -
83 -
84 -@pytest.mark.parametrize("tp_size", [2])
85 -@pytest.mark.parametrize("test_target",
86 - [all_reduce_test_worker, all_gather_test_worker])
87 -@reinitialize_neuron_runtime
88 -def test_neuron_multi_process_tensor_parallel(monkeypatch, tp_size,
89 - test_target):
90 -
91 - with patch('torch_xla._XLAC._xla_runtime_is_initialized',
92 - return_value=False):
93 - distributed_init_method = get_distributed_init_method(
94 - "127.0.0.1", get_open_port())
95 -
96 - monkeypatch.setenv("VLLM_USE_V1", "1")
97 - monkeypatch.setenv("NEURONCORE_NUM_DEVICES", str(tp_size))
98 - monkeypatch.setenv("NEURON_PJRT_PROCESSES_NUM_DEVICES",
99 - ','.join(['1' for _ in range(tp_size)]))
100 -
101 - xmp.spawn(test_target, args=(tp_size, distributed_init_method))
tests/neuron/2_core/test_eagle.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 json
5 -import os
6 -import shutil
7 -import tempfile
8 -
9 -import torch
10 -from huggingface_hub import snapshot_download
11 -from safetensors import safe_open
12 -
13 -from vllm import LLM, SamplingParams
14 -
15 -
16 -def patch_eagle_draft_with_lm_head(target_model_id: str,
17 - draft_model_id: str) -> str:
18 - # In NxDI, draft model checkpoint must include lm_head weights from target
19 - # model. For more details see https://awsdocs-neuron.readthedocs-hosted.com
20 - # /en/latest/libraries/nxd-inference/developer_guides/feature-guide.html
21 - # #eagle-checkpoint-compatibility
22 - final_draft_dir = "/tmp/patched_eagle_draft"
23 -
24 - with tempfile.TemporaryDirectory() as tmp_dir:
25 - target_dir = snapshot_download(repo_id=target_model_id,
26 - local_dir=os.path.join(
27 - tmp_dir, "target"))
28 - draft_dir = snapshot_download(repo_id=draft_model_id,
29 - local_dir=os.path.join(tmp_dir, "draft"))
30 -
31 - lm_head_key = "lm_head.weight"
32 - index_path = os.path.join(target_dir, "model.safetensors.index.json")
33 - with open(index_path) as f:
34 - index = json.load(f)
35 - shard_name = index["weight_map"][lm_head_key]
36 - target_safetensor_path = os.path.join(target_dir, shard_name)
37 -
38 - with safe_open(target_safetensor_path, framework="pt") as f:
39 - target_lm_head = f.get_tensor(lm_head_key)
40 -
41 - draft_path = os.path.join(draft_dir, "pytorch_model.bin")
42 - draft_state_dict = torch.load(draft_path, map_location="cpu")
43 - draft_state_dict[lm_head_key] = target_lm_head.to(torch.float16)
44 - torch.save(draft_state_dict, draft_path)
45 -
46 - shutil.copytree(draft_dir, final_draft_dir, dirs_exist_ok=True)
47 -
48 - return final_draft_dir
49 -
50 -
51 -def test_eagle():
52 - patched_draft_path = patch_eagle_draft_with_lm_head(
53 - target_model_id="meta-llama/Llama-2-7b-hf",
54 - draft_model_id="yuhuili/EAGLE-llama2-chat-7B")
55 - llm = LLM(
56 - model="meta-llama/Llama-2-7b-hf",
57 - speculative_config={
58 - "model": patched_draft_path,
59 - "num_speculative_tokens": 5,
60 - "max_model_len": 128
61 - },
62 - max_num_seqs=1,
63 - max_model_len=128,
64 - tensor_parallel_size=2,
65 - override_neuron_config={
66 - "enable_eagle_speculation": True,
67 - "enable_fused_speculation": True,
68 - "fused_qkv": True
69 - },
70 - )
71 - prompts = [
72 - "The president of the United States is",
73 - ]
74 - outputs = llm.generate(prompts, SamplingParams(top_k=1))
75 - expected_output = " the head of state and head of government of " \
76 - "the United States. The president direct"
77 -
78 - for output in outputs:
79 - generated_text = output.outputs[0].text
80 - print(f"Prompt: {output.prompt!r}, Generated text: {generated_text!r}")
81 - assert (expected_output == generated_text)
82 -
83 - print("Neuron Eagle speculation test passed.")
tests/neuron/2_core/test_mistral.py deleted
-64
@@ -1,64 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -from vllm import LLM, SamplingParams
5 -
6 -
7 -def test_mistral():
8 - llm = LLM(model="mistralai/Mistral-7B-v0.1",
9 - tensor_parallel_size=2,
10 - max_num_seqs=4,
11 - max_model_len=128,
12 - override_neuron_config={
13 - "sequence_parallel_enabled": False,
14 - "skip_warmup": True
15 - })
16 -
17 - # Send more prompts than the compiled batch size (4) and request
18 - # varying generation lengths to test accuracy related to Neuron
19 - # specific sequence id sorting.
20 - prompts = [
21 - "The president of the United States is",
22 - "The capital of France is",
23 - "What is Annapurna labs?",
24 - "I believe the meaning of life is",
25 - "Tell me a story about a brave knight",
26 - "Hello, my name is Llama",
27 - ]
28 -
29 - sampling_params = [
30 - SamplingParams(top_k=1, max_tokens=10),
31 - SamplingParams(top_k=1, max_tokens=20),
32 - SamplingParams(top_k=1, max_tokens=30),
33 - SamplingParams(top_k=1, max_tokens=40),
34 - SamplingParams(top_k=1, max_tokens=50),
35 - SamplingParams(top_k=1, max_tokens=60)
36 - ]
37 -
38 - outputs = llm.generate(prompts, sampling_params)
39 -
40 - expected_outputs = [
41 - " the most powerful person in the world. He is",
42 - " a city of many faces. It is a city of history, culture, art, "
43 - "fashion, and",
44 - "\n\nAnnapurna Labs is a semiconductor company that was founded "
45 - "in 2013 by Amazon. The company is",
46 - " to be happy.\n\nI believe that happiness is a choice.\n\nI "
47 - "believe that happiness is a state of mind.\n\nI believe that "
48 - "happiness is a journey.\n\nI believe",
49 - " who rescued a princess from a dragon.\n\nTell me a story about"
50 - " a princess who rescued herself from a dragon.\n\nTell me a "
51 - "story about a princess who rescued herself from a dragon and "
52 - "then rescued a knight from",
53 - " and I am a 10 year old male. I am a very friendly and "
54 - "affectionate boy who loves to be around people. I am a very "
55 - "active boy who loves to play and run around. I am a very smart "
56 - "boy who loves to learn new things. I am a very loyal boy"
57 - ]
58 -
59 - for expected_output, output in zip(expected_outputs, outputs):
60 - generated_text = output.outputs[0].text
61 - print(f"Prompt: {output.prompt!r}, Generated text: {generated_text!r}")
62 - assert (expected_output == generated_text)
63 -
64 - print("Neuron Mistral test passed.")
tests/neuron/2_core/test_multi_lora.py deleted
-97
@@ -1,97 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -from huggingface_hub import snapshot_download
5 -
6 -from vllm import LLM, SamplingParams
7 -from vllm.lora.request import LoRARequest
8 -
9 -
10 -def test_llama_single_lora():
11 - sql_lora_files = snapshot_download(
12 - repo_id="yard1/llama-2-7b-sql-lora-test")
13 - llm = LLM(model="meta-llama/Llama-2-7b-hf",
14 - tensor_parallel_size=2,
15 - max_num_seqs=4,
16 - max_model_len=512,
17 - override_neuron_config={
18 - "sequence_parallel_enabled": False,
19 - "skip_warmup": True,
20 - "lora_modules": [{
21 - "name": "lora_id_1",
22 - "path": sql_lora_files
23 - }]
24 - },
25 - enable_lora=True,
26 - max_loras=1,
27 - max_lora_rank=256,
28 - device="neuron")
29 - """For multi-lora requests using NxDI as the backend, only the lora_name
30 - needs to be specified. The lora_id and lora_path are supplied at the LLM
31 - class/server initialization, after which the paths are handled by NxDI"""
32 - lora_req_1 = LoRARequest("lora_id_1", 0, " ")
33 - prompts = [
34 - "The president of the United States is",
35 - "The capital of France is",
36 - ]
37 - outputs = llm.generate(prompts,
38 - SamplingParams(top_k=1),
39 - lora_request=[lora_req_1, lora_req_1])
40 -
41 - expected_outputs = [
42 - " the head of state and head of government of the United States. "
43 - "The president direct",
44 - " a city of contrasts. The city is home to the Eiffel Tower"
45 - ]
46 -
47 - for expected_output, output in zip(expected_outputs, outputs):
48 - generated_text = output.outputs[0].text
49 - assert (expected_output == generated_text)
50 -
51 -
52 -def test_llama_multiple_lora():
53 - sql_lora_files = snapshot_download(
54 - repo_id="yard1/llama-2-7b-sql-lora-test")
55 - llm = LLM(model="meta-llama/Llama-2-7b-hf",
56 - tensor_parallel_size=2,
57 - max_num_seqs=4,
58 - max_model_len=512,
59 - override_neuron_config={
60 - "sequence_parallel_enabled":
61 - False,
62 - "skip_warmup":
63 - True,
64 - "lora_modules": [{
65 - "name": "lora_id_1",
66 - "path": sql_lora_files
67 - }, {
68 - "name": "lora_id_2",
69 - "path": sql_lora_files
70 - }]
71 - },
72 - enable_lora=True,
73 - max_loras=2,
74 - max_lora_rank=256,
75 - device="neuron")
76 - """For multi-lora requests using NxDI as the backend, only the lora_name
77 - needs to be specified. The lora_id and lora_path are supplied at the LLM
78 - class/server initialization, after which the paths are handled by NxDI"""
79 - lora_req_1 = LoRARequest("lora_id_1", 0, " ")
80 - lora_req_2 = LoRARequest("lora_id_2", 1, " ")
81 - prompts = [
82 - "The president of the United States is",
83 - "The capital of France is",
84 - ]
85 - outputs = llm.generate(prompts,
86 - SamplingParams(top_k=1),
87 - lora_request=[lora_req_1, lora_req_2])
88 -
89 - expected_outputs = [
90 - " the head of state and head of government of the United States. "
91 - "The president direct",
92 - " a city of contrasts. The city is home to the Eiffel Tower"
93 - ]
94 -
95 - for expected_output, output in zip(expected_outputs, outputs):
96 - generated_text = output.outputs[0].text
97 - assert (expected_output == generated_text)
vllm/attention/ops/nki_flash_attn.py deleted
-903
@@ -1,903 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -import neuronxcc.nki.isa as nisa
5 -import neuronxcc.nki.language as nl
6 -import numpy as np
7 -import torch
8 -from neuronxcc import nki
9 -from neuronxcc.nki.language import par_dim
10 -
11 -from vllm.utils import cdiv
12 -
13 -
14 -def is_power_of_2(x):
15 - return x > 0 and (x & (x - 1)) == 0
16 -
17 -
18 -@nki.jit
19 -def load_block_tables(block_tables_hbm, num_tiles, num_blocks_per_tile):
20 - """
21 - Load block tables from HBM into SRAM
22 -
23 - `block_tables_hbm` has shape `(num_tiles * num_blocks_per_tile, )`.
24 - In case `num_tiles > B_P_SIZE`, we need further tile `num_tile` dimension.
25 - """
26 - B_P_SIZE = 128
27 -
28 - # reshape as `(num_tiles, num_blocks_per_tile)`
29 - assert len(block_tables_hbm.shape) == 1
30 - (num_total_blocks, ) = block_tables_hbm.shape
31 - assert num_blocks_per_tile * num_tiles == num_total_blocks
32 - block_tables_hbm = block_tables_hbm.reshape(
33 - (num_tiles, num_blocks_per_tile))
34 -
35 - block_tables_sbuf = nl.zeros(
36 - (cdiv(num_tiles, B_P_SIZE), par_dim(B_P_SIZE), num_blocks_per_tile),
37 - dtype=nl.int32,
38 - )
39 - for i in nl.affine_range(cdiv(num_tiles, B_P_SIZE)):
40 - i_p = nl.arange(B_P_SIZE)[:, None]
41 - i_f = nl.arange(num_blocks_per_tile)[None, :]
42 - block_tables_sbuf[i, i_p, i_f] = nl.load(
43 - block_tables_hbm[i_p + i * B_P_SIZE, i_f],
44 - dtype=nl.int32,
45 - mask=(i_p + i * B_P_SIZE < num_tiles),
46 - )
47 - return block_tables_sbuf
48 -
49 -
50 -@nki.jit
51 -def transform_block_tables_for_indirect_load(
52 - block_tables,
53 - block_size_tiling_factor,
54 - num_head,
55 - head_id,
56 -):
57 - """
58 - This function does two things:
59 - 1. calculate new `block_tables` for a `head_id` after flattening
60 - `num_block`, `num_head`, and `block_size_tiling_factor` dimensions
61 - 2. transpose the result so that `block_table` for each tile is mapped to
62 - SBUF Partition dimension for vectorized DMA
63 -
64 - Tiling trick to further improve DMA performance:
65 - Given KV cache shape `(num_block, num_head, block_size, D)`, when loading M
66 - blocks of a given `head_id` from HBM, the load `cache[block_tables,
67 - head_id]` has shape `(M, block_size, D)`. If M < B_P_SIZE = 128, DMA may not
68 - fully utilize hardware parallelization. The solution is to tile `block_size`
69 - into `(block_size_tiling_factor, tiled_block_size)` s.t. `M *
70 - block_size_tiling_factor = B_P_SIZE`. After tiling, KV cache has shape
71 - `(num_block, num_head, block_size_tiling_factor, tiled_block_size, D)`.
72 -
73 - Note:
74 - We don't further tile D dimension as small DMA size also hurts performance.
75 - """
76 - B_P_SIZE = 128
77 - num_partitions, num_tiles_per_partition, num_blocks_per_tile = (
78 - block_tables.shape)
79 - assert num_tiles_per_partition == B_P_SIZE
80 - assert is_power_of_2(
81 - num_blocks_per_tile), f"{num_blocks_per_tile=} is not power of 2"
82 -
83 - num_loads = cdiv(num_blocks_per_tile, B_P_SIZE)
84 - block_tables_transposed = nl.ndarray(
85 - (
86 - num_loads,
87 - par_dim(B_P_SIZE),
88 - num_partitions * num_tiles_per_partition,
89 - ),
90 - dtype=nl.int32,
91 - )
92 -
93 - # prepare iota ahead of time to avoid repeatedly using Gpsimd
94 - if num_head > 1:
95 - head_id = nisa.iota(head_id, dtype=nl.int32).reshape((1, 1))
96 - head_id = nl.transpose(
97 - head_id.broadcast_to((1, num_tiles_per_partition)))
98 - if num_blocks_per_tile > 1:
99 - head_id = head_id.broadcast_to(
100 - (num_tiles_per_partition, num_blocks_per_tile))
101 -
102 - if block_size_tiling_factor > 1:
103 - broadcast_shape = (
104 - num_tiles_per_partition,
105 - num_blocks_per_tile,
106 - block_size_tiling_factor,
107 - )
108 - offset = nisa.iota(nl.arange(block_size_tiling_factor)[None, None, :],
109 - dtype=nl.int32).broadcast_to(broadcast_shape)
110 -
111 - for partition_id in nl.affine_range(num_partitions):
112 - block_tables_partition = block_tables[partition_id]
113 - if num_head > 1:
114 - # fuse num_block and num_head dimension
115 - block_tables_partition = block_tables_partition * num_head + head_id
116 -
117 - if block_size_tiling_factor > 1:
118 - # need to apply block size tiling trick
119 - assert num_blocks_per_tile * block_size_tiling_factor == B_P_SIZE
120 - block_tables_partition = ((block_tables_partition *
121 - block_size_tiling_factor).reshape(
122 - (num_tiles_per_partition,
123 - num_blocks_per_tile,
124 - 1)).broadcast_to(broadcast_shape))
125 - new_block_tables = block_tables_partition + offset
126 - new_block_tables = new_block_tables.reshape(
127 - (num_tiles_per_partition, B_P_SIZE))
128 - else:
129 - new_block_tables = block_tables_partition
130 -
131 - # transpose the block table so that it can be used by vector DGE
132 - for i in nl.affine_range(num_loads):
133 - i_p = nl.arange(B_P_SIZE)[:, None]
134 - i_f = (partition_id * num_tiles_per_partition +
135 - nl.arange(num_tiles_per_partition)[None, :])
136 - block_tables_transposed[i, i_p, i_f] = nl.transpose(
137 - new_block_tables[:, nl.ds(i * B_P_SIZE, B_P_SIZE)])
138 - return block_tables_transposed
139 -
140 -
141 -@nki.jit
142 -def load_kv_tile_from_cache(
143 - cur_k_tile,
144 - cur_v_tile,
145 - kv_cache,
146 - block_tables,
147 - large_k_tile_idx,
148 - num_blocks_per_large_tile,
149 - tiled_block_size,
150 - B_P_SIZE,
151 - B_D_SIZE,
152 -):
153 - """
154 - Load KV cache and transform Key and Value into layout required by Matmul
155 -
156 - Vectorized DMA Load layout:
157 - Key and Value: (par_dim(B_P_SIZE), seqlen_kv // B_P_SIZE * B_D_SIZE)
158 -
159 - Layout used by attention matmuls:
160 - Key: (par_dim(B_D_SIZE), seqlen_kv)
161 - Value: (seqlen_kv // B_P_SIZE, par_dim(B_P_SIZE), B_D_SIZE)
162 - equivalent to (par_dim(B_P_SIZE), seqlen_kv // B_P_SIZE * B_D_SIZE)
163 - """
164 - # load key cache
165 - num_loads = cdiv(num_blocks_per_large_tile, B_P_SIZE)
166 - for load_idx in nl.affine_range(num_loads):
167 - i_p = nl.arange(B_P_SIZE)[:, None]
168 - i_f = nl.arange(tiled_block_size * B_D_SIZE)[None, :]
169 - loaded = nl.load(kv_cache[0, block_tables[load_idx, i_p,
170 - large_k_tile_idx], i_f])
171 - if cur_k_tile.dtype != loaded.dtype:
172 - loaded = nl.copy(loaded, dtype=cur_k_tile.dtype)
173 - # Transpose SBUF tensor using PE
174 - for tb_i in nl.affine_range(tiled_block_size):
175 - cur_k_tile[
176 - :,
177 - nl.ds(
178 - load_idx * B_P_SIZE * tiled_block_size + tb_i * B_P_SIZE,
179 - B_P_SIZE,
180 - ),
181 - ] = nl.transpose(loaded[:, nl.ds(tb_i * B_D_SIZE, B_D_SIZE)])
182 -
183 - # load value cache
184 - for load_idx in nl.affine_range(num_loads):
185 - loaded = nl.load(kv_cache[1, block_tables[load_idx, i_p,
186 - large_k_tile_idx], i_f])
187 - if cur_v_tile.dtype != loaded.dtype:
188 - loaded = nl.copy(loaded, dtype=cur_v_tile.dtype)
189 - i_p = nl.arange(B_P_SIZE)[:, None]
190 - i_f = nl.arange(tiled_block_size * B_D_SIZE)[None, :]
191 - cur_v_tile[
192 - :,
193 - nl.ds(
194 - load_idx * tiled_block_size * B_D_SIZE,
195 - tiled_block_size * B_D_SIZE,
196 - ),
197 - ] = loaded
198 -
199 -
200 -@nki.jit
201 -def transpose_p_local(p_local_transposed,
202 - p_local,
203 - LARGE_TILE_SZ,
204 - B_F_SIZE=512):
205 - for i in nl.affine_range(LARGE_TILE_SZ // B_F_SIZE):
206 - if nisa.get_nc_version() == nisa.nc_version.gen3:
207 - p_local_t_tmp = nl.ndarray((par_dim(128), B_F_SIZE),
208 - buffer=nl.sbuf,
209 - dtype=p_local.dtype)
210 - else:
211 - p_local_t_tmp = nl.ndarray((par_dim(128), B_F_SIZE),
212 - buffer=nl.psum,
213 - dtype=np.float32)
214 -
215 - for j in nl.affine_range(B_F_SIZE // 128):
216 - j_128_slice = nl.ds(j * 128, 128)
217 - i_j_128_slice = nl.ds(i * B_F_SIZE + j * 128, 128)
218 -
219 - if nisa.get_nc_version() == nisa.nc_version.gen3:
220 - p_local_t_tmp[:, j_128_slice] = nisa.dma_transpose(
221 - p_local[:, i_j_128_slice])
222 - else:
223 - p_local_t_tmp[:, j_128_slice] = nisa.nc_transpose(
224 - p_local[:, i_j_128_slice])
225 -
226 - p_local_transposed[:, nl.ds(i * B_F_SIZE, B_F_SIZE)] = nl.copy(
227 - p_local_t_tmp, dtype=p_local_transposed.dtype)
228 -
229 -
230 -@nki.jit
231 -def _flash_attention_core(
232 - q_local_tile,
233 - k,
234 - v,
235 - o_buffer,
236 - l_buffer,
237 - m_buffer,
238 - kernel_dtype,
239 - acc_type,
240 - tile_mask,
241 - use_causal_mask,
242 - q_tile_idx=None,
243 - initialize=False,
244 - LARGE_TILE_SZ=2048,
245 - B_P_SIZE=128,
246 - B_F_SIZE=512,
247 - B_D_SIZE=128,
248 - qk_res_buffer=None,
249 -):
250 - """
251 - The flash attention core function to calculate self attention between a tile
252 - of q and a block of K and V.
253 - The q_local_tile has (B_P_SIZE, B_D_SIZE)
254 - The K and V have shape (B_D_SIZE, LARGE_TILE_SZ), whose free dimension will
255 - be split into size B_F_SIZE tiles
256 -
257 - The results are stored in the following three buffers
258 - o_buffer: (B_P_SIZE, d)
259 - l_buffer: (B_P_SIZE, 1)
260 - m_buffer: (B_P_SIZE, 1)
261 -
262 - All IO buffers are in SBUF.
263 - """
264 - num_k_tile_per_large_tile = LARGE_TILE_SZ // B_F_SIZE
265 -
266 - qk_res_buf = nl.ndarray((par_dim(B_P_SIZE), LARGE_TILE_SZ),
267 - buffer=nl.sbuf,
268 - dtype=acc_type)
269 - max_local = nl.ndarray((par_dim(B_P_SIZE), num_k_tile_per_large_tile),
270 - dtype=acc_type)
271 - for k_i in nl.affine_range(num_k_tile_per_large_tile):
272 - k_i_b_f_slice = nl.ds(k_i * B_F_SIZE, B_F_SIZE)
273 -
274 - if use_causal_mask:
275 - # mask are used to only apply computation to the lower half of the
276 - # matrix, which reduce the arithmetic intensity by up to 50%
277 - multiplication_required_selection = (q_tile_idx * B_P_SIZE
278 - >= k_i * B_F_SIZE)
279 - else:
280 - multiplication_required_selection = True
281 -
282 - if multiplication_required_selection:
283 - qk_psum = nl.ndarray((par_dim(B_P_SIZE), B_F_SIZE),
284 - dtype=np.float32,
285 - buffer=nl.psum) # (128, 512)
286 - qk_psum[:, :] = nl.matmul(q_local_tile,
287 - k[:, k_i_b_f_slice],
288 - transpose_x=True) # (p(128), 512)
289 - qk_res_buf[:, k_i_b_f_slice] = nl.where(
290 - tile_mask[:, k_i_b_f_slice],
291 - qk_psum[:, nl.ds(0, B_F_SIZE)],
292 - -9984.0,
293 - dtype=acc_type,
294 - )
295 - else:
296 - qk_res_buf[:, k_i_b_f_slice] = -9984.0
297 -
298 - # Calculate max of the current tile
299 - max_local[:, k_i] = nisa.tensor_reduce(
300 - np.max,
301 - qk_res_buf[:, k_i_b_f_slice],
302 - axis=(1, ),
303 - dtype=acc_type,
304 - negate=False,
305 - )
306 -
307 - if qk_res_buffer is not None:
308 - qk_res_buffer[:, :] = nl.copy(qk_res_buf[:, :])
309 -
310 - max_ = nisa.tensor_reduce(
311 - np.max,
312 - max_local[:, :],
313 - axis=(1, ),
314 - dtype=acc_type,
315 - negate=False,
316 - )
317 -
318 - o_previous_scaled = nl.ndarray((par_dim(B_P_SIZE), B_D_SIZE),
319 - dtype=o_buffer.dtype)
320 -
321 - if initialize:
322 - m_buffer[:, 0] = nl.copy(max_)
323 - m_current = max_
324 - else:
325 - m_previous = nl.copy(m_buffer[:, 0])
326 - m_buffer[:, 0] = nl.maximum(m_previous, max_) # (128,1)
327 -
328 - m_current = m_buffer[:, 0]
329 - # Compute scaling factor
330 - alpha = nisa.activation(
331 - np.exp,
332 - m_previous,
333 - bias=-1 * m_current,
334 - scale=1.0,
335 - )
336 - o_previous_scaled[...] = nl.multiply(o_buffer[:, :], alpha)
337 -
338 - p_local = nl.ndarray((par_dim(B_P_SIZE), LARGE_TILE_SZ),
339 - dtype=kernel_dtype)
340 - REDUCTION_TILE = min(2048, LARGE_TILE_SZ // 2)
341 -
342 - p_partial_sum = nl.ndarray(
343 - (par_dim(B_P_SIZE), LARGE_TILE_SZ // REDUCTION_TILE),
344 - dtype=acc_type,
345 - )
346 -
347 - for k_r_i in nl.affine_range(LARGE_TILE_SZ // REDUCTION_TILE):
348 - k_r_i_reduce_slice = nl.ds(k_r_i * REDUCTION_TILE, REDUCTION_TILE)
349 -
350 - # compute exp(qk - max)
351 - # Compute partial row - tile sum of exp(qk - max))
352 - # FIXME : Use activation accumulate to accumulate over k_r_i loop ?
353 - p_local[:, k_r_i_reduce_slice] = nisa.activation_reduce(
354 - np.exp,
355 - qk_res_buf[:, k_r_i_reduce_slice],
356 - bias=-1 * m_current,
357 - scale=1.0,
358 - reduce_op=nl.add,
359 - reduce_res=p_partial_sum[:, k_r_i],
360 - dtype=kernel_dtype,
361 - )
362 -
363 - ps = nl.sum(p_partial_sum, axis=1, dtype=acc_type)
364 -
365 - p_local_transposed = nl.ndarray((par_dim(B_P_SIZE), LARGE_TILE_SZ),
366 - dtype=kernel_dtype)
367 - transpose_p_local(
368 - p_local_transposed=p_local_transposed,
369 - p_local=p_local,
370 - LARGE_TILE_SZ=LARGE_TILE_SZ,
371 - B_F_SIZE=B_F_SIZE,
372 - )
373 -
374 - pv_psum = nl.zeros(
375 - (par_dim(B_P_SIZE), B_D_SIZE),
376 - dtype=np.float32,
377 - buffer=nl.psum,
378 - )
379 - for k_i in nl.affine_range(LARGE_TILE_SZ // B_P_SIZE):
380 - pv_psum[:, :] += nl.matmul(
381 - p_local_transposed[:, nl.ds(k_i * B_P_SIZE, B_P_SIZE)],
382 - v[:, nl.ds(k_i * B_D_SIZE, B_D_SIZE)],
383 - transpose_x=True,
384 - ) # (128, 128) (p(Br), d)
385 -
386 - if initialize:
387 - o_buffer[:, :] = nl.copy(pv_psum[:, :])
388 - l_buffer[:, 0] = nl.add(nl.log(ps), max_)
389 - else:
390 - o_buffer[:, :] = nl.add(o_previous_scaled, pv_psum)
391 -
392 - l_prev = l_buffer[:, 0]
393 - l_exp = nl.add(
394 - nl.exp(nl.subtract(l_prev, m_current)),
395 - ps,
396 - )
397 - l_buffer[:, 0] = nl.add(m_current, nl.log(l_exp))
398 -
399 -
400 -@nki.jit
401 -def load_v_tile(v_hbm_tile, cur_v_tile, large_tile_idx, v_i, LARGE_TILE_SZ):
402 - B_P_SIZE = 128
403 - B_D_SIZE = v_hbm_tile.shape[-1]
404 - loaded = nl.load(v_hbm_tile[
405 - nl.ds(large_tile_idx * LARGE_TILE_SZ + B_P_SIZE * v_i, B_P_SIZE),
406 - :,
407 - ])
408 - if cur_v_tile.dtype != loaded.dtype:
409 - loaded = nl.copy(loaded, dtype=cur_v_tile.dtype)
410 - cur_v_tile[:, nl.ds(v_i * B_D_SIZE, B_D_SIZE)] = loaded
411 -
412 -
413 -@nki.jit
414 -def flash_paged_attention(
415 - query,
416 - key,
417 - value,
418 - kv_cache,
419 - block_tables,
420 - mask,
421 - softmax_scale=None,
422 - mixed_precision=True,
423 - LARGE_TILE_SZ=2048,
424 - return_debug_tensors=False,
425 -):
426 - """
427 - Flash PagedAttention Forward Kernel.
428 -
429 - IO tensor layouts:
430 - - query: shape (1, n_heads, d, seq_q)
431 - - key: shape (1, n_kv_heads, d, seq_k)
432 - - value: shape (1, n_kv_heads, seq_v, d)
433 - - kv_cache: (2, num_blocks, n_kv_heads, block_size, d)
434 - - block_tables: (num_active_blocks, )
435 - - mask: (seq_q, num_active_blocks * block_size + seq_q)
436 - - o: shape (1, n_heads, seq_q, d)
437 -
438 - - This kernel requires seq_k == seq_v
439 - - We use continuous batching by default, so the batch dimension is
440 - always 1, and different requests are concatenated along sequence
441 - dimension.
442 - - We use paged cache blocks (kv_cache) to store KV cache.
443 -
444 - IO tensor dtypes:
445 - - This kernel assumes all IO tensors have the same dtype except for
446 - block_tables (int32) and mask (int32)
447 - - If mixed_precision is True, then all Tensor Engine operation will be
448 - performed in bfloat16 and accumulation will be performed in float32.
449 - Otherwise the intermediates will be in the same type as the inputs.
450 -
451 - Compile-time Constants:
452 - - softmax_scale: scaling for softmax, is None, default is `1.0/(d**0.5)`
453 - - mixed_precision: flag to set non-matmul ops in fp32 precision, default
454 - is set to `true`, if false, we use same precision as input types
455 - - LARGE_TILE_SZ: `default=2048`, size of the kv tile size for attention
456 - computation reduction
457 -
458 - GQA support Notes:
459 - the spmd kernel for launching kernel should be on kv_heads instead of
460 - nheads
461 -
462 - Example usage:
463 - MHA: q: [b, h, d, s], k: [b, h, d, s], v: [b, h, s, d]
464 - usage: `flash_fwd[b, h](q, k, v, ...)`
465 - GQA: q: [b, h, d, s], k: [b, kv_h, d, s], v: [b, kv_h, s, d]
466 - usage: `flash_fwd[b, kv_h](q, k, v, ...)`
467 - """
468 - B_F_SIZE = 512
469 - B_P_SIZE = 128
470 - b, h, d, seqlen_q = query.shape
471 - B_D_SIZE = d
472 - n_tile_q = seqlen_q // B_P_SIZE # since q will be loaded on tensor engine
473 - _, num_blocks, k_h, block_size, _ = kv_cache.shape
474 - q_h_per_k_h = h // k_h
475 - assert b == 1, f"invalid batch size {b=}"
476 - assert d <= 128, f" we do not support head_dim > 128, got head dim {d=}"
477 - cache_shape = (2, num_blocks, k_h, block_size, d)
478 - assert (tuple(kv_cache.shape) == cache_shape
479 - ), f"{kv_cache.shape=} mismatch, expect {cache_shape}"
480 - assert key is None or tuple(key.shape) == (
481 - 1,
482 - k_h,
483 - d,
484 - seqlen_q,
485 - ), f"key shape {key.shape} mismatch!"
486 - assert value is None or tuple(value.shape) == (
487 - 1,
488 - k_h,
489 - seqlen_q,
490 - d,
491 - ), f"value shape {value.shape} mismatch!"
492 -
493 - assert (
494 - nl.program_ndim() == 2
495 - ), f"Expect spmd grid with 2 dimensions, got {nl.program_ndim()} instead!"
496 - batch_id = nl.program_id(axis=0)
497 - head_id = nl.program_id(axis=1)
498 -
499 - (num_active_blocks, ) = block_tables.shape
500 - context_kv_len = num_active_blocks * block_size
501 - assert (
502 - LARGE_TILE_SZ % B_F_SIZE == 0
503 - ), f"Need {LARGE_TILE_SZ=} to be divisible by {B_F_SIZE=} in transpose_p"
504 - assert (context_kv_len % LARGE_TILE_SZ == 0
505 - ), f"Need {context_kv_len=} to be divisible by {LARGE_TILE_SZ=}"
506 -
507 - num_blocks_per_large_tile = LARGE_TILE_SZ // block_size
508 - assert is_power_of_2(
509 - num_blocks_per_large_tile
510 - ), f"{num_blocks_per_large_tile=} is expected of be power of 2"
511 - if seqlen_q > B_F_SIZE:
512 - MAX_REDUCTION_TILE = 2048
513 - if seqlen_q // 2 > MAX_REDUCTION_TILE:
514 - assert (
515 - seqlen_q % MAX_REDUCTION_TILE == 0
516 - ), f"{seqlen_q=} should be divisible by {MAX_REDUCTION_TILE=}"
517 - else:
518 - assert (seqlen_q % B_F_SIZE == 0
519 - ), f"{seqlen_q=} should be divisible by {B_F_SIZE=})"
520 -
521 - kernel_dtype = nl.bfloat16 if mixed_precision else query.dtype
522 - acc_type = np.dtype(np.float32) if mixed_precision else kernel_dtype
523 - softmax_scale = softmax_scale or (1.0 / (d**0.5))
524 - num_large_k_tile = context_kv_len // LARGE_TILE_SZ
525 -
526 - o = nl.ndarray((b, h, seqlen_q, d),
527 - dtype=query.dtype,
528 - buffer=nl.shared_hbm)
529 - hbm_l_buffer, hbm_m_buffer, hbm_qk_res, qk_res_buffer = (
530 - None,
531 - None,
532 - None,
533 - None,
534 - )
535 - if return_debug_tensors:
536 - hbm_l_buffer = nl.ndarray((b, h, seqlen_q),
537 - dtype=acc_type,
538 - buffer=nl.shared_hbm)
539 - hbm_m_buffer = nl.ndarray((b, h, seqlen_q),
540 - dtype=acc_type,
541 - buffer=nl.shared_hbm)
542 - hbm_qk_res = nl.ndarray((b, h, B_P_SIZE, seqlen_q),
543 - dtype=acc_type,
544 - buffer=nl.shared_hbm)
545 - qk_res_buffer = nl.zeros(
546 - (n_tile_q, q_h_per_k_h, par_dim(B_P_SIZE), seqlen_q),
547 - dtype=acc_type,
548 - buffer=nl.sbuf,
549 - lazy_initialization=True,
550 - )
551 - block_tables_sbuf = load_block_tables(
552 - block_tables_hbm=block_tables,
553 - num_tiles=num_large_k_tile,
554 - num_blocks_per_tile=num_blocks_per_large_tile,
555 - )
556 -
557 - # On Neuron, we need B_P_SIZE = 128 blocks to make DMA efficient
558 - if num_blocks_per_large_tile < B_P_SIZE:
559 - # we checked num_blocks_per_tile is a power of 2
560 - assert B_P_SIZE % num_blocks_per_large_tile == 0
561 - block_size_tiling_factor = B_P_SIZE // num_blocks_per_large_tile
562 - # We assume block_size >= block_size_tiling_factor
563 - assert block_size % block_size_tiling_factor == 0
564 - else:
565 - block_size_tiling_factor = 1
566 - tiled_block_size = block_size // block_size_tiling_factor
567 -
568 - # Indirect DMA load must be placed along Partition Dimension
569 - block_tables_sbuf = transform_block_tables_for_indirect_load(
570 - block_tables_sbuf,
571 - block_size_tiling_factor=block_size_tiling_factor,
572 - num_head=k_h,
573 - head_id=head_id,
574 - )
575 -
576 - # Flatten KV cache to be 3D for loading into SBUF
577 - new_cache_shape = (
578 - 2,
579 - num_blocks * k_h * block_size_tiling_factor,
580 - tiled_block_size * d,
581 - )
582 - kv_cache = kv_cache.reshape(new_cache_shape)
583 -
584 - # Global Flash Attention accumulators
585 - o_buffer = nl.zeros(
586 - (n_tile_q, q_h_per_k_h, par_dim(B_P_SIZE), d),
587 - dtype=acc_type,
588 - buffer=nl.sbuf,
589 - lazy_initialization=True,
590 - )
591 - l_buffer = nl.zeros(
592 - (n_tile_q, q_h_per_k_h, par_dim(B_P_SIZE), 1),
593 - dtype=acc_type,
594 - buffer=nl.sbuf,
595 - lazy_initialization=True,
596 - )
597 - m_buffer = nl.zeros(
598 - (n_tile_q, q_h_per_k_h, par_dim(B_P_SIZE), 1),
599 - dtype=acc_type,
600 - buffer=nl.sbuf,
601 - lazy_initialization=True,
602 - )
603 -
604 - for large_k_tile_idx in nl.sequential_range(0, num_large_k_tile):
605 - num_loads = cdiv(num_blocks_per_large_tile, B_P_SIZE)
606 - cur_k_tile = nl.ndarray(
607 - (par_dim(B_D_SIZE), LARGE_TILE_SZ),
608 - dtype=kernel_dtype,
609 - )
610 - cur_v_tile = nl.ndarray(
611 - (par_dim(B_P_SIZE), num_loads * tiled_block_size * B_D_SIZE),
612 - dtype=kernel_dtype,
613 - )
614 - load_kv_tile_from_cache(
615 - cur_k_tile=cur_k_tile,
616 - cur_v_tile=cur_v_tile,
617 - kv_cache=kv_cache,
618 - block_tables=block_tables_sbuf,
619 - large_k_tile_idx=large_k_tile_idx,
620 - num_blocks_per_large_tile=num_blocks_per_large_tile,
621 - tiled_block_size=tiled_block_size,
622 - B_P_SIZE=B_P_SIZE,
623 - B_D_SIZE=B_D_SIZE,
624 - )
625 -
626 - for i in nl.affine_range(n_tile_q):
627 - cur_mask = nl.load(mask[
628 - nl.ds(i * B_P_SIZE, B_P_SIZE),
629 - nl.ds(large_k_tile_idx * LARGE_TILE_SZ, LARGE_TILE_SZ),
630 - ])
631 - for i_q_h in nl.affine_range(q_h_per_k_h):
632 - q_tile = nl.ndarray((B_D_SIZE, B_P_SIZE), dtype=kernel_dtype)
633 - q_hbm_tile = query[batch_id, head_id * q_h_per_k_h + i_q_h]
634 - q_sbuf_tile = nl.load(q_hbm_tile[:,
635 - nl.ds(i *
636 - B_P_SIZE, B_P_SIZE)])
637 - if q_sbuf_tile.dtype != kernel_dtype:
638 - q_sbuf_tile = nl.copy(q_sbuf_tile, dtype=kernel_dtype)
639 - q_tile[:, :] = q_sbuf_tile * softmax_scale
640 -
641 - _flash_attention_core(
642 - q_local_tile=q_tile,
643 - k=cur_k_tile,
644 - v=cur_v_tile,
645 - o_buffer=o_buffer[i, i_q_h],
646 - l_buffer=l_buffer[i, i_q_h],
647 - m_buffer=m_buffer[i, i_q_h],
648 - kernel_dtype=kernel_dtype,
649 - acc_type=acc_type,
650 - tile_mask=cur_mask,
651 - use_causal_mask=False,
652 - q_tile_idx=i,
653 - initialize=large_k_tile_idx == 0,
654 - LARGE_TILE_SZ=LARGE_TILE_SZ,
655 - B_P_SIZE=B_P_SIZE,
656 - B_F_SIZE=B_F_SIZE,
657 - B_D_SIZE=B_D_SIZE,
658 - )
659 -
660 - # compute attention between input query, key and value
661 - if key is not None and value is not None:
662 - B_F_SIZE = min(seqlen_q, B_F_SIZE)
663 - LARGE_TILE_SZ = seqlen_q
664 -
665 - cur_k_tile = nl.ndarray((par_dim(B_D_SIZE), LARGE_TILE_SZ),
666 - dtype=kernel_dtype)
667 - cur_v_tile = nl.ndarray(
668 - (par_dim(B_P_SIZE), LARGE_TILE_SZ // B_P_SIZE * B_D_SIZE),
669 - dtype=kernel_dtype,
670 - )
671 -
672 - loaded = nl.load(key[batch_id, head_id, :, :])
673 - if loaded.dtype != kernel_dtype:
674 - loaded = nl.copy(loaded, dtype=kernel_dtype)
675 - cur_k_tile[:, :] = loaded
676 -
677 - v_hbm_tile = value[batch_id, head_id]
678 - for v_i in nl.affine_range(LARGE_TILE_SZ // B_P_SIZE):
679 - load_v_tile(
680 - v_hbm_tile=v_hbm_tile,
681 - cur_v_tile=cur_v_tile,
682 - large_tile_idx=0,
683 - v_i=v_i,
684 - LARGE_TILE_SZ=LARGE_TILE_SZ,
685 - )
686 -
687 - for i in nl.affine_range(n_tile_q):
688 - cur_mask = nl.load(mask[
689 - nl.ds(i * B_P_SIZE, B_P_SIZE),
690 - nl.ds(context_kv_len, LARGE_TILE_SZ),
691 - ])
692 - for i_q_h in nl.affine_range(q_h_per_k_h):
693 -
694 - q_tile = nl.ndarray((B_D_SIZE, B_P_SIZE), dtype=kernel_dtype)
695 - q_hbm_tile = query[batch_id, head_id * q_h_per_k_h + i_q_h]
696 - q_sbuf_tile = nl.load(q_hbm_tile[:,
697 - nl.ds(i *
698 - B_P_SIZE, B_P_SIZE)])
699 - if q_sbuf_tile.dtype != kernel_dtype:
700 - q_sbuf_tile = nl.copy(q_sbuf_tile, dtype=kernel_dtype)
701 - q_tile[:, :] = q_sbuf_tile * softmax_scale
702 - _flash_attention_core(
703 - q_local_tile=q_tile,
704 - k=cur_k_tile,
705 - v=cur_v_tile,
706 - o_buffer=o_buffer[i, i_q_h],
707 - l_buffer=l_buffer[i, i_q_h],
708 - m_buffer=m_buffer[i, i_q_h],
709 - kernel_dtype=kernel_dtype,
710 - acc_type=acc_type,
711 - tile_mask=cur_mask,
712 - use_causal_mask=True,
713 - q_tile_idx=i,
714 - initialize=False,
715 - LARGE_TILE_SZ=LARGE_TILE_SZ,
716 - B_P_SIZE=B_P_SIZE,
717 - B_F_SIZE=B_F_SIZE,
718 - B_D_SIZE=B_D_SIZE,
719 - qk_res_buffer=(qk_res_buffer[i, i_q_h]
720 - if qk_res_buffer is not None else None),
721 - )
722 -
723 - # -- -- -- -- write output to buffer on HBM -- -- -- -- -- -- #
724 - for i_q_h in nl.affine_range(q_h_per_k_h):
725 - for i in nl.affine_range(n_tile_q):
726 - out = nl.multiply(
727 - o_buffer[i, i_q_h],
728 - nl.exp(m_buffer[i, i_q_h] - l_buffer[i, i_q_h]),
729 - dtype=kernel_dtype,
730 - )
731 -
732 - nl.store(
733 - o[
734 - batch_id,
735 - head_id * q_h_per_k_h + i_q_h,
736 - nl.ds(i * B_P_SIZE, B_P_SIZE),
737 - :,
738 - ],
739 - out,
740 - )
741 - # maximum and summation statistics
742 - if return_debug_tensors:
743 - nl.store(
744 - hbm_m_buffer[
745 - batch_id,
746 - head_id * q_h_per_k_h + i_q_h,
747 - nl.ds(i * B_P_SIZE, B_P_SIZE),
748 - ],
749 - m_buffer[i, i_q_h, :, :],
750 - )
751 - nl.store(
752 - hbm_l_buffer[
753 - batch_id,
754 - head_id * q_h_per_k_h + i_q_h,
755 - nl.ds(i * B_P_SIZE, B_P_SIZE),
756 - ],
757 - l_buffer[i, i_q_h],
758 - )
759 - nl.store(
760 - hbm_qk_res[batch_id, head_id * q_h_per_k_h + i_q_h, :, :],
761 - qk_res_buffer[batch_id, i_q_h, :, :],
762 - )
763 -
764 - if return_debug_tensors:
765 - return o, hbm_m_buffer, hbm_l_buffer, hbm_qk_res
766 - return o
767 -
768 -
769 -def reorder_context_mask(mask, LARGE_TILE_SZ, block_size):
770 - """
771 - Reorder the mask to make it compatible with the flash attention kernel.
772 -
773 - We vectorize KV cache read to improve DMA utilization. However, the layout
774 - that maximizes DMA bandwidth changes the order tokens are consumed.
775 -
776 - The token layout (inner 2 dimensions) after vectorized load is (B_P_SIZE,
777 - tiled_block_size) in a tile of `B_P_SIZE * tiled_block_size` tokens. And
778 - each step the engine consumes a column (rather than a row) of B_P_SIZE
779 - tokens. Therefore, the tokens are visited in a strided way.
780 -
781 - To make sure mask matches the order tokens are consumed, we need to properly
782 - transpose mask.
783 - """
784 - total_query_len, total_seq_len = mask.shape
785 - context_kv_len = total_seq_len - total_query_len
786 -
787 - B_P_SIZE = 128
788 - assert (LARGE_TILE_SZ
789 - >= B_P_SIZE), f"{LARGE_TILE_SZ=} must be larger than {B_P_SIZE=}"
790 - num_tiled_blocks = max(B_P_SIZE, LARGE_TILE_SZ // block_size)
791 - tiled_block_size = LARGE_TILE_SZ // num_tiled_blocks
792 - if tiled_block_size > 1:
793 - # Mask reordering is needed when tiled_block_size > 1
794 - device = mask.device
795 - mask = mask.cpu()
796 - context_mask = mask[:, :context_kv_len]
797 - context_mask = context_mask.view(
798 - total_query_len,
799 - context_kv_len // LARGE_TILE_SZ,
800 - num_tiled_blocks // B_P_SIZE,
801 - B_P_SIZE,
802 - tiled_block_size,
803 - )
804 - context_mask = context_mask.transpose(3, 4).reshape(
805 - total_query_len, context_kv_len)
806 - new_mask = mask[:, context_kv_len:]
807 - return torch.concat([context_mask, new_mask], dim=1).to(device)
808 - else:
809 - return mask
810 -
811 -
812 -def flash_attn_varlen_nkifunc(
813 - query,
814 - key,
815 - value,
816 - kv_cache,
817 - block_table,
818 - attn_mask,
819 - n_kv_head=None,
820 - head_size=None,
821 - LARGE_TILE_SZ=2048,
822 - mixed_precision=True,
823 -):
824 - """
825 - Compute flash paged attention for variable length sequences.
826 -
827 - This function is a wrapper around the flash attention NKI kernel. It takes
828 - in the following arguments:
829 - - query: (1, n_heads, d, seq_q)
830 - - key: (1, n_kv_heads, d, seq_k)
831 - - value: (1, n_kv_heads, seq_v, d)
832 - - kv_cache: (2, n_blocks, n_kv_heads, block_size, d)
833 - - block_tables: (n_active_blocks, )
834 - - attn_mask: (seq_q, n_active_blocks * block_size + seq_q)
835 -
836 - Notes:
837 - - attn_mask must be reordered outside using `reorder_context_mask`
838 - - Key/value cache layout must be (n_blocks, n_kv_heads, block_size, d)
839 - for better DMA throughput
840 - """
841 - if n_kv_head is None:
842 - n_kv_head = kv_cache.shape[2]
843 - assert kv_cache.shape[0] == 2
844 - assert kv_cache.shape[2] == n_kv_head
845 - if head_size is None:
846 - head_size = kv_cache.shape[-1]
847 -
848 - kwargs = dict(
849 - query=query,
850 - key=key,
851 - value=value,
852 - kv_cache=kv_cache,
853 - block_tables=block_table,
854 - mask=attn_mask,
855 - softmax_scale=1.0 / (head_size**0.5),
856 - mixed_precision=mixed_precision,
857 - LARGE_TILE_SZ=LARGE_TILE_SZ,
858 - )
859 -
860 - o = flash_paged_attention[1, n_kv_head](**kwargs)
861 - return o
862 -
863 -
864 -def reshape_and_cache(
865 - key: torch.Tensor,
866 - value: torch.Tensor,
867 - kv_cache: torch.Tensor,
868 - slot_mapping: torch.Tensor,
869 -) -> None:
870 - """
871 - Writes key-value pairs to the KV cache at specified positions.
872 -
873 - Args:
874 - key (torch.Tensor): Key tensor with shape
875 - (num_tokens, n_kv_head, d_head)
876 - value (torch.Tensor): Value tensor with shape
877 - (num_tokens, n_kv_head, d_head)
878 - kv_cache (torch.Tensor): Key/value cache tensor with shape
879 - (2, num_blocks, n_kv_head, block_size, d_head)
880 - slot_mapping (torch.Tensor): Mapping tensor indicating cache positions
881 - with shape (num_tokens)
882 -
883 - Returns:
884 - None: Updates the kv_cache tensor in-place
885 - """
886 - block_size = kv_cache.size(3)
887 - n_kv_head = key.size(1)
888 -
889 - # Calculate indices with explicit floor division
890 - block_indices = torch.div(slot_mapping, block_size, rounding_mode="floor")
891 - block_offsets = slot_mapping % block_size
892 -
893 - # Create the head indices tensor
894 - head_indices = torch.arange(n_kv_head, device=key.device)
895 -
896 - # Update caches using index_put_
897 - kv_cache.index_put_(
898 - (torch.tensor([0], device=key.device), block_indices[:, None],
899 - head_indices[None, :], block_offsets[:, None]), key)
900 -
901 - kv_cache.index_put_(
902 - (torch.tensor([1], device=key.device), block_indices[:, None],
903 - head_indices[None, :], block_offsets[:, None]), value)
vllm/collect_env.py
+1 -15
@@ -54,7 +54,6 @@ SystemEnv = namedtuple(
54 'is_xnnpack_available',
55 'cpu_info',
56 'rocm_version', # vllm specific field
57 - 'neuron_sdk_version', # vllm specific field
57 'vllm_version', # vllm specific field
58 'vllm_build_flags', # vllm specific field
59 'gpu_topo', # vllm specific field
@@ -275,15 +274,6 @@ def get_rocm_version(run_lambda):
274 r'HIP version: (\S+)')
275
276
278 -def get_neuron_sdk_version(run_lambda):
279 - # Adapted from your install script
280 - try:
281 - result = run_lambda(["neuron-ls"])
282 - return result if result[0] == 0 else 'N/A'
283 - except Exception:
284 - return 'N/A'
285 -
286 -
277 def get_vllm_version():
278 from vllm import __version__, __version_tuple__
279
@@ -306,10 +296,9 @@ def get_vllm_version():
296
297 def summarize_vllm_build_flags():
298 # This could be a static method if the flags are constant, or dynamic if you need to check environment variables, etc.
309 - return 'CUDA Archs: {}; ROCm: {}; Neuron: {}'.format(
299 + return 'CUDA Archs: {}; ROCm: {}'.format(
300 os.environ.get('TORCH_CUDA_ARCH_LIST', 'Not Set'),
301 'Enabled' if os.environ.get('ROCM_HOME') else 'Disabled',
312 - 'Enabled' if os.environ.get('NEURON_CORES') else 'Disabled',
302 )
303
304
@@ -601,7 +590,6 @@ def get_env_info():
590 conda_packages = get_conda_packages(run_lambda)
591
592 rocm_version = get_rocm_version(run_lambda)
604 - neuron_sdk_version = get_neuron_sdk_version(run_lambda)
593 vllm_version = get_vllm_version()
594 vllm_build_flags = summarize_vllm_build_flags()
595 gpu_topo = get_gpu_topo(run_lambda)
@@ -635,7 +623,6 @@ def get_env_info():
623 is_xnnpack_available=is_xnnpack_available(),
624 cpu_info=get_cpu_info(run_lambda),
625 rocm_version=rocm_version,
638 - neuron_sdk_version=neuron_sdk_version,
626 vllm_version=vllm_version,
627 vllm_build_flags=vllm_build_flags,
628 gpu_topo=gpu_topo,
@@ -702,7 +689,6 @@ env_info_fmt += """
689 vLLM Info
690 ==============================
691 ROCM Version : {rocm_version}
705 -Neuron SDK Version : {neuron_sdk_version}
692 vLLM Version : {vllm_version}
693 vLLM Build Flags:
694 {vllm_build_flags}
vllm/config/__init__.py
+2 -20
@@ -461,11 +461,6 @@ class ModelConfig:
461 DP (which is controlled by `--data-parallel-size`).
462 This is only supported on a per-model basis and falls back to
463 `"weights"` if the encoder does not support DP."""
464 - override_neuron_config: dict[str, Any] = field(default_factory=dict)
465 - """Initialize non-default neuron config or override default neuron config
466 - that are specific to Neuron devices, this argument will be used to
467 - configure the neuron config that can not be gathered from the vllm
468 - arguments. e.g. `{"cast_logits_dtype": "bfloat16"}`."""
464 pooler_config: Optional["PoolerConfig"] = field(init=False)
465 """Pooler config which controls the behaviour of output pooling in pooling
466 models."""
@@ -785,10 +780,6 @@ class ModelConfig:
780 if not self.skip_tokenizer_init:
781 self._verify_tokenizer_mode()
782
788 - if (not current_platform.is_neuron() and self.override_neuron_config):
789 - raise ValueError(
790 - "`override_neuron_config` is only supported on Neuron.")
791 -
783 # Avoid running try_verify_and_update_config multiple times
784 self.config_updated = False
785
@@ -1696,13 +1687,7 @@ class ModelConfig:
1687 """
1688 For Mllama, VLLM overrides HF's is_encoder_decoder flag and sets it to
1689 True to enable cross-attention
1699 - Neuron needs all multimodal data to be in the decoder and does not
1700 - need to explicitly enable cross-attention
1690 """
1702 - if (current_platform.is_neuron()
1703 - and self.hf_config.model_type == "mllama"):
1704 - return False
1705 -
1691 return is_encoder_decoder(self.hf_config)
1692
1693 @property
@@ -1871,7 +1856,7 @@ class LoadConfig:
1856 self.ignore_patterns = ["original/**/*"]
1857
1858
1874 -Device = Literal["auto", "cuda", "neuron", "cpu", "tpu", "xpu"]
1859 +Device = Literal["auto", "cuda", "cpu", "tpu", "xpu"]
1860
1861
1862 @config
@@ -1927,9 +1912,7 @@ class DeviceConfig:
1912 self.device_type = self.device.type
1913
1914 # Some device types require processing inputs on CPU
1930 - if self.device_type in ["neuron"]:
1931 - self.device = torch.device("cpu")
1932 - elif self.device_type in ["tpu"]:
1915 + if self.device_type in ["tpu"]:
1916 self.device = None
1917 else:
1918 # Set device with device type
@@ -3941,7 +3924,6 @@ class VllmConfig:
3924 f"skip_tokenizer_init={self.model_config.skip_tokenizer_init}, "
3925 f"tokenizer_mode={self.model_config.tokenizer_mode}, "
3926 f"revision={self.model_config.revision}, "
3944 - f"override_neuron_config={self.model_config.override_neuron_config}, " # noqa
3927 f"tokenizer_revision={self.model_config.tokenizer_revision}, "
3928 f"trust_remote_code={self.model_config.trust_remote_code}, "
3929 f"dtype={self.model_config.dtype}, "
vllm/config/cache.py
+2 -3
@@ -33,9 +33,8 @@ class CacheConfig:
33 """Configuration for the KV cache."""
34
35 block_size: SkipValidation[BlockSize] = None # type: ignore
36 - """Size of a contiguous cache block in number of tokens. This is ignored on
37 - neuron devices and set to `--max-model-len`. On CUDA devices, only block
38 - sizes up to 32 are supported. On HPU devices, block size defaults to 128.
36 + """Size of a contiguous cache block in number of tokens. On CUDA devices,
37 + only block sizes up to 32 are supported.
38
39 This config has no static default. If left unspecified by the user, it will
40 be set in `Platform.check_and_update_config()` based on the current
vllm/config/parallel.py
+1 -4
@@ -377,10 +377,7 @@ class ParallelConfig:
377 from vllm.executor import ray_utils
378 backend: DistributedExecutorBackend = "mp"
379 ray_found = ray_utils.ray_is_available()
380 - if current_platform.is_neuron():
381 - # neuron uses single process to control multiple devices
382 - backend = "uni"
383 - elif current_platform.is_tpu() and envs.VLLM_XLA_USE_SPMD:
380 + if current_platform.is_tpu() and envs.VLLM_XLA_USE_SPMD:
381 backend = "uni"
382 elif (current_platform.is_cuda()
383 and cuda_device_count_stateless() < self.world_size):
vllm/distributed/device_communicators/neuron_communicator.py deleted
-20
@@ -1,20 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -import torch
4 -
5 -from vllm.distributed.device_communicators.base_device_communicator import (
6 - DeviceCommunicatorBase)
7 -from vllm.platforms import current_platform
8 -
9 -if current_platform.is_neuron():
10 - import torch_xla.core.xla_model as xm
11 -
12 -
13 -class NeuronCommunicator(DeviceCommunicatorBase):
14 -
15 - def all_reduce(self, x: torch.Tensor) -> torch.Tensor:
16 - return xm.all_reduce(xm.REDUCE_SUM, x)
17 -
18 - def all_gather(self, x: torch.Tensor, dim: int = -1) -> torch.Tensor:
19 - assert dim == -1, "Neuron only supports dim=-1 for all-gather."
20 - return xm.all_gather(x, dim=dim)
vllm/engine/arg_utils.py
-5
@@ -419,8 +419,6 @@ class EngineArgs:
419 scheduling_policy: SchedulerPolicy = SchedulerConfig.policy
420 scheduler_cls: Union[str, Type[object]] = SchedulerConfig.scheduler_cls
421
422 - override_neuron_config: dict[str, Any] = \
423 - get_field(ModelConfig, "override_neuron_config")
422 override_pooler_config: Optional[Union[dict, PoolerConfig]] = \
423 ModelConfig.override_pooler_config
424 compilation_config: CompilationConfig = \
@@ -561,8 +559,6 @@ class EngineArgs:
559 help=model_kwargs["hf_token"]["help"])
560 model_group.add_argument("--hf-overrides",
561 **model_kwargs["hf_overrides"])
564 - model_group.add_argument("--override-neuron-config",
565 - **model_kwargs["override_neuron_config"])
562 model_group.add_argument("--override-pooler-config",
563 **model_kwargs["override_pooler_config"])
564 model_group.add_argument("--logits-processor-pattern",
@@ -992,7 +988,6 @@ class EngineArgs:
988 mm_processor_kwargs=self.mm_processor_kwargs,
989 mm_processor_cache_gb=self.mm_processor_cache_gb,
990 mm_encoder_tp_mode=self.mm_encoder_tp_mode,
995 - override_neuron_config=self.override_neuron_config,
991 override_pooler_config=self.override_pooler_config,
992 logits_processor_pattern=self.logits_processor_pattern,
993 generation_config=self.generation_config,
vllm/envs.py
+1 -1
@@ -236,7 +236,7 @@ environment_variables: dict[str, Callable[[], Any]] = {
236 # ================== Installation Time Env Vars ==================
237
238 # Target device of vLLM, supporting [cuda (by default),
239 - # rocm, neuron, cpu]
239 + # rocm, cpu]
240 "VLLM_TARGET_DEVICE":
241 lambda: os.getenv("VLLM_TARGET_DEVICE", "cuda").lower(),
242
vllm/model_executor/custom_op.py
-7
@@ -73,11 +73,6 @@ class CustomOp(nn.Module):
73 # NOTE(woosuk): This is a placeholder for future extensions.
74 return self.forward_native(*args, **kwargs)
75
76 - def forward_neuron(self, *args, **kwargs):
77 - # By default, we assume that Neuron ops are compatible with the
78 - # PyTorch-native implementation.
79 - return self.forward_native(*args, **kwargs)
80 -
76 def forward_oot(self, *args, **kwargs):
77 # By default, we assume that OOT ops are compatible with the
78 # PyTorch-native implementation.
@@ -105,8 +100,6 @@ class CustomOp(nn.Module):
100 return self.forward_tpu
101 elif current_platform.is_xpu():
102 return self.forward_xpu
108 - elif current_platform.is_neuron():
109 - return self.forward_neuron
103 elif current_platform.is_out_of_tree():
104 return self.forward_oot
105 else:
vllm/model_executor/layers/activation.py
-7
@@ -95,13 +95,6 @@ class SiluAndMul(CustomOp):
95 self.op(out, x)
96 return out
97
98 - def forward_neuron(self, x: torch.Tensor) -> torch.Tensor:
99 - d = x.shape[-1] // 2
100 - x_reshaped = x.view(-1, x.shape[-1])
101 - s = x_reshaped[:, :d] * F.sigmoid(x_reshaped[:, :d])
102 - result = s * x_reshaped[:, d:]
103 - return result.view(*x.shape[:-1], d)
104 -
98
99 @CustomOp.register("mul_and_silu")
100 class MulAndSilu(CustomOp):
vllm/model_executor/layers/quantization/__init__.py
-3
@@ -26,7 +26,6 @@ QuantizationMethods = Literal[
26 "bitsandbytes",
27 "hqq",
28 "experts_int8",
29 - "neuron_quant",
29 "ipex",
30 "quark",
31 "moe_wna16",
@@ -108,7 +107,6 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
107 from .modelopt import ModelOptFp8Config, ModelOptNvFp4Config
108 from .moe_wna16 import MoeWNA16Config
109 from .mxfp4 import Mxfp4Config
111 - from .neuron_quant import NeuronQuantConfig
110 from .petit import PetitNvFp4Config
111 from .ptpc_fp8 import PTPCFp8Config
112 from .rtn import RTNConfig
@@ -135,7 +133,6 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
133 "ptpc_fp8": PTPCFp8Config,
134 "hqq": HQQMarlinConfig,
135 "experts_int8": ExpertsInt8Config,
138 - "neuron_quant": NeuronQuantConfig,
136 "ipex": IPEXConfig,
137 "quark": QuarkConfig,
138 "moe_wna16": MoeWNA16Config,
vllm/model_executor/layers/quantization/neuron_quant.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 os
5 -from importlib.util import find_spec
6 -from typing import Any, Optional
7 -
8 -from torch.nn import Module
9 -
10 -from vllm.model_executor.layers.quantization import QuantizationMethods
11 -from vllm.model_executor.layers.quantization.base_config import (
12 - QuantizationConfig)
13 -
14 -SUPPORTED_QUANT_DTYPE_LIST = ['s8', 'f8e4m3fn']
15 -
16 -
17 -class AlwaysSupportedDtypes(list):
18 -
19 - def __contains__(self, item):
20 - return True
21 -
22 -
23 -class NeuronQuantConfig(QuantizationConfig):
24 - """Int8 Quantization Config class for Neuron Backend."""
25 -
26 - def __init__(
27 - self,
28 - dequant_dtype: str = "f16",
29 - quantize_method: str = "vector_dynamic",
30 - ) -> None:
31 - super().__init__()
32 - self.quant_dtype = os.getenv("NEURON_QUANT_DTYPE", "s8")
33 - if self.quant_dtype not in SUPPORTED_QUANT_DTYPE_LIST:
34 - raise ValueError(
35 - f"Neuron quantization datatype {self.quant_dtype} is not valid,"
36 - f" the quantization datatype should match one of the below "
37 - f"types {SUPPORTED_QUANT_DTYPE_LIST}")
38 - self.dequant_dtype = dequant_dtype
39 - self.quantize_method = quantize_method
40 -
41 - def get_name(self) -> QuantizationMethods:
42 - return "neuron_quant"
43 -
44 - def get_supported_act_dtypes(self) -> list[str]:
45 - # Neuron implements custom handling logic for quantization support
46 - return AlwaysSupportedDtypes()
47 -
48 - @classmethod
49 - def get_min_capability(cls) -> int:
50 - raise NotImplementedError(
51 - "This function should not be called with Neuron Backend")
52 -
53 - @staticmethod
54 - def get_config_filenames() -> list[str]:
55 - return []
56 -
57 - @classmethod
58 - def from_config(cls, config: dict[str, Any]) -> "NeuronQuantConfig":
59 - quantize_method = cls.get_from_keys(config, ["quantize_method"])
60 - dequant_dtype = cls.get_from_keys(config, ["dequant_dtype"])
61 - return cls(dequant_dtype=dequant_dtype,
62 - quantize_method=quantize_method)
63 -
64 - def get_quant_method(self, layer: Module, prefix: str) -> Optional[Any]:
65 - if find_spec("transformers_neuronx") is not None:
66 - return self.get_quantization_config()
67 - else:
68 - raise NotImplementedError(
69 - "Neuron Quantization is only supported through"
70 - " transformers_neuronx.")
71 -
72 - def get_quantization_config(self):
73 - from transformers_neuronx.config import QuantizationConfig
74 - return QuantizationConfig(quant_dtype=self.quant_dtype,
75 - dequant_dtype=self.dequant_dtype,
76 - quantize_method=self.quantize_method)
vllm/model_executor/layers/rotary_embedding/base.py
+1 -82
@@ -7,7 +7,7 @@ import torch
7
8 from vllm.model_executor.custom_op import CustomOp
9
10 -from .common import apply_rotary_emb_dispatch, apply_rotary_emb_torch
10 +from .common import apply_rotary_emb_torch
11
12
13 @CustomOp.register("rotary_embedding")
@@ -149,87 +149,6 @@ class RotaryEmbedding(CustomOp):
149 self.cos_sin_cache, self.is_neox_style)
150 return query, key
151
152 - def forward_neuron(
153 - self,
154 - positions: torch.Tensor,
155 - query: torch.Tensor,
156 - key: Optional[torch.Tensor] = None,
157 - offsets: Optional[torch.Tensor] = None,
158 - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
159 -
160 - def _apply_rotary_emb_neuron(
161 - x: torch.Tensor,
162 - cos: torch.Tensor,
163 - sin: torch.Tensor,
164 - is_neox_style: bool,
165 - ) -> torch.Tensor:
166 - cos = cos.unsqueeze(-2).to(x.dtype)
167 - sin = sin.unsqueeze(-2).to(x.dtype)
168 - if is_neox_style:
169 - x1, x2 = torch.chunk(x, 2, dim=-1)
170 - else:
171 - # x1 = x[..., ::2]
172 -
173 - # x2 = x[..., 1::2]
174 - d = x.shape[-1] // 2
175 - x_reshaped = x.view(-1, x.shape[-1])
176 - x1 = x_reshaped[:, ::2].view(*x.shape[:-1], d)
177 - x2 = x_reshaped[:, 1::2].view(*x.shape[:-1], d)
178 - o1 = x1 * cos - x2 * sin
179 - o2 = x2 * cos + x1 * sin
180 - if is_neox_style:
181 - return torch.cat((o1, o2), dim=-1)
182 - else:
183 - return torch.stack((o1, o2), dim=-1).flatten(-2)
184 -
185 - if offsets is not None:
186 - positions = positions + offsets
187 -
188 - self.cos_sin_cache = self.cos_sin_cache.to(query.device,
189 - dtype=query.dtype)
190 -
191 - positions = positions.flatten()
192 - num_tokens = positions.shape[0]
193 - cos_sin = self.cos_sin_cache.index_select(0, positions)
194 - cos, sin = cos_sin.chunk(2, dim=-1)
195 -
196 - query_shape = query.shape
197 - query = query.view(num_tokens, -1, self.head_size)
198 - if key is not None:
199 - key_shape = key.shape
200 - key = key.view(num_tokens, -1, self.head_size)
201 -
202 - if self.rotary_dim == self.head_size:
203 - query = apply_rotary_emb_dispatch(query, cos, sin,
204 - self.is_neox_style)
205 - query = query.reshape(query_shape)
206 - if key is not None:
207 - key = apply_rotary_emb_dispatch(key, cos, sin,
208 - self.is_neox_style)
209 - key = key.reshape(key_shape)
210 - else:
211 - head_size = query.shape[-1]
212 - query_reshaped = query.view(-1, head_size)
213 - query_pass = query_reshaped[:, self.rotary_dim:].view(
214 - *query.shape[:-1], head_size - self.rotary_dim)
215 - query_rot = query_reshaped[:, :self.rotary_dim].view(
216 - *query.shape[:-1], self.rotary_dim)
217 - query_rot = _apply_rotary_emb_neuron(query_rot, cos, sin,
218 - self.is_neox_style)
219 - query = torch.cat((query_rot, query_pass),
220 - dim=-1).reshape(query_shape)
221 -
222 - if key is not None:
223 - key_reshaped = key.view(-1, head_size)
224 - key_pass = key_reshaped[:, self.rotary_dim:].view(
225 - *key.shape[:-1], head_size - self.rotary_dim)
226 - key_rot = key_reshaped[:, :self.rotary_dim].view(
227 - *key.shape[:-1], self.rotary_dim)
228 - key_rot = _apply_rotary_emb_neuron(key_rot, cos, sin,
229 - self.is_neox_style)
230 - key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)
231 - return query, key
232 -
152 def extra_repr(self) -> str:
153 s = f"head_size={self.head_size}, rotary_dim={self.rotary_dim}"
154 s += f", max_position_embeddings={self.max_position_embeddings}"
vllm/model_executor/model_loader/neuron.py deleted
-476
@@ -1,476 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""Utilities for selecting and loading Neuron models in transformers-neuronx
4 -framework."""
5 -import ast
6 -import copy
7 -import importlib
8 -import os
9 -from typing import Optional
10 -
11 -import torch
12 -import torch.nn as nn
13 -from transformers import PretrainedConfig
14 -
15 -from vllm.config import (ModelConfig, ParallelConfig, SchedulerConfig,
16 - SpeculativeConfig)
17 -from vllm.logprobs import Logprob
18 -from vllm.model_executor.layers.logits_processor import LogitsProcessor
19 -from vllm.model_executor.layers.quantization import get_quantization_config
20 -from vllm.model_executor.layers.sampler import Sampler, SamplerOutput
21 -from vllm.model_executor.sampling_metadata import SamplingMetadata
22 -from vllm.sequence import CompletionSequenceGroupOutput, SequenceOutput
23 -
24 -TORCH_DTYPE_TO_NEURON_AMP = {
25 - "auto": "f32",
26 - "half": "f16",
27 - "float16": "f16",
28 - "bfloat16": "bf16",
29 - "float": "f32",
30 - "float32": "f32",
31 - torch.float16: "f16",
32 - torch.bfloat16: "bf16",
33 - torch.float32: "f32",
34 -}
35 -
36 -# Models supported by Neuron.
37 -_NEURON_SUPPORTED_MODELS: dict[str, tuple[str, str, str]] = {
38 - "LlamaForCausalLM": ("transformers_neuronx.llama.model",
39 - "LlamaForSampling", "LlamaForCausalLM"),
40 - "MistralForCausalLM": ("transformers_neuronx.mistral.model",
41 - "MistralForSampling", "MistralForCausalLM")
42 -}
43 -
44 -
45 -class NeuronCausalLM(nn.Module):
46 -
47 - def __init__(self,
48 - config: PretrainedConfig,
49 - on_device_sampling_disabled: bool = False) -> None:
50 - super().__init__()
51 - self.config = config
52 - self.logits_processor = LogitsProcessor(config.vocab_size,
53 - logits_as_input=True)
54 -
55 - self.on_device_sampling_disabled = on_device_sampling_disabled
56 - if self.on_device_sampling_disabled:
57 - # Use default sampler
58 - self.sampler = Sampler()
59 -
60 - # Lazy initialized
61 - self.model: nn.Module
62 -
63 - def forward(
64 - self,
65 - input_ids: torch.Tensor,
66 - positions: torch.Tensor,
67 - input_block_ids: torch.Tensor,
68 - ) -> torch.Tensor:
69 - logits = self.model(input_ids,
70 - cache_ids=positions,
71 - start_ids=input_block_ids)
72 - return logits
73 -
74 - def compute_logits(self, hidden_states: torch.Tensor,
75 - sampling_metadata: SamplingMetadata) -> torch.Tensor:
76 - logits = self.logits_processor(None, hidden_states, sampling_metadata)
77 - return logits
78 -
79 - def sample(
80 - self,
81 - logits: torch.Tensor,
82 - sampling_metadata: SamplingMetadata,
83 - ) -> Optional[SamplerOutput]:
84 -
85 - if self.on_device_sampling_disabled:
86 - next_tokens = self.sampler(logits, sampling_metadata)
87 - return next_tokens
88 -
89 - # On-device sampling outputs the token ids directly.
90 - sampled_token_ids = logits.flatten()
91 - next_tokens = []
92 - sample_idx = 0
93 - for seq_group in sampling_metadata.seq_groups:
94 - samples = []
95 - for seq_id in seq_group.seq_ids:
96 - token_id = sampled_token_ids[sample_idx].item()
97 - samples.append(
98 - SequenceOutput(parent_seq_id=seq_id,
99 - output_token=token_id,
100 - logprobs={token_id: Logprob(token_id)}))
101 - sample_idx += 1
102 - next_tokens.append(
103 - CompletionSequenceGroupOutput(samples=samples,
104 - prompt_logprobs=None))
105 -
106 - return SamplerOutput(outputs=next_tokens)
107 -
108 - def load_weights(self, model_name_or_path: str, **kwargs):
109 - arch = _get_model_architecture(self.config)
110 - neuronx_module_path, neuronx_model_cls_name, hf_model_cls_name = (
111 - _NEURON_SUPPORTED_MODELS[arch])
112 - neuronx_module = importlib.import_module(neuronx_module_path)
113 - neuronx_model_cls = getattr(neuronx_module, neuronx_model_cls_name)
114 -
115 - self.model = neuronx_model_cls.from_pretrained(model_name_or_path,
116 - **kwargs)
117 - self.model.to_neuron()
118 -
119 -
120 -class NeuronSpeculationCausalLM(nn.Module):
121 - """A Neuron-optimized causal language model with speculative decoding."""
122 -
123 - SPECULATION_TERMINATION_ID = -1
124 -
125 - def __init__(self, speculation_model) -> None:
126 - super().__init__()
127 - self.model = speculation_model
128 -
129 - def forward(
130 - self,
131 - input_ids: torch.Tensor,
132 - positions: torch.Tensor,
133 - input_block_ids: torch.Tensor,
134 - ) -> torch.Tensor:
135 - tokens, counts = self.model.speculative_iteration(
136 - input_ids, positions, input_block_ids)
137 -
138 - # Mark the end of accepted speculative tokens for each sequence with the
139 - # speculation termination id.
140 - batch_size, steps = tokens.shape
141 - mask = torch.arange(steps).expand(batch_size, -1) >= counts
142 - tokens[mask] = self.SPECULATION_TERMINATION_ID
143 -
144 - return tokens
145 -
146 - def sample(
147 - self,
148 - logits: torch.Tensor,
149 - sampling_metadata: SamplingMetadata,
150 - ) -> Optional[list[SamplerOutput]]:
151 - batch_size, num_steps = logits.shape
152 - seq_ids = [
153 - seq_id for sg in sampling_metadata.seq_groups
154 - for seq_id in sg.seq_ids
155 - ]
156 - # Organize input tensors by step instead of by sequence.
157 - accepted_token_ids_by_step = logits.transpose(0, 1)
158 - accepted_token_ids_by_step = accepted_token_ids_by_step.tolist()
159 -
160 - sampler_output_list = []
161 - for step_index in range(num_steps):
162 - if all(token_id == self.SPECULATION_TERMINATION_ID
163 - for token_id in accepted_token_ids_by_step[step_index]):
164 - break
165 - step_output_token_ids = []
166 - for sequence_index in range(batch_size):
167 - token_id = accepted_token_ids_by_step[step_index][
168 - sequence_index]
169 - step_output_token_ids.append(
170 - CompletionSequenceGroupOutput(samples=[
171 - SequenceOutput(parent_seq_id=seq_ids[sequence_index],
172 - output_token=token_id,
173 - logprobs={token_id: Logprob(token_id)})
174 - ],
175 - prompt_logprobs=None))
176 - sampler_output_list.append(
177 - SamplerOutput(outputs=step_output_token_ids))
178 - return sampler_output_list
179 -
180 -
181 -def _get_model_architecture(config: PretrainedConfig) -> str:
182 - architectures = getattr(config, "architectures", [])
183 - for arch in architectures:
184 - if arch in _NEURON_SUPPORTED_MODELS:
185 - return arch
186 - raise ValueError(
187 - f"Model architectures {architectures} are not supported on Neuron "
188 - f"for now. Supported architectures: "
189 - f"{list(_NEURON_SUPPORTED_MODELS.keys())}")
190 -
191 -
192 -def _get_buckets(env: str, default_value: list[int]) -> list[int]:
193 - env_value = os.getenv(env)
194 - if env_value is None:
195 - return default_value
196 - buckets_remove_empty = filter(
197 - lambda x: x is not None and len(x.strip()) > 0, env_value.split(","))
198 - buckets_int = map(int, buckets_remove_empty)
199 - buckets_list = list(buckets_int)
200 - return buckets_list
201 -
202 -
203 -def _get_default_neuron_config(model_config: ModelConfig,
204 - parallel_config: ParallelConfig,
205 - scheduler_config: SchedulerConfig):
206 - """Generate a neuron config based on vllm config args."""
207 - from transformers_neuronx.config import ContinuousBatchingConfig
208 - from transformers_neuronx.constants import LAYOUT_BSH
209 -
210 - continuous_batching_config = ContinuousBatchingConfig(
211 - batch_size_for_shared_caches=scheduler_config.max_num_seqs)
212 - quant_config = dict(
213 - dequant_dtype=TORCH_DTYPE_TO_NEURON_AMP[model_config.dtype],
214 - quantize_method="vector_dynamic")
215 - neuron_quantization_config_builder = lambda quant: get_quantization_config(
216 - quant).from_config(quant_config).get_quant_method(None, "")
217 - # TODO: Add Paged attention config to the default neuron arguments.
218 - default_neuron_args = dict(
219 - collectives_layout=LAYOUT_BSH,
220 - attention_layout=LAYOUT_BSH,
221 - fuse_qkv=True,
222 - quant=neuron_quantization_config_builder(model_config.quantization)
223 - if model_config.quantization else None,
224 - continuous_batching=continuous_batching_config,
225 - weight_tiling=bool(model_config.quantization),
226 - on_device_generation=_get_neuron_on_device_generation_config(
227 - model_config))
228 - return default_neuron_args
229 -
230 -
231 -def _get_default_neuron_config_for_speculation(
232 - model_config: ModelConfig, parallel_config: ParallelConfig,
233 - scheduler_config: SchedulerConfig):
234 - """Generate a neuron config for speculative decoding based on
235 - vllm config args."""
236 - from transformers_neuronx.config import ContinuousBatchingConfig
237 - from transformers_neuronx.constants import LAYOUT_BSH
238 -
239 - continuous_batching_config = ContinuousBatchingConfig(
240 - batch_size_for_shared_caches=scheduler_config.max_num_seqs)
241 -
242 - default_neuron_args = dict(collectives_layout=LAYOUT_BSH,
243 - attention_layout=LAYOUT_BSH,
244 - fuse_qkv=True,
245 - on_device_embedding=True,
246 - continuous_batching=continuous_batching_config,
247 - on_device_generation=copy.deepcopy(
248 - model_config.neuron_sampling_params))
249 - return default_neuron_args
250 -
251 -
252 -def _get_neuron_on_device_generation_config(model_config: ModelConfig):
253 - if not _is_neuron_on_device_sampling_disabled(model_config):
254 - return copy.deepcopy(model_config.neuron_sampling_params)
255 - return None
256 -
257 -
258 -def _is_neuron_on_device_sampling_disabled(model_config: ModelConfig) -> bool:
259 - return not getattr(model_config, "neuron_sampling_params", None)
260 -
261 -
262 -def _get_neuron_config_after_override(default_neuron_config,
263 - overridden_neuron_config):
264 - from transformers_neuronx.config import (ContinuousBatchingConfig,
265 - GenerationConfig,
266 - KVCacheQuantizationConfig,
267 - NeuronConfig, QuantizationConfig,
268 - SparseAttnConfig)
269 -
270 - sparse_attn = overridden_neuron_config.pop("sparse_attn", {})
271 - if sparse_attn:
272 - overridden_neuron_config["sparse_attn"] = SparseAttnConfig(
273 - **sparse_attn)
274 -
275 - kv_cache_quant = overridden_neuron_config.pop("kv_cache_quant", {})
276 - if kv_cache_quant:
277 - overridden_neuron_config["kv_cache_quant"] = KVCacheQuantizationConfig(
278 - **kv_cache_quant)
279 -
280 - continuous_batching = overridden_neuron_config.pop("continuous_batching",
281 - {})
282 - if continuous_batching:
283 - overridden_neuron_config[
284 - "continuous_batching"] = ContinuousBatchingConfig(
285 - **continuous_batching)
286 -
287 - quant = overridden_neuron_config.pop("quant", {})
288 - if quant:
289 - overridden_neuron_config["quant"] = QuantizationConfig(**quant)
290 -
291 - on_device_generation = overridden_neuron_config.pop(
292 - "on_device_generation", {})
293 - if on_device_generation:
294 - overridden_neuron_config["on_device_generation"] = GenerationConfig(
295 - **on_device_generation)
296 - default_neuron_config.update(overridden_neuron_config)
297 - return NeuronConfig(**default_neuron_config)
298 -
299 -
300 -def get_neuron_model(model_config: ModelConfig,
301 - parallel_config: ParallelConfig,
302 - scheduler_config: SchedulerConfig) -> nn.Module:
303 - """Initializes a neuron-optimized model for inference."""
304 - # Create a model instance.
305 - model = NeuronCausalLM(
306 - model_config.hf_config,
307 - _is_neuron_on_device_sampling_disabled(model_config))
308 -
309 - default_neuron_config_args = _get_default_neuron_config(
310 - model_config, parallel_config, scheduler_config)
311 -
312 - neuron_config = _get_neuron_config_after_override(
313 - default_neuron_config_args, model_config.override_neuron_config)
314 -
315 - context_length_estimates = _get_buckets("NEURON_CONTEXT_LENGTH_BUCKETS",
316 - [scheduler_config.max_model_len])
317 - n_positions = _get_buckets("NEURON_TOKEN_GEN_BUCKETS",
318 - [scheduler_config.max_model_len])
319 -
320 - model.load_weights(model_config.model,
321 - tp_degree=parallel_config.tensor_parallel_size,
322 - amp=TORCH_DTYPE_TO_NEURON_AMP[model_config.dtype],
323 - neuron_config=neuron_config,
324 - context_length_estimate=context_length_estimates,
325 - n_positions=n_positions,
326 - batch_size=scheduler_config.max_num_seqs)
327 -
328 - return model.eval()
329 -
330 -
331 -def get_neuron_speculation_model(model_config: ModelConfig,
332 - parallel_config: ParallelConfig,
333 - scheduler_config: SchedulerConfig,
334 - speculation_config: SpeculativeConfig):
335 - """Initializes a neuron-optimized speculation model for inference.
336 -
337 - This method is only applicable for speculation with a standalone draft model
338 - """
339 - from transformers_neuronx.fused_speculation import FusedSpeculativeDecoder
340 -
341 - # For Eagle SD, we need to pass in additional parameters in neuron config.
342 - is_eagle = getattr(speculation_config.draft_model_config.hf_config,
343 - "is_eagle", False)
344 -
345 - # Create target model instance.
346 - target_model = NeuronCausalLM(model_config.hf_config)
347 -
348 - default_neuron_config_args = _get_default_neuron_config_for_speculation(
349 - model_config, parallel_config, scheduler_config)
350 - if is_eagle:
351 - default_neuron_config_args['is_eagle_target'] = True
352 -
353 - neuron_config = _get_neuron_config_after_override(
354 - default_neuron_config_args, model_config.override_neuron_config)
355 -
356 - context_length_estimates = _get_buckets("NEURON_CONTEXT_LENGTH_BUCKETS",
357 - [scheduler_config.max_model_len])
358 - n_positions = _get_buckets("NEURON_TOKEN_GEN_BUCKETS",
359 - [scheduler_config.max_model_len])
360 -
361 - target_model.load_weights(
362 - model_config.model,
363 - tp_degree=parallel_config.tensor_parallel_size,
364 - amp=TORCH_DTYPE_TO_NEURON_AMP[model_config.dtype],
365 - neuron_config=neuron_config,
366 - context_length_estimate=context_length_estimates,
367 - n_positions=n_positions,
368 - batch_size=scheduler_config.max_num_seqs)
369 -
370 - target_model.eval()
371 -
372 - # Create draft model instance.
373 - draft_model = NeuronCausalLM(
374 - speculation_config.draft_model_config.hf_config)
375 -
376 - default_draft_neuron_config_args = (
377 - _get_default_neuron_config_for_speculation(
378 - speculation_config.draft_model_config, parallel_config,
379 - scheduler_config))
380 - if is_eagle:
381 - default_draft_neuron_config_args['is_eagle_draft'] = True
382 - default_draft_neuron_config_args['has_pre_attention_norm'] = False
383 -
384 - draft_neuron_config = _get_neuron_config_after_override(
385 - default_draft_neuron_config_args,
386 - speculation_config.draft_model_config.override_neuron_config)
387 -
388 - draft_model.load_weights(speculation_config.draft_model_config.model,
389 - tp_degree=speculation_config.
390 - draft_parallel_config.tensor_parallel_size,
391 - amp=TORCH_DTYPE_TO_NEURON_AMP[
392 - speculation_config.draft_model_config.dtype],
393 - neuron_config=draft_neuron_config,
394 - context_length_estimate=context_length_estimates,
395 - n_positions=n_positions,
396 - batch_size=scheduler_config.max_num_seqs)
397 -
398 - draft_model.eval()
399 -
400 - num_speculative_tokens = speculation_config.num_speculative_tokens
401 - # Create speculation model instance.
402 - speculation_model = FusedSpeculativeDecoder(draft_model.model,
403 - target_model.model,
404 - num_speculative_tokens)
405 - speculation_model.to_neuron()
406 -
407 - return NeuronSpeculationCausalLM(speculation_model)
408 -
409 -
410 -def get_neuron_eagle_speculation_model(model_config: ModelConfig,
411 - parallel_config: ParallelConfig,
412 - scheduler_config: SchedulerConfig,
413 - speculation_config: SpeculativeConfig):
414 - """Initializes a neuron-optimized EAGLE speculation model for inference."""
415 - from transformers_neuronx.eagle_speculation import EagleSpeculativeDecoder
416 -
417 - # Create target model instance.
418 - target_model = NeuronCausalLM(model_config.hf_config)
419 -
420 - default_neuron_config_args = _get_default_neuron_config_for_speculation(
421 - model_config, parallel_config, scheduler_config)
422 - default_neuron_config_args['is_eagle_target'] = True
423 - neuron_config = _get_neuron_config_after_override(
424 - default_neuron_config_args, model_config.override_neuron_config)
425 -
426 - context_length_estimates = _get_buckets("NEURON_CONTEXT_LENGTH_BUCKETS",
427 - [scheduler_config.max_model_len])
428 - n_positions = _get_buckets("NEURON_TOKEN_GEN_BUCKETS",
429 - [scheduler_config.max_model_len])
430 -
431 - target_model.load_weights(
432 - model_config.model,
433 - tp_degree=parallel_config.tensor_parallel_size,
434 - amp=TORCH_DTYPE_TO_NEURON_AMP[model_config.dtype],
435 - neuron_config=neuron_config,
436 - context_length_estimate=context_length_estimates,
437 - n_positions=n_positions,
438 - batch_size=scheduler_config.max_num_seqs)
439 -
440 - target_model.eval()
441 -
442 - # Create draft model instance.
443 - draft_model = NeuronCausalLM(
444 - speculation_config.draft_model_config.hf_config)
445 -
446 - default_draft_neuron_config_args = (
447 - _get_default_neuron_config_for_speculation(
448 - speculation_config.draft_model_config, parallel_config,
449 - scheduler_config))
450 - default_draft_neuron_config_args['is_eagle_draft'] = True
451 - default_draft_neuron_config_args['has_pre_attention_norm'] = False
452 - draft_neuron_config = _get_neuron_config_after_override(
453 - default_draft_neuron_config_args,
454 - speculation_config.draft_model_config.override_neuron_config)
455 -
456 - draft_model.load_weights(speculation_config.draft_model_config.model,
457 - tp_degree=speculation_config.
458 - draft_parallel_config.tensor_parallel_size,
459 - amp=TORCH_DTYPE_TO_NEURON_AMP[
460 - speculation_config.draft_model_config.dtype],
461 - neuron_config=draft_neuron_config,
462 - context_length_estimate=context_length_estimates,
463 - n_positions=n_positions,
464 - batch_size=scheduler_config.max_num_seqs)
465 -
466 - draft_model.eval()
467 -
468 - token_tree: dict[int, list[int]] = ast.literal_eval(
469 - speculation_config.speculative_token_tree)
470 -
471 - speculation_model = EagleSpeculativeDecoder(draft_model.model,
472 - target_model.model,
473 - token_tree=token_tree)
474 - speculation_model.to_neuron()
475 -
476 - return NeuronSpeculationCausalLM(speculation_model)
vllm/model_executor/model_loader/neuronx_distributed.py deleted
-685
@@ -1,685 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""Utilities for selecting and loading Neuron models in
4 -neuronx-distributed-inference framework."""
5 -# Disabling yapf because yapf and isort have conflicts for the below imports
6 -# yapf: disable
7 -import copy
8 -import hashlib
9 -import importlib
10 -import multiprocessing
11 -import os
12 -import shutil
13 -from typing import Optional
14 -
15 -import torch
16 -import torch.nn as nn
17 -from neuronx_distributed_inference.models.config import (
18 - FusedSpecNeuronConfig, OnDeviceSamplingConfig)
19 -from neuronx_distributed_inference.models.mllama.utils import (
20 - create_vision_mask)
21 -from neuronx_distributed_inference.modules.lora_serving import (
22 - LoraServingConfig)
23 -from neuronx_distributed_inference.utils.hf_adapter import (
24 - load_pretrained_config)
25 -from transformers import AutoModelForCausalLM, AutoTokenizer, PretrainedConfig
26 -
27 -from vllm.config import (ModelConfig, ParallelConfig, SchedulerConfig,
28 - SpeculativeConfig)
29 -from vllm.logger import init_logger
30 -from vllm.logprobs import Logprob
31 -from vllm.model_executor.layers.logits_processor import LogitsProcessor
32 -from vllm.model_executor.layers.sampler import Sampler, SamplerOutput
33 -from vllm.model_executor.sampling_metadata import SamplingMetadata
34 -from vllm.sequence import CompletionSequenceGroupOutput, SequenceOutput
35 -
36 -# yapf: enable
37 -logger = init_logger(__name__)
38 -
39 -TORCH_DTYPE_TO_NEURON_AMP = {
40 - "auto": "float32",
41 - "half": "float16",
42 - "float16": "float16",
43 - "bfloat16": "bfloat16",
44 - "float": "float32",
45 - "float32": "float32",
46 - torch.float16: "float16",
47 - torch.bfloat16: "bfloat16",
48 - torch.float32: "float32",
49 -}
50 -
51 -# Models supported by Neuronx distributed for inference.
52 -_NEURON_SUPPORTED_MODELS: dict[str, tuple[str, str]] = {
53 - "LlamaForCausalLM":
54 - ("neuronx_distributed_inference.models.llama.modeling_llama",
55 - "NeuronLlamaForCausalLM"),
56 - "MistralForCausalLM":
57 - ("neuronx_distributed_inference.models.llama.modeling_llama",
58 - "NeuronLlamaForCausalLM"),
59 - "DbrxForCausalLM":
60 - ("neuronx_distributed_inference.models.dbrx.modeling_dbrx",
61 - "NeuronDbrxForCausalLM"),
62 - "MixtralForCausalLM":
63 - ("neuronx_distributed_inference.models.mixtral.modeling_mixtral",
64 - "NeuronMixtralForCausalLM"),
65 - "MllamaForConditionalGeneration":
66 - ("neuronx_distributed_inference.models.mllama.modeling_mllama",
67 - "NeuronMllamaForCausalLM"),
68 -}
69 -
70 -
71 -class NeuronCausalLM(nn.Module):
72 -
73 - def __init__(
74 - self,
75 - config: PretrainedConfig,
76 - ) -> None:
77 - super().__init__()
78 - self.config = config
79 - self.logits_processor = LogitsProcessor(config.vocab_size,
80 - logits_as_input=True)
81 - self.sampler = Sampler()
82 -
83 - # Lazy initialized
84 - self.model: nn.Module
85 -
86 - def forward(self,
87 - input_ids: torch.Tensor,
88 - positions: torch.Tensor,
89 - input_block_ids: torch.Tensor,
90 - sampling_params: torch.Tensor,
91 - prev_hidden: Optional[torch.Tensor] = None,
92 - adapter_ids: Optional[torch.Tensor] = None) -> torch.Tensor:
93 - # sort block ids sequentially for perf/neuron support reasons
94 - sorted_input_block_ids, sorted_indices = torch.sort(input_block_ids)
95 - input_ids = torch.index_select(input_ids, 0, sorted_indices)
96 - positions = torch.index_select(positions, 0, sorted_indices)
97 - sampling_params = torch.index_select(sampling_params, 0,
98 - sorted_indices)
99 - output = self.model(input_ids,
100 - attention_mask=None,
101 - position_ids=positions,
102 - seq_ids=sorted_input_block_ids,
103 - sampling_params=sampling_params,
104 - prev_hidden=prev_hidden,
105 - adapter_ids=adapter_ids)
106 - # on-device sampling
107 - if self.config.neuron_config.on_device_sampling_config:
108 - output = output.hidden_states
109 - else:
110 - output = output.logits[:, -1, :]
111 -
112 - restored_indices = torch.argsort(sorted_indices)
113 - if input_block_ids.shape[0] != 1:
114 - output = torch.index_select(output, 0, restored_indices)
115 -
116 - return output
117 -
118 - def compute_logits(self, hidden_states: torch.Tensor,
119 - sampling_metadata: SamplingMetadata) -> torch.Tensor:
120 - logits = self.logits_processor(None, hidden_states, sampling_metadata)
121 - return logits
122 -
123 - def sample(
124 - self,
125 - logits: torch.Tensor,
126 - sampling_metadata: SamplingMetadata,
127 - ) -> Optional[SamplerOutput]:
128 - # on-device sampling
129 - if self.config.neuron_config.on_device_sampling_config:
130 - batch_size = logits.shape
131 - seq_ids = [
132 - seq_id for sg in sampling_metadata.seq_groups
133 - for seq_id in sg.seq_ids
134 - ]
135 - assert len(seq_ids) == list(batch_size)[0], "batch size mismatch"
136 - # Organize input tensors by step instead of by sequence.
137 - accepted_token_ids_by_step = logits.flatten()
138 - accepted_token_ids_by_step = accepted_token_ids_by_step.tolist()
139 -
140 - step_output_token_ids = []
141 - for i, seq_id in enumerate(seq_ids):
142 - token_id = accepted_token_ids_by_step[i]
143 - step_output_token_ids.append(
144 - CompletionSequenceGroupOutput(samples=[
145 - SequenceOutput(parent_seq_id=seq_id,
146 - output_token=token_id,
147 - logprobs={token_id: Logprob(token_id)})
148 - ],
149 - prompt_logprobs=None))
150 - return SamplerOutput(outputs=step_output_token_ids)
151 - else:
152 - return self.sampler(logits, sampling_metadata)
153 -
154 - def load_weights(self, model_name_or_path: str, **kwargs):
155 - arch = _get_model_architecture(self.config)
156 - neuronx_module_path, neuronx_model_cls_name = (
157 - _NEURON_SUPPORTED_MODELS[arch])
158 - neuronx_module = importlib.import_module(neuronx_module_path)
159 - neuronx_model_cls = getattr(neuronx_module, neuronx_model_cls_name)
160 - neuron_config = neuronx_model_cls.get_neuron_config_cls()(
161 - **kwargs['neuron_config'])
162 - self.config.neuron_config = neuron_config
163 - config = neuronx_model_cls.get_config_cls()(
164 - neuron_config,
165 - load_config=load_pretrained_config(model_name_or_path))
166 - hashed_config = hashlib.md5(config.to_json_string().encode('utf-8'),
167 - usedforsecurity=False).hexdigest()
168 - if os.getenv("NEURON_COMPILED_ARTIFACTS") is not None:
169 - compiled_model_path = os.getenv("NEURON_COMPILED_ARTIFACTS")
170 - elif os.path.exists(model_name_or_path):
171 - compiled_model_path = os.path.join(model_name_or_path,
172 - "neuron-compiled-artifacts",
173 - hashed_config)
174 - shutil.rmtree(compiled_model_path, ignore_errors=True)
175 - else:
176 - compiled_model_path = os.path.join("local-models",
177 - model_name_or_path,
178 - "neuron-compiled-artifacts",
179 - hashed_config)
180 - shutil.rmtree(compiled_model_path, ignore_errors=True)
181 - try:
182 - self.model = neuronx_model_cls(compiled_model_path)
183 - override_neuron_config = kwargs["override_neuron_config"]
184 - for k, v in override_neuron_config.items():
185 - setattr(self.model.config.neuron_config, k, v)
186 - self.model.load(compiled_model_path)
187 - return
188 - except (FileNotFoundError, ValueError) as e:
189 - logger.warning("Exception: %s", e)
190 - logger.warning("Failed to load the model from %s, Recompiling...",
191 - compiled_model_path)
192 - if not os.path.exists(model_name_or_path):
193 - hf_model = AutoModelForCausalLM.from_pretrained(model_name_or_path)
194 - saved_path = os.path.join("local-models", model_name_or_path)
195 - hf_model.save_pretrained(saved_path)
196 - model_name_or_path = saved_path
197 - self.model = neuronx_model_cls(model_name_or_path, config)
198 - self.model.compile(compiled_model_path)
199 - self.model.load(compiled_model_path)
200 -
201 -
202 -class NeuronMllamaForCausalLM(nn.Module):
203 -
204 - def __init__(self,
205 - config: PretrainedConfig,
206 - on_device_sampling_disabled: bool = False) -> None:
207 - super().__init__()
208 - # has_image is the only multimodal input that is used in
209 - # token-generation
210 - # This is a cache (on CPU) that saves has_image data per sequence id
211 - # The number of entries in this cache is <= Batch-Size
212 - self.has_image_cache: dict[int, torch.Tensor] = {}
213 - self.config = config
214 - self.logits_processor = LogitsProcessor(
215 - config.get_text_config().vocab_size, logits_as_input=True)
216 -
217 - self.on_device_sampling_disabled = on_device_sampling_disabled
218 - if self.on_device_sampling_disabled:
219 - # Use default sampler
220 - self.sampler = Sampler()
221 -
222 - # Lazy initialized
223 - self.model: nn.Module
224 - self.is_reorder_needed: bool = True
225 -
226 - def read_from_has_image_cache(self, seq_ids: torch.Tensor):
227 - has_image_list = []
228 - for index in range(len(seq_ids)):
229 - seq_id = seq_ids[index].item()
230 - if seq_id in self.has_image_cache:
231 - has_image_list.append(self.has_image_cache[seq_id])
232 - else:
233 - has_image_list.append(torch.tensor([0]))
234 - return torch.tensor(has_image_list)
235 -
236 - def write_to_has_image_cache(self, seq_ids: torch.Tensor,
237 - has_image: torch.Tensor):
238 - for index in range(len(seq_ids)):
239 - seq_id = seq_ids[index].item()
240 - if index < len(has_image):
241 - self.has_image_cache[seq_id] = has_image[index]
242 - else:
243 - self.has_image_cache[seq_id] = torch.zeros(1)
244 -
245 - def forward(self, input_ids: torch.Tensor, positions: torch.Tensor,
246 - seq_ids: torch.Tensor, pixel_values: torch.Tensor,
247 - aspect_ratios: torch.Tensor, num_chunks: torch.Tensor,
248 - has_image: torch.Tensor, sampling_params) -> torch.Tensor:
249 -
250 - # We update the has_image cache during prefill
251 - # and read the has_image cache during decode
252 - if input_ids.shape[-1] > 1: # prefill
253 - self.write_to_has_image_cache(seq_ids, has_image)
254 - else:
255 - has_image = self.read_from_has_image_cache(seq_ids)
256 - bs = input_ids.shape[0]
257 - num_chunks = torch.zeros((bs, 1))
258 - aspect_ratios = torch.zeros((bs, 1, 2))
259 -
260 - input_block_ids = seq_ids
261 - origin_input_block_ids = seq_ids
262 - if self.is_reorder_needed:
263 - # sort block ids sequentially for perf/neuron support reasons
264 - input_block_ids, sorted_indices = torch.sort(input_block_ids)
265 - input_ids = torch.index_select(input_ids, 0, sorted_indices)
266 - positions = torch.index_select(positions, 0, sorted_indices)
267 - sampling_params = torch.index_select(sampling_params, 0,
268 - sorted_indices)
269 - pixel_values = torch.index_select(pixel_values, 0, sorted_indices)
270 - aspect_ratios = torch.index_select(aspect_ratios, 0,
271 - sorted_indices)
272 - num_chunks = torch.index_select(num_chunks, 0, sorted_indices)
273 - has_image = torch.index_select(has_image, 0, sorted_indices)
274 -
275 - self.vision_mask = create_vision_mask(input_ids, self.vision_token_id)
276 - output = self.model(
277 - input_ids.to(torch.int32),
278 - attention_mask=None,
279 - position_ids=positions.to(torch.int32),
280 - seq_ids=seq_ids.flatten().to(torch.int32),
281 - pixel_values=pixel_values.to(
282 - self.config.vision_config.torch_dtype),
283 - aspect_ratios=aspect_ratios.to(torch.int32),
284 - vision_mask=self.vision_mask.to(torch.int32),
285 - sampling_params=sampling_params,
286 - num_chunks=num_chunks.to(torch.int32),
287 - has_image=has_image.to(torch.int32),
288 - )
289 - if self.config.neuron_config.on_device_sampling_config:
290 - output = output.hidden_states
291 - else:
292 - output = output.logits[:, -1, :]
293 -
294 - if self.is_reorder_needed and origin_input_block_ids.shape[0] != 1:
295 - restored_indices = torch.argsort(sorted_indices)
296 - output = torch.index_select(output, 0, restored_indices)
297 - return output
298 -
299 - def compute_logits(self, hidden_states: torch.Tensor,
300 - sampling_metadata: SamplingMetadata) -> torch.Tensor:
301 - logits = self.logits_processor(None, hidden_states, sampling_metadata)
302 - return logits
303 -
304 - def sample(self, hidden_states, sampling_metadata):
305 - if not self.on_device_sampling_disabled:
306 - with torch.profiler.record_function("sample"):
307 - hidden_states = hidden_states.flatten()
308 - res = []
309 - sample_idx = 0
310 - for seq_group in sampling_metadata.seq_groups:
311 - seq_ids = seq_group.seq_ids
312 - samples = []
313 - for seq_id in seq_ids:
314 - token_id = hidden_states[sample_idx].item()
315 - samples.append(
316 - SequenceOutput(
317 - parent_seq_id=seq_id,
318 - output_token=token_id,
319 - logprobs={token_id: Logprob(token_id)}))
320 - sample_idx += 1
321 - res.append(
322 - CompletionSequenceGroupOutput(samples=samples,
323 - prompt_logprobs=None))
324 - next_tokens = SamplerOutput(outputs=res)
325 - else:
326 - next_tokens = self.sampler(None, hidden_states, sampling_metadata)
327 - return next_tokens
328 -
329 - def load_weights(self, model_name_or_path: str, **kwargs):
330 - arch = _get_model_architecture(self.config)
331 - neuronx_module_path, neuronx_model_cls_name = (
332 - _NEURON_SUPPORTED_MODELS[arch])
333 - neuronx_module = importlib.import_module(neuronx_module_path)
334 - neuronx_model_cls = getattr(neuronx_module, neuronx_model_cls_name)
335 - neuron_config = neuronx_model_cls.get_neuron_config_cls()(
336 - **kwargs['neuron_config'])
337 - self.config.neuron_config = neuron_config
338 - logger.info("neuron_config buckets: %s",
339 - self.config.neuron_config.buckets)
340 - config = neuronx_model_cls.get_config_cls()(
341 - neuron_config,
342 - load_config=load_pretrained_config(model_name_or_path))
343 - hashed_config = hashlib.md5(config.to_json_string().encode('utf-8'),
344 - usedforsecurity=False).hexdigest()
345 - if os.getenv("NEURON_COMPILED_ARTIFACTS") is not None:
346 - compiled_model_path = os.getenv("NEURON_COMPILED_ARTIFACTS")
347 - elif os.path.exists(model_name_or_path):
348 - compiled_model_path = os.path.join(model_name_or_path,
349 - "neuron-compiled-artifacts",
350 - hashed_config)
351 - else:
352 - compiled_model_path = os.path.join("local-models",
353 - model_name_or_path,
354 - "neuron-compiled-artifacts",
355 - hashed_config)
356 - try:
357 - self.model = neuronx_model_cls(compiled_model_path)
358 - tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
359 - self.vision_token_id = tokenizer(
360 - "<|image|>", add_special_tokens=False).input_ids[0]
361 - self.model.load(compiled_model_path)
362 - return
363 - except (FileNotFoundError, ValueError):
364 - logger.warning("Failed to load the model from %s, Recompiling...",
365 - compiled_model_path)
366 - if not os.path.exists(model_name_or_path):
367 - hf_model = AutoModelForCausalLM.from_pretrained(model_name_or_path)
368 - saved_path = os.path.join("local-models", model_name_or_path)
369 - hf_model.save_pretrained(saved_path)
370 - model_name_or_path = saved_path
371 - self.model = neuronx_model_cls(model_name_or_path, config)
372 -
373 - logger.info("\nCompiling and saving model to %s", model_name_or_path)
374 -
375 - p = multiprocessing.Process(target=compile_model,
376 - args=(self, compiled_model_path))
377 - p.start()
378 - p.join()
379 -
380 - tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
381 - tokenizer.save_pretrained(compiled_model_path)
382 - logger.info("Successfully compiled and saved the model in %s",
383 - compiled_model_path)
384 -
385 - # Read "<|image|>" token_id from the tokenizer
386 - self.vision_token_id = tokenizer("<|image|>",
387 - add_special_tokens=False).input_ids[0]
388 - logger.info("\nLoading model from compiled checkpoint...")
389 - self.model.load(compiled_model_path)
390 -
391 -
392 -def compile_model(neuron_model, traced_model_path):
393 - neuron_model.model.compile(traced_model_path)
394 -
395 -
396 -class NeuronSpeculationCausalLM(nn.Module):
397 - """A Neuron-optimized causal language model with speculative decoding."""
398 -
399 - def __init__(
400 - self,
401 - config: PretrainedConfig,
402 - ) -> None:
403 - super().__init__()
404 - self.config = config
405 - self.logits_processor = LogitsProcessor(config.vocab_size,
406 - logits_as_input=True)
407 - # Lazy initialized
408 - self.model: nn.Module
409 -
410 - def forward(
411 - self,
412 - input_ids: torch.Tensor,
413 - positions: torch.Tensor,
414 - input_block_ids: torch.Tensor,
415 - sampling_params: torch.Tensor,
416 - ) -> torch.Tensor:
417 - # sort block ids sequentially for perf/neuron support reasons
418 - sorted_input_block_ids, sorted_indices = torch.sort(input_block_ids)
419 - input_ids = torch.index_select(input_ids, 0, sorted_indices)
420 - positions = torch.index_select(positions, 0, sorted_indices)
421 - sampling_params = torch.index_select(sampling_params, 0,
422 - sorted_indices)
423 -
424 - output = self.model(input_ids,
425 - attention_mask=None,
426 - position_ids=positions,
427 - seq_ids=sorted_input_block_ids,
428 - sampling_params=sampling_params)
429 - restored_indices = torch.argsort(sorted_indices)
430 -
431 - # CTX encoding
432 - if (positions[:, 0]).sum().item() == 0:
433 - output = output.fused_outputs[0][:, 0:1]
434 - if input_block_ids.shape[0] != 1:
435 - output = torch.index_select(output, 0, restored_indices)
436 - return output
437 -
438 - # Fused Spec (Generation)
439 - accepted_tokens_with_padding = output.fused_outputs[0]
440 - next_pos_ids = output.fused_outputs[-1]
441 - generated_token_counts = next_pos_ids - positions
442 -
443 - assert torch.any(generated_token_counts == 0).item() is False, \
444 - "NxDI model generated no output for one or more sequences."
445 -
446 - batch_size, steps = accepted_tokens_with_padding.shape
447 - mask = torch.arange(steps).expand(batch_size,
448 - -1) >= generated_token_counts
449 - accepted_tokens_with_padding[mask] = -1
450 -
451 - if input_block_ids.shape[0] != 1:
452 - accepted_tokens_with_padding = torch.index_select(
453 - accepted_tokens_with_padding, 0, restored_indices)
454 -
455 - return accepted_tokens_with_padding
456 -
457 - def sample(
458 - self,
459 - logits: torch.Tensor,
460 - sampling_metadata: SamplingMetadata,
461 - ) -> Optional[list[SamplerOutput]]:
462 - batch_size, num_steps = logits.shape
463 - seq_ids = [
464 - seq_id for sg in sampling_metadata.seq_groups
465 - for seq_id in sg.seq_ids
466 - ]
467 - # Organize input tensors by step instead of by sequence.
468 - accepted_token_ids_by_step = logits.transpose(0, 1)
469 - accepted_token_ids_by_step = accepted_token_ids_by_step.tolist()
470 -
471 - sampler_output_list = []
472 - for step_index in range(num_steps):
473 - if all(token_id == -1
474 - for token_id in accepted_token_ids_by_step[step_index]):
475 - break
476 - step_output_token_ids = []
477 - for sequence_index in range(batch_size):
478 - token_id = accepted_token_ids_by_step[step_index][
479 - sequence_index]
480 - step_output_token_ids.append(
481 - CompletionSequenceGroupOutput(samples=[
482 - SequenceOutput(parent_seq_id=seq_ids[sequence_index],
483 - output_token=token_id,
484 - logprobs={token_id: Logprob(token_id)})
485 - ],
486 - prompt_logprobs=None))
487 - sampler_output_list.append(
488 - SamplerOutput(outputs=step_output_token_ids))
489 - return sampler_output_list
490 -
491 - def load_weights(self, model_name_or_path: str,
492 - draft_model_name_or_path: str, **kwargs):
493 - arch = _get_model_architecture(self.config)
494 - neuronx_module_path, neuronx_model_cls_name = (
495 - _NEURON_SUPPORTED_MODELS[arch])
496 - neuronx_module = importlib.import_module(neuronx_module_path)
497 - neuronx_model_cls = getattr(neuronx_module, neuronx_model_cls_name)
498 - neuron_config = neuronx_model_cls.get_neuron_config_cls()(
499 - **kwargs['neuron_config'])
500 - config = neuronx_model_cls.get_config_cls()(
501 - neuron_config,
502 - load_config=load_pretrained_config(model_name_or_path))
503 -
504 - draft_neuron_config = copy.deepcopy(config.neuron_config)
505 - if not config.neuron_config.enable_eagle_speculation:
506 - draft_neuron_config.speculation_length = 0
507 - draft_neuron_config.trace_tokengen_model = True
508 - draft_neuron_config.enable_fused_speculation = False
509 - if getattr(config.neuron_config, "draft_model_modules_to_not_convert",
510 - None):
511 - draft_neuron_config.modules_to_not_convert = (
512 - draft_neuron_config.draft_model_modules_to_not_convert)
513 - if config.neuron_config.enable_eagle_speculation:
514 - draft_neuron_config.is_eagle_draft = True
515 - draft_neuron_config.sequence_parallel_enabled = False
516 - draft_config = neuronx_model_cls.get_config_cls()(
517 - draft_neuron_config,
518 - load_config=load_pretrained_config(draft_model_name_or_path))
519 - fused_spec_config = (FusedSpecNeuronConfig(
520 - neuronx_model_cls._model_cls,
521 - draft_config=draft_config,
522 - draft_model_path=draft_model_name_or_path))
523 - config.fused_spec_config = fused_spec_config
524 - self.config.neuron_config = neuron_config
525 -
526 - hashed_config = hashlib.md5(config.to_json_string().encode('utf-8'),
527 - usedforsecurity=False).hexdigest()
528 - if os.getenv("NEURON_COMPILED_ARTIFACTS") is not None:
529 - compiled_model_path = os.getenv("NEURON_COMPILED_ARTIFACTS")
530 - elif os.path.exists(model_name_or_path):
531 - compiled_model_path = os.path.join(model_name_or_path,
532 - "neuron-compiled-artifacts",
533 - hashed_config)
534 - shutil.rmtree(compiled_model_path, ignore_errors=True)
535 - else:
536 - compiled_model_path = os.path.join("local-models",
537 - model_name_or_path,
538 - "neuron-compiled-artifacts",
539 - hashed_config)
540 - shutil.rmtree(compiled_model_path, ignore_errors=True)
541 - try:
542 - self.model = neuronx_model_cls(compiled_model_path)
543 - override_neuron_config = kwargs["override_neuron_config"]
544 - for k, v in override_neuron_config.items():
545 - setattr(self.model.config.neuron_config, k, v)
546 - self.model.load(compiled_model_path)
547 - return
548 - except (FileNotFoundError, ValueError) as e:
549 - logger.warning("Exception: %s", e)
550 - logger.warning("Failed to load the model from %s Recompiling...",
551 - compiled_model_path)
552 - if not os.path.exists(model_name_or_path):
553 - hf_model = AutoModelForCausalLM.from_pretrained(model_name_or_path)
554 - saved_path = os.path.join("local-models", model_name_or_path)
555 - hf_model.save_pretrained(saved_path)
556 - model_name_or_path = saved_path
557 - if not os.path.exists(draft_model_name_or_path):
558 - if draft_model_name_or_path != model_name_or_path:
559 - hf_model = AutoModelForCausalLM.from_pretrained(
560 - draft_model_name_or_path)
561 - saved_path = os.path.join("local-models",
562 - draft_model_name_or_path)
563 - hf_model.save_pretrained(saved_path)
564 - draft_model_name_or_path = saved_path
565 - else:
566 - draft_model_name_or_path = model_name_or_path
567 - config.fused_spec_config.draft_model_path = draft_model_name_or_path
568 - self.model = neuronx_model_cls(model_name_or_path, config)
569 - self.model.compile(compiled_model_path)
570 - self.model.load(compiled_model_path)
571 -
572 -
573 -def _get_model_architecture(config: PretrainedConfig) -> str:
574 - architectures = getattr(config, "architectures", [])
575 - for arch in architectures:
576 - if arch in _NEURON_SUPPORTED_MODELS:
577 - return arch
578 - raise ValueError(
579 - f"Model architectures {architectures} are not supported on Neuron "
580 - f"for now. Supported architectures: "
581 - f"{list(_NEURON_SUPPORTED_MODELS.keys())}")
582 -
583 -
584 -def _get_default_neuron_config(model_config: ModelConfig,
585 - parallel_config: ParallelConfig,
586 - scheduler_config: SchedulerConfig,
587 - lora_serving_config: LoraServingConfig):
588 - """Generate a neuron config based on vllm config args."""
589 - on_device_sampling_config = OnDeviceSamplingConfig(dynamic=True,
590 - deterministic=False)
591 - batch_size = scheduler_config.max_num_seqs
592 -
593 - neuron_config = dict(
594 - tp_degree=parallel_config.tensor_parallel_size,
595 - ctx_batch_size=1,
596 - batch_size=batch_size,
597 - max_context_length=scheduler_config.max_model_len,
598 - seq_len=scheduler_config.max_model_len,
599 - enable_bucketing=True,
600 - is_continuous_batching=True,
601 - quantized=False,
602 - torch_dtype=TORCH_DTYPE_TO_NEURON_AMP[model_config.dtype],
603 - padding_side="right",
604 - on_device_sampling_config=on_device_sampling_config,
605 - sequence_parallel_enabled=True,
606 - lora_serving_config=lora_serving_config)
607 - return neuron_config
608 -
609 -
610 -def _get_default_speculation_config(model_config: ModelConfig,
611 - parallel_config: ParallelConfig,
612 - scheduler_config: SchedulerConfig,
613 - speculation_config: SpeculativeConfig):
614 - """Generate a neuron config for speculative decoding based on vllm config
615 - args."""
616 - neuron_config = dict(
617 - tp_degree=parallel_config.tensor_parallel_size,
618 - ctx_batch_size=1,
619 - batch_size=scheduler_config.max_num_seqs,
620 - max_context_length=scheduler_config.max_model_len,
621 - seq_len=scheduler_config.max_model_len,
622 - speculation_length=speculation_config.num_speculative_tokens,
623 - trace_tokengen_model=False,
624 - enable_fused_speculation=True,
625 - enable_bucketing=True,
626 - is_continuous_batching=True,
627 - quantized=False,
628 - torch_dtype=TORCH_DTYPE_TO_NEURON_AMP[model_config.dtype],
629 - on_device_sampling_config=dict(
630 - top_k=1,
631 - do_sample=False,
632 - ))
633 - return neuron_config
634 -
635 -
636 -def _get_neuron_config_after_override(default_neuron_config,
637 - overridden_neuron_config):
638 - """Update default neuron config values with override args"""
639 - overridden_neuron_config = overridden_neuron_config or {}
640 - default_neuron_config.update(overridden_neuron_config)
641 - return default_neuron_config
642 -
643 -
644 -def get_neuron_model(model_config: ModelConfig,
645 - parallel_config: ParallelConfig,
646 - scheduler_config: SchedulerConfig,
647 - lora_serving_config: LoraServingConfig) -> nn.Module:
648 - """Initializes a neuron-optimized model for inference."""
649 - model_arch = _get_model_architecture(model_config.hf_config)
650 - if model_arch == "MllamaForConditionalGeneration":
651 - model = NeuronMllamaForCausalLM(model_config.hf_config)
652 - else:
653 - model = NeuronCausalLM(model_config.hf_config)
654 - default_neuron_config_args = _get_default_neuron_config(
655 - model_config, parallel_config, scheduler_config, lora_serving_config)
656 - neuron_config = _get_neuron_config_after_override(
657 - default_neuron_config_args, model_config.override_neuron_config)
658 -
659 - override_neuron_config = model_config.override_neuron_config
660 - model.load_weights(model_config.model,
661 - neuron_config=neuron_config,
662 - override_neuron_config=override_neuron_config)
663 - return model.eval()
664 -
665 -
666 -def get_neuron_speculation_model(model_config: ModelConfig,
667 - parallel_config: ParallelConfig,
668 - scheduler_config: SchedulerConfig,
669 - speculation_config: SpeculativeConfig):
670 - """Initializes a neuron-optimized speculation model for inference.
671 -
672 - This model handles speculation using both a draft model and an EAGLE draft.
673 - """
674 - model = NeuronSpeculationCausalLM(model_config.hf_config)
675 - default_neuron_config_args = _get_default_speculation_config(
676 - model_config, parallel_config, scheduler_config, speculation_config)
677 - neuron_config = _get_neuron_config_after_override(
678 - default_neuron_config_args, model_config.override_neuron_config)
679 -
680 - override_neuron_config = model_config.override_neuron_config
681 - model.load_weights(model_config.model,
682 - speculation_config.draft_model_config.model,
683 - neuron_config=neuron_config,
684 - override_neuron_config=override_neuron_config)
685 - return model.eval()
vllm/platforms/__init__.py
-25
@@ -169,37 +169,12 @@ def cpu_platform_plugin() -> Optional[str]:
169 return "vllm.platforms.cpu.CpuPlatform" if is_cpu else None
170
171
172 -def neuron_platform_plugin() -> Optional[str]:
173 - tnx_installed = False
174 - nxd_installed = False
175 - logger.debug("Checking if Neuron platform is available.")
176 - try:
177 - import transformers_neuronx # noqa: F401
178 - tnx_installed = True
179 - logger.debug("Confirmed Neuron platform is available because"
180 - " transformers_neuronx is found.")
181 - except ImportError:
182 - pass
183 -
184 - try:
185 - import neuronx_distributed_inference # noqa: F401
186 - nxd_installed = True
187 - logger.debug("Confirmed Neuron platform is available because"
188 - " neuronx_distributed_inference is found.")
189 - except ImportError:
190 - pass
191 -
192 - is_neuron = tnx_installed or nxd_installed
193 - return "vllm.platforms.neuron.NeuronPlatform" if is_neuron else None
194 -
195 -
172 builtin_platform_plugins = {
173 'tpu': tpu_platform_plugin,
174 'cuda': cuda_platform_plugin,
175 'rocm': rocm_platform_plugin,
176 'xpu': xpu_platform_plugin,
177 'cpu': cpu_platform_plugin,
202 - 'neuron': neuron_platform_plugin,
178 }
179
180
vllm/platforms/interface.py
-4
@@ -73,7 +73,6 @@ class PlatformEnum(enum.Enum):
73 TPU = enum.auto()
74 XPU = enum.auto()
75 CPU = enum.auto()
76 - NEURON = enum.auto()
76 OOT = enum.auto()
77 UNSPECIFIED = enum.auto()
78
@@ -164,9 +163,6 @@ class Platform:
163 def is_cpu(self) -> bool:
164 return self._enum == PlatformEnum.CPU
165
167 - def is_neuron(self) -> bool:
168 - return self._enum == PlatformEnum.NEURON
169 -
166 def is_out_of_tree(self) -> bool:
167 return self._enum == PlatformEnum.OOT
168
vllm/platforms/neuron.py deleted
-151
@@ -1,151 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -import enum
4 -import os
5 -from functools import lru_cache
6 -from typing import TYPE_CHECKING, Optional
7 -
8 -from vllm import envs
9 -from vllm.logger import init_logger
10 -from vllm.utils import DEFAULT_MAX_NUM_BATCHED_TOKENS
11 -
12 -from .interface import Platform, PlatformEnum
13 -
14 -if TYPE_CHECKING:
15 - from vllm.config import VllmConfig
16 -else:
17 - VllmConfig = None
18 -
19 -logger = init_logger(__name__)
20 -
21 -
22 -class NeuronFramework(enum.Enum):
23 - TRANSFORMERS_NEURONX = "transformers-neuronx"
24 - NEURONX_DISTRIBUTED_INFERENCE = "neuronx-distributed-inference"
25 -
26 -
27 -class NeuronPlatform(Platform):
28 - _enum = PlatformEnum.NEURON
29 - device_name: str = "neuron"
30 - device_type: str = "neuron"
31 - ray_device_key: str = "neuron_cores"
32 - supported_quantization: list[str] = ["neuron_quant", "fbgemm_fp8"]
33 - dist_backend: str = "gloo"
34 - device_control_env_var: str = "NEURON_RT_VISIBLE_CORES"
35 -
36 - @classmethod
37 - def get_device_name(cls, device_id: int = 0) -> str:
38 - return "neuron"
39 -
40 - @classmethod
41 - def is_async_output_supported(cls, enforce_eager: Optional[bool]) -> bool:
42 - return False
43 -
44 - @classmethod
45 - def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
46 - parallel_config = vllm_config.parallel_config
47 - if parallel_config.worker_cls == "auto":
48 - parallel_config.worker_cls = \
49 - "vllm.worker.neuron_worker.NeuronWorker"
50 -
51 - if parallel_config.world_size > 1:
52 - parallel_config.distributed_executor_backend = "uni"
53 -
54 - if vllm_config.cache_config and vllm_config.model_config:
55 - # neuron needs block_size = max_model_len
56 - vllm_config.cache_config.block_size = \
57 - vllm_config.model_config.max_model_len # type: ignore
58 -
59 - if vllm_config.model_config and vllm_config.model_config.use_mla:
60 - logger.info(
61 - "MLA is enabled on a non-GPU platform; forcing chunked "
62 - "prefill and prefix caching to be disabled.")
63 - vllm_config.scheduler_config.enable_chunked_prefill = False
64 - vllm_config.scheduler_config.chunked_prefill_enabled = False
65 - vllm_config.scheduler_config.max_num_batched_tokens = max(
66 - vllm_config.scheduler_config.max_model_len,
67 - DEFAULT_MAX_NUM_BATCHED_TOKENS)
68 -
69 - @classmethod
70 - def is_pin_memory_available(cls) -> bool:
71 - logger.warning("Pin memory is not supported on Neuron.")
72 - return False
73 -
74 - @classmethod
75 - def get_device_communicator_cls(cls) -> str:
76 - if envs.VLLM_USE_V1:
77 - return "vllm.distributed.device_communicators.neuron_communicator.NeuronCommunicator" # noqa
78 - else:
79 - return Platform.get_device_communicator_cls()
80 -
81 - @classmethod
82 - def use_all_gather(cls) -> bool:
83 - return True
84 -
85 - @classmethod
86 - @lru_cache
87 - def is_neuronx_distributed_inference(cls) -> bool:
88 - try:
89 - import neuronx_distributed_inference
90 - except ImportError:
91 - neuronx_distributed_inference = None
92 - return neuronx_distributed_inference is not None
93 -
94 - @classmethod
95 - @lru_cache
96 - def is_transformers_neuronx(cls) -> bool:
97 - try:
98 - import transformers_neuronx
99 - except ImportError:
100 - transformers_neuronx = None
101 - return transformers_neuronx is not None
102 -
103 - def get_neuron_framework_to_use(self):
104 - """Return the specified framework if corresponding installations are
105 - available.
106 -
107 - If no framework is specified, use neuronx-distributed-inference by
108 - default.
109 - If that's unavailable, check and switch to transformers-neuronx.
110 - """
111 - if not self.is_neuron():
112 - raise AssertionError(
113 - f"Neuron Framework unavailable for platform: {self}")
114 -
115 - tnx_installed = self.is_transformers_neuronx()
116 - nxd_installed = self.is_neuronx_distributed_inference()
117 -
118 - specified_framework = os.environ.get("VLLM_NEURON_FRAMEWORK")
119 - tnx_framework = NeuronFramework.TRANSFORMERS_NEURONX.value
120 - nxd_framework = NeuronFramework.NEURONX_DISTRIBUTED_INFERENCE.value
121 - if specified_framework == tnx_framework and tnx_installed:
122 - return self.TRANSFORMERS_NEURONX
123 -
124 - if ((specified_framework == nxd_framework and nxd_installed)
125 - or (specified_framework is None and nxd_installed)):
126 - return NeuronFramework.NEURONX_DISTRIBUTED_INFERENCE
127 -
128 - if specified_framework is None and tnx_installed:
129 - return NeuronFramework.TRANSFORMERS_NEURONX
130 -
131 - return None
132 -
133 - def use_neuronx_distributed(self):
134 - """
135 - Return True if the framework determined in get_neuron_framework_to_use()
136 - is NeuronFramework.NEURONX_DISTRIBUTED_INFERENCE, False otherwise. This
137 - is used to select the Neuron model framework and framework-specific
138 - configuration to apply during model compilation.
139 - """
140 - nxd_framework = NeuronFramework.NEURONX_DISTRIBUTED_INFERENCE
141 - return self.get_neuron_framework_to_use() == nxd_framework
142 -
143 - def use_transformers_neuronx(self):
144 - """
145 - Return True if the framework determined in get_neuron_framework_to_use()
146 - is NeuronFramework.TRANSFORMERS_NEURONX, False otherwise. This is used
147 - to select the Neuron model framework and framework-specific
148 - configuration to apply during model compilation.
149 - """
150 - return self.get_neuron_framework_to_use(
151 - ) == NeuronFramework.TRANSFORMERS_NEURONX
vllm/worker/neuron_model_runner.py deleted
-455
@@ -1,455 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -import os
5 -from dataclasses import dataclass
6 -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Tuple, Union
7 -
8 -import torch
9 -from torch import nn
10 -
11 -from vllm.config import DeviceConfig, VllmConfig
12 -from vllm.logger import init_logger
13 -from vllm.lora.layers import LoRAMapping
14 -from vllm.lora.request import LoRARequest
15 -from vllm.model_executor import SamplingMetadata
16 -from vllm.model_executor.layers.sampler import SamplerOutput
17 -from vllm.model_executor.model_loader.neuron import get_neuron_model
18 -from vllm.multimodal import BatchedTensorInputs, MultiModalKwargs
19 -from vllm.platforms import current_platform
20 -from vllm.sampling_params import SamplingParams
21 -from vllm.sequence import IntermediateTensors, SequenceGroupMetadata
22 -from vllm.utils import is_pin_memory_available, make_tensor_with_pad
23 -from vllm.worker.model_runner_base import ModelRunnerBase, ModelRunnerInputBase
24 -
25 -if TYPE_CHECKING:
26 - from vllm.attention.backends.abstract import AttentionBackend
27 -
28 -logger = init_logger(__name__)
29 -
30 -
31 -@dataclass(frozen=True)
32 -class ModelInputForNeuron(ModelRunnerInputBase):
33 - """
34 - Used by the NeuronModelRunner.
35 - """
36 - input_tokens: Optional[torch.Tensor] = None
37 - input_positions: Optional[torch.Tensor] = None
38 - input_block_ids: Optional[torch.Tensor] = None
39 - sampling_metadata: SamplingMetadata = None
40 - multi_modal_kwargs: BatchedTensorInputs = None
41 - adapter_ids: Optional[str] = None
42 -
43 - def as_broadcastable_tensor_dict(
44 - self) -> Dict[str, Union[int, torch.Tensor]]:
45 - return {
46 - "input_tokens": self.input_tokens,
47 - "input_positions": self.input_positions,
48 - "input_block_ids": self.input_block_ids,
49 - "sampling_metadata": self.sampling_metadata,
50 - "multi_modal_kwargs": self.multi_modal_kwargs,
51 - }
52 -
53 - @classmethod
54 - def from_broadcasted_tensor_dict(
55 - cls,
56 - tensor_dict: Dict[str, Any],
57 - attn_backend: Optional["AttentionBackend"] = None,
58 - ) -> "ModelInputForNeuron":
59 - return ModelInputForNeuron(
60 - input_tokens=tensor_dict["input_tokens"],
61 - input_positions=tensor_dict["input_positions"],
62 - input_block_ids=tensor_dict["input_block_ids"],
63 - sampling_metadata=tensor_dict["sampling_metadata"],
64 - multi_modal_kwargs=tensor_dict["multi_modal_kwargs"],
65 - )
66 -
67 -
68 -class NeuronModelRunner(ModelRunnerBase[ModelInputForNeuron]):
69 - """A model runner for AWS Neuron hardware"""
70 -
71 - # NEURON has an upper limit on the top_k
72 - _MAX_NEURON_SAMPLING_TOP_K = 256
73 -
74 - def __init__(
75 - self,
76 - vllm_config: VllmConfig,
77 - ):
78 - ModelRunnerBase.__init__(self, vllm_config)
79 -
80 - if (self.model_config is not None
81 - and self.model_config.get_sliding_window()):
82 - logger.warning("Sliding window is not supported on Neuron. "
83 - "The model will run without sliding window.")
84 - self.device_config = (self.device_config if self.device_config
85 - is not None else DeviceConfig())
86 - self.lora_config = vllm_config.lora_config
87 - self.device = self.device_config.device
88 - self.pin_memory = is_pin_memory_available()
89 -
90 - # Lazy initialization.
91 - self.model: nn.Module # initialize after load_model.
92 -
93 - # Once NEURON_ON_DEVICE_SAMPLING_DISABLED is set to a non-zero value,
94 - # turn off on-device sampling.
95 - self._on_device_sampling_disabled = int(
96 - os.getenv("NEURON_ON_DEVICE_SAMPLING_DISABLED", "0"))
97 -
98 - # NEURON needs to update sampling parameters when request IDs change
99 - # across batches. This variable stores the previous batch's request IDs
100 - # to determine if an update is needed.
101 - self._previous_batch_request_ids: List[str] = []
102 -
103 - if not self._on_device_sampling_disabled:
104 - self._init_neuron_sampling()
105 -
106 - def _init_neuron_sampling(self) -> None:
107 - if current_platform.use_transformers_neuronx():
108 - from transformers_neuronx.config import GenerationConfig
109 - else:
110 - from transformers import GenerationConfig
111 - logger.warning(
112 - "On-device sampling is turned on in Neuron by default, only "
113 - "top_k, top_p, and temperature are current supported sampling "
114 - "parameters. To turn off the on-device sampling, please set "
115 - "the environment variable NEURON_ON_DEVICE_SAMPLING_DISABLED=1.")
116 - self.model_config.neuron_sampling_params = GenerationConfig(
117 - max_length=self.scheduler_config.max_model_len,
118 - do_sample=True,
119 - per_batch_line=True,
120 - top_k=[self._MAX_NEURON_SAMPLING_TOP_K] \
121 - * self.scheduler_config.max_num_seqs,
122 - top_p=[1.0] * self.scheduler_config.max_num_seqs,
123 - temperature=[1.0] * self.scheduler_config.max_num_seqs,
124 - dynamic=True,
125 - global_top_k=self._MAX_NEURON_SAMPLING_TOP_K)
126 -
127 - def load_model(self) -> None:
128 - self.model = get_neuron_model(self.model_config,
129 - parallel_config=self.parallel_config,
130 - scheduler_config=self.scheduler_config)
131 -
132 - def get_model(self) -> nn.Module:
133 - return self.model
134 -
135 - def _prepare_prompt(
136 - self,
137 - seq_group_metadata_list: List[SequenceGroupMetadata],
138 - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, List[int],
139 - BatchedTensorInputs]:
140 - assert len(seq_group_metadata_list) > 0
141 - input_tokens: List[List[int]] = []
142 - input_positions: List[List[int]] = []
143 - input_block_ids: List[int] = []
144 -
145 - seq_lens: List[int] = []
146 - multi_modal_kwargs_list: List[MultiModalKwargs] = []
147 - for seq_group_metadata in seq_group_metadata_list:
148 - assert seq_group_metadata.is_prompt
149 - seq_ids = list(seq_group_metadata.seq_data.keys())
150 - assert len(seq_ids) == 1
151 - seq_id = seq_ids[0]
152 -
153 - seq_data = seq_group_metadata.seq_data[seq_id]
154 - prompt_tokens = seq_data.get_token_ids()
155 - seq_len = len(prompt_tokens)
156 - seq_lens.append(seq_len)
157 -
158 - input_tokens.append(prompt_tokens)
159 - input_positions.append(list(range(seq_len)))
160 -
161 - assert seq_group_metadata.block_tables is not None
162 - block_table = seq_group_metadata.block_tables[seq_id]
163 - assert len(block_table) == 1
164 - input_block_ids.append(block_table[0])
165 -
166 - mm_kwargs = seq_group_metadata.multi_modal_data
167 - if mm_kwargs:
168 - mm_kwargs = self.process_multi_modal_data_neuron(mm_kwargs)
169 - multi_modal_kwargs_list.append(mm_kwargs)
170 -
171 - max_seq_len = max(seq_lens)
172 - assert max_seq_len > 0
173 - input_tokens = make_tensor_with_pad(input_tokens,
174 - pad=0,
175 - max_len=max_seq_len,
176 - dtype=torch.long,
177 - device=self.device)
178 - input_positions = make_tensor_with_pad(input_positions,
179 - pad=0,
180 - max_len=max_seq_len,
181 - dtype=torch.long,
182 - device=self.device)
183 - input_block_ids = torch.tensor(input_block_ids,
184 - dtype=torch.long,
185 - device=self.device)
186 -
187 - multi_modal_kwargs = MultiModalKwargs.batch(multi_modal_kwargs_list)
188 -
189 - return (input_tokens, input_positions, input_block_ids, seq_lens,
190 - multi_modal_kwargs)
191 -
192 - def _prepare_decode(
193 - self,
194 - seq_group_metadata_list: List[SequenceGroupMetadata],
195 - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
196 - assert len(seq_group_metadata_list) > 0
197 - input_tokens: List[List[int]] = []
198 - input_positions: List[List[int]] = []
199 - input_block_ids: List[int] = []
200 - context_lens: List[int] = []
201 -
202 - for seq_group_metadata in seq_group_metadata_list:
203 - assert not seq_group_metadata.is_prompt
204 -
205 - seq_ids = list(seq_group_metadata.seq_data.keys())
206 -
207 - for seq_id in seq_ids:
208 - seq_data = seq_group_metadata.seq_data[seq_id]
209 - generation_token = seq_data.get_last_token_id()
210 - input_tokens.append([generation_token])
211 -
212 - seq_len = seq_data.get_len()
213 - position = seq_len - 1
214 - input_positions.append([position])
215 - context_lens.append(seq_len)
216 -
217 - assert seq_group_metadata.block_tables is not None
218 - block_table = seq_group_metadata.block_tables[seq_id]
219 - assert len(block_table) == 1
220 - input_block_ids.append(block_table[0])
221 -
222 - input_tokens = make_tensor_with_pad(input_tokens,
223 - pad=0,
224 - max_len=1,
225 - dtype=torch.long,
226 - device=self.device)
227 - input_positions = make_tensor_with_pad(input_positions,
228 - pad=0,
229 - max_len=1,
230 - dtype=torch.long,
231 - device=self.device)
232 - context_lens = torch.tensor(context_lens,
233 - dtype=torch.int,
234 - device=self.device)
235 - input_block_ids = torch.tensor(input_block_ids,
236 - dtype=torch.long,
237 - device=self.device)
238 -
239 - return input_tokens, input_positions, input_block_ids
240 -
241 - def make_model_input_from_broadcasted_tensor_dict(
242 - self, tensor_dict: Dict[str, Any]) -> ModelInputForNeuron:
243 - return ModelInputForNeuron.from_broadcasted_tensor_dict(tensor_dict)
244 -
245 - def prepare_model_input(
246 - self,
247 - seq_group_metadata_list: List[SequenceGroupMetadata],
248 - virtual_engine: int = 0,
249 - finished_requests_ids: Optional[List[str]] = None
250 - ) -> ModelInputForNeuron:
251 - multi_modal_kwargs = None
252 - # NOTE: We assume that all sequences in the group are all prompts or
253 - # all decodes.
254 - is_prompt = seq_group_metadata_list[0].is_prompt
255 - # Prepare input tensors.
256 - if is_prompt:
257 - (input_tokens, input_positions, input_block_ids, seq_lens,
258 - multi_modal_kwargs
259 - ) = self._prepare_prompt(seq_group_metadata_list)
260 - else:
261 - (input_tokens, input_positions,
262 - input_block_ids) = self._prepare_decode(seq_group_metadata_list)
263 - seq_lens = None
264 -
265 - if not self._on_device_sampling_disabled:
266 - for seq_group_metadata in seq_group_metadata_list:
267 - sampling_params = seq_group_metadata.sampling_params
268 - top_k, top_p, temperature = (
269 - self._convert_to_neuron_sampling_params(sampling_params))
270 - sampling_params.top_k = top_k
271 - sampling_params.top_p = top_p
272 - sampling_params.temperature = temperature
273 -
274 - # we need multi_modal_data for later tokens as well
275 - multi_modal_kwargs_list: List[MultiModalKwargs] = []
276 - for seq_group_metadata in seq_group_metadata_list:
277 - mm_data = seq_group_metadata.multi_modal_data
278 - if mm_data:
279 - multi_modal_kwargs_list.append(mm_data)
280 - multi_modal_kwargs = MultiModalKwargs.batch(multi_modal_kwargs_list)
281 -
282 - sampling_metadata = SamplingMetadata.prepare(
283 - seq_group_metadata_list,
284 - seq_lens,
285 - # query_lens is not needed if chunked prefill is not
286 - # supported. Since neuron worker doesn't support chunked prefill
287 - # just use seq_lens instead.
288 - seq_lens,
289 - self.device,
290 - self.pin_memory,
291 - generators=self.get_generators(finished_requests_ids))
292 -
293 - if current_platform.use_transformers_neuronx(
294 - ) and not self._on_device_sampling_disabled:
295 - # Once the request IDs are changed in current iteration, we will
296 - # update the on-device sampling parameters.
297 - current_batch_request_ids = [
298 - seq_group_meta_data.request_id
299 - for seq_group_meta_data in seq_group_metadata_list
300 - ]
301 - if current_batch_request_ids != self._previous_batch_request_ids:
302 - self._update_neuron_sampling_params(seq_group_metadata_list)
303 - self._previous_batch_request_ids = current_batch_request_ids
304 -
305 - return ModelInputForNeuron(input_tokens=input_tokens,
306 - input_positions=input_positions,
307 - input_block_ids=input_block_ids,
308 - sampling_metadata=sampling_metadata,
309 - multi_modal_kwargs=multi_modal_kwargs)
310 -
311 - def _update_neuron_sampling_params(
312 - self, seq_group_metadata_list: List[SequenceGroupMetadata]):
313 - # Update Neuron sampling parameters (GenerationConfig in Neuron)
314 - current_sampling_params = self.model_config.neuron_sampling_params
315 - assert current_sampling_params is not None, (
316 - f"Failed to update sampling_params, "
317 - f"current sampling params is {current_sampling_params}")
318 -
319 - is_update_needed = False
320 -
321 - top_k = current_sampling_params.top_k
322 - top_p = current_sampling_params.top_p
323 - temperature = current_sampling_params.temperature
324 -
325 - # The index of a sequence's sampling parameters in neuron is equal to
326 - # its index in `input_block_ids`.
327 - for seq_group_metadata in seq_group_metadata_list:
328 - seq_ids = list(seq_group_metadata.seq_data.keys())
329 - sampling_params = seq_group_metadata.sampling_params
330 -
331 - seq_group_top_k = sampling_params.top_k
332 - seq_group_top_p = sampling_params.top_p
333 - seq_group_temperature = sampling_params.temperature
334 -
335 - for seq_id in seq_ids:
336 - index = seq_group_metadata.block_tables[seq_id][0]
337 - if (top_k[index] != seq_group_top_k
338 - or top_p[index] != seq_group_top_p
339 - or temperature[index] != seq_group_temperature):
340 - is_update_needed = True
341 -
342 - top_k[index] = seq_group_top_k
343 - top_p[index] = seq_group_top_p
344 - temperature[index] = seq_group_temperature
345 -
346 - # update_generation_config is only available in transformers-neuronx
347 - if is_update_needed and current_platform.use_transformers_neuronx():
348 - self.model.model.update_generation_config(current_sampling_params)
349 -
350 - def _convert_to_neuron_sampling_params(
351 - self, sampling_params: SamplingParams) -> Tuple[int, float, float]:
352 - # Returns the top_k, top_p and temperature parameters for neuron.
353 - top_k = sampling_params.top_k
354 - top_p = sampling_params.top_p
355 - temperature = sampling_params.temperature
356 -
357 - if temperature == 0.0:
358 - # Enable greedy sampling on zero temperature
359 - return (1, 1.0, 1.0)
360 - if top_k < 1 or top_k > self._MAX_NEURON_SAMPLING_TOP_K:
361 - top_k = self._MAX_NEURON_SAMPLING_TOP_K
362 -
363 - return (top_k, top_p, temperature)
364 -
365 - @torch.inference_mode()
366 - def execute_model(
367 - self,
368 - model_input: ModelInputForNeuron,
369 - kv_caches: Optional[List[torch.Tensor]] = None,
370 - intermediate_tensors: Optional[IntermediateTensors] = None,
371 - num_steps: int = 1,
372 - ) -> Optional[List[SamplerOutput]]:
373 - if num_steps > 1:
374 - raise ValueError(
375 - "NeuronModelRunner does not support multi-step execution.")
376 -
377 - # extract top_k, top_p and temperature from model_input for neuron
378 - # forward call
379 - sampling_params = (torch.tensor([[
380 - seq_group.sampling_params.top_k, seq_group.sampling_params.top_p,
381 - seq_group.sampling_params.temperature
382 - ] for seq_group in model_input.sampling_metadata.seq_groups]))
383 -
384 - if current_platform.use_neuronx_distributed():
385 - hidden_states = self.model(
386 - input_ids=model_input.input_tokens,
387 - positions=model_input.input_positions,
388 - input_block_ids=model_input.input_block_ids,
389 - sampling_params=sampling_params,
390 - adapter_ids=model_input.adapter_ids,
391 - **MultiModalKwargs.as_kwargs(
392 - model_input.multi_modal_kwargs or {},
393 - device=self.device,
394 - ),
395 - )
396 - elif current_platform.use_transformers_neuronx():
397 - # [TODO] validate on-device sampling
398 - # The model signature may need change for on-device sampling
399 - hidden_states = self.model(
400 - input_ids=model_input.input_tokens,
401 - positions=model_input.input_positions,
402 - input_block_ids=model_input.input_block_ids,
403 - **MultiModalKwargs.as_kwargs(
404 - model_input.multi_modal_kwargs or {},
405 - device=self.device,
406 - ),
407 - )
408 -
409 - # Compute the logits only if the on-device sampling is turned off as
410 - # on-device sampling outputs the token ids.
411 - if self._on_device_sampling_disabled:
412 - logits = self.model.compute_logits(hidden_states,
413 - model_input.sampling_metadata)
414 - else:
415 - logits = hidden_states
416 -
417 - # Sample the next token.
418 - output = self.model.sample(
419 - logits=logits,
420 - sampling_metadata=model_input.sampling_metadata,
421 - )
422 - return [output]
423 -
424 - @property
425 - def vocab_size(self) -> int:
426 - return self.model_config.get_vocab_size()
427 -
428 - def process_multi_modal_data_neuron(self, mm_data):
429 - # this is a no-op for NeuronModelRunner
430 - return mm_data
431 -
432 - def remove_all_loras(self):
433 - raise NotImplementedError(
434 - "LoRAs are not supported for Transformers NeuronX framework")
435 -
436 - def set_active_loras(self, lora_requests: Set[LoRARequest],
437 - lora_mapping: LoRAMapping) -> None:
438 - raise NotImplementedError(
439 - "LoRAs are not supported for Transformers NeuronX framework")
440 -
441 - def add_lora(self, lora_request: LoRARequest):
442 - raise NotImplementedError(
443 - "LoRAs are not supported for Transformers NeuronX framework")
444 -
445 - def remove_lora(self, lora_id: int) -> bool:
446 - raise NotImplementedError(
447 - "LoRAs are not supported for Transformers NeuronX framework")
448 -
449 - def pin_lora(self, lora_id: int) -> bool:
450 - raise NotImplementedError(
451 - "LoRAs are not supported for Transformers NeuronX framework")
452 -
453 - def list_loras(self) -> Set[int]:
454 - raise NotImplementedError(
455 - "LoRAs are not supported for Transformers NeuronX framework")
vllm/worker/neuron_worker.py deleted
-189
@@ -1,189 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""A Neuron worker class."""
4 -import os
5 -from typing import List, Optional, Set, Tuple
6 -
7 -import torch.distributed
8 -
9 -from vllm.config import VllmConfig
10 -from vllm.distributed import (ensure_model_parallel_initialized,
11 - init_distributed_environment)
12 -from vllm.logger import init_logger
13 -from vllm.lora.request import LoRARequest
14 -from vllm.model_executor import set_random_seed
15 -from vllm.platforms import current_platform
16 -from vllm.platforms.neuron import NeuronFramework
17 -from vllm.sequence import ExecuteModelRequest
18 -from vllm.worker.neuron_model_runner import NeuronModelRunner
19 -from vllm.worker.worker_base import (LocalOrDistributedWorkerBase, WorkerBase,
20 - WorkerInput)
21 -
22 -logger = init_logger(__name__)
23 -
24 -
25 -class NeuronWorker(LocalOrDistributedWorkerBase):
26 - """A worker class that executes the model on a group of neuron cores.
27 - """
28 -
29 - model_runner: NeuronModelRunner
30 -
31 - def __init__(self,
32 - vllm_config: VllmConfig,
33 - local_rank: int,
34 - rank: int,
35 - distributed_init_method: str,
36 - is_driver_worker: bool = False) -> None:
37 - WorkerBase.__init__(self, vllm_config=vllm_config)
38 - self.local_rank = local_rank
39 - self.rank = rank
40 - self.distributed_init_method = distributed_init_method
41 - self.is_driver_worker = is_driver_worker
42 - self.lora_config = vllm_config.lora_config
43 -
44 - if self.model_config.trust_remote_code:
45 - # note: lazy import to avoid importing torch before initializing
46 - from vllm.utils import init_cached_hf_modules
47 - init_cached_hf_modules()
48 -
49 - neuron_framework = current_platform.get_neuron_framework_to_use()
50 - if neuron_framework == NeuronFramework.TRANSFORMERS_NEURONX:
51 - self.model_runner = self.get_tnx_model_runner(vllm_config)
52 - elif neuron_framework == NeuronFramework.NEURONX_DISTRIBUTED_INFERENCE:
53 - self.model_runner = self.get_neuronx_distributed_model_runner(
54 - vllm_config)
55 - else:
56 - raise NotImplementedError(
57 - "Specified framework" +
58 - f" {os.environ.get('VLLM_NEURON_FRAMEWORK')}" +
59 - " is either not installed or not supported." +
60 - " Supported frameworks: " +
61 - "[transformers-neuronx, neuronx-distributed-inference]")
62 -
63 - def get_tnx_model_runner(self, vllm_config):
64 - assert (self.lora_config
65 - is None), ("LoRA is not supported for TransformersNeuronX "
66 - "framework.")
67 - if self.speculative_config is not None:
68 - raise NotImplementedError(
69 - "Speculative decoding is not supported for TransformersNeuronX"
70 - )
71 - return NeuronModelRunner(vllm_config=vllm_config)
72 -
73 - def get_neuronx_distributed_model_runner(self, vllm_config):
74 - from vllm.worker.neuronx_distributed_model_runner import (
75 - NeuronxDistributedModelRunner)
76 - if self.speculative_config is not None:
77 - assert (self.lora_config is None), (
78 - "LoRA is not supported for Speculative Decoding")
79 - raise NotImplementedError(
80 - "Speculative decoding is not supported for NeuronxDistributed")
81 - return NeuronxDistributedModelRunner(vllm_config=vllm_config)
82 -
83 - def init_device(self) -> None:
84 - self.init_distributed_environment()
85 -
86 - # Set random seed.
87 - set_random_seed(self.model_config.seed)
88 -
89 - def load_model(self):
90 - self.model_runner.load_model()
91 -
92 - def determine_num_available_blocks(self) -> Tuple[int, int]:
93 - """Determine the number of available KV blocks.
94 -
95 - Swapping is not yet supported, so always return num_cpu_blocks=0.
96 -
97 - We configure num_gpu_blocks to be equal to max_num_seqs.
98 - """
99 - # Set the number of GPU blocks to be the same as the maximum number of
100 - # sequences that can be processed in a single batch. This is equivalent
101 - # to schedule without PagedAttention.
102 - num_gpu_blocks = self.scheduler_config.max_num_seqs + 1
103 -
104 - # Swap not yet supported with Neuron backend.
105 - num_cpu_blocks = 0
106 -
107 - return num_gpu_blocks, num_cpu_blocks
108 -
109 - def initialize_cache(self, num_gpu_blocks: int,
110 - num_cpu_blocks: int) -> None:
111 - """Initialize the KV cache.
112 - """
113 -
114 - # Different values are not tested.
115 - assert num_cpu_blocks == 0
116 - assert num_gpu_blocks == self.scheduler_config.max_num_seqs + 1
117 -
118 - self.cache_config.num_gpu_blocks = num_gpu_blocks
119 - self.cache_config.num_cpu_blocks = num_cpu_blocks
120 -
121 - @property
122 - def do_metadata_broadcast(self) -> bool:
123 - return False
124 -
125 - @property
126 - def kv_cache(self) -> Optional[List[List[torch.Tensor]]]:
127 - return None
128 -
129 - @torch.inference_mode()
130 - def prepare_worker_input(
131 - self, execute_model_req: ExecuteModelRequest) -> WorkerInput:
132 - return WorkerInput(num_seq_groups=len(
133 - execute_model_req.seq_group_metadata_list), )
134 -
135 - def execute_worker(self, worker_input: WorkerInput) -> None:
136 - pass
137 -
138 - def get_cache_block_size_bytes(self) -> int:
139 - """Determine the size in bytes of a cache block.
140 -
141 - This is required for speculative decoding; it is not yet implemented.
142 - """
143 - raise NotImplementedError
144 -
145 - def init_distributed_environment(self):
146 - """Neuron uses transformers-neuronx for tensor parallelism.
147 -
148 - vLLM still needs the environment initialized when TP/PP > 1
149 - """
150 - init_distributed_environment(
151 - world_size=1,
152 - rank=self.rank,
153 - local_rank=self.local_rank,
154 - distributed_init_method=self.distributed_init_method,
155 - backend=current_platform.dist_backend,
156 - )
157 -
158 - ensure_model_parallel_initialized(
159 - 1,
160 - 1,
161 - )
162 -
163 - def add_lora(self, lora_request: LoRARequest) -> bool:
164 - if current_platform.use_transformers_neuronx():
165 - raise NotImplementedError(
166 - f"{type(self)} does not support LoRA with Neuron Framework "
167 - f"Transformers NeuronX")
168 - return self.model_runner.add_lora(lora_request)
169 -
170 - def remove_lora(self, lora_id: int) -> bool:
171 - if current_platform.use_transformers_neuronx():
172 - raise NotImplementedError(
173 - f"{type(self)} does not support LoRA with Neuron Framework "
174 - f"Transformers NeuronX")
175 - return self.model_runner.remove_lora(lora_id)
176 -
177 - def pin_lora(self, lora_id: int) -> bool:
178 - if current_platform.use_transformers_neuronx():
179 - raise NotImplementedError(
180 - f"{type(self)} does not support LoRA with Neuron Framework "
181 - f"Transformers NeuronX")
182 - return self.model_runner.pin_lora(lora_id)
183 -
184 - def list_loras(self) -> Set[int]:
185 - if current_platform.use_transformers_neuronx():
186 - raise NotImplementedError(
187 - f"{type(self)} does not support LoRA with Neuron Framework "
188 - f"Transformers NeuronX")
189 - return self.model_runner.list_loras()
vllm/worker/neuronx_distributed_model_runner.py deleted
-294
@@ -1,294 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -from typing import List, Optional, Set
5 -
6 -import torch
7 -from neuronx_distributed_inference.models.mllama.aspect_ratio_utils import (
8 - get_all_supported_aspect_ratios)
9 -from neuronx_distributed_inference.modules.generation.sampling import (
10 - prepare_sampling_params)
11 -from neuronx_distributed_inference.modules.lora_serving import (
12 - LoraCheckpoint, LoraServingConfig)
13 -
14 -from vllm.config import VllmConfig
15 -from vllm.entrypoints.openai.serving_models import LoRAModulePath
16 -from vllm.logger import init_logger
17 -from vllm.lora.layers import LoRAMapping
18 -from vllm.lora.request import LoRARequest
19 -from vllm.model_executor import SamplingMetadata
20 -from vllm.model_executor.layers.sampler import SamplerOutput
21 -from vllm.model_executor.model_loader.neuronx_distributed import (
22 - _get_model_architecture, get_neuron_model)
23 -from vllm.multimodal import MultiModalKwargs
24 -from vllm.sequence import IntermediateTensors, SequenceGroupMetadata
25 -from vllm.worker.neuron_model_runner import (ModelInputForNeuron,
26 - NeuronModelRunner)
27 -
28 -logger = init_logger(__name__)
29 -
30 -
31 -class NeuronxDistributedModelRunner(NeuronModelRunner):
32 -
33 - def __init__(
34 - self,
35 - vllm_config: VllmConfig,
36 - ):
37 - super().__init__(vllm_config)
38 - self.lora_checkpoint = None
39 - self.model = None
40 - self.lora_serving_config = None
41 -
42 - @staticmethod
43 - def _get_lora_paths_strings(lora_modules: List[LoRAModulePath]):
44 - if not lora_modules:
45 - return None
46 - return {_.get("name"): _.get("path") for _ in lora_modules}
47 -
48 - def _get_nxdi_lora_config(self):
49 - override_neuron_config = self.model_config.override_neuron_config
50 - lora_modules = override_neuron_config.pop("lora_modules", None)
51 - target_modules = override_neuron_config.pop("target_modules", None)
52 - lora_ckpt_paths = self._get_lora_paths_strings(lora_modules)
53 - if self.lora_config.max_loras < len(lora_ckpt_paths):
54 - raise ValueError(
55 - "Number of LoRAs (%s) exceeds maximum "
56 - "allowed (%s)", len(lora_ckpt_paths),
57 - self.lora_config.max_loras)
58 -
59 - return LoraServingConfig(
60 - max_loras=self.lora_config.max_loras,
61 - max_lora_rank=self.lora_config.max_lora_rank,
62 - target_modules=target_modules,
63 - lora_ckpt_paths=lora_ckpt_paths,
64 - )
65 -
66 - def load_model(self) -> None:
67 - # Update LoRA config
68 - if self.lora_config is not None:
69 - self.lora_serving_config = self._get_nxdi_lora_config()
70 - self.lora_checkpoint = LoraCheckpoint(self.lora_serving_config)
71 - self.model = get_neuron_model(
72 - self.model_config,
73 - parallel_config=self.parallel_config,
74 - scheduler_config=self.scheduler_config,
75 - lora_serving_config=self.lora_serving_config)
76 -
77 - def get_nxd_sampling_params(self, sampling_metadata):
78 - if self.model.config.neuron_config.on_device_sampling_config:
79 - max_topk = (self.model.config.neuron_config.
80 - on_device_sampling_config.global_topk)
81 - else:
82 - max_topk = self.model.config.vocab_size
83 -
84 - top_k = [1] * self.scheduler_config.max_num_seqs
85 - top_p = [1.0] * self.scheduler_config.max_num_seqs
86 - temperature = [1.0] * self.scheduler_config.max_num_seqs
87 -
88 - for index, sequenceGroupToSample in enumerate(
89 - sampling_metadata.seq_groups):
90 - top_k[index] = (sequenceGroupToSample.sampling_params.top_k
91 - if sequenceGroupToSample.sampling_params.top_k > 0
92 - else max_topk)
93 - top_p[index] = sequenceGroupToSample.sampling_params.top_p
94 - temperature[index] = (
95 - sequenceGroupToSample.sampling_params.temperature)
96 -
97 - sampling_params = prepare_sampling_params(
98 - batch_size=self.scheduler_config.max_num_seqs,
99 - top_k=top_k,
100 - top_p=top_p,
101 - temperature=temperature)
102 - return sampling_params
103 -
104 - def get_multi_modal_data_neuron(self, input_images):
105 - raise NotImplementedError("need to restore multi-modal support")
106 -
107 - @torch.inference_mode()
108 - def execute_model(
109 - self,
110 - model_input: ModelInputForNeuron,
111 - kv_caches: Optional[List[torch.Tensor]] = None,
112 - intermediate_tensors: Optional[IntermediateTensors] = None,
113 - num_steps: int = 1,
114 - ) -> Optional[List[SamplerOutput]]:
115 - if num_steps > 1:
116 - raise ValueError(
117 - "NeuronModelRunner does not support multi-step execution.")
118 -
119 - if _get_model_architecture(
120 - self.model.config) != "MllamaForConditionalGeneration":
121 - return super().execute_model(model_input, kv_caches,
122 - intermediate_tensors, num_steps)
123 -
124 - sampling_params = self.get_nxd_sampling_params(
125 - model_input.sampling_metadata)
126 -
127 - if model_input.multi_modal_kwargs.get('pixel_values') is not None:
128 - hidden_states = self.model(
129 - input_ids=model_input.input_tokens,
130 - positions=model_input.input_positions,
131 - seq_ids=model_input.input_block_ids,
132 - pixel_values=model_input.multi_modal_kwargs.get(
133 - 'pixel_values'),
134 - aspect_ratios=model_input.multi_modal_kwargs.get(
135 - 'aspect_ratios'),
136 - sampling_params=sampling_params,
137 - num_chunks=model_input.multi_modal_kwargs.get('num_chunks'),
138 - has_image=model_input.multi_modal_kwargs.get(
139 - 'has_image').squeeze(1),
140 - )
141 - else:
142 - bs = model_input.input_tokens.shape[0] if (model_input.input_tokens
143 - is not None) else 1
144 - empty_pixel_values = torch.zeros([bs, 1, 4, 3, 560, 560],
145 - dtype=torch.bfloat16)
146 - empty_aspect_ratios = torch.ones([bs, 1, 2], dtype=torch.int64)
147 - num_chunks = torch.zeros((bs, 1), dtype=torch.int32)
148 - has_image = torch.zeros([bs], dtype=torch.int32)
149 - hidden_states = self.model(
150 - input_ids=model_input.input_tokens,
151 - positions=model_input.input_positions,
152 - seq_ids=model_input.input_block_ids,
153 - pixel_values=empty_pixel_values,
154 - aspect_ratios=empty_aspect_ratios,
155 - sampling_params=sampling_params,
156 - num_chunks=num_chunks,
157 - has_image=has_image,
158 - )
159 -
160 - output = self.model.sample(
161 - hidden_states=hidden_states,
162 - sampling_metadata=model_input.sampling_metadata,
163 - )
164 -
165 - return [output]
166 -
167 - def process_multi_modal_data_neuron(self, mm_data):
168 - # Neuron uses aspect_ratios instead of aspect_ratio_ids
169 - all_supported_aspect_ratios = get_all_supported_aspect_ratios(
170 - self.model.config.vision_config.max_num_tiles)
171 - aspect_ratio_ids = mm_data.get("aspect_ratio_ids")
172 - mm_data["aspect_ratios"] = torch.tensor(
173 - all_supported_aspect_ratios[aspect_ratio_ids]).unsqueeze(0)
174 -
175 - # Neuron's num_chunks is HF's num_tiles
176 - mm_data["num_chunks"] = mm_data.get("num_tiles")
177 -
178 - # Input has an image if it has pixel_values
179 - bs = mm_data["num_chunks"].shape[0]
180 - pixel_values = mm_data.get("pixel_values")
181 - if pixel_values is not None and not torch.all(pixel_values == 0):
182 - mm_data["has_image"] = torch.ones(bs)
183 -
184 - else:
185 - mm_data["has_image"] = torch.zeros(bs)
186 - return mm_data
187 -
188 - def _get_lora_adapter_ids(self, seq_group_metadata_list):
189 - # set LoRA adapter IDs for multi-lora serving
190 - batch_size = len(seq_group_metadata_list)
191 - if self.lora_checkpoint is not None:
192 - # "0" indicates NxDI to use the base model for inference
193 - adapter_ids = ["0"] * batch_size
194 - for idx, seq_group_metadata in enumerate(seq_group_metadata_list):
195 - if seq_group_metadata.lora_request is not None:
196 - adapter_ids[
197 - idx] = seq_group_metadata.lora_request.lora_name
198 -
199 - # convert adapter_ids from strings to integers
200 - adapter_ids = self.lora_checkpoint.convert_adapter_ids_to_indices(
201 - adapter_ids, batch_size)
202 - else:
203 - adapter_ids = torch.zeros((batch_size), dtype=torch.int32)
204 -
205 - return adapter_ids
206 -
207 - def prepare_model_input(
208 - self,
209 - seq_group_metadata_list: List[SequenceGroupMetadata],
210 - virtual_engine: int = 0,
211 - finished_requests_ids: Optional[List[str]] = None
212 - ) -> ModelInputForNeuron:
213 - # NOTE: We assume that all sequences in the group are all prompts or
214 - # all decodes.
215 - is_prompt = seq_group_metadata_list[0].is_prompt
216 - # Prepare input tensors.
217 - if is_prompt:
218 - (input_tokens, input_positions, input_block_ids, seq_lens,
219 - multi_modal_kwargs
220 - ) = self._prepare_prompt(seq_group_metadata_list)
221 - else:
222 - (input_tokens, input_positions,
223 - input_block_ids) = self._prepare_decode(seq_group_metadata_list)
224 - seq_lens = None
225 -
226 - if not self._on_device_sampling_disabled:
227 - for seq_group_metadata in seq_group_metadata_list:
228 - sampling_params = seq_group_metadata.sampling_params
229 - top_k, top_p, temperature = (
230 - self._convert_to_neuron_sampling_params(sampling_params))
231 - sampling_params.top_k = top_k
232 - sampling_params.top_p = top_p
233 - sampling_params.temperature = temperature
234 -
235 - # we need multi_modal_data for later tokens as well
236 - multi_modal_kwargs_list: List[MultiModalKwargs] = []
237 - for seq_group_metadata in seq_group_metadata_list:
238 - mm_data = seq_group_metadata.multi_modal_data
239 - if mm_data:
240 - multi_modal_kwargs_list.append(mm_data)
241 - multi_modal_kwargs = MultiModalKwargs.batch(multi_modal_kwargs_list)
242 -
243 - lora_adapter_ids = self._get_lora_adapter_ids(seq_group_metadata_list)
244 -
245 - sampling_metadata = SamplingMetadata.prepare(
246 - seq_group_metadata_list,
247 - seq_lens,
248 - # query_lens is not needed if chunked prefill is not
249 - # supported. Since neuron worker doesn't support chunked prefill
250 - # just use seq_lens instead.
251 - seq_lens,
252 - self.device,
253 - self.pin_memory,
254 - generators=self.get_generators(finished_requests_ids))
255 -
256 - return ModelInputForNeuron(input_tokens=input_tokens,
257 - input_positions=input_positions,
258 - input_block_ids=input_block_ids,
259 - sampling_metadata=sampling_metadata,
260 - multi_modal_kwargs=multi_modal_kwargs,
261 - adapter_ids=lora_adapter_ids)
262 -
263 - def remove_all_loras(self):
264 - raise NotImplementedError(
265 - "Managing LoRAs is only supported through the "
266 - "lora_modules parameter in override_neuron_config")
267 -
268 - def set_active_loras(self, lora_requests: Set[LoRARequest],
269 - lora_mapping: LoRAMapping) -> None:
270 - raise NotImplementedError(
271 - "Managing LoRAs is only supported through the "
272 - "lora_modules parameter in override_neuron_config")
273 -
274 - def add_lora(self, lora_request: LoRARequest):
275 - logger.warning(
276 - "Adding LoRAs is only supported through the "
277 - "lora_modules parameter in override_neuron_config. If you supplied "
278 - "the parameter, you can ignore this warning. Ignoring"
279 - "lora request: ", lora_request)
280 -
281 - def remove_lora(self, lora_id: int) -> bool:
282 - raise NotImplementedError(
283 - "Managing LoRAs is only supported through the "
284 - "lora_modules parameter in override_neuron_config")
285 -
286 - def pin_lora(self, lora_id: int) -> bool:
287 - raise NotImplementedError(
288 - "Managing LoRAs is only supported through the "
289 - "lora_modules parameter in override_neuron_config")
290 -
291 - def list_loras(self) -> Set[int]:
292 - raise NotImplementedError(
293 - "Managing LoRAs is only supported through the "
294 - "lora_modules parameter in override_neuron_config")