forked from Karylab-cklius/vllm
Compare commits
12
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
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,
|
||||
)
|
||||
+82
-25
@@ -11,29 +11,74 @@
|
||||
namespace vllm {
|
||||
|
||||
template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&),
|
||||
bool act_first>
|
||||
bool act_first, bool HAS_CLAMP>
|
||||
__device__ __forceinline__ scalar_t compute(const scalar_t& x,
|
||||
const scalar_t& y) {
|
||||
return act_first ? ACT_FN(x) * y : x * ACT_FN(y);
|
||||
const scalar_t& y,
|
||||
const float limit) {
|
||||
if constexpr (act_first) {
|
||||
scalar_t gate = x;
|
||||
scalar_t up = y;
|
||||
if constexpr (HAS_CLAMP) {
|
||||
gate = (scalar_t)fminf((float)gate, limit);
|
||||
up = (scalar_t)fmaxf(fminf((float)up, limit), -limit);
|
||||
}
|
||||
return ACT_FN(gate) * up;
|
||||
} else {
|
||||
scalar_t gate = x;
|
||||
scalar_t up = y;
|
||||
if constexpr (HAS_CLAMP) {
|
||||
gate = (scalar_t)fmaxf(fminf((float)gate, limit), -limit);
|
||||
up = (scalar_t)fminf((float)up, limit);
|
||||
}
|
||||
return gate * ACT_FN(up);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename packed_t, packed_t (*PACKED_ACT_FN)(const packed_t&),
|
||||
bool act_first>
|
||||
bool act_first, bool HAS_CLAMP>
|
||||
__device__ __forceinline__ packed_t packed_compute(const packed_t& x,
|
||||
const packed_t& y) {
|
||||
return act_first ? packed_mul(PACKED_ACT_FN(x), y)
|
||||
: packed_mul(x, PACKED_ACT_FN(y));
|
||||
const packed_t& y,
|
||||
const float limit) {
|
||||
if constexpr (act_first) {
|
||||
packed_t gate = x;
|
||||
packed_t up = y;
|
||||
if constexpr (HAS_CLAMP) {
|
||||
float2 g = cast_to_float2(gate);
|
||||
float2 u = cast_to_float2(up);
|
||||
g.x = fminf(g.x, limit);
|
||||
g.y = fminf(g.y, limit);
|
||||
u.x = fmaxf(fminf(u.x, limit), -limit);
|
||||
u.y = fmaxf(fminf(u.y, limit), -limit);
|
||||
gate = cast_to_packed<packed_t>(g);
|
||||
up = cast_to_packed<packed_t>(u);
|
||||
}
|
||||
return packed_mul(PACKED_ACT_FN(gate), up);
|
||||
} else {
|
||||
packed_t gate = x;
|
||||
packed_t up = y;
|
||||
if constexpr (HAS_CLAMP) {
|
||||
float2 g = cast_to_float2(gate);
|
||||
float2 u = cast_to_float2(up);
|
||||
g.x = fmaxf(fminf(g.x, limit), -limit);
|
||||
g.y = fmaxf(fminf(g.y, limit), -limit);
|
||||
u.x = fminf(u.x, limit);
|
||||
u.y = fminf(u.y, limit);
|
||||
gate = cast_to_packed<packed_t>(g);
|
||||
up = cast_to_packed<packed_t>(u);
|
||||
}
|
||||
return packed_mul(gate, PACKED_ACT_FN(up));
|
||||
}
|
||||
}
|
||||
|
||||
// Activation and gating kernel template.
|
||||
template <typename scalar_t, typename packed_t,
|
||||
scalar_t (*ACT_FN)(const scalar_t&),
|
||||
packed_t (*PACKED_ACT_FN)(const packed_t&), bool act_first,
|
||||
bool use_vec, bool use_256b = false>
|
||||
bool use_vec, bool HAS_CLAMP, bool use_256b = false>
|
||||
__global__ void act_and_mul_kernel(
|
||||
scalar_t* __restrict__ out, // [..., d]
|
||||
const scalar_t* __restrict__ input, // [..., 2, d]
|
||||
const int d) {
|
||||
const int d, const float limit) {
|
||||
const scalar_t* x_ptr = input + blockIdx.x * 2 * d;
|
||||
const scalar_t* y_ptr = x_ptr + d;
|
||||
scalar_t* out_ptr = out + blockIdx.x * d;
|
||||
@@ -58,8 +103,9 @@ __global__ void act_and_mul_kernel(
|
||||
}
|
||||
#pragma unroll
|
||||
for (int j = 0; j < pvec_t::NUM_ELTS; j++) {
|
||||
x.elts[j] = packed_compute<packed_t, PACKED_ACT_FN, act_first>(
|
||||
x.elts[j], y.elts[j]);
|
||||
x.elts[j] =
|
||||
packed_compute<packed_t, PACKED_ACT_FN, act_first, HAS_CLAMP>(
|
||||
x.elts[j], y.elts[j], limit);
|
||||
}
|
||||
if constexpr (use_256b) {
|
||||
st256(x, &out_vec[i]);
|
||||
@@ -72,7 +118,8 @@ __global__ void act_and_mul_kernel(
|
||||
for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
|
||||
const scalar_t x = VLLM_LDG(&x_ptr[idx]);
|
||||
const scalar_t y = VLLM_LDG(&y_ptr[idx]);
|
||||
out_ptr[idx] = compute<scalar_t, ACT_FN, act_first>(x, y);
|
||||
out_ptr[idx] =
|
||||
compute<scalar_t, ACT_FN, act_first, HAS_CLAMP>(x, y, limit);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -151,8 +198,11 @@ packed_gelu_tanh_kernel(const packed_t& val) {
|
||||
|
||||
// Launch activation and gating kernel.
|
||||
// Use ACT_FIRST (bool) indicating whether to apply the activation function
|
||||
// first.
|
||||
#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, PACKED_KERNEL, ACT_FIRST) \
|
||||
// first. HAS_CLAMP (bool) enables pre-activation clamping: gate input is
|
||||
// clamped (max only) and up input is clamped (both sides) before the
|
||||
// activation function is applied.
|
||||
#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, PACKED_KERNEL, ACT_FIRST, \
|
||||
HAS_CLAMP, LIMIT) \
|
||||
auto dtype = input.scalar_type(); \
|
||||
int d = input.size(-1) / 2; \
|
||||
int64_t num_tokens = input.numel() / input.size(-1); \
|
||||
@@ -177,8 +227,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
|
||||
scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type, \
|
||||
KERNEL<scalar_t>, \
|
||||
PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
|
||||
ACT_FIRST, true, true><<<grid, block, 0, stream>>>( \
|
||||
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d); \
|
||||
ACT_FIRST, true, HAS_CLAMP, true><<<grid, block, 0, stream>>>( \
|
||||
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT); \
|
||||
}); \
|
||||
} else { \
|
||||
VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] { \
|
||||
@@ -186,8 +236,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
|
||||
scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type, \
|
||||
KERNEL<scalar_t>, \
|
||||
PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
|
||||
ACT_FIRST, true, false><<<grid, block, 0, stream>>>( \
|
||||
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d); \
|
||||
ACT_FIRST, true, HAS_CLAMP, false><<<grid, block, 0, stream>>>( \
|
||||
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT); \
|
||||
}); \
|
||||
} \
|
||||
} else { \
|
||||
@@ -197,8 +247,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
|
||||
scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type, \
|
||||
KERNEL<scalar_t>, \
|
||||
PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
|
||||
ACT_FIRST, false><<<grid, block, 0, stream>>>( \
|
||||
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d); \
|
||||
ACT_FIRST, false, HAS_CLAMP><<<grid, block, 0, stream>>>( \
|
||||
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT); \
|
||||
}); \
|
||||
}
|
||||
|
||||
@@ -206,7 +256,14 @@ void silu_and_mul(torch::Tensor& out, // [..., d]
|
||||
torch::Tensor& input) // [..., 2 * d]
|
||||
{
|
||||
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
|
||||
true);
|
||||
true, false, 0.0f);
|
||||
}
|
||||
|
||||
void silu_and_mul_clamp(torch::Tensor& out, // [..., d]
|
||||
torch::Tensor& input, // [..., 2 * d]
|
||||
double limit) {
|
||||
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
|
||||
true, true, (float)limit);
|
||||
}
|
||||
|
||||
void mul_and_silu(torch::Tensor& out, // [..., d]
|
||||
@@ -215,21 +272,21 @@ void mul_and_silu(torch::Tensor& out, // [..., d]
|
||||
// The difference between mul_and_silu and silu_and_mul is that mul_and_silu
|
||||
// applies the silu to the latter half of the input.
|
||||
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
|
||||
false);
|
||||
false, false, 0.0f);
|
||||
}
|
||||
|
||||
void gelu_and_mul(torch::Tensor& out, // [..., d]
|
||||
torch::Tensor& input) // [..., 2 * d]
|
||||
{
|
||||
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_kernel, vllm::packed_gelu_kernel,
|
||||
true);
|
||||
true, false, 0.0f);
|
||||
}
|
||||
|
||||
void gelu_tanh_and_mul(torch::Tensor& out, // [..., d]
|
||||
torch::Tensor& input) // [..., 2 * d]
|
||||
{
|
||||
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_tanh_kernel,
|
||||
vllm::packed_gelu_tanh_kernel, true);
|
||||
LAUNCH_ACTIVATION_GATE_KERNEL(
|
||||
vllm::gelu_tanh_kernel, vllm::packed_gelu_tanh_kernel, true, false, 0.0f);
|
||||
}
|
||||
|
||||
namespace vllm {
|
||||
|
||||
@@ -163,6 +163,8 @@ void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
|
||||
|
||||
void silu_and_mul(torch::Tensor& out, torch::Tensor& input);
|
||||
|
||||
void silu_and_mul_clamp(torch::Tensor& out, torch::Tensor& input, double limit);
|
||||
|
||||
void silu_and_mul_quant(torch::Tensor& out, torch::Tensor& input,
|
||||
torch::Tensor& scale);
|
||||
|
||||
|
||||
@@ -106,6 +106,12 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
ops.def("silu_and_mul(Tensor! result, Tensor input) -> ()");
|
||||
ops.impl("silu_and_mul", torch::kCUDA, &silu_and_mul);
|
||||
|
||||
// SwiGLU activation with input clamping.
|
||||
ops.def(
|
||||
"silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit) "
|
||||
"-> ()");
|
||||
ops.impl("silu_and_mul_with_clamp", torch::kCUDA, &silu_and_mul_clamp);
|
||||
|
||||
ops.def(
|
||||
"silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
|
||||
ops.impl("silu_and_mul_quant", torch::kCUDA, &silu_and_mul_quant);
|
||||
|
||||
+1
-1
@@ -538,7 +538,7 @@ 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 && \
|
||||
# Fixes nccl_allocator requiring nccl.h at runtime
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
@@ -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")
|
||||
@@ -16,6 +16,7 @@ from vllm.model_executor.layers.activation import (
|
||||
NewGELU,
|
||||
QuickGELU,
|
||||
SiluAndMul,
|
||||
SiluAndMulWithClamp,
|
||||
SwigluOAIAndMul,
|
||||
SwigluStepAndMul,
|
||||
swiglustep_and_mul_triton,
|
||||
@@ -116,6 +117,85 @@ def test_act_and_mul(
|
||||
opcheck(fn, (out, x))
|
||||
|
||||
|
||||
SWIGLU_LIMITS = [3.0, 7.0, 15.0]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("swiglu_limit", SWIGLU_LIMITS)
|
||||
@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
|
||||
@pytest.mark.parametrize("d", D)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@pytest.mark.parametrize("device", CUDA_DEVICES)
|
||||
@torch.inference_mode()
|
||||
def test_silu_and_mul_with_clamp(
|
||||
default_vllm_config,
|
||||
swiglu_limit: float,
|
||||
num_tokens: int,
|
||||
d: int,
|
||||
dtype: torch.dtype,
|
||||
seed: int,
|
||||
device: str,
|
||||
) -> None:
|
||||
"""SiluAndMulWithClamp: cuda kernel must match native reference."""
|
||||
set_random_seed(seed)
|
||||
torch.set_default_device(device)
|
||||
# Use large values to ensure clamping is exercised.
|
||||
x = torch.randn(num_tokens, 2 * d, dtype=dtype) * swiglu_limit * 2
|
||||
|
||||
layer = SiluAndMulWithClamp(swiglu_limit, compile_native=False)
|
||||
out = layer(x)
|
||||
ref_out = layer.forward_native(x)
|
||||
|
||||
rtol = {
|
||||
torch.float16: 2e-3,
|
||||
torch.bfloat16: 2e-2,
|
||||
torch.float: 1.3e-6,
|
||||
}
|
||||
torch.testing.assert_close(
|
||||
out, ref_out, atol=get_default_atol(out), rtol=rtol[out.dtype]
|
||||
)
|
||||
|
||||
# Verify clamping is actually being applied: the clamped output should
|
||||
# differ from the unclamped SiluAndMul output when inputs are large.
|
||||
unclamped_out = SiluAndMul.forward_native(x)
|
||||
assert not torch.equal(ref_out.float(), unclamped_out.float()), (
|
||||
"Input was not large enough to exercise the clamp; increase scale"
|
||||
)
|
||||
|
||||
# Verify gate clamping semantics with a controlled scalar case.
|
||||
# gate=large_val is clamped to limit first, then silu(limit) * 1.0.
|
||||
x_gate = torch.tensor(
|
||||
[[swiglu_limit * 20.0, 1.0]], dtype=torch.float32, device=device
|
||||
)
|
||||
out_gate = SiluAndMulWithClamp(swiglu_limit, compile_native=False)(x_gate)
|
||||
expected_gate = torch.nn.functional.silu(
|
||||
torch.tensor(swiglu_limit, dtype=torch.float32)
|
||||
).item()
|
||||
torch.testing.assert_close(
|
||||
out_gate,
|
||||
torch.tensor([[expected_gate]], dtype=torch.float32, device=device),
|
||||
atol=1e-3,
|
||||
rtol=1e-3,
|
||||
)
|
||||
|
||||
# Verify up clamping semantics: up >> limit gets clamped to limit.
|
||||
x_up = torch.tensor(
|
||||
[[1.0, swiglu_limit * 20.0]], dtype=torch.float32, device=device
|
||||
)
|
||||
out_up = SiluAndMulWithClamp(swiglu_limit, compile_native=False)(x_up)
|
||||
silu_1 = torch.nn.functional.silu(torch.tensor(1.0)).item()
|
||||
torch.testing.assert_close(
|
||||
out_up,
|
||||
torch.tensor([[silu_1 * swiglu_limit]], dtype=torch.float32, device=device),
|
||||
atol=1e-3,
|
||||
rtol=1e-3,
|
||||
)
|
||||
|
||||
# opcheck
|
||||
out_buf = torch.empty(x.shape[:-1] + (d,), dtype=dtype, device=device)
|
||||
opcheck(torch.ops._C.silu_and_mul_with_clamp, (out_buf, x, swiglu_limit))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"activation",
|
||||
[
|
||||
|
||||
@@ -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
|
||||
@@ -961,6 +961,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",
|
||||
|
||||
@@ -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]}'"
|
||||
)
|
||||
@@ -2512,3 +2512,111 @@ def test_block_lookup_cache_multi_blocks_per_key():
|
||||
assert cache.pop(key1, 11) is block11
|
||||
assert cache.get_one_block(key1) is None
|
||||
assert cache.pop(key1, 12) is None
|
||||
|
||||
|
||||
def test_can_fit_full_sequence_swa_cap_admits_long_prompt():
|
||||
"""Hybrid full+SWA model with a pool sized at the startup minimum should
|
||||
admit a prompt longer than the SWA cap, because SlidingWindowManager
|
||||
recycles blocks during chunked prefill (issue #39734)."""
|
||||
block_size = 16
|
||||
sliding_window = 4 * block_size # 64 tokens
|
||||
max_num_batched_tokens = 8 * block_size # 128 tokens
|
||||
max_model_len = 64 * block_size # 1024 tokens — much larger than the SWA cap
|
||||
# Startup pool sizing: full demands cdiv(max_model_len, bs) = 64 blocks,
|
||||
# SWA demands cdiv(SW-1+max_batched, bs) + 1 = cdiv(191, 16) + 1 = 13.
|
||||
# Pool minimum = 64 + 13 = 77; +1 for the null block.
|
||||
num_blocks = 64 + 13 + 1
|
||||
|
||||
config = KVCacheConfig(
|
||||
num_blocks=num_blocks,
|
||||
kv_cache_tensors=[],
|
||||
kv_cache_groups=[
|
||||
KVCacheGroupSpec(
|
||||
["layer_full"],
|
||||
FullAttentionSpec(
|
||||
block_size=block_size,
|
||||
num_kv_heads=1,
|
||||
head_size=1,
|
||||
dtype=torch.float32,
|
||||
),
|
||||
),
|
||||
KVCacheGroupSpec(
|
||||
["layer_swa"],
|
||||
SlidingWindowSpec(
|
||||
block_size=block_size,
|
||||
num_kv_heads=1,
|
||||
head_size=1,
|
||||
dtype=torch.float32,
|
||||
sliding_window=sliding_window,
|
||||
),
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
manager = KVCacheManager(
|
||||
config,
|
||||
max_model_len=max_model_len,
|
||||
max_num_batched_tokens=max_num_batched_tokens,
|
||||
enable_caching=True,
|
||||
hash_block_size=block_size,
|
||||
)
|
||||
|
||||
# A prompt that is shorter than max_model_len but longer than SW + chunk:
|
||||
# cdiv(prompt_len, bs) = 32 blocks. Without the cap, admission would
|
||||
# demand 32 (full) + 32 (SWA) = 64 blocks. With the cap, SWA contributes
|
||||
# only 13, so total = 32 + 13 = 45 ≤ pool size.
|
||||
prompt_len = 32 * block_size
|
||||
req = make_request("long", list(range(prompt_len)), block_size, sha256)
|
||||
|
||||
assert manager.can_fit_full_sequence(req)
|
||||
|
||||
|
||||
def test_can_fit_full_sequence_full_attention_still_gates_oversized():
|
||||
"""The cap only loosens the SWA group; a prompt that exceeds the
|
||||
full-attention pool capacity must still be rejected."""
|
||||
block_size = 16
|
||||
sliding_window = 4 * block_size
|
||||
max_num_batched_tokens = 8 * block_size
|
||||
max_model_len = 64 * block_size
|
||||
# Provide a tiny pool — even a small prompt should be rejected.
|
||||
num_blocks = 5
|
||||
|
||||
config = KVCacheConfig(
|
||||
num_blocks=num_blocks,
|
||||
kv_cache_tensors=[],
|
||||
kv_cache_groups=[
|
||||
KVCacheGroupSpec(
|
||||
["layer_full"],
|
||||
FullAttentionSpec(
|
||||
block_size=block_size,
|
||||
num_kv_heads=1,
|
||||
head_size=1,
|
||||
dtype=torch.float32,
|
||||
),
|
||||
),
|
||||
KVCacheGroupSpec(
|
||||
["layer_swa"],
|
||||
SlidingWindowSpec(
|
||||
block_size=block_size,
|
||||
num_kv_heads=1,
|
||||
head_size=1,
|
||||
dtype=torch.float32,
|
||||
sliding_window=sliding_window,
|
||||
),
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
manager = KVCacheManager(
|
||||
config,
|
||||
max_model_len=max_model_len,
|
||||
max_num_batched_tokens=max_num_batched_tokens,
|
||||
enable_caching=True,
|
||||
hash_block_size=block_size,
|
||||
)
|
||||
|
||||
# 16 blocks of full attention demand alone exceeds the 5-block pool.
|
||||
prompt_len = 16 * block_size
|
||||
req = make_request("oversized", list(range(prompt_len)), block_size, sha256)
|
||||
|
||||
assert not manager.can_fit_full_sequence(req)
|
||||
|
||||
@@ -22,11 +22,13 @@ pytestmark = pytest.mark.cpu_test
|
||||
|
||||
|
||||
def get_sliding_window_manager(sliding_window_spec, block_pool, enable_caching=True):
|
||||
# Tests don't exercise admission gating; pass a large cap that is a no-op.
|
||||
return SlidingWindowManager(
|
||||
sliding_window_spec,
|
||||
block_pool=block_pool,
|
||||
enable_caching=enable_caching,
|
||||
kv_cache_group_id=0,
|
||||
max_admission_blocks_per_request=10**9,
|
||||
)
|
||||
|
||||
|
||||
@@ -38,6 +40,7 @@ def get_chunked_local_attention_manager(
|
||||
block_pool=block_pool,
|
||||
enable_caching=enable_caching,
|
||||
kv_cache_group_id=0,
|
||||
max_admission_blocks_per_request=10**9,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -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,
|
||||
@@ -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
|
||||
|
||||
@@ -63,7 +63,8 @@ SpeculativeMethod = Literal[
|
||||
EagleModelTypes,
|
||||
NgramGPUTypes,
|
||||
]
|
||||
RejectionSampleMethod = Literal["strict", "probabilistic", "synthetic"]
|
||||
RejectionSampleMethod = Literal["standard", "synthetic"]
|
||||
DraftSampleMethod = Literal["greedy", "gumbel"]
|
||||
|
||||
|
||||
@config
|
||||
@@ -183,11 +184,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 +249,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,
|
||||
|
||||
@@ -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:
|
||||
"""
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
@@ -151,6 +151,46 @@ class SiluAndMul(CustomOp):
|
||||
return self.forward_cuda(x)
|
||||
|
||||
|
||||
@CustomOp.register("silu_and_mul_with_clamp")
|
||||
class SiluAndMulWithClamp(CustomOp):
|
||||
"""SwiGLU activation with input clamping (used by some MoE shared experts).
|
||||
|
||||
Computes:
|
||||
gate = clamp(x[..., :d], max=swiglu_limit)
|
||||
up = clamp(x[..., d:], min=-swiglu_limit, max=swiglu_limit)
|
||||
out = silu(gate) * up
|
||||
where d = x.shape[-1] // 2.
|
||||
|
||||
Shapes:
|
||||
x: (num_tokens, 2 * d) or (batch_size, seq_len, 2 * d)
|
||||
return: (num_tokens, d) or (batch_size, seq_len, d)
|
||||
"""
|
||||
|
||||
def __init__(self, swiglu_limit: float, *, compile_native: bool = True):
|
||||
super().__init__(compile_native=compile_native)
|
||||
self.swiglu_limit = float(swiglu_limit)
|
||||
if current_platform.is_cuda_alike() or current_platform.is_xpu():
|
||||
self.op = torch.ops._C.silu_and_mul_with_clamp
|
||||
elif current_platform.is_cpu():
|
||||
self._forward_method = self.forward_native
|
||||
|
||||
def forward_native(self, x: torch.Tensor) -> torch.Tensor:
|
||||
d = x.shape[-1] // 2
|
||||
gate = torch.clamp(x[..., :d], max=self.swiglu_limit)
|
||||
up = torch.clamp(x[..., d:], min=-self.swiglu_limit, max=self.swiglu_limit)
|
||||
return F.silu(gate) * up
|
||||
|
||||
def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
|
||||
d = x.shape[-1] // 2
|
||||
output_shape = x.shape[:-1] + (d,)
|
||||
out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
|
||||
self.op(out, x, self.swiglu_limit)
|
||||
return out
|
||||
|
||||
def forward_xpu(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return self.forward_cuda(x)
|
||||
|
||||
|
||||
# --8<-- [start:mul_and_silu]
|
||||
@CustomOp.register("mul_and_silu")
|
||||
class MulAndSilu(CustomOp):
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -45,7 +45,7 @@ def _gelu_and_mul(
|
||||
# Uses static methods or standalone functions to avoid instantiating CustomOp
|
||||
# classes, which would call get_current_vllm_config() before config is set.
|
||||
_CPU_MOE_ACT_FN: dict[MoEActivation, Callable[[torch.Tensor], torch.Tensor]] = {
|
||||
MoEActivation.SILU: SiluAndMul.forward_native,
|
||||
MoEActivation.SILU: lambda x: SiluAndMul(compile_native=False).forward_native(x),
|
||||
MoEActivation.SWIGLUOAI: _swigluoai_forward_native,
|
||||
MoEActivation.GELU: _gelu_and_mul,
|
||||
}
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -17,6 +17,7 @@ from vllm.distributed import (
|
||||
get_tensor_model_parallel_world_size,
|
||||
)
|
||||
from vllm.forward_context import get_forward_context
|
||||
from vllm.model_executor.layers.activation import SiluAndMul, SiluAndMulWithClamp
|
||||
from vllm.model_executor.layers.deepseek_v4_attention import (
|
||||
DeepseekV4Indexer,
|
||||
DeepseekV4MLAModules,
|
||||
@@ -34,7 +35,10 @@ from vllm.model_executor.layers.linear import (
|
||||
RowParallelLinear,
|
||||
)
|
||||
from vllm.model_executor.layers.logits_processor import LogitsProcessor
|
||||
from vllm.model_executor.layers.quantization import QuantizationMethods
|
||||
from vllm.model_executor.layers.quantization import (
|
||||
QuantizationConfig,
|
||||
QuantizationMethods,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.fp8 import Fp8Config
|
||||
from vllm.model_executor.layers.quantization.mxfp4 import Mxfp4MoEMethod
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
@@ -46,7 +50,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
|
||||
VocabParallelEmbedding,
|
||||
)
|
||||
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
|
||||
from vllm.model_executor.models.deepseek_v2 import DeepseekV2MLP
|
||||
from vllm.model_executor.utils import set_weight_attrs
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.sequence import IntermediateTensors
|
||||
@@ -63,6 +66,57 @@ from .utils import (
|
||||
)
|
||||
|
||||
|
||||
class DeepseekV4MLP(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
hidden_size: int,
|
||||
intermediate_size: int,
|
||||
hidden_act: str,
|
||||
swiglu_limit: float | None = None,
|
||||
quant_config: QuantizationConfig | None = None,
|
||||
reduce_results: bool = True,
|
||||
is_sequence_parallel: bool = False,
|
||||
prefix: str = "",
|
||||
) -> None:
|
||||
super().__init__()
|
||||
|
||||
# If is_sequence_parallel, the input and output tensors are sharded
|
||||
# across the ranks within the tp_group. In this case the weights are
|
||||
# replicated and no collective ops are needed.
|
||||
# Otherwise we use standard TP with an allreduce at the end.
|
||||
self.gate_up_proj = MergedColumnParallelLinear(
|
||||
hidden_size,
|
||||
[intermediate_size] * 2,
|
||||
bias=False,
|
||||
quant_config=quant_config,
|
||||
disable_tp=is_sequence_parallel,
|
||||
prefix=f"{prefix}.gate_up_proj",
|
||||
)
|
||||
self.down_proj = RowParallelLinear(
|
||||
intermediate_size,
|
||||
hidden_size,
|
||||
bias=False,
|
||||
quant_config=quant_config,
|
||||
reduce_results=reduce_results,
|
||||
disable_tp=is_sequence_parallel,
|
||||
prefix=f"{prefix}.down_proj",
|
||||
)
|
||||
if hidden_act != "silu":
|
||||
raise ValueError(
|
||||
f"Unsupported activation: {hidden_act}. Only silu is supported for now."
|
||||
)
|
||||
if swiglu_limit is not None:
|
||||
self.act_fn = SiluAndMulWithClamp(swiglu_limit)
|
||||
else:
|
||||
self.act_fn = SiluAndMul()
|
||||
|
||||
def forward(self, x):
|
||||
gate_up, _ = self.gate_up_proj(x)
|
||||
x = self.act_fn(gate_up)
|
||||
x, _ = self.down_proj(x)
|
||||
return x
|
||||
|
||||
|
||||
class DeepseekV4FP8Config(Fp8Config):
|
||||
"""FP8 config that routes MoE layers to MXFP4 quantization.
|
||||
|
||||
@@ -672,10 +726,11 @@ class DeepseekV4MoE(nn.Module):
|
||||
else:
|
||||
intermediate_size = config.moe_intermediate_size * config.n_shared_experts
|
||||
|
||||
self.shared_experts = DeepseekV2MLP(
|
||||
self.shared_experts = DeepseekV4MLP(
|
||||
hidden_size=config.hidden_size,
|
||||
intermediate_size=intermediate_size,
|
||||
hidden_act=config.hidden_act,
|
||||
swiglu_limit=self.swiglu_limit,
|
||||
quant_config=quant_config,
|
||||
reduce_results=self.use_mega_moe,
|
||||
prefix=f"{prefix}.shared_experts",
|
||||
|
||||
@@ -1778,11 +1778,6 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
|
||||
return None # If vision tower is skipped
|
||||
if getattr(self, "deepstack_input_embeds_num_tokens", 0) == 0:
|
||||
return None
|
||||
if num_tokens > self.deepstack_input_embeds_num_tokens:
|
||||
raise ValueError(
|
||||
"Requested more deepstack tokens than available in buffer: "
|
||||
f"{num_tokens=} > {self.deepstack_input_embeds_num_tokens=}"
|
||||
)
|
||||
|
||||
# get deepstack_input_embeds from buffer, and clear the buffer
|
||||
return IntermediateTensors(
|
||||
@@ -1824,12 +1819,6 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
|
||||
|
||||
# clear deepstack_input_embeds in buffer
|
||||
if num_tokens > 0:
|
||||
if num_tokens > self.deepstack_input_embeds_num_tokens:
|
||||
raise ValueError(
|
||||
"Requested to clear more deepstack tokens than available in "
|
||||
"buffer: "
|
||||
f"{num_tokens=} > {self.deepstack_input_embeds_num_tokens=}"
|
||||
)
|
||||
for idx in range(self.deepstack_num_level):
|
||||
self.deepstack_input_embeds[idx][:num_tokens].zero_()
|
||||
self.deepstack_input_embeds_num_tokens = 0
|
||||
|
||||
@@ -136,6 +136,7 @@ from .utils import (
|
||||
maybe_prefix,
|
||||
)
|
||||
from .vision import (
|
||||
get_fp8_padded_hidden_size,
|
||||
get_vit_attn_backend,
|
||||
is_vit_use_data_parallel,
|
||||
run_dp_sharded_mrope_vision_model,
|
||||
@@ -562,6 +563,13 @@ class Qwen3_VisionTransformer(nn.Module):
|
||||
|
||||
norm_layer = partial(nn.LayerNorm, eps=norm_eps)
|
||||
head_dim = self.hidden_size // self.num_heads
|
||||
|
||||
# FP8 attention: Q/K/V become independent contiguous tensors
|
||||
# after quantization, so cu_seqlens uses uniform stride (no 3x V).
|
||||
self.fp8_padded_hidden_size = get_fp8_padded_hidden_size(
|
||||
self.num_heads, head_dim
|
||||
)
|
||||
|
||||
self.rotary_pos_emb = get_rope(
|
||||
head_size=head_dim,
|
||||
max_position=8192,
|
||||
@@ -776,6 +784,7 @@ class Qwen3_VisionTransformer(nn.Module):
|
||||
self.hidden_size,
|
||||
self.tp_size,
|
||||
device,
|
||||
fp8_padded_hidden_size=self.fp8_padded_hidden_size,
|
||||
)
|
||||
|
||||
return metadata
|
||||
@@ -1707,11 +1716,6 @@ class Qwen3VLForConditionalGeneration(
|
||||
return None # If vision tower is skipped
|
||||
if getattr(self, "deepstack_input_embeds_num_tokens", 0) == 0:
|
||||
return None
|
||||
if num_tokens > self.deepstack_input_embeds_num_tokens:
|
||||
raise ValueError(
|
||||
"Requested more deepstack tokens than available in buffer: "
|
||||
f"{num_tokens=} > {self.deepstack_input_embeds_num_tokens=}"
|
||||
)
|
||||
|
||||
# get deepstack_input_embeds from buffer, and clear the buffer
|
||||
return IntermediateTensors(
|
||||
@@ -1753,12 +1757,6 @@ class Qwen3VLForConditionalGeneration(
|
||||
|
||||
# clear deepstack_input_embeds in buffer
|
||||
if num_tokens > 0:
|
||||
if num_tokens > self.deepstack_input_embeds_num_tokens:
|
||||
raise ValueError(
|
||||
"Requested to clear more deepstack tokens than available in "
|
||||
"buffer: "
|
||||
f"{num_tokens=} > {self.deepstack_input_embeds_num_tokens=}"
|
||||
)
|
||||
for idx in range(self.deepstack_num_level):
|
||||
self.deepstack_input_embeds[idx][:num_tokens].zero_()
|
||||
self.deepstack_input_embeds_num_tokens = 0
|
||||
|
||||
@@ -10,7 +10,7 @@ from typing import Final, Generic, Literal, Protocol, TypeAlias, TypeVar
|
||||
import torch
|
||||
from transformers import PretrainedConfig
|
||||
|
||||
from vllm.config import MultiModalConfig, VllmConfig, get_current_vllm_config
|
||||
from vllm.config import MultiModalConfig, get_current_vllm_config_or_none
|
||||
from vllm.distributed import (
|
||||
get_tensor_model_parallel_rank,
|
||||
get_tensor_model_parallel_world_size,
|
||||
@@ -18,6 +18,7 @@ from vllm.distributed import (
|
||||
)
|
||||
from vllm.logger import init_logger
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.math_utils import round_up
|
||||
from vllm.v1.attention.backends.registry import AttentionBackendEnum
|
||||
|
||||
logger = init_logger(__name__)
|
||||
@@ -102,45 +103,48 @@ def get_vit_attn_backend(
|
||||
"""
|
||||
Get the attention backend for Vision Transformer.
|
||||
"""
|
||||
try:
|
||||
vllm_config: VllmConfig = get_current_vllm_config()
|
||||
model_config = vllm_config.model_config
|
||||
multimodal_config: MultiModalConfig | None = (
|
||||
model_config.multimodal_config if model_config is not None else None
|
||||
)
|
||||
except (AssertionError, AttributeError):
|
||||
multimodal_config = None
|
||||
|
||||
mm_cfg = get_multimodal_config()
|
||||
attn_backend_override = (
|
||||
multimodal_config.mm_encoder_attn_backend
|
||||
if multimodal_config is not None
|
||||
else None
|
||||
mm_cfg.mm_encoder_attn_backend if mm_cfg is not None else None
|
||||
)
|
||||
attn_backend = _get_vit_attn_backend(
|
||||
return _get_vit_attn_backend(
|
||||
head_size,
|
||||
dtype,
|
||||
attn_backend_override=attn_backend_override,
|
||||
)
|
||||
return attn_backend
|
||||
|
||||
|
||||
def get_multimodal_config() -> MultiModalConfig | None:
|
||||
"""Return the current ``MultiModalConfig``, or ``None`` when no engine
|
||||
config context is active (e.g., during unit tests) or when the current
|
||||
``model_config`` does not carry a ``multimodal_config`` (e.g., minimal
|
||||
stubs used in tests)."""
|
||||
vllm_config = get_current_vllm_config_or_none()
|
||||
if vllm_config is None or vllm_config.model_config is None:
|
||||
return None
|
||||
return getattr(vllm_config.model_config, "multimodal_config", None)
|
||||
|
||||
|
||||
def get_fp8_padded_hidden_size(num_heads: int, head_dim: int) -> int | None:
|
||||
"""Return the padded hidden size for FP8 ViT encoder attention, or
|
||||
``None`` when FP8 is not enabled.
|
||||
|
||||
cuDNN FP8 prefill attention requires ``head_dim`` to be a multiple of
|
||||
16. For non-aligned ``head_dim`` (e.g. 72), Q/K/V are padded to the
|
||||
nearest multiple of 16.
|
||||
"""
|
||||
mm_cfg = get_multimodal_config()
|
||||
if mm_cfg is None or mm_cfg.mm_encoder_attn_dtype != "fp8":
|
||||
return None
|
||||
return num_heads * round_up(head_dim, 16)
|
||||
|
||||
|
||||
def is_vit_use_data_parallel():
|
||||
"""
|
||||
Get the tensor parallel type for Vision Transformer.
|
||||
"""
|
||||
try:
|
||||
vllm_config: VllmConfig = get_current_vllm_config()
|
||||
model_config = vllm_config.model_config
|
||||
multimodal_config: MultiModalConfig | None = (
|
||||
model_config.multimodal_config if model_config is not None else None
|
||||
)
|
||||
except (AssertionError, AttributeError):
|
||||
multimodal_config = None
|
||||
|
||||
mm_encoder_tp_mode = (
|
||||
multimodal_config.mm_encoder_tp_mode if multimodal_config is not None else None
|
||||
)
|
||||
return mm_encoder_tp_mode == "data"
|
||||
mm_cfg = get_multimodal_config()
|
||||
return mm_cfg is not None and mm_cfg.mm_encoder_tp_mode == "data"
|
||||
|
||||
|
||||
VisionFeatureSelectStrategyStr = Literal["class", "default", "full"]
|
||||
|
||||
@@ -772,6 +772,40 @@ def should_use_flashinfer_for_blockscale_fp8_gemm(
|
||||
return should_use_flashinfer
|
||||
|
||||
|
||||
_MIN_CUDNN_FP8 = 91701 # cuDNN >= 9.17.1 required for FP8 attention
|
||||
|
||||
|
||||
@functools.cache
|
||||
def is_flashinfer_cudnn_fp8_prefill_attn_supported() -> bool:
|
||||
"""Check if FP8 ViT attention is supported on this platform.
|
||||
|
||||
Requires native FP8 hardware support, the FlashInfer cuDNN backend,
|
||||
and cuDNN >= 9.17.1.
|
||||
"""
|
||||
from vllm.v1.attention.backends.registry import AttentionBackendEnum
|
||||
|
||||
# cuDNN SDPA FP8 requires Hopper (SM 90) or newer.
|
||||
if not current_platform.has_device_capability(90):
|
||||
return False
|
||||
|
||||
try:
|
||||
supported = current_platform.get_supported_vit_attn_backends()
|
||||
if AttentionBackendEnum.FLASHINFER not in supported:
|
||||
return False
|
||||
except (ImportError, AttributeError):
|
||||
return False
|
||||
|
||||
try:
|
||||
import torch.backends.cudnn as cudnn
|
||||
|
||||
if cudnn.is_available() and cudnn.version() < _MIN_CUDNN_FP8:
|
||||
return False
|
||||
except (ImportError, AttributeError):
|
||||
pass
|
||||
|
||||
return True
|
||||
|
||||
|
||||
__all__ = [
|
||||
"has_flashinfer",
|
||||
"flashinfer_trtllm_fp8_block_scale_moe",
|
||||
@@ -803,4 +837,5 @@ __all__ = [
|
||||
"flashinfer_fp8_blockscale_gemm",
|
||||
"should_use_flashinfer_for_blockscale_fp8_gemm",
|
||||
"is_flashinfer_fp8_blockscale_gemm_supported",
|
||||
"is_flashinfer_cudnn_fp8_prefill_attn_supported",
|
||||
]
|
||||
|
||||
@@ -26,6 +26,9 @@ import torch.nn.functional as F
|
||||
|
||||
from vllm.config import get_current_vllm_config
|
||||
from vllm.config.cache import CacheDType
|
||||
from vllm.model_executor.layers.quantization.turboquant.centroids import (
|
||||
get_centroids,
|
||||
)
|
||||
from vllm.triton_utils import triton
|
||||
from vllm.v1.attention.backend import (
|
||||
AttentionBackend,
|
||||
@@ -49,6 +52,10 @@ from vllm.v1.attention.ops.triton_turboquant_decode import (
|
||||
triton_turboquant_decode_attention,
|
||||
)
|
||||
from vllm.v1.attention.ops.triton_turboquant_store import triton_turboquant_store
|
||||
from vllm.v1.worker.workspace import (
|
||||
current_workspace_manager,
|
||||
is_workspace_manager_initialized,
|
||||
)
|
||||
|
||||
_HAS_FLASH_ATTN = is_flash_attn_varlen_func_available()
|
||||
if _HAS_FLASH_ATTN:
|
||||
@@ -335,9 +342,15 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):
|
||||
H = _build_hadamard(D, str(device))
|
||||
layer._tq_PiT = H
|
||||
layer._tq_Pi = H
|
||||
# fp16 copy for rotation in continuation prefill path
|
||||
layer._tq_Pi_half = H.to(torch.float16)
|
||||
|
||||
c = layer._tq_centroids.to(device=device, dtype=torch.float32)
|
||||
c_sorted, _ = c.sort()
|
||||
# Centroids for Lloyd-Max quantization.
|
||||
layer._tq_centroids = get_centroids(D, self.tq_config.centroid_bits).to(
|
||||
device=device, dtype=torch.float32
|
||||
)
|
||||
|
||||
c_sorted, _ = layer._tq_centroids.sort()
|
||||
layer._tq_midpoints = (c_sorted[:-1] + c_sorted[1:]) / 2
|
||||
layer._tq_cached = True
|
||||
|
||||
@@ -572,7 +585,17 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):
|
||||
|
||||
# Pre-allocate cu_seqlens for single-request flash_attn calls
|
||||
# to avoid per-request host→device tensor creation.
|
||||
_cu_2 = torch.zeros(2, device=query.device, dtype=torch.int32)
|
||||
if not hasattr(self, "_cu_2"):
|
||||
self._cu_2 = torch.zeros(2, device=query.device, dtype=torch.int32)
|
||||
# Cache arange on self (avoid per-call kernel launch).
|
||||
_max_seq = attn_metadata.max_seq_len
|
||||
_ac: torch.Tensor | None = getattr(self, "_arange_cache", None)
|
||||
if _ac is None or _ac.shape[0] <= _max_seq:
|
||||
_ac = torch.arange(
|
||||
0, _max_seq + 1, device=query.device, dtype=attn_metadata.seq_lens.dtype
|
||||
)
|
||||
self._arange_cache = _ac
|
||||
_arange_cache: torch.Tensor = _ac
|
||||
|
||||
for i in range(num_reqs):
|
||||
q_start = qsl[i]
|
||||
@@ -589,8 +612,8 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):
|
||||
if q_len == seq_len:
|
||||
# First-chunk prefill: all K/V are in the current batch.
|
||||
if _HAS_FLASH_ATTN:
|
||||
_cu_2[1] = q_len
|
||||
cu = _cu_2
|
||||
self._cu_2[1] = q_len
|
||||
cu = self._cu_2
|
||||
out = self._flash_attn_varlen(
|
||||
q=q_seq,
|
||||
k=k_seq,
|
||||
@@ -622,12 +645,8 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):
|
||||
if q_len <= _CONTINUATION_DECODE_THRESHOLD:
|
||||
# Fast path: treat each query as a decode request
|
||||
# with incremental seq_lens for causal masking.
|
||||
synth_seq_lens = torch.arange(
|
||||
cached_len + 1,
|
||||
seq_len + 1,
|
||||
device=query.device,
|
||||
dtype=attn_metadata.seq_lens.dtype,
|
||||
)
|
||||
# Slice from pre-built arange (no kernel launch)
|
||||
synth_seq_lens = _arange_cache[cached_len + 1 : seq_len + 1]
|
||||
synth_bt = attn_metadata.block_table[i : i + 1].expand(q_len, -1)
|
||||
out = triton_turboquant_decode_attention(
|
||||
query=q_seq,
|
||||
@@ -695,16 +714,17 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):
|
||||
# Reuse cached buffers to avoid per-call allocation (~16MB at 8K).
|
||||
alloc_len = math.ceil(cached_len / block_size) * block_size
|
||||
buf_shape = (1, Hk, alloc_len, D)
|
||||
k_buf = getattr(layer, "_tq_k_dequant_buf", None)
|
||||
if k_buf is None or k_buf.shape[2] < alloc_len:
|
||||
k_buf = torch.empty(buf_shape, dtype=torch.float16, device=device)
|
||||
v_buf = torch.empty(buf_shape, dtype=torch.float16, device=device)
|
||||
layer._tq_k_dequant_buf = k_buf
|
||||
layer._tq_v_dequant_buf = v_buf
|
||||
else:
|
||||
v_buf = layer._tq_v_dequant_buf
|
||||
k_cached = k_buf[:, :, :alloc_len, :].zero_()
|
||||
v_cached = v_buf[:, :, :alloc_len, :].zero_()
|
||||
# Use WorkspaceManager for dequant buffers.
|
||||
# Shared across all layers — saves 60× memory at long context.
|
||||
# Required for CUDA Graph capture (per-layer growth incompatible with CG).
|
||||
k_buf, v_buf = current_workspace_manager().get_simultaneous(
|
||||
(buf_shape, torch.float16),
|
||||
(buf_shape, torch.float16),
|
||||
)
|
||||
# Skip .zero_() — kernel writes all positions up to cached_len,
|
||||
# and we only read [:cached_len] afterwards.
|
||||
k_cached = k_buf[:, :, :alloc_len, :]
|
||||
v_cached = v_buf[:, :, :alloc_len, :]
|
||||
|
||||
grid = (alloc_len, 1 * Hk)
|
||||
_tq_full_dequant_kv[grid](
|
||||
@@ -740,29 +760,41 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):
|
||||
|
||||
# Inverse-rotate MSE keys back to original space
|
||||
if not self.tq_config.key_fp8:
|
||||
k_flat = k_cached[0, :, :cached_len, :].reshape(-1, D).float()
|
||||
k_flat = k_flat @ Pi
|
||||
k_cached_trim = (
|
||||
k_flat.to(torch.float16).reshape(Hk, cached_len, D).transpose(0, 1)
|
||||
) # (cached_len, Hk, D)
|
||||
# fp16 matmul for rotation (2× less bandwidth, uses fp16 tensor cores)
|
||||
Pi_half = layer._tq_Pi_half
|
||||
k_flat = k_cached[0, :, :cached_len, :].reshape(-1, D)
|
||||
k_flat = k_flat @ Pi_half
|
||||
k_cached_trim = k_flat.reshape(Hk, cached_len, D).transpose(
|
||||
0, 1
|
||||
) # (cached_len, Hk, D) — already fp16
|
||||
else:
|
||||
k_cached_trim = (
|
||||
k_cached[0, :, :cached_len, :].transpose(0, 1).contiguous()
|
||||
k_cached_trim = k_cached[0, :, :cached_len, :].transpose(
|
||||
0, 1
|
||||
) # (cached_len, Hk, D)
|
||||
|
||||
v_cached_trim = (
|
||||
v_cached[0, :, :cached_len, :].transpose(0, 1).contiguous()
|
||||
) # (cached_len, Hk, D)
|
||||
# Skip .contiguous() — the copy into k_full/v_full handles layout
|
||||
v_cached_trim = v_cached[0, :, :cached_len, :].transpose(0, 1)
|
||||
|
||||
# Concatenate cached + current chunk K/V (match query dtype)
|
||||
# Pre-allocate full K/V buffer, copy into slices (no cat alloc)
|
||||
qdtype = query.dtype
|
||||
k_full = torch.cat([k_cached_trim.to(qdtype), key_chunk], dim=0)
|
||||
v_full = torch.cat([v_cached_trim.to(qdtype), val_chunk], dim=0)
|
||||
k_full = torch.empty(seq_len, Hk, D, dtype=qdtype, device=device)
|
||||
v_full = torch.empty(seq_len, Hk, D, dtype=qdtype, device=device)
|
||||
k_full[:cached_len] = k_cached_trim.to(qdtype)
|
||||
k_full[cached_len:] = key_chunk
|
||||
v_full[:cached_len] = v_cached_trim.to(qdtype)
|
||||
v_full[cached_len:] = val_chunk
|
||||
|
||||
# Attention: q_len queries attending to seq_len K/V with causal mask
|
||||
if _HAS_FLASH_ATTN:
|
||||
cu_seqlens_q = torch.tensor([0, q_len], device=device, dtype=torch.int32)
|
||||
cu_seqlens_k = torch.tensor([0, seq_len], device=device, dtype=torch.int32)
|
||||
# Reuse pre-allocated cu_seqlens (avoid host→device transfer)
|
||||
if not hasattr(self, "_cu_2_q"):
|
||||
self._cu_2_q = torch.zeros(2, device=device, dtype=torch.int32)
|
||||
self._cu_2_k = torch.zeros(2, device=device, dtype=torch.int32)
|
||||
self._cu_2_q[1] = q_len
|
||||
self._cu_2_k[1] = seq_len
|
||||
cu_seqlens_q = self._cu_2_q
|
||||
cu_seqlens_k = self._cu_2_k
|
||||
return self._flash_attn_varlen(
|
||||
q=query,
|
||||
k=k_full,
|
||||
@@ -805,12 +837,23 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):
|
||||
PiT: torch.Tensor | None = None,
|
||||
layer: torch.nn.Module | None = None,
|
||||
) -> torch.Tensor:
|
||||
# Grab cached decode buffers from the layer (lazily allocated).
|
||||
# Acquire shared decode scratch buffers from WorkspaceManager.
|
||||
# Layers execute sequentially so one set of buffers is sufficient.
|
||||
# Falls back to kernel-internal allocation if workspace unavailable.
|
||||
B = query.shape[0]
|
||||
D = self.head_size
|
||||
S = self.max_num_kv_splits
|
||||
Hq = self.num_heads
|
||||
mid_o_buf = output_buf = lse_buf = None
|
||||
if layer is not None:
|
||||
mid_o_buf = getattr(layer, "_tq_mid_o_buf", None)
|
||||
output_buf = getattr(layer, "_tq_output_buf", None)
|
||||
lse_buf = getattr(layer, "_tq_lse_buf", None)
|
||||
if is_workspace_manager_initialized():
|
||||
# output_buf in query dtype — matches the in-kernel fp16 cast in stage2.
|
||||
mid_o_buf, output_buf, lse_buf = (
|
||||
current_workspace_manager().get_simultaneous(
|
||||
((B, Hq, S, D + 1), torch.float32),
|
||||
((B, Hq, D), query.dtype),
|
||||
((B, Hq), torch.float32),
|
||||
)
|
||||
)
|
||||
|
||||
result = triton_turboquant_decode_attention(
|
||||
query=query,
|
||||
|
||||
@@ -551,6 +551,7 @@ def _fwd_kernel_stage2(
|
||||
NUM_KV_SPLITS: tl.constexpr,
|
||||
BLOCK_DV: tl.constexpr,
|
||||
Lv: tl.constexpr,
|
||||
OUTPUT_FP16: tl.constexpr = 0,
|
||||
):
|
||||
cur_batch = tl.program_id(0)
|
||||
cur_head = tl.program_id(1)
|
||||
@@ -587,9 +588,12 @@ def _fwd_kernel_stage2(
|
||||
e_sum = e_sum * old_scale + exp_logic
|
||||
e_max = n_e_max
|
||||
|
||||
result = acc / e_sum
|
||||
if OUTPUT_FP16:
|
||||
result = result.to(tl.float16)
|
||||
tl.store(
|
||||
o + cur_batch * stride_obs + cur_head * stride_oh + offs_d,
|
||||
acc / e_sum,
|
||||
result,
|
||||
mask=mask_d,
|
||||
)
|
||||
lse_val = e_max + tl.log(e_sum)
|
||||
|
||||
@@ -588,10 +588,16 @@ def triton_turboquant_decode_attention(
|
||||
)
|
||||
|
||||
# Stage 2: Reduce across KV splits
|
||||
if output_buf is not None and output_buf.shape[0] >= B:
|
||||
# Output in query dtype — eliminates float16_copy kernel after stage2
|
||||
out_dtype = query.dtype
|
||||
if (
|
||||
output_buf is not None
|
||||
and output_buf.shape[0] >= B
|
||||
and output_buf.dtype == out_dtype
|
||||
):
|
||||
output = output_buf[:B, :Hq, :D]
|
||||
else:
|
||||
output = torch.empty(B, Hq, D, dtype=torch.float32, device=device)
|
||||
output = torch.empty(B, Hq, D, dtype=out_dtype, device=device)
|
||||
if buf_holder is not None:
|
||||
buf_holder._tq_output_buf = output
|
||||
if lse_buf is not None and lse_buf.shape[0] >= B:
|
||||
@@ -616,8 +622,9 @@ def triton_turboquant_decode_attention(
|
||||
NUM_KV_SPLITS=NUM_KV_SPLITS,
|
||||
BLOCK_DV=cfg["BLOCK_D"],
|
||||
Lv=D,
|
||||
OUTPUT_FP16=1 if out_dtype == torch.float16 else 0,
|
||||
num_warps=4,
|
||||
num_stages=2,
|
||||
)
|
||||
|
||||
return output.to(query.dtype)
|
||||
return output # already in query dtype
|
||||
|
||||
@@ -279,6 +279,10 @@ def flashinfer_wrapper(
|
||||
cu_seqlens: torch.Tensor | None = None,
|
||||
max_seqlen: torch.Tensor | None = None,
|
||||
sequence_lengths: torch.Tensor | None = None,
|
||||
q_scale: torch.Tensor | None = None,
|
||||
k_scale: torch.Tensor | None = None,
|
||||
v_scale: torch.Tensor | None = None,
|
||||
o_data_type: torch.dtype | None = None,
|
||||
) -> torch.Tensor:
|
||||
from flashinfer.prefill import cudnn_batch_prefill_with_kv_cache
|
||||
|
||||
@@ -318,6 +322,10 @@ def flashinfer_wrapper(
|
||||
batch_offsets_k=batch_offsets_qko,
|
||||
batch_offsets_v=batch_offsets_v,
|
||||
batch_offsets_o=batch_offsets_qko,
|
||||
q_scale=q_scale,
|
||||
k_scale=k_scale,
|
||||
v_scale=v_scale,
|
||||
o_data_type=o_data_type,
|
||||
)
|
||||
|
||||
if is_reshaped:
|
||||
@@ -335,8 +343,12 @@ def vit_flashinfer_wrapper_fake(
|
||||
cu_seqlens: torch.Tensor | None = None,
|
||||
max_seqlen: torch.Tensor | None = None,
|
||||
sequence_lengths: torch.Tensor | None = None,
|
||||
q_scale: torch.Tensor | None = None,
|
||||
k_scale: torch.Tensor | None = None,
|
||||
v_scale: torch.Tensor | None = None,
|
||||
o_data_type: torch.dtype | None = None,
|
||||
) -> torch.Tensor:
|
||||
return torch.empty_like(q)
|
||||
return torch.empty_like(q, dtype=o_data_type or q.dtype)
|
||||
|
||||
|
||||
direct_register_custom_op(
|
||||
@@ -355,7 +367,22 @@ def vit_flashinfer_wrapper(
|
||||
cu_seqlens: torch.Tensor | None = None,
|
||||
max_seqlen: torch.Tensor | None = None,
|
||||
sequence_lengths: torch.Tensor | None = None,
|
||||
q_scale: torch.Tensor | None = None,
|
||||
k_scale: torch.Tensor | None = None,
|
||||
v_scale: torch.Tensor | None = None,
|
||||
o_data_type: torch.dtype | None = None,
|
||||
) -> torch.Tensor:
|
||||
return torch.ops.vllm.flashinfer_wrapper(
|
||||
q, k, v, scale, workspace_buffer, cu_seqlens, max_seqlen, sequence_lengths
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
scale,
|
||||
workspace_buffer,
|
||||
cu_seqlens,
|
||||
max_seqlen,
|
||||
sequence_lengths,
|
||||
q_scale,
|
||||
k_scale,
|
||||
v_scale,
|
||||
o_data_type,
|
||||
)
|
||||
|
||||
@@ -34,6 +34,7 @@ class KVCacheCoordinator(ABC):
|
||||
self,
|
||||
kv_cache_config: KVCacheConfig,
|
||||
max_model_len: int,
|
||||
max_num_batched_tokens: int,
|
||||
use_eagle: bool,
|
||||
enable_caching: bool,
|
||||
enable_kv_cache_events: bool,
|
||||
@@ -65,6 +66,8 @@ class KVCacheCoordinator(ABC):
|
||||
self.single_type_managers = tuple(
|
||||
get_manager_for_kv_cache_spec(
|
||||
kv_cache_spec=kv_cache_group.kv_cache_spec,
|
||||
max_num_batched_tokens=max_num_batched_tokens,
|
||||
max_model_len=max_model_len,
|
||||
block_pool=self.block_pool,
|
||||
enable_caching=enable_caching,
|
||||
kv_cache_group_id=i,
|
||||
@@ -271,6 +274,7 @@ class KVCacheCoordinatorNoPrefixCache(KVCacheCoordinator):
|
||||
self,
|
||||
kv_cache_config: KVCacheConfig,
|
||||
max_model_len: int,
|
||||
max_num_batched_tokens: int,
|
||||
use_eagle: bool,
|
||||
enable_kv_cache_events: bool,
|
||||
dcp_world_size: int,
|
||||
@@ -281,6 +285,7 @@ class KVCacheCoordinatorNoPrefixCache(KVCacheCoordinator):
|
||||
super().__init__(
|
||||
kv_cache_config,
|
||||
max_model_len,
|
||||
max_num_batched_tokens,
|
||||
use_eagle,
|
||||
False,
|
||||
enable_kv_cache_events,
|
||||
@@ -316,6 +321,7 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator):
|
||||
self,
|
||||
kv_cache_config: KVCacheConfig,
|
||||
max_model_len: int,
|
||||
max_num_batched_tokens: int,
|
||||
use_eagle: bool,
|
||||
enable_caching: bool,
|
||||
enable_kv_cache_events: bool,
|
||||
@@ -327,6 +333,7 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator):
|
||||
super().__init__(
|
||||
kv_cache_config,
|
||||
max_model_len,
|
||||
max_num_batched_tokens,
|
||||
use_eagle,
|
||||
enable_caching,
|
||||
enable_kv_cache_events,
|
||||
@@ -381,6 +388,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
|
||||
self,
|
||||
kv_cache_config: KVCacheConfig,
|
||||
max_model_len: int,
|
||||
max_num_batched_tokens: int,
|
||||
use_eagle: bool,
|
||||
enable_caching: bool,
|
||||
enable_kv_cache_events: bool,
|
||||
@@ -392,6 +400,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
|
||||
super().__init__(
|
||||
kv_cache_config,
|
||||
max_model_len,
|
||||
max_num_batched_tokens,
|
||||
use_eagle,
|
||||
enable_caching,
|
||||
enable_kv_cache_events,
|
||||
@@ -574,6 +583,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
|
||||
def get_kv_cache_coordinator(
|
||||
kv_cache_config: KVCacheConfig,
|
||||
max_model_len: int,
|
||||
max_num_batched_tokens: int,
|
||||
use_eagle: bool,
|
||||
enable_caching: bool,
|
||||
enable_kv_cache_events: bool,
|
||||
@@ -586,6 +596,7 @@ def get_kv_cache_coordinator(
|
||||
return KVCacheCoordinatorNoPrefixCache(
|
||||
kv_cache_config,
|
||||
max_model_len,
|
||||
max_num_batched_tokens,
|
||||
use_eagle,
|
||||
enable_kv_cache_events,
|
||||
dcp_world_size=dcp_world_size,
|
||||
@@ -597,6 +608,7 @@ def get_kv_cache_coordinator(
|
||||
return UnitaryKVCacheCoordinator(
|
||||
kv_cache_config,
|
||||
max_model_len,
|
||||
max_num_batched_tokens,
|
||||
use_eagle,
|
||||
enable_caching,
|
||||
enable_kv_cache_events,
|
||||
@@ -608,6 +620,7 @@ def get_kv_cache_coordinator(
|
||||
return HybridKVCacheCoordinator(
|
||||
kv_cache_config,
|
||||
max_model_len,
|
||||
max_num_batched_tokens,
|
||||
use_eagle,
|
||||
enable_caching,
|
||||
enable_kv_cache_events,
|
||||
|
||||
@@ -109,6 +109,7 @@ class KVCacheManager:
|
||||
kv_cache_config: KVCacheConfig,
|
||||
max_model_len: int,
|
||||
hash_block_size: int,
|
||||
max_num_batched_tokens: int | None = None,
|
||||
enable_caching: bool = True,
|
||||
use_eagle: bool = False,
|
||||
log_stats: bool = False,
|
||||
@@ -118,6 +119,11 @@ class KVCacheManager:
|
||||
metrics_collector: KVCacheMetricsCollector | None = None,
|
||||
) -> None:
|
||||
self.max_model_len = max_model_len
|
||||
# When unset, fall back to `max_model_len` so the recycling-aware cap
|
||||
# collapses to the prior (uncapped) admission behavior. The scheduler
|
||||
# always supplies the real value at runtime.
|
||||
if max_num_batched_tokens is None:
|
||||
max_num_batched_tokens = max_model_len
|
||||
|
||||
self.enable_caching = enable_caching
|
||||
self.use_eagle = use_eagle
|
||||
@@ -131,6 +137,7 @@ class KVCacheManager:
|
||||
self.coordinator = get_kv_cache_coordinator(
|
||||
kv_cache_config=kv_cache_config,
|
||||
max_model_len=self.max_model_len,
|
||||
max_num_batched_tokens=max_num_batched_tokens,
|
||||
use_eagle=self.use_eagle,
|
||||
enable_caching=self.enable_caching,
|
||||
enable_kv_cache_events=enable_kv_cache_events,
|
||||
|
||||
@@ -228,6 +228,7 @@ class Scheduler(SchedulerInterface):
|
||||
self.kv_cache_manager = KVCacheManager(
|
||||
kv_cache_config=kv_cache_config,
|
||||
max_model_len=self.max_model_len,
|
||||
max_num_batched_tokens=self.scheduler_config.max_num_batched_tokens,
|
||||
enable_caching=self.cache_config.enable_prefix_caching,
|
||||
use_eagle=self.use_eagle,
|
||||
log_stats=self.log_stats,
|
||||
|
||||
@@ -41,6 +41,7 @@ class SingleTypeKVCacheManager(ABC):
|
||||
kv_cache_group_id: int,
|
||||
dcp_world_size: int = 1,
|
||||
pcp_world_size: int = 1,
|
||||
max_admission_blocks_per_request: int | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Initializes the SingleTypeKVCacheManager.
|
||||
@@ -48,6 +49,12 @@ class SingleTypeKVCacheManager(ABC):
|
||||
kv_cache_spec: The kv_cache_spec for this manager.
|
||||
block_pool: The block pool.
|
||||
kv_cache_group_id: The id of the kv cache group of this manager.
|
||||
max_admission_blocks_per_request: Recycling-aware per-request
|
||||
block cap used by `get_num_blocks_to_allocate`. Only set for
|
||||
spec types that recycle blocks across chunks (SWA,
|
||||
chunked-local); `None` (the default) means no cap, which is
|
||||
correct for full-attention-style specs that hold every
|
||||
block until the request finishes.
|
||||
"""
|
||||
self.block_size = kv_cache_spec.block_size
|
||||
self.dcp_world_size = dcp_world_size
|
||||
@@ -57,6 +64,7 @@ class SingleTypeKVCacheManager(ABC):
|
||||
self.kv_cache_spec = kv_cache_spec
|
||||
self.block_pool = block_pool
|
||||
self.enable_caching = enable_caching
|
||||
self._max_admission_blocks_per_request = max_admission_blocks_per_request
|
||||
self.new_block_ids: list[int] = []
|
||||
|
||||
# Mapping from request ID to blocks to track the blocks allocated
|
||||
@@ -105,6 +113,19 @@ class SingleTypeKVCacheManager(ABC):
|
||||
"""
|
||||
|
||||
num_required_blocks = cdiv(num_tokens, self.block_size)
|
||||
if self._max_admission_blocks_per_request is not None:
|
||||
# Recycling-aware specs (SWA, chunked-local) cap the per-request
|
||||
# reservation here so admission matches the startup pool sizer
|
||||
# (`SlidingWindowSpec.max_admission_blocks_per_request` / its
|
||||
# chunked-local counterpart). `remove_skipped_blocks` runs from
|
||||
# `allocate_slots` before each chunk's `get_num_blocks_to_allocate`,
|
||||
# so per-request peak real-held blocks <= this cap, which keeps
|
||||
# `sum(reservations) <= pool` <=> `sum(peak_real_held) <= pool`.
|
||||
# Drift between the two would re-introduce the deadlock from
|
||||
# issue #39734 or, worse, mid-prefill OOM.
|
||||
num_required_blocks = min(
|
||||
num_required_blocks, self._max_admission_blocks_per_request
|
||||
)
|
||||
num_req_blocks = len(self.req_to_blocks.get(request_id, ()))
|
||||
|
||||
if request_id in self.num_cached_block:
|
||||
@@ -1126,8 +1147,21 @@ spec_manager_map: dict[type[KVCacheSpec], type[SingleTypeKVCacheManager]] = {
|
||||
|
||||
|
||||
def get_manager_for_kv_cache_spec(
|
||||
kv_cache_spec: KVCacheSpec, **kwargs
|
||||
kv_cache_spec: KVCacheSpec,
|
||||
max_num_batched_tokens: int,
|
||||
max_model_len: int,
|
||||
**kwargs,
|
||||
) -> SingleTypeKVCacheManager:
|
||||
manager_class = spec_manager_map[type(kv_cache_spec)]
|
||||
# SlidingWindow / ChunkedLocalAttention managers recycle blocks across
|
||||
# chunks; the runtime admission cap must match the recycling-aware bound
|
||||
# the startup pool sizer uses (single source of truth: the spec method).
|
||||
if isinstance(kv_cache_spec, (SlidingWindowSpec, ChunkedLocalAttentionSpec)):
|
||||
kwargs["max_admission_blocks_per_request"] = (
|
||||
kv_cache_spec.max_admission_blocks_per_request(
|
||||
max_num_batched_tokens=max_num_batched_tokens,
|
||||
max_model_len=max_model_len,
|
||||
)
|
||||
)
|
||||
manager = manager_class(kv_cache_spec, **kwargs)
|
||||
return manager
|
||||
|
||||
@@ -376,19 +376,28 @@ class MLAAttentionSpec(FullAttentionSpec):
|
||||
class ChunkedLocalAttentionSpec(AttentionSpec):
|
||||
attention_chunk_size: int
|
||||
|
||||
def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
|
||||
max_model_len = vllm_config.model_config.max_model_len
|
||||
max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens
|
||||
def max_admission_blocks_per_request(
|
||||
self, max_num_batched_tokens: int, max_model_len: int
|
||||
) -> int:
|
||||
"""Per-request admission cap, in blocks.
|
||||
|
||||
# During chunked prefill, we allocate KV cache for at most
|
||||
# `self.attention_chunk_size` computed tokens plus the newly scheduled
|
||||
# tokens. And we won't allocate KV cache for more than `max_model_len`
|
||||
# tokens.
|
||||
Single source of truth for both startup pool sizing
|
||||
(`max_memory_usage_bytes`) and the runtime admission gate, so requests
|
||||
admitted by startup can also be admitted at runtime.
|
||||
"""
|
||||
# During chunked prefill, we hold KV for at most one chunk window.
|
||||
num_tokens = min(
|
||||
self.attention_chunk_size + max_num_batched_tokens, max_model_len
|
||||
)
|
||||
return cdiv(num_tokens, self.block_size)
|
||||
|
||||
return cdiv(num_tokens, self.block_size) * self.page_size_bytes
|
||||
def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
|
||||
max_model_len = vllm_config.model_config.max_model_len
|
||||
max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens
|
||||
max_blocks = self.max_admission_blocks_per_request(
|
||||
max_num_batched_tokens=max_num_batched_tokens, max_model_len=max_model_len
|
||||
)
|
||||
return max_blocks * self.page_size_bytes
|
||||
|
||||
|
||||
@dataclass(frozen=True, kw_only=True)
|
||||
@@ -409,26 +418,38 @@ class SlidingWindowSpec(AttentionSpec):
|
||||
* get_dtype_size(self.dtype)
|
||||
)
|
||||
|
||||
def max_admission_blocks_per_request(
|
||||
self, max_num_batched_tokens: int, max_model_len: int
|
||||
) -> int:
|
||||
"""Per-request admission cap, in blocks.
|
||||
|
||||
Single source of truth for both startup pool sizing
|
||||
(`max_memory_usage_bytes`) and the runtime admission gate. Per-request
|
||||
real-held blocks plateau at this bound because
|
||||
`SlidingWindowManager.remove_skipped_blocks` runs from `allocate_slots`
|
||||
before each chunk's `get_num_blocks_to_allocate`.
|
||||
"""
|
||||
# During chunked prefill, we hold KV for the last `sliding_window-1`
|
||||
# computed tokens plus the newly scheduled tokens, and never more
|
||||
# than `max_model_len`.
|
||||
num_tokens = min(
|
||||
self.sliding_window - 1 + max_num_batched_tokens, max_model_len
|
||||
)
|
||||
# +1 because the sliding window may not start from the beginning of
|
||||
# the block. E.g. block size 4 and num_token 4 needs two blocks
|
||||
# [XXCD][EF] to store the 6-token window [CDEF].
|
||||
return cdiv(num_tokens, self.block_size) + 1
|
||||
|
||||
def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
|
||||
assert vllm_config.parallel_config.decode_context_parallel_size == 1, (
|
||||
"DCP not support sliding window."
|
||||
)
|
||||
max_model_len = vllm_config.model_config.max_model_len
|
||||
max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens
|
||||
|
||||
# During chunked prefill, we allocate KV cache for the last
|
||||
# `self.sliding_window-1` computed tokens plus the newly scheduled
|
||||
# tokens. And we won't allocate KV cache for more than `max_model_len`
|
||||
# tokens.
|
||||
num_tokens = min(
|
||||
self.sliding_window - 1 + max_num_batched_tokens, max_model_len
|
||||
max_blocks = self.max_admission_blocks_per_request(
|
||||
max_num_batched_tokens=max_num_batched_tokens, max_model_len=max_model_len
|
||||
)
|
||||
|
||||
# +1 here because the sliding window may not start from the beginning
|
||||
# of the block. For example, if the block size is 4 and num_token
|
||||
# is 4, we need two blocks [XXCD] [EF] to store the sliding
|
||||
# window [CDEF] of 6 tokens.
|
||||
return (cdiv(num_tokens, self.block_size) + 1) * self.page_size_bytes
|
||||
return max_blocks * self.page_size_bytes
|
||||
|
||||
|
||||
@dataclass(frozen=True, kw_only=True)
|
||||
|
||||
@@ -110,6 +110,9 @@ class SimpleCPUOffloadScheduler:
|
||||
self.cpu_coordinator: KVCacheCoordinator = get_kv_cache_coordinator(
|
||||
kv_cache_config=self.cpu_kv_cache_config,
|
||||
max_model_len=vllm_config.model_config.max_model_len,
|
||||
max_num_batched_tokens=(
|
||||
vllm_config.scheduler_config.max_num_batched_tokens
|
||||
),
|
||||
use_eagle=False,
|
||||
enable_caching=True,
|
||||
enable_kv_cache_events=self.enable_kv_cache_events,
|
||||
|
||||
@@ -99,7 +99,7 @@ class EagleSpeculator:
|
||||
)
|
||||
|
||||
self.draft_logits: torch.Tensor | None = None
|
||||
if self.speculative_config.rejection_sample_method == "probabilistic":
|
||||
if self.speculative_config.draft_sample_method == "gumbel":
|
||||
self.draft_logits = torch.zeros(
|
||||
self.max_num_reqs,
|
||||
self.num_speculative_steps,
|
||||
@@ -220,6 +220,28 @@ class EagleSpeculator:
|
||||
last_hidden_states, hidden_states = ret_hidden_states
|
||||
return last_hidden_states, hidden_states
|
||||
|
||||
def _sample_draft(
|
||||
self,
|
||||
logits: torch.Tensor,
|
||||
idx_mapping: torch.Tensor,
|
||||
pos: torch.Tensor,
|
||||
step: int,
|
||||
) -> torch.Tensor:
|
||||
if self.draft_logits is not None:
|
||||
# NOTE(woosuk): We must add 1 to the positions to match the Gumbel noise
|
||||
# used for draft and target sampling.
|
||||
return gumbel_sample(
|
||||
logits,
|
||||
idx_mapping,
|
||||
self.temperature,
|
||||
self.seeds,
|
||||
pos + 1,
|
||||
apply_temperature=True,
|
||||
processed_logits_out=self.draft_logits[:, step],
|
||||
)
|
||||
else:
|
||||
return logits.argmax(dim=-1)
|
||||
|
||||
def prefill(
|
||||
self,
|
||||
num_reqs: int,
|
||||
@@ -245,18 +267,11 @@ class EagleSpeculator:
|
||||
sample_hidden_states = last_hidden_states[last_token_indices]
|
||||
logits = self.model.compute_logits(sample_hidden_states)
|
||||
|
||||
# NOTE(woosuk): We must add 1 to the positions to match the Gumbel noise
|
||||
# used for draft and target sampling.
|
||||
self.draft_tokens[:num_reqs, 0] = gumbel_sample(
|
||||
self.draft_tokens[:num_reqs, 0] = self._sample_draft(
|
||||
logits,
|
||||
idx_mapping,
|
||||
self.temperature,
|
||||
self.seeds,
|
||||
pos + 1,
|
||||
apply_temperature=True,
|
||||
processed_logits_out=self.draft_logits[:, 0]
|
||||
if self.draft_logits is not None
|
||||
else None,
|
||||
pos,
|
||||
step=0,
|
||||
)
|
||||
self.hidden_states[:num_reqs] = hidden_states[last_token_indices]
|
||||
self.input_buffers.positions[:num_reqs] = pos
|
||||
@@ -286,18 +301,11 @@ class EagleSpeculator:
|
||||
hidden_states = hidden_states[:num_reqs]
|
||||
logits = self.model.compute_logits(last_hidden_states)
|
||||
|
||||
# NOTE(woosuk): We must add 1 to the positions to match the Gumbel noise
|
||||
# used for draft and target sampling.
|
||||
draft_tokens = gumbel_sample(
|
||||
draft_tokens = self._sample_draft(
|
||||
logits,
|
||||
idx_mapping,
|
||||
self.temperature,
|
||||
self.seeds,
|
||||
pos + 1,
|
||||
apply_temperature=True,
|
||||
processed_logits_out=self.draft_logits[:, step]
|
||||
if self.draft_logits is not None
|
||||
else None,
|
||||
pos,
|
||||
step=step,
|
||||
)
|
||||
self.draft_tokens[:num_reqs, step] = draft_tokens
|
||||
|
||||
|
||||
@@ -45,7 +45,7 @@ def _compute_global_lse(
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _compute_block_max_and_sumexp_kernel(
|
||||
def _compute_block_stats_kernel(
|
||||
# [num_logits, num_blocks]
|
||||
target_local_argmax_ptr,
|
||||
target_local_argmax_stride,
|
||||
@@ -77,6 +77,7 @@ def _compute_block_max_and_sumexp_kernel(
|
||||
vocab_size,
|
||||
num_speculative_steps,
|
||||
BLOCK_SIZE: tl.constexpr,
|
||||
HAS_DRAFT_LOGITS: tl.constexpr,
|
||||
):
|
||||
logit_idx = tl.program_id(0)
|
||||
draft_step_idx = tl.load(expanded_local_pos_ptr + logit_idx)
|
||||
@@ -112,24 +113,6 @@ def _compute_block_max_and_sumexp_kernel(
|
||||
value,
|
||||
)
|
||||
else:
|
||||
# Get local draft max and summed exponentials.
|
||||
draft_logits = tl.load(
|
||||
draft_logits_ptr
|
||||
+ req_state_idx * draft_logits_stride_0
|
||||
+ draft_step_idx * draft_logits_stride_1
|
||||
+ block_offsets,
|
||||
mask=mask,
|
||||
other=float("-inf"),
|
||||
).to(tl.float32)
|
||||
draft_max, draft_sumexp = _compute_block_max_and_sumexp(draft_logits)
|
||||
tl.store(
|
||||
draft_local_max_ptr + logit_idx * draft_local_max_stride + block_idx,
|
||||
draft_max,
|
||||
)
|
||||
tl.store(
|
||||
draft_local_sumexp_ptr + logit_idx * draft_local_sumexp_stride + block_idx,
|
||||
draft_sumexp,
|
||||
)
|
||||
# Get local target max and summed exponentials.
|
||||
target_logits = tl.load(
|
||||
target_logits_ptr + logit_idx * target_logits_stride + block_offsets,
|
||||
@@ -147,6 +130,27 @@ def _compute_block_max_and_sumexp_kernel(
|
||||
+ block_idx,
|
||||
target_sumexp,
|
||||
)
|
||||
if HAS_DRAFT_LOGITS:
|
||||
# Get local draft max and summed exponentials.
|
||||
draft_logits = tl.load(
|
||||
draft_logits_ptr
|
||||
+ req_state_idx * draft_logits_stride_0
|
||||
+ draft_step_idx * draft_logits_stride_1
|
||||
+ block_offsets,
|
||||
mask=mask,
|
||||
other=float("-inf"),
|
||||
).to(tl.float32)
|
||||
draft_max, draft_sumexp = _compute_block_max_and_sumexp(draft_logits)
|
||||
tl.store(
|
||||
draft_local_max_ptr + logit_idx * draft_local_max_stride + block_idx,
|
||||
draft_max,
|
||||
)
|
||||
tl.store(
|
||||
draft_local_sumexp_ptr
|
||||
+ logit_idx * draft_local_sumexp_stride
|
||||
+ block_idx,
|
||||
draft_sumexp,
|
||||
)
|
||||
|
||||
|
||||
@triton.jit
|
||||
@@ -196,6 +200,7 @@ def _probabilistic_rejection_kernel(
|
||||
pos_ptr,
|
||||
vocab_num_blocks,
|
||||
PADDED_VOCAB_NUM_BLOCKS: tl.constexpr,
|
||||
HAS_DRAFT_LOGITS: tl.constexpr,
|
||||
):
|
||||
req_idx = tl.program_id(0)
|
||||
req_state_idx = tl.load(idx_mapping_ptr + req_idx)
|
||||
@@ -238,12 +243,6 @@ def _probabilistic_rejection_kernel(
|
||||
target_logit = tl.load(
|
||||
target_logits_ptr + logit_idx * target_logits_stride + draft_sampled
|
||||
).to(tl.float32)
|
||||
draft_logit = tl.load(
|
||||
draft_logits_ptr
|
||||
+ req_state_idx * draft_logits_stride_0
|
||||
+ i * draft_logits_stride_1
|
||||
+ draft_sampled
|
||||
).to(tl.float32)
|
||||
target_lse = _compute_global_lse(
|
||||
target_local_max_ptr,
|
||||
target_local_max_stride,
|
||||
@@ -253,19 +252,29 @@ def _probabilistic_rejection_kernel(
|
||||
vocab_num_blocks,
|
||||
PADDED_VOCAB_NUM_BLOCKS,
|
||||
)
|
||||
draft_lse = _compute_global_lse(
|
||||
draft_local_max_ptr,
|
||||
draft_local_max_stride,
|
||||
draft_local_sumexp_ptr,
|
||||
draft_local_sumexp_stride,
|
||||
logit_idx,
|
||||
vocab_num_blocks,
|
||||
PADDED_VOCAB_NUM_BLOCKS,
|
||||
)
|
||||
target_log_prob = target_logit - target_lse
|
||||
draft_log_prob = draft_logit - draft_lse
|
||||
pos = tl.load(pos_ptr + logit_idx)
|
||||
u = tl_rand64(seed, pos, includes_zero=False)
|
||||
if HAS_DRAFT_LOGITS:
|
||||
draft_logit = tl.load(
|
||||
draft_logits_ptr
|
||||
+ req_state_idx * draft_logits_stride_0
|
||||
+ i * draft_logits_stride_1
|
||||
+ draft_sampled
|
||||
).to(tl.float32)
|
||||
draft_lse = _compute_global_lse(
|
||||
draft_local_max_ptr,
|
||||
draft_local_max_stride,
|
||||
draft_local_sumexp_ptr,
|
||||
draft_local_sumexp_stride,
|
||||
logit_idx,
|
||||
vocab_num_blocks,
|
||||
PADDED_VOCAB_NUM_BLOCKS,
|
||||
)
|
||||
draft_log_prob = draft_logit - draft_lse
|
||||
else:
|
||||
# One-hot draft: q(draft_token) = 1, log_q = 0.
|
||||
draft_log_prob = 0
|
||||
# Probability ratio test: p(x) > u * q(x)
|
||||
# Equivalent log form: log_p(x) > log(u) + log_q(x)
|
||||
accepted &= target_log_prob > tl.log(u) + draft_log_prob
|
||||
@@ -301,6 +310,8 @@ def _resample_kernel(
|
||||
cu_num_logits_ptr,
|
||||
# [num_logits]
|
||||
expanded_idx_mapping_ptr,
|
||||
# [num_logits]
|
||||
draft_sampled_ptr,
|
||||
# [max_num_reqs]
|
||||
temp_ptr,
|
||||
# [max_num_reqs]
|
||||
@@ -309,6 +320,7 @@ def _resample_kernel(
|
||||
pos_ptr,
|
||||
vocab_size,
|
||||
BLOCK_SIZE: tl.constexpr,
|
||||
HAS_DRAFT_LOGITS: tl.constexpr,
|
||||
):
|
||||
req_idx = tl.program_id(0)
|
||||
resample_idx = tl.load(rejected_step_ptr + req_idx)
|
||||
@@ -327,22 +339,17 @@ def _resample_kernel(
|
||||
block_idx = tl.program_id(1)
|
||||
block = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
|
||||
mask = block < vocab_size
|
||||
target_logits = tl.load(
|
||||
target_logits_ptr + resample_token_idx * target_logits_stride + block,
|
||||
mask=mask,
|
||||
other=float("-inf"),
|
||||
).to(tl.float32)
|
||||
|
||||
# Compute the residual logits to resample the rejected token
|
||||
# from. In the case of no rejections (bonus token), we directly
|
||||
# use the target logits.
|
||||
# Compute the residual logits to resample the rejected token from.
|
||||
if is_bonus:
|
||||
residual_logits = tl.load(
|
||||
target_logits_ptr + resample_token_idx * target_logits_stride + block,
|
||||
mask=mask,
|
||||
other=float("-inf"),
|
||||
).to(tl.float32)
|
||||
else:
|
||||
target_logits = tl.load(
|
||||
target_logits_ptr + resample_token_idx * target_logits_stride + block,
|
||||
mask=mask,
|
||||
other=float("-inf"),
|
||||
).to(tl.float32)
|
||||
# Bonus token (no rejections). Directly use the target logits.
|
||||
residual_logits = target_logits
|
||||
elif HAS_DRAFT_LOGITS:
|
||||
draft_logits = tl.load(
|
||||
draft_logits_ptr
|
||||
+ req_state_idx * draft_logits_stride_0
|
||||
@@ -365,6 +372,15 @@ def _resample_kernel(
|
||||
target_log_probs + tl.log(1 - ratio),
|
||||
float("-inf"),
|
||||
).to(tl.float32)
|
||||
else:
|
||||
# One-hot draft. The residual is just the target distribution with
|
||||
# the rejected draft token probability zeroed out.
|
||||
rejected_draft_token = tl.load(draft_sampled_ptr + resample_token_idx + 1)
|
||||
residual_logits = tl.where(
|
||||
block != rejected_draft_token,
|
||||
target_logits,
|
||||
float("-inf"),
|
||||
).to(tl.float32)
|
||||
|
||||
# Resample the rejected/bonus token.
|
||||
value, idx = gumbel_block_argmax(
|
||||
@@ -456,7 +472,7 @@ def probabilistic_rejection_sample(
|
||||
# [num_logits, V]
|
||||
target_logits: torch.Tensor,
|
||||
# [max_num_reqs, num_speculative_steps, V]
|
||||
draft_logits: torch.Tensor,
|
||||
draft_logits: torch.Tensor | None,
|
||||
# [num_logits]
|
||||
draft_sampled: torch.Tensor,
|
||||
# [num_reqs + 1]
|
||||
@@ -477,9 +493,17 @@ def probabilistic_rejection_sample(
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
num_reqs = cu_num_logits.shape[0] - 1
|
||||
num_logits, vocab_size = target_logits.shape
|
||||
has_draft_logits = draft_logits is not None
|
||||
|
||||
# Gather draft logits, compute target argmax for greedy, and
|
||||
# compute per-block LSE and max for non-greedy requests.
|
||||
if draft_logits is None:
|
||||
# When draft_logits is None, create a dummy tensor so that Triton
|
||||
# kernel signatures receive valid pointers/strides. The kernels
|
||||
# will never read from it when HAS_DRAFT_LOGITS=False.
|
||||
draft_logits = target_logits.new_empty(1, 1, 1)
|
||||
|
||||
# Compute the block-level logits stats, such as target argmax
|
||||
# (for greedy requests), and target max + softmax exponential
|
||||
# (for non-greedy requests).
|
||||
VOCAB_BLOCK_SIZE = 8192
|
||||
vocab_num_blocks = triton.cdiv(vocab_size, VOCAB_BLOCK_SIZE)
|
||||
padded_vocab_num_blocks = triton.next_power_of_2(vocab_num_blocks)
|
||||
@@ -498,7 +522,7 @@ def probabilistic_rejection_sample(
|
||||
draft_local_sumexp = target_logits.new_empty(
|
||||
num_logits, vocab_num_blocks, dtype=torch.float32
|
||||
)
|
||||
_compute_block_max_and_sumexp_kernel[(num_logits, vocab_num_blocks)](
|
||||
_compute_block_stats_kernel[(num_logits, vocab_num_blocks)](
|
||||
target_local_argmax,
|
||||
target_local_argmax.stride(0),
|
||||
target_local_max,
|
||||
@@ -520,6 +544,7 @@ def probabilistic_rejection_sample(
|
||||
vocab_size,
|
||||
num_speculative_steps,
|
||||
BLOCK_SIZE=VOCAB_BLOCK_SIZE,
|
||||
HAS_DRAFT_LOGITS=has_draft_logits,
|
||||
)
|
||||
|
||||
# Sample up until the first rejected/bonus token, and store
|
||||
@@ -559,6 +584,7 @@ def probabilistic_rejection_sample(
|
||||
pos,
|
||||
vocab_num_blocks,
|
||||
PADDED_VOCAB_NUM_BLOCKS=padded_vocab_num_blocks,
|
||||
HAS_DRAFT_LOGITS=has_draft_logits,
|
||||
num_warps=1,
|
||||
)
|
||||
|
||||
@@ -587,11 +613,13 @@ def probabilistic_rejection_sample(
|
||||
num_sampled,
|
||||
cu_num_logits,
|
||||
expanded_idx_mapping,
|
||||
draft_sampled,
|
||||
temperature,
|
||||
seed,
|
||||
pos,
|
||||
vocab_size,
|
||||
BLOCK_SIZE=RESAMPLE_BLOCK_SIZE,
|
||||
HAS_DRAFT_LOGITS=has_draft_logits,
|
||||
)
|
||||
|
||||
# Insert the resampled tokens into the output sampled.
|
||||
|
||||
@@ -20,63 +20,6 @@ from vllm.v1.worker.gpu.spec_decode.synthetic_rejection_sampler_utils import (
|
||||
)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _strict_rejection_sample_kernel(
|
||||
sampled_ptr, # [num_reqs, num_speculative_steps + 1]
|
||||
sampled_stride,
|
||||
num_sampled_ptr, # [num_reqs]
|
||||
target_sampled_ptr, # [num_draft_tokens + num_reqs]
|
||||
input_ids_ptr, # [num_draft_tokens + num_reqs]
|
||||
cu_num_logits_ptr, # [num_reqs + 1]
|
||||
):
|
||||
req_idx = tl.program_id(0)
|
||||
start_idx = tl.load(cu_num_logits_ptr + req_idx)
|
||||
end_idx = tl.load(cu_num_logits_ptr + req_idx + 1)
|
||||
num_tokens = end_idx - start_idx
|
||||
|
||||
num_sampled = 0
|
||||
rejected = False
|
||||
for i in range(num_tokens - 1):
|
||||
if not rejected:
|
||||
target_sampled = tl.load(target_sampled_ptr + start_idx + i)
|
||||
draft_sampled = tl.load(input_ids_ptr + start_idx + i + 1)
|
||||
tl.store(sampled_ptr + req_idx * sampled_stride + i, target_sampled)
|
||||
num_sampled += 1
|
||||
if target_sampled != draft_sampled:
|
||||
rejected = True
|
||||
if not rejected:
|
||||
target_sampled = tl.load(target_sampled_ptr + start_idx + num_tokens - 1)
|
||||
tl.store(
|
||||
sampled_ptr + req_idx * sampled_stride + num_tokens - 1, target_sampled
|
||||
)
|
||||
num_sampled += 1
|
||||
tl.store(num_sampled_ptr + req_idx, num_sampled)
|
||||
|
||||
|
||||
def strict_rejection_sample(
|
||||
# [num_draft_tokens + num_reqs]
|
||||
target_sampled: torch.Tensor,
|
||||
# [num_draft_tokens + num_reqs]
|
||||
draft_sampled: torch.Tensor,
|
||||
# [num_reqs + 1]
|
||||
cu_num_logits: torch.Tensor,
|
||||
num_speculative_steps,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
num_reqs = cu_num_logits.shape[0] - 1
|
||||
sampled = target_sampled.new_empty(num_reqs, num_speculative_steps + 1)
|
||||
num_sampled = target_sampled.new_empty(num_reqs, dtype=torch.int32)
|
||||
_strict_rejection_sample_kernel[(num_reqs,)](
|
||||
sampled,
|
||||
sampled.stride(0),
|
||||
num_sampled,
|
||||
target_sampled,
|
||||
draft_sampled,
|
||||
cu_num_logits,
|
||||
num_warps=1,
|
||||
)
|
||||
return sampled, num_sampled
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _flatten_sampled_kernel(
|
||||
# [num_logits]
|
||||
@@ -163,17 +106,7 @@ class RejectionSampler:
|
||||
# that num_nans is computed before applying penalties and temperature.
|
||||
num_nans = get_num_nans(logits) if self.sampler.compute_nans else None
|
||||
|
||||
if self.rejection_sample_method == "strict":
|
||||
sampler_output = self.sampler(logits, input_batch)
|
||||
logprobs_tensors = sampler_output.logprobs_tensors
|
||||
sampled, num_sampled = strict_rejection_sample(
|
||||
sampler_output.sampled_token_ids.view(-1),
|
||||
draft_sampled,
|
||||
input_batch.cu_num_logits,
|
||||
self.num_speculative_steps,
|
||||
)
|
||||
elif self.rejection_sample_method == "probabilistic":
|
||||
assert draft_logits is not None
|
||||
if self.rejection_sample_method == "standard":
|
||||
pos = input_batch.positions[input_batch.logits_indices]
|
||||
processed_logits = self.sampler.apply_sampling_params(
|
||||
logits,
|
||||
|
||||
Reference in New Issue
Block a user