forked from Karylab-cklius/vllm
Compare commits
18
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
e0dd86f6c3 | ||
|
|
01f19dee7f | ||
|
|
2c06cf3486 | ||
|
|
e6f710a87f | ||
|
|
c245d35ff4 | ||
|
|
f8ac0c7cf0 | ||
|
|
ebf862c351 | ||
|
|
8d8062d0a7 | ||
|
|
985961345a | ||
|
|
706a04d34b | ||
|
|
22631f80a0 | ||
|
|
2cc008e7b4 | ||
|
|
5d5c776444 | ||
|
|
592ae6805c | ||
|
|
7b1bc0a3eb | ||
|
|
c0879d9483 | ||
|
|
f5f9878514 | ||
|
|
2ce95a761b |
+14
-14
@@ -388,10 +388,10 @@ steps:
|
||||
- python3 basic/offline_inference/embed.py
|
||||
- python3 basic/offline_inference/score.py
|
||||
# Multi-modal models
|
||||
- python3 offline_inference/audio_language.py --seed 0
|
||||
- python3 offline_inference/vision_language.py --seed 0
|
||||
- python3 offline_inference/vision_language_multi_image.py --seed 0
|
||||
- python3 offline_inference/encoder_decoder_multimodal.py --model-type whisper --seed 0
|
||||
- python3 generate/multimodal/audio_language_offline.py --seed 0
|
||||
- python3 generate/multimodal/vision_language_offline.py --seed 0
|
||||
- python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0
|
||||
- python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0
|
||||
# Pooling models
|
||||
- python3 pooling/embed/vision_embedding_offline.py --seed 0
|
||||
# Features demo
|
||||
@@ -1647,10 +1647,10 @@ steps:
|
||||
- python3 basic/offline_inference/embed.py
|
||||
- python3 basic/offline_inference/score.py
|
||||
# Multi-modal models
|
||||
- python3 offline_inference/audio_language.py --seed 0
|
||||
- python3 offline_inference/vision_language.py --seed 0
|
||||
- python3 offline_inference/vision_language_multi_image.py --seed 0
|
||||
- python3 offline_inference/encoder_decoder_multimodal.py --model-type whisper --seed 0
|
||||
- python3 generate/multimodal/audio_language_offline.py --seed 0
|
||||
- python3 generate/multimodal/vision_language_offline.py --seed 0
|
||||
- python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0
|
||||
- python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0
|
||||
# Pooling models
|
||||
- python3 pooling/embed/vision_embedding_offline.py --seed 0
|
||||
# Features demo
|
||||
@@ -1951,8 +1951,8 @@ steps:
|
||||
- pytest -v -s tests/models/multimodal/processing/
|
||||
- pytest -v -s tests/models/multimodal/test_mapping.py
|
||||
- python3 examples/basic/offline_inference/chat.py
|
||||
- python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper
|
||||
- python3 examples/generate/multimodal/vision_language_offline.py --model-type qwen2_5_vl
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/generate/multimodal/audio_language_offline.py --model-type whisper
|
||||
|
||||
#------------------------------------------------------- mi300 · quantization --------------------------------------------------------#
|
||||
|
||||
@@ -2930,10 +2930,10 @@ steps:
|
||||
- python3 basic/offline_inference/embed.py
|
||||
- python3 basic/offline_inference/score.py
|
||||
# Multi-modal models
|
||||
- python3 offline_inference/audio_language.py --seed 0
|
||||
- python3 offline_inference/vision_language.py --seed 0
|
||||
- python3 offline_inference/vision_language_multi_image.py --seed 0
|
||||
- python3 offline_inference/encoder_decoder_multimodal.py --model-type whisper --seed 0
|
||||
- python3 generate/multimodal/audio_language_offline.py --seed 0
|
||||
- python3 generate/multimodal/vision_language_offline.py --seed 0
|
||||
- python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0
|
||||
- python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0
|
||||
# Pooling models
|
||||
- python3 pooling/embed/vision_embedding_offline.py --seed 0
|
||||
# Features demo
|
||||
|
||||
@@ -95,11 +95,13 @@ steps:
|
||||
- tests/kernels/moe/test_deepgemm.py
|
||||
- tests/kernels/moe/test_batched_deepgemm.py
|
||||
- tests/kernels/attention/test_deepgemm_attention.py
|
||||
- tests/quantization/test_cutlass_w4a16.py
|
||||
commands:
|
||||
- pytest -v -s kernels/quantization/test_block_fp8.py
|
||||
- pytest -v -s kernels/moe/test_deepgemm.py
|
||||
- pytest -v -s kernels/moe/test_batched_deepgemm.py
|
||||
- pytest -v -s kernels/attention/test_deepgemm_attention.py
|
||||
- pytest -v -s quantization/test_cutlass_w4a16.py
|
||||
|
||||
- label: Kernels (B200)
|
||||
timeout_in_minutes: 30
|
||||
|
||||
@@ -113,10 +113,10 @@ steps:
|
||||
- python3 basic/offline_inference/embed.py
|
||||
- python3 basic/offline_inference/score.py
|
||||
# for multi-modal models
|
||||
- python3 offline_inference/audio_language.py --seed 0
|
||||
- python3 offline_inference/vision_language.py --seed 0
|
||||
- python3 offline_inference/vision_language_multi_image.py --seed 0
|
||||
- python3 offline_inference/encoder_decoder_multimodal.py --model-type whisper --seed 0
|
||||
- python3 generate/multimodal/audio_language_offline.py --seed 0
|
||||
- python3 generate/multimodal/vision_language_offline.py --seed 0
|
||||
- python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0
|
||||
- python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0
|
||||
# for pooling models
|
||||
- python3 pooling/embed/vision_embedding_offline.py --seed 0
|
||||
# for features demo
|
||||
|
||||
@@ -44,10 +44,10 @@ steps:
|
||||
#- python3 basic/offline_inference/generate.py --model meta-llama/Llama-2-13b-chat-hf --cpu-offload-gb 10 # TODO
|
||||
#- python3 basic/offline_inference/embed.py # TODO
|
||||
# for multi-modal models
|
||||
- python3 offline_inference/audio_language.py --seed 0
|
||||
- python3 offline_inference/vision_language.py --seed 0
|
||||
- python3 offline_inference/vision_language_multi_image.py --seed 0
|
||||
- python3 offline_inference/encoder_decoder_multimodal.py --model-type whisper --seed 0
|
||||
- python3 generate/multimodal/audio_language_offline.py --seed 0
|
||||
- python3 generate/multimodal/vision_language_offline.py --seed 0
|
||||
- python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0
|
||||
- python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0
|
||||
# for pooling models
|
||||
- python3 pooling/embed/vision_embedding_offline.py --seed 0
|
||||
# for features demo
|
||||
|
||||
@@ -69,9 +69,9 @@ steps:
|
||||
- pytest -v -s tests/models/multimodal/processing/
|
||||
- pytest -v -s tests/models/multimodal/test_mapping.py
|
||||
- python3 examples/basic/offline_inference/chat.py
|
||||
- python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl
|
||||
- python3 examples/generate/multimodal/vision_language_offline.py --model-type qwen2_5_vl
|
||||
# Whisper needs spawn method to avoid deadlock
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/generate/multimodal/audio_language_offline.py --model-type whisper
|
||||
|
||||
- label: Transformers Backward Compatibility Models Test
|
||||
working_dir: "/vllm-workspace/"
|
||||
@@ -83,7 +83,7 @@ steps:
|
||||
- pytest -v -s tests/models/test_transformers.py
|
||||
- pytest -v -s tests/models/multimodal/processing/
|
||||
- pytest -v -s tests/models/multimodal/test_mapping.py
|
||||
- python3 examples/offline_inference/basic/chat.py
|
||||
- python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl
|
||||
- python3 examples/basic/offline_inference/chat.py
|
||||
- python3 examples/generate/multimodal/vision_language_offline.py --model-type qwen2_5_vl
|
||||
# Whisper needs spawn method to avoid deadlock
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/generate/multimodal/audio_language_offline.py --model-type whisper
|
||||
|
||||
+1
-5
@@ -389,11 +389,7 @@ pull_request_rules:
|
||||
- files~=^tests/entrypoints/anthropic/.*tool.*
|
||||
- files~=^vllm/tool_parsers/
|
||||
- files=docs/features/tool_calling.md
|
||||
- files~=^examples/tool_chat_*
|
||||
- files=examples/offline_inference/chat_with_tools.py
|
||||
- files=examples/online_serving/openai_chat_completion_client_with_tools_required.py
|
||||
- files=examples/online_serving/openai_chat_completion_tool_calls_with_reasoning.py
|
||||
- files=examples/online_serving/openai_chat_completion_client_with_tools.py
|
||||
- files~=^examples/tool_calling/
|
||||
actions:
|
||||
label:
|
||||
add:
|
||||
|
||||
@@ -0,0 +1,324 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
# Benchmarks FP8 vs BF16 ViT attention via FlashInfer cuDNN backend.
|
||||
#
|
||||
# == Usage Examples ==
|
||||
#
|
||||
# Benchmark mode (default, FlashInfer CUDAGraph Bench)
|
||||
# python3 benchmark_vit_fp8_attn.py
|
||||
#
|
||||
# Profile mode (PyTorch profiler, saves TensorBoard traces):
|
||||
# python3 benchmark_vit_fp8_attn.py --profile
|
||||
# python3 benchmark_vit_fp8_attn.py --profile --profile-output-dir ./profile_traces
|
||||
#
|
||||
# Custom seq_lens:
|
||||
# python3 benchmark_vit_fp8_attn.py --seq-lens 4096 8192 16384
|
||||
|
||||
from functools import partial
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch.profiler import ProfilerActivity, profile, record_function
|
||||
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
# Qwen3-VL defaults
|
||||
NUM_HEADS = 16
|
||||
HEAD_DIM = 72
|
||||
DEFAULT_SEQ_LENS = [2304, 4096, 8192, 16384]
|
||||
|
||||
|
||||
def _setup_fp8_attention(num_heads: int, head_dim: int) -> tuple:
|
||||
"""Create FP8 and BF16 attention modules + workspace."""
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
from vllm.config import VllmConfig, set_current_vllm_config
|
||||
from vllm.config.multimodal import MultiModalConfig
|
||||
from vllm.model_executor.layers.attention.mm_encoder_attention import (
|
||||
MMEncoderAttention,
|
||||
_get_flashinfer_workspace_buffer,
|
||||
)
|
||||
from vllm.v1.attention.backends.registry import AttentionBackendEnum
|
||||
|
||||
old_dtype = torch.get_default_dtype()
|
||||
torch.set_default_dtype(torch.bfloat16)
|
||||
|
||||
backend_patch = patch(
|
||||
"vllm.model_executor.layers.attention.mm_encoder_attention"
|
||||
".get_vit_attn_backend",
|
||||
return_value=AttentionBackendEnum.FLASHINFER,
|
||||
)
|
||||
|
||||
# FP8 attention
|
||||
mm_config_fp8 = MultiModalConfig(mm_encoder_attn_dtype="fp8")
|
||||
vllm_config_fp8 = VllmConfig()
|
||||
vllm_config_fp8.model_config = SimpleNamespace(multimodal_config=mm_config_fp8)
|
||||
with set_current_vllm_config(vllm_config_fp8), backend_patch:
|
||||
attn_fp8 = MMEncoderAttention(
|
||||
num_heads=num_heads,
|
||||
head_size=head_dim,
|
||||
prefix="visual.blocks.0.attn",
|
||||
).to("cuda")
|
||||
|
||||
# BF16 attention (no FP8)
|
||||
with set_current_vllm_config(VllmConfig()), backend_patch:
|
||||
attn_bf16 = MMEncoderAttention(
|
||||
num_heads=num_heads,
|
||||
head_size=head_dim,
|
||||
prefix="visual.blocks.0.attn",
|
||||
).to("cuda")
|
||||
|
||||
torch.set_default_dtype(old_dtype)
|
||||
|
||||
workspace = _get_flashinfer_workspace_buffer()
|
||||
return attn_fp8, attn_bf16, workspace
|
||||
|
||||
|
||||
def _build_meta(
|
||||
seq_len: int,
|
||||
num_heads: int,
|
||||
head_dim: int,
|
||||
fp8: bool,
|
||||
):
|
||||
"""Build cu_seqlens, max_seqlen, sequence_lengths."""
|
||||
from vllm.model_executor.layers.attention.mm_encoder_attention import (
|
||||
MMEncoderAttention,
|
||||
)
|
||||
from vllm.utils.math_utils import round_up
|
||||
from vllm.v1.attention.backends.registry import AttentionBackendEnum
|
||||
|
||||
cu_np = np.array([0, seq_len], dtype=np.int32)
|
||||
fp8_padded = num_heads * round_up(head_dim, 16) if fp8 else None
|
||||
|
||||
seq_lengths = MMEncoderAttention.maybe_compute_seq_lens(
|
||||
AttentionBackendEnum.FLASHINFER, cu_np, torch.device("cuda")
|
||||
)
|
||||
max_seqlen = torch.tensor(
|
||||
MMEncoderAttention.compute_max_seqlen(AttentionBackendEnum.FLASHINFER, cu_np),
|
||||
dtype=torch.int32,
|
||||
)
|
||||
cu_seqlens = MMEncoderAttention.maybe_recompute_cu_seqlens(
|
||||
AttentionBackendEnum.FLASHINFER,
|
||||
cu_np,
|
||||
num_heads * head_dim,
|
||||
1,
|
||||
torch.device("cuda"),
|
||||
fp8_padded_hidden_size=fp8_padded,
|
||||
)
|
||||
return cu_seqlens, max_seqlen, seq_lengths
|
||||
|
||||
|
||||
def run_benchmark(
|
||||
seq_lens: list[int],
|
||||
num_heads: int,
|
||||
head_dim: int,
|
||||
method: str,
|
||||
):
|
||||
"""Benchmark FP8 vs BF16 attention across seq_lens.
|
||||
|
||||
Uses FlashInfer GPU-level timing to measure pure kernel time,
|
||||
excluding CPU launch overhead.
|
||||
"""
|
||||
if method == "cupti":
|
||||
from flashinfer.testing import bench_gpu_time_with_cupti as bench_fn
|
||||
|
||||
bench_fn = partial(bench_fn, use_cuda_graph=True, cold_l2_cache=False)
|
||||
elif method == "cudagraph":
|
||||
from flashinfer.testing import (
|
||||
bench_gpu_time_with_cudagraph as bench_fn,
|
||||
)
|
||||
|
||||
bench_fn = partial(bench_fn, cold_l2_cache=False)
|
||||
else:
|
||||
raise ValueError(f"Invalid method: {method}")
|
||||
|
||||
attn_fp8, attn_bf16, workspace = _setup_fp8_attention(num_heads, head_dim)
|
||||
|
||||
print(f"Timing method: {method}")
|
||||
print(f"{'seq_len':>8} {'BF16 (us)':>12} {'FP8 (us)':>12} {'Speedup':>10}")
|
||||
print("-" * 46)
|
||||
|
||||
for seq_len in seq_lens:
|
||||
torch.manual_seed(42)
|
||||
|
||||
q = torch.randn(
|
||||
seq_len,
|
||||
num_heads,
|
||||
head_dim,
|
||||
device="cuda",
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
k = torch.randn_like(q)
|
||||
v = torch.randn_like(q)
|
||||
|
||||
cu_fp8, max_s, seq_l = _build_meta(seq_len, num_heads, head_dim, fp8=True)
|
||||
# we can reuse cu_fp8 for cu_bf16 since q, k, and v are contiguous
|
||||
cu_bf16 = cu_fp8.clone()
|
||||
|
||||
def bf16_fn(q=q, k=k, v=v, cu=cu_bf16, ms=max_s, sl=seq_l):
|
||||
attn_bf16._forward_flashinfer(q, k, v, cu, ms, sl)
|
||||
|
||||
def fp8_fn(q=q, k=k, v=v, cu=cu_fp8, ms=max_s, sl=seq_l):
|
||||
attn_fp8._forward_flashinfer(q, k, v, cu, ms, sl)
|
||||
|
||||
# bench_fn returns List[float] of per-iteration times in ms
|
||||
bf16_times = bench_fn(bf16_fn)
|
||||
fp8_times = bench_fn(fp8_fn)
|
||||
|
||||
bf16_us = np.median(bf16_times) * 1e3 # ms -> us
|
||||
fp8_us = np.median(fp8_times) * 1e3
|
||||
speedup = bf16_us / fp8_us if fp8_us > 0 else float("inf")
|
||||
|
||||
print(f"{seq_len:>8} {bf16_us:>12.1f} {fp8_us:>12.1f} {speedup:>9.2f}x")
|
||||
|
||||
|
||||
def _make_trace_handler(output_dir: str, worker_name: str, label: str):
|
||||
"""Create a trace handler that saves to TensorBoard and prints summary."""
|
||||
|
||||
def handler(prof):
|
||||
torch.profiler.tensorboard_trace_handler(output_dir, worker_name)(prof)
|
||||
print(f"\n{'=' * 80}")
|
||||
print(label)
|
||||
print(f"{'=' * 80}")
|
||||
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=20))
|
||||
|
||||
return handler
|
||||
|
||||
|
||||
def run_profile(
|
||||
seq_len: int,
|
||||
num_heads: int,
|
||||
head_dim: int,
|
||||
warmup: int,
|
||||
output_dir: str,
|
||||
):
|
||||
"""Profile FP8 vs BF16 attention with PyTorch profiler."""
|
||||
attn_fp8, attn_bf16, workspace = _setup_fp8_attention(num_heads, head_dim)
|
||||
|
||||
torch.manual_seed(42)
|
||||
q = torch.randn(
|
||||
seq_len,
|
||||
num_heads,
|
||||
head_dim,
|
||||
device="cuda",
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
k = torch.randn_like(q)
|
||||
v = torch.randn_like(q)
|
||||
|
||||
cu_fp8, max_s, seq_l = _build_meta(seq_len, num_heads, head_dim, fp8=True)
|
||||
# we can reuse cu_fp8 for cu_bf16 since q, k, and v are contiguous
|
||||
cu_bf16 = cu_fp8.clone()
|
||||
|
||||
sched = torch.profiler.schedule(wait=0, warmup=warmup, active=1)
|
||||
|
||||
# Profile BF16 (warmup handled by profiler schedule)
|
||||
with profile(
|
||||
activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
|
||||
schedule=sched,
|
||||
on_trace_ready=_make_trace_handler(
|
||||
output_dir,
|
||||
f"bf16_h{head_dim}_s{seq_len}",
|
||||
f"BF16 Attention (seq_len={seq_len}, heads={num_heads}, "
|
||||
f"head_dim={head_dim})",
|
||||
),
|
||||
) as prof_bf16:
|
||||
for _ in range(warmup + 1):
|
||||
with record_function("bf16_attention"):
|
||||
attn_bf16._forward_flashinfer(
|
||||
q.clone(), k.clone(), v.clone(), cu_bf16, max_s, seq_l
|
||||
)
|
||||
torch.accelerator.synchronize()
|
||||
prof_bf16.step()
|
||||
|
||||
# Profile FP8 (warmup handled by profiler schedule)
|
||||
with profile(
|
||||
activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
|
||||
schedule=sched,
|
||||
on_trace_ready=_make_trace_handler(
|
||||
output_dir,
|
||||
f"fp8_h{head_dim}_s{seq_len}",
|
||||
f"FP8 Attention (seq_len={seq_len}, heads={num_heads}, "
|
||||
f"head_dim={head_dim})",
|
||||
),
|
||||
) as prof_fp8:
|
||||
for _ in range(warmup + 1):
|
||||
with record_function("fp8_attention"):
|
||||
attn_fp8._forward_flashinfer(
|
||||
q.clone(), k.clone(), v.clone(), cu_fp8, max_s, seq_l
|
||||
)
|
||||
torch.accelerator.synchronize()
|
||||
prof_fp8.step()
|
||||
|
||||
print(f"\nTensorBoard traces saved to: {output_dir}")
|
||||
print(f"View with: tensorboard --logdir={output_dir}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = FlexibleArgumentParser(description="Benchmark FP8 vs BF16 ViT attention.")
|
||||
parser.add_argument(
|
||||
"--seq-lens",
|
||||
type=int,
|
||||
nargs="+",
|
||||
default=DEFAULT_SEQ_LENS,
|
||||
help="Sequence lengths to benchmark",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-heads",
|
||||
type=int,
|
||||
default=NUM_HEADS,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--head-dim",
|
||||
type=int,
|
||||
default=HEAD_DIM,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--method",
|
||||
choices=["cupti", "cudagraph"],
|
||||
default="cudagraph",
|
||||
help="GPU timing method: cupti (CUPTI kernel timing) or "
|
||||
"cudagraph (CUDA graph capture/replay). Default: cudagraph",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--warmup",
|
||||
type=int,
|
||||
default=10,
|
||||
help="Warmup iterations (profile mode only)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--profile",
|
||||
action="store_true",
|
||||
help="Run PyTorch profiler instead of benchmark",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--profile-seq-len",
|
||||
type=int,
|
||||
default=8192,
|
||||
help="Sequence length for profiling (default: 8192)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--profile-output-dir",
|
||||
type=str,
|
||||
default="./profile_traces",
|
||||
help="Output directory for TensorBoard traces (default: ./profile_traces)",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.profile:
|
||||
run_profile(
|
||||
args.profile_seq_len,
|
||||
args.num_heads,
|
||||
args.head_dim,
|
||||
args.warmup,
|
||||
args.profile_output_dir,
|
||||
)
|
||||
else:
|
||||
run_benchmark(
|
||||
args.seq_lens,
|
||||
args.num_heads,
|
||||
args.head_dim,
|
||||
args.method,
|
||||
)
|
||||
+2
-4
@@ -538,11 +538,9 @@ RUN CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-') && \
|
||||
cuda-nvrtc-${CUDA_VERSION_DASH} \
|
||||
cuda-cuobjdump-${CUDA_VERSION_DASH} \
|
||||
libcurand-dev-${CUDA_VERSION_DASH} \
|
||||
libcublas-${CUDA_VERSION_DASH} \
|
||||
libcublas-dev-${CUDA_VERSION_DASH} \
|
||||
# Required by fastsafetensors (fixes #20384)
|
||||
libnuma-dev \
|
||||
# numactl CLI for NUMA binding at runtime
|
||||
numactl && \
|
||||
libnuma-dev && \
|
||||
# Fixes nccl_allocator requiring nccl.h at runtime
|
||||
# https://github.com/vllm-project/vllm/blob/1336a1ea244fa8bfd7e72751cabbdb5b68a0c11a/vllm/distributed/device_communicators/pynccl_allocator.py#L22
|
||||
# NCCL packages don't use the cuda-MAJOR-MINOR naming convention,
|
||||
|
||||
@@ -68,7 +68,7 @@ You can pass a single image to the `'image'` field of the multi-modal dictionary
|
||||
print(generated_text)
|
||||
```
|
||||
|
||||
Full example: [examples/offline_inference/vision_language.py](../../examples/offline_inference/vision_language.py)
|
||||
Full example: [examples/generate/multimodal/vision_language_offline.py](../../examples/generate/multimodal/vision_language_offline.py)
|
||||
|
||||
To substitute multiple images inside the same text prompt, you can pass in a list of images instead:
|
||||
|
||||
@@ -101,7 +101,7 @@ To substitute multiple images inside the same text prompt, you can pass in a lis
|
||||
print(generated_text)
|
||||
```
|
||||
|
||||
Full example: [examples/offline_inference/vision_language_multi_image.py](../../examples/offline_inference/vision_language_multi_image.py)
|
||||
Full example: [examples/generate/multimodal/vision_language_multi_image_offline.py](../../examples/generate/multimodal/vision_language_multi_image_offline.py)
|
||||
|
||||
If using the [LLM.chat](../models/generative_models.md#llmchat) method, you can pass images directly in the message content using various formats: image URLs, PIL Image objects, or pre-computed embeddings:
|
||||
|
||||
@@ -287,13 +287,13 @@ Instead of NumPy arrays, you can also pass `'torch.Tensor'` instances, as shown
|
||||
!!! note
|
||||
'process_vision_info' is only applicable to Qwen2.5-VL and similar models.
|
||||
|
||||
Full example: [examples/offline_inference/vision_language.py](../../examples/offline_inference/vision_language.py)
|
||||
Full example: [examples/generate/multimodal/vision_language_offline.py](../../examples/generate/multimodal/vision_language_offline.py)
|
||||
|
||||
### Audio Inputs
|
||||
|
||||
You can pass a tuple `(array, sampling_rate)` to the `'audio'` field of the multi-modal dictionary.
|
||||
|
||||
Full example: [examples/offline_inference/audio_language.py](../../examples/offline_inference/audio_language.py)
|
||||
Full example: [examples/generate/multimodal/audio_language_offline.py](../../examples/generate/multimodal/audio_language_offline.py)
|
||||
|
||||
#### Chunking Long Audio for Transcription
|
||||
|
||||
@@ -674,7 +674,7 @@ Then, you can use the OpenAI client as follows:
|
||||
print("Chat completion output:", chat_response.choices[0].message.content)
|
||||
```
|
||||
|
||||
Full example: [examples/online_serving/openai_chat_completion_client_for_multimodal.py](../../examples/online_serving/openai_chat_completion_client_for_multimodal.py)
|
||||
Full example: [examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py](../../examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py)
|
||||
|
||||
!!! tip
|
||||
Loading from local file paths is also supported on vLLM: You can specify the allowed local media path via `--allowed-local-media-path` when launching the API server/engine,
|
||||
@@ -745,7 +745,7 @@ Then, you can use the OpenAI client as follows:
|
||||
print("Chat completion output from image url:", result)
|
||||
```
|
||||
|
||||
Full example: [examples/online_serving/openai_chat_completion_client_for_multimodal.py](../../examples/online_serving/openai_chat_completion_client_for_multimodal.py)
|
||||
Full example: [examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py](../../examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py)
|
||||
|
||||
!!! note
|
||||
By default, the timeout for fetching videos through HTTP URL is `30` seconds.
|
||||
@@ -958,7 +958,7 @@ Alternatively, you can pass `audio_url`, which is the audio counterpart of `imag
|
||||
print("Chat completion output from audio url:", result)
|
||||
```
|
||||
|
||||
Full example: [examples/online_serving/openai_chat_completion_client_for_multimodal.py](../../examples/online_serving/openai_chat_completion_client_for_multimodal.py)
|
||||
Full example: [examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py](../../examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py)
|
||||
|
||||
!!! note
|
||||
By default, the timeout for fetching audios through HTTP URL is `10` seconds.
|
||||
|
||||
@@ -20,6 +20,7 @@ The following are the supported quantization formats for vLLM:
|
||||
- [AMD Quark](quark.md)
|
||||
- [Quantized KV Cache](quantized_kvcache.md)
|
||||
- [TorchAO](torchao.md)
|
||||
- [FP8 ViT Encoder Attention](fp8_vit_attn.md)
|
||||
|
||||
## Supported Hardware
|
||||
|
||||
|
||||
@@ -0,0 +1,109 @@
|
||||
# FP8 ViT Encoder Attention
|
||||
|
||||
For visual understanding workloads with large images (e.g. QHD, 4K) and relatively
|
||||
short text prompts/generation, the ViT encoder attention can become a significant
|
||||
bottleneck, especially when the text model is quantized (e.g. NVFP4). vLLM
|
||||
supports optional FP8 quantization for the ViT encoder attention via the
|
||||
FlashInfer cuDNN backend. Q/K/V are quantized on-the-fly to FP8 before the
|
||||
cuDNN attention call.
|
||||
|
||||
!!! note
|
||||
- Currently supports Qwen3-VL family models only (`qwen3_vl`, `qwen3_vl_moe`,
|
||||
`qwen3_5`, `qwen3_5_moe`, and other models using Qwen3 ViT).
|
||||
- Dynamic scaling is not compatible with ViT full CUDA graphs.
|
||||
- Performance gains are mostly visible at QHD/4K resolutions or multi-image
|
||||
requests. Smaller images may see no speedup due to quantization overhead
|
||||
(3 quantization kernel launches + un-padding).
|
||||
- FP8 tensor-core speedup is more pronounced on GB300 than GB200.
|
||||
|
||||
## Requirements
|
||||
|
||||
- FlashInfer cuDNN backend with cuDNN >= 9.17.1.
|
||||
|
||||
## Usage
|
||||
|
||||
Enable FP8 ViT attention by passing `--mm-encoder-attn-dtype fp8` together
|
||||
with `--mm-encoder-attn-backend FLASHINFER`:
|
||||
|
||||
```bash
|
||||
vllm serve $MODEL \
|
||||
--mm-encoder-attn-backend FLASHINFER \
|
||||
--mm-encoder-attn-dtype fp8
|
||||
```
|
||||
|
||||
By default (no scale file), **dynamic scaling** is used: a 16-entry circular
|
||||
buffer of observed Q/K/V amax values drives per-forward scale updates. This
|
||||
matches BF16 accuracy without any calibration but adds a small per-forward
|
||||
overhead.
|
||||
|
||||
## Calibrate-Once, Reuse Workflow (Recommended)
|
||||
|
||||
For production, calibrate static scales on a representative dataset once and
|
||||
reuse them to avoid the dynamic overhead:
|
||||
|
||||
```bash
|
||||
# Step 1: calibrate and save scales (runs dynamic scaling for 16 passes,
|
||||
# then dumps the learned scales to JSON).
|
||||
vllm bench mm-processor \
|
||||
--model $MODEL --mm-encoder-attn-backend FLASHINFER \
|
||||
--mm-encoder-attn-dtype fp8 \
|
||||
--mm-encoder-fp8-scale-save-path /path/to/scales.json \
|
||||
--dataset-name hf --dataset-path lmarena-ai/VisionArena-Chat \
|
||||
--num-prompts 100
|
||||
|
||||
# Step 2: serve with static scales (no dynamic overhead).
|
||||
vllm serve $MODEL \
|
||||
--mm-encoder-attn-backend FLASHINFER \
|
||||
--mm-encoder-attn-dtype fp8 \
|
||||
--mm-encoder-fp8-scale-path /path/to/scales.json
|
||||
```
|
||||
|
||||
Saved scales are multiplied by `--mm-encoder-fp8-scale-save-margin` (default
|
||||
`1.5`) to leave headroom against activation outliers not present in the
|
||||
calibration set. The default has been validated to generalize across datasets
|
||||
(e.g. VisionArena-Chat calibration maintains BF16 accuracy on ChartQA).
|
||||
|
||||
## Scale File Format
|
||||
|
||||
```json
|
||||
{
|
||||
"visual.blocks.0.attn.attn": {"q": 224.0, "k": 198.0, "v": 210.0},
|
||||
"visual.blocks.1.attn.attn": {"q": 218.0, "k": 195.0, "v": 207.0}
|
||||
}
|
||||
```
|
||||
|
||||
Keys `q_scale` / `k_scale` / `v_scale` are accepted as aliases.
|
||||
|
||||
## Performance
|
||||
|
||||
**Core cuDNN attention kernel** (PyTorch profiler, `cudnn_generated_fort_native_sdpa_sm100_flash_fprop`, head_dim=128, seq_len=8192):
|
||||
|
||||
| Hardware | BF16 | FP8 | Speedup |
|
||||
| -------- | ---- | ---- | ------- |
|
||||
| GB200 | 350 us | 312 us | **1.12x** |
|
||||
| GB300 | 300 us | 211 us | **1.42x** |
|
||||
|
||||
**End-to-end encoder forward time** (Qwen3-VL-30B-A3B-Instruct on GB200, 3 images/request):
|
||||
|
||||
| Resolution | BF16 median | FP8 median | Speedup |
|
||||
| ---------- | ----------- | ---------- | ------- |
|
||||
| HD (720x1280) | 31.77 ms | 36.39 ms | 0.87x |
|
||||
| FullHD (1080x1920) | 57.99 ms | 58.73 ms | ~same |
|
||||
| QHD (1440x2560) | 131.83 ms | 122.30 ms | **1.08x** |
|
||||
| 4K (2160x3840) | 543.44 ms | 460.31 ms | **1.18x** |
|
||||
|
||||
Crossover is around FullHD with 3 images/request. At QHD and above, FP8 wins.
|
||||
|
||||
## Accuracy
|
||||
|
||||
ChartQA, Qwen3-VL-8B-Instruct, 500 samples. FP8 static uses scales calibrated
|
||||
on VisionArena-Chat (with default 1.5x margin):
|
||||
|
||||
| Metric | BF16 | FP8 dynamic | FP8 static |
|
||||
| ------ | ---- | ----------- | ---------- |
|
||||
| relaxed_accuracy | 0.780 | 0.776 | 0.780 |
|
||||
| anywhere_accuracy | 0.806 | 0.816 | 0.814 |
|
||||
| exact_match | 0.584 | 0.582 | 0.578 |
|
||||
|
||||
All three configurations match within statistical noise, confirming that
|
||||
static scales calibrated on one dataset generalize to another.
|
||||
@@ -202,7 +202,7 @@ The reasoning content is also available when both tool calling and the reasoning
|
||||
print(f"Arguments: {tool_call.arguments}")
|
||||
```
|
||||
|
||||
For more examples, please refer to [examples/online_serving/openai_chat_completion_tool_calls_with_reasoning.py](../../examples/online_serving/openai_chat_completion_tool_calls_with_reasoning.py).
|
||||
For more examples, please refer to [examples/reasoning/openai_chat_completion_tool_calls_with_reasoning.py](../../examples/reasoning/openai_chat_completion_tool_calls_with_reasoning.py).
|
||||
|
||||
## Server-Level Default Chat Template Kwargs
|
||||
|
||||
|
||||
@@ -292,10 +292,10 @@ Pooling models now support token-wise task.
|
||||
|
||||
### Score task
|
||||
|
||||
`score` task is deprecated and will be removed in v0.20. Please use `classify` instead. Only when a
|
||||
classification model outputs num_labels equal to 1 can it be used as a scoring model and have its scoring API enabled.
|
||||
`score` task have has been removed in v0.21, use `classify` instead. Only when a classification model outputs num_labels
|
||||
equal to 1 can it be used as a scoring model and have its scoring API enabled.
|
||||
|
||||
### Pooling multitask support
|
||||
|
||||
Pooling multitask support is deprecated and will be removed in v0.20. When the default pooling task is not what you want,
|
||||
Pooling multitask support has been removed in v0.21. When the default pooling task is not what you want,
|
||||
you need to manually specify it via `PoolerConfig(task=<task>)` offline or `--pooler-config.task <task>` online.
|
||||
|
||||
@@ -439,6 +439,7 @@ th {
|
||||
| `Mamba2ForCausalLM` | Mamba2 | `mistralai/Mamba-Codestral-7B-v0.1`, etc. | | ✅︎ |
|
||||
| `MiMoForCausalLM` | MiMo | `XiaomiMiMo/MiMo-7B-RL`, etc. | ✅︎ | ✅︎ |
|
||||
| `MiMoV2FlashForCausalLM` | MiMoV2Flash | `XiaomiMiMo/MiMo-V2-Flash`, etc. | | ✅︎ |
|
||||
| `MiMoV2ProForCausalLM` | MiMoV2Pro | `XiaomiMiMo/MiMo-V2.5-Pro`, etc. | | ✅︎ |
|
||||
| `MiniCPMForCausalLM` | MiniCPM | `openbmb/MiniCPM-2B-sft-bf16`, `openbmb/MiniCPM-2B-dpo-bf16`, `openbmb/MiniCPM-S-1B-sft`, etc. | ✅︎ | ✅︎ |
|
||||
| `MiniCPM3ForCausalLM` | MiniCPM3 | `openbmb/MiniCPM3-4B`, etc. | ✅︎ | ✅︎ |
|
||||
| `MiniMaxForCausalLM` | MiniMax-Text | `MiniMaxAI/MiniMax-Text-01-hf`, etc. | | |
|
||||
@@ -590,6 +591,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
|
||||
| `LlavaNextVideoForConditionalGeneration` | LLaVA-NeXT-Video | T + V | `llava-hf/LLaVA-NeXT-Video-7B-hf`, etc. | | ✅︎ |
|
||||
| `LlavaOnevisionForConditionalGeneration` | LLaVA-Onevision | T + I<sup>+</sup> + V<sup>+</sup> | `llava-hf/llava-onevision-qwen2-7b-ov-hf`, `llava-hf/llava-onevision-qwen2-0.5b-ov-hf`, etc. | | ✅︎ |
|
||||
| `MiDashengLMModel` | MiDashengLM | T + A<sup>+</sup> | `mispeech/midashenglm-7b` | | ✅︎ |
|
||||
| `MiMoV2OmniForCausalLM` | MiMo-V2.5-Omni | T + I<sup>E+</sup> + V<sup>E+</sup> + A<sup>+</sup> | `XiaomiMiMo/MiMo-V2.5-Omni` | | ✅︎ |
|
||||
| `MiniCPMO` | MiniCPM-O | T + I<sup>E+</sup> + V<sup>E+</sup> + A<sup>E+</sup> | `openbmb/MiniCPM-o-2_6`, etc. | ✅︎ | ✅︎ |
|
||||
| `MiniCPMV` | MiniCPM-V | T + I<sup>E+</sup> + V<sup>E+</sup> | `openbmb/MiniCPM-V-2` (see note), `openbmb/MiniCPM-Llama3-V-2_5`, `openbmb/MiniCPM-V-2_6`, `openbmb/MiniCPM-V-4`, `openbmb/MiniCPM-V-4_5`, etc. | ✅︎ | |
|
||||
| `MiniMaxVL01ForConditionalGeneration` | MiniMax-VL | T + I<sup>E+</sup> | `MiniMaxAI/MiniMax-VL-01`, etc. | | ✅︎ |
|
||||
|
||||
@@ -251,7 +251,7 @@ The following extra parameters are supported:
|
||||
Our Responses API is compatible with [OpenAI's Responses API](https://platform.openai.com/docs/api-reference/responses);
|
||||
you can use the [official OpenAI Python client](https://github.com/openai/openai-python) to interact with it.
|
||||
|
||||
Code example: [examples/online_serving/openai_responses_client_with_tools.py](../../examples/online_serving/openai_responses_client_with_tools.py)
|
||||
Code example: [examples/online_serving/openai_responses_client_with_tools.py](../../examples/tool_calling/openai_responses_client_with_tools.py)
|
||||
|
||||
#### Extra parameters
|
||||
|
||||
@@ -279,7 +279,7 @@ you can use the [official OpenAI Python client](https://github.com/openai/openai
|
||||
!!! note
|
||||
To use the Transcriptions API, please install with extra audio dependencies using `pip install vllm[audio]`.
|
||||
|
||||
Code example: [examples/online_serving/openai_transcription_client.py](../../examples/online_serving/openai_transcription_client.py)
|
||||
Code example: [examples/speech_to_text/openai/openai_transcription_client.py](../../examples/speech_to_text/openai/openai_transcription_client.py)
|
||||
|
||||
NOTE: beam search is currently supported in the transcriptions endpoint for encoder-decoder multimodal models, e.g., whisper, but highly inefficient as work for handling the encoder/decoder cache is actively ongoing. This is an active point of ongoing optimization and will be handled properly in the very near future.
|
||||
|
||||
@@ -397,7 +397,7 @@ Please mind that the popular `openai/whisper-large-v3-turbo` model does not supp
|
||||
!!! note
|
||||
To use the Translation API, please install with extra audio dependencies using `pip install vllm[audio]`.
|
||||
|
||||
Code example: [examples/online_serving/openai_translation_client.py](../../examples/online_serving/openai_translation_client.py)
|
||||
Code example: [examples/speech_to_text/openai/openai_translation_client.py](../../examples/speech_to_text/openai/openai_translation_client.py)
|
||||
|
||||
#### Extra Parameters
|
||||
|
||||
|
||||
Executable → Regular
+6
-6
@@ -6,15 +6,15 @@ This folder provides several example scripts on how to inference Qwen2.5-Omni of
|
||||
|
||||
```bash
|
||||
# Audio + image + video
|
||||
python examples/offline_inference/qwen2_5_omni/only_thinker.py \
|
||||
python examples/generate/multimodal/qwen2_5_omni/only_thinker.py \
|
||||
-q mixed_modalities
|
||||
|
||||
# Read vision and audio inputs from a single video file
|
||||
python examples/offline_inference/qwen2_5_omni/only_thinker.py \
|
||||
python examples/generate/multimodal/qwen2_5_omni/only_thinker.py \
|
||||
-q use_audio_in_video
|
||||
|
||||
# Multiple audios
|
||||
python examples/offline_inference/qwen2_5_omni/only_thinker.py \
|
||||
python examples/generate/multimodal/qwen2_5_omni/only_thinker.py \
|
||||
-q multi_audios
|
||||
```
|
||||
|
||||
@@ -24,16 +24,16 @@ You can also test Qwen2.5-Omni on a single modality:
|
||||
|
||||
```bash
|
||||
# Process audio inputs
|
||||
python examples/offline_inference/audio_language.py \
|
||||
python examples/generate/multimodal/audio_language_offline.py \
|
||||
--model-type qwen2_5_omni
|
||||
|
||||
# Process image inputs
|
||||
python examples/offline_inference/vision_language.py \
|
||||
python examples/generate/multimodal/vision_language_offline.py \
|
||||
--modality image \
|
||||
--model-type qwen2_5_omni
|
||||
|
||||
# Process video inputs
|
||||
python examples/offline_inference/vision_language.py \
|
||||
python examples/generate/multimodal/vision_language_offline.py \
|
||||
--modality video \
|
||||
--model-type qwen2_5_omni
|
||||
```
|
||||
Executable → Regular
Executable → Regular
+1
-1
@@ -1402,7 +1402,7 @@ def run_mantis(questions: list[str], modality: str) -> ModelRequestData:
|
||||
# MiniCPM-V
|
||||
def run_minicpmv_base(questions: list[str], modality: str, model_name):
|
||||
assert modality in ["image", "video", "image+video"]
|
||||
# If you want to use `MiniCPM-o-2_6` with audio inputs, check `audio_language.py` # noqa
|
||||
# If you want to use `MiniCPM-o-2_6` with audio inputs, check `audio_language_offline.py` # noqa
|
||||
|
||||
# 2.0
|
||||
# The official repo doesn't work yet, so we need to use a fork for now
|
||||
@@ -4,68 +4,74 @@
|
||||
import torch
|
||||
|
||||
from vllm import LLM
|
||||
from vllm.config import PoolerConfig
|
||||
from vllm.inputs import TextPrompt
|
||||
from vllm.multimodal.utils import fetch_image
|
||||
|
||||
# Initialize model
|
||||
model = LLM(
|
||||
model="jinaai/jina-embeddings-v4-vllm-text-matching",
|
||||
runner="pooling",
|
||||
max_model_len=1024,
|
||||
gpu_memory_utilization=0.8,
|
||||
)
|
||||
|
||||
# Create text prompts
|
||||
text1 = "Ein wunderschöner Sonnenuntergang am Strand"
|
||||
text1_prompt = TextPrompt(prompt=f"Query: {text1}")
|
||||
def main():
|
||||
# Initialize model
|
||||
model = LLM(
|
||||
model="jinaai/jina-embeddings-v4-vllm-text-matching",
|
||||
pooler_config=PoolerConfig(task="token_embed"),
|
||||
runner="pooling",
|
||||
max_model_len=1024,
|
||||
gpu_memory_utilization=0.8,
|
||||
)
|
||||
|
||||
text2 = "浜辺に沈む美しい夕日"
|
||||
text2_prompt = TextPrompt(prompt=f"Query: {text2}")
|
||||
# Create text prompts
|
||||
text1 = "Ein wunderschöner Sonnenuntergang am Strand"
|
||||
text1_prompt = TextPrompt(prompt=f"Query: {text1}")
|
||||
|
||||
# Create image prompt
|
||||
image = fetch_image(
|
||||
"https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/eskimo.jpg" # noqa: E501
|
||||
)
|
||||
image_prompt = TextPrompt(
|
||||
prompt="<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Describe the image.<|im_end|>\n", # noqa: E501
|
||||
multi_modal_data={"image": image},
|
||||
)
|
||||
text2 = "浜辺に沈む美しい夕日"
|
||||
text2_prompt = TextPrompt(prompt=f"Query: {text2}")
|
||||
|
||||
# Encode all prompts
|
||||
prompts = [text1_prompt, text2_prompt, image_prompt]
|
||||
outputs = model.encode(prompts, pooling_task="token_embed")
|
||||
# Create image prompt
|
||||
image = fetch_image(
|
||||
"https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/eskimo.jpg" # noqa: E501
|
||||
)
|
||||
image_prompt = TextPrompt(
|
||||
prompt="<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Describe the image.<|im_end|>\n", # noqa: E501
|
||||
multi_modal_data={"image": image},
|
||||
)
|
||||
|
||||
# Encode all prompts
|
||||
prompts = [text1_prompt, text2_prompt, image_prompt]
|
||||
outputs = model.encode(prompts, pooling_task="token_embed")
|
||||
|
||||
def get_embeddings(outputs):
|
||||
VISION_START_TOKEN_ID, VISION_END_TOKEN_ID = 151652, 151653
|
||||
|
||||
embeddings = []
|
||||
for output in outputs:
|
||||
if VISION_START_TOKEN_ID in output.prompt_token_ids:
|
||||
# Gather only vision tokens
|
||||
img_start_pos = torch.where(
|
||||
torch.tensor(output.prompt_token_ids) == VISION_START_TOKEN_ID
|
||||
)[0][0]
|
||||
img_end_pos = torch.where(
|
||||
torch.tensor(output.prompt_token_ids) == VISION_END_TOKEN_ID
|
||||
)[0][0]
|
||||
embeddings_tensor = output.outputs.data.detach().clone()[
|
||||
img_start_pos : img_end_pos + 1
|
||||
]
|
||||
else:
|
||||
# Use all tokens for text-only prompts
|
||||
embeddings_tensor = output.outputs.data.detach().clone()
|
||||
|
||||
# Pool and normalize embeddings
|
||||
pooled_output = (
|
||||
embeddings_tensor.sum(dim=0, dtype=torch.float32)
|
||||
/ embeddings_tensor.shape[0]
|
||||
)
|
||||
embeddings.append(torch.nn.functional.normalize(pooled_output, dim=-1))
|
||||
return embeddings
|
||||
|
||||
embeddings = get_embeddings(outputs)
|
||||
|
||||
for embedding in embeddings:
|
||||
print(embedding.shape)
|
||||
|
||||
|
||||
def get_embeddings(outputs):
|
||||
VISION_START_TOKEN_ID, VISION_END_TOKEN_ID = 151652, 151653
|
||||
|
||||
embeddings = []
|
||||
for output in outputs:
|
||||
if VISION_START_TOKEN_ID in output.prompt_token_ids:
|
||||
# Gather only vision tokens
|
||||
img_start_pos = torch.where(
|
||||
torch.tensor(output.prompt_token_ids) == VISION_START_TOKEN_ID
|
||||
)[0][0]
|
||||
img_end_pos = torch.where(
|
||||
torch.tensor(output.prompt_token_ids) == VISION_END_TOKEN_ID
|
||||
)[0][0]
|
||||
embeddings_tensor = output.outputs.data.detach().clone()[
|
||||
img_start_pos : img_end_pos + 1
|
||||
]
|
||||
else:
|
||||
# Use all tokens for text-only prompts
|
||||
embeddings_tensor = output.outputs.data.detach().clone()
|
||||
|
||||
# Pool and normalize embeddings
|
||||
pooled_output = (
|
||||
embeddings_tensor.sum(dim=0, dtype=torch.float32)
|
||||
/ embeddings_tensor.shape[0]
|
||||
)
|
||||
embeddings.append(torch.nn.functional.normalize(pooled_output, dim=-1))
|
||||
return embeddings
|
||||
|
||||
|
||||
embeddings = get_embeddings(outputs)
|
||||
|
||||
for embedding in embeddings:
|
||||
print(embedding.shape)
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
@@ -4,6 +4,7 @@
|
||||
from argparse import Namespace
|
||||
|
||||
from vllm import LLM, EngineArgs
|
||||
from vllm.config import PoolerConfig
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
|
||||
@@ -13,6 +14,7 @@ def parse_args():
|
||||
# Set example specific arguments
|
||||
parser.set_defaults(
|
||||
model="BAAI/bge-m3",
|
||||
pooler_config=PoolerConfig(task="token_embed"),
|
||||
runner="pooling",
|
||||
enforce_eager=True,
|
||||
)
|
||||
@@ -32,15 +34,6 @@ def main(args: Namespace):
|
||||
# You should pass runner="pooling" for embedding models
|
||||
llm = LLM(**vars(args))
|
||||
|
||||
# Generate embedding. The output is a list of EmbeddingRequestOutputs.
|
||||
outputs = llm.embed(prompts)
|
||||
|
||||
# Print the outputs.
|
||||
print("\nGenerated Outputs:\n" + "-" * 60)
|
||||
for prompt, output in zip(prompts, outputs):
|
||||
embeds = output.outputs.embedding
|
||||
print(len(embeds))
|
||||
|
||||
# Generate embedding for each token. The output is a list of PoolingRequestOutput.
|
||||
outputs = llm.encode(prompts, pooling_task="token_embed")
|
||||
|
||||
@@ -50,6 +43,20 @@ def main(args: Namespace):
|
||||
multi_vector = output.outputs.data
|
||||
print(multi_vector.shape)
|
||||
|
||||
query = "What is the capital of France?"
|
||||
documents = [
|
||||
"The capital of Brazil is Brasilia.",
|
||||
"The capital of France is Paris.",
|
||||
]
|
||||
# Generate scores.
|
||||
outputs = llm.score(query, documents)
|
||||
# Print the outputs.
|
||||
print("\nGenerated Outputs:\n" + "-" * 60)
|
||||
for document, output in zip(documents, outputs):
|
||||
score = output.outputs.score
|
||||
print(f"Pair: {[query, document]!r} \nScore: {score}")
|
||||
print("-" * 60)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
|
||||
@@ -7,10 +7,11 @@ Example online usage of Pooling API for multi vector retrieval.
|
||||
Run `vllm serve <model> --runner pooling`
|
||||
to start up the server in vLLM. e.g.
|
||||
|
||||
vllm serve BAAI/bge-m3
|
||||
vllm serve BAAI/bge-m3 --pooler-config.task token_embed
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import pprint
|
||||
|
||||
import requests
|
||||
import torch
|
||||
@@ -32,7 +33,8 @@ def parse_args():
|
||||
|
||||
|
||||
def main(args):
|
||||
api_url = f"http://{args.host}:{args.port}/pooling"
|
||||
pooling_url = f"http://{args.host}:{args.port}/pooling"
|
||||
score_url = f"http://{args.host}:{args.port}/score"
|
||||
model_name = args.model
|
||||
|
||||
prompts = [
|
||||
@@ -43,11 +45,23 @@ def main(args):
|
||||
]
|
||||
prompt = {"model": model_name, "input": prompts}
|
||||
|
||||
pooling_response = post_http_request(prompt=prompt, api_url=api_url)
|
||||
pooling_response = post_http_request(prompt=prompt, api_url=pooling_url)
|
||||
for output in pooling_response.json()["data"]:
|
||||
multi_vector = torch.tensor(output["data"])
|
||||
print(multi_vector.shape)
|
||||
|
||||
queries = "What is the capital of France?"
|
||||
documents = [
|
||||
"The capital of Brazil is Brasilia.",
|
||||
"The capital of France is Paris.",
|
||||
]
|
||||
prompt = {"model": model_name, "queries": queries, "documents": documents}
|
||||
score_response = post_http_request(prompt=prompt, api_url=score_url)
|
||||
print("\nPrompt when queries is string and documents is a list:")
|
||||
pprint.pprint(prompt)
|
||||
print("\nScore Response:")
|
||||
pprint.pprint(score_response.json())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
|
||||
@@ -41,3 +41,21 @@ def test_language_model_only_affects_model_hash():
|
||||
base_hash = ModelConfig(model).compute_hash()
|
||||
lm_only_hash = ModelConfig(model, language_model_only=True).compute_hash()
|
||||
assert base_hash != lm_only_hash
|
||||
|
||||
|
||||
def test_mm_encoder_fp8_scale_path_requires_fp8():
|
||||
with pytest.raises(ValueError, match="mm_encoder_attn_dtype"):
|
||||
MultiModalConfig(mm_encoder_fp8_scale_path="/tmp/scales.json")
|
||||
|
||||
|
||||
def test_mm_encoder_attn_dtype_hash_updates(tmp_path):
|
||||
scale_file = tmp_path / "scales.json"
|
||||
scale_file.write_text("{}")
|
||||
base_hash = MultiModalConfig().compute_hash()
|
||||
fp8_hash = MultiModalConfig(mm_encoder_attn_dtype="fp8").compute_hash()
|
||||
fp8_static_hash = MultiModalConfig(
|
||||
mm_encoder_attn_dtype="fp8",
|
||||
mm_encoder_fp8_scale_path=str(scale_file),
|
||||
).compute_hash()
|
||||
assert base_hash != fp8_hash
|
||||
assert fp8_hash != fp8_static_hash
|
||||
|
||||
@@ -0,0 +1,76 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for ``system_fingerprint`` construction."""
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.entrypoints.openai import fingerprint as fp
|
||||
|
||||
|
||||
def _cfg(tp=1, pp=1, dp=1, ep=False, digest="a3b21f94deadbeef"):
|
||||
c = SimpleNamespace(
|
||||
parallel_config=SimpleNamespace(
|
||||
tensor_parallel_size=tp,
|
||||
pipeline_parallel_size=pp,
|
||||
data_parallel_size=dp,
|
||||
enable_expert_parallel=ep,
|
||||
)
|
||||
)
|
||||
c.compute_hash = lambda: digest # type: ignore[attr-defined]
|
||||
return c
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _reset():
|
||||
fp.set_default_fingerprint_mode("full")
|
||||
yield
|
||||
fp.set_default_fingerprint_mode("full")
|
||||
|
||||
|
||||
def test_four_modes_produce_expected_shapes():
|
||||
from vllm import __version__ as v
|
||||
|
||||
cfg = _cfg(tp=8, ep=True)
|
||||
|
||||
assert fp.build_system_fingerprint(cfg, "full") == (f"vllm-{v}-tp8-ep-a3b21f94")
|
||||
assert fp.build_system_fingerprint(cfg, "hash") == f"vllm-{v}-a3b21f94"
|
||||
assert fp.build_system_fingerprint(cfg, "custom", "my-fp") == "my-fp"
|
||||
assert fp.build_system_fingerprint(cfg, "none") is None
|
||||
|
||||
|
||||
def test_full_mode_emits_only_non_trivial_parallelism():
|
||||
from vllm import __version__ as v
|
||||
|
||||
# Single-GPU: nothing between version and hash.
|
||||
assert fp.build_system_fingerprint(_cfg(), "full") == f"vllm-{v}-a3b21f94"
|
||||
# All parallelism axes.
|
||||
assert (
|
||||
fp.build_system_fingerprint(_cfg(tp=8, pp=2, dp=4, ep=True), "full")
|
||||
== f"vllm-{v}-tp8-pp2-dp4-ep-a3b21f94"
|
||||
)
|
||||
|
||||
|
||||
def test_get_respects_set_default():
|
||||
cfg = _cfg(tp=8)
|
||||
full = fp.get_system_fingerprint(cfg)
|
||||
assert full == fp.get_system_fingerprint(cfg)
|
||||
|
||||
fp.set_default_fingerprint_mode("hash")
|
||||
hashed = fp.get_system_fingerprint(cfg)
|
||||
assert hashed != full
|
||||
assert "tp8" not in hashed
|
||||
|
||||
fp.set_default_fingerprint_mode("custom", "deploy-42")
|
||||
assert fp.get_system_fingerprint(cfg) == "deploy-42"
|
||||
|
||||
fp.set_default_fingerprint_mode("none")
|
||||
assert fp.get_system_fingerprint(cfg) is None
|
||||
|
||||
|
||||
def test_compute_hash_failure_does_not_raise():
|
||||
cfg = _cfg()
|
||||
cfg.compute_hash = lambda: (_ for _ in ()).throw(RuntimeError("boom"))
|
||||
assert fp.build_system_fingerprint(cfg, "full").endswith("-nohash")
|
||||
assert fp.build_system_fingerprint(cfg, "hash").endswith("-nohash")
|
||||
@@ -1,13 +1,12 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import logging
|
||||
import weakref
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.models.utils import softmax
|
||||
from vllm import LLM, ClassificationRequestOutput, PoolingParams, PoolingRequestOutput
|
||||
from vllm import LLM, ClassificationRequestOutput, PoolingParams
|
||||
from vllm.distributed import cleanup_dist_env_and_memory
|
||||
from vllm.tasks import PoolingTask
|
||||
|
||||
@@ -66,18 +65,6 @@ def test_list_prompts(llm: LLM):
|
||||
assert len(outputs[i].outputs.probs) == num_labels
|
||||
|
||||
|
||||
@pytest.mark.skip_global_cleanup
|
||||
def test_token_classify(llm: LLM, caplog_vllm):
|
||||
with caplog_vllm.at_level(level=logging.WARNING, logger="vllm"):
|
||||
outputs = llm.encode(prompt, pooling_task="token_classify", use_tqdm=False)
|
||||
assert "deprecated" in caplog_vllm.text
|
||||
|
||||
assert len(outputs) == 1
|
||||
assert isinstance(outputs[0], PoolingRequestOutput)
|
||||
assert outputs[0].prompt_token_ids == prompt_token_ids
|
||||
assert outputs[0].outputs.data.shape == (len(prompt_token_ids), num_labels)
|
||||
|
||||
|
||||
@pytest.mark.skip_global_cleanup
|
||||
def test_pooling_params(llm: LLM):
|
||||
def get_outputs(use_activation):
|
||||
@@ -110,10 +97,12 @@ def test_score_api(llm: LLM):
|
||||
llm.score("ping", "pong", use_tqdm=False)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
|
||||
@pytest.mark.parametrize("task", ["embed", "token_embed", "token_classify", "plugin"])
|
||||
def test_unsupported_tasks(llm: LLM, task: PoolingTask):
|
||||
if task == "plugin":
|
||||
err_msg = "No IOProcessor plugin installed."
|
||||
elif task == "token_classify":
|
||||
err_msg = "Try switching the model's pooling_task via.+"
|
||||
else:
|
||||
err_msg = "Embedding API is not supported by this model.+"
|
||||
with pytest.raises(ValueError, match=err_msg):
|
||||
|
||||
@@ -436,26 +436,7 @@ async def test_pooling_classify(server: RemoteOpenAIServer, model_name: str):
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
async def test_pooling_token_classify(server: RemoteOpenAIServer, model_name: str):
|
||||
task = "token_classify"
|
||||
response = requests.post(
|
||||
server.url_for("pooling"),
|
||||
json={
|
||||
"model": model_name,
|
||||
"input": input_text,
|
||||
"encoding_format": "float",
|
||||
"task": task,
|
||||
},
|
||||
)
|
||||
poolings = PoolingResponse.model_validate(response.json())
|
||||
assert len(poolings.data) == 1
|
||||
assert len(poolings.data[0].data) == 8
|
||||
assert len(poolings.data[0].data[0]) == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
|
||||
@pytest.mark.parametrize("task", ["embed", "token_embed", "token_classify", "plugin"])
|
||||
async def test_pooling_not_supported(
|
||||
server: RemoteOpenAIServer, model_name: str, task: str
|
||||
):
|
||||
@@ -469,8 +450,11 @@ async def test_pooling_not_supported(
|
||||
},
|
||||
)
|
||||
assert response.json()["error"]["type"] == "BadRequestError"
|
||||
|
||||
if task == "plugin":
|
||||
err_msg = "No IOProcessor plugin installed."
|
||||
elif task == "token_classify":
|
||||
err_msg = "Try switching the model's pooling_task via"
|
||||
else:
|
||||
err_msg = f"Unsupported task: {task!r}"
|
||||
assert response.json()["error"]["message"].startswith(err_msg)
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import logging
|
||||
import weakref
|
||||
|
||||
import pytest
|
||||
@@ -38,11 +37,11 @@ def llm():
|
||||
seed=0,
|
||||
attention_config=attention_config,
|
||||
)
|
||||
assert embedding_size == llm.model_config.embedding_size
|
||||
|
||||
yield weakref.proxy(llm)
|
||||
|
||||
del llm
|
||||
|
||||
cleanup_dist_env_and_memory()
|
||||
|
||||
|
||||
@@ -74,16 +73,6 @@ def test_list_prompts(llm: LLM):
|
||||
assert len(outputs[i].outputs.embedding) == embedding_size
|
||||
|
||||
|
||||
@pytest.mark.skip_global_cleanup
|
||||
def test_token_embed(llm: LLM, caplog_vllm):
|
||||
with caplog_vllm.at_level(level=logging.WARNING, logger="vllm"):
|
||||
outputs = llm.encode(prompt, pooling_task="token_embed", use_tqdm=False)
|
||||
assert "deprecated" in caplog_vllm.text
|
||||
|
||||
multi_vector = outputs[0].outputs.data
|
||||
assert multi_vector.shape == (11, 384)
|
||||
|
||||
|
||||
@pytest.mark.skip_global_cleanup
|
||||
def test_pooling_params(llm: LLM):
|
||||
def get_outputs(normalize):
|
||||
@@ -107,10 +96,14 @@ def test_pooling_params(llm: LLM):
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("task", ["token_classify", "classify", "plugin"])
|
||||
@pytest.mark.parametrize(
|
||||
"task", ["token_classify", "classify", "token_embed", "plugin"]
|
||||
)
|
||||
def test_unsupported_tasks(llm: LLM, task: PoolingTask):
|
||||
if task == "plugin":
|
||||
err_msg = "No IOProcessor plugin installed."
|
||||
elif task == "token_embed":
|
||||
err_msg = "Try switching the model's pooling_task via.+"
|
||||
else:
|
||||
err_msg = "Classification API is not supported by this model.+"
|
||||
with pytest.raises(ValueError, match=err_msg):
|
||||
|
||||
@@ -732,28 +732,9 @@ async def test_pooling_embed(server: RemoteOpenAIServer, model_name: str):
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
async def test_pooling_token_embed(server: RemoteOpenAIServer, model_name: str):
|
||||
task = "token_embed"
|
||||
response = requests.post(
|
||||
server.url_for("pooling"),
|
||||
json={
|
||||
"model": model_name,
|
||||
"input": input_text,
|
||||
"encoding_format": "float",
|
||||
"task": task,
|
||||
},
|
||||
)
|
||||
|
||||
poolings = PoolingResponse.model_validate(response.json())
|
||||
|
||||
assert len(poolings.data) == 1
|
||||
assert len(poolings.data[0].data) == len(input_tokens)
|
||||
assert len(poolings.data[0].data[0]) == 384
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
@pytest.mark.parametrize("task", ["classify", "token_classify", "plugin"])
|
||||
@pytest.mark.parametrize(
|
||||
"task", ["classify", "token_classify", "token_embed", "plugin"]
|
||||
)
|
||||
async def test_pooling_not_supported(
|
||||
server: RemoteOpenAIServer, model_name: str, task: str
|
||||
):
|
||||
@@ -769,6 +750,8 @@ async def test_pooling_not_supported(
|
||||
assert response.json()["error"]["type"] == "BadRequestError"
|
||||
if task == "plugin":
|
||||
err_msg = "No IOProcessor plugin installed."
|
||||
elif task == "token_embed":
|
||||
err_msg = "Try switching the model's pooling_task via"
|
||||
else:
|
||||
err_msg = f"Unsupported task: {task!r}"
|
||||
assert response.json()["error"]["message"].startswith(err_msg)
|
||||
|
||||
@@ -452,25 +452,6 @@ async def test_pooling_classify(server: RemoteOpenAIServer):
|
||||
assert len(poolings.data[0].data) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pooling_token_classify(server: RemoteOpenAIServer):
|
||||
response = requests.post(
|
||||
server.url_for("pooling"),
|
||||
json={
|
||||
"model": MODEL_NAME,
|
||||
"task": "token_classify",
|
||||
"input": input_text,
|
||||
"encoding_format": "float",
|
||||
},
|
||||
)
|
||||
|
||||
poolings = PoolingResponse.model_validate(response.json())
|
||||
|
||||
assert len(poolings.data) == 1
|
||||
assert len(poolings.data[0].data) == len(input_tokens)
|
||||
assert len(poolings.data[0].data[0]) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rerank_max_tokens_per_doc(
|
||||
server: RemoteOpenAIServer,
|
||||
@@ -544,7 +525,7 @@ async def test_rerank_max_tokens_per_doc_validation(
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
|
||||
@pytest.mark.parametrize("task", ["embed", "token_embed", "token_classify", "plugin"])
|
||||
async def test_pooling_not_supported(server: RemoteOpenAIServer, task: str):
|
||||
response = requests.post(
|
||||
server.url_for("pooling"),
|
||||
@@ -558,6 +539,8 @@ async def test_pooling_not_supported(server: RemoteOpenAIServer, task: str):
|
||||
assert response.json()["error"]["type"] == "BadRequestError"
|
||||
if task == "plugin":
|
||||
err_msg = "No IOProcessor plugin installed."
|
||||
elif task == "token_classify":
|
||||
err_msg = "Try switching the model's pooling_task via"
|
||||
else:
|
||||
err_msg = f"Unsupported task: {task!r}"
|
||||
assert response.json()["error"]["message"].startswith(err_msg)
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import logging
|
||||
import weakref
|
||||
|
||||
import pytest
|
||||
@@ -60,22 +59,19 @@ def test_token_ids_prompts(llm: LLM):
|
||||
|
||||
@pytest.mark.skip_global_cleanup
|
||||
def test_score_api(llm: LLM):
|
||||
err_msg = "Scoring API is only enabled for num_labels == 1."
|
||||
err_msg = "This model does not support the Scoring API."
|
||||
with pytest.raises(ValueError, match=err_msg):
|
||||
llm.score("ping", "pong", use_tqdm=False)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("task", ["classify", "embed", "token_embed", "plugin"])
|
||||
def test_unsupported_tasks(llm: LLM, task: PoolingTask, caplog_vllm):
|
||||
if task == "classify":
|
||||
with caplog_vllm.at_level(level=logging.WARNING, logger="vllm"):
|
||||
llm.encode(prompt, pooling_task=task, use_tqdm=False)
|
||||
assert "deprecated" in caplog_vllm.text
|
||||
if task == "plugin":
|
||||
err_msg = "No IOProcessor plugin installed."
|
||||
elif task == "classify":
|
||||
err_msg = "Try switching the model's pooling_task via.+"
|
||||
else:
|
||||
if task == "plugin":
|
||||
err_msg = "No IOProcessor plugin installed."
|
||||
else:
|
||||
err_msg = "Embedding API is not supported by this model.+"
|
||||
err_msg = "Embedding API is not supported by this model.+"
|
||||
|
||||
with pytest.raises(ValueError, match=err_msg):
|
||||
llm.encode(prompt, pooling_task=task, use_tqdm=False)
|
||||
with pytest.raises(ValueError, match=err_msg):
|
||||
llm.encode(prompt, pooling_task=task, use_tqdm=False)
|
||||
|
||||
@@ -50,7 +50,7 @@ async def test_pooling_token_classify(server: RemoteOpenAIServer, model_name: st
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
|
||||
@pytest.mark.parametrize("task", ["classify", "embed", "token_embed", "plugin"])
|
||||
async def test_pooling_not_supported(
|
||||
server: RemoteOpenAIServer, model_name: str, task: str
|
||||
):
|
||||
@@ -63,9 +63,12 @@ async def test_pooling_not_supported(
|
||||
"task": task,
|
||||
},
|
||||
)
|
||||
assert response.json()["error"]["type"] == "BadRequestError"
|
||||
|
||||
if task == "plugin":
|
||||
err_msg = "No IOProcessor plugin installed."
|
||||
elif task == "classify":
|
||||
err_msg = "Try switching the model's pooling_task via"
|
||||
else:
|
||||
err_msg = f"Unsupported task: {task!r}"
|
||||
assert response.json()["error"]["message"].startswith(err_msg)
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import logging
|
||||
import weakref
|
||||
|
||||
import pytest
|
||||
@@ -64,15 +63,12 @@ def test_token_ids_prompts(llm: LLM):
|
||||
|
||||
@pytest.mark.parametrize("task", ["embed", "classify", "token_classify", "plugin"])
|
||||
def test_unsupported_tasks(llm: LLM, task: PoolingTask, caplog_vllm):
|
||||
if task == "embed":
|
||||
with caplog_vllm.at_level(level=logging.WARNING, logger="vllm"):
|
||||
llm.encode(prompt, pooling_task=task, use_tqdm=False)
|
||||
assert "deprecated" in caplog_vllm.text
|
||||
if task == "plugin":
|
||||
err_msg = "No IOProcessor plugin installed."
|
||||
elif task == "embed":
|
||||
err_msg = "Try switching the model's pooling_task via.+"
|
||||
else:
|
||||
if task == "plugin":
|
||||
err_msg = "No IOProcessor plugin installed."
|
||||
else:
|
||||
err_msg = "Classification API is not supported by this model.+"
|
||||
err_msg = "Classification API is not supported by this model.+"
|
||||
|
||||
with pytest.raises(ValueError, match=err_msg):
|
||||
llm.encode(prompt, pooling_task=task, use_tqdm=False)
|
||||
with pytest.raises(ValueError, match=err_msg):
|
||||
llm.encode(prompt, pooling_task=task, use_tqdm=False)
|
||||
|
||||
@@ -73,7 +73,7 @@ async def test_pooling_token_embed(server: RemoteOpenAIServer, model_name: str):
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
@pytest.mark.parametrize("task", ["classify", "token_classify", "plugin"])
|
||||
@pytest.mark.parametrize("task", ["embed", "classify", "token_classify", "plugin"])
|
||||
async def test_pooling_not_supported(
|
||||
server: RemoteOpenAIServer, model_name: str, task: str
|
||||
):
|
||||
@@ -86,9 +86,12 @@ async def test_pooling_not_supported(
|
||||
"task": task,
|
||||
},
|
||||
)
|
||||
assert response.json()["error"]["type"] == "BadRequestError"
|
||||
|
||||
if task == "plugin":
|
||||
err_msg = "No IOProcessor plugin installed."
|
||||
elif task == "embed":
|
||||
err_msg = "Try switching the model's pooling_task via"
|
||||
else:
|
||||
err_msg = f"Unsupported task: {task!r}"
|
||||
assert response.json()["error"]["message"].startswith(err_msg)
|
||||
|
||||
@@ -0,0 +1,279 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for the full FP8 ViT attention path (quantize -> cuDNN -> un-pad)."""
|
||||
|
||||
import contextlib
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.triton_utils import HAS_TRITON
|
||||
from vllm.utils.flashinfer import (
|
||||
is_flashinfer_cudnn_fp8_prefill_attn_supported,
|
||||
)
|
||||
from vllm.v1.attention.backends.registry import AttentionBackendEnum
|
||||
|
||||
|
||||
def _has_flashinfer_cudnn() -> bool:
|
||||
"""Check if FlashInfer cuDNN backend is available."""
|
||||
try:
|
||||
from flashinfer.prefill import (
|
||||
cudnn_batch_prefill_with_kv_cache, # noqa: F401
|
||||
)
|
||||
|
||||
return True
|
||||
except ImportError:
|
||||
return False
|
||||
|
||||
|
||||
HEAD_DIMS = [72, 80]
|
||||
SEQ_LENS = [256]
|
||||
NUM_HEADS = [16]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def _fp8_attention():
|
||||
"""Create FP8-enabled MMEncoderAttention via config."""
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
from vllm.config import VllmConfig, set_current_vllm_config
|
||||
from vllm.config.multimodal import MultiModalConfig
|
||||
|
||||
if not is_flashinfer_cudnn_fp8_prefill_attn_supported():
|
||||
pytest.skip("FlashInfer cuDNN FP8 prefill attention not supported")
|
||||
|
||||
mm_config = MultiModalConfig(mm_encoder_attn_dtype="fp8")
|
||||
vllm_config = VllmConfig()
|
||||
vllm_config.model_config = SimpleNamespace(multimodal_config=mm_config)
|
||||
|
||||
# MMEncoderAttention reads torch.get_default_dtype() during init
|
||||
# to determine the output dtype. In real model loading this is bf16.
|
||||
old_dtype = torch.get_default_dtype()
|
||||
torch.set_default_dtype(torch.bfloat16)
|
||||
|
||||
with (
|
||||
set_current_vllm_config(vllm_config),
|
||||
patch(
|
||||
"vllm.model_executor.layers.attention.mm_encoder_attention"
|
||||
".get_vit_attn_backend",
|
||||
return_value=AttentionBackendEnum.FLASHINFER,
|
||||
),
|
||||
):
|
||||
yield
|
||||
|
||||
torch.set_default_dtype(old_dtype)
|
||||
|
||||
|
||||
def _build_cu_seqlens_and_meta(
|
||||
seq_len: int,
|
||||
num_heads: int,
|
||||
head_dim: int,
|
||||
fp8_padded_hidden_size: int | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""Build cu_seqlens, max_seqlen, sequence_lengths for a single sequence."""
|
||||
import numpy as np
|
||||
|
||||
from vllm.model_executor.layers.attention.mm_encoder_attention import (
|
||||
MMEncoderAttention,
|
||||
)
|
||||
|
||||
cu_seqlens_np = np.array([0, seq_len], dtype=np.int32)
|
||||
|
||||
sequence_lengths = MMEncoderAttention.maybe_compute_seq_lens(
|
||||
AttentionBackendEnum.FLASHINFER,
|
||||
cu_seqlens_np,
|
||||
torch.device("cuda"),
|
||||
)
|
||||
|
||||
max_seqlen = torch.tensor(
|
||||
MMEncoderAttention.compute_max_seqlen(
|
||||
AttentionBackendEnum.FLASHINFER, cu_seqlens_np
|
||||
),
|
||||
dtype=torch.int32,
|
||||
)
|
||||
|
||||
cu_seqlens = MMEncoderAttention.maybe_recompute_cu_seqlens(
|
||||
AttentionBackendEnum.FLASHINFER,
|
||||
cu_seqlens_np,
|
||||
num_heads * head_dim,
|
||||
1, # tp_size
|
||||
torch.device("cuda"),
|
||||
fp8_padded_hidden_size=fp8_padded_hidden_size,
|
||||
)
|
||||
|
||||
return cu_seqlens, max_seqlen, sequence_lengths
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not (HAS_TRITON and _has_flashinfer_cudnn()),
|
||||
reason="Triton and FlashInfer cuDNN required",
|
||||
)
|
||||
@pytest.mark.parametrize("head_dim", HEAD_DIMS)
|
||||
@pytest.mark.parametrize("seq_len", SEQ_LENS)
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
def test_fp8_attn_output_shape(
|
||||
head_dim: int,
|
||||
seq_len: int,
|
||||
num_heads: int,
|
||||
_fp8_attention,
|
||||
) -> None:
|
||||
"""Verify FP8 attention produces correct output shape after un-padding."""
|
||||
from vllm.model_executor.layers.attention.mm_encoder_attention import (
|
||||
MMEncoderAttention,
|
||||
)
|
||||
from vllm.utils.math_utils import round_up
|
||||
|
||||
attn = None
|
||||
with contextlib.suppress(ValueError, ImportError):
|
||||
attn = MMEncoderAttention(
|
||||
num_heads=num_heads,
|
||||
head_size=head_dim,
|
||||
prefix="visual.blocks.0.attn",
|
||||
).to("cuda")
|
||||
|
||||
if attn is None or not attn.fp8_enabled:
|
||||
pytest.skip("FP8 MMEncoderAttention not available")
|
||||
assert attn is not None # mypy narrowing
|
||||
|
||||
# FP8 always needs fp8_padded_hidden_size for correct cu_seqlens
|
||||
fp8_padded_hidden_size = num_heads * round_up(head_dim, 16)
|
||||
|
||||
cu_seqlens, max_seqlen, sequence_lengths = _build_cu_seqlens_and_meta(
|
||||
seq_len, num_heads, head_dim, fp8_padded_hidden_size=fp8_padded_hidden_size
|
||||
)
|
||||
|
||||
q = torch.randn(
|
||||
seq_len,
|
||||
num_heads,
|
||||
head_dim,
|
||||
device="cuda",
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
k = torch.randn_like(q)
|
||||
v = torch.randn_like(q)
|
||||
|
||||
output = attn._forward_flashinfer(q, k, v, cu_seqlens, max_seqlen, sequence_lengths)
|
||||
|
||||
# Output should have original head_dim (un-padded)
|
||||
assert output.shape[-1] == head_dim
|
||||
assert output.dtype == torch.bfloat16
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not (HAS_TRITON and _has_flashinfer_cudnn()),
|
||||
reason="Triton and FlashInfer cuDNN required",
|
||||
)
|
||||
@pytest.mark.parametrize("head_dim", HEAD_DIMS)
|
||||
@pytest.mark.parametrize("seq_len", SEQ_LENS)
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
def test_fp8_vs_bf16_close(
|
||||
head_dim: int, seq_len: int, num_heads: int, _fp8_attention
|
||||
) -> None:
|
||||
"""FP8 attention output should be reasonably close to BF16 baseline."""
|
||||
from vllm.model_executor.layers.attention.mm_encoder_attention import (
|
||||
MMEncoderAttention,
|
||||
)
|
||||
from vllm.utils.math_utils import round_up
|
||||
|
||||
torch.manual_seed(42)
|
||||
q = torch.randn(
|
||||
1,
|
||||
seq_len,
|
||||
num_heads,
|
||||
head_dim,
|
||||
device="cuda",
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
k = torch.randn_like(q)
|
||||
v = torch.randn_like(q)
|
||||
|
||||
# FP8 path
|
||||
attn_fp8 = None
|
||||
with contextlib.suppress(ValueError, ImportError):
|
||||
attn_fp8 = MMEncoderAttention(
|
||||
num_heads=num_heads,
|
||||
head_size=head_dim,
|
||||
prefix="visual.blocks.0.attn",
|
||||
).to("cuda")
|
||||
|
||||
if attn_fp8 is None or not attn_fp8.fp8_enabled:
|
||||
pytest.skip("FP8 MMEncoderAttention not available")
|
||||
assert attn_fp8 is not None # mypy narrowing
|
||||
|
||||
fp8_padded_hidden_size = num_heads * round_up(head_dim, 16)
|
||||
cu_seqlens, max_seqlen, seq_lengths = _build_cu_seqlens_and_meta(
|
||||
seq_len,
|
||||
num_heads,
|
||||
head_dim,
|
||||
fp8_padded_hidden_size=fp8_padded_hidden_size,
|
||||
)
|
||||
|
||||
out_fp8 = attn_fp8._forward_flashinfer(
|
||||
q.clone(),
|
||||
k.clone(),
|
||||
v.clone(),
|
||||
cu_seqlens,
|
||||
max_seqlen,
|
||||
seq_lengths,
|
||||
)
|
||||
|
||||
# BF16 baseline (create non-FP8 attention by using scale=attn_fp8.scale
|
||||
# and calling the wrapper directly without FP8 quantization)
|
||||
from vllm.model_executor.layers.attention.mm_encoder_attention import (
|
||||
_get_flashinfer_workspace_buffer,
|
||||
)
|
||||
from vllm.v1.attention.ops.vit_attn_wrappers import (
|
||||
vit_flashinfer_wrapper,
|
||||
)
|
||||
|
||||
out_bf16 = vit_flashinfer_wrapper(
|
||||
q=q.clone(),
|
||||
k=k.clone(),
|
||||
v=v.clone(),
|
||||
scale=attn_fp8.scale,
|
||||
workspace_buffer=_get_flashinfer_workspace_buffer(),
|
||||
cu_seqlens=cu_seqlens,
|
||||
max_seqlen=max_seqlen,
|
||||
sequence_lengths=seq_lengths,
|
||||
)
|
||||
|
||||
out_fp8_f = out_fp8.float()
|
||||
out_bf16_f = out_bf16.float()
|
||||
|
||||
abs_diff = (out_fp8_f - out_bf16_f).abs()
|
||||
abs_diff_flat = abs_diff.flatten()
|
||||
|
||||
# Relative diff (avoid division by zero)
|
||||
denom = out_bf16_f.abs().clamp(min=1e-6)
|
||||
rel_diff_flat = (abs_diff / denom).flatten()
|
||||
|
||||
cosine_sim = torch.nn.functional.cosine_similarity(
|
||||
out_fp8_f.flatten().unsqueeze(0),
|
||||
out_bf16_f.flatten().unsqueeze(0),
|
||||
).item()
|
||||
|
||||
pcts = [50, 90, 95, 99, 99.9]
|
||||
abs_pct = {p: torch.quantile(abs_diff_flat, p / 100).item() for p in pcts}
|
||||
rel_pct = {p: torch.quantile(rel_diff_flat, p / 100).item() for p in pcts}
|
||||
|
||||
print(f"\nFP8 vs BF16 (head_dim={head_dim}, seq_len={seq_len}):")
|
||||
print(f" cosine_sim={cosine_sim:.6f}")
|
||||
print(
|
||||
f" abs_diff: max={abs_diff_flat.max().item():.6f}, "
|
||||
f"mean={abs_diff_flat.mean().item():.6f}, "
|
||||
+ ", ".join(f"p{p}={abs_pct[p]:.6f}" for p in pcts)
|
||||
)
|
||||
print(
|
||||
f" rel_diff: max={rel_diff_flat.max().item():.6f}, "
|
||||
f"mean={rel_diff_flat.mean().item():.6f}, "
|
||||
+ ", ".join(f"p{p}={rel_pct[p]:.6f}" for p in pcts)
|
||||
)
|
||||
|
||||
assert abs_diff_flat.max().item() < 0.3, (
|
||||
f"FP8 vs BF16 max abs diff too large: {abs_diff_flat.max().item()}"
|
||||
)
|
||||
assert abs_diff_flat.mean().item() < 0.03, (
|
||||
f"FP8 vs BF16 mean abs diff too large: {abs_diff_flat.mean().item()}"
|
||||
)
|
||||
assert cosine_sim > 0.99, f"Cosine similarity too low: {cosine_sim:.6f}"
|
||||
@@ -0,0 +1,124 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for the stride-aware FP8 quantization kernel with head_dim padding."""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.triton_utils import HAS_TRITON
|
||||
|
||||
if HAS_TRITON:
|
||||
from vllm.kernels.triton.qkv_padded_fp8_quant import (
|
||||
quantize_fp8_pad_head_dim_triton,
|
||||
)
|
||||
|
||||
HEAD_DIMS = [72, 80, 128]
|
||||
SEQ_LENS = [64, 256]
|
||||
NUM_HEADS = [16]
|
||||
SCALES = [0.01, 0.1, 1.0]
|
||||
|
||||
|
||||
def _naive_fp8_quantize(
|
||||
tensor: torch.Tensor, scale: torch.Tensor, skip_scale: bool
|
||||
) -> torch.Tensor:
|
||||
"""Reference FP8 quantization in PyTorch."""
|
||||
fp8_dtype = current_platform.fp8_dtype()
|
||||
fp8_max = torch.finfo(fp8_dtype).max
|
||||
fp8_min = -fp8_max
|
||||
|
||||
x = tensor.float()
|
||||
if not skip_scale:
|
||||
x = x / scale.item()
|
||||
x = x.clamp(fp8_min, fp8_max)
|
||||
return x.to(fp8_dtype)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not HAS_TRITON, reason="Triton not available")
|
||||
@pytest.mark.parametrize("head_dim", HEAD_DIMS)
|
||||
@pytest.mark.parametrize("seq_len", SEQ_LENS)
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
@pytest.mark.parametrize("scale_val", SCALES)
|
||||
def test_quantize_contiguous(
|
||||
head_dim: int, seq_len: int, num_heads: int, scale_val: float
|
||||
) -> None:
|
||||
"""Test quantization of contiguous 3D tensors."""
|
||||
torch.manual_seed(42)
|
||||
tensor = torch.randn(
|
||||
seq_len, num_heads, head_dim, device="cuda", dtype=torch.bfloat16
|
||||
)
|
||||
scale = torch.tensor([scale_val], dtype=torch.float32, device="cuda").view(
|
||||
1, 1, 1, 1
|
||||
)
|
||||
|
||||
result = quantize_fp8_pad_head_dim_triton(tensor, scale)
|
||||
|
||||
padded_dim = (head_dim + 15) // 16 * 16
|
||||
assert result.shape == (seq_len, num_heads, padded_dim)
|
||||
assert result.is_contiguous()
|
||||
assert result.dtype == current_platform.fp8_dtype()
|
||||
|
||||
# Compare unpadded portion against reference
|
||||
ref = _naive_fp8_quantize(tensor, scale, skip_scale=False)
|
||||
torch.testing.assert_close(result[:, :, :head_dim].float(), ref.float())
|
||||
|
||||
# Padded region should be zero
|
||||
if padded_dim > head_dim:
|
||||
assert (result[:, :, head_dim:].float() == 0).all()
|
||||
|
||||
|
||||
@pytest.mark.skipif(not HAS_TRITON, reason="Triton not available")
|
||||
@pytest.mark.parametrize("head_dim", [72, 80])
|
||||
def test_quantize_non_contiguous(head_dim: int) -> None:
|
||||
"""Test quantization from non-contiguous QKV views (interleaved buffer)."""
|
||||
seq_len, num_heads = 64, 16
|
||||
# Simulate interleaved QKV buffer: shape (seq_len, 3 * num_heads, head_dim)
|
||||
qkv = torch.randn(
|
||||
seq_len, 3 * num_heads, head_dim, device="cuda", dtype=torch.bfloat16
|
||||
)
|
||||
# Q is every 3rd head slice - non-contiguous view
|
||||
q = qkv[:, 0::3, :]
|
||||
assert not q.is_contiguous()
|
||||
|
||||
scale = torch.tensor([0.1], dtype=torch.float32, device="cuda").view(1, 1, 1, 1)
|
||||
result = quantize_fp8_pad_head_dim_triton(q, scale)
|
||||
|
||||
padded_dim = (head_dim + 15) // 16 * 16
|
||||
assert result.shape == (seq_len, num_heads, padded_dim)
|
||||
assert result.is_contiguous()
|
||||
|
||||
# Compare against contiguous reference
|
||||
ref = _naive_fp8_quantize(q.contiguous(), scale, skip_scale=False)
|
||||
torch.testing.assert_close(result[:, :, :head_dim].float(), ref.float())
|
||||
|
||||
|
||||
@pytest.mark.skipif(not HAS_TRITON, reason="Triton not available")
|
||||
def test_skip_scale() -> None:
|
||||
"""Test skip_scale=True produces cast-only output (no division)."""
|
||||
seq_len, num_heads, head_dim = 32, 8, 80
|
||||
tensor = torch.randn(
|
||||
seq_len, num_heads, head_dim, device="cuda", dtype=torch.bfloat16
|
||||
)
|
||||
scale = torch.tensor([0.5], dtype=torch.float32, device="cuda").view(1, 1, 1, 1)
|
||||
|
||||
result_skip = quantize_fp8_pad_head_dim_triton(tensor, scale, skip_scale=True)
|
||||
result_noskip = quantize_fp8_pad_head_dim_triton(tensor, scale, skip_scale=False)
|
||||
|
||||
# skip_scale should just cast, not divide
|
||||
ref_cast = _naive_fp8_quantize(tensor, scale, skip_scale=True)
|
||||
torch.testing.assert_close(result_skip[:, :, :head_dim].float(), ref_cast.float())
|
||||
|
||||
# With scale != 1.0, skip and no-skip should differ
|
||||
assert not torch.equal(result_skip.float(), result_noskip.float())
|
||||
|
||||
|
||||
@pytest.mark.skipif(not HAS_TRITON, reason="Triton not available")
|
||||
def test_4d_input() -> None:
|
||||
"""Test that 4D input (B, S, H, D) is handled correctly."""
|
||||
B, S, H, D = 2, 32, 8, 72
|
||||
tensor = torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16)
|
||||
scale = torch.tensor([0.1], dtype=torch.float32, device="cuda").view(1, 1, 1, 1)
|
||||
|
||||
result = quantize_fp8_pad_head_dim_triton(tensor, scale)
|
||||
padded_dim = (D + 15) // 16 * 16
|
||||
assert result.shape == (B, S, H, padded_dim)
|
||||
@@ -0,0 +1,251 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for FP8 scaling (dynamic and static) in MMEncoderAttention."""
|
||||
|
||||
import contextlib
|
||||
import json
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.model_executor.layers.attention.mm_encoder_attention import (
|
||||
_FP8_AMAX_HISTORY_LEN,
|
||||
_FP8_MAX,
|
||||
)
|
||||
from vllm.utils.flashinfer import (
|
||||
is_flashinfer_cudnn_fp8_prefill_attn_supported,
|
||||
)
|
||||
|
||||
LAYER_0 = "visual.blocks.0.attn.attn"
|
||||
LAYER_1 = "visual.blocks.1.attn.attn"
|
||||
NUM_HEADS = 16
|
||||
HEAD_DIM = 72
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _build_attention(mm_config):
|
||||
"""Yield an MMEncoderAttention with the given multimodal config.
|
||||
|
||||
The VllmConfig context stays active while the test runs so that
|
||||
``get_multimodal_config()`` calls during the forward path resolve. Also
|
||||
invokes ``process_weights_after_loading`` to simulate the model loader's
|
||||
auto-scan. Yields ``None`` if FlashInfer cuDNN is not available.
|
||||
"""
|
||||
from vllm.config import VllmConfig, set_current_vllm_config
|
||||
from vllm.model_executor.layers.attention.mm_encoder_attention import (
|
||||
MMEncoderAttention,
|
||||
)
|
||||
from vllm.v1.attention.backends.registry import AttentionBackendEnum
|
||||
|
||||
if not is_flashinfer_cudnn_fp8_prefill_attn_supported():
|
||||
yield None
|
||||
return
|
||||
|
||||
vllm_config = VllmConfig()
|
||||
vllm_config.model_config = SimpleNamespace(multimodal_config=mm_config)
|
||||
|
||||
with (
|
||||
set_current_vllm_config(vllm_config),
|
||||
patch(
|
||||
"vllm.model_executor.layers.attention.mm_encoder_attention"
|
||||
".get_vit_attn_backend",
|
||||
return_value=AttentionBackendEnum.FLASHINFER,
|
||||
),
|
||||
):
|
||||
attn = MMEncoderAttention(
|
||||
num_heads=NUM_HEADS,
|
||||
head_size=HEAD_DIM,
|
||||
prefix=LAYER_0,
|
||||
)
|
||||
attn.process_weights_after_loading(torch.bfloat16)
|
||||
yield attn
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def _make_attention():
|
||||
"""Create an MMEncoderAttention with dynamic FP8 scaling."""
|
||||
from vllm.config.multimodal import MultiModalConfig
|
||||
|
||||
with _build_attention(MultiModalConfig(mm_encoder_attn_dtype="fp8")) as attn:
|
||||
yield attn
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def _make_static_attention(tmp_path):
|
||||
"""Create an MMEncoderAttention with static FP8 scales from a file."""
|
||||
from vllm.config.multimodal import MultiModalConfig
|
||||
|
||||
scale_file = tmp_path / "scales.json"
|
||||
scale_file.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
LAYER_0: {"q": 224.0, "k": 198.0, "v": 210.0},
|
||||
LAYER_1: {"q": 100.0, "k": 110.0, "v": 120.0},
|
||||
}
|
||||
)
|
||||
)
|
||||
with _build_attention(
|
||||
MultiModalConfig(
|
||||
mm_encoder_attn_dtype="fp8",
|
||||
mm_encoder_fp8_scale_path=str(scale_file),
|
||||
)
|
||||
) as attn:
|
||||
yield attn
|
||||
|
||||
|
||||
def test_dynamic_scaling_updates_scales(_make_attention) -> None:
|
||||
"""Verify that _record_amax_and_update_scales updates scale buffers."""
|
||||
attn = _make_attention
|
||||
if attn is None or not attn.fp8_enabled:
|
||||
pytest.skip("FP8 attention not available (FlashInfer backend required)")
|
||||
|
||||
attn = attn.to("cuda")
|
||||
|
||||
S, H, D = 32, NUM_HEADS, HEAD_DIM
|
||||
q = torch.full((S, H, D), 2.0, device="cuda", dtype=torch.bfloat16)
|
||||
k = torch.full((S, H, D), 3.0, device="cuda", dtype=torch.bfloat16)
|
||||
v = torch.full((S, H, D), 4.0, device="cuda", dtype=torch.bfloat16)
|
||||
|
||||
attn._record_amax_and_update_scales(q, k, v)
|
||||
|
||||
expected_q_scale = 2.0 / _FP8_MAX
|
||||
expected_k_scale = 3.0 / _FP8_MAX
|
||||
expected_v_scale = 4.0 / _FP8_MAX
|
||||
|
||||
torch.testing.assert_close(attn._fp8_q_scale.item(), expected_q_scale)
|
||||
torch.testing.assert_close(attn._fp8_k_scale.item(), expected_k_scale)
|
||||
torch.testing.assert_close(attn._fp8_v_scale.item(), expected_v_scale)
|
||||
|
||||
|
||||
def test_circular_buffer_wraps(_make_attention) -> None:
|
||||
"""Verify the amax circular buffer wraps at HISTORY_LEN."""
|
||||
attn = _make_attention
|
||||
if attn is None or not attn.fp8_enabled:
|
||||
pytest.skip("FP8 attention not available (FlashInfer backend required)")
|
||||
|
||||
attn = attn.to("cuda")
|
||||
S, H, D = 16, NUM_HEADS, HEAD_DIM
|
||||
|
||||
for i in range(_FP8_AMAX_HISTORY_LEN + 2):
|
||||
mag = float(i + 1)
|
||||
q = torch.full((S, H, D), mag, device="cuda", dtype=torch.bfloat16)
|
||||
k = torch.full((S, H, D), mag, device="cuda", dtype=torch.bfloat16)
|
||||
v = torch.full((S, H, D), mag, device="cuda", dtype=torch.bfloat16)
|
||||
attn._record_amax_and_update_scales(q, k, v)
|
||||
|
||||
assert attn._fp8_amax_pos == 2
|
||||
|
||||
expected_max = float(_FP8_AMAX_HISTORY_LEN + 2)
|
||||
expected_scale = expected_max / _FP8_MAX
|
||||
torch.testing.assert_close(attn._fp8_q_scale.item(), expected_scale)
|
||||
|
||||
|
||||
def test_static_scales_loaded(_make_static_attention) -> None:
|
||||
"""Verify static scales are loaded from the JSON file."""
|
||||
attn = _make_static_attention
|
||||
if attn is None or not attn.fp8_enabled:
|
||||
pytest.skip("FP8 attention not available (FlashInfer backend required)")
|
||||
|
||||
assert attn.fp8_enabled
|
||||
assert not attn._fp8_dynamic_scale
|
||||
|
||||
# Layer 0 scales (the layer this attention was created with).
|
||||
assert attn._fp8_q_scale.item() == 224.0
|
||||
assert attn._fp8_k_scale.item() == 198.0
|
||||
assert attn._fp8_v_scale.item() == 210.0
|
||||
|
||||
assert not attn.skip_scale_q
|
||||
assert not attn.skip_scale_k
|
||||
assert not attn.skip_scale_v
|
||||
|
||||
# No amax history buffers for static scaling.
|
||||
assert not hasattr(attn, "_fp8_q_amax")
|
||||
|
||||
|
||||
def test_static_scales_missing_layer(tmp_path) -> None:
|
||||
"""Verify error when requested layer is not in the scale file."""
|
||||
from vllm.config import VllmConfig, set_current_vllm_config
|
||||
from vllm.config.multimodal import MultiModalConfig
|
||||
from vllm.v1.attention.backends.registry import AttentionBackendEnum
|
||||
|
||||
if not is_flashinfer_cudnn_fp8_prefill_attn_supported():
|
||||
pytest.skip("FlashInfer cuDNN not available")
|
||||
|
||||
scale_file = tmp_path / "wrong_layer.json"
|
||||
scale_file.write_text(
|
||||
json.dumps({"visual.blocks.99.attn": {"q": 1.0, "k": 1.0, "v": 1.0}})
|
||||
)
|
||||
mm_config = MultiModalConfig(
|
||||
mm_encoder_attn_dtype="fp8",
|
||||
mm_encoder_fp8_scale_path=str(scale_file),
|
||||
)
|
||||
vllm_config = VllmConfig()
|
||||
vllm_config.model_config = SimpleNamespace(multimodal_config=mm_config)
|
||||
|
||||
from vllm.model_executor.layers.attention.mm_encoder_attention import (
|
||||
MMEncoderAttention,
|
||||
)
|
||||
|
||||
with (
|
||||
set_current_vllm_config(vllm_config),
|
||||
patch(
|
||||
"vllm.model_executor.layers.attention.mm_encoder_attention"
|
||||
".get_vit_attn_backend",
|
||||
return_value=AttentionBackendEnum.FLASHINFER,
|
||||
),
|
||||
):
|
||||
attn = MMEncoderAttention(
|
||||
num_heads=NUM_HEADS,
|
||||
head_size=HEAD_DIM,
|
||||
prefix=LAYER_0,
|
||||
)
|
||||
with pytest.raises(ValueError, match="scales not found for layer"):
|
||||
attn.process_weights_after_loading(torch.bfloat16)
|
||||
|
||||
|
||||
def test_dynamic_scales_auto_save(tmp_path) -> None:
|
||||
"""Verify scales are saved to disk after the amax buffer fills."""
|
||||
import vllm.model_executor.layers.attention.mm_encoder_attention as _mod
|
||||
from vllm.config.multimodal import MultiModalConfig
|
||||
|
||||
if not is_flashinfer_cudnn_fp8_prefill_attn_supported():
|
||||
pytest.skip("FlashInfer cuDNN not available")
|
||||
|
||||
# Reset module-level state between runs (other tests may have left
|
||||
# state behind after triggering a save).
|
||||
_mod._fp8_scale_save_path = None
|
||||
_mod._fp8_saved_scale_refs.clear()
|
||||
|
||||
save_file = tmp_path / "auto_scales.json"
|
||||
with _build_attention(
|
||||
MultiModalConfig(
|
||||
mm_encoder_attn_dtype="fp8",
|
||||
mm_encoder_fp8_scale_save_path=str(save_file),
|
||||
)
|
||||
) as attn:
|
||||
if attn is None or not attn.fp8_enabled:
|
||||
pytest.skip("FP8 attention not available")
|
||||
|
||||
attn = attn.to("cuda")
|
||||
S, H, D = 16, NUM_HEADS, HEAD_DIM
|
||||
|
||||
# Run exactly _FP8_AMAX_HISTORY_LEN forward passes.
|
||||
for i in range(_FP8_AMAX_HISTORY_LEN):
|
||||
mag = float(i + 1)
|
||||
q = torch.full((S, H, D), mag, device="cuda", dtype=torch.bfloat16)
|
||||
k = torch.full((S, H, D), mag * 0.5, device="cuda", dtype=torch.bfloat16)
|
||||
v = torch.full((S, H, D), mag * 0.3, device="cuda", dtype=torch.bfloat16)
|
||||
attn._record_amax_and_update_scales(q, k, v)
|
||||
|
||||
# File should have been written on the 16th call (buffer wrap).
|
||||
assert save_file.is_file(), "Scale file was not saved"
|
||||
scales = json.loads(save_file.read_text())
|
||||
assert LAYER_0 in scales
|
||||
assert set(scales[LAYER_0].keys()) == {"q", "k", "v"}
|
||||
for val in scales[LAYER_0].values():
|
||||
assert isinstance(val, float) and val > 0
|
||||
|
||||
# Path is cleared after the one-shot save fires.
|
||||
assert _mod._fp8_scale_save_path is None
|
||||
@@ -6,6 +6,7 @@ from transformers import AutoModel
|
||||
|
||||
from tests.models.utils import check_embeddings_close
|
||||
from vllm import TokensPrompt
|
||||
from vllm.config import PoolerConfig
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
@@ -21,6 +22,7 @@ def test_embed_models(hf_runner, vllm_runner, model: str):
|
||||
with vllm_runner(
|
||||
model,
|
||||
runner="pooling",
|
||||
pooler_config=PoolerConfig(task="token_embed"),
|
||||
max_model_len=128,
|
||||
max_num_batched_tokens=chunk_size,
|
||||
enforce_eager=True,
|
||||
|
||||
@@ -3,7 +3,6 @@
|
||||
import httpx
|
||||
import openai
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
import torch
|
||||
|
||||
from ....utils import RemoteOpenAIServer
|
||||
@@ -25,29 +24,42 @@ sentences_2 = [
|
||||
similarity_reference = [[0.6259, 0.3474], [0.3309, 0.6734]]
|
||||
lexical_score_reference = [0.19554901123046875, 0.0]
|
||||
colbert_score_reference = [0.7797, 0.4620]
|
||||
SUPPORTED_TASKS = ["embed", "token_embed", "token_classify"]
|
||||
|
||||
|
||||
@pytest.fixture(scope="module", params=SUPPORTED_TASKS)
|
||||
def pooling_task(request):
|
||||
yield request.param
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def server():
|
||||
def server(pooling_task):
|
||||
args = [
|
||||
"--max-model-len",
|
||||
str(MAX_MODEL_LEN),
|
||||
"--hf-overrides",
|
||||
'{"architectures": ["BgeM3EmbeddingModel"]}',
|
||||
"--pooler-config.task",
|
||||
pooling_task,
|
||||
]
|
||||
|
||||
with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
|
||||
yield remote_server
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def client(server):
|
||||
async with server.get_async_client() as async_client:
|
||||
yield async_client
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bge_m3_api_server_embedding(client: openai.AsyncOpenAI):
|
||||
async def test_bge_m3_api_server_embedding(server, pooling_task):
|
||||
client = server.get_async_client()
|
||||
|
||||
if pooling_task != "embed":
|
||||
with pytest.raises(openai.InternalServerError):
|
||||
await run_client_embeddings(
|
||||
client,
|
||||
MODEL_NAME,
|
||||
sentences_1,
|
||||
)
|
||||
return
|
||||
|
||||
embeddings_list_1 = await run_client_embeddings(
|
||||
client,
|
||||
MODEL_NAME,
|
||||
@@ -117,7 +129,14 @@ def compute_lexical_matching_score(
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bge_m3_api_server_sparse_embedding(client: openai.AsyncOpenAI):
|
||||
async def test_bge_m3_api_server_sparse_embedding(server, pooling_task):
|
||||
client = server.get_async_client()
|
||||
|
||||
if pooling_task != "token_classify":
|
||||
with pytest.raises(openai.BadRequestError):
|
||||
await sparse_embeddings(client, sentences_1)
|
||||
return
|
||||
|
||||
embeddings_1 = await sparse_embeddings(client, sentences_1)
|
||||
embeddings_2 = await sparse_embeddings(client, sentences_2)
|
||||
|
||||
@@ -137,9 +156,11 @@ async def test_bge_m3_api_server_sparse_embedding(client: openai.AsyncOpenAI):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bge_m3_api_server_sparse_embedding_corner_case(
|
||||
client: openai.AsyncOpenAI,
|
||||
):
|
||||
async def test_bge_m3_api_server_sparse_embedding_corner_case(server, pooling_task):
|
||||
if pooling_task != "token_classify":
|
||||
return
|
||||
|
||||
client = server.get_async_client()
|
||||
embeddings = await sparse_embeddings(client, ["Hi"])
|
||||
assert len(embeddings) == 1
|
||||
assert 2673 in embeddings[0]
|
||||
@@ -155,7 +176,18 @@ def colbert_score(q_reps: torch.Tensor, p_reps: torch.Tensor) -> torch.Tensor:
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bge_m3_api_server_multi_vector(client: openai.AsyncOpenAI):
|
||||
async def test_bge_m3_api_server_multi_vector(server, pooling_task):
|
||||
client = server.get_async_client()
|
||||
|
||||
if pooling_task != "token_embed":
|
||||
with pytest.raises(openai.BadRequestError):
|
||||
await client.post(
|
||||
"../pooling",
|
||||
body={"model": MODEL_NAME, "input": sentences_1, "task": "token_embed"},
|
||||
cast_to=httpx.Response,
|
||||
)
|
||||
return
|
||||
|
||||
result_1 = await client.post(
|
||||
"../pooling",
|
||||
body={"model": MODEL_NAME, "input": sentences_1, "task": "token_embed"},
|
||||
|
||||
@@ -4,6 +4,7 @@ import pytest
|
||||
import torch
|
||||
|
||||
from vllm import TokensPrompt
|
||||
from vllm.config import PoolerConfig
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
@@ -20,6 +21,7 @@ def test_extract_hidden_states(hf_runner, vllm_runner, model: str):
|
||||
max_model_len=128,
|
||||
enforce_eager=True,
|
||||
runner="pooling",
|
||||
pooler_config=PoolerConfig(task="token_embed"),
|
||||
enable_prefix_caching=True,
|
||||
) as vllm_model:
|
||||
pooling_outputs = vllm_model.llm.encode(
|
||||
@@ -44,14 +46,3 @@ def test_extract_hidden_states(hf_runner, vllm_runner, model: str):
|
||||
assert len(output.prompt_token_ids) == n
|
||||
assert len(output.outputs.data) == n
|
||||
assert output.num_cached_tokens == 0
|
||||
|
||||
# skip_reading_prefix_cache can still write to cache
|
||||
# to accelerate following requests
|
||||
pooling_outputs = vllm_model.llm.encode(
|
||||
[TokensPrompt(prompt_token_ids=t) for t in token_prompts],
|
||||
pooling_task="embed",
|
||||
)
|
||||
|
||||
for n, output in zip(n_prompt_tokens, pooling_outputs):
|
||||
assert len(output.prompt_token_ids) == n
|
||||
assert output.num_cached_tokens > 0
|
||||
|
||||
@@ -5,6 +5,7 @@ import torch
|
||||
from transformers import AutoModel
|
||||
|
||||
from tests.models.utils import check_embeddings_close
|
||||
from vllm.config import PoolerConfig
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
@@ -17,6 +18,7 @@ def test_embed_models(hf_runner, vllm_runner, example_prompts, model: str, dtype
|
||||
with vllm_runner(
|
||||
model,
|
||||
runner="pooling",
|
||||
pooler_config=PoolerConfig(task="token_embed"),
|
||||
max_model_len=None,
|
||||
) as vllm_model:
|
||||
vllm_outputs = vllm_model.token_embed(example_prompts)
|
||||
|
||||
@@ -146,7 +146,7 @@ def test_multi_vector_retrieval_models_using_normalize(
|
||||
model,
|
||||
max_model_len=512,
|
||||
dtype=dtype,
|
||||
pooler_config=PoolerConfig(use_activation=False),
|
||||
pooler_config=PoolerConfig(use_activation=False, task="token_embed"),
|
||||
) as vllm_model:
|
||||
wo_normalize = vllm_model.token_embed(example_prompts)
|
||||
|
||||
@@ -154,7 +154,7 @@ def test_multi_vector_retrieval_models_using_normalize(
|
||||
model,
|
||||
max_model_len=512,
|
||||
dtype=dtype,
|
||||
pooler_config=PoolerConfig(use_activation=True),
|
||||
pooler_config=PoolerConfig(use_activation=True, task="token_embed"),
|
||||
) as vllm_model:
|
||||
w_normalize = vllm_model.token_embed(example_prompts)
|
||||
|
||||
|
||||
@@ -594,6 +594,9 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
|
||||
"MiMoV2FlashForCausalLM": _HfExamplesInfo(
|
||||
"XiaomiMiMo/MiMo-V2-Flash", trust_remote_code=True
|
||||
),
|
||||
"MiMoV2ProForCausalLM": _HfExamplesInfo(
|
||||
"XiaomiMiMo/MiMo-V2.5-Pro", trust_remote_code=True, is_available_online=False
|
||||
),
|
||||
"Dots1ForCausalLM": _HfExamplesInfo("rednote-hilab/dots.llm1.inst"),
|
||||
}
|
||||
|
||||
@@ -961,6 +964,18 @@ _MULTIMODAL_EXAMPLE_MODELS = {
|
||||
"PerceptronAI/Isaac-0.1",
|
||||
trust_remote_code=True,
|
||||
extras={"0.2-2B-Preview": "PerceptronAI/Isaac-0.2-2B-Preview"},
|
||||
max_transformers_version="4.57",
|
||||
transformers_version_reason={
|
||||
"vllm": (
|
||||
"Custom Isaac code is not compatible with Transformers v5. "
|
||||
"The model should be upstreamed to Transformers for "
|
||||
"long-term support."
|
||||
),
|
||||
"hf": (
|
||||
"Isaac's remote model and processor code import or configure "
|
||||
"APIs that changed in Transformers v5."
|
||||
),
|
||||
},
|
||||
),
|
||||
"InternS1ForConditionalGeneration": _HfExamplesInfo(
|
||||
"internlm/Intern-S1",
|
||||
@@ -1057,6 +1072,9 @@ _MULTIMODAL_EXAMPLE_MODELS = {
|
||||
"MiDashengLMModel": _HfExamplesInfo(
|
||||
"mispeech/midashenglm-7b", trust_remote_code=True
|
||||
),
|
||||
"MiMoV2OmniForCausalLM": _HfExamplesInfo(
|
||||
"XiaomiMiMo/MiMo-V2.5-Omni", trust_remote_code=True, is_available_online=False
|
||||
),
|
||||
"MiniCPMO": _HfExamplesInfo(
|
||||
"openbmb/MiniCPM-o-2_6",
|
||||
trust_remote_code=True,
|
||||
@@ -1540,6 +1558,18 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
|
||||
trust_remote_code=True,
|
||||
speculative_model="XiaomiMiMo/MiMo-7B-RL",
|
||||
),
|
||||
"MiMoV2MTPModel": _HfExamplesInfo(
|
||||
"XiaomiMiMo/MiMo-V2.5-Pro",
|
||||
trust_remote_code=True,
|
||||
speculative_model="XiaomiMiMo/MiMo-V2.5-Pro",
|
||||
is_available_online=False,
|
||||
),
|
||||
"MiMoV2OmniMTPModel": _HfExamplesInfo(
|
||||
"XiaomiMiMo/MiMo-V2.5-Omni",
|
||||
trust_remote_code=True,
|
||||
speculative_model="XiaomiMiMo/MiMo-V2.5-Omni",
|
||||
is_available_online=False,
|
||||
),
|
||||
"NemotronHMTPModel": _HfExamplesInfo(
|
||||
"nvidia/Nemotron-Super-Placeholder",
|
||||
speculative_model="nvidia/Nemotron-Super-Placeholder",
|
||||
|
||||
@@ -5,6 +5,8 @@ Unit tests for MultiModalRegistry.supports_multimodal_inputs and
|
||||
Qwen2.5-VL visual component loading behavior.
|
||||
"""
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.multimodal import MULTIMODAL_REGISTRY
|
||||
@@ -32,3 +34,17 @@ def test_supports_multimodal_inputs(model_id, limit_mm_per_prompt, expected):
|
||||
limit_mm_per_prompt=limit_mm_per_prompt,
|
||||
)
|
||||
assert MULTIMODAL_REGISTRY.supports_multimodal_inputs(ctx.model_config) is expected
|
||||
|
||||
|
||||
def test_create_processor_error_uses_served_model_name():
|
||||
model_config = SimpleNamespace(
|
||||
is_multimodal_model=False,
|
||||
model="/path/to/model/weights",
|
||||
served_model_name="friendly-model-name",
|
||||
)
|
||||
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match="friendly-model-name is not a multimodal model",
|
||||
):
|
||||
MULTIMODAL_REGISTRY.create_processor(model_config)
|
||||
|
||||
@@ -0,0 +1,185 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for Cutlass W4A16 (Machete) kernel on Hopper.
|
||||
|
||||
Verifies that W4A16 quantized models loaded through vllm select the
|
||||
MacheteLinearKernel on sm_90 GPUs, that weights are correctly repacked,
|
||||
and that inference produces valid output.
|
||||
|
||||
Run `pytest tests/quantization/test_cutlass_w4a16.py`.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
if not current_platform.has_device_capability(90):
|
||||
pytest.skip(
|
||||
"Machete W4A16 requires Hopper (sm_90).",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
from vllm.model_executor.kernels.linear import (
|
||||
MPLinearLayerConfig,
|
||||
choose_mp_linear_kernel,
|
||||
)
|
||||
from vllm.model_executor.kernels.linear.mixed_precision import (
|
||||
MacheteLinearKernel,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501
|
||||
CompressedTensorsLinearMethod,
|
||||
CompressedTensorsWNA16,
|
||||
)
|
||||
from vllm.scalar_type import scalar_types
|
||||
|
||||
|
||||
@pytest.fixture(scope="function", autouse=True)
|
||||
def enable_pickle(monkeypatch):
|
||||
"""`LLM.apply_model` requires pickling a function."""
|
||||
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"act_type,weight_type,group_size,zero_points",
|
||||
[
|
||||
(torch.float16, scalar_types.uint4b8, 128, False),
|
||||
(torch.bfloat16, scalar_types.uint4b8, 128, False),
|
||||
(torch.float16, scalar_types.uint4, 128, True),
|
||||
(torch.float16, scalar_types.uint4b8, -1, False),
|
||||
],
|
||||
ids=[
|
||||
"fp16-gptq-g128",
|
||||
"bf16-gptq-g128",
|
||||
"fp16-awq-g128",
|
||||
"fp16-channelwise",
|
||||
],
|
||||
)
|
||||
def test_machete_kernel_selected(act_type, weight_type, group_size, zero_points):
|
||||
"""Verify choose_mp_linear_kernel picks MacheteLinearKernel."""
|
||||
config = MPLinearLayerConfig(
|
||||
full_weight_shape=(4096, 4096),
|
||||
partition_weight_shape=(4096, 4096),
|
||||
act_type=act_type,
|
||||
weight_type=weight_type,
|
||||
group_size=group_size,
|
||||
zero_points=zero_points,
|
||||
has_g_idx=False,
|
||||
)
|
||||
kernel = choose_mp_linear_kernel(config)
|
||||
assert kernel is MacheteLinearKernel, (
|
||||
f"Expected MacheteLinearKernel, got {kernel.__name__}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"full_shape,part_shape,weight_type,group_size,has_g_idx,expected_reason",
|
||||
[
|
||||
((4096, 4096), (2048, 4096), scalar_types.uint4b8, 128, True, "Act reordering"),
|
||||
(
|
||||
(4096, 4096),
|
||||
(4096, 4096),
|
||||
scalar_types.float6_e3m2f,
|
||||
128,
|
||||
False,
|
||||
"Quant type",
|
||||
),
|
||||
((4096, 4096), (4096, 4096), scalar_types.uint4b8, 32, False, "Group size"),
|
||||
],
|
||||
ids=["partitioned-g_idx", "unsupported-quant-type", "unsupported-group-size"],
|
||||
)
|
||||
def test_machete_rejects_invalid_config(
|
||||
full_shape, part_shape, weight_type, group_size, has_g_idx, expected_reason
|
||||
):
|
||||
"""Verify Machete rejects unsupported configurations."""
|
||||
config = MPLinearLayerConfig(
|
||||
full_weight_shape=full_shape,
|
||||
partition_weight_shape=part_shape,
|
||||
act_type=torch.float16,
|
||||
weight_type=weight_type,
|
||||
group_size=group_size,
|
||||
zero_points=False,
|
||||
has_g_idx=has_g_idx,
|
||||
)
|
||||
can_impl, reason = MacheteLinearKernel.can_implement(config)
|
||||
assert not can_impl
|
||||
assert expected_reason in reason
|
||||
|
||||
|
||||
def test_kernel_selection_with_disabled_machete(monkeypatch):
|
||||
"""Verify kernel selection falls back when Machete is disabled."""
|
||||
monkeypatch.setattr("vllm.envs.VLLM_DISABLED_KERNELS", ["MacheteLinearKernel"])
|
||||
|
||||
config = MPLinearLayerConfig(
|
||||
full_weight_shape=(4096, 4096),
|
||||
partition_weight_shape=(4096, 4096),
|
||||
act_type=torch.float16,
|
||||
weight_type=scalar_types.uint4b8,
|
||||
group_size=128,
|
||||
zero_points=False,
|
||||
has_g_idx=False,
|
||||
)
|
||||
kernel = choose_mp_linear_kernel(config)
|
||||
assert kernel is not MacheteLinearKernel, "MacheteLinearKernel should be disabled"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model_name",
|
||||
[
|
||||
"nm-testing/tinyllama-oneshot-w4a16-channel-v2",
|
||||
"nm-testing/TinyLlama-1.1B-Chat-v1.0-W4A16-G128-Asym-Updated-ActOrder",
|
||||
],
|
||||
)
|
||||
def test_w4a16_machete_e2e(vllm_runner, model_name):
|
||||
"""Load a W4A16 model, verify Machete kernel is used, and generate."""
|
||||
with vllm_runner(model_name, enforce_eager=True, gpu_memory_utilization=0.5) as llm:
|
||||
|
||||
def check_model(model):
|
||||
layer = model.model.layers[0]
|
||||
qkv_proj = layer.self_attn.qkv_proj
|
||||
|
||||
assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod)
|
||||
assert isinstance(qkv_proj.scheme, CompressedTensorsWNA16)
|
||||
assert isinstance(qkv_proj.scheme.kernel, MacheteLinearKernel), (
|
||||
f"Expected MacheteLinearKernel on Hopper, "
|
||||
f"got {type(qkv_proj.scheme.kernel).__name__}"
|
||||
)
|
||||
|
||||
assert hasattr(qkv_proj, "weight_packed")
|
||||
assert hasattr(qkv_proj, "weight_scale")
|
||||
assert qkv_proj.weight_packed.dtype == torch.int32
|
||||
|
||||
llm.apply_model(check_model)
|
||||
|
||||
output = llm.generate_greedy("Hello my name is", max_tokens=10)
|
||||
assert output
|
||||
assert len(output[0][1]) > 0
|
||||
|
||||
|
||||
def test_w4a16_machete_bfloat16_deterministic(vllm_runner):
|
||||
"""Verify Machete works with bf16 activations and is deterministic."""
|
||||
model_name = "nm-testing/tinyllama-oneshot-w4a16-channel-v2"
|
||||
prompt = "The capital of France is"
|
||||
|
||||
with vllm_runner(
|
||||
model_name,
|
||||
enforce_eager=True,
|
||||
dtype="bfloat16",
|
||||
gpu_memory_utilization=0.5,
|
||||
) as llm:
|
||||
|
||||
def check_kernel_type(model):
|
||||
layer = model.model.layers[0]
|
||||
scheme = layer.self_attn.qkv_proj.scheme
|
||||
assert isinstance(scheme.kernel, MacheteLinearKernel), (
|
||||
f"Expected MacheteLinearKernel with bf16, "
|
||||
f"got {type(scheme.kernel).__name__}"
|
||||
)
|
||||
|
||||
llm.apply_model(check_kernel_type)
|
||||
|
||||
out1 = llm.generate_greedy(prompt, max_tokens=10)
|
||||
out2 = llm.generate_greedy(prompt, max_tokens=10)
|
||||
assert out1[0][1] == out2[0][1], (
|
||||
f"Non-deterministic: '{out1[0][1]}' vs '{out2[0][1]}'"
|
||||
)
|
||||
+13
-1
@@ -326,6 +326,10 @@ class ModelConfig:
|
||||
mm_encoder_only: InitVar[bool | None] = None
|
||||
mm_encoder_tp_mode: InitVar[MMEncoderTPMode | None] = None
|
||||
mm_encoder_attn_backend: InitVar[AttentionBackendEnum | str | None] = None
|
||||
mm_encoder_attn_dtype: InitVar[str | None] = None
|
||||
mm_encoder_fp8_scale_path: InitVar[str | None] = None
|
||||
mm_encoder_fp8_scale_save_path: InitVar[str | None] = None
|
||||
mm_encoder_fp8_scale_save_margin: InitVar[float | None] = None
|
||||
interleave_mm_strings: InitVar[bool | None] = None
|
||||
skip_mm_profiling: InitVar[bool | None] = None
|
||||
video_pruning_rate: InitVar[float | None] = None
|
||||
@@ -447,6 +451,10 @@ class ModelConfig:
|
||||
mm_encoder_only: bool | None,
|
||||
mm_encoder_tp_mode: MMEncoderTPMode | None,
|
||||
mm_encoder_attn_backend: AttentionBackendEnum | str | None,
|
||||
mm_encoder_attn_dtype: str | None,
|
||||
mm_encoder_fp8_scale_path: str | None,
|
||||
mm_encoder_fp8_scale_save_path: str | None,
|
||||
mm_encoder_fp8_scale_save_margin: float | None,
|
||||
interleave_mm_strings: bool | None,
|
||||
skip_mm_profiling: bool | None,
|
||||
video_pruning_rate: float | None,
|
||||
@@ -513,6 +521,7 @@ class ModelConfig:
|
||||
if dict_overrides:
|
||||
self._apply_dict_overrides(hf_config, dict_overrides)
|
||||
self.hf_text_config = get_hf_text_config(self.hf_config)
|
||||
self.model_arch_config = self.get_model_arch_config()
|
||||
self.attention_chunk_size = getattr(
|
||||
self.hf_text_config, "attention_chunk_size", None
|
||||
)
|
||||
@@ -520,7 +529,6 @@ class ModelConfig:
|
||||
self.hf_image_processor_config = get_hf_image_processor_config(
|
||||
self.model, hf_token=self.hf_token, revision=self.revision
|
||||
)
|
||||
self.model_arch_config = self.get_model_arch_config()
|
||||
|
||||
architectures = self.architectures
|
||||
registry = self.registry
|
||||
@@ -643,6 +651,10 @@ class ModelConfig:
|
||||
mm_encoder_only=mm_encoder_only,
|
||||
mm_encoder_tp_mode=mm_encoder_tp_mode,
|
||||
mm_encoder_attn_backend=mm_encoder_attn_backend,
|
||||
mm_encoder_attn_dtype=mm_encoder_attn_dtype,
|
||||
mm_encoder_fp8_scale_path=mm_encoder_fp8_scale_path,
|
||||
mm_encoder_fp8_scale_save_path=mm_encoder_fp8_scale_save_path,
|
||||
mm_encoder_fp8_scale_save_margin=mm_encoder_fp8_scale_save_margin,
|
||||
interleave_mm_strings=interleave_mm_strings,
|
||||
skip_mm_profiling=skip_mm_profiling,
|
||||
video_pruning_rate=video_pruning_rate,
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from collections.abc import Mapping
|
||||
from pathlib import Path
|
||||
from typing import Any, Literal, TypeAlias, TypedDict, final
|
||||
|
||||
from pydantic import ConfigDict, Field, field_validator, model_validator
|
||||
@@ -158,6 +159,24 @@ class MultiModalConfig:
|
||||
"""Optional override for the multi-modal encoder attention backend when
|
||||
using vision transformers. Accepts any value from
|
||||
`vllm.v1.attention.backends.registry.AttentionBackendEnum` (e.g. `FLASH_ATTN`)."""
|
||||
mm_encoder_attn_dtype: Literal["fp8"] | None = None
|
||||
"""Optional dtype override for ViT encoder attention. Set to `"fp8"` to
|
||||
enable FP8 quantization via the FlashInfer cuDNN backend. When set to
|
||||
`"fp8"` without a scale file, dynamic scaling is used automatically.
|
||||
See docs/features/quantization/fp8_vit_attn.md for details."""
|
||||
mm_encoder_fp8_scale_path: str | None = None
|
||||
"""Path to a JSON file containing per-layer FP8 Q/K/V scales for ViT
|
||||
encoder attention. When provided (with `mm_encoder_attn_dtype="fp8"`),
|
||||
static scaling is used. When omitted, dynamic scaling is used."""
|
||||
mm_encoder_fp8_scale_save_path: str | None = None
|
||||
"""When set with dynamic FP8 scaling (`mm_encoder_attn_dtype="fp8"`
|
||||
and no `mm_encoder_fp8_scale_path`), saves the calibrated scales to
|
||||
this file after the amax history buffer is full. The saved file can
|
||||
then be used as `mm_encoder_fp8_scale_path` in subsequent runs."""
|
||||
mm_encoder_fp8_scale_save_margin: float = Field(default=1.5, gt=0.0)
|
||||
"""Safety margin multiplied onto scales when auto-saving. A value > 1
|
||||
leaves headroom so that inputs with larger activations than the
|
||||
calibration set do not overflow FP8 range. Default 1.5."""
|
||||
interleave_mm_strings: bool = False
|
||||
"""Enable fully interleaved support for multimodal prompts, while using
|
||||
--chat-template-content-format=string."""
|
||||
@@ -233,6 +252,36 @@ class MultiModalConfig:
|
||||
"'mm_shm_cache_max_object_size_mb' should only be set when "
|
||||
"'mm_processor_cache_type' is 'shm'."
|
||||
)
|
||||
# Validate FP8 scale path combinations.
|
||||
if self.mm_encoder_attn_dtype != "fp8" and (
|
||||
self.mm_encoder_fp8_scale_path is not None
|
||||
or self.mm_encoder_fp8_scale_save_path is not None
|
||||
):
|
||||
raise ValueError(
|
||||
"'mm_encoder_fp8_scale_path' and "
|
||||
"'mm_encoder_fp8_scale_save_path' require "
|
||||
"'mm_encoder_attn_dtype' to be 'fp8'."
|
||||
)
|
||||
if (
|
||||
self.mm_encoder_fp8_scale_path is not None
|
||||
and self.mm_encoder_fp8_scale_save_path is not None
|
||||
):
|
||||
raise ValueError(
|
||||
"'mm_encoder_fp8_scale_save_path' cannot be used with "
|
||||
"'mm_encoder_fp8_scale_path' (saving requires dynamic scaling)."
|
||||
)
|
||||
|
||||
# Validate file paths exist.
|
||||
if self.mm_encoder_fp8_scale_path is not None:
|
||||
scale_path = Path(self.mm_encoder_fp8_scale_path)
|
||||
if not scale_path.is_file():
|
||||
raise FileNotFoundError(f"FP8 scale file not found: {scale_path}")
|
||||
if self.mm_encoder_fp8_scale_save_path is not None:
|
||||
save_parent = Path(self.mm_encoder_fp8_scale_save_path).parent
|
||||
if not save_parent.is_dir():
|
||||
raise FileNotFoundError(
|
||||
f"Parent directory for FP8 scale save path not found: {save_parent}"
|
||||
)
|
||||
return self
|
||||
|
||||
def compute_hash(self) -> str:
|
||||
@@ -252,6 +301,8 @@ class MultiModalConfig:
|
||||
if self.mm_encoder_attn_backend is not None
|
||||
else None,
|
||||
self.mm_encoder_tp_mode,
|
||||
self.mm_encoder_attn_dtype,
|
||||
self.mm_encoder_fp8_scale_path,
|
||||
]
|
||||
hash_str = safe_hash(str(factors).encode(), usedforsecurity=False).hexdigest()
|
||||
return hash_str
|
||||
|
||||
@@ -34,6 +34,7 @@ logger = init_logger(__name__)
|
||||
MTPModelTypes = Literal[
|
||||
"deepseek_mtp",
|
||||
"mimo_mtp",
|
||||
"mimo_v2_mtp",
|
||||
"glm4_moe_mtp",
|
||||
"glm4_moe_lite_mtp",
|
||||
"glm_ocr_mtp",
|
||||
@@ -63,7 +64,8 @@ SpeculativeMethod = Literal[
|
||||
EagleModelTypes,
|
||||
NgramGPUTypes,
|
||||
]
|
||||
RejectionSampleMethod = Literal["strict", "probabilistic", "synthetic"]
|
||||
RejectionSampleMethod = Literal["standard", "synthetic"]
|
||||
DraftSampleMethod = Literal["greedy", "gumbel"]
|
||||
|
||||
|
||||
@config
|
||||
@@ -183,11 +185,11 @@ class SpeculativeConfig:
|
||||
"""Load config for the draft model. If not specified, will use the load
|
||||
config from the target model."""
|
||||
|
||||
rejection_sample_method: RejectionSampleMethod = "strict"
|
||||
"""Whether to use strict (target and draft sampled tokens match exactly)
|
||||
or probabilistic rejection sampling. Both respect the target model
|
||||
distribution, but the latter yields a higher acceptance rate at the cost
|
||||
of more memory to cache draft logits."""
|
||||
rejection_sample_method: RejectionSampleMethod = "standard"
|
||||
"""The rejection sampling method to use. 'standard' uses probabilistic
|
||||
rejection sampling (with or without cached draft logits, controlled by
|
||||
draft_sample_method). 'synthetic' accepts draft tokens with a decaying
|
||||
probability calibrated to synthetic_acceptance_rate."""
|
||||
|
||||
synthetic_acceptance_rates: list[float] | None = None
|
||||
"""Per-position *unconditional* acceptance rates for synthetic rejection
|
||||
@@ -248,6 +250,14 @@ class SpeculativeConfig:
|
||||
)
|
||||
return SpeculativeConfig._acceptance_length_to_rates(length, n)
|
||||
|
||||
draft_sample_method: DraftSampleMethod = "greedy"
|
||||
"""How the draft model samples tokens. 'greedy' always picks the argmax
|
||||
token, and the draft probabilities are treated as one-hot during rejection
|
||||
sampling. 'gumbel' adds Gumbel noise for stochastic sampling, and the full
|
||||
draft logits are used for the probability ratio test during rejection
|
||||
sampling. This comes at the cost of additional GPU memory usage. This
|
||||
parameter currently only applies to Model Runner V2."""
|
||||
|
||||
def compute_hash(self) -> str:
|
||||
"""
|
||||
WARNING: Whenever a new field is added to this config,
|
||||
@@ -323,6 +333,48 @@ class SpeculativeConfig:
|
||||
}
|
||||
)
|
||||
|
||||
if (arch := hf_config.architectures[0]) in (
|
||||
"MiMoV2ProForCausalLM",
|
||||
"MiMoV2OmniForCausalLM",
|
||||
):
|
||||
from vllm.model_executor.models.mimo_v2_mtp import (
|
||||
_MIMO_V2_PRO_NUM_MTP_LAYERS,
|
||||
)
|
||||
|
||||
mtp_arch_maps = {
|
||||
"MiMoV2ProForCausalLM": "MiMoV2MTPModel",
|
||||
"MiMoV2OmniForCausalLM": "MiMoV2OmniMTPModel",
|
||||
}
|
||||
|
||||
hf_config.model_type = "mimo_v2_mtp"
|
||||
# vLLM currently supports only the first MiMo-V2 MTP layer.
|
||||
n_predict = _MIMO_V2_PRO_NUM_MTP_LAYERS
|
||||
hf_config.update(
|
||||
{
|
||||
"num_hidden_layers": 0,
|
||||
"n_predict": n_predict,
|
||||
"num_nextn_predict_layers": n_predict,
|
||||
"architectures": [mtp_arch_maps[arch]],
|
||||
}
|
||||
)
|
||||
|
||||
if hf_config.architectures[0] == "MiMoV2FlashForCausalLM":
|
||||
from vllm.model_executor.models.mimo_v2_mtp import (
|
||||
_MIMO_V2_FLASH_NUM_MTP_LAYERS,
|
||||
)
|
||||
|
||||
hf_config.model_type = "mimo_v2_mtp"
|
||||
# vLLM currently supports only the first MiMo-V2 MTP layer.
|
||||
n_predict = _MIMO_V2_FLASH_NUM_MTP_LAYERS
|
||||
hf_config.update(
|
||||
{
|
||||
"num_hidden_layers": 0,
|
||||
"n_predict": n_predict,
|
||||
"num_nextn_predict_layers": n_predict,
|
||||
"architectures": ["MiMoV2MTPModel"],
|
||||
}
|
||||
)
|
||||
|
||||
if hf_config.architectures[0] == "Glm4MoeForCausalLM":
|
||||
hf_config.model_type = "glm4_moe_mtp"
|
||||
n_predict = getattr(hf_config, "num_nextn_predict_layers", None)
|
||||
|
||||
@@ -128,13 +128,6 @@ class CuMemAllocator:
|
||||
return CuMemAllocator.instance
|
||||
|
||||
def __init__(self):
|
||||
conf = os.environ.get("PYTORCH_CUDA_ALLOC_CONF", "")
|
||||
assert "expandable_segments:True" not in conf, (
|
||||
"Expandable segments are not compatible with memory pool. "
|
||||
"Please track https://github.com/pytorch/pytorch/issues/147851 "
|
||||
"for the latest updates."
|
||||
)
|
||||
|
||||
self.pointer_to_data: dict[int, AllocationData] = {}
|
||||
self.current_tag: str = CuMemAllocator.default_tag
|
||||
self.allocator_and_pools: dict[str, Any] = {}
|
||||
@@ -264,34 +257,49 @@ class CuMemAllocator:
|
||||
|
||||
assert isinstance(tag, str)
|
||||
|
||||
# Expandable segments are incompatible with the memory pool used for
|
||||
# sleep mode (see https://github.com/pytorch/pytorch/issues/147851).
|
||||
# If the user has enabled expandable segments via
|
||||
# PYTORCH_CUDA_ALLOC_CONF, temporarily disable them for the duration
|
||||
# of the memory pool context and restore on exit.
|
||||
conf = os.environ.get("PYTORCH_CUDA_ALLOC_CONF", "")
|
||||
expandable_was_enabled = "expandable_segments:True" in conf
|
||||
if expandable_was_enabled:
|
||||
torch.cuda.memory._set_allocator_settings("expandable_segments:False")
|
||||
|
||||
old_tag = self.current_tag
|
||||
self.current_tag = tag
|
||||
with use_memory_pool_with_allocator(
|
||||
self.python_malloc_callback, self.python_free_callback
|
||||
) as data:
|
||||
# start to hit another PyTorch bug in PyTorch 2.6,
|
||||
# possibly because of gc-related issue w.r.t. the allocator and
|
||||
# the memory pool.
|
||||
# to avoid the issue, we keep a reference of the data.
|
||||
# see https://github.com/pytorch/pytorch/issues/146431 .
|
||||
self.allocator_and_pools[tag] = data
|
||||
yield
|
||||
# PyTorch's bug, calling torch.cuda.empty_cache() will error
|
||||
# when using pluggable allocator, see
|
||||
# https://github.com/pytorch/pytorch/issues/145168 .
|
||||
# if we have some memory allocated and then freed,
|
||||
# the memory will not be released, e.g. in online quantization,
|
||||
# where the model is created in higher precision, and then
|
||||
# quantized in lower precision.
|
||||
# Find all unused allocations and manually release them.
|
||||
# TODO: we should expose `empty_cache` method in the memory pool.
|
||||
# TODO: ask for help from PyTorch team to expose this method.
|
||||
allocations = data[0].snapshot()
|
||||
for allocation in allocations:
|
||||
if allocation["allocated_size"] == 0:
|
||||
handle = self._python_free_callback(allocation["address"])
|
||||
unmap_and_release(handle)
|
||||
try:
|
||||
with use_memory_pool_with_allocator(
|
||||
self.python_malloc_callback, self.python_free_callback
|
||||
) as data:
|
||||
# start to hit another PyTorch bug in PyTorch 2.6,
|
||||
# possibly because of gc-related issue w.r.t. the allocator
|
||||
# and the memory pool.
|
||||
# to avoid the issue, we keep a reference of the data.
|
||||
# see https://github.com/pytorch/pytorch/issues/146431 .
|
||||
self.allocator_and_pools[tag] = data
|
||||
yield
|
||||
# PyTorch's bug, calling torch.cuda.empty_cache() will error
|
||||
# when using pluggable allocator, see
|
||||
# https://github.com/pytorch/pytorch/issues/145168 .
|
||||
# if we have some memory allocated and then freed,
|
||||
# the memory will not be released, e.g. in online
|
||||
# quantization, where the model is created in higher
|
||||
# precision, and then quantized in lower precision.
|
||||
# Find all unused allocations and manually release them.
|
||||
# TODO: we should expose `empty_cache` method in the memory
|
||||
# pool.
|
||||
# TODO: ask for help from PyTorch team to expose this method.
|
||||
allocations = data[0].snapshot()
|
||||
for allocation in allocations:
|
||||
if allocation["allocated_size"] == 0:
|
||||
handle = self._python_free_callback(allocation["address"])
|
||||
unmap_and_release(handle)
|
||||
finally:
|
||||
self.current_tag = old_tag
|
||||
if expandable_was_enabled:
|
||||
torch.cuda.memory._set_allocator_settings("expandable_segments:True")
|
||||
|
||||
def get_current_usage(self) -> int:
|
||||
"""
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Scheduler-side logic for the NIXL connector."""
|
||||
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
from typing import TYPE_CHECKING, Any
|
||||
@@ -88,6 +89,12 @@ class NixlConnectorScheduler:
|
||||
if vllm_config.scheduler_config.disable_hybrid_kv_cache_manager:
|
||||
logger.info("Hybrid Memory Allocator is enabled with NIXL")
|
||||
|
||||
if os.environ.get("VLLM_NIXL_ABORT_REQUEST_TIMEOUT") is not None:
|
||||
logger.warning(
|
||||
"VLLM_NIXL_ABORT_REQUEST_TIMEOUT is deprecated and will be "
|
||||
"removed in release 0.22.0."
|
||||
)
|
||||
|
||||
# Background thread for handling new handshake requests.
|
||||
self._nixl_handshake_listener_t: threading.Thread | None = None
|
||||
self._stop_event = threading.Event()
|
||||
|
||||
@@ -542,6 +542,14 @@ class EngineArgs:
|
||||
mm_encoder_attn_backend: AttentionBackendEnum | str | None = (
|
||||
MultiModalConfig.mm_encoder_attn_backend
|
||||
)
|
||||
mm_encoder_attn_dtype: str | None = MultiModalConfig.mm_encoder_attn_dtype
|
||||
mm_encoder_fp8_scale_path: str | None = MultiModalConfig.mm_encoder_fp8_scale_path
|
||||
mm_encoder_fp8_scale_save_path: str | None = (
|
||||
MultiModalConfig.mm_encoder_fp8_scale_save_path
|
||||
)
|
||||
mm_encoder_fp8_scale_save_margin: float = (
|
||||
MultiModalConfig.mm_encoder_fp8_scale_save_margin
|
||||
)
|
||||
io_processor_plugin: str | None = None
|
||||
renderer_num_workers: int = 1
|
||||
skip_mm_profiling: bool = MultiModalConfig.skip_mm_profiling
|
||||
@@ -1179,6 +1187,22 @@ class EngineArgs:
|
||||
"--mm-encoder-attn-backend",
|
||||
**multimodal_kwargs["mm_encoder_attn_backend"],
|
||||
)
|
||||
multimodal_group.add_argument(
|
||||
"--mm-encoder-attn-dtype",
|
||||
**multimodal_kwargs["mm_encoder_attn_dtype"],
|
||||
)
|
||||
multimodal_group.add_argument(
|
||||
"--mm-encoder-fp8-scale-path",
|
||||
**multimodal_kwargs["mm_encoder_fp8_scale_path"],
|
||||
)
|
||||
multimodal_group.add_argument(
|
||||
"--mm-encoder-fp8-scale-save-path",
|
||||
**multimodal_kwargs["mm_encoder_fp8_scale_save_path"],
|
||||
)
|
||||
multimodal_group.add_argument(
|
||||
"--mm-encoder-fp8-scale-save-margin",
|
||||
**multimodal_kwargs["mm_encoder_fp8_scale_save_margin"],
|
||||
)
|
||||
multimodal_group.add_argument(
|
||||
"--interleave-mm-strings", **multimodal_kwargs["interleave_mm_strings"]
|
||||
)
|
||||
@@ -1517,6 +1541,10 @@ class EngineArgs:
|
||||
mm_encoder_only=self.mm_encoder_only,
|
||||
mm_encoder_tp_mode=self.mm_encoder_tp_mode,
|
||||
mm_encoder_attn_backend=self.mm_encoder_attn_backend,
|
||||
mm_encoder_attn_dtype=self.mm_encoder_attn_dtype,
|
||||
mm_encoder_fp8_scale_path=self.mm_encoder_fp8_scale_path,
|
||||
mm_encoder_fp8_scale_save_path=self.mm_encoder_fp8_scale_save_path,
|
||||
mm_encoder_fp8_scale_save_margin=self.mm_encoder_fp8_scale_save_margin,
|
||||
pooler_config=self.pooler_config,
|
||||
generation_config=self.generation_config,
|
||||
override_generation_config=self.override_generation_config,
|
||||
|
||||
@@ -79,7 +79,7 @@ from vllm.renderers.inputs.preprocess import (
|
||||
prompt_to_seq,
|
||||
)
|
||||
from vllm.sampling_params import BeamSearchParams, RequestOutputKind, SamplingParams
|
||||
from vllm.tasks import PoolingTask
|
||||
from vllm.tasks import SCORE_TYPE_MAP, PoolingTask
|
||||
from vllm.tokenizers import TokenizerLike
|
||||
from vllm.usage.usage_lib import UsageContext
|
||||
from vllm.utils.counter import Counter
|
||||
@@ -1204,12 +1204,9 @@ class LLM:
|
||||
f"Supported tasks: {self.supported_tasks}"
|
||||
)
|
||||
else:
|
||||
logger.warning_once(
|
||||
"Pooling multitask support is deprecated and will "
|
||||
"be removed in v0.20. When the default pooling task is "
|
||||
"not what you want, you need to manually specify it "
|
||||
'via PoolerConfig(task="%s"). ',
|
||||
pooling_task,
|
||||
raise ValueError(
|
||||
f"Try switching the model's pooling_task "
|
||||
f'via `PoolerConfig(task="{pooling_task}")`'
|
||||
)
|
||||
|
||||
if pooling_task == "plugin" and "plugin" not in self.pooling_io_processors:
|
||||
@@ -1412,7 +1409,7 @@ class LLM:
|
||||
"pooling model."
|
||||
)
|
||||
|
||||
score_type = self.model_config.score_type
|
||||
score_type: str | None = SCORE_TYPE_MAP.get(self.pooling_task, None) # type: ignore[arg-type]
|
||||
if (
|
||||
score_type == "cross-encoder"
|
||||
and getattr(self.model_config.hf_config, "num_labels", 0) != 1
|
||||
|
||||
@@ -317,4 +317,5 @@ class OpenAIServingChatBatch(OpenAIServingChat):
|
||||
model=model_name,
|
||||
choices=choices,
|
||||
usage=usage,
|
||||
system_fingerprint=self.system_fingerprint,
|
||||
)
|
||||
|
||||
@@ -129,6 +129,9 @@ class ChatCompletionStreamResponse(OpenAIBaseModel):
|
||||
model: str
|
||||
choices: list[ChatCompletionResponseStreamChoice]
|
||||
usage: UsageInfo | None = Field(default=None)
|
||||
# Set only on the final chunk of a stream to mirror non-streaming responses
|
||||
# without the per-chunk serialization overhead.
|
||||
system_fingerprint: str | None = None
|
||||
# not part of the OpenAI spec but for tracing the tokens
|
||||
prompt_token_ids: list[int] | None = None
|
||||
|
||||
|
||||
@@ -1195,6 +1195,16 @@ class OpenAIServingChat(OpenAIServing):
|
||||
choices=[choice_data],
|
||||
model=model_name,
|
||||
)
|
||||
# Stamp the fingerprint on terminal chunks only (those with
|
||||
# finish_reason set). When ``include_usage`` is on, the
|
||||
# trailing usage chunk below overrides this as the true
|
||||
# final message.
|
||||
if (
|
||||
not include_usage
|
||||
and self.system_fingerprint is not None
|
||||
and choice_data.finish_reason is not None
|
||||
):
|
||||
chunk.system_fingerprint = self.system_fingerprint
|
||||
|
||||
# handle usage stats if requested & if continuous
|
||||
if include_continuous_usage:
|
||||
@@ -1229,6 +1239,7 @@ class OpenAIServingChat(OpenAIServing):
|
||||
choices=[],
|
||||
model=model_name,
|
||||
usage=final_usage,
|
||||
system_fingerprint=self.system_fingerprint,
|
||||
)
|
||||
final_usage_data = final_usage_chunk.model_dump_json(
|
||||
exclude_unset=True, exclude_none=True
|
||||
@@ -1637,6 +1648,7 @@ class OpenAIServingChat(OpenAIServing):
|
||||
model=model_name,
|
||||
choices=choices,
|
||||
usage=usage,
|
||||
system_fingerprint=self.system_fingerprint,
|
||||
prompt_logprobs=clamp_prompt_logprobs(final_res.prompt_logprobs),
|
||||
prompt_token_ids=(
|
||||
final_res.prompt_token_ids if request.return_token_ids else None
|
||||
|
||||
@@ -153,9 +153,21 @@ class BaseFrontendArgs:
|
||||
"""If set to True, log the stack trace of error responses"""
|
||||
tokens_only: bool = False
|
||||
"""
|
||||
If set to True, only enable the Tokens In<>Out endpoint.
|
||||
If set to True, only enable the Tokens In<>Out endpoint.
|
||||
This is intended for use in a Disaggregated Everything setup.
|
||||
"""
|
||||
fingerprint_mode: Literal["full", "hash", "custom", "none"] = "full"
|
||||
"""Controls the ``system_fingerprint`` field on responses.
|
||||
|
||||
- ``full`` (default): ``vllm-<version>[-<parallelism>]-<hash8>``. Encodes
|
||||
server version, non-trivial parallelism degrees (tp/pp/dp/ep), and an
|
||||
8-char config hash.
|
||||
- ``hash``: ``vllm-<version>-<hash8>``. Parallelism stripped.
|
||||
- ``custom``: emits the literal string from ``--fingerprint-value``.
|
||||
- ``none``: the field is omitted (serialized as ``null``).
|
||||
"""
|
||||
fingerprint_value: str | None = None
|
||||
"""Literal fingerprint string used when ``--fingerprint-mode=custom``."""
|
||||
|
||||
@classmethod
|
||||
def _customize_cli_kwargs(
|
||||
|
||||
@@ -512,3 +512,6 @@ class CompletionStreamResponse(OpenAIBaseModel):
|
||||
model: str
|
||||
choices: list[CompletionResponseStreamChoice]
|
||||
usage: UsageInfo | None = Field(default=None)
|
||||
# Set only on the final chunk of a stream to mirror non-streaming responses
|
||||
# without the per-chunk serialization overhead.
|
||||
system_fingerprint: str | None = None
|
||||
|
||||
@@ -383,6 +383,7 @@ class OpenAIServingCompletion(OpenAIServing):
|
||||
|
||||
chunk = CompletionStreamResponse(
|
||||
id=request_id,
|
||||
object="text_completion",
|
||||
created=created_time,
|
||||
model=model_name,
|
||||
choices=[
|
||||
@@ -401,6 +402,14 @@ class OpenAIServingCompletion(OpenAIServing):
|
||||
)
|
||||
],
|
||||
)
|
||||
# Stamp on terminal chunk only when no trailing usage chunk
|
||||
# will follow (that one is the true final message).
|
||||
if (
|
||||
not include_usage
|
||||
and self.system_fingerprint is not None
|
||||
and finish_reason is not None
|
||||
):
|
||||
chunk.system_fingerprint = self.system_fingerprint
|
||||
if include_continuous_usage:
|
||||
prompt_tokens = num_prompt_tokens[prompt_idx]
|
||||
completion_tokens = previous_num_tokens[i]
|
||||
@@ -410,7 +419,7 @@ class OpenAIServingCompletion(OpenAIServing):
|
||||
total_tokens=prompt_tokens + completion_tokens,
|
||||
)
|
||||
|
||||
response_json = chunk.model_dump_json(exclude_unset=False)
|
||||
response_json = chunk.model_dump_json(exclude_unset=True)
|
||||
yield f"data: {response_json}\n\n"
|
||||
|
||||
total_prompt_tokens = sum(num_prompt_tokens)
|
||||
@@ -433,6 +442,7 @@ class OpenAIServingCompletion(OpenAIServing):
|
||||
model=model_name,
|
||||
choices=[],
|
||||
usage=final_usage_info,
|
||||
system_fingerprint=self.system_fingerprint,
|
||||
)
|
||||
final_usage_data = final_usage_chunk.model_dump_json(
|
||||
exclude_unset=False, exclude_none=True
|
||||
@@ -562,6 +572,7 @@ class OpenAIServingCompletion(OpenAIServing):
|
||||
model=model_name,
|
||||
choices=choices,
|
||||
usage=usage,
|
||||
system_fingerprint=self.system_fingerprint,
|
||||
kv_transfer_params=kv_transfer_params,
|
||||
)
|
||||
|
||||
|
||||
@@ -157,6 +157,19 @@ class OpenAIServing:
|
||||
self.renderer = engine_client.renderer
|
||||
self.input_processor = engine_client.input_processor
|
||||
|
||||
# Computed once at startup (cached by ``vllm_config`` identity) and
|
||||
# stamped on non-streaming responses. Streaming chunks deliberately
|
||||
# omit it to avoid per-chunk overhead.
|
||||
from vllm.entrypoints.openai.fingerprint import get_system_fingerprint
|
||||
|
||||
try:
|
||||
self.system_fingerprint: str | None = get_system_fingerprint(
|
||||
engine_client.vllm_config
|
||||
)
|
||||
except Exception:
|
||||
# Never fail server startup over the fingerprint.
|
||||
self.system_fingerprint = None
|
||||
|
||||
async def beam_search(
|
||||
self,
|
||||
prompt: EngineInput,
|
||||
|
||||
@@ -0,0 +1,84 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Build the ``system_fingerprint`` string returned by the OpenAI-compatible
|
||||
server.
|
||||
|
||||
Four modes, configured via ``--fingerprint-mode``:
|
||||
|
||||
* ``full`` (default): ``vllm-<version>[-<parallelism>]-<hash8>`` — encodes
|
||||
server version, any non-trivial parallelism degree (tp/pp/dp/ep), and an
|
||||
8-char prefix of ``vllm_config.compute_hash()`` (covers model identity,
|
||||
quant config, speculative, attention backend, etc.).
|
||||
* ``hash``: ``vllm-<version>-<hash8>`` — parallelism stripped.
|
||||
* ``custom``: user-provided literal via ``--fingerprint-value``.
|
||||
* ``none``: the field is omitted (serialized as ``null``).
|
||||
|
||||
``get_system_fingerprint`` is only called at serving-class init (a handful
|
||||
of times per server); each subclass caches the returned string on
|
||||
``self.system_fingerprint``, so per-request cost is one attribute read.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Literal
|
||||
|
||||
FingerprintMode = Literal["full", "hash", "custom", "none"]
|
||||
|
||||
_DEFAULT_MODE: FingerprintMode = "full"
|
||||
_CUSTOM_VALUE: str | None = None
|
||||
|
||||
|
||||
def set_default_fingerprint_mode(
|
||||
mode: FingerprintMode,
|
||||
custom_value: str | None = None,
|
||||
) -> None:
|
||||
"""Configure the fingerprint mode for subsequent ``get_system_fingerprint``
|
||||
calls. Called once at server startup."""
|
||||
global _DEFAULT_MODE, _CUSTOM_VALUE
|
||||
_DEFAULT_MODE = mode
|
||||
_CUSTOM_VALUE = custom_value
|
||||
|
||||
|
||||
def get_system_fingerprint(vllm_config: Any) -> str | None:
|
||||
"""Return the fingerprint for ``vllm_config`` using the mode configured by
|
||||
``set_default_fingerprint_mode``."""
|
||||
return build_system_fingerprint(vllm_config, _DEFAULT_MODE, _CUSTOM_VALUE)
|
||||
|
||||
|
||||
def build_system_fingerprint(
|
||||
vllm_config: Any,
|
||||
mode: FingerprintMode = "full",
|
||||
custom_value: str | None = None,
|
||||
) -> str | None:
|
||||
if mode == "none":
|
||||
return None
|
||||
if mode == "custom":
|
||||
return custom_value
|
||||
|
||||
from vllm import __version__ as vllm_version
|
||||
|
||||
try:
|
||||
hash8 = vllm_config.compute_hash()[:8]
|
||||
except Exception:
|
||||
hash8 = "nohash"
|
||||
|
||||
if mode == "hash":
|
||||
return f"vllm-{vllm_version}-{hash8}"
|
||||
|
||||
# mode == "full"
|
||||
parts: list[str] = [f"vllm-{vllm_version}"]
|
||||
pc = getattr(vllm_config, "parallel_config", None)
|
||||
if pc is not None:
|
||||
tp = getattr(pc, "tensor_parallel_size", 1)
|
||||
if tp > 1:
|
||||
parts.append(f"tp{tp}")
|
||||
pp = getattr(pc, "pipeline_parallel_size", 1)
|
||||
if pp > 1:
|
||||
parts.append(f"pp{pp}")
|
||||
dp = getattr(pc, "data_parallel_size", 1)
|
||||
if dp > 1:
|
||||
parts.append(f"dp{dp}")
|
||||
if getattr(pc, "enable_expert_parallel", False):
|
||||
parts.append("ep")
|
||||
parts.append(hash8)
|
||||
return "-".join(parts)
|
||||
@@ -61,9 +61,17 @@ async def init_generate_state(
|
||||
)
|
||||
from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat
|
||||
from vllm.entrypoints.openai.completion.serving import OpenAIServingCompletion
|
||||
from vllm.entrypoints.openai.fingerprint import set_default_fingerprint_mode
|
||||
from vllm.entrypoints.openai.responses.serving import OpenAIServingResponses
|
||||
from vllm.entrypoints.serve.disagg.serving import ServingTokens
|
||||
|
||||
# Applied before any serving class is constructed so that each one picks
|
||||
# up the chosen mode on its first cache miss.
|
||||
set_default_fingerprint_mode(
|
||||
getattr(args, "fingerprint_mode", "full"),
|
||||
getattr(args, "fingerprint_value", None),
|
||||
)
|
||||
|
||||
if args.tool_server == "demo":
|
||||
tool_server: ToolServer | None = DemoToolServer()
|
||||
assert isinstance(tool_server, DemoToolServer)
|
||||
|
||||
@@ -15,10 +15,7 @@ from starlette.datastructures import Headers
|
||||
from vllm import PoolingParams, PoolingRequestOutput, envs
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.engine.protocol import EngineClient
|
||||
from vllm.entrypoints.chat_utils import (
|
||||
ChatTemplateConfig,
|
||||
ChatTemplateContentFormatOption,
|
||||
)
|
||||
from vllm.entrypoints.chat_utils import ChatTemplateConfig
|
||||
from vllm.entrypoints.logger import RequestLogger
|
||||
from vllm.entrypoints.openai.engine.protocol import ErrorResponse
|
||||
from vllm.entrypoints.openai.models.serving import OpenAIServingModels
|
||||
@@ -48,9 +45,7 @@ class PoolingServingBase(ABC):
|
||||
models: OpenAIServingModels,
|
||||
*,
|
||||
request_logger: RequestLogger | None,
|
||||
chat_template: str | None = None,
|
||||
chat_template_content_format: ChatTemplateContentFormatOption = "auto",
|
||||
trust_request_chat_template: bool = False,
|
||||
chat_template_config: ChatTemplateConfig,
|
||||
return_tokens_as_token_ids: bool = False,
|
||||
log_error_stack: bool = False,
|
||||
):
|
||||
@@ -63,11 +58,7 @@ class PoolingServingBase(ABC):
|
||||
self.request_logger = request_logger
|
||||
self.return_tokens_as_token_ids = return_tokens_as_token_ids
|
||||
self.log_error_stack = log_error_stack
|
||||
self.chat_template_config = ChatTemplateConfig(
|
||||
chat_template=chat_template,
|
||||
chat_template_content_format=chat_template_content_format,
|
||||
trust_request_chat_template=trust_request_chat_template,
|
||||
)
|
||||
self.chat_template_config = chat_template_config
|
||||
|
||||
# Shared thread pool executor for preprocessing and postprocessing.
|
||||
self._executor: Executor = models.renderer._executor
|
||||
|
||||
@@ -10,7 +10,7 @@ from vllm.entrypoints.chat_utils import ChatTemplateConfig
|
||||
from vllm.logger import init_logger
|
||||
from vllm.plugins.io_processors import has_io_processor
|
||||
from vllm.renderers import BaseRenderer
|
||||
from vllm.tasks import POOLING_TASKS, SupportedTask
|
||||
from vllm.tasks import POOLING_TASKS, SCORE_TYPE_MAP, SupportedTask
|
||||
|
||||
from .base.io_processor import PoolingIOProcessor
|
||||
from .utils import enable_scoring_api
|
||||
@@ -43,23 +43,24 @@ def init_pooling_io_processors(
|
||||
) -> dict[str, PoolingIOProcessor]:
|
||||
model_config = vllm_config.model_config
|
||||
processors: dict[str, type[PoolingIOProcessor]] = {}
|
||||
pooling_task = model_config.get_pooling_task(supported_tasks)
|
||||
|
||||
if "classify" in supported_tasks:
|
||||
if pooling_task == "classify":
|
||||
from .classify.io_processor import ClassifyIOProcessor
|
||||
|
||||
processors["classify"] = ClassifyIOProcessor
|
||||
|
||||
if "token_classify" in supported_tasks:
|
||||
if pooling_task == "token_classify":
|
||||
from .classify.io_processor import TokenClassifyIOProcessor
|
||||
|
||||
processors["token_classify"] = TokenClassifyIOProcessor
|
||||
|
||||
if "embed" in supported_tasks:
|
||||
if pooling_task == "embed":
|
||||
from .embed.io_processor import EmbedIOProcessor
|
||||
|
||||
processors["embed"] = EmbedIOProcessor
|
||||
|
||||
if "token_embed" in supported_tasks:
|
||||
if pooling_task == "token_embed":
|
||||
from .embed.io_processor import TokenEmbedIOProcessor
|
||||
|
||||
processors["token_embed"] = TokenEmbedIOProcessor
|
||||
@@ -71,15 +72,15 @@ def init_pooling_io_processors(
|
||||
from .pooling.io_processor import PluginWithIOProcessorPlugins
|
||||
|
||||
processors["plugin"] = PluginWithIOProcessorPlugins
|
||||
elif "plugin" in supported_tasks:
|
||||
elif pooling_task == "plugin":
|
||||
from .pooling.io_processor import PluginWithoutIOProcessorPlugins
|
||||
|
||||
processors["plugin"] = PluginWithoutIOProcessorPlugins
|
||||
|
||||
if enable_scoring_api(supported_tasks, model_config):
|
||||
score_type = model_config.score_type
|
||||
from .scoring.io_processor import ScoringIOProcessors
|
||||
|
||||
score_type: str | None = SCORE_TYPE_MAP.get(pooling_task, None) # type: ignore[arg-type]
|
||||
if score_type is not None and score_type in ScoringIOProcessors:
|
||||
processors[score_type] = ScoringIOProcessors[score_type]
|
||||
|
||||
@@ -140,6 +141,10 @@ def init_pooling_state(
|
||||
request_logger: RequestLogger | None,
|
||||
supported_tasks: tuple["SupportedTask", ...],
|
||||
):
|
||||
model_config = engine_client.model_config
|
||||
if model_config is None:
|
||||
return
|
||||
|
||||
from vllm.entrypoints.chat_utils import load_chat_template
|
||||
from vllm.tasks import POOLING_TASKS
|
||||
|
||||
@@ -148,8 +153,14 @@ def init_pooling_state(
|
||||
from .pooling.serving import ServingPooling
|
||||
from .scoring.serving import ServingScores
|
||||
|
||||
model_config = engine_client.model_config
|
||||
resolved_chat_template = load_chat_template(args.chat_template)
|
||||
pooling_task = model_config.get_pooling_task(supported_tasks)
|
||||
|
||||
chat_template_config = ChatTemplateConfig(
|
||||
chat_template=resolved_chat_template,
|
||||
chat_template_content_format=args.chat_template_content_format,
|
||||
trust_request_chat_template=args.trust_request_chat_template,
|
||||
)
|
||||
|
||||
state.serving_pooling = (
|
||||
(
|
||||
@@ -158,9 +169,7 @@ def init_pooling_state(
|
||||
state.openai_serving_models,
|
||||
supported_tasks=supported_tasks,
|
||||
request_logger=request_logger,
|
||||
chat_template=resolved_chat_template,
|
||||
chat_template_content_format=args.chat_template_content_format,
|
||||
trust_request_chat_template=args.trust_request_chat_template,
|
||||
chat_template_config=chat_template_config,
|
||||
)
|
||||
)
|
||||
if any(t in supported_tasks for t in POOLING_TASKS)
|
||||
@@ -171,11 +180,9 @@ def init_pooling_state(
|
||||
engine_client,
|
||||
state.openai_serving_models,
|
||||
request_logger=request_logger,
|
||||
chat_template=resolved_chat_template,
|
||||
chat_template_content_format=args.chat_template_content_format,
|
||||
trust_request_chat_template=args.trust_request_chat_template,
|
||||
chat_template_config=chat_template_config,
|
||||
)
|
||||
if "embed" in supported_tasks
|
||||
if pooling_task == "embed"
|
||||
else None
|
||||
)
|
||||
state.serving_classification = (
|
||||
@@ -183,21 +190,18 @@ def init_pooling_state(
|
||||
engine_client,
|
||||
state.openai_serving_models,
|
||||
request_logger=request_logger,
|
||||
chat_template=resolved_chat_template,
|
||||
chat_template_content_format=args.chat_template_content_format,
|
||||
trust_request_chat_template=args.trust_request_chat_template,
|
||||
chat_template_config=chat_template_config,
|
||||
)
|
||||
if "classify" in supported_tasks
|
||||
if pooling_task == "classify"
|
||||
else None
|
||||
)
|
||||
state.serving_scores = (
|
||||
ServingScores(
|
||||
engine_client,
|
||||
state.openai_serving_models,
|
||||
supported_tasks=supported_tasks,
|
||||
request_logger=request_logger,
|
||||
chat_template=resolved_chat_template,
|
||||
chat_template_content_format=args.chat_template_content_format,
|
||||
trust_request_chat_template=args.trust_request_chat_template,
|
||||
chat_template_config=chat_template_config,
|
||||
enable_flash_late_interaction=getattr(
|
||||
args, "enable_flash_late_interaction", True
|
||||
),
|
||||
@@ -214,7 +218,12 @@ def get_pooling_invocation_types(
|
||||
# NOTE: Items defined earlier take higher priority
|
||||
invocation_types: list[tuple[RequestType, tuple[GetHandlerFn, EndpointFn]]] = []
|
||||
|
||||
if "embed" in supported_tasks:
|
||||
if model_config is None:
|
||||
return invocation_types
|
||||
|
||||
pooling_task = model_config.get_pooling_task(supported_tasks)
|
||||
|
||||
if pooling_task == "embed":
|
||||
from .embed.api_router import create_embedding, embedding
|
||||
from .embed.protocol import EmbeddingRequest
|
||||
|
||||
@@ -222,7 +231,7 @@ def get_pooling_invocation_types(
|
||||
(EmbeddingRequest, (embedding, create_embedding)),
|
||||
]
|
||||
|
||||
if "classify" in supported_tasks:
|
||||
if pooling_task == "classify":
|
||||
from .classify.api_router import classify, create_classify
|
||||
from .classify.protocol import ClassificationRequest
|
||||
|
||||
|
||||
@@ -78,17 +78,15 @@ class ServingPooling(PoolingServingBase):
|
||||
|
||||
# plugin task uses io_processor.parse_request to verify inputs
|
||||
if pooling_task != "plugin" and pooling_task != self.pooling_task:
|
||||
if pooling_task not in self.io_processors:
|
||||
if pooling_task not in self.supported_tasks:
|
||||
raise ValueError(
|
||||
f"Unsupported task: {pooling_task!r} "
|
||||
f"Supported tasks: {self.supported_tasks}"
|
||||
)
|
||||
else:
|
||||
logger.warning_once(
|
||||
"Pooling multitask support is deprecated and will be removed "
|
||||
"in v0.20. When the default pooling task is not what you want, you "
|
||||
"need to manually specify it via --pooler-config.task %s. ",
|
||||
pooling_task,
|
||||
raise ValueError(
|
||||
"Try switching the model's pooling_task "
|
||||
f"via --pooler-config.task {request.task}."
|
||||
)
|
||||
|
||||
if pooling_task == "plugin" and "plugin" not in self.io_processors:
|
||||
|
||||
@@ -8,6 +8,7 @@ from vllm.engine.protocol import EngineClient
|
||||
from vllm.entrypoints.openai.engine.protocol import UsageInfo
|
||||
from vllm.logger import init_logger
|
||||
from vllm.outputs import PoolingRequestOutput, ScoringRequestOutput
|
||||
from vllm.tasks import SCORE_TYPE_MAP, SupportedTask
|
||||
from vllm.v1.pool.late_interaction import (
|
||||
build_late_interaction_doc_params,
|
||||
build_late_interaction_query_params,
|
||||
@@ -38,10 +39,15 @@ class ServingScores(PoolingServing):
|
||||
self,
|
||||
engine_client: EngineClient,
|
||||
*args,
|
||||
supported_tasks: tuple[SupportedTask, ...],
|
||||
enable_flash_late_interaction: bool = True,
|
||||
**kwargs,
|
||||
):
|
||||
self.io_processor_name: str = engine_client.model_config.score_type
|
||||
pooling_task = engine_client.model_config.get_pooling_task(supported_tasks)
|
||||
score_type = SCORE_TYPE_MAP.get(pooling_task, None) # type: ignore[arg-type]
|
||||
assert score_type is not None
|
||||
|
||||
self.io_processor_name: str = score_type
|
||||
self.enable_flash_late_interaction = (
|
||||
self.io_processor_name == "late-interaction"
|
||||
and enable_flash_late_interaction
|
||||
|
||||
@@ -141,10 +141,14 @@ def enable_scoring_api(
|
||||
supported_tasks: tuple["SupportedTask", ...],
|
||||
model_config: ModelConfig | None = None,
|
||||
) -> bool:
|
||||
if any(t in supported_tasks for t in ("embed", "token_embed")):
|
||||
if model_config is None:
|
||||
return False
|
||||
|
||||
pooling_task = model_config.get_pooling_task(supported_tasks)
|
||||
if pooling_task in ("embed", "token_embed"):
|
||||
return True
|
||||
|
||||
if model_config is not None and "classify" in supported_tasks:
|
||||
if pooling_task == "classify":
|
||||
num_labels = getattr(model_config.hf_config, "num_labels", 0)
|
||||
if num_labels != 1:
|
||||
logger.debug_once("Scoring API is only enabled for num_labels == 1.")
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Triton kernel implementations."""
|
||||
@@ -0,0 +1,180 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Stride-aware FP8 quantization with head_dim padding for ViT attention.
|
||||
|
||||
Reads directly from non-contiguous QKV views using 3D strides and pads
|
||||
head_dim to a multiple of 16 for cuDNN compatibility.
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
get_fp8_min_max,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.triton_utils import HAS_TRITON, tl, triton
|
||||
from vllm.utils.math_utils import round_up
|
||||
|
||||
_FP8_MIN, _FP8_MAX = get_fp8_min_max()
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _quantize_pad_fp8_kernel(
|
||||
x_ptr,
|
||||
y_ptr,
|
||||
scale_ptr,
|
||||
stride_xs,
|
||||
stride_xh,
|
||||
stride_xd,
|
||||
stride_ys,
|
||||
stride_yh,
|
||||
stride_yd,
|
||||
num_heads,
|
||||
n_rows,
|
||||
n_cols,
|
||||
n_cols_padded,
|
||||
fp8_min,
|
||||
fp8_max,
|
||||
SKIP_SCALE: tl.constexpr,
|
||||
BLOCK_M: tl.constexpr,
|
||||
BLOCK_N: tl.constexpr,
|
||||
):
|
||||
pid_m = tl.program_id(0)
|
||||
pid_n = tl.program_id(1)
|
||||
|
||||
offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
|
||||
offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
|
||||
mask_m = offs_m < n_rows
|
||||
mask_out = mask_m[:, None] & (offs_n[None, :] < n_cols_padded)
|
||||
mask_in = mask_m[:, None] & (offs_n[None, :] < n_cols)
|
||||
|
||||
# Decompose flattened row into (token, head) for 3D stride indexing.
|
||||
s = offs_m // num_heads
|
||||
h = offs_m % num_heads
|
||||
|
||||
x_ptrs = (
|
||||
x_ptr
|
||||
+ s[:, None] * stride_xs
|
||||
+ h[:, None] * stride_xh
|
||||
+ offs_n[None, :] * stride_xd
|
||||
)
|
||||
x = tl.load(x_ptrs, mask=mask_in, other=0.0).to(tl.float32)
|
||||
if SKIP_SCALE:
|
||||
x_q = x
|
||||
else:
|
||||
scale = tl.load(scale_ptr)
|
||||
x_q = x / scale
|
||||
x_q = tl.clamp(x_q, fp8_min, fp8_max).to(y_ptr.dtype.element_ty)
|
||||
|
||||
y_ptrs = (
|
||||
y_ptr
|
||||
+ s[:, None] * stride_ys
|
||||
+ h[:, None] * stride_yh
|
||||
+ offs_n[None, :] * stride_yd
|
||||
)
|
||||
tl.store(y_ptrs, x_q, mask=mask_out)
|
||||
|
||||
|
||||
def _get_fp8_pad_quant_config(padded_head_dim: int) -> tuple[int, int, int]:
|
||||
block_n = triton.next_power_of_2(padded_head_dim)
|
||||
block_n = max(16, min(block_n, 128))
|
||||
block_m = 16
|
||||
num_warps = 4
|
||||
return block_m, block_n, num_warps
|
||||
|
||||
|
||||
def quantize_fp8_pad_head_dim_triton(
|
||||
tensor: torch.Tensor,
|
||||
scale: torch.Tensor,
|
||||
skip_scale: bool = False,
|
||||
block_m: int | None = None,
|
||||
block_n: int | None = None,
|
||||
num_warps: int | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Quantize a 3D/4D tensor to FP8, padding head_dim to a multiple of 16.
|
||||
|
||||
Reads directly from the input using its 3D strides, so non-contiguous
|
||||
views (e.g. Q/K/V slices from an interleaved QKV buffer) are handled
|
||||
without an extra copy. Output is always a fresh contiguous tensor
|
||||
with shape (S, H, padded_D).
|
||||
"""
|
||||
if not HAS_TRITON:
|
||||
raise RuntimeError("Triton is required to quantize with head_dim padding.")
|
||||
|
||||
original_shape = tensor.shape
|
||||
if tensor.dim() == 4:
|
||||
tensor = tensor.view(-1, tensor.shape[-2], tensor.shape[-1])
|
||||
assert tensor.dim() == 3, f"Expected 3D input (S, H, D), got {tensor.dim()}D"
|
||||
S, H, D = tensor.shape
|
||||
padded_head_dim = round_up(D, 16)
|
||||
out_dtype = current_platform.fp8_dtype()
|
||||
output = torch.empty(
|
||||
(S, H, padded_head_dim),
|
||||
device=tensor.device,
|
||||
dtype=out_dtype,
|
||||
)
|
||||
|
||||
scale_1d = scale.reshape(-1)
|
||||
n_rows = S * H
|
||||
|
||||
if block_m is None or block_n is None or num_warps is None:
|
||||
block_m, block_n, num_warps = _get_fp8_pad_quant_config(padded_head_dim)
|
||||
|
||||
grid = (
|
||||
triton.cdiv(n_rows, block_m),
|
||||
triton.cdiv(padded_head_dim, block_n),
|
||||
)
|
||||
|
||||
_quantize_pad_fp8_kernel[grid](
|
||||
tensor,
|
||||
output,
|
||||
scale_1d,
|
||||
tensor.stride(0),
|
||||
tensor.stride(1),
|
||||
tensor.stride(2),
|
||||
output.stride(0),
|
||||
output.stride(1),
|
||||
output.stride(2),
|
||||
H,
|
||||
n_rows,
|
||||
D,
|
||||
padded_head_dim,
|
||||
_FP8_MIN,
|
||||
_FP8_MAX,
|
||||
SKIP_SCALE=skip_scale,
|
||||
BLOCK_M=block_m,
|
||||
BLOCK_N=block_n,
|
||||
num_warps=num_warps,
|
||||
)
|
||||
|
||||
return output.view((*original_shape[:-1], padded_head_dim))
|
||||
|
||||
|
||||
def quantize_fp8_maybe_pad_head_dim(
|
||||
tensor: torch.Tensor,
|
||||
scale: torch.Tensor,
|
||||
fp8_quant: QuantFP8,
|
||||
skip_scale: bool = False,
|
||||
) -> torch.Tensor:
|
||||
"""Quantize a 3D/4D tensor to FP8, padding head_dim to a multiple of 16
|
||||
only when needed.
|
||||
|
||||
Accepts (S, H, D) or (B, S, H, D) input. Uses ``fp8_quant`` (a
|
||||
:class:`QuantFP8` CustomOp) when head_dim is already aligned to 16
|
||||
(no padding); otherwise falls back to a stride-aware Triton kernel
|
||||
that pads head_dim to a multiple of 16.
|
||||
"""
|
||||
head_dim = tensor.shape[-1]
|
||||
if head_dim % 16 != 0:
|
||||
return quantize_fp8_pad_head_dim_triton(tensor, scale, skip_scale=skip_scale)
|
||||
|
||||
if skip_scale:
|
||||
return tensor.to(current_platform.fp8_dtype())
|
||||
|
||||
# QuantFP8 expects 2D: flatten all dims except (H, D).
|
||||
orig_shape = tensor.shape
|
||||
total_tokens = tensor.numel() // (orig_shape[-1] * orig_shape[-2])
|
||||
tensor_2d = tensor.reshape(total_tokens, -1)
|
||||
fp8_tensor, _ = fp8_quant(tensor_2d, scale=scale)
|
||||
return fp8_tensor.reshape(orig_shape)
|
||||
@@ -386,10 +386,6 @@ class Attention(nn.Module, AttentionLayerBase):
|
||||
# Initialize KV cache quantization attributes
|
||||
_init_kv_cache_quant(self, quant_config, prefix)
|
||||
|
||||
# Initialize TurboQuant buffers (Pi, S, centroids) if tq cache dtype
|
||||
if kv_cache_dtype.startswith("turboquant_"):
|
||||
self._init_turboquant_buffers(kv_cache_dtype, head_size, prefix)
|
||||
|
||||
# for attn backends supporting query quantization
|
||||
self.query_quant = None
|
||||
if (
|
||||
@@ -410,50 +406,6 @@ class Attention(nn.Module, AttentionLayerBase):
|
||||
else GroupShape.PER_TENSOR,
|
||||
)
|
||||
|
||||
def _init_turboquant_buffers(
|
||||
self, cache_dtype: str, head_size: int, prefix: str
|
||||
) -> None:
|
||||
"""Initialize TurboQuant centroids for Lloyd-Max quantization."""
|
||||
from vllm.model_executor.layers.quantization.turboquant.centroids import (
|
||||
get_centroids,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.turboquant.config import (
|
||||
TurboQuantConfig,
|
||||
)
|
||||
|
||||
tq_config = TurboQuantConfig.from_cache_dtype(cache_dtype, head_size)
|
||||
|
||||
self.register_buffer(
|
||||
"_tq_centroids",
|
||||
get_centroids(head_size, tq_config.centroid_bits),
|
||||
)
|
||||
self._tq_config = tq_config
|
||||
|
||||
# Pre-allocate decode intermediate buffers so model.to(device) moves
|
||||
# them to GPU *before* the memory profiler runs. Without this the
|
||||
# profiler gives all free memory to KV cache blocks and the first
|
||||
# decode OOMs when these buffers are lazily allocated.
|
||||
_vllm_cfg = get_current_vllm_config()
|
||||
B = _vllm_cfg.scheduler_config.max_num_seqs
|
||||
Hq = self.num_heads
|
||||
S = _vllm_cfg.attention_config.tq_max_kv_splits_for_cuda_graph
|
||||
D = head_size
|
||||
self.register_buffer(
|
||||
"_tq_mid_o_buf",
|
||||
torch.empty(B, Hq, S, D + 1, dtype=torch.float32),
|
||||
persistent=False,
|
||||
)
|
||||
self.register_buffer(
|
||||
"_tq_output_buf",
|
||||
torch.empty(B, Hq, D, dtype=torch.float32),
|
||||
persistent=False,
|
||||
)
|
||||
self.register_buffer(
|
||||
"_tq_lse_buf",
|
||||
torch.empty(B, Hq, dtype=torch.float32),
|
||||
persistent=False,
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
|
||||
@@ -1,13 +1,32 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import functools
|
||||
import json
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from vllm.config import MultiModalConfig
|
||||
from vllm.kernels.triton.qkv_padded_fp8_quant import (
|
||||
quantize_fp8_maybe_pad_head_dim,
|
||||
)
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.custom_op import CustomOp, maybe_get_oot_by_class
|
||||
from vllm.model_executor.models.vision import get_vit_attn_backend
|
||||
from vllm.model_executor.layers.quantization.input_quant_fp8 import (
|
||||
QuantFP8,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
GroupShape,
|
||||
get_fp8_min_max,
|
||||
)
|
||||
from vllm.model_executor.models.vision import (
|
||||
get_multimodal_config,
|
||||
get_vit_attn_backend,
|
||||
)
|
||||
from vllm.utils.flashinfer import (
|
||||
is_flashinfer_cudnn_fp8_prefill_attn_supported,
|
||||
)
|
||||
from vllm.utils.math_utils import round_up
|
||||
from vllm.v1.attention.backends.fa_utils import get_flash_attn_version
|
||||
from vllm.v1.attention.backends.registry import AttentionBackendEnum
|
||||
@@ -20,6 +39,108 @@ from vllm.v1.attention.ops.vit_attn_wrappers import (
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
_, _FP8_MAX = get_fp8_min_max()
|
||||
_FP8_AMAX_HISTORY_LEN = 16
|
||||
|
||||
# Module-level state for auto-saving dynamic scales. The save is a one-shot
|
||||
# triggered by the first layer whose amax buffer wraps. Path and margin are
|
||||
# captured during layer init (set_current_vllm_config context only lives
|
||||
# across model init, not forward passes).
|
||||
_fp8_scale_save_path: str | None = None
|
||||
_fp8_scale_save_margin: float = MultiModalConfig.mm_encoder_fp8_scale_save_margin
|
||||
_fp8_saved_scale_refs: dict[str, tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = {}
|
||||
|
||||
|
||||
@functools.cache
|
||||
def _load_fp8_scales_file(path: str | None) -> dict[str, dict[str, float]]:
|
||||
"""Load per-layer FP8 Q/K/V scales from a JSON file. Results are cached.
|
||||
|
||||
Expected format (keys ``q_scale`` / ``k_scale`` / ``v_scale`` also accepted)::
|
||||
|
||||
{
|
||||
"visual.blocks.0.attn.attn": {"q": 224.0, "k": 198.0, "v": 210.0},
|
||||
"visual.blocks.1.attn.attn": {"q": 218.0, "k": 195.0, "v": 207.0},
|
||||
}
|
||||
|
||||
To produce such a file, run with ``mm_encoder_fp8_scale_save_path`` set.
|
||||
"""
|
||||
if path is None:
|
||||
return {}
|
||||
|
||||
with open(path, encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
|
||||
# Handle nested "layers" format
|
||||
if "layers" in data and isinstance(data["layers"], dict):
|
||||
data = data["layers"]
|
||||
|
||||
scales: dict[str, dict[str, float]] = {}
|
||||
for layer_name, layer_scales in data.items():
|
||||
if not isinstance(layer_scales, dict):
|
||||
continue
|
||||
q = layer_scales.get("q", layer_scales.get("q_scale"))
|
||||
k = layer_scales.get("k", layer_scales.get("k_scale"))
|
||||
v = layer_scales.get("v", layer_scales.get("v_scale"))
|
||||
if q is not None and k is not None and v is not None:
|
||||
q_f, k_f, v_f = float(q), float(k), float(v)
|
||||
if q_f <= 0 or k_f <= 0 or v_f <= 0:
|
||||
raise ValueError(
|
||||
f"FP8 scales must be positive, got q={q_f}, "
|
||||
f"k={k_f}, v={v_f} for layer '{layer_name}'"
|
||||
)
|
||||
scales[layer_name] = {"q": q_f, "k": k_f, "v": v_f}
|
||||
|
||||
logger.info_once(
|
||||
"Loaded FP8 attention scales from %s (%d layers)", path, len(scales)
|
||||
)
|
||||
return scales
|
||||
|
||||
|
||||
def _maybe_save_fp8_scales(
|
||||
layer_name: str,
|
||||
q_scale: torch.Tensor,
|
||||
k_scale: torch.Tensor,
|
||||
v_scale: torch.Tensor,
|
||||
buffer_wrapped: bool,
|
||||
) -> None:
|
||||
"""Accumulate a layer's scale tensors; on the first amax buffer wrap,
|
||||
dump all accumulated scales to ``mm_encoder_fp8_scale_save_path``.
|
||||
|
||||
No-op unless auto-save is configured. Tensor references are stored on
|
||||
every call (no GPU->CPU sync); ``.item()`` is only called at the single
|
||||
save point to avoid stalling the forward path.
|
||||
"""
|
||||
global _fp8_scale_save_path
|
||||
# Fast path: auto-save either disabled or already finished. Path is
|
||||
# captured at layer init and cleared once the save fires.
|
||||
if _fp8_scale_save_path is None:
|
||||
return
|
||||
|
||||
# Stash scale tensor refs (no GPU->CPU sync yet); wait until the amax
|
||||
# history has seen a full cycle before committing scales to disk.
|
||||
_fp8_saved_scale_refs[layer_name] = (q_scale, k_scale, v_scale)
|
||||
if not buffer_wrapped:
|
||||
return
|
||||
|
||||
# Buffer just wrapped for the first time: materialize scales (with
|
||||
# safety margin) and dump to disk. Clearing _fp8_scale_save_path
|
||||
# makes this a one-shot across all layers.
|
||||
path, margin = _fp8_scale_save_path, _fp8_scale_save_margin
|
||||
scales = {
|
||||
name: {
|
||||
"q": q.item() * margin,
|
||||
"k": k.item() * margin,
|
||||
"v": v.item() * margin,
|
||||
}
|
||||
for name, (q, k, v) in _fp8_saved_scale_refs.items()
|
||||
}
|
||||
_fp8_scale_save_path = None
|
||||
_fp8_saved_scale_refs.clear()
|
||||
with open(path, "w", encoding="utf-8") as f:
|
||||
json.dump(scales, f, indent=2)
|
||||
logger.info("Saved FP8 scales (%d layers) to %s", len(scales), path)
|
||||
|
||||
|
||||
# Batch buckets for cuDNN graph caching.
|
||||
# Graphs use batch size and max sequence length as cache key.
|
||||
# This avoids creating a new graph for each unique set of
|
||||
@@ -148,27 +269,47 @@ class MMEncoderAttention(CustomOp):
|
||||
hidden_size: int,
|
||||
tp_size: int,
|
||||
device: torch.device,
|
||||
fp8_padded_hidden_size: int | None = None,
|
||||
) -> torch.Tensor:
|
||||
if (oot_class := maybe_get_oot_by_class(cls)) is not cls:
|
||||
return oot_class.maybe_recompute_cu_seqlens( # type: ignore[attr-defined]
|
||||
attn_backend, cu_seqlens, hidden_size, tp_size, device
|
||||
attn_backend,
|
||||
cu_seqlens,
|
||||
hidden_size,
|
||||
tp_size,
|
||||
device,
|
||||
fp8_padded_hidden_size=fp8_padded_hidden_size,
|
||||
)
|
||||
|
||||
if attn_backend == AttentionBackendEnum.FLASHINFER:
|
||||
batch_size = len(cu_seqlens) - 1
|
||||
scale = hidden_size // tp_size
|
||||
cu_seqlens = cu_seqlens * scale
|
||||
|
||||
cu_seqlens_qko = cu_seqlens
|
||||
cu_seqlens_v = cu_seqlens * 3
|
||||
if fp8_padded_hidden_size is not None:
|
||||
# FP8 path: after quantization Q/K/V are each independent
|
||||
# contiguous tensors with stride H * padded_D per token.
|
||||
# All sections use the same element stride.
|
||||
scale = fp8_padded_hidden_size // tp_size
|
||||
cu_seqlens = cu_seqlens * scale
|
||||
cu_seqlens_padded = add_padding_to_seqlens(
|
||||
cu_seqlens, batch_size, cu_seqlens[-1]
|
||||
)
|
||||
cu_seqlens = np.concatenate([cu_seqlens_padded, cu_seqlens_padded])
|
||||
else:
|
||||
# BF16 path: Q/K/V are non-contiguous views into shared
|
||||
# buffers. V section has 3x stride from interleaved QKV.
|
||||
scale = hidden_size // tp_size
|
||||
cu_seqlens = cu_seqlens * scale
|
||||
|
||||
cu_seqlens_qko = add_padding_to_seqlens(
|
||||
cu_seqlens_qko, batch_size, cu_seqlens_qko[-1]
|
||||
)
|
||||
cu_seqlens_v = add_padding_to_seqlens(
|
||||
cu_seqlens_v, batch_size, cu_seqlens_v[-1]
|
||||
)
|
||||
cu_seqlens = np.concatenate([cu_seqlens_qko, cu_seqlens_v])
|
||||
cu_seqlens_qko = cu_seqlens
|
||||
cu_seqlens_v = cu_seqlens * 3
|
||||
|
||||
cu_seqlens_qko = add_padding_to_seqlens(
|
||||
cu_seqlens_qko, batch_size, cu_seqlens_qko[-1]
|
||||
)
|
||||
cu_seqlens_v = add_padding_to_seqlens(
|
||||
cu_seqlens_v, batch_size, cu_seqlens_v[-1]
|
||||
)
|
||||
cu_seqlens = np.concatenate([cu_seqlens_qko, cu_seqlens_v])
|
||||
|
||||
cu_seqlens = torch.from_numpy(cu_seqlens).to(device, non_blocking=True)
|
||||
return cu_seqlens
|
||||
@@ -206,6 +347,7 @@ class MMEncoderAttention(CustomOp):
|
||||
# During model initialization, the default dtype is set as the model
|
||||
# weight and activation dtype.
|
||||
dtype = torch.get_default_dtype()
|
||||
self.dtype = dtype
|
||||
|
||||
# Get device-specific vision attention backend.
|
||||
self.attn_backend = get_vit_attn_backend(
|
||||
@@ -229,6 +371,113 @@ class MMEncoderAttention(CustomOp):
|
||||
|
||||
logger.info_once(f"Using {self.attn_backend} for MMEncoderAttention.")
|
||||
|
||||
self._init_fp8_state()
|
||||
|
||||
def _init_fp8_state(self) -> None:
|
||||
"""Initialize FP8 attention state from multimodal config.
|
||||
|
||||
No-op if FP8 is not requested. Raises ``ValueError`` if FP8 is
|
||||
requested but the platform does not support it.
|
||||
"""
|
||||
# Populate defaults so ``_forward_flashinfer`` can
|
||||
# check ``self.fp8_enabled`` and others without AttributeError.
|
||||
self.fp8_enabled = False
|
||||
self._fp8_dynamic_scale = False
|
||||
self.fp8_quant: QuantFP8 | None = None
|
||||
self.skip_scale_q = False
|
||||
self.skip_scale_k = False
|
||||
self.skip_scale_v = False
|
||||
|
||||
mm_cfg = get_multimodal_config()
|
||||
if mm_cfg is None or mm_cfg.mm_encoder_attn_dtype != "fp8":
|
||||
return
|
||||
|
||||
# FP8 path
|
||||
if not is_flashinfer_cudnn_fp8_prefill_attn_supported():
|
||||
raise ValueError(
|
||||
"mm_encoder_attn_dtype='fp8' requires the FlashInfer "
|
||||
"cuDNN backend with cuDNN >= 9.17.1 on a GPU with native "
|
||||
"FP8 support."
|
||||
)
|
||||
|
||||
self.fp8_enabled = True
|
||||
self._fp8_dynamic_scale = mm_cfg.mm_encoder_fp8_scale_path is None
|
||||
self.fp8_quant = QuantFP8(static=True, group_shape=GroupShape.PER_TENSOR)
|
||||
|
||||
# Register buffers pre-device-move; values populated in
|
||||
# process_weights_after_loading. Shape (1, 1, 1, 1) is required by cuDNN.
|
||||
for attr in ("_fp8_q_scale", "_fp8_k_scale", "_fp8_v_scale"):
|
||||
self.register_buffer(
|
||||
attr, torch.ones(1, dtype=torch.float32).view(1, 1, 1, 1)
|
||||
)
|
||||
if self._fp8_dynamic_scale:
|
||||
for attr in ("_fp8_q_amax", "_fp8_k_amax", "_fp8_v_amax"):
|
||||
self.register_buffer(
|
||||
attr,
|
||||
torch.zeros(_FP8_AMAX_HISTORY_LEN, dtype=torch.float32),
|
||||
persistent=False,
|
||||
)
|
||||
self._fp8_amax_pos = 0
|
||||
|
||||
# Capture auto-save config now: the VllmConfig context only lives
|
||||
# across model init, not forward passes, so ``_maybe_save_fp8_scales``
|
||||
# reads these globals instead of re-querying ``get_multimodal_config``.
|
||||
if (
|
||||
mm_cfg.mm_encoder_fp8_scale_save_path is not None
|
||||
and self._fp8_dynamic_scale
|
||||
):
|
||||
global _fp8_scale_save_path, _fp8_scale_save_margin
|
||||
_fp8_scale_save_path = mm_cfg.mm_encoder_fp8_scale_save_path
|
||||
_fp8_scale_save_margin = mm_cfg.mm_encoder_fp8_scale_save_margin
|
||||
|
||||
def process_weights_after_loading(self, act_dtype: torch.dtype) -> None:
|
||||
"""Populate FP8 scale buffers after weights are loaded.
|
||||
|
||||
``act_dtype`` matches the signature used by :class:`Attention` and
|
||||
:class:`MLAAttention` for the loader auto-scan but is unused:
|
||||
FP8 scales are always float32.
|
||||
"""
|
||||
if not self.fp8_enabled:
|
||||
return
|
||||
|
||||
mm_cfg = get_multimodal_config()
|
||||
scale_path = mm_cfg.mm_encoder_fp8_scale_path if mm_cfg is not None else None
|
||||
if scale_path is None:
|
||||
logger.info_once(
|
||||
"FP8 attention enabled with dynamic scaling "
|
||||
"(no scale file provided). Scales will adapt from "
|
||||
"observed Q/K/V amax values (history_len=%d).",
|
||||
_FP8_AMAX_HISTORY_LEN,
|
||||
)
|
||||
return
|
||||
|
||||
all_scales = _load_fp8_scales_file(scale_path)
|
||||
layer_scales = all_scales.get(self.layer_name)
|
||||
if layer_scales is None:
|
||||
raise ValueError(
|
||||
"FP8 attention enabled but scales not found for layer "
|
||||
f"'{self.layer_name}' in {scale_path}. "
|
||||
f"Available layers: {list(all_scales.keys())}"
|
||||
)
|
||||
|
||||
for attr, key in (
|
||||
("_fp8_q_scale", "q"),
|
||||
("_fp8_k_scale", "k"),
|
||||
("_fp8_v_scale", "v"),
|
||||
):
|
||||
getattr(self, attr).fill_(layer_scales[key])
|
||||
self.skip_scale_q = layer_scales["q"] == 1.0
|
||||
self.skip_scale_k = layer_scales["k"] == 1.0
|
||||
self.skip_scale_v = layer_scales["v"] == 1.0
|
||||
|
||||
logger.debug(
|
||||
"FP8 attention enabled for %s: q=%.4f, k=%.4f, v=%.4f",
|
||||
self.layer_name if self.layer_name else "MMEncoderAttention",
|
||||
layer_scales["q"],
|
||||
layer_scales["k"],
|
||||
layer_scales["v"],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def enabled(cls) -> bool:
|
||||
return True
|
||||
@@ -353,6 +602,44 @@ class MMEncoderAttention(CustomOp):
|
||||
output = output.reshape(bsz, q_len, -1)
|
||||
return output
|
||||
|
||||
@torch.no_grad()
|
||||
def _record_amax_and_update_scales(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
) -> None:
|
||||
"""Record Q/K/V amax into circular history and recompute scales.
|
||||
|
||||
All work stays on GPU with no device-to-host sync. The Python-side
|
||||
history position counter is mutated, so this method must NOT be
|
||||
called inside CUDA graph capture/replay. When CUDA graphs are
|
||||
used for the encoder, dynamic scaling should be disabled by
|
||||
providing a static scale file via --mm-encoder-fp8-scale-path.
|
||||
"""
|
||||
pos = self._fp8_amax_pos
|
||||
self._fp8_amax_pos = (pos + 1) % _FP8_AMAX_HISTORY_LEN
|
||||
|
||||
for tensor, amax_buf, scale_buf in (
|
||||
(query, self._fp8_q_amax, self._fp8_q_scale),
|
||||
(key, self._fp8_k_amax, self._fp8_k_scale),
|
||||
(value, self._fp8_v_amax, self._fp8_v_scale),
|
||||
):
|
||||
amax_buf[pos] = tensor.amax()
|
||||
max_amax = amax_buf.max()
|
||||
scale_buf.fill_(
|
||||
torch.clamp(max_amax, min=torch.finfo(torch.float32).tiny) / _FP8_MAX
|
||||
)
|
||||
|
||||
buffer_wrapped = self._fp8_amax_pos == 0 and pos == _FP8_AMAX_HISTORY_LEN - 1
|
||||
_maybe_save_fp8_scales(
|
||||
self.layer_name,
|
||||
self._fp8_q_scale,
|
||||
self._fp8_k_scale,
|
||||
self._fp8_v_scale,
|
||||
buffer_wrapped,
|
||||
)
|
||||
|
||||
def _forward_flashinfer(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
@@ -363,7 +650,32 @@ class MMEncoderAttention(CustomOp):
|
||||
sequence_lengths: torch.Tensor
|
||||
| None = None, # Only used for FlashInfer CuDNN backend
|
||||
) -> torch.Tensor:
|
||||
return vit_flashinfer_wrapper(
|
||||
if self.fp8_enabled:
|
||||
assert self.fp8_quant is not None
|
||||
|
||||
if self._fp8_dynamic_scale:
|
||||
self._record_amax_and_update_scales(query, key, value)
|
||||
|
||||
query = quantize_fp8_maybe_pad_head_dim(
|
||||
query,
|
||||
self._fp8_q_scale,
|
||||
skip_scale=self.skip_scale_q,
|
||||
fp8_quant=self.fp8_quant,
|
||||
)
|
||||
key = quantize_fp8_maybe_pad_head_dim(
|
||||
key,
|
||||
self._fp8_k_scale,
|
||||
skip_scale=self.skip_scale_k,
|
||||
fp8_quant=self.fp8_quant,
|
||||
)
|
||||
value = quantize_fp8_maybe_pad_head_dim(
|
||||
value,
|
||||
self._fp8_v_scale,
|
||||
skip_scale=self.skip_scale_v,
|
||||
fp8_quant=self.fp8_quant,
|
||||
)
|
||||
|
||||
output = vit_flashinfer_wrapper(
|
||||
q=query,
|
||||
k=key,
|
||||
v=value,
|
||||
@@ -372,8 +684,17 @@ class MMEncoderAttention(CustomOp):
|
||||
cu_seqlens=cu_seqlens,
|
||||
max_seqlen=max_seqlen,
|
||||
sequence_lengths=sequence_lengths,
|
||||
q_scale=self._fp8_q_scale if self.fp8_enabled else None,
|
||||
k_scale=self._fp8_k_scale if self.fp8_enabled else None,
|
||||
v_scale=self._fp8_v_scale if self.fp8_enabled else None,
|
||||
o_data_type=self.dtype if self.fp8_enabled else None,
|
||||
)
|
||||
|
||||
if self.fp8_enabled and output.shape[-1] != self.head_size:
|
||||
output = output[..., : self.head_size].contiguous()
|
||||
|
||||
return output
|
||||
|
||||
def forward_native(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
|
||||
@@ -1120,7 +1120,12 @@ class QKVParallelLinear(ColumnParallelLinear):
|
||||
# Special case for Quantization.
|
||||
# If quantized, we need to adjust the offset and size to account
|
||||
# for the packing.
|
||||
if (
|
||||
if isinstance(param, BlockQuantScaleParameter):
|
||||
weight_block_size = getattr(self, "weight_block_size", None)
|
||||
shard_size, shard_offset = adjust_block_scale_shard(
|
||||
weight_block_size, shard_size, shard_offset
|
||||
)
|
||||
elif (
|
||||
isinstance(param, (PackedColumnParameter, PackedvLLMParameter))
|
||||
and param.packed_dim == param.output_dim
|
||||
):
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
|
||||
import math
|
||||
from collections.abc import Callable
|
||||
from functools import partial
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
@@ -33,28 +34,46 @@ from vllm.v1.attention.backend import AttentionMetadata
|
||||
from vllm.v1.attention.backends.linear_attn import LinearAttentionMetadata
|
||||
|
||||
|
||||
@CustomOp.register("minimax_text01_rmsnorm_tp")
|
||||
class MiniMaxText01RMSNormTP(CustomOp):
|
||||
name = "MiniMaxText01RMSNormTP"
|
||||
|
||||
def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
|
||||
def __init__(
|
||||
self,
|
||||
hidden_size: int,
|
||||
eps: float = 1e-6,
|
||||
*,
|
||||
weight_shard_world_size: int | None = None,
|
||||
weight_shard_rank: int | None = None,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.tp_world = get_tensor_model_parallel_world_size()
|
||||
self.tp_rank = get_tensor_model_parallel_rank()
|
||||
self.weight = nn.Parameter(torch.ones(int(hidden_size / self.tp_world)))
|
||||
self.weight_shard_world = weight_shard_world_size or self.tp_world
|
||||
self.weight_shard_rank = (
|
||||
self.tp_rank if weight_shard_rank is None else weight_shard_rank
|
||||
)
|
||||
|
||||
self.weight.weight_loader = self.weight_loader
|
||||
self.weight = nn.Parameter(torch.ones(hidden_size // self.weight_shard_world))
|
||||
self.weight.weight_loader = partial(
|
||||
self.weight_loader,
|
||||
shard_world_size=self.weight_shard_world,
|
||||
shard_rank=self.weight_shard_rank,
|
||||
)
|
||||
self.variance_epsilon = eps
|
||||
|
||||
@staticmethod
|
||||
def weight_loader(
|
||||
param: nn.Parameter,
|
||||
loaded_weight: torch.Tensor,
|
||||
shard_world_size: int | None = None,
|
||||
shard_rank: int | None = None,
|
||||
) -> None:
|
||||
tp_world = get_tensor_model_parallel_world_size()
|
||||
tp_rank = get_tensor_model_parallel_rank()
|
||||
if shard_world_size is None:
|
||||
shard_world_size = get_tensor_model_parallel_world_size()
|
||||
if shard_rank is None:
|
||||
shard_rank = get_tensor_model_parallel_rank()
|
||||
|
||||
shard_size = loaded_weight.shape[0] // tp_world
|
||||
shard = slice(tp_rank * shard_size, (tp_rank + 1) * shard_size)
|
||||
shard_size = loaded_weight.shape[0] // shard_world_size
|
||||
shard = slice(shard_rank * shard_size, (shard_rank + 1) * shard_size)
|
||||
param.data.copy_(loaded_weight[shard])
|
||||
|
||||
def _forward(
|
||||
|
||||
@@ -15,7 +15,11 @@ from typing_extensions import assert_never
|
||||
import vllm.envs as envs
|
||||
from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.attention import Attention, MLAAttention
|
||||
from vllm.model_executor.layers.attention import (
|
||||
Attention,
|
||||
MLAAttention,
|
||||
MMEncoderAttention,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.base_config import (
|
||||
QuantizationConfig,
|
||||
QuantizeMethodBase,
|
||||
@@ -106,12 +110,12 @@ def process_weights_after_loading(
|
||||
with device_loading_context(module, target_device):
|
||||
quant_method.process_weights_after_loading(module)
|
||||
|
||||
# Initialize post-load attention weights for both Attention and MLA.
|
||||
# Initialize post-load attention weights for Attention, MLA, and MM encoder.
|
||||
# NOTE: Happens after other modules so we can easily decompress weights.
|
||||
for _, module in model.named_modules():
|
||||
if isinstance(module, (Attention, MLAAttention)) and hasattr(
|
||||
module, "process_weights_after_loading"
|
||||
):
|
||||
if isinstance(
|
||||
module, (Attention, MLAAttention, MMEncoderAttention)
|
||||
) and hasattr(module, "process_weights_after_loading"):
|
||||
# TODO(lucas): see if there is a way to unify the signatures
|
||||
# of process_weights_after_loading
|
||||
with device_loading_context(module, target_device):
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
+22
-2
@@ -6,6 +6,7 @@ from itertools import islice
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from vllm.compilation.decorators import support_torch_compile
|
||||
from vllm.config import (
|
||||
CacheConfig,
|
||||
VllmConfig,
|
||||
@@ -268,7 +269,7 @@ class MiMoV2Attention(nn.Module):
|
||||
self.total_num_heads * self.v_head_dim,
|
||||
hidden_size,
|
||||
bias=False,
|
||||
quant_config=quant_config,
|
||||
quant_config=quant_config if "mtp.layers" not in prefix else None,
|
||||
reduce_results=True,
|
||||
prefix=f"{prefix}.o_proj",
|
||||
)
|
||||
@@ -440,6 +441,7 @@ class MiMoV2FlashDecoderLayer(nn.Module):
|
||||
return self.config.hybrid_layer_pattern[self.layer_id] == 1
|
||||
|
||||
|
||||
@support_torch_compile
|
||||
class MiMoV2Model(nn.Module):
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
||||
super().__init__()
|
||||
@@ -603,7 +605,13 @@ class MiMoV2Model(nn.Module):
|
||||
|
||||
if expert_matched:
|
||||
continue
|
||||
|
||||
# Support fused qkv_proj checkpoint (Pro format)
|
||||
if "qkv_proj" in name:
|
||||
if name in params_dict:
|
||||
param = params_dict[name]
|
||||
loaded_weight = loaded_weight.chunk(tp_size, dim=0)[tp_rank]
|
||||
default_weight_loader(param, loaded_weight)
|
||||
continue
|
||||
stacked_matched = False
|
||||
for param_name, weight_name, shard_id in stacked_params_mapping:
|
||||
if weight_name not in name:
|
||||
@@ -662,6 +670,11 @@ class MiMoV2Model(nn.Module):
|
||||
|
||||
|
||||
class MiMoV2FlashForCausalLM(nn.Module, SupportsPP, MixtureOfExperts):
|
||||
packed_modules_mapping = {
|
||||
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
|
||||
"gate_up_proj": ["gate_proj", "up_proj"],
|
||||
}
|
||||
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
||||
super().__init__()
|
||||
config = vllm_config.model_config.hf_config
|
||||
@@ -718,3 +731,10 @@ class MiMoV2FlashForCausalLM(nn.Module, SupportsPP, MixtureOfExperts):
|
||||
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
||||
loader = AutoWeightsLoader(self)
|
||||
return loader.load_weights(weights)
|
||||
|
||||
|
||||
class MiMoV2ProForCausalLM(MiMoV2FlashForCausalLM):
|
||||
packed_modules_mapping = {
|
||||
"qkv_proj": ["qkv_proj"],
|
||||
"gate_up_proj": ["gate_proj", "up_proj"],
|
||||
}
|
||||
@@ -0,0 +1,373 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""Inference-only MiMo-V2 MTP (Multi-Token Prediction) draft model.
|
||||
|
||||
Supports both MiMo-V2-Pro and MiMo-V2-Flash checkpoints.
|
||||
|
||||
Checkpoint weight layout (model.mtp.layers.{idx}.*):
|
||||
enorm - RMSNorm for token embeddings
|
||||
hnorm - RMSNorm for previous hidden states
|
||||
eh_proj - ReplicatedLinear(hidden*2 -> hidden)
|
||||
input_layernorm - pre-attention RMSNorm
|
||||
self_attn.* - attention weights; format differs by variant:
|
||||
Pro: fused qkv_proj [Q;K;V] concatenated
|
||||
Flash: separate q_proj, k_proj, v_proj
|
||||
pre_mlp_layernorm - post-attention / pre-MLP RMSNorm
|
||||
mlp.* - dense MLP (gate_proj / up_proj / down_proj)
|
||||
final_layernorm - norm applied before logit computation
|
||||
"""
|
||||
|
||||
from collections.abc import Iterable
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from transformers import PretrainedConfig
|
||||
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.distributed import (
|
||||
get_tensor_model_parallel_rank,
|
||||
get_tensor_model_parallel_world_size,
|
||||
)
|
||||
from vllm.model_executor.layers.layernorm import RMSNorm
|
||||
from vllm.model_executor.layers.linear import ReplicatedLinear
|
||||
from vllm.model_executor.layers.logits_processor import LogitsProcessor
|
||||
from vllm.model_executor.layers.quantization import QuantizationConfig
|
||||
from vllm.model_executor.layers.vocab_parallel_embedding import (
|
||||
ParallelLMHead,
|
||||
VocabParallelEmbedding,
|
||||
)
|
||||
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
|
||||
from vllm.sequence import IntermediateTensors
|
||||
|
||||
from .interfaces import (
|
||||
MultiModalEmbeddings,
|
||||
SupportsMultiModal,
|
||||
_require_is_multimodal,
|
||||
)
|
||||
from .mimo_v2 import MiMoV2Attention, MiMoV2MLP
|
||||
from .utils import _merge_multimodal_embeddings, maybe_prefix
|
||||
|
||||
# MiMo-V2 checkpoints contain multiple MTP layers, but vLLM currently supports
|
||||
# only the first layer and only one speculative token.
|
||||
_MIMO_V2_PRO_NUM_MTP_LAYERS = 1
|
||||
_MIMO_V2_FLASH_NUM_MTP_LAYERS = 1
|
||||
|
||||
|
||||
class MiMoV2MTPLayer(nn.Module):
|
||||
"""Single MTP predictor layer for MiMo-V2 (Pro and Flash).
|
||||
|
||||
Mirrors the single-layer MiMo-V2 nextn reference implementation.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: PretrainedConfig,
|
||||
prefix: str,
|
||||
quant_config: QuantizationConfig | None = None,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
|
||||
# Predictor head components
|
||||
self.enorm = RMSNorm(config.hidden_size, eps=config.layernorm_epsilon)
|
||||
self.hnorm = RMSNorm(config.hidden_size, eps=config.layernorm_epsilon)
|
||||
self.eh_proj = ReplicatedLinear(
|
||||
config.hidden_size * 2, config.hidden_size, bias=False
|
||||
)
|
||||
|
||||
# MTP uses the SWA attention configuration
|
||||
# implementation.
|
||||
swa_rope_theta = getattr(
|
||||
config,
|
||||
"swa_rope_theta",
|
||||
getattr(config, "rope_theta", 1000000),
|
||||
)
|
||||
sliding_window_size = getattr(config, "sliding_window_size", -1)
|
||||
|
||||
self.input_layernorm = RMSNorm(config.hidden_size, eps=config.layernorm_epsilon)
|
||||
self.self_attn = MiMoV2Attention(
|
||||
hidden_size=config.hidden_size,
|
||||
num_heads=config.swa_num_attention_heads,
|
||||
num_kv_heads=config.swa_num_key_value_heads,
|
||||
head_dim=config.swa_head_dim,
|
||||
v_head_dim=getattr(config, "swa_v_head_dim", None),
|
||||
v_scale=getattr(config, "attention_value_scale", None),
|
||||
sliding_window_size=sliding_window_size,
|
||||
attention_bias=config.attention_bias,
|
||||
add_swa_attention_sink_bias=getattr(
|
||||
config, "add_swa_attention_sink_bias", False
|
||||
),
|
||||
layer_id=0,
|
||||
rope_theta=swa_rope_theta,
|
||||
max_position_embeddings=getattr(config, "max_position_embeddings", 32768),
|
||||
quant_config=quant_config,
|
||||
partial_rotary_factor=getattr(config, "partial_rotary_factor", 1.0),
|
||||
prefix=f"{prefix}.self_attn",
|
||||
)
|
||||
self.pre_mlp_layernorm = RMSNorm(
|
||||
config.hidden_size, eps=config.layernorm_epsilon
|
||||
)
|
||||
self.mlp = MiMoV2MLP(
|
||||
hidden_size=config.hidden_size,
|
||||
intermediate_size=config.intermediate_size,
|
||||
hidden_act=config.hidden_act,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.mlp",
|
||||
)
|
||||
self.final_layernorm = RMSNorm(config.hidden_size, eps=config.layernorm_epsilon)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
inputs_embeds: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
previous_hidden_states: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
# Combine token embedding and previous hidden state
|
||||
h, _ = self.eh_proj(
|
||||
torch.cat(
|
||||
[self.enorm(inputs_embeds), self.hnorm(previous_hidden_states)], dim=-1
|
||||
)
|
||||
)
|
||||
|
||||
# Transformer block with fused residual norms
|
||||
residual = h
|
||||
h = self.input_layernorm(h)
|
||||
h = self.self_attn(positions=positions, hidden_states=h)
|
||||
h, residual = self.pre_mlp_layernorm(h, residual)
|
||||
h = self.mlp(h)
|
||||
h = h + residual
|
||||
|
||||
return self.final_layernorm(h)
|
||||
|
||||
|
||||
class _MiMoV2MTPLayers(nn.Module):
|
||||
"""Thin wrapper so parameter paths match checkpoint: model.mtp.layers.*"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: PretrainedConfig,
|
||||
num_mtp_layers: int,
|
||||
quant_config: QuantizationConfig | None,
|
||||
prefix: str,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.layers = nn.ModuleDict(
|
||||
{
|
||||
str(i): MiMoV2MTPLayer(
|
||||
config=config,
|
||||
prefix=f"{prefix}.{i}",
|
||||
quant_config=quant_config,
|
||||
)
|
||||
for i in range(num_mtp_layers)
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
class MiMoV2MultiTokenPredictor(nn.Module):
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
|
||||
super().__init__()
|
||||
|
||||
config = vllm_config.model_config.hf_config
|
||||
spec_cfg = vllm_config.speculative_config
|
||||
assert spec_cfg is not None
|
||||
if spec_cfg.num_speculative_tokens != 1:
|
||||
raise ValueError(
|
||||
"MiMo-V2 MTP in vLLM only supports num_speculative_tokens=1."
|
||||
)
|
||||
num_mtp_layers = 1
|
||||
|
||||
self.num_mtp_layers = num_mtp_layers
|
||||
|
||||
self.embed_tokens = VocabParallelEmbedding(
|
||||
config.vocab_size,
|
||||
config.hidden_size,
|
||||
)
|
||||
|
||||
self.mtp = _MiMoV2MTPLayers(
|
||||
config=config,
|
||||
num_mtp_layers=num_mtp_layers,
|
||||
quant_config=vllm_config.quant_config,
|
||||
prefix=maybe_prefix(prefix, "mtp.layers"),
|
||||
)
|
||||
|
||||
self.logits_processor = LogitsProcessor(config.vocab_size)
|
||||
|
||||
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
|
||||
return self.embed_tokens(input_ids)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
previous_hidden_states: torch.Tensor,
|
||||
inputs_embeds: torch.Tensor | None = None,
|
||||
spec_step_idx: int = 0,
|
||||
) -> torch.Tensor:
|
||||
assert spec_step_idx == 0, "MiMo-V2 MTP only supports one speculative token."
|
||||
if inputs_embeds is None:
|
||||
inputs_embeds = self.embed_input_ids(input_ids)
|
||||
return self.mtp.layers[str(spec_step_idx)](
|
||||
inputs_embeds, positions, previous_hidden_states
|
||||
)
|
||||
|
||||
def compute_logits(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
lm_head: ParallelLMHead,
|
||||
spec_step_idx: int = 0,
|
||||
) -> torch.Tensor:
|
||||
assert spec_step_idx == 0, "MiMo-V2 MTP only supports one speculative token."
|
||||
return self.logits_processor(lm_head, hidden_states)
|
||||
|
||||
|
||||
class MiMoV2MTP(nn.Module):
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
|
||||
super().__init__()
|
||||
self.config = vllm_config.model_config.hf_config
|
||||
self.model = MiMoV2MultiTokenPredictor(
|
||||
vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
|
||||
)
|
||||
self.lm_head = ParallelLMHead(
|
||||
self.config.vocab_size,
|
||||
self.config.hidden_size,
|
||||
prefix=maybe_prefix(prefix, "lm_head"),
|
||||
)
|
||||
|
||||
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
|
||||
return self.model.embed_input_ids(input_ids)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.Tensor | None,
|
||||
positions: torch.Tensor,
|
||||
hidden_states: torch.Tensor,
|
||||
intermediate_tensors: IntermediateTensors | None = None,
|
||||
inputs_embeds: torch.Tensor | None = None,
|
||||
spec_step_idx: int = 0,
|
||||
) -> torch.Tensor:
|
||||
assert spec_step_idx == 0, "MiMo-V2 MTP only supports one speculative token."
|
||||
return self.model(
|
||||
input_ids, positions, hidden_states, inputs_embeds, spec_step_idx
|
||||
)
|
||||
|
||||
def compute_logits(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
spec_step_idx: int = 0,
|
||||
) -> torch.Tensor | None:
|
||||
assert spec_step_idx == 0, "MiMo-V2 MTP only supports one speculative token."
|
||||
return self.model.compute_logits(hidden_states, self.lm_head, spec_step_idx)
|
||||
|
||||
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
||||
tp_rank = get_tensor_model_parallel_rank()
|
||||
tp_size = get_tensor_model_parallel_world_size()
|
||||
|
||||
stacked_params_mapping = [
|
||||
("gate_up_proj", "gate_proj", 0),
|
||||
("gate_up_proj", "up_proj", 1),
|
||||
# Flash format: separate projections → fused qkv_proj
|
||||
("qkv_proj", "q_proj", "q"),
|
||||
("qkv_proj", "k_proj", "k"),
|
||||
("qkv_proj", "v_proj", "v"),
|
||||
]
|
||||
|
||||
params_dict = dict(self.named_parameters())
|
||||
loaded_params: set[str] = set()
|
||||
|
||||
for name, loaded_weight in weights:
|
||||
if "rotary_emb.inv_freq" in name:
|
||||
continue
|
||||
|
||||
# Only load MTP-related weights, shared embeddings, and lm_head
|
||||
if (
|
||||
"model.mtp" not in name
|
||||
and "model.embed_tokens" not in name
|
||||
and not name.startswith("lm_head")
|
||||
):
|
||||
continue
|
||||
|
||||
# Support fused qkv_proj checkpoint (Pro format).
|
||||
# The checkpoint is stored pre-sharded for TP=8 as
|
||||
# [Q_rank0, K_rank0, V_rank0, Q_rank1, ...], so splitting along
|
||||
# dim 0 with chunk(tp_size) gives each rank its Q+K+V slice for
|
||||
# both the FP8 weight and the block weight_scale_inv. This matches
|
||||
# how the main model loads the same layout.
|
||||
if "qkv_proj" in name:
|
||||
if name in params_dict:
|
||||
param = params_dict[name]
|
||||
loaded_weight = loaded_weight.chunk(tp_size, dim=0)[tp_rank]
|
||||
default_weight_loader(param, loaded_weight)
|
||||
loaded_params.add(name)
|
||||
continue
|
||||
|
||||
# gate_proj/up_proj → gate_up_proj stacking (both formats);
|
||||
# Flash: q_proj/k_proj/v_proj → qkv_proj merging.
|
||||
stacked_matched = False
|
||||
for param_name, weight_name, shard_id in stacked_params_mapping:
|
||||
if weight_name not in name:
|
||||
continue
|
||||
name_rewritten = name.replace(weight_name, param_name)
|
||||
if (
|
||||
name_rewritten.endswith(".bias")
|
||||
and name_rewritten not in params_dict
|
||||
):
|
||||
continue
|
||||
if name_rewritten not in params_dict:
|
||||
continue
|
||||
param = params_dict[name_rewritten]
|
||||
weight_loader = getattr(param, "weight_loader", default_weight_loader)
|
||||
weight_loader(param, loaded_weight, shard_id)
|
||||
loaded_params.add(name_rewritten)
|
||||
stacked_matched = True
|
||||
break
|
||||
|
||||
if stacked_matched:
|
||||
continue
|
||||
|
||||
if name.endswith(".bias") and name not in params_dict:
|
||||
continue
|
||||
if name not in params_dict:
|
||||
continue
|
||||
|
||||
param = params_dict[name]
|
||||
# attention_sink_bias is head-parallel; slice by tp
|
||||
if "attention_sink_bias" in name:
|
||||
total_heads = loaded_weight.shape[0]
|
||||
heads_per_rank = total_heads // tp_size
|
||||
loaded_weight = loaded_weight.narrow(
|
||||
0, tp_rank * heads_per_rank, heads_per_rank
|
||||
)
|
||||
|
||||
weight_loader = getattr(param, "weight_loader", default_weight_loader)
|
||||
weight_loader(param, loaded_weight)
|
||||
loaded_params.add(name)
|
||||
|
||||
return loaded_params
|
||||
|
||||
|
||||
class MiMoV2OmniMTP(MiMoV2MTP, SupportsMultiModal):
|
||||
def embed_input_ids(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
multimodal_embeddings: MultiModalEmbeddings | None = None,
|
||||
*,
|
||||
is_multimodal: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
inputs_embeds = self._embed_text_input_ids(
|
||||
input_ids,
|
||||
self.model.embed_input_ids,
|
||||
is_multimodal=is_multimodal,
|
||||
)
|
||||
|
||||
if multimodal_embeddings is None or len(multimodal_embeddings) == 0:
|
||||
return inputs_embeds
|
||||
|
||||
is_multimodal = _require_is_multimodal(is_multimodal)
|
||||
|
||||
inputs_embeds = _merge_multimodal_embeddings(
|
||||
inputs_embeds=inputs_embeds,
|
||||
multimodal_embeddings=multimodal_embeddings,
|
||||
is_multimodal=is_multimodal,
|
||||
)
|
||||
|
||||
return inputs_embeds
|
||||
File diff suppressed because it is too large
Load Diff
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user