forked from Karylab-cklius/vllm
Compare commits
74
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
54a62a79f7 | ||
|
|
384dc7f77b | ||
|
|
f04d5226f8 | ||
|
|
0a0a1a198b | ||
|
|
6c1cfbad32 | ||
|
|
45f526d652 | ||
|
|
5db91f0aaf | ||
|
|
061980c36a | ||
|
|
7a49742b88 | ||
|
|
3e6a1e1686 | ||
|
|
7961486a9b | ||
|
|
4f9b14c21c | ||
|
|
31a458c091 | ||
|
|
a3a51d20e7 | ||
|
|
e5b807607c | ||
|
|
fd4d96302a | ||
|
|
c0f011918d | ||
|
|
e6ae4b1be1 | ||
|
|
2dccb38f73 | ||
|
|
d157216093 | ||
|
|
93f3c8e531 | ||
|
|
2cc26c3a99 | ||
|
|
dfa8852db2 | ||
|
|
714c6e0eab | ||
|
|
0fefd00e6c | ||
|
|
f5c081d432 | ||
|
|
c88ea8338b | ||
|
|
9f9ecff4cd | ||
|
|
ca1954d58c | ||
|
|
55e6d3d5c0 | ||
|
|
6682c231fa | ||
|
|
5ae685c1c8 | ||
|
|
ce8cf9161d | ||
|
|
18be11fd59 | ||
|
|
8d8855fdae | ||
|
|
e855d380fa | ||
|
|
0e5a9382af | ||
|
|
04bf5a35fa | ||
|
|
43a73f853b | ||
|
|
ffbc2e5bdb | ||
|
|
f9e6db3034 | ||
|
|
d61d2b08e9 | ||
|
|
f5e59ee7a6 | ||
|
|
9b005edc48 | ||
|
|
bf9a185395 | ||
|
|
ad041c79db | ||
|
|
747b068136 | ||
|
|
122f75d939 | ||
|
|
d8f8a7aad2 | ||
|
|
0115e957d4 | ||
|
|
116ed130f4 | ||
|
|
8374387bd8 | ||
|
|
912fbe9555 | ||
|
|
52131f88d9 | ||
|
|
821eb80c0d | ||
|
|
a2956a0f8e | ||
|
|
911355e216 | ||
|
|
8d3f8f485e | ||
|
|
96efb91480 | ||
|
|
2754231ba3 | ||
|
|
2390d44209 | ||
|
|
7362b4450a | ||
|
|
57a314d155 | ||
|
|
d4c57863f7 | ||
|
|
68e1b711f1 | ||
|
|
0024f39a32 | ||
|
|
e9163b536e | ||
|
|
7acaea634c | ||
|
|
697e4ff352 | ||
|
|
a3e2e250f0 | ||
|
|
143e4dccdf | ||
|
|
6590a3ecda | ||
|
|
b3debb7e77 | ||
|
|
458c1a4b2d |
@@ -333,15 +333,15 @@ apply_rocm_test_overrides() {
|
||||
# --- Entrypoint ignores ---
|
||||
if [[ $cmds == *" entrypoints/openai "* ]]; then
|
||||
cmds=${cmds//" entrypoints/openai "/" entrypoints/openai \
|
||||
--ignore=entrypoints/openai/test_audio.py \
|
||||
--ignore=entrypoints/openai/test_shutdown.py \
|
||||
--ignore=entrypoints/openai/chat_completion/test_audio.py \
|
||||
--ignore=entrypoints/openai/completion/test_shutdown.py \
|
||||
--ignore=entrypoints/openai/test_completion.py \
|
||||
--ignore=entrypoints/openai/test_models.py \
|
||||
--ignore=entrypoints/openai/test_lora_adapters.py \
|
||||
--ignore=entrypoints/openai/test_return_tokens_as_ids.py \
|
||||
--ignore=entrypoints/openai/test_root_path.py \
|
||||
--ignore=entrypoints/openai/chat_completion/test_root_path.py \
|
||||
--ignore=entrypoints/openai/test_tokenization.py \
|
||||
--ignore=entrypoints/openai/test_prompt_validation.py "}
|
||||
--ignore=entrypoints/openai/completion/test_prompt_validation.py "}
|
||||
fi
|
||||
|
||||
if [[ $cmds == *" entrypoints/llm "* ]]; then
|
||||
|
||||
+12
-12
@@ -162,7 +162,7 @@ steps:
|
||||
- tests/entrypoints/test_chat_utils
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/test_oot_registration.py --ignore=entrypoints/openai/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses
|
||||
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses
|
||||
- pytest -v -s entrypoints/test_chat_utils.py
|
||||
|
||||
- label: Entrypoints Integration Test (API Server 2)
|
||||
@@ -674,12 +674,12 @@ steps:
|
||||
- vllm/config/model.py
|
||||
- vllm/model_executor
|
||||
- tests/model_executor
|
||||
- tests/entrypoints/openai/test_tensorizer_entrypoint.py
|
||||
- tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
|
||||
commands:
|
||||
- apt-get update && apt-get install -y curl libsodium23
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s model_executor
|
||||
- pytest -v -s entrypoints/openai/test_tensorizer_entrypoint.py
|
||||
- pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py
|
||||
|
||||
- label: Benchmarks # 11min
|
||||
timeout_in_minutes: 20
|
||||
@@ -1143,7 +1143,7 @@ steps:
|
||||
- pytest -v -s plugins_tests/test_scheduler_plugins.py
|
||||
- pip install -e ./plugins/vllm_add_dummy_model
|
||||
- pytest -v -s distributed/test_distributed_oot.py
|
||||
- pytest -v -s entrypoints/openai/test_oot_registration.py
|
||||
- pytest -v -s entrypoints/openai/chat_completion/test_oot_registration.py
|
||||
- pytest -v -s models/test_oot_registration.py
|
||||
- pytest -v -s plugins/lora_resolvers
|
||||
|
||||
@@ -1502,7 +1502,7 @@ steps:
|
||||
- tests/entrypoints/test_chat_utils
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/test_oot_registration.py --ignore=entrypoints/openai/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses
|
||||
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses
|
||||
- pytest -v -s entrypoints/test_chat_utils.py
|
||||
|
||||
- label: Entrypoints Integration Test (API Server 2)
|
||||
@@ -2133,12 +2133,12 @@ steps:
|
||||
- vllm/config/model.py
|
||||
- vllm/model_executor
|
||||
- tests/model_executor
|
||||
- tests/entrypoints/openai/test_tensorizer_entrypoint.py
|
||||
- tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
|
||||
commands:
|
||||
- apt-get update && apt-get install -y curl libsodium23
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s model_executor
|
||||
- pytest -v -s entrypoints/openai/test_tensorizer_entrypoint.py
|
||||
- pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py
|
||||
|
||||
- label: Benchmarks # 11min
|
||||
timeout_in_minutes: 20
|
||||
@@ -2735,7 +2735,7 @@ steps:
|
||||
- pytest -v -s plugins_tests/test_scheduler_plugins.py
|
||||
- pip install -e ./plugins/vllm_add_dummy_model
|
||||
- pytest -v -s distributed/test_distributed_oot.py
|
||||
- pytest -v -s entrypoints/openai/test_oot_registration.py # it needs a clean process
|
||||
- pytest -v -s entrypoints/openai/chat_completion/test_oot_registration.py # it needs a clean process
|
||||
- pytest -v -s models/test_oot_registration.py # it needs a clean process
|
||||
- pytest -v -s plugins/lora_resolvers # unit tests for in-tree lora resolver plugins
|
||||
|
||||
@@ -3257,7 +3257,7 @@ steps:
|
||||
- tests/entrypoints/test_chat_utils
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/test_oot_registration.py --ignore=entrypoints/openai/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses
|
||||
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses
|
||||
- pytest -v -s entrypoints/test_chat_utils.py
|
||||
|
||||
- label: Entrypoints Integration Test (API Server 2)
|
||||
@@ -3872,12 +3872,12 @@ steps:
|
||||
- vllm/config/model.py
|
||||
- vllm/model_executor
|
||||
- tests/model_executor
|
||||
- tests/entrypoints/openai/test_tensorizer_entrypoint.py
|
||||
- tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
|
||||
commands:
|
||||
- apt-get update && apt-get install -y curl libsodium23
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s model_executor
|
||||
- pytest -v -s entrypoints/openai/test_tensorizer_entrypoint.py
|
||||
- pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py
|
||||
|
||||
- label: Benchmarks # 11min
|
||||
timeout_in_minutes: 20
|
||||
@@ -4508,7 +4508,7 @@ steps:
|
||||
- pytest -v -s plugins_tests/test_scheduler_plugins.py
|
||||
- pip install -e ./plugins/vllm_add_dummy_model
|
||||
- pytest -v -s distributed/test_distributed_oot.py
|
||||
- pytest -v -s entrypoints/openai/test_oot_registration.py # it needs a clean process
|
||||
- pytest -v -s entrypoints/openai/chat_completion/test_oot_registration.py # it needs a clean process
|
||||
- pytest -v -s models/test_oot_registration.py # it needs a clean process
|
||||
- pytest -v -s plugins/lora_resolvers # unit tests for in-tree lora resolver plugins
|
||||
|
||||
|
||||
@@ -34,7 +34,7 @@ steps:
|
||||
- tests/entrypoints/test_chat_utils
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/test_oot_registration.py --ignore=entrypoints/openai/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses
|
||||
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses
|
||||
- pytest -v -s entrypoints/test_chat_utils.py
|
||||
mirror:
|
||||
amd:
|
||||
|
||||
@@ -9,9 +9,9 @@ steps:
|
||||
- vllm/config/model.py
|
||||
- vllm/model_executor
|
||||
- tests/model_executor
|
||||
- tests/entrypoints/openai/test_tensorizer_entrypoint.py
|
||||
- tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
|
||||
commands:
|
||||
- apt-get update && apt-get install -y curl libsodium23
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s model_executor
|
||||
- pytest -v -s entrypoints/openai/test_tensorizer_entrypoint.py
|
||||
- pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py
|
||||
|
||||
@@ -36,6 +36,6 @@ steps:
|
||||
- pytest -v -s plugins_tests/test_scheduler_plugins.py
|
||||
- pip install -e ./plugins/vllm_add_dummy_model
|
||||
- pytest -v -s distributed/test_distributed_oot.py
|
||||
- pytest -v -s entrypoints/openai/test_oot_registration.py # it needs a clean process
|
||||
- pytest -v -s entrypoints/openai/chat_completion/test_oot_registration.py # it needs a clean process
|
||||
- pytest -v -s models/test_oot_registration.py # it needs a clean process
|
||||
- pytest -v -s plugins/lora_resolvers # unit tests for in-tree lora resolver plugins
|
||||
|
||||
+2
-2
@@ -27,7 +27,7 @@ pull_request_rules:
|
||||
Hi @{{author}}, the pre-commit checks have failed. Please run:
|
||||
|
||||
```bash
|
||||
uv pip install pre-commit
|
||||
uv pip install pre-commit>=4.5.1
|
||||
pre-commit install
|
||||
pre-commit run --all-files
|
||||
```
|
||||
@@ -381,7 +381,7 @@ pull_request_rules:
|
||||
- or:
|
||||
- files~=^vllm/model_executor/model_loader/tensorizer.py
|
||||
- files~=^vllm/model_executor/model_loader/tensorizer_loader.py
|
||||
- files~=^tests/entrypoints/openai/test_tensorizer_entrypoint.py
|
||||
- files~=^tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
|
||||
- files~=^tests/model_executor/model_loader/tensorizer_loader/
|
||||
actions:
|
||||
assign:
|
||||
|
||||
@@ -47,6 +47,8 @@ from common import (
|
||||
is_mla_backend,
|
||||
)
|
||||
|
||||
from vllm.v1.worker.workspace import init_workspace_manager
|
||||
|
||||
|
||||
def run_standard_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:
|
||||
"""Run standard attention benchmark (Flash/Triton/FlashInfer)."""
|
||||
@@ -462,7 +464,7 @@ def main():
|
||||
parser.add_argument(
|
||||
"--batch-specs",
|
||||
nargs="+",
|
||||
default=["q2k", "8q1s1k"],
|
||||
default=None,
|
||||
help="Batch specifications using extended grammar",
|
||||
)
|
||||
|
||||
@@ -478,6 +480,21 @@ def main():
|
||||
parser.add_argument("--repeats", type=int, default=1, help="Repetitions")
|
||||
parser.add_argument("--warmup-iters", type=int, default=3, help="Warmup iterations")
|
||||
parser.add_argument("--profile-memory", action="store_true", help="Profile memory")
|
||||
parser.add_argument(
|
||||
"--kv-cache-dtype",
|
||||
default="auto",
|
||||
choices=["auto", "fp8"],
|
||||
help="KV cache dtype: auto or fp8",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--cuda-graphs",
|
||||
action=argparse.BooleanOptionalAction,
|
||||
default=True,
|
||||
help=(
|
||||
"Launch kernels with CUDA graphs to eliminate CPU overhead"
|
||||
"in measurements (default: True)"
|
||||
),
|
||||
)
|
||||
|
||||
# Parameter sweep (use YAML config for advanced sweeps)
|
||||
parser.add_argument(
|
||||
@@ -536,21 +553,24 @@ def main():
|
||||
|
||||
# Batch specs and sizes
|
||||
# Support both explicit batch_specs and generated batch_spec_ranges
|
||||
if "batch_spec_ranges" in yaml_config:
|
||||
# Generate batch specs from ranges
|
||||
generated_specs = generate_batch_specs_from_ranges(
|
||||
yaml_config["batch_spec_ranges"]
|
||||
)
|
||||
# Combine with any explicit batch_specs
|
||||
if "batch_specs" in yaml_config:
|
||||
args.batch_specs = yaml_config["batch_specs"] + generated_specs
|
||||
else:
|
||||
args.batch_specs = generated_specs
|
||||
console.print(
|
||||
f"[dim]Generated {len(generated_specs)} batch specs from ranges[/]"
|
||||
)
|
||||
elif "batch_specs" in yaml_config:
|
||||
args.batch_specs = yaml_config["batch_specs"]
|
||||
# CLI --batch-specs takes precedence over YAML when provided.
|
||||
cli_batch_specs_provided = args.batch_specs is not None
|
||||
if not cli_batch_specs_provided:
|
||||
if "batch_spec_ranges" in yaml_config:
|
||||
# Generate batch specs from ranges
|
||||
generated_specs = generate_batch_specs_from_ranges(
|
||||
yaml_config["batch_spec_ranges"]
|
||||
)
|
||||
# Combine with any explicit batch_specs
|
||||
if "batch_specs" in yaml_config:
|
||||
args.batch_specs = yaml_config["batch_specs"] + generated_specs
|
||||
else:
|
||||
args.batch_specs = generated_specs
|
||||
console.print(
|
||||
f"[dim]Generated {len(generated_specs)} batch specs from ranges[/]"
|
||||
)
|
||||
elif "batch_specs" in yaml_config:
|
||||
args.batch_specs = yaml_config["batch_specs"]
|
||||
|
||||
if "batch_sizes" in yaml_config:
|
||||
args.batch_sizes = yaml_config["batch_sizes"]
|
||||
@@ -575,6 +595,10 @@ def main():
|
||||
args.warmup_iters = yaml_config["warmup_iters"]
|
||||
if "profile_memory" in yaml_config:
|
||||
args.profile_memory = yaml_config["profile_memory"]
|
||||
if "kv_cache_dtype" in yaml_config:
|
||||
args.kv_cache_dtype = yaml_config["kv_cache_dtype"]
|
||||
if "cuda_graphs" in yaml_config:
|
||||
args.cuda_graphs = yaml_config["cuda_graphs"]
|
||||
|
||||
# Parameter sweep configuration
|
||||
if "parameter_sweep" in yaml_config:
|
||||
@@ -629,12 +653,18 @@ def main():
|
||||
# Determine backends
|
||||
backends = args.backends or ([args.backend] if args.backend else ["flash"])
|
||||
prefill_backends = getattr(args, "prefill_backends", None)
|
||||
if not args.batch_specs:
|
||||
args.batch_specs = ["q2k", "8q1s1k"]
|
||||
console.print(f"Backends: {', '.join(backends)}")
|
||||
if prefill_backends:
|
||||
console.print(f"Prefill backends: {', '.join(prefill_backends)}")
|
||||
console.print(f"Batch specs: {', '.join(args.batch_specs)}")
|
||||
console.print(f"KV cache dtype: {args.kv_cache_dtype}")
|
||||
console.print(f"CUDA graphs: {args.cuda_graphs}")
|
||||
console.print()
|
||||
|
||||
init_workspace_manager(args.device)
|
||||
|
||||
# Run benchmarks
|
||||
all_results = []
|
||||
|
||||
@@ -687,6 +717,8 @@ def main():
|
||||
repeats=args.repeats,
|
||||
warmup_iters=args.warmup_iters,
|
||||
profile_memory=args.profile_memory,
|
||||
kv_cache_dtype=args.kv_cache_dtype,
|
||||
use_cuda_graphs=args.cuda_graphs,
|
||||
)
|
||||
|
||||
# Add decode pipeline config
|
||||
@@ -839,6 +871,8 @@ def main():
|
||||
"repeats": args.repeats,
|
||||
"warmup_iters": args.warmup_iters,
|
||||
"profile_memory": args.profile_memory,
|
||||
"kv_cache_dtype": args.kv_cache_dtype,
|
||||
"use_cuda_graphs": args.cuda_graphs,
|
||||
}
|
||||
all_results = run_model_parameter_sweep(
|
||||
backends,
|
||||
@@ -861,6 +895,8 @@ def main():
|
||||
"repeats": args.repeats,
|
||||
"warmup_iters": args.warmup_iters,
|
||||
"profile_memory": args.profile_memory,
|
||||
"kv_cache_dtype": args.kv_cache_dtype,
|
||||
"use_cuda_graphs": args.cuda_graphs,
|
||||
}
|
||||
all_results = run_parameter_sweep(
|
||||
backends, args.batch_specs, base_config_args, args.parameter_sweep, console
|
||||
@@ -891,6 +927,8 @@ def main():
|
||||
repeats=args.repeats,
|
||||
warmup_iters=args.warmup_iters,
|
||||
profile_memory=args.profile_memory,
|
||||
kv_cache_dtype=args.kv_cache_dtype,
|
||||
use_cuda_graphs=args.cuda_graphs,
|
||||
)
|
||||
|
||||
result = run_benchmark(config)
|
||||
|
||||
@@ -213,6 +213,9 @@ class BenchmarkConfig:
|
||||
profile_memory: bool = False
|
||||
use_cuda_graphs: bool = False
|
||||
|
||||
# "auto" or "fp8"
|
||||
kv_cache_dtype: str = "auto"
|
||||
|
||||
# MLA-specific
|
||||
prefill_backend: str | None = None
|
||||
kv_lora_rank: int | None = None
|
||||
@@ -369,6 +372,7 @@ class ResultsFormatter:
|
||||
"backend",
|
||||
"batch_spec",
|
||||
"num_layers",
|
||||
"kv_cache_dtype",
|
||||
"mean_time",
|
||||
"std_time",
|
||||
"throughput",
|
||||
@@ -382,6 +386,7 @@ class ResultsFormatter:
|
||||
"backend": r.config.backend,
|
||||
"batch_spec": r.config.batch_spec,
|
||||
"num_layers": r.config.num_layers,
|
||||
"kv_cache_dtype": r.config.kv_cache_dtype,
|
||||
"mean_time": r.mean_time,
|
||||
"std_time": r.std_time,
|
||||
"throughput": r.throughput_tokens_per_sec or 0,
|
||||
|
||||
@@ -30,9 +30,9 @@ batch_specs:
|
||||
- "2q16k_32q1s4k" # 2 very large prefill + 32 decode
|
||||
|
||||
# Context extension + decode
|
||||
- "2q1kkv2k_16q1s1k" # 2 extend + 16 decode
|
||||
- "4q2kkv4k_32q1s2k" # 4 extend + 32 decode
|
||||
- "2q1kkv8k_32q1s2k" # 2 large extend + 32 decode
|
||||
- "2q1ks2k_16q1s1k" # 2 extend + 16 decode
|
||||
- "4q2ks4k_32q1s2k" # 4 extend + 32 decode
|
||||
- "2q1ks8k_32q1s2k" # 2 large extend + 32 decode
|
||||
|
||||
# Explicitly chunked prefill
|
||||
- "q8k" # 8k prefill with chunking hint
|
||||
|
||||
@@ -0,0 +1,58 @@
|
||||
# MLA decode-only benchmark configuration
|
||||
|
||||
model:
|
||||
name: "deepseek-v3"
|
||||
num_layers: 60
|
||||
num_q_heads: 128 # Base value, can be swept for TP simulation
|
||||
num_kv_heads: 1 # MLA uses single latent KV
|
||||
head_dim: 576
|
||||
kv_lora_rank: 512
|
||||
qk_nope_head_dim: 128
|
||||
qk_rope_head_dim: 64
|
||||
v_head_dim: 128
|
||||
block_size: 128 # CUTLASS MLA and FlashAttn MLA use 128
|
||||
|
||||
# Model parameter sweep: simulate tensor parallelism by varying num_q_heads
|
||||
# TP=1: 128 heads, TP=2: 64 heads, TP=4: 32 heads, TP=8: 16 heads
|
||||
model_parameter_sweep:
|
||||
param_name: "num_q_heads"
|
||||
values: [128, 64, 32, 16]
|
||||
label_format: "{backend}_{value}h"
|
||||
|
||||
batch_specs:
|
||||
# Small batches, varying sequence lengths
|
||||
- "16q1s512" # 16 requests, 512 KV cache
|
||||
- "16q1s1k" # 16 requests, 1k KV cache
|
||||
- "16q1s2k" # 16 requests, 2k KV cache
|
||||
- "16q1s4k" # 16 requests, 4k KV cache
|
||||
|
||||
# Medium batches
|
||||
- "32q1s1k" # 32 requests, 1k KV cache
|
||||
- "32q1s2k" # 32 requests, 2k KV cache
|
||||
- "32q1s4k" # 32 requests, 4k KV cache
|
||||
- "32q1s8k" # 32 requests, 8k KV cache
|
||||
|
||||
# Large batches
|
||||
- "64q1s1k" # 64 requests, 1k KV cache
|
||||
- "64q1s2k" # 64 requests, 2k KV cache
|
||||
- "64q1s4k" # 64 requests, 4k KV cache
|
||||
- "64q1s8k" # 64 requests, 8k KV cache
|
||||
|
||||
# Very large batches
|
||||
- "128q1s1k" # 128 requests, 1k KV cache
|
||||
- "128q1s2k" # 128 requests, 2k KV cache
|
||||
- "128q1s4k" # 128 requests, 4k KV cache
|
||||
- "128q1s8k" # 128 requests, 8k KV cache
|
||||
|
||||
# Long context
|
||||
- "32q1s16k" # 32 requests, 16k KV cache
|
||||
- "32q1s32k" # 32 requests, 32k KV cache
|
||||
|
||||
backends:
|
||||
- FLASHMLA_SPARSE
|
||||
- FLASHINFER_MLA_SPARSE
|
||||
|
||||
device: "cuda:0"
|
||||
repeats: 100
|
||||
warmup_iters: 10
|
||||
profile_memory: true
|
||||
@@ -60,9 +60,11 @@ def create_minimal_vllm_config(
|
||||
model_name: str = "deepseek-v3",
|
||||
block_size: int = 128,
|
||||
max_num_seqs: int = 256,
|
||||
max_num_batched_tokens: int = 8192,
|
||||
mla_dims: dict | None = None,
|
||||
index_topk: int | None = None,
|
||||
prefill_backend: str | None = None,
|
||||
kv_cache_dtype: str = "auto",
|
||||
) -> VllmConfig:
|
||||
"""
|
||||
Create minimal VllmConfig for MLA benchmarks.
|
||||
@@ -149,13 +151,13 @@ def create_minimal_vllm_config(
|
||||
cache_config = CacheConfig(
|
||||
block_size=block_size,
|
||||
gpu_memory_utilization=0.9,
|
||||
cache_dtype="auto",
|
||||
cache_dtype=kv_cache_dtype,
|
||||
enable_prefix_caching=False,
|
||||
)
|
||||
|
||||
scheduler_config = SchedulerConfig(
|
||||
max_num_seqs=max_num_seqs,
|
||||
max_num_batched_tokens=8192,
|
||||
max_num_batched_tokens=max(max_num_batched_tokens, max_num_seqs),
|
||||
max_model_len=32768,
|
||||
is_encoder_decoder=False,
|
||||
enable_chunked_prefill=True,
|
||||
@@ -535,6 +537,7 @@ def _create_backend_impl(
|
||||
device: torch.device,
|
||||
max_num_tokens: int = 8192,
|
||||
index_topk: int | None = None,
|
||||
kv_cache_dtype: str = "auto",
|
||||
):
|
||||
"""
|
||||
Create backend implementation instance.
|
||||
@@ -583,7 +586,7 @@ def _create_backend_impl(
|
||||
"num_kv_heads": mla_dims["num_kv_heads"],
|
||||
"alibi_slopes": None,
|
||||
"sliding_window": None,
|
||||
"kv_cache_dtype": "auto",
|
||||
"kv_cache_dtype": kv_cache_dtype,
|
||||
"logits_soft_cap": None,
|
||||
"attn_type": "decoder",
|
||||
"kv_sharing_target_layer_name": None,
|
||||
@@ -701,6 +704,7 @@ def _run_single_benchmark(
|
||||
mla_dims: dict,
|
||||
device: torch.device,
|
||||
indexer=None,
|
||||
kv_cache_dtype: str | None = None,
|
||||
) -> BenchmarkResult:
|
||||
"""
|
||||
Run a single benchmark iteration.
|
||||
@@ -734,49 +738,124 @@ def _run_single_benchmark(
|
||||
)
|
||||
|
||||
# Create KV cache
|
||||
kv_cache = torch.zeros(
|
||||
num_blocks,
|
||||
block_size,
|
||||
mla_dims["kv_lora_rank"] + mla_dims["qk_rope_head_dim"],
|
||||
device=device,
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
if kv_cache_dtype is None:
|
||||
kv_cache_dtype = getattr(config, "kv_cache_dtype", "auto")
|
||||
head_size = mla_dims["kv_lora_rank"] + mla_dims["qk_rope_head_dim"]
|
||||
if kv_cache_dtype == "fp8_ds_mla":
|
||||
# FlashMLA sparse custom format: 656 bytes per token, stored as uint8.
|
||||
# Layout: kv_lora_rank fp8 bytes + 4 float32 tile scales
|
||||
# + 2*rope_dim bf16 bytes
|
||||
# = 512 + 16 + 128 = 656 bytes for DeepSeek dims.
|
||||
kv_cache = torch.zeros(
|
||||
num_blocks,
|
||||
block_size,
|
||||
656,
|
||||
device=device,
|
||||
dtype=torch.uint8,
|
||||
)
|
||||
elif kv_cache_dtype == "fp8":
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
# Create input tensors for both decode and prefill modes
|
||||
decode_inputs, prefill_inputs = _create_input_tensors(
|
||||
total_q,
|
||||
mla_dims,
|
||||
backend_cfg["query_format"],
|
||||
device,
|
||||
torch.bfloat16,
|
||||
)
|
||||
kv_cache = torch.zeros(
|
||||
num_blocks,
|
||||
block_size,
|
||||
head_size,
|
||||
device=device,
|
||||
dtype=torch.uint8,
|
||||
).view(current_platform.fp8_dtype())
|
||||
else:
|
||||
kv_cache = torch.zeros(
|
||||
num_blocks,
|
||||
block_size,
|
||||
head_size,
|
||||
device=device,
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
|
||||
# Fill indexer with random indices for sparse backends
|
||||
is_sparse = backend_cfg.get("is_sparse", False)
|
||||
if is_sparse and indexer is not None:
|
||||
indexer.fill_random_indices(total_q, max_kv_len)
|
||||
|
||||
# Determine which forward method to use based on metadata
|
||||
if metadata.decode is not None:
|
||||
forward_fn = lambda: impl.forward_mqa(decode_inputs, kv_cache, metadata, layer)
|
||||
elif metadata.prefill is not None:
|
||||
forward_fn = lambda: impl.forward_mha(
|
||||
prefill_inputs["q"],
|
||||
prefill_inputs["k_c_normed"],
|
||||
prefill_inputs["k_pe"],
|
||||
kv_cache,
|
||||
metadata,
|
||||
prefill_inputs["k_scale"],
|
||||
prefill_inputs["output"],
|
||||
)
|
||||
else:
|
||||
# Determine which forward methods to use based on metadata.
|
||||
# Sparse MLA backends always use forward_mqa
|
||||
has_decode = is_sparse or getattr(metadata, "decode", None) is not None
|
||||
has_prefill = not is_sparse and getattr(metadata, "prefill", None) is not None
|
||||
if not has_decode and not has_prefill:
|
||||
raise RuntimeError("Metadata has neither decode nor prefill metadata")
|
||||
|
||||
num_decode = (
|
||||
metadata.num_decode_tokens
|
||||
if (has_decode and has_prefill)
|
||||
else total_q
|
||||
if has_decode
|
||||
else 0
|
||||
)
|
||||
num_prefill = total_q - num_decode
|
||||
|
||||
# Some backends requires fp8 queries when using fp8 KV cache.
|
||||
is_fp8_kvcache = kv_cache_dtype.startswith("fp8")
|
||||
quantize_query = is_fp8_kvcache and getattr(
|
||||
impl, "supports_quant_query_input", False
|
||||
)
|
||||
|
||||
# quantize_query forces concat format
|
||||
query_fmt = "concat" if quantize_query else backend_cfg["query_format"]
|
||||
|
||||
# Create decode query tensors
|
||||
if has_decode:
|
||||
decode_inputs, _ = _create_input_tensors(
|
||||
num_decode, mla_dims, query_fmt, device, torch.bfloat16
|
||||
)
|
||||
# Cast decode query to fp8 if the backend supports it
|
||||
if quantize_query:
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
if isinstance(decode_inputs, tuple):
|
||||
decode_inputs = torch.cat(list(decode_inputs), dim=-1)
|
||||
decode_inputs = decode_inputs.to(current_platform.fp8_dtype())
|
||||
|
||||
# Create prefill input tensors
|
||||
if has_prefill:
|
||||
_, prefill_inputs = _create_input_tensors(
|
||||
num_prefill, mla_dims, query_fmt, device, torch.bfloat16
|
||||
)
|
||||
|
||||
# Build forward function
|
||||
def forward_fn():
|
||||
results = []
|
||||
if has_decode:
|
||||
results.append(impl.forward_mqa(decode_inputs, kv_cache, metadata, layer))
|
||||
if has_prefill:
|
||||
results.append(
|
||||
impl.forward_mha(
|
||||
prefill_inputs["q"],
|
||||
prefill_inputs["k_c_normed"],
|
||||
prefill_inputs["k_pe"],
|
||||
kv_cache,
|
||||
metadata,
|
||||
prefill_inputs["k_scale"],
|
||||
prefill_inputs["output"],
|
||||
)
|
||||
)
|
||||
return results[0] if len(results) == 1 else tuple(results)
|
||||
|
||||
# Warmup
|
||||
for _ in range(config.warmup_iters):
|
||||
forward_fn()
|
||||
torch.accelerator.synchronize()
|
||||
|
||||
# Optionally capture a CUDA graph after warmup.
|
||||
# Graph replay eliminates CPU launch overhead so timings reflect pure
|
||||
# kernel time.
|
||||
if config.use_cuda_graphs:
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(graph):
|
||||
forward_fn()
|
||||
benchmark_fn = graph.replay
|
||||
else:
|
||||
benchmark_fn = forward_fn
|
||||
|
||||
# Benchmark
|
||||
times = []
|
||||
for _ in range(config.repeats):
|
||||
@@ -785,7 +864,7 @@ def _run_single_benchmark(
|
||||
|
||||
start.record()
|
||||
for _ in range(config.num_layers):
|
||||
forward_fn()
|
||||
benchmark_fn()
|
||||
end.record()
|
||||
|
||||
torch.accelerator.synchronize()
|
||||
@@ -852,13 +931,30 @@ def _run_mla_benchmark_batched(
|
||||
# Determine if this is a sparse backend
|
||||
is_sparse = backend_cfg.get("is_sparse", False)
|
||||
|
||||
# Extract kv_cache_dtype from the first config
|
||||
kv_cache_dtype = getattr(first_config, "kv_cache_dtype", "auto")
|
||||
|
||||
# FlashMLA sparse only supports "fp8_ds_mla" internally (not generic "fp8").
|
||||
# Remap here so the user can pass --kv-cache-dtype fp8 regardless of backend.
|
||||
if backend.upper() == "FLASHMLA_SPARSE" and kv_cache_dtype == "fp8":
|
||||
kv_cache_dtype = "fp8_ds_mla"
|
||||
|
||||
# Compute max total_q across all configs so the metadata builder buffer
|
||||
# and scheduler config are large enough for all batch specs.
|
||||
max_total_q = max(
|
||||
sum(r.q_len for r in parse_batch_spec(cfg.batch_spec))
|
||||
for cfg, *_ in configs_with_params
|
||||
)
|
||||
|
||||
# Create and set vLLM config for MLA (reused across all benchmarks)
|
||||
vllm_config = create_minimal_vllm_config(
|
||||
model_name="deepseek-v3", # Used only for model path
|
||||
block_size=block_size,
|
||||
max_num_batched_tokens=max_total_q,
|
||||
mla_dims=mla_dims, # Use custom dims from config or default
|
||||
index_topk=index_topk if is_sparse else None,
|
||||
prefill_backend=prefill_backend,
|
||||
kv_cache_dtype=kv_cache_dtype,
|
||||
)
|
||||
|
||||
results = []
|
||||
@@ -883,7 +979,9 @@ def _run_mla_benchmark_batched(
|
||||
mla_dims,
|
||||
vllm_config,
|
||||
device,
|
||||
max_num_tokens=max_total_q,
|
||||
index_topk=index_topk if is_sparse else None,
|
||||
kv_cache_dtype=kv_cache_dtype,
|
||||
)
|
||||
|
||||
# Verify the actual prefill backend matches what was requested
|
||||
@@ -942,6 +1040,7 @@ def _run_mla_benchmark_batched(
|
||||
mla_dims,
|
||||
device,
|
||||
indexer=indexer,
|
||||
kv_cache_dtype=kv_cache_dtype,
|
||||
)
|
||||
results.append(result)
|
||||
|
||||
|
||||
@@ -140,7 +140,7 @@ def _create_vllm_config(
|
||||
|
||||
cache_config = CacheConfig(
|
||||
block_size=config.block_size,
|
||||
cache_dtype="auto",
|
||||
cache_dtype=config.kv_cache_dtype,
|
||||
)
|
||||
cache_config.num_gpu_blocks = max_num_blocks
|
||||
cache_config.num_cpu_blocks = 0
|
||||
@@ -215,7 +215,7 @@ def _create_backend_impl(
|
||||
num_kv_heads=config.num_kv_heads,
|
||||
alibi_slopes=None,
|
||||
sliding_window=None,
|
||||
kv_cache_dtype="auto",
|
||||
kv_cache_dtype=config.kv_cache_dtype,
|
||||
)
|
||||
|
||||
kv_cache_spec = FullAttentionSpec(
|
||||
@@ -288,12 +288,22 @@ def _create_input_tensors(
|
||||
total_q: int,
|
||||
device: torch.device,
|
||||
dtype: torch.dtype,
|
||||
quantize_query: bool = False,
|
||||
) -> tuple:
|
||||
"""Create Q, K, V input tensors for all layers."""
|
||||
"""Create Q, K, V input tensors for all layers.
|
||||
|
||||
When quantize_query is True, queries are cast to fp8 to match backends
|
||||
that require query/key/value dtype consistency.
|
||||
"""
|
||||
q_dtype = dtype
|
||||
if quantize_query:
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
q_dtype = current_platform.fp8_dtype()
|
||||
q_list = [
|
||||
torch.randn(
|
||||
total_q, config.num_q_heads, config.head_dim, device=device, dtype=dtype
|
||||
)
|
||||
).to(q_dtype)
|
||||
for _ in range(config.num_layers)
|
||||
]
|
||||
k_list = [
|
||||
@@ -344,10 +354,17 @@ def _create_kv_cache(
|
||||
# Compute inverse permutation to get back to logical view
|
||||
inv_order = [stride_order.index(i) for i in range(len(stride_order))]
|
||||
|
||||
# Use fp8 dtype for cache when requested.
|
||||
cache_dtype = dtype
|
||||
if config.kv_cache_dtype == "fp8":
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
cache_dtype = current_platform.fp8_dtype()
|
||||
|
||||
cache_list = []
|
||||
for _ in range(config.num_layers):
|
||||
# Allocate in physical layout order (contiguous in memory)
|
||||
cache = torch.zeros(*physical_shape, device=device, dtype=dtype)
|
||||
cache = torch.zeros(*physical_shape, device=device, dtype=cache_dtype)
|
||||
# Permute to logical view
|
||||
cache = cache.permute(*inv_order)
|
||||
cache_list.append(cache)
|
||||
@@ -392,6 +409,37 @@ def _run_single_benchmark(
|
||||
)
|
||||
torch.accelerator.synchronize()
|
||||
|
||||
# Optionally capture a CUDA graph after warmup.
|
||||
# Graph replay eliminates CPU launch overhead so timings reflect pure
|
||||
# kernel time.
|
||||
if config.use_cuda_graphs:
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(graph):
|
||||
for i in range(config.num_layers):
|
||||
impl.forward(
|
||||
layer,
|
||||
q_list[i],
|
||||
k_list[i],
|
||||
v_list[i],
|
||||
cache_list[i],
|
||||
attn_metadata,
|
||||
output=out,
|
||||
)
|
||||
benchmark_fn = graph.replay
|
||||
else:
|
||||
|
||||
def benchmark_fn():
|
||||
for i in range(config.num_layers):
|
||||
impl.forward(
|
||||
layer,
|
||||
q_list[i],
|
||||
k_list[i],
|
||||
v_list[i],
|
||||
cache_list[i],
|
||||
attn_metadata,
|
||||
output=out,
|
||||
)
|
||||
|
||||
# Benchmark
|
||||
times = []
|
||||
for _ in range(config.repeats):
|
||||
@@ -399,16 +447,7 @@ def _run_single_benchmark(
|
||||
end = torch.cuda.Event(enable_timing=True)
|
||||
|
||||
start.record()
|
||||
for i in range(config.num_layers):
|
||||
impl.forward(
|
||||
layer,
|
||||
q_list[i],
|
||||
k_list[i],
|
||||
v_list[i],
|
||||
cache_list[i],
|
||||
attn_metadata,
|
||||
output=out,
|
||||
)
|
||||
benchmark_fn()
|
||||
end.record()
|
||||
|
||||
torch.accelerator.synchronize()
|
||||
@@ -418,8 +457,8 @@ def _run_single_benchmark(
|
||||
mem_stats = {}
|
||||
if config.profile_memory:
|
||||
mem_stats = {
|
||||
"allocated_mb": torch.cuda.memory_allocated(device) / 1024**2,
|
||||
"reserved_mb": torch.cuda.memory_reserved(device) / 1024**2,
|
||||
"allocated_mb": torch.accelerator.memory_allocated(device) / 1024**2,
|
||||
"reserved_mb": torch.accelerator.memory_reserved(device) / 1024**2,
|
||||
}
|
||||
|
||||
return times, mem_stats
|
||||
@@ -502,8 +541,12 @@ def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:
|
||||
common_attn_metadata=common_metadata,
|
||||
)
|
||||
|
||||
# Only quantize queries when the impl supports it
|
||||
quantize_query = config.kv_cache_dtype.startswith("fp8") and getattr(
|
||||
impl, "supports_quant_query_input", False
|
||||
)
|
||||
q_list, k_list, v_list = _create_input_tensors(
|
||||
config, total_q, device, dtype
|
||||
config, total_q, device, dtype, quantize_query=quantize_query
|
||||
)
|
||||
|
||||
cache_list = _create_kv_cache(
|
||||
|
||||
@@ -95,13 +95,16 @@ def create_logits(
|
||||
def measure_memory() -> tuple[int, int]:
|
||||
"""Return (allocated, reserved) memory in bytes."""
|
||||
torch.accelerator.synchronize()
|
||||
return torch.cuda.memory_allocated(), torch.cuda.max_memory_allocated()
|
||||
return (
|
||||
torch.accelerator.memory_allocated(),
|
||||
torch.accelerator.max_memory_allocated(),
|
||||
)
|
||||
|
||||
|
||||
def reset_memory_stats():
|
||||
"""Reset peak memory statistics."""
|
||||
reset_buffer_cache()
|
||||
torch.cuda.reset_peak_memory_stats()
|
||||
torch.accelerator.reset_peak_memory_stats()
|
||||
torch.accelerator.empty_cache()
|
||||
gc.collect()
|
||||
|
||||
|
||||
@@ -196,6 +196,7 @@ __forceinline__ __device__ u32x8_t ld256_cs(const u32x8_t* addr) {
|
||||
return val;
|
||||
#else
|
||||
assert(false && "ld256_cs requires SM100+ with CUDA 12.9+");
|
||||
return u32x8_t{};
|
||||
#endif
|
||||
}
|
||||
|
||||
|
||||
@@ -109,16 +109,18 @@ void create_and_map(unsigned long long device, ssize_t size, CUdeviceptr d_mem,
|
||||
|
||||
#ifndef USE_ROCM
|
||||
int flag = 0;
|
||||
CUDA_CHECK(cuDeviceGetAttribute(
|
||||
CUresult rdma_result = cuDeviceGetAttribute(
|
||||
&flag, CU_DEVICE_ATTRIBUTE_GPU_DIRECT_RDMA_WITH_CUDA_VMM_SUPPORTED,
|
||||
device));
|
||||
if (flag) { // support GPUDirect RDMA if possible
|
||||
device);
|
||||
if (rdma_result == CUDA_SUCCESS &&
|
||||
flag) { // support GPUDirect RDMA if possible
|
||||
prop.allocFlags.gpuDirectRDMACapable = 1;
|
||||
}
|
||||
int fab_flag = 0;
|
||||
CUDA_CHECK(cuDeviceGetAttribute(
|
||||
&fab_flag, CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED, device));
|
||||
if (fab_flag) { // support fabric handle if possible
|
||||
CUresult fab_result = cuDeviceGetAttribute(
|
||||
&fab_flag, CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED, device);
|
||||
if (fab_result == CUDA_SUCCESS &&
|
||||
fab_flag) { // support fabric handle if possible
|
||||
prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_FABRIC;
|
||||
}
|
||||
#endif
|
||||
|
||||
@@ -73,10 +73,9 @@ void moe_permute(
|
||||
MOE_DISPATCH(input.scalar_type(), [&] {
|
||||
expandInputRowsKernelLauncher<scalar_t>(
|
||||
get_ptr<scalar_t>(input), get_ptr<scalar_t>(permuted_input),
|
||||
get_ptr<int>(permuted_experts_id), get_ptr<int>(sorted_row_idx),
|
||||
get_ptr<int>(inv_permuted_idx), get_ptr<int>(permuted_idx),
|
||||
get_ptr<int64_t>(expert_first_token_offset), n_token, valid_num_ptr,
|
||||
n_hidden, topk, n_local_expert, stream);
|
||||
get_ptr<int>(sorted_row_idx), get_ptr<int>(inv_permuted_idx),
|
||||
get_ptr<int>(permuted_idx), get_ptr<int64_t>(expert_first_token_offset),
|
||||
n_token, valid_num_ptr, n_hidden, topk, n_local_expert, stream);
|
||||
});
|
||||
}
|
||||
|
||||
|
||||
@@ -57,7 +57,7 @@ void sortAndScanExpert(const int* expert_for_source_row, const int* source_rows,
|
||||
|
||||
template <typename T>
|
||||
void expandInputRowsKernelLauncher(
|
||||
T const* unpermuted_input, T* permuted_output, int* sorted_experts,
|
||||
T const* unpermuted_input, T* permuted_output,
|
||||
int const* expanded_dest_row_to_expanded_source_row,
|
||||
int* expanded_source_row_to_expanded_dest_row, int* permuted_idx,
|
||||
int64_t const* expert_first_token_offset, int64_t const num_rows,
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
|
||||
template <typename T, bool CHECK_SKIPPED>
|
||||
__global__ void expandInputRowsKernel(
|
||||
T const* unpermuted_input, T* permuted_output, int* sorted_experts,
|
||||
T const* unpermuted_input, T* permuted_output,
|
||||
int const* expanded_dest_row_to_expanded_source_row,
|
||||
int* expanded_source_row_to_expanded_dest_row, int* permuted_idx,
|
||||
int64_t const* expert_first_token_offset, int64_t const num_rows,
|
||||
@@ -16,7 +16,6 @@ __global__ void expandInputRowsKernel(
|
||||
int64_t expanded_dest_row = blockIdx.x;
|
||||
int64_t const expanded_source_row =
|
||||
expanded_dest_row_to_expanded_source_row[expanded_dest_row];
|
||||
int expert_id = sorted_experts[expanded_dest_row];
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
assert(expanded_dest_row <= INT32_MAX);
|
||||
@@ -54,7 +53,7 @@ __global__ void expandInputRowsKernel(
|
||||
|
||||
template <typename T>
|
||||
void expandInputRowsKernelLauncher(
|
||||
T const* unpermuted_input, T* permuted_output, int* sorted_experts,
|
||||
T const* unpermuted_input, T* permuted_output,
|
||||
int const* expanded_dest_row_to_expanded_source_row,
|
||||
int* expanded_source_row_to_expanded_dest_row, int* permuted_idx,
|
||||
int64_t const* expert_first_token_offset, int64_t const num_rows,
|
||||
@@ -70,12 +69,12 @@ void expandInputRowsKernelLauncher(
|
||||
bool is_check_skip = num_valid_tokens_ptr != nullptr;
|
||||
auto func = func_map[is_check_skip];
|
||||
|
||||
func<<<blocks, threads, 0, stream>>>(
|
||||
unpermuted_input, permuted_output, sorted_experts,
|
||||
expanded_dest_row_to_expanded_source_row,
|
||||
expanded_source_row_to_expanded_dest_row, permuted_idx,
|
||||
expert_first_token_offset, num_rows, num_valid_tokens_ptr, cols, k,
|
||||
num_local_experts);
|
||||
func<<<blocks, threads, 0, stream>>>(unpermuted_input, permuted_output,
|
||||
expanded_dest_row_to_expanded_source_row,
|
||||
expanded_source_row_to_expanded_dest_row,
|
||||
permuted_idx, expert_first_token_offset,
|
||||
num_rows, num_valid_tokens_ptr, cols, k,
|
||||
num_local_experts);
|
||||
}
|
||||
|
||||
template <class T, class U>
|
||||
|
||||
+8
-4
@@ -295,10 +295,14 @@ void cutlass_scaled_sparse_mm(torch::Tensor& out, torch::Tensor const& a,
|
||||
|
||||
std::vector<torch::Tensor> cutlass_sparse_compress(torch::Tensor const& a);
|
||||
|
||||
void scaled_fp4_quant(torch::Tensor& output, torch::Tensor const& input,
|
||||
torch::Tensor& output_scale,
|
||||
torch::Tensor const& input_scale,
|
||||
bool is_sf_swizzled_layout);
|
||||
std::tuple<torch::Tensor, torch::Tensor> scaled_fp4_quant_func(
|
||||
torch::Tensor const& input, torch::Tensor const& input_scale,
|
||||
bool is_sf_swizzled_layout);
|
||||
|
||||
void scaled_fp4_quant_out(torch::Tensor const& input,
|
||||
torch::Tensor const& input_scale,
|
||||
bool is_sf_swizzled_layout, torch::Tensor& output,
|
||||
torch::Tensor& output_scale);
|
||||
|
||||
void scaled_fp4_experts_quant(
|
||||
torch::Tensor& output, torch::Tensor& output_scale,
|
||||
|
||||
@@ -16,6 +16,8 @@
|
||||
|
||||
#include <torch/all.h>
|
||||
|
||||
#include "nvfp4_utils.cuh"
|
||||
|
||||
#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
|
||||
(defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
|
||||
void scaled_fp4_quant_sm1xxa(torch::Tensor const& output,
|
||||
@@ -51,9 +53,10 @@ void silu_and_mul_scaled_fp4_experts_quant_sm1xxa(
|
||||
torch::Tensor const& output_scale_offset_by_experts);
|
||||
#endif
|
||||
|
||||
void scaled_fp4_quant(torch::Tensor& output, torch::Tensor const& input,
|
||||
torch::Tensor& output_sf, torch::Tensor const& input_sf,
|
||||
bool is_sf_swizzled_layout) {
|
||||
void scaled_fp4_quant_out(torch::Tensor const& input,
|
||||
torch::Tensor const& input_sf,
|
||||
bool is_sf_swizzled_layout, torch::Tensor& output,
|
||||
torch::Tensor& output_sf) {
|
||||
#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
|
||||
(defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
|
||||
return scaled_fp4_quant_sm1xxa(output, input, output_sf, input_sf,
|
||||
@@ -62,6 +65,34 @@ void scaled_fp4_quant(torch::Tensor& output, torch::Tensor const& input,
|
||||
TORCH_CHECK_NOT_IMPLEMENTED(false, "No compiled nvfp4 quantization kernel");
|
||||
}
|
||||
|
||||
std::tuple<torch::Tensor, torch::Tensor> scaled_fp4_quant_func(
|
||||
torch::Tensor const& input, torch::Tensor const& input_sf,
|
||||
bool is_sf_swizzled_layout) {
|
||||
int64_t n = input.size(-1);
|
||||
int64_t m = input.numel() / n;
|
||||
auto device = input.device();
|
||||
|
||||
// Two fp4 values packed into a uint8
|
||||
auto output = torch::empty(
|
||||
{m, n / 2}, torch::TensorOptions().device(device).dtype(torch::kUInt8));
|
||||
|
||||
torch::Tensor output_sf;
|
||||
if (is_sf_swizzled_layout) {
|
||||
auto [sf_m, sf_n] = vllm::computeSwizzledSFShape(m, n);
|
||||
output_sf = torch::empty(
|
||||
{sf_m, sf_n},
|
||||
torch::TensorOptions().device(device).dtype(torch::kInt32));
|
||||
} else {
|
||||
output_sf = torch::empty(
|
||||
{m, n / CVT_FP4_SF_VEC_SIZE},
|
||||
torch::TensorOptions().device(device).dtype(torch::kUInt8));
|
||||
}
|
||||
|
||||
scaled_fp4_quant_out(input, input_sf, is_sf_swizzled_layout, output,
|
||||
output_sf);
|
||||
return {output, output_sf};
|
||||
}
|
||||
|
||||
void scaled_fp4_experts_quant(
|
||||
torch::Tensor& output, torch::Tensor& output_scale,
|
||||
torch::Tensor const& input, torch::Tensor const& input_global_scale,
|
||||
|
||||
@@ -18,6 +18,7 @@
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
#include <cuda_fp8.h>
|
||||
#include <utility>
|
||||
|
||||
#include "../../cuda_vec_utils.cuh"
|
||||
|
||||
@@ -54,6 +55,18 @@ inline int computeEffectiveRows(int m) {
|
||||
return round_up(m, ROW_TILE);
|
||||
}
|
||||
|
||||
// Compute the shape of the swizzled SF output tensor.
|
||||
// Returns (rounded_m, rounded_n / 4) where:
|
||||
// rounded_m = round_up(m, 128)
|
||||
// rounded_n = round_up(n / CVT_FP4_SF_VEC_SIZE, 4)
|
||||
inline std::pair<int64_t, int64_t> computeSwizzledSFShape(int64_t m,
|
||||
int64_t n) {
|
||||
int64_t rounded_m = round_up(m, static_cast<int64_t>(128));
|
||||
int64_t scale_n = n / CVT_FP4_SF_VEC_SIZE;
|
||||
int64_t rounded_n = round_up(scale_n, static_cast<int64_t>(4));
|
||||
return {rounded_m, rounded_n / 4};
|
||||
}
|
||||
|
||||
// Convert 8 float32 values into 8 e2m1 values (represented as one uint32_t).
|
||||
inline __device__ uint32_t fp32_vec8_to_e2m1(float (&array)[8]) {
|
||||
uint32_t val;
|
||||
|
||||
@@ -286,6 +286,15 @@ void rms_norm_per_block_quant(torch::Tensor& out, torch::Tensor const& input,
|
||||
"Outer scale stride must be 1 when scales are not transposed");
|
||||
}
|
||||
|
||||
int64_t hidden_size = input.size(-1);
|
||||
TORCH_CHECK(hidden_size > 0 && hidden_size % group_size == 0,
|
||||
"hidden_size must be a positive multiple of group_size");
|
||||
int64_t num_tokens = input.numel() / hidden_size;
|
||||
int64_t num_groups = hidden_size / group_size;
|
||||
TORCH_CHECK(scales.numel() >= num_tokens * num_groups,
|
||||
"scales buffer too small: need ", num_tokens * num_groups,
|
||||
" elements, got ", scales.numel());
|
||||
|
||||
rms_norm_per_block_quant_dispatch(out, input, weight, scales, group_size,
|
||||
var_epsilon, scale_ub, residual,
|
||||
is_scale_transposed);
|
||||
|
||||
+1
-1
@@ -575,7 +575,7 @@ static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode(
|
||||
// The range of logits within the row.
|
||||
int rowStart = 0;
|
||||
int seq_len = seqLens[rowIdx / next_n];
|
||||
int rowEnd = seq_len - next_n + (rowIdx % next_n) + 1;
|
||||
int rowEnd = max(0, seq_len - next_n + (rowIdx % next_n) + 1);
|
||||
|
||||
// Local pointers to this block
|
||||
if constexpr (!multipleBlocksPerRow && !mergeBlocks) {
|
||||
|
||||
+15
-4
@@ -564,10 +564,21 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
|
||||
// Compute NVFP4 block quantized tensor.
|
||||
ops.def(
|
||||
"scaled_fp4_quant(Tensor! output, Tensor input,"
|
||||
" Tensor! output_scale, Tensor input_scale, bool "
|
||||
"is_sf_swizzled_layout) -> ()");
|
||||
ops.impl("scaled_fp4_quant", torch::kCUDA, &scaled_fp4_quant);
|
||||
"scaled_fp4_quant(Tensor input,"
|
||||
" Tensor input_scale, bool "
|
||||
"is_sf_swizzled_layout) -> (Tensor, Tensor)");
|
||||
ops.impl("scaled_fp4_quant", torch::kCUDA, &scaled_fp4_quant_func);
|
||||
|
||||
// Out variant
|
||||
// TODO: Add {at::Tag::out_variant} tag and update all call sites
|
||||
// to use the functional variant once vLLM upgrades PyTorch.
|
||||
// See pytorch/pytorch#176117.
|
||||
ops.def(
|
||||
"scaled_fp4_quant.out(Tensor input,"
|
||||
" Tensor input_scale, bool "
|
||||
"is_sf_swizzled_layout, *, Tensor(a!) output, Tensor(b!) output_scale) "
|
||||
"-> ()");
|
||||
ops.impl("scaled_fp4_quant.out", torch::kCUDA, &scaled_fp4_quant_out);
|
||||
|
||||
// Compute NVFP4 experts quantization.
|
||||
ops.def(
|
||||
|
||||
+2
-2
@@ -620,7 +620,7 @@ RUN set -eux; \
|
||||
ARG BITSANDBYTES_VERSION_X86=0.46.1
|
||||
ARG BITSANDBYTES_VERSION_ARM64=0.42.0
|
||||
ARG TIMM_VERSION=">=1.0.17"
|
||||
ARG RUNAI_MODEL_STREAMER_VERSION=">=0.15.3"
|
||||
ARG RUNAI_MODEL_STREAMER_VERSION=">=0.15.7"
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
if [ "$TARGETPLATFORM" = "linux/arm64" ]; then \
|
||||
BITSANDBYTES_VERSION="${BITSANDBYTES_VERSION_ARM64}"; \
|
||||
@@ -628,7 +628,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
BITSANDBYTES_VERSION="${BITSANDBYTES_VERSION_X86}"; \
|
||||
fi; \
|
||||
uv pip install --system accelerate hf_transfer modelscope \
|
||||
"bitsandbytes>=${BITSANDBYTES_VERSION}" "timm${TIMM_VERSION}" "runai-model-streamer[s3,gcs]${RUNAI_MODEL_STREAMER_VERSION}"
|
||||
"bitsandbytes>=${BITSANDBYTES_VERSION}" "timm${TIMM_VERSION}" "runai-model-streamer[s3,gcs,azure]${RUNAI_MODEL_STREAMER_VERSION}"
|
||||
|
||||
# ============================================================
|
||||
# VLLM INSTALLATION (depends on build stage)
|
||||
|
||||
@@ -9,6 +9,7 @@
|
||||
#
|
||||
# Build targets:
|
||||
# vllm-openai (default): used for serving deployment
|
||||
# vllm-openai-zen: vLLM from source + zentorch from PyPI via vllm[zen]
|
||||
# vllm-test: used for CI tests
|
||||
# vllm-dev: used for development
|
||||
#
|
||||
@@ -222,3 +223,19 @@ LABEL ai.vllm.build.cpu-arm-bf16="${VLLM_CPU_ARM_BF16:-false}"
|
||||
LABEL ai.vllm.build.python-version="${PYTHON_VERSION:-3.12}"
|
||||
|
||||
ENTRYPOINT ["vllm", "serve"]
|
||||
|
||||
|
||||
######################### ZEN CPU PYPI IMAGE #########################
|
||||
FROM vllm-openai AS vllm-openai-zen
|
||||
|
||||
ARG TARGETARCH
|
||||
|
||||
RUN if [ "$TARGETARCH" != "amd64" ]; then \
|
||||
echo "ERROR: vllm-openai-amd only supports --platform=linux/amd64"; \
|
||||
exit 1; \
|
||||
fi
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install "vllm[zen]"
|
||||
|
||||
ENTRYPOINT ["vllm", "serve"]
|
||||
|
||||
@@ -83,7 +83,7 @@
|
||||
"default": ">=1.0.17"
|
||||
},
|
||||
"RUNAI_MODEL_STREAMER_VERSION": {
|
||||
"default": ">=0.15.3"
|
||||
"default": ">=0.15.7"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -75,7 +75,7 @@ For an optimized workflow when iterating on C++/CUDA kernels, see the [Increment
|
||||
vLLM uses `pre-commit` to lint and format the codebase. See <https://pre-commit.com/#usage> if `pre-commit` is new to you. Setting up `pre-commit` is as easy as:
|
||||
|
||||
```bash
|
||||
uv pip install pre-commit
|
||||
uv pip install pre-commit>=4.5.1
|
||||
pre-commit install
|
||||
```
|
||||
|
||||
|
||||
@@ -164,18 +164,18 @@ Priority is **1 = highest** (tried first).
|
||||
| Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | MM Prefix | DCP | Attention Types | Compute Cap. |
|
||||
| ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | --------- | --- | --------------- | ------------ |
|
||||
| `CPU_ATTN` | | fp16, bf16, fp32 | `auto` | Any | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | All | N/A |
|
||||
| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
|
||||
| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ✅ | Decoder | 10.x |
|
||||
| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥8.0 |
|
||||
| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ❌ | ✅ | All | 9.x |
|
||||
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥10.0 |
|
||||
| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
|
||||
| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ✅ | Decoder | 10.x |
|
||||
| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥8.0 |
|
||||
| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ❌ | ✅ | All | 9.x |
|
||||
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥10.0 |
|
||||
| `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ✅ | Decoder | Any |
|
||||
| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `bfloat16` | Any | Any | ❌ | ✅ | ❌ | Decoder, Encoder Only | Any |
|
||||
| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ❌ | ❌ | ❌ | Decoder, Enc-Dec | N/A |
|
||||
| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | Any | Any | ❌ | ✅ | ❌ | Decoder, Encoder Only | Any |
|
||||
| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ❌ | ❌ | ❌ | Decoder, Enc-Dec | N/A |
|
||||
| `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto` | %16 | Any | ✅ | ✅ | ❌ | All | N/A |
|
||||
| `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ✅ | ✅ | ❌ | All | N/A |
|
||||
| `TREE_ATTN` | | fp16, bf16 | `auto` | %16 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | Decoder | Any |
|
||||
| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | All | Any |
|
||||
| `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ✅ | ✅ | ❌ | All | N/A |
|
||||
| `TREE_ATTN` | | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | Decoder | Any |
|
||||
| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | All | Any |
|
||||
|
||||
> **†** FlashInfer uses TRTLLM attention on Blackwell (SM100), which supports sinks. Disable via `--attention-config.use_trtllm_attention=0`.
|
||||
>
|
||||
@@ -204,14 +204,14 @@ configuration.
|
||||
|
||||
| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. |
|
||||
| ------- | ------ | --------- | ----------- | ---------- | ---- | ------ | --------- | --- | --------------- | ------------ |
|
||||
| `CUTLASS_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
|
||||
| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
|
||||
| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x |
|
||||
| `FLASHMLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |
|
||||
| `CUTLASS_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
|
||||
| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
|
||||
| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x |
|
||||
| `FLASHMLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |
|
||||
| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
|
||||
| `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x |
|
||||
| `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 1 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
|
||||
| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `bfloat16` | 1 | Any | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
|
||||
| `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x |
|
||||
| `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 1 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
|
||||
| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | 1 | Any | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
|
||||
| `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
|
||||
| `TRITON_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
|
||||
| `XPU_MLA_SPARSE` | fp16, bf16 | `auto`, `bfloat16` | Any | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | Any |
|
||||
| `TRITON_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
|
||||
| `XPU_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | Any | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | Any |
|
||||
|
||||
@@ -35,7 +35,8 @@ th {
|
||||
| naive | standard | all<sup>1</sup> | G,A,T | N | <sup>6</sup> | [layer.py][vllm.model_executor.layers.fused_moe.layer.FusedMoE] |
|
||||
| deepep_high_throughput | standard | fp8 | G(128),A,T<sup>2</sup> | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize.DeepEPHTPrepareAndFinalize] |
|
||||
| deepep_low_latency | batched | fp8 | G(128),A,T<sup>3</sup> | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ll_prepare_finalize.DeepEPLLPrepareAndFinalize] |
|
||||
| flashinfer_all2allv | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferA2APrepareAndFinalize`][vllm.model_executor.layers.fused_moe.flashinfer_a2a_prepare_finalize.FlashInferA2APrepareAndFinalize] |
|
||||
| flashinfer_nvlink_two_sided | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferNVLinkTwoSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.flashinfer_nvlink_two_sided_prepare_finalize.FlashInferNVLinkTwoSidedPrepareAndFinalize] |
|
||||
| flashinfer_nvlink_one_sided | standard | nvfp4 | G,A,T | N | N | [`FlashInferNVLinkOneSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.flashinfer_nvlink_one_sided_prepare_finalize.FlashInferNVLinkOneSidedPrepareAndFinalize] |
|
||||
|
||||
!!! info "Table key"
|
||||
1. All types: mxfp4, nvfp4, int4, int8, fp8
|
||||
|
||||
@@ -34,9 +34,6 @@ relies on caching artifacts to reduce start time, we must properly propagate the
|
||||
with the LLM text-backbone, or other instances of the same artifact (as is the case with vision block). `is_encoder=True` is also needed for encoder
|
||||
components (see Compile Range Integration).
|
||||
|
||||
3. `with set_forward_context` context manager should be used around the nn.Module's forward call. This will properly forward the vllm_config which is needed
|
||||
for torch.compile integration.
|
||||
|
||||
### CompilationConfig
|
||||
|
||||
With the exception of `compile_mm_encoder: true`, the multimodal encoder will inherit from the same compilation config as the text LLM. We may extend
|
||||
|
||||
@@ -107,6 +107,27 @@ vLLM supports the `tool_choice='none'` option in the chat completion API. When t
|
||||
!!! note
|
||||
When tools are specified in the request, vLLM includes tool definitions in the prompt by default, regardless of the `tool_choice` setting. To exclude tool definitions when `tool_choice='none'`, use the `--exclude-tools-when-tool-choice-none` option.
|
||||
|
||||
## Constrained Decoding Behavior
|
||||
|
||||
Whether vLLM enforces the tool parameter schema during generation depends on the `tool_choice` mode:
|
||||
|
||||
| `tool_choice` value | Schema-constrained decoding | Behavior |
|
||||
| --- | --- | --- |
|
||||
| Named function | Yes (via structured outputs backend) | Arguments are guaranteed to be valid JSON conforming to the function's parameter schema. |
|
||||
| `"required"` | Yes (via structured outputs backend) | Same as named function. The model must produce at least one tool call. |
|
||||
| `"auto"` | No | The model generates freely. A tool-call parser extracts tool calls from the raw text. Arguments may be malformed or not match the schema. |
|
||||
| `"none"` | N/A | No tool calls are produced. |
|
||||
|
||||
When schema conformance matters, prefer `tool_choice="required"` or named function calling over `"auto"`.
|
||||
|
||||
### Strict Mode (`strict` parameter)
|
||||
|
||||
The [OpenAI API](https://platform.openai.com/docs/guides/function-calling#strict-mode) supports a `strict` field on function definitions. When set to `true`, OpenAI uses constrained decoding to guarantee that tool-call arguments match the function schema, even in `tool_choice="auto"` mode.
|
||||
|
||||
vLLM **does not implement** `strict` mode today. The `strict` field is accepted in requests (to avoid breaking clients that set it), but it has no effect on decoding behavior. In auto mode, argument validity depends entirely on the model's output quality and the parser's extraction logic.
|
||||
|
||||
Tracking issues: [#15526](https://github.com/vllm-project/vllm/issues/15526), [#16313](https://github.com/vllm-project/vllm/issues/16313).
|
||||
|
||||
## Automatic Function Calling
|
||||
|
||||
To enable this feature, you should set the following flags:
|
||||
@@ -124,6 +145,9 @@ from HuggingFace; and you can find an example of this in a `tokenizer_config.jso
|
||||
|
||||
If your favorite tool-calling model is not supported, please feel free to contribute a parser & tool use chat template!
|
||||
|
||||
!!! note
|
||||
With `tool_choice="auto"`, tool-call arguments are extracted from the model's raw text output by the selected parser. No schema-level constraint is applied during decoding, so arguments may occasionally be malformed or violate the function's parameter schema. See [Constrained Decoding Behavior](#constrained-decoding-behavior) for details.
|
||||
|
||||
### Hermes Models (`hermes`)
|
||||
|
||||
All Nous Research Hermes-series models newer than Hermes 2 Pro should be supported.
|
||||
@@ -219,7 +243,7 @@ Supported models:
|
||||
|
||||
* `ibm-granite/granite-4.0-h-small` and other Granite 4.0 models
|
||||
|
||||
Recommended flags: `--tool-call-parser hermes`
|
||||
Recommended flags: `--tool-call-parser granite4`
|
||||
|
||||
* `ibm-granite/granite-3.0-8b-instruct`
|
||||
|
||||
|
||||
@@ -16,4 +16,6 @@ vLLM supports the following hardware platforms:
|
||||
|
||||
vLLM supports third-party hardware plugins that live **outside** the main `vllm` repository. These follow the [Hardware-Pluggable RFC](../../design/plugin_system.md).
|
||||
|
||||
A list of all supported hardware can be found on the [vllm.ai website](https://vllm.ai/#compatibility). If you want to add new hardware, please contact us on [Slack](https://slack.vllm.ai/) or [Email](mailto:collaboration@vllm.ai).
|
||||
A list of all supported hardware can be found on the vLLM website, see [Universal Compatibility - Hardware](https://vllm.ai/#compatibility).
|
||||
|
||||
If you want to add new hardware, please contact us on [Slack](https://slack.vllm.ai/) or [Email](mailto:collaboration@vllm.ai).
|
||||
|
||||
@@ -31,6 +31,16 @@ vllm serve gs://core-llm/Llama-3-8b \
|
||||
--load-format runai_streamer
|
||||
```
|
||||
|
||||
To run model from Azure Blob Storage run:
|
||||
|
||||
```bash
|
||||
AZURE_STORAGE_ACCOUNT_NAME=<account> \
|
||||
vllm serve az://<container>/<model-path> \
|
||||
--load-format runai_streamer
|
||||
```
|
||||
|
||||
Authentication uses `DefaultAzureCredential`, which supports `az login`, managed identity, environment variables (`AZURE_CLIENT_ID`, `AZURE_TENANT_ID`, `AZURE_CLIENT_SECRET`), and other methods.
|
||||
|
||||
To run model from a S3 compatible object store run:
|
||||
|
||||
```bash
|
||||
|
||||
@@ -418,6 +418,7 @@ th {
|
||||
| `Grok1ForCausalLM` | Grok2 | `xai-org/grok-2` | ✅︎ | ✅︎ |
|
||||
| `HunYuanDenseV1ForCausalLM` | Hunyuan Dense | `tencent/Hunyuan-7B-Instruct` | ✅︎ | ✅︎ |
|
||||
| `HunYuanMoEV1ForCausalLM` | Hunyuan-A13B | `tencent/Hunyuan-A13B-Instruct`, `tencent/Hunyuan-A13B-Pretrain`, `tencent/Hunyuan-A13B-Instruct-FP8`, etc. | ✅︎ | ✅︎ |
|
||||
| `HyperCLOVAXForCausalLM` | HyperCLOVAX-SEED-Think-14B | `naver-hyperclovax/HyperCLOVAX-SEED-Think-14B` | ✅︎ | ✅︎ |
|
||||
| `InternLMForCausalLM` | InternLM | `internlm/internlm-7b`, `internlm/internlm-chat-7b`, etc. | ✅︎ | ✅︎ |
|
||||
| `InternLM2ForCausalLM` | InternLM2 | `internlm/internlm2-7b`, `internlm/internlm2-chat-7b`, etc. | ✅︎ | ✅︎ |
|
||||
| `InternLM3ForCausalLM` | InternLM3 | `internlm/internlm3-8b-instruct`, etc. | ✅︎ | ✅︎ |
|
||||
|
||||
@@ -21,7 +21,8 @@ vLLM provides multiple communication backends for EP. Use `--all2all-backend` to
|
||||
| `allgather_reducescatter` | Default backend | Standard all2all using allgather/reducescatter primitives | General purpose, works with any EP+DP configuration |
|
||||
| `deepep_high_throughput` | Multi-node prefill | Grouped GEMM with continuous layout, optimized for prefill | Prefill-dominated workloads, high-throughput scenarios |
|
||||
| `deepep_low_latency` | Multi-node decode | CUDA graph support, masked layout, optimized for decode | Decode-dominated workloads, low-latency scenarios |
|
||||
| `flashinfer_all2allv` | MNNVL systems | FlashInfer alltoallv kernels for multi-node NVLink | Systems with NVLink across nodes |
|
||||
| `flashinfer_nvlink_one_sided` | MNNVL systems | FlashInfer's one-sided A2A strategy for multi-node NVLink | High-throughput workloads |
|
||||
| `flashinfer_nvlink_two_sided` | MNNVL systems | FlashInfer's two-sided A2A strategy for multi-node NVLink | Systems with NVLink across nodes |
|
||||
| `naive` | Testing/debugging | Simple broadcast-based implementation | Debugging, not recommended for production |
|
||||
|
||||
## Single Node Deployment
|
||||
|
||||
@@ -72,6 +72,9 @@ In addition, we have the following custom APIs:
|
||||
- Only applicable to [classification models](../models/pooling_models.md).
|
||||
- [Score API](#score-api) (`/score`)
|
||||
- Applicable to [embedding models and cross-encoder models](../models/pooling_models.md).
|
||||
- [Cohere Embed API](#cohere-embed-api) (`/v2/embed`)
|
||||
- Compatible with [Cohere's Embed API](https://docs.cohere.com/reference/embed)
|
||||
- Works with any [embedding model](../models/pooling_models.md), including multimodal models.
|
||||
- [Re-rank API](#re-rank-api) (`/rerank`, `/v1/rerank`, `/v2/rerank`)
|
||||
- Implements [Jina AI's v1 re-rank API](https://jina.ai/reranker/)
|
||||
- Also compatible with [Cohere's v1 & v2 re-rank APIs](https://docs.cohere.com/v2/reference/rerank)
|
||||
@@ -429,6 +432,137 @@ these extra parameters are supported instead:
|
||||
--8<-- "vllm/entrypoints/pooling/base/protocol.py:embed-extra-params"
|
||||
```
|
||||
|
||||
### Cohere Embed API
|
||||
|
||||
Our API is also compatible with [Cohere's Embed v2 API](https://docs.cohere.com/reference/embed) which adds support for some modern embedding feature such as truncation, output dimensions, embedding types, and input types. This endpoint works with any embedding model (including multimodal models).
|
||||
|
||||
#### Cohere Embed API request parameters
|
||||
|
||||
| Parameter | Type | Required | Description |
|
||||
| --------- | ---- | -------- | ----------- |
|
||||
| `model` | string | Yes | Model name |
|
||||
| `input_type` | string | No | Prompt prefix key (model-dependent, see below) |
|
||||
| `texts` | list[string] | No | Text inputs (use one of `texts`, `images`, or `inputs`) |
|
||||
| `images` | list[string] | No | Base64 data URI images |
|
||||
| `inputs` | list[object] | No | Mixed text and image content objects |
|
||||
| `embedding_types` | list[string] | No | Output types (default: `["float"]`) |
|
||||
| `output_dimension` | int | No | Truncate embeddings to this dimension (Matryoshka) |
|
||||
| `truncate` | string | No | `END`, `START`, or `NONE` (default: `END`) |
|
||||
|
||||
#### Text embedding
|
||||
|
||||
```bash
|
||||
curl -X POST "http://localhost:8000/v2/embed" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "Snowflake/snowflake-arctic-embed-m-v1.5",
|
||||
"input_type": "query",
|
||||
"texts": ["Hello world", "How are you?"],
|
||||
"embedding_types": ["float"]
|
||||
}'
|
||||
```
|
||||
|
||||
??? console "Response"
|
||||
|
||||
```json
|
||||
{
|
||||
"id": "embd-...",
|
||||
"embeddings": {
|
||||
"float": [
|
||||
[0.012, -0.034, ...],
|
||||
[0.056, 0.078, ...]
|
||||
]
|
||||
},
|
||||
"texts": ["Hello world", "How are you?"],
|
||||
"meta": {
|
||||
"api_version": {"version": "2"},
|
||||
"billed_units": {"input_tokens": 12}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
#### Mixed text and image inputs
|
||||
|
||||
For multimodal models, you can embed images by passing base64 data URIs. The `inputs` field accepts a list of objects with mixed text and image content:
|
||||
|
||||
```bash
|
||||
curl -X POST "http://localhost:8000/v2/embed" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "google/siglip-so400m-patch14-384",
|
||||
"inputs": [
|
||||
{
|
||||
"content": [
|
||||
{"type": "text", "text": "A photo of a cat"},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBOR..."}}
|
||||
]
|
||||
}
|
||||
],
|
||||
"embedding_types": ["float"]
|
||||
}'
|
||||
```
|
||||
|
||||
#### Embedding types
|
||||
|
||||
The `embedding_types` parameter controls the output format. Multiple types can be requested in a single call:
|
||||
|
||||
| Type | Description |
|
||||
| ---- | ----------- |
|
||||
| `float` | Raw float32 embeddings (default) |
|
||||
| `binary` | Bit-packed signed binary |
|
||||
| `ubinary` | Bit-packed unsigned binary |
|
||||
| `base64` | Little-endian float32 encoded as base64 |
|
||||
|
||||
```bash
|
||||
curl -X POST "http://localhost:8000/v2/embed" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "Snowflake/snowflake-arctic-embed-m-v1.5",
|
||||
"input_type": "query",
|
||||
"texts": ["What is machine learning?"],
|
||||
"embedding_types": ["float", "binary"]
|
||||
}'
|
||||
```
|
||||
|
||||
??? console "Response"
|
||||
|
||||
```json
|
||||
{
|
||||
"id": "embd-...",
|
||||
"embeddings": {
|
||||
"float": [[0.012, -0.034, ...]],
|
||||
"binary": [[42, -117, ...]]
|
||||
},
|
||||
"texts": ["What is machine learning?"],
|
||||
"meta": {
|
||||
"api_version": {"version": "2"},
|
||||
"billed_units": {"input_tokens": 8}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
#### Truncation
|
||||
|
||||
The `truncate` parameter controls how inputs exceeding the model's maximum sequence length are handled:
|
||||
|
||||
| Value | Behavior |
|
||||
| ----- | --------- |
|
||||
| `END` (default) | Keep the first tokens, drop the end |
|
||||
| `START` | Keep the last tokens, drop the beginning |
|
||||
| `NONE` | Return an error if the input is too long |
|
||||
|
||||
#### Input type and prompt prefixes
|
||||
|
||||
The `input_type` field selects a prompt prefix to prepend to each text input. The available values
|
||||
depend on the model:
|
||||
|
||||
- **Models with `task_instructions` in `config.json`**: The keys from the `task_instructions` dict are
|
||||
the valid `input_type` values and the corresponding value is prepended to each text.
|
||||
- **Models with `config_sentence_transformers.json` prompts**: The keys from the `prompts` dict are
|
||||
the valid `input_type` values. For example, `Snowflake/snowflake-arctic-embed-xs` defines `"query"`,
|
||||
so setting `input_type: "query"` prepends `"Represent this sentence for searching relevant passages: "`.
|
||||
- **Other models**: `input_type` is not accepted and will raise a validation error if passed.
|
||||
|
||||
### Transcriptions API
|
||||
|
||||
Our Transcriptions API is compatible with [OpenAI's Transcriptions API](https://platform.openai.com/docs/api-reference/audio/createTranscription);
|
||||
|
||||
@@ -14,6 +14,10 @@ import regex as re
|
||||
import zmq
|
||||
from quart import Quart, make_response, request
|
||||
|
||||
from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import (
|
||||
MoRIIOConstants,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.setLevel(logging.DEBUG)
|
||||
prefill_instances: list[dict] = []
|
||||
@@ -213,6 +217,8 @@ async def handle_request():
|
||||
|
||||
dip, dport = extract_ip_port_fast(decode_instance_endpoint["request_address"])
|
||||
|
||||
transfer_id = f"{MoRIIOConstants.TRANSFER_PREFIX}-{str(uuid.uuid4())}"
|
||||
|
||||
req_data_to_prefill = copy.deepcopy(req_data)
|
||||
req_data_to_prefill["kv_transfer_params"] = {}
|
||||
req_data["kv_transfer_params"] = {}
|
||||
@@ -222,6 +228,7 @@ async def handle_request():
|
||||
req_data_to_prefill["kv_transfer_params"]["remote_tp_size"] = (
|
||||
decode_instance_endpoint["tp_size"]
|
||||
)
|
||||
req_data_to_prefill["kv_transfer_params"]["transfer_id"] = transfer_id
|
||||
|
||||
send_prefill_task = asyncio.create_task(
|
||||
send_request_to_prefill(
|
||||
@@ -267,6 +274,7 @@ async def handle_request():
|
||||
|
||||
if selected_prefill_dp_rank is not None:
|
||||
req_data["kv_transfer_params"]["remote_dp_rank"] = selected_prefill_dp_rank
|
||||
req_data["kv_transfer_params"]["transfer_id"] = transfer_id
|
||||
|
||||
decode_request_task = asyncio.create_task(
|
||||
start_decode_request(
|
||||
|
||||
@@ -24,7 +24,7 @@ outlines_core == 0.2.11
|
||||
# required for outlines backend disk cache
|
||||
diskcache == 5.6.3
|
||||
lark == 1.2.2
|
||||
xgrammar == 0.1.29; platform_machine == "x86_64" or platform_machine == "aarch64" or platform_machine == "arm64" or platform_machine == "s390x" or platform_machine == "ppc64le"
|
||||
xgrammar >= 0.1.32, < 1.0.0; platform_machine == "x86_64" or platform_machine == "aarch64" or platform_machine == "arm64" or platform_machine == "s390x" or platform_machine == "ppc64le"
|
||||
typing_extensions >= 4.10
|
||||
filelock >= 3.16.1 # need to contain https://github.com/tox-dev/filelock/pull/317
|
||||
partial-json-parser # used for parsing partial JSON outputs
|
||||
|
||||
@@ -1,2 +1,2 @@
|
||||
# formatting
|
||||
pre-commit==4.0.1
|
||||
pre-commit>=4.5.1
|
||||
|
||||
@@ -42,7 +42,7 @@ tritonclient>=2.51.0
|
||||
|
||||
numba == 0.61.2 # Required for N-gram speculative decoding
|
||||
numpy
|
||||
runai-model-streamer[s3,gcs]==0.15.3
|
||||
runai-model-streamer[s3,gcs,azure]==0.15.7
|
||||
fastsafetensors>=0.2.2
|
||||
instanttensor>=0.1.5
|
||||
pydantic>=2.12 # 2.11 leads to error on python 3.13
|
||||
|
||||
@@ -45,10 +45,12 @@ pystemmer==3.0.0
|
||||
# via mteb
|
||||
|
||||
# Multi-modal processing
|
||||
av==16.1.0
|
||||
# required for audio_in_video tests
|
||||
blobfile==3.0.0
|
||||
# Multi-Modal Models Test
|
||||
decord==0.6.0
|
||||
# video processing, required by entrypoints/openai/test_video.py
|
||||
# video processing, required by entrypoints/openai/chat_completion/test_video.py
|
||||
rapidfuzz==3.12.1
|
||||
|
||||
# OpenAI compatibility and testing
|
||||
|
||||
@@ -15,7 +15,7 @@ tensorizer==2.10.1
|
||||
packaging>=24.2
|
||||
setuptools>=77.0.3,<80.0.0
|
||||
setuptools-scm>=8
|
||||
runai-model-streamer[s3,gcs]==0.15.3
|
||||
runai-model-streamer[s3,gcs,azure]==0.15.7
|
||||
conch-triton-kernels==1.2.1
|
||||
timm>=1.0.17
|
||||
# amd-quark: required for Quark quantization on ROCm
|
||||
|
||||
@@ -10,6 +10,7 @@ pytest-cov
|
||||
|
||||
# testing utils
|
||||
albumentations # required for Nemotron Parse in test_common.py
|
||||
av # required for audio_in_video tests
|
||||
backoff # required for phi4mm test
|
||||
blobfile # required for kimi-vl test
|
||||
einops # required for MPT, qwen-vl
|
||||
@@ -55,7 +56,7 @@ grpcio-reflection==1.78.0
|
||||
arctic-inference == 0.1.1 # Required for suffix decoding test
|
||||
numba == 0.61.2 # Required for N-gram speculative decoding
|
||||
numpy
|
||||
runai-model-streamer[s3,gcs]==0.15.3
|
||||
runai-model-streamer[s3,gcs,azure]==0.15.7
|
||||
fastsafetensors>=0.2.2 # 0.2.2 contains important fixes for multi-GPU mem usage
|
||||
instanttensor>=0.1.5
|
||||
pydantic>=2.12 # 2.11 leads to error on python 3.13
|
||||
|
||||
+40
-5
@@ -62,6 +62,16 @@ attrs==24.2.0
|
||||
# referencing
|
||||
audioread==3.0.1
|
||||
# via librosa
|
||||
av==16.1.0
|
||||
# via -r requirements/test.in
|
||||
azure-core==1.38.2
|
||||
# via
|
||||
# azure-identity
|
||||
# azure-storage-blob
|
||||
azure-identity==1.25.2
|
||||
# via runai-model-streamer-azure
|
||||
azure-storage-blob==12.28.0
|
||||
# via runai-model-streamer-azure
|
||||
backoff==2.2.1
|
||||
# via
|
||||
# -r requirements/test.in
|
||||
@@ -101,8 +111,10 @@ certifi==2024.8.30
|
||||
# rasterio
|
||||
# requests
|
||||
# sentry-sdk
|
||||
cffi==1.17.1
|
||||
# via soundfile
|
||||
cffi==2.0.0
|
||||
# via
|
||||
# cryptography
|
||||
# soundfile
|
||||
chardet==5.2.0
|
||||
# via mbstrdecoder
|
||||
charset-normalizer==3.4.0
|
||||
@@ -146,6 +158,12 @@ coverage==7.10.6
|
||||
# via pytest-cov
|
||||
cramjam==2.9.0
|
||||
# via fastparquet
|
||||
cryptography==46.0.5
|
||||
# via
|
||||
# azure-identity
|
||||
# azure-storage-blob
|
||||
# msal
|
||||
# pyjwt
|
||||
cuda-bindings==12.9.4
|
||||
# via torch
|
||||
cuda-pathfinder==1.3.3
|
||||
@@ -377,6 +395,8 @@ iniconfig==2.0.0
|
||||
# via pytest
|
||||
instanttensor==0.1.5
|
||||
# via -r requirements/test.in
|
||||
isodate==0.7.2
|
||||
# via azure-storage-blob
|
||||
isoduration==20.11.0
|
||||
# via jsonschema
|
||||
isort==5.13.2
|
||||
@@ -490,6 +510,12 @@ more-itertools==10.5.0
|
||||
# via lm-eval
|
||||
mpmath==1.3.0
|
||||
# via sympy
|
||||
msal==1.34.0
|
||||
# via
|
||||
# azure-identity
|
||||
# msal-extensions
|
||||
msal-extensions==1.3.1
|
||||
# via azure-identity
|
||||
msgpack==1.1.0
|
||||
# via
|
||||
# librosa
|
||||
@@ -826,6 +852,8 @@ pydantic-extra-types==2.10.5
|
||||
# via mistral-common
|
||||
pygments==2.18.0
|
||||
# via rich
|
||||
pyjwt==2.11.0
|
||||
# via msal
|
||||
pyogrio==0.11.0
|
||||
# via geopandas
|
||||
pyparsing==3.2.0
|
||||
@@ -943,6 +971,7 @@ regex==2024.9.11
|
||||
# transformers
|
||||
requests==2.32.3
|
||||
# via
|
||||
# azure-core
|
||||
# buildkite-test-collector
|
||||
# datasets
|
||||
# diffusers
|
||||
@@ -955,6 +984,7 @@ requests==2.32.3
|
||||
# lightly
|
||||
# lm-eval
|
||||
# mistral-common
|
||||
# msal
|
||||
# mteb
|
||||
# pooch
|
||||
# ray
|
||||
@@ -991,11 +1021,13 @@ rsa==4.9.1
|
||||
# via google-auth
|
||||
rtree==1.4.0
|
||||
# via torchgeo
|
||||
runai-model-streamer==0.15.3
|
||||
runai-model-streamer==0.15.7
|
||||
# via -r requirements/test.in
|
||||
runai-model-streamer-gcs==0.15.3
|
||||
runai-model-streamer-azure==0.15.7
|
||||
# via runai-model-streamer
|
||||
runai-model-streamer-s3==0.15.3
|
||||
runai-model-streamer-gcs==0.15.7
|
||||
# via runai-model-streamer
|
||||
runai-model-streamer-s3==0.15.7
|
||||
# via runai-model-streamer
|
||||
s3transfer==0.10.3
|
||||
# via boto3
|
||||
@@ -1264,6 +1296,9 @@ typing-extensions==4.15.0
|
||||
# aiosignal
|
||||
# albumentations
|
||||
# alembic
|
||||
# azure-core
|
||||
# azure-identity
|
||||
# azure-storage-blob
|
||||
# chz
|
||||
# fastapi
|
||||
# grpcio
|
||||
|
||||
@@ -657,13 +657,18 @@ class precompiled_wheel_utils:
|
||||
def get_base_commit_in_main_branch() -> str:
|
||||
try:
|
||||
# Get the latest commit hash of the upstream main branch.
|
||||
resp_json = subprocess.check_output(
|
||||
[
|
||||
"curl",
|
||||
"-s",
|
||||
"https://api.github.com/repos/vllm-project/vllm/commits/main",
|
||||
curl_cmd = [
|
||||
"curl",
|
||||
"-s",
|
||||
"https://api.github.com/repos/vllm-project/vllm/commits/main",
|
||||
]
|
||||
github_token = os.getenv("GH_TOKEN", os.getenv("GITHUB_TOKEN"))
|
||||
if github_token:
|
||||
curl_cmd += [
|
||||
"-H",
|
||||
f"Authorization: token {github_token}",
|
||||
]
|
||||
).decode("utf-8")
|
||||
resp_json = subprocess.check_output(curl_cmd).decode("utf-8")
|
||||
upstream_main_commit = json.loads(resp_json)["sha"]
|
||||
print(f"Upstream main branch latest commit: {upstream_main_commit}")
|
||||
|
||||
@@ -966,18 +971,19 @@ setup(
|
||||
ext_modules=ext_modules,
|
||||
install_requires=get_requirements(),
|
||||
extras_require={
|
||||
# AMD Zen CPU optimizations via zentorch
|
||||
"zen": ["zentorch"],
|
||||
"bench": ["pandas", "matplotlib", "seaborn", "datasets", "scipy", "plotly"],
|
||||
"tensorizer": ["tensorizer==2.10.1"],
|
||||
"fastsafetensors": ["fastsafetensors >= 0.2.2"],
|
||||
"instanttensor": ["instanttensor >= 0.1.5"],
|
||||
"runai": ["runai-model-streamer[s3,gcs] >= 0.15.3"],
|
||||
"runai": ["runai-model-streamer[s3,gcs,azure] >= 0.15.7"],
|
||||
"audio": [
|
||||
"librosa",
|
||||
"scipy",
|
||||
"soundfile",
|
||||
"mistral_common[audio]",
|
||||
"av",
|
||||
"torchcodec",
|
||||
], # Required for audio processing
|
||||
"video": [], # Kept for backwards compatibility
|
||||
"flashinfer": [], # Kept for backwards compatibility
|
||||
|
||||
@@ -179,7 +179,7 @@ class TestAllReduceFusedAddRMSNormStaticQuantFP4Model(torch.nn.Module):
|
||||
def ops_in_model_before(self):
|
||||
return [
|
||||
torch.ops.vllm.all_reduce.default,
|
||||
torch.ops._C.scaled_fp4_quant.default,
|
||||
torch.ops._C.scaled_fp4_quant.out,
|
||||
]
|
||||
|
||||
|
||||
|
||||
@@ -1,7 +1,9 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from ..entrypoints.openai.test_oot_registration import run_and_test_dummy_opt_api_server
|
||||
from tests.entrypoints.openai.chat_completion.test_oot_registration import (
|
||||
run_and_test_dummy_opt_api_server,
|
||||
)
|
||||
|
||||
|
||||
def test_distributed_oot(dummy_opt_path: str):
|
||||
|
||||
@@ -4,6 +4,9 @@
|
||||
|
||||
Tests the image source handling and tool_result content parsing in
|
||||
AnthropicServingMessages._convert_anthropic_to_openai_request().
|
||||
|
||||
Also covers extended-thinking edge cases such as ``redacted_thinking``
|
||||
blocks echoed back by Anthropic clients.
|
||||
"""
|
||||
|
||||
from vllm.entrypoints.anthropic.protocol import (
|
||||
@@ -373,3 +376,262 @@ class TestAttributionHeaderStripping:
|
||||
result = _convert(request)
|
||||
system_msg = result.messages[0]
|
||||
assert system_msg["content"] == "You are a helpful assistant."
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# Thinking block conversion (Anthropic → OpenAI)
|
||||
# ======================================================================
|
||||
|
||||
|
||||
class TestThinkingBlockConversion:
|
||||
"""Verify that thinking blocks in assistant messages are correctly
|
||||
moved to the ``reasoning`` field and stripped from ``content`` during
|
||||
the Anthropic→OpenAI conversion.
|
||||
|
||||
This is the Anthropic-endpoint path: the client echoes back the full
|
||||
assistant message (including thinking blocks emitted by vllm) in
|
||||
subsequent requests.
|
||||
"""
|
||||
|
||||
def test_thinking_plus_text_in_assistant_message(self):
|
||||
"""thinking + text → reasoning field + plain-string content."""
|
||||
request = _make_request(
|
||||
[
|
||||
{"role": "user", "content": "Write me some code."},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"type": "thinking",
|
||||
"thinking": "I should write a simple example.",
|
||||
"signature": "sig_abc123",
|
||||
},
|
||||
{"type": "text", "text": "Sure! Here is the code."},
|
||||
],
|
||||
},
|
||||
{"role": "user", "content": "Can you fix the bug?"},
|
||||
]
|
||||
)
|
||||
result = _convert(request)
|
||||
|
||||
# Find the assistant message in the converted output.
|
||||
asst_msgs = [m for m in result.messages if m.get("role") == "assistant"]
|
||||
assert len(asst_msgs) == 1
|
||||
asst = asst_msgs[0]
|
||||
|
||||
# Thinking content must be in reasoning, NOT in content.
|
||||
assert asst.get("reasoning") == "I should write a simple example."
|
||||
assert asst.get("content") == "Sure! Here is the code."
|
||||
|
||||
def test_thinking_only_in_assistant_message(self):
|
||||
"""Assistant message with only a thinking block (no visible text).
|
||||
|
||||
This can happen when the model emits reasoning but no final answer
|
||||
yet (e.g. a mid-turn reasoning step). Content should be None.
|
||||
"""
|
||||
request = _make_request(
|
||||
[
|
||||
{"role": "user", "content": "Hello"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"type": "thinking",
|
||||
"thinking": "Just thinking...",
|
||||
"signature": "sig_xyz",
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "user", "content": "Go on."},
|
||||
]
|
||||
)
|
||||
result = _convert(request)
|
||||
|
||||
asst_msgs = [m for m in result.messages if m.get("role") == "assistant"]
|
||||
assert len(asst_msgs) == 1
|
||||
asst = asst_msgs[0]
|
||||
|
||||
assert asst.get("reasoning") == "Just thinking..."
|
||||
# No visible text → content should be absent or None.
|
||||
assert asst.get("content") is None
|
||||
|
||||
def test_thinking_plus_tool_use_in_assistant_message(self):
|
||||
"""thinking + tool_use: reasoning field set, tool_calls populated."""
|
||||
request = _make_request(
|
||||
[
|
||||
{"role": "user", "content": "What is 2+2?"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"type": "thinking",
|
||||
"thinking": "I need to call the calculator.",
|
||||
"signature": "sig_tool",
|
||||
},
|
||||
{
|
||||
"type": "tool_use",
|
||||
"id": "call_001",
|
||||
"name": "calculator",
|
||||
"input": {"expression": "2+2"},
|
||||
},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "call_001",
|
||||
"content": "4",
|
||||
}
|
||||
],
|
||||
},
|
||||
]
|
||||
)
|
||||
result = _convert(request)
|
||||
|
||||
asst_msgs = [m for m in result.messages if m.get("role") == "assistant"]
|
||||
assert len(asst_msgs) == 1
|
||||
asst = asst_msgs[0]
|
||||
|
||||
assert asst.get("reasoning") == "I need to call the calculator."
|
||||
tool_calls = list(asst.get("tool_calls", []))
|
||||
assert len(tool_calls) == 1
|
||||
assert tool_calls[0]["function"]["name"] == "calculator"
|
||||
# No text content alongside reasoning + tool_use.
|
||||
assert asst.get("content") is None
|
||||
|
||||
def test_multiple_thinking_blocks_concatenated(self):
|
||||
"""Multiple thinking blocks should be joined in order."""
|
||||
request = _make_request(
|
||||
[
|
||||
{"role": "user", "content": "Think hard."},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"type": "thinking",
|
||||
"thinking": "First thought. ",
|
||||
"signature": "s1",
|
||||
},
|
||||
{
|
||||
"type": "thinking",
|
||||
"thinking": "Second thought.",
|
||||
"signature": "s2",
|
||||
},
|
||||
{"type": "text", "text": "Done."},
|
||||
],
|
||||
},
|
||||
]
|
||||
)
|
||||
result = _convert(request)
|
||||
|
||||
asst_msgs = [m for m in result.messages if m.get("role") == "assistant"]
|
||||
assert len(asst_msgs) == 1
|
||||
asst = asst_msgs[0]
|
||||
|
||||
assert asst.get("reasoning") == "First thought. Second thought."
|
||||
assert asst.get("content") == "Done."
|
||||
|
||||
def test_no_thinking_blocks_unchanged(self):
|
||||
"""Messages without thinking blocks must not be modified."""
|
||||
request = _make_request(
|
||||
[
|
||||
{"role": "user", "content": "Hi"},
|
||||
{"role": "assistant", "content": "Hello!"},
|
||||
]
|
||||
)
|
||||
result = _convert(request)
|
||||
|
||||
asst_msgs = [m for m in result.messages if m.get("role") == "assistant"]
|
||||
assert len(asst_msgs) == 1
|
||||
asst = asst_msgs[0]
|
||||
|
||||
assert asst.get("content") == "Hello!"
|
||||
assert "reasoning" not in asst
|
||||
|
||||
def test_multi_turn_with_thinking_blocks(self):
|
||||
"""Full multi-turn conversation: previous assistant messages that
|
||||
include thinking blocks must all be converted without a 400 error.
|
||||
|
||||
This is the primary regression scenario from the bug report:
|
||||
upgrading vllm from v0.15.1 → v0.17.0 introduced thinking-block
|
||||
support in responses, but echoing those responses back in subsequent
|
||||
requests caused a Pydantic validation failure.
|
||||
"""
|
||||
request = _make_request(
|
||||
[
|
||||
{"role": "user", "content": "Turn 1 question"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"type": "thinking",
|
||||
"thinking": "Reasoning for turn 1.",
|
||||
"signature": "s_t1",
|
||||
},
|
||||
{"type": "text", "text": "Answer for turn 1."},
|
||||
],
|
||||
},
|
||||
{"role": "user", "content": "Turn 2 question"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"type": "thinking",
|
||||
"thinking": "Reasoning for turn 2.",
|
||||
"signature": "s_t2",
|
||||
},
|
||||
{"type": "text", "text": "Answer for turn 2."},
|
||||
],
|
||||
},
|
||||
{"role": "user", "content": "Turn 3 question"},
|
||||
]
|
||||
)
|
||||
# Must not raise a ValidationError / 400.
|
||||
result = _convert(request)
|
||||
|
||||
asst_msgs = [m for m in result.messages if m.get("role") == "assistant"]
|
||||
assert len(asst_msgs) == 2
|
||||
|
||||
assert asst_msgs[0].get("reasoning") == "Reasoning for turn 1."
|
||||
assert asst_msgs[0].get("content") == "Answer for turn 1."
|
||||
assert asst_msgs[1].get("reasoning") == "Reasoning for turn 2."
|
||||
assert asst_msgs[1].get("content") == "Answer for turn 2."
|
||||
|
||||
def test_redacted_thinking_block_is_accepted(self):
|
||||
"""Anthropic clients may echo back redacted thinking blocks.
|
||||
|
||||
vLLM should accept these blocks (to avoid 400 validation errors)
|
||||
and ignore them when constructing the OpenAI-format prompt.
|
||||
"""
|
||||
request = _make_request(
|
||||
[
|
||||
{"role": "user", "content": "Hello"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"type": "thinking",
|
||||
"thinking": "Thinking...",
|
||||
"signature": "sig_think",
|
||||
},
|
||||
{
|
||||
"type": "redacted_thinking",
|
||||
"data": "BASE64_OR_OTHER_OPAQUE_DATA",
|
||||
},
|
||||
{"type": "text", "text": "Hi!"},
|
||||
],
|
||||
},
|
||||
{"role": "user", "content": "Continue"},
|
||||
]
|
||||
)
|
||||
result = _convert(request)
|
||||
|
||||
asst_msgs = [m for m in result.messages if m.get("role") == "assistant"]
|
||||
assert len(asst_msgs) == 1
|
||||
asst = asst_msgs[0]
|
||||
|
||||
# Redacted thinking is ignored, normal thinking still becomes reasoning.
|
||||
assert asst.get("reasoning") == "Thinking..."
|
||||
assert asst.get("content") == "Hi!"
|
||||
|
||||
@@ -4,12 +4,11 @@ import weakref
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.entrypoints.openai.chat_completion.test_vision import TEST_IMAGE_ASSETS
|
||||
from vllm import LLM
|
||||
from vllm.distributed import cleanup_dist_env_and_memory
|
||||
from vllm.sampling_params import SamplingParams
|
||||
|
||||
from ..openai.test_vision import TEST_IMAGE_ASSETS
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def text_llm():
|
||||
|
||||
@@ -6,13 +6,12 @@ import logging
|
||||
import pytest
|
||||
import regex as re
|
||||
|
||||
from tests.entrypoints.openai.chat_completion.test_vision import TEST_IMAGE_ASSETS
|
||||
from vllm import LLM
|
||||
from vllm.entrypoints.chat_utils import ChatCompletionMessageParam
|
||||
from vllm.v1.metrics import loggers as stat_loggers
|
||||
from vllm.v1.metrics.reader import Counter, Metric
|
||||
|
||||
from ..openai.test_vision import TEST_IMAGE_ASSETS
|
||||
|
||||
|
||||
def _make_messages(image_url: str) -> list[ChatCompletionMessageParam]:
|
||||
return [
|
||||
|
||||
+1
-2
@@ -7,11 +7,10 @@ import openai
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
from vllm.assets.audio import AudioAsset
|
||||
from vllm.multimodal.utils import encode_audio_base64, encode_audio_url, fetch_audio
|
||||
|
||||
from ...utils import RemoteOpenAIServer
|
||||
|
||||
MODEL_NAME = "fixie-ai/ultravox-v0_5-llama-3_2-1b"
|
||||
TEST_AUDIO_URLS = [
|
||||
AudioAsset("winning_call").url,
|
||||
@@ -0,0 +1,175 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import base64
|
||||
import json
|
||||
|
||||
import openai
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
|
||||
from tests.conftest import VideoTestAssets
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
|
||||
MODEL_NAME = "Qwen/Qwen2.5-Omni-3B"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def server():
|
||||
args = [
|
||||
"--max-model-len",
|
||||
"16384",
|
||||
"--enforce-eager",
|
||||
"--limit-mm-per-prompt",
|
||||
json.dumps({"audio": 3, "video": 3}),
|
||||
]
|
||||
|
||||
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.core_model
|
||||
@pytest.mark.asyncio
|
||||
async def test_online_audio_in_video(
|
||||
client: openai.AsyncOpenAI, video_assets: VideoTestAssets
|
||||
):
|
||||
"""Test video input with `audio_in_video=True`"""
|
||||
|
||||
# we don't use video_urls above because they missed audio stream.
|
||||
video_path = video_assets[0].video_path
|
||||
with open(video_path, "rb") as f:
|
||||
video_base64 = base64.b64encode(f.read()).decode("utf-8")
|
||||
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "What's in this video?"},
|
||||
{
|
||||
"type": "video_url",
|
||||
"video_url": {"url": f"data:video/mp4;base64,{video_base64}"},
|
||||
},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
# multi-turn to test mm processor cache as well
|
||||
for _ in range(2):
|
||||
chat_completion = await client.chat.completions.create(
|
||||
model=MODEL_NAME,
|
||||
messages=messages,
|
||||
max_tokens=16,
|
||||
extra_body={
|
||||
"mm_processor_kwargs": {
|
||||
"use_audio_in_video": True,
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
assert len(chat_completion.choices) == 1
|
||||
choice = chat_completion.choices[0]
|
||||
assert choice.finish_reason == "length"
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.asyncio
|
||||
async def test_online_audio_in_video_multi_videos(
|
||||
client: openai.AsyncOpenAI, video_assets: VideoTestAssets
|
||||
):
|
||||
"""Test multi-video input with `audio_in_video=True`"""
|
||||
|
||||
# we don't use video_urls above because they missed audio stream.
|
||||
video_path = video_assets[0].video_path
|
||||
with open(video_path, "rb") as f:
|
||||
video_base64 = base64.b64encode(f.read()).decode("utf-8")
|
||||
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "What's in these two videos?"},
|
||||
{
|
||||
"type": "video_url",
|
||||
"video_url": {"url": f"data:video/mp4;base64,{video_base64}"},
|
||||
},
|
||||
{
|
||||
"type": "video_url",
|
||||
"video_url": {"url": f"data:video/mp4;base64,{video_base64}"},
|
||||
},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
# multi-turn to test mm processor cache as well
|
||||
for _ in range(2):
|
||||
chat_completion = await client.chat.completions.create(
|
||||
model=MODEL_NAME,
|
||||
messages=messages,
|
||||
max_tokens=16,
|
||||
extra_body={
|
||||
"mm_processor_kwargs": {
|
||||
"use_audio_in_video": True,
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
assert len(chat_completion.choices) == 1
|
||||
choice = chat_completion.choices[0]
|
||||
assert choice.finish_reason == "length"
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.asyncio
|
||||
async def test_online_audio_in_video_interleaved(
|
||||
client: openai.AsyncOpenAI, video_assets: VideoTestAssets
|
||||
):
|
||||
"""Test interleaved video/audio input with `audio_in_video=True`"""
|
||||
|
||||
# we don't use video_urls above because they missed audio stream.
|
||||
video_path = video_assets[0].video_path
|
||||
with open(video_path, "rb") as f:
|
||||
video_base64 = base64.b64encode(f.read()).decode("utf-8")
|
||||
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "What's in these two videos?"},
|
||||
{
|
||||
"type": "video_url",
|
||||
"video_url": {"url": f"data:video/mp4;base64,{video_base64}"},
|
||||
},
|
||||
{
|
||||
"type": "audio_url",
|
||||
"audio_url": {"url": f"data:audio/mp4;base64,{video_base64}"},
|
||||
},
|
||||
{
|
||||
"type": "video_url",
|
||||
"video_url": {"url": f"data:video/mp4;base64,{video_base64}"},
|
||||
},
|
||||
],
|
||||
}
|
||||
]
|
||||
with pytest.raises(
|
||||
openai.BadRequestError,
|
||||
match="use_audio_in_video requires equal number of audio and video items",
|
||||
):
|
||||
await client.chat.completions.create(
|
||||
model=MODEL_NAME,
|
||||
messages=messages,
|
||||
max_tokens=16,
|
||||
extra_body={
|
||||
"mm_processor_kwargs": {
|
||||
"use_audio_in_video": True,
|
||||
}
|
||||
},
|
||||
)
|
||||
+2
-2
@@ -8,8 +8,8 @@ import pytest
|
||||
import pytest_asyncio
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
from ...conftest import AudioTestAssets
|
||||
from ...utils import RemoteOpenAIServer
|
||||
from tests.conftest import AudioTestAssets
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
|
||||
# NOTE - the tests in this module are currently analogous to test_chat, but are
|
||||
# separated to avoid OOM killing due to module-scoped servers, since we
|
||||
+1
-1
@@ -1,7 +1,7 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from ...utils import VLLM_PATH, RemoteOpenAIServer
|
||||
from tests.utils import VLLM_PATH, RemoteOpenAIServer
|
||||
|
||||
chatml_jinja_path = VLLM_PATH / "examples/template_chatml.jinja"
|
||||
assert chatml_jinja_path.exists()
|
||||
+1
-1
@@ -8,7 +8,7 @@ from typing import Any, NamedTuple
|
||||
import openai # use the official client for correctness check
|
||||
import pytest
|
||||
|
||||
from ...utils import RemoteOpenAIServer
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
|
||||
# # any model with a chat template should work here
|
||||
MODEL_NAME = "Qwen/Qwen2-1.5B-Instruct"
|
||||
+1
-2
@@ -7,11 +7,10 @@ import openai
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
from vllm.multimodal.utils import encode_video_url, fetch_video
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
from ...utils import RemoteOpenAIServer
|
||||
|
||||
MODEL_NAME = "llava-hf/llava-onevision-qwen2-0.5b-ov-hf"
|
||||
MAXIMUM_VIDEOS = 3
|
||||
|
||||
+1
-2
@@ -8,12 +8,11 @@ import pytest
|
||||
import pytest_asyncio
|
||||
from transformers import AutoProcessor
|
||||
|
||||
from tests.utils import ROCM_ENV_OVERRIDES, ROCM_EXTRA_ARGS, RemoteOpenAIServer
|
||||
from vllm.multimodal.media import MediaWithBytes
|
||||
from vllm.multimodal.utils import encode_image_url, fetch_image
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
from ...utils import ROCM_ENV_OVERRIDES, ROCM_EXTRA_ARGS, RemoteOpenAIServer
|
||||
|
||||
MODEL_NAME = "microsoft/Phi-3.5-vision-instruct"
|
||||
MAXIMUM_IMAGES = 2
|
||||
|
||||
+1
-2
@@ -8,10 +8,9 @@ import pytest
|
||||
import requests
|
||||
import torch
|
||||
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
from vllm.utils.serial_utils import tensor2base64
|
||||
|
||||
from ...utils import RemoteOpenAIServer
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model_name", ["ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11"]
|
||||
+1
-1
@@ -14,7 +14,7 @@ import torch
|
||||
from openai import BadRequestError
|
||||
from transformers import AutoConfig
|
||||
|
||||
from ...utils import RemoteOpenAIServer
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
|
||||
# any model with a chat template should work here
|
||||
MODEL_NAME = "facebook/opt-125m"
|
||||
+1
-2
@@ -11,11 +11,10 @@ import pytest
|
||||
import regex as re
|
||||
import torch
|
||||
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
from vllm.config import ModelConfig
|
||||
from vllm.renderers.embed_utils import safe_load_prompt_embeds
|
||||
|
||||
from ...utils import RemoteOpenAIServer
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_empty_prompt():
|
||||
+1
-2
@@ -9,6 +9,7 @@ import pytest
|
||||
import pytest_asyncio
|
||||
import torch.cuda
|
||||
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
from vllm.engine.arg_utils import EngineArgs
|
||||
from vllm.model_executor.model_loader.tensorizer import (
|
||||
TensorizerConfig,
|
||||
@@ -17,8 +18,6 @@ from vllm.model_executor.model_loader.tensorizer import (
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
from ...utils import RemoteOpenAIServer
|
||||
|
||||
MODEL_NAME = "unsloth/llama-3.2-1b-Instruct"
|
||||
LORA_PATH = "davzoku/finqa_adapter_1b"
|
||||
|
||||
+1
-2
@@ -6,11 +6,10 @@ import tempfile
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
from vllm.model_executor.model_loader.weight_utils import download_weights_from_hf
|
||||
from vllm.tokenizers import get_tokenizer
|
||||
|
||||
from ...utils import RemoteOpenAIServer
|
||||
|
||||
MODEL_NAME = "Qwen/Qwen3-0.6B"
|
||||
MODEL_PATH = os.path.join(tempfile.gettempdir(), "qwen3_06b")
|
||||
|
||||
@@ -137,6 +137,59 @@ async def test_streaming_output_consistency(client: OpenAI, model_name: str):
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
async def test_streaming_logprobs(client: OpenAI, model_name: str):
|
||||
"""Test that streaming with logprobs returns valid logprob data on
|
||||
output_text.delta events and that top_logprobs has the requested count."""
|
||||
response = await client.responses.create(
|
||||
model=model_name,
|
||||
input="Say hello.",
|
||||
stream=True,
|
||||
top_logprobs=3,
|
||||
include=["message.output_text.logprobs"],
|
||||
)
|
||||
|
||||
events = []
|
||||
async for event in response:
|
||||
events.append(event)
|
||||
|
||||
assert len(events) > 0
|
||||
|
||||
# Collect all output_text.delta events that carry logprobs
|
||||
text_delta_events = [e for e in events if e.type == "response.output_text.delta"]
|
||||
assert len(text_delta_events) > 0, "Expected at least one text delta event"
|
||||
|
||||
for delta_event in text_delta_events:
|
||||
logprobs = delta_event.logprobs
|
||||
assert logprobs is not None, "logprobs should be present on text delta events"
|
||||
assert len(logprobs) > 0, "logprobs list should not be empty"
|
||||
for lp in logprobs:
|
||||
# Each logprob entry must have a token and a logprob value
|
||||
assert lp.token is not None
|
||||
assert isinstance(lp.logprob, float)
|
||||
assert lp.logprob <= 0.0, f"logprob should be <= 0, got {lp.logprob}"
|
||||
# top_logprobs should have up to 3 entries
|
||||
assert lp.top_logprobs is not None
|
||||
assert len(lp.top_logprobs) <= 3
|
||||
for tl in lp.top_logprobs:
|
||||
assert tl.token is not None
|
||||
assert isinstance(tl.logprob, float)
|
||||
|
||||
# Verify that top_logprobs are actually populated, not always empty
|
||||
all_top_logprobs = [
|
||||
tl for e in text_delta_events for lp in e.logprobs for tl in lp.top_logprobs
|
||||
]
|
||||
assert len(all_top_logprobs) > 0, (
|
||||
"Expected at least one top_logprobs entry across all delta events"
|
||||
)
|
||||
|
||||
# Verify the completed event still has valid output
|
||||
completed = events[-1]
|
||||
assert completed.type == "response.completed"
|
||||
assert completed.response.status == "completed"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
async def test_streaming_reasoning_tokens_e2e(client: OpenAI, model_name: str):
|
||||
|
||||
@@ -1,279 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""Integration tests for GPT-OSS structural tags functionality (PR #25515)."""
|
||||
|
||||
import json
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.entrypoints.mcp.tool_server import ToolServer
|
||||
from vllm.reasoning.gptoss_reasoning_parser import (
|
||||
GptOssReasoningParser,
|
||||
)
|
||||
from vllm.sampling_params import StructuredOutputsParams
|
||||
|
||||
|
||||
class TestGptOssStructuralTagsIntegration:
|
||||
"""Integration tests for structural tags in GPT-OSS tool calls."""
|
||||
|
||||
@pytest.fixture
|
||||
def mock_tokenizer(self):
|
||||
"""Create a mock tokenizer."""
|
||||
tokenizer = Mock()
|
||||
tokenizer.encode = Mock(return_value=[1, 2, 3, 4, 5])
|
||||
tokenizer.get_vocab = Mock(return_value={"<|end|>": 6})
|
||||
return tokenizer
|
||||
|
||||
@pytest.fixture
|
||||
def gptoss_parser(self, mock_tokenizer):
|
||||
"""Create a real GptOssReasoningParser instance."""
|
||||
return GptOssReasoningParser(mock_tokenizer)
|
||||
|
||||
@pytest.fixture
|
||||
def tool_server_with_python(self):
|
||||
"""Create a tool server with Python tool enabled."""
|
||||
tool_server = Mock(spec=ToolServer)
|
||||
tool_server.has_tool = Mock(side_effect=lambda tool: tool == "python")
|
||||
return tool_server
|
||||
|
||||
@pytest.fixture
|
||||
def tool_server_empty(self):
|
||||
"""Create a tool server with no tools."""
|
||||
tool_server = Mock(spec=ToolServer)
|
||||
tool_server.has_tool = Mock(return_value=False)
|
||||
return tool_server
|
||||
|
||||
def test_end_to_end_no_tools(self, gptoss_parser):
|
||||
"""Test end-to-end flow when no tools are available."""
|
||||
# Test the parser directly
|
||||
result = gptoss_parser.prepare_structured_tag(None, None)
|
||||
parsed_result = json.loads(result)
|
||||
|
||||
# Verify basic structure
|
||||
assert parsed_result["type"] == "structural_tag"
|
||||
assert parsed_result["format"]["type"] == "triggered_tags"
|
||||
assert len(parsed_result["format"]["tags"]) == 1
|
||||
|
||||
# Verify only analysis channel is allowed
|
||||
analysis_tag = parsed_result["format"]["tags"][0]
|
||||
assert analysis_tag["begin"] == "<|channel|>analysis<|message|>"
|
||||
assert analysis_tag["content"]["type"] == "any_text"
|
||||
assert analysis_tag["end"] == "<|end|>"
|
||||
|
||||
# Verify triggers
|
||||
assert parsed_result["format"]["triggers"] == ["<|channel|>analysis"]
|
||||
assert parsed_result["format"]["stop_after_first"] is False
|
||||
|
||||
def test_end_to_end_with_python_tool(self, gptoss_parser, tool_server_with_python):
|
||||
"""Test end-to-end flow with Python tool enabled."""
|
||||
result = gptoss_parser.prepare_structured_tag(None, tool_server_with_python)
|
||||
parsed_result = json.loads(result)
|
||||
|
||||
# Should have analysis tag + 2 python tags
|
||||
assert len(parsed_result["format"]["tags"]) == 3
|
||||
|
||||
# Verify all expected tags are present
|
||||
tag_begins = [tag["begin"] for tag in parsed_result["format"]["tags"]]
|
||||
expected_begins = [
|
||||
"<|channel|>analysis<|message|>",
|
||||
"<|channel|>commentary to=python",
|
||||
"<|channel|>analysis to=python",
|
||||
]
|
||||
|
||||
for expected in expected_begins:
|
||||
assert expected in tag_begins
|
||||
|
||||
# Verify triggers include commentary
|
||||
assert "<|channel|>analysis" in parsed_result["format"]["triggers"]
|
||||
assert "<|channel|>commentary to=" in parsed_result["format"]["triggers"]
|
||||
|
||||
def test_structured_outputs_params_integration(
|
||||
self, gptoss_parser, tool_server_with_python
|
||||
):
|
||||
"""Test integration with StructuredOutputsParams."""
|
||||
# Generate structural tag
|
||||
structural_tag = gptoss_parser.prepare_structured_tag(
|
||||
None, tool_server_with_python
|
||||
)
|
||||
|
||||
# Create StructuredOutputsParams
|
||||
params = StructuredOutputsParams(structural_tag=structural_tag)
|
||||
|
||||
# Verify the tag is properly stored and accessible
|
||||
assert params.structural_tag == structural_tag
|
||||
|
||||
# Verify the tag is valid JSON
|
||||
parsed_tag = json.loads(params.structural_tag)
|
||||
assert parsed_tag["type"] == "structural_tag"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"browser, python, container, expected_tags",
|
||||
[
|
||||
# No tools
|
||||
(False, False, False, 1),
|
||||
# Single tool
|
||||
(True, False, False, 3),
|
||||
# Multiple tools
|
||||
(True, True, False, 5),
|
||||
# All tools
|
||||
(True, True, True, 7),
|
||||
],
|
||||
)
|
||||
def test_tool_server_interaction_flow(
|
||||
self, gptoss_parser, browser, python, container, expected_tags
|
||||
):
|
||||
"""Test the complete tool server interaction flow."""
|
||||
|
||||
# Create a mock ToolServer
|
||||
tool_server = Mock(spec=ToolServer)
|
||||
|
||||
# Simulate tool availability based on parameters
|
||||
tool_server.has_tool = Mock(
|
||||
side_effect=lambda tool: {
|
||||
"browser": browser,
|
||||
"python": python,
|
||||
"container": container,
|
||||
}.get(tool, False)
|
||||
)
|
||||
|
||||
# Run the parser and verify results
|
||||
result = gptoss_parser.prepare_structured_tag(None, tool_server)
|
||||
parsed_result = json.loads(result)
|
||||
|
||||
# Validate number of tags
|
||||
assert len(parsed_result["format"]["tags"]) == expected_tags
|
||||
|
||||
# Verify tool-specific tags exist for enabled tools
|
||||
tag_begins = [tag["begin"] for tag in parsed_result["format"]["tags"]]
|
||||
for tool, enabled in {
|
||||
"browser": browser,
|
||||
"python": python,
|
||||
"container": container,
|
||||
}.items():
|
||||
if enabled:
|
||||
assert f"<|channel|>commentary to={tool}" in tag_begins
|
||||
assert f"<|channel|>analysis to={tool}" in tag_begins
|
||||
|
||||
def test_original_tag_preservation(self, gptoss_parser, tool_server_with_python):
|
||||
"""Test that original tags are preserved when provided."""
|
||||
original_tag = '{"type": "custom_tag", "data": "preserved"}'
|
||||
|
||||
result = gptoss_parser.prepare_structured_tag(
|
||||
original_tag, tool_server_with_python
|
||||
)
|
||||
|
||||
# Should return original tag unchanged
|
||||
assert result == original_tag
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"tools",
|
||||
[
|
||||
[],
|
||||
["browser"],
|
||||
["python"],
|
||||
["container"],
|
||||
["browser", "python"],
|
||||
["browser", "container"],
|
||||
["python", "container"],
|
||||
["browser", "python", "container"],
|
||||
],
|
||||
)
|
||||
def test_json_validity_comprehensive(self, gptoss_parser, tools):
|
||||
"""Test JSON validity across all possible tool combinations."""
|
||||
|
||||
tool_server = Mock(spec=ToolServer)
|
||||
tool_server.has_tool = Mock(side_effect=lambda tool: tool in tools)
|
||||
|
||||
result = gptoss_parser.prepare_structured_tag(None, tool_server)
|
||||
|
||||
# Should be valid JSON
|
||||
parsed_result = json.loads(result)
|
||||
|
||||
# Should have correct structure
|
||||
assert parsed_result["type"] == "structural_tag"
|
||||
assert "format" in parsed_result
|
||||
assert "tags" in parsed_result["format"]
|
||||
assert "triggers" in parsed_result["format"]
|
||||
|
||||
# Tag count should be: 1 (analysis) + 2 * len(tools)
|
||||
expected_tag_count = 1 + (2 * len(tools))
|
||||
assert len(parsed_result["format"]["tags"]) == expected_tag_count
|
||||
|
||||
def test_error_handling_invalid_tool_server(self, gptoss_parser):
|
||||
"""Test error handling with invalid tool server."""
|
||||
# Tool server that raises exceptions
|
||||
tool_server = Mock(spec=ToolServer)
|
||||
tool_server.has_tool = Mock(side_effect=Exception("Tool server error"))
|
||||
|
||||
# Should handle gracefully and still return a valid tag
|
||||
with pytest.raises(Exception, match="Tool server error"):
|
||||
gptoss_parser.prepare_structured_tag(None, tool_server)
|
||||
|
||||
def test_concurrent_requests_isolation(self, gptoss_parser):
|
||||
"""Test that concurrent requests don't interfere with each other."""
|
||||
# Simulate concurrent requests with different tool servers
|
||||
tool_server_1 = Mock(spec=ToolServer)
|
||||
tool_server_1.has_tool = Mock(side_effect=lambda tool: tool == "python")
|
||||
|
||||
tool_server_2 = Mock(spec=ToolServer)
|
||||
tool_server_2.has_tool = Mock(side_effect=lambda tool: tool == "browser")
|
||||
|
||||
# Generate tags concurrently
|
||||
result_1 = gptoss_parser.prepare_structured_tag(None, tool_server_1)
|
||||
result_2 = gptoss_parser.prepare_structured_tag(None, tool_server_2)
|
||||
|
||||
# Parse results
|
||||
parsed_1 = json.loads(result_1)
|
||||
parsed_2 = json.loads(result_2)
|
||||
|
||||
# Verify they have different tool configurations
|
||||
tags_1 = [tag["begin"] for tag in parsed_1["format"]["tags"]]
|
||||
tags_2 = [tag["begin"] for tag in parsed_2["format"]["tags"]]
|
||||
|
||||
# Result 1 should have python tags
|
||||
assert "<|channel|>commentary to=python" in tags_1
|
||||
assert "<|channel|>commentary to=browser" not in tags_1
|
||||
|
||||
# Result 2 should have browser tags
|
||||
assert "<|channel|>commentary to=browser" in tags_2
|
||||
assert "<|channel|>commentary to=python" not in tags_2
|
||||
|
||||
def test_tag_format_consistency(self, gptoss_parser):
|
||||
"""Test that all generated tags follow consistent format."""
|
||||
tool_server = Mock(spec=ToolServer)
|
||||
tool_server.has_tool = Mock(
|
||||
side_effect=lambda tool: tool in ["python", "browser"]
|
||||
)
|
||||
|
||||
result = gptoss_parser.prepare_structured_tag(None, tool_server)
|
||||
parsed_result = json.loads(result)
|
||||
|
||||
# Verify all tags have required fields
|
||||
for tag in parsed_result["format"]["tags"]:
|
||||
assert "begin" in tag
|
||||
assert "content" in tag
|
||||
assert "end" in tag
|
||||
assert tag["content"]["type"] == "any_text"
|
||||
assert tag["end"] == "<|end|>"
|
||||
|
||||
# Verify begin format
|
||||
assert tag["begin"].startswith("<|channel|>")
|
||||
|
||||
def test_trigger_configuration(self, gptoss_parser):
|
||||
"""Test trigger configuration for different tool setups."""
|
||||
# Test with no tools
|
||||
result_no_tools = gptoss_parser.prepare_structured_tag(None, None)
|
||||
parsed_no_tools = json.loads(result_no_tools)
|
||||
assert parsed_no_tools["format"]["triggers"] == ["<|channel|>analysis"]
|
||||
|
||||
# Test with tools
|
||||
tool_server = Mock(spec=ToolServer)
|
||||
tool_server.has_tool = Mock(side_effect=lambda tool: tool == "python")
|
||||
|
||||
result_with_tools = gptoss_parser.prepare_structured_tag(None, tool_server)
|
||||
parsed_with_tools = json.loads(result_with_tools)
|
||||
|
||||
expected_triggers = ["<|channel|>analysis", "<|channel|>commentary to="]
|
||||
assert set(parsed_with_tools["format"]["triggers"]) == set(expected_triggers)
|
||||
@@ -5,7 +5,7 @@ import json
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.entrypoints.openai.tool_parsers.utils import (
|
||||
from tests.tool_parsers.utils import (
|
||||
run_tool_extraction,
|
||||
run_tool_extraction_streaming,
|
||||
)
|
||||
|
||||
@@ -0,0 +1,360 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import json
|
||||
import random
|
||||
from typing import Any
|
||||
|
||||
import openai
|
||||
import pytest
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
|
||||
from vllm.entrypoints.openai.engine.protocol import (
|
||||
DeltaMessage,
|
||||
)
|
||||
from vllm.tool_parsers.granite4_tool_parser import Granite4ToolParser
|
||||
|
||||
from ....utils import RemoteOpenAIServer
|
||||
|
||||
MODEL = "ibm-granite/granite-4.0-h-tiny"
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def server():
|
||||
model = MODEL
|
||||
args_for_model = [
|
||||
"--enforce-eager",
|
||||
"--enable-auto-tool-choice",
|
||||
"--tool-call-parser",
|
||||
"granite4",
|
||||
"--tokenizer",
|
||||
"ibm-granite/granite-4.0-h-tiny",
|
||||
"--max-model-len",
|
||||
"4096",
|
||||
"--max-num-seqs",
|
||||
"2",
|
||||
]
|
||||
with RemoteOpenAIServer(model, args_for_model, max_wait_seconds=480) as server:
|
||||
yield server
|
||||
|
||||
|
||||
def create_complex_input(create_string_args: bool):
|
||||
coord_arg: dict | str = {
|
||||
"coordinates": [[23.54, 43.1], [-12.2, 54.3], [4, 5]],
|
||||
"coordinate_type": "latlong",
|
||||
}
|
||||
if create_string_args:
|
||||
# test granite behavior
|
||||
coord_arg = json.dumps(coord_arg)
|
||||
return [
|
||||
{"name": "find_bbox", "arguments": coord_arg},
|
||||
{
|
||||
"name": "get_stock_price",
|
||||
"arguments": {
|
||||
"symbol": "AAPL",
|
||||
"start_date": "2021-01-01",
|
||||
"end_date": "2021-12-31",
|
||||
},
|
||||
},
|
||||
{"name": "find_bbox", "arguments": coord_arg},
|
||||
]
|
||||
|
||||
|
||||
def random_chunks(s: str, min_len: int, max_len: int):
|
||||
chunks = []
|
||||
i = 0
|
||||
n = len(s)
|
||||
|
||||
while i < n:
|
||||
size = random.randint(min_len, max_len)
|
||||
chunks.append(s[i : i + size])
|
||||
i += size
|
||||
|
||||
return chunks
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def tokenizer():
|
||||
return AutoTokenizer.from_pretrained(MODEL)
|
||||
|
||||
|
||||
# create a variety of input chunk sizes
|
||||
@pytest.mark.parametrize(
|
||||
"min_chunk, max_chunk",
|
||||
[
|
||||
(1, 1),
|
||||
(1, 2),
|
||||
(5, 7),
|
||||
(6, 20),
|
||||
],
|
||||
)
|
||||
def test_tool_call_parser_complex(min_chunk: int, max_chunk: int, tokenizer):
|
||||
input_dicts = create_complex_input(True)
|
||||
|
||||
formatted_tcs = [
|
||||
"<tool_call> " + json.dumps(call) + " </tool_call>" for call in input_dicts
|
||||
]
|
||||
|
||||
text_messages = [
|
||||
"Here goes the bbox call: \n",
|
||||
" Now the stock price call: \n ",
|
||||
" Now another bbox call: \n ",
|
||||
" See? I'm a helpful assistant.",
|
||||
]
|
||||
|
||||
test_input = (
|
||||
text_messages[0]
|
||||
+ formatted_tcs[0]
|
||||
+ text_messages[1]
|
||||
+ formatted_tcs[1]
|
||||
+ text_messages[2]
|
||||
+ formatted_tcs[2]
|
||||
+ text_messages[3]
|
||||
)
|
||||
|
||||
any_chat_request = ChatCompletionRequest(
|
||||
seed=42,
|
||||
model=MODEL,
|
||||
messages=[],
|
||||
)
|
||||
|
||||
parser = Granite4ToolParser(tokenizer=tokenizer)
|
||||
|
||||
delta_messages = list[DeltaMessage]()
|
||||
for text in random_chunks(test_input, min_chunk, max_chunk):
|
||||
delta = parser.extract_tool_calls_streaming(
|
||||
previous_text="",
|
||||
current_text="",
|
||||
delta_text=text,
|
||||
previous_token_ids=[],
|
||||
current_token_ids=[],
|
||||
delta_token_ids=[],
|
||||
request=any_chat_request,
|
||||
)
|
||||
if delta is not None:
|
||||
delta_messages.append(delta)
|
||||
|
||||
content = ""
|
||||
tool_calls = list[dict[str, Any]]()
|
||||
|
||||
current_name = "__start__"
|
||||
current_args = ""
|
||||
|
||||
for msg in delta_messages:
|
||||
if msg.content:
|
||||
content += msg.content
|
||||
for tool_call in msg.tool_calls:
|
||||
if delta_func := tool_call.function:
|
||||
if delta_func.name is not None:
|
||||
if current_name == "__start__":
|
||||
current_name = delta_func.name
|
||||
|
||||
if delta_func.name != current_name:
|
||||
tool_calls.append(
|
||||
{
|
||||
"name": current_name,
|
||||
"arguments": json.loads(current_args),
|
||||
}
|
||||
)
|
||||
current_name = delta_func.name
|
||||
current_args = ""
|
||||
|
||||
if delta_func.arguments:
|
||||
current_args += delta_func.arguments
|
||||
|
||||
if current_name != "__start__":
|
||||
tool_calls.append({"name": current_name, "arguments": json.loads(current_args)})
|
||||
|
||||
assert content == "".join(text_messages)
|
||||
assert tool_calls == create_complex_input(False)
|
||||
|
||||
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_acme_region_name_for_transaction_id",
|
||||
"description": "Returns ACME transaction/transaction ID information"
|
||||
" including ACME regions\n\nArgs:\n start_time "
|
||||
"(str): Start date and time in datetime format "
|
||||
'"%Y-%m-%dT%H:%M:%S.%f"\n end_time (str): End '
|
||||
"date and time in datetime format "
|
||||
'"%Y-%m-%dT%H:%M:%S.%f"\n size (int, optional): '
|
||||
"Number of ACME Transaction IDs to return\n "
|
||||
"order (str, optional): Sort by most run "
|
||||
"transaction IDs. The value can be 'asc' for "
|
||||
"ascending or 'desc' for descending\n "
|
||||
"transaction_id (str, optional): ACME Transaction "
|
||||
"ID to filter on\n acme_region (str, optional): "
|
||||
"ACME Region to filter on\nReturns:\n - A "
|
||||
"dictionary containing a list of ACME transaction "
|
||||
"ids and the ACME regions they run in:\n {\n"
|
||||
' "Number of transaction IDs" : int,\n'
|
||||
' "Total transaction IDs available": int'
|
||||
',\n "ACME Transaction IDs": [\n '
|
||||
' {\n "Transaction ID": '
|
||||
'str,\n "Number of runs": int,\n'
|
||||
' "ACME Regions": [str],\n '
|
||||
" },\n ...\n ],"
|
||||
'\n "Start time" : datetime,\n '
|
||||
' "End time" : datetime,\n '
|
||||
' "Order" : str\n }\n '
|
||||
" - If no ACME region found for transaction id, "
|
||||
'returns:\n {"Success": "No ACME region '
|
||||
'found for transaction id."}\n - If an error '
|
||||
'occurs, returns:\n {"Error": "{exception'
|
||||
' message}"}',
|
||||
"parameters": {
|
||||
"properties": {
|
||||
"start_time": {},
|
||||
"end_time": {},
|
||||
"size": {"default": 500},
|
||||
"order": {"default": "desc"},
|
||||
"transaction_id": {"default": None},
|
||||
"acme_region": {"default": None},
|
||||
},
|
||||
"required": ["start_time", "end_time"],
|
||||
"type": "object",
|
||||
},
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
tools2 = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_current_weather",
|
||||
"description": "Get the current weather",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"location": {
|
||||
"description": "The city and state, e.g. San Francisco, CA",
|
||||
"type": "string",
|
||||
}
|
||||
},
|
||||
"required": ["location"],
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_stock_price",
|
||||
"description": "Retrieves the current stock price for a given "
|
||||
"ticker symbol. The ticker symbol must be a valid "
|
||||
"symbol for a publicly traded company on a major US"
|
||||
" stock exchange like NYSE or NASDAQ. The tool will"
|
||||
" return the latest trade price in USD. It should "
|
||||
"be used when the user asks about the current or "
|
||||
"most recent price of a specific stock. It will not"
|
||||
" provide any other information about the stock or"
|
||||
" company.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"ticker": {
|
||||
"description": "The stock ticker symbol, e.g."
|
||||
" AAPL for Apple Inc.",
|
||||
"type": "string",
|
||||
}
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
messages = [
|
||||
{
|
||||
"content": "\n\nSystem: You are a helpful, precise, and methodical AI"
|
||||
" assistant that uses tool outputs provided inline.\nAlways"
|
||||
" assume the current datetime is 2026-01-29T13:59:09.238901"
|
||||
"+00:00.\n\nIf you receive a ToolMessage with `tool_call_id"
|
||||
'` equal to "get_time_range" (or "time_range_tool"), you '
|
||||
"MUST:\n 1. Parse that JSON and use the values `start` and"
|
||||
" `end` directly when calling other tools.\n 2. Do not "
|
||||
"re-call or re-compute the time range.\n 3. Pass resolved "
|
||||
"values (ISO strings) as arguments to any subsequent tool "
|
||||
"(do not pass function metadata or placeholders).\n 4. If "
|
||||
"a tool requires datetime objects rather than strings, "
|
||||
"convert the ISO strings into language-native datetime "
|
||||
"objects before invoking.\n\nAlways return fully resolved "
|
||||
"arguments in correct types (e.g., ISO datetime strings or"
|
||||
" datetime objects) and never include placeholders like "
|
||||
'"<start>".\n\n',
|
||||
"role": "system",
|
||||
},
|
||||
{
|
||||
"content": "What are the transaction IDs that ran in the"
|
||||
" ACME region A9345 over the last two months?",
|
||||
"role": "user",
|
||||
},
|
||||
{
|
||||
"content": '["2026-01-26T09: 51: 55.467722Z", "2026-01-27T09: 51: 55.467722Z"]',
|
||||
"role": "tool",
|
||||
"tool_call_id": "time_range_tool",
|
||||
},
|
||||
]
|
||||
messages2 = [{"role": "user", "content": "What's stock price for IBM?"}]
|
||||
|
||||
messages3 = [{"role": "user", "content": "What's the current weather in New York?"}]
|
||||
|
||||
|
||||
def get_args(client: openai.OpenAI, _tools, _messages, _stop):
|
||||
response = client.chat.completions.create(
|
||||
model=MODEL,
|
||||
messages=_messages,
|
||||
temperature=0,
|
||||
tools=_tools,
|
||||
max_tokens=200,
|
||||
stop=_stop,
|
||||
tool_choice="auto",
|
||||
)
|
||||
|
||||
return response.choices[0].message.tool_calls[0].function.arguments
|
||||
|
||||
|
||||
async def get_args_streaming(
|
||||
async_client: openai.AsyncOpenAI, _tools, _messages, _stop
|
||||
):
|
||||
stream = await async_client.chat.completions.create(
|
||||
model=MODEL,
|
||||
messages=_messages,
|
||||
temperature=0,
|
||||
tools=_tools,
|
||||
max_tokens=200,
|
||||
stop=_stop,
|
||||
tool_choice="auto",
|
||||
stream=True,
|
||||
)
|
||||
full_call = []
|
||||
async for chunk in stream:
|
||||
tc = chunk.choices[0].delta.tool_calls
|
||||
if tc and tc[0].function.arguments:
|
||||
full_call.append(tc[0].function.arguments)
|
||||
return "".join(full_call)
|
||||
|
||||
|
||||
async def run_scenario(server: RemoteOpenAIServer, _tools, _messages, _stop):
|
||||
non_streaming = get_args(server.get_client(), _tools, _messages, _stop)
|
||||
json.loads(non_streaming) # verify that it is json loadable
|
||||
streaming = await get_args_streaming(
|
||||
server.get_async_client(), _tools, _messages, _stop
|
||||
)
|
||||
json.loads(streaming)
|
||||
assert non_streaming == streaming, f"{non_streaming=}, {streaming=}"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stop_sequence_interference(server: RemoteOpenAIServer):
|
||||
print("Testing scenario 1")
|
||||
await run_scenario(server, tools, messages, "veroniqueprattyushveroniqueprattyush")
|
||||
|
||||
print("Testing scenario 2")
|
||||
await run_scenario(
|
||||
server, tools2, messages2, "veroniqueprattyushveroniqueprattyush"
|
||||
)
|
||||
|
||||
print("Testing scenario 3")
|
||||
await run_scenario(server, tools2, messages3, "prattyush")
|
||||
@@ -3,29 +3,22 @@
|
||||
|
||||
import json
|
||||
|
||||
import openai
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from huggingface_hub import snapshot_download
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
|
||||
from vllm.tokenizers import TokenizerLike
|
||||
from vllm.tool_parsers.abstract_tool_parser import ToolParser
|
||||
from vllm.tool_parsers.granite4_tool_parser import Granite4ToolParser
|
||||
from vllm.tool_parsers.hermes_tool_parser import Hermes2ProToolParser
|
||||
|
||||
from ....utils import RemoteOpenAIServer
|
||||
|
||||
MODEL_NAME = "meta-llama/Llama-3.2-1B-Instruct"
|
||||
LORA_MODEL = "minpeter/LoRA-Llama-3.2-1B-tool-vllm-ci"
|
||||
|
||||
SERVER_ARGS = [
|
||||
"--enforce-eager",
|
||||
"--enable-auto-tool-choice",
|
||||
"--tool-call-parser",
|
||||
"hermes",
|
||||
"--enable-lora",
|
||||
"--lora-modules",
|
||||
f"{LORA_MODEL}={LORA_MODEL}",
|
||||
"--tokenizer",
|
||||
f"{LORA_MODEL}",
|
||||
]
|
||||
|
||||
TOOLS = [
|
||||
{
|
||||
"type": "function",
|
||||
@@ -50,6 +43,75 @@ TOOLS = [
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
class ServerConfig(TypedDict, total=False):
|
||||
model: str
|
||||
arguments: list[str]
|
||||
model_arg: str
|
||||
tool_parser: ToolParser
|
||||
|
||||
|
||||
CONFIGS: dict[str, ServerConfig] = {
|
||||
"llama": {
|
||||
"model": "meta-llama/Llama-3.2-1B-Instruct",
|
||||
"arguments": [
|
||||
"--enforce-eager",
|
||||
"--enable-auto-tool-choice",
|
||||
"--tool-call-parser",
|
||||
"hermes",
|
||||
"--enable-lora",
|
||||
"--lora-modules",
|
||||
f"{LORA_MODEL}={LORA_MODEL}",
|
||||
"--tokenizer",
|
||||
f"{LORA_MODEL}",
|
||||
],
|
||||
"model_arg": LORA_MODEL,
|
||||
"tool_parser": Hermes2ProToolParser,
|
||||
},
|
||||
"granite4": {
|
||||
"model": "ibm-granite/granite-4.0-h-tiny",
|
||||
"arguments": [
|
||||
"--enforce-eager",
|
||||
"--enable-auto-tool-choice",
|
||||
"--tool-call-parser",
|
||||
"granite4",
|
||||
"--tokenizer",
|
||||
"ibm-granite/granite-4.0-h-tiny",
|
||||
"--max-model-len",
|
||||
"4096",
|
||||
"--max-num-seqs",
|
||||
"2",
|
||||
],
|
||||
"model_arg": "ibm-granite/granite-4.0-h-tiny",
|
||||
"tool_parser": Granite4ToolParser,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
# for each server config, download the model and return the config
|
||||
@pytest.fixture(scope="session", params=CONFIGS.keys())
|
||||
def server_config(request):
|
||||
config = CONFIGS[request.param]
|
||||
|
||||
# download model and tokenizer using transformers
|
||||
snapshot_download(config["model"])
|
||||
yield CONFIGS[request.param]
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def server(request, server_config: ServerConfig):
|
||||
model = server_config["model"]
|
||||
args_for_model = server_config["arguments"]
|
||||
with RemoteOpenAIServer(model, args_for_model, max_wait_seconds=480) as server:
|
||||
yield server
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def client(server: RemoteOpenAIServer):
|
||||
async with server.get_async_client() as async_client:
|
||||
yield async_client
|
||||
|
||||
|
||||
PRODUCT_TOOLS = [
|
||||
{
|
||||
"type": "function",
|
||||
@@ -87,186 +149,182 @@ PRODUCT_MESSAGES = [
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_streaming_tool_call():
|
||||
async def test_non_streaming_tool_call(
|
||||
client: openai.AsyncOpenAI, server_config: ServerConfig
|
||||
):
|
||||
"""Test tool call in non-streaming mode."""
|
||||
with RemoteOpenAIServer(MODEL_NAME, SERVER_ARGS) as server:
|
||||
client = server.get_async_client()
|
||||
|
||||
response = await client.chat.completions.create(
|
||||
model=LORA_MODEL,
|
||||
messages=MESSAGES,
|
||||
tools=TOOLS,
|
||||
tool_choice="auto",
|
||||
temperature=0.0,
|
||||
)
|
||||
response = await client.chat.completions.create(
|
||||
model=server_config["model_arg"],
|
||||
messages=MESSAGES,
|
||||
tools=TOOLS,
|
||||
tool_choice="auto",
|
||||
temperature=0.0,
|
||||
)
|
||||
|
||||
assert response.choices
|
||||
choice = response.choices[0]
|
||||
message = choice.message
|
||||
assert response.choices
|
||||
choice = response.choices[0]
|
||||
message = choice.message
|
||||
|
||||
assert choice.finish_reason == "tool_calls"
|
||||
assert message.tool_calls is not None
|
||||
assert choice.finish_reason == "tool_calls"
|
||||
assert message.tool_calls is not None
|
||||
|
||||
tool_call = message.tool_calls[0]
|
||||
assert tool_call.type == "function"
|
||||
assert tool_call.function.name == "get_current_weather"
|
||||
tool_call = message.tool_calls[0]
|
||||
assert tool_call.type == "function"
|
||||
assert tool_call.function.name == "get_current_weather"
|
||||
|
||||
arguments = json.loads(tool_call.function.arguments)
|
||||
assert "location" in arguments
|
||||
assert "Boston" in arguments["location"]
|
||||
print("\n[Non-Streaming Test Passed]")
|
||||
print(f"Tool Call: {tool_call.function.name}")
|
||||
print(f"Arguments: {arguments}")
|
||||
arguments = json.loads(tool_call.function.arguments)
|
||||
assert "location" in arguments
|
||||
assert "Boston" in arguments["location"]
|
||||
print("\n[Non-Streaming Test Passed]")
|
||||
print(f"Tool Call: {tool_call.function.name}")
|
||||
print(f"Arguments: {arguments}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_tool_call():
|
||||
async def test_streaming_tool_call(
|
||||
client: openai.AsyncOpenAI, server_config: ServerConfig
|
||||
):
|
||||
"""Test tool call in streaming mode."""
|
||||
with RemoteOpenAIServer(MODEL_NAME, SERVER_ARGS) as server:
|
||||
client = server.get_async_client()
|
||||
|
||||
stream = await client.chat.completions.create(
|
||||
model=LORA_MODEL,
|
||||
messages=MESSAGES,
|
||||
tools=TOOLS,
|
||||
tool_choice="auto",
|
||||
temperature=0.0,
|
||||
stream=True,
|
||||
)
|
||||
stream = await client.chat.completions.create(
|
||||
model=server_config["model_arg"],
|
||||
messages=MESSAGES,
|
||||
tools=TOOLS,
|
||||
tool_choice="auto",
|
||||
temperature=0.0,
|
||||
stream=True,
|
||||
)
|
||||
|
||||
tool_call_chunks = {}
|
||||
async for chunk in stream:
|
||||
if not chunk.choices:
|
||||
continue
|
||||
tool_call_chunks = {}
|
||||
async for chunk in stream:
|
||||
if not chunk.choices:
|
||||
continue
|
||||
|
||||
delta = chunk.choices[0].delta
|
||||
if not delta or not delta.tool_calls:
|
||||
continue
|
||||
delta = chunk.choices[0].delta
|
||||
if not delta or not delta.tool_calls:
|
||||
continue
|
||||
|
||||
for tool_chunk in delta.tool_calls:
|
||||
index = tool_chunk.index
|
||||
if index not in tool_call_chunks:
|
||||
tool_call_chunks[index] = {"name": "", "arguments": ""}
|
||||
for tool_chunk in delta.tool_calls:
|
||||
index = tool_chunk.index
|
||||
if index not in tool_call_chunks:
|
||||
tool_call_chunks[index] = {"name": "", "arguments": ""}
|
||||
|
||||
if tool_chunk.function.name:
|
||||
tool_call_chunks[index]["name"] += tool_chunk.function.name
|
||||
if tool_chunk.function.arguments:
|
||||
tool_call_chunks[index]["arguments"] += (
|
||||
tool_chunk.function.arguments
|
||||
)
|
||||
if tool_chunk.function.name:
|
||||
tool_call_chunks[index]["name"] += tool_chunk.function.name
|
||||
if tool_chunk.function.arguments:
|
||||
tool_call_chunks[index]["arguments"] += tool_chunk.function.arguments
|
||||
|
||||
assert len(tool_call_chunks) == 1
|
||||
reconstructed_tool_call = tool_call_chunks[0]
|
||||
assert len(tool_call_chunks) == 1
|
||||
reconstructed_tool_call = tool_call_chunks[0]
|
||||
|
||||
assert reconstructed_tool_call["name"] == "get_current_weather"
|
||||
assert reconstructed_tool_call["name"] == "get_current_weather"
|
||||
|
||||
arguments = json.loads(reconstructed_tool_call["arguments"])
|
||||
assert "location" in arguments
|
||||
assert "Boston" in arguments["location"]
|
||||
print("\n[Streaming Test Passed]")
|
||||
print(f"Reconstructed Tool Call: {reconstructed_tool_call['name']}")
|
||||
print(f"Reconstructed Arguments: {arguments}")
|
||||
arguments = json.loads(reconstructed_tool_call["arguments"])
|
||||
assert "location" in arguments
|
||||
assert "Boston" in arguments["location"]
|
||||
print("\n[Streaming Test Passed]")
|
||||
print(f"Reconstructed Tool Call: {reconstructed_tool_call['name']}")
|
||||
print(f"Reconstructed Arguments: {arguments}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_streaming_product_tool_call():
|
||||
async def test_non_streaming_product_tool_call(
|
||||
client: openai.AsyncOpenAI, server_config: ServerConfig
|
||||
):
|
||||
"""Test tool call integer and boolean parameters in non-streaming mode."""
|
||||
with RemoteOpenAIServer(MODEL_NAME, SERVER_ARGS) as server:
|
||||
client = server.get_async_client()
|
||||
|
||||
response = await client.chat.completions.create(
|
||||
model=LORA_MODEL,
|
||||
messages=PRODUCT_MESSAGES,
|
||||
tools=PRODUCT_TOOLS,
|
||||
tool_choice="auto",
|
||||
temperature=0.66,
|
||||
)
|
||||
response = await client.chat.completions.create(
|
||||
model=server_config["model_arg"],
|
||||
messages=PRODUCT_MESSAGES,
|
||||
tools=PRODUCT_TOOLS,
|
||||
tool_choice="auto",
|
||||
temperature=0.66,
|
||||
)
|
||||
|
||||
assert response.choices
|
||||
choice = response.choices[0]
|
||||
message = choice.message
|
||||
assert response.choices
|
||||
choice = response.choices[0]
|
||||
message = choice.message
|
||||
|
||||
assert choice.finish_reason == "tool_calls"
|
||||
assert message.tool_calls is not None
|
||||
assert choice.finish_reason == "tool_calls"
|
||||
assert message.tool_calls is not None
|
||||
|
||||
tool_call = message.tool_calls[0]
|
||||
assert tool_call.type == "function"
|
||||
assert tool_call.function.name == "get_product_info"
|
||||
tool_call = message.tool_calls[0]
|
||||
assert tool_call.type == "function"
|
||||
assert tool_call.function.name == "get_product_info"
|
||||
|
||||
arguments = json.loads(tool_call.function.arguments)
|
||||
assert "product_id" in arguments
|
||||
assert "inserted" in arguments
|
||||
arguments = json.loads(tool_call.function.arguments)
|
||||
assert "product_id" in arguments
|
||||
assert "inserted" in arguments
|
||||
|
||||
product_id = arguments.get("product_id")
|
||||
inserted = arguments.get("inserted")
|
||||
product_id = arguments.get("product_id")
|
||||
inserted = arguments.get("inserted")
|
||||
|
||||
assert isinstance(product_id, int)
|
||||
assert product_id == 7355608
|
||||
assert isinstance(inserted, bool)
|
||||
assert inserted is True
|
||||
assert isinstance(product_id, int)
|
||||
assert product_id == 7355608
|
||||
assert isinstance(inserted, bool)
|
||||
assert inserted is True
|
||||
|
||||
print("\n[Non-Streaming Product Test Passed]")
|
||||
print(f"Tool Call: {tool_call.function.name}")
|
||||
print(f"Arguments: {arguments}")
|
||||
print("\n[Non-Streaming Product Test Passed]")
|
||||
print(f"Tool Call: {tool_call.function.name}")
|
||||
print(f"Arguments: {arguments}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_product_tool_call():
|
||||
async def test_streaming_product_tool_call(
|
||||
client: openai.AsyncOpenAI, server_config: ServerConfig
|
||||
):
|
||||
"""Test tool call integer and boolean parameters in streaming mode."""
|
||||
with RemoteOpenAIServer(MODEL_NAME, SERVER_ARGS) as server:
|
||||
client = server.get_async_client()
|
||||
|
||||
stream = await client.chat.completions.create(
|
||||
model=LORA_MODEL,
|
||||
messages=PRODUCT_MESSAGES,
|
||||
tools=PRODUCT_TOOLS,
|
||||
tool_choice="auto",
|
||||
temperature=0.66,
|
||||
stream=True,
|
||||
)
|
||||
stream = await client.chat.completions.create(
|
||||
model=server_config["model_arg"],
|
||||
messages=PRODUCT_MESSAGES,
|
||||
tools=PRODUCT_TOOLS,
|
||||
tool_choice="auto",
|
||||
temperature=0.66,
|
||||
stream=True,
|
||||
)
|
||||
|
||||
tool_call_chunks = {}
|
||||
async for chunk in stream:
|
||||
if not chunk.choices:
|
||||
continue
|
||||
tool_call_chunks = {}
|
||||
async for chunk in stream:
|
||||
if not chunk.choices:
|
||||
continue
|
||||
|
||||
delta = chunk.choices[0].delta
|
||||
if not delta or not delta.tool_calls:
|
||||
continue
|
||||
delta = chunk.choices[0].delta
|
||||
if not delta or not delta.tool_calls:
|
||||
continue
|
||||
|
||||
for tool_chunk in delta.tool_calls:
|
||||
index = tool_chunk.index
|
||||
if index not in tool_call_chunks:
|
||||
tool_call_chunks[index] = {"name": "", "arguments": ""}
|
||||
for tool_chunk in delta.tool_calls:
|
||||
index = tool_chunk.index
|
||||
if index not in tool_call_chunks:
|
||||
tool_call_chunks[index] = {"name": "", "arguments": ""}
|
||||
|
||||
if tool_chunk.function.name:
|
||||
tool_call_chunks[index]["name"] += tool_chunk.function.name
|
||||
if tool_chunk.function.arguments:
|
||||
tool_call_chunks[index]["arguments"] += (
|
||||
tool_chunk.function.arguments
|
||||
)
|
||||
if tool_chunk.function.name:
|
||||
tool_call_chunks[index]["name"] += tool_chunk.function.name
|
||||
if tool_chunk.function.arguments:
|
||||
tool_call_chunks[index]["arguments"] += tool_chunk.function.arguments
|
||||
|
||||
assert len(tool_call_chunks) == 1
|
||||
reconstructed_tool_call = tool_call_chunks[0]
|
||||
assert len(tool_call_chunks) == 1
|
||||
reconstructed_tool_call = tool_call_chunks[0]
|
||||
|
||||
assert reconstructed_tool_call["name"] == "get_product_info"
|
||||
assert reconstructed_tool_call["name"] == "get_product_info"
|
||||
|
||||
arguments = json.loads(reconstructed_tool_call["arguments"])
|
||||
assert "product_id" in arguments
|
||||
assert "inserted" in arguments
|
||||
arguments = json.loads(reconstructed_tool_call["arguments"])
|
||||
assert "product_id" in arguments
|
||||
assert "inserted" in arguments
|
||||
|
||||
# Handle type coercion for streaming test as well
|
||||
product_id = arguments.get("product_id")
|
||||
inserted = arguments.get("inserted")
|
||||
# Handle type coercion for streaming test as well
|
||||
product_id = arguments.get("product_id")
|
||||
inserted = arguments.get("inserted")
|
||||
|
||||
assert isinstance(product_id, int)
|
||||
assert product_id == 7355608
|
||||
assert isinstance(inserted, bool)
|
||||
assert inserted is True
|
||||
assert isinstance(product_id, int)
|
||||
assert product_id == 7355608
|
||||
assert isinstance(inserted, bool)
|
||||
assert inserted is True
|
||||
|
||||
print("\n[Streaming Product Test Passed]")
|
||||
print(f"Reconstructed Tool Call: {reconstructed_tool_call['name']}")
|
||||
print(f"Reconstructed Arguments: {arguments}")
|
||||
print("\n[Streaming Product Test Passed]")
|
||||
print(f"Reconstructed Tool Call: {reconstructed_tool_call['name']}")
|
||||
print(f"Reconstructed Arguments: {arguments}")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
@@ -276,9 +334,10 @@ def qwen_tokenizer() -> TokenizerLike:
|
||||
return get_tokenizer("Qwen/Qwen3-32B")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def hermes_parser(qwen_tokenizer: TokenizerLike) -> Hermes2ProToolParser:
|
||||
return Hermes2ProToolParser(qwen_tokenizer)
|
||||
@pytest.fixture(params=CONFIGS.keys())
|
||||
def hermes_parser(request, qwen_tokenizer: TokenizerLike) -> ToolParser:
|
||||
config = CONFIGS[request.param]
|
||||
return config["tool_parser"](qwen_tokenizer)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
@@ -292,7 +351,7 @@ def any_chat_request() -> ChatCompletionRequest:
|
||||
|
||||
def test_hermes_parser_streaming_just_forward_text(
|
||||
qwen_tokenizer: TokenizerLike,
|
||||
hermes_parser: Hermes2ProToolParser,
|
||||
hermes_parser: ToolParser,
|
||||
any_chat_request: ChatCompletionRequest,
|
||||
) -> None:
|
||||
text = """This is some prior text that has nothing to do with tool calling."""
|
||||
@@ -324,7 +383,7 @@ def test_hermes_parser_streaming_just_forward_text(
|
||||
|
||||
def test_hermes_parser_streaming_failure_case_bug_19056(
|
||||
qwen_tokenizer: TokenizerLike,
|
||||
hermes_parser: Hermes2ProToolParser,
|
||||
hermes_parser: ToolParser,
|
||||
any_chat_request: ChatCompletionRequest,
|
||||
) -> None:
|
||||
text = """<tool_call>
|
||||
@@ -358,7 +417,7 @@ def test_hermes_parser_streaming_failure_case_bug_19056(
|
||||
|
||||
def test_hermes_parser_streaming(
|
||||
qwen_tokenizer: TokenizerLike,
|
||||
hermes_parser: Hermes2ProToolParser,
|
||||
hermes_parser: ToolParser,
|
||||
any_chat_request: ChatCompletionRequest,
|
||||
) -> None:
|
||||
text = '<tool_call>\
|
||||
@@ -387,16 +446,20 @@ def test_hermes_parser_streaming(
|
||||
delta_messages.append(delta)
|
||||
print(delta_messages)
|
||||
assert delta_messages[0].tool_calls[0].function.name == "get_current_temperature"
|
||||
tool_call_args = "".join(
|
||||
delta.tool_calls[0].function.arguments or "" for delta in delta_messages
|
||||
)
|
||||
assert tool_call_args == (
|
||||
'{"location":"San Francisco, California, United States", "unit": "celsius"}'
|
||||
# load to normalize whitespace
|
||||
tool_call_args = json.loads(
|
||||
"".join(
|
||||
delta.tool_calls[0].function.arguments or "" for delta in delta_messages
|
||||
)
|
||||
)
|
||||
assert tool_call_args == {
|
||||
"location": "San Francisco, California, United States",
|
||||
"unit": "celsius",
|
||||
}
|
||||
|
||||
|
||||
def test_hermes_parser_non_streaming_no_tool_call(
|
||||
hermes_parser: Hermes2ProToolParser,
|
||||
hermes_parser: ToolParser,
|
||||
any_chat_request: ChatCompletionRequest,
|
||||
) -> None:
|
||||
text = """This is not a tool call."""
|
||||
@@ -410,7 +473,7 @@ def test_hermes_parser_non_streaming_no_tool_call(
|
||||
|
||||
|
||||
def test_hermes_parser_non_streaming_tool_call_between_tags(
|
||||
hermes_parser: Hermes2ProToolParser,
|
||||
hermes_parser: ToolParser,
|
||||
any_chat_request: ChatCompletionRequest,
|
||||
) -> None:
|
||||
text = """<tool_call>
|
||||
@@ -428,9 +491,12 @@ def test_hermes_parser_non_streaming_tool_call_between_tags(
|
||||
|
||||
|
||||
def test_hermes_parser_non_streaming_tool_call_until_eos(
|
||||
hermes_parser: Hermes2ProToolParser,
|
||||
hermes_parser: ToolParser,
|
||||
any_chat_request: ChatCompletionRequest,
|
||||
) -> None:
|
||||
if isinstance(hermes_parser, Granite4ToolParser):
|
||||
pytest.skip(reason="The Granite4 tool parser enforces a complete response")
|
||||
|
||||
text = """<tool_call>
|
||||
{"name": "final_answer", "arguments": {"trigger": true}}"""
|
||||
tool_call = hermes_parser.extract_tool_calls(
|
||||
@@ -445,7 +511,7 @@ def test_hermes_parser_non_streaming_tool_call_until_eos(
|
||||
|
||||
|
||||
def test_hermes_parser_non_streaming_tool_call_invalid_json(
|
||||
hermes_parser: Hermes2ProToolParser,
|
||||
hermes_parser: ToolParser,
|
||||
any_chat_request: ChatCompletionRequest,
|
||||
) -> None:
|
||||
# Missing closing brace to trigger exception
|
||||
|
||||
@@ -7,7 +7,7 @@ from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.entrypoints.openai.tool_parsers.utils import (
|
||||
from tests.tool_parsers.utils import (
|
||||
run_tool_extraction,
|
||||
run_tool_extraction_streaming,
|
||||
)
|
||||
|
||||
@@ -5,7 +5,7 @@ from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.entrypoints.openai.tool_parsers.utils import (
|
||||
from tests.tool_parsers.utils import (
|
||||
run_tool_extraction,
|
||||
run_tool_extraction_streaming,
|
||||
)
|
||||
|
||||
@@ -5,7 +5,7 @@ from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.entrypoints.openai.tool_parsers.utils import (
|
||||
from tests.tool_parsers.utils import (
|
||||
run_tool_extraction,
|
||||
run_tool_extraction_streaming,
|
||||
)
|
||||
|
||||
@@ -5,7 +5,7 @@ from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.entrypoints.openai.tool_parsers.utils import (
|
||||
from tests.tool_parsers.utils import (
|
||||
run_tool_extraction,
|
||||
run_tool_extraction_streaming,
|
||||
)
|
||||
|
||||
@@ -0,0 +1,310 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for the Cohere /v2/embed API with generic (non-Cohere) models.
|
||||
|
||||
Validates that the Cohere v2 embed endpoint works correctly with standard
|
||||
embedding models, covering text embedding, embedding type conversions,
|
||||
response structure, batching, normalisation, and semantic similarity.
|
||||
"""
|
||||
|
||||
import base64
|
||||
import struct
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import requests
|
||||
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
|
||||
DTYPE = "bfloat16"
|
||||
|
||||
MODELS: list[tuple[str, list[str]]] = [
|
||||
("intfloat/multilingual-e5-small", []),
|
||||
(
|
||||
"Snowflake/snowflake-arctic-embed-m-v1.5",
|
||||
[
|
||||
"--trust_remote_code",
|
||||
"--hf_overrides",
|
||||
'{"matryoshka_dimensions":[256]}',
|
||||
],
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture(scope="module", params=MODELS, ids=lambda m: m[0])
|
||||
def model_config(request):
|
||||
return request.param
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def model_name(model_config):
|
||||
return model_config[0]
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def server(model_config):
|
||||
name, extra_args = model_config
|
||||
args = [
|
||||
"--runner",
|
||||
"pooling",
|
||||
"--dtype",
|
||||
DTYPE,
|
||||
"--enforce-eager",
|
||||
"--max-model-len",
|
||||
"512",
|
||||
"--gpu-memory-utilization",
|
||||
"0.02",
|
||||
] + extra_args
|
||||
with RemoteOpenAIServer(name, args) as remote_server:
|
||||
yield remote_server
|
||||
|
||||
|
||||
def _cohere_embed(
|
||||
server: RemoteOpenAIServer,
|
||||
model_name: str,
|
||||
texts: list[str] | None = None,
|
||||
images: list[str] | None = None,
|
||||
input_type: str | None = None,
|
||||
embedding_types: list[str] | None = None,
|
||||
) -> dict:
|
||||
body: dict = {"model": model_name}
|
||||
if input_type is not None:
|
||||
body["input_type"] = input_type
|
||||
if texts is not None:
|
||||
body["texts"] = texts
|
||||
if images is not None:
|
||||
body["images"] = images
|
||||
if embedding_types is not None:
|
||||
body["embedding_types"] = embedding_types
|
||||
resp = requests.post(server.url_for("/v2/embed"), json=body)
|
||||
resp.raise_for_status()
|
||||
return resp.json()
|
||||
|
||||
|
||||
def _openai_embed(
|
||||
server: RemoteOpenAIServer, model_name: str, texts: list[str]
|
||||
) -> dict:
|
||||
body = {"model": model_name, "input": texts, "encoding_format": "float"}
|
||||
resp = requests.post(server.url_for("/v1/embeddings"), json=body)
|
||||
resp.raise_for_status()
|
||||
return resp.json()
|
||||
|
||||
|
||||
def _cosine_sim(a: list[float], b: list[float]) -> float:
|
||||
va, vb = np.array(a), np.array(b)
|
||||
return float(np.dot(va, vb) / (np.linalg.norm(va) * np.linalg.norm(vb)))
|
||||
|
||||
|
||||
# -----------------------------------------------------------
|
||||
# Text embedding tests
|
||||
# -----------------------------------------------------------
|
||||
|
||||
|
||||
def test_basic_embed(server: RemoteOpenAIServer, model_name: str):
|
||||
r = _cohere_embed(
|
||||
server, model_name, texts=["hello world"], embedding_types=["float"]
|
||||
)
|
||||
assert "embeddings" in r
|
||||
assert len(r["embeddings"]["float"]) == 1
|
||||
assert len(r["embeddings"]["float"][0]) > 0
|
||||
|
||||
|
||||
def test_unsupported_input_type_rejected(server: RemoteOpenAIServer, model_name: str):
|
||||
"""An input_type not defined in the model's prompt config should be
|
||||
rejected with a 400 error."""
|
||||
body = {
|
||||
"model": model_name,
|
||||
"input_type": "nonexistent_type",
|
||||
"texts": ["hello world"],
|
||||
"embedding_types": ["float"],
|
||||
}
|
||||
resp = requests.post(server.url_for("/v2/embed"), json=body)
|
||||
assert resp.status_code == 400
|
||||
assert "Unsupported input_type" in resp.json()["error"]["message"]
|
||||
|
||||
|
||||
def test_omitted_input_type_accepted(server: RemoteOpenAIServer, model_name: str):
|
||||
"""Omitting input_type should always work (no prompt prefix applied)."""
|
||||
body = {
|
||||
"model": model_name,
|
||||
"texts": ["hello world"],
|
||||
"embedding_types": ["float"],
|
||||
}
|
||||
resp = requests.post(server.url_for("/v2/embed"), json=body)
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert len(data["embeddings"]["float"]) == 1
|
||||
|
||||
|
||||
def test_v1_v2_parity(server: RemoteOpenAIServer, model_name: str):
|
||||
"""v1 (OpenAI) and v2 (Cohere) endpoints should produce the same
|
||||
float embeddings for a generic model."""
|
||||
texts = ["hello world"]
|
||||
v2 = _cohere_embed(server, model_name, texts=texts, embedding_types=["float"])
|
||||
v1 = _openai_embed(server, model_name, texts)
|
||||
cos = _cosine_sim(v2["embeddings"]["float"][0], v1["data"][0]["embedding"])
|
||||
assert cos > 0.9999, f"v1/v2 parity failed, cosine={cos}"
|
||||
|
||||
|
||||
def test_embedding_types(server: RemoteOpenAIServer, model_name: str):
|
||||
r = _cohere_embed(
|
||||
server,
|
||||
model_name,
|
||||
texts=["test"],
|
||||
embedding_types=["float", "binary", "ubinary"],
|
||||
)
|
||||
dim = len(r["embeddings"]["float"][0])
|
||||
assert len(r["embeddings"]["binary"][0]) == dim // 8
|
||||
assert len(r["embeddings"]["ubinary"][0]) == dim // 8
|
||||
|
||||
|
||||
def test_response_structure(server: RemoteOpenAIServer, model_name: str):
|
||||
r = _cohere_embed(server, model_name, texts=["test"], embedding_types=["float"])
|
||||
assert "id" in r
|
||||
assert "embeddings" in r
|
||||
assert "texts" in r
|
||||
assert r["texts"] == ["test"]
|
||||
assert "meta" in r
|
||||
assert r["meta"]["api_version"]["version"] == "2"
|
||||
assert "billed_units" in r["meta"]
|
||||
assert r["meta"]["billed_units"]["input_tokens"] > 0
|
||||
assert r["meta"]["billed_units"]["image_tokens"] == 0
|
||||
|
||||
|
||||
def test_batch(server: RemoteOpenAIServer, model_name: str):
|
||||
texts = ["apple", "banana", "cherry"]
|
||||
r = _cohere_embed(server, model_name, texts=texts, embedding_types=["float"])
|
||||
assert len(r["embeddings"]["float"]) == 3
|
||||
dim = len(r["embeddings"]["float"][0])
|
||||
for emb in r["embeddings"]["float"]:
|
||||
assert len(emb) == dim
|
||||
|
||||
|
||||
def test_l2_normalized(server: RemoteOpenAIServer, model_name: str):
|
||||
r = _cohere_embed(
|
||||
server, model_name, texts=["hello world"], embedding_types=["float"]
|
||||
)
|
||||
emb = np.array(r["embeddings"]["float"][0])
|
||||
assert abs(float(np.linalg.norm(emb)) - 1.0) < 0.01
|
||||
|
||||
|
||||
def test_semantic_similarity(server: RemoteOpenAIServer, model_name: str):
|
||||
r = _cohere_embed(
|
||||
server,
|
||||
model_name,
|
||||
texts=["machine learning", "deep learning", "chocolate cake recipe"],
|
||||
embedding_types=["float"],
|
||||
)
|
||||
embs = r["embeddings"]["float"]
|
||||
cos_related = _cosine_sim(embs[0], embs[1])
|
||||
cos_unrelated = _cosine_sim(embs[0], embs[2])
|
||||
assert cos_related > cos_unrelated
|
||||
|
||||
|
||||
def test_missing_input_returns_error(server: RemoteOpenAIServer, model_name: str):
|
||||
body = {"model": model_name}
|
||||
resp = requests.post(server.url_for("/v2/embed"), json=body)
|
||||
assert resp.status_code == 400
|
||||
|
||||
|
||||
def test_base64_embedding_type(server: RemoteOpenAIServer, model_name: str):
|
||||
r = _cohere_embed(
|
||||
server,
|
||||
model_name,
|
||||
texts=["test encoding"],
|
||||
embedding_types=["float", "base64"],
|
||||
)
|
||||
float_emb = r["embeddings"]["float"][0]
|
||||
b64_str = r["embeddings"]["base64"][0]
|
||||
decoded = struct.unpack(f"<{len(float_emb)}f", base64.b64decode(b64_str))
|
||||
np.testing.assert_allclose(float_emb, decoded, rtol=1e-5)
|
||||
|
||||
|
||||
# -----------------------------------------------------------
|
||||
# Truncation tests
|
||||
# -----------------------------------------------------------
|
||||
|
||||
|
||||
def _cohere_embed_raw(
|
||||
server: RemoteOpenAIServer,
|
||||
body: dict,
|
||||
) -> requests.Response:
|
||||
return requests.post(server.url_for("/v2/embed"), json=body)
|
||||
|
||||
|
||||
def test_truncate_end_succeeds(server: RemoteOpenAIServer, model_name: str):
|
||||
"""truncate=END should silently truncate long input."""
|
||||
long_text = " ".join(["word"] * 2000)
|
||||
body = {
|
||||
"model": model_name,
|
||||
"texts": [long_text],
|
||||
"embedding_types": ["float"],
|
||||
"truncate": "END",
|
||||
}
|
||||
resp = _cohere_embed_raw(server, body)
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert len(data["embeddings"]["float"]) == 1
|
||||
|
||||
|
||||
def test_truncate_start_succeeds(server: RemoteOpenAIServer, model_name: str):
|
||||
"""truncate=START should silently truncate long input from the start."""
|
||||
long_text = " ".join(["word"] * 2000)
|
||||
body = {
|
||||
"model": model_name,
|
||||
"texts": [long_text],
|
||||
"embedding_types": ["float"],
|
||||
"truncate": "START",
|
||||
}
|
||||
resp = _cohere_embed_raw(server, body)
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert len(data["embeddings"]["float"]) == 1
|
||||
|
||||
|
||||
def test_truncate_none_rejects_long_input(server: RemoteOpenAIServer, model_name: str):
|
||||
"""truncate=NONE should error when input exceeds model context."""
|
||||
long_text = " ".join(["word"] * 2000)
|
||||
body = {
|
||||
"model": model_name,
|
||||
"texts": [long_text],
|
||||
"embedding_types": ["float"],
|
||||
"truncate": "NONE",
|
||||
}
|
||||
resp = _cohere_embed_raw(server, body)
|
||||
assert resp.status_code == 400
|
||||
|
||||
|
||||
def test_truncate_start_vs_end_differ(server: RemoteOpenAIServer, model_name: str):
|
||||
"""START and END truncation should produce different embeddings
|
||||
when the input is long enough to actually be truncated.
|
||||
|
||||
We construct input with distinct tokens at the start vs end
|
||||
so that keeping different halves produces different embeddings.
|
||||
"""
|
||||
start_words = " ".join([f"alpha{i}" for i in range(300)])
|
||||
end_words = " ".join([f"omega{i}" for i in range(300)])
|
||||
long_text = start_words + " " + end_words
|
||||
|
||||
body_end = {
|
||||
"model": model_name,
|
||||
"texts": [long_text],
|
||||
"embedding_types": ["float"],
|
||||
"truncate": "END",
|
||||
}
|
||||
body_start = {
|
||||
"model": model_name,
|
||||
"texts": [long_text],
|
||||
"embedding_types": ["float"],
|
||||
"truncate": "START",
|
||||
}
|
||||
r_end = _cohere_embed_raw(server, body_end).json()
|
||||
r_start = _cohere_embed_raw(server, body_start).json()
|
||||
|
||||
emb_end = r_end["embeddings"]["float"][0]
|
||||
emb_start = r_start["embeddings"]["float"][0]
|
||||
cos = _cosine_sim(emb_end, emb_start)
|
||||
assert cos < 0.99, (
|
||||
f"START and END truncation should produce different embeddings "
|
||||
f"for long input, but cosine similarity was {cos}"
|
||||
)
|
||||
@@ -0,0 +1,135 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for the Cohere /v2/embed API with a multimodal model (SigLIP).
|
||||
|
||||
Validates image embedding, batching, normalisation, and embedding type
|
||||
conversions through the /v2/embed endpoint.
|
||||
"""
|
||||
|
||||
import base64
|
||||
import struct
|
||||
import zlib
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import requests
|
||||
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
|
||||
MODEL_NAME = "google/siglip-so400m-patch14-384"
|
||||
DTYPE = "bfloat16"
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def server():
|
||||
args = [
|
||||
"--runner",
|
||||
"pooling",
|
||||
"--dtype",
|
||||
DTYPE,
|
||||
"--enforce-eager",
|
||||
"--max-model-len",
|
||||
"64",
|
||||
"--gpu-memory-utilization",
|
||||
"0.3",
|
||||
]
|
||||
with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
|
||||
yield remote_server
|
||||
|
||||
|
||||
def _make_tiny_png(r: int, g: int, b: int, w: int = 2, h: int = 2) -> str:
|
||||
raw = b""
|
||||
for _ in range(h):
|
||||
raw += b"\x00" + bytes([r, g, b]) * w
|
||||
compressed = zlib.compress(raw)
|
||||
|
||||
def chunk(ctype: bytes, cdata: bytes) -> bytes:
|
||||
c = ctype + cdata
|
||||
return (
|
||||
struct.pack(">I", len(cdata))
|
||||
+ c
|
||||
+ struct.pack(">I", zlib.crc32(c) & 0xFFFFFFFF)
|
||||
)
|
||||
|
||||
ihdr = struct.pack(">IIBBBBB", w, h, 8, 2, 0, 0, 0)
|
||||
png = (
|
||||
b"\x89PNG\r\n\x1a\n"
|
||||
+ chunk(b"IHDR", ihdr)
|
||||
+ chunk(b"IDAT", compressed)
|
||||
+ chunk(b"IEND", b"")
|
||||
)
|
||||
return "data:image/png;base64," + base64.b64encode(png).decode()
|
||||
|
||||
|
||||
def _cohere_embed(
|
||||
server: RemoteOpenAIServer,
|
||||
texts: list[str] | None = None,
|
||||
images: list[str] | None = None,
|
||||
embedding_types: list[str] | None = None,
|
||||
) -> dict:
|
||||
body: dict = {"model": MODEL_NAME}
|
||||
if texts is not None:
|
||||
body["texts"] = texts
|
||||
if images is not None:
|
||||
body["images"] = images
|
||||
if embedding_types is not None:
|
||||
body["embedding_types"] = embedding_types
|
||||
resp = requests.post(server.url_for("/v2/embed"), json=body)
|
||||
resp.raise_for_status()
|
||||
return resp.json()
|
||||
|
||||
|
||||
def test_image_embed(server: RemoteOpenAIServer):
|
||||
img_uri = _make_tiny_png(255, 0, 0)
|
||||
r = _cohere_embed(
|
||||
server,
|
||||
images=[img_uri],
|
||||
embedding_types=["float"],
|
||||
)
|
||||
assert "embeddings" in r
|
||||
assert len(r["embeddings"]["float"]) == 1
|
||||
assert len(r["embeddings"]["float"][0]) > 0
|
||||
assert r["meta"]["billed_units"]["image_tokens"] > 0
|
||||
assert r["meta"]["billed_units"]["input_tokens"] == 0
|
||||
|
||||
|
||||
def test_image_batch(server: RemoteOpenAIServer):
|
||||
red = _make_tiny_png(255, 0, 0)
|
||||
blue = _make_tiny_png(0, 0, 255)
|
||||
r = _cohere_embed(
|
||||
server,
|
||||
images=[red, blue],
|
||||
embedding_types=["float"],
|
||||
)
|
||||
assert len(r["embeddings"]["float"]) == 2
|
||||
|
||||
|
||||
def test_image_l2_normalized(server: RemoteOpenAIServer):
|
||||
img_uri = _make_tiny_png(0, 255, 0)
|
||||
r = _cohere_embed(
|
||||
server,
|
||||
images=[img_uri],
|
||||
embedding_types=["float"],
|
||||
)
|
||||
emb = np.array(r["embeddings"]["float"][0])
|
||||
assert abs(float(np.linalg.norm(emb)) - 1.0) < 0.01
|
||||
|
||||
|
||||
def test_image_embedding_types(server: RemoteOpenAIServer):
|
||||
img_uri = _make_tiny_png(128, 128, 128)
|
||||
r = _cohere_embed(
|
||||
server,
|
||||
images=[img_uri],
|
||||
embedding_types=["float", "binary", "ubinary"],
|
||||
)
|
||||
dim = len(r["embeddings"]["float"][0])
|
||||
assert len(r["embeddings"]["binary"][0]) == dim // 8
|
||||
assert len(r["embeddings"]["ubinary"][0]) == dim // 8
|
||||
|
||||
|
||||
def test_text_embed_on_multimodal(server: RemoteOpenAIServer):
|
||||
"""SigLIP also supports text-only embedding via /v2/embed."""
|
||||
r = _cohere_embed(server, texts=["hello world"], embedding_types=["float"])
|
||||
assert "embeddings" in r
|
||||
assert len(r["embeddings"]["float"]) == 1
|
||||
assert len(r["embeddings"]["float"][0]) > 0
|
||||
@@ -0,0 +1,102 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Parity test between Cohere /v2/embed and OpenAI /v1/embeddings.
|
||||
|
||||
Verifies that both endpoints produce identical float embeddings when
|
||||
no prompt prefix is applied (input_type omitted for Cohere /v2/embed).
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import requests
|
||||
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
|
||||
MODEL_NAME = "BAAI/bge-base-en-v1.5"
|
||||
DTYPE = "bfloat16"
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def server():
|
||||
args = [
|
||||
"--runner",
|
||||
"pooling",
|
||||
"--dtype",
|
||||
DTYPE,
|
||||
"--enforce-eager",
|
||||
"--max-model-len",
|
||||
"512",
|
||||
"--gpu-memory-utilization",
|
||||
"0.02",
|
||||
]
|
||||
with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
|
||||
yield remote_server
|
||||
|
||||
|
||||
def _cohere_embed(
|
||||
server: RemoteOpenAIServer,
|
||||
texts: list[str],
|
||||
) -> list[list[float]]:
|
||||
body = {
|
||||
"model": MODEL_NAME,
|
||||
"texts": texts,
|
||||
"embedding_types": ["float"],
|
||||
}
|
||||
resp = requests.post(server.url_for("/v2/embed"), json=body)
|
||||
resp.raise_for_status()
|
||||
return resp.json()["embeddings"]["float"]
|
||||
|
||||
|
||||
def _openai_embed(
|
||||
server: RemoteOpenAIServer,
|
||||
texts: list[str],
|
||||
) -> list[list[float]]:
|
||||
body = {"model": MODEL_NAME, "input": texts, "encoding_format": "float"}
|
||||
resp = requests.post(server.url_for("/v1/embeddings"), json=body)
|
||||
resp.raise_for_status()
|
||||
return [item["embedding"] for item in resp.json()["data"]]
|
||||
|
||||
|
||||
def test_single_text_parity(server: RemoteOpenAIServer):
|
||||
"""A single text should produce identical embeddings via both APIs."""
|
||||
texts = ["the quick brown fox jumps over the lazy dog"]
|
||||
v2 = _cohere_embed(server, texts)
|
||||
v1 = _openai_embed(server, texts)
|
||||
np.testing.assert_allclose(v2[0], v1[0], rtol=1e-5)
|
||||
|
||||
|
||||
def test_batch_parity(server: RemoteOpenAIServer):
|
||||
"""A batch of texts should produce identical embeddings via both APIs,
|
||||
in the same order."""
|
||||
texts = [
|
||||
"machine learning",
|
||||
"deep learning",
|
||||
"natural language processing",
|
||||
]
|
||||
v2 = _cohere_embed(server, texts)
|
||||
v1 = _openai_embed(server, texts)
|
||||
assert len(v2) == len(v1) == 3
|
||||
for i in range(3):
|
||||
np.testing.assert_allclose(v2[i], v1[i], rtol=1e-5, err_msg=f"index {i}")
|
||||
|
||||
|
||||
def test_token_count_parity(server: RemoteOpenAIServer):
|
||||
"""Both APIs should report the same prompt token count."""
|
||||
texts = ["hello world"]
|
||||
v2_resp = requests.post(
|
||||
server.url_for("/v2/embed"),
|
||||
json={
|
||||
"model": MODEL_NAME,
|
||||
"texts": texts,
|
||||
"embedding_types": ["float"],
|
||||
},
|
||||
)
|
||||
v1_resp = requests.post(
|
||||
server.url_for("/v1/embeddings"),
|
||||
json={"model": MODEL_NAME, "input": texts, "encoding_format": "float"},
|
||||
)
|
||||
v2_resp.raise_for_status()
|
||||
v1_resp.raise_for_status()
|
||||
v2_tokens = v2_resp.json()["meta"]["billed_units"]["input_tokens"]
|
||||
v1_tokens = v1_resp.json()["usage"]["prompt_tokens"]
|
||||
assert v2_tokens == v1_tokens
|
||||
@@ -0,0 +1,208 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for EmbedIOProcessor."""
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.entrypoints.pooling.embed.io_processor import EmbedIOProcessor
|
||||
from vllm.entrypoints.pooling.embed.protocol import (
|
||||
CohereEmbedRequest,
|
||||
)
|
||||
|
||||
|
||||
class TestResolveTruncation:
|
||||
"""Unit tests for EmbedIOProcessor._resolve_cohere_truncation."""
|
||||
|
||||
@staticmethod
|
||||
def _make_request(**kwargs) -> CohereEmbedRequest:
|
||||
defaults = {
|
||||
"model": "test",
|
||||
"input_type": "search_document",
|
||||
"texts": ["hello"],
|
||||
}
|
||||
return CohereEmbedRequest(**(defaults | kwargs))
|
||||
|
||||
def test_truncate_end_default(self):
|
||||
req = self._make_request()
|
||||
tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
|
||||
assert tokens == -1
|
||||
assert side is None
|
||||
|
||||
def test_truncate_end_explicit(self):
|
||||
req = self._make_request(truncate="END")
|
||||
tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
|
||||
assert tokens == -1
|
||||
assert side is None
|
||||
|
||||
def test_truncate_end_with_max_tokens(self):
|
||||
req = self._make_request(truncate="END", max_tokens=128)
|
||||
tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
|
||||
assert tokens == 128
|
||||
assert side is None
|
||||
|
||||
def test_truncate_none(self):
|
||||
req = self._make_request(truncate="NONE")
|
||||
tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
|
||||
assert tokens is None
|
||||
assert side is None
|
||||
|
||||
def test_truncate_none_with_max_tokens(self):
|
||||
"""truncate=NONE should NOT set truncate_prompt_tokens; the
|
||||
max_tokens limit is enforced separately via _check_max_tokens."""
|
||||
req = self._make_request(truncate="NONE", max_tokens=10)
|
||||
tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
|
||||
assert tokens is None
|
||||
assert side is None
|
||||
|
||||
def test_truncate_start(self):
|
||||
req = self._make_request(truncate="START")
|
||||
tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
|
||||
assert tokens == -1
|
||||
assert side == "left"
|
||||
|
||||
def test_truncate_start_with_max_tokens(self):
|
||||
req = self._make_request(truncate="START", max_tokens=64)
|
||||
tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
|
||||
assert tokens == 64
|
||||
assert side == "left"
|
||||
|
||||
|
||||
class TestApplyStPrompt:
|
||||
"""Unit tests for EmbedIOProcessor._apply_task_instruction."""
|
||||
|
||||
@staticmethod
|
||||
def _make_handler(task_instructions: dict[str, str] | None):
|
||||
handler = object.__new__(EmbedIOProcessor)
|
||||
handler.task_instructions = task_instructions
|
||||
return handler
|
||||
|
||||
def test_no_prompts_configured(self):
|
||||
handler = self._make_handler(None)
|
||||
texts = ["hello", "world"]
|
||||
assert handler._apply_task_instruction(texts, "query") is texts
|
||||
|
||||
def test_matching_input_type(self):
|
||||
handler = self._make_handler({"query": "search_query: "})
|
||||
result = handler._apply_task_instruction(["hello"], "query")
|
||||
assert result == ["search_query: hello"]
|
||||
|
||||
def test_non_matching_input_type(self):
|
||||
handler = self._make_handler({"query": "search_query: "})
|
||||
texts = ["hello"]
|
||||
assert handler._apply_task_instruction(texts, "document") is texts
|
||||
|
||||
def test_multiple_texts(self):
|
||||
handler = self._make_handler(
|
||||
{"query": "Represent this sentence for searching: "}
|
||||
)
|
||||
result = handler._apply_task_instruction(["a", "b", "c"], "query")
|
||||
assert result == [
|
||||
"Represent this sentence for searching: a",
|
||||
"Represent this sentence for searching: b",
|
||||
"Represent this sentence for searching: c",
|
||||
]
|
||||
|
||||
def test_empty_prefix_returns_unchanged(self):
|
||||
handler = self._make_handler({"passage": ""})
|
||||
texts = ["hello"]
|
||||
assert handler._apply_task_instruction(texts, "passage") is texts
|
||||
|
||||
|
||||
class TestLoadTaskInstructions:
|
||||
"""Unit tests for EmbedIOProcessor._load_task_instructions."""
|
||||
|
||||
def test_no_attribute(self):
|
||||
class FakeConfig:
|
||||
pass
|
||||
|
||||
assert EmbedIOProcessor._load_task_instructions(FakeConfig()) is None
|
||||
|
||||
def test_with_task_instructions(self):
|
||||
class FakeConfig:
|
||||
task_instructions = {
|
||||
"retrieval.query": "Represent the query: ",
|
||||
"retrieval.passage": "",
|
||||
}
|
||||
|
||||
result = EmbedIOProcessor._load_task_instructions(FakeConfig())
|
||||
assert result == {
|
||||
"retrieval.query": "Represent the query: ",
|
||||
"retrieval.passage": "",
|
||||
}
|
||||
|
||||
def test_empty_dict(self):
|
||||
class FakeConfig:
|
||||
task_instructions = {}
|
||||
|
||||
assert EmbedIOProcessor._load_task_instructions(FakeConfig()) is None
|
||||
|
||||
def test_non_dict(self):
|
||||
class FakeConfig:
|
||||
task_instructions = "not a dict"
|
||||
|
||||
assert EmbedIOProcessor._load_task_instructions(FakeConfig()) is None
|
||||
|
||||
|
||||
class TestCheckMaxTokens:
|
||||
"""Unit tests for EmbedIOProcessor._check_cohere_max_tokens."""
|
||||
|
||||
@staticmethod
|
||||
def _fake_output(n_tokens: int):
|
||||
class _Out:
|
||||
def __init__(self, n: int):
|
||||
self.prompt_token_ids = list(range(n))
|
||||
|
||||
return _Out(n_tokens)
|
||||
|
||||
def test_none_check_is_noop(self):
|
||||
outs = [self._fake_output(100)]
|
||||
EmbedIOProcessor._check_cohere_max_tokens(outs, None)
|
||||
|
||||
def test_within_limit(self):
|
||||
outs = [self._fake_output(5), self._fake_output(3)]
|
||||
EmbedIOProcessor._check_cohere_max_tokens(outs, 5)
|
||||
|
||||
def test_exceeds_limit(self):
|
||||
outs = [self._fake_output(3), self._fake_output(10)]
|
||||
with pytest.raises(ValueError, match="exceeds max_tokens=5"):
|
||||
EmbedIOProcessor._check_cohere_max_tokens(outs, 5)
|
||||
|
||||
def test_exact_limit(self):
|
||||
outs = [self._fake_output(5)]
|
||||
EmbedIOProcessor._check_cohere_max_tokens(outs, 5)
|
||||
|
||||
|
||||
class TestValidateInputType:
|
||||
"""Unit tests for EmbedIOProcessor._validate_input_type."""
|
||||
|
||||
@staticmethod
|
||||
def _make_handler(task_instructions: dict[str, str] | None):
|
||||
handler = object.__new__(EmbedIOProcessor)
|
||||
handler.task_instructions = task_instructions
|
||||
return handler
|
||||
|
||||
def test_none_input_type_always_accepted(self):
|
||||
handler = self._make_handler(None)
|
||||
handler._validate_input_type(None)
|
||||
handler_with = self._make_handler({"query": "q: "})
|
||||
handler_with._validate_input_type(None)
|
||||
|
||||
def test_no_prompts_rejects(self):
|
||||
handler = self._make_handler(None)
|
||||
with pytest.raises(ValueError, match="does not define any input_type"):
|
||||
handler._validate_input_type("anything")
|
||||
|
||||
def test_known_type_accepted(self):
|
||||
handler = self._make_handler({"query": "q: ", "document": "d: "})
|
||||
handler._validate_input_type("query")
|
||||
handler._validate_input_type("document")
|
||||
|
||||
def test_unknown_type_rejected(self):
|
||||
handler = self._make_handler({"query": "q: ", "document": "d: "})
|
||||
with pytest.raises(ValueError, match="Unsupported input_type 'other'"):
|
||||
handler._validate_input_type("other")
|
||||
|
||||
def test_error_lists_supported(self):
|
||||
handler = self._make_handler({"a": "", "b": ""})
|
||||
with pytest.raises(ValueError, match="Supported values: a, b"):
|
||||
handler._validate_input_type("z")
|
||||
@@ -0,0 +1,129 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for Cohere embed protocol: build_typed_embeddings and its
|
||||
underlying packing helpers, plus Cohere-specific serving helpers."""
|
||||
|
||||
import base64
|
||||
import struct
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from vllm.entrypoints.pooling.embed.protocol import (
|
||||
build_typed_embeddings,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_embeddings() -> list[list[float]]:
|
||||
return [
|
||||
[0.1, -0.2, 0.3, -0.4, 0.5, -0.6, 0.7, -0.8],
|
||||
[-0.05, 0.15, -0.25, 0.35, -0.45, 0.55, -0.65, 0.75],
|
||||
]
|
||||
|
||||
|
||||
class TestBuildTypedEmbeddingsFloat:
|
||||
def test_float_passthrough(self, sample_embeddings: list[list[float]]):
|
||||
result = build_typed_embeddings(sample_embeddings, ["float"])
|
||||
assert result.float == sample_embeddings
|
||||
assert result.binary is None
|
||||
|
||||
def test_empty_input(self):
|
||||
result = build_typed_embeddings([], ["float"])
|
||||
assert result.float == []
|
||||
|
||||
|
||||
class TestBuildTypedEmbeddingsBinary:
|
||||
def test_binary_packing(self):
|
||||
# 8 values: positive->1, negative->0 => bits: 10101010 = 0xAA = 170
|
||||
# signed: 170 - 128 = 42
|
||||
embs = [[1.0, -1.0, 1.0, -1.0, 1.0, -1.0, 1.0, -1.0]]
|
||||
result = build_typed_embeddings(embs, ["binary"])
|
||||
assert result.binary is not None
|
||||
assert result.binary[0] == [42]
|
||||
|
||||
def test_ubinary_packing(self):
|
||||
embs = [[1.0, -1.0, 1.0, -1.0, 1.0, -1.0, 1.0, -1.0]]
|
||||
result = build_typed_embeddings(embs, ["ubinary"])
|
||||
assert result.ubinary is not None
|
||||
assert result.ubinary[0] == [170] # 0b10101010
|
||||
|
||||
def test_binary_all_positive(self):
|
||||
embs = [[0.1] * 8]
|
||||
result = build_typed_embeddings(embs, ["binary"])
|
||||
assert result.binary is not None
|
||||
# all bits = 1 => 0xFF = 255, signed: 255 - 128 = 127
|
||||
assert result.binary[0] == [127]
|
||||
|
||||
def test_binary_all_negative(self):
|
||||
embs = [[-0.1] * 8]
|
||||
result = build_typed_embeddings(embs, ["binary"])
|
||||
assert result.binary is not None
|
||||
# all bits = 0, signed: 0 - 128 = -128
|
||||
assert result.binary[0] == [-128]
|
||||
|
||||
def test_binary_dimension_is_eighth(self, sample_embeddings: list[list[float]]):
|
||||
result = build_typed_embeddings(sample_embeddings, ["binary"])
|
||||
assert result.binary is not None
|
||||
for orig, packed in zip(sample_embeddings, result.binary):
|
||||
assert len(packed) == len(orig) // 8
|
||||
|
||||
def test_zero_treated_as_positive(self):
|
||||
embs = [[0.0] * 8]
|
||||
result = build_typed_embeddings(embs, ["binary"])
|
||||
assert result.binary is not None
|
||||
# 0.0 >= 0 is True, so bit=1 for all => 127 (signed)
|
||||
assert result.binary[0] == [127]
|
||||
|
||||
def test_non_multiple_of_8_raises(self):
|
||||
embs = [[0.1] * 7]
|
||||
with pytest.raises(ValueError, match="multiple of 8"):
|
||||
build_typed_embeddings(embs, ["binary"])
|
||||
|
||||
def test_ubinary_non_multiple_of_8_raises(self):
|
||||
embs = [[0.1] * 10]
|
||||
with pytest.raises(ValueError, match="multiple of 8"):
|
||||
build_typed_embeddings(embs, ["ubinary"])
|
||||
|
||||
|
||||
class TestBuildTypedEmbeddingsBase64:
|
||||
def test_base64_roundtrip(self, sample_embeddings: list[list[float]]):
|
||||
result = build_typed_embeddings(sample_embeddings, ["base64"])
|
||||
assert result.base64 is not None
|
||||
assert len(result.base64) == 2
|
||||
|
||||
for orig, b64_str in zip(sample_embeddings, result.base64):
|
||||
decoded = base64.b64decode(b64_str)
|
||||
n = len(orig)
|
||||
values = struct.unpack(f"<{n}f", decoded)
|
||||
np.testing.assert_allclose(orig, values, rtol=1e-5)
|
||||
|
||||
def test_base64_byte_length(self):
|
||||
embs = [[0.1, 0.2, 0.3]]
|
||||
result = build_typed_embeddings(embs, ["base64"])
|
||||
assert result.base64 is not None
|
||||
raw = base64.b64decode(result.base64[0])
|
||||
assert len(raw) == 3 * 4 # 3 floats * 4 bytes each
|
||||
|
||||
|
||||
class TestBuildTypedEmbeddingsMultiple:
|
||||
def test_all_types_at_once(self, sample_embeddings: list[list[float]]):
|
||||
result = build_typed_embeddings(
|
||||
sample_embeddings,
|
||||
["float", "binary", "ubinary", "base64"],
|
||||
)
|
||||
assert result.float is not None
|
||||
assert result.binary is not None
|
||||
assert result.ubinary is not None
|
||||
assert result.base64 is not None
|
||||
|
||||
def test_subset_types(self, sample_embeddings: list[list[float]]):
|
||||
result = build_typed_embeddings(sample_embeddings, ["float", "binary"])
|
||||
assert result.float is not None
|
||||
assert result.binary is not None
|
||||
assert result.ubinary is None
|
||||
assert result.base64 is None
|
||||
|
||||
def test_unknown_type_ignored(self, sample_embeddings: list[list[float]]):
|
||||
result = build_typed_embeddings(sample_embeddings, ["float", "unknown_type"])
|
||||
assert result.float is not None
|
||||
@@ -0,0 +1,434 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Standalone unit tests for trtllm_prefill_attn_kvfp8_dequant.
|
||||
|
||||
Tests both contiguous and non-contiguous (cross-layer unified) KV cache
|
||||
layouts against a pure-PyTorch reference implementation.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
FP8_DTYPE = current_platform.fp8_dtype()
|
||||
|
||||
NUM_BLOCKS = 128
|
||||
|
||||
|
||||
def to_float8(x, dtype=None):
|
||||
if dtype is None:
|
||||
dtype = FP8_DTYPE
|
||||
finfo = torch.finfo(dtype)
|
||||
min_val, max_val = x.aminmax()
|
||||
amax = torch.maximum(min_val.abs(), max_val.abs()).clamp(min=1e-12)
|
||||
scale = finfo.max / amax * 0.1
|
||||
x_scl_sat = (x * scale).clamp(min=finfo.min, max=finfo.max)
|
||||
return x_scl_sat.to(dtype), scale.float().reciprocal()
|
||||
|
||||
|
||||
def make_contiguous_kv_cache(num_blocks, num_kv_heads, block_size, head_size):
|
||||
"""Create a standard contiguous fp8 KV cache (HND layout)."""
|
||||
raw = torch.randn(
|
||||
num_blocks,
|
||||
2,
|
||||
num_kv_heads,
|
||||
block_size,
|
||||
head_size,
|
||||
dtype=torch.bfloat16,
|
||||
device="cuda",
|
||||
)
|
||||
kv_cache, scale = to_float8(raw)
|
||||
return kv_cache, scale
|
||||
|
||||
|
||||
def make_cross_layer_kv_cache(
|
||||
num_blocks,
|
||||
num_kv_heads,
|
||||
block_size,
|
||||
head_size,
|
||||
num_layers=4,
|
||||
):
|
||||
"""
|
||||
Create a non-contiguous per-layer view mimicking cross-layer allocation.
|
||||
|
||||
Physical layout: (num_blocks, 2, num_kv_heads, num_layers, block_size, head_size)
|
||||
Returned view: (num_blocks, 2, num_kv_heads, block_size, head_size)
|
||||
with non-contiguous strides on dims 0, 1, 2 (they skip over num_layers).
|
||||
"""
|
||||
raw = torch.randn(
|
||||
num_blocks,
|
||||
2,
|
||||
num_kv_heads,
|
||||
num_layers,
|
||||
block_size,
|
||||
head_size,
|
||||
dtype=torch.bfloat16,
|
||||
device="cuda",
|
||||
)
|
||||
fp8_full, scale = to_float8(raw)
|
||||
layer_view = fp8_full[:, :, :, 0, :, :]
|
||||
assert not layer_view.is_contiguous(), (
|
||||
f"Expected non-contiguous view, got strides {layer_view.stride()}"
|
||||
)
|
||||
return layer_view, scale
|
||||
|
||||
|
||||
def ref_dequant(kv_cache, block_tables, k_scale, v_scale, dequant_dtype):
|
||||
"""Pure PyTorch reference: gather pages and dequantize fp8 -> dequant_dtype."""
|
||||
batch_size, num_pages_per_seq = block_tables.shape
|
||||
s = kv_cache.shape
|
||||
out = torch.zeros(
|
||||
batch_size * num_pages_per_seq + 1,
|
||||
s[1],
|
||||
s[2],
|
||||
s[3],
|
||||
s[4],
|
||||
dtype=dequant_dtype,
|
||||
device=kv_cache.device,
|
||||
)
|
||||
for b in range(batch_size):
|
||||
for p in range(num_pages_per_seq):
|
||||
page_idx = block_tables[b, p].item()
|
||||
if page_idx <= 0:
|
||||
continue
|
||||
mock_idx = b * num_pages_per_seq + p + 1
|
||||
out[mock_idx, 0] = (kv_cache[page_idx, 0].float() * k_scale.item()).to(
|
||||
dequant_dtype
|
||||
)
|
||||
out[mock_idx, 1] = (kv_cache[page_idx, 1].float() * v_scale.item()).to(
|
||||
dequant_dtype
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_kv_heads", [1, 8])
|
||||
@pytest.mark.parametrize("head_size", [64, 128])
|
||||
@pytest.mark.parametrize("block_size", [16, 32])
|
||||
@pytest.mark.parametrize("batch_size", [1, 4])
|
||||
@pytest.mark.parametrize("num_pages_per_seq", [3, 8])
|
||||
@pytest.mark.parametrize("contiguous", [True, False])
|
||||
@torch.inference_mode()
|
||||
def test_trtllm_kvfp8_dequant(
|
||||
num_kv_heads: int,
|
||||
head_size: int,
|
||||
block_size: int,
|
||||
batch_size: int,
|
||||
num_pages_per_seq: int,
|
||||
contiguous: bool,
|
||||
):
|
||||
from vllm.v1.attention.backends.flashinfer import (
|
||||
trtllm_prefill_attn_kvfp8_dequant,
|
||||
)
|
||||
|
||||
torch.set_default_device("cuda")
|
||||
|
||||
if contiguous:
|
||||
kv_cache, scale = make_contiguous_kv_cache(
|
||||
NUM_BLOCKS,
|
||||
num_kv_heads,
|
||||
block_size,
|
||||
head_size,
|
||||
)
|
||||
else:
|
||||
kv_cache, scale = make_cross_layer_kv_cache(
|
||||
NUM_BLOCKS,
|
||||
num_kv_heads,
|
||||
block_size,
|
||||
head_size,
|
||||
)
|
||||
|
||||
k_scale = scale.clone()
|
||||
v_scale = scale.clone()
|
||||
|
||||
block_tables = torch.randint(
|
||||
1,
|
||||
NUM_BLOCKS,
|
||||
(batch_size, num_pages_per_seq),
|
||||
dtype=torch.int32,
|
||||
)
|
||||
|
||||
mock_kv_cache, mock_block_table = trtllm_prefill_attn_kvfp8_dequant(
|
||||
kv_cache,
|
||||
block_tables,
|
||||
k_scale,
|
||||
v_scale,
|
||||
torch.bfloat16,
|
||||
)
|
||||
|
||||
ref = ref_dequant(kv_cache, block_tables, k_scale, v_scale, torch.bfloat16)
|
||||
|
||||
expected_bt = torch.arange(
|
||||
1,
|
||||
batch_size * num_pages_per_seq + 1,
|
||||
dtype=torch.int32,
|
||||
device="cuda",
|
||||
).reshape(batch_size, num_pages_per_seq)
|
||||
torch.testing.assert_close(mock_block_table, expected_bt)
|
||||
|
||||
# Page 0 is padding (never written), compare only pages 1+
|
||||
torch.testing.assert_close(mock_kv_cache[1:], ref[1:], atol=1e-3, rtol=1e-3)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def test_block_tables_with_zero_pages():
|
||||
"""Pages with index <= 0 must be skipped (early return in kernel)."""
|
||||
from vllm.v1.attention.backends.flashinfer import (
|
||||
trtllm_prefill_attn_kvfp8_dequant,
|
||||
)
|
||||
|
||||
torch.set_default_device("cuda")
|
||||
num_kv_heads, block_size, head_size = 8, 16, 64
|
||||
|
||||
kv_cache, scale = make_contiguous_kv_cache(
|
||||
NUM_BLOCKS,
|
||||
num_kv_heads,
|
||||
block_size,
|
||||
head_size,
|
||||
)
|
||||
k_scale = v_scale = scale.clone()
|
||||
|
||||
# Mix of valid pages and zeros (padding)
|
||||
block_tables = torch.tensor(
|
||||
[[5, 0, 10], [0, 0, 0], [3, 7, 0]],
|
||||
dtype=torch.int32,
|
||||
device="cuda",
|
||||
)
|
||||
|
||||
mock_kv_cache, _ = trtllm_prefill_attn_kvfp8_dequant(
|
||||
kv_cache,
|
||||
block_tables,
|
||||
k_scale,
|
||||
v_scale,
|
||||
torch.bfloat16,
|
||||
)
|
||||
ref = ref_dequant(kv_cache, block_tables, k_scale, v_scale, torch.bfloat16)
|
||||
|
||||
# Only compare pages that were actually written (non-zero page indices)
|
||||
for b in range(block_tables.shape[0]):
|
||||
for p in range(block_tables.shape[1]):
|
||||
if block_tables[b, p].item() > 0:
|
||||
idx = b * block_tables.shape[1] + p + 1
|
||||
torch.testing.assert_close(
|
||||
mock_kv_cache[idx],
|
||||
ref[idx],
|
||||
atol=1e-3,
|
||||
rtol=1e-3,
|
||||
)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def test_all_zero_block_tables():
|
||||
"""All-zero block_tables: kernel should write nothing."""
|
||||
from vllm.v1.attention.backends.flashinfer import (
|
||||
trtllm_prefill_attn_kvfp8_dequant,
|
||||
)
|
||||
|
||||
torch.set_default_device("cuda")
|
||||
num_kv_heads, block_size, head_size = 4, 16, 64
|
||||
|
||||
kv_cache, scale = make_contiguous_kv_cache(
|
||||
NUM_BLOCKS,
|
||||
num_kv_heads,
|
||||
block_size,
|
||||
head_size,
|
||||
)
|
||||
k_scale = v_scale = scale.clone()
|
||||
|
||||
block_tables = torch.zeros(2, 4, dtype=torch.int32, device="cuda")
|
||||
|
||||
# Should not crash even though no pages are valid
|
||||
mock_kv_cache, mock_block_table = trtllm_prefill_attn_kvfp8_dequant(
|
||||
kv_cache,
|
||||
block_tables,
|
||||
k_scale,
|
||||
v_scale,
|
||||
torch.bfloat16,
|
||||
)
|
||||
assert mock_kv_cache.shape[0] == 2 * 4 + 1
|
||||
assert mock_block_table.shape == (2, 4)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def test_different_k_v_scales():
|
||||
"""Verify K and V are dequantized with independent scales."""
|
||||
from vllm.v1.attention.backends.flashinfer import (
|
||||
trtllm_prefill_attn_kvfp8_dequant,
|
||||
)
|
||||
|
||||
torch.set_default_device("cuda")
|
||||
num_kv_heads, block_size, head_size = 8, 16, 64
|
||||
|
||||
kv_cache, _ = make_contiguous_kv_cache(
|
||||
NUM_BLOCKS,
|
||||
num_kv_heads,
|
||||
block_size,
|
||||
head_size,
|
||||
)
|
||||
k_scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
|
||||
v_scale = torch.tensor([2.0], dtype=torch.float32, device="cuda")
|
||||
|
||||
block_tables = torch.tensor([[1, 2]], dtype=torch.int32, device="cuda")
|
||||
|
||||
mock_kv_cache, _ = trtllm_prefill_attn_kvfp8_dequant(
|
||||
kv_cache,
|
||||
block_tables,
|
||||
k_scale,
|
||||
v_scale,
|
||||
torch.bfloat16,
|
||||
)
|
||||
ref = ref_dequant(kv_cache, block_tables, k_scale, v_scale, torch.bfloat16)
|
||||
|
||||
torch.testing.assert_close(mock_kv_cache[1:], ref[1:], atol=1e-3, rtol=1e-3)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def test_single_page_per_seq():
|
||||
"""Minimum grid dim 1 = 1 page per sequence."""
|
||||
from vllm.v1.attention.backends.flashinfer import (
|
||||
trtllm_prefill_attn_kvfp8_dequant,
|
||||
)
|
||||
|
||||
torch.set_default_device("cuda")
|
||||
num_kv_heads, block_size, head_size = 8, 16, 128
|
||||
|
||||
kv_cache, scale = make_contiguous_kv_cache(
|
||||
NUM_BLOCKS,
|
||||
num_kv_heads,
|
||||
block_size,
|
||||
head_size,
|
||||
)
|
||||
k_scale = v_scale = scale.clone()
|
||||
|
||||
block_tables = torch.tensor([[5], [10], [20]], dtype=torch.int32, device="cuda")
|
||||
|
||||
mock_kv_cache, _ = trtllm_prefill_attn_kvfp8_dequant(
|
||||
kv_cache,
|
||||
block_tables,
|
||||
k_scale,
|
||||
v_scale,
|
||||
torch.bfloat16,
|
||||
)
|
||||
ref = ref_dequant(kv_cache, block_tables, k_scale, v_scale, torch.bfloat16)
|
||||
|
||||
torch.testing.assert_close(mock_kv_cache[1:], ref[1:], atol=1e-3, rtol=1e-3)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def test_large_page_indices():
|
||||
"""Page indices near the top of the buffer stress offset arithmetic."""
|
||||
from vllm.v1.attention.backends.flashinfer import (
|
||||
trtllm_prefill_attn_kvfp8_dequant,
|
||||
)
|
||||
|
||||
torch.set_default_device("cuda")
|
||||
num_kv_heads, block_size, head_size = 8, 16, 128
|
||||
large_num_blocks = 32768
|
||||
|
||||
kv_cache, scale = make_contiguous_kv_cache(
|
||||
large_num_blocks,
|
||||
num_kv_heads,
|
||||
block_size,
|
||||
head_size,
|
||||
)
|
||||
k_scale = v_scale = scale.clone()
|
||||
|
||||
# Use page indices near the top of the buffer
|
||||
block_tables = torch.tensor(
|
||||
[[large_num_blocks - 1, large_num_blocks - 2, 1]],
|
||||
dtype=torch.int32,
|
||||
device="cuda",
|
||||
)
|
||||
|
||||
mock_kv_cache, _ = trtllm_prefill_attn_kvfp8_dequant(
|
||||
kv_cache,
|
||||
block_tables,
|
||||
k_scale,
|
||||
v_scale,
|
||||
torch.bfloat16,
|
||||
)
|
||||
ref = ref_dequant(kv_cache, block_tables, k_scale, v_scale, torch.bfloat16)
|
||||
|
||||
torch.testing.assert_close(mock_kv_cache[1:], ref[1:], atol=1e-3, rtol=1e-3)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def test_large_block_size():
|
||||
"""block_size=64 -> HEAD_STRIDE=8192, large tl.arange per thread block."""
|
||||
from vllm.v1.attention.backends.flashinfer import (
|
||||
trtllm_prefill_attn_kvfp8_dequant,
|
||||
)
|
||||
|
||||
torch.set_default_device("cuda")
|
||||
num_kv_heads, block_size, head_size = 4, 64, 128
|
||||
|
||||
kv_cache, scale = make_contiguous_kv_cache(
|
||||
NUM_BLOCKS,
|
||||
num_kv_heads,
|
||||
block_size,
|
||||
head_size,
|
||||
)
|
||||
k_scale = v_scale = scale.clone()
|
||||
|
||||
block_tables = torch.randint(
|
||||
1,
|
||||
NUM_BLOCKS,
|
||||
(2, 4),
|
||||
dtype=torch.int32,
|
||||
device="cuda",
|
||||
)
|
||||
|
||||
mock_kv_cache, _ = trtllm_prefill_attn_kvfp8_dequant(
|
||||
kv_cache,
|
||||
block_tables,
|
||||
k_scale,
|
||||
v_scale,
|
||||
torch.bfloat16,
|
||||
)
|
||||
ref = ref_dequant(kv_cache, block_tables, k_scale, v_scale, torch.bfloat16)
|
||||
|
||||
torch.testing.assert_close(mock_kv_cache[1:], ref[1:], atol=1e-3, rtol=1e-3)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def test_cross_layer_many_layers():
|
||||
"""
|
||||
Non-contiguous with 36 layers -- matches real gpt-oss-120b.
|
||||
Strides are far from contiguous (factor of 36 in the gaps).
|
||||
"""
|
||||
from vllm.v1.attention.backends.flashinfer import (
|
||||
trtllm_prefill_attn_kvfp8_dequant,
|
||||
)
|
||||
|
||||
torch.set_default_device("cuda")
|
||||
num_kv_heads, block_size, head_size = 8, 16, 64
|
||||
num_layers = 36
|
||||
|
||||
kv_cache, scale = make_cross_layer_kv_cache(
|
||||
NUM_BLOCKS,
|
||||
num_kv_heads,
|
||||
block_size,
|
||||
head_size,
|
||||
num_layers=num_layers,
|
||||
)
|
||||
k_scale = v_scale = scale.clone()
|
||||
|
||||
block_tables = torch.randint(
|
||||
1,
|
||||
NUM_BLOCKS,
|
||||
(4, 6),
|
||||
dtype=torch.int32,
|
||||
device="cuda",
|
||||
)
|
||||
|
||||
mock_kv_cache, _ = trtllm_prefill_attn_kvfp8_dequant(
|
||||
kv_cache,
|
||||
block_tables,
|
||||
k_scale,
|
||||
v_scale,
|
||||
torch.bfloat16,
|
||||
)
|
||||
ref = ref_dequant(kv_cache, block_tables, k_scale, v_scale, torch.bfloat16)
|
||||
|
||||
torch.testing.assert_close(mock_kv_cache[1:], ref[1:], atol=1e-3, rtol=1e-3)
|
||||
@@ -280,21 +280,22 @@ def test_rms_norm(
|
||||
assert torch.allclose(ref_residual, ops_residual)
|
||||
|
||||
output = torch.empty(x.shape, dtype=quant_dtype, device=x.device)
|
||||
scales = torch.empty(
|
||||
(x.numel() // x.shape[-1], 1), device=x.device, dtype=torch.float32
|
||||
)
|
||||
|
||||
if group_size is None:
|
||||
scales = torch.empty(
|
||||
(x.numel() // x.shape[-1], 1), device=x.device, dtype=torch.float32
|
||||
)
|
||||
opcheck(
|
||||
torch.ops._C.rms_norm_dynamic_per_token_quant,
|
||||
(output, x, layer.weight, scales, 1e-5, scale_ub, residual),
|
||||
)
|
||||
else:
|
||||
# TODO(luka/eliza) opcheck is broken?
|
||||
# Somehow the cloned args are getting mutated in-place,
|
||||
# which causes the opcheck to fail.
|
||||
# https://github.com/vllm-project/vllm/issues/36688
|
||||
return
|
||||
assert hidden_size % group_size[1] == 0
|
||||
num_groups = hidden_size // group_size[1]
|
||||
scales = torch.empty(
|
||||
(num_groups, num_tokens),
|
||||
device=x.device,
|
||||
dtype=torch.float32,
|
||||
).transpose(0, 1)
|
||||
opcheck(
|
||||
torch.ops._C.rms_norm_per_block_quant,
|
||||
(
|
||||
|
||||
@@ -33,7 +33,10 @@ from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.deep_gemm import is_deep_gemm_supported
|
||||
from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
|
||||
from vllm.utils.flashinfer import (
|
||||
has_flashinfer_cutlass_fused_moe,
|
||||
has_flashinfer_nvlink_one_sided,
|
||||
)
|
||||
from vllm.utils.import_utils import (
|
||||
has_aiter,
|
||||
has_deep_ep,
|
||||
@@ -234,15 +237,15 @@ if has_mori():
|
||||
)
|
||||
|
||||
if has_flashinfer_cutlass_fused_moe() and current_platform.has_device_capability(100):
|
||||
from vllm.model_executor.layers.fused_moe.flashinfer_a2a_prepare_finalize import ( # noqa: E501
|
||||
FlashInferA2APrepareAndFinalize,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (
|
||||
FlashInferExperts,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.flashinfer_nvlink_two_sided_prepare_finalize import ( # noqa: E501
|
||||
FlashInferNVLinkTwoSidedPrepareAndFinalize,
|
||||
)
|
||||
|
||||
register_prepare_and_finalize(
|
||||
FlashInferA2APrepareAndFinalize,
|
||||
FlashInferNVLinkTwoSidedPrepareAndFinalize,
|
||||
standard_format,
|
||||
nvfp4_types + fp8_types,
|
||||
blocked_quantization_support=True,
|
||||
@@ -263,6 +266,36 @@ else:
|
||||
FlashInferCutlassMoEPrepareAndFinalize = None
|
||||
FlashInferExperts = None
|
||||
|
||||
if (
|
||||
has_flashinfer_nvlink_one_sided()
|
||||
and has_flashinfer_cutlass_fused_moe()
|
||||
and current_platform.has_device_capability(100)
|
||||
):
|
||||
from vllm.model_executor.layers.fused_moe.flashinfer_nvlink_one_sided_prepare_finalize import ( # noqa: E501
|
||||
FlashInferNVLinkOneSidedPrepareAndFinalize,
|
||||
)
|
||||
|
||||
register_prepare_and_finalize(
|
||||
FlashInferNVLinkOneSidedPrepareAndFinalize,
|
||||
standard_format,
|
||||
nvfp4_types,
|
||||
blocked_quantization_support=False,
|
||||
backend="flashinfer_nvlink_one_sided",
|
||||
supports_apply_weight_on_input=False,
|
||||
)
|
||||
|
||||
if has_flashinfer_cutlass_fused_moe() and current_platform.has_device_capability(100):
|
||||
from vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe import (
|
||||
TrtLlmNvFp4ExpertsModular,
|
||||
)
|
||||
|
||||
register_experts(
|
||||
TrtLlmNvFp4ExpertsModular,
|
||||
standard_format,
|
||||
nvfp4_types,
|
||||
blocked_quantization_support=False,
|
||||
supports_expert_map=True,
|
||||
)
|
||||
|
||||
if has_aiter():
|
||||
from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
|
||||
|
||||
@@ -159,6 +159,52 @@ def test_quantize_to_fp4(
|
||||
torch.testing.assert_close(scale_ans, scale_ref)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"shape",
|
||||
[(32, 4096), (128, 4096), (1, 64), (127, 1024), (256, 16384)],
|
||||
)
|
||||
@pytest.mark.parametrize("is_sf_swizzled_layout", [True, False])
|
||||
@torch.inference_mode()
|
||||
def test_python_util_matches_cpp_allocation(
|
||||
shape: tuple[int, int],
|
||||
is_sf_swizzled_layout: bool,
|
||||
) -> None:
|
||||
"""
|
||||
Verify that the Python utility (create_fp4_output_tensors) allocates
|
||||
tensors with the same shapes and dtypes as the C++ functional variant
|
||||
(scaled_fp4_quant_func).
|
||||
"""
|
||||
from vllm._custom_ops import create_fp4_output_tensors
|
||||
|
||||
torch.set_default_device("cuda:0")
|
||||
m, n = shape
|
||||
input_tensor = torch.randn((m, n), dtype=torch.bfloat16)
|
||||
input_scale = torch.tensor([1.0], dtype=torch.float32, device="cuda:0")
|
||||
|
||||
# C++ functional variant allocates internally
|
||||
cpp_out, cpp_scale = torch.ops._C.scaled_fp4_quant(
|
||||
input_tensor, input_scale, is_sf_swizzled_layout
|
||||
)
|
||||
|
||||
# Python utility
|
||||
py_out, py_scale = create_fp4_output_tensors(
|
||||
m, n, torch.device("cuda:0"), is_sf_swizzled_layout
|
||||
)
|
||||
|
||||
assert py_out.shape == cpp_out.shape, (
|
||||
f"Output shape mismatch: Python {py_out.shape} vs C++ {cpp_out.shape}"
|
||||
)
|
||||
assert py_out.dtype == cpp_out.dtype, (
|
||||
f"Output dtype mismatch: Python {py_out.dtype} vs C++ {cpp_out.dtype}"
|
||||
)
|
||||
assert py_scale.shape == cpp_scale.shape, (
|
||||
f"Scale shape mismatch: Python {py_scale.shape} vs C++ {cpp_scale.shape}"
|
||||
)
|
||||
assert py_scale.dtype == cpp_scale.dtype, (
|
||||
f"Scale dtype mismatch: Python {py_scale.dtype} vs C++ {cpp_scale.dtype}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("pad_shape", PAD_SHAPES)
|
||||
@torch.inference_mode()
|
||||
def test_quantize_to_fp4_padded(pad_shape: tuple[int, int]) -> None:
|
||||
|
||||
@@ -19,6 +19,7 @@ from vllm.transformers_utils.runai_utils import (
|
||||
def test_is_runai_obj_uri():
|
||||
assert is_runai_obj_uri("gs://some-gcs-bucket/path")
|
||||
assert is_runai_obj_uri("s3://some-s3-bucket/path")
|
||||
assert is_runai_obj_uri("az://some-azure-container/path")
|
||||
assert not is_runai_obj_uri("nfs://some-nfs-path")
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,361 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for EP weight filtering during model loading."""
|
||||
|
||||
import glob
|
||||
import tempfile
|
||||
|
||||
import huggingface_hub.constants
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.model_executor.model_loader.ep_weight_filter import (
|
||||
compute_local_expert_ids,
|
||||
parse_expert_id,
|
||||
should_skip_weight,
|
||||
)
|
||||
from vllm.model_executor.model_loader.weight_utils import (
|
||||
safetensors_weights_iterator,
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Unit tests for parse_expert_id
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestParseExpertId:
|
||||
def test_routed_expert(self):
|
||||
name = "model.layers.0.mlp.experts.42.gate_proj.weight"
|
||||
assert parse_expert_id(name) == 42
|
||||
|
||||
def test_large_expert_id(self):
|
||||
name = "model.layers.60.mlp.experts.383.down_proj.weight"
|
||||
assert parse_expert_id(name) == 383
|
||||
|
||||
def test_shared_expert(self):
|
||||
# Shared experts use a different naming convention in most models
|
||||
name = "model.layers.0.mlp.shared_experts.gate_proj.weight"
|
||||
assert parse_expert_id(name) is None
|
||||
|
||||
def test_attention_weight(self):
|
||||
name = "model.layers.0.self_attn.q_proj.weight"
|
||||
assert parse_expert_id(name) is None
|
||||
|
||||
def test_embedding(self):
|
||||
name = "model.embed_tokens.weight"
|
||||
assert parse_expert_id(name) is None
|
||||
|
||||
def test_layernorm(self):
|
||||
name = "model.layers.0.input_layernorm.weight"
|
||||
assert parse_expert_id(name) is None
|
||||
|
||||
def test_fused_3d_expert(self):
|
||||
# 3D fused-expert tensors (e.g. gpt-oss) have no numeric expert id.
|
||||
# They must NOT be filtered — slicing happens later in weight_loader.
|
||||
name = "model.layers.0.mlp.experts.gate_proj.weight"
|
||||
assert parse_expert_id(name) is None
|
||||
|
||||
def test_fused_3d_expert_down_proj(self):
|
||||
name = "model.layers.10.mlp.experts.down_proj.weight"
|
||||
assert parse_expert_id(name) is None
|
||||
|
||||
def test_expert_scale(self):
|
||||
# NVFP4 quantized models have scale tensors for experts
|
||||
name = "model.layers.5.mlp.experts.100.gate_proj.weight_scale"
|
||||
assert parse_expert_id(name) == 100
|
||||
|
||||
def test_expert_zero_id(self):
|
||||
name = "model.layers.0.mlp.experts.0.up_proj.weight"
|
||||
assert parse_expert_id(name) == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Unit tests for compute_local_expert_ids
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestComputeLocalExpertIds:
|
||||
def test_ep_disabled(self):
|
||||
assert compute_local_expert_ids(64, ep_size=1, ep_rank=0) is None
|
||||
|
||||
def test_even_split(self):
|
||||
# 64 experts, EP=8 → 8 per rank
|
||||
ids = compute_local_expert_ids(64, ep_size=8, ep_rank=0)
|
||||
assert ids == set(range(0, 8))
|
||||
|
||||
ids = compute_local_expert_ids(64, ep_size=8, ep_rank=7)
|
||||
assert ids == set(range(56, 64))
|
||||
|
||||
def test_uneven_split(self):
|
||||
# 10 experts, EP=3 → ranks get 4, 3, 3
|
||||
ids_0 = compute_local_expert_ids(10, ep_size=3, ep_rank=0)
|
||||
ids_1 = compute_local_expert_ids(10, ep_size=3, ep_rank=1)
|
||||
ids_2 = compute_local_expert_ids(10, ep_size=3, ep_rank=2)
|
||||
|
||||
assert len(ids_0) == 4
|
||||
assert len(ids_1) == 3
|
||||
assert len(ids_2) == 3
|
||||
# All experts covered, no overlap
|
||||
assert ids_0 | ids_1 | ids_2 == set(range(10))
|
||||
assert ids_0.isdisjoint(ids_1)
|
||||
assert ids_1.isdisjoint(ids_2)
|
||||
|
||||
def test_384_experts_ep8(self):
|
||||
# Kimi-K2.5 config: 384 experts, EP=8
|
||||
for rank in range(8):
|
||||
ids = compute_local_expert_ids(384, ep_size=8, ep_rank=rank)
|
||||
assert len(ids) == 48
|
||||
|
||||
# All experts covered
|
||||
all_ids = set()
|
||||
for rank in range(8):
|
||||
ids = compute_local_expert_ids(384, ep_size=8, ep_rank=rank)
|
||||
all_ids |= ids
|
||||
assert all_ids == set(range(384))
|
||||
|
||||
def test_384_experts_ep16(self):
|
||||
for rank in range(16):
|
||||
ids = compute_local_expert_ids(384, ep_size=16, ep_rank=rank)
|
||||
assert len(ids) == 24
|
||||
|
||||
def test_384_experts_ep24(self):
|
||||
# 384 / 24 = 16 exactly
|
||||
for rank in range(24):
|
||||
ids = compute_local_expert_ids(384, ep_size=24, ep_rank=rank)
|
||||
assert len(ids) == 16
|
||||
|
||||
# round_robin placement tests
|
||||
|
||||
def test_round_robin_basic(self):
|
||||
# 8 experts, EP=2: rank 0 → {0,2,4,6}, rank 1 → {1,3,5,7}
|
||||
rr = "round_robin"
|
||||
ids_0 = compute_local_expert_ids(8, 2, 0, placement=rr)
|
||||
ids_1 = compute_local_expert_ids(8, 2, 1, placement=rr)
|
||||
assert ids_0 == {0, 2, 4, 6}
|
||||
assert ids_1 == {1, 3, 5, 7}
|
||||
|
||||
def test_round_robin_full_coverage(self):
|
||||
# 384 experts, EP=8: all experts covered, no overlap
|
||||
rr = "round_robin"
|
||||
all_ids: set[int] = set()
|
||||
for rank in range(8):
|
||||
ids = compute_local_expert_ids(384, 8, rank, placement=rr)
|
||||
assert ids is not None and len(ids) == 48
|
||||
assert all_ids.isdisjoint(ids)
|
||||
all_ids |= ids
|
||||
assert all_ids == set(range(384))
|
||||
|
||||
def test_round_robin_uneven(self):
|
||||
# 10 experts, EP=3: rank 0→{0,3,6,9}, rank 1→{1,4,7}, rank 2→{2,5,8}
|
||||
rr = "round_robin"
|
||||
ids_0 = compute_local_expert_ids(10, 3, 0, placement=rr)
|
||||
ids_1 = compute_local_expert_ids(10, 3, 1, placement=rr)
|
||||
ids_2 = compute_local_expert_ids(10, 3, 2, placement=rr)
|
||||
assert ids_0 == {0, 3, 6, 9}
|
||||
assert ids_1 == {1, 4, 7}
|
||||
assert ids_2 == {2, 5, 8}
|
||||
assert ids_0 | ids_1 | ids_2 == set(range(10))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Unit tests for should_skip_weight
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestShouldSkipWeight:
|
||||
def setup_method(self):
|
||||
# Simulate EP=8, rank=0 → experts 0-47
|
||||
self.local_ids = compute_local_expert_ids(384, ep_size=8, ep_rank=0)
|
||||
|
||||
def test_no_filter(self):
|
||||
assert not should_skip_weight("anything", None)
|
||||
|
||||
def test_dense_not_skipped(self):
|
||||
assert not should_skip_weight(
|
||||
"model.layers.0.self_attn.q_proj.weight", self.local_ids
|
||||
)
|
||||
|
||||
def test_local_expert_not_skipped(self):
|
||||
assert not should_skip_weight(
|
||||
"model.layers.0.mlp.experts.10.gate_proj.weight", self.local_ids
|
||||
)
|
||||
|
||||
def test_remote_expert_skipped(self):
|
||||
assert should_skip_weight(
|
||||
"model.layers.0.mlp.experts.200.gate_proj.weight", self.local_ids
|
||||
)
|
||||
|
||||
def test_boundary_expert(self):
|
||||
# Expert 47 is local (last one), 48 is not
|
||||
assert not should_skip_weight(
|
||||
"model.layers.0.mlp.experts.47.gate_proj.weight", self.local_ids
|
||||
)
|
||||
assert should_skip_weight(
|
||||
"model.layers.0.mlp.experts.48.gate_proj.weight", self.local_ids
|
||||
)
|
||||
|
||||
def test_shared_expert_not_skipped(self):
|
||||
assert not should_skip_weight(
|
||||
"model.layers.0.mlp.shared_experts.gate_proj.weight", self.local_ids
|
||||
)
|
||||
|
||||
def test_embedding_not_skipped(self):
|
||||
assert not should_skip_weight("model.embed_tokens.weight", self.local_ids)
|
||||
|
||||
def test_fused_3d_expert_not_skipped(self):
|
||||
# 3D fused-expert tensors (gpt-oss style) have no numeric id.
|
||||
# Must not be skipped — weight_loader handles slicing later.
|
||||
assert not should_skip_weight(
|
||||
"model.layers.0.mlp.experts.gate_proj.weight", self.local_ids
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Integration test: safetensors_weights_iterator with EP filtering
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestSafetensorsWeightsIteratorWithEpFilter:
|
||||
"""Verify that EP filtering produces a strict subset of unfiltered loading
|
||||
and that all expected dense + local expert weights are present."""
|
||||
|
||||
@pytest.fixture(scope="class")
|
||||
def gpt2_files(self):
|
||||
"""Download GPT-2 safetensors to a temp dir (shared across class)."""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
huggingface_hub.constants.HF_HUB_OFFLINE = False
|
||||
from vllm.model_executor.model_loader.weight_utils import (
|
||||
download_weights_from_hf,
|
||||
)
|
||||
|
||||
download_weights_from_hf(
|
||||
"openai-community/gpt2",
|
||||
allow_patterns=["*.safetensors"],
|
||||
cache_dir=tmpdir,
|
||||
)
|
||||
files = glob.glob(f"{tmpdir}/**/*.safetensors", recursive=True)
|
||||
assert len(files) > 0
|
||||
yield files
|
||||
|
||||
def test_no_filter_returns_all(self, gpt2_files):
|
||||
"""With local_expert_ids=None, all weights are returned (no MoE)."""
|
||||
all_weights = dict(safetensors_weights_iterator(gpt2_files, False))
|
||||
filtered_weights = dict(
|
||||
safetensors_weights_iterator(gpt2_files, False, local_expert_ids=None)
|
||||
)
|
||||
assert set(all_weights.keys()) == set(filtered_weights.keys())
|
||||
|
||||
def test_empty_filter_skips_experts_only(self, gpt2_files):
|
||||
"""GPT-2 has no expert weights, so even an empty local_expert_ids
|
||||
set should return all weights (all are dense)."""
|
||||
all_weights = dict(safetensors_weights_iterator(gpt2_files, False))
|
||||
filtered_weights = dict(
|
||||
safetensors_weights_iterator(gpt2_files, False, local_expert_ids=set())
|
||||
)
|
||||
# GPT-2 has no experts, so nothing should be filtered
|
||||
assert set(all_weights.keys()) == set(filtered_weights.keys())
|
||||
|
||||
|
||||
class TestEpFilterOnSyntheticMoeWeights:
|
||||
"""Create synthetic safetensors files with expert-like naming and verify
|
||||
that the filter correctly skips non-local experts."""
|
||||
|
||||
@pytest.fixture
|
||||
def synthetic_moe_files(self, tmp_path):
|
||||
"""Create synthetic safetensors with expert-patterned tensor names."""
|
||||
from safetensors.torch import save_file
|
||||
|
||||
tensors = {}
|
||||
# Dense weights
|
||||
tensors["model.embed_tokens.weight"] = torch.randn(100, 64)
|
||||
tensors["model.layers.0.self_attn.q_proj.weight"] = torch.randn(64, 64)
|
||||
tensors["model.layers.0.input_layernorm.weight"] = torch.randn(64)
|
||||
# Expert weights: 8 experts
|
||||
for expert_id in range(8):
|
||||
tensors[f"model.layers.0.mlp.experts.{expert_id}.gate_proj.weight"] = (
|
||||
torch.randn(128, 64)
|
||||
)
|
||||
tensors[f"model.layers.0.mlp.experts.{expert_id}.up_proj.weight"] = (
|
||||
torch.randn(128, 64)
|
||||
)
|
||||
tensors[f"model.layers.0.mlp.experts.{expert_id}.down_proj.weight"] = (
|
||||
torch.randn(64, 128)
|
||||
)
|
||||
# Shared expert (should never be filtered)
|
||||
tensors["model.layers.0.mlp.shared_experts.gate_proj.weight"] = torch.randn(
|
||||
128, 64
|
||||
)
|
||||
|
||||
filepath = str(tmp_path / "model-00001-of-00001.safetensors")
|
||||
save_file(tensors, filepath)
|
||||
return [filepath], tensors
|
||||
|
||||
def test_no_filter_returns_all(self, synthetic_moe_files):
|
||||
files, expected = synthetic_moe_files
|
||||
loaded = dict(safetensors_weights_iterator(files, False))
|
||||
assert set(loaded.keys()) == set(expected.keys())
|
||||
|
||||
def test_ep2_rank0_gets_half_experts(self, synthetic_moe_files):
|
||||
files, expected = synthetic_moe_files
|
||||
# EP=2, rank=0 → experts 0-3
|
||||
local_ids = compute_local_expert_ids(8, ep_size=2, ep_rank=0)
|
||||
loaded = dict(
|
||||
safetensors_weights_iterator(files, False, local_expert_ids=local_ids)
|
||||
)
|
||||
|
||||
# Should have all dense + shared + experts 0-3 only
|
||||
for name in loaded:
|
||||
eid = parse_expert_id(name)
|
||||
if eid is not None:
|
||||
assert eid in local_ids, f"Non-local expert {eid} was loaded"
|
||||
|
||||
# Check expert count: 4 experts × 3 weights = 12
|
||||
expert_names = [n for n in loaded if parse_expert_id(n) is not None]
|
||||
assert len(expert_names) == 4 * 3
|
||||
|
||||
# Check all dense weights present
|
||||
assert "model.embed_tokens.weight" in loaded
|
||||
assert "model.layers.0.self_attn.q_proj.weight" in loaded
|
||||
assert "model.layers.0.input_layernorm.weight" in loaded
|
||||
assert "model.layers.0.mlp.shared_experts.gate_proj.weight" in loaded
|
||||
|
||||
def test_ep2_rank1_gets_other_half(self, synthetic_moe_files):
|
||||
files, expected = synthetic_moe_files
|
||||
local_ids = compute_local_expert_ids(8, ep_size=2, ep_rank=1)
|
||||
loaded = dict(
|
||||
safetensors_weights_iterator(files, False, local_expert_ids=local_ids)
|
||||
)
|
||||
|
||||
expert_names = [n for n in loaded if parse_expert_id(n) is not None]
|
||||
assert len(expert_names) == 4 * 3
|
||||
for name in expert_names:
|
||||
assert parse_expert_id(name) in local_ids
|
||||
|
||||
def test_ep8_each_rank_gets_one_expert(self, synthetic_moe_files):
|
||||
files, _ = synthetic_moe_files
|
||||
all_expert_names = set()
|
||||
for rank in range(8):
|
||||
local_ids = compute_local_expert_ids(8, ep_size=8, ep_rank=rank)
|
||||
loaded = dict(
|
||||
safetensors_weights_iterator(files, False, local_expert_ids=local_ids)
|
||||
)
|
||||
expert_names = {n for n in loaded if parse_expert_id(n) is not None}
|
||||
# 1 expert × 3 weights
|
||||
assert len(expert_names) == 3
|
||||
all_expert_names |= expert_names
|
||||
|
||||
# All 8 experts × 3 weights covered across ranks
|
||||
assert len(all_expert_names) == 24
|
||||
|
||||
def test_tensor_values_match(self, synthetic_moe_files):
|
||||
"""Filtered tensors have identical values to unfiltered ones."""
|
||||
files, _ = synthetic_moe_files
|
||||
all_weights = dict(safetensors_weights_iterator(files, False))
|
||||
|
||||
local_ids = compute_local_expert_ids(8, ep_size=2, ep_rank=0)
|
||||
filtered = dict(
|
||||
safetensors_weights_iterator(files, False, local_expert_ids=local_ids)
|
||||
)
|
||||
|
||||
for name, tensor in filtered.items():
|
||||
assert torch.equal(tensor, all_weights[name]), f"Tensor mismatch for {name}"
|
||||
@@ -0,0 +1,68 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for CPU unquantized GEMM dispatch behavior."""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.model_executor.layers import utils
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def _mock_zentorch_linear_unary():
|
||||
"""Register a mock zentorch_linear_unary op when zentorch is not installed.
|
||||
|
||||
Allows the dispatch tests to run in CI without a real zentorch build.
|
||||
Skips registration when zentorch is already available.
|
||||
"""
|
||||
if hasattr(torch.ops.zentorch, "zentorch_linear_unary"):
|
||||
yield
|
||||
return
|
||||
|
||||
lib_def = torch.library.Library("zentorch", "DEF")
|
||||
lib_def.define(
|
||||
"zentorch_linear_unary("
|
||||
"Tensor input, "
|
||||
"Tensor weight, "
|
||||
"Tensor? bias, "
|
||||
"bool is_weight_prepacked=False"
|
||||
") -> Tensor"
|
||||
)
|
||||
|
||||
lib_impl = torch.library.Library("zentorch", "IMPL", "CPU")
|
||||
lib_impl.impl(
|
||||
"zentorch_linear_unary",
|
||||
lambda input, weight, bias, is_weight_prepacked=False: (
|
||||
torch.nn.functional.linear(input, weight, bias)
|
||||
),
|
||||
)
|
||||
|
||||
yield
|
||||
|
||||
lib_impl._destroy()
|
||||
lib_def._destroy()
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("_mock_zentorch_linear_unary")
|
||||
def test_dispatch_cpu_unquantized_gemm_uses_zentorch_on_zen(monkeypatch):
|
||||
monkeypatch.setattr(current_platform, "is_zen_cpu", lambda: True)
|
||||
|
||||
layer = torch.nn.Linear(16, 8, bias=True)
|
||||
x = torch.randn(4, 16)
|
||||
expected = torch.nn.functional.linear(x, layer.weight, layer.bias)
|
||||
|
||||
utils.dispatch_cpu_unquantized_gemm(layer, remove_weight=False)
|
||||
output = layer.cpu_linear(x, layer.weight, layer.bias)
|
||||
|
||||
torch.testing.assert_close(output, expected)
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("_mock_zentorch_linear_unary")
|
||||
def test_dispatch_cpu_unquantized_gemm_zen_remove_weight(monkeypatch):
|
||||
monkeypatch.setattr(current_platform, "is_zen_cpu", lambda: True)
|
||||
|
||||
layer = torch.nn.Linear(16, 8, bias=True)
|
||||
utils.dispatch_cpu_unquantized_gemm(layer, remove_weight=True)
|
||||
|
||||
assert layer.weight.numel() == 0
|
||||
@@ -103,6 +103,10 @@ AITER_MODEL_LIST = [
|
||||
marks=[pytest.mark.core_model, pytest.mark.cpu_model],
|
||||
),
|
||||
pytest.param("swiss-ai/Apertus-8B-Instruct-2509"), # apertus
|
||||
pytest.param(
|
||||
"naver-hyperclovax/HyperCLOVAX-SEED-Think-14B", # hyperclovax
|
||||
marks=[large_gpu_mark(min_gb=32)],
|
||||
),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("max_tokens", [32])
|
||||
|
||||
@@ -1,8 +1,10 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import contextlib
|
||||
from dataclasses import asdict
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from mistral_common.audio import Audio
|
||||
from mistral_common.protocol.instruct.chunk import RawAudio
|
||||
from mistral_common.protocol.transcription.request import (
|
||||
@@ -17,18 +19,21 @@ from vllm.assets.audio import AudioAsset
|
||||
from vllm.engine.arg_utils import AsyncEngineArgs
|
||||
from vllm.v1.engine.async_llm import AsyncLLM
|
||||
|
||||
from ....utils import ROCM_ENGINE_KWARGS
|
||||
|
||||
MODEL_NAME = "mistralai/Voxtral-Mini-4B-Realtime-2602"
|
||||
ENGINE_CONFIG = dict(
|
||||
model=MODEL_NAME,
|
||||
max_model_len=8192,
|
||||
max_num_seqs=4,
|
||||
limit_mm_per_prompt={"audio": 1},
|
||||
config_format="mistral",
|
||||
load_format="mistral",
|
||||
tokenizer_mode="mistral",
|
||||
enforce_eager=True,
|
||||
gpu_memory_utilization=0.9,
|
||||
)
|
||||
ENGINE_CONFIG = {
|
||||
"model": MODEL_NAME,
|
||||
"max_model_len": 8192,
|
||||
"max_num_seqs": 4,
|
||||
"limit_mm_per_prompt": {"audio": 1},
|
||||
"config_format": "mistral",
|
||||
"load_format": "mistral",
|
||||
"tokenizer_mode": "mistral",
|
||||
"enforce_eager": True,
|
||||
"gpu_memory_utilization": 0.9,
|
||||
**ROCM_ENGINE_KWARGS,
|
||||
}
|
||||
|
||||
|
||||
EXPECTED_TEXT = [
|
||||
@@ -49,6 +54,14 @@ EXPECTED_TEXT = [
|
||||
]
|
||||
|
||||
|
||||
def _normalize(texts: list[str]) -> list[str]:
|
||||
# The model occasionally transcribes "OBS" as "a base hit" and
|
||||
# "oh, my" as "oh my", but both are acoustically valid. Normalise so
|
||||
# the assertion is stable across runs and hardware.
|
||||
texts[1] = texts[1].replace("a base hit", "OBS").replace("oh my", "oh, my")
|
||||
return texts
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def audio_assets() -> list[AudioAsset]:
|
||||
return [AudioAsset("mary_had_lamb"), AudioAsset("winning_call")]
|
||||
@@ -60,15 +73,27 @@ def tokenizer() -> MistralTokenizer:
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def engine() -> LLM:
|
||||
def engine():
|
||||
engine_args = EngineArgs(**ENGINE_CONFIG)
|
||||
return LLM(**asdict(engine_args))
|
||||
llm = LLM(**asdict(engine_args))
|
||||
try:
|
||||
yield llm
|
||||
finally:
|
||||
with contextlib.suppress(Exception):
|
||||
llm.llm_engine.engine_core.shutdown()
|
||||
import torch
|
||||
|
||||
torch.accelerator.empty_cache()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def async_engine() -> AsyncLLM:
|
||||
@pytest_asyncio.fixture
|
||||
async def async_engine():
|
||||
engine_args = AsyncEngineArgs(**ENGINE_CONFIG)
|
||||
return AsyncLLM.from_engine_args(engine_args)
|
||||
llm = AsyncLLM.from_engine_args(engine_args)
|
||||
try:
|
||||
yield llm
|
||||
finally:
|
||||
llm.shutdown()
|
||||
|
||||
|
||||
def test_voxtral_realtime_forward(audio_assets, tokenizer, engine):
|
||||
@@ -108,8 +133,13 @@ def test_voxtral_realtime_forward(audio_assets, tokenizer, engine):
|
||||
sampling_params=sampling_params,
|
||||
)
|
||||
|
||||
texts = [out.outputs[0].text for out in outputs]
|
||||
assert texts == EXPECTED_TEXT
|
||||
texts = _normalize([out.outputs[0].text for out in outputs])
|
||||
for i, (got, expected) in enumerate(zip(texts, EXPECTED_TEXT)):
|
||||
assert got == expected, (
|
||||
f"Output mismatch at index {i}:\n"
|
||||
f" got: {got!r}\n"
|
||||
f" expected: {expected!r}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -149,9 +179,17 @@ async def test_voxtral_realtime_generator(audio_assets, tokenizer, async_engine)
|
||||
|
||||
output_tokens_list.append(output_tokens)
|
||||
|
||||
texts = [
|
||||
tokenizer.decode(output_tokens, special_token_policy=SpecialTokenPolicy.IGNORE)
|
||||
for output_tokens in output_tokens_list
|
||||
]
|
||||
texts[1] = texts[1].replace("a base hit", "OBS").replace("oh my", "oh, my")
|
||||
assert texts == EXPECTED_TEXT
|
||||
texts = _normalize(
|
||||
[
|
||||
tokenizer.decode(
|
||||
output_tokens, special_token_policy=SpecialTokenPolicy.IGNORE
|
||||
)
|
||||
for output_tokens in output_tokens_list
|
||||
]
|
||||
)
|
||||
for i, (got, expected) in enumerate(zip(texts, EXPECTED_TEXT)):
|
||||
assert got == expected, (
|
||||
f"Output mismatch at index {i}:\n"
|
||||
f" got: {got!r}\n"
|
||||
f" expected: {expected!r}"
|
||||
)
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from transformers import CLIPModel
|
||||
|
||||
from ....conftest import IMAGE_ASSETS, HfRunner, PromptImageInput, VllmRunner
|
||||
@@ -50,13 +51,16 @@ def _run_test(
|
||||
if "pixel_values" in inputs:
|
||||
pooled_output = hf_model.model.get_image_features(
|
||||
pixel_values=inputs.pixel_values,
|
||||
).squeeze(0)
|
||||
)
|
||||
else:
|
||||
pooled_output = hf_model.model.get_text_features(
|
||||
input_ids=inputs.input_ids,
|
||||
attention_mask=inputs.attention_mask,
|
||||
).squeeze(0)
|
||||
)
|
||||
|
||||
if not isinstance(pooled_output, torch.Tensor):
|
||||
pooled_output = pooled_output.pooler_output
|
||||
pooled_output = pooled_output.squeeze(0)
|
||||
all_outputs.append(pooled_output.tolist())
|
||||
|
||||
hf_outputs = all_outputs
|
||||
|
||||
@@ -4,6 +4,7 @@
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from transformers import SiglipModel
|
||||
|
||||
from ....conftest import IMAGE_ASSETS, HfRunner, PromptImageInput, VllmRunner
|
||||
@@ -68,12 +69,15 @@ def _run_test(
|
||||
if "pixel_values" in inputs:
|
||||
pooled_output = hf_model.model.get_image_features(
|
||||
pixel_values=inputs.pixel_values,
|
||||
).squeeze(0)
|
||||
)
|
||||
else:
|
||||
pooled_output = hf_model.model.get_text_features(
|
||||
input_ids=inputs.input_ids,
|
||||
).squeeze(0)
|
||||
)
|
||||
|
||||
if not isinstance(pooled_output, torch.Tensor):
|
||||
pooled_output = pooled_output.pooler_output
|
||||
pooled_output = pooled_output.squeeze(0)
|
||||
all_outputs.append(pooled_output.tolist())
|
||||
|
||||
hf_outputs = all_outputs
|
||||
|
||||
@@ -34,8 +34,22 @@ MODELS = [
|
||||
]
|
||||
|
||||
|
||||
def create_mm_data(num_videos: int) -> dict[str, list]:
|
||||
# Small video (8 frames, 64×64) and ~0.5 s of audio at 16 kHz so the test
|
||||
# stays fast even without a GPU.
|
||||
mm_data = dict[str, list](video=[], audio=[])
|
||||
for i in range(num_videos):
|
||||
rng = np.random.RandomState(i)
|
||||
video = random_video(rng, min_frames=8, max_frames=9, min_wh=64, max_wh=65)
|
||||
audio, sr = random_audio(rng, min_len=8000, max_len=8001, sr=16000)
|
||||
mm_data["video"].append(video)
|
||||
mm_data["audio"].append((audio, sr))
|
||||
return mm_data
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_id", MODELS)
|
||||
def test_audio_in_video_cache_correctness(model_id: str) -> None:
|
||||
@pytest.mark.parametrize("num_videos", [1, 2])
|
||||
def test_audio_in_video_cache_correctness(model_id: str, num_videos: int) -> None:
|
||||
"""
|
||||
Regression test for https://github.com/vllm-project/vllm/pull/36800
|
||||
|
||||
@@ -47,7 +61,7 @@ def test_audio_in_video_cache_correctness(model_id: str) -> None:
|
||||
"""
|
||||
ctx = build_model_context(
|
||||
model_id,
|
||||
limit_mm_per_prompt={"audio": 1, "image": 0, "video": 1},
|
||||
limit_mm_per_prompt={"audio": num_videos, "image": 0, "video": num_videos},
|
||||
mm_processor_cache_gb=1,
|
||||
)
|
||||
|
||||
@@ -65,17 +79,12 @@ def test_audio_in_video_cache_correctness(model_id: str) -> None:
|
||||
|
||||
video_token_id = baseline_processor.info.get_hf_config().video_token_id
|
||||
|
||||
rng = np.random.RandomState(0)
|
||||
# Small video (8 frames, 64×64) and ~0.5 s of audio at 16 kHz so the test
|
||||
# stays fast even without a GPU.
|
||||
video = random_video(rng, min_frames=8, max_frames=9, min_wh=64, max_wh=65)
|
||||
audio, sr = random_audio(rng, min_len=8000, max_len=8001, sr=16000)
|
||||
mm_data = {"video": [video], "audio": [(audio, sr)]}
|
||||
mm_data = create_mm_data(num_videos)
|
||||
hf_processor_mm_kwargs = {"use_audio_in_video": True}
|
||||
|
||||
def run(processor):
|
||||
return processor(
|
||||
[video_token_id],
|
||||
[video_token_id] * num_videos,
|
||||
mm_items=baseline_processor.info.parse_mm_data(mm_data),
|
||||
hf_processor_mm_kwargs=hf_processor_mm_kwargs,
|
||||
)["prompt_token_ids"]
|
||||
|
||||
@@ -0,0 +1,104 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""E2E tests for online MXFP8 quantization.
|
||||
|
||||
Loads a BF16 model with ``--quantization mxfp8`` (online quantization) and
|
||||
compares log-probabilities against the same model served in BF16 without
|
||||
quantization. This exercises the full pipeline: config parsing,
|
||||
``Mxfp8OnlineLinearMethod``, ``Mxfp8OnlineMoEMethod``, weight loading,
|
||||
online quantization / shuffling, and inference through ``apply_monolithic``.
|
||||
|
||||
Layer skipping (``modules_to_not_convert``) is configured in the model's
|
||||
``config.json`` under ``quantization_config`` and is not tested here.
|
||||
|
||||
``example_prompts`` is a pytest fixture (from conftest.py) that loads 8
|
||||
diverse prompts from ``tests/prompts/example.txt``.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.quantization.utils import is_quant_method_supported
|
||||
|
||||
from ..utils import check_logprobs_close
|
||||
|
||||
# A small MoE model that fits on a single GPU and has both linear + MoE layers.
|
||||
MOE_MODEL = "Qwen/Qwen3-30B-A3B"
|
||||
# A small dense model (no MoE) to validate the linear-only path.
|
||||
DENSE_MODEL = "Qwen/Qwen3-0.6B"
|
||||
|
||||
MAX_MODEL_LEN = 1024
|
||||
MAX_TOKENS = 4
|
||||
NUM_LOG_PROBS = 8
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not is_quant_method_supported("mxfp8"),
|
||||
reason="mxfp8 is not supported on this GPU type (requires sm_100+).",
|
||||
)
|
||||
@pytest.mark.quant_model
|
||||
@pytest.mark.parametrize("model", [DENSE_MODEL, MOE_MODEL], ids=["dense", "moe"])
|
||||
def test_mxfp8_logprobs(
|
||||
vllm_runner,
|
||||
example_prompts,
|
||||
model: str,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Compare BF16 baseline logprobs against online MXFP8-quantized model.
|
||||
|
||||
Runs the same model twice -- once in BF16 (baseline) and once with
|
||||
online MXFP8 quantization -- then checks that the top log-probabilities
|
||||
are close. Only 4 tokens are generated to keep the test fast while
|
||||
still catching numerical divergence.
|
||||
"""
|
||||
with monkeypatch.context() as m:
|
||||
m.setenv("TOKENIZERS_PARALLELISM", "true")
|
||||
|
||||
with vllm_runner(
|
||||
model,
|
||||
max_model_len=MAX_MODEL_LEN,
|
||||
enforce_eager=True,
|
||||
) as vllm_model:
|
||||
baseline_outputs = vllm_model.generate_greedy_logprobs(
|
||||
example_prompts, MAX_TOKENS, NUM_LOG_PROBS
|
||||
)
|
||||
|
||||
with vllm_runner(
|
||||
model,
|
||||
max_model_len=MAX_MODEL_LEN,
|
||||
enforce_eager=True,
|
||||
quantization="mxfp8",
|
||||
) as vllm_model:
|
||||
test_outputs = vllm_model.generate_greedy_logprobs(
|
||||
example_prompts, MAX_TOKENS, NUM_LOG_PROBS
|
||||
)
|
||||
|
||||
check_logprobs_close(
|
||||
outputs_0_lst=baseline_outputs,
|
||||
outputs_1_lst=test_outputs,
|
||||
name_0="bf16",
|
||||
name_1="mxfp8",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not is_quant_method_supported("mxfp8"),
|
||||
reason="mxfp8 is not supported on this GPU type (requires sm_100+).",
|
||||
)
|
||||
@pytest.mark.quant_model
|
||||
@pytest.mark.parametrize("model", [DENSE_MODEL, MOE_MODEL], ids=["dense", "moe"])
|
||||
def test_mxfp8_generation(vllm_runner, model: str) -> None:
|
||||
"""Smoke test: verify online MXFP8 model generates coherent text."""
|
||||
prompt = "1 2 3 4 5"
|
||||
with vllm_runner(
|
||||
model,
|
||||
enforce_eager=True,
|
||||
quantization="mxfp8",
|
||||
max_model_len=MAX_MODEL_LEN,
|
||||
) as vllm_model:
|
||||
output = vllm_model.generate_greedy([prompt], max_tokens=5)
|
||||
|
||||
generated = output[0][1]
|
||||
assert len(generated) > len(prompt), (
|
||||
f"MXFP8 model produced no new tokens. Output: {generated!r}"
|
||||
)
|
||||
@@ -320,7 +320,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
|
||||
"tencent/Hunyuan-A13B-Instruct", trust_remote_code=True
|
||||
),
|
||||
"HyperCLOVAXForCausalLM": _HfExamplesInfo(
|
||||
"naver-hyperclovax/HyperCLOVAX-SEED-Think-32B",
|
||||
"naver-hyperclovax/HyperCLOVAX-SEED-Think-14B",
|
||||
trust_remote_code=True,
|
||||
),
|
||||
"InternLMForCausalLM": _HfExamplesInfo(
|
||||
|
||||
@@ -4,6 +4,7 @@ import base64
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
import librosa
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
@@ -71,3 +72,13 @@ def test_audio_media_io_encode_base64(dummy_audio):
|
||||
decoded = base64.b64decode(out)
|
||||
assert decoded == b"dummy_wav_data"
|
||||
mock_write.assert_called_once()
|
||||
|
||||
|
||||
def test_audio_media_io_from_video(video_assets):
|
||||
audio_io = AudioMediaIO()
|
||||
video_path = video_assets[0].video_path
|
||||
with open(video_path, "rb") as f:
|
||||
audio, sr = audio_io.load_bytes(f.read())
|
||||
audio_ref, sr_ref = librosa.load(video_path, sr=None)
|
||||
assert sr == sr_ref
|
||||
np.testing.assert_allclose(audio_ref, audio, atol=1e-4)
|
||||
|
||||
@@ -1,11 +1,19 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import json
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
from vllm.entrypoints.mcp.tool_server import ToolServer
|
||||
from vllm.reasoning import ReasoningParser
|
||||
from vllm.reasoning.gptoss_reasoning_parser import GptOssReasoningParser
|
||||
from vllm.reasoning.gptoss_reasoning_parser import (
|
||||
GptOssReasoningParser,
|
||||
from_builtin_tool_to_tag,
|
||||
no_func_reasoning_tag,
|
||||
)
|
||||
|
||||
REASONING_MODEL_NAME = "openai/gpt-oss-120b"
|
||||
|
||||
@@ -142,3 +150,133 @@ def test_gptoss_is_reasoning_end(
|
||||
output_ids = gpt_oss_tokenizer.convert_tokens_to_ids(output)
|
||||
actual_is_reasoning_end = parser.is_reasoning_end(output_ids)
|
||||
assert is_reasoning_end == actual_is_reasoning_end
|
||||
|
||||
|
||||
class TestGptOssStructuralTags:
|
||||
"""Test cases for GptOssReasoningParser structural tag functionality."""
|
||||
|
||||
@pytest.fixture
|
||||
def mock_tokenizer(self):
|
||||
"""Create a mock tokenizer for testing."""
|
||||
tokenizer = Mock()
|
||||
tokenizer.encode = Mock(return_value=[1, 2, 3, 4, 5])
|
||||
tokenizer.get_vocab = Mock(return_value={"<|end|>": 6})
|
||||
return tokenizer
|
||||
|
||||
@pytest.fixture
|
||||
def reasoning_parser(self, mock_tokenizer):
|
||||
"""Create a GptOssReasoningParser instance."""
|
||||
return GptOssReasoningParser(mock_tokenizer)
|
||||
|
||||
def test_prepare_structured_tag_no_tool_server(self, reasoning_parser):
|
||||
"""Test prepare_structured_tag with no tool server."""
|
||||
result = reasoning_parser.prepare_structured_tag(None, None)
|
||||
expected = json.dumps(no_func_reasoning_tag)
|
||||
|
||||
assert result == expected
|
||||
|
||||
# Verify the structure is correct
|
||||
parsed = json.loads(result)
|
||||
assert parsed["type"] == "structural_tag"
|
||||
assert parsed["format"]["type"] == "triggered_tags"
|
||||
assert len(parsed["format"]["tags"]) == 1
|
||||
assert parsed["format"]["tags"][0]["begin"] == "<|channel|>analysis<|message|>"
|
||||
assert parsed["format"]["triggers"] == ["<|channel|>analysis"]
|
||||
|
||||
def test_prepare_structured_tag_with_original_tag(self, reasoning_parser):
|
||||
"""Test prepare_structured_tag when original_tag is provided."""
|
||||
original_tag = '{"custom": "tag"}'
|
||||
result = reasoning_parser.prepare_structured_tag(original_tag, None)
|
||||
|
||||
# Should return the original tag unchanged
|
||||
assert result == original_tag
|
||||
|
||||
def test_from_builtin_tool_to_tag(self):
|
||||
"""Test from_builtin_tool_to_tag function."""
|
||||
tags = from_builtin_tool_to_tag("python")
|
||||
|
||||
assert len(tags) == 2
|
||||
assert tags[0]["begin"] == "<|channel|>commentary to=python"
|
||||
assert tags[0]["content"]["type"] == "any_text"
|
||||
assert tags[0]["end"] == "<|end|>"
|
||||
|
||||
assert tags[1]["begin"] == "<|channel|>analysis to=python"
|
||||
assert tags[1]["content"]["type"] == "any_text"
|
||||
assert tags[1]["end"] == "<|end|>"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"tools",
|
||||
[
|
||||
[],
|
||||
["browser"],
|
||||
["python"],
|
||||
["container"],
|
||||
["browser", "python"],
|
||||
["browser", "container"],
|
||||
["python", "container"],
|
||||
["browser", "python", "container"],
|
||||
],
|
||||
)
|
||||
def test_json_validity_comprehensive(self, reasoning_parser, tools):
|
||||
"""Test JSON validity across all possible tool combinations."""
|
||||
tool_server = Mock(spec=ToolServer)
|
||||
tool_server.has_tool = Mock(side_effect=lambda tool: tool in tools)
|
||||
|
||||
result = reasoning_parser.prepare_structured_tag(None, tool_server)
|
||||
parsed_result = json.loads(result)
|
||||
|
||||
assert parsed_result["type"] == "structural_tag"
|
||||
assert "format" in parsed_result
|
||||
assert "tags" in parsed_result["format"]
|
||||
assert "triggers" in parsed_result["format"]
|
||||
|
||||
# Tag count should be: 1 (analysis) + 2 * len(tools)
|
||||
expected_tag_count = 1 + (2 * len(tools))
|
||||
assert len(parsed_result["format"]["tags"]) == expected_tag_count
|
||||
|
||||
# Verify triggers are correctly configured
|
||||
expected_triggers = ["<|channel|>analysis"]
|
||||
if tools:
|
||||
expected_triggers.append("<|channel|>commentary to=")
|
||||
assert set(parsed_result["format"]["triggers"]) == set(expected_triggers)
|
||||
|
||||
def test_no_cross_request_state_pollution(self, reasoning_parser):
|
||||
"""Test that sequential calls with different tool servers produce
|
||||
independent results, guarding against shared mutable state
|
||||
(e.g. missing deepcopy in tag_with_builtin_funcs)."""
|
||||
tool_server_1 = Mock(spec=ToolServer)
|
||||
tool_server_1.has_tool = Mock(side_effect=lambda tool: tool == "python")
|
||||
|
||||
tool_server_2 = Mock(spec=ToolServer)
|
||||
tool_server_2.has_tool = Mock(side_effect=lambda tool: tool == "browser")
|
||||
|
||||
result_1 = reasoning_parser.prepare_structured_tag(None, tool_server_1)
|
||||
result_2 = reasoning_parser.prepare_structured_tag(None, tool_server_2)
|
||||
|
||||
tags_1 = [tag["begin"] for tag in json.loads(result_1)["format"]["tags"]]
|
||||
tags_2 = [tag["begin"] for tag in json.loads(result_2)["format"]["tags"]]
|
||||
|
||||
assert "<|channel|>commentary to=python" in tags_1
|
||||
assert "<|channel|>commentary to=browser" not in tags_1
|
||||
|
||||
assert "<|channel|>commentary to=browser" in tags_2
|
||||
assert "<|channel|>commentary to=python" not in tags_2
|
||||
|
||||
def test_tag_format_consistency(self, reasoning_parser):
|
||||
"""Test that all generated tags follow consistent format,
|
||||
catching malformed tags from from_builtin_tool_to_tag."""
|
||||
tool_server = Mock(spec=ToolServer)
|
||||
tool_server.has_tool = Mock(
|
||||
side_effect=lambda tool: tool in ["python", "browser"]
|
||||
)
|
||||
|
||||
result = reasoning_parser.prepare_structured_tag(None, tool_server)
|
||||
parsed_result = json.loads(result)
|
||||
|
||||
for tag in parsed_result["format"]["tags"]:
|
||||
assert "begin" in tag
|
||||
assert "content" in tag
|
||||
assert "end" in tag
|
||||
assert tag["content"]["type"] == "any_text"
|
||||
assert tag["end"] == "<|end|>"
|
||||
assert tag["begin"].startswith("<|channel|>")
|
||||
|
||||
@@ -55,7 +55,7 @@ def test_gc():
|
||||
# The memory allocated for model and KV cache should be released.
|
||||
# The memory allocated for PyTorch and others should be less than 50MB.
|
||||
# Usually, it's around 10MB.
|
||||
allocated = torch.cuda.memory_allocated()
|
||||
allocated = torch.accelerator.memory_allocated()
|
||||
assert allocated < 50 * 1024 * 1024
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,37 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
from unittest.mock import mock_open, patch
|
||||
|
||||
from vllm.platforms import _is_amd_zen_cpu
|
||||
|
||||
|
||||
def test_is_amd_zen_cpu_detects_amd_with_avx512():
|
||||
cpuinfo = "vendor_id: AuthenticAMD\nflags: avx avx2 avx512f avx512bw"
|
||||
with (
|
||||
patch("os.path.exists", return_value=True),
|
||||
patch("builtins.open", mock_open(read_data=cpuinfo)),
|
||||
):
|
||||
assert _is_amd_zen_cpu()
|
||||
|
||||
|
||||
def test_is_amd_zen_cpu_returns_false_for_amd_without_avx512():
|
||||
cpuinfo = "vendor_id: AuthenticAMD\nflags: avx avx2"
|
||||
with (
|
||||
patch("os.path.exists", return_value=True),
|
||||
patch("builtins.open", mock_open(read_data=cpuinfo)),
|
||||
):
|
||||
assert not _is_amd_zen_cpu()
|
||||
|
||||
|
||||
def test_is_amd_zen_cpu_returns_false_for_intel_with_avx512():
|
||||
cpuinfo = "vendor_id: GenuineIntel\nflags: avx avx2 avx512f"
|
||||
with (
|
||||
patch("os.path.exists", return_value=True),
|
||||
patch("builtins.open", mock_open(read_data=cpuinfo)),
|
||||
):
|
||||
assert not _is_amd_zen_cpu()
|
||||
|
||||
|
||||
def test_is_amd_zen_cpu_returns_false_when_cpuinfo_missing():
|
||||
with patch("os.path.exists", return_value=False):
|
||||
assert not _is_amd_zen_cpu()
|
||||
@@ -0,0 +1,378 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import json
|
||||
from dataclasses import dataclass, field
|
||||
from types import NoneType
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.tool_parsers.utils import run_tool_extraction
|
||||
from vllm.tokenizers import TokenizerLike
|
||||
from vllm.tool_parsers import ToolParserManager
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolParserTestConfig:
|
||||
"""Configuration for a tool parser's common tests.
|
||||
|
||||
This dataclass contains all the test data and expected results needed
|
||||
to run the common test suite for a parser. Each parser test file
|
||||
creates one instance of this config with parser-specific values.
|
||||
|
||||
Attributes:
|
||||
parser_name: Name used with ToolParserManager (e.g., "mistral")
|
||||
|
||||
Test data (model outputs):
|
||||
no_tool_calls_output: Plain text without any tool syntax
|
||||
single_tool_call_output: One tool call with simple arguments
|
||||
parallel_tool_calls_output: Multiple tool calls in one response
|
||||
various_data_types_output: Tool with various data types
|
||||
empty_arguments_output: Tool call with no parameters
|
||||
surrounding_text_output: Tool call mixed with regular text
|
||||
escaped_strings_output: Tool call with escaped chars
|
||||
malformed_input_outputs: List of invalid inputs
|
||||
|
||||
Expected results:
|
||||
single_tool_call_expected_name: Expected function name
|
||||
single_tool_call_expected_args: Expected arguments dict
|
||||
parallel_tool_calls_count: Number of tools in parallel test
|
||||
parallel_tool_calls_names: Function names in order
|
||||
single_tool_call_expected_content: Content field when tool called
|
||||
parallel_tool_calls_expected_content: Content for parallel test
|
||||
|
||||
xfail markers:
|
||||
xfail_streaming: Mapping test name to xfail reason (streaming only)
|
||||
xfail_nonstreaming: Mapping test name to xfail reason (non-streaming)
|
||||
|
||||
Special flags:
|
||||
allow_empty_or_json_empty_args: True if "" or "{}" both valid for empty args
|
||||
supports_typed_arguments: True if the parser supports typed function arguments
|
||||
"""
|
||||
|
||||
# Parser identification
|
||||
parser_name: str
|
||||
|
||||
# Test data - model outputs for each common test
|
||||
no_tool_calls_output: str
|
||||
single_tool_call_output: str
|
||||
parallel_tool_calls_output: str
|
||||
various_data_types_output: str
|
||||
empty_arguments_output: str
|
||||
surrounding_text_output: str
|
||||
escaped_strings_output: str
|
||||
malformed_input_outputs: list[str]
|
||||
|
||||
# Expected results for specific tests (optional overrides)
|
||||
single_tool_call_expected_name: str = "get_weather"
|
||||
single_tool_call_expected_args: dict[str, Any] = field(
|
||||
default_factory=lambda: {"city": "Tokyo"}
|
||||
)
|
||||
parallel_tool_calls_count: int = 2
|
||||
parallel_tool_calls_names: list[str] = field(
|
||||
default_factory=lambda: ["get_weather", "get_time"]
|
||||
)
|
||||
|
||||
# xfail configuration - maps test name to xfail reason
|
||||
xfail_streaming: dict[str, str] = field(default_factory=dict)
|
||||
xfail_nonstreaming: dict[str, str] = field(default_factory=dict)
|
||||
|
||||
# Content expectations (some parsers strip content, others don't)
|
||||
single_tool_call_expected_content: str | None = None
|
||||
parallel_tool_calls_expected_content: str | None = None
|
||||
|
||||
# Special assertions for edge cases
|
||||
allow_empty_or_json_empty_args: bool = True # "{}" or "" for empty args
|
||||
supports_typed_arguments: bool = True
|
||||
|
||||
|
||||
class ToolParserTests:
|
||||
"""Mixin class providing common test suite for tool parsers.
|
||||
|
||||
To use this mixin in a parser test file:
|
||||
|
||||
1. Create a test_config fixture that returns a ToolParserTestConfig instance
|
||||
2. Inherit from this class
|
||||
3. Add parser-specific tests as additional methods
|
||||
|
||||
Example:
|
||||
class TestMistralToolParser(ToolParserTests):
|
||||
@pytest.fixture
|
||||
def test_config(self) -> ToolParserTestConfig:
|
||||
return ToolParserTestConfig(
|
||||
parser_name="mistral",
|
||||
no_tool_calls_output="Plain text...",
|
||||
# ... other config ...
|
||||
)
|
||||
|
||||
# Parser-specific tests
|
||||
def test_mistral_specific_feature(self, tool_parser):
|
||||
# Custom test logic
|
||||
pass
|
||||
"""
|
||||
|
||||
@pytest.fixture
|
||||
def test_config(self) -> ToolParserTestConfig:
|
||||
"""Override this to provide parser-specific configuration."""
|
||||
raise NotImplementedError(
|
||||
"Subclass must provide test_config fixture returning ToolParserTestConfig"
|
||||
)
|
||||
|
||||
@pytest.fixture
|
||||
def tokenizer(self, default_tokenizer: TokenizerLike) -> TokenizerLike:
|
||||
"""Override this to provide parser-specific tokenizer."""
|
||||
return default_tokenizer
|
||||
|
||||
@pytest.fixture
|
||||
def tool_parser(self, test_config: ToolParserTestConfig, tokenizer: TokenizerLike):
|
||||
return ToolParserManager.get_tool_parser(test_config.parser_name)(tokenizer)
|
||||
|
||||
@pytest.fixture(params=[True, False])
|
||||
def streaming(self, request: pytest.FixtureRequest) -> bool:
|
||||
return request.param
|
||||
|
||||
def test_no_tool_calls(
|
||||
self,
|
||||
request: pytest.FixtureRequest,
|
||||
tool_parser: Any,
|
||||
test_config: ToolParserTestConfig,
|
||||
streaming: bool,
|
||||
):
|
||||
"""Verify parser handles plain text without tool syntax."""
|
||||
# Apply xfail markers if configured
|
||||
test_name = "test_no_tool_calls"
|
||||
self.apply_xfail_mark(request, test_config, test_name, streaming)
|
||||
|
||||
content, tool_calls = run_tool_extraction(
|
||||
tool_parser, test_config.no_tool_calls_output, streaming=streaming
|
||||
)
|
||||
assert content == test_config.no_tool_calls_output, (
|
||||
f"Expected content to match input, got {content}"
|
||||
)
|
||||
assert len(tool_calls) == 0, f"Expected no tool calls, got {len(tool_calls)}"
|
||||
|
||||
def test_single_tool_call_simple_args(
|
||||
self,
|
||||
request: pytest.FixtureRequest,
|
||||
tool_parser: Any,
|
||||
test_config: ToolParserTestConfig,
|
||||
streaming: bool,
|
||||
):
|
||||
"""Verify parser extracts one tool with simple arguments."""
|
||||
# Apply xfail markers if configured
|
||||
test_name = "test_single_tool_call_simple_args"
|
||||
self.apply_xfail_mark(request, test_config, test_name, streaming)
|
||||
|
||||
content, tool_calls = run_tool_extraction(
|
||||
tool_parser, test_config.single_tool_call_output, streaming=streaming
|
||||
)
|
||||
|
||||
# Content check (some parsers strip it)
|
||||
if test_config.single_tool_call_expected_content is not None:
|
||||
assert content == test_config.single_tool_call_expected_content
|
||||
|
||||
assert len(tool_calls) == 1, f"Expected 1 tool call, got {len(tool_calls)}"
|
||||
assert tool_calls[0].type == "function"
|
||||
assert tool_calls[0].function.name == test_config.single_tool_call_expected_name
|
||||
|
||||
args = json.loads(tool_calls[0].function.arguments)
|
||||
for key, value in test_config.single_tool_call_expected_args.items():
|
||||
assert args.get(key) == value, (
|
||||
f"Expected {key}={value}, got {args.get(key)}"
|
||||
)
|
||||
|
||||
def test_parallel_tool_calls(
|
||||
self,
|
||||
request: pytest.FixtureRequest,
|
||||
tool_parser: Any,
|
||||
test_config: ToolParserTestConfig,
|
||||
streaming: bool,
|
||||
):
|
||||
"""Verify parser handles multiple tools in one response."""
|
||||
# Apply xfail markers if configured
|
||||
test_name = "test_parallel_tool_calls"
|
||||
self.apply_xfail_mark(request, test_config, test_name, streaming)
|
||||
|
||||
content, tool_calls = run_tool_extraction(
|
||||
tool_parser,
|
||||
test_config.parallel_tool_calls_output,
|
||||
streaming=streaming,
|
||||
)
|
||||
|
||||
assert len(tool_calls) == test_config.parallel_tool_calls_count, (
|
||||
f"Expected {test_config.parallel_tool_calls_count} "
|
||||
f"tool calls, got {len(tool_calls)}"
|
||||
)
|
||||
|
||||
# Verify tool names match expected
|
||||
for i, expected_name in enumerate(test_config.parallel_tool_calls_names):
|
||||
assert tool_calls[i].type == "function"
|
||||
assert tool_calls[i].function.name == expected_name
|
||||
|
||||
# Verify unique IDs
|
||||
ids = [tc.id for tc in tool_calls]
|
||||
assert len(ids) == len(set(ids)), "Tool call IDs should be unique"
|
||||
|
||||
def test_various_data_types(
|
||||
self,
|
||||
request: pytest.FixtureRequest,
|
||||
tool_parser: Any,
|
||||
test_config: ToolParserTestConfig,
|
||||
streaming: bool,
|
||||
):
|
||||
"""Verify parser handles all JSON types in arguments."""
|
||||
# Apply xfail markers if configured
|
||||
test_name = "test_various_data_types"
|
||||
self.apply_xfail_mark(request, test_config, test_name, streaming)
|
||||
|
||||
content, tool_calls = run_tool_extraction(
|
||||
tool_parser,
|
||||
test_config.various_data_types_output,
|
||||
streaming=streaming,
|
||||
)
|
||||
assert len(tool_calls) == 1, f"Expected 1 tool call, got {len(tool_calls)}"
|
||||
|
||||
args = json.loads(tool_calls[0].function.arguments)
|
||||
# Verify all expected fields present
|
||||
required_fields_types = {
|
||||
"string_field": str,
|
||||
"int_field": int,
|
||||
"float_field": float,
|
||||
"bool_field": bool,
|
||||
"null_field": NoneType,
|
||||
"array_field": list,
|
||||
"object_field": dict,
|
||||
}
|
||||
for required_field, expected_type in required_fields_types.items():
|
||||
assert required_field in args, (
|
||||
f"Expected field '{required_field}' in arguments"
|
||||
)
|
||||
if test_config.supports_typed_arguments:
|
||||
found_type = type(args[required_field])
|
||||
assert found_type is expected_type, (
|
||||
f"Expected field '{required_field}' to have type {expected_type}, "
|
||||
f"got {found_type}"
|
||||
)
|
||||
|
||||
def test_empty_arguments(
|
||||
self,
|
||||
request: pytest.FixtureRequest,
|
||||
tool_parser: Any,
|
||||
test_config: ToolParserTestConfig,
|
||||
streaming: bool,
|
||||
):
|
||||
"""Verify parser handles parameterless tool calls."""
|
||||
# Apply xfail markers if configured
|
||||
test_name = "test_empty_arguments"
|
||||
self.apply_xfail_mark(request, test_config, test_name, streaming)
|
||||
|
||||
content, tool_calls = run_tool_extraction(
|
||||
tool_parser, test_config.empty_arguments_output, streaming=streaming
|
||||
)
|
||||
assert len(tool_calls) == 1, f"Expected 1 tool call, got {len(tool_calls)}"
|
||||
|
||||
args = tool_calls[0].function.arguments
|
||||
if test_config.allow_empty_or_json_empty_args:
|
||||
assert args in ["{}", ""], f"Expected empty args, got {args}"
|
||||
else:
|
||||
assert args == "{}", f"Expected {{}}, got {args}"
|
||||
|
||||
def test_surrounding_text(
|
||||
self,
|
||||
request: pytest.FixtureRequest,
|
||||
tool_parser: Any,
|
||||
test_config: ToolParserTestConfig,
|
||||
streaming: bool,
|
||||
):
|
||||
"""Verify parser extracts tools from mixed content."""
|
||||
# Apply xfail markers if configured
|
||||
test_name = "test_surrounding_text"
|
||||
self.apply_xfail_mark(request, test_config, test_name, streaming)
|
||||
|
||||
content, tool_calls = run_tool_extraction(
|
||||
tool_parser, test_config.surrounding_text_output, streaming=streaming
|
||||
)
|
||||
assert len(tool_calls) >= 1, (
|
||||
f"Expected at least 1 tool call, got {len(tool_calls)}"
|
||||
)
|
||||
|
||||
def test_escaped_strings(
|
||||
self,
|
||||
request: pytest.FixtureRequest,
|
||||
tool_parser: Any,
|
||||
test_config: ToolParserTestConfig,
|
||||
streaming: bool,
|
||||
):
|
||||
"""Verify parser handles escaped characters in arguments."""
|
||||
# Apply xfail markers if configured
|
||||
test_name = "test_escaped_strings"
|
||||
self.apply_xfail_mark(request, test_config, test_name, streaming)
|
||||
|
||||
content, tool_calls = run_tool_extraction(
|
||||
tool_parser, test_config.escaped_strings_output, streaming=streaming
|
||||
)
|
||||
assert len(tool_calls) == 1, f"Expected 1 tool call, got {len(tool_calls)}"
|
||||
|
||||
args = json.loads(tool_calls[0].function.arguments)
|
||||
# At minimum, verify we can parse and have expected fields
|
||||
# Exact escaping behavior varies by parser
|
||||
assert len(args) > 0, "Expected some arguments with escaped strings"
|
||||
|
||||
def test_malformed_input(
|
||||
self,
|
||||
request: pytest.FixtureRequest,
|
||||
tool_parser: Any,
|
||||
test_config: ToolParserTestConfig,
|
||||
streaming: bool,
|
||||
):
|
||||
"""Verify parser gracefully handles invalid syntax."""
|
||||
# Apply xfail markers if configured
|
||||
test_name = "test_malformed_input"
|
||||
self.apply_xfail_mark(request, test_config, test_name, streaming)
|
||||
|
||||
for malformed_input in test_config.malformed_input_outputs:
|
||||
# Should not raise exception
|
||||
content, tool_calls = run_tool_extraction(
|
||||
tool_parser, malformed_input, streaming=streaming
|
||||
)
|
||||
# Parser should handle gracefully (exact behavior varies)
|
||||
|
||||
def test_streaming_reconstruction(
|
||||
self,
|
||||
request: pytest.FixtureRequest,
|
||||
tool_parser: Any,
|
||||
test_config: ToolParserTestConfig,
|
||||
):
|
||||
"""Verify streaming produces same result as non-streaming."""
|
||||
test_name = "test_streaming_reconstruction"
|
||||
self.apply_xfail_mark(request, test_config, test_name, True)
|
||||
|
||||
test_output = test_config.single_tool_call_output
|
||||
|
||||
# Non-streaming result
|
||||
content_non, tools_non = run_tool_extraction(
|
||||
tool_parser, test_output, streaming=False
|
||||
)
|
||||
|
||||
# Streaming result
|
||||
content_stream, tools_stream = run_tool_extraction(
|
||||
tool_parser, test_output, streaming=True
|
||||
)
|
||||
|
||||
# Compare results
|
||||
assert content_non == content_stream, "Content should match between modes"
|
||||
assert len(tools_non) == len(tools_stream), "Tool count should match"
|
||||
if len(tools_non) > 0:
|
||||
assert tools_non[0].function.name == tools_stream[0].function.name
|
||||
assert tools_non[0].function.arguments == tools_stream[0].function.arguments
|
||||
|
||||
def apply_xfail_mark(self, request, test_config, test_name, streaming):
|
||||
reason = None
|
||||
if streaming and test_name in test_config.xfail_streaming:
|
||||
reason = test_config.xfail_streaming[test_name]
|
||||
elif not streaming and test_name in test_config.xfail_nonstreaming:
|
||||
reason = test_config.xfail_nonstreaming[test_name]
|
||||
if reason is not None:
|
||||
mark = pytest.mark.xfail(reason=reason, strict=True)
|
||||
request.node.add_marker(mark)
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user