forked from Karylab-cklius/vllm
Compare commits
4
Commits
tml-inkling
...
v0.20.0
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
88d34c6409 | ||
|
|
b8160878f0 | ||
|
|
84c276d7ea | ||
|
|
5eb3657578 |
+82
-25
@@ -11,29 +11,74 @@
|
||||
namespace vllm {
|
||||
|
||||
template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&),
|
||||
bool act_first>
|
||||
bool act_first, bool HAS_CLAMP>
|
||||
__device__ __forceinline__ scalar_t compute(const scalar_t& x,
|
||||
const scalar_t& y) {
|
||||
return act_first ? ACT_FN(x) * y : x * ACT_FN(y);
|
||||
const scalar_t& y,
|
||||
const float limit) {
|
||||
if constexpr (act_first) {
|
||||
scalar_t gate = x;
|
||||
scalar_t up = y;
|
||||
if constexpr (HAS_CLAMP) {
|
||||
gate = (scalar_t)fminf((float)gate, limit);
|
||||
up = (scalar_t)fmaxf(fminf((float)up, limit), -limit);
|
||||
}
|
||||
return ACT_FN(gate) * up;
|
||||
} else {
|
||||
scalar_t gate = x;
|
||||
scalar_t up = y;
|
||||
if constexpr (HAS_CLAMP) {
|
||||
gate = (scalar_t)fmaxf(fminf((float)gate, limit), -limit);
|
||||
up = (scalar_t)fminf((float)up, limit);
|
||||
}
|
||||
return gate * ACT_FN(up);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename packed_t, packed_t (*PACKED_ACT_FN)(const packed_t&),
|
||||
bool act_first>
|
||||
bool act_first, bool HAS_CLAMP>
|
||||
__device__ __forceinline__ packed_t packed_compute(const packed_t& x,
|
||||
const packed_t& y) {
|
||||
return act_first ? packed_mul(PACKED_ACT_FN(x), y)
|
||||
: packed_mul(x, PACKED_ACT_FN(y));
|
||||
const packed_t& y,
|
||||
const float limit) {
|
||||
if constexpr (act_first) {
|
||||
packed_t gate = x;
|
||||
packed_t up = y;
|
||||
if constexpr (HAS_CLAMP) {
|
||||
float2 g = cast_to_float2(gate);
|
||||
float2 u = cast_to_float2(up);
|
||||
g.x = fminf(g.x, limit);
|
||||
g.y = fminf(g.y, limit);
|
||||
u.x = fmaxf(fminf(u.x, limit), -limit);
|
||||
u.y = fmaxf(fminf(u.y, limit), -limit);
|
||||
gate = cast_to_packed<packed_t>(g);
|
||||
up = cast_to_packed<packed_t>(u);
|
||||
}
|
||||
return packed_mul(PACKED_ACT_FN(gate), up);
|
||||
} else {
|
||||
packed_t gate = x;
|
||||
packed_t up = y;
|
||||
if constexpr (HAS_CLAMP) {
|
||||
float2 g = cast_to_float2(gate);
|
||||
float2 u = cast_to_float2(up);
|
||||
g.x = fmaxf(fminf(g.x, limit), -limit);
|
||||
g.y = fmaxf(fminf(g.y, limit), -limit);
|
||||
u.x = fminf(u.x, limit);
|
||||
u.y = fminf(u.y, limit);
|
||||
gate = cast_to_packed<packed_t>(g);
|
||||
up = cast_to_packed<packed_t>(u);
|
||||
}
|
||||
return packed_mul(gate, PACKED_ACT_FN(up));
|
||||
}
|
||||
}
|
||||
|
||||
// Activation and gating kernel template.
|
||||
template <typename scalar_t, typename packed_t,
|
||||
scalar_t (*ACT_FN)(const scalar_t&),
|
||||
packed_t (*PACKED_ACT_FN)(const packed_t&), bool act_first,
|
||||
bool use_vec, bool use_256b = false>
|
||||
bool use_vec, bool HAS_CLAMP, bool use_256b = false>
|
||||
__global__ void act_and_mul_kernel(
|
||||
scalar_t* __restrict__ out, // [..., d]
|
||||
const scalar_t* __restrict__ input, // [..., 2, d]
|
||||
const int d) {
|
||||
const int d, const float limit) {
|
||||
const scalar_t* x_ptr = input + blockIdx.x * 2 * d;
|
||||
const scalar_t* y_ptr = x_ptr + d;
|
||||
scalar_t* out_ptr = out + blockIdx.x * d;
|
||||
@@ -58,8 +103,9 @@ __global__ void act_and_mul_kernel(
|
||||
}
|
||||
#pragma unroll
|
||||
for (int j = 0; j < pvec_t::NUM_ELTS; j++) {
|
||||
x.elts[j] = packed_compute<packed_t, PACKED_ACT_FN, act_first>(
|
||||
x.elts[j], y.elts[j]);
|
||||
x.elts[j] =
|
||||
packed_compute<packed_t, PACKED_ACT_FN, act_first, HAS_CLAMP>(
|
||||
x.elts[j], y.elts[j], limit);
|
||||
}
|
||||
if constexpr (use_256b) {
|
||||
st256(x, &out_vec[i]);
|
||||
@@ -72,7 +118,8 @@ __global__ void act_and_mul_kernel(
|
||||
for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
|
||||
const scalar_t x = VLLM_LDG(&x_ptr[idx]);
|
||||
const scalar_t y = VLLM_LDG(&y_ptr[idx]);
|
||||
out_ptr[idx] = compute<scalar_t, ACT_FN, act_first>(x, y);
|
||||
out_ptr[idx] =
|
||||
compute<scalar_t, ACT_FN, act_first, HAS_CLAMP>(x, y, limit);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -151,8 +198,11 @@ packed_gelu_tanh_kernel(const packed_t& val) {
|
||||
|
||||
// Launch activation and gating kernel.
|
||||
// Use ACT_FIRST (bool) indicating whether to apply the activation function
|
||||
// first.
|
||||
#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, PACKED_KERNEL, ACT_FIRST) \
|
||||
// first. HAS_CLAMP (bool) enables pre-activation clamping: gate input is
|
||||
// clamped (max only) and up input is clamped (both sides) before the
|
||||
// activation function is applied.
|
||||
#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, PACKED_KERNEL, ACT_FIRST, \
|
||||
HAS_CLAMP, LIMIT) \
|
||||
auto dtype = input.scalar_type(); \
|
||||
int d = input.size(-1) / 2; \
|
||||
int64_t num_tokens = input.numel() / input.size(-1); \
|
||||
@@ -177,8 +227,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
|
||||
scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type, \
|
||||
KERNEL<scalar_t>, \
|
||||
PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
|
||||
ACT_FIRST, true, true><<<grid, block, 0, stream>>>( \
|
||||
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d); \
|
||||
ACT_FIRST, true, HAS_CLAMP, true><<<grid, block, 0, stream>>>( \
|
||||
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT); \
|
||||
}); \
|
||||
} else { \
|
||||
VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] { \
|
||||
@@ -186,8 +236,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
|
||||
scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type, \
|
||||
KERNEL<scalar_t>, \
|
||||
PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
|
||||
ACT_FIRST, true, false><<<grid, block, 0, stream>>>( \
|
||||
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d); \
|
||||
ACT_FIRST, true, HAS_CLAMP, false><<<grid, block, 0, stream>>>( \
|
||||
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT); \
|
||||
}); \
|
||||
} \
|
||||
} else { \
|
||||
@@ -197,8 +247,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
|
||||
scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type, \
|
||||
KERNEL<scalar_t>, \
|
||||
PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
|
||||
ACT_FIRST, false><<<grid, block, 0, stream>>>( \
|
||||
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d); \
|
||||
ACT_FIRST, false, HAS_CLAMP><<<grid, block, 0, stream>>>( \
|
||||
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT); \
|
||||
}); \
|
||||
}
|
||||
|
||||
@@ -206,7 +256,14 @@ void silu_and_mul(torch::Tensor& out, // [..., d]
|
||||
torch::Tensor& input) // [..., 2 * d]
|
||||
{
|
||||
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
|
||||
true);
|
||||
true, false, 0.0f);
|
||||
}
|
||||
|
||||
void silu_and_mul_clamp(torch::Tensor& out, // [..., d]
|
||||
torch::Tensor& input, // [..., 2 * d]
|
||||
double limit) {
|
||||
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
|
||||
true, true, (float)limit);
|
||||
}
|
||||
|
||||
void mul_and_silu(torch::Tensor& out, // [..., d]
|
||||
@@ -215,21 +272,21 @@ void mul_and_silu(torch::Tensor& out, // [..., d]
|
||||
// The difference between mul_and_silu and silu_and_mul is that mul_and_silu
|
||||
// applies the silu to the latter half of the input.
|
||||
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
|
||||
false);
|
||||
false, false, 0.0f);
|
||||
}
|
||||
|
||||
void gelu_and_mul(torch::Tensor& out, // [..., d]
|
||||
torch::Tensor& input) // [..., 2 * d]
|
||||
{
|
||||
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_kernel, vllm::packed_gelu_kernel,
|
||||
true);
|
||||
true, false, 0.0f);
|
||||
}
|
||||
|
||||
void gelu_tanh_and_mul(torch::Tensor& out, // [..., d]
|
||||
torch::Tensor& input) // [..., 2 * d]
|
||||
{
|
||||
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_tanh_kernel,
|
||||
vllm::packed_gelu_tanh_kernel, true);
|
||||
LAUNCH_ACTIVATION_GATE_KERNEL(
|
||||
vllm::gelu_tanh_kernel, vllm::packed_gelu_tanh_kernel, true, false, 0.0f);
|
||||
}
|
||||
|
||||
namespace vllm {
|
||||
|
||||
@@ -163,6 +163,8 @@ void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
|
||||
|
||||
void silu_and_mul(torch::Tensor& out, torch::Tensor& input);
|
||||
|
||||
void silu_and_mul_clamp(torch::Tensor& out, torch::Tensor& input, double limit);
|
||||
|
||||
void silu_and_mul_quant(torch::Tensor& out, torch::Tensor& input,
|
||||
torch::Tensor& scale);
|
||||
|
||||
|
||||
@@ -106,6 +106,12 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
ops.def("silu_and_mul(Tensor! result, Tensor input) -> ()");
|
||||
ops.impl("silu_and_mul", torch::kCUDA, &silu_and_mul);
|
||||
|
||||
// SwiGLU activation with input clamping.
|
||||
ops.def(
|
||||
"silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit) "
|
||||
"-> ()");
|
||||
ops.impl("silu_and_mul_with_clamp", torch::kCUDA, &silu_and_mul_clamp);
|
||||
|
||||
ops.def(
|
||||
"silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
|
||||
ops.impl("silu_and_mul_quant", torch::kCUDA, &silu_and_mul_quant);
|
||||
|
||||
+3
-1
@@ -540,7 +540,9 @@ RUN CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-') && \
|
||||
libcurand-dev-${CUDA_VERSION_DASH} \
|
||||
libcublas-${CUDA_VERSION_DASH} \
|
||||
# Required by fastsafetensors (fixes #20384)
|
||||
libnuma-dev && \
|
||||
libnuma-dev \
|
||||
# numactl CLI for NUMA binding at runtime
|
||||
numactl && \
|
||||
# Fixes nccl_allocator requiring nccl.h at runtime
|
||||
# https://github.com/vllm-project/vllm/blob/1336a1ea244fa8bfd7e72751cabbdb5b68a0c11a/vllm/distributed/device_communicators/pynccl_allocator.py#L22
|
||||
# NCCL packages don't use the cuda-MAJOR-MINOR naming convention,
|
||||
|
||||
@@ -292,10 +292,10 @@ Pooling models now support token-wise task.
|
||||
|
||||
### Score task
|
||||
|
||||
`score` task have has been removed in v0.21, use `classify` instead. Only when a classification model outputs num_labels
|
||||
equal to 1 can it be used as a scoring model and have its scoring API enabled.
|
||||
`score` task is deprecated and will be removed in v0.20. Please use `classify` instead. Only when a
|
||||
classification model outputs num_labels equal to 1 can it be used as a scoring model and have its scoring API enabled.
|
||||
|
||||
### Pooling multitask support
|
||||
|
||||
Pooling multitask support has been removed in v0.21. When the default pooling task is not what you want,
|
||||
Pooling multitask support is deprecated and will be removed in v0.20. When the default pooling task is not what you want,
|
||||
you need to manually specify it via `PoolerConfig(task=<task>)` offline or `--pooler-config.task <task>` online.
|
||||
|
||||
@@ -4,74 +4,68 @@
|
||||
import torch
|
||||
|
||||
from vllm import LLM
|
||||
from vllm.config import PoolerConfig
|
||||
from vllm.inputs import TextPrompt
|
||||
from vllm.multimodal.utils import fetch_image
|
||||
|
||||
# Initialize model
|
||||
model = LLM(
|
||||
model="jinaai/jina-embeddings-v4-vllm-text-matching",
|
||||
runner="pooling",
|
||||
max_model_len=1024,
|
||||
gpu_memory_utilization=0.8,
|
||||
)
|
||||
|
||||
def main():
|
||||
# Initialize model
|
||||
model = LLM(
|
||||
model="jinaai/jina-embeddings-v4-vllm-text-matching",
|
||||
pooler_config=PoolerConfig(task="token_embed"),
|
||||
runner="pooling",
|
||||
max_model_len=1024,
|
||||
gpu_memory_utilization=0.8,
|
||||
)
|
||||
# Create text prompts
|
||||
text1 = "Ein wunderschöner Sonnenuntergang am Strand"
|
||||
text1_prompt = TextPrompt(prompt=f"Query: {text1}")
|
||||
|
||||
# Create text prompts
|
||||
text1 = "Ein wunderschöner Sonnenuntergang am Strand"
|
||||
text1_prompt = TextPrompt(prompt=f"Query: {text1}")
|
||||
text2 = "浜辺に沈む美しい夕日"
|
||||
text2_prompt = TextPrompt(prompt=f"Query: {text2}")
|
||||
|
||||
text2 = "浜辺に沈む美しい夕日"
|
||||
text2_prompt = TextPrompt(prompt=f"Query: {text2}")
|
||||
# Create image prompt
|
||||
image = fetch_image(
|
||||
"https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/eskimo.jpg" # noqa: E501
|
||||
)
|
||||
image_prompt = TextPrompt(
|
||||
prompt="<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Describe the image.<|im_end|>\n", # noqa: E501
|
||||
multi_modal_data={"image": image},
|
||||
)
|
||||
|
||||
# Create image prompt
|
||||
image = fetch_image(
|
||||
"https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/eskimo.jpg" # noqa: E501
|
||||
)
|
||||
image_prompt = TextPrompt(
|
||||
prompt="<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Describe the image.<|im_end|>\n", # noqa: E501
|
||||
multi_modal_data={"image": image},
|
||||
)
|
||||
|
||||
# Encode all prompts
|
||||
prompts = [text1_prompt, text2_prompt, image_prompt]
|
||||
outputs = model.encode(prompts, pooling_task="token_embed")
|
||||
|
||||
def get_embeddings(outputs):
|
||||
VISION_START_TOKEN_ID, VISION_END_TOKEN_ID = 151652, 151653
|
||||
|
||||
embeddings = []
|
||||
for output in outputs:
|
||||
if VISION_START_TOKEN_ID in output.prompt_token_ids:
|
||||
# Gather only vision tokens
|
||||
img_start_pos = torch.where(
|
||||
torch.tensor(output.prompt_token_ids) == VISION_START_TOKEN_ID
|
||||
)[0][0]
|
||||
img_end_pos = torch.where(
|
||||
torch.tensor(output.prompt_token_ids) == VISION_END_TOKEN_ID
|
||||
)[0][0]
|
||||
embeddings_tensor = output.outputs.data.detach().clone()[
|
||||
img_start_pos : img_end_pos + 1
|
||||
]
|
||||
else:
|
||||
# Use all tokens for text-only prompts
|
||||
embeddings_tensor = output.outputs.data.detach().clone()
|
||||
|
||||
# Pool and normalize embeddings
|
||||
pooled_output = (
|
||||
embeddings_tensor.sum(dim=0, dtype=torch.float32)
|
||||
/ embeddings_tensor.shape[0]
|
||||
)
|
||||
embeddings.append(torch.nn.functional.normalize(pooled_output, dim=-1))
|
||||
return embeddings
|
||||
|
||||
embeddings = get_embeddings(outputs)
|
||||
|
||||
for embedding in embeddings:
|
||||
print(embedding.shape)
|
||||
# Encode all prompts
|
||||
prompts = [text1_prompt, text2_prompt, image_prompt]
|
||||
outputs = model.encode(prompts, pooling_task="token_embed")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
def get_embeddings(outputs):
|
||||
VISION_START_TOKEN_ID, VISION_END_TOKEN_ID = 151652, 151653
|
||||
|
||||
embeddings = []
|
||||
for output in outputs:
|
||||
if VISION_START_TOKEN_ID in output.prompt_token_ids:
|
||||
# Gather only vision tokens
|
||||
img_start_pos = torch.where(
|
||||
torch.tensor(output.prompt_token_ids) == VISION_START_TOKEN_ID
|
||||
)[0][0]
|
||||
img_end_pos = torch.where(
|
||||
torch.tensor(output.prompt_token_ids) == VISION_END_TOKEN_ID
|
||||
)[0][0]
|
||||
embeddings_tensor = output.outputs.data.detach().clone()[
|
||||
img_start_pos : img_end_pos + 1
|
||||
]
|
||||
else:
|
||||
# Use all tokens for text-only prompts
|
||||
embeddings_tensor = output.outputs.data.detach().clone()
|
||||
|
||||
# Pool and normalize embeddings
|
||||
pooled_output = (
|
||||
embeddings_tensor.sum(dim=0, dtype=torch.float32)
|
||||
/ embeddings_tensor.shape[0]
|
||||
)
|
||||
embeddings.append(torch.nn.functional.normalize(pooled_output, dim=-1))
|
||||
return embeddings
|
||||
|
||||
|
||||
embeddings = get_embeddings(outputs)
|
||||
|
||||
for embedding in embeddings:
|
||||
print(embedding.shape)
|
||||
|
||||
@@ -4,7 +4,6 @@
|
||||
from argparse import Namespace
|
||||
|
||||
from vllm import LLM, EngineArgs
|
||||
from vllm.config import PoolerConfig
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
|
||||
@@ -14,7 +13,6 @@ def parse_args():
|
||||
# Set example specific arguments
|
||||
parser.set_defaults(
|
||||
model="BAAI/bge-m3",
|
||||
pooler_config=PoolerConfig(task="token_embed"),
|
||||
runner="pooling",
|
||||
enforce_eager=True,
|
||||
)
|
||||
@@ -34,6 +32,15 @@ def main(args: Namespace):
|
||||
# You should pass runner="pooling" for embedding models
|
||||
llm = LLM(**vars(args))
|
||||
|
||||
# Generate embedding. The output is a list of EmbeddingRequestOutputs.
|
||||
outputs = llm.embed(prompts)
|
||||
|
||||
# Print the outputs.
|
||||
print("\nGenerated Outputs:\n" + "-" * 60)
|
||||
for prompt, output in zip(prompts, outputs):
|
||||
embeds = output.outputs.embedding
|
||||
print(len(embeds))
|
||||
|
||||
# Generate embedding for each token. The output is a list of PoolingRequestOutput.
|
||||
outputs = llm.encode(prompts, pooling_task="token_embed")
|
||||
|
||||
@@ -43,20 +50,6 @@ def main(args: Namespace):
|
||||
multi_vector = output.outputs.data
|
||||
print(multi_vector.shape)
|
||||
|
||||
query = "What is the capital of France?"
|
||||
documents = [
|
||||
"The capital of Brazil is Brasilia.",
|
||||
"The capital of France is Paris.",
|
||||
]
|
||||
# Generate scores.
|
||||
outputs = llm.score(query, documents)
|
||||
# Print the outputs.
|
||||
print("\nGenerated Outputs:\n" + "-" * 60)
|
||||
for document, output in zip(documents, outputs):
|
||||
score = output.outputs.score
|
||||
print(f"Pair: {[query, document]!r} \nScore: {score}")
|
||||
print("-" * 60)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
|
||||
@@ -7,11 +7,10 @@ Example online usage of Pooling API for multi vector retrieval.
|
||||
Run `vllm serve <model> --runner pooling`
|
||||
to start up the server in vLLM. e.g.
|
||||
|
||||
vllm serve BAAI/bge-m3 --pooler-config.task token_embed
|
||||
vllm serve BAAI/bge-m3
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import pprint
|
||||
|
||||
import requests
|
||||
import torch
|
||||
@@ -33,8 +32,7 @@ def parse_args():
|
||||
|
||||
|
||||
def main(args):
|
||||
pooling_url = f"http://{args.host}:{args.port}/pooling"
|
||||
score_url = f"http://{args.host}:{args.port}/score"
|
||||
api_url = f"http://{args.host}:{args.port}/pooling"
|
||||
model_name = args.model
|
||||
|
||||
prompts = [
|
||||
@@ -45,23 +43,11 @@ def main(args):
|
||||
]
|
||||
prompt = {"model": model_name, "input": prompts}
|
||||
|
||||
pooling_response = post_http_request(prompt=prompt, api_url=pooling_url)
|
||||
pooling_response = post_http_request(prompt=prompt, api_url=api_url)
|
||||
for output in pooling_response.json()["data"]:
|
||||
multi_vector = torch.tensor(output["data"])
|
||||
print(multi_vector.shape)
|
||||
|
||||
queries = "What is the capital of France?"
|
||||
documents = [
|
||||
"The capital of Brazil is Brasilia.",
|
||||
"The capital of France is Paris.",
|
||||
]
|
||||
prompt = {"model": model_name, "queries": queries, "documents": documents}
|
||||
score_response = post_http_request(prompt=prompt, api_url=score_url)
|
||||
print("\nPrompt when queries is string and documents is a list:")
|
||||
pprint.pprint(prompt)
|
||||
print("\nScore Response:")
|
||||
pprint.pprint(score_response.json())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
|
||||
@@ -1,12 +1,13 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import logging
|
||||
import weakref
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.models.utils import softmax
|
||||
from vllm import LLM, ClassificationRequestOutput, PoolingParams
|
||||
from vllm import LLM, ClassificationRequestOutput, PoolingParams, PoolingRequestOutput
|
||||
from vllm.distributed import cleanup_dist_env_and_memory
|
||||
from vllm.tasks import PoolingTask
|
||||
|
||||
@@ -65,6 +66,18 @@ def test_list_prompts(llm: LLM):
|
||||
assert len(outputs[i].outputs.probs) == num_labels
|
||||
|
||||
|
||||
@pytest.mark.skip_global_cleanup
|
||||
def test_token_classify(llm: LLM, caplog_vllm):
|
||||
with caplog_vllm.at_level(level=logging.WARNING, logger="vllm"):
|
||||
outputs = llm.encode(prompt, pooling_task="token_classify", use_tqdm=False)
|
||||
assert "deprecated" in caplog_vllm.text
|
||||
|
||||
assert len(outputs) == 1
|
||||
assert isinstance(outputs[0], PoolingRequestOutput)
|
||||
assert outputs[0].prompt_token_ids == prompt_token_ids
|
||||
assert outputs[0].outputs.data.shape == (len(prompt_token_ids), num_labels)
|
||||
|
||||
|
||||
@pytest.mark.skip_global_cleanup
|
||||
def test_pooling_params(llm: LLM):
|
||||
def get_outputs(use_activation):
|
||||
@@ -97,12 +110,10 @@ def test_score_api(llm: LLM):
|
||||
llm.score("ping", "pong", use_tqdm=False)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("task", ["embed", "token_embed", "token_classify", "plugin"])
|
||||
@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
|
||||
def test_unsupported_tasks(llm: LLM, task: PoolingTask):
|
||||
if task == "plugin":
|
||||
err_msg = "No IOProcessor plugin installed."
|
||||
elif task == "token_classify":
|
||||
err_msg = "Try switching the model's pooling_task via.+"
|
||||
else:
|
||||
err_msg = "Embedding API is not supported by this model.+"
|
||||
with pytest.raises(ValueError, match=err_msg):
|
||||
|
||||
@@ -436,7 +436,26 @@ async def test_pooling_classify(server: RemoteOpenAIServer, model_name: str):
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
@pytest.mark.parametrize("task", ["embed", "token_embed", "token_classify", "plugin"])
|
||||
async def test_pooling_token_classify(server: RemoteOpenAIServer, model_name: str):
|
||||
task = "token_classify"
|
||||
response = requests.post(
|
||||
server.url_for("pooling"),
|
||||
json={
|
||||
"model": model_name,
|
||||
"input": input_text,
|
||||
"encoding_format": "float",
|
||||
"task": task,
|
||||
},
|
||||
)
|
||||
poolings = PoolingResponse.model_validate(response.json())
|
||||
assert len(poolings.data) == 1
|
||||
assert len(poolings.data[0].data) == 8
|
||||
assert len(poolings.data[0].data[0]) == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
|
||||
async def test_pooling_not_supported(
|
||||
server: RemoteOpenAIServer, model_name: str, task: str
|
||||
):
|
||||
@@ -450,11 +469,8 @@ async def test_pooling_not_supported(
|
||||
},
|
||||
)
|
||||
assert response.json()["error"]["type"] == "BadRequestError"
|
||||
|
||||
if task == "plugin":
|
||||
err_msg = "No IOProcessor plugin installed."
|
||||
elif task == "token_classify":
|
||||
err_msg = "Try switching the model's pooling_task via"
|
||||
else:
|
||||
err_msg = f"Unsupported task: {task!r}"
|
||||
assert response.json()["error"]["message"].startswith(err_msg)
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import logging
|
||||
import weakref
|
||||
|
||||
import pytest
|
||||
@@ -37,11 +38,11 @@ def llm():
|
||||
seed=0,
|
||||
attention_config=attention_config,
|
||||
)
|
||||
assert embedding_size == llm.model_config.embedding_size
|
||||
|
||||
yield weakref.proxy(llm)
|
||||
|
||||
del llm
|
||||
|
||||
cleanup_dist_env_and_memory()
|
||||
|
||||
|
||||
@@ -73,6 +74,16 @@ def test_list_prompts(llm: LLM):
|
||||
assert len(outputs[i].outputs.embedding) == embedding_size
|
||||
|
||||
|
||||
@pytest.mark.skip_global_cleanup
|
||||
def test_token_embed(llm: LLM, caplog_vllm):
|
||||
with caplog_vllm.at_level(level=logging.WARNING, logger="vllm"):
|
||||
outputs = llm.encode(prompt, pooling_task="token_embed", use_tqdm=False)
|
||||
assert "deprecated" in caplog_vllm.text
|
||||
|
||||
multi_vector = outputs[0].outputs.data
|
||||
assert multi_vector.shape == (11, 384)
|
||||
|
||||
|
||||
@pytest.mark.skip_global_cleanup
|
||||
def test_pooling_params(llm: LLM):
|
||||
def get_outputs(normalize):
|
||||
@@ -96,14 +107,10 @@ def test_pooling_params(llm: LLM):
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"task", ["token_classify", "classify", "token_embed", "plugin"]
|
||||
)
|
||||
@pytest.mark.parametrize("task", ["token_classify", "classify", "plugin"])
|
||||
def test_unsupported_tasks(llm: LLM, task: PoolingTask):
|
||||
if task == "plugin":
|
||||
err_msg = "No IOProcessor plugin installed."
|
||||
elif task == "token_embed":
|
||||
err_msg = "Try switching the model's pooling_task via.+"
|
||||
else:
|
||||
err_msg = "Classification API is not supported by this model.+"
|
||||
with pytest.raises(ValueError, match=err_msg):
|
||||
|
||||
@@ -732,9 +732,28 @@ async def test_pooling_embed(server: RemoteOpenAIServer, model_name: str):
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
@pytest.mark.parametrize(
|
||||
"task", ["classify", "token_classify", "token_embed", "plugin"]
|
||||
)
|
||||
async def test_pooling_token_embed(server: RemoteOpenAIServer, model_name: str):
|
||||
task = "token_embed"
|
||||
response = requests.post(
|
||||
server.url_for("pooling"),
|
||||
json={
|
||||
"model": model_name,
|
||||
"input": input_text,
|
||||
"encoding_format": "float",
|
||||
"task": task,
|
||||
},
|
||||
)
|
||||
|
||||
poolings = PoolingResponse.model_validate(response.json())
|
||||
|
||||
assert len(poolings.data) == 1
|
||||
assert len(poolings.data[0].data) == len(input_tokens)
|
||||
assert len(poolings.data[0].data[0]) == 384
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
@pytest.mark.parametrize("task", ["classify", "token_classify", "plugin"])
|
||||
async def test_pooling_not_supported(
|
||||
server: RemoteOpenAIServer, model_name: str, task: str
|
||||
):
|
||||
@@ -750,8 +769,6 @@ async def test_pooling_not_supported(
|
||||
assert response.json()["error"]["type"] == "BadRequestError"
|
||||
if task == "plugin":
|
||||
err_msg = "No IOProcessor plugin installed."
|
||||
elif task == "token_embed":
|
||||
err_msg = "Try switching the model's pooling_task via"
|
||||
else:
|
||||
err_msg = f"Unsupported task: {task!r}"
|
||||
assert response.json()["error"]["message"].startswith(err_msg)
|
||||
|
||||
@@ -452,6 +452,25 @@ async def test_pooling_classify(server: RemoteOpenAIServer):
|
||||
assert len(poolings.data[0].data) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pooling_token_classify(server: RemoteOpenAIServer):
|
||||
response = requests.post(
|
||||
server.url_for("pooling"),
|
||||
json={
|
||||
"model": MODEL_NAME,
|
||||
"task": "token_classify",
|
||||
"input": input_text,
|
||||
"encoding_format": "float",
|
||||
},
|
||||
)
|
||||
|
||||
poolings = PoolingResponse.model_validate(response.json())
|
||||
|
||||
assert len(poolings.data) == 1
|
||||
assert len(poolings.data[0].data) == len(input_tokens)
|
||||
assert len(poolings.data[0].data[0]) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rerank_max_tokens_per_doc(
|
||||
server: RemoteOpenAIServer,
|
||||
@@ -525,7 +544,7 @@ async def test_rerank_max_tokens_per_doc_validation(
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("task", ["embed", "token_embed", "token_classify", "plugin"])
|
||||
@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
|
||||
async def test_pooling_not_supported(server: RemoteOpenAIServer, task: str):
|
||||
response = requests.post(
|
||||
server.url_for("pooling"),
|
||||
@@ -539,8 +558,6 @@ async def test_pooling_not_supported(server: RemoteOpenAIServer, task: str):
|
||||
assert response.json()["error"]["type"] == "BadRequestError"
|
||||
if task == "plugin":
|
||||
err_msg = "No IOProcessor plugin installed."
|
||||
elif task == "token_classify":
|
||||
err_msg = "Try switching the model's pooling_task via"
|
||||
else:
|
||||
err_msg = f"Unsupported task: {task!r}"
|
||||
assert response.json()["error"]["message"].startswith(err_msg)
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import logging
|
||||
import weakref
|
||||
|
||||
import pytest
|
||||
@@ -59,19 +60,22 @@ def test_token_ids_prompts(llm: LLM):
|
||||
|
||||
@pytest.mark.skip_global_cleanup
|
||||
def test_score_api(llm: LLM):
|
||||
err_msg = "This model does not support the Scoring API."
|
||||
err_msg = "Scoring API is only enabled for num_labels == 1."
|
||||
with pytest.raises(ValueError, match=err_msg):
|
||||
llm.score("ping", "pong", use_tqdm=False)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("task", ["classify", "embed", "token_embed", "plugin"])
|
||||
def test_unsupported_tasks(llm: LLM, task: PoolingTask, caplog_vllm):
|
||||
if task == "plugin":
|
||||
err_msg = "No IOProcessor plugin installed."
|
||||
elif task == "classify":
|
||||
err_msg = "Try switching the model's pooling_task via.+"
|
||||
if task == "classify":
|
||||
with caplog_vllm.at_level(level=logging.WARNING, logger="vllm"):
|
||||
llm.encode(prompt, pooling_task=task, use_tqdm=False)
|
||||
assert "deprecated" in caplog_vllm.text
|
||||
else:
|
||||
err_msg = "Embedding API is not supported by this model.+"
|
||||
if task == "plugin":
|
||||
err_msg = "No IOProcessor plugin installed."
|
||||
else:
|
||||
err_msg = "Embedding API is not supported by this model.+"
|
||||
|
||||
with pytest.raises(ValueError, match=err_msg):
|
||||
llm.encode(prompt, pooling_task=task, use_tqdm=False)
|
||||
with pytest.raises(ValueError, match=err_msg):
|
||||
llm.encode(prompt, pooling_task=task, use_tqdm=False)
|
||||
|
||||
@@ -50,7 +50,7 @@ async def test_pooling_token_classify(server: RemoteOpenAIServer, model_name: st
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
@pytest.mark.parametrize("task", ["classify", "embed", "token_embed", "plugin"])
|
||||
@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
|
||||
async def test_pooling_not_supported(
|
||||
server: RemoteOpenAIServer, model_name: str, task: str
|
||||
):
|
||||
@@ -63,12 +63,9 @@ async def test_pooling_not_supported(
|
||||
"task": task,
|
||||
},
|
||||
)
|
||||
assert response.json()["error"]["type"] == "BadRequestError"
|
||||
|
||||
if task == "plugin":
|
||||
err_msg = "No IOProcessor plugin installed."
|
||||
elif task == "classify":
|
||||
err_msg = "Try switching the model's pooling_task via"
|
||||
else:
|
||||
err_msg = f"Unsupported task: {task!r}"
|
||||
assert response.json()["error"]["message"].startswith(err_msg)
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import logging
|
||||
import weakref
|
||||
|
||||
import pytest
|
||||
@@ -63,12 +64,15 @@ def test_token_ids_prompts(llm: LLM):
|
||||
|
||||
@pytest.mark.parametrize("task", ["embed", "classify", "token_classify", "plugin"])
|
||||
def test_unsupported_tasks(llm: LLM, task: PoolingTask, caplog_vllm):
|
||||
if task == "plugin":
|
||||
err_msg = "No IOProcessor plugin installed."
|
||||
elif task == "embed":
|
||||
err_msg = "Try switching the model's pooling_task via.+"
|
||||
if task == "embed":
|
||||
with caplog_vllm.at_level(level=logging.WARNING, logger="vllm"):
|
||||
llm.encode(prompt, pooling_task=task, use_tqdm=False)
|
||||
assert "deprecated" in caplog_vllm.text
|
||||
else:
|
||||
err_msg = "Classification API is not supported by this model.+"
|
||||
if task == "plugin":
|
||||
err_msg = "No IOProcessor plugin installed."
|
||||
else:
|
||||
err_msg = "Classification API is not supported by this model.+"
|
||||
|
||||
with pytest.raises(ValueError, match=err_msg):
|
||||
llm.encode(prompt, pooling_task=task, use_tqdm=False)
|
||||
with pytest.raises(ValueError, match=err_msg):
|
||||
llm.encode(prompt, pooling_task=task, use_tqdm=False)
|
||||
|
||||
@@ -73,7 +73,7 @@ async def test_pooling_token_embed(server: RemoteOpenAIServer, model_name: str):
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
@pytest.mark.parametrize("task", ["embed", "classify", "token_classify", "plugin"])
|
||||
@pytest.mark.parametrize("task", ["classify", "token_classify", "plugin"])
|
||||
async def test_pooling_not_supported(
|
||||
server: RemoteOpenAIServer, model_name: str, task: str
|
||||
):
|
||||
@@ -86,12 +86,9 @@ async def test_pooling_not_supported(
|
||||
"task": task,
|
||||
},
|
||||
)
|
||||
assert response.json()["error"]["type"] == "BadRequestError"
|
||||
|
||||
if task == "plugin":
|
||||
err_msg = "No IOProcessor plugin installed."
|
||||
elif task == "embed":
|
||||
err_msg = "Try switching the model's pooling_task via"
|
||||
else:
|
||||
err_msg = f"Unsupported task: {task!r}"
|
||||
assert response.json()["error"]["message"].startswith(err_msg)
|
||||
|
||||
@@ -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",
|
||||
[
|
||||
|
||||
@@ -6,7 +6,6 @@ from transformers import AutoModel
|
||||
|
||||
from tests.models.utils import check_embeddings_close
|
||||
from vllm import TokensPrompt
|
||||
from vllm.config import PoolerConfig
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
@@ -22,7 +21,6 @@ def test_embed_models(hf_runner, vllm_runner, model: str):
|
||||
with vllm_runner(
|
||||
model,
|
||||
runner="pooling",
|
||||
pooler_config=PoolerConfig(task="token_embed"),
|
||||
max_model_len=128,
|
||||
max_num_batched_tokens=chunk_size,
|
||||
enforce_eager=True,
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
import httpx
|
||||
import openai
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
import torch
|
||||
|
||||
from ....utils import RemoteOpenAIServer
|
||||
@@ -24,42 +25,29 @@ sentences_2 = [
|
||||
similarity_reference = [[0.6259, 0.3474], [0.3309, 0.6734]]
|
||||
lexical_score_reference = [0.19554901123046875, 0.0]
|
||||
colbert_score_reference = [0.7797, 0.4620]
|
||||
SUPPORTED_TASKS = ["embed", "token_embed", "token_classify"]
|
||||
|
||||
|
||||
@pytest.fixture(scope="module", params=SUPPORTED_TASKS)
|
||||
def pooling_task(request):
|
||||
yield request.param
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def server(pooling_task):
|
||||
def server():
|
||||
args = [
|
||||
"--max-model-len",
|
||||
str(MAX_MODEL_LEN),
|
||||
"--hf-overrides",
|
||||
'{"architectures": ["BgeM3EmbeddingModel"]}',
|
||||
"--pooler-config.task",
|
||||
pooling_task,
|
||||
]
|
||||
|
||||
with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
|
||||
yield remote_server
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def client(server):
|
||||
async with server.get_async_client() as async_client:
|
||||
yield async_client
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bge_m3_api_server_embedding(server, pooling_task):
|
||||
client = server.get_async_client()
|
||||
|
||||
if pooling_task != "embed":
|
||||
with pytest.raises(openai.InternalServerError):
|
||||
await run_client_embeddings(
|
||||
client,
|
||||
MODEL_NAME,
|
||||
sentences_1,
|
||||
)
|
||||
return
|
||||
|
||||
async def test_bge_m3_api_server_embedding(client: openai.AsyncOpenAI):
|
||||
embeddings_list_1 = await run_client_embeddings(
|
||||
client,
|
||||
MODEL_NAME,
|
||||
@@ -129,14 +117,7 @@ def compute_lexical_matching_score(
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bge_m3_api_server_sparse_embedding(server, pooling_task):
|
||||
client = server.get_async_client()
|
||||
|
||||
if pooling_task != "token_classify":
|
||||
with pytest.raises(openai.BadRequestError):
|
||||
await sparse_embeddings(client, sentences_1)
|
||||
return
|
||||
|
||||
async def test_bge_m3_api_server_sparse_embedding(client: openai.AsyncOpenAI):
|
||||
embeddings_1 = await sparse_embeddings(client, sentences_1)
|
||||
embeddings_2 = await sparse_embeddings(client, sentences_2)
|
||||
|
||||
@@ -156,11 +137,9 @@ async def test_bge_m3_api_server_sparse_embedding(server, pooling_task):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bge_m3_api_server_sparse_embedding_corner_case(server, pooling_task):
|
||||
if pooling_task != "token_classify":
|
||||
return
|
||||
|
||||
client = server.get_async_client()
|
||||
async def test_bge_m3_api_server_sparse_embedding_corner_case(
|
||||
client: openai.AsyncOpenAI,
|
||||
):
|
||||
embeddings = await sparse_embeddings(client, ["Hi"])
|
||||
assert len(embeddings) == 1
|
||||
assert 2673 in embeddings[0]
|
||||
@@ -176,18 +155,7 @@ def colbert_score(q_reps: torch.Tensor, p_reps: torch.Tensor) -> torch.Tensor:
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bge_m3_api_server_multi_vector(server, pooling_task):
|
||||
client = server.get_async_client()
|
||||
|
||||
if pooling_task != "token_embed":
|
||||
with pytest.raises(openai.BadRequestError):
|
||||
await client.post(
|
||||
"../pooling",
|
||||
body={"model": MODEL_NAME, "input": sentences_1, "task": "token_embed"},
|
||||
cast_to=httpx.Response,
|
||||
)
|
||||
return
|
||||
|
||||
async def test_bge_m3_api_server_multi_vector(client: openai.AsyncOpenAI):
|
||||
result_1 = await client.post(
|
||||
"../pooling",
|
||||
body={"model": MODEL_NAME, "input": sentences_1, "task": "token_embed"},
|
||||
|
||||
@@ -4,7 +4,6 @@ import pytest
|
||||
import torch
|
||||
|
||||
from vllm import TokensPrompt
|
||||
from vllm.config import PoolerConfig
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
@@ -21,7 +20,6 @@ def test_extract_hidden_states(hf_runner, vllm_runner, model: str):
|
||||
max_model_len=128,
|
||||
enforce_eager=True,
|
||||
runner="pooling",
|
||||
pooler_config=PoolerConfig(task="token_embed"),
|
||||
enable_prefix_caching=True,
|
||||
) as vllm_model:
|
||||
pooling_outputs = vllm_model.llm.encode(
|
||||
@@ -46,3 +44,14 @@ def test_extract_hidden_states(hf_runner, vllm_runner, model: str):
|
||||
assert len(output.prompt_token_ids) == n
|
||||
assert len(output.outputs.data) == n
|
||||
assert output.num_cached_tokens == 0
|
||||
|
||||
# skip_reading_prefix_cache can still write to cache
|
||||
# to accelerate following requests
|
||||
pooling_outputs = vllm_model.llm.encode(
|
||||
[TokensPrompt(prompt_token_ids=t) for t in token_prompts],
|
||||
pooling_task="embed",
|
||||
)
|
||||
|
||||
for n, output in zip(n_prompt_tokens, pooling_outputs):
|
||||
assert len(output.prompt_token_ids) == n
|
||||
assert output.num_cached_tokens > 0
|
||||
|
||||
@@ -5,7 +5,6 @@ import torch
|
||||
from transformers import AutoModel
|
||||
|
||||
from tests.models.utils import check_embeddings_close
|
||||
from vllm.config import PoolerConfig
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
@@ -18,7 +17,6 @@ def test_embed_models(hf_runner, vllm_runner, example_prompts, model: str, dtype
|
||||
with vllm_runner(
|
||||
model,
|
||||
runner="pooling",
|
||||
pooler_config=PoolerConfig(task="token_embed"),
|
||||
max_model_len=None,
|
||||
) as vllm_model:
|
||||
vllm_outputs = vllm_model.token_embed(example_prompts)
|
||||
|
||||
@@ -146,7 +146,7 @@ def test_multi_vector_retrieval_models_using_normalize(
|
||||
model,
|
||||
max_model_len=512,
|
||||
dtype=dtype,
|
||||
pooler_config=PoolerConfig(use_activation=False, task="token_embed"),
|
||||
pooler_config=PoolerConfig(use_activation=False),
|
||||
) as vllm_model:
|
||||
wo_normalize = vllm_model.token_embed(example_prompts)
|
||||
|
||||
@@ -154,7 +154,7 @@ def test_multi_vector_retrieval_models_using_normalize(
|
||||
model,
|
||||
max_model_len=512,
|
||||
dtype=dtype,
|
||||
pooler_config=PoolerConfig(use_activation=True, task="token_embed"),
|
||||
pooler_config=PoolerConfig(use_activation=True),
|
||||
) as vllm_model:
|
||||
w_normalize = vllm_model.token_embed(example_prompts)
|
||||
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -79,7 +79,7 @@ from vllm.renderers.inputs.preprocess import (
|
||||
prompt_to_seq,
|
||||
)
|
||||
from vllm.sampling_params import BeamSearchParams, RequestOutputKind, SamplingParams
|
||||
from vllm.tasks import SCORE_TYPE_MAP, PoolingTask
|
||||
from vllm.tasks import PoolingTask
|
||||
from vllm.tokenizers import TokenizerLike
|
||||
from vllm.usage.usage_lib import UsageContext
|
||||
from vllm.utils.counter import Counter
|
||||
@@ -1204,9 +1204,12 @@ class LLM:
|
||||
f"Supported tasks: {self.supported_tasks}"
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Try switching the model's pooling_task "
|
||||
f'via `PoolerConfig(task="{pooling_task}")`'
|
||||
logger.warning_once(
|
||||
"Pooling multitask support is deprecated and will "
|
||||
"be removed in v0.20. When the default pooling task is "
|
||||
"not what you want, you need to manually specify it "
|
||||
'via PoolerConfig(task="%s"). ',
|
||||
pooling_task,
|
||||
)
|
||||
|
||||
if pooling_task == "plugin" and "plugin" not in self.pooling_io_processors:
|
||||
@@ -1409,7 +1412,7 @@ class LLM:
|
||||
"pooling model."
|
||||
)
|
||||
|
||||
score_type: str | None = SCORE_TYPE_MAP.get(self.pooling_task, None) # type: ignore[arg-type]
|
||||
score_type = self.model_config.score_type
|
||||
if (
|
||||
score_type == "cross-encoder"
|
||||
and getattr(self.model_config.hf_config, "num_labels", 0) != 1
|
||||
|
||||
@@ -15,7 +15,10 @@ from starlette.datastructures import Headers
|
||||
from vllm import PoolingParams, PoolingRequestOutput, envs
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.engine.protocol import EngineClient
|
||||
from vllm.entrypoints.chat_utils import ChatTemplateConfig
|
||||
from vllm.entrypoints.chat_utils import (
|
||||
ChatTemplateConfig,
|
||||
ChatTemplateContentFormatOption,
|
||||
)
|
||||
from vllm.entrypoints.logger import RequestLogger
|
||||
from vllm.entrypoints.openai.engine.protocol import ErrorResponse
|
||||
from vllm.entrypoints.openai.models.serving import OpenAIServingModels
|
||||
@@ -45,7 +48,9 @@ class PoolingServingBase(ABC):
|
||||
models: OpenAIServingModels,
|
||||
*,
|
||||
request_logger: RequestLogger | None,
|
||||
chat_template_config: ChatTemplateConfig,
|
||||
chat_template: str | None = None,
|
||||
chat_template_content_format: ChatTemplateContentFormatOption = "auto",
|
||||
trust_request_chat_template: bool = False,
|
||||
return_tokens_as_token_ids: bool = False,
|
||||
log_error_stack: bool = False,
|
||||
):
|
||||
@@ -58,7 +63,11 @@ class PoolingServingBase(ABC):
|
||||
self.request_logger = request_logger
|
||||
self.return_tokens_as_token_ids = return_tokens_as_token_ids
|
||||
self.log_error_stack = log_error_stack
|
||||
self.chat_template_config = chat_template_config
|
||||
self.chat_template_config = ChatTemplateConfig(
|
||||
chat_template=chat_template,
|
||||
chat_template_content_format=chat_template_content_format,
|
||||
trust_request_chat_template=trust_request_chat_template,
|
||||
)
|
||||
|
||||
# Shared thread pool executor for preprocessing and postprocessing.
|
||||
self._executor: Executor = models.renderer._executor
|
||||
|
||||
@@ -10,7 +10,7 @@ from vllm.entrypoints.chat_utils import ChatTemplateConfig
|
||||
from vllm.logger import init_logger
|
||||
from vllm.plugins.io_processors import has_io_processor
|
||||
from vllm.renderers import BaseRenderer
|
||||
from vllm.tasks import POOLING_TASKS, SCORE_TYPE_MAP, SupportedTask
|
||||
from vllm.tasks import POOLING_TASKS, SupportedTask
|
||||
|
||||
from .base.io_processor import PoolingIOProcessor
|
||||
from .utils import enable_scoring_api
|
||||
@@ -43,24 +43,23 @@ def init_pooling_io_processors(
|
||||
) -> dict[str, PoolingIOProcessor]:
|
||||
model_config = vllm_config.model_config
|
||||
processors: dict[str, type[PoolingIOProcessor]] = {}
|
||||
pooling_task = model_config.get_pooling_task(supported_tasks)
|
||||
|
||||
if pooling_task == "classify":
|
||||
if "classify" in supported_tasks:
|
||||
from .classify.io_processor import ClassifyIOProcessor
|
||||
|
||||
processors["classify"] = ClassifyIOProcessor
|
||||
|
||||
if pooling_task == "token_classify":
|
||||
if "token_classify" in supported_tasks:
|
||||
from .classify.io_processor import TokenClassifyIOProcessor
|
||||
|
||||
processors["token_classify"] = TokenClassifyIOProcessor
|
||||
|
||||
if pooling_task == "embed":
|
||||
if "embed" in supported_tasks:
|
||||
from .embed.io_processor import EmbedIOProcessor
|
||||
|
||||
processors["embed"] = EmbedIOProcessor
|
||||
|
||||
if pooling_task == "token_embed":
|
||||
if "token_embed" in supported_tasks:
|
||||
from .embed.io_processor import TokenEmbedIOProcessor
|
||||
|
||||
processors["token_embed"] = TokenEmbedIOProcessor
|
||||
@@ -72,15 +71,15 @@ def init_pooling_io_processors(
|
||||
from .pooling.io_processor import PluginWithIOProcessorPlugins
|
||||
|
||||
processors["plugin"] = PluginWithIOProcessorPlugins
|
||||
elif pooling_task == "plugin":
|
||||
elif "plugin" in supported_tasks:
|
||||
from .pooling.io_processor import PluginWithoutIOProcessorPlugins
|
||||
|
||||
processors["plugin"] = PluginWithoutIOProcessorPlugins
|
||||
|
||||
if enable_scoring_api(supported_tasks, model_config):
|
||||
score_type = model_config.score_type
|
||||
from .scoring.io_processor import ScoringIOProcessors
|
||||
|
||||
score_type: str | None = SCORE_TYPE_MAP.get(pooling_task, None) # type: ignore[arg-type]
|
||||
if score_type is not None and score_type in ScoringIOProcessors:
|
||||
processors[score_type] = ScoringIOProcessors[score_type]
|
||||
|
||||
@@ -141,10 +140,6 @@ def init_pooling_state(
|
||||
request_logger: RequestLogger | None,
|
||||
supported_tasks: tuple["SupportedTask", ...],
|
||||
):
|
||||
model_config = engine_client.model_config
|
||||
if model_config is None:
|
||||
return
|
||||
|
||||
from vllm.entrypoints.chat_utils import load_chat_template
|
||||
from vllm.tasks import POOLING_TASKS
|
||||
|
||||
@@ -153,14 +148,8 @@ def init_pooling_state(
|
||||
from .pooling.serving import ServingPooling
|
||||
from .scoring.serving import ServingScores
|
||||
|
||||
model_config = engine_client.model_config
|
||||
resolved_chat_template = load_chat_template(args.chat_template)
|
||||
pooling_task = model_config.get_pooling_task(supported_tasks)
|
||||
|
||||
chat_template_config = ChatTemplateConfig(
|
||||
chat_template=resolved_chat_template,
|
||||
chat_template_content_format=args.chat_template_content_format,
|
||||
trust_request_chat_template=args.trust_request_chat_template,
|
||||
)
|
||||
|
||||
state.serving_pooling = (
|
||||
(
|
||||
@@ -169,7 +158,9 @@ def init_pooling_state(
|
||||
state.openai_serving_models,
|
||||
supported_tasks=supported_tasks,
|
||||
request_logger=request_logger,
|
||||
chat_template_config=chat_template_config,
|
||||
chat_template=resolved_chat_template,
|
||||
chat_template_content_format=args.chat_template_content_format,
|
||||
trust_request_chat_template=args.trust_request_chat_template,
|
||||
)
|
||||
)
|
||||
if any(t in supported_tasks for t in POOLING_TASKS)
|
||||
@@ -180,9 +171,11 @@ def init_pooling_state(
|
||||
engine_client,
|
||||
state.openai_serving_models,
|
||||
request_logger=request_logger,
|
||||
chat_template_config=chat_template_config,
|
||||
chat_template=resolved_chat_template,
|
||||
chat_template_content_format=args.chat_template_content_format,
|
||||
trust_request_chat_template=args.trust_request_chat_template,
|
||||
)
|
||||
if pooling_task == "embed"
|
||||
if "embed" in supported_tasks
|
||||
else None
|
||||
)
|
||||
state.serving_classification = (
|
||||
@@ -190,18 +183,21 @@ def init_pooling_state(
|
||||
engine_client,
|
||||
state.openai_serving_models,
|
||||
request_logger=request_logger,
|
||||
chat_template_config=chat_template_config,
|
||||
chat_template=resolved_chat_template,
|
||||
chat_template_content_format=args.chat_template_content_format,
|
||||
trust_request_chat_template=args.trust_request_chat_template,
|
||||
)
|
||||
if pooling_task == "classify"
|
||||
if "classify" in supported_tasks
|
||||
else None
|
||||
)
|
||||
state.serving_scores = (
|
||||
ServingScores(
|
||||
engine_client,
|
||||
state.openai_serving_models,
|
||||
supported_tasks=supported_tasks,
|
||||
request_logger=request_logger,
|
||||
chat_template_config=chat_template_config,
|
||||
chat_template=resolved_chat_template,
|
||||
chat_template_content_format=args.chat_template_content_format,
|
||||
trust_request_chat_template=args.trust_request_chat_template,
|
||||
enable_flash_late_interaction=getattr(
|
||||
args, "enable_flash_late_interaction", True
|
||||
),
|
||||
@@ -218,12 +214,7 @@ def get_pooling_invocation_types(
|
||||
# NOTE: Items defined earlier take higher priority
|
||||
invocation_types: list[tuple[RequestType, tuple[GetHandlerFn, EndpointFn]]] = []
|
||||
|
||||
if model_config is None:
|
||||
return invocation_types
|
||||
|
||||
pooling_task = model_config.get_pooling_task(supported_tasks)
|
||||
|
||||
if pooling_task == "embed":
|
||||
if "embed" in supported_tasks:
|
||||
from .embed.api_router import create_embedding, embedding
|
||||
from .embed.protocol import EmbeddingRequest
|
||||
|
||||
@@ -231,7 +222,7 @@ def get_pooling_invocation_types(
|
||||
(EmbeddingRequest, (embedding, create_embedding)),
|
||||
]
|
||||
|
||||
if pooling_task == "classify":
|
||||
if "classify" in supported_tasks:
|
||||
from .classify.api_router import classify, create_classify
|
||||
from .classify.protocol import ClassificationRequest
|
||||
|
||||
|
||||
@@ -78,15 +78,17 @@ class ServingPooling(PoolingServingBase):
|
||||
|
||||
# plugin task uses io_processor.parse_request to verify inputs
|
||||
if pooling_task != "plugin" and pooling_task != self.pooling_task:
|
||||
if pooling_task not in self.supported_tasks:
|
||||
if pooling_task not in self.io_processors:
|
||||
raise ValueError(
|
||||
f"Unsupported task: {pooling_task!r} "
|
||||
f"Supported tasks: {self.supported_tasks}"
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
"Try switching the model's pooling_task "
|
||||
f"via --pooler-config.task {request.task}."
|
||||
logger.warning_once(
|
||||
"Pooling multitask support is deprecated and will be removed "
|
||||
"in v0.20. When the default pooling task is not what you want, you "
|
||||
"need to manually specify it via --pooler-config.task %s. ",
|
||||
pooling_task,
|
||||
)
|
||||
|
||||
if pooling_task == "plugin" and "plugin" not in self.io_processors:
|
||||
|
||||
@@ -8,7 +8,6 @@ from vllm.engine.protocol import EngineClient
|
||||
from vllm.entrypoints.openai.engine.protocol import UsageInfo
|
||||
from vllm.logger import init_logger
|
||||
from vllm.outputs import PoolingRequestOutput, ScoringRequestOutput
|
||||
from vllm.tasks import SCORE_TYPE_MAP, SupportedTask
|
||||
from vllm.v1.pool.late_interaction import (
|
||||
build_late_interaction_doc_params,
|
||||
build_late_interaction_query_params,
|
||||
@@ -39,15 +38,10 @@ class ServingScores(PoolingServing):
|
||||
self,
|
||||
engine_client: EngineClient,
|
||||
*args,
|
||||
supported_tasks: tuple[SupportedTask, ...],
|
||||
enable_flash_late_interaction: bool = True,
|
||||
**kwargs,
|
||||
):
|
||||
pooling_task = engine_client.model_config.get_pooling_task(supported_tasks)
|
||||
score_type = SCORE_TYPE_MAP.get(pooling_task, None) # type: ignore[arg-type]
|
||||
assert score_type is not None
|
||||
|
||||
self.io_processor_name: str = score_type
|
||||
self.io_processor_name: str = engine_client.model_config.score_type
|
||||
self.enable_flash_late_interaction = (
|
||||
self.io_processor_name == "late-interaction"
|
||||
and enable_flash_late_interaction
|
||||
|
||||
@@ -141,14 +141,10 @@ def enable_scoring_api(
|
||||
supported_tasks: tuple["SupportedTask", ...],
|
||||
model_config: ModelConfig | None = None,
|
||||
) -> bool:
|
||||
if model_config is None:
|
||||
return False
|
||||
|
||||
pooling_task = model_config.get_pooling_task(supported_tasks)
|
||||
if pooling_task in ("embed", "token_embed"):
|
||||
if any(t in supported_tasks for t in ("embed", "token_embed")):
|
||||
return True
|
||||
|
||||
if pooling_task == "classify":
|
||||
if model_config is not None and "classify" in supported_tasks:
|
||||
num_labels = getattr(model_config.hf_config, "num_labels", 0)
|
||||
if num_labels != 1:
|
||||
logger.debug_once("Scoring API is only enabled for num_labels == 1.")
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -87,6 +87,13 @@ class PoolingParams(
|
||||
return deepcopy(self)
|
||||
|
||||
def verify(self, model_config: ModelConfig) -> None:
|
||||
if self.task == "score":
|
||||
logger.warning_once(
|
||||
"`score` task is deprecated and will be removed in v0.20. "
|
||||
"Please use `classify` instead."
|
||||
)
|
||||
self.task = "classify"
|
||||
|
||||
# plugin task uses io_processor.parse_request to verify inputs,
|
||||
# skipping PoolingParams verify
|
||||
if self.task == "plugin":
|
||||
|
||||
@@ -16,11 +16,6 @@ PoolingTask = Literal[
|
||||
POOLING_TASKS: tuple[PoolingTask, ...] = get_args(PoolingTask)
|
||||
|
||||
ScoreType = Literal["bi-encoder", "cross-encoder", "late-interaction"]
|
||||
SCORE_TYPE_MAP: dict[PoolingTask, ScoreType] = {
|
||||
"embed": "bi-encoder",
|
||||
"classify": "cross-encoder",
|
||||
"token_embed": "late-interaction",
|
||||
}
|
||||
|
||||
FrontendTask = Literal["render"]
|
||||
FRONTEND_TASKS: tuple[FrontendTask, ...] = get_args(FrontendTask)
|
||||
|
||||
@@ -34,6 +34,7 @@ class KVCacheCoordinator(ABC):
|
||||
self,
|
||||
kv_cache_config: KVCacheConfig,
|
||||
max_model_len: int,
|
||||
max_num_batched_tokens: int,
|
||||
use_eagle: bool,
|
||||
enable_caching: bool,
|
||||
enable_kv_cache_events: bool,
|
||||
@@ -65,6 +66,8 @@ class KVCacheCoordinator(ABC):
|
||||
self.single_type_managers = tuple(
|
||||
get_manager_for_kv_cache_spec(
|
||||
kv_cache_spec=kv_cache_group.kv_cache_spec,
|
||||
max_num_batched_tokens=max_num_batched_tokens,
|
||||
max_model_len=max_model_len,
|
||||
block_pool=self.block_pool,
|
||||
enable_caching=enable_caching,
|
||||
kv_cache_group_id=i,
|
||||
@@ -271,6 +274,7 @@ class KVCacheCoordinatorNoPrefixCache(KVCacheCoordinator):
|
||||
self,
|
||||
kv_cache_config: KVCacheConfig,
|
||||
max_model_len: int,
|
||||
max_num_batched_tokens: int,
|
||||
use_eagle: bool,
|
||||
enable_kv_cache_events: bool,
|
||||
dcp_world_size: int,
|
||||
@@ -281,6 +285,7 @@ class KVCacheCoordinatorNoPrefixCache(KVCacheCoordinator):
|
||||
super().__init__(
|
||||
kv_cache_config,
|
||||
max_model_len,
|
||||
max_num_batched_tokens,
|
||||
use_eagle,
|
||||
False,
|
||||
enable_kv_cache_events,
|
||||
@@ -316,6 +321,7 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator):
|
||||
self,
|
||||
kv_cache_config: KVCacheConfig,
|
||||
max_model_len: int,
|
||||
max_num_batched_tokens: int,
|
||||
use_eagle: bool,
|
||||
enable_caching: bool,
|
||||
enable_kv_cache_events: bool,
|
||||
@@ -327,6 +333,7 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator):
|
||||
super().__init__(
|
||||
kv_cache_config,
|
||||
max_model_len,
|
||||
max_num_batched_tokens,
|
||||
use_eagle,
|
||||
enable_caching,
|
||||
enable_kv_cache_events,
|
||||
@@ -381,6 +388,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
|
||||
self,
|
||||
kv_cache_config: KVCacheConfig,
|
||||
max_model_len: int,
|
||||
max_num_batched_tokens: int,
|
||||
use_eagle: bool,
|
||||
enable_caching: bool,
|
||||
enable_kv_cache_events: bool,
|
||||
@@ -392,6 +400,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
|
||||
super().__init__(
|
||||
kv_cache_config,
|
||||
max_model_len,
|
||||
max_num_batched_tokens,
|
||||
use_eagle,
|
||||
enable_caching,
|
||||
enable_kv_cache_events,
|
||||
@@ -574,6 +583,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
|
||||
def get_kv_cache_coordinator(
|
||||
kv_cache_config: KVCacheConfig,
|
||||
max_model_len: int,
|
||||
max_num_batched_tokens: int,
|
||||
use_eagle: bool,
|
||||
enable_caching: bool,
|
||||
enable_kv_cache_events: bool,
|
||||
@@ -586,6 +596,7 @@ def get_kv_cache_coordinator(
|
||||
return KVCacheCoordinatorNoPrefixCache(
|
||||
kv_cache_config,
|
||||
max_model_len,
|
||||
max_num_batched_tokens,
|
||||
use_eagle,
|
||||
enable_kv_cache_events,
|
||||
dcp_world_size=dcp_world_size,
|
||||
@@ -597,6 +608,7 @@ def get_kv_cache_coordinator(
|
||||
return UnitaryKVCacheCoordinator(
|
||||
kv_cache_config,
|
||||
max_model_len,
|
||||
max_num_batched_tokens,
|
||||
use_eagle,
|
||||
enable_caching,
|
||||
enable_kv_cache_events,
|
||||
@@ -608,6 +620,7 @@ def get_kv_cache_coordinator(
|
||||
return HybridKVCacheCoordinator(
|
||||
kv_cache_config,
|
||||
max_model_len,
|
||||
max_num_batched_tokens,
|
||||
use_eagle,
|
||||
enable_caching,
|
||||
enable_kv_cache_events,
|
||||
|
||||
@@ -109,6 +109,7 @@ class KVCacheManager:
|
||||
kv_cache_config: KVCacheConfig,
|
||||
max_model_len: int,
|
||||
hash_block_size: int,
|
||||
max_num_batched_tokens: int | None = None,
|
||||
enable_caching: bool = True,
|
||||
use_eagle: bool = False,
|
||||
log_stats: bool = False,
|
||||
@@ -118,6 +119,11 @@ class KVCacheManager:
|
||||
metrics_collector: KVCacheMetricsCollector | None = None,
|
||||
) -> None:
|
||||
self.max_model_len = max_model_len
|
||||
# When unset, fall back to `max_model_len` so the recycling-aware cap
|
||||
# collapses to the prior (uncapped) admission behavior. The scheduler
|
||||
# always supplies the real value at runtime.
|
||||
if max_num_batched_tokens is None:
|
||||
max_num_batched_tokens = max_model_len
|
||||
|
||||
self.enable_caching = enable_caching
|
||||
self.use_eagle = use_eagle
|
||||
@@ -131,6 +137,7 @@ class KVCacheManager:
|
||||
self.coordinator = get_kv_cache_coordinator(
|
||||
kv_cache_config=kv_cache_config,
|
||||
max_model_len=self.max_model_len,
|
||||
max_num_batched_tokens=max_num_batched_tokens,
|
||||
use_eagle=self.use_eagle,
|
||||
enable_caching=self.enable_caching,
|
||||
enable_kv_cache_events=enable_kv_cache_events,
|
||||
|
||||
@@ -228,6 +228,7 @@ class Scheduler(SchedulerInterface):
|
||||
self.kv_cache_manager = KVCacheManager(
|
||||
kv_cache_config=kv_cache_config,
|
||||
max_model_len=self.max_model_len,
|
||||
max_num_batched_tokens=self.scheduler_config.max_num_batched_tokens,
|
||||
enable_caching=self.cache_config.enable_prefix_caching,
|
||||
use_eagle=self.use_eagle,
|
||||
log_stats=self.log_stats,
|
||||
|
||||
@@ -41,6 +41,7 @@ class SingleTypeKVCacheManager(ABC):
|
||||
kv_cache_group_id: int,
|
||||
dcp_world_size: int = 1,
|
||||
pcp_world_size: int = 1,
|
||||
max_admission_blocks_per_request: int | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Initializes the SingleTypeKVCacheManager.
|
||||
@@ -48,6 +49,12 @@ class SingleTypeKVCacheManager(ABC):
|
||||
kv_cache_spec: The kv_cache_spec for this manager.
|
||||
block_pool: The block pool.
|
||||
kv_cache_group_id: The id of the kv cache group of this manager.
|
||||
max_admission_blocks_per_request: Recycling-aware per-request
|
||||
block cap used by `get_num_blocks_to_allocate`. Only set for
|
||||
spec types that recycle blocks across chunks (SWA,
|
||||
chunked-local); `None` (the default) means no cap, which is
|
||||
correct for full-attention-style specs that hold every
|
||||
block until the request finishes.
|
||||
"""
|
||||
self.block_size = kv_cache_spec.block_size
|
||||
self.dcp_world_size = dcp_world_size
|
||||
@@ -57,6 +64,7 @@ class SingleTypeKVCacheManager(ABC):
|
||||
self.kv_cache_spec = kv_cache_spec
|
||||
self.block_pool = block_pool
|
||||
self.enable_caching = enable_caching
|
||||
self._max_admission_blocks_per_request = max_admission_blocks_per_request
|
||||
self.new_block_ids: list[int] = []
|
||||
|
||||
# Mapping from request ID to blocks to track the blocks allocated
|
||||
@@ -105,6 +113,19 @@ class SingleTypeKVCacheManager(ABC):
|
||||
"""
|
||||
|
||||
num_required_blocks = cdiv(num_tokens, self.block_size)
|
||||
if self._max_admission_blocks_per_request is not None:
|
||||
# Recycling-aware specs (SWA, chunked-local) cap the per-request
|
||||
# reservation here so admission matches the startup pool sizer
|
||||
# (`SlidingWindowSpec.max_admission_blocks_per_request` / its
|
||||
# chunked-local counterpart). `remove_skipped_blocks` runs from
|
||||
# `allocate_slots` before each chunk's `get_num_blocks_to_allocate`,
|
||||
# so per-request peak real-held blocks <= this cap, which keeps
|
||||
# `sum(reservations) <= pool` <=> `sum(peak_real_held) <= pool`.
|
||||
# Drift between the two would re-introduce the deadlock from
|
||||
# issue #39734 or, worse, mid-prefill OOM.
|
||||
num_required_blocks = min(
|
||||
num_required_blocks, self._max_admission_blocks_per_request
|
||||
)
|
||||
num_req_blocks = len(self.req_to_blocks.get(request_id, ()))
|
||||
|
||||
if request_id in self.num_cached_block:
|
||||
@@ -1126,8 +1147,21 @@ spec_manager_map: dict[type[KVCacheSpec], type[SingleTypeKVCacheManager]] = {
|
||||
|
||||
|
||||
def get_manager_for_kv_cache_spec(
|
||||
kv_cache_spec: KVCacheSpec, **kwargs
|
||||
kv_cache_spec: KVCacheSpec,
|
||||
max_num_batched_tokens: int,
|
||||
max_model_len: int,
|
||||
**kwargs,
|
||||
) -> SingleTypeKVCacheManager:
|
||||
manager_class = spec_manager_map[type(kv_cache_spec)]
|
||||
# SlidingWindow / ChunkedLocalAttention managers recycle blocks across
|
||||
# chunks; the runtime admission cap must match the recycling-aware bound
|
||||
# the startup pool sizer uses (single source of truth: the spec method).
|
||||
if isinstance(kv_cache_spec, (SlidingWindowSpec, ChunkedLocalAttentionSpec)):
|
||||
kwargs["max_admission_blocks_per_request"] = (
|
||||
kv_cache_spec.max_admission_blocks_per_request(
|
||||
max_num_batched_tokens=max_num_batched_tokens,
|
||||
max_model_len=max_model_len,
|
||||
)
|
||||
)
|
||||
manager = manager_class(kv_cache_spec, **kwargs)
|
||||
return manager
|
||||
|
||||
@@ -376,19 +376,28 @@ class MLAAttentionSpec(FullAttentionSpec):
|
||||
class ChunkedLocalAttentionSpec(AttentionSpec):
|
||||
attention_chunk_size: int
|
||||
|
||||
def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
|
||||
max_model_len = vllm_config.model_config.max_model_len
|
||||
max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens
|
||||
def max_admission_blocks_per_request(
|
||||
self, max_num_batched_tokens: int, max_model_len: int
|
||||
) -> int:
|
||||
"""Per-request admission cap, in blocks.
|
||||
|
||||
# During chunked prefill, we allocate KV cache for at most
|
||||
# `self.attention_chunk_size` computed tokens plus the newly scheduled
|
||||
# tokens. And we won't allocate KV cache for more than `max_model_len`
|
||||
# tokens.
|
||||
Single source of truth for both startup pool sizing
|
||||
(`max_memory_usage_bytes`) and the runtime admission gate, so requests
|
||||
admitted by startup can also be admitted at runtime.
|
||||
"""
|
||||
# During chunked prefill, we hold KV for at most one chunk window.
|
||||
num_tokens = min(
|
||||
self.attention_chunk_size + max_num_batched_tokens, max_model_len
|
||||
)
|
||||
return cdiv(num_tokens, self.block_size)
|
||||
|
||||
return cdiv(num_tokens, self.block_size) * self.page_size_bytes
|
||||
def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
|
||||
max_model_len = vllm_config.model_config.max_model_len
|
||||
max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens
|
||||
max_blocks = self.max_admission_blocks_per_request(
|
||||
max_num_batched_tokens=max_num_batched_tokens, max_model_len=max_model_len
|
||||
)
|
||||
return max_blocks * self.page_size_bytes
|
||||
|
||||
|
||||
@dataclass(frozen=True, kw_only=True)
|
||||
@@ -409,26 +418,38 @@ class SlidingWindowSpec(AttentionSpec):
|
||||
* get_dtype_size(self.dtype)
|
||||
)
|
||||
|
||||
def max_admission_blocks_per_request(
|
||||
self, max_num_batched_tokens: int, max_model_len: int
|
||||
) -> int:
|
||||
"""Per-request admission cap, in blocks.
|
||||
|
||||
Single source of truth for both startup pool sizing
|
||||
(`max_memory_usage_bytes`) and the runtime admission gate. Per-request
|
||||
real-held blocks plateau at this bound because
|
||||
`SlidingWindowManager.remove_skipped_blocks` runs from `allocate_slots`
|
||||
before each chunk's `get_num_blocks_to_allocate`.
|
||||
"""
|
||||
# During chunked prefill, we hold KV for the last `sliding_window-1`
|
||||
# computed tokens plus the newly scheduled tokens, and never more
|
||||
# than `max_model_len`.
|
||||
num_tokens = min(
|
||||
self.sliding_window - 1 + max_num_batched_tokens, max_model_len
|
||||
)
|
||||
# +1 because the sliding window may not start from the beginning of
|
||||
# the block. E.g. block size 4 and num_token 4 needs two blocks
|
||||
# [XXCD][EF] to store the 6-token window [CDEF].
|
||||
return cdiv(num_tokens, self.block_size) + 1
|
||||
|
||||
def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
|
||||
assert vllm_config.parallel_config.decode_context_parallel_size == 1, (
|
||||
"DCP not support sliding window."
|
||||
)
|
||||
max_model_len = vllm_config.model_config.max_model_len
|
||||
max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens
|
||||
|
||||
# During chunked prefill, we allocate KV cache for the last
|
||||
# `self.sliding_window-1` computed tokens plus the newly scheduled
|
||||
# tokens. And we won't allocate KV cache for more than `max_model_len`
|
||||
# tokens.
|
||||
num_tokens = min(
|
||||
self.sliding_window - 1 + max_num_batched_tokens, max_model_len
|
||||
max_blocks = self.max_admission_blocks_per_request(
|
||||
max_num_batched_tokens=max_num_batched_tokens, max_model_len=max_model_len
|
||||
)
|
||||
|
||||
# +1 here because the sliding window may not start from the beginning
|
||||
# of the block. For example, if the block size is 4 and num_token
|
||||
# is 4, we need two blocks [XXCD] [EF] to store the sliding
|
||||
# window [CDEF] of 6 tokens.
|
||||
return (cdiv(num_tokens, self.block_size) + 1) * self.page_size_bytes
|
||||
return max_blocks * self.page_size_bytes
|
||||
|
||||
|
||||
@dataclass(frozen=True, kw_only=True)
|
||||
|
||||
@@ -110,6 +110,9 @@ class SimpleCPUOffloadScheduler:
|
||||
self.cpu_coordinator: KVCacheCoordinator = get_kv_cache_coordinator(
|
||||
kv_cache_config=self.cpu_kv_cache_config,
|
||||
max_model_len=vllm_config.model_config.max_model_len,
|
||||
max_num_batched_tokens=(
|
||||
vllm_config.scheduler_config.max_num_batched_tokens
|
||||
),
|
||||
use_eagle=False,
|
||||
enable_caching=True,
|
||||
enable_kv_cache_events=self.enable_kv_cache_events,
|
||||
|
||||
Reference in New Issue
Block a user