Merge branch 'main' into patches
This commit is contained in:
@@ -8,7 +8,7 @@ ARG TRITON_VERSION=3.8.0+git669b31ac
|
||||
ARG TRITON_REPO="https://github.com/ROCm/triton.git"
|
||||
ARG FA_BRANCH="0e60e394"
|
||||
ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git"
|
||||
ARG AITER_BRANCH="v0.1.23"
|
||||
ARG AITER_BRANCH="v0.1.24.post1"
|
||||
ARG AITER_REPO="https://github.com/ROCm/aiter.git"
|
||||
ARG MORI_BRANCH="v1.2.3.post1"
|
||||
ARG MORI_REPO="https://github.com/ROCm/mori.git"
|
||||
|
||||
@@ -9,7 +9,7 @@ ARG PYTORCH_AUDIO_BRANCH="v2.11.0"
|
||||
ARG PYTORCH_AUDIO_REPO="https://github.com/pytorch/audio.git"
|
||||
ARG FA_BRANCH="0e60e394"
|
||||
ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git"
|
||||
ARG AITER_BRANCH="v0.1.23"
|
||||
ARG AITER_BRANCH="v0.1.24.post1"
|
||||
ARG AITER_REPO="https://github.com/ROCm/aiter.git"
|
||||
ARG MORI_BRANCH="v1.2.3.post1"
|
||||
ARG MORI_REPO="https://github.com/ROCm/mori.git"
|
||||
|
||||
@@ -19,12 +19,10 @@ vLLM also supports model implementations that are available in Transformers. We
|
||||
|
||||
Currently, the Transformers modeling backend works for the following:
|
||||
|
||||
- Modalities: embedding models, language models, vision-language models* and audio-language models
|
||||
- Modalities: embedding models, language models, vision-language models and audio-language models
|
||||
- Architectures: encoder-only, decoder-only, mixture-of-experts
|
||||
- Attention types: full attention and/or sliding attention
|
||||
|
||||
_*Vision-language models currently accept only image inputs. Support for video inputs will be added in a future release._
|
||||
|
||||
If the Transformers model implementation follows all the steps in [writing a custom model](#writing-custom-models) then, when used with the Transformers modeling backend, it will be compatible with the following features of vLLM:
|
||||
|
||||
- All the features listed in the [compatibility matrix](../features/README.md#feature-x-feature)
|
||||
|
||||
@@ -455,6 +455,31 @@ def test_per_head_quant_scales_backend_selection(
|
||||
assert backend_name in str(exc_info.value)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
CudaPlatform is None, reason="CUDA platform is required for this test"
|
||||
)
|
||||
@pytest.mark.parametrize("head_size", [64, 80, 96, 128, 192, 256])
|
||||
def test_int4_per_token_head_head_size_selection(head_size: int):
|
||||
"""INT4 per-token-head KV cache is selectable at any even head size.
|
||||
|
||||
The Hadamard rotation the INT4 write path applies is now defined for
|
||||
non-power-of-two rows as well (it runs block-diagonally, see
|
||||
``fast_hadamard_transform``), so a model such as Phi-3 (head_size=96) must
|
||||
keep selecting TRITON_ATTN instead of being rejected.
|
||||
"""
|
||||
vllm_config = VllmConfig(cache_config=CacheConfig(block_size=64))
|
||||
with (
|
||||
set_current_vllm_config(vllm_config),
|
||||
patch("vllm.platforms.current_platform", CudaPlatform()),
|
||||
):
|
||||
backend = get_attn_backend(
|
||||
head_size=head_size,
|
||||
dtype=torch.float16,
|
||||
kv_cache_dtype="int4_per_token_head",
|
||||
)
|
||||
assert backend.get_name() == "TRITON_ATTN"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"backend_name,use_non_causal,should_succeed",
|
||||
[
|
||||
|
||||
@@ -713,8 +713,8 @@ def test_low_latency_table_capability_routing(
|
||||
def test_qwen4_exp_hopper_plans_are_valid() -> None:
|
||||
plans = qwen4_exp_gemm.QWEN4_EXP_SM90_GEMM_PLANS
|
||||
|
||||
assert len(plans) == 9
|
||||
assert sum(map(len, plans.values())) == 31
|
||||
assert len(plans) == 7
|
||||
assert sum(map(len, plans.values())) == 24
|
||||
assert (320, 10240) in plans
|
||||
assert (10240, 320) not in plans
|
||||
for (n, k), shape_plans in plans.items():
|
||||
@@ -728,10 +728,11 @@ def test_qwen4_exp_hopper_plans_are_valid() -> None:
|
||||
@pytest.mark.parametrize(
|
||||
"capability,expected_plans",
|
||||
[
|
||||
((10, 3), qwen4_exp_gemm.QWEN4_EXP_GEMM_PLANS),
|
||||
((10, 3), qwen4_exp_gemm.QWEN4_EXP_SM103_GEMM_PLANS),
|
||||
((10, 0), qwen4_exp_gemm.QWEN4_EXP_SM100_GEMM_PLANS),
|
||||
((9, 0), qwen4_exp_gemm.QWEN4_EXP_SM90_GEMM_PLANS),
|
||||
((12, 1), qwen4_exp_gemm.QWEN4_EXP_SM121_GEMM_PLANS),
|
||||
((12, 0), {}),
|
||||
((8, 0), {}),
|
||||
],
|
||||
)
|
||||
@@ -742,11 +743,16 @@ def test_qwen4_exp_gemm_capability_routing(
|
||||
) -> None:
|
||||
monkeypatch.setattr(
|
||||
qwen4_exp_gemm.current_platform,
|
||||
"is_device_capability",
|
||||
lambda target: capability == target,
|
||||
"get_device_capability",
|
||||
lambda: capability,
|
||||
)
|
||||
|
||||
assert qwen4_exp_gemm._gemm_plans() == expected_plans
|
||||
assert (
|
||||
qwen4_exp_gemm.QWEN4_EXP_GEMM_PLANS_BY_CAPABILITY.get(
|
||||
qwen4_exp_gemm.current_platform.get_device_capability(), {}
|
||||
)
|
||||
== expected_plans
|
||||
)
|
||||
|
||||
|
||||
def test_installation_is_shape_specific_and_unquantized(
|
||||
|
||||
@@ -69,13 +69,8 @@ AUDIO_MODEL_SETTINGS: dict[str, dict[str, Any]] = {
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_id", list(AUDIO_MODEL_SETTINGS))
|
||||
def test_transformers_audio_generation(
|
||||
hf_runner: type[HfRunner],
|
||||
vllm_runner: type[VllmRunner],
|
||||
monkeypatch,
|
||||
model_id: str,
|
||||
):
|
||||
@pytest.fixture(autouse=True)
|
||||
def use_spawn_for_audio_models(monkeypatch):
|
||||
"""Single-process workaround for V1 fork safety deadlock issue
|
||||
(vllm-project/vllm/issues/17676). Running multiple audio models together
|
||||
under pytest can cause (possibly flaky) hangs, so they are grouped under
|
||||
@@ -87,6 +82,13 @@ def test_transformers_audio_generation(
|
||||
disable_envs_cache()
|
||||
monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_id", list(AUDIO_MODEL_SETTINGS))
|
||||
def test_transformers_audio_generation(
|
||||
hf_runner: type[HfRunner],
|
||||
vllm_runner: type[VllmRunner],
|
||||
model_id: str,
|
||||
):
|
||||
settings = AUDIO_MODEL_SETTINGS[model_id]
|
||||
audio_lora_path = settings.get("audio_lora_path")
|
||||
|
||||
|
||||
@@ -0,0 +1,100 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from transformers import AutoModelForImageTextToText
|
||||
|
||||
from vllm.assets.video import VideoAsset
|
||||
from vllm.envs import disable_envs_cache
|
||||
|
||||
from ....conftest import HfRunner, VllmRunner
|
||||
from ...utils import check_logprobs_close
|
||||
from .vlm_utils import model_utils
|
||||
|
||||
VIDEO_ASSET = VideoAsset("baby_reading", num_frames=8)
|
||||
|
||||
VIDEO_MODEL_SETTINGS: dict[str, dict[str, Any]] = {
|
||||
"Qwen/Qwen3-VL-2B-Instruct": {
|
||||
"prompt": (
|
||||
"<|im_start|>user\n"
|
||||
"<|vision_start|><|video_pad|><|vision_end|>"
|
||||
"Describe this video.<|im_end|>\n"
|
||||
"<|im_start|>assistant\n"
|
||||
),
|
||||
},
|
||||
"llava-hf/llava-onevision-qwen2-0.5b-ov-hf": {
|
||||
"prompt": (
|
||||
"<|im_start|>user <video>\nDescribe this video.<|im_end|>"
|
||||
"<|im_start|>assistant\n"
|
||||
),
|
||||
},
|
||||
"CohereLabs/North-Micro-Vision-Instruct": {
|
||||
"prompt": (
|
||||
"<BOS_TOKEN><|START_OF_TURN_TOKEN|><|USER_TOKEN|>"
|
||||
"<|VISION_START|><|VIDEO_PAD|><|VISION_END|>Describe this video."
|
||||
"<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>"
|
||||
),
|
||||
},
|
||||
"Qwen/Qwen2.5-VL-3B-Instruct": {
|
||||
"prompt": (
|
||||
"<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
|
||||
"<|im_start|>user\n"
|
||||
"<|vision_start|><|video_pad|><|vision_end|>"
|
||||
"Describe this video.<|im_end|>\n"
|
||||
"<|im_start|>assistant\n"
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def use_spawn_for_video_models(monkeypatch):
|
||||
"""Single-process workaround for V1 fork safety deadlock issue
|
||||
(vllm-project/vllm/issues/17676). Running multiple video models together
|
||||
under pytest can cause (possibly flaky) hangs, so they are grouped under
|
||||
the same config. Using VLLM_WORKER_MULTIPROC_METHOD=spawn avoids the
|
||||
deadlock and allows worker processes to terminate cleanly, and release
|
||||
GPU memory between test runs until the issue is fixed."""
|
||||
# TODO: Remove monkeypatch once
|
||||
# https://github.com/vllm-project/vllm/issues/17676 is fixed.
|
||||
disable_envs_cache()
|
||||
monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_id", list(VIDEO_MODEL_SETTINGS))
|
||||
def test_transformers_video_generation(
|
||||
hf_runner: type[HfRunner],
|
||||
vllm_runner: type[VllmRunner],
|
||||
model_id: str,
|
||||
):
|
||||
prompt = VIDEO_MODEL_SETTINGS[model_id]["prompt"]
|
||||
video = (VIDEO_ASSET.np_ndarrays, VIDEO_ASSET.metadata)
|
||||
|
||||
with vllm_runner(
|
||||
model_id,
|
||||
model_impl="transformers",
|
||||
dtype="bfloat16",
|
||||
max_model_len=8192,
|
||||
enforce_eager=True,
|
||||
limit_mm_per_prompt={"video": 1},
|
||||
) as vllm_model:
|
||||
vllm_outputs = vllm_model.generate_greedy_logprobs(
|
||||
[prompt], 128, num_logprobs=10, videos=[video]
|
||||
)
|
||||
|
||||
with hf_runner(
|
||||
model_id, dtype="bfloat16", auto_cls=AutoModelForImageTextToText
|
||||
) as hf_model:
|
||||
hf_model = model_utils.qwen3_vl_patch_hf_runner(hf_model)
|
||||
hf_outputs = hf_model.generate_greedy_logprobs_limit(
|
||||
[prompt], 128, num_logprobs=10, videos=[video]
|
||||
)
|
||||
|
||||
check_logprobs_close(
|
||||
outputs_0_lst=hf_outputs,
|
||||
outputs_1_lst=vllm_outputs,
|
||||
name_0="hf",
|
||||
name_1="vllm",
|
||||
)
|
||||
@@ -0,0 +1,277 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import numpy as np
|
||||
import pytest
|
||||
from transformers import Qwen3VLVideoProcessor
|
||||
|
||||
from vllm.config import ModelConfig
|
||||
from vllm.multimodal import MULTIMODAL_REGISTRY
|
||||
from vllm.multimodal.cache import MultiModalProcessorOnlyCache
|
||||
|
||||
VIDEO_MODEL_SETTINGS = {
|
||||
"Qwen/Qwen3-VL-2B-Instruct": {
|
||||
"prompt": (
|
||||
"<|im_start|>user\n"
|
||||
"<|vision_start|><|video_pad|><|vision_end|>"
|
||||
"Describe this video.<|im_end|>\n"
|
||||
"<|im_start|>assistant\n"
|
||||
),
|
||||
},
|
||||
"llava-hf/llava-onevision-qwen2-0.5b-ov-hf": {
|
||||
"prompt": (
|
||||
"<|im_start|>user <video>\nDescribe this video.<|im_end|>"
|
||||
"<|im_start|>assistant\n"
|
||||
),
|
||||
},
|
||||
"CohereLabs/North-Micro-Vision-Instruct": {
|
||||
"prompt": (
|
||||
"<BOS_TOKEN><|START_OF_TURN_TOKEN|><|USER_TOKEN|>"
|
||||
"<|VISION_START|><|VIDEO_PAD|><|VISION_END|>Describe this video."
|
||||
"<|END_OF_TURN_TOKEN|><|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>"
|
||||
),
|
||||
},
|
||||
"Qwen/Qwen2.5-VL-3B-Instruct": {
|
||||
"prompt": (
|
||||
"<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
|
||||
"<|im_start|>user\n"
|
||||
"<|vision_start|><|video_pad|><|vision_end|>"
|
||||
"Describe this video.<|im_end|>\n"
|
||||
"<|im_start|>assistant\n"
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _video(
|
||||
num_frames: int = 8,
|
||||
seed: int = 0,
|
||||
fps: float = 1.0,
|
||||
do_sample_frames: bool = False,
|
||||
) -> tuple[np.ndarray, dict]:
|
||||
rng = np.random.default_rng(seed)
|
||||
frames = rng.integers(256, size=(num_frames, 64, 64, 3), dtype=np.uint8)
|
||||
metadata = {
|
||||
"fps": fps,
|
||||
"total_num_frames": num_frames,
|
||||
"frames_indices": list(range(num_frames)),
|
||||
"do_sample_frames": do_sample_frames,
|
||||
}
|
||||
return frames, metadata
|
||||
|
||||
|
||||
def _hf_video_kwargs(*videos):
|
||||
return {
|
||||
"videos": [frames for frames, _ in videos],
|
||||
"video_metadata": [
|
||||
{k: v for k, v in metadata.items() if k != "do_sample_frames"}
|
||||
for _, metadata in videos
|
||||
],
|
||||
"do_sample_frames": videos[0][1]["do_sample_frames"],
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_id", list(VIDEO_MODEL_SETTINGS))
|
||||
def test_video_multimodal_processor(model_id):
|
||||
mm_processor = MULTIMODAL_REGISTRY.create_processor(
|
||||
ModelConfig(model=model_id, model_impl="transformers")
|
||||
)
|
||||
hf_processor = mm_processor.info.get_hf_processor()
|
||||
|
||||
video = _video()
|
||||
frames, _ = video
|
||||
result = mm_processor(
|
||||
prompt=VIDEO_MODEL_SETTINGS[model_id]["prompt"],
|
||||
mm_items=mm_processor.info.parse_mm_data({"video": video}),
|
||||
hf_processor_mm_kwargs={},
|
||||
)
|
||||
|
||||
hf_inputs = hf_processor(
|
||||
text=VIDEO_MODEL_SETTINGS[model_id]["prompt"],
|
||||
**_hf_video_kwargs(video),
|
||||
return_mm_token_type_ids=True,
|
||||
return_tensors="pt",
|
||||
)
|
||||
assert result["prompt_token_ids"] == hf_inputs["input_ids"][0].tolist()
|
||||
|
||||
(placeholder,) = result["mm_placeholders"]["video"]
|
||||
(num_video_tokens,) = hf_processor._get_num_multimodal_tokens(
|
||||
video_sizes=[frames.shape[:3]]
|
||||
)["num_video_tokens"]
|
||||
assert placeholder.get_num_embeds() == num_video_tokens
|
||||
assert placeholder.get_num_embeds() == int(
|
||||
(hf_inputs["mm_token_type_ids"] == 2).sum()
|
||||
)
|
||||
|
||||
(item,) = result["mm_kwargs"]["video"]
|
||||
rows = item["pixel_values_videos"].data.shape[0]
|
||||
if "video_grid_thw" in item:
|
||||
assert rows == int(item["num_video_patches"].data)
|
||||
else:
|
||||
assert rows == len(frames)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_id", list(VIDEO_MODEL_SETTINGS))
|
||||
def test_video_multiple_inputs(model_id):
|
||||
"""Multiple videos per prompt are each detected as a separate placeholder
|
||||
and multi-modal item by the Transformers modeling backend."""
|
||||
mm_processor = MULTIMODAL_REGISTRY.create_processor(
|
||||
ModelConfig(model=model_id, model_impl="transformers")
|
||||
)
|
||||
hf_processor = mm_processor.info.get_hf_processor()
|
||||
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "video"},
|
||||
{"type": "video"},
|
||||
{"type": "text", "text": "Compare these videos."},
|
||||
],
|
||||
}
|
||||
]
|
||||
prompt = hf_processor.apply_chat_template(
|
||||
messages, add_generation_prompt=True, tokenize=False
|
||||
)
|
||||
videos = [_video(8, seed=0), _video(8, seed=1)]
|
||||
result = mm_processor(
|
||||
prompt=prompt,
|
||||
mm_items=mm_processor.info.parse_mm_data({"video": videos}),
|
||||
hf_processor_mm_kwargs={},
|
||||
)
|
||||
|
||||
assert len(result["mm_placeholders"]["video"]) == 2
|
||||
assert len(result["mm_kwargs"]["video"]) == 2
|
||||
hf_inputs = hf_processor(
|
||||
text=prompt, **_hf_video_kwargs(*videos), return_tensors="pt"
|
||||
)
|
||||
assert result["prompt_token_ids"] == hf_inputs["input_ids"][0].tolist()
|
||||
|
||||
|
||||
def test_video_fields_not_claimed_by_image():
|
||||
"""A prompt holding an image and a video keeps the fields of each on its own
|
||||
modality, so the image branch does not claim the video processor's outputs."""
|
||||
model_id = "Qwen/Qwen3-VL-2B-Instruct"
|
||||
mm_processor = MULTIMODAL_REGISTRY.create_processor(
|
||||
ModelConfig(model=model_id, model_impl="transformers")
|
||||
)
|
||||
|
||||
result = mm_processor(
|
||||
prompt=(
|
||||
"<|vision_start|><|image_pad|><|vision_end|> and "
|
||||
"<|vision_start|><|video_pad|><|vision_end|> What do these show?"
|
||||
),
|
||||
mm_items=mm_processor.info.parse_mm_data(
|
||||
{"image": _video()[0][0], "video": _video()}
|
||||
),
|
||||
hf_processor_mm_kwargs={},
|
||||
)
|
||||
|
||||
(image_item,) = result["mm_kwargs"]["image"]
|
||||
(video_item,) = result["mm_kwargs"]["video"]
|
||||
assert {"pixel_values", "image_grid_thw"} <= set(image_item.keys())
|
||||
assert {"pixel_values_videos", "video_grid_thw"} <= set(video_item.keys())
|
||||
assert len(result["mm_placeholders"]["image"]) == 1
|
||||
assert len(result["mm_placeholders"]["video"]) == 1
|
||||
|
||||
|
||||
def test_repeated_video_hits_the_processor_cache():
|
||||
"""Check that mm caching is actually working."""
|
||||
model_config = ModelConfig(
|
||||
model="Qwen/Qwen3-VL-2B-Instruct", model_impl="transformers"
|
||||
)
|
||||
model_config.multimodal_config.mm_processor_cache_gb = 4
|
||||
mm_processor = MULTIMODAL_REGISTRY.create_processor(model_config)
|
||||
cache = MultiModalProcessorOnlyCache(model_config)
|
||||
video = _video()
|
||||
|
||||
def process():
|
||||
return mm_processor(
|
||||
prompt=VIDEO_MODEL_SETTINGS["Qwen/Qwen3-VL-2B-Instruct"]["prompt"],
|
||||
mm_items=mm_processor.info.parse_mm_data({"video": video}),
|
||||
hf_processor_mm_kwargs={},
|
||||
cache=cache,
|
||||
)
|
||||
|
||||
first, second = process(), process()
|
||||
|
||||
assert cache.make_stats().hits > 0
|
||||
assert first["prompt_token_ids"] == second["prompt_token_ids"]
|
||||
assert first["mm_hashes"] == second["mm_hashes"]
|
||||
|
||||
|
||||
def test_video_unsupported_when_processor_cannot_count_tokens(monkeypatch):
|
||||
"""A processor that cannot count video tokens is served as image-only."""
|
||||
monkeypatch.delattr(Qwen3VLVideoProcessor, "get_num_of_video_patches")
|
||||
mm_processor = MULTIMODAL_REGISTRY.create_processor(
|
||||
ModelConfig(model="Qwen/Qwen3-VL-2B-Instruct", model_impl="transformers")
|
||||
)
|
||||
|
||||
assert "video" not in mm_processor.info.get_supported_mm_limits()
|
||||
|
||||
|
||||
def test_merged_token_fields_split_per_video():
|
||||
"""VideoLLaMA3's compression mask has one row per merged token, not per patch."""
|
||||
mm_processor = MULTIMODAL_REGISTRY.create_processor(
|
||||
ModelConfig(model="lkhl/VideoLLaMA3-2B-Image-HF", model_impl="transformers")
|
||||
)
|
||||
|
||||
result = mm_processor(
|
||||
prompt="<|im_start|>user\n<|video_pad|>\n<|video_pad|>\nDescribe.<|im_end|>\n",
|
||||
mm_items=mm_processor.info.parse_mm_data({"video": [_video(8), _video(4)]}),
|
||||
hf_processor_mm_kwargs={},
|
||||
)
|
||||
|
||||
items = result["mm_kwargs"]["video"]
|
||||
assert [len(item["video_compression_mask"].data) for item in items] == [128, 64]
|
||||
|
||||
|
||||
def test_video_sampled_by_processor():
|
||||
"""Frames left for the processor to sample yield the tokens HF does."""
|
||||
model_id = "Qwen/Qwen3-VL-2B-Instruct"
|
||||
mm_processor = MULTIMODAL_REGISTRY.create_processor(
|
||||
ModelConfig(model=model_id, model_impl="transformers")
|
||||
)
|
||||
hf_processor = mm_processor.info.get_hf_processor()
|
||||
prompt = VIDEO_MODEL_SETTINGS[model_id]["prompt"]
|
||||
video = _video(64, fps=30.0, do_sample_frames=True)
|
||||
|
||||
result = mm_processor(
|
||||
prompt=prompt,
|
||||
mm_items=mm_processor.info.parse_mm_data({"video": video}),
|
||||
hf_processor_mm_kwargs={},
|
||||
)
|
||||
hf_inputs = hf_processor(
|
||||
text=prompt,
|
||||
**_hf_video_kwargs(video),
|
||||
return_mm_token_type_ids=True,
|
||||
return_tensors="pt",
|
||||
)
|
||||
|
||||
assert result["prompt_token_ids"] == hf_inputs["input_ids"][0].tolist()
|
||||
(placeholder,) = result["mm_placeholders"]["video"]
|
||||
assert placeholder.get_num_embeds() == int(
|
||||
(hf_inputs["mm_token_type_ids"] == 2).sum()
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_id", list(VIDEO_MODEL_SETTINGS))
|
||||
def test_dummy_video_has_the_most_tokens(model_id):
|
||||
"""Profiling must build a dummy video with exactly the most tokens a video
|
||||
can produce, or memory is under- or over-reserved."""
|
||||
mm_processor = MULTIMODAL_REGISTRY.create_processor(
|
||||
ModelConfig(model=model_id, model_impl="transformers")
|
||||
)
|
||||
info = mm_processor.info
|
||||
seq_len = info.ctx.model_config.max_model_len
|
||||
|
||||
result = mm_processor.get_dummy_mm_inputs({"video": 1})
|
||||
|
||||
(placeholder,) = result["mm_placeholders"]["video"]
|
||||
max_tokens = info.get_max_video_tokens(seq_len, {"video": 1})
|
||||
assert placeholder.get_num_embeds() == max_tokens
|
||||
# Shrinking the dummy frames must never cost tokens
|
||||
num_frames = info.get_num_frames_with_most_features(seq_len, {"video": 1})
|
||||
(huge_frame_tokens,) = info.get_hf_processor()._get_num_multimodal_tokens(
|
||||
video_sizes=[[num_frames, 10_000, 10_000]]
|
||||
)["num_video_tokens"]
|
||||
assert max_tokens >= huge_frame_tokens
|
||||
@@ -74,13 +74,14 @@ def _make_main_inputs(num_tokens: int, fp8_cache: bool) -> dict:
|
||||
slots = torch.randperm(num_blocks * MAIN_PAGE, device="cuda")[:num_tokens]
|
||||
slots[-1] = -1
|
||||
norm_weights = torch.randn(2, MAIN_D, dtype=torch.bfloat16, device="cuda") * 0.2
|
||||
qkv_width = 2 * (MAIN_HQ + MAIN_HK) * MAIN_D
|
||||
return dict(
|
||||
main_qkv=torch.randn(
|
||||
num_tokens,
|
||||
2 * (MAIN_HQ + MAIN_HK) * MAIN_D,
|
||||
qkv_width + (HQ + 1) * D,
|
||||
dtype=torch.bfloat16,
|
||||
device="cuda",
|
||||
),
|
||||
)[:, :qkv_width],
|
||||
main_q_norm_weight=norm_weights[0],
|
||||
main_k_norm_weight=norm_weights[1],
|
||||
main_eps=EPS,
|
||||
@@ -316,9 +317,12 @@ def test_qsa_fused_prepare_matches_unfused(
|
||||
)
|
||||
unfused_compressed = fused_compressed.clone()
|
||||
|
||||
# TP=4 QKVG width makes the indexer slice row-strided, as in the merged
|
||||
# projection. The exact width is unimportant; the larger stride is tested.
|
||||
main_qkvg_width = 3584
|
||||
projected_qk = torch.randn(
|
||||
num_tokens, (HQ + 1) * D, dtype=torch.bfloat16, device=device
|
||||
)
|
||||
num_tokens, main_qkvg_width + (HQ + 1) * D, dtype=torch.bfloat16, device=device
|
||||
)[:, main_qkvg_width:]
|
||||
q_weight = torch.randn(D, dtype=torch.bfloat16, device=device) * 0.2
|
||||
k_weight = torch.randn(D, dtype=torch.bfloat16, device=device) * 0.2
|
||||
|
||||
|
||||
@@ -587,12 +587,14 @@ class MarkingStub(SupportsMultiModal, nn.Module):
|
||||
_pre_trained_model_classes = Base._pre_trained_model_classes
|
||||
|
||||
|
||||
def build_marked_model(image_limit: int, skip_tokenizer_init: bool = False):
|
||||
def build_marked_model(
|
||||
image_limit: int, video_limit: int = 0, skip_tokenizer_init: bool = False
|
||||
):
|
||||
"""Build the HF model inside the marking context and return it."""
|
||||
model_config = ModelConfig(
|
||||
model=MULTIMODAL_MODEL,
|
||||
model_impl="transformers",
|
||||
limit_mm_per_prompt={"image": image_limit},
|
||||
limit_mm_per_prompt={"image": image_limit, "video": video_limit},
|
||||
)
|
||||
# Set after construction: building the config itself needs the tokenizer
|
||||
model_config.skip_tokenizer_init = skip_tokenizer_init
|
||||
@@ -606,10 +608,15 @@ def build_marked_model(image_limit: int, skip_tokenizer_init: bool = False):
|
||||
return stub.model
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("image_limit", "skipped"), [(0, True), (4, False)])
|
||||
def test_tower_weights_skipped_when_modality_disabled(image_limit, skipped):
|
||||
"""`--limit-mm-per-prompt image=0` should drop the vision tower's weights."""
|
||||
vision_tower = build_marked_model(image_limit).vision_tower
|
||||
@pytest.mark.parametrize(
|
||||
("image_limit", "video_limit", "skipped"),
|
||||
[(0, 0, True), (4, 0, False), (0, 4, False)],
|
||||
)
|
||||
def test_tower_weights_skipped_when_modality_disabled(
|
||||
image_limit, video_limit, skipped
|
||||
):
|
||||
"""Only `--limit-mm-per-prompt image=0,video=0` drops the tower they share."""
|
||||
vision_tower = build_marked_model(image_limit, video_limit).vision_tower
|
||||
assert isinstance(vision_tower, StageMissingLayer) is skipped
|
||||
|
||||
|
||||
|
||||
@@ -1552,6 +1552,53 @@ def test_chunked_prefill_reads_live_block_ids() -> None:
|
||||
assert len(meta2.store_gpu_blocks) == 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize("use_v2_model_runner", [False, True])
|
||||
@pytest.mark.parametrize("same_output", [False, True])
|
||||
def test_eager_store_resets_block_cursor_after_preemption(
|
||||
use_v2_model_runner: bool,
|
||||
same_output: bool,
|
||||
) -> None:
|
||||
"""A resumed request stores its confirmed tail after preemption."""
|
||||
fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=8, lazy=False)
|
||||
sched = fix.scheduler
|
||||
req = make_request(num_blocks=3)
|
||||
|
||||
initial_blocks = _alloc_and_register(fix, req, num_blocks=2)
|
||||
sched.update_state_after_alloc(req, initial_blocks, num_external_tokens=0)
|
||||
initial_output = make_scheduler_output(
|
||||
{req.request_id: 2 * BLOCK_SIZE},
|
||||
new_reqs={req.request_id: initial_blocks.get_block_ids()},
|
||||
)
|
||||
initial_meta = sched.build_connector_meta(initial_output)
|
||||
assert len(initial_meta.store_gpu_blocks) == 2
|
||||
simulate_store_completion(sched, initial_meta.store_event)
|
||||
|
||||
if not same_output:
|
||||
preemption_output = make_scheduler_output({})
|
||||
preemption_output.preempted_req_ids = {req.request_id}
|
||||
sched.build_connector_meta(preemption_output)
|
||||
fix.gpu_block_pool.free_blocks(initial_blocks.blocks[0])
|
||||
|
||||
resumed_blocks = _alloc_and_register(fix, req, num_blocks=3)
|
||||
sched.update_state_after_alloc(req, resumed_blocks, num_external_tokens=0)
|
||||
if use_v2_model_runner:
|
||||
resumed_output = make_scheduler_output(
|
||||
{req.request_id: 3 * BLOCK_SIZE},
|
||||
new_reqs={req.request_id: resumed_blocks.get_block_ids()},
|
||||
)
|
||||
else:
|
||||
resumed_output = make_scheduler_output(
|
||||
{req.request_id: 3 * BLOCK_SIZE},
|
||||
cached_req_new_blocks={req.request_id: resumed_blocks.get_block_ids()},
|
||||
)
|
||||
resumed_output.scheduled_cached_reqs.resumed_req_ids.add(req.request_id)
|
||||
if same_output:
|
||||
resumed_output.preempted_req_ids = {req.request_id}
|
||||
|
||||
resumed_meta = sched.build_connector_meta(resumed_output)
|
||||
assert resumed_meta.store_gpu_blocks == [resumed_blocks.blocks[0][2].block_id]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 10: Partial GPU prefix hit + CPU load + new compute blocks
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -152,6 +152,46 @@ def test_scoped_max_split_ignores_async_allocator(monkeypatch):
|
||||
set_settings.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kv_cache_memory_bytes", [None, 1])
|
||||
def test_memory_profile_bounds_decode_logits_rows(monkeypatch, kv_cache_memory_bytes):
|
||||
"""ROCm decode workspace uses max_num_seqs during startup profiling."""
|
||||
from vllm.config import get_current_vllm_config_or_none
|
||||
from vllm.v1.attention.ops.rocm_aiter_mla_sparse import _max_decode_logits_rows
|
||||
|
||||
config = SimpleNamespace(
|
||||
scheduler_config=SimpleNamespace(max_num_seqs=16),
|
||||
speculative_config=SimpleNamespace(num_speculative_tokens=5),
|
||||
)
|
||||
|
||||
class Profiled(Exception):
|
||||
pass
|
||||
|
||||
def profile_run(*, randomize_inputs=False):
|
||||
assert randomize_inputs is True
|
||||
# Without the config context, the helper falls back to all 32768
|
||||
# batched tokens, which would reserve 128 GiB at 1M context.
|
||||
assert _max_decode_logits_rows(32768) == 16 * (1 + 5)
|
||||
raise Profiled
|
||||
|
||||
worker = SimpleNamespace(
|
||||
vllm_config=config,
|
||||
randomize_dummy_inputs=True,
|
||||
cache_config=SimpleNamespace(kv_cache_memory_bytes=kv_cache_memory_bytes),
|
||||
model_runner=SimpleNamespace(profile_run=profile_run, model_memory_usage=0),
|
||||
init_snapshot=SimpleNamespace(free_memory=2),
|
||||
_scoped_allocator_max_split=lambda **kwargs: nullcontext(),
|
||||
)
|
||||
monkeypatch.setattr(gpu_worker, "maybe_apply_startup_plan", lambda _: None)
|
||||
monkeypatch.setattr(
|
||||
gpu_worker, "memory_profiling", lambda *args, **kwargs: nullcontext()
|
||||
)
|
||||
|
||||
assert get_current_vllm_config_or_none() is None
|
||||
with pytest.raises(Profiled):
|
||||
gpu_worker.Worker.determine_available_memory(worker)
|
||||
assert get_current_vllm_config_or_none() is None
|
||||
|
||||
|
||||
# Startup-plan persistence (vllm/v1/worker/startup_plan.py), applied and
|
||||
# saved by Worker.determine_available_memory / compile_or_warm_up_model.
|
||||
|
||||
|
||||
+6
-18
@@ -37,7 +37,7 @@ from vllm.transformers_utils.config import (
|
||||
get_pooling_config,
|
||||
get_sentence_transformer_tokenizer_config,
|
||||
is_encoder_decoder,
|
||||
is_rope_parameters_nested,
|
||||
iter_rope_parameters,
|
||||
mrope_num_dims,
|
||||
try_get_dense_modules,
|
||||
try_get_generation_config,
|
||||
@@ -2566,24 +2566,14 @@ def _get_and_verify_max_len(
|
||||
)
|
||||
derived_max_model_len = default_max_len
|
||||
|
||||
# In Transformers v5 rope_parameters could be TypedDict or dict[str, TypedDict].
|
||||
# To simplify the verification, we convert it to dict[str, TypedDict].
|
||||
rope_parameters = getattr(hf_config, "rope_parameters", None)
|
||||
if rope_parameters and not is_rope_parameters_nested(rope_parameters):
|
||||
rope_parameters = {"": rope_parameters}
|
||||
if rope_parameters is not None:
|
||||
# Layers without RoPE do not contribute to context length scaling.
|
||||
rope_parameters = {
|
||||
layer_type: rp
|
||||
for layer_type, rp in rope_parameters.items()
|
||||
if rp is not None
|
||||
}
|
||||
# Layers without RoPE do not contribute to context length scaling.
|
||||
rope_parameters = list(iter_rope_parameters(hf_config))
|
||||
|
||||
# NOTE(woosuk): Gemma3's max_model_len (128K) is already scaled by RoPE
|
||||
# scaling, so we skip applying the scaling factor again.
|
||||
if rope_parameters is not None and "gemma3" not in hf_config.model_type:
|
||||
if rope_parameters and "gemma3" not in hf_config.model_type:
|
||||
scaling_factor = 1.0
|
||||
for rp in rope_parameters.values():
|
||||
for rp in rope_parameters:
|
||||
# No need to consider "type" key because of patch_rope_parameters when
|
||||
# loading HF config
|
||||
rope_type = rp["rope_type"]
|
||||
@@ -2623,9 +2613,7 @@ def _get_and_verify_max_len(
|
||||
if max_model_len is None or max_model_len == -1:
|
||||
# For LongRoPE, default to original_max_position_embeddings to avoid
|
||||
# performance degradation for shorter sequences
|
||||
if rope_parameters is not None and any(
|
||||
rp["rope_type"] == "longrope" for rp in rope_parameters.values()
|
||||
):
|
||||
if any(rp["rope_type"] == "longrope" for rp in rope_parameters):
|
||||
max_model_len = int(
|
||||
getattr(
|
||||
hf_config, "original_max_position_embeddings", derived_max_model_len
|
||||
|
||||
@@ -16,14 +16,15 @@
|
||||
# limitations under the License.
|
||||
"""Transformers modeling backend mixin for multi-modal models."""
|
||||
|
||||
import math
|
||||
from collections import defaultdict
|
||||
from collections.abc import Iterable, Mapping, Sequence
|
||||
from collections.abc import Callable, Iterable, Mapping, Sequence
|
||||
from contextlib import ExitStack, contextmanager
|
||||
from functools import cached_property
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import torch
|
||||
import transformers
|
||||
from transformers.utils.generic import ModelOutput
|
||||
|
||||
from vllm.compilation.decorators import should_torch_compile_mm_encoder
|
||||
from vllm.config.utils import getattr_iter
|
||||
@@ -47,6 +48,7 @@ from vllm.multimodal.processing import (
|
||||
PromptUpdateDetails,
|
||||
cached_encode,
|
||||
)
|
||||
from vllm.multimodal.processing.processor import HFMultiModalInputs
|
||||
from vllm.sequence import IntermediateTensors
|
||||
from vllm.utils.gpu_sync_debug import gpu_sync_allowed
|
||||
from vllm.utils.torch_utils import async_tensor_h2d
|
||||
@@ -62,7 +64,19 @@ if TYPE_CHECKING:
|
||||
logger = init_logger(__name__)
|
||||
|
||||
_MODALITY_TO_TOKEN_TYPE_ID = {"image": 1, "video": 2, "audio": 3}
|
||||
_MODALITY_SIZE_KEYS = {"audio": "num_audio_tokens", "image": "num_image_patches"}
|
||||
_MODALITY_SIZE_KEYS = {
|
||||
"audio": "num_audio_tokens",
|
||||
"image": "num_image_patches",
|
||||
"video": "num_video_patches",
|
||||
}
|
||||
# NOTE: Profiling cap as in llava_onevision._MAX_FRAMES_PER_VIDEO
|
||||
# past which a pixel budget only shrinks the frames, so Qwen3-VL's most is
|
||||
# 12168 tokens at 16 frames and 9600 at its `max_frames` of 768
|
||||
_MAX_FRAMES_PER_VIDEO = 16
|
||||
# Arbitrary large size, bounded only by the processor's own resizing
|
||||
_MAX_DUMMY_SIDE = 10_000
|
||||
# Tiny inputs trip processors' channel-axis inference
|
||||
_MIN_DUMMY_SIDE = 224
|
||||
|
||||
|
||||
def _get_embed_token_id(replacement_ids: torch.Tensor) -> int:
|
||||
@@ -70,6 +84,15 @@ def _get_embed_token_id(replacement_ids: torch.Tensor) -> int:
|
||||
return int(replacement_ids.mode().values)
|
||||
|
||||
|
||||
def _count_embed_tokens(seqs: list[list[int]]) -> torch.Tensor:
|
||||
"""Number of embedding tokens in each item's replacement."""
|
||||
counts = []
|
||||
for seq in seqs:
|
||||
ids = torch.tensor(seq)
|
||||
counts.append(int(ids.eq(_get_embed_token_id(ids)).sum()))
|
||||
return torch.tensor(counts)
|
||||
|
||||
|
||||
class MultiModalProcessingInfo(BaseProcessingInfo):
|
||||
def _get_audio_processor(self) -> Any:
|
||||
# TODO: drop feature_extractor branch once huggingface/transformers#44394 lands.
|
||||
@@ -77,23 +100,62 @@ class MultiModalProcessingInfo(BaseProcessingInfo):
|
||||
self.get_hf_processor(), ("audio_processor", "feature_extractor")
|
||||
)
|
||||
|
||||
@cached_property
|
||||
def _is_audio_model(self) -> bool:
|
||||
return self._get_audio_processor() is not None
|
||||
|
||||
@cached_property
|
||||
def _is_image_model(self) -> bool:
|
||||
return hasattr(self.get_hf_processor(), "image_processor")
|
||||
|
||||
@cached_property
|
||||
def _is_video_model(self) -> bool:
|
||||
if not hasattr(self.get_hf_processor(), "video_processor"):
|
||||
return False
|
||||
try:
|
||||
Base.check_version("5.18.0", "video inputs")
|
||||
except ImportError as e:
|
||||
logger.info_once("%s, so video inputs are disabled.", e)
|
||||
return False
|
||||
num_frames = self._get_min_video_frames()
|
||||
side = _MIN_DUMMY_SIDE
|
||||
# TODO: Drop the except branch once every video processor can count, see
|
||||
# https://github.com/huggingface/transformers/issues/43329
|
||||
try:
|
||||
mm_tokens = self._get_num_mm_tokens(video_sizes=([num_frames, side, side],))
|
||||
except AttributeError:
|
||||
logger.info_once(
|
||||
"%s cannot count video tokens yet, so the Transformers modeling "
|
||||
"backend serves this model without video inputs. See "
|
||||
"https://github.com/huggingface/transformers/issues/43329",
|
||||
type(self.get_hf_processor()).__name__,
|
||||
)
|
||||
return False
|
||||
return mm_tokens["num_video_tokens"] is not None
|
||||
|
||||
@cached_property
|
||||
def _video_needs_metadata(self) -> bool:
|
||||
video_processor = getattr(self.get_hf_processor(), "video_processor", None)
|
||||
return getattr(video_processor, "do_sample_frames", False)
|
||||
|
||||
def _get_min_video_frames(self) -> int:
|
||||
video_processor = self.get_hf_processor().video_processor
|
||||
# A processor without `temporal_patch_size` groups no frames, so one is enough
|
||||
return getattr(video_processor, "temporal_patch_size", 1)
|
||||
|
||||
def _get_supported_modalities(self) -> list[str]:
|
||||
modalities = []
|
||||
if self._is_audio_model():
|
||||
if self._is_audio_model:
|
||||
modalities.append("audio")
|
||||
if self._is_image_model():
|
||||
if self._is_image_model:
|
||||
modalities.append("image")
|
||||
if self._is_video_model:
|
||||
modalities.append("video")
|
||||
if not modalities:
|
||||
raise ValueError(
|
||||
f"{type(self.get_hf_processor()).__name__} exposes neither an image "
|
||||
"processor nor an audio processor, so the Transformers modeling "
|
||||
"backend cannot serve this model as multi-modal."
|
||||
f"{type(self.get_hf_processor()).__name__} exposes no image, video "
|
||||
"or audio processor, so the Transformers modeling backend cannot "
|
||||
"serve this model as multi-modal."
|
||||
)
|
||||
return modalities
|
||||
|
||||
@@ -104,12 +166,13 @@ class MultiModalProcessingInfo(BaseProcessingInfo):
|
||||
return 16000.0
|
||||
|
||||
def get_data_parser(self) -> MultiModalDataParser:
|
||||
if self._is_audio_model():
|
||||
return MultiModalDataParser(
|
||||
target_sr=self._get_audio_sampling_rate(),
|
||||
expected_hidden_size=self._get_expected_hidden_size(),
|
||||
)
|
||||
return super().get_data_parser()
|
||||
target_sr = self._get_audio_sampling_rate() if self._is_audio_model else None
|
||||
return MultiModalDataParser(
|
||||
target_sr=target_sr,
|
||||
video_needs_metadata=self._video_needs_metadata,
|
||||
expected_hidden_size=self._get_expected_hidden_size(),
|
||||
allow_missing_mm_embeddings=self.allow_missing_mm_embeddings,
|
||||
)
|
||||
|
||||
def get_supported_mm_limits(self):
|
||||
return dict.fromkeys(self._get_supported_modalities())
|
||||
@@ -121,6 +184,8 @@ class MultiModalProcessingInfo(BaseProcessingInfo):
|
||||
max_tokens["audio"] = self.get_max_audio_tokens()
|
||||
if "image" in modalities:
|
||||
max_tokens["image"] = self.get_max_image_tokens()
|
||||
if "video" in modalities:
|
||||
max_tokens["video"] = self.get_max_video_tokens(seq_len, mm_counts)
|
||||
return max_tokens
|
||||
|
||||
def get_max_audio_tokens(self) -> int:
|
||||
@@ -136,30 +201,127 @@ class MultiModalProcessingInfo(BaseProcessingInfo):
|
||||
f"The following attribute names were checked: {names}."
|
||||
)
|
||||
|
||||
def get_max_image_tokens(self) -> int:
|
||||
width, height = self.get_image_size_with_most_features()
|
||||
def _get_num_mm_tokens(self, **sizes: Sequence[Sequence[int]]) -> Any:
|
||||
processor = self.get_hf_processor()
|
||||
multimodal_config = self.ctx.model_config.get_multimodal_config()
|
||||
mm_processor_kwargs = multimodal_config.mm_processor_kwargs or {}
|
||||
mm_tokens = processor._get_num_multimodal_tokens(
|
||||
image_sizes=([height, width],), **mm_processor_kwargs
|
||||
)
|
||||
image_tokens = mm_tokens["num_image_tokens"][0]
|
||||
return image_tokens
|
||||
return processor._get_num_multimodal_tokens(**sizes, **mm_processor_kwargs)
|
||||
|
||||
def get_max_image_tokens(self) -> int:
|
||||
size = self.get_image_size_with_most_features()
|
||||
return self._get_num_image_tokens(size)
|
||||
|
||||
def get_max_video_tokens(self, seq_len: int, mm_counts: Mapping[str, int]) -> int:
|
||||
num_frames = self.get_num_frames_with_most_features(seq_len, mm_counts)
|
||||
return self._get_max_video_tokens_for_frames(num_frames)
|
||||
|
||||
def _get_max_video_tokens_for_frames(self, num_frames: int) -> int:
|
||||
size = self.get_video_size_with_most_features(num_frames)
|
||||
return self._get_num_video_tokens(num_frames, size)
|
||||
|
||||
def _get_num_image_tokens(self, size: ImageSize) -> int:
|
||||
mm_tokens = self._get_num_mm_tokens(image_sizes=([size.height, size.width],))
|
||||
return mm_tokens["num_image_tokens"][0]
|
||||
|
||||
def _get_num_video_tokens(self, num_frames: int, size: ImageSize) -> int:
|
||||
video_sizes = ([num_frames, size.height, size.width],)
|
||||
mm_tokens = self._get_num_mm_tokens(video_sizes=video_sizes)
|
||||
return mm_tokens["num_video_tokens"][0]
|
||||
|
||||
def _get_size_candidates(
|
||||
self, sub_processor: Any, divisors: tuple[int, ...]
|
||||
) -> list[ImageSize]:
|
||||
"""Candidate sizes read off a sub-processor's `size`.
|
||||
|
||||
The keys are one of `VALID_SIZE_DICT_KEYS`, so the bound is either exact
|
||||
or an area budget, which `divisors` splits over its items. `shortest_edge`
|
||||
bounds only the small side, so it yields no candidate.
|
||||
|
||||
`longest_edge` is read as the area budget Qwen and GLM use it for. A
|
||||
processor using it as an edge length (SmolVLM) yields a tiny candidate,
|
||||
which is only chosen if it ties on tokens, so it is still correct.
|
||||
"""
|
||||
size = getattr(sub_processor, "size", None) or {}
|
||||
height = size.get("height", size.get("max_height"))
|
||||
width = size.get("width", size.get("max_width"))
|
||||
if height and width:
|
||||
return [ImageSize(width=width, height=height)]
|
||||
if max_pixels := size.get("max_pixels", size.get("longest_edge")):
|
||||
sides = {math.isqrt(max_pixels // divisor) for divisor in divisors}
|
||||
return [
|
||||
ImageSize(width=side, height=side) for side in sorted(sides) if side
|
||||
]
|
||||
return []
|
||||
|
||||
def _get_size_with_most_tokens(
|
||||
self, candidates: list[ImageSize], count: Callable[[ImageSize], int]
|
||||
) -> ImageSize:
|
||||
"""Smallest size found that yields the most tokens.
|
||||
|
||||
`size` bounds the resized output, not the token count, so the candidates
|
||||
are compared by token count against an arbitrary large fallback, which
|
||||
tiling processors need. The fallback goes last so a candidate wins a tie.
|
||||
Resize rounding can give a large input fewer tokens (Qwen2.5-VL), so the
|
||||
candidates matter for the count, not only for memory.
|
||||
|
||||
The winner is then halved while the count holds, so processors with no
|
||||
usable bound aren't profiled on huge dummy inputs.
|
||||
"""
|
||||
sizes = [*candidates, ImageSize(width=_MAX_DUMMY_SIDE, height=_MAX_DUMMY_SIDE)]
|
||||
num_tokens = [count(size) for size in sizes]
|
||||
most_tokens = max(num_tokens)
|
||||
size = sizes[num_tokens.index(most_tokens)]
|
||||
|
||||
while min(size.width, size.height) // 2 >= _MIN_DUMMY_SIDE:
|
||||
smaller = ImageSize(width=size.width // 2, height=size.height // 2)
|
||||
if count(smaller) < most_tokens:
|
||||
break
|
||||
size = smaller
|
||||
return size
|
||||
|
||||
def get_image_size_with_most_features(self) -> ImageSize:
|
||||
return ImageSize(width=10_000, height=10_000) # arbitrary very large size
|
||||
return self._get_size_with_most_tokens(
|
||||
self._get_size_candidates(self.get_hf_processor().image_processor, (1,)),
|
||||
self._get_num_image_tokens,
|
||||
)
|
||||
|
||||
def get_video_size_with_most_features(self, num_frames: int) -> ImageSize:
|
||||
"""Per-frame size that yields the most video tokens.
|
||||
|
||||
The video processor's `size` budget may cover one frame, every frame or
|
||||
every temporal patch, so all three splits are tried.
|
||||
"""
|
||||
video_processor = self.get_hf_processor().video_processor
|
||||
grid_t = max(num_frames // self._get_min_video_frames(), 1)
|
||||
return self._get_size_with_most_tokens(
|
||||
self._get_size_candidates(video_processor, (num_frames, grid_t, 1)),
|
||||
lambda size: self._get_num_video_tokens(num_frames, size),
|
||||
)
|
||||
|
||||
def get_num_frames_with_most_features(
|
||||
self, seq_len: int, mm_counts: Mapping[str, int]
|
||||
) -> int:
|
||||
max_videos = max(mm_counts.get("video", 0), 1)
|
||||
step = self._get_min_video_frames()
|
||||
num_frames = step
|
||||
while (
|
||||
num_frames + step <= _MAX_FRAMES_PER_VIDEO
|
||||
and self._get_max_video_tokens_for_frames(num_frames + step) * max_videos
|
||||
<= seq_len
|
||||
):
|
||||
num_frames += step
|
||||
return num_frames
|
||||
|
||||
|
||||
class MultiModalDummyInputsBuilder(BaseDummyInputsBuilder[MultiModalProcessingInfo]):
|
||||
def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
|
||||
text = ""
|
||||
if self.info._is_audio_model() and (num_audios := mm_counts.get("audio", 0)):
|
||||
if self.info._is_audio_model and (num_audios := mm_counts.get("audio", 0)):
|
||||
processor = self.info.get_hf_processor()
|
||||
audio_token = getattr(processor, "audio_token", "")
|
||||
# Separated so that adjacent placeholders stay distinguishable
|
||||
text += " ".join([audio_token] * num_audios)
|
||||
if self.info._is_image_model() and (num_images := mm_counts.get("image", 0)):
|
||||
if self.info._is_image_model and (num_images := mm_counts.get("image", 0)):
|
||||
processor = self.info.get_hf_processor()
|
||||
if "gemma3" in processor.__class__.__name__.lower():
|
||||
image_token = processor.boi_token
|
||||
@@ -171,6 +333,10 @@ class MultiModalDummyInputsBuilder(BaseDummyInputsBuilder[MultiModalProcessingIn
|
||||
end_token = getattr(processor, "image_end_token", "")
|
||||
image_token = f"{start_token}{image_token}{end_token}"
|
||||
text += image_token * num_images
|
||||
if self.info._is_video_model and (num_videos := mm_counts.get("video", 0)):
|
||||
processor = self.info.get_hf_processor()
|
||||
video_token = getattr(processor, "video_token", "")
|
||||
text += video_token * num_videos
|
||||
return text
|
||||
|
||||
def get_dummy_mm_data(
|
||||
@@ -180,7 +346,7 @@ class MultiModalDummyInputsBuilder(BaseDummyInputsBuilder[MultiModalProcessingIn
|
||||
mm_options: "MultiModalDummyOptions",
|
||||
) -> MultiModalDataDict:
|
||||
data = MultiModalDataBuiltins()
|
||||
if self.info._is_audio_model() and (num_audios := mm_counts.get("audio", 0)):
|
||||
if self.info._is_audio_model and (num_audios := mm_counts.get("audio", 0)):
|
||||
sampling_rate = self.info._get_audio_sampling_rate()
|
||||
sub = self.info._get_audio_processor()
|
||||
chunk_length = getattr(sub, "chunk_length", None) if sub else None
|
||||
@@ -191,7 +357,7 @@ class MultiModalDummyInputsBuilder(BaseDummyInputsBuilder[MultiModalProcessingIn
|
||||
num_audios=num_audios,
|
||||
overrides=mm_options.get("audio"),
|
||||
)
|
||||
if self.info._is_image_model() and (num_images := mm_counts.get("image", 0)):
|
||||
if self.info._is_image_model and (num_images := mm_counts.get("image", 0)):
|
||||
width, height = self.info.get_image_size_with_most_features()
|
||||
data["image"] = self._get_dummy_images(
|
||||
width=width,
|
||||
@@ -199,6 +365,36 @@ class MultiModalDummyInputsBuilder(BaseDummyInputsBuilder[MultiModalProcessingIn
|
||||
num_images=num_images,
|
||||
overrides=mm_options.get("image"),
|
||||
)
|
||||
if self.info._is_video_model and (num_videos := mm_counts.get("video", 0)):
|
||||
num_frames = self.info.get_num_frames_with_most_features(seq_len, mm_counts)
|
||||
width, height = self.info.get_video_size_with_most_features(num_frames)
|
||||
videos = self._get_dummy_videos(
|
||||
width=width,
|
||||
height=height,
|
||||
num_frames=num_frames,
|
||||
num_videos=num_videos,
|
||||
overrides=mm_options.get("video"),
|
||||
)
|
||||
# Only processors that sample frames need metadata, which has the
|
||||
# dummy frames consumed verbatim
|
||||
if self.info._video_needs_metadata:
|
||||
video_processor = self.info.get_hf_processor().video_processor
|
||||
fps = getattr(video_processor, "fps", None)
|
||||
videos = [
|
||||
(
|
||||
video,
|
||||
{
|
||||
"fps": fps,
|
||||
"duration": len(video) / fps if fps else None,
|
||||
"total_num_frames": len(video),
|
||||
"frames_indices": list(range(len(video))),
|
||||
"video_backend": "opencv",
|
||||
"do_sample_frames": False,
|
||||
},
|
||||
)
|
||||
for video in videos
|
||||
]
|
||||
data["video"] = videos
|
||||
return data
|
||||
|
||||
|
||||
@@ -211,6 +407,40 @@ class MultiModalProcessor(BaseMultiModalProcessor[MultiModalProcessingInfo]):
|
||||
and with it the multi-modal processor cache.
|
||||
"""
|
||||
|
||||
def _get_hf_mm_inputs(
|
||||
self,
|
||||
mm_items: MultiModalDataItems,
|
||||
hf_kwargs: Mapping[str, object],
|
||||
) -> HFMultiModalInputs:
|
||||
"""Pass the video metadata along to the HF processor."""
|
||||
hf_inputs = super()._get_hf_mm_inputs(mm_items, hf_kwargs)
|
||||
hf_data = hf_inputs.hf_data
|
||||
|
||||
if "videos" not in hf_data or not self.info._video_needs_metadata:
|
||||
return hf_inputs
|
||||
|
||||
videos = hf_data["videos"]
|
||||
assert isinstance(videos, Sequence)
|
||||
videos, metadata = zip(*videos)
|
||||
hf_data["videos"] = list(videos)
|
||||
hf_data["video_metadata"] = [
|
||||
{k: v for k, v in item.items() if k != "do_sample_frames"}
|
||||
for item in metadata
|
||||
]
|
||||
do_sample_frames = {item.get("do_sample_frames", False) for item in metadata}
|
||||
if len(do_sample_frames) > 1:
|
||||
raise ValueError(
|
||||
"The videos in a request must agree on `do_sample_frames`, "
|
||||
"since the HF processor takes one value for all of them."
|
||||
)
|
||||
|
||||
return hf_inputs._replace(
|
||||
hf_kwargs={
|
||||
"do_sample_frames": do_sample_frames.pop(),
|
||||
**hf_inputs.hf_kwargs,
|
||||
}
|
||||
)
|
||||
|
||||
def _get_modality_field_names(self, modality: str) -> set[str]:
|
||||
"""Names of the fields the sub-processor for `modality` produces."""
|
||||
# TODO: use else branch only once huggingface/transformers#44394 lands.
|
||||
@@ -287,23 +517,18 @@ class MultiModalProcessor(BaseMultiModalProcessor[MultiModalProcessingInfo]):
|
||||
modalities = [m for m, size in sizes.items() if size is not None]
|
||||
|
||||
# Keys we wrote ourselves, rather than ones a sub-processor produced
|
||||
own_keys = set(_MODALITY_SIZE_KEYS.values()) | {
|
||||
own_keys = {*_MODALITY_SIZE_KEYS.values(), "num_video_tokens"} | {
|
||||
f"{modality}_replacement_{suffix}"
|
||||
for modality in modalities
|
||||
for suffix in ("ids", "sizes")
|
||||
}
|
||||
# Registered by name below, so no sub-processor has to claim them
|
||||
own_keys |= {"image_grid_thw", "video_grid_thw", "second_per_grid_ts"}
|
||||
keys = [key for key in hf_inputs if key not in own_keys]
|
||||
owned = self._partition_keys_by_modality(keys, modalities)
|
||||
|
||||
# Un-padded fields are already one entry per item, so index rather than slice
|
||||
mm_fields: dict[str, MultiModalFieldConfig] = {
|
||||
key: MultiModalFieldConfig.batched(modality)
|
||||
if modality == "audio" or isinstance(hf_inputs[key], list)
|
||||
else MultiModalFieldConfig.flat_from_sizes(
|
||||
modality,
|
||||
sizes[modality],
|
||||
dim=self._get_slice_dim(hf_inputs[key], int(sizes[modality].sum())),
|
||||
)
|
||||
mm_fields = {
|
||||
key: self._get_field_config(modality, key, hf_inputs, sizes[modality])
|
||||
for modality in modalities
|
||||
for key in owned[modality]
|
||||
}
|
||||
@@ -326,13 +551,51 @@ class MultiModalProcessor(BaseMultiModalProcessor[MultiModalProcessingInfo]):
|
||||
mm_fields["image_grid_thw"] = MultiModalFieldConfig.batched(
|
||||
"image", keep_on_cpu=True
|
||||
)
|
||||
# TODO: route to "video" once the video modality is supported
|
||||
if "video" in modalities:
|
||||
mm_fields["video_grid_thw"] = MultiModalFieldConfig.batched(
|
||||
"image", keep_on_cpu=True
|
||||
"video", keep_on_cpu=True
|
||||
)
|
||||
mm_fields["second_per_grid_ts"] = MultiModalFieldConfig.batched(
|
||||
"video", keep_on_cpu=True
|
||||
)
|
||||
|
||||
return mm_fields
|
||||
|
||||
def _get_field_config(
|
||||
self,
|
||||
modality: str,
|
||||
key: str,
|
||||
hf_inputs: "BatchFeature",
|
||||
sizes: torch.Tensor,
|
||||
) -> MultiModalFieldConfig:
|
||||
data = hf_inputs[key]
|
||||
# Un-padded fields are already one entry per item, so index rather than slice
|
||||
if modality == "audio" or isinstance(data, list):
|
||||
return MultiModalFieldConfig.batched(modality)
|
||||
|
||||
total = int(sizes.sum())
|
||||
dim = self._get_slice_dim(data, total)
|
||||
if modality != "video":
|
||||
return MultiModalFieldConfig.flat_from_sizes(modality, sizes, dim=dim)
|
||||
|
||||
rows = data.shape[dim]
|
||||
if rows == len(sizes):
|
||||
return MultiModalFieldConfig.batched("video")
|
||||
if rows == total:
|
||||
return MultiModalFieldConfig.flat_from_sizes("video", sizes, dim=dim)
|
||||
# VideoLLaMA3's compression mask has one row per token the processor counts
|
||||
num_video_tokens = hf_inputs["num_video_tokens"]
|
||||
if rows == int(num_video_tokens.sum()):
|
||||
return MultiModalFieldConfig.flat_from_sizes(
|
||||
"video", num_video_tokens, dim=dim
|
||||
)
|
||||
# NOTE: Any other layout would need a per-model guess, which we'd like to avoid
|
||||
raise ValueError(
|
||||
f"{type(self.info.get_hf_processor()).__name__} returned {rows} row(s) of "
|
||||
f"`{key}` for {len(sizes)} video(s) with {total} patch(es), so the rows "
|
||||
"cannot be attributed to a video."
|
||||
)
|
||||
|
||||
def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str:
|
||||
return self.dummy_inputs.get_dummy_text(mm_counts)
|
||||
|
||||
@@ -569,11 +832,15 @@ class MultiModalProcessor(BaseMultiModalProcessor[MultiModalProcessingInfo]):
|
||||
hf_inputs, hf_data, len(seqs)
|
||||
)
|
||||
elif modality == "audio":
|
||||
counts = []
|
||||
for seq in seqs:
|
||||
ids = torch.tensor(seq)
|
||||
counts.append(int(ids.eq(_get_embed_token_id(ids)).sum()))
|
||||
hf_inputs["num_audio_tokens"] = torch.tensor(counts)
|
||||
hf_inputs["num_audio_tokens"] = _count_embed_tokens(seqs)
|
||||
elif modality == "video":
|
||||
grid = hf_inputs.get("video_grid_thw")
|
||||
hf_inputs["num_video_patches"] = (
|
||||
grid.prod(-1)
|
||||
if grid is not None
|
||||
else torch.ones(len(seqs), dtype=torch.long)
|
||||
)
|
||||
hf_inputs["num_video_tokens"] = _count_embed_tokens(seqs)
|
||||
|
||||
hf_inputs.pop("input_ids")
|
||||
|
||||
@@ -833,25 +1100,47 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE, Base):
|
||||
if pixel_values is None:
|
||||
return None
|
||||
|
||||
num_image_patches = kwargs.pop("num_image_patches")
|
||||
return self._process_vision_input(
|
||||
"image", pixel_values, kwargs.pop("num_image_patches"), **kwargs
|
||||
)
|
||||
|
||||
split_sizes = num_image_patches.flatten().tolist()
|
||||
def _process_video_input(self, **kwargs) -> list[torch.Tensor] | None:
|
||||
pixel_values_videos: torch.Tensor | None = kwargs.pop(
|
||||
"pixel_values_videos", None
|
||||
)
|
||||
if pixel_values_videos is None:
|
||||
return None
|
||||
|
||||
kwargs.pop("second_per_grid_ts", None)
|
||||
return self._process_vision_input(
|
||||
"video", pixel_values_videos, kwargs.pop("num_video_patches"), **kwargs
|
||||
)
|
||||
|
||||
def _process_vision_input(
|
||||
self,
|
||||
modality: str,
|
||||
pixel_values: torch.Tensor,
|
||||
num_patches: torch.Tensor,
|
||||
**kwargs,
|
||||
) -> list[torch.Tensor]:
|
||||
split_sizes = num_patches.flatten().tolist()
|
||||
if isinstance(pixel_values, torch.Tensor):
|
||||
vision_embeddings = self._get_image_features(pixel_values, **kwargs)
|
||||
vision_embeddings = self._get_features(modality, pixel_values, **kwargs)
|
||||
if isinstance(vision_embeddings, torch.Tensor):
|
||||
return self._split_embeddings(vision_embeddings, split_sizes)
|
||||
return list(vision_embeddings)
|
||||
|
||||
# Images the processor left un-padded arrive as a list once their
|
||||
# Items the processor left un-padded arrive as a list once their
|
||||
# shapes differ. Encode them one at a time so that none of them is
|
||||
# padded to match another.
|
||||
embeddings: list[torch.Tensor] = []
|
||||
for index, image in enumerate(pixel_values):
|
||||
features = self._get_image_features(
|
||||
image.unsqueeze(0),
|
||||
for index, item in enumerate(pixel_values):
|
||||
features = self._get_features(
|
||||
modality,
|
||||
item.unsqueeze(0),
|
||||
**self._select_item_kwargs(kwargs, index, len(pixel_values)),
|
||||
)
|
||||
# Encoders which return one entry per image return a single entry
|
||||
# Encoders which return one entry per item return a single entry
|
||||
if not isinstance(features, torch.Tensor):
|
||||
features = torch.cat(list(features))
|
||||
embeddings.extend(self._split_embeddings(features, [split_sizes[index]]))
|
||||
@@ -870,35 +1159,31 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE, Base):
|
||||
for key, value in kwargs.items()
|
||||
}
|
||||
|
||||
def _get_image_features(self, pixel_values: torch.Tensor, **kwargs) -> Any:
|
||||
def _get_features(self, modality: str, pixel_values: torch.Tensor, **kwargs) -> Any:
|
||||
# grid_thw fields are registered keep_on_cpu; restore the on-device
|
||||
# placement that HF get_image_features implementations expect.
|
||||
# placement that HF get_*_features implementations expect.
|
||||
for key, value in kwargs.items():
|
||||
if isinstance(value, torch.Tensor) and value.is_cpu:
|
||||
kwargs[key] = async_tensor_h2d(value, pixel_values.device)
|
||||
|
||||
# The underlying HuggingFace `get_image_features` implementations
|
||||
# The underlying HuggingFace `get_*_features` implementations
|
||||
# contain model-internal syncs (e.g. Idefics3 filters all-zero
|
||||
# padding images via boolean-mask indexing, LlavaOnevision
|
||||
# branches on per-sample batch counts).
|
||||
with gpu_sync_allowed():
|
||||
features = self.model.get_image_features(pixel_values, **kwargs)
|
||||
|
||||
# Transformers `v5`, `self.get_image_features` returns a tuple
|
||||
# containing the features and optionally attentions/hidden_states
|
||||
# After v5 is settled, we can enable qwen3-vl with several outputs
|
||||
# from `self.get_image_features`
|
||||
if isinstance(features, tuple):
|
||||
return features[0]
|
||||
if isinstance(features, ModelOutput):
|
||||
return features.pooler_output
|
||||
return features
|
||||
get_modality_features = getattr(self.model, f"get_{modality}_features")
|
||||
features = get_modality_features(pixel_values, return_dict=True, **kwargs)
|
||||
return features.pooler_output
|
||||
|
||||
def embed_multimodal(self, **kwargs) -> MultiModalEmbeddings:
|
||||
# Each helper detects its own inputs. We are called once per modality, so the
|
||||
# leftovers a helper forwards to the HF model can't belong to the other one.
|
||||
embeddings: list[torch.Tensor] = []
|
||||
for process_input in (self._process_audio_input, self._process_image_input):
|
||||
for process_input in (
|
||||
self._process_audio_input,
|
||||
self._process_image_input,
|
||||
self._process_video_input,
|
||||
):
|
||||
embeddings.extend(process_input(**kwargs) or [])
|
||||
return embeddings
|
||||
|
||||
@@ -917,16 +1202,21 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE, Base):
|
||||
"use_audio_in_video",
|
||||
},
|
||||
)
|
||||
if any(v for k, v in kwargs.items() if k not in {"image_grid_thw"}):
|
||||
if kwargs.get("audio_feature_lengths") or kwargs.get("use_audio_in_video"):
|
||||
raise NotImplementedError(
|
||||
"Transformers modeling backend only supports images."
|
||||
"Transformers modeling backend does not support audio inputs in "
|
||||
"M-RoPE models."
|
||||
)
|
||||
|
||||
image_grid_thw = kwargs.get("image_grid_thw", [])
|
||||
video_grid_thw = kwargs.get("video_grid_thw", [])
|
||||
second_per_grid_ts = kwargs.get("second_per_grid_ts", [])
|
||||
|
||||
image_grid_thw = torch.stack(image_grid_thw) if image_grid_thw else None
|
||||
video_grid_thw = torch.stack(video_grid_thw) if video_grid_thw else None
|
||||
second_per_grid_ts = (
|
||||
torch.stack(second_per_grid_ts) if second_per_grid_ts else None
|
||||
)
|
||||
|
||||
# `get_rope_index` doesn't always accept arbitrary `kwargs`
|
||||
if not hasattr(self, "_get_rope_index_kwarg_names"):
|
||||
@@ -950,6 +1240,7 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE, Base):
|
||||
for name, value in (
|
||||
("image_grid_thw", image_grid_thw),
|
||||
("video_grid_thw", video_grid_thw),
|
||||
("second_per_grid_ts", second_per_grid_ts),
|
||||
)
|
||||
if value is not None or accepts_kwarg(name)
|
||||
}
|
||||
@@ -958,8 +1249,11 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE, Base):
|
||||
for feature in mm_features:
|
||||
position = feature.mm_position
|
||||
offset, length = position.offset, position.length
|
||||
is_embed = position.is_embed
|
||||
if is_embed is None:
|
||||
is_embed = slice(None)
|
||||
mm_token_type_id = _MODALITY_TO_TOKEN_TYPE_ID[feature.modality]
|
||||
mm_token_type_ids[offset : offset + length] = mm_token_type_id
|
||||
mm_token_type_ids[offset : offset + length][is_embed] = mm_token_type_id
|
||||
kwargs["mm_token_type_ids"] = mm_token_type_ids.unsqueeze(0)
|
||||
|
||||
mrope_positions, mrope_position_delta = self.model.get_rope_index(
|
||||
|
||||
@@ -36,7 +36,7 @@ from vllm.model_executor.layers.linear import (
|
||||
)
|
||||
from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
|
||||
from vllm.model_executor.models.utils import maybe_prefix
|
||||
from vllm.transformers_utils.config import is_rope_parameters_nested
|
||||
from vllm.transformers_utils.config import iter_rope_parameters
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from vllm.config import VllmConfig
|
||||
@@ -366,10 +366,4 @@ def can_enable_torch_compile(vllm_config: "VllmConfig") -> bool:
|
||||
"""
|
||||
text_config = vllm_config.model_config.hf_config.get_text_config()
|
||||
# Dynamic rope scaling is not compatible with torch.compile
|
||||
rope_parameters: dict | None = getattr(text_config, "rope_parameters", None) or {}
|
||||
if rope_parameters:
|
||||
# Nest rope_parameters if not nested already to simplify logic
|
||||
if not is_rope_parameters_nested(rope_parameters):
|
||||
rope_parameters = {"": rope_parameters}
|
||||
return all(rp["rope_type"] != "dynamic" for rp in rope_parameters.values())
|
||||
return True
|
||||
return all(rp["rope_type"] != "dynamic" for rp in iter_rope_parameters(text_config))
|
||||
|
||||
@@ -23,7 +23,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.torch_utils import direct_register_custom_op
|
||||
|
||||
QWEN4_EXP_GEMM_PLANS: dict[tuple[int, int], dict[int, SkinnyGemmConfig]] = {
|
||||
QWEN4_EXP_SM103_GEMM_PLANS: dict[tuple[int, int], dict[int, SkinnyGemmConfig]] = {
|
||||
# GDN fused QKVZ projection, TP=4.
|
||||
(4096, 2560): {
|
||||
1: SkinnyGemmConfig(1, 64, 4, k_unroll=4),
|
||||
@@ -43,18 +43,11 @@ QWEN4_EXP_GEMM_PLANS: dict[tuple[int, int], dict[int, SkinnyGemmConfig]] = {
|
||||
8: SkinnyGemmConfig(8, 128, 1, k_unroll=4, vector_width=4),
|
||||
16: SkinnyGemmConfig(16, 128, 1, k_unroll=4, vector_width=4),
|
||||
},
|
||||
# QSA fused QKV/gate projection, TP=4.
|
||||
(3584, 2560): {
|
||||
1: SkinnyGemmConfig(1, 128, 4, k_unroll=4, vector_width=4),
|
||||
2: SkinnyGemmConfig(2, 64, 2, k_unroll=2),
|
||||
4: SkinnyGemmConfig(4, 64, 2, k_unroll=2),
|
||||
},
|
||||
# QSA indexer Q/K projection, replicated in a TP=4 deployment.
|
||||
(640, 2560): {
|
||||
1: SkinnyGemmConfig(1, 128, 1, k_unroll=4, vector_width=4),
|
||||
2: SkinnyGemmConfig(2, 128, 1, k_unroll=4, vector_width=4),
|
||||
4: SkinnyGemmConfig(4, 128, 1, k_unroll=4, vector_width=4),
|
||||
8: SkinnyGemmConfig(8, 128, 1, k_unroll=4, vector_width=4),
|
||||
# QSA fused QKV/gate + replicated indexer Q/K, TP=4 (GB300).
|
||||
(4224, 2560): {
|
||||
1: SkinnyGemmConfig(1, 128, 4, vector_width=4, static_k=2560),
|
||||
2: SkinnyGemmConfig(2, 128, 2, vector_width=4, static_k=2560),
|
||||
4: SkinnyGemmConfig(4, 64, 4, vector_width=4, static_k=2560),
|
||||
},
|
||||
# Shared-expert fused gate/up projection, TP=4.
|
||||
(320, 2560): {
|
||||
@@ -69,13 +62,6 @@ QWEN4_EXP_GEMM_PLANS: dict[tuple[int, int], dict[int, SkinnyGemmConfig]] = {
|
||||
1: SkinnyGemmConfig(1, 64, 4, k_unroll=2),
|
||||
2: SkinnyGemmConfig(2, 32, 4, k_unroll=2),
|
||||
},
|
||||
# HC merged down/injection projection, replicated in a TP=4 deployment.
|
||||
(336, 10240): {
|
||||
1: SkinnyGemmConfig(1, 128, 1, static_k=10240),
|
||||
2: SkinnyGemmConfig(2, 128, 1, static_k=10240),
|
||||
4: SkinnyGemmConfig(4, 128, 2, static_k=10240),
|
||||
8: SkinnyGemmConfig(8, 128, 1, k_unroll=4),
|
||||
},
|
||||
}
|
||||
|
||||
# H200 plans selected by exhaustive CUDA graph replay measurements over
|
||||
@@ -88,6 +74,12 @@ QWEN4_EXP_SM90_GEMM_PLANS: dict[tuple[int, int], dict[int, SkinnyGemmConfig]] =
|
||||
1: SkinnyGemmConfig(1, 128, 2, vector_width=4, static_k=2560),
|
||||
2: SkinnyGemmConfig(2, 64, 4, vector_width=4, static_k=2560),
|
||||
},
|
||||
# QSA fused QKV/gate + replicated indexer Q/K, TP=4.
|
||||
(4224, 2560): {
|
||||
1: SkinnyGemmConfig(1, 128, 2, k_unroll=6, vector_width=2, static_k=2560),
|
||||
2: SkinnyGemmConfig(2, 128, 2, k_unroll=3, vector_width=2, static_k=2560),
|
||||
4: SkinnyGemmConfig(4, 64, 4, k_unroll=2, vector_width=4),
|
||||
},
|
||||
# GDN and QSA output projections, TP=4.
|
||||
(2560, 1536): {
|
||||
1: SkinnyGemmConfig(1, 128, 4, vector_width=2, static_k=1536),
|
||||
@@ -102,18 +94,6 @@ QWEN4_EXP_SM90_GEMM_PLANS: dict[tuple[int, int], dict[int, SkinnyGemmConfig]] =
|
||||
8: SkinnyGemmConfig(8, 128, 1, vector_width=4, static_k=2560),
|
||||
16: SkinnyGemmConfig(16, 128, 1, vector_width=4, static_k=2560),
|
||||
},
|
||||
# QSA fused QKV/gate projection, TP=4.
|
||||
(3584, 2560): {
|
||||
1: SkinnyGemmConfig(1, 128, 4, vector_width=2, static_k=2560),
|
||||
2: SkinnyGemmConfig(2, 64, 4, k_unroll=5),
|
||||
},
|
||||
# QSA indexer Q/K projection, replicated in a TP=4 deployment.
|
||||
(640, 2560): {
|
||||
1: SkinnyGemmConfig(1, 256, 2, vector_width=2, static_k=2560),
|
||||
2: SkinnyGemmConfig(2, 128, 2, vector_width=4, static_k=2560),
|
||||
4: SkinnyGemmConfig(4, 128, 1, vector_width=2, static_k=2560),
|
||||
8: SkinnyGemmConfig(8, 128, 2, vector_width=4, static_k=2560),
|
||||
},
|
||||
# Shared-expert fused gate/up projection, TP=4.
|
||||
(320, 2560): {
|
||||
1: SkinnyGemmConfig(1, 64, 2, vector_width=4, static_k=2560),
|
||||
@@ -127,13 +107,6 @@ QWEN4_EXP_SM90_GEMM_PLANS: dict[tuple[int, int], dict[int, SkinnyGemmConfig]] =
|
||||
1: SkinnyGemmConfig(1, 64, 2, vector_width=2, static_k=2560),
|
||||
2: SkinnyGemmConfig(2, 64, 2, vector_width=2, static_k=2560),
|
||||
},
|
||||
# HC merged down/injection projection, replicated in a TP=4 deployment.
|
||||
(336, 10240): {
|
||||
1: SkinnyGemmConfig(1, 256, 1, k_unroll=5),
|
||||
2: SkinnyGemmConfig(2, 256, 3, static_k=10240),
|
||||
4: SkinnyGemmConfig(4, 256, 3, static_k=10240),
|
||||
8: SkinnyGemmConfig(8, 256, 3, static_k=10240),
|
||||
},
|
||||
# Final HC down projection, replicated in a TP=4 deployment.
|
||||
(320, 10240): {
|
||||
1: SkinnyGemmConfig(1, 256, 1, static_k=10240),
|
||||
@@ -169,30 +142,17 @@ QWEN4_EXP_SM100_GEMM_PLANS: dict[tuple[int, int], dict[int, SkinnyGemmConfig]] =
|
||||
4: SkinnyGemmConfig(4, 256, 1, k_unroll=2),
|
||||
8: SkinnyGemmConfig(8, 256, 1),
|
||||
},
|
||||
# HC merged down/injection projection, replicated in a TP=4 deployment.
|
||||
(336, 10240): {
|
||||
1: SkinnyGemmConfig(1, 256, 1, k_unroll=4, static_k=10240),
|
||||
2: SkinnyGemmConfig(2, 256, 1, k_unroll=4, static_k=10240),
|
||||
8: SkinnyGemmConfig(8, 256, 3, static_k=10240),
|
||||
},
|
||||
# QSA indexer Q/K projection, replicated in a TP=4 deployment.
|
||||
(640, 2560): {
|
||||
1: SkinnyGemmConfig(1, 128, 2, k_unroll=4, vector_width=4, static_k=2560),
|
||||
2: SkinnyGemmConfig(2, 128, 1, k_unroll=4, vector_width=4, static_k=2560),
|
||||
4: SkinnyGemmConfig(4, 128, 1, k_unroll=4, vector_width=4, static_k=2560),
|
||||
8: SkinnyGemmConfig(8, 128, 1, k_unroll=4, vector_width=4, static_k=2560),
|
||||
},
|
||||
# GDN and QSA output projections, TP=4.
|
||||
(2560, 1536): {
|
||||
1: SkinnyGemmConfig(1, 128, 2, k_unroll=2, vector_width=4, static_k=1536),
|
||||
2: SkinnyGemmConfig(2, 64, 2, k_unroll=2, static_k=1536),
|
||||
4: SkinnyGemmConfig(4, 64, 2, k_unroll=2, static_k=1536),
|
||||
},
|
||||
# QSA fused QKV/gate projection, TP=4.
|
||||
(3584, 2560): {
|
||||
1: SkinnyGemmConfig(1, 128, 4, k_unroll=2, vector_width=4, static_k=2560),
|
||||
2: SkinnyGemmConfig(2, 128, 4, k_unroll=2, vector_width=4, static_k=2560),
|
||||
4: SkinnyGemmConfig(4, 64, 2, k_unroll=4, vector_width=4, static_k=2560),
|
||||
# QSA fused QKV/gate + replicated indexer Q/K, TP=4.
|
||||
(4224, 2560): {
|
||||
1: SkinnyGemmConfig(1, 128, 4, vector_width=4, static_k=2560),
|
||||
2: SkinnyGemmConfig(2, 128, 2, vector_width=4, static_k=2560),
|
||||
4: SkinnyGemmConfig(4, 64, 4, vector_width=4, static_k=2560),
|
||||
},
|
||||
# GDN fused QKVZ projection, TP=4.
|
||||
(4096, 2560): {
|
||||
@@ -207,7 +167,7 @@ QWEN4_EXP_SM100_GEMM_PLANS: dict[tuple[int, int], dict[int, SkinnyGemmConfig]] =
|
||||
}
|
||||
|
||||
|
||||
# DGX Spark (GB10) plans, TP=1.
|
||||
# DGX Spark (GB10) plans, TP=1 unless noted.
|
||||
QWEN4_EXP_SM121_GEMM_PLANS: dict[tuple[int, int], dict[int, SkinnyGemmConfig]] = {
|
||||
# Shared-expert gate.
|
||||
(1, 2560): {
|
||||
@@ -225,14 +185,6 @@ QWEN4_EXP_SM121_GEMM_PLANS: dict[tuple[int, int], dict[int, SkinnyGemmConfig]] =
|
||||
8: SkinnyGemmConfig(8, 128, 4, k_unroll=2, vector_width=4, static_k=2560),
|
||||
16: SkinnyGemmConfig(16, 64, 1),
|
||||
},
|
||||
# HC merged down/injection projection.
|
||||
(336, 10240): {
|
||||
1: SkinnyGemmConfig(1, 256, 1, vector_width=4, static_k=10240),
|
||||
2: SkinnyGemmConfig(2, 256, 1, k_unroll=4, vector_width=4, static_k=10240),
|
||||
4: SkinnyGemmConfig(4, 128, 1, k_unroll=4, static_k=10240),
|
||||
8: SkinnyGemmConfig(8, 128, 4, k_unroll=4, static_k=10240),
|
||||
16: SkinnyGemmConfig(16, 128, 4, k_unroll=4, static_k=10240),
|
||||
},
|
||||
# Router.
|
||||
(512, 2560): {
|
||||
1: SkinnyGemmConfig(1, 256, 1, k_unroll=4, vector_width=2, static_k=2560),
|
||||
@@ -241,7 +193,7 @@ QWEN4_EXP_SM121_GEMM_PLANS: dict[tuple[int, int], dict[int, SkinnyGemmConfig]] =
|
||||
8: SkinnyGemmConfig(8, 64, 2, static_k=2560),
|
||||
16: SkinnyGemmConfig(16, 64, 2, k_unroll=2, static_k=2560),
|
||||
},
|
||||
# QSA indexer Q/K projection.
|
||||
# Shared-expert fused gate/up, TP=2; retain after indexer fusion.
|
||||
(640, 2560): {
|
||||
1: SkinnyGemmConfig(1, 256, 1, vector_width=2, static_k=2560),
|
||||
2: SkinnyGemmConfig(2, 256, 1, vector_width=2, static_k=2560),
|
||||
@@ -280,13 +232,13 @@ QWEN4_EXP_SM121_GEMM_PLANS: dict[tuple[int, int], dict[int, SkinnyGemmConfig]] =
|
||||
4: SkinnyGemmConfig(4, 32, 1, k_unroll=2, vector_width=2, static_k=320),
|
||||
8: SkinnyGemmConfig(8, 32, 1, vector_width=2, static_k=320),
|
||||
},
|
||||
# QSA fused QKV/gate projection.
|
||||
(13312, 2560): {
|
||||
# QSA fused QKV/gate + replicated indexer Q/K projection.
|
||||
(13952, 2560): {
|
||||
1: SkinnyGemmConfig(1, 32, 1, vector_width=4, static_k=2560),
|
||||
2: SkinnyGemmConfig(2, 32, 1, vector_width=4, static_k=2560),
|
||||
2: SkinnyGemmConfig(2, 32, 1, k_unroll=2, vector_width=4, static_k=2560),
|
||||
4: SkinnyGemmConfig(4, 32, 1, vector_width=4, static_k=2560),
|
||||
8: SkinnyGemmConfig(8, 64, 1, vector_width=2, static_k=2560),
|
||||
16: SkinnyGemmConfig(16, 32, 1, vector_width=4, static_k=2560),
|
||||
16: SkinnyGemmConfig(16, 32, 1, k_unroll=2, vector_width=4, static_k=2560),
|
||||
},
|
||||
# GDN fused QKVZ projection.
|
||||
(16384, 2560): {
|
||||
@@ -307,32 +259,14 @@ QWEN4_EXP_SM121_GEMM_PLANS: dict[tuple[int, int], dict[int, SkinnyGemmConfig]] =
|
||||
}
|
||||
|
||||
|
||||
def _is_sm121() -> bool:
|
||||
return current_platform.is_device_capability((12, 1))
|
||||
|
||||
|
||||
def _is_sm100() -> bool:
|
||||
return current_platform.is_device_capability((10, 0))
|
||||
|
||||
|
||||
def _is_sm103() -> bool:
|
||||
return current_platform.is_device_capability((10, 3))
|
||||
|
||||
|
||||
def _is_sm90() -> bool:
|
||||
return current_platform.is_device_capability((9, 0))
|
||||
|
||||
|
||||
def _gemm_plans() -> dict[tuple[int, int], dict[int, SkinnyGemmConfig]]:
|
||||
if _is_sm103():
|
||||
return QWEN4_EXP_GEMM_PLANS
|
||||
if _is_sm100():
|
||||
return QWEN4_EXP_SM100_GEMM_PLANS
|
||||
if _is_sm90():
|
||||
return QWEN4_EXP_SM90_GEMM_PLANS
|
||||
if _is_sm121():
|
||||
return QWEN4_EXP_SM121_GEMM_PLANS
|
||||
return {}
|
||||
QWEN4_EXP_GEMM_PLANS_BY_CAPABILITY: dict[
|
||||
tuple[int, int] | None, dict[tuple[int, int], dict[int, SkinnyGemmConfig]]
|
||||
] = {
|
||||
(10, 3): QWEN4_EXP_SM103_GEMM_PLANS,
|
||||
(10, 0): QWEN4_EXP_SM100_GEMM_PLANS,
|
||||
(9, 0): QWEN4_EXP_SM90_GEMM_PLANS,
|
||||
(12, 1): QWEN4_EXP_SM121_GEMM_PLANS,
|
||||
}
|
||||
|
||||
|
||||
def _is_packed_row_major(tensor: torch.Tensor) -> bool:
|
||||
@@ -376,7 +310,10 @@ class Qwen4ExpLowLatencyEmbeddingMethod(
|
||||
|
||||
|
||||
def _qwen4_exp_low_latency_gemm(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
|
||||
plan = _gemm_plans().get((weight.shape[0], weight.shape[1]))
|
||||
plans = QWEN4_EXP_GEMM_PLANS_BY_CAPABILITY.get(
|
||||
current_platform.get_device_capability(), {}
|
||||
)
|
||||
plan = plans.get((weight.shape[0], weight.shape[1]))
|
||||
config = None if plan is None else plan.get(x.shape[0])
|
||||
if (
|
||||
config is not None
|
||||
@@ -404,7 +341,9 @@ def enable_qwen4_exp_low_latency_gemm(
|
||||
module: nn.Module,
|
||||
dtype: torch.dtype,
|
||||
) -> None:
|
||||
plans = _gemm_plans()
|
||||
plans = QWEN4_EXP_GEMM_PLANS_BY_CAPABILITY.get(
|
||||
current_platform.get_device_capability(), {}
|
||||
)
|
||||
if dtype != torch.bfloat16 or not plans:
|
||||
return
|
||||
if not shape_dynamic_skinny_gemm.is_available():
|
||||
|
||||
@@ -612,7 +612,7 @@ class Qwen4ExpForCausalLM(
|
||||
SupportsReplaySSM,
|
||||
):
|
||||
packed_modules_mapping = {
|
||||
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
|
||||
"qkv_proj": ["q_proj", "k_proj", "v_proj", "indexer.index_qk_proj"],
|
||||
"gate_up_proj": ["gate_proj", "up_proj"],
|
||||
"kv_proj": ["key_proj", "value_proj"],
|
||||
"in_proj_qkvz": ["in_proj_qkv", "in_proj_z"],
|
||||
@@ -631,6 +631,12 @@ class Qwen4ExpForCausalLM(
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
|
||||
super().__init__()
|
||||
config: Qwen4ExpTextConfig = vllm_config.model_config.hf_text_config
|
||||
if vllm_config.lora_config is not None:
|
||||
# LoRA does not support the merged QKV/indexer projection, so its
|
||||
# packed mapping must keep the indexer separate.
|
||||
self.packed_modules_mapping = self.packed_modules_mapping | {
|
||||
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
|
||||
}
|
||||
self.vllm_config = vllm_config
|
||||
self.model_config = vllm_config.model_config
|
||||
self.quant_config = vllm_config.quant_config
|
||||
@@ -887,6 +893,7 @@ class Qwen4ExpForConditionalGeneration(
|
||||
requires_raw_input_tokens = True
|
||||
|
||||
packed_modules_mapping = Qwen3_5ForConditionalGeneration.packed_modules_mapping | {
|
||||
"qkv_proj": Qwen4ExpForCausalLM.packed_modules_mapping["qkv_proj"],
|
||||
"kv_proj": ["key_proj", "value_proj"],
|
||||
"input_mix_weight_down_block_inject": [
|
||||
"input_mix_weight_down",
|
||||
@@ -904,6 +911,12 @@ class Qwen4ExpForConditionalGeneration(
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model") -> None:
|
||||
nn.Module.__init__(self)
|
||||
config: Qwen4ExpConfig = vllm_config.model_config.hf_config
|
||||
if vllm_config.lora_config is not None:
|
||||
# LoRA does not support the merged QKV/indexer projection, so its
|
||||
# packed mapping must keep the indexer separate.
|
||||
self.packed_modules_mapping = self.packed_modules_mapping | {
|
||||
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
|
||||
}
|
||||
quant_config = vllm_config.quant_config
|
||||
multimodal_config = vllm_config.model_config.multimodal_config
|
||||
if multimodal_config is None:
|
||||
|
||||
@@ -154,6 +154,7 @@ def _qsa_prepare_kernel(
|
||||
MROPE_H: tl.constexpr,
|
||||
MROPE_W: tl.constexpr,
|
||||
main_qkv_ptr,
|
||||
main_qkv_stride_token,
|
||||
main_q_norm_weight_ptr,
|
||||
main_k_norm_weight_ptr,
|
||||
main_eps,
|
||||
@@ -178,12 +179,12 @@ def _qsa_prepare_kernel(
|
||||
# indexer work. RoPE covers the first D // 2 dims, the width of the
|
||||
# shared cos/sin table.
|
||||
main_pid = pid - num_index_work
|
||||
token = main_pid // (MAIN_HQ + MAIN_HK)
|
||||
token = (main_pid // (MAIN_HQ + MAIN_HK)).to(tl.int64)
|
||||
head = main_pid % (MAIN_HQ + MAIN_HK)
|
||||
HALF: tl.constexpr = D // 4
|
||||
dims = tl.arange(0, MAIN_D)
|
||||
rot = tl.arange(0, HALF)
|
||||
row = main_qkv_ptr + token * (2 * (MAIN_HQ + MAIN_HK) * MAIN_D)
|
||||
row = main_qkv_ptr + token * main_qkv_stride_token
|
||||
is_k = head >= MAIN_HQ
|
||||
kv_head = head - MAIN_HQ
|
||||
if is_k:
|
||||
@@ -241,7 +242,7 @@ def _qsa_prepare_kernel(
|
||||
num_head_tiles: tl.constexpr = tl.cdiv(HQ, TILE_H_Q)
|
||||
token_tile = q_pid // num_head_tiles
|
||||
head_tile = q_pid % num_head_tiles
|
||||
tokens = token_tile * TILE_T_Q + tl.arange(0, TILE_T_Q)
|
||||
tokens = (token_tile * TILE_T_Q + tl.arange(0, TILE_T_Q)).to(tl.int64)
|
||||
heads = head_tile * TILE_H_Q + tl.arange(0, TILE_H_Q)
|
||||
valid_tokens = tokens < num_tokens
|
||||
valid_heads = heads < HQ
|
||||
@@ -342,7 +343,8 @@ def _qsa_prepare_kernel(
|
||||
& (source_tokens < num_tokens)
|
||||
)
|
||||
current_base = (
|
||||
k_ptr + tl.maximum(source_tokens, 0)[:, None] * k_stride_token
|
||||
k_ptr
|
||||
+ tl.maximum(source_tokens, 0).to(tl.int64)[:, None] * k_stride_token
|
||||
)
|
||||
cached_base = (
|
||||
state_cache_ptr
|
||||
@@ -449,7 +451,7 @@ def _qsa_prepare_kernel(
|
||||
# One CTA per request commits only the suffix retained by the ring.
|
||||
num_state_rows = tl.minimum(query_len, STATE_SIZE)
|
||||
for state_offset in tl.range(0, num_state_rows):
|
||||
token = query_end - num_state_rows + state_offset
|
||||
token = (query_end - num_state_rows + state_offset).to(tl.int64)
|
||||
valid_token = (
|
||||
(token >= query_start) & (token < query_end) & (token < num_tokens)
|
||||
)
|
||||
@@ -567,7 +569,8 @@ def qsa_prepare(
|
||||
section = mrope_section if mrope_section is not None else (0, 0, 0)
|
||||
assert len(section) == 3
|
||||
qkv_width = 2 * (num_main_q_heads + num_main_kv_heads) * main_head_dim
|
||||
assert main_qkv.shape == (num_tokens, qkv_width) and main_qkv.is_contiguous()
|
||||
assert main_qkv.shape == (num_tokens, qkv_width)
|
||||
assert main_qkv.stride(-1) == 1
|
||||
assert main_slot_mapping.shape == (num_tokens,)
|
||||
|
||||
if num_tokens <= 4096:
|
||||
@@ -622,6 +625,7 @@ def qsa_prepare(
|
||||
MROPE_H=section[1],
|
||||
MROPE_W=section[2],
|
||||
main_qkv_ptr=main_qkv,
|
||||
main_qkv_stride_token=main_qkv.stride(0),
|
||||
main_q_norm_weight_ptr=main_q_norm_weight,
|
||||
main_k_norm_weight_ptr=main_k_norm_weight,
|
||||
main_eps=main_eps,
|
||||
|
||||
@@ -4,6 +4,7 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable, Iterable
|
||||
from typing import ClassVar, cast
|
||||
|
||||
import torch
|
||||
@@ -20,10 +21,15 @@ from vllm.model_executor.layers.attention.attention import (
|
||||
)
|
||||
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
|
||||
from vllm.model_executor.layers.layernorm import GemmaRMSNorm
|
||||
from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear
|
||||
from vllm.model_executor.layers.linear import (
|
||||
MergedColumnParallelLinear,
|
||||
QKVParallelLinear,
|
||||
RowParallelLinear,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization import QuantizationConfig
|
||||
from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding, get_rope
|
||||
from vllm.model_executor.models.qwen3_next import Qwen3NextAttention
|
||||
from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.platforms.interface import DeviceCapability
|
||||
from vllm.utils.torch_utils import (
|
||||
@@ -53,6 +59,72 @@ from . import model
|
||||
from .indexer_qsa import QSAIndexer
|
||||
|
||||
|
||||
class Qwen4ExpQSAQKVIndexerLinear(MergedColumnParallelLinear):
|
||||
"""Pack sharded Q/gate, K, V and replicated indexer Q/K in one GEMM."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
qkv_proj: QKVParallelLinear,
|
||||
index_size: int,
|
||||
quant_config: QuantizationConfig | None,
|
||||
) -> None:
|
||||
self.num_kv_head_replicas = qkv_proj.num_kv_head_replicas
|
||||
super().__init__(
|
||||
input_size=qkv_proj.input_size,
|
||||
output_sizes=[*qkv_proj.output_sizes, index_size * qkv_proj.tp_size],
|
||||
bias=False,
|
||||
params_dtype=qkv_proj.params_dtype,
|
||||
quant_config=quant_config,
|
||||
prefix=qkv_proj.prefix,
|
||||
)
|
||||
|
||||
def _load_shard(
|
||||
self,
|
||||
loader: Callable[..., None],
|
||||
param: nn.Parameter,
|
||||
loaded_weight: torch.Tensor,
|
||||
loaded_shard_id: int | tuple[int, ...] | None,
|
||||
) -> None:
|
||||
tp_rank = self.tp_rank
|
||||
param_tp_rank = getattr(param, "tp_rank", None)
|
||||
if loaded_shard_id == 3:
|
||||
shard_rank = 0
|
||||
elif loaded_shard_id in (1, 2):
|
||||
shard_rank = tp_rank // self.num_kv_head_replicas
|
||||
else:
|
||||
shard_rank = tp_rank
|
||||
self.tp_rank = shard_rank
|
||||
if param_tp_rank is not None:
|
||||
param.tp_rank = shard_rank
|
||||
try:
|
||||
loader(param, loaded_weight, loaded_shard_id)
|
||||
finally:
|
||||
self.tp_rank = tp_rank
|
||||
if param_tp_rank is not None:
|
||||
param.tp_rank = param_tp_rank
|
||||
|
||||
def weight_loader(self, param, loaded_weight, loaded_shard_id=None) -> None:
|
||||
self._load_shard(super().weight_loader, param, loaded_weight, loaded_shard_id)
|
||||
|
||||
def weight_loader_v2(self, param, loaded_weight, loaded_shard_id=None) -> None:
|
||||
self._load_shard(
|
||||
super().weight_loader_v2, param, loaded_weight, loaded_shard_id
|
||||
)
|
||||
|
||||
def load_weights(
|
||||
self, weights: Iterable[tuple[str, torch.Tensor]]
|
||||
) -> Iterable[str]:
|
||||
def remap_shards():
|
||||
for name, weight in weights:
|
||||
shard_id = getattr(weight, "shard_id", None)
|
||||
if isinstance(shard_id, str):
|
||||
weight = weight.detach()
|
||||
weight.shard_id = {"q": 0, "k": 1, "v": 2}[shard_id]
|
||||
yield name, weight
|
||||
|
||||
return super().load_weights(remap_shards())
|
||||
|
||||
|
||||
class Qwen4ExpQSAMetadataBuilder(FlashAttentionMetadataBuilder):
|
||||
"""Flash metadata supporting uniform decode and target-verify graphs."""
|
||||
|
||||
@@ -411,6 +483,16 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase):
|
||||
self.use_fused_qsa_prepare = (
|
||||
self.use_fused_qk_norm_rope_gate and self.indexer.use_fused_pre_indexer
|
||||
)
|
||||
self.fuse_indexer_projection = vllm_config.lora_config is None
|
||||
if self.fuse_indexer_projection:
|
||||
self.index_qk_size = self.indexer.index_qk_proj.output_size
|
||||
self.qkv_proj = Qwen4ExpQSAQKVIndexerLinear(
|
||||
self.qkv_proj,
|
||||
self.index_qk_size,
|
||||
model.without_modelopt_fp4(quant_config),
|
||||
)
|
||||
del self.indexer.index_qk_proj
|
||||
|
||||
max_tokens = vllm_config.scheduler_config.max_num_batched_tokens
|
||||
# PACKED selection buffer: the trailing column holds each row's
|
||||
# valid-entry count (written by the expand kernel) — never a token
|
||||
@@ -432,6 +514,16 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase):
|
||||
raise ValueError(f"Duplicate layer name: {self.layer_name}")
|
||||
static_context[self.layer_name] = self
|
||||
|
||||
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
||||
mapper = None
|
||||
if self.fuse_indexer_projection:
|
||||
mapper = WeightsMapper(
|
||||
orig_to_new_stacked={
|
||||
"indexer.index_qk_proj.": ("qkv_proj.", 3),
|
||||
}
|
||||
)
|
||||
return AutoWeightsLoader(self).load_weights(weights, mapper=mapper)
|
||||
|
||||
def get_attn_backend(self) -> type[AttentionBackend]:
|
||||
return self.attn_backend
|
||||
|
||||
@@ -518,6 +610,12 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase):
|
||||
hidden_states: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
qkv, _ = self.qkv_proj(hidden_states)
|
||||
if self.fuse_indexer_projection:
|
||||
qkv, projected_qk = qkv.split(
|
||||
[2 * self.q_size + 2 * self.kv_size, self.index_qk_size], dim=-1
|
||||
)
|
||||
else:
|
||||
projected_qk, _ = self.indexer.index_qk_proj(hidden_states)
|
||||
num_tokens = hidden_states.shape[0]
|
||||
if not self.use_fused_qsa_prepare:
|
||||
q, k, v, gate = self._project_qkv_gate(qkv, positions)
|
||||
@@ -529,8 +627,6 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase):
|
||||
# Norm/RoPE/gate and the K/V cache write happen inside _run_qsa.
|
||||
query = key = value = gate = None
|
||||
attn_output = qkv.new_empty(num_tokens, self.num_heads, self.head_dim)
|
||||
# Keep the index projection outside the eager break.
|
||||
projected_qk, _ = self.indexer.index_qk_proj(hidden_states)
|
||||
self._run_qsa(
|
||||
projected_qk,
|
||||
positions,
|
||||
|
||||
@@ -683,15 +683,15 @@ def patch_rope_parameters(config: PreTrainedConfig) -> None:
|
||||
config.validate_rope()
|
||||
|
||||
|
||||
def _iter_rope_parameters(config: PreTrainedConfig) -> Iterator[dict[str, Any]]:
|
||||
def iter_rope_parameters(config: PreTrainedConfig) -> Iterator[dict[str, Any]]:
|
||||
"""Yield a config's rope parameters, one dict per layer type if nested."""
|
||||
rope_parameters = getattr(config, "rope_parameters", None)
|
||||
if not isinstance(rope_parameters, dict):
|
||||
return
|
||||
|
||||
if is_rope_parameters_nested(rope_parameters):
|
||||
yield from (p for p in rope_parameters.values() if isinstance(p, dict))
|
||||
else:
|
||||
yield from (p for p in rope_parameters.values() if isinstance(p, dict) and p)
|
||||
elif rope_parameters:
|
||||
yield rope_parameters
|
||||
|
||||
|
||||
@@ -705,7 +705,7 @@ def _mrope_section(config: PreTrainedConfig) -> Sequence[int] | None:
|
||||
|
||||
names = ("mrope_section", "xdrope_section")
|
||||
|
||||
for params in _iter_rope_parameters(config):
|
||||
for params in iter_rope_parameters(config):
|
||||
for i, name in enumerate(names):
|
||||
section = params.get(name)
|
||||
if isinstance(section, (list, tuple)):
|
||||
|
||||
@@ -18,6 +18,7 @@ import torch
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.triton_utils import tl, triton
|
||||
from vllm.utils.math_utils import largest_power_of_2_divisor
|
||||
from vllm.v1.attention.ops.triton_attention_helpers import (
|
||||
apply_alibi_to_score,
|
||||
apply_softcap,
|
||||
@@ -908,14 +909,15 @@ def unified_attention_int4(
|
||||
) -> None:
|
||||
"""Paged attention over the INT4 packed cache, writing into *out*.
|
||||
|
||||
The forward RHT has norm ``sqrt(head_size)``, so ``softmax_scale`` is
|
||||
divided by ``head_size`` and the inverse RHT divides the output by
|
||||
``head_size`` as well.
|
||||
The forward RHT has norm ``sqrt(b)`` for ``b`` the largest power-of-two
|
||||
divisor of ``head_size`` (``b == head_size`` for power-of-two sizes), so
|
||||
``softmax_scale`` is divided by ``b`` and the inverse RHT divides the output
|
||||
by ``b`` as well.
|
||||
"""
|
||||
q_orig_dtype = q.dtype
|
||||
q = single_rht(q.float()).to(q_orig_dtype)
|
||||
head_size = q.shape[2]
|
||||
softmax_scale = softmax_scale / head_size
|
||||
rht_norm = largest_power_of_2_divisor(q.shape[2])
|
||||
softmax_scale = softmax_scale / rht_norm
|
||||
|
||||
_launch_packed_attn(
|
||||
q=q,
|
||||
@@ -945,7 +947,7 @@ def unified_attention_int4(
|
||||
packing_factor=_INT4_PACKING_FACTOR,
|
||||
)
|
||||
|
||||
out_f = single_rht(out.float(), inverse=True) / head_size
|
||||
out_f = single_rht(out.float(), inverse=True) / rht_norm
|
||||
out.copy_(out_f.to(q_orig_dtype))
|
||||
|
||||
|
||||
@@ -1148,6 +1150,16 @@ def single_rht(x: torch.Tensor, inverse: bool = False) -> torch.Tensor:
|
||||
before asymmetric quantization.
|
||||
"""
|
||||
d = x.shape[-1]
|
||||
block = largest_power_of_2_divisor(d)
|
||||
if block < d:
|
||||
# H is only defined for power-of-two rows, so apply it
|
||||
# block-diagonally: H_b inside each of the d / b blocks, for b the
|
||||
# largest power-of-two divisor of d. That keeps H @ H.T = b * I
|
||||
# without padding the row, hence the norm is sqrt(b), not sqrt(d).
|
||||
lead = x.shape[:-1]
|
||||
return single_rht(x.reshape(*lead, d // block, block), inverse).reshape(
|
||||
*lead, d
|
||||
)
|
||||
d1 = _get_rht_signs(d, 0, x.device)
|
||||
if inverse:
|
||||
return fast_hadamard_transform(x) * d1
|
||||
|
||||
@@ -738,14 +738,19 @@ class SimpleCPUOffloadScheduler:
|
||||
if scheduled_for_req:
|
||||
req_ids.append(req_id)
|
||||
|
||||
for req_id, new_block_id_groups, preempted in yield_req_data(scheduler_output):
|
||||
# Preemption frees the request's blocks. It resumes with a full new
|
||||
# block table (MRV1 resumed-cached, MRV2 NewRequestData).
|
||||
for req_id in preempted_req_ids:
|
||||
state = self._reqs_to_store.get(req_id)
|
||||
if state is not None:
|
||||
state.block_ids = tuple([] for _ in range(num_groups))
|
||||
state.num_stored_blocks = [0] * num_groups
|
||||
|
||||
for req_id, new_block_id_groups, _ in yield_req_data(scheduler_output):
|
||||
state = self._reqs_to_store.get(req_id)
|
||||
if state is None or state.finished:
|
||||
continue
|
||||
|
||||
if preempted:
|
||||
state.block_ids = tuple([] for _ in range(num_groups))
|
||||
state.num_stored_blocks = [0] * num_groups
|
||||
if new_block_id_groups:
|
||||
for g in range(min(num_groups, len(new_block_id_groups))):
|
||||
if new_block_id_groups[g] is not None:
|
||||
|
||||
@@ -615,7 +615,10 @@ class Worker(WorkerBase):
|
||||
if kv_cache_memory_bytes := self.cache_config.kv_cache_memory_bytes:
|
||||
# still need a profile run which compiles the model for
|
||||
# max_num_batched_tokens
|
||||
self.model_runner.profile_run(randomize_inputs=self.randomize_dummy_inputs)
|
||||
with set_current_vllm_config(self.vllm_config):
|
||||
self.model_runner.profile_run(
|
||||
randomize_inputs=self.randomize_dummy_inputs
|
||||
)
|
||||
|
||||
msg = (
|
||||
f"Initial free memory {format_gib(self.init_snapshot.free_memory)} "
|
||||
@@ -651,6 +654,7 @@ class Worker(WorkerBase):
|
||||
# limit are never split, so they stay releasable. Exits before
|
||||
# memory_profiling measures, restoring the original limit.
|
||||
self._scoped_allocator_max_split(max_split_size_mb=20),
|
||||
set_current_vllm_config(self.vllm_config),
|
||||
):
|
||||
self.model_runner.profile_run(randomize_inputs=self.randomize_dummy_inputs)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user