Compare commits

...
Author SHA1 Message Date
ebf862c351 Add system_fingerprint field to OpenAI-compatible API responses (#40537)
Co-authored-by: Claude <noreply@anthropic.com>
2026-04-27 16:17:52 +08:00
wang.yuqiandGitHub 8d8062d0a7 [Examples] Resettle generate examples. (#36464)
Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io>
2026-04-27 07:48:37 +00:00
985961345a [Bugfix] Install libcublas-dev in Dockerfile for FlashInfer CuTe DSL JIT (#39855)
Signed-off-by: esmeetu <jasonailu87@gmail.com>
Co-authored-by: Roger Wang <hey@rogerw.io>
2026-04-27 15:47:39 +08:00
Yongye ZhuandGitHub 706a04d34b [DSV4] Add silu clamp limit to shared expert (#40950)
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
2026-04-27 00:37:43 -07:00
Isotr0pyandGitHub 22631f80a0 [Bugfix] Remove invalid deepstack boundary check for Qwen3-VL (#40932)
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
2026-04-27 07:27:06 +00:00
BhoomitandGitHub 2cc008e7b4 [Attention][TurboQuant] Share dequant buffers, eliminate float16_copy (#40941)
Signed-off-by: Bhoomit Vasani <bhoomit.2010@gmail.com>
Signed-off-by: Vasani Bhoomit <bhoomit.2010@gmail.com>
2026-04-27 13:48:36 +08:00
5d5c776444 [Perf] FP8 FlashInfer Attn for ViT (#38065)
Signed-off-by: Zhanda Zhu <zhandazhu@gmail.com>
Co-authored-by: Yubo Gao <ybgao-nvidia@users.noreply.github.com>
2026-04-27 13:44:15 +08:00
ojhaanshikaandGitHub 592ae6805c Cutlass W4A16 (Machete) Tests (#35450)
Signed-off-by: Anshika Ojha <anshikao@nvidia.com>
2026-04-27 05:15:29 +00:00
7b1bc0a3eb [Bugfix] Cap SWA/chunked-local runtime admission to startup pool-sizing bound (#40946)
Signed-off-by: Dao Le <Dao007forever@gmail.com>
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Nick Hill <nickhill123@gmail.com>
2026-04-27 04:33:13 +00:00
Silu PandaandGitHub c0879d9483 [Tests] Gate Isaac under Transformers v5 (#40907)
Signed-off-by: Silu Panda <31051721+SiluPanda@users.noreply.github.com>
2026-04-26 19:26:51 -07:00
Giancarlo DelfinandGitHub f5f9878514 [Model Runner V2] Fix rejection sampling acceptance rate gap vs MRV1 (#40651)
Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai>
2026-04-26 19:12:08 -07:00
youkaichaoGitHubClaudegemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
2ce95a761b Auto-disable expandable_segments around cumem memory pool (#40812)
Signed-off-by: youkaichao <youkaichao@gmail.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
2026-04-27 09:37:22 +08:00
93 changed files with 3045 additions and 447 deletions
+14 -14
View File
@@ -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
+2
View File
@@ -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
+4 -4
View File
@@ -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
+4 -4
View File
@@ -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
+5 -5
View File
@@ -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
View File
@@ -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
View File
@@ -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 {
+2
View File
@@ -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);
+6
View File
@@ -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
View File
@@ -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
+7 -7
View File
@@ -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.
+1
View File
@@ -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
+109
View File
@@ -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.
+1 -1
View File
@@ -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
+3 -3
View File
@@ -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
@@ -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
```
@@ -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
+18
View File
@@ -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")
+80
View File
@@ -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",
[
+279
View File
@@ -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}"
+124
View File
@@ -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)
+251
View File
@@ -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
+12
View File
@@ -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",
+185
View File
@@ -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]}'"
)
+108
View File
@@ -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,
)
+12
View File
@@ -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,
+51
View File
@@ -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
+15 -6
View File
@@ -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,
+40 -32
View File
@@ -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:
"""
+28
View File
@@ -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
+13 -1
View File
@@ -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
+12 -1
View File
@@ -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,
)
+13
View File
@@ -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,
+84
View File
@@ -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)
+3
View File
@@ -0,0 +1,3 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Triton kernel implementations."""
+180
View File
@@ -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)
+40
View File
@@ -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,
}
+9 -5
View File
@@ -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):
+58 -3
View File
@@ -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
+9 -11
View File
@@ -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
+32 -28
View File
@@ -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"]
+35
View File
@@ -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",
]
+83 -40
View File
@@ -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
+29 -2
View File
@@ -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,
)
+13
View File
@@ -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,
+7
View File
@@ -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,
+1
View File
@@ -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,
+35 -1
View File
@@ -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
+42 -21
View File
@@ -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)
+3
View File
@@ -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,