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

Upstream Llama4 Support to Main (#16113)

Signed-off-by: Aston Zhang <22279212+astonzhang@users.noreply.github.com> Signed-off-by: Chris Thi <chris.c.thi@gmail.com> Signed-off-by: drisspg <drisspguessous@gmail.com> Signed-off-by: Jon Swenson <jmswen@gmail.com> Signed-off-by: Keyun Tong <tongkeyun@gmail.com> Signed-off-by: Lu Fang <fanglu@meta.com> Signed-off-by: Xiaodong Wang <xdwang@meta.com> Signed-off-by: Yang Chen <yangche@fb.com> Signed-off-by: Ye (Charlotte) Qi <yeq@meta.com> Signed-off-by: Yong Hoon Shin <yhshin@meta.com> Signed-off-by: Zijing Liu <liuzijing2014@gmail.com> Signed-off-by: Lu Fang <lufang@fb.com> Signed-off-by: Lu Fang <fanglu@fb.com> Signed-off-by: Lucia Fang <fanglu@fb.com> Signed-off-by: Roger Wang <ywang@roblox.com> Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk> Co-authored-by: Lu Fang <fanglu@fb.com> Co-authored-by: Roger Wang <ywang@roblox.com> Co-authored-by: DarkLight1337 <tlleungac@connect.ust.hk>

Lu Fang committed Apr 7, 2025 at 08:06 UTC 55dcce91df150f576c28520d987eaf1498fcb0bd
43 files changed +2436 -155
.buildkite/test-pipeline.yaml
+2 -1
@@ -389,7 +389,8 @@ steps:
389 - pytest -v -s models/test_transformers.py
390 - pytest -v -s models/test_registry.py
391 # V1 Test: https://github.com/vllm-project/vllm/issues/14531
392 - - VLLM_USE_V1=0 pytest -v -s models/test_initialization.py
392 + - VLLM_USE_V1=0 pytest -v -s models/test_initialization.py -k 'not llama4'
393 + - VLLM_USE_V1=0 pytest -v -s models/test_initialization.py -k 'llama4'
394
395 - label: Language Models Test (Standard) # 32min
396 #mirror_hardwares: [amd]
benchmarks/kernels/benchmark_moe.py
+3
@@ -553,6 +553,9 @@ def main(args: argparse.Namespace):
553 intermediate_size = config.moe_intermediate_size
554 shard_intermediate_size = 2 * intermediate_size // args.tp_size
555 else:
556 + if not hasattr(config, "hidden_size"):
557 + # Support for llama4
558 + config = config.text_config
559 # Default: Mixtral.
560 E = config.num_local_experts
561 topk = config.num_experts_per_tok
docs/source/models/supported_models.md
+9 -2
@@ -24,7 +24,7 @@ vLLM also supports model implementations that are available in Transformers. Thi
24
25 To check if the modeling backend is Transformers, you can simply do this:
26
27 -```python
27 +```python
28 from vllm import LLM
29 llm = LLM(model=..., task="generate") # Name or path of your model
30 llm.apply_model(lambda model: print(type(model)))
@@ -55,7 +55,7 @@ If your model is neither supported natively by vLLM or Transformers, you can sti
55 Simply set `trust_remote_code=True` and vLLM will run any model on the Model Hub that is compatible with Transformers.
56 Provided that the model writer implements their model in a compatible way, this means that you can run new models before they are officially supported in Transformers or vLLM!
57
58 -```python
58 +```python
59 from vllm import LLM
60 llm = LLM(model=..., task="generate", trust_remote_code=True) # Name or path of your model
61 llm.apply_model(lambda model: print(model.__class__))
@@ -850,6 +850,13 @@ See [this page](#generative-models) for more information on how to use generativ
850 *
851 * ✅︎
852 * ✅︎
853 +- * `Llama4ForConditionalGeneration`
854 + * Llama-4-17B-Omni-Instruct
855 + * T + I<sup>+</sup>
856 + * `meta-llama/Llama-4-Scout-17B-16E-Instruct`, `meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8`, `meta-llama/Llama-4-Maverick-17B-128E-Instruct`, etc.
857 + *
858 + * ✅︎
859 + * ✅︎
860 - * `LlavaForConditionalGeneration`
861 * LLaVA-1.5
862 * T + I<sup>E+</sup>
examples/offline_inference/audio_language.py
+1 -1
@@ -47,7 +47,7 @@ def run_minicpmo(question: str, audio_count: int) -> ModelRequestData:
47 model=model_name,
48 trust_remote_code=True,
49 max_model_len=4096,
50 - max_num_seqs=5,
50 + max_num_seqs=2,
51 limit_mm_per_prompt={"audio": audio_count},
52 )
53
examples/offline_inference/vision_language.py
+37
@@ -582,6 +582,42 @@ def run_mllama(questions: list[str], modality: str) -> ModelRequestData:
582 )
583
584
585 +def run_llama4(questions: list[str], modality: str):
586 + assert modality == "image"
587 +
588 + model_name = "meta-llama/Llama-4-Scout-17B-16E-Instruct"
589 +
590 + engine_args = EngineArgs(
591 + model=model_name,
592 + max_model_len=8192,
593 + max_num_seqs=4,
594 + tensor_parallel_size=8,
595 + disable_mm_preprocessor_cache=args.disable_mm_preprocessor_cache,
596 + gpu_memory_utilization=0.4,
597 + )
598 +
599 + tokenizer = AutoTokenizer.from_pretrained(model_name)
600 + messages = [[{
601 + "role":
602 + "user",
603 + "content": [{
604 + "type": "image"
605 + }, {
606 + "type": "text",
607 + "text": f"{question}"
608 + }]
609 + }] for question in questions]
610 + prompts = tokenizer.apply_chat_template(messages,
611 + add_generation_prompt=True,
612 + tokenize=False)
613 + stop_token_ids = None
614 + return ModelRequestData(
615 + engine_args=engine_args,
616 + prompts=prompts,
617 + stop_token_ids=stop_token_ids,
618 + )
619 +
620 +
621 # Molmo
622 def run_molmo(questions: list[str], modality: str) -> ModelRequestData:
623 assert modality == "image"
@@ -907,6 +943,7 @@ model_example_map = {
943 "minicpmv": run_minicpmv,
944 "mistral3": run_mistral3,
945 "mllama": run_mllama,
946 + "llama4": run_llama4,
947 "molmo": run_molmo,
948 "NVLM_D": run_nvlm_d,
949 "paligemma": run_paligemma,
examples/offline_inference/vision_language_multi_image.py
+38
@@ -253,6 +253,43 @@ def load_internvl(question: str, image_urls: list[str]) -> ModelRequestData:
253 )
254
255
256 +def load_llama4(question: str, image_urls: list[str]) -> ModelRequestData:
257 + model_name = "meta-llama/Llama-4-Scout-17B-16E-Instruct"
258 +
259 + engine_args = EngineArgs(
260 + model=model_name,
261 + max_model_len=8192,
262 + max_num_seqs=4,
263 + tensor_parallel_size=8,
264 + limit_mm_per_prompt={"image": len(image_urls)},
265 + )
266 +
267 + placeholders = [{"type": "image", "image": url} for url in image_urls]
268 + messages = [{
269 + "role":
270 + "user",
271 + "content": [
272 + *placeholders,
273 + {
274 + "type": "text",
275 + "text": question
276 + },
277 + ],
278 + }]
279 +
280 + processor = AutoProcessor.from_pretrained(model_name)
281 +
282 + prompt = processor.apply_chat_template(messages,
283 + tokenize=False,
284 + add_generation_prompt=True)
285 +
286 + return ModelRequestData(
287 + engine_args=engine_args,
288 + prompt=prompt,
289 + image_data=[fetch_image(url) for url in image_urls],
290 + )
291 +
292 +
293 def load_mistral3(question: str, image_urls: list[str]) -> ModelRequestData:
294 model_name = "mistralai/Mistral-Small-3.1-24B-Instruct-2503"
295
@@ -567,6 +604,7 @@ model_example_map = {
604 "h2ovl_chat": load_h2ovl,
605 "idefics3": load_idefics3,
606 "internvl_chat": load_internvl,
607 + "llama4": load_llama4,
608 "mistral3": load_mistral3,
609 "mllama": load_mllama,
610 "NVLM_D": load_nvlm_d,
requirements/common.txt
+1 -1
@@ -6,7 +6,7 @@ requests >= 2.26.0
6 tqdm
7 blake3
8 py-cpuinfo
9 -transformers >= 4.50.3
9 +transformers >= 4.51.0
10 huggingface-hub[hf_xet] >= 0.30.0 # Required for Xet downloads.
11 tokenizers >= 0.19.1 # Required for Llama 3.
12 protobuf # Required by LlamaTokenizer.
requirements/test.in
+1 -1
@@ -30,7 +30,7 @@ mistral_common[opencv] >= 1.5.4 # required for pixtral test
30 opencv-python-headless >= 4.11.0 # required for video test
31 datamodel_code_generator # required for minicpm3 test
32 lm-eval[api]==0.4.8 # required for model evaluation test
33 -transformers==4.50.3
33 +transformers==4.51.0
34 huggingface-hub[hf_xet]>=0.30.0 # Required for Xet downloads.
35 # quantization
36 bitsandbytes>=0.45.3
requirements/test.txt
+1 -1
@@ -645,7 +645,7 @@ tqdm==4.66.6
645 # transformers
646 tqdm-multiprocess==0.0.11
647 # via lm-eval
648 -transformers==4.50.3
648 +transformers==4.51.0
649 # via
650 # -r requirements/test.in
651 # genai-perf
tests/models/decoder_only/audio_language/test_ultravox.py
+13 -1
@@ -12,6 +12,7 @@ from vllm.sequence import SampleLogprobs
12
13 from ....conftest import HfRunner, VllmRunner
14 from ....utils import RemoteOpenAIServer
15 +from ...registry import HF_EXAMPLE_MODELS
16 from ...utils import check_logprobs_close
17
18 MODEL_NAME = "fixie-ai/ultravox-v0_5-llama-3_2-1b"
@@ -55,7 +56,10 @@ def server(request, audio_assets):
56 for key, value in request.param.items()
57 ]
58
58 - with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
59 + with RemoteOpenAIServer(MODEL_NAME,
60 + args,
61 + env_dict={"VLLM_AUDIO_FETCH_TIMEOUT":
62 + "30"}) as remote_server:
63 yield remote_server
64
65
@@ -106,6 +110,10 @@ def run_test(
110 **kwargs,
111 ):
112 """Inference result should be the same between hf and vllm."""
113 + model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
114 + model_info.check_available_online(on_fail="skip")
115 + model_info.check_transformers_version(on_fail="skip")
116 +
117 # NOTE: take care of the order. run vLLM first, and then run HF.
118 # vLLM needs a fresh new process without cuda initialization.
119 # if we run HF first, the cuda initialization will be done and it
@@ -156,6 +164,10 @@ def run_multi_audio_test(
164 num_logprobs: int,
165 **kwargs,
166 ):
167 + model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
168 + model_info.check_available_online(on_fail="skip")
169 + model_info.check_transformers_version(on_fail="skip")
170 +
171 with vllm_runner(model,
172 dtype=dtype,
173 enforce_eager=True,
tests/models/decoder_only/vision_language/test_models.py
+34 -3
@@ -160,17 +160,32 @@ VLM_TEST_SETTINGS = {
160 ),
161 "aya_vision": VLMTestInfo(
162 models=["CohereForAI/aya-vision-8b"],
163 - test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE),
163 + test_type=(VLMTestType.IMAGE),
164 prompt_formatter=lambda img_prompt: f"<|START_OF_TURN_TOKEN|><|USER_TOKEN|>{img_prompt}<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>", # noqa: E501
165 single_image_prompts=IMAGE_ASSETS.prompts({
166 "stop_sign": "<image>What's the content in the center of the image?", # noqa: E501
167 "cherry_blossom": "<image>What is the season?", # noqa: E501
168 }),
169 multi_image_prompt="<image><image>Describe the two images in detail.", # noqa: E501
170 - max_model_len=8192,
170 + max_model_len=4096,
171 + max_num_seqs=2,
172 + auto_cls=AutoModelForImageTextToText,
173 + vllm_runner_kwargs={"mm_processor_kwargs": {"crop_to_patches": True}},
174 + ),
175 + "aya_vision-multi_image": VLMTestInfo(
176 + models=["CohereForAI/aya-vision-8b"],
177 + test_type=(VLMTestType.MULTI_IMAGE),
178 + prompt_formatter=lambda img_prompt: f"<|START_OF_TURN_TOKEN|><|USER_TOKEN|>{img_prompt}<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>", # noqa: E501
179 + single_image_prompts=IMAGE_ASSETS.prompts({
180 + "stop_sign": "<image>What's the content in the center of the image?", # noqa: E501
181 + "cherry_blossom": "<image>What is the season?", # noqa: E501
182 + }),
183 + multi_image_prompt="<image><image>Describe the two images in detail.", # noqa: E501
184 + max_model_len=4096,
185 max_num_seqs=2,
186 auto_cls=AutoModelForImageTextToText,
173 - vllm_runner_kwargs={"mm_processor_kwargs": {"crop_to_patches": True}}
187 + vllm_runner_kwargs={"mm_processor_kwargs": {"crop_to_patches": True}},
188 + marks=[large_gpu_mark(min_gb=32)],
189 ),
190 "blip2": VLMTestInfo(
191 # TODO: Change back to 2.7b once head_dim = 80 is supported
@@ -303,6 +318,22 @@ VLM_TEST_SETTINGS = {
318 use_tokenizer_eos=True,
319 patch_hf_runner=model_utils.internvl_patch_hf_runner,
320 ),
321 + "llama4": VLMTestInfo(
322 + models=["meta-llama/Llama-4-Scout-17B-16E-Instruct"],
323 + prompt_formatter=lambda img_prompt: f"<|begin_of_text|><|header_start|>user<|header_end|>\n\n{img_prompt}<|eot|><|header_start|>assistant<|header_end|>\n\n", # noqa: E501
324 + img_idx_to_prompt=lambda _: "<|image|>",
325 + test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE),
326 + distributed_executor_backend="mp",
327 + image_size_factors=[(.25, 0.5, 1.0)],
328 + hf_model_kwargs={"device_map": "auto"},
329 + max_model_len=8192,
330 + max_num_seqs=4,
331 + dtype="bfloat16",
332 + auto_cls=AutoModelForImageTextToText,
333 + tensor_parallel_size=8,
334 + vllm_runner_kwargs={"gpu_memory_utilization": 0.8},
335 + marks=multi_gpu_marks(num_gpus=8),
336 + ),
337 "llava_next": VLMTestInfo(
338 models=["llava-hf/llava-v1.6-mistral-7b-hf"],
339 test_type=(VLMTestType.IMAGE, VLMTestType.CUSTOM_INPUTS),
tests/models/decoder_only/vision_language/test_phi3v.py
+9
@@ -5,7 +5,9 @@ import re
5 from typing import Optional
6
7 import pytest
8 +from packaging.version import Version
9 from transformers import AutoTokenizer
10 +from transformers import __version__ as TRANSFORMERS_VERSION
11
12 from vllm.multimodal.image import rescale_image_size
13 from vllm.platforms import current_platform
@@ -81,6 +83,13 @@ def run_test(
83 from transformers import AutoImageProcessor # noqa: F401
84 from transformers import AutoProcessor # noqa: F401
85
86 + # Once the model repo is updated to 4.49, we should be able to run the
87 + # test in `test_models.py` without the above workaround
88 + if Version(TRANSFORMERS_VERSION) >= Version("4.49"):
89 + pytest.skip(f"`transformers=={TRANSFORMERS_VERSION}` installed, "
90 + "but `transformers<=4.49` is required to run this model. "
91 + "Reason: Cannot run HF implementation")
92 +
93 # NOTE: take care of the order. run vLLM first, and then run HF.
94 # vLLM needs a fresh new process without cuda initialization.
95 # if we run HF first, the cuda initialization will be done and it
tests/models/decoder_only/vision_language/test_pixtral.py
+2
@@ -176,6 +176,8 @@ def test_chat(
176 model,
177 dtype=dtype,
178 tokenizer_mode="mistral",
179 + load_format="mistral",
180 + config_format="mistral",
181 max_model_len=max_model_len,
182 limit_mm_per_prompt=LIMIT_MM_PER_PROMPT,
183 ) as vllm_model:
tests/models/multimodal/processing/test_common.py
+1
@@ -257,6 +257,7 @@ def _test_processing_correctness_mistral(
257 "h2oai/h2ovl-mississippi-800m",
258 "OpenGVLab/InternVL2-1B",
259 "HuggingFaceM4/Idefics3-8B-Llama3",
260 + "meta-llama/Llama-4-Scout-17B-16E-Instruct",
261 "llava-hf/llava-1.5-7b-hf",
262 "llava-hf/llava-v1.6-mistral-7b-hf",
263 "llava-hf/LLaVA-NeXT-Video-7B-hf",
tests/models/multimodal/processing/test_llama4.py new
+99
@@ -0,0 +1,99 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +"""Tests for Llama4's multimodal preprocessing kwargs."""
3 +
4 +import pytest
5 +
6 +from vllm.multimodal import MULTIMODAL_REGISTRY
7 +from vllm.transformers_utils.tokenizer import encode_tokens
8 +
9 +from ....conftest import _ImageAssets
10 +from ...utils import build_model_context
11 +
12 +
13 +@pytest.mark.parametrize("model_id",
14 + ["meta-llama/Llama-4-Scout-17B-16E-Instruct"])
15 +@pytest.mark.parametrize("mm_processor_kwargs", [{}])
16 +@pytest.mark.parametrize("num_imgs", [1, 5])
17 +@pytest.mark.parametrize("disable_mm_preprocessor_cache", [True, False])
18 +@pytest.mark.parametrize("tokenized_prompt", [True, False])
19 +def test_processor_override(
20 + image_assets: _ImageAssets,
21 + model_id: str,
22 + mm_processor_kwargs: dict,
23 + num_imgs: int,
24 + disable_mm_preprocessor_cache: bool,
25 + tokenized_prompt: bool,
26 +):
27 + """Ensure llama4 processor works properly."""
28 + ctx = build_model_context(
29 + model_id,
30 + mm_processor_kwargs=mm_processor_kwargs,
31 + limit_mm_per_prompt={"image": num_imgs},
32 + disable_mm_preprocessor_cache=disable_mm_preprocessor_cache,
33 + )
34 + processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
35 + config = processor.info.get_hf_config()
36 + tokenizer = processor.info.get_tokenizer()
37 + hf_processor = processor.info.get_hf_processor()
38 + vocab = tokenizer.get_vocab()
39 +
40 + prompt = "<|begin_of_text|><|header_start|>user<|header_end|>" \
41 + + "<|image|>" * num_imgs \
42 + + "<|eot|><|header_start|>assistant<|header_end|>"
43 + mm_data = {
44 + "image": [
45 + image_assets[(i % len(image_assets))].pil_image
46 + for i in range(num_imgs)
47 + ]
48 + }
49 + if tokenized_prompt:
50 + prompt = encode_tokens(tokenizer, prompt)
51 +
52 + processed_inputs = processor.apply(prompt, mm_data, mm_processor_kwargs)
53 + mm_kwargs = processed_inputs["mm_kwargs"]
54 +
55 + # place holder replacements
56 + prompt_token_ids = processed_inputs["prompt_token_ids"]
57 + assert prompt_token_ids.count(config.boi_token_index) == num_imgs
58 + assert prompt_token_ids.count(config.eoi_token_index) == num_imgs
59 + assert prompt_token_ids.count(vocab[hf_processor.image_token]) == num_imgs
60 + aspect_ratios = mm_kwargs["aspect_ratios"]
61 + num_x_separators = num_y_separators = 0
62 + for tiles_y, tiles_x in aspect_ratios:
63 + if tiles_x * tiles_y > 1:
64 + num_x_separators += (tiles_x - 1) * tiles_y
65 + num_y_separators += tiles_y
66 + assert prompt_token_ids.count(vocab[hf_processor.tile_token]) \
67 + == num_x_separators
68 + assert prompt_token_ids.count(vocab[hf_processor.tile_global_token]) \
69 + == num_y_separators
70 +
71 + # image token offsets
72 + img_locs = processed_inputs["mm_placeholders"].get("image", [])
73 + assert len(img_locs) == num_imgs
74 + assert [img_loc["offset"] for img_loc in img_locs] == \
75 + [i for i, v in enumerate(prompt_token_ids) \
76 + if v == config.boi_token_index]
77 +
78 + # patch sizes and masks
79 + assert prompt_token_ids.count(config.image_token_index) \
80 + == sum(img_patch.sum() for img_patch in mm_kwargs["embed_is_patch"])
81 + patch_token_id = vocab[hf_processor.img_patch_token]
82 + num_patches = processed_inputs["prompt_token_ids"].count(patch_token_id)
83 + mm_counts = {"image": num_imgs}
84 + assert num_patches / num_imgs <= \
85 + processor.info.get_mm_max_tokens_per_item(32768, mm_counts)["image"]
86 + num_patches_per_chunk = processor.info.get_patch_per_chunk(
87 + config.vision_config)
88 + assert prompt_token_ids.count(config.image_token_index) \
89 + == mm_kwargs["patches_per_image"].sum() * num_patches_per_chunk
90 + assert mm_kwargs["pixel_values"].shape[0] \
91 + == mm_kwargs["patches_per_image"].sum()
92 +
93 + for embed_is_patch, aspect_ratio in zip(mm_kwargs["embed_is_patch"],
94 + mm_kwargs["aspect_ratios"]):
95 + assert embed_is_patch.shape[0] == \
96 + len(tokenizer.encode(
97 + hf_processor._prompt_split_image(
98 + aspect_ratio, num_patches_per_chunk),
99 + add_special_tokens=False))
tests/models/registry.py
+5 -1
@@ -287,12 +287,16 @@ _MULTIMODAL_EXAMPLE_MODELS = {
287 trust_remote_code=True,
288 hf_overrides={"architectures": ["GLM4VForCausalLM"]}), # noqa: E501
289 "H2OVLChatModel": _HfExamplesInfo("h2oai/h2ovl-mississippi-800m",
290 - extras={"2b": "h2oai/h2ovl-mississippi-2b"}), # noqa: E501
290 + extras={"2b": "h2oai/h2ovl-mississippi-2b"}, # noqa: E501
291 + max_transformers_version="4.48", # noqa: E501
292 + transformers_version_reason="HF model is not compatible."), # noqa: E501
293 "InternVLChatModel": _HfExamplesInfo("OpenGVLab/InternVL2-1B",
294 extras={"2B": "OpenGVLab/InternVL2-2B"}, # noqa: E501
295 trust_remote_code=True),
296 "Idefics3ForConditionalGeneration": _HfExamplesInfo("HuggingFaceM4/Idefics3-8B-Llama3", # noqa: E501
297 {"tiny": "HuggingFaceTB/SmolVLM-256M-Instruct"}), # noqa: E501
298 + "Llama4ForConditionalGeneration": _HfExamplesInfo("meta-llama/Llama-4-Scout-17B-16E-Instruct", # noqa: E501
299 + min_transformers_version="4.51"),
300 "LlavaForConditionalGeneration": _HfExamplesInfo("llava-hf/llava-1.5-7b-hf",
301 extras={"mistral": "mistral-community/pixtral-12b", # noqa: E501
302 "mistral-fp8": "nm-testing/pixtral-12b-FP8-dynamic"}), # noqa: E501
tests/models/test_initialization.py
+13 -4
@@ -7,6 +7,8 @@ from transformers import PretrainedConfig
7
8 from vllm import LLM
9 from vllm.engine.llm_engine import LLMEngine as V0LLMEngine
10 +from vllm.utils import GiB_bytes
11 +from vllm.v1.core.kv_cache_utils import get_kv_cache_config
12 from vllm.v1.engine.core import EngineCore as V1EngineCore
13
14 from .registry import HF_EXAMPLE_MODELS
@@ -42,14 +44,21 @@ def test_can_initialize(model_arch):
44 self.cache_config.num_gpu_blocks = 0
45 self.cache_config.num_cpu_blocks = 0
46
45 - def _initalize_kv_caches_v1(self, vllm_config):
46 - # gpu_blocks (> 0), cpu_blocks
47 - return 1, 0
47 + def _initialize_kv_caches_v1(self, vllm_config):
48 + kv_cache_specs = self.model_executor.get_kv_cache_specs()
49 + scheduler_kv_cache_config = get_kv_cache_config(
50 + vllm_config,
51 + kv_cache_specs[0],
52 + 20 * GiB_bytes,
53 + )
54 +
55 + # gpu_blocks (> 0), cpu_blocks, scheduler_kv_cache_config
56 + return 1, 0, scheduler_kv_cache_config
57
58 with (patch.object(V0LLMEngine, "_initialize_kv_caches",
59 _initialize_kv_caches_v0),
60 patch.object(V1EngineCore, "_initialize_kv_caches",
52 - _initalize_kv_caches_v1)):
61 + _initialize_kv_caches_v1)):
62 LLM(
63 model_info.default,
64 tokenizer=model_info.tokenizer,
vllm/config.py
+2
@@ -358,6 +358,8 @@ class ModelConfig:
358 self.hf_config = hf_config
359
360 self.hf_text_config = get_hf_text_config(self.hf_config)
361 + self.attention_chunk_size = getattr(self.hf_text_config,
362 + "attention_chunk_size", None)
363 self.encoder_config = self._get_encoder_config()
364 self.hf_image_processor_config = get_hf_image_processor_config(
365 self.model, hf_token=hf_token, revision=revision)
vllm/entrypoints/chat_utils.py
+1 -1
@@ -500,7 +500,7 @@ class BaseMultiModalItemTracker(ABC, Generic[_T]):
500 "internvl_chat", "skywork_chat", "NVLM_D",
501 "h2ovl_chat", "idefics3"):
502 return "<image>"
503 - if model_type == "mllama":
503 + if model_type in ("mllama", "llama4"):
504 return "<|image|>"
505 if model_type in ("qwen2_vl", "qwen2_5_vl"):
506 return "<|vision_start|><|image_pad|><|vision_end|>"
vllm/model_executor/layers/fused_moe/configs/E=16,N=1024,device_name=AMD_Instinct_MI300X.json new
+200
@@ -0,0 +1,200 @@
1 +{
2 + "1": {
3 + "BLOCK_SIZE_M": 16,
4 + "BLOCK_SIZE_N": 16,
5 + "BLOCK_SIZE_K": 256,
6 + "GROUP_SIZE_M": 1,
7 + "num_warps": 4,
8 + "num_stages": 2,
9 + "waves_per_eu": 0,
10 + "matrix_instr_nonkdim": 16,
11 + "kpack": 1
12 + },
13 + "2": {
14 + "BLOCK_SIZE_M": 16,
15 + "BLOCK_SIZE_N": 16,
16 + "BLOCK_SIZE_K": 256,
17 + "GROUP_SIZE_M": 1,
18 + "num_warps": 4,
19 + "num_stages": 2,
20 + "waves_per_eu": 0,
21 + "matrix_instr_nonkdim": 16,
22 + "kpack": 1
23 + },
24 + "4": {
25 + "BLOCK_SIZE_M": 16,
26 + "BLOCK_SIZE_N": 16,
27 + "BLOCK_SIZE_K": 256,
28 + "GROUP_SIZE_M": 1,
29 + "num_warps": 2,
30 + "num_stages": 2,
31 + "waves_per_eu": 0,
32 + "matrix_instr_nonkdim": 16,
33 + "kpack": 1
34 + },
35 + "8": {
36 + "BLOCK_SIZE_M": 16,
37 + "BLOCK_SIZE_N": 64,
38 + "BLOCK_SIZE_K": 128,
39 + "GROUP_SIZE_M": 1,
40 + "num_warps": 4,
41 + "num_stages": 2,
42 + "waves_per_eu": 0,
43 + "matrix_instr_nonkdim": 16,
44 + "kpack": 2
45 + },
46 + "16": {
47 + "BLOCK_SIZE_M": 16,
48 + "BLOCK_SIZE_N": 32,
49 + "BLOCK_SIZE_K": 128,
50 + "GROUP_SIZE_M": 1,
51 + "num_warps": 2,
52 + "num_stages": 2,
53 + "waves_per_eu": 0,
54 + "matrix_instr_nonkdim": 16,
55 + "kpack": 1
56 + },
57 + "24": {
58 + "BLOCK_SIZE_M": 16,
59 + "BLOCK_SIZE_N": 64,
60 + "BLOCK_SIZE_K": 256,
61 + "GROUP_SIZE_M": 1,
62 + "num_warps": 4,
63 + "num_stages": 2,
64 + "waves_per_eu": 0,
65 + "matrix_instr_nonkdim": 16,
66 + "kpack": 2
67 + },
68 + "32": {
69 + "BLOCK_SIZE_M": 32,
70 + "BLOCK_SIZE_N": 64,
71 + "BLOCK_SIZE_K": 128,
72 + "GROUP_SIZE_M": 1,
73 + "num_warps": 4,
74 + "num_stages": 2,
75 + "waves_per_eu": 0,
76 + "matrix_instr_nonkdim": 16,
77 + "kpack": 1
78 + },
79 + "48": {
80 + "BLOCK_SIZE_M": 16,
81 + "BLOCK_SIZE_N": 64,
82 + "BLOCK_SIZE_K": 128,
83 + "GROUP_SIZE_M": 1,
84 + "num_warps": 4,
85 + "num_stages": 2,
86 + "waves_per_eu": 0,
87 + "matrix_instr_nonkdim": 16,
88 + "kpack": 1
89 + },
90 + "64": {
91 + "BLOCK_SIZE_M": 16,
92 + "BLOCK_SIZE_N": 64,
93 + "BLOCK_SIZE_K": 128,
94 + "GROUP_SIZE_M": 1,
95 + "num_warps": 2,
96 + "num_stages": 2,
97 + "waves_per_eu": 0,
98 + "matrix_instr_nonkdim": 16,
99 + "kpack": 2
100 + },
101 + "96": {
102 + "BLOCK_SIZE_M": 16,
103 + "BLOCK_SIZE_N": 64,
104 + "BLOCK_SIZE_K": 256,
105 + "GROUP_SIZE_M": 1,
106 + "num_warps": 4,
107 + "num_stages": 2,
108 + "waves_per_eu": 0,
109 + "matrix_instr_nonkdim": 16,
110 + "kpack": 2
111 + },
112 + "128": {
113 + "BLOCK_SIZE_M": 16,
114 + "BLOCK_SIZE_N": 64,
115 + "BLOCK_SIZE_K": 256,
116 + "GROUP_SIZE_M": 1,
117 + "num_warps": 4,
118 + "num_stages": 2,
119 + "waves_per_eu": 0,
120 + "matrix_instr_nonkdim": 16,
121 + "kpack": 2
122 + },
123 + "256": {
124 + "BLOCK_SIZE_M": 32,
125 + "BLOCK_SIZE_N": 32,
126 + "BLOCK_SIZE_K": 256,
127 + "GROUP_SIZE_M": 8,
128 + "num_warps": 4,
129 + "num_stages": 2,
130 + "waves_per_eu": 0,
131 + "matrix_instr_nonkdim": 16,
132 + "kpack": 2
133 + },
134 + "512": {
135 + "BLOCK_SIZE_M": 64,
136 + "BLOCK_SIZE_N": 64,
137 + "BLOCK_SIZE_K": 128,
138 + "GROUP_SIZE_M": 1,
139 + "num_warps": 8,
140 + "num_stages": 2,
141 + "waves_per_eu": 0,
142 + "matrix_instr_nonkdim": 16,
143 + "kpack": 2
144 + },
145 + "1024": {
146 + "BLOCK_SIZE_M": 64,
147 + "BLOCK_SIZE_N": 64,
148 + "BLOCK_SIZE_K": 128,
149 + "GROUP_SIZE_M": 1,
150 + "num_warps": 8,
151 + "num_stages": 2,
152 + "waves_per_eu": 0,
153 + "matrix_instr_nonkdim": 16,
154 + "kpack": 2
155 + },
156 + "1536": {
157 + "BLOCK_SIZE_M": 128,
158 + "BLOCK_SIZE_N": 128,
159 + "BLOCK_SIZE_K": 128,
160 + "GROUP_SIZE_M": 8,
161 + "num_warps": 8,
162 + "num_stages": 2,
163 + "waves_per_eu": 0,
164 + "matrix_instr_nonkdim": 16,
165 + "kpack": 2
166 + },
167 + "2048": {
168 + "BLOCK_SIZE_M": 64,
169 + "BLOCK_SIZE_N": 64,
170 + "BLOCK_SIZE_K": 128,
171 + "GROUP_SIZE_M": 1,
172 + "num_warps": 8,
173 + "num_stages": 2,
174 + "waves_per_eu": 0,
175 + "matrix_instr_nonkdim": 16,
176 + "kpack": 2
177 + },
178 + "3072": {
179 + "BLOCK_SIZE_M": 128,
180 + "BLOCK_SIZE_N": 128,
181 + "BLOCK_SIZE_K": 64,
182 + "GROUP_SIZE_M": 1,
183 + "num_warps": 8,
184 + "num_stages": 2,
185 + "waves_per_eu": 0,
186 + "matrix_instr_nonkdim": 16,
187 + "kpack": 2
188 + },
189 + "4096": {
190 + "BLOCK_SIZE_M": 64,
191 + "BLOCK_SIZE_N": 128,
192 + "BLOCK_SIZE_K": 64,
193 + "GROUP_SIZE_M": 1,
194 + "num_warps": 4,
195 + "num_stages": 2,
196 + "waves_per_eu": 0,
197 + "matrix_instr_nonkdim": 16,
198 + "kpack": 2
199 + }
200 +}
vllm/model_executor/layers/fused_moe/cutlass_moe.py
+12 -3
@@ -23,6 +23,7 @@ def cutlass_moe_fp8(
23 a1_scale: Optional[torch.Tensor] = None,
24 a2_scale: Optional[torch.Tensor] = None,
25 out_dtype: torch.dtype = torch.half,
26 + apply_router_weight_on_input: bool = False,
27 ) -> torch.Tensor:
28 """
29 This function computes a a8w8-quantized Mixture of Experts (MoE) layer
@@ -96,8 +97,14 @@ def cutlass_moe_fp8(
97 n = w2_q.size(1)
98
99 topk = topk_ids.size(1)
100 +
101 per_act_token = a1_scale.numel() != 1 if a1_scale is not None else (
102 a2_scale.numel() != 1 if a2_scale is not None else False)
103 + if apply_router_weight_on_input:
104 + assert topk == 1, \
105 + "apply_router_weight_on_input is only implemented for topk=1"
106 + # TODO: this only works for topK=1, will need to update for topK>1
107 + a = a * topk_weights.to(out_dtype)
108
109 a_q, a1_scale = ops.scaled_fp8_quant(
110 a, a1_scale, use_per_token_if_dynamic=per_act_token)
@@ -139,6 +146,8 @@ def cutlass_moe_fp8(
146 ops.cutlass_moe_mm(c2, intemediate_q, w2_q, a2_scale, w2_scale,
147 expert_offsets[:-1], problem_sizes2, ab_strides2,
148 ab_strides2, c_strides2)
142 -
143 - return (c2[c_map].view(m, topk, k) *
144 - topk_weights.view(m, topk, 1).to(out_dtype)).sum(dim=1)
149 + # Gather tokens
150 + c2 = c2[c_map].view(m, topk, k)
151 + if not apply_router_weight_on_input:
152 + c2 = c2 * topk_weights.view(m, topk, 1).to(out_dtype)
153 + return c2.sum(dim=1)
vllm/model_executor/layers/fused_moe/fused_moe.py
+18 -10
@@ -954,6 +954,7 @@ def inplace_fused_experts(hidden_states: torch.Tensor,
954 topk_weights: torch.Tensor,
955 topk_ids: torch.Tensor,
956 activation: str = "silu",
957 + apply_router_weight_on_input: bool = False,
958 use_fp8_w8a8: bool = False,
959 use_int8_w8a16: bool = False,
960 use_int4_w4a16: bool = False,
@@ -967,10 +968,10 @@ def inplace_fused_experts(hidden_states: torch.Tensor,
968 a2_scale: Optional[torch.Tensor] = None,
969 block_shape: Optional[List[int]] = None) -> None:
970 fused_experts_impl(hidden_states, w1, w2, topk_weights, topk_ids, True,
970 - activation, use_fp8_w8a8, use_int8_w8a16,
971 - use_int4_w4a16, global_num_experts, expert_map,
972 - w1_scale, w2_scale, w1_zp, w2_zp, a1_scale, a2_scale,
973 - block_shape)
971 + activation, apply_router_weight_on_input, use_fp8_w8a8,
972 + use_int8_w8a16, use_int4_w4a16, global_num_experts,
973 + expert_map, w1_scale, w2_scale, w1_zp, w2_zp, a1_scale,
974 + a2_scale, block_shape)
975
976
977 def inplace_fused_experts_fake(
@@ -980,6 +981,7 @@ def inplace_fused_experts_fake(
981 topk_weights: torch.Tensor,
982 topk_ids: torch.Tensor,
983 activation: str = "silu",
984 + apply_router_weight_on_input: bool = False,
985 use_fp8_w8a8: bool = False,
986 use_int8_w8a16: bool = False,
987 use_int4_w4a16: bool = False,
@@ -1010,6 +1012,7 @@ def outplace_fused_experts(
1012 topk_weights: torch.Tensor,
1013 topk_ids: torch.Tensor,
1014 activation: str = "silu",
1015 + apply_router_weight_on_input: bool = False,
1016 use_fp8_w8a8: bool = False,
1017 use_int8_w8a16: bool = False,
1018 use_int4_w4a16: bool = False,
@@ -1023,10 +1026,11 @@ def outplace_fused_experts(
1026 a2_scale: Optional[torch.Tensor] = None,
1027 block_shape: Optional[List[int]] = None) -> torch.Tensor:
1028 return fused_experts_impl(hidden_states, w1, w2, topk_weights, topk_ids,
1026 - False, activation, use_fp8_w8a8, use_int8_w8a16,
1027 - use_int4_w4a16, global_num_experts, expert_map,
1028 - w1_scale, w2_scale, w1_zp, w2_zp, a1_scale,
1029 - a2_scale, block_shape)
1029 + False, activation, apply_router_weight_on_input,
1030 + use_fp8_w8a8, use_int8_w8a16, use_int4_w4a16,
1031 + global_num_experts, expert_map, w1_scale,
1032 + w2_scale, w1_zp, w2_zp, a1_scale, a2_scale,
1033 + block_shape)
1034
1035
1036 def outplace_fused_experts_fake(
@@ -1084,6 +1088,7 @@ def fused_experts(hidden_states: torch.Tensor,
1088 topk_ids: torch.Tensor,
1089 inplace: bool = False,
1090 activation: str = "silu",
1091 + apply_router_weight_on_input: bool = False,
1092 use_fp8_w8a8: bool = False,
1093 use_int8_w8a16: bool = False,
1094 use_int4_w4a16: bool = False,
@@ -1099,6 +1104,7 @@ def fused_experts(hidden_states: torch.Tensor,
1104 allow_deep_gemm: bool = False) -> torch.Tensor:
1105 if (allow_deep_gemm and use_fp8_w8a8
1106 and _valid_deep_gemm(hidden_states, w1, w2, expert_map)):
1107 + assert apply_router_weight_on_input is False
1108 return deep_gemm_moe_fp8(
1109 hidden_states=hidden_states,
1110 w1=w1,
@@ -1122,6 +1128,7 @@ def fused_experts(hidden_states: torch.Tensor,
1128 topk_weights=topk_weights,
1129 topk_ids=topk_ids,
1130 activation=activation,
1131 + apply_router_weight_on_input=apply_router_weight_on_input,
1132 use_fp8_w8a8=use_fp8_w8a8,
1133 use_int8_w8a16=use_int8_w8a16,
1134 use_int4_w4a16=use_int4_w4a16,
@@ -1143,6 +1150,7 @@ def fused_experts_impl(hidden_states: torch.Tensor,
1150 topk_ids: torch.Tensor,
1151 inplace: bool = False,
1152 activation: str = "silu",
1153 + apply_router_weight_on_input: bool = False,
1154 use_fp8_w8a8: bool = False,
1155 use_int8_w8a16: bool = False,
1156 use_int4_w4a16: bool = False,
@@ -1270,7 +1278,7 @@ def fused_experts_impl(hidden_states: torch.Tensor,
1278 sorted_token_ids,
1279 expert_ids,
1280 num_tokens_post_padded,
1273 - False,
1281 + apply_router_weight_on_input,
1282 top_k_num,
1283 config,
1284 compute_type=compute_type,
@@ -1307,7 +1315,7 @@ def fused_experts_impl(hidden_states: torch.Tensor,
1315 sorted_token_ids,
1316 expert_ids,
1317 num_tokens_post_padded,
1310 - True,
1318 + not apply_router_weight_on_input,
1319 1,
1320 config,
1321 compute_type=compute_type,
vllm/model_executor/layers/fused_moe/layer.py
+41 -24
@@ -65,7 +65,9 @@ class FusedMoEMethodBase(QuantizeMethodBase):
65 expert_map: Optional[torch.Tensor] = None,
66 custom_routing_function: Optional[Callable] = None,
67 scoring_func: str = "softmax",
68 - e_score_correction_bias: Optional[torch.Tensor] = None
68 + e_score_correction_bias: Optional[torch.Tensor] = None,
69 + apply_router_weight_on_input: bool = False,
70 + activation: str = "silu",
71 ) -> torch.Tensor:
72 raise NotImplementedError
73
@@ -156,22 +158,25 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
158 custom_routing_function: Optional[Callable] = None,
159 scoring_func: str = "softmax",
160 e_score_correction_bias: Optional[torch.Tensor] = None,
161 + apply_router_weight_on_input: bool = False,
162 activation: str = "silu",
163 ) -> torch.Tensor:
161 - return self.forward(x=x,
162 - layer=layer,
163 - router_logits=router_logits,
164 - top_k=top_k,
165 - renormalize=renormalize,
166 - use_grouped_topk=use_grouped_topk,
167 - topk_group=topk_group,
168 - num_expert_group=num_expert_group,
169 - global_num_experts=global_num_experts,
170 - expert_map=expert_map,
171 - custom_routing_function=custom_routing_function,
172 - scoring_func=scoring_func,
173 - e_score_correction_bias=e_score_correction_bias,
174 - activation=activation)
164 + return self.forward(
165 + x=x,
166 + layer=layer,
167 + router_logits=router_logits,
168 + top_k=top_k,
169 + renormalize=renormalize,
170 + use_grouped_topk=use_grouped_topk,
171 + topk_group=topk_group,
172 + num_expert_group=num_expert_group,
173 + global_num_experts=global_num_experts,
174 + expert_map=expert_map,
175 + custom_routing_function=custom_routing_function,
176 + scoring_func=scoring_func,
177 + e_score_correction_bias=e_score_correction_bias,
178 + activation=activation,
179 + apply_router_weight_on_input=apply_router_weight_on_input)
180
181 def forward_cuda(
182 self,
@@ -188,6 +193,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
193 custom_routing_function: Optional[Callable] = None,
194 scoring_func: str = "softmax",
195 e_score_correction_bias: Optional[torch.Tensor] = None,
196 + apply_router_weight_on_input: bool = False,
197 activation: str = "silu",
198 ) -> torch.Tensor:
199 topk_weights, topk_ids = FusedMoE.select_experts(
@@ -202,15 +208,17 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
208 scoring_func=scoring_func,
209 e_score_correction_bias=e_score_correction_bias)
210
205 - return fused_experts(hidden_states=x,
206 - w1=layer.w13_weight,
207 - w2=layer.w2_weight,
208 - topk_weights=topk_weights,
209 - topk_ids=topk_ids,
210 - inplace=True,
211 - activation=activation,
212 - global_num_experts=global_num_experts,
213 - expert_map=expert_map)
211 + return fused_experts(
212 + hidden_states=x,
213 + w1=layer.w13_weight,
214 + w2=layer.w2_weight,
215 + topk_weights=topk_weights,
216 + topk_ids=topk_ids,
217 + inplace=True,
218 + activation=activation,
219 + apply_router_weight_on_input=apply_router_weight_on_input,
220 + global_num_experts=global_num_experts,
221 + expert_map=expert_map)
222
223 def forward_cpu(
224 self,
@@ -228,9 +236,11 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
236 scoring_func: str = "softmax",
237 e_score_correction_bias: Optional[torch.Tensor] = None,
238 activation: str = "silu",
239 + apply_router_weight_on_input: bool = False,
240 **kwargs,
241 ):
242 assert activation == "silu", f"{activation} is not supported."
243 + assert apply_router_weight_on_input is False
244 return layer.ipex_fusion(
245 x,
246 use_grouped_topk,
@@ -259,6 +269,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
269 custom_routing_function: Optional[Callable] = None,
270 scoring_func: str = "softmax",
271 e_score_correction_bias: Optional[torch.Tensor] = None,
272 + apply_router_weight_on_input: bool = False,
273 activation: str = "silu",
274 ) -> torch.Tensor:
275 assert not use_grouped_topk
@@ -266,6 +277,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
277 assert topk_group is None
278 assert custom_routing_function is None
279 assert layer is not None
280 + assert apply_router_weight_on_input is False
281 if scoring_func != "softmax":
282 raise NotImplementedError(
283 "Only softmax scoring function is supported for HPU.")
@@ -290,12 +302,14 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
302 custom_routing_function: Optional[Callable] = None,
303 scoring_func: str = "softmax",
304 e_score_correction_bias: Optional[torch.Tensor] = None,
305 + apply_router_weight_on_input: bool = False,
306 activation: str = "silu",
307 ) -> torch.Tensor:
308 assert not use_grouped_topk
309 assert num_expert_group is None
310 assert topk_group is None
311 assert custom_routing_function is None
312 + assert apply_router_weight_on_input is False
313 if scoring_func != "softmax":
314 raise NotImplementedError(
315 "Only softmax scoring function is supported for TPU.")
@@ -401,6 +415,7 @@ class FusedMoE(torch.nn.Module):
415 custom_routing_function: Optional[Callable] = None,
416 scoring_func: str = "softmax",
417 e_score_correction_bias: Optional[torch.Tensor] = None,
418 + apply_router_weight_on_input: bool = False,
419 activation: str = "silu",
420 ):
421 super().__init__()
@@ -486,6 +501,7 @@ class FusedMoE(torch.nn.Module):
501 self.quant_method = quant_config.get_quant_method(self, prefix)
502 assert self.quant_method is not None
503
504 + self.apply_router_weight_on_input = apply_router_weight_on_input
505 moe_quant_params = {
506 "num_experts": self.local_num_experts,
507 "hidden_size": hidden_size,
@@ -853,6 +869,7 @@ class FusedMoE(torch.nn.Module):
869 scoring_func=self.scoring_func,
870 e_score_correction_bias=self.e_score_correction_bias,
871 activation=self.activation,
872 + apply_router_weight_on_input=self.apply_router_weight_on_input,
873 )
874
875 if self.dp_size > 1:
vllm/model_executor/layers/layernorm.py
+5 -2
@@ -92,6 +92,7 @@ class RMSNorm(CustomOp):
92 eps: float = 1e-6,
93 var_hidden_size: Optional[int] = None,
94 has_weight: bool = True,
95 + dtype: Optional[torch.dtype] = None,
96 ) -> None:
97 super().__init__()
98
@@ -100,8 +101,10 @@ class RMSNorm(CustomOp):
101 self.variance_size_override = (None if var_hidden_size == hidden_size
102 else var_hidden_size)
103 self.has_weight = has_weight
103 -
104 - self.weight = torch.ones(hidden_size)
104 + if dtype is not None:
105 + self.weight = torch.ones(hidden_size, dtype=dtype)
106 + else:
107 + self.weight = torch.ones(hidden_size)
108 if self.has_weight:
109 self.weight = nn.Parameter(self.weight)
110
vllm/model_executor/layers/quantization/awq_marlin.py
+5
@@ -469,6 +469,7 @@ class AWQMoEMethod(FusedMoEMethodBase):
469 custom_routing_function: Optional[Callable] = None,
470 scoring_func: str = "softmax",
471 e_score_correction_bias: Optional[torch.Tensor] = None,
472 + apply_router_weight_on_input: bool = False,
473 activation: str = "silu",
474 ) -> torch.Tensor:
475 assert activation == "silu", "Only SiLU activation is supported."
@@ -476,6 +477,10 @@ class AWQMoEMethod(FusedMoEMethodBase):
477 raise NotImplementedError(
478 "Expert Parallelism is not supported for "
479 "fused Marlin MoE method.")
480 + if apply_router_weight_on_input:
481 + raise NotImplementedError(
482 + "Apply router weight on input is not supported for"
483 + "fused Marlin MoE method.")
484
485 topk_weights, topk_ids = FusedMoE.select_experts(
486 hidden_states=x,
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+24 -14
@@ -224,6 +224,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
224 custom_routing_function: Optional[Callable] = None,
225 scoring_func: str = "softmax",
226 e_score_correction_bias: Optional[torch.Tensor] = None,
227 + apply_router_weight_on_input: bool = False,
228 activation: str = "silu",
229 ) -> torch.Tensor:
230 from vllm.model_executor.layers.fused_moe import fused_experts
@@ -240,20 +241,22 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
241 scoring_func=scoring_func,
242 e_score_correction_bias=e_score_correction_bias)
243
243 - return fused_experts(x,
244 - layer.w13_weight,
245 - layer.w2_weight,
246 - topk_weights=topk_weights,
247 - topk_ids=topk_ids,
248 - inplace=True,
249 - activation=activation,
250 - use_fp8_w8a8=True,
251 - global_num_experts=global_num_experts,
252 - expert_map=expert_map,
253 - w1_scale=layer.w13_weight_scale,
254 - w2_scale=layer.w2_weight_scale,
255 - a1_scale=layer.w13_input_scale,
256 - a2_scale=layer.w2_input_scale)
244 + return fused_experts(
245 + x,
246 + layer.w13_weight,
247 + layer.w2_weight,
248 + topk_weights=topk_weights,
249 + topk_ids=topk_ids,
250 + inplace=True,
251 + activation=activation,
252 + apply_router_weight_on_input=apply_router_weight_on_input,
253 + use_fp8_w8a8=True,
254 + global_num_experts=global_num_experts,
255 + expert_map=expert_map,
256 + w1_scale=layer.w13_weight_scale,
257 + w2_scale=layer.w2_weight_scale,
258 + a1_scale=layer.w13_input_scale,
259 + a2_scale=layer.w2_input_scale)
260
261
262 class CompressedTensorsW8A8Fp8MoECutlassMethod(CompressedTensorsMoEMethod):
@@ -438,6 +441,7 @@ class CompressedTensorsW8A8Fp8MoECutlassMethod(CompressedTensorsMoEMethod):
441 custom_routing_function: Optional[Callable] = None,
442 scoring_func: str = "softmax",
443 e_score_correction_bias: Optional[torch.Tensor] = None,
444 + apply_router_weight_on_input: bool = False,
445 activation: str = "silu",
446 ) -> torch.Tensor:
447
@@ -474,6 +478,7 @@ class CompressedTensorsW8A8Fp8MoECutlassMethod(CompressedTensorsMoEMethod):
478 a1_scale=layer.w13_input_scale,
479 a2_scale=layer.w2_input_scale,
480 out_dtype=x.dtype,
481 + apply_router_weight_on_input=apply_router_weight_on_input,
482 )
483
484
@@ -778,6 +783,7 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
783 custom_routing_function: Optional[Callable] = None,
784 scoring_func: str = "softmax",
785 e_score_correction_bias: Optional[torch.Tensor] = None,
786 + apply_router_weight_on_input: bool = False,
787 activation: str = "silu",
788 ) -> torch.Tensor:
789 assert activation == "silu", "Only SiLU activation is supported."
@@ -785,6 +791,10 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
791 raise NotImplementedError(
792 "Expert Parallelism is not supported for "
793 "fused Marlin MoE method.")
794 + if apply_router_weight_on_input:
795 + raise NotImplementedError(
796 + "Apply router weight on input is not supported for "
797 + "fused Marlin MoE method.")
798
799 topk_weights, topk_ids = FusedMoE.select_experts(
800 hidden_states=x,
vllm/model_executor/layers/quantization/experts_int8.py
+15 -12
@@ -113,6 +113,7 @@ class ExpertsInt8MoEMethod(FusedMoEMethodBase):
113 custom_routing_function: Optional[Callable] = None,
114 scoring_func: str = "softmax",
115 e_score_correction_bias: Optional[torch.Tensor] = None,
116 + apply_router_weight_on_input: bool = False,
117 activation: str = "silu",
118 ) -> torch.Tensor:
119 from vllm.model_executor.layers.fused_moe import fused_experts
@@ -129,18 +130,20 @@ class ExpertsInt8MoEMethod(FusedMoEMethodBase):
130 scoring_func=scoring_func,
131 e_score_correction_bias=e_score_correction_bias)
132
132 - return fused_experts(x,
133 - layer.w13_weight,
134 - layer.w2_weight,
135 - topk_weights=topk_weights,
136 - topk_ids=topk_ids,
137 - inplace=True,
138 - activation=activation,
139 - use_int8_w8a16=True,
140 - global_num_experts=global_num_experts,
141 - expert_map=expert_map,
142 - w1_scale=layer.w13_scale,
143 - w2_scale=layer.w2_scale)
133 + return fused_experts(
134 + x,
135 + layer.w13_weight,
136 + layer.w2_weight,
137 + topk_weights=topk_weights,
138 + topk_ids=topk_ids,
139 + inplace=True,
140 + activation=activation,
141 + use_int8_w8a16=True,
142 + global_num_experts=global_num_experts,
143 + apply_router_weight_on_input=apply_router_weight_on_input,
144 + expert_map=expert_map,
145 + w1_scale=layer.w13_scale,
146 + w2_scale=layer.w2_scale)
147
148 @staticmethod
149 def quantizing_weight_loader(layer, weight_loader):
vllm/model_executor/layers/quantization/fp8.py
+2
@@ -773,6 +773,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
773 custom_routing_function: Optional[Callable] = None,
774 scoring_func: str = "softmax",
775 e_score_correction_bias: Optional[torch.Tensor] = None,
776 + apply_router_weight_on_input: bool = False,
777 activation: str = "silu",
778 ) -> torch.Tensor:
779 from vllm.model_executor.layers.fused_moe import fused_experts
@@ -800,6 +801,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
801 activation=activation,
802 use_fp8_w8a8=True,
803 global_num_experts=global_num_experts,
804 + apply_router_weight_on_input=apply_router_weight_on_input,
805 expert_map=expert_map,
806 w1_scale=(layer.w13_weight_scale_inv
807 if self.block_quant else layer.w13_weight_scale),
vllm/model_executor/layers/quantization/gguf.py
+6
@@ -338,9 +338,15 @@ class GGUFMoEMethod(FusedMoEMethodBase):
338 custom_routing_function: Optional[Callable] = None,
339 scoring_func: str = "softmax",
340 e_score_correction_bias: Optional[torch.Tensor] = None,
341 + apply_router_weight_on_input: bool = False,
342 activation: str = "silu",
343 ):
344 assert activation == "silu", "Only SiLU activation is supported."
345 + if apply_router_weight_on_input:
346 + raise NotImplementedError(
347 + "Apply router weight on input is not supported for"
348 + "fused GGUF MoE method.")
349 +
350 topk_weights, topk_ids = FusedMoE.select_experts(
351 hidden_states=x,
352 router_logits=router_logits,
vllm/model_executor/layers/quantization/gptq_marlin.py
+5
@@ -592,9 +592,14 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
592 custom_routing_function: Optional[Callable] = None,
593 scoring_func: str = "softmax",
594 e_score_correction_bias: Optional[torch.Tensor] = None,
595 + apply_router_weight_on_input: bool = False,
596 activation: str = "silu",
597 ) -> torch.Tensor:
598 assert activation == "silu", "Only SiLU activation is supported."
599 + if apply_router_weight_on_input is not None:
600 + raise NotImplementedError(
601 + "Apply router weight on input is not supported for"
602 + "fused Marlin MoE method.")
603
604 # The input must currently be float16
605 orig_dtype = x.dtype
vllm/model_executor/layers/quantization/moe_wna16.py
+18 -15
@@ -293,6 +293,7 @@ class MoeWNA16Method(FusedMoEMethodBase):
293 custom_routing_function: Optional[Callable] = None,
294 scoring_func: str = "softmax",
295 e_score_correction_bias: Optional[torch.Tensor] = None,
296 + apply_router_weight_on_input: bool = False,
297 activation: str = "silu",
298 ) -> torch.Tensor:
299 from vllm.model_executor.layers.fused_moe import fused_experts
@@ -312,21 +313,23 @@ class MoeWNA16Method(FusedMoEMethodBase):
313 weight_bits = self.quant_config.weight_bits
314 has_zp = self.quant_config.has_zp
315
315 - return fused_experts(x,
316 - layer.w13_qweight,
317 - layer.w2_qweight,
318 - topk_weights=topk_weights,
319 - topk_ids=topk_ids,
320 - inplace=True,
321 - use_int4_w4a16=weight_bits == 4,
322 - use_int8_w8a16=weight_bits == 8,
323 - global_num_experts=global_num_experts,
324 - expert_map=expert_map,
325 - w1_scale=layer.w13_scales,
326 - w2_scale=layer.w2_scales,
327 - w1_zp=layer.w13_qzeros if has_zp else None,
328 - w2_zp=layer.w2_qzeros if has_zp else None,
329 - block_shape=[0, layer.group_size])
316 + return fused_experts(
317 + x,
318 + layer.w13_qweight,
319 + layer.w2_qweight,
320 + topk_weights=topk_weights,
321 + topk_ids=topk_ids,
322 + inplace=True,
323 + use_int4_w4a16=weight_bits == 4,
324 + use_int8_w8a16=weight_bits == 8,
325 + global_num_experts=global_num_experts,
326 + apply_router_weight_on_input=apply_router_weight_on_input,
327 + expert_map=expert_map,
328 + w1_scale=layer.w13_scales,
329 + w2_scale=layer.w2_scales,
330 + w1_zp=layer.w13_qzeros if has_zp else None,
331 + w2_zp=layer.w2_qzeros if has_zp else None,
332 + block_shape=[0, layer.group_size])
333
334 @staticmethod
335 def get_weight_loader(layer, weight_loader):
vllm/model_executor/layers/quantization/quark/quark_moe.py
+17 -13
@@ -202,6 +202,8 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
202 custom_routing_function: Optional[Callable] = None,
203 scoring_func: str = "softmax",
204 e_score_correction_bias: Optional[torch.Tensor] = None,
205 + apply_router_weight_on_input: bool = False,
206 + activation: str = "silu",
207 ) -> torch.Tensor:
208 from vllm.model_executor.layers.fused_moe import fused_experts
209
@@ -217,16 +219,18 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
219 scoring_func=scoring_func,
220 e_score_correction_bias=e_score_correction_bias)
221
220 - return fused_experts(x,
221 - layer.w13_weight,
222 - layer.w2_weight,
223 - topk_weights=topk_weights,
224 - topk_ids=topk_ids,
225 - inplace=True,
226 - use_fp8_w8a8=True,
227 - global_num_experts=global_num_experts,
228 - expert_map=expert_map,
229 - w1_scale=layer.w13_weight_scale,
230 - w2_scale=layer.w2_weight_scale,
231 - a1_scale=layer.w13_input_scale,
232 - a2_scale=layer.w2_input_scale)
222 + return fused_experts(
223 + x,
224 + layer.w13_weight,
225 + layer.w2_weight,
226 + topk_weights=topk_weights,
227 + topk_ids=topk_ids,
228 + inplace=True,
229 + use_fp8_w8a8=True,
230 + global_num_experts=global_num_experts,
231 + apply_router_weight_on_input=apply_router_weight_on_input,
232 + expert_map=expert_map,
233 + w1_scale=layer.w13_weight_scale,
234 + w2_scale=layer.w2_weight_scale,
235 + a1_scale=layer.w13_input_scale,
236 + a2_scale=layer.w2_input_scale)
vllm/model_executor/layers/rotary_embedding.py
+68
@@ -851,6 +851,70 @@ class Llama3RotaryEmbedding(RotaryEmbedding):
851 return new_freqs
852
853
854 +class Llama4VisionRotaryEmbedding(RotaryEmbedding):
855 +
856 + def __init__(
857 + self,
858 + head_size: int,
859 + rotary_dim: int,
860 + max_position_embeddings: int,
861 + base: int,
862 + is_neox_style: bool,
863 + dtype: torch.dtype,
864 + ):
865 + super().__init__(head_size, rotary_dim, max_position_embeddings, base,
866 + is_neox_style, dtype)
867 +
868 + def _compute_inv_freq(self, base: Union[int, float]) -> torch.Tensor:
869 + inv_freqs = super()._compute_inv_freq(base)
870 + inv_freqs = inv_freqs[:(self.rotary_dim // 2)]
871 + return inv_freqs
872 +
873 + def _compute_cos_sin_cache(self) -> torch.Tensor:
874 + inv_freq = self._compute_inv_freq(self.base)
875 +
876 + # self.max_position_embeddings here is number of image patches
877 + # i.e. (image_size // patch_size) ** 2
878 + num_patches = self.max_position_embeddings
879 + img_idx = torch.arange(num_patches,
880 + dtype=torch.int32) \
881 + .reshape(num_patches, 1)
882 + img_idx = torch.cat([img_idx, img_idx[:1]], dim=0)
883 + img_idx[-1, -1] = -2 # set to ID_CLS_TOKEN
884 + num_patches_single_dim = int(math.sqrt(num_patches))
885 + frequencies_x = img_idx % num_patches_single_dim
886 + frequencies_y = img_idx // num_patches_single_dim
887 + freqs_x = ((frequencies_x + 1)[..., None] *
888 + inv_freq[None, None, :]).repeat_interleave(2, dim=-1)
889 + freqs_y = ((frequencies_y + 1)[..., None] *
890 + inv_freq[None, None, :]).repeat_interleave(2, dim=-1)
891 + freqs = torch.cat([freqs_x, freqs_y],
892 + dim=-1).float().contiguous()[..., ::2]
893 + freqs = freqs.masked_fill(img_idx.reshape(-1, 1, 1) < 0, 0)
894 + cache = torch.view_as_complex(
895 + torch.stack([torch.cos(freqs), torch.sin(freqs)], dim=-1))
896 + return cache
897 +
898 + def forward(
899 + self,
900 + query: torch.Tensor,
901 + key: torch.Tensor,
902 + ) -> Tuple[torch.Tensor, torch.Tensor]:
903 + self.cos_sin_cache: torch.Tensor = self.cos_sin_cache.to(query.device)
904 + query_ = torch.view_as_complex(query.float().reshape(
905 + *query.shape[:-1], -1, 2))
906 + key_ = torch.view_as_complex(key.float().reshape(
907 + *key.shape[:-1], -1, 2))
908 + broadcast_shape = [
909 + d if i == 1 or i == (query_.ndim - 1) else 1
910 + for i, d in enumerate(query_.shape)
911 + ]
912 + freqs_ci = self.cos_sin_cache.view(*broadcast_shape)
913 + query_out = torch.view_as_real(query_ * freqs_ci).flatten(3)
914 + key_out = torch.view_as_real(key_ * freqs_ci).flatten(3)
915 + return query_out.type_as(query), key_out.type_as(key)
916 +
917 +
918 class MRotaryEmbedding(RotaryEmbedding):
919 """Rotary Embedding with Multimodal Sections."""
920
@@ -1130,6 +1194,10 @@ def get_rope(
1194 scaling_factor, low_freq_factor,
1195 high_freq_factor,
1196 original_max_position)
1197 + elif scaling_type == "mllama4":
1198 + rotary_emb = Llama4VisionRotaryEmbedding(head_size, rotary_dim,
1199 + max_position, base,
1200 + is_neox_style, dtype)
1201 elif scaling_type == "default":
1202 if "mrope_section" in rope_scaling:
1203 rotary_emb = MRotaryEmbedding(
vllm/model_executor/model_loader/loader.py
+3 -1
@@ -111,10 +111,12 @@ def _initialize_model(
111 vllm_config: VllmConfig,
112 *,
113 prefix: str = "",
114 + model_class: Optional[type[nn.Module]] = None,
115 ) -> nn.Module:
116 """Initialize a model with the given configurations."""
117 model_config = vllm_config.model_config
117 - model_class, _ = get_model_architecture(model_config)
118 + if model_class is None:
119 + model_class, _ = get_model_architecture(model_config)
120
121 if vllm_config.quant_config is not None:
122 configure_quant_config(vllm_config.quant_config, model_class)
vllm/model_executor/models/llama.py
+19 -7
@@ -22,7 +22,7 @@
22 # See the License for the specific language governing permissions and
23 # limitations under the License.
24 """Inference-only LLaMA model compatible with HuggingFace weights."""
25 -from typing import Any, Dict, Iterable, Optional, Set, Tuple, Type, Union
25 +from typing import Any, Dict, Iterable, Optional, Set, Tuple, Union
26
27 import torch
28 from torch import nn
@@ -65,6 +65,7 @@ class LlamaMLP(nn.Module):
65 quant_config: Optional[QuantizationConfig] = None,
66 bias: bool = False,
67 prefix: str = "",
68 + reduce_results: bool = True,
69 ) -> None:
70 super().__init__()
71 self.gate_up_proj = MergedColumnParallelLinear(
@@ -79,6 +80,7 @@ class LlamaMLP(nn.Module):
80 output_size=hidden_size,
81 bias=bias,
82 quant_config=quant_config,
83 + reduce_results=reduce_results,
84 prefix=f"{prefix}.down_proj",
85 )
86 if hidden_act != "silu":
@@ -292,7 +294,7 @@ class LlamaModel(nn.Module):
294 *,
295 vllm_config: VllmConfig,
296 prefix: str = "",
295 - layer_type: Type[LlamaDecoderLayer] = LlamaDecoderLayer):
297 + layer_type: type[nn.Module] = LlamaDecoderLayer):
298 super().__init__()
299
300 config = vllm_config.model_config.hf_config
@@ -466,10 +468,14 @@ class LlamaForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
468 "ffn_norm": "post_attention_layernorm",
469 "tok_embeddings": "model.embed_tokens",
470 "output": "lm_head",
469 - "norm": "model.norm"
471 + "norm": "model.norm",
472 }
473
472 - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
474 + def __init__(self,
475 + *,
476 + vllm_config: VllmConfig,
477 + prefix: str = "",
478 + layer_type: type[nn.Module] = LlamaDecoderLayer):
479 super().__init__()
480 config = vllm_config.model_config.hf_config
481 quant_config = vllm_config.quant_config
@@ -478,7 +484,8 @@ class LlamaForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
484 self.lora_config = lora_config
485
486 self.model = self._init_model(vllm_config=vllm_config,
481 - prefix=maybe_prefix(prefix, "model"))
487 + prefix=maybe_prefix(prefix, "model"),
488 + layer_type=layer_type)
489
490 if get_pp_group().is_last_rank:
491 self.unpadded_vocab_size = config.vocab_size
@@ -513,8 +520,13 @@ class LlamaForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
520 self.make_empty_intermediate_tensors = (
521 self.model.make_empty_intermediate_tensors)
522
516 - def _init_model(self, vllm_config: VllmConfig, prefix: str = ""):
517 - return LlamaModel(vllm_config=vllm_config, prefix=prefix)
523 + def _init_model(self,
524 + vllm_config: VllmConfig,
525 + prefix: str = "",
526 + layer_type: type[nn.Module] = LlamaDecoderLayer):
527 + return LlamaModel(vllm_config=vllm_config,
528 + prefix=prefix,
529 + layer_type=layer_type)
530
531 def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
532 return self.model.get_input_embeddings(input_ids)
vllm/model_executor/models/llama4.py new
+531
@@ -0,0 +1,531 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +#
3 +# Copyright 2025 the LLAMA4, Meta Inc., vLLM, and HuggingFace Inc. team.
4 +# All rights reserved.
5 +#
6 +#
7 +# Licensed under the Apache License, Version 2.0 (the "License");
8 +# you may not use this file except in compliance with the License.
9 +# You may obtain a copy of the License at
10 +#
11 +# http://www.apache.org/licenses/LICENSE-2.0
12 +#
13 +# Unless required by applicable law or agreed to in writing, software
14 +# distributed under the License is distributed on an "AS IS" BASIS,
15 +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
16 +# See the License for the specific language governing permissions and
17 +# limitations under the License.
18 +"""Inference-only LLaMA model compatible with HuggingFace weights."""
19 +from typing import Any, Dict, Iterable, List, Optional, Set, Tuple
20 +
21 +import torch
22 +from torch import nn
23 +from transformers import Llama4TextConfig
24 +
25 +from vllm.attention import Attention
26 +from vllm.compilation.decorators import support_torch_compile
27 +from vllm.config import CacheConfig, VllmConfig
28 +from vllm.distributed import (get_tensor_model_parallel_world_size,
29 + tensor_model_parallel_all_reduce)
30 +from vllm.model_executor.layers.fused_moe import FusedMoE
31 +from vllm.model_executor.layers.layernorm import RMSNorm
32 +from vllm.model_executor.layers.linear import (QKVParallelLinear,
33 + ReplicatedLinear,
34 + RowParallelLinear)
35 +from vllm.model_executor.layers.quantization import QuantizationConfig
36 +from vllm.model_executor.layers.rotary_embedding import get_rope
37 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader
38 +
39 +from .llama import LlamaForCausalLM, LlamaMLP, LlamaModel
40 +from .utils import (AutoWeightsLoader, extract_layer_index,
41 + is_pp_missing_parameter)
42 +
43 +
44 +class Llama4MoE(nn.Module):
45 +
46 + @staticmethod
47 + def custom_routing_function(
48 + hidden_states: torch.Tensor,
49 + gating_output: torch.Tensor,
50 + topk: int,
51 + renormalize: bool,
52 + ) -> Tuple[torch.Tensor, torch.Tensor]:
53 + router_scores, router_indices = torch.topk(gating_output, topk, dim=-1)
54 + router_scores = torch.sigmoid(router_scores.float()).to(
55 + hidden_states.dtype)
56 + return (router_scores, router_indices.to(torch.int32))
57 +
58 + def __init__(self,
59 + config: Llama4TextConfig,
60 + quant_config: Optional[QuantizationConfig] = None,
61 + prefix: str = ""):
62 + super().__init__()
63 + self.tp_size = get_tensor_model_parallel_world_size()
64 + self.top_k = config.num_experts_per_tok
65 +
66 + intermediate_size_moe = config.intermediate_size
67 + self.router = ReplicatedLinear(config.hidden_size,
68 + config.num_local_experts,
69 + bias=False,
70 + quant_config=None,
71 + prefix=f"{prefix}.router")
72 +
73 + self.experts = FusedMoE(
74 + num_experts=config.num_local_experts,
75 + top_k=config.num_experts_per_tok,
76 + hidden_size=config.hidden_size,
77 + custom_routing_function=Llama4MoE.custom_routing_function,
78 + intermediate_size=intermediate_size_moe,
79 + apply_router_weight_on_input=True,
80 + reduce_results=False,
81 + renormalize=False,
82 + quant_config=quant_config,
83 + prefix=f"{prefix}.experts")
84 +
85 + self.shared_expert = LlamaMLP(
86 + hidden_size=config.hidden_size,
87 + intermediate_size=intermediate_size_moe,
88 + hidden_act="silu",
89 + quant_config=quant_config,
90 + bias=False,
91 + prefix=f"{prefix}.shared_expert",
92 + reduce_results=False, # We need to do scatter before reduce
93 + )
94 +
95 + def forward(self, hidden_states):
96 + router_logits, _ = self.router(hidden_states)
97 + shared_out = self.shared_expert(hidden_states)
98 + routed_out = self.experts(
99 + hidden_states=hidden_states,
100 + router_logits=router_logits,
101 + )
102 + experts_out = routed_out + shared_out
103 +
104 + if self.tp_size > 1:
105 + experts_out = tensor_model_parallel_all_reduce(experts_out)
106 +
107 + return experts_out
108 +
109 +
110 +class Llama4Attention(nn.Module):
111 +
112 + def __init__(self,
113 + config: Llama4TextConfig,
114 + hidden_size: int,
115 + num_heads: int,
116 + num_kv_heads: int,
117 + rope_theta: float = 10000,
118 + rope_scaling: Optional[Dict[str, Any]] = None,
119 + max_position_embeddings: int = 8192,
120 + quant_config: Optional[QuantizationConfig] = None,
121 + bias: bool = False,
122 + bias_o_proj: bool = False,
123 + cache_config: Optional[CacheConfig] = None,
124 + prefix: str = "") -> None:
125 + super().__init__()
126 + self.layer_idx = extract_layer_index(prefix)
127 + self.hidden_size = hidden_size
128 + self.no_rope_layers = config.no_rope_layers
129 + self.nope = self.no_rope_layers[self.layer_idx] == 0
130 + self.use_qk_norm = config.use_qk_norm and not self.nope
131 + tp_size = get_tensor_model_parallel_world_size()
132 + self.total_num_heads = num_heads
133 + assert self.total_num_heads % tp_size == 0
134 + self.num_heads = self.total_num_heads // tp_size
135 + self.total_num_kv_heads = num_kv_heads
136 + if self.total_num_kv_heads >= tp_size:
137 + # Number of KV heads is greater than TP size, so we partition
138 + # the KV heads across multiple tensor parallel GPUs.
139 + assert self.total_num_kv_heads % tp_size == 0
140 + else:
141 + # Number of KV heads is less than TP size, so we replicate
142 + # the KV heads across multiple tensor parallel GPUs.
143 + assert tp_size % self.total_num_kv_heads == 0
144 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
145 + self.head_dim = config.head_dim
146 + self.q_size = self.num_heads * self.head_dim
147 + self.kv_size = self.num_kv_heads * self.head_dim
148 + self.scaling = self.head_dim**-0.5
149 + # TODO: attn_temperature_tuning should be a bool in huggingface
150 + self.attn_temperature_tuning = self.nope and \
151 + config.attn_temperature_tuning > 0
152 +
153 + self.floor_scale = getattr(config, "floor_scale", 8192.0)
154 + self.attn_scale = getattr(config, "attn_scale", 0.1)
155 + self.rope_theta = rope_theta
156 + self.max_position_embeddings = max_position_embeddings
157 + self.n_rep = self.num_heads // self.num_kv_heads
158 + self.q_norm = RMSNorm(
159 + hidden_size=self.q_size,
160 + eps=config.rms_norm_eps,
161 + has_weight=False,
162 + dtype=torch.float32,
163 + ) if self.use_qk_norm else None
164 + self.k_norm = RMSNorm(
165 + hidden_size=self.kv_size,
166 + eps=config.rms_norm_eps,
167 + has_weight=False,
168 + dtype=torch.float32,
169 + ) if self.use_qk_norm else None
170 + self.qkv_proj = QKVParallelLinear(
171 + hidden_size=hidden_size,
172 + head_size=self.head_dim,
173 + total_num_heads=self.total_num_heads,
174 + total_num_kv_heads=self.total_num_kv_heads,
175 + bias=bias,
176 + quant_config=quant_config,
177 + prefix=f"{prefix}.qkv_proj",
178 + )
179 +
180 + self.o_proj = RowParallelLinear(
181 + input_size=self.total_num_heads * self.head_dim,
182 + output_size=hidden_size,
183 + bias=bias_o_proj,
184 + quant_config=quant_config,
185 + prefix=f"{prefix}.o_proj",
186 + )
187 + is_neox_style = True
188 + is_gguf = quant_config and quant_config.get_name() == "gguf"
189 + if is_gguf and config.model_type == "llama":
190 + is_neox_style = False
191 +
192 + self.rotary_emb = get_rope(
193 + self.head_dim,
194 + rotary_dim=self.head_dim,
195 + max_position=max_position_embeddings,
196 + base=int(rope_theta),
197 + rope_scaling=rope_scaling if rope_scaling != "default" else None,
198 + is_neox_style=is_neox_style,
199 + ) if not self.nope else None
200 +
201 + self.attn = Attention(
202 + self.num_heads,
203 + self.head_dim,
204 + self.scaling,
205 + num_kv_heads=self.num_kv_heads,
206 + cache_config=cache_config,
207 + quant_config=quant_config,
208 + per_layer_sliding_window=None,
209 + use_irope=not self.nope,
210 + prefix=f"{prefix}.attn",
211 + )
212 +
213 + def _get_attn_scale(self, positions: torch.Tensor) -> torch.Tensor:
214 + floor = torch.floor((positions + 1.0) / self.floor_scale)
215 + attn_scale = torch.log(floor + 1.0) * self.attn_scale + 1.0
216 +
217 + return attn_scale.unsqueeze(-1)
218 +
219 + def forward(
220 + self,
221 + positions: torch.Tensor,
222 + hidden_states: torch.Tensor,
223 + ) -> torch.Tensor:
224 + qkv, _ = self.qkv_proj(hidden_states)
225 + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
226 +
227 + if self.rotary_emb is not None:
228 + q, k = self.rotary_emb(positions, q, k)
229 + if self.q_norm is not None:
230 + q = self.q_norm(q.float()).to(q.dtype)
231 + if self.k_norm is not None:
232 + k = self.k_norm(k.float()).to(k.dtype)
233 +
234 + # We are applying temperature tuning (https://arxiv.org/abs/2501.19399)
235 + # to NoPE layers, where the inference-time temperature tuning function
236 + # is customized to not affect short context
237 + # while working at very long context
238 + # https://arxiv.org/abs/2501.19399
239 + #
240 + # We should apply temperature tuning between (after) rotary / QK norm
241 + # and (before) attention.
242 + if self.attn_temperature_tuning and self.nope:
243 + attn_scale = self._get_attn_scale(positions)
244 + q = (q * attn_scale).to(q.dtype)
245 + attn_output = self.attn(q, k, v)
246 + output, _ = self.o_proj(attn_output)
247 + return output
248 +
249 +
250 +class Llama4DecoderLayer(nn.Module):
251 +
252 + def __init__(
253 + self,
254 + config: Llama4TextConfig,
255 + cache_config: Optional[CacheConfig] = None,
256 + quant_config: Optional[QuantizationConfig] = None,
257 + prefix: str = "",
258 + ) -> None:
259 + super().__init__()
260 +
261 + self.layer_idx = extract_layer_index(prefix)
262 + self.hidden_size = config.hidden_size
263 + rope_theta = config.rope_theta
264 + rope_scaling = config.rope_scaling
265 + max_position_embeddings = config.max_position_embeddings
266 +
267 + self.self_attn = Llama4Attention(
268 + config=config,
269 + hidden_size=self.hidden_size,
270 + num_heads=config.num_attention_heads,
271 + num_kv_heads=config.num_key_value_heads,
272 + rope_theta=rope_theta,
273 + rope_scaling=rope_scaling,
274 + max_position_embeddings=max_position_embeddings,
275 + quant_config=quant_config,
276 + bias=False,
277 + bias_o_proj=False,
278 + cache_config=cache_config,
279 + prefix=f"{prefix}.self_attn",
280 + )
281 + is_moe_layer = (self.layer_idx +
282 + 1) % config.interleave_moe_layer_step == 0
283 + if is_moe_layer:
284 + self.feed_forward = Llama4MoE(
285 + config=config,
286 + quant_config=quant_config,
287 + prefix=f"{prefix}.feed_forward",
288 + )
289 + else:
290 + self.feed_forward = LlamaMLP(
291 + hidden_size=self.hidden_size,
292 + intermediate_size=config.intermediate_size_mlp,
293 + hidden_act="silu",
294 + quant_config=quant_config,
295 + bias=False,
296 + prefix=f"{prefix}.feed_forward",
297 + )
298 + self.input_layernorm = RMSNorm(config.hidden_size,
299 + eps=config.rms_norm_eps)
300 + self.post_attention_layernorm = RMSNorm(config.hidden_size,
301 + eps=config.rms_norm_eps)
302 +
303 + def forward(
304 + self,
305 + positions: torch.Tensor,
306 + hidden_states: torch.Tensor,
307 + residual: Optional[torch.Tensor],
308 + ) -> Tuple[torch.Tensor, torch.Tensor]:
309 + # Self Attention
310 + if residual is None:
311 + residual = hidden_states
312 + hidden_states = self.input_layernorm(hidden_states)
313 + else:
314 + hidden_states, residual = self.input_layernorm(
315 + hidden_states, residual)
316 + hidden_states = self.self_attn(positions=positions,
317 + hidden_states=hidden_states)
318 +
319 + # Fully Connected
320 + hidden_states, residual = self.post_attention_layernorm(
321 + hidden_states, residual)
322 + hidden_states = self.feed_forward(hidden_states)
323 + return hidden_states, residual
324 +
325 +
326 +@support_torch_compile
327 +class Llama4Model(LlamaModel):
328 +
329 + def __init__(self,
330 + *,
331 + vllm_config: VllmConfig,
332 + prefix: str = "",
333 + layer_type: type[Llama4DecoderLayer] = Llama4DecoderLayer):
334 + self.num_experts = vllm_config.model_config.hf_config.num_local_experts
335 + super().__init__(vllm_config=vllm_config,
336 + prefix=prefix,
337 + layer_type=layer_type)
338 +
339 + def load_moe_expert_weights(
340 + self,
341 + name: str,
342 + loaded_weight: torch.Tensor,
343 + params_dict: Dict[str, nn.Parameter],
344 + loaded_params: Set[str],
345 + expert_params_mapping: List[Tuple[str, str, int, str]],
346 + fused: bool = True,
347 + ) -> bool:
348 + expert_param_loaded = False
349 + if "experts.gate_up_proj" in name:
350 + loaded_weight = loaded_weight.chunk(2, dim=-1)
351 + for (param_name, weight_name, expert_id,
352 + shard_id) in expert_params_mapping:
353 + new_loaded_weight = loaded_weight
354 + if fused:
355 + e_str, _, proj_str, _ = weight_name.split('.')
356 + weight_name = f"{e_str}.{proj_str}"
357 + param_name = f"{param_name}weight"
358 + if weight_name not in name:
359 + continue
360 + full_param_name = name.replace(weight_name, param_name)
361 + # Skip layers on other devices.
362 + if is_pp_missing_parameter(name, self):
363 + continue
364 + if ((name.endswith(".bias") or name.endswith("_bias"))
365 + and name not in params_dict):
366 + continue
367 + param = params_dict[full_param_name]
368 + weight_loader = param.weight_loader
369 + if fused:
370 + if "w13" in full_param_name:
371 + shard_idx = 0 if shard_id == "w1" else 1
372 + new_loaded_weight = new_loaded_weight[shard_idx]
373 + new_loaded_weight = new_loaded_weight.transpose(-1, -2)
374 + layer_idx = extract_layer_index(name)
375 + # EP mapping
376 + expert_map = self.layers[
377 + layer_idx].feed_forward.experts.expert_map
378 + if expert_map is not None:
379 + local_expert_indices = (expert_map != -1) \
380 + .nonzero() \
381 + .flatten() \
382 + .to(new_loaded_weight.device)
383 + new_loaded_weight = new_loaded_weight[local_expert_indices]
384 + expert_id = local_expert_indices[0].item()
385 + else:
386 + # TODO: add EP support for non fused weights
387 + pass
388 + weight_loader(param,
389 + new_loaded_weight,
390 + full_param_name,
391 + shard_id=shard_id,
392 + expert_id=expert_id)
393 +
394 + loaded_params.add(full_param_name)
395 + expert_param_loaded = True
396 + return expert_param_loaded
397 +
398 + def load_weights(self, weights: Iterable[Tuple[str,
399 + torch.Tensor]]) -> Set[str]:
400 + stacked_params_mapping = [
401 + # (param_name, shard_name, shard_id)
402 + (".qkv_proj", ".q_proj", "q"),
403 + (".qkv_proj", ".k_proj", "k"),
404 + (".qkv_proj", ".v_proj", "v"),
405 + (".gate_up_proj", ".gate_proj", 0),
406 + (".gate_up_proj", ".up_proj", 1),
407 + ]
408 + fused_experts_params = False
409 + expert_params_mapping = FusedMoE.make_expert_params_mapping(
410 + ckpt_gate_proj_name="gate_proj",
411 + ckpt_down_proj_name="down_proj",
412 + ckpt_up_proj_name="up_proj",
413 + num_experts=self.num_experts)
414 + expert_params_mapping_fused = FusedMoE.make_expert_params_mapping(
415 + ckpt_gate_proj_name="gate_up_proj",
416 + ckpt_down_proj_name="down_proj",
417 + ckpt_up_proj_name="gate_up_proj",
418 + num_experts=1)
419 + params_dict = dict(self.named_parameters())
420 + loaded_params: Set[str] = set()
421 + for name, loaded_weight in weights:
422 + if "experts.gate_up_proj" in name or "experts.down_proj" in name:
423 + fused_experts_params = True
424 + expert_params_mapping = expert_params_mapping_fused
425 + if (self.quant_config is not None and
426 + (scale_name := self.quant_config.get_cache_scale(name))):
427 + # Loading kv cache quantization scales
428 + param = params_dict[scale_name]
429 + weight_loader = getattr(param, "weight_loader",
430 + default_weight_loader)
431 + loaded_weight = (loaded_weight if loaded_weight.dim() == 0 else
432 + loaded_weight[0])
433 + weight_loader(param, loaded_weight)
434 + loaded_params.add(scale_name)
435 + continue
436 + for param_name, weight_name, shard_id in stacked_params_mapping:
437 + if weight_name not in name or "experts" in name:
438 + continue
439 + name = name.replace(weight_name, param_name)
440 + if is_pp_missing_parameter(name, self):
441 + continue
442 + param = params_dict[name]
443 + weight_loader = param.weight_loader
444 + weight_loader(param, loaded_weight, shard_id)
445 + loaded_params.add(name)
446 + break
447 + else:
448 + moe_loaded = self.load_moe_expert_weights(
449 + name,
450 + loaded_weight,
451 + params_dict,
452 + loaded_params,
453 + expert_params_mapping,
454 + fused=fused_experts_params)
455 +
456 + if not moe_loaded:
457 + if is_pp_missing_parameter(name, self):
458 + continue
459 + param = params_dict[name]
460 + weight_loader = getattr(param, "weight_loader",
461 + default_weight_loader)
462 + weight_loader(param, loaded_weight)
463 + loaded_params.add(name)
464 + return loaded_params
465 +
466 +
467 +class Llama4ForCausalLM(LlamaForCausalLM):
468 +
469 + packed_modules_mapping = {
470 + "qkv_proj": ["q_proj", "k_proj", "v_proj"],
471 + "gate_up_proj": ["gate_proj", "up_proj"],
472 + }
473 +
474 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
475 + # Update temperature tuning config from generation config
476 + gen_config = vllm_config.model_config.try_get_generation_config()
477 + gen_config.update(vllm_config.model_config.override_generation_config)
478 + vllm_config.model_config.hf_config.attn_temperature_tuning \
479 + = gen_config.get("attn_temperature_tuning", False)
480 +
481 + super().__init__(vllm_config=vllm_config,
482 + prefix=prefix,
483 + layer_type=Llama4DecoderLayer)
484 +
485 + def _init_model(self,
486 + vllm_config: VllmConfig,
487 + prefix: str = "",
488 + layer_type: type[Llama4DecoderLayer] = Llama4DecoderLayer):
489 + return Llama4Model(vllm_config=vllm_config,
490 + prefix=prefix,
491 + layer_type=layer_type)
492 +
493 + def load_weights(self, weights: Iterable[Tuple[str,
494 + torch.Tensor]]) -> Set[str]:
495 + loader = AutoWeightsLoader(
496 + self,
497 + skip_prefixes=(["lm_head."]
498 + if self.config.tie_word_embeddings else None),
499 + )
500 + weights = [
501 + self.permute_qk_weight_for_rotary(name, loaded_weight)
502 + for name, loaded_weight in weights
503 + ]
504 + return loader.load_weights(weights)
505 +
506 + def permute_qk_weight_for_rotary(
507 + self,
508 + name: str,
509 + loaded_weight: torch.Tensor,
510 + ) -> Tuple[str, torch.Tensor]:
511 +
512 + def permute(w: torch.Tensor, n_heads: int):
513 + attn_in = self.config.head_dim * n_heads
514 + attn_out = self.config.hidden_size
515 +
516 + return w.view(n_heads, attn_in // n_heads // 2, 2,
517 + attn_out).transpose(1, 2).reshape(attn_in, attn_out)
518 +
519 + modules = name.split(".")
520 +
521 + # rotary embeds should be sliced
522 + if ("wk" in modules or "k_proj" in modules) \
523 + and modules[-1] == "weight":
524 + loaded_weight = permute(loaded_weight,
525 + self.config.num_key_value_heads)
526 + elif ("wq" in modules or "q_proj" in modules) \
527 + and modules[-1] == "weight":
528 + loaded_weight = permute(loaded_weight,
529 + self.config.num_attention_heads)
530 +
531 + return name, loaded_weight
vllm/model_executor/models/mllama4.py new
+895
@@ -0,0 +1,895 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +#
3 +# Copyright 2025 the LLAMA4, Meta Inc., vLLM, and HuggingFace Inc. team.
4 +# All rights reserved.
5 +#
6 +#
7 +# Licensed under the Apache License, Version 2.0 (the "License");
8 +# you may not use this file except in compliance with the License.
9 +# You may obtain a copy of the License at
10 +#
11 +# http://www.apache.org/licenses/LICENSE-2.0
12 +#
13 +# Unless required by applicable law or agreed to in writing, software
14 +# distributed under the License is distributed on an "AS IS" BASIS,
15 +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
16 +# See the License for the specific language governing permissions and
17 +# limitations under the License.
18 +import math
19 +from collections.abc import Iterable, Mapping
20 +from functools import cached_property
21 +from itertools import tee
22 +from typing import List, Literal, Optional, Set, Tuple, TypedDict, Union
23 +
24 +import torch
25 +from torch import nn
26 +from transformers import BatchFeature, Llama4Config, Llama4VisionConfig
27 +from transformers.image_utils import SizeDict
28 +from transformers.models.llama4 import Llama4Processor
29 +from transformers.models.llama4.image_processing_llama4_fast import (
30 + find_supported_resolutions, get_best_fit)
31 +
32 +from vllm.attention.layer import MultiHeadAttention
33 +from vllm.config import VllmConfig
34 +from vllm.distributed import get_tensor_model_parallel_world_size
35 +from vllm.inputs import InputProcessingContext
36 +from vllm.logger import init_logger
37 +from vllm.model_executor.layers.linear import (ColumnParallelLinear,
38 + QKVParallelLinear,
39 + RowParallelLinear)
40 +from vllm.model_executor.layers.quantization import QuantizationConfig
41 +from vllm.model_executor.layers.rotary_embedding import get_rope
42 +from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
43 +from vllm.model_executor.model_loader.loader import _initialize_model
44 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader
45 +from vllm.model_executor.sampling_metadata import SamplingMetadata
46 +from vllm.multimodal import MULTIMODAL_REGISTRY
47 +from vllm.multimodal.inputs import (MultiModalFieldConfig, MultiModalKwargs,
48 + NestedTensors)
49 +from vllm.multimodal.parse import (ImageProcessorItems, ImageSize,
50 + MultiModalDataItems)
51 +from vllm.multimodal.processing import (BaseMultiModalProcessor,
52 + BaseProcessingInfo, PromptReplacement,
53 + PromptUpdate)
54 +from vllm.multimodal.profiling import BaseDummyInputsBuilder, ProcessorInputs
55 +from vllm.sequence import IntermediateTensors
56 +
57 +from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
58 +from .llama4 import Llama4ForCausalLM
59 +from .utils import (AutoWeightsLoader, flatten_bn, maybe_prefix,
60 + merge_multimodal_embeddings)
61 +from .vision import scatter_patch_features, select_patch_features
62 +
63 +logger = init_logger(__name__)
64 +
65 +
66 +class Llama4ImagePatchInputs(TypedDict):
67 + type: Literal["pixel_values"]
68 + flat_data: torch.Tensor
69 + """
70 + Shape:
71 + `(batch_size * num_chunks, num_channels, image size, image size)`
72 + """
73 + patches_per_image: torch.Tensor
74 + """
75 + The number of total patches for each image in the batch.
76 +
77 + This is used to split the embeddings which has the first two dimensions
78 + flattened just like `flat_data`.
79 + """
80 + embed_is_patch: Union[torch.Tensor, list[torch.Tensor]]
81 + """
82 + A boolean mask indicating which image embeddings correspond
83 + to patch tokens.
84 + """
85 + aspect_ratios: Union[torch.Tensor, list[torch.Tensor]]
86 + """
87 + A list of aspect ratios corresponding to the number of tiles
88 + in each dimension that each image in the batch corresponds to.
89 +
90 + Shape:
91 + `(batch_size, ratio)` where ratio is a pair `(ratio_h, ratio_w)`
92 + """
93 +
94 +
95 +class Llama4VisionMLP(nn.Module):
96 +
97 + def __init__(self,
98 + input_size: int,
99 + intermediate_size: int,
100 + output_size: int,
101 + bias: bool,
102 + output_activation: bool,
103 + quant_config: Optional[QuantizationConfig] = None,
104 + prefix: str = ""):
105 + super().__init__()
106 + self.fc1 = ColumnParallelLinear(
107 + input_size=input_size,
108 + output_size=intermediate_size,
109 + bias=bias,
110 + quant_config=quant_config,
111 + prefix=f"{prefix}.fc1",
112 + )
113 + self.fc2 = RowParallelLinear(
114 + input_size=intermediate_size,
115 + output_size=output_size,
116 + bias=bias,
117 + quant_config=quant_config,
118 + prefix=f"{prefix}.fc2",
119 + )
120 + self.activation_fn = nn.GELU()
121 + self.output_activation = output_activation
122 +
123 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
124 + hidden_states, _ = self.fc1(hidden_states)
125 + hidden_states = self.activation_fn(hidden_states)
126 + hidden_states, _ = self.fc2(hidden_states)
127 + if self.output_activation:
128 + return self.activation_fn(hidden_states)
129 + return hidden_states
130 +
131 +
132 +class Llama4MultiModalProjector(nn.Module):
133 +
134 + def __init__(
135 + self,
136 + config,
137 + quant_config: Optional[QuantizationConfig] = None,
138 + prefix: str = "",
139 + ):
140 + super().__init__()
141 + self.linear_1 = ColumnParallelLinear(
142 + input_size=config.vision_config.vision_output_dim,
143 + output_size=config.text_config.hidden_size,
144 + bias=False,
145 + quant_config=quant_config,
146 + gather_output=True,
147 + prefix=f"{prefix}.linear_1",
148 + )
149 +
150 + def forward(self, image_features):
151 + hidden_states, _ = self.linear_1(image_features)
152 + return hidden_states
153 +
154 +
155 +def pixel_shuffle(input_tensor, shuffle_ratio):
156 + # input_tensor: [batch_size, num_patches, channels]
157 + batch_size, num_patches, channels = input_tensor.shape
158 + patch_size = int(math.sqrt(num_patches))
159 +
160 + input_tensor = input_tensor.view(batch_size, patch_size, patch_size, -1)
161 + batch_size, height, width, channels = input_tensor.size()
162 +
163 + reshaped_tensor = input_tensor.view(batch_size, height,
164 + int(width * shuffle_ratio),
165 + int(channels / shuffle_ratio))
166 + reshaped_tensor = reshaped_tensor.permute(0, 2, 1, 3).contiguous()
167 +
168 + reshaped_tensor = reshaped_tensor.view(batch_size,
169 + int(height * shuffle_ratio),
170 + int(width * shuffle_ratio),
171 + int(channels / (shuffle_ratio**2)))
172 + reshaped_tensor = reshaped_tensor.permute(0, 2, 1, 3).contiguous()
173 +
174 + output_tensor = reshaped_tensor.view(batch_size, -1,
175 + reshaped_tensor.shape[-1])
176 + return output_tensor
177 +
178 +
179 +class Llama4VisionPixelShuffleMLP(nn.Module):
180 +
181 + def __init__(
182 + self,
183 + config,
184 + quant_config: Optional[QuantizationConfig] = None,
185 + prefix: str = "",
186 + ):
187 + super().__init__()
188 + self.pixel_shuffle_ratio = config.pixel_shuffle_ratio
189 + self.inner_dim = int(config.projector_input_dim //
190 + (self.pixel_shuffle_ratio**2))
191 + self.output_dim = config.projector_output_dim
192 + self.mlp = Llama4VisionMLP(
193 + input_size=config.intermediate_size,
194 + intermediate_size=config.projector_input_dim,
195 + output_size=config.projector_output_dim,
196 + bias=config.multi_modal_projector_bias,
197 + output_activation=True,
198 + quant_config=quant_config,
199 + prefix=f"{prefix}.mlp")
200 +
201 + def forward(self, encoded_patches: torch.Tensor) -> torch.Tensor:
202 + encoded_patches = pixel_shuffle(encoded_patches,
203 + self.pixel_shuffle_ratio)
204 + return self.mlp(encoded_patches)
205 +
206 +
207 +class Llama4VisionAttention(nn.Module):
208 +
209 + def __init__(
210 + self,
211 + config: Llama4VisionConfig,
212 + quant_config: Optional[QuantizationConfig],
213 + prefix: str = "",
214 + ):
215 + super().__init__()
216 + self.config = config
217 + self.tp_size = get_tensor_model_parallel_world_size()
218 + self.embed_dim = config.hidden_size
219 + self.num_heads = config.num_attention_heads
220 + self.head_dim = config.hidden_size // self.num_heads
221 + assert self.num_heads % self.tp_size == 0
222 + self.num_local_heads = self.num_heads // self.tp_size
223 + self.q_size = self.num_local_heads * self.head_dim
224 + self.kv_size = self.num_local_heads * self.head_dim
225 + self.attention_dropout = config.attention_dropout
226 + self.scaling = self.head_dim**-0.5
227 +
228 + self.attn = MultiHeadAttention(self.num_local_heads, self.head_dim,
229 + self.scaling)
230 + self.qkv_proj = QKVParallelLinear(
231 + self.embed_dim,
232 + self.head_dim,
233 + self.num_heads,
234 + bias=True,
235 + quant_config=quant_config,
236 + prefix=f"{prefix}.qkv_proj",
237 + )
238 + self.o_proj = RowParallelLinear(
239 + self.num_heads * self.head_dim,
240 + self.embed_dim,
241 + bias=True,
242 + input_is_parallel=True,
243 + quant_config=quant_config,
244 + prefix=f"{prefix}.o_proj",
245 + )
246 +
247 + self.rotary_emb = get_rope(
248 + head_size=self.head_dim,
249 + rotary_dim=config.hidden_size // config.num_attention_heads // 2,
250 + # number of image patches
251 + max_position=(config.image_size // config.patch_size)**2,
252 + base=config.rope_theta,
253 + rope_scaling={"rope_type": "mllama4"},
254 + is_neox_style=False,
255 + dtype=torch.complex64, # important
256 + )
257 +
258 + def forward(
259 + self,
260 + hidden_states: torch.Tensor,
261 + ) -> torch.Tensor:
262 + input_shape = hidden_states.shape[:-1]
263 +
264 + qkv, _ = self.qkv_proj(hidden_states)
265 + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
266 +
267 + q = q.view(q.shape[0], q.shape[1], self.num_local_heads, self.head_dim)
268 + k = k.view(k.shape[0], k.shape[1], self.num_local_heads, self.head_dim)
269 + q, k = self.rotary_emb(q, k)
270 +
271 + q = q.view(q.shape[0], q.shape[1], -1)
272 + k = k.view(k.shape[0], k.shape[1], -1)
273 +
274 + attn_output = self.attn(q, k, v)
275 + attn_output = attn_output.reshape(*input_shape, -1).contiguous()
276 + attn_output, _ = self.o_proj(attn_output)
277 +
278 + return attn_output
279 +
280 +
281 +class Llama4VisionEncoderLayer(nn.Module):
282 +
283 + def __init__(
284 + self,
285 + config: Llama4VisionConfig,
286 + quant_config: Optional[QuantizationConfig],
287 + prefix: str = "",
288 + ):
289 + super().__init__()
290 + self.hidden_size = config.hidden_size
291 + self.num_attention_heads = config.num_attention_heads
292 + self.intermediate_size = config.intermediate_size
293 +
294 + self.self_attn = Llama4VisionAttention(config,
295 + quant_config=quant_config,
296 + prefix=f"{prefix}.self_attn")
297 + self.mlp = Llama4VisionMLP(input_size=config.hidden_size,
298 + intermediate_size=config.intermediate_size,
299 + output_size=config.hidden_size,
300 + bias=True,
301 + output_activation=False,
302 + quant_config=quant_config,
303 + prefix=f"{prefix}.mlp")
304 +
305 + self.input_layernorm = nn.LayerNorm(config.hidden_size)
306 + self.post_attention_layernorm = nn.LayerNorm(config.hidden_size)
307 +
308 + def forward(
309 + self,
310 + hidden_state: torch.Tensor,
311 + ):
312 + # Self Attention
313 + residual = hidden_state
314 + hidden_state = self.input_layernorm(hidden_state)
315 + hidden_state = self.self_attn(hidden_state)
316 + hidden_state = residual + hidden_state
317 +
318 + # Feed forward
319 + residual = hidden_state
320 + hidden_state = self.post_attention_layernorm(hidden_state)
321 + hidden_state = self.mlp(hidden_state)
322 + hidden_state = residual + hidden_state
323 +
324 + outputs = (hidden_state, )
325 + return outputs
326 +
327 +
328 +class Llama4VisionEncoder(nn.Module):
329 +
330 + def __init__(
331 + self,
332 + config: Llama4VisionConfig,
333 + quant_config: Optional[QuantizationConfig],
334 + prefix: str = "",
335 + ):
336 + super().__init__()
337 + self.config = config
338 + self.layers = nn.ModuleList([
339 + Llama4VisionEncoderLayer(
340 + config,
341 + quant_config=quant_config,
342 + prefix=f"{prefix}.layers.{layer_idx}",
343 + ) for layer_idx in range(config.num_hidden_layers)
344 + ])
345 +
346 + def forward(
347 + self,
348 + hidden_states: torch.Tensor,
349 + ) -> torch.Tensor:
350 + r"""
351 + Args:
352 + inputs_embeds (`torch.FloatTensor` of shape
353 + `(batch_size, sequence_length, hidden_size)`):
354 + Optionally, instead of passing `input_ids` you can choose to
355 + directly pass an embedded representation. This is useful if you
356 + want more control over how to convert `input_ids` indices into
357 + associated vectors than the model's internal embedding
358 + lookup matrix.
359 + """
360 +
361 + for encoder_layer in self.layers:
362 + layer_outputs = encoder_layer(hidden_states)
363 + hidden_states = layer_outputs[0]
364 +
365 + return hidden_states
366 +
367 +
368 +class Llama4UnfoldConvolution(nn.Module):
369 +
370 + def __init__(self,
371 + config: Llama4VisionConfig,
372 + quant_config: Optional[QuantizationConfig] = None,
373 + prefix: str = ""):
374 + super().__init__()
375 + kernel_size = config.patch_size
376 + if isinstance(kernel_size, int):
377 + kernel_size = (kernel_size, kernel_size)
378 + self.unfold = torch.nn.Unfold(kernel_size=kernel_size,
379 + stride=config.patch_size)
380 + self.linear = ColumnParallelLinear(config.num_channels *
381 + kernel_size[0] * kernel_size[1],
382 + config.hidden_size,
383 + bias=False,
384 + quant_config=quant_config,
385 + gather_output=True,
386 + prefix=f"{prefix}.linear")
387 +
388 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
389 + hidden_states = self.unfold(hidden_states)
390 + hidden_states = hidden_states.permute(0, 2, 1)
391 + hidden_states, _ = self.linear(hidden_states)
392 + return hidden_states
393 +
394 +
395 +class Llama4VisionModel(nn.Module):
396 +
397 + def __init__(
398 + self,
399 + config: Llama4VisionConfig,
400 + quant_config: Optional[QuantizationConfig] = None,
401 + prefix: str = "",
402 + ):
403 + super().__init__()
404 + self.config = config
405 + self.image_size = config.image_size
406 + self.patch_size = config.patch_size
407 + self.hidden_size = config.hidden_size
408 + self.num_channels = config.num_channels
409 +
410 + self.num_patches = (self.image_size // self.patch_size)**2 + 1
411 + self.scale = config.hidden_size**-0.5
412 +
413 + self.patch_embedding = Llama4UnfoldConvolution(
414 + config,
415 + quant_config=quant_config,
416 + prefix=f"{prefix}.patch_embedding")
417 +
418 + self.class_embedding = nn.Parameter(self.scale *
419 + torch.randn(self.hidden_size))
420 + self.positional_embedding_vlm = nn.Parameter(
421 + self.scale * torch.randn(self.num_patches, self.hidden_size))
422 +
423 + # layer norms
424 + self.layernorm_pre = nn.LayerNorm(self.hidden_size, eps=1e-5)
425 + self.layernorm_post = nn.LayerNorm(self.hidden_size, eps=1e-5)
426 +
427 + # encoders
428 + self.model = Llama4VisionEncoder(config,
429 + quant_config=quant_config,
430 + prefix=f"{prefix}.model")
431 + self.vision_adapter = Llama4VisionPixelShuffleMLP(
432 + config, quant_config, prefix=f"{prefix}.vision_adapter")
433 +
434 + def forward(
435 + self,
436 + images_flattened: torch.Tensor,
437 + ) -> torch.Tensor:
438 + # Patch embedding
439 + hidden_state = self.patch_embedding(images_flattened)
440 + num_tiles, num_patches, hidden_dim = hidden_state.shape
441 +
442 + # Add cls token
443 + class_embedding = self.class_embedding.expand(hidden_state.shape[0], 1,
444 + hidden_state.shape[-1])
445 + hidden_state = torch.cat([hidden_state, class_embedding], dim=1)
446 + num_patches += 1
447 +
448 + # Position embeddings
449 + hidden_state = hidden_state.reshape(
450 + num_tiles,
451 + 1,
452 + num_patches,
453 + hidden_dim,
454 + )
455 + positional_embedding = self.positional_embedding_vlm.to(
456 + dtype=hidden_state.dtype, device=hidden_state.device)
457 + hidden_state = hidden_state + positional_embedding
458 + hidden_state = self.layernorm_pre(hidden_state)
459 + hidden_state = hidden_state.view(num_tiles, -1, hidden_dim)
460 +
461 + # Apply encoder
462 + hidden_state = self.model(hidden_state)
463 + hidden_state = self.layernorm_post(hidden_state)
464 +
465 + # Remove CLS token output
466 + hidden_state = hidden_state[:, :-1, :]
467 +
468 + # now, we use Llama4VisionPixelShuffle + mlp to project embeddings
469 + hidden_state = self.vision_adapter(hidden_state)
470 +
471 + return hidden_state
472 +
473 +
474 +class Mllama4ProcessingInfo(BaseProcessingInfo):
475 +
476 + def __init__(self, ctx: InputProcessingContext) -> None:
477 + super().__init__(ctx)
478 +
479 + def get_hf_config(self) -> Llama4Config:
480 + return self.ctx.get_hf_config(Llama4Config)
481 +
482 + def get_hf_processor(self, **kwargs: object) -> Llama4Processor:
483 + return self.ctx.get_hf_processor(Llama4Processor,
484 + use_fast=True,
485 + **kwargs)
486 +
487 + def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:
488 + return {"image": 10}
489 +
490 + @staticmethod
491 + def get_patch_per_chunk(vision_config: Llama4VisionConfig) -> int:
492 + image_size = vision_config.image_size
493 + patch_size = vision_config.patch_size
494 +
495 + assert (
496 + image_size %
497 + patch_size == 0), f"chunk size {image_size} should be multiple of "
498 + f"patch_size {patch_size}"
499 +
500 + ds_ratio = int(round(1.0 / (vision_config.pixel_shuffle_ratio**2)))
501 + return (image_size // patch_size)**2 // ds_ratio
502 +
503 + def get_max_num_tiles(self) -> int:
504 + image_processor = self.get_hf_processor().image_processor
505 + return image_processor.max_patches
506 +
507 + def get_mm_max_tokens_per_item(
508 + self,
509 + seq_len: int,
510 + mm_counts: Mapping[str, int],
511 + ) -> Mapping[str, int]:
512 + vision_config = self.get_hf_config().vision_config
513 + # image_start + local tiles * (patches + 1 x separator) +
514 + # 1 global tile * (image x 1 + patches) + image_end
515 + token_per_chunk = self.get_patch_per_chunk(vision_config) + 1
516 + mm_max_tokens = (self.get_max_num_tiles() + 1) * token_per_chunk + 2
517 + return {"image": mm_max_tokens}
518 +
519 + def get_image_size_with_most_features(self) -> ImageSize:
520 + vision_config = self.get_hf_config().vision_config
521 + image_size = vision_config.image_size
522 + # Result in the max possible feature size (h:w = 16:1)
523 + return ImageSize(height=self.get_max_num_tiles() * image_size,
524 + width=image_size)
525 +
526 +
527 +class Mllama4MultiModalProcessor(BaseMultiModalProcessor[Mllama4ProcessingInfo]
528 + ):
529 +
530 + def _call_hf_processor(
531 + self,
532 + prompt: str,
533 + mm_data: Mapping[str, object],
534 + mm_kwargs: Mapping[str, object],
535 + ) -> BatchFeature:
536 + tokenizer = self.info.get_tokenizer()
537 +
538 + if mm_data is None:
539 + return tokenizer(prompt, add_special_tokens=False) # exclude bos
540 + processed_outputs = super()._call_hf_processor(
541 + prompt=prompt,
542 + mm_data=mm_data,
543 + mm_kwargs=mm_kwargs,
544 + )
545 +
546 + processor = self.info.get_hf_processor(**mm_kwargs)
547 + image_processor = processor.image_processor
548 + vision_config = self.info.get_hf_config().vision_config
549 +
550 + if processed_outputs.get("pixel_values") is not None:
551 + assert "images" in mm_data, \
552 + "images expected to be in mm_data when pixel_values is present"
553 +
554 + images = mm_data["images"]
555 + parsed_images = (self._get_data_parser().parse_mm_data({
556 + "image":
557 + images
558 + }).get_items("image", ImageProcessorItems))
559 +
560 + tile_size = vision_config.image_size
561 + possible_resolutions = find_supported_resolutions(
562 + max_num_chunks=self.info.get_max_num_tiles(),
563 + patch_size=SizeDict(height=tile_size, width=tile_size),
564 + )
565 + best_fit_sizes = [
566 + get_best_fit(
567 + (image.size[1], image.size[0]),
568 + torch.tensor(possible_resolutions),
569 + resize_to_max_canvas=image_processor.resize_to_max_canvas)
570 + for image in parsed_images
571 + ]
572 + # TODO tile height/width do not necessarily need to match
573 + aspect_ratios = [(image_size[0] // tile_size,
574 + image_size[1] // tile_size)
575 + for image_size in best_fit_sizes]
576 + patches_per_image = [
577 + 1 if r_h * r_w == 1 else 1 + r_h * r_w
578 + for (r_h, r_w) in aspect_ratios
579 + ]
580 +
581 + # embed_is_patch should have one feature per image-related token:
582 + # <|image_start|>, <|tile_*_separator|>, <|image|>, <|image_end|>
583 + # -> False
584 + # <|patch|> -> True
585 + # embed_is_patch has no entries corresponding to non-image-related
586 + # tokens.
587 + patch_id = tokenizer.get_vocab()[processor.img_patch_token]
588 + num_patches_per_chunk = self.info.get_patch_per_chunk(
589 + vision_config)
590 + expanded_image_tokens_list = [
591 + processor._prompt_split_image(aspect_ratio,
592 + num_patches_per_chunk)
593 + for aspect_ratio in aspect_ratios
594 + ]
595 + expanded_image_token_ids = [
596 + tokenizer.encode(image_tokens, add_special_tokens=False)
597 + for image_tokens in expanded_image_tokens_list
598 + ]
599 + embed_is_patch = [
600 + torch.tensor(tokens) == patch_id
601 + for tokens in expanded_image_token_ids
602 + ]
603 +
604 + processed_outputs["aspect_ratios"] = aspect_ratios
605 + processed_outputs["patches_per_image"] = torch.tensor(
606 + patches_per_image)
607 + processed_outputs["embed_is_patch"] = embed_is_patch
608 +
609 + return processed_outputs
610 +
611 + def _get_mm_fields_config(
612 + self,
613 + hf_inputs: BatchFeature,
614 + hf_processor_mm_kwargs: Mapping[str, object],
615 + ) -> Mapping[str, MultiModalFieldConfig]:
616 + patches_per_image = hf_inputs.get("patches_per_image", torch.empty(0))
617 + return dict(
618 + pixel_values=MultiModalFieldConfig.flat_from_sizes(
619 + "image", patches_per_image),
620 + patches_per_image=MultiModalFieldConfig.batched("image"),
621 + aspect_ratios=MultiModalFieldConfig.batched("image"),
622 + embed_is_patch=MultiModalFieldConfig.batched("image"),
623 + )
624 +
625 + def _get_prompt_updates(
626 + self,
627 + mm_items: MultiModalDataItems,
628 + hf_processor_mm_kwargs: Mapping[str, object],
629 + out_mm_kwargs: MultiModalKwargs,
630 + ) -> List[PromptUpdate]:
631 + assert (
632 + mm_items.get_count("image", strict=False) == 0
633 + or "aspect_ratios" in out_mm_kwargs
634 + ), "Transformers expect to include aspect_ratios in out_mm_kwargs"
635 +
636 + config = self.info.get_hf_config()
637 + vision_config = config.vision_config
638 +
639 + num_patches_per_chunk = self.info.get_patch_per_chunk(vision_config)
640 + hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
641 + image_token = hf_processor.image_token
642 +
643 + def get_replacement(item_idx: int):
644 + aspect_ratio = out_mm_kwargs["aspect_ratios"][item_idx]
645 + return hf_processor._prompt_split_image(
646 + aspect_ratio=aspect_ratio,
647 + num_patches_per_chunk=num_patches_per_chunk)
648 +
649 + return [
650 + PromptReplacement(
651 + modality="image",
652 + target=image_token,
653 + replacement=get_replacement,
654 + )
655 + ]
656 +
657 +
658 +class Mllama4DummyInputsBuilder(BaseDummyInputsBuilder[Mllama4ProcessingInfo]):
659 +
660 + def get_dummy_processor_inputs(
661 + self,
662 + seq_len: int,
663 + mm_counts: Mapping[str, int],
664 + ) -> ProcessorInputs:
665 + num_images = mm_counts.get("image", 0)
666 +
667 + (target_width,
668 + target_height) = self.info.get_image_size_with_most_features()
669 +
670 + mm_data = {
671 + "image":
672 + self._get_dummy_images(width=target_width,
673 + height=target_height,
674 + num_images=num_images)
675 + }
676 +
677 + image_token = self.info.get_hf_processor().fake_image_token
678 + return ProcessorInputs(
679 + prompt_text=image_token * num_images,
680 + mm_data=mm_data,
681 + )
682 +
683 +
684 +@MULTIMODAL_REGISTRY.register_processor(
685 + Mllama4MultiModalProcessor,
686 + info=Mllama4ProcessingInfo,
687 + dummy_inputs=Mllama4DummyInputsBuilder,
688 +)
689 +class Llama4ForConditionalGeneration(nn.Module, SupportsMultiModal,
690 + SupportsPP):
691 + packed_modules_mapping = {
692 + "qkv_proj": ["q_proj", "k_proj", "v_proj"],
693 + }
694 +
695 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
696 + super().__init__()
697 + config = vllm_config.model_config.hf_config
698 + quant_config = vllm_config.quant_config
699 + multimodal_config = vllm_config.model_config.multimodal_config
700 + self.config = config
701 + self.quant_config = quant_config
702 + self.multimodal_config = multimodal_config
703 + self.vision_model = Llama4VisionModel(config.vision_config,
704 + None,
705 + prefix=maybe_prefix(
706 + prefix, "vision_model"))
707 + self.multi_modal_projector = Llama4MultiModalProjector(
708 + self.config,
709 + None,
710 + prefix=maybe_prefix(prefix, "multi_modal_projector"))
711 +
712 + self.language_model = _initialize_model(
713 + vllm_config=vllm_config.with_hf_config(config.text_config),
714 + prefix=maybe_prefix(prefix, "language_model"),
715 + model_class=Llama4ForCausalLM,
716 + )
717 +
718 + self.make_empty_intermediate_tensors = (
719 + self.language_model.make_empty_intermediate_tensors)
720 +
721 + @cached_property
722 + def sampler(self):
723 + if hasattr(self.language_model, "sampler"):
724 + return self.language_model.sampler
725 +
726 + return get_sampler()
727 +
728 + def _parse_and_validate_image_input(
729 + self, **kwargs: object) -> Optional[Llama4ImagePatchInputs]:
730 + # num_images, 1, num_chunks, channel, image_size, image_size
731 + pixel_values = kwargs.pop("pixel_values", None)
732 + if pixel_values is None:
733 + return None
734 +
735 + # num_images x num_chunks, channel, image_size, image_size
736 + # TODO: confirm handling for variable lengths
737 + flat_pixel_values = flatten_bn(pixel_values, concat=True)
738 + patches_per_image = flatten_bn(kwargs.pop("patches_per_image"))
739 +
740 + embed_is_patch = kwargs.pop("embed_is_patch", None)
741 + if not isinstance(embed_is_patch, (torch.Tensor, list)):
742 + raise ValueError("Incorrect type of embed_is_patch. "
743 + f"Got type: {type(embed_is_patch)}")
744 +
745 + aspect_ratios = kwargs.pop("aspect_ratios", None)
746 + if not isinstance(aspect_ratios, (torch.Tensor, list)):
747 + raise ValueError("Incorrect type of aspect_ratios. "
748 + f"Got type: {type(aspect_ratios)}")
749 +
750 + return Llama4ImagePatchInputs(
751 + type="pixel_values",
752 + flat_data=flat_pixel_values,
753 + patches_per_image=patches_per_image,
754 + embed_is_patch=embed_is_patch,
755 + aspect_ratios=aspect_ratios,
756 + )
757 +
758 + def _process_image_input(
759 + self, image_input: Llama4ImagePatchInputs) -> MultiModalEmbeddings:
760 + flat_data = image_input["flat_data"]
761 + patches_per_image = image_input["patches_per_image"].tolist()
762 + vision_embeddings_flat = self.vision_model(flat_data)
763 + return vision_embeddings_flat.split(patches_per_image, dim=0)
764 +
765 + def get_multimodal_embeddings(self,
766 + **kwargs) -> Optional[MultiModalEmbeddings]:
767 + image_input = self._parse_and_validate_image_input(**kwargs)
768 + if image_input is None:
769 + return None
770 +
771 + # num_images x [num_chunks, num_patches, hidden_dim]
772 + image_features = self._process_image_input(image_input)
773 + # num_images x [num_chunks x num_patches, hidden_dim]
774 + image_features_flat = [img.flatten(0, 1) for img in image_features]
775 + # num_images x [1, input_len] -> num_images x [input_len]
776 + embed_is_patch_flat = [
777 + is_patch.flatten(0, 1)
778 + for is_patch in image_input["embed_is_patch"]
779 + ]
780 +
781 + return scatter_patch_features(
782 + image_features_flat,
783 + embed_is_patch_flat,
784 + )
785 +
786 + def get_input_embeddings(
787 + self,
788 + input_ids: torch.Tensor,
789 + multimodal_embeddings: Optional[NestedTensors] = None,
790 + ) -> torch.Tensor:
791 + inputs_embeds = self.language_model.get_input_embeddings(input_ids)
792 +
793 + if multimodal_embeddings is not None:
794 + multimodal_embeddings = torch.cat(multimodal_embeddings)
795 + mm_embeddings = self.multi_modal_projector(multimodal_embeddings)
796 + inputs_embeds = merge_multimodal_embeddings(
797 + input_ids, inputs_embeds, select_patch_features(mm_embeddings),
798 + self.config.image_token_index)
799 +
800 + return inputs_embeds
801 +
802 + def forward(
803 + self,
804 + input_ids: torch.Tensor,
805 + positions: torch.Tensor,
806 + intermediate_tensors: Optional[IntermediateTensors] = None,
807 + inputs_embeds: Optional[torch.Tensor] = None,
808 + **kwargs: object,
809 + ) -> Union[torch.Tensor, IntermediateTensors]:
810 + if intermediate_tensors is not None:
811 + inputs_embeds = None
812 +
813 + # NOTE: In v1, inputs_embeds is always generated at model runner,
814 + # this condition is for v0 compatibility.
815 + elif inputs_embeds is None:
816 + vision_embeddings = self.get_multimodal_embeddings(**kwargs)
817 + inputs_embeds = self.get_input_embeddings(input_ids,
818 + vision_embeddings)
819 + input_ids = None
820 +
821 + return self.language_model(input_ids, positions, intermediate_tensors,
822 + inputs_embeds)
823 +
824 + def compute_logits(
825 + self,
826 + hidden_states: torch.Tensor,
827 + sampling_metadata: SamplingMetadata,
828 + ) -> Optional[torch.Tensor]:
829 + return self.language_model.compute_logits(hidden_states,
830 + sampling_metadata)
831 +
832 + def sample(self, logits: torch.Tensor,
833 + sampling_metadata: SamplingMetadata) -> Optional[SamplerOutput]:
834 + return self.language_model.sample(logits, sampling_metadata)
835 +
836 + def separate_weights(
837 + self,
838 + weights: Iterable[Tuple[str, torch.Tensor]],
839 + prefix: str,
840 + ) -> Tuple[Iterable[Tuple[str, torch.Tensor]], Iterable[Tuple[
841 + str, torch.Tensor]]]:
842 + weights1, weights2 = tee(weights, 2)
843 +
844 + def get_prefix_weights() -> Iterable[Tuple[str, torch.Tensor]]:
845 + for name, data in weights1:
846 + if name.startswith(prefix):
847 + yield (name, data)
848 +
849 + def get_other_weights() -> Iterable[Tuple[str, torch.Tensor]]:
850 + for name, data in weights2:
851 + if not name.startswith(prefix):
852 + yield (name, data)
853 +
854 + return get_prefix_weights(), get_other_weights()
855 +
856 + def load_weights(self, weights: Iterable[Tuple[str,
857 + torch.Tensor]]) -> Set[str]:
858 +
859 + stacked_params_mapping = [
860 + # (param_name, shard_name, shard_id)
861 + (".self_attn.qkv_proj", ".self_attn.q_proj", "q"),
862 + (".self_attn.qkv_proj", ".self_attn.k_proj", "k"),
863 + (".self_attn.qkv_proj", ".self_attn.v_proj", "v"),
864 + ]
865 + params_dict = dict(self.named_parameters())
866 + updated_params: Set[str] = set()
867 +
868 + # language_model is an Llama4ForCausalLM instance. We load it's
869 + # using llama4's load_weights routine.
870 + language_model_weights, other_weights = self.separate_weights(
871 + weights, prefix="language_model.model.")
872 + loader = AutoWeightsLoader(self)
873 + loaded_language_model_params = loader.load_weights(
874 + language_model_weights)
875 + assert loaded_language_model_params is not None
876 + updated_params.update(loaded_language_model_params)
877 +
878 + for name, loaded_weight in other_weights:
879 + for param_name, weight_name, shard_id in stacked_params_mapping:
880 + if weight_name not in name:
881 + continue
882 + name = name.replace(weight_name, param_name)
883 + param = params_dict[name]
884 + updated_params.add(name)
885 + weight_loader = param.weight_loader
886 + weight_loader(param, loaded_weight, shard_id)
887 + break
888 + else:
889 + param = params_dict[name]
890 + weight_loader = getattr(param, "weight_loader",
891 + default_weight_loader)
892 +
893 + weight_loader(param, loaded_weight)
894 + updated_params.add(name)
895 + return updated_params
vllm/model_executor/models/registry.py
+1
@@ -196,6 +196,7 @@ _MULTIMODAL_MODELS = {
196 # [Encoder-decoder]
197 "Florence2ForConditionalGeneration": ("florence2", "Florence2ForConditionalGeneration"), # noqa: E501
198 "MllamaForConditionalGeneration": ("mllama", "MllamaForConditionalGeneration"), # noqa: E501
199 + "Llama4ForConditionalGeneration": ("mllama4", "Llama4ForConditionalGeneration"), # noqa: E501
200 "SkyworkR1VChatModel": ("skyworkr1v", "SkyworkR1VChatModel"),
201 "WhisperForConditionalGeneration": ("whisper", "WhisperForConditionalGeneration"), # noqa: E501
202 }
vllm/model_executor/models/telechat2.py
+3 -2
@@ -19,9 +19,10 @@
19 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
20 # See the License for the specific language governing permissions and
21 # limitations under the License.
22 -from typing import Iterable, Set, Tuple, Type
22 +from typing import Iterable, Set, Tuple
23
24 import torch
25 +import torch.nn as nn
26
27 from vllm.config import VllmConfig
28 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
@@ -124,7 +125,7 @@ class TeleChat2ForCausalLM(LlamaForCausalLM):
125 def _init_model(self,
126 vllm_config: VllmConfig,
127 prefix: str = "",
127 - layer_type: Type[LlamaDecoderLayer] = LlamaDecoderLayer):
128 + layer_type: type[nn.Module] = LlamaDecoderLayer):
129 return TeleChat2Model(vllm_config=vllm_config, prefix=prefix)
130
131 def load_weights(self, weights: Iterable[Tuple[str,
vllm/model_executor/models/teleflm.py
+2 -3
@@ -22,9 +22,8 @@
22 # See the License for the specific language governing permissions and
23 # limitations under the License.
24
25 -from typing import Type
26 -
25 import torch
26 +import torch.nn as nn
27
28 from vllm.config import VllmConfig
29 from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -39,7 +38,7 @@ class TeleFLMModel(LlamaModel):
38 *,
39 vllm_config: VllmConfig,
40 prefix: str = "",
42 - layer_type: Type[LlamaDecoderLayer] = LlamaDecoderLayer,
41 + layer_type: type[nn.Module] = LlamaDecoderLayer,
42 ):
43 super().__init__(vllm_config=vllm_config,
44 prefix=prefix,
vllm/v1/attention/backends/flash_attn.py
+236 -14
@@ -96,6 +96,183 @@ class FlashAttentionMetadata:
96 # For logging.
97 num_input_tokens: int = 0 # Number of tokens including padding.
98
99 + # for local attention
100 + @dataclass
101 + class LocalAttentionMetadata:
102 + local_query_start_loc: torch.Tensor
103 + local_seqused_k: torch.Tensor
104 + local_block_table: torch.Tensor
105 + local_max_query_len: int
106 + local_max_seq_len: int
107 +
108 + local_attn_metadata: Optional[LocalAttentionMetadata] = None
109 +
110 +
111 +#
112 +# Take in `query_start_loc_np` and `seq_lens_np` and break the sequences into
113 +# local attention blocks, where each block is passed to the attention kernel
114 +# as an independent local ("virtual") batch item.
115 +#
116 +# For example, if are performing a chunked prefill a batch of 3 sequences:
117 +# q_seqlens = [4, 10, 5]
118 +# kv_seqlens = [6, 17, 9]
119 +# Then normally for regular attention we would compute with an attention mask
120 +# for batch idx 0 (q_seqlens = 4, kv_seqlens = 6) like:
121 +# batch idx: 0 (q_seqlens = 4, kv_seqlens = 6)
122 +# k_toks > 0 1 2 3 4 5
123 +# q_toks v _____________
124 +# 0 | 1 1 1
125 +# 1 | 1 1 1 1
126 +# 2 | 1 1 1 1 1
127 +# 3 | 1 1 1 1 1 1
128 +#
129 +# for local attention (with attn_chunk_size = 4) we would compute with an
130 +# attention mask like:
131 +# batch idx: 0 (q_seqlens = 4, kv_seqlens = 6, attn_chunk_size = 4)
132 +# k_toks > 0 1 2 3 4 5
133 +# q_toks v _____________
134 +# 0 | 1 1 1
135 +# 1 | 1 1 1 1
136 +# 2 | 1
137 +# 3 | 1 1
138 +#
139 +# We can simulate this mask using standard flash-attention by breaking the
140 +# sequences into local ("virtual") batches, where each local batch item is a
141 +# local attention block, so in this case batch idx 0 would be broken up into:
142 +#
143 +# local-batch idx: 0 (q_seqlens = 2, kv_seqlens = 4) (batch 0)
144 +# k_toks > 0 1 2 3
145 +# q_toks v _____________
146 +# 0 | 1 1 1
147 +# 1 | 1 1 1 1
148 +# local-batch idx: 1 (q_seqlens = 2, kv_seqlens = 2) (batch 0)
149 +# k_toks > 4 5
150 +# q_toks v _____________
151 +# 2 | 1
152 +# 3 | 1 1
153 +#
154 +# e.g. if we have:
155 +# attn_chunk_size = 4
156 +# query_start_loc_np = [0, 4, 14, 19] (q_seqlens = [4, 10, 5])
157 +# Then this function would return:
158 +# __b0__ ______b1______ __b2__ < orig batch indices
159 +# q_seqlens_local = [ 2, 2, 1, 4, 4, 1, 4, 1]
160 +# cu_seqlens_q_local = [0, 4, 6, 10, 14, 18, 19, 23, 24]
161 +# seqlens_k_local = [ 4, 2, 4, 4, 4, 1, 4, 1]
162 +# block_table_local : shape[local_virtual_batches, pages_per_local_batch]
163 +def make_local_attention_virtual_batches(
164 + attn_chunk_size: int,
165 + query_start_loc_np: np.ndarray,
166 + seq_lens_np: np.ndarray,
167 + block_table: torch.Tensor,
168 + page_size: int = 0,
169 +) -> tuple[np.ndarray, np.ndarray, np.ndarray, torch.Tensor]:
170 + q_seqlens = query_start_loc_np[1:] - query_start_loc_np[:-1]
171 + actual_batch_size = seq_lens_np.shape[0]
172 +
173 + # Handle if we are starting in the middle of a local attention block,
174 + # we assume q_seqlens > 0 (for all elements), for each batch idx we compute
175 + # the number of tokens that are not in the first local attention block and
176 + # then we can simply use a cdiv for the rest.
177 + # For example if we have:
178 + # attn_chunk_size = 4
179 + # q_seqlens = [4, 10, 5]
180 + # k_seqlens = [6, 17, 9]
181 + # Then we would get:
182 + # new_tokens_in_first_block = [2, 1, 4]
183 + # local_blocks = [2, 4, 2]
184 + q_tokens_in_first_block = np.minimum(
185 + attn_chunk_size - ((seq_lens_np - q_seqlens) % attn_chunk_size),
186 + q_seqlens).astype(np.int32)
187 + tokens_in_last_block = attn_chunk_size + (seq_lens_np % -attn_chunk_size)
188 + local_blocks = 1 + cdiv(q_seqlens - q_tokens_in_first_block,
189 + attn_chunk_size)
190 +
191 + # Once we know the number of local blocks we can compute the request spans
192 + # for each batch idx, we can figure out the number of "virtual" requests we
193 + # have to make,
194 + # For the above example we would get:
195 + # seqlens_q_local = [2, 2, 1, 4, 4, 1, 4, 1]
196 + #
197 + # First Get batched arange. (E.g., [2, 4, 2] -> [0, 1, 0, 1, 2, 3, 0, 1])
198 + # (TODO: max a utility to share this code with _prepare_inputs)
199 + # arange step 1. [2, 4, 2] -> [2, 6, 8]
200 + cu_num_blocks = np.cumsum(local_blocks)
201 + virtual_batches = cu_num_blocks[-1]
202 + # arange step 2. [2, 6, 8] -> [0, 0, 2, 2, 2, 2, 6, 6]
203 + block_offsets = np.repeat(cu_num_blocks - local_blocks, local_blocks)
204 + # arange step 3. [0, 1, 0, 1, 2, 3, 0, 1]
205 + arange = np.arange(virtual_batches, dtype=np.int32) - block_offsets
206 + # also compute reverse arange (i.e. [1, 0, 3, 2, 1, 0, 1, 0])
207 + rarange = np.repeat(local_blocks, local_blocks) - arange - 1
208 + # Then we can compute the seqlens_q_local, handling the fact that the
209 + # first and last blocks could be partial
210 + seqlens_q_local = \
211 + np.repeat(q_seqlens - q_tokens_in_first_block, local_blocks)
212 + # set the first block since this may be a partial block
213 + seqlens_q_local[arange == 0] = q_tokens_in_first_block
214 + # set the remaining blocks
215 + seqlens_q_local[arange > 0] = np.minimum(
216 + seqlens_q_local - attn_chunk_size * (arange - 1),
217 + attn_chunk_size)[arange > 0]
218 +
219 + # convert from q_seqlens to cu_seqlens_q
220 + cu_seqlens_q_local = np.pad(np.cumsum(seqlens_q_local), (1, 0))\
221 + .astype(np.int32)
222 +
223 + # compute the seqlens_k_local,
224 + # basically a full local attention block for all but the last block in each
225 + # batch
226 + # For our example this will be:
227 + # seqlens_k_local = [4, 2, 4, 4, 4, 1, 4, 1]
228 + seqlens_k_local = np.full(cu_num_blocks[-1],
229 + attn_chunk_size,
230 + dtype=np.int32)
231 + seqlens_k_local[cu_num_blocks - 1] = tokens_in_last_block
232 +
233 + k_seqstarts_absolute = np.repeat(seq_lens_np, local_blocks) - \
234 + (rarange * attn_chunk_size + \
235 + np.repeat(tokens_in_last_block, local_blocks))
236 + # For the example the local attention blocks start at:
237 + # _b0_ _____b1_____ _b2_
238 + # k_seqstarts_absolute = [0, 4, 4, 8, 12, 16, 4, 8]
239 + block_starts = k_seqstarts_absolute // page_size
240 + assert attn_chunk_size % page_size == 0, \
241 + f"attn_chunk_size {attn_chunk_size} is not " \
242 + f"divisible by page_size {page_size}"
243 + pages_per_local_batch = attn_chunk_size // page_size
244 +
245 + # Create a block_table for the local attention blocks
246 + # For out example if we have a block-table like (assuming page_size=2):
247 + # block_table = [
248 + # [ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9], < batch 0
249 + # [10, 11, 12, 13, 14, 15, 16, 17, 18, 19], < batch 1
250 + # [20, 21, 22, 23, 24, 25, 26, 27, 28, 29], < batch 2
251 + # ]
252 + # Then for the local batches we would want a block-table like
253 + # block_table_local = [
254 + # [ 0, 1 ], < local-batch 0, (batch 0, starting from k[0])
255 + # [ 2, 3 ], < local-batch 1, (batch 0, starting from k[4])
256 + # [ 12, 13 ], < local-batch 2, (batch 1, starting from k[4])
257 + # [ 14, 15 ], < local-batch 3, (batch 1, starting from k[8])
258 + # [ 16, 17 ], < local-batch 4, (batch 1, starting from k[12])
259 + # [ 18, 19 ], < local-batch 5, (batch 1, starting from k[16])
260 + # [ 22, 23 ], < local-batch 6, (batch 2, starting from k[4])
261 + # [ 24, 25 ], < local-batch 7, (batch 2, starting from k[8])
262 + # ]
263 + block_indices= np.broadcast_to(
264 + np.arange(pages_per_local_batch, dtype=np.int32),
265 + (virtual_batches, pages_per_local_batch)) \
266 + + np.expand_dims(block_starts, axis=1)
267 + block_indices = block_indices.flatten()
268 + batch_indices = np.repeat(np.arange(actual_batch_size, dtype=np.int32),
269 + local_blocks * pages_per_local_batch)
270 + block_table_local = block_table[batch_indices, block_indices]\
271 + .view(virtual_batches, -1)
272 +
273 + return seqlens_q_local, cu_seqlens_q_local, seqlens_k_local, \
274 + block_table_local
275 +
276
277 class FlashAttentionMetadataBuilder:
278
@@ -109,18 +286,40 @@ class FlashAttentionMetadataBuilder:
286 def build(self, num_reqs: int, num_actual_tokens: int, max_query_len: int,
287 common_prefix_len: int):
288 max_seq_len = self.runner.seq_lens_np[:num_reqs].max()
112 - query_start_loc = self.runner.query_start_loc_cpu[:num_reqs + 1].to(
113 - self.runner.device, non_blocking=True)
114 - seq_lens = self.runner.seq_lens_cpu[:num_reqs].to(self.runner.device,
115 - non_blocking=True)
289 + query_start_loc_cpu = self.runner.query_start_loc_cpu[:num_reqs + 1]
290 + query_start_loc = query_start_loc_cpu.to(self.runner.device,
291 + non_blocking=True)
292 + seq_lens_cpu = self.runner.seq_lens_cpu[:num_reqs]
293 + seq_lens = seq_lens_cpu.to(self.runner.device, non_blocking=True)
294 block_table = (
295 self.runner.input_batch.block_table.get_device_tensor()[:num_reqs])
296 slot_mapping = self.runner.slot_mapping_cpu[:num_actual_tokens].to(
297 self.runner.device, non_blocking=True).long()
298
299 + # for local attention
300 + local_attn_metadata = None
301 + if self.runner.attention_chunk_size is not None:
302 + seqlens_q_local_np, virt_q_cu_seqlens_np, virt_k_seqlens_np, \
303 + virt_block_table = make_local_attention_virtual_batches(
304 + self.runner.attention_chunk_size,
305 + self.runner.query_start_loc_np[:num_reqs + 1],
306 + self.runner.seq_lens_np[:num_reqs],
307 + block_table,
308 + self.runner.block_size,
309 + )
310 + local_attn_metadata = FlashAttentionMetadata.LocalAttentionMetadata(
311 + local_query_start_loc=torch.from_numpy(
312 + virt_q_cu_seqlens_np).to(self.runner.device,
313 + non_blocking=True),
314 + local_seqused_k=torch.from_numpy(virt_k_seqlens_np).to(
315 + self.runner.device, non_blocking=True),
316 + local_block_table=virt_block_table,
317 + local_max_query_len=seqlens_q_local_np.max(),
318 + local_max_seq_len=virt_k_seqlens_np.max(),
319 + )
320 +
321 use_cascade = common_prefix_len > 0
322 if use_cascade:
123 - # TODO: Optimize.
323 cu_prefix_query_lens = torch.tensor([0, num_actual_tokens],
324 dtype=torch.int32,
325 device=self.runner.device)
@@ -149,6 +348,7 @@ class FlashAttentionMetadataBuilder:
348 cu_prefix_query_lens=cu_prefix_query_lens,
349 prefix_kv_lens=prefix_kv_lens,
350 suffix_kv_lens=suffix_kv_lens,
351 + local_attn_metadata=local_attn_metadata,
352 )
353 return attn_metadata
354
@@ -167,6 +367,7 @@ class FlashAttentionImpl(AttentionImpl):
367 blocksparse_params: Optional[dict[str, Any]] = None,
368 logits_soft_cap: Optional[float] = None,
369 attn_type: AttentionType = AttentionType.DECODER,
370 + use_irope: bool = False,
371 ) -> None:
372 if blocksparse_params is not None:
373 raise ValueError(
@@ -203,6 +404,7 @@ class FlashAttentionImpl(AttentionImpl):
404 "encoder/decoder cross-attention "
405 "are not implemented for "
406 "FlashAttentionImpl")
407 + self.use_irope = use_irope
408 self.vllm_flash_attn_version = get_flash_attn_version()
409 if is_quantized_kv_cache(self.kv_cache_dtype) \
410 and not flash_attn_supports_fp8():
@@ -265,8 +467,7 @@ class FlashAttentionImpl(AttentionImpl):
467 layer._k_scale,
468 layer._v_scale,
469 )
268 - descale_shape = (attn_metadata.query_start_loc.shape[0] - 1,
269 - key.shape[1])
470 +
471 if self.kv_cache_dtype.startswith("fp8"):
472 key_cache = key_cache.view(torch.float8_e4m3fn)
473 value_cache = value_cache.view(torch.float8_e4m3fn)
@@ -278,22 +479,41 @@ class FlashAttentionImpl(AttentionImpl):
479 query = query.reshape((num_tokens, num_heads, head_size))
480
481 # Compute attention and update output up to `num_actual_tokens`.
281 - if not attn_metadata.use_cascade:
282 - # Regular attention (common case).
482 + use_local_attn = \
483 + (self.use_irope and attn_metadata.local_attn_metadata is not None)
484 +
485 + if not attn_metadata.use_cascade or use_local_attn:
486 + if use_local_attn:
487 + assert attn_metadata.local_attn_metadata is not None
488 + local_metadata = attn_metadata.local_attn_metadata
489 + cu_seqlens_q = local_metadata.local_query_start_loc
490 + seqused_k = local_metadata.local_seqused_k
491 + max_seqlen_q = local_metadata.local_max_query_len
492 + max_seqlen_k = local_metadata.local_max_seq_len
493 + block_table = local_metadata.local_block_table
494 + else:
495 + cu_seqlens_q = attn_metadata.query_start_loc
496 + seqused_k = attn_metadata.seq_lens
497 + max_seqlen_q = attn_metadata.max_query_len
498 + max_seqlen_k = attn_metadata.max_seq_len
499 + block_table = attn_metadata.block_table
500 +
501 + descale_shape = (cu_seqlens_q.shape[0] - 1, key.shape[1])
502 +
503 flash_attn_varlen_func(
504 q=query[:num_actual_tokens],
505 k=key_cache,
506 v=value_cache,
507 out=output[:num_actual_tokens],
288 - cu_seqlens_q=attn_metadata.query_start_loc,
289 - max_seqlen_q=attn_metadata.max_query_len,
290 - seqused_k=attn_metadata.seq_lens,
291 - max_seqlen_k=attn_metadata.max_seq_len,
508 + cu_seqlens_q=cu_seqlens_q,
509 + max_seqlen_q=max_seqlen_q,
510 + seqused_k=seqused_k,
511 + max_seqlen_k=max_seqlen_k,
512 softmax_scale=self.scale,
513 causal=True,
514 alibi_slopes=self.alibi_slopes,
515 window_size=self.sliding_window,
296 - block_table=attn_metadata.block_table,
516 + block_table=block_table,
517 softcap=self.logits_soft_cap,
518 fa_version=self.vllm_flash_attn_version,
519 q_descale=layer._q_scale.expand(descale_shape),
@@ -302,6 +522,8 @@ class FlashAttentionImpl(AttentionImpl):
522 )
523 return output
524
525 + assert not use_local_attn, (
526 + "Cascade attention does not support local attention.")
527 # Cascade attention (rare case).
528 cascade_attention(
529 output[:num_actual_tokens],
vllm/v1/attention/backends/triton_attn.py
+37 -18
@@ -70,6 +70,7 @@ class TritonAttentionImpl(AttentionImpl):
70 blocksparse_params: Optional[dict[str, Any]] = None,
71 logits_soft_cap: Optional[float] = None,
72 attn_type: AttentionType = AttentionType.DECODER,
73 + use_irope: bool = False,
74 ) -> None:
75 if blocksparse_params is not None:
76 raise ValueError(
@@ -86,6 +87,7 @@ class TritonAttentionImpl(AttentionImpl):
87 else:
88 self.sliding_window = (sliding_window - 1, 0)
89 self.kv_cache_dtype = kv_cache_dtype
90 + self.use_irope = use_irope
91
92 assert self.num_heads % self.num_kv_heads == 0
93 self.num_queries_per_kv = self.num_heads // self.num_kv_heads
@@ -156,24 +158,41 @@ class TritonAttentionImpl(AttentionImpl):
158 layer._v_scale,
159 )
160
161 + use_local_attn = \
162 + (self.use_irope and attn_metadata.local_attn_metadata is not None)
163 +
164 + if use_local_attn:
165 + assert attn_metadata.local_attn_metadata is not None
166 + local_metadata = attn_metadata.local_attn_metadata
167 + cu_seqlens_q = local_metadata.local_query_start_loc
168 + sequesd_k = local_metadata.local_seqused_k
169 + max_seqlen_q = local_metadata.local_max_query_len
170 + max_seqlen_k = local_metadata.local_max_seq_len
171 + block_table = local_metadata.local_block_table
172 + else:
173 + cu_seqlens_q = attn_metadata.query_start_loc
174 + sequesd_k = attn_metadata.seq_lens
175 + max_seqlen_q = attn_metadata.max_query_len
176 + max_seqlen_k = attn_metadata.max_seq_len
177 + block_table = attn_metadata.block_table
178 +
179 # Compute attention and update output up to `num_actual_tokens`.
160 - chunked_prefill_paged_decode(
161 - query=query[:num_actual_tokens],
162 - key=key[:num_actual_tokens],
163 - value=value[:num_actual_tokens],
164 - output=output[:num_actual_tokens],
165 - kv_cache_dtype=self.kv_cache_dtype,
166 - key_cache=key_cache,
167 - value_cache=value_cache,
168 - block_table=attn_metadata.block_table,
169 - query_start_loc=attn_metadata.query_start_loc,
170 - seq_lens=attn_metadata.seq_lens,
171 - max_seq_len=attn_metadata.max_seq_len,
172 - max_query_len=attn_metadata.max_query_len,
173 - k_scale=layer._k_scale,
174 - v_scale=layer._v_scale,
175 - alibi_slopes=self.alibi_slopes,
176 - sliding_window=self.sliding_window[0],
177 - sm_scale=self.scale)
180 + chunked_prefill_paged_decode(query=query[:num_actual_tokens],
181 + key=key[:num_actual_tokens],
182 + value=value[:num_actual_tokens],
183 + output=output[:num_actual_tokens],
184 + kv_cache_dtype=self.kv_cache_dtype,
185 + key_cache=key_cache,
186 + value_cache=value_cache,
187 + block_table=block_table,
188 + query_start_loc=cu_seqlens_q,
189 + seq_lens=sequesd_k,
190 + max_seq_len=max_seqlen_k,
191 + max_query_len=max_seqlen_q,
192 + k_scale=layer._k_scale,
193 + v_scale=layer._v_scale,
194 + alibi_slopes=self.alibi_slopes,
195 + sliding_window=self.sliding_window[0],
196 + sm_scale=self.scale)
197
198 return output
vllm/v1/worker/gpu_model_runner.py
+1
@@ -113,6 +113,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
113 self.num_kv_heads = model_config.get_num_kv_heads(parallel_config)
114 self.head_size = model_config.get_head_size()
115 self.hidden_size = model_config.get_hidden_size()
116 + self.attention_chunk_size = model_config.attention_chunk_size
117
118 self.attn_backend = get_attn_backend(
119 self.head_size,