karylab_agents/vllm
1
0
Fork 0
forked from Karylab-cklius/vllm
Code Pull Requests 2 Actions 1 Packages Activity

Merge branch 'main' into zhuohan/remove-unnecessary-instance_id-setup

Browse Source
This commit is contained in:
Cyrus Leung
2026-03-12 15:52:39 +08:00
committed by GitHub
parent 72ee63dd34 3e64fe4a18
commit 423ff4ebaa
321 changed files with 42382 additions and 33118 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.buildkite/performance-benchmarks/scripts/compare-json-results.py
+300 -89
View File
@@ -7,12 +7,12 @@ import argparse
import html as _html
import json
import os
from contextlib import nullcontext
from dataclasses import dataclass
from importlib import util
from pathlib import Path
import pandas as pd
import regex as re
pd.options.display.float_format = "{:.2f}".format
plotly_found = util.find_spec("plotly.express") is not None
@@ -33,6 +33,45 @@ pd.set_option("display.precision", 2)
pd.set_option("display.float_format", lambda x: f"{x:.2f}")
# -----------------------------
# Concurrency normalization (NEW, small)
# -----------------------------
def _find_concurrency_col(df: pd.DataFrame) -> str:
for c in [
"# of max concurrency.",
"# of max concurrency",
"Max Concurrency",
"max_concurrency",
"Concurrency",
]:
if c in df.columns:
return c
for c in df.columns:
if "concurr" in str(c).lower():
s = df[c]
if s.dtype.kind in "iu" and s.nunique() > 1 and s.min() >= 1:
return c
raise ValueError(
"Cannot infer concurrency column. "
"Please rename the column to one of the known names "
"or add an explicit override (e.g., --concurrency-col)."
)
def _normalize_concurrency_in_df(
df: pd.DataFrame, canonical: str = "# of max concurrency."
) -> pd.DataFrame:
if canonical in df.columns:
return df
detected = _find_concurrency_col(df)
if detected in df.columns and detected != canonical:
return df.rename(columns={detected: canonical})
df[canonical] = pd.NA
return df
# -----------------------------
# Core data compare
# -----------------------------
@@ -52,19 +91,25 @@ def compare_data_columns(
- Concat along axis=1 (indexes align), then reset_index so callers can
group by columns.
- If --debug, add a <file_label>_name column per file.
Minimal fix to support different max_concurrency lists across files:
- normalize concurrency column naming to "# of max concurrency."
- align on UNION of keys (missing points become NaN)
- BUGFIX: don't drop throughput rows based on P99/Median presence
"""
print("\ncompare_data_column:", data_column)
frames = []
raw_data_cols: list[str] = []
compare_frames = []
# Determine key cols after normalizing concurrency
cols_per_file: list[set] = []
for f in files:
try:
df_tmp = pd.read_json(f, orient="records")
except Exception as err:
raise ValueError(f"Failed to read {f}") from err
df_tmp = _normalize_concurrency_in_df(df_tmp, canonical="# of max concurrency.")
cols_per_file.append(set(df_tmp.columns))
key_cols = [c for c in info_cols if all(c in cset for cset in cols_per_file)]
@@ -75,12 +120,25 @@ def compare_data_columns(
"No common key columns found from info_cols across the input files."
)
meta_added = False
union_index = None
metas: list[pd.DataFrame] = []
staged: list[tuple[str, pd.Series, pd.Series | None]] = []
for file in files:
df = pd.read_json(file, orient="records")
df = _normalize_concurrency_in_df(df, canonical="# of max concurrency.")
if drop_column in df.columns:
# BUGFIX: only drop rows for latency-like metrics; throughput rows may have
# NaN in P99/Median columns even if the column exists in the JSON.
metric_lc = str(data_column).lower()
is_latency_metric = (
"ttft" in metric_lc
or "tpot" in metric_lc
or "p99" in metric_lc
or "median" in metric_lc
or metric_lc.strip() in {"p99", "median"}
)
if is_latency_metric and drop_column in df.columns:
df = df.dropna(subset=[drop_column], ignore_index=True)
for c in (
@@ -105,35 +163,61 @@ def compare_data_columns(
meta = meta.groupby(level=key_cols, dropna=False).first()
file_label = "/".join(file.split("/")[:-1]) or os.path.basename(file)
s = df_idx[data_column]
if not s.index.is_unique:
s = s.groupby(level=key_cols, dropna=False).mean()
if data_column in df_idx.columns:
s = df_idx[data_column]
if not s.index.is_unique:
s = s.groupby(level=key_cols, dropna=False).mean()
else:
# keep NA series to preserve meta keys for union_index
s = pd.Series(pd.NA, index=meta.index)
s.name = file_label
if not meta_added:
frames.append(meta)
meta_added = True
name_s = None
if debug and name_column in df_idx.columns:
name_s = df_idx[name_column]
if not name_s.index.is_unique:
name_s = name_s.groupby(level=key_cols, dropna=False).first()
name_s.name = f"{file_label}_name"
frames.append(name_s)
frames.append(s)
if union_index is None:
union_index = meta.index
else:
union_index = union_index.union(meta.index)
metas.append(meta)
staged.append((file_label, s, name_s))
if union_index is None:
raise ValueError("No data found after loading inputs.")
# meta first (union-aligned): build UNION meta across all files
if metas:
meta_union = pd.concat(metas, axis=0)
# Collapse duplicates on the MultiIndex; keep first non-null per column
meta_union = meta_union.groupby(level=key_cols, dropna=False).first()
frames.append(meta_union.reindex(union_index))
# values + ratios (union-aligned)
metric_series_aligned: list[pd.Series] = []
for file_label, s, name_s in staged:
s_aligned = s.reindex(union_index)
frames.append(s_aligned)
raw_data_cols.append(file_label)
compare_frames.append(s)
metric_series_aligned.append(s_aligned)
if len(compare_frames) >= 2:
base = compare_frames[0]
current = compare_frames[-1]
if "P99" in data_column or "Median" in data_column:
if debug and name_s is not None:
frames.append(name_s.reindex(union_index))
if len(metric_series_aligned) >= 2:
base = metric_series_aligned[0]
current = metric_series_aligned[-1]
if "P99" in str(data_column) or "Median" in str(data_column):
ratio = base / current
else:
ratio = current / base
ratio = ratio.mask(base == 0)
ratio.name = f"Ratio 1 vs {len(compare_frames)}"
ratio.name = f"Ratio 1 vs {len(metric_series_aligned)}"
frames.append(ratio)
concat_df = pd.concat(frames, axis=1).reset_index(drop=True)
@@ -204,24 +288,10 @@ def split_json_by_tp_pp(
# -----------------------------
# Styling helpers
# -----------------------------
def _find_concurrency_col(df: pd.DataFrame) -> str:
for c in [
"# of max concurrency.",
"# of max concurrency",
"Max Concurrency",
"max_concurrency",
"Concurrency",
]:
if c in df.columns:
return c
for c in df.columns:
if df[c].dtype.kind in "iu" and df[c].nunique() > 1 and df[c].min() >= 1:
return c
return "# of max concurrency."
def _highlight_threshold(
df: pd.DataFrame, threshold: float
df: pd.DataFrame,
threshold: float,
slack_pct: float = 0.0,
) -> pd.io.formats.style.Styler:
conc_col = _find_concurrency_col(df)
key_cols = [
@@ -234,12 +304,24 @@ def _highlight_threshold(
]
conf_cols = [c for c in conf_cols if pd.api.types.is_numeric_dtype(df[c])]
return df.style.map(
lambda v: "background-color:#e6ffe6;font-weight:bold;"
if pd.notna(v) and v <= threshold
else "",
subset=conf_cols,
)
try:
slack_pct = float(slack_pct or 0.0)
except Exception:
slack_pct = 0.0
slack_limit = threshold * (1.0 + slack_pct / 100.0)
def _cell(v):
if pd.isna(v):
return ""
if v <= threshold:
# Strict SLA
return "background-color:#e6ffe6;font-weight:bold;"
if v <= slack_limit:
# Within slack range
return "background-color:#ffe5cc;font-weight:bold;"
return ""
return df.style.map(_cell, subset=conf_cols)
def highlight_ratio_columns(styler: pd.io.formats.style.Styler):
@@ -286,11 +368,30 @@ def _sanitize_sheet_name(name: str) -> str:
- max 31 chars
- cannot contain: : \ / ? * [ ]
- cannot be empty
NOTE: Use fast, non-regex operations here to avoid the third-party `regex`
module's compile overhead/edge-cases on some systems.
"""
name = "sheet" if name is None else str(name)
name = re.sub(r"[:\\/?*\[\]]", "_", name)
# Replace illegal characters with underscore.
trans = str.maketrans(
{
":": "_",
"\\": "_",
"/": "_",
"?": "_",
"*": "_",
"[": "_",
"]": "_",
}
)
name = name.translate(trans)
# Strip quotes/spaces and collapse whitespace.
name = name.strip().strip("'")
name = re.sub(r"\s+", " ", name)
name = " ".join(name.split())
if not name:
name = "sheet"
return name[:31]
@@ -298,30 +399,57 @@ def _sanitize_sheet_name(name: str) -> str:
def _group_to_sheet_base(group_cols: list[str], gkey_tuple) -> str:
d = dict(zip(group_cols, gkey_tuple))
model = d.get("Model", "model")
model_short = str(model).split("/")[-1]
# Always keep input/output lengths (these are important).
ilen = d.get("Input Len", "")
olen = d.get("Output Len", "")
lens = f"_{ilen}x{olen}" if ilen != "" and olen != "" else ""
# Shorten model name aggressively to make room for lens.
model = d.get("Model", "model")
leaf = str(model).split("/")[-1]
max_model_len = max(1, 31 - len(lens))
model_short = leaf[:max_model_len]
return _sanitize_sheet_name(f"{model_short}{lens}")
def _write_tables_to_excel_sheet(
writer: pd.ExcelWriter, sheet: str, blocks: list[tuple[str, pd.DataFrame]]
):
startrow = 0
"""Write all blocks to a sheet with a single to_excel() call.
Pandas+openpyxl can be extremely slow when called many times per sheet.
We flatten blocks into one table with a 'Section' column to keep structure
while making Excel generation fast and deterministic.
"""
if not blocks:
pd.DataFrame().to_excel(writer, sheet_name=sheet, index=False)
return
combined_parts: list[pd.DataFrame] = []
for title, df in blocks:
pd.DataFrame([[title]]).to_excel(
writer, sheet_name=sheet, index=False, header=False, startrow=startrow
)
startrow += 1
df.to_excel(writer, sheet_name=sheet, index=False, startrow=startrow)
startrow += len(df) + 3
df2 = df.copy()
# Put the section label as the first column for readability.
df2.insert(0, "Section", title)
combined_parts.append(df2)
combined = pd.concat(combined_parts, axis=0, ignore_index=True, sort=False)
combined.to_excel(writer, sheet_name=sheet, index=False)
def _safe_filename(s: str) -> str:
s = re.sub(r"[^\w\-.]+", "_", str(s).strip())
return s[:180] if len(s) > 180 else s
# Fast path without the third-party `regex` module.
s = " ".join(str(s).strip().split())
allowed = []
for ch in s:
if ch.isalnum() or ch in "._-":
allowed.append(ch)
else:
allowed.append("_")
out = "".join(allowed)
return out[:180] if len(out) > 180 else out
# -----------------------------
@@ -428,7 +556,11 @@ def _config_value_columns(df: pd.DataFrame, conc_col: str) -> list[str]:
def _max_concurrency_ok(
df: pd.DataFrame, conc_col: str, cfg_col: str, threshold: float
df: pd.DataFrame,
conc_col: str,
cfg_col: str,
threshold: float,
slack_pct: float = 0.0,
):
if df is None or conc_col not in df.columns or cfg_col not in df.columns:
return pd.NA
@@ -441,7 +573,14 @@ def _max_concurrency_ok(
if d.empty:
return pd.NA
ok = d[d[cfg_col] <= threshold]
# Accept values up to (1 + slack_pct%) above the SLA.
try:
slack_pct = float(slack_pct or 0.0)
except Exception:
slack_pct = 0.0
effective_limit = float(threshold) * (1.0 + slack_pct / 100.0)
ok = d[d[cfg_col] <= effective_limit]
if ok.empty:
return pd.NA
@@ -507,15 +646,25 @@ def build_valid_max_concurrency_summary_html(
if not cfg_cols:
cfg_cols = sorted(set(ttft_cols) | set(tpot_cols) | set(tput_cols), key=str)
# Display SLA ranges in the table header (SLA .. SLA*(1+slack))
ttft_hi = args.ttft_max_ms * (1.0 + args.ttft_slack_pct / 100.0)
tpot_hi = args.tpot_max_ms * (1.0 + args.tpot_slack_pct / 100.0)
ttft_range = f"{args.ttft_max_ms:g}–{ttft_hi:g} ms (+{args.ttft_slack_pct:g}%)"
tpot_range = f"{args.tpot_max_ms:g}–{tpot_hi:g} ms (+{args.tpot_slack_pct:g}%)"
rows = []
for cfg in cfg_cols:
ttft_max = (
_max_concurrency_ok(ttft_group_df, conc_col, cfg, args.ttft_max_ms)
_max_concurrency_ok(
ttft_group_df, conc_col, cfg, args.ttft_max_ms, args.ttft_slack_pct
)
if ttft_group_df is not None
else pd.NA
)
tpot_max = (
_max_concurrency_ok(tpot_group_df, conc_col, cfg, args.tpot_max_ms)
_max_concurrency_ok(
tpot_group_df, conc_col, cfg, args.tpot_max_ms, args.tpot_slack_pct
)
if tpot_group_df is not None
else pd.NA
)
@@ -544,8 +693,8 @@ def build_valid_max_concurrency_summary_html(
rows.append(
{
"Configuration": cfg,
f"Max {conc_col} (TTFT ≤ {args.ttft_max_ms:g} ms)": ttft_max,
f"Max {conc_col} (TPOT ≤ {args.tpot_max_ms:g} ms)": tpot_max,
f"Max {conc_col} (TTFT ≤ {ttft_range})": ttft_max,
f"Max {conc_col} (TPOT ≤ {tpot_range})": tpot_max,
f"Max {conc_col} (Both)": both,
"Output Tput @ Both (tok/s)": tput_at_both,
"TTFT @ Both (ms)": ttft_at_both,
@@ -620,15 +769,24 @@ def build_valid_max_concurrency_summary_df(
if not cfg_cols:
cfg_cols = sorted(set(ttft_cols) | set(tpot_cols) | set(tput_cols), key=str)
ttft_hi = args.ttft_max_ms * (1.0 + args.ttft_slack_pct / 100.0)
tpot_hi = args.tpot_max_ms * (1.0 + args.tpot_slack_pct / 100.0)
ttft_range = f"{args.ttft_max_ms:g}–{ttft_hi:g} ms (+{args.ttft_slack_pct:g}%)"
tpot_range = f"{args.tpot_max_ms:g}–{tpot_hi:g} ms (+{args.tpot_slack_pct:g}%)"
rows = []
for cfg in cfg_cols:
ttft_max = (
_max_concurrency_ok(ttft_group_df, conc_col, cfg, args.ttft_max_ms)
_max_concurrency_ok(
ttft_group_df, conc_col, cfg, args.ttft_max_ms, args.ttft_slack_pct
)
if ttft_group_df is not None
else pd.NA
)
tpot_max = (
_max_concurrency_ok(tpot_group_df, conc_col, cfg, args.tpot_max_ms)
_max_concurrency_ok(
tpot_group_df, conc_col, cfg, args.tpot_max_ms, args.tpot_slack_pct
)
if tpot_group_df is not None
else pd.NA
)
@@ -657,8 +815,8 @@ def build_valid_max_concurrency_summary_df(
rows.append(
{
"Configuration": cfg,
f"Max {conc_col} (TTFT ≤ {args.ttft_max_ms:g} ms)": ttft_max,
f"Max {conc_col} (TPOT ≤ {args.tpot_max_ms:g} ms)": tpot_max,
f"Max {conc_col} (TTFT ≤ {ttft_range})": ttft_max,
f"Max {conc_col} (TPOT ≤ {tpot_range})": tpot_max,
f"Max {conc_col} (Both)": both,
"Output Tput @ Both (tok/s)": tput_at_both,
"TTFT @ Both (ms)": ttft_at_both,
@@ -751,7 +909,21 @@ def build_parser() -> argparse.ArgumentParser:
help="Reference limit for TPOT plots (ms)",
)
# ---- NEW: export options ----
# ---- SLA tolerance (slack) options ----
parser.add_argument(
"--ttft-slack-pct",
type=float,
default=5.0,
help="Allowed percentage above TTFT SLA (default: 5).",
)
parser.add_argument(
"--tpot-slack-pct",
type=float,
default=5.0,
help="Allowed percentage above TPOT SLA (default: 5).",
)
# ---- export options ----
parser.add_argument(
"--excel-out",
type=str,
@@ -843,9 +1015,13 @@ def render_metric_table_html(
metric_name = metric_label.lower()
if "ttft" in metric_name:
styler = _highlight_threshold(display_group, args.ttft_max_ms)
styler = _highlight_threshold(
display_group, args.ttft_max_ms, args.ttft_slack_pct
)
elif ("tpot" in metric_name) or ("median" in metric_name) or ("p99" in metric_name):
styler = _highlight_threshold(display_group, args.tpot_max_ms)
styler = _highlight_threshold(
display_group, args.tpot_max_ms, args.tpot_slack_pct
)
else:
styler = display_group.style
@@ -962,22 +1138,46 @@ def write_report_group_first(
csv_dir.mkdir(parents=True, exist_ok=True)
excel_path = args.excel_out or "perf_comparison.xlsx"
with pd.ExcelWriter(excel_path, engine="openpyxl") as xw:
disable_excel = os.getenv("VLLM_COMPARE_DISABLE_EXCEL", "0") == "1"
# Prefer xlsxwriter for speed; fallback to openpyxl if unavailable.
excel_engine = (
os.getenv("VLLM_COMPARE_EXCEL_ENGINE", "xlsxwriter").strip() or "xlsxwriter"
)
if excel_engine == "xlsxwriter" and util.find_spec("xlsxwriter") is None:
excel_engine = "openpyxl"
excel_engine_kwargs = {}
if excel_engine == "xlsxwriter":
# Reduce memory pressure & usually faster writes.
excel_engine_kwargs = {"options": {"constant_memory": True}}
xw_ctx = (
nullcontext(None)
if disable_excel
else pd.ExcelWriter(
excel_path, engine=excel_engine, engine_kwargs=excel_engine_kwargs
)
)
with xw_ctx as xw:
used_sheets: set[str] = set()
# ---- Environment sheet (first) ----
env_sheet = _sanitize_sheet_name("Environment")
env_df = _load_env_df_for_inputs(args, files)
if env_df is None or env_df.empty:
pd.DataFrame(
[
{
"Section": "Environment",
"Key": "vllm_env.txt",
"Value": "NOT FOUND (or empty)",
}
]
).to_excel(xw, sheet_name=env_sheet, index=False)
else:
env_df.to_excel(xw, sheet_name=env_sheet, index=False)
if xw is not None:
if env_df is None or env_df.empty:
pd.DataFrame(
[
{
"Section": "Environment",
"Key": "vllm_env.txt",
"Value": "NOT FOUND (or empty)",
}
]
).to_excel(xw, sheet_name=env_sheet, index=False)
else:
env_df.to_excel(xw, sheet_name=env_sheet, index=False)
used_sheets.add(env_sheet)
with open("perf_comparison.html", "w", encoding="utf-8") as main_fh:
main_fh.write('<meta charset="utf-8">\n')
for gkey in group_keys:
@@ -993,12 +1193,19 @@ def write_report_group_first(
main_fh.write(group_header)
do_excel = xw is not None
sheet = _group_to_sheet_base(group_cols_canonical, gkey_tuple)
sheet_base = sheet
dedup_i = 1
while sheet in xw.sheets:
dedup_i += 1
sheet = _sanitize_sheet_name(f"{sheet_base}_{dedup_i}")
if do_excel:
dedup_i = 1
while sheet in used_sheets:
dedup_i += 1
suffix = f"_{dedup_i}"
# Ensure uniqueness even when sheet names are truncated.
base = str(sheet_base)
keep = max(1, 31 - len(suffix))
sheet = _sanitize_sheet_name(base[:keep] + suffix)
used_sheets.add(sheet)
excel_blocks: list[tuple[str, pd.DataFrame]] = []
@@ -1059,7 +1266,7 @@ def write_report_group_first(
)
excel_blocks.append(
(metric_label, display_group.reset_index(drop=True))
(metric_label, group_df.reset_index(drop=True))
)
if csv_dir:
fn = _safe_filename(
@@ -1067,7 +1274,7 @@ def write_report_group_first(
"/", "_"
)
)
display_group.to_csv(csv_dir / f"{fn}.csv", index=False)
group_df.to_csv(csv_dir / f"{fn}.csv", index=False)
summary_html = build_valid_max_concurrency_summary_html(
tput_group_df=tput_group_df,
@@ -1097,9 +1304,13 @@ def write_report_group_first(
)
summary_df.to_csv(csv_dir / f"{fn}.csv", index=False)
_write_tables_to_excel_sheet(xw, sheet, excel_blocks)
if do_excel:
_write_tables_to_excel_sheet(xw, sheet, excel_blocks)
print(f"Wrote Excel: {excel_path}")
if disable_excel:
print("Skipped Excel generation (VLLM_COMPARE_DISABLE_EXCEL=1).")
else:
print(f"Wrote Excel: {excel_path}")
if csv_dir:
print(f"Wrote CSVs under: {csv_dir}")
.buildkite/performance-benchmarks/scripts/run-performance-benchmarks.sh
Executable → Regular
+361 -4
View File
@@ -12,6 +12,13 @@ DRY_RUN="${DRY_RUN:-0}"
MODEL_FILTER="${MODEL_FILTER:-}"
DTYPE_FILTER="${DTYPE_FILTER:-}"
# Adaptive search controls
ENABLE_ADAPTIVE_CONCURRENCY="${ENABLE_ADAPTIVE_CONCURRENCY:-0}"
SLA_TTFT_MS="${SLA_TTFT_MS:-3000}"
SLA_TPOT_MS="${SLA_TPOT_MS:-100}"
ADAPTIVE_MAX_PROBES="${ADAPTIVE_MAX_PROBES:-8}"
ADAPTIVE_MAX_CONCURRENCY="${ADAPTIVE_MAX_CONCURRENCY:-1024}"
check_gpus() {
if command -v nvidia-smi; then
# check the number of GPUs and GPU type.
@@ -183,6 +190,304 @@ upload_to_buildkite() {
$BUILDKITE_AGENT_COMMAND artifact upload "$RESULTS_FOLDER/*"
}
# -------------------------------
# Adaptive concurrency helpers
# -------------------------------
result_json_path_for_serving() {
local test_name=$1
local qps=$2
local max_concurrency=$3
echo "$RESULTS_FOLDER/${test_name}_qps_${qps}_concurrency_${max_concurrency}.json"
}
extract_metric_ms() {
local metric_name=$1
local json_file=$2
[[ -f "$json_file" ]] || return 0
if [[ "$metric_name" == "ttft" ]]; then
jq -r '
[
.ttft_ms.p99?,
.metrics.ttft_ms.p99?,
.ttft.p99?,
.metrics.ttft.p99?,
.p99_ttft_ms?,
.ttft_ms.mean?,
.metrics.ttft_ms.mean?,
.ttft.mean?,
.metrics.ttft.mean?,
.mean_ttft_ms?
] | map(select(. != null)) | .[0] // empty
' "$json_file"
else
jq -r '
[
.tpot_ms.p99?,
.metrics.tpot_ms.p99?,
.tpot.p99?,
.metrics.tpot.p99?,
.p99_tpot_ms?,
.itl_ms.p99?,
.metrics.itl_ms.p99?,
.inter_token_latency_ms.p99?,
.tpot_ms.mean?,
.metrics.tpot_ms.mean?,
.tpot.mean?,
.metrics.tpot.mean?,
.itl_ms.mean?,
.metrics.itl_ms.mean?,
.mean_tpot_ms?,
.mean_itl_ms?
] | map(select(. != null)) | .[0] // empty
' "$json_file"
fi
}
evaluate_sla_from_json() {
local json_file=$1
local ttft
local tpot
local pass
[[ -f "$json_file" ]] || return 2
ttft=$(extract_metric_ms ttft "$json_file")
tpot=$(extract_metric_ms tpot "$json_file")
[[ -n "$ttft" && -n "$tpot" ]] || return 2
pass=$(jq -n \
--argjson ttft "$ttft" \
--argjson tpot "$tpot" \
--argjson sla_ttft "$SLA_TTFT_MS" \
--argjson sla_tpot "$SLA_TPOT_MS" \
'($ttft <= $sla_ttft) and ($tpot <= $sla_tpot)')
[[ "$pass" == "true" ]]
}
write_adaptive_summary_json() {
local summary_file=$1
local test_name=$2
local qps=$3
local static_last_pass=$4
local static_first_fail=$5
local final_last_pass=$6
local final_first_fail=$7
jq -n \
--arg test_name "$test_name" \
--arg qps "$qps" \
--argjson sla_ttft "$SLA_TTFT_MS" \
--argjson sla_tpot "$SLA_TPOT_MS" \
--arg static_last_pass "${static_last_pass:-}" \
--arg static_first_fail "${static_first_fail:-}" \
--arg final_last_pass "${final_last_pass:-}" \
--arg final_first_fail "${final_first_fail:-}" \
'{
test_name: $test_name,
qps: $qps,
sla_ttft_ms: $sla_ttft,
sla_tpot_ms: $sla_tpot,
static_last_pass: (if $static_last_pass == "" then null else ($static_last_pass | tonumber) end),
static_first_fail: (if $static_first_fail == "" then null else ($static_first_fail | tonumber) end),
final_last_pass: (if $final_last_pass == "" then null else ($final_last_pass | tonumber) end),
final_first_fail: (if $final_first_fail == "" then null else ($final_first_fail | tonumber) end)
}' > "$summary_file"
}
run_single_serving_probe() {
local test_name=$1
local qps=$2
local max_concurrency=$3
local tp=$4
local compilation_config_mode=$5
local optimization_level=$6
local client_args_effective=$7
local client_remote_args=$8
local server_command=$9
local new_test_name="${test_name}_qps_${qps}_concurrency_${max_concurrency}"
local result_json
local num_prompts_arg=""
local client_command
result_json=$(result_json_path_for_serving "$test_name" "$qps" "$max_concurrency")
if [[ -f "$result_json" ]]; then
evaluate_sla_from_json "$result_json"
return $?
fi
if [[ -n "${PROMPTS_PER_CONCURRENCY}" ]]; then
num_prompts=$(( max_concurrency * PROMPTS_PER_CONCURRENCY ))
if (( num_prompts < MIN_NUM_PROMPTS )); then num_prompts=$MIN_NUM_PROMPTS; fi
if (( num_prompts > MAX_NUM_PROMPTS )); then num_prompts=$MAX_NUM_PROMPTS; fi
num_prompts_arg="--num-prompts $num_prompts"
fi
client_command="vllm bench serve \
--save-result \
--result-dir $RESULTS_FOLDER \
--result-filename ${new_test_name}.json \
--request-rate $qps \
--max-concurrency $max_concurrency \
$num_prompts_arg \
--metadata tensor_parallel_size=$tp compilation_config.mode=$compilation_config_mode optimization_level=$optimization_level adaptive_search=1 \
$client_args_effective $client_remote_args "
echo "Adaptive probe: $client_command"
if [[ "${DRY_RUN:-0}" != "1" ]]; then
bash -c "$client_command"
fi
jq_output=$(jq -n \
--arg server "$server_command" \
--arg client "$client_command" \
--arg gpu "$gpu_type" \
'{
server_command: $server,
client_command: $client,
gpu_type: $gpu,
adaptive_search: true
}')
echo "$jq_output" > "$RESULTS_FOLDER/${new_test_name}.commands"
evaluate_sla_from_json "$result_json"
}
adaptive_refine_from_static_results() {
local test_name=$1
local qps=$2
local max_concurrency_list_raw=$3
local tp=$4
local compilation_config_mode=$5
local optimization_level=$6
local client_args_effective=$7
local client_remote_args=$8
local server_command=$9
local sorted_points
local point
local rc
local static_last_pass=""
local static_first_fail=""
local largest_static=""
local step_hint=1
local previous_point=""
local low
local high
local mid
local probes=0
local summary_file="$RESULTS_FOLDER/${test_name}_qps_${qps}_sla_summary.json"
[[ "${ENABLE_ADAPTIVE_CONCURRENCY}" == "1" ]] || return 0
[[ "${DRY_RUN:-0}" != "1" ]] || return 0
sorted_points=$(for point in $max_concurrency_list_raw; do printf '%s\n' "$point"; done | tr -d "'" | awk '/^[0-9]+$/' | sort -n | uniq)
[[ -n "$sorted_points" ]] || return 0
while read -r point; do
[[ -z "$point" ]] && continue
largest_static="$point"
evaluate_sla_from_json "$(result_json_path_for_serving "$test_name" "$qps" "$point")"
rc=$?
if (( rc == 0 )); then
static_last_pass="$point"
elif (( rc == 1 )); then
if [[ -n "$static_last_pass" ]]; then
static_first_fail="$point"
break
fi
fi
if [[ -n "$previous_point" ]]; then
step_hint=$(( point - previous_point ))
if (( step_hint < 1 )); then step_hint=1; fi
fi
previous_point="$point"
done <<< "$sorted_points"
if [[ -z "$static_last_pass" ]]; then
write_adaptive_summary_json "$summary_file" "$test_name" "$qps" "" "$static_first_fail" "" "$static_first_fail"
return 0
fi
if [[ -n "$static_first_fail" ]]; then
low=$static_last_pass
high=$static_first_fail
while (( low + 1 < high )) && (( probes < ADAPTIVE_MAX_PROBES )); do
mid=$(( (low + high) / 2 ))
probes=$(( probes + 1 ))
run_single_serving_probe \
"$test_name" "$qps" "$mid" "$tp" \
"$compilation_config_mode" "$optimization_level" \
"$client_args_effective" "$client_remote_args" "$server_command"
rc=$?
if (( rc == 0 )); then
low=$mid
elif (( rc == 1 )); then
high=$mid
else
break
fi
done
write_adaptive_summary_json "$summary_file" "$test_name" "$qps" "$static_last_pass" "$static_first_fail" "$low" "$high"
return 0
fi
low=$largest_static
high=""
while (( probes < ADAPTIVE_MAX_PROBES )); do
point=$(( low + step_hint ))
if (( point > ADAPTIVE_MAX_CONCURRENCY )); then
point=$ADAPTIVE_MAX_CONCURRENCY
fi
(( point > low )) || break
probes=$(( probes + 1 ))
run_single_serving_probe \
"$test_name" "$qps" "$point" "$tp" \
"$compilation_config_mode" "$optimization_level" \
"$client_args_effective" "$client_remote_args" "$server_command"
rc=$?
if (( rc == 0 )); then
low=$point
(( point == ADAPTIVE_MAX_CONCURRENCY )) && break
step_hint=$(( step_hint * 2 ))
if (( step_hint < 1 )); then step_hint=1; fi
elif (( rc == 1 )); then
high=$point
break
else
break
fi
done
if [[ -n "$high" ]]; then
while (( low + 1 < high )) && (( probes < ADAPTIVE_MAX_PROBES )); do
mid=$(( (low + high) / 2 ))
probes=$(( probes + 1 ))
run_single_serving_probe \
"$test_name" "$qps" "$mid" "$tp" \
"$compilation_config_mode" "$optimization_level" \
"$client_args_effective" "$client_remote_args" "$server_command"
rc=$?
if (( rc == 0 )); then
low=$mid
elif (( rc == 1 )); then
high=$mid
else
break
fi
done
fi
write_adaptive_summary_json "$summary_file" "$test_name" "$qps" "$static_last_pass" "" "$low" "$high"
}
run_benchmark_tests() {
# run benchmark tests using `vllm bench <test_type>` command
# $1: test type (latency or throughput)
@@ -347,10 +652,48 @@ run_serving_tests() {
server_envs=$(echo "$params" | jq -r '.server_environment_variables')
client_params=$(echo "$params" | jq -r '.client_parameters')
server_args=$(json2args "$server_params")
# vLLM serve CLI: model must be positional (no --model). Convert server_parameters accordingly.
server_model=$(echo "$server_params" | jq -r '.model // empty')
if [[ -z "$server_model" || "$server_model" == "null" ]]; then
echo "Error: serving test '$test_name' is missing server_parameters.model" >&2
exit 1
fi
server_params_no_model=$(echo "$server_params" | jq -c 'del(.model)')
server_args=$(json2args "$server_params_no_model")
server_envs=$(json2envs "$server_envs")
client_args=$(json2args "$client_params")
# ------------------------------------------------------------
# Option 1: Dynamic num-prompts scaling based on max_concurrency
#
# If PROMPTS_PER_CONCURRENCY is set, override JSON num_prompts with:
# num_prompts = max_concurrency * PROMPTS_PER_CONCURRENCY
#
# If PROMPTS_PER_CONCURRENCY is NOT set, keep JSON num_prompts behavior
# unchanged (i.e., whatever is in serving-tests-*.json).
# ------------------------------------------------------------
PROMPTS_PER_CONCURRENCY="${PROMPTS_PER_CONCURRENCY-}" # no default on purpose
MIN_NUM_PROMPTS="${MIN_NUM_PROMPTS:-1}"
MAX_NUM_PROMPTS="${MAX_NUM_PROMPTS:-1000000}"
if [[ -n "${PROMPTS_PER_CONCURRENCY}" ]]; then
# Remove any fixed --num-prompts from JSON-derived args (avoid duplicates)
# Remove any fixed --num-prompts from JSON-derived args (avoid duplicates)
# Handles: --num-prompts 123 and --num-prompts=123
client_args_no_np="$(
printf ' %s ' "$client_args" \
| sed -E \
-e 's/[[:space:]]--num-prompts=([^[:space:]]+)([[:space:]]|$)/ /g' \
-e 's/[[:space:]]--num-prompts[[:space:]]+([^[:space:]]+)([[:space:]]|$)/ /g'
)"
# normalize whitespace
client_args_no_np="$(echo "$client_args_no_np" | tr -s ' ' | sed -E 's/^ //; s/ $//')"
client_args_no_np="$(echo "$client_args_no_np" | xargs)"
client_args_effective="$client_args_no_np"
else
client_args_effective="$client_args"
fi
# qps_list
qps_list=$(echo "$params" | jq -r '.qps_list')
qps_list=$(echo "$qps_list" | jq -r '.[] | @sh')
@@ -382,14 +725,13 @@ run_serving_tests() {
fi
# check if server model and client model is aligned
server_model=$(echo "$server_params" | jq -r '.model')
client_model=$(echo "$client_params" | jq -r '.model')
if [[ $server_model != "$client_model" ]]; then
echo "Server model and client model must be the same. Skip testcase $test_name."
continue
fi
server_command="$server_envs vllm serve \
server_command="$server_envs vllm serve $server_model \
$server_args"
# run the server
@@ -436,6 +778,14 @@ run_serving_tests() {
for max_concurrency in $max_concurrency_list; do
new_test_name="${test_name}_qps_${qps}_concurrency_${max_concurrency}"
echo " new test name $new_test_name"
# If PROMPTS_PER_CONCURRENCY is set, compute per-concurrency --num-prompts.
num_prompts_arg=""
if [[ -n "${PROMPTS_PER_CONCURRENCY}" ]]; then
num_prompts=$(( max_concurrency * PROMPTS_PER_CONCURRENCY ))
if (( num_prompts < MIN_NUM_PROMPTS )); then num_prompts=$MIN_NUM_PROMPTS; fi
if (( num_prompts > MAX_NUM_PROMPTS )); then num_prompts=$MAX_NUM_PROMPTS; fi
num_prompts_arg="--num-prompts $num_prompts"
fi
# pass the tensor parallel size, the compilation mode, and the optimization
# level to the client so that they can be used on the benchmark dashboard
client_command="vllm bench serve \
@@ -444,8 +794,9 @@ run_serving_tests() {
--result-filename ${new_test_name}.json \
--request-rate $qps \
--max-concurrency $max_concurrency \
$num_prompts_arg \
--metadata tensor_parallel_size=$tp compilation_config.mode=$compilation_config_mode optimization_level=$optimization_level \
$client_args $client_remote_args "
$client_args_effective $client_remote_args "
echo "Running test case $test_name with qps $qps"
echo "Client command: $client_command"
@@ -467,6 +818,11 @@ run_serving_tests() {
echo "$jq_output" >"$RESULTS_FOLDER/${new_test_name}.commands"
done
adaptive_refine_from_static_results \
"$test_name" "$qps" "$max_concurrency_list" "$tp" \
"$compilation_config_mode" "$optimization_level" \
"$client_args_effective" "$client_remote_args" "$server_command"
done
# clean up
@@ -532,6 +888,7 @@ main() {
# postprocess benchmarking results
pip install tabulate pandas
python3 $QUICK_BENCHMARK_ROOT/scripts/convert-results-json-to-markdown.py
python3 $QUICK_BENCHMARK_ROOT/scripts/compare-json-results.py -f $RESULTS_FOLDER/benchmark_results.json
upload_to_buildkite
}
.buildkite/performance-benchmarks/tests/serving-tests-cpu-asr.json
+37
View File
@@ -0,0 +1,37 @@
{
"defaults": {
"qps_list": [
"inf"
],
"max_concurrency_list": [12, 16, 24, 32, 64, 128, 200],
"server_environment_variables": {
"VLLM_RPC_TIMEOUT": 100000,
"VLLM_ENGINE_ITERATION_TIMEOUT_S": 120
},
"server_parameters": {
"dtype": "bfloat16",
"model": "openai/whisper-large-v3-turbo"
},
"client_parameters": {
"model": "openai/whisper-large-v3-turbo",
"backend": "openai-audio",
"endpoint": "/v1/audio/transcriptions",
"dataset_name": "hf",
"dataset_path": "openslr/librispeech_asr",
"hf_subset": "clean",
"hf_split": "test",
"no_stream": "",
"no_oversample": "",
"num_prompts": 200
}
},
"tests": [
{
"test_name": "serving_whisper_large_v3_turbo_librispeech_clean_tp1",
"server_parameters": {
"tensor_parallel_size": 1
},
"client_parameters": {}
}
]
}
.buildkite/performance-benchmarks/tests/serving-tests-cpu-text.json
+72
View File
@@ -149,6 +149,39 @@
"random-output-len": 128
}
},
{
"test_name": "serving_llama8B_tp1_random_2048_2048",
"server_parameters": {
"tensor_parallel_size": 1
},
"client_parameters": {
"dataset_name": "random",
"random-input-len": 2048,
"random-output-len": 2048
}
},
{
"test_name": "serving_llama8B_tp2_random_2048_2048",
"server_parameters": {
"tensor_parallel_size": 2
},
"client_parameters": {
"dataset_name": "random",
"random-input-len": 2048,
"random-output-len": 2048
}
},
{
"test_name": "serving_llama8B_tp4_random_2048_2048",
"server_parameters": {
"tensor_parallel_size": 4
},
"client_parameters": {
"dataset_name": "random",
"random-input-len": 2048,
"random-output-len": 2048
}
},
{
"test_name": "serving_llama8B_int4_tp1_random_128_128",
"server_parameters": {
@@ -188,6 +221,45 @@
"random-output-len": 128
}
},
{
"test_name": "serving_llama8B_int8_tp1_random_128_128",
"server_parameters": {
"model": "RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w8a8",
"tensor_parallel_size": 1
},
"client_parameters": {
"model": "RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w8a8",
"dataset_name": "random",
"random-input-len": 128,
"random-output-len": 128
}
},
{
"test_name": "serving_llama8B_int8_tp2_random_128_128",
"server_parameters": {
"model": "RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w8a8",
"tensor_parallel_size": 2
},
"client_parameters": {
"model": "RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w8a8",
"dataset_name": "random",
"random-input-len": 128,
"random-output-len": 128
}
},
{
"test_name": "serving_llama8B_int8_tp4_random_128_128",
"server_parameters": {
"model": "RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w8a8",
"tensor_parallel_size": 4
},
"client_parameters": {
"model": "RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w8a8",
"dataset_name": "random",
"random-input-len": 128,
"random-output-len": 128
}
},
{
"test_name": "serving_llama3B_tp1_random_128_128",
"server_parameters": {
.buildkite/performance-benchmarks/tests/serving-tests-cpu.json
+14 -25
View File
@@ -72,17 +72,6 @@
"random-output-len": 128
}
},
{
"test_name": "serving_llama8B_tp4_random_128_128",
"server_parameters": {
"tensor_parallel_size": 4
},
"client_parameters": {
"dataset_name": "random",
"random-input-len": 128,
"random-output-len": 128
}
},
{
"test_name": "serving_llama8B_tp1_random_128_2048",
"server_parameters": {
@@ -105,17 +94,6 @@
"random-output-len": 2048
}
},
{
"test_name": "serving_llama8B_tp4_random_128_2048",
"server_parameters": {
"tensor_parallel_size": 4
},
"client_parameters": {
"dataset_name": "random",
"random-input-len": 128,
"random-output-len": 2048
}
},
{
"test_name": "serving_llama8B_tp1_random_2048_128",
"server_parameters": {
@@ -139,14 +117,25 @@
}
},
{
"test_name": "serving_llama8B_tp4_random_2048_128",
"test_name": "serving_llama8B_tp1_random_2048_2048",
"server_parameters": {
"tensor_parallel_size": 4
"tensor_parallel_size": 1
},
"client_parameters": {
"dataset_name": "random",
"random-input-len": 2048,
"random-output-len": 128
"random-output-len": 2048
}
},
{
"test_name": "serving_llama8B_tp2_random_2048_2048",
"server_parameters": {
"tensor_parallel_size": 2
},
"client_parameters": {
"dataset_name": "random",
"random-input-len": 2048,
"random-output-len": 2048
}
}
]
.buildkite/scripts/tool_call/run-bfcl-eval.sh
Executable
+248
View File
@@ -0,0 +1,248 @@
#!/bin/bash
# Run BFCL (Berkeley Function Call Leaderboard) tool-calling correctness
# evaluation against a local vLLM server.
#
# Usage:
# # Run with defaults (gpt-oss-20b, multi_turn)
# bash .buildkite/scripts/tool_call/run-bfcl-eval.sh
#
# # Run with gpt-oss-120b and multiple test categories
# BFCL_MODEL="openai/gpt-oss-120b" BFCL_TP_SIZE=4 \
# BFCL_TEST_CATEGORY="live_simple, multiple, parallel_multiple" \
# bash .buildkite/scripts/tool_call/run-bfcl-eval.sh
#
# # Chain both API types (use BFCL_OUTPUT_DIR to avoid overwriting results)
# BFCL_OUTPUT_DIR=./bfcl-chat-completions BFCL_API_TYPE=chat_completions \
# bash .buildkite/scripts/tool_call/run-bfcl-eval.sh && \
# BFCL_OUTPUT_DIR=./bfcl-responses BFCL_API_TYPE=responses \
# bash .buildkite/scripts/tool_call/run-bfcl-eval.sh
#
# Environment variables (all optional, with defaults):
# BFCL_MODEL - HF model name (default: openai/gpt-oss-20b)
# BFCL_API_TYPE - API type: "chat_completions" or "responses" (default: chat_completions)
# BFCL_OUTPUT_DIR - Directory for BFCL results (default: current working directory)
# BFCL_TEST_CATEGORY - BFCL test categories (default: multi_turn)
# BFCL_TOOL_CALL_PARSER - Tool call parser name (default: openai)
# BFCL_NUM_THREADS - Threads for BFCL generate (default: 8)
# BFCL_TP_SIZE - Tensor parallel size (default: 1)
# BFCL_MAX_MODEL_LEN - Max model length (default: 4096)
# BFCL_PORT - Server port (default: 8000)
# BFCL_REASONING_PARSER - Reasoning parser name (default: disabled)
# BFCL_EXTRA_ARGS - Additional vLLM server args
set -euo pipefail
# ---- Configuration ----
MODEL="${BFCL_MODEL:-openai/gpt-oss-20b}"
API_TYPE="${BFCL_API_TYPE:-chat_completions}"
OUTPUT_DIR="${BFCL_OUTPUT_DIR:-}"
TEST_CATEGORY="${BFCL_TEST_CATEGORY:-multi_turn}"
TOOL_CALL_PARSER="${BFCL_TOOL_CALL_PARSER:-openai}"
NUM_THREADS="${BFCL_NUM_THREADS:-8}"
TP_SIZE="${BFCL_TP_SIZE:-1}"
MAX_MODEL_LEN="${BFCL_MAX_MODEL_LEN:-4096}"
PORT="${BFCL_PORT:-8000}"
REASONING_PARSER="${BFCL_REASONING_PARSER:-}"
EXTRA_ARGS="${BFCL_EXTRA_ARGS:-}"
# Set up output directory
if [ -n "$OUTPUT_DIR" ]; then
mkdir -p "$OUTPUT_DIR"
OUTPUT_DIR="$(cd "$OUTPUT_DIR" && pwd)"
fi
echo "============================================"
echo "BFCL Tool Call Correctness Evaluation"
echo "============================================"
echo "Model: $MODEL"
echo "Tool parser: $TOOL_CALL_PARSER"
echo "API type: $API_TYPE"
echo "Output dir: ${OUTPUT_DIR:-<cwd>}"
echo "Test category: $TEST_CATEGORY"
echo "TP size: $TP_SIZE"
echo "Max model len: $MAX_MODEL_LEN"
echo "Port: $PORT"
echo "Num threads: $NUM_THREADS"
echo "============================================"
# ---- Install bfcl-eval if missing ----
if ! python3 -c "import bfcl_eval" 2>/dev/null; then
echo "Installing bfcl-eval..."
pip install "bfcl-eval>=2025.10.20.1,<2026"
fi
# ---- Cleanup handler ----
SERVER_PID=""
cleanup() {
if [ -n "$SERVER_PID" ]; then
echo "Stopping vLLM server (pid=$SERVER_PID)..."
kill "$SERVER_PID" 2>/dev/null || true
wait "$SERVER_PID" 2>/dev/null || true
fi
# Remove BFCL lock files (created by filelock for thread-safe writes)
rm -rf .file_locks/
if [ -n "${OUTPUT_DIR:-}" ]; then
rm -rf "$OUTPUT_DIR/.file_locks/"
fi
}
trap cleanup EXIT
# ---- Start vLLM server ----
echo "Starting vLLM server..."
SERVE_ARGS=(
"$MODEL"
--port "$PORT"
--enable-auto-tool-choice
--tool-call-parser "$TOOL_CALL_PARSER"
--tensor-parallel-size "$TP_SIZE"
--max-model-len "$MAX_MODEL_LEN"
--enforce-eager
--no-enable-prefix-caching
)
# Append reasoning parser if specified
if [ -n "$REASONING_PARSER" ]; then
SERVE_ARGS+=(--reasoning-parser "$REASONING_PARSER")
fi
# Append any extra args
if [ -n "$EXTRA_ARGS" ]; then
read -ra EXTRA_ARGS_ARRAY <<< "$EXTRA_ARGS"
SERVE_ARGS+=("${EXTRA_ARGS_ARRAY[@]}")
fi
echo "Command: vllm serve ${SERVE_ARGS[*]}"
vllm serve "${SERVE_ARGS[@]}" &
SERVER_PID=$!
# ---- Wait for server to be ready ----
echo "Waiting for vLLM server to start (timeout: 600s)..."
SECONDS_WAITED=0
until curl -sf "http://localhost:${PORT}/health" > /dev/null 2>&1; do
if [ $SECONDS_WAITED -ge 600 ]; then
echo ""
echo "ERROR: vLLM server failed to start within 600s"
exit 1
fi
if (( SECONDS_WAITED % 30 == 0 && SECONDS_WAITED > 0 )); then
echo " Still waiting... (${SECONDS_WAITED}s elapsed)"
fi
sleep 2
SECONDS_WAITED=$((SECONDS_WAITED + 2))
done
echo "vLLM server is ready. (started in ${SECONDS_WAITED}s)"
# ---- Run BFCL evaluation ----
# bfcl-eval has no CLI entry point; generate() and evaluate() are Typer
# functions that must be called from Python. The MODEL_CONFIG_MAPPING must
# be patched in-process so BFCL knows to use the OpenAI-compatible handler
# against our local vLLM server.
bfcl_exit_code=0
python3 - "$MODEL" "$TEST_CATEGORY" "$NUM_THREADS" "$PORT" "$API_TYPE" "$OUTPUT_DIR" << 'PYEOF' || bfcl_exit_code=$?
import os
import sys
model = sys.argv[1]
test_category = sys.argv[2]
num_threads = int(sys.argv[3])
port = sys.argv[4]
api_type = sys.argv[5]
output_dir = sys.argv[6] if len(sys.argv) > 6 and sys.argv[6] else os.getcwd()
os.environ["OPENAI_BASE_URL"] = f"http://localhost:{port}/v1"
os.environ["OPENAI_API_KEY"] = "dummy"
os.environ["BFCL_PROJECT_ROOT"] = output_dir
import bfcl_eval.constants.model_config as bfcl_model_config
from bfcl_eval.constants.model_config import ModelConfig
from bfcl_eval.model_handler.api_inference.openai_completion import (
OpenAICompletionsHandler,
)
from bfcl_eval.model_handler.api_inference.openai_response import (
OpenAIResponsesHandler,
)
if api_type == "responses":
handler = OpenAIResponsesHandler
else:
handler = OpenAICompletionsHandler
bfcl_model_config.MODEL_CONFIG_MAPPING[model] = ModelConfig(
model_name=model,
display_name=f"{model} (FC) (vLLM)",
url=f"https://huggingface.co/{model}",
org="",
license="apache-2.0",
model_handler=handler,
input_price=None,
output_price=None,
is_fc_model=True,
underscore_to_dot=True,
)
from bfcl_eval.__main__ import evaluate, generate
import inspect
import typer
def _get_default_kwargs(function):
kwargs = {}
for k, v in inspect.signature(function).parameters.items():
if v.default is not inspect.Parameter.empty:
default = v.default
if isinstance(default, typer.models.OptionInfo):
default = default.default
kwargs[k] = default
return kwargs
# ---- generate ----
print(f"=== BFCL generate: model={model} test_category={test_category} ===")
gen_kwargs = _get_default_kwargs(generate)
gen_kwargs["model"] = [model]
gen_kwargs["test_category"] = [c.strip() for c in test_category.split(",")]
gen_kwargs["skip_server_setup"] = True
gen_kwargs["num_threads"] = num_threads
generate(**gen_kwargs)
# ---- evaluate ----
print(f"=== BFCL evaluate: model={model} test_category={test_category} ===")
eval_kwargs = _get_default_kwargs(evaluate)
eval_kwargs["model"] = [model]
eval_kwargs["test_category"] = [c.strip() for c in test_category.split(",")]
evaluate(**eval_kwargs)
print("=== BFCL evaluation completed successfully ===")
PYEOF
# ---- Upload results to buildkite ----
if command -v buildkite-agent &>/dev/null; then
if [ $bfcl_exit_code -eq 0 ]; then
STYLE="success"
STATUS="PASSED"
else
STYLE="error"
STATUS="FAILED"
fi
buildkite-agent annotate --style "$STYLE" --context "bfcl-results" <<EOF
### BFCL Tool Call Correctness - ${STATUS}
- **Model:** \`${MODEL}\`
- **Parser:** \`${TOOL_CALL_PARSER}\`
- **API type:** \`${API_TYPE}\`
- **Test category:** \`${TEST_CATEGORY}\`
EOF
# BFCL writes results to $BFCL_PROJECT_ROOT/result/ and scores to
# $BFCL_PROJECT_ROOT/score/
RESULTS_ROOT="${OUTPUT_DIR:-.}"
if [ -d "$RESULTS_ROOT/result" ]; then
buildkite-agent artifact upload "$RESULTS_ROOT/result/**/*"
fi
if [ -d "$RESULTS_ROOT/score" ]; then
buildkite-agent artifact upload "$RESULTS_ROOT/score/**/*"
fi
fi
exit $bfcl_exit_code
.buildkite/test-amd.yaml
+117 -4
View File
@@ -42,6 +42,7 @@ steps:
mirror_hardwares: [amdexperimental, amdproduction, amdtentative]
agent_pool: mi325_1
grade: Blocking
optional: true
soft_fail: true
source_file_dependencies:
- requirements/nightly_torch_test.txt
@@ -67,6 +68,7 @@ steps:
timeout_in_minutes: 30
mirror_hardwares: [amdexperimental, amdproduction, amdtentative]
agent_pool: mi325_1
optional: true
# grade: Blocking
source_file_dependencies:
- vllm/
@@ -97,6 +99,7 @@ steps:
timeout_in_minutes: 20
mirror_hardwares: [amdexperimental]
agent_pool: mi325_1
optional: true
# grade: Blocking
source_file_dependencies:
- tests/standalone_tests/python_only_compile.sh
@@ -140,6 +143,7 @@ steps:
timeout_in_minutes: 40
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi325_1
optional: true
# grade: Blocking
working_dir: "/vllm-workspace/tests"
fast_check: true
@@ -503,6 +507,7 @@ steps:
mirror_hardwares: [amdexperimental, amdproduction, amdtentative]
agent_pool: mi325_1
grade: Blocking
optional: true
source_file_dependencies:
- vllm/
- tests/v1
@@ -520,6 +525,7 @@ steps:
timeout_in_minutes: 45
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi325_1
optional: true
# grade: Blocking
working_dir: "/vllm-workspace/examples"
source_file_dependencies:
@@ -823,6 +829,7 @@ steps:
timeout_in_minutes: 90
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi325_1
optional: true
# grade: Blocking
source_file_dependencies:
- csrc/
@@ -936,6 +943,7 @@ steps:
timeout_in_minutes: 25
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi325_1
optional: true
# grade: Blocking
torch_nightly: true
source_file_dependencies:
@@ -1046,6 +1054,7 @@ steps:
timeout_in_minutes: 60
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi325_1
optional: true
source_file_dependencies:
- vllm/
- tests/models/multimodal
@@ -1059,6 +1068,7 @@ steps:
timeout_in_minutes: 60
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi325_1
optional: true
# grade: Blocking
source_file_dependencies:
- vllm/
@@ -1072,6 +1082,7 @@ steps:
timeout_in_minutes: 100
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi325_1
optional: true
# grade: Blocking
torch_nightly: true
source_file_dependencies:
@@ -1090,6 +1101,7 @@ steps:
timeout_in_minutes: 10
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi325_1
optional: true
# grade: Blocking
working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
source_file_dependencies:
@@ -1355,6 +1367,7 @@ steps:
timeout_in_minutes: 60
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi325_2
optional: true
# grade: Blocking
working_dir: "/vllm-workspace/tests"
num_gpus: 2
@@ -1393,6 +1406,7 @@ steps:
timeout_in_minutes: 60
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi325_4
optional: true
# grade: Blocking
working_dir: "/vllm-workspace/tests"
num_gpus: 4
@@ -1410,6 +1424,7 @@ steps:
timeout_in_minutes: 30
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi325_4
optional: true
# grade: Blocking
num_gpus: 4
source_file_dependencies:
@@ -1461,6 +1476,7 @@ steps:
- label: NixlConnector PD accuracy tests (Distributed) # 30min
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi325_4
optional: true
# grade: Blocking
timeout_in_minutes: 30
working_dir: "/vllm-workspace/tests"
@@ -1475,6 +1491,7 @@ steps:
- label: DP EP NixlConnector PD accuracy tests (Distributed) # 15min
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi325_4
optional: true
# grade: Blocking
timeout_in_minutes: 15
working_dir: "/vllm-workspace/tests"
@@ -1779,6 +1796,7 @@ steps:
# in /vllm/tools/pre_commit/generate_nightly_torch_test.py
mirror_hardwares: [amdexperimental, amdproduction, amdtentative]
agent_pool: mi355_1
optional: true
soft_fail: true
source_file_dependencies:
- requirements/nightly_torch_test.txt
@@ -1789,6 +1807,7 @@ steps:
timeout_in_minutes: 15
mirror_hardwares: [amdexperimental, amdproduction, amdtentative]
agent_pool: mi355_1
optional: true
source_file_dependencies:
- vllm/
- tests/multimodal
@@ -1801,6 +1820,7 @@ steps:
timeout_in_minutes: 30
mirror_hardwares: [amdexperimental, amdproduction, amdtentative]
agent_pool: mi355_1
optional: true
source_file_dependencies:
- vllm/
- tests/test_inputs.py
@@ -1830,6 +1850,7 @@ steps:
timeout_in_minutes: 20
mirror_hardwares: [amdexperimental]
agent_pool: mi355_1
optional: true
source_file_dependencies:
- tests/standalone_tests/python_only_compile.sh
- setup.py
@@ -1840,6 +1861,7 @@ steps:
timeout_in_minutes: 30
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi355_1
optional: true
fast_check: true
torch_nightly: true
source_file_dependencies:
@@ -1870,6 +1892,7 @@ steps:
timeout_in_minutes: 40
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi355_1
optional: true
working_dir: "/vllm-workspace/tests"
fast_check: true
torch_nightly: true
@@ -1887,6 +1910,7 @@ steps:
timeout_in_minutes: 130
mirror_hardwares: [amdexperimental]
agent_pool: mi355_1
optional: true
working_dir: "/vllm-workspace/tests"
fast_check: true
torch_nightly: true
@@ -1903,6 +1927,7 @@ steps:
timeout_in_minutes: 50
mirror_hardwares: [amdexperimental]
agent_pool: mi355_1
optional: true
working_dir: "/vllm-workspace/tests"
fast_check: true
torch_nightly: true
@@ -1921,6 +1946,7 @@ steps:
timeout_in_minutes: 50
mirror_hardwares: [amdexperimental]
agent_pool: mi355_1
optional: true
working_dir: "/vllm-workspace/tests"
fast_check: true
torch_nightly: true
@@ -1935,6 +1961,7 @@ steps:
timeout_in_minutes: 50
mirror_hardwares: [amdexperimental]
agent_pool: mi355_1
optional: true
working_dir: "/vllm-workspace/tests"
fast_check: true
torch_nightly: true
@@ -2013,6 +2040,7 @@ steps:
timeout_in_minutes: 10
mirror_hardwares: [amdexperimental]
agent_pool: mi355_8
optional: true
gpu: h100
num_gpus: 8
working_dir: "/vllm-workspace/tests"
@@ -2033,6 +2061,7 @@ steps:
- label: EPLB Algorithm Test # 5min
mirror_hardwares: [amdexperimental, amdproduction, amdtentative]
agent_pool: mi355_1
optional: true
timeout_in_minutes: 15
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
@@ -2044,6 +2073,7 @@ steps:
- label: EPLB Execution Test # 10min
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi355_4
optional: true
timeout_in_minutes: 20
working_dir: "/vllm-workspace/tests"
num_gpus: 4
@@ -2058,6 +2088,7 @@ steps:
timeout_in_minutes: 20
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi355_2
optional: true
num_gpus: 2
source_file_dependencies:
- vllm/
@@ -2099,12 +2130,13 @@ steps:
commands:
- pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py
- label: V1 Test e2e + engine # 65min
timeout_in_minutes: 90
mirror_hardwares: [amdexperimental]
# The test uses 4 GPUs, but we schedule it on 8-GPU machines for stability.
# See discussion here: https://github.com/vllm-project/vllm/pull/31040
agent_pool: mi355_8
agent_pool: mi355_1
optional: true
# grade: Blocking
source_file_dependencies:
- vllm/
- tests/v1
@@ -2114,10 +2146,39 @@ steps:
- pytest -v -s v1/e2e
- pytest -v -s v1/engine
- label: V1 Test e2e (2 GPUs) # 65min
timeout_in_minutes: 90
mirror_hardwares: [amdexperimental]
agent_pool: mi355_2
optional: true
# grade: Blocking
source_file_dependencies:
- vllm/
- tests/v1
commands:
# Only run tests that need exactly 2 GPUs
- pytest -v -s v1/e2e/test_spec_decode.py -k "tensor_parallelism"
- label: V1 Test e2e (4 GPUs) # 65min
timeout_in_minutes: 90
mirror_hardwares: [amdexperimental]
# The test uses 4 GPUs, but we schedule it on 8-GPU machines for stability.
# See discussion here: https://github.com/vllm-project/vllm/pull/31040
agent_pool: mi355_4
optional: true
# grade: Blocking
source_file_dependencies:
- vllm/
- tests/v1
commands:
# Only run tests that need 4 GPUs
- pytest -v -s v1/e2e/test_spec_decode.py -k "eagle_correctness_heavy"
- label: V1 Test entrypoints # 35min
timeout_in_minutes: 50
mirror_hardwares: [amdexperimental, amdproduction, amdtentative]
agent_pool: mi355_1
optional: true
source_file_dependencies:
- vllm/
- tests/v1
@@ -2128,6 +2189,7 @@ steps:
timeout_in_minutes: 60
mirror_hardwares: [amdexperimental]
agent_pool: mi355_1
optional: true
source_file_dependencies:
- vllm/
- tests/v1
@@ -2150,7 +2212,19 @@ steps:
- pip install -U git+https://github.com/robertgshaw2-redhat/lm-evaluation-harness.git@streaming-api
- pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
# TODO: Add the "V1 Test attention (MI300)" test group
- label: V1 Test attention (H100) # 10min
mirror_hardwares: [amdexperimental]
agent_pool: mi355_1
optional: true
timeout_in_minutes: 30
gpu: h100
source_file_dependencies:
- vllm/config/attention.py
- vllm/model_executor/layers/attention
- vllm/v1/attention
- tests/v1/attention
commands:
- pytest -v -s v1/attention
- label: Batch Invariance Tests (H100) # 10min
mirror_hardwares: [amdexperimental]
@@ -2200,6 +2274,7 @@ steps:
timeout_in_minutes: 45
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi355_1
optional: true
working_dir: "/vllm-workspace/examples"
source_file_dependencies:
- vllm/entrypoints
@@ -2234,6 +2309,7 @@ steps:
timeout_in_minutes: 15
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi355_1
optional: true
source_file_dependencies:
- vllm/
- tests/cuda
@@ -2245,6 +2321,7 @@ steps:
timeout_in_minutes: 75
mirror_hardwares: [amdexperimental]
agent_pool: mi355_1
optional: true
source_file_dependencies:
- vllm/model_executor/layers
- vllm/sampling_metadata.py
@@ -2277,6 +2354,7 @@ steps:
timeout_in_minutes: 30
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi355_1
optional: true
torch_nightly: true
source_file_dependencies:
- vllm/
@@ -2293,6 +2371,7 @@ steps:
timeout_in_minutes: 30
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi355_1
optional: true
torch_nightly: true
source_file_dependencies:
- vllm/
@@ -2308,6 +2387,7 @@ steps:
timeout_in_minutes: 40
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi355_1
optional: true
# grade: Blocking
torch_nightly: true
source_file_dependencies:
@@ -2325,6 +2405,7 @@ steps:
timeout_in_minutes: 20
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi355_1
optional: true
source_file_dependencies:
- tests/v1/cudagraph
- vllm/v1/cudagraph_dispatcher.py
@@ -2338,6 +2419,7 @@ steps:
timeout_in_minutes: 75
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi355_1
optional: true
source_file_dependencies:
- csrc/
- tests/kernels/core
@@ -2349,6 +2431,7 @@ steps:
timeout_in_minutes: 35
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi355_1
optional: true
source_file_dependencies:
- csrc/attention/
- vllm/v1/attention
@@ -2363,6 +2446,7 @@ steps:
timeout_in_minutes: 90
mirror_hardwares: [amdexperimental]
agent_pool: mi355_1
optional: true
source_file_dependencies:
- csrc/quantization/
- vllm/model_executor/layers/quantization
@@ -2375,6 +2459,7 @@ steps:
timeout_in_minutes: 60
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi355_1
optional: true
source_file_dependencies:
- csrc/quantization/cutlass_w8a8/moe/
- csrc/moe/
@@ -2391,6 +2476,7 @@ steps:
timeout_in_minutes: 45
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi355_1
optional: true
source_file_dependencies:
- csrc/mamba/
- tests/kernels/mamba
@@ -2422,6 +2508,7 @@ steps:
timeout_in_minutes: 30
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi355_1
optional: true
source_file_dependencies:
- vllm/utils/import_utils.py
- tests/kernels/helion/
@@ -2434,6 +2521,7 @@ steps:
torch_nightly: true
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi355_1
optional: true
source_file_dependencies:
- vllm/engine/arg_utils.py
- vllm/config/model.py
@@ -2450,6 +2538,7 @@ steps:
timeout_in_minutes: 20
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi355_1
optional: true
working_dir: "/vllm-workspace/.buildkite"
source_file_dependencies:
- benchmarks/
@@ -2460,6 +2549,7 @@ steps:
timeout_in_minutes: 20
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi355_1
optional: true
source_file_dependencies:
- vllm/
- tests/benchmarks/
@@ -2470,6 +2560,7 @@ steps:
timeout_in_minutes: 90
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi355_1
optional: true
source_file_dependencies:
- csrc/
- vllm/model_executor/layers/quantization
@@ -2490,6 +2581,7 @@ steps:
timeout_in_minutes: 75
mirror_hardwares: [amdexperimental]
agent_pool: mi355_1
optional: true
source_file_dependencies:
- csrc/
- vllm/model_executor/layers/quantization
@@ -2501,6 +2593,7 @@ steps:
timeout_in_minutes: 15
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi355_1
optional: true
source_file_dependencies:
- csrc/
- vllm/entrypoints/openai/
@@ -2517,6 +2610,7 @@ steps:
timeout_in_minutes: 45
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi355_1
optional: true
torch_nightly: true
source_file_dependencies:
- vllm/
@@ -2529,6 +2623,7 @@ steps:
timeout_in_minutes: 45
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi355_1
optional: true
torch_nightly: true
source_file_dependencies:
- vllm/model_executor/models/
@@ -2548,6 +2643,7 @@ steps:
timeout_in_minutes: 45
mirror_hardwares: [amdexperimental]
agent_pool: mi355_1
optional: true
torch_nightly: true
source_file_dependencies:
- vllm/
@@ -2560,6 +2656,7 @@ steps:
- label: Basic Models Test (Other CPU) # 5min
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi355_1
optional: true
timeout_in_minutes: 10
torch_nightly: true
source_file_dependencies:
@@ -2574,6 +2671,7 @@ steps:
timeout_in_minutes: 25
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi355_1
optional: true
torch_nightly: true
source_file_dependencies:
- vllm/
@@ -2587,6 +2685,7 @@ steps:
timeout_in_minutes: 45
mirror_hardwares: [amdexperimental]
agent_pool: mi355_1
optional: true
torch_nightly: true
source_file_dependencies:
- vllm/model_executor/models/
@@ -2607,6 +2706,7 @@ steps:
timeout_in_minutes: 75
mirror_hardwares: [amdexperimental]
agent_pool: mi355_1
optional: true
torch_nightly: true
source_file_dependencies:
- vllm/
@@ -2676,6 +2776,7 @@ steps:
timeout_in_minutes: 60
mirror_hardwares: [amdexperimental]
agent_pool: mi355_1
optional: true
source_file_dependencies:
- vllm/
- tests/models/multimodal
@@ -2688,6 +2789,7 @@ steps:
timeout_in_minutes: 60
mirror_hardwares: [amdexperimental]
agent_pool: mi355_1
optional: true
source_file_dependencies:
- vllm/
- tests/models/multimodal
@@ -2699,6 +2801,7 @@ steps:
timeout_in_minutes: 100
mirror_hardwares: [amdexperimental]
agent_pool: mi355_1
optional: true
torch_nightly: true
source_file_dependencies:
- vllm/
@@ -2716,6 +2819,7 @@ steps:
timeout_in_minutes: 10
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi355_1
optional: true
working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
source_file_dependencies:
- vllm/multimodal/
@@ -2772,6 +2876,7 @@ steps:
timeout_in_minutes: 60
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi355_1
optional: true
source_file_dependencies:
- vllm/model_executor/layers/quantization
- tests/models/quantization
@@ -2923,6 +3028,7 @@ steps:
timeout_in_minutes: 20
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi355_2
optional: true
working_dir: "/vllm-workspace/tests"
num_gpus: 2
source_file_dependencies:
@@ -3005,6 +3111,7 @@ steps:
timeout_in_minutes: 50
mirror_hardwares: [amdexperimental]
agent_pool: mi355_2
optional: true
working_dir: "/vllm-workspace/tests"
num_gpus: 2
source_file_dependencies:
@@ -3026,6 +3133,7 @@ steps:
timeout_in_minutes: 60
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi355_2
optional: true
working_dir: "/vllm-workspace/tests"
num_gpus: 2
source_file_dependencies:
@@ -3063,6 +3171,7 @@ steps:
timeout_in_minutes: 60
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi355_4
optional: true
working_dir: "/vllm-workspace/tests"
num_gpus: 4
source_file_dependencies:
@@ -3079,6 +3188,7 @@ steps:
timeout_in_minutes: 30
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi355_4
optional: true
num_gpus: 4
source_file_dependencies:
- vllm/lora
@@ -3127,6 +3237,7 @@ steps:
- label: NixlConnector PD accuracy tests (Distributed) # 30min
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi355_4
optional: true
timeout_in_minutes: 30
working_dir: "/vllm-workspace/tests"
num_gpus: 4
@@ -3140,6 +3251,7 @@ steps:
- label: DP EP NixlConnector PD accuracy tests (Distributed) # 15min
mirror_hardwares: [amdexperimental, amdproduction]
agent_pool: mi355_4
optional: true
timeout_in_minutes: 15
working_dir: "/vllm-workspace/tests"
num_gpus: 4
@@ -3278,6 +3390,7 @@ steps:
- label: ROCm LM Eval Large Models (8 Card)
mirror_hardwares: [amdproduction]
agent_pool: mi355_8
optional: true
num_gpus: 8
working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
commands:
.buildkite/test_areas/basic_correctness.yaml
-5
View File
@@ -14,8 +14,3 @@ steps:
- pytest -v -s basic_correctness/test_cumem.py
- pytest -v -s basic_correctness/test_basic_correctness.py
- pytest -v -s basic_correctness/test_cpu_offload.py
mirror:
amd:
device: mi325_1
depends_on:
- image-build-amd
.buildkite/test_areas/compile.yaml
+9 -9
View File
@@ -101,8 +101,8 @@ steps:
- nvidia-smi
# Run all models and attn backends but only Inductor partition and native custom ops
- pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "inductor_partition and not +rms_norm and not +quant_fp8"
# Qwen requires +quant_fp8 as -quant_fp8 rms+quant fusion is not supported
- pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "inductor_partition and not +rms_norm and +quant_fp8 and qwen3"
# Qwen/Deepseek requires +quant_fp8 as -quant_fp8 rms+quant fusion is not supported
- pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "inductor_partition and not +rms_norm and +quant_fp8 and (qwen3 or deepseek)"
- label: Fusion E2E Config Sweep (H100)
timeout_in_minutes: 30
@@ -132,9 +132,9 @@ steps:
commands:
- nvidia-smi
# Run all models but only FLASHINFER, Inductor partition and native custom ops
# Qwen requires +quant_fp8 as -quant_fp8 rms+quant fusion is not supported
# Qwen/Deepseek requires +quant_fp8 as -quant_fp8 rms+quant fusion is not supported
# Run just llama3 (fp8 & fp4) for all config combinations (only inductor partition)
- pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "inductor_partition and (FLASHINFER and not +rms_norm and (not +quant_fp8 or +quant_fp8 and qwen3) or llama-3)"
- pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "inductor_partition and (FLASHINFER and not +rms_norm and (not +quant_fp8 or +quant_fp8 and (qwen3 or deepseek)) or llama-3)"
- label: Fusion E2E TP2 Quick (H100)
timeout_in_minutes: 20
@@ -150,8 +150,8 @@ steps:
commands:
- nvidia-smi
# Run all models and attn backends but only Inductor partition and native custom ops
- pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k "inductor_partition and not +rms_norm and not +quant_fp8"
- pytest -v -s tests/compile/fusions_e2e/test_tp2_async_tp.py -k "inductor_partition and not +rms_norm and not +quant_fp8"
- pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k "inductor_partition and not +rms_norm and (not +quant_fp8 or +quant_fp8 and (qwen3 or deepseek))"
- pytest -v -s tests/compile/fusions_e2e/test_tp2_async_tp.py -k "inductor_partition and not +rms_norm and (not +quant_fp8 or +quant_fp8 and (qwen3 or deepseek))"
- label: Fusion E2E TP2 AR-RMS Config Sweep (H100)
timeout_in_minutes: 40
@@ -205,7 +205,7 @@ steps:
commands:
- nvidia-smi
# Run all models but only FLASHINFER, Inductor partition and native custom ops
# include qwen with +quant_fp8 as -quant_fp8 rms+quant fusion is not supported
# include qwen/deepseek with +quant_fp8 as -quant_fp8 rms+quant fusion is not supported
# for ar-rms-quant-fp4, also sweep llama3
- pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k "(FLASHINFER and inductor_partition and not +rms_norm and (not +quant_fp8 or +quant_fp8 and qwen3)) or Llama-3.1-8B-Instruct-FP4"
- pytest -v -s tests/compile/fusions_e2e/test_tp2_async_tp.py -k "FLASHINFER and inductor_partition and not +rms_norm and (not +quant_fp8 or +quant_fp8 and qwen3)"
- pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k "(FLASHINFER and inductor_partition and not +rms_norm and (not +quant_fp8 or +quant_fp8 and (qwen3 or deepseek))) or Llama-3.1-8B-Instruct-FP4"
- pytest -v -s tests/compile/fusions_e2e/test_tp2_async_tp.py -k "FLASHINFER and inductor_partition and not +rms_norm and (not +quant_fp8 or +quant_fp8 and (qwen3 or deepseek))"
.buildkite/test_areas/distributed.yaml
+1 -1
View File
@@ -149,7 +149,7 @@ steps:
num_devices: 2
commands:
- pytest -v -s tests/distributed/test_context_parallel.py
# - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/offline_inference/new_weight_syncing/rlhf_async_new_apis.py --- failing, need to re-enable
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/offline_inference/new_weight_syncing/rlhf_async_new_apis.py
- VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
- pytest -v -s tests/v1/distributed/test_dbo.py
.buildkite/test_areas/entrypoints.yaml
-15
View File
@@ -24,11 +24,6 @@ steps:
- pytest -v -s entrypoints/llm --ignore=entrypoints/llm/test_generate.py --ignore=entrypoints/llm/test_collective_rpc.py
- pytest -v -s entrypoints/llm/test_generate.py # it needs a clean process
- pytest -v -s entrypoints/offline_mode # Needs to avoid interference with other tests
mirror:
amd:
device: mi325_1
depends_on:
- image-build-amd
- label: Entrypoints Integration (API Server 1)
timeout_in_minutes: 130
@@ -60,11 +55,6 @@ steps:
- pytest -v -s entrypoints/instrumentator
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc
- pytest -v -s tool_use
mirror:
amd:
device: mi325_1
depends_on:
- image-build-amd
- label: Entrypoints Integration (Pooling)
timeout_in_minutes: 50
@@ -75,11 +65,6 @@ steps:
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/pooling
mirror:
amd:
device: mi325_1
depends_on:
- image-build-amd
- label: Entrypoints Integration (Responses API)
timeout_in_minutes: 50
.buildkite/test_areas/misc.yaml
-5
View File
@@ -88,11 +88,6 @@ steps:
- python3 offline_inference/spec_decode.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
# https://github.com/vllm-project/vllm/pull/26682 uses slightly more memory in PyTorch 2.9+ causing this test to OOM in 1xL4 GPU
- python3 offline_inference/spec_decode.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536
mirror:
amd:
device: mi325_1
depends_on:
- image-build-amd
- label: Metrics, Tracing (2 GPUs)
timeout_in_minutes: 20
.buildkite/test_areas/model_runner_v2.yaml
+110
View File
@@ -0,0 +1,110 @@
group: Model Runner V2
depends_on:
- image-build
steps:
- label: Model Runner V2 Core Tests
timeout_in_minutes: 45
source_file_dependencies:
- vllm/v1/worker/gpu/
- vllm/v1/worker/gpu_worker.py
- vllm/v1/core/sched/
- vllm/v1/attention/
- tests/v1/engine/test_llm_engine.py
- tests/v1/e2e/
- tests/v1/entrypoints/llm/test_struct_output_generate.py
commands:
- set -x
- export VLLM_USE_V2_MODEL_RUNNER=1
- pytest -v -s v1/engine/test_llm_engine.py -k "not test_engine_metrics"
# This requires eager until we sort out CG correctness issues.
# TODO: remove ENFORCE_EAGER here after https://github.com/vllm-project/vllm/pull/32936 is merged.
- ENFORCE_EAGER=1 pytest -v -s v1/e2e/test_async_scheduling.py -k "not ngram"
- pytest -v -s v1/e2e/test_context_length.py
- pytest -v -s v1/e2e/test_min_tokens.py
# Temporary hack filter to exclude ngram spec decoding based tests.
- pytest -v -s v1/entrypoints/llm/test_struct_output_generate.py -k "xgrammar and not speculative_config6 and not speculative_config7 and not speculative_config8 and not speculative_config0"
- label: Model Runner V2 Examples
timeout_in_minutes: 45
working_dir: "/vllm-workspace/examples"
source_file_dependencies:
- vllm/v1/worker/gpu/
- vllm/v1/core/sched/
- vllm/v1/worker/gpu_worker.py
- examples/offline_inference/
- examples/basic/offline_inference/
- examples/pooling/embed/vision_embedding_offline.py
- examples/others/tensorize_vllm_model.py
commands:
- set -x
- export VLLM_USE_V2_MODEL_RUNNER=1
- pip install tensorizer # for tensorizer test
- python3 basic/offline_inference/chat.py # for basic
- python3 basic/offline_inference/generate.py --model facebook/opt-125m
#- python3 basic/offline_inference/generate.py --model meta-llama/Llama-2-13b-chat-hf --cpu-offload-gb 10 # TODO
#- python3 basic/offline_inference/embed.py # TODO
# for multi-modal models
- python3 offline_inference/audio_language.py --seed 0
- python3 offline_inference/vision_language.py --seed 0
- python3 offline_inference/vision_language_multi_image.py --seed 0
- python3 offline_inference/encoder_decoder_multimodal.py --model-type whisper --seed 0
# for pooling models
- python3 pooling/embed/vision_embedding_offline.py --seed 0
# for features demo
- python3 offline_inference/prefix_caching.py
- python3 offline_inference/llm_engine_example.py
- python3 others/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 others/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
- python3 offline_inference/spec_decode.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
# https://github.com/vllm-project/vllm/pull/26682 uses slightly more memory in PyTorch 2.9+ causing this test to OOM in 1xL4 GPU
- python3 offline_inference/spec_decode.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536
- label: Model Runner V2 Distributed (2 GPUs)
timeout_in_minutes: 45
working_dir: "/vllm-workspace/tests"
num_devices: 2
source_file_dependencies:
- vllm/v1/worker/gpu/
- vllm/v1/worker/gpu_worker.py
- tests/basic_correctness/test_basic_correctness.py
- tests/v1/distributed/test_async_llm_dp.py
- tests/v1/distributed/test_eagle_dp.py
commands:
- set -x
- export VLLM_USE_V2_MODEL_RUNNER=1
# The "and not True" here is a hacky way to exclude the prompt_embeds cases which aren't yet supported.
- TARGET_TEST_SUITE=L4 pytest -v -s basic_correctness/test_basic_correctness.py -m 'distributed(num_gpus=2)' -k "not ray and not True"
# https://github.com/NVIDIA/nccl/issues/1838
- export NCCL_CUMEM_HOST_ENABLE=0
- TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py -k "not ray"
- TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
# These require fix https://github.com/vllm-project/vllm/pull/36280
- label: Model Runner V2 Pipeline Parallelism (4 GPUs)
timeout_in_minutes: 60
working_dir: "/vllm-workspace/tests"
num_devices: 4
source_file_dependencies:
- vllm/v1/worker/gpu/
- vllm/v1/worker/gpu_worker.py
- tests/distributed/test_pipeline_parallel.py
#- tests/distributed/test_pp_cudagraph.py
commands:
- set -x
- export VLLM_USE_V2_MODEL_RUNNER=1
- pytest -v -s distributed/test_pipeline_parallel.py -k "not ray and not Jamba"
# TODO: Uncomment once https://github.com/vllm-project/vllm/pull/35162 is merged.
#- pytest -v -s distributed/test_pp_cudagraph.py -k "not ray"
- label: Model Runner V2 Spec Decode
timeout_in_minutes: 30
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/v1/worker/gpu/
- vllm/v1/worker/gpu_worker.py
- tests/v1/spec_decode/test_max_len.py
- tests/v1/e2e/test_spec_decode.py
commands:
- set -x
- export VLLM_USE_V2_MODEL_RUNNER=1
- pytest -v -s v1/spec_decode/test_max_len.py -k "eagle or mtp"
- pytest -v -s v1/e2e/test_spec_decode.py -k "eagle or mtp"
.buildkite/test_areas/plugins.yaml
-5
View File
@@ -39,8 +39,3 @@ steps:
- pytest -v -s entrypoints/openai/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
mirror:
amd:
device: mi325_2
depends_on:
- image-build-amd
.pre-commit-config.yaml
+1 -1
View File
@@ -55,7 +55,7 @@ repos:
language: python
types_or: [python, pyi]
require_serial: true
additional_dependencies: ["mypy[faster-cache]==1.15.0", regex, types-cachetools, types-setuptools, types-PyYAML, types-requests, types-torch, pydantic]
additional_dependencies: ["mypy[faster-cache]==1.19.1", regex, types-cachetools, types-setuptools, types-PyYAML, types-requests, types-torch, pydantic]
- id: mypy-3.10 # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup is less awkward
name: Run mypy for Python 3.10
entry: python tools/pre_commit/mypy.py 1 "3.10"
.readthedocs.yaml
+1 -1
View File
@@ -9,7 +9,7 @@ build:
python: "3.12"
jobs:
post_checkout:
- bash docs/maybe_skip_pr_build.sh
# - bash docs/maybe_skip_pr_build.sh
- git fetch origin main --unshallow --no-tags --filter=blob:none || true
pre_create_environment:
- pip install uv
CMakeLists.txt
+1 -1
View File
@@ -37,7 +37,7 @@ install(CODE "set(CMAKE_INSTALL_LOCAL_ONLY TRUE)" ALL_COMPONENTS)
set(PYTHON_SUPPORTED_VERSIONS "3.10" "3.11" "3.12" "3.13")
# Supported AMD GPU architectures.
set(HIP_SUPPORTED_ARCHS "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1030;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151")
set(HIP_SUPPORTED_ARCHS "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1030;gfx1100;gfx1101;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201")
# ROCm installation prefix. Default to /opt/rocm but allow override via
# -DROCM_PATH=/your/rocm/path when invoking cmake.
benchmarks/kernels/benchmark_moe.py
+5 -1
View File
@@ -626,7 +626,11 @@ class BenchmarkWorker:
if visible_device != f"{self.device_id}":
need_device_guard = True
with torch.cuda.device(self.device_id) if need_device_guard else nullcontext():
with (
torch.accelerator.device_index(self.device_id)
if need_device_guard
else nullcontext()
):
for idx, config in enumerate(tqdm(search_space)):
try:
kernel_time = benchmark_config(
cmake/cpu_extension.cmake
+12 -3
View File
@@ -79,7 +79,8 @@ else()
find_isa(${CPUINFO} "asimd" ASIMD_FOUND) # Check for ARM NEON support
find_isa(${CPUINFO} "bf16" ARM_BF16_FOUND) # Check for ARM BF16 support
find_isa(${CPUINFO} "S390" S390_FOUND)
find_isa(${CPUINFO} "v" RVV_FOUND) # Check for RISC-V RVV support
find_isa(${CPUINFO} "zvfhmin" RVV_FP16_FOUND) # Check for RISC-V Vector FP16 support
find_isa(${CPUINFO} "zvfbfmin" RVV_BF16_FOUND) # Check for RISC-V Vector BF16 support
# Support cross-compilation by allowing override via environment variables
if (ENABLE_ARM_BF16)
@@ -142,11 +143,19 @@ elseif (S390_FOUND)
"-march=native"
"-mtune=native")
elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64")
if(RVV_FOUND)
message(FAIL_ERROR "Can't support rvv now.")
message(STATUS "RISC-V detected")
if(RVV_BF16_FOUND)
message(STATUS "BF16 extension detected")
set(MARCH_FLAGS -march=rv64gcv_zvfh_zfbfmin_zvfbfmin_zvl128b -mrvv-vector-bits=zvl -mabi=lp64d)
add_compile_definitions(RISCV_BF16_SUPPORT)
elseif (RVV_FP16_FOUND)
message(WARNING "BF16 functionality is not available")
set(MARCH_FLAGS -march=rv64gcv_zvfh_zvl128b -mrvv-vector-bits=zvl -mabi=lp64d)
else()
message(STATUS "compile riscv with scalar")
list(APPEND CXX_COMPILE_FLAGS "-march=rv64gc")
endif()
list(APPEND CXX_COMPILE_FLAGS ${MARCH_FLAGS})
else()
message(FATAL_ERROR "vLLM CPU backend requires X86, Power9+ ISA, S390X ISA, ARMv8 or RISC-V support.")
endif()
csrc/cache_kernels.cu
+19 -6
View File
@@ -919,8 +919,8 @@ __global__ void gather_and_maybe_dequant_cache(
// SCALAR_T is the data type of the destination tensor.
// CACHE_T is the stored data type of kv-cache.
// KV_DTYPE is the real data type of kv-cache.
#define CALL_GATHER_CACHE(SCALAR_T, CACHE_T, KV_DTYPE) \
vllm::gather_and_maybe_dequant_cache<SCALAR_T, CACHE_T, KV_DTYPE, 576, \
#define CALL_GATHER_CACHE(SCALAR_T, CACHE_T, KV_DTYPE, ENTRY_SZ) \
vllm::gather_and_maybe_dequant_cache<SCALAR_T, CACHE_T, KV_DTYPE, ENTRY_SZ, \
thread_block_size> \
<<<grid, block, 0, stream>>>( \
reinterpret_cast<CACHE_T*>(src_cache.data_ptr()), \
@@ -931,6 +931,12 @@ __global__ void gather_and_maybe_dequant_cache(
dst_entry_stride, reinterpret_cast<const float*>(scale.data_ptr()), \
seq_starts_ptr);
#define CALL_GATHER_CACHE_576(SCALAR_T, CACHE_T, KV_DTYPE) \
CALL_GATHER_CACHE(SCALAR_T, CACHE_T, KV_DTYPE, 576)
#define CALL_GATHER_CACHE_320(SCALAR_T, CACHE_T, KV_DTYPE) \
CALL_GATHER_CACHE(SCALAR_T, CACHE_T, KV_DTYPE, 320)
// Gather sequences from the cache into the destination tensor.
// - cu_seq_lens contains the cumulative sequence lengths for each batch
// - block_table contains the cache block indices for each sequence
@@ -960,9 +966,10 @@ void gather_and_maybe_dequant_cache(
TORCH_CHECK(seq_starts.value().dtype() == torch::kInt32,
"seq_starts must be int32");
}
TORCH_CHECK(head_dim == 576,
"gather_and_maybe_dequant_cache only support the head_dim to 576 "
"for better performance")
TORCH_CHECK(
head_dim == 320 || head_dim == 576,
"gather_and_maybe_dequant_cache only support the head_dim to 320 or 576 "
"for better performance")
TORCH_CHECK(src_cache.device() == dst.device(),
"src_cache and dst must be on the same device");
@@ -987,7 +994,13 @@ void gather_and_maybe_dequant_cache(
const int32_t* seq_starts_ptr =
seq_starts.has_value() ? seq_starts.value().data_ptr<int32_t>() : nullptr;
DISPATCH_BY_KV_CACHE_DTYPE(dst.dtype(), kv_cache_dtype, CALL_GATHER_CACHE);
if (head_dim == 576) {
DISPATCH_BY_KV_CACHE_DTYPE(dst.dtype(), kv_cache_dtype,
CALL_GATHER_CACHE_576);
} else {
DISPATCH_BY_KV_CACHE_DTYPE(dst.dtype(), kv_cache_dtype,
CALL_GATHER_CACHE_320);
}
}
namespace vllm {
csrc/cpu/cpu_types.hpp
+3
View File
@@ -13,6 +13,9 @@
#elif defined(__aarch64__)
// arm implementation
#include "cpu_types_arm.hpp"
#elif defined(__riscv_v)
// riscv implementation
#include "cpu_types_riscv.hpp"
#else
#warning "unsupported vLLM cpu implementation, vLLM will compile with scalar"
#include "cpu_types_scalar.hpp"
csrc/cpu/cpu_types_riscv.hpp
+832
View File
@@ -0,0 +1,832 @@
#ifndef CPU_TYPES_RISCV_HPP
#define CPU_TYPES_RISCV_HPP
#include <algorithm>
#include <cmath>
#include <cstring>
#include <iostream>
#include <limits>
#include <riscv_vector.h>
#include <torch/all.h>
// ============================================================================
// Vector Register Type Definitions (VLEN=128 bits)
// ============================================================================
typedef vfloat16m1_t fixed_vfloat16m1_t
__attribute__((riscv_rvv_vector_bits(128)));
typedef vfloat16m2_t fixed_vfloat16m2_t
__attribute__((riscv_rvv_vector_bits(256)));
typedef vfloat32m1_t fixed_vfloat32m1_t
__attribute__((riscv_rvv_vector_bits(128)));
typedef vfloat32m2_t fixed_vfloat32m2_t
__attribute__((riscv_rvv_vector_bits(256)));
typedef vfloat32m4_t fixed_vfloat32m4_t
__attribute__((riscv_rvv_vector_bits(512)));
typedef vfloat32m8_t fixed_vfloat32m8_t
__attribute__((riscv_rvv_vector_bits(1024)));
typedef vint32m2_t fixed_vint32m2_t __attribute__((riscv_rvv_vector_bits(256)));
typedef vint32m4_t fixed_vint32m4_t __attribute__((riscv_rvv_vector_bits(512)));
typedef vuint16m1_t fixed_vuint16m1_t
__attribute__((riscv_rvv_vector_bits(128)));
typedef vuint16m2_t fixed_vuint16m2_t
__attribute__((riscv_rvv_vector_bits(256)));
typedef vuint16m4_t fixed_vuint16m4_t
__attribute__((riscv_rvv_vector_bits(512)));
#ifdef RISCV_BF16_SUPPORT
typedef vbfloat16m1_t fixed_vbfloat16m1_t
__attribute__((riscv_rvv_vector_bits(128)));
typedef vbfloat16m2_t fixed_vbfloat16m2_t
__attribute__((riscv_rvv_vector_bits(256)));
typedef vbfloat16m4_t fixed_vbfloat16m4_t
__attribute__((riscv_rvv_vector_bits(512)));
#endif
namespace vec_op {
#ifdef RISCV_BF16_SUPPORT
#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \
AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \
AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__)
#else
#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \
AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__)
#endif
#define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
#define FORCE_INLINE __attribute__((always_inline)) inline
namespace {
template <typename T, T... indexes, typename F>
constexpr void unroll_loop_item(std::integer_sequence<T, indexes...>, F&& f) {
(f(std::integral_constant<T, indexes>{}), ...);
};
} // namespace
template <typename T, T count, typename F,
typename = std::enable_if_t<std::is_invocable_v<F, T>>>
constexpr void unroll_loop(F&& f) {
unroll_loop_item(std::make_integer_sequence<T, count>{}, std::forward<F>(f));
}
template <typename T>
struct Vec {
constexpr static int get_elem_num() { return T::VEC_ELEM_NUM; };
};
struct FP32Vec8;
struct FP32Vec16;
// ============================================================================
// FP16 Implementation
// ============================================================================
struct FP16Vec8 : public Vec<FP16Vec8> {
constexpr static int VEC_ELEM_NUM = 8;
fixed_vfloat16m1_t reg;
explicit FP16Vec8(const void* ptr)
: reg(__riscv_vle16_v_f16m1(static_cast<const _Float16*>(ptr),
VEC_ELEM_NUM)) {};
explicit FP16Vec8(const FP32Vec8&);
void save(void* ptr) const {
__riscv_vse16_v_f16m1(static_cast<_Float16*>(ptr), reg, VEC_ELEM_NUM);
}
void save(void* ptr, int elem_num) const {
__riscv_vse16_v_f16m1(static_cast<_Float16*>(ptr), reg, elem_num);
}
void save_strided(void* ptr, ptrdiff_t stride) const {
ptrdiff_t byte_stride = stride * sizeof(_Float16);
__riscv_vsse16_v_f16m1(static_cast<_Float16*>(ptr), byte_stride, reg,
VEC_ELEM_NUM);
}
};
struct FP16Vec16 : public Vec<FP16Vec16> {
constexpr static int VEC_ELEM_NUM = 16;
fixed_vfloat16m2_t reg;
explicit FP16Vec16(const void* ptr)
: reg(__riscv_vle16_v_f16m2(static_cast<const _Float16*>(ptr),
VEC_ELEM_NUM)) {};
explicit FP16Vec16(const FP32Vec16& vec);
void save(void* ptr) const {
__riscv_vse16_v_f16m2(static_cast<_Float16*>(ptr), reg, VEC_ELEM_NUM);
}
void save(void* ptr, int elem_num) const {
__riscv_vse16_v_f16m2(static_cast<_Float16*>(ptr), reg, elem_num);
}
void save_strided(void* ptr, ptrdiff_t stride) const {
ptrdiff_t byte_stride = stride * sizeof(_Float16);
__riscv_vsse16_v_f16m2(static_cast<_Float16*>(ptr), byte_stride, reg,
VEC_ELEM_NUM);
}
};
// ============================================================================
// BF16 Implementation
// ============================================================================
#ifdef RISCV_BF16_SUPPORT
FORCE_INLINE fixed_vuint16m1_t bf16_to_u16(fixed_vbfloat16m1_t v) {
return __riscv_vreinterpret_v_bf16m1_u16m1(v);
}
FORCE_INLINE fixed_vuint16m2_t bf16_to_u16(fixed_vbfloat16m2_t v) {
return __riscv_vreinterpret_v_bf16m2_u16m2(v);
}
FORCE_INLINE fixed_vuint16m4_t bf16_to_u16(fixed_vbfloat16m4_t v) {
return __riscv_vreinterpret_v_bf16m4_u16m4(v);
}
struct BF16Vec8 : public Vec<BF16Vec8> {
constexpr static int VEC_ELEM_NUM = 8;
fixed_vbfloat16m1_t reg;
explicit BF16Vec8(const void* ptr)
: reg(__riscv_vreinterpret_v_u16m1_bf16m1(__riscv_vle16_v_u16m1(
reinterpret_cast<const uint16_t*>(ptr), VEC_ELEM_NUM))) {};
explicit BF16Vec8(fixed_vbfloat16m1_t data) : reg(data) {};
explicit BF16Vec8(const FP32Vec8&);
void save(void* ptr) const {
__riscv_vse16_v_u16m1(reinterpret_cast<uint16_t*>(ptr), bf16_to_u16(reg),
VEC_ELEM_NUM);
}
void save(void* ptr, int elem_num) const {
__riscv_vse16_v_u16m1(reinterpret_cast<uint16_t*>(ptr), bf16_to_u16(reg),
elem_num);
}
void save_strided(void* ptr, ptrdiff_t stride) const {
ptrdiff_t byte_stride = stride * sizeof(uint16_t);
__riscv_vsse16_v_u16m1(reinterpret_cast<uint16_t*>(ptr), byte_stride,
bf16_to_u16(reg), VEC_ELEM_NUM);
}
};
struct BF16Vec16 : public Vec<BF16Vec16> {
constexpr static int VEC_ELEM_NUM = 16;
fixed_vbfloat16m2_t reg;
explicit BF16Vec16(const void* ptr)
: reg(__riscv_vreinterpret_v_u16m2_bf16m2(__riscv_vle16_v_u16m2(
reinterpret_cast<const uint16_t*>(ptr), VEC_ELEM_NUM))) {};
explicit BF16Vec16(fixed_vbfloat16m2_t data) : reg(data) {};
explicit BF16Vec16(const FP32Vec16&);
void save(void* ptr) const {
__riscv_vse16_v_u16m2(reinterpret_cast<uint16_t*>(ptr), bf16_to_u16(reg),
VEC_ELEM_NUM);
}
void save(void* ptr, int elem_num) const {
__riscv_vse16_v_u16m2(reinterpret_cast<uint16_t*>(ptr), bf16_to_u16(reg),
elem_num);
}
void save_strided(void* ptr, ptrdiff_t stride) const {
ptrdiff_t byte_stride = stride * sizeof(uint16_t);
__riscv_vsse16_v_u16m2(reinterpret_cast<uint16_t*>(ptr), byte_stride,
bf16_to_u16(reg), VEC_ELEM_NUM);
}
};
struct BF16Vec32 : public Vec<BF16Vec32> {
constexpr static int VEC_ELEM_NUM = 32;
fixed_vbfloat16m4_t reg;
explicit BF16Vec32(const void* ptr)
: reg(__riscv_vreinterpret_v_u16m4_bf16m4(__riscv_vle16_v_u16m4(
reinterpret_cast<const uint16_t*>(ptr), VEC_ELEM_NUM))) {};
explicit BF16Vec32(fixed_vbfloat16m4_t data) : reg(data) {};
explicit BF16Vec32(const BF16Vec8& v) {
fixed_vuint16m1_t u16_val = bf16_to_u16(v.reg);
fixed_vuint16m4_t u16_combined =
__riscv_vcreate_v_u16m1_u16m4(u16_val, u16_val, u16_val, u16_val);
reg = __riscv_vreinterpret_v_u16m4_bf16m4(u16_combined);
};
void save(void* ptr) const {
__riscv_vse16_v_u16m4(reinterpret_cast<uint16_t*>(ptr), bf16_to_u16(reg),
VEC_ELEM_NUM);
}
void save(void* ptr, int elem_num) const {
__riscv_vse16_v_u16m4(reinterpret_cast<uint16_t*>(ptr), bf16_to_u16(reg),
elem_num);
}
void save_strided(void* ptr, ptrdiff_t stride) const {
ptrdiff_t byte_stride = stride * sizeof(uint16_t);
__riscv_vsse16_v_u16m4(reinterpret_cast<uint16_t*>(ptr), byte_stride,
bf16_to_u16(reg), VEC_ELEM_NUM);
}
};
#else
// ============================================================================
// BF16 Fallback Implementation (FP32 Simulation)
// ============================================================================
struct BF16Vec8 : public Vec<BF16Vec8> {
constexpr static int VEC_ELEM_NUM = 8;
fixed_vfloat32m2_t reg_fp32;
explicit BF16Vec8(const void* ptr) {
const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
float tmp[8];
for (int i = 0; i < 8; ++i) {
uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
std::memcpy(&tmp[i], &v, 4);
}
reg_fp32 = __riscv_vle32_v_f32m2(tmp, 8);
}
explicit BF16Vec8(const FP32Vec8&);
void save(void* ptr) const {
float tmp[8];
__riscv_vse32_v_f32m2(tmp, reg_fp32, 8);
uint16_t* u16 = static_cast<uint16_t*>(ptr);
for (int i = 0; i < 8; ++i) {
uint32_t v;
std::memcpy(&v, &tmp[i], 4);
u16[i] = static_cast<uint16_t>(v >> 16);
}
}
void save(void* ptr, int elem_num) const {
float tmp[8];
__riscv_vse32_v_f32m2(tmp, reg_fp32, 8);
uint16_t* u16 = static_cast<uint16_t*>(ptr);
for (int i = 0; i < elem_num; ++i) {
uint32_t v;
std::memcpy(&v, &tmp[i], 4);
u16[i] = static_cast<uint16_t>(v >> 16);
}
}
void save_strided(void* ptr, ptrdiff_t stride) const {
float tmp[8];
__riscv_vse32_v_f32m2(tmp, reg_fp32, 8);
uint8_t* u8 = static_cast<uint8_t*>(ptr);
ptrdiff_t byte_stride = stride * sizeof(uint16_t);
for (int i = 0; i < 8; ++i) {
uint32_t v;
std::memcpy(&v, &tmp[i], 4);
uint16_t val = static_cast<uint16_t>(v >> 16);
*reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
}
}
};
struct BF16Vec16 : public Vec<BF16Vec16> {
constexpr static int VEC_ELEM_NUM = 16;
fixed_vfloat32m4_t reg_fp32;
explicit BF16Vec16(const void* ptr) {
const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
float tmp[16];
for (int i = 0; i < 16; ++i) {
uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
std::memcpy(&tmp[i], &v, 4);
}
reg_fp32 = __riscv_vle32_v_f32m4(tmp, 16);
}
explicit BF16Vec16(const FP32Vec16&);
void save(void* ptr) const {
float tmp[16];
__riscv_vse32_v_f32m4(tmp, reg_fp32, 16);
uint16_t* u16 = static_cast<uint16_t*>(ptr);
for (int i = 0; i < 16; ++i) {
uint32_t v;
std::memcpy(&v, &tmp[i], 4);
u16[i] = static_cast<uint16_t>(v >> 16);
}
}
void save(void* ptr, int elem_num) const {
float tmp[16];
__riscv_vse32_v_f32m4(tmp, reg_fp32, 16);
uint16_t* u16 = static_cast<uint16_t*>(ptr);
for (int i = 0; i < elem_num; ++i) {
uint32_t v;
std::memcpy(&v, &tmp[i], 4);
u16[i] = static_cast<uint16_t>(v >> 16);
}
}
void save_strided(void* ptr, ptrdiff_t stride) const {
float tmp[16];
__riscv_vse32_v_f32m4(tmp, reg_fp32, 16);
uint8_t* u8 = static_cast<uint8_t*>(ptr);
ptrdiff_t byte_stride = stride * sizeof(uint16_t);
for (int i = 0; i < 16; ++i) {
uint32_t v;
std::memcpy(&v, &tmp[i], 4);
uint16_t val = static_cast<uint16_t>(v >> 16);
*reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
}
}
};
struct BF16Vec32 : public Vec<BF16Vec32> {
constexpr static int VEC_ELEM_NUM = 32;
fixed_vfloat32m8_t reg_fp32;
explicit BF16Vec32(const void* ptr) {
const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
float tmp[32];
for (int i = 0; i < 32; ++i) {
uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
std::memcpy(&tmp[i], &v, 4);
}
reg_fp32 = __riscv_vle32_v_f32m8(tmp, 32);
}
explicit BF16Vec32(const BF16Vec8& v) {
float tmp_small[8];
__riscv_vse32_v_f32m2(tmp_small, v.reg_fp32, 8);
float tmp_large[32];
for (int i = 0; i < 4; ++i) {
std::memcpy(tmp_large + (i * 8), tmp_small, 8 * sizeof(float));
}
reg_fp32 = __riscv_vle32_v_f32m8(tmp_large, 32);
}
void save(void* ptr) const {
float tmp[32];
__riscv_vse32_v_f32m8(tmp, reg_fp32, 32);
uint16_t* u16 = static_cast<uint16_t*>(ptr);
for (int i = 0; i < 32; ++i) {
uint32_t v;
std::memcpy(&v, &tmp[i], 4);
u16[i] = static_cast<uint16_t>(v >> 16);
}
}
void save(void* ptr, int elem_num) const {
float tmp[32];
__riscv_vse32_v_f32m8(tmp, reg_fp32, 32);
uint16_t* u16 = static_cast<uint16_t*>(ptr);
for (int i = 0; i < elem_num; ++i) {
uint32_t v;
std::memcpy(&v, &tmp[i], 4);
u16[i] = static_cast<uint16_t>(v >> 16);
}
}
void save_strided(void* ptr, ptrdiff_t stride) const {
float tmp[32];
__riscv_vse32_v_f32m8(tmp, reg_fp32, 32);
uint8_t* u8 = static_cast<uint8_t*>(ptr);
ptrdiff_t byte_stride = stride * sizeof(uint16_t);
for (int i = 0; i < 32; ++i) {
uint32_t v;
std::memcpy(&v, &tmp[i], 4);
uint16_t val = static_cast<uint16_t>(v >> 16);
*reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
}
}
};
#endif
// ============================================================================
// FP32 Implementation
// ============================================================================
struct FP32Vec4 : public Vec<FP32Vec4> {
constexpr static int VEC_ELEM_NUM = 4;
fixed_vfloat32m1_t reg;
explicit FP32Vec4(float v) : reg(__riscv_vfmv_v_f_f32m1(v, VEC_ELEM_NUM)) {};
explicit FP32Vec4() : reg(__riscv_vfmv_v_f_f32m1(0.0f, VEC_ELEM_NUM)) {};
explicit FP32Vec4(const float* ptr)
: reg(__riscv_vle32_v_f32m1(ptr, VEC_ELEM_NUM)) {};
explicit FP32Vec4(fixed_vfloat32m1_t data) : reg(data) {};
explicit FP32Vec4(const FP32Vec4& data) : reg(data.reg) {};
void save(float* ptr) const { __riscv_vse32_v_f32m1(ptr, reg, VEC_ELEM_NUM); }
void save(float* ptr, int elem_num) const {
__riscv_vse32_v_f32m1(ptr, reg, elem_num);
}
};
struct FP32Vec8 : public Vec<FP32Vec8> {
constexpr static int VEC_ELEM_NUM = 8;
fixed_vfloat32m2_t reg;
explicit FP32Vec8(float v) : reg(__riscv_vfmv_v_f_f32m2(v, VEC_ELEM_NUM)) {};
explicit FP32Vec8() : reg(__riscv_vfmv_v_f_f32m2(0.0f, VEC_ELEM_NUM)) {};
explicit FP32Vec8(const float* ptr)
: reg(__riscv_vle32_v_f32m2(ptr, VEC_ELEM_NUM)) {};
explicit FP32Vec8(fixed_vfloat32m2_t data) : reg(data) {};
explicit FP32Vec8(const FP32Vec8& data) : reg(data.reg) {};
explicit FP32Vec8(const FP16Vec8& v)
: reg(__riscv_vfwcvt_f_f_v_f32m2(v.reg, VEC_ELEM_NUM)) {};
explicit FP32Vec8(fixed_vfloat16m1_t v)
: reg(__riscv_vfwcvt_f_f_v_f32m2(v, VEC_ELEM_NUM)) {};
#ifdef RISCV_BF16_SUPPORT
explicit FP32Vec8(fixed_vbfloat16m1_t v)
: reg(__riscv_vfwcvtbf16_f_f_v_f32m2(v, VEC_ELEM_NUM)) {};
explicit FP32Vec8(const BF16Vec8& v)
: reg(__riscv_vfwcvtbf16_f_f_v_f32m2(v.reg, VEC_ELEM_NUM)) {};
#else
explicit FP32Vec8(const BF16Vec8& v) : reg(v.reg_fp32) {};
#endif
float reduce_sum() const {
fixed_vfloat32m1_t scalar = __riscv_vfmv_s_f_f32m1(0.0f, 1);
scalar = __riscv_vfredusum_vs_f32m2_f32m1(reg, scalar, VEC_ELEM_NUM);
return __riscv_vfmv_f_s_f32m1_f32(scalar);
}
FP32Vec8 operator*(const FP32Vec8& b) const {
return FP32Vec8(__riscv_vfmul_vv_f32m2(reg, b.reg, VEC_ELEM_NUM));
}
FP32Vec8 operator+(const FP32Vec8& b) const {
return FP32Vec8(__riscv_vfadd_vv_f32m2(reg, b.reg, VEC_ELEM_NUM));
}
FP32Vec8 operator-(const FP32Vec8& b) const {
return FP32Vec8(__riscv_vfsub_vv_f32m2(reg, b.reg, VEC_ELEM_NUM));
}
FP32Vec8 operator/(const FP32Vec8& b) const {
return FP32Vec8(__riscv_vfdiv_vv_f32m2(reg, b.reg, VEC_ELEM_NUM));
}
FP32Vec8 min(const FP32Vec8& b) const {
return FP32Vec8(__riscv_vfmin_vv_f32m2(reg, b.reg, VEC_ELEM_NUM));
}
FP32Vec8 max(const FP32Vec8& b) const {
return FP32Vec8(__riscv_vfmax_vv_f32m2(reg, b.reg, VEC_ELEM_NUM));
}
FP32Vec8 abs() const {
return FP32Vec8(__riscv_vfabs_v_f32m2(reg, VEC_ELEM_NUM));
}
FP32Vec8 min(const FP32Vec8& b, int elem_num) const {
return FP32Vec8(__riscv_vfmin_vv_f32m2(reg, b.reg, elem_num));
}
FP32Vec8 max(const FP32Vec8& b, int elem_num) const {
return FP32Vec8(__riscv_vfmax_vv_f32m2(reg, b.reg, elem_num));
}
FP32Vec8 clamp(const FP32Vec8& min_v, const FP32Vec8& max_v) const {
fixed_vfloat32m2_t temp =
__riscv_vfmax_vv_f32m2(min_v.reg, reg, VEC_ELEM_NUM);
return FP32Vec8(__riscv_vfmin_vv_f32m2(max_v.reg, temp, VEC_ELEM_NUM));
}
void save(float* ptr) const { __riscv_vse32_v_f32m2(ptr, reg, VEC_ELEM_NUM); }
void save(float* ptr, int elem_num) const {
__riscv_vse32_v_f32m2(ptr, reg, elem_num);
}
void save_strided(float* ptr, ptrdiff_t stride) const {
ptrdiff_t byte_stride = stride * sizeof(float);
__riscv_vsse32_v_f32m2(ptr, byte_stride, reg, VEC_ELEM_NUM);
}
FP32Vec8 exp() const {
const float inv_ln2 = 1.44269504088896341f;
fixed_vfloat32m2_t x_scaled =
__riscv_vfmul_vf_f32m2(reg, inv_ln2, VEC_ELEM_NUM);
fixed_vint32m2_t n_int = __riscv_vfcvt_x_f_v_i32m2(x_scaled, VEC_ELEM_NUM);
fixed_vfloat32m2_t n_float = __riscv_vfcvt_f_x_v_f32m2(n_int, VEC_ELEM_NUM);
fixed_vfloat32m2_t r =
__riscv_vfsub_vv_f32m2(x_scaled, n_float, VEC_ELEM_NUM);
fixed_vfloat32m2_t poly =
__riscv_vfmv_v_f_f32m2(0.001333355810164f, VEC_ELEM_NUM);
poly = __riscv_vfmul_vv_f32m2(poly, r, VEC_ELEM_NUM);
poly = __riscv_vfadd_vf_f32m2(poly, 0.009618129107628f, VEC_ELEM_NUM);
poly = __riscv_vfmul_vv_f32m2(poly, r, VEC_ELEM_NUM);
poly = __riscv_vfadd_vf_f32m2(poly, 0.055504108664821f, VEC_ELEM_NUM);
poly = __riscv_vfmul_vv_f32m2(poly, r, VEC_ELEM_NUM);
poly = __riscv_vfadd_vf_f32m2(poly, 0.240226506959101f, VEC_ELEM_NUM);
poly = __riscv_vfmul_vv_f32m2(poly, r, VEC_ELEM_NUM);
poly = __riscv_vfadd_vf_f32m2(poly, 0.693147180559945f, VEC_ELEM_NUM);
poly = __riscv_vfmul_vv_f32m2(poly, r, VEC_ELEM_NUM);
poly = __riscv_vfadd_vf_f32m2(poly, 1.0f, VEC_ELEM_NUM);
fixed_vint32m2_t biased_exp =
__riscv_vadd_vx_i32m2(n_int, 127, VEC_ELEM_NUM);
biased_exp = __riscv_vmax_vx_i32m2(biased_exp, 0, VEC_ELEM_NUM);
fixed_vint32m2_t exponent_bits =
__riscv_vsll_vx_i32m2(biased_exp, 23, VEC_ELEM_NUM);
fixed_vfloat32m2_t scale =
__riscv_vreinterpret_v_i32m2_f32m2(exponent_bits);
return FP32Vec8(__riscv_vfmul_vv_f32m2(poly, scale, VEC_ELEM_NUM));
}
FP32Vec8 tanh() const {
fixed_vfloat32m2_t x_clamped = __riscv_vfmin_vf_f32m2(
__riscv_vfmax_vf_f32m2(reg, -9.0f, VEC_ELEM_NUM), 9.0f, VEC_ELEM_NUM);
fixed_vfloat32m2_t x2 =
__riscv_vfmul_vf_f32m2(x_clamped, 2.0f, VEC_ELEM_NUM);
FP32Vec8 exp_val = FP32Vec8(x2).exp();
fixed_vfloat32m2_t num =
__riscv_vfsub_vf_f32m2(exp_val.reg, 1.0f, VEC_ELEM_NUM);
fixed_vfloat32m2_t den =
__riscv_vfadd_vf_f32m2(exp_val.reg, 1.0f, VEC_ELEM_NUM);
return FP32Vec8(__riscv_vfdiv_vv_f32m2(num, den, VEC_ELEM_NUM));
}
FP32Vec8 er() const {
const float p = 0.3275911f, a1 = 0.254829592f, a2 = -0.284496736f,
a3 = 1.421413741f, a4 = -1.453152027f, a5 = 1.061405429f;
fixed_vfloat32m2_t abs_x = __riscv_vfabs_v_f32m2(reg, VEC_ELEM_NUM);
fixed_vfloat32m2_t t = __riscv_vfadd_vf_f32m2(
__riscv_vfmul_vf_f32m2(abs_x, p, VEC_ELEM_NUM), 1.0f, VEC_ELEM_NUM);
t = __riscv_vfrdiv_vf_f32m2(t, 1.0f, VEC_ELEM_NUM);
fixed_vfloat32m2_t poly = __riscv_vfmv_v_f_f32m2(a5, VEC_ELEM_NUM);
poly = __riscv_vfadd_vf_f32m2(__riscv_vfmul_vv_f32m2(poly, t, VEC_ELEM_NUM),
a4, VEC_ELEM_NUM);
poly = __riscv_vfadd_vf_f32m2(__riscv_vfmul_vv_f32m2(poly, t, VEC_ELEM_NUM),
a3, VEC_ELEM_NUM);
poly = __riscv_vfadd_vf_f32m2(__riscv_vfmul_vv_f32m2(poly, t, VEC_ELEM_NUM),
a2, VEC_ELEM_NUM);
poly = __riscv_vfadd_vf_f32m2(__riscv_vfmul_vv_f32m2(poly, t, VEC_ELEM_NUM),
a1, VEC_ELEM_NUM);
poly = __riscv_vfmul_vv_f32m2(poly, t, VEC_ELEM_NUM);
fixed_vfloat32m2_t exp_val =
FP32Vec8(__riscv_vfneg_v_f32m2(
__riscv_vfmul_vv_f32m2(abs_x, abs_x, VEC_ELEM_NUM),
VEC_ELEM_NUM))
.exp()
.reg;
fixed_vfloat32m2_t res = __riscv_vfrsub_vf_f32m2(
__riscv_vfmul_vv_f32m2(poly, exp_val, VEC_ELEM_NUM), 1.0f,
VEC_ELEM_NUM);
vbool16_t mask = __riscv_vmflt_vf_f32m2_b16(reg, 0.0f, VEC_ELEM_NUM);
return FP32Vec8(__riscv_vfneg_v_f32m2_m(mask, res, VEC_ELEM_NUM));
}
};
struct FP32Vec16 : public Vec<FP32Vec16> {
constexpr static int VEC_ELEM_NUM = 16;
fixed_vfloat32m4_t reg;
explicit FP32Vec16(float v) : reg(__riscv_vfmv_v_f_f32m4(v, VEC_ELEM_NUM)) {};
explicit FP32Vec16() : reg(__riscv_vfmv_v_f_f32m4(0.0f, VEC_ELEM_NUM)) {};
explicit FP32Vec16(const float* ptr)
: reg(__riscv_vle32_v_f32m4(ptr, VEC_ELEM_NUM)) {};
explicit FP32Vec16(fixed_vfloat32m4_t data) : reg(data) {};
explicit FP32Vec16(const FP32Vec8& data)
: reg(__riscv_vcreate_v_f32m2_f32m4(data.reg, data.reg)) {};
explicit FP32Vec16(const FP32Vec16& data) : reg(data.reg) {};
explicit FP32Vec16(const FP16Vec16& v);
#ifdef RISCV_BF16_SUPPORT
explicit FP32Vec16(fixed_vbfloat16m2_t v)
: reg(__riscv_vfwcvtbf16_f_f_v_f32m4(v, VEC_ELEM_NUM)) {};
explicit FP32Vec16(const BF16Vec16& v)
: reg(__riscv_vfwcvtbf16_f_f_v_f32m4(v.reg, VEC_ELEM_NUM)) {};
#else
explicit FP32Vec16(const BF16Vec16& v) : reg(v.reg_fp32) {};
#endif
FP32Vec16 operator+(const FP32Vec16& b) const {
return FP32Vec16(__riscv_vfadd_vv_f32m4(reg, b.reg, VEC_ELEM_NUM));
}
FP32Vec16 operator-(const FP32Vec16& b) const {
return FP32Vec16(__riscv_vfsub_vv_f32m4(reg, b.reg, VEC_ELEM_NUM));
}
FP32Vec16 operator*(const FP32Vec16& b) const {
return FP32Vec16(__riscv_vfmul_vv_f32m4(reg, b.reg, VEC_ELEM_NUM));
}
FP32Vec16 operator/(const FP32Vec16& b) const {
return FP32Vec16(__riscv_vfdiv_vv_f32m4(reg, b.reg, VEC_ELEM_NUM));
}
FP32Vec16 fma(const FP32Vec16& a, const FP32Vec16& b) const {
return FP32Vec16(__riscv_vfmacc_vv_f32m4(reg, a.reg, b.reg, VEC_ELEM_NUM));
}
float reduce_sum() const {
fixed_vfloat32m1_t scalar = __riscv_vfmv_s_f_f32m1(0.0f, 1);
scalar = __riscv_vfredusum_vs_f32m4_f32m1(reg, scalar, VEC_ELEM_NUM);
return __riscv_vfmv_f_s_f32m1_f32(scalar);
}
float reduce_max() const {
fixed_vfloat32m1_t scalar =
__riscv_vfmv_s_f_f32m1(std::numeric_limits<float>::lowest(), 1);
scalar = __riscv_vfredmax_vs_f32m4_f32m1(reg, scalar, VEC_ELEM_NUM);
return __riscv_vfmv_f_s_f32m1_f32(scalar);
}
float reduce_min() const {
fixed_vfloat32m1_t scalar =
__riscv_vfmv_s_f_f32m1(std::numeric_limits<float>::max(), 1);
scalar = __riscv_vfredmin_vs_f32m4_f32m1(reg, scalar, VEC_ELEM_NUM);
return __riscv_vfmv_f_s_f32m1_f32(scalar);
}
template <int group_size>
float reduce_sub_sum(int idx) {
static_assert(VEC_ELEM_NUM % group_size == 0);
const int start = idx * group_size;
vuint32m4_t indices = __riscv_vid_v_u32m4(VEC_ELEM_NUM);
vbool8_t mask = __riscv_vmand_mm_b8(
__riscv_vmsgeu_vx_u32m4_b8(indices, start, VEC_ELEM_NUM),
__riscv_vmsltu_vx_u32m4_b8(indices, start + group_size, VEC_ELEM_NUM),
VEC_ELEM_NUM);
fixed_vfloat32m1_t scalar = __riscv_vfmv_s_f_f32m1(0.0f, 1);
scalar =
__riscv_vfredusum_vs_f32m4_f32m1_m(mask, reg, scalar, VEC_ELEM_NUM);
return __riscv_vfmv_f_s_f32m1_f32(scalar);
};
FP32Vec16 max(const FP32Vec16& b) const {
return FP32Vec16(__riscv_vfmax_vv_f32m4(reg, b.reg, VEC_ELEM_NUM));
}
FP32Vec16 min(const FP32Vec16& b) const {
return FP32Vec16(__riscv_vfmin_vv_f32m4(reg, b.reg, VEC_ELEM_NUM));
}
FP32Vec16 abs() const {
return FP32Vec16(__riscv_vfabs_v_f32m4(reg, VEC_ELEM_NUM));
}
FP32Vec16 clamp(const FP32Vec16& min_v, const FP32Vec16& max_v) const {
return FP32Vec16(__riscv_vfmin_vv_f32m4(
max_v.reg, __riscv_vfmax_vv_f32m4(min_v.reg, reg, VEC_ELEM_NUM),
VEC_ELEM_NUM));
}
void save(float* ptr) const { __riscv_vse32_v_f32m4(ptr, reg, VEC_ELEM_NUM); }
void save(float* ptr, int elem_num) const {
__riscv_vse32_v_f32m4(ptr, reg, elem_num);
}
void save_strided(float* ptr, ptrdiff_t stride) const {
ptrdiff_t byte_stride = stride * sizeof(float);
__riscv_vsse32_v_f32m4(ptr, byte_stride, reg, VEC_ELEM_NUM);
}
FP32Vec16 exp() const {
const float inv_ln2 = 1.44269504088896341f;
fixed_vfloat32m4_t x_scaled =
__riscv_vfmul_vf_f32m4(reg, inv_ln2, VEC_ELEM_NUM);
fixed_vint32m4_t n_int = __riscv_vfcvt_x_f_v_i32m4(x_scaled, VEC_ELEM_NUM);
fixed_vfloat32m4_t n_float = __riscv_vfcvt_f_x_v_f32m4(n_int, VEC_ELEM_NUM);
fixed_vfloat32m4_t r =
__riscv_vfsub_vv_f32m4(x_scaled, n_float, VEC_ELEM_NUM);
fixed_vfloat32m4_t poly =
__riscv_vfmv_v_f_f32m4(0.001333355810164f, VEC_ELEM_NUM);
poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, r, VEC_ELEM_NUM),
0.009618129107628f, VEC_ELEM_NUM);
poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, r, VEC_ELEM_NUM),
0.055504108664821f, VEC_ELEM_NUM);
poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, r, VEC_ELEM_NUM),
0.240226506959101f, VEC_ELEM_NUM);
poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, r, VEC_ELEM_NUM),
0.693147180559945f, VEC_ELEM_NUM);
poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, r, VEC_ELEM_NUM),
1.0f, VEC_ELEM_NUM);
fixed_vint32m4_t biased_exp = __riscv_vmax_vx_i32m4(
__riscv_vadd_vx_i32m4(n_int, 127, VEC_ELEM_NUM), 0, VEC_ELEM_NUM);
fixed_vfloat32m4_t scale = __riscv_vreinterpret_v_i32m4_f32m4(
__riscv_vsll_vx_i32m4(biased_exp, 23, VEC_ELEM_NUM));
return FP32Vec16(__riscv_vfmul_vv_f32m4(poly, scale, VEC_ELEM_NUM));
}
FP32Vec16 tanh() const {
fixed_vfloat32m4_t x_clamped = __riscv_vfmin_vf_f32m4(
__riscv_vfmax_vf_f32m4(reg, -9.0f, VEC_ELEM_NUM), 9.0f, VEC_ELEM_NUM);
FP32Vec16 exp_val =
FP32Vec16(__riscv_vfmul_vf_f32m4(x_clamped, 2.0f, VEC_ELEM_NUM)).exp();
return FP32Vec16(__riscv_vfdiv_vv_f32m4(
__riscv_vfsub_vf_f32m4(exp_val.reg, 1.0f, VEC_ELEM_NUM),
__riscv_vfadd_vf_f32m4(exp_val.reg, 1.0f, VEC_ELEM_NUM), VEC_ELEM_NUM));
}
FP32Vec16 er() const {
const float p = 0.3275911f, a1 = 0.254829592f, a2 = -0.284496736f,
a3 = 1.421413741f, a4 = -1.453152027f, a5 = 1.061405429f;
fixed_vfloat32m4_t abs_x = __riscv_vfabs_v_f32m4(reg, VEC_ELEM_NUM);
fixed_vfloat32m4_t t = __riscv_vfrdiv_vf_f32m4(
__riscv_vfadd_vf_f32m4(__riscv_vfmul_vf_f32m4(abs_x, p, VEC_ELEM_NUM),
1.0f, VEC_ELEM_NUM),
1.0f, VEC_ELEM_NUM);
fixed_vfloat32m4_t poly = __riscv_vfmv_v_f_f32m4(a5, VEC_ELEM_NUM);
poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, t, VEC_ELEM_NUM),
a4, VEC_ELEM_NUM);
poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, t, VEC_ELEM_NUM),
a3, VEC_ELEM_NUM);
poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, t, VEC_ELEM_NUM),
a2, VEC_ELEM_NUM);
poly = __riscv_vfadd_vf_f32m4(__riscv_vfmul_vv_f32m4(poly, t, VEC_ELEM_NUM),
a1, VEC_ELEM_NUM);
poly = __riscv_vfmul_vv_f32m4(poly, t, VEC_ELEM_NUM);
fixed_vfloat32m4_t exp_val =
FP32Vec16(__riscv_vfneg_v_f32m4(
__riscv_vfmul_vv_f32m4(abs_x, abs_x, VEC_ELEM_NUM),
VEC_ELEM_NUM))
.exp()
.reg;
fixed_vfloat32m4_t res = __riscv_vfrsub_vf_f32m4(
__riscv_vfmul_vv_f32m4(poly, exp_val, VEC_ELEM_NUM), 1.0f,
VEC_ELEM_NUM);
vbool8_t mask = __riscv_vmflt_vf_f32m4_b8(reg, 0.0f, VEC_ELEM_NUM);
return FP32Vec16(__riscv_vfneg_v_f32m4_m(mask, res, VEC_ELEM_NUM));
}
};
// ============================================================================
// Type Traits & Global Helpers
// ============================================================================
template <typename T>
struct VecType {
using vec_type = void;
using vec_t = void;
};
template <typename T>
using vec_t = typename VecType<T>::vec_type;
template <>
struct VecType<float> {
using vec_type = FP32Vec8;
using vec_t = FP32Vec8;
};
template <>
struct VecType<c10::Half> {
using vec_type = FP16Vec8;
using vec_t = FP16Vec8;
};
template <>
struct VecType<c10::BFloat16> {
using vec_type = BF16Vec8;
using vec_t = BF16Vec8;
};
template <typename T>
void storeFP32(float v, T* ptr) {
*ptr = v;
}
template <>
inline void storeFP32<c10::Half>(float v, c10::Half* ptr) {
*reinterpret_cast<_Float16*>(ptr) = static_cast<_Float16>(v);
}
inline FP16Vec16::FP16Vec16(const FP32Vec16& v) {
reg = __riscv_vfncvt_f_f_w_f16m2(v.reg, VEC_ELEM_NUM);
}
inline FP16Vec8::FP16Vec8(const FP32Vec8& v) {
reg = __riscv_vfncvt_f_f_w_f16m1(v.reg, VEC_ELEM_NUM);
}
inline FP32Vec16::FP32Vec16(const FP16Vec16& v) {
reg = __riscv_vfwcvt_f_f_v_f32m4(v.reg, VEC_ELEM_NUM);
}
inline void fma(FP32Vec16& acc, const FP32Vec16& a, const FP32Vec16& b) {
acc = acc.fma(a, b);
}
#ifdef RISCV_BF16_SUPPORT
template <>
inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
*ptr = static_cast<__bf16>(v);
};
inline BF16Vec8::BF16Vec8(const FP32Vec8& v)
: reg(__riscv_vfncvtbf16_f_f_w_bf16m1(v.reg, VEC_ELEM_NUM)) {};
inline BF16Vec16::BF16Vec16(const FP32Vec16& v)
: reg(__riscv_vfncvtbf16_f_f_w_bf16m2(v.reg, VEC_ELEM_NUM)) {};
#else
template <>
inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
uint32_t val;
std::memcpy(&val, &v, 4);
*reinterpret_cast<uint16_t*>(ptr) = static_cast<uint16_t>(val >> 16);
}
inline BF16Vec8::BF16Vec8(const FP32Vec8& v) : reg_fp32(v.reg) {}
inline BF16Vec16::BF16Vec16(const FP32Vec16& v) : reg_fp32(v.reg) {}
#endif
inline void prefetch(const void* addr) { __builtin_prefetch(addr, 0, 1); }
} // namespace vec_op
#ifndef CPU_KERNEL_GUARD_IN
#define CPU_KERNEL_GUARD_IN(NAME)
#endif
#ifndef CPU_KERNEL_GUARD_OUT
#define CPU_KERNEL_GUARD_OUT(NAME)
#endif
#endif // CPU_TYPES_RISCV_HPP
csrc/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu
+45 -24
View File
@@ -15,31 +15,33 @@ __device__ void rms_norm_dynamic_per_token_quant_vec(
scalar_t const* __restrict__ input, // [..., hidden_size]
scalar_t const* __restrict__ weight, // [hidden_size]
float const* scale_ub, float const var_epsilon, int32_t const hidden_size,
scalar_t* __restrict__ residual = nullptr) {
int32_t const input_stride, scalar_t* __restrict__ residual = nullptr) {
float rms = 0.0f;
float token_scale = 0.0f;
// Compute rms
vllm::vectorized::compute_rms<scalar_t, has_residual>(
&rms, input, hidden_size, var_epsilon, residual);
&rms, input, hidden_size, input_stride, var_epsilon, residual);
// Compute scale
vllm::vectorized::compute_dynamic_per_token_scales<scalar_t, scalar_out_t,
has_residual>(
&token_scale, scales, input, weight, rms, scale_ub, hidden_size,
residual);
input_stride, residual);
// RMS Norm + Quant
if constexpr (std::is_same_v<scalar_out_t, int8_t>) {
token_scale = 1.0f / token_scale;
vllm::vectorized::norm_and_quant<scalar_t, scalar_out_t, true,
has_residual>(
out, input, weight, rms, &token_scale, hidden_size, residual);
has_residual>(out, input, weight, rms,
&token_scale, hidden_size,
input_stride, residual);
} else {
// FP8 - Do not invert token_scale for exact match with FBGemm
vllm::vectorized::norm_and_quant<scalar_t, scalar_out_t, false,
has_residual>(
out, input, weight, rms, &token_scale, hidden_size, residual);
has_residual>(out, input, weight, rms,
&token_scale, hidden_size,
input_stride, residual);
}
}
@@ -51,38 +53,40 @@ __global__ void rms_norm_dynamic_per_token_quant_kernel(
scalar_t const* __restrict__ input, // [..., hidden_size]
scalar_t const* __restrict__ weight, // [hidden_size]
float const* scale_ub, float const var_epsilon, int32_t const hidden_size,
scalar_t* __restrict__ residual = nullptr) {
int32_t const input_stride, scalar_t* __restrict__ residual = nullptr) {
// For vectorization, token_input and token_output pointers need to be
// aligned at 8-byte and 4-byte addresses respectively.
bool const can_vectorize = hidden_size % 4 == 0;
bool const can_vectorize = hidden_size % 4 == 0 and input_stride % 4 == 0;
if (can_vectorize) {
return rms_norm_dynamic_per_token_quant_vec<scalar_t, scalar_out_t,
has_residual>(
out, scales, input, weight, scale_ub, var_epsilon, hidden_size,
residual);
input_stride, residual);
}
float rms = 0.0f;
float token_scale = 0.0f;
// Compute RMS
vllm::compute_rms<scalar_t, has_residual>(&rms, input, hidden_size,
var_epsilon, residual);
vllm::compute_rms<scalar_t, has_residual>(
&rms, input, hidden_size, input_stride, var_epsilon, residual);
// Compute Scale
vllm::compute_dynamic_per_token_scales<scalar_t, scalar_out_t, has_residual>(
&token_scale, scales, input, weight, rms, scale_ub, hidden_size,
residual);
input_stride, residual);
// RMS Norm + Quant
if constexpr (std::is_same_v<scalar_out_t, int8_t>) {
token_scale = 1.0f / token_scale;
vllm::norm_and_quant<scalar_t, scalar_out_t, true, has_residual>(
out, input, weight, rms, &token_scale, hidden_size, residual);
out, input, weight, rms, &token_scale, hidden_size, input_stride,
residual);
} else {
// FP8 - Do not invert s_token_scale for exact match with FBGemm
vllm::norm_and_quant<scalar_t, scalar_out_t, false, has_residual>(
out, input, weight, rms, &token_scale, hidden_size, residual);
out, input, weight, rms, &token_scale, hidden_size, input_stride,
residual);
}
}
@@ -97,19 +101,20 @@ __global__ void rms_norm_per_block_quant_kernel(
scalar_t const* __restrict__ input, // [..., hidden_size]
scalar_t const* __restrict__ weight, // [hidden_size]
float const* scale_ub, float const var_epsilon, int32_t const hidden_size,
scalar_t* __restrict__ residual = nullptr, int64_t outer_scale_stride = 1) {
int32_t const input_stride, scalar_t* __restrict__ residual = nullptr,
int64_t outer_scale_stride = 1) {
float rms;
// Compute RMS
// Always able to vectorize due to constraints on hidden_size
vllm::vectorized::compute_rms<scalar_t, has_residual>(
&rms, input, hidden_size, var_epsilon, residual);
&rms, input, hidden_size, input_stride, var_epsilon, residual);
// Compute Scale
// Always able to vectorize due to constraints on hidden_size and group_size
vllm::vectorized::compute_dynamic_per_token_scales<
scalar_t, scalar_out_t, has_residual, is_scale_transposed, group_size>(
nullptr, scales, input, weight, rms, scale_ub, hidden_size, residual,
outer_scale_stride);
nullptr, scales, input, weight, rms, scale_ub, hidden_size, input_stride,
residual, outer_scale_stride);
// RMS Norm + Quant
// Always able to vectorize due to constraints on hidden_size
@@ -120,7 +125,7 @@ __global__ void rms_norm_per_block_quant_kernel(
vllm::vectorized::norm_and_quant<
scalar_t, scalar_out_t, std::is_same_v<scalar_out_t, int8_t>,
has_residual, is_scale_transposed, group_size>(
out, input, weight, rms, scales, hidden_size, residual,
out, input, weight, rms, scales, hidden_size, input_stride, residual,
outer_scale_stride);
}
@@ -137,6 +142,7 @@ void rms_norm_dynamic_per_token_quant_dispatch(
std::optional<at::Tensor> const& scale_ub,
std::optional<at::Tensor>& residual) {
int32_t hidden_size = input.size(-1);
int32_t input_stride = input.view({-1, hidden_size}).stride(0);
auto num_tokens = input.numel() / hidden_size;
dim3 grid(num_tokens);
@@ -153,7 +159,7 @@ void rms_norm_dynamic_per_token_quant_dispatch(
out.data_ptr<scalar_t>(), scales.data_ptr<float>(),
input.data_ptr<scalar_in_t>(), weight.data_ptr<scalar_in_t>(),
scale_ub.has_value() ? scale_ub->data_ptr<float>() : nullptr,
var_epsilon, hidden_size,
var_epsilon, hidden_size, input_stride,
has_residual ? residual->data_ptr<scalar_in_t>() : nullptr);
});
});
@@ -170,7 +176,9 @@ void rms_norm_dynamic_per_token_quant(
? c10::ScalarType::Float8_e4m3fn
: c10::ScalarType::Float8_e4m3fnuz;
TORCH_CHECK(out.dtype() == kFp8Type || out.dtype() == torch::kInt8);
TORCH_CHECK(out.is_contiguous() && input.is_contiguous());
TORCH_CHECK(out.is_contiguous());
TORCH_CHECK(input.stride(-1) == 1,
"Input must be contiguous in the last dimension");
if (scale_ub.has_value()) {
TORCH_CHECK(out.dtype() == kFp8Type);
@@ -179,6 +187,7 @@ void rms_norm_dynamic_per_token_quant(
TORCH_CHECK(scales.dtype() == torch::kFloat32);
if (residual) {
TORCH_CHECK(residual->scalar_type() == input.scalar_type());
TORCH_CHECK(residual->is_contiguous());
}
VLLM_DISPATCH_FLOATING_TYPES(
@@ -200,6 +209,15 @@ void rms_norm_per_block_quant_dispatch(
std::optional<at::Tensor> const& scale_ub,
std::optional<at::Tensor>& residual, bool is_scale_transposed) {
int32_t hidden_size = input.size(-1);
int32_t input_stride = input.view({-1, hidden_size}).stride(0);
TORCH_CHECK(hidden_size % 4 == 0,
"Hidden size must be divisible by 4 for vectorized access");
TORCH_CHECK(input_stride % 4 == 0,
"Input stride must be divisible by 4 for vectorized access");
TORCH_CHECK(group_size % 4 == 0,
"Group size must be divisible by 4 for vectorized access");
auto num_tokens = input.numel() / hidden_size;
dim3 grid(num_tokens);
@@ -225,7 +243,7 @@ void rms_norm_per_block_quant_dispatch(
weight.data_ptr<scalar_in_t>(),
scale_ub.has_value() ? scale_ub->data_ptr<float>()
: nullptr,
var_epsilon, hidden_size,
var_epsilon, hidden_size, input_stride,
has_residual ? residual->data_ptr<scalar_in_t>()
: nullptr,
scales.stride(1));
@@ -246,7 +264,9 @@ void rms_norm_per_block_quant(torch::Tensor& out, torch::Tensor const& input,
? c10::ScalarType::Float8_e4m3fn
: c10::ScalarType::Float8_e4m3fnuz;
TORCH_CHECK(out.dtype() == kFp8Type || out.dtype() == torch::kInt8);
TORCH_CHECK(out.is_contiguous() && input.is_contiguous());
TORCH_CHECK(out.is_contiguous());
TORCH_CHECK(input.stride(-1) == 1,
"Input must be contiguous in the last dimension");
if (scale_ub.has_value()) {
TORCH_CHECK(out.dtype() == kFp8Type);
@@ -255,6 +275,7 @@ void rms_norm_per_block_quant(torch::Tensor& out, torch::Tensor const& input,
TORCH_CHECK(scales.dtype() == torch::kFloat32);
if (residual) {
TORCH_CHECK(residual->scalar_type() == input.scalar_type());
TORCH_CHECK(residual->is_contiguous());
}
TORCH_CHECK(group_size == 128 || group_size == 64,
csrc/quantization/fused_kernels/layernorm_utils.cuh
+44 -27
View File
@@ -16,14 +16,17 @@ namespace vllm {
// has_residual must be true, if residual is not a nullptr
template <typename scalar_t, bool has_residual = false>
__device__ void compute_rms(float* rms, scalar_t const* __restrict__ input,
int32_t const hidden_size, float const epsilon,
int32_t const hidden_size,
int32_t const input_stride, float const epsilon,
scalar_t const* __restrict__ residual = nullptr) {
int64_t const input_token_offset =
blockIdx.x * static_cast<int64_t>(input_stride);
int64_t const token_offset = blockIdx.x * static_cast<int64_t>(hidden_size);
// sum of squares
float ss = 0.0f;
for (auto i = threadIdx.x; i < hidden_size; i += blockDim.x) {
float x = static_cast<float>(input[token_offset + i]);
float x = static_cast<float>(input[input_token_offset + i]);
if constexpr (has_residual) {
x += static_cast<float>(residual[token_offset + i]);
}
@@ -73,15 +76,20 @@ __device__ void compute_dynamic_per_token_scales(
float* __restrict__ token_scale, float* __restrict__ all_token_scales,
scalar_t const* __restrict__ input, scalar_t const* __restrict__ weight,
float const rms, float const* __restrict__ scale_ub,
int32_t const hidden_size, scalar_t const* __restrict__ residual = nullptr,
int32_t const hidden_size, int32_t const input_stride,
scalar_t const* __restrict__ residual = nullptr,
int32_t const group_size = 0, int64_t outer_scale_stride = 1) {
float block_absmax_val_maybe = 0.0f;
constexpr scalar_out_t qmax{quant_type_max_v<scalar_out_t>};
__syncthreads();
int64_t const input_token_offset =
blockIdx.x * static_cast<int64_t>(input_stride);
int64_t const token_offset = blockIdx.x * static_cast<int64_t>(hidden_size);
if (group_size > 0) {
__shared__ float s_max_vals[1024];
int64_t const token_offset = blockIdx.x * static_cast<int64_t>(hidden_size);
int64_t num_groups = hidden_size / group_size;
__shared__ float s_max_vals[1024];
int64_t const threads_per_group = blockDim.x / num_groups;
int64_t const thread_in_group = threadIdx.x % threads_per_group;
int64_t const group_offset = threadIdx.x / threads_per_group * group_size;
@@ -89,7 +97,7 @@ __device__ void compute_dynamic_per_token_scales(
int64_t const thread_end =
min(group_offset + group_size, static_cast<int64_t>(hidden_size));
for (auto i = thread_offset; i < thread_end; i += threads_per_group) {
float x = static_cast<float>(input[token_offset + i]);
float x = static_cast<float>(input[input_token_offset + i]);
if constexpr (has_residual) {
x += static_cast<float>(residual[token_offset + i]);
}
@@ -144,10 +152,8 @@ __device__ void compute_dynamic_per_token_scales(
}
__syncthreads();
} else {
int64_t const token_offset = blockIdx.x * static_cast<int64_t>(hidden_size);
for (auto i = threadIdx.x; i < hidden_size; i += blockDim.x) {
float x = static_cast<float>(input[token_offset + i]);
float x = static_cast<float>(input[input_token_offset + i]);
if constexpr (has_residual) {
x += static_cast<float>(residual[token_offset + i]);
}
@@ -185,12 +191,15 @@ template <typename scalar_t, typename scalar_out_t, bool is_scale_inverted,
__device__ void norm_and_quant(
scalar_out_t* __restrict__ output, scalar_t const* __restrict__ input,
scalar_t const* __restrict__ weight, float const rms, float* const scale,
int32_t const hidden_size, scalar_t* __restrict__ residual = nullptr,
int32_t const group_size = 0, int64_t outer_scale_stride = 1) {
int32_t const hidden_size, int32_t const input_stride,
scalar_t* __restrict__ residual = nullptr, int32_t const group_size = 0,
int64_t outer_scale_stride = 1) {
int64_t const input_token_offset =
blockIdx.x * static_cast<int64_t>(input_stride);
int64_t const token_offset = blockIdx.x * static_cast<int64_t>(hidden_size);
for (auto i = threadIdx.x; i < hidden_size; i += blockDim.x) {
float x = static_cast<float>(input[token_offset + i]);
float x = static_cast<float>(input[input_token_offset + i]);
if constexpr (has_residual) {
x += static_cast<float>(residual[token_offset + i]);
residual[token_offset + i] = static_cast<scalar_t>(x);
@@ -224,13 +233,16 @@ namespace vectorized {
// hidden_size must be a multiple of 4
template <typename scalar_t, bool has_residual = false>
__device__ void compute_rms(float* rms, scalar_t const* __restrict__ input,
int32_t const hidden_size, float const epsilon,
int32_t const hidden_size,
int32_t const input_stride, float const epsilon,
scalar_t const* __restrict__ residual = nullptr) {
int64_t const input_token_offset =
blockIdx.x * static_cast<int64_t>(input_stride);
int64_t const token_offset = blockIdx.x * static_cast<int64_t>(hidden_size);
// Vectorized input/output to better utilize memory bandwidth.
vec4_t<scalar_t> const* vec_input =
reinterpret_cast<vec4_t<scalar_t> const*>(&input[token_offset]);
reinterpret_cast<vec4_t<scalar_t> const*>(&input[input_token_offset]);
vec4_t<scalar_t> const* vec_residual = nullptr;
if constexpr (has_residual) {
vec_residual =
@@ -288,7 +300,8 @@ __device__ void compute_dynamic_per_token_scales(
float* __restrict__ token_scale, float* __restrict__ all_token_scales,
scalar_t const* __restrict__ input, scalar_t const* __restrict__ weight,
float const rms, float const* __restrict__ scale_ub,
int32_t const hidden_size, scalar_t const* __restrict__ residual = nullptr,
int32_t const hidden_size, int32_t const input_stride,
scalar_t const* __restrict__ residual = nullptr,
int64_t outer_scale_stride = 1) {
constexpr scalar_out_t qmax{quant_type_max_v<scalar_out_t>};
@@ -300,10 +313,13 @@ __device__ void compute_dynamic_per_token_scales(
vec4_t<scalar_t> const* vec_weight = nullptr;
vec4_t<scalar_t> const* vec_residual = nullptr;
int64_t const input_token_offset =
blockIdx.x * static_cast<int64_t>(input_stride);
int64_t const token_offset = blockIdx.x * static_cast<int64_t>(hidden_size);
if constexpr (group_size > 0) {
__shared__ float s_max_vals[1024];
int64_t const token_offset = blockIdx.x * static_cast<int64_t>(hidden_size);
int64_t const num_groups = hidden_size / group_size;
int64_t const threads_per_group = blockDim.x / num_groups;
int64_t const thread_in_group = threadIdx.x % threads_per_group;
@@ -312,7 +328,8 @@ __device__ void compute_dynamic_per_token_scales(
int64_t const thread_offset = group_offset + thread_in_group;
int64_t const thread_end = min(group_offset + (group_size >> 2),
static_cast<int64_t>(hidden_size >> 2));
vec_input = reinterpret_cast<vec4_t<scalar_t> const*>(&input[token_offset]);
vec_input =
reinterpret_cast<vec4_t<scalar_t> const*>(&input[input_token_offset]);
vec_weight = reinterpret_cast<vec4_t<scalar_t> const*>(weight);
if constexpr (has_residual) {
vec_residual =
@@ -396,8 +413,8 @@ __device__ void compute_dynamic_per_token_scales(
__syncthreads();
} else {
int64_t const token_offset = blockIdx.x * static_cast<int64_t>(hidden_size);
vec_input = reinterpret_cast<vec4_t<scalar_t> const*>(&input[token_offset]);
vec_input =
reinterpret_cast<vec4_t<scalar_t> const*>(&input[input_token_offset]);
vec_weight = reinterpret_cast<vec4_t<scalar_t> const*>(weight);
if constexpr (has_residual) {
vec_residual =
@@ -462,18 +479,18 @@ __device__ void compute_dynamic_per_token_scales(
template <typename scalar_t, typename scalar_out_t, bool is_scale_inverted,
bool has_residual = false, bool is_scale_transposed = false,
int32_t group_size = 0>
__device__ void norm_and_quant(scalar_out_t* __restrict__ output,
scalar_t const* __restrict__ input,
scalar_t const* __restrict__ weight,
float const rms, float* const scale,
int32_t const hidden_size,
scalar_t* __restrict__ residual = nullptr,
int64_t outer_scale_stride = 1) {
__device__ void norm_and_quant(
scalar_out_t* __restrict__ output, scalar_t const* __restrict__ input,
scalar_t const* __restrict__ weight, float const rms, float* const scale,
int32_t const hidden_size, int32_t const input_stride,
scalar_t* __restrict__ residual = nullptr, int64_t outer_scale_stride = 1) {
int64_t const input_token_offset =
blockIdx.x * static_cast<int64_t>(input_stride);
int64_t const token_offset = blockIdx.x * static_cast<int64_t>(hidden_size);
// Vectorized input/output/weight/residual to better utilize memory bandwidth.
vec4_t<scalar_t> const* vec_input =
reinterpret_cast<vec4_t<scalar_t> const*>(&input[token_offset]);
reinterpret_cast<vec4_t<scalar_t> const*>(&input[input_token_offset]);
vec4_t<scalar_t> const* vec_weight =
reinterpret_cast<vec4_t<scalar_t> const*>(weight);
q8x4_t<scalar_out_t>* vec_output =
csrc/rocm/skinny_gemms.cu
+151 -85
View File
@@ -12,6 +12,7 @@
#include "../cuda_compat.h"
#include "dispatch_utils.h"
#include "quantization/w8a8/fp8/common.cuh"
#include "core/batch_invariant.hpp"
// TODO(rasmith): The kernels in this file are susceptible to integer overflow
// issues, do not take strides, and are unable to handle PyTorch tensors that
@@ -1224,17 +1225,14 @@ torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b,
#if defined(__gfx950__)
#define WVSPLITKRC_1KPASS
template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
int UNRL, int N, int GrpsShrB, int CHUNKK>
int UNRL, int N, int GrpsShrB, int CHUNKK, int DTRMNSTC>
__global__ void __launch_bounds__(WvPrGrp* THRDS)
__attribute__((amdgpu_waves_per_eu(1, 1)))
wvSplitKrc_(const int actlN, const int K, const int M, const int Bx,
const int By, const scalar_t* __restrict__ B,
const scalar_t* __restrict__ A,
const scalar_t* __restrict__ BIAS, float* glbl, scalar_t* C,
const int CuCount) {
// Use upper half of glbl buffer for atomic reduce counting
int* cntr = (int*)(&glbl[M * N]);
wvSplitKrc_(const int actlN, const int K, const int Kap, const int M,
const int Bx, const int By, const scalar_t* __restrict__ A,
const scalar_t* __restrict__ B,
const scalar_t* __restrict__ BIAS, float* glbl, int* cntr,
scalar_t* C, const int CuCount) {
constexpr int NTILE = 16;
constexpr int APAD = 1;
constexpr int ASTRD = 64;
@@ -1425,11 +1423,11 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
unsigned int kOffcp = min__(K - A_CHUNK, k_str + kOff);
for (unsigned int n = 0; n < N; n += CHUNKK * sprdN) {
__builtin_amdgcn_global_load_lds(
(int*)(&A[min__(
K * actlN - A_CHUNK,
kOffcp + K * (n / CHUNKK +
(N / CHUNKK) * (threadIdx.x / (64 / CHUNKK)) +
(threadIdx.y % sprdN)))]),
(int*)(&A[min__(Kap * actlN - A_CHUNK,
kOffcp + Kap * (n / CHUNKK +
(N / CHUNKK) * (threadIdx.x /
(64 / CHUNKK)) +
(threadIdx.y % sprdN)))]),
(int*)(&s[(k +
kFitPdd * ((n / CHUNKK) + (threadIdx.y % sprdN)))]),
16, 0, 0);
@@ -1533,45 +1531,98 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
}
}
union flt4 {
scalar8 s8;
float2 f2[2];
float4 f4;
};
if (m + (threadIdx.x % 16) < M) {
int my_cntr;
int mindx = m + (threadIdx.x % 16);
int g_mindx = m * 4 + (threadIdx.x % 64); // coalesced atomic reduction
scalar_t biases[N / NTILE / GrpsShrB][4] = {};
// Atomic add the output, read biases
for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++)
for (uint32_t j = 0; j < 4; j++) {
// int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE +
// (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
// int adr = mindx + M * nindx;
int g_nindx =
j + (nt * NTILE + (N / GrpsShrB) * (threadIdx.y % GrpsShrB)) / 4;
int g_adr = g_mindx + M * g_nindx * 4;
atomicAdd(&glbl[g_adr], sum4[nt][0][j]);
for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
int g_nindx =
(nt * NTILE + (N / GrpsShrB) * (threadIdx.y % GrpsShrB)) / 4;
int g_adr = g_mindx * 4 + 0 + M * g_nindx * 4;
if (DTRMNSTC) {
flt4 flt4_ = {.s8 = sum4[nt][0]};
__hip_atomic_store((float2*)&glbl[g_adr + M * N * (m0 / Mmod)],
flt4_.f2[0], __ATOMIC_RELAXED,
__HIP_MEMORY_SCOPE_AGENT);
__hip_atomic_store((float2*)&glbl[g_adr + 2 + M * N * (m0 / Mmod)],
flt4_.f2[1], __ATOMIC_RELAXED,
__HIP_MEMORY_SCOPE_AGENT);
} else {
for (uint32_t j = 0; j < 4; j++)
atomicAdd((&glbl[g_adr + j]), sum4[nt][0][j]);
}
}
__atomic_signal_fence(__ATOMIC_SEQ_CST);
asm volatile("s_waitcnt vmcnt(0)" ::: "memory");
__atomic_signal_fence(__ATOMIC_SEQ_CST);
int nindx_ = (0 + (threadIdx.x / 16) * 4) + 0 * NTILE +
(N / GrpsShrB) * (threadIdx.y % GrpsShrB);
int adr_ = mindx + M * nindx_ / 4;
// Update the complete counter
my_cntr = atomicAdd(&cntr[adr_], 1);
float vals[N / NTILE / GrpsShrB][4] = {};
// make sure LDS is free for write out staging
if (DTRMNSTC) __syncthreads();
// Update the complete counter
flt4 vals[N / NTILE / GrpsShrB] = {};
// If we're the last k-shard, read back the value and convert...
if (my_cntr + 1 == k_rnd) {
if (BIAS)
for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
for (uint32_t j = 0; j < 4; j++) {
int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE +
(N / GrpsShrB) * (threadIdx.y % GrpsShrB);
biases[nt][j] = BIAS[(mindx % Bx) + (nindx % By) * Bx];
cntr[adr_] = 0; // clear for next round
if constexpr (DTRMNSTC) {
#pragma unroll
for (int ks = 0; ks < k_rnd; ks++) {
for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
int g_nindx =
(nt * NTILE + (N / GrpsShrB) * (threadIdx.y % GrpsShrB)) / 4;
int g_adr = g_mindx * 4 + 0 + M * g_nindx * 4;
__builtin_amdgcn_global_load_lds(
(float4*)(&glbl[g_adr + M * N * ks]),
&(((float4*)s)[(threadIdx.y * THRDS) + ks * THRDS * 4 +
nt * THRDS * 4 * k_rnd]),
16, 0, 0);
}
}
for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
for (uint32_t j = 0; j < 4; j++) {
int g_nindx =
j + (nt * NTILE + (N / GrpsShrB) * (threadIdx.y % GrpsShrB)) / 4;
int g_adr = g_mindx + M * g_nindx * 4;
vals[nt][j] = glbl[g_adr];
if (BIAS)
for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
for (uint32_t j = 0; j < 4; j++) {
int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE +
(N / GrpsShrB) * (threadIdx.y % GrpsShrB);
biases[nt][j] = BIAS[(mindx % Bx) + (nindx % By) * Bx];
}
}
asm volatile("s_waitcnt 0");
for (int ks = 0; ks < k_rnd; ks++) {
for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
float4 eval = ((float4*)s)[(threadIdx.x + threadIdx.y * THRDS) +
ks * THRDS * 4 + nt * THRDS * 4 * k_rnd];
vals[nt].f4 += eval;
}
}
} else {
for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
int g_nindx =
(nt * NTILE + (N / GrpsShrB) * (threadIdx.y % GrpsShrB)) / 4;
int g_adr = g_mindx * 4 + 0 + M * g_nindx * 4;
vals[nt].f4 = *(float4*)(&glbl[g_adr]);
*(float4*)(&glbl[g_adr]) = {}; // clear out for next round
}
if (BIAS)
for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
for (uint32_t j = 0; j < 4; j++) {
int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE +
(N / GrpsShrB) * (threadIdx.y % GrpsShrB);
biases[nt][j] = BIAS[(mindx % Bx) + (nindx % By) * Bx];
}
}
}
__builtin_amdgcn_sched_barrier(0);
for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
@@ -1581,11 +1632,11 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
if (nindx < actlN) {
int adr = mindx + M * nindx;
if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) {
vals[nt][j] += __bfloat162float(biases[nt][j]);
C[adr] = __float2bfloat16(vals[nt][j]);
vals[nt].s8[j] += __bfloat162float(biases[nt][j]);
C[adr] = __float2bfloat16(vals[nt].s8[j]);
} else {
vals[nt][j] += __half2float(biases[nt][j]);
C[adr] = __float2half(vals[nt][j]);
vals[nt].s8[j] += __half2float(biases[nt][j]);
C[adr] = __float2half(vals[nt].s8[j]);
}
}
}
@@ -1604,21 +1655,25 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
}
#else // !defined(__HIP__GFX9__) TODO: Add NAVI support
template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
int UNRL, int N, int GrpsShrB, int CHUNKK>
__global__ void wvSplitKrc_(const int actlN, const int K, const int M,
const int Bx, const int By, const scalar_t* B,
const scalar_t* __restrict__ A,
int UNRL, int N, int GrpsShrB, int CHUNKK, int DTRMNSTC>
__global__ void wvSplitKrc_(const int actlN, const int K, const int Kap,
const int M, const int Bx, const int By,
const scalar_t* B, const scalar_t* __restrict__ A,
const scalar_t* __restrict__ BIAS, float* glbl,
// int* cntr,
scalar_t* C, const int CuCount){UNREACHABLE_CODE}
int* cntr, scalar_t* C,
const int CuCount){UNREACHABLE_CODE}
#endif // defined(__HIP__GFX9__) TODO: Add NAVI support
torch::Tensor wvSplitKrc(const at::Tensor& in_a, const at::Tensor& in_b,
const std::optional<at::Tensor>& in_bias,
const int64_t CuCount) {
auto M_in = in_a.size(0);
auto N_in = in_b.size(0);
auto K_in = in_a.size(1);
int _DTRMNSTC = 1; // vllm::vllm_is_batch_invariant();
auto M_in = in_b.size(0);
auto N_in = in_a.size(0);
auto K_in = in_b.size(1);
auto Kap_in = in_a.stride(0);
auto Bx_in =
(in_bias.has_value() && in_bias->numel() > 0)
? (in_bias->sizes().size() == 2) ? in_bias->size(1) : in_bias->size(0)
@@ -1635,13 +1690,9 @@ torch::Tensor wvSplitKrc(const at::Tensor& in_a, const at::Tensor& in_b,
auto out_c = torch::empty(
{N_in, M_in},
torch::TensorOptions().dtype(in_b.dtype()).device(in_b.device()));
torch::TensorOptions().dtype(in_a.dtype()).device(in_a.device()));
auto N_p2 = 1U << (32 - __builtin_clz(N_in - 1));
auto axl_glbl = torch::empty(
{N_p2 + N_p2 / 4, M_in + M_in / 4},
torch::TensorOptions().dtype(torch::kFloat32).device(in_b.device()));
axl_glbl.zero_(); // disable for FAST_UNSAFE_RDC_INIT
dim3 grid(CuCount);
@@ -1649,55 +1700,70 @@ torch::Tensor wvSplitKrc(const at::Tensor& in_a, const at::Tensor& in_b,
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
// const int max_lds_len = get_lds_size() / 2;
// With 64 Ms per CU (each of 4 SIMDs working on a 16x16 tile),
// and each working on a 512-shard of K, how many CUs would we need?
int rndup_cus = ((M_in + 64 - 1) / 64) * ((K_in + 512 - 1) / 512);
// How many of 4 waves in a group can work on same 16 Ms at same time? First
// try to maximize this. This reduces the Ms each group works on, i.e.
// increasing the number of CUs needed.
int GrpsShrB = min(N_p2 / 16, 4);
// Given the above, how many CUs would we need?
int CuNeeded = rndup_cus * GrpsShrB;
if (CuNeeded > CuCount) throw std::runtime_error("Invalid wvSplitKrc size");
// Can we increase SplitK by shrinking the K-shared to 256?
int chunkk = (CuNeeded * 2 <= CuCount) ? 2 : 1;
static torch::Tensor axl_glbl =
torch::zeros(
128 * 1024 * (_DTRMNSTC ? 12 : 1),
torch::TensorOptions().dtype(torch::kFloat32).device(in_a.device()))
.detach();
static torch::Tensor axl_cntr =
torch::zeros(
128 * 1024 * (_DTRMNSTC ? 12 : 1) / 4,
torch::TensorOptions().dtype(torch::kInt).device(in_a.device()))
.detach();
auto glbl = axl_glbl.data_ptr<float>();
auto cntr = axl_cntr.data_ptr<int>();
#define WVSPLITKrc(_N, _GrpsShrB, _CHUNKK) \
{ \
dim3 block(64, 4); \
wvSplitKrc_<fptype, 64, 16, 4, 8, 1, _N, _GrpsShrB, _CHUNKK> \
<<<grid, block, 0, stream>>>(N_in, K_in, M_in, Bx_in, By_in, af4, bf4, \
biasf4, glbl, c, CuCount); \
if (_DTRMNSTC) \
wvSplitKrc_<fptype, 64, 16, 4, 8, 1, _N, _GrpsShrB, _CHUNKK, 1> \
<<<grid, block, 0, stream>>>(N_in, K_in, Kap_in, M_in, Bx_in, By_in, \
af4, bf4, biasf4, glbl, cntr, c, \
CuCount); \
else \
wvSplitKrc_<fptype, 64, 16, 4, 8, 1, _N, _GrpsShrB, _CHUNKK, 0> \
<<<grid, block, 0, stream>>>(N_in, K_in, Kap_in, M_in, Bx_in, By_in, \
af4, bf4, biasf4, glbl, cntr, c, \
CuCount); \
}
AT_DISPATCH_REDUCED_FLOATING_TYPES(in_b.scalar_type(), "wvSplitKrc", [&] {
AT_DISPATCH_REDUCED_FLOATING_TYPES(in_a.scalar_type(), "wvSplitKrc", [&] {
using fptype = typename scalar<scalar_t>::type;
fptype* af4 = reinterpret_cast<fptype*>(in_a.data_ptr());
const fptype* af4 = reinterpret_cast<const fptype*>(in_a.data_ptr());
const fptype* bf4 = reinterpret_cast<const fptype*>(in_b.data_ptr());
const fptype* biasf4 =
(in_bias.has_value() && in_bias->numel() > 0)
? reinterpret_cast<const fptype*>(in_bias->data_ptr())
: nullptr;
fptype* c = reinterpret_cast<fptype*>(out_c.data_ptr());
auto glbl = axl_glbl.data_ptr<float>();
// With 64 Ms per CU (each of 4 SIMDs working on a 16x16 tile),
// and each working on a 512-shard of K, how many CUs would we need?
int rndup_cus = ((M_in + 64 - 1) / 64) * ((K_in + 512 - 1) / 512);
// How many of 4 waves in a group can work on same 16 Ms at same time? First
// try to maximize this. This reduces the Ms each group works on, i.e.
// increasing the number of CUs needed.
int GrpsShrB = min(N_p2 / 16, 4);
// Given the above, how many CUs would we need?
int CuNeeded = rndup_cus * GrpsShrB;
if (CuNeeded > CuCount) std::runtime_error("Invalid wvSplitKrc size");
// Can we increase SplitK by shrinking the K-shared to 256?
int chunkk = (CuNeeded * 2 <= CuCount) ? 2 : 1;
switch (N_p2) {
case 16:
WVSPLITKrc(16, 1, 1) break;
case 32:
if (chunkk == 2)
WVSPLITKrc(32, 2, 2) else if (chunkk == 1) WVSPLITKrc(32, 2, 1) break;
if (chunkk == 2) WVSPLITKrc(32, 2, 2) else WVSPLITKrc(32, 2, 1) break;
case 64:
if (chunkk == 2)
WVSPLITKrc(64, 4, 2) else if (chunkk == 1) WVSPLITKrc(64, 4, 1) break;
if (chunkk == 2) WVSPLITKrc(64, 4, 2) else WVSPLITKrc(64, 4, 1) break;
case 128:
if (chunkk == 2)
WVSPLITKrc(128, 4, 2) else if (chunkk == 1)
WVSPLITKrc(128, 4, 1) break;
if (chunkk == 2) WVSPLITKrc(128, 4, 2) else WVSPLITKrc(128, 4, 1) break;
default:
throw std::runtime_error(
"Unsupported N value: " + std::to_string(M_in) + "," +
docs/benchmarking/dashboard.md
+6
View File
@@ -39,6 +39,12 @@ When run, benchmark script generates results under **benchmark/results** folder,
- `THROUGHPUT_JSON`: JSON file to use for the throughout tests. Default value is empty string (use default file).
- `REMOTE_HOST`: IP for the remote vLLM service to benchmark. Default value is empty string.
- `REMOTE_PORT`: Port for the remote vLLM service to benchmark. Default value is empty string.
- `PROMPTS_PER_CONCURRENCY`: Multiplier to compute `num_prompts` for serving tests (`num_prompts = max_concurrency × value`). Overrides JSON `num_prompts`. Default is NULL.
- `ENABLE_ADAPTIVE_CONCURRENCY`: set the value to '1' to enable adaptive SLA-based concurrency search after the static serving max_concurrency sweep. Default value is 0.
- `SLA_TTFT_MS`: default TTFT SLA threshold in milliseconds for adaptive concurrency search. Default value is 3000.
- `SLA_TPOT_MS`: default TPOT SLA threshold in milliseconds for adaptive concurrency search. Default value is 100.
- `ADAPTIVE_MAX_PROBES`: maximum number of extra adaptive search probes. Default value is 8.
- `ADAPTIVE_MAX_CONCURRENCY`: maximum allowed concurrency during adaptive search. Default value is 1024.
### Visualization
docs/design/attention_backends.md
+2 -1
View File
@@ -173,7 +173,7 @@ Priority is **1 = highest** (tried first).
| `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 |
| `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, 544 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ✅ | ✅ | ❌ | 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 |
@@ -214,3 +214,4 @@ configuration.
| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `bfloat16` | 1 | Any | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
| `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
| `TRITON_MLA` | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
| `XPU_MLA_SPARSE` | fp16, bf16 | `auto`, `bfloat16` | Any | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | Any |
docs/features/disagg_prefill.md
+6
View File
@@ -44,6 +44,12 @@ For NixlConnector, you may also specify one or multiple NIXL_Backend. Such as:
--kv-transfer-config '{"kv_connector":"OffloadingConnector","kv_role":"kv_both","kv_connector_extra_config":{"block_size": 64, "cpu_bytes_to_use": 1000000000}}'
```
- **FlexKVConnectorV1**: refer to [examples/offline_inference/prefix_caching_flexkv.py](../../examples/offline_inference/prefix_caching_flexkv.py) for the example usage of FlexKVConnectorV1. FlexKV is a distributed KV Store and multi-level cache management system for ultra-large-scale LLM inference.
```bash
--kv-transfer-config '{"kv_connector":"FlexKVConnectorV1","kv_role":"kv_both"}'
```
## Benchmarks
Please refer to [benchmarks/disagg_benchmarks](../../benchmarks/disagg_benchmarks) for disaggregated prefilling benchmarks.
docs/getting_started/installation/README.md
+1 -1
View File
@@ -16,4 +16,4 @@ 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/#hardware). 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.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).
docs/getting_started/installation/gpu.xpu.inc.md
+2 -3
View File
@@ -7,7 +7,6 @@ vLLM initially supports basic model inference and serving on Intel GPU platform.
--8<-- [start:requirements]
- Supported Hardware: Intel Data Center GPU, Intel ARC GPU
- OneAPI requirements: oneAPI 2025.3
- Dependency: [vllm-xpu-kernels](https://github.com/vllm-project/vllm-xpu-kernels): a package provide all necessary vllm custom kernel when running vLLM on Intel GPU platform,
- Python: 3.12
!!! warning
@@ -26,8 +25,8 @@ Currently, there are no pre-built XPU wheels.
--8<-- [end:pre-built-wheels]
--8<-- [start:build-wheel-from-source]
- First, install required [driver](https://dgpu-docs.intel.com/driver/installation.html#installing-gpu-drivers) and [Intel OneAPI](https://www.intel.com/content/www/us/en/developer/tools/oneapi/base-toolkit.html) 2025.3 or later.
- Second, install Python packages for vLLM XPU backend building:
- First, install required [driver](https://dgpu-docs.intel.com/driver/installation.html#installing-gpu-drivers).
- Second, install Python packages for vLLM XPU backend building (Intel OneAPI dependencies are installed automatically as part of `torch-xpu`, see [PyTorch XPU get started](https://docs.pytorch.org/docs/stable/notes/get_start_xpu.html)):
```bash
git clone https://github.com/vllm-project/vllm.git
docs/models/supported_models.md
+3 -1
View File
@@ -701,6 +701,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
| `GlmOcrForConditionalGeneration` | GLM-OCR | T + I<sup>E+</sup> | `zai-org/GLM-OCR`, etc. | ✅︎ | ✅︎ |
| `GraniteSpeechForConditionalGeneration` | Granite Speech | T + A | `ibm-granite/granite-speech-3.3-8b` | ✅︎ | ✅︎ |
| `HCXVisionForCausalLM` | HyperCLOVAX-SEED-Vision-Instruct-3B | T + I<sup>+</sup> + V<sup>+</sup> | `naver-hyperclovax/HyperCLOVAX-SEED-Vision-Instruct-3B` | | |
| `HCXVisionV2ForCausalLM` | HyperCLOVAX-SEED-Think-32B | T + I<sup>+</sup> + V<sup>+</sup> | `naver-hyperclovax/HyperCLOVAX-SEED-Think-32B` | | |
| `H2OVLChatModel` | H2OVL | T + I<sup>E+</sup> | `h2oai/h2ovl-mississippi-800m`, `h2oai/h2ovl-mississippi-2b`, etc. | | ✅︎ |
| `HunYuanVLForConditionalGeneration` | HunyuanOCR | T + I<sup>E+</sup> | `tencent/HunyuanOCR`, etc. | ✅︎ | ✅︎ |
| `Idefics3ForConditionalGeneration` | Idefics3 | T + I | `HuggingFaceM4/Idefics3-8B-Llama3`, etc. | ✅︎ | |
@@ -712,8 +713,9 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
| `KananaVForConditionalGeneration` | Kanana-V | T + I<sup>+</sup> | `kakaocorp/kanana-1.5-v-3b-instruct`, etc. | | ✅︎ |
| `KeyeForConditionalGeneration` | Keye-VL-8B-Preview | T + I<sup>E+</sup> + V<sup>E+</sup> | `Kwai-Keye/Keye-VL-8B-Preview` | ✅︎ | ✅︎ |
| `KeyeVL1_5ForConditionalGeneration` | Keye-VL-1_5-8B | T + I<sup>E+</sup> + V<sup>E+</sup> | `Kwai-Keye/Keye-VL-1_5-8B` | ✅︎ | ✅︎ |
| `KimiVLForConditionalGeneration` | Kimi-VL-A3B-Instruct, Kimi-VL-A3B-Thinking | T + I<sup>+</sup> | `moonshotai/Kimi-VL-A3B-Instruct`, `moonshotai/Kimi-VL-A3B-Thinking` | | ✅︎ |
| `KimiAudioForConditionalGeneration` | Kimi-Audio | T + A<sup>+</sup> | `moonshotai/Kimi-Audio-7B-Instruct` | | ✅︎ |
| `KimiK25ForConditionalGeneration` | Kimi-K2.5 | T + I<sup>+</sup> | `moonshotai/Kimi-K2.5` | | ✅︎ |
| `KimiVLForConditionalGeneration` | Kimi-VL-A3B-Instruct, Kimi-VL-A3B-Thinking | T + I<sup>+</sup> | `moonshotai/Kimi-VL-A3B-Instruct`, `moonshotai/Kimi-VL-A3B-Thinking` | | ✅︎ |
| `LightOnOCRForConditionalGeneration` | LightOnOCR-1B | T + I<sup>+</sup> | `lightonai/LightOnOCR-1B`, etc | ✅︎ | ✅︎ |
| `Lfm2VlForConditionalGeneration` | LFM2-VL | T + I<sup>+</sup> | `LiquidAI/LFM2-VL-450M`, `LiquidAI/LFM2-VL-3B`, `LiquidAI/LFM2-VL-8B-A1B`, etc. | ✅︎ | ✅︎ |
| `Llama4ForConditionalGeneration` | Llama 4 | T + I<sup>+</sup> | `meta-llama/Llama-4-Scout-17B-16E-Instruct`, `meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8`, `meta-llama/Llama-4-Maverick-17B-128E-Instruct`, etc. | ✅︎ | ✅︎ |
docs/serving/integrations/claude_code.md
+3
View File
@@ -60,6 +60,9 @@ The environment variables:
!!! tip
You can add these environment variables to your shell profile (e.g., `.bashrc`, `.zshrc`), Claude Code configuration file (`~/.claude/settings.json`), or create a wrapper script for convenience.
!!! warning
Claude Code recently started injecting a per-request hash in the system prompt, which can defeat [prefix caching](../../design/prefix_caching.md) because the prompt changes on every request, causing greatly reduced performance. This is addressed automatically in vLLM versions > 0.17.1 but for older versions `"CLAUDE_CODE_ATTRIBUTION_HEADER": "0"` should be added to the `"env"` section of `~/.claude/settings.json` (see this [blog post](https://unsloth.ai/docs/basics/claude-code#fixing-90-slower-inference-in-claude-code) from Unsloth).
## Testing the Setup
Once Claude Code launches, try a simple prompt to verify the connection:
examples/offline_inference/audio_language.py
+29
View File
@@ -201,6 +201,34 @@ def run_granite_speech(question: str, audio_count: int) -> ModelRequestData:
)
# Kimi-Audio-7B-Instruct
def run_kimi_audio(question: str, audio_count: int) -> ModelRequestData:
"""Kimi-Audio-7B-Instruct for audio transcription and understanding."""
model_name = "moonshotai/Kimi-Audio-7B-Instruct"
engine_args = EngineArgs(
model=model_name,
trust_remote_code=True,
max_model_len=4096,
max_num_seqs=2,
limit_mm_per_prompt={"audio": audio_count},
)
# Kimi-Audio uses <|im_kimia_text_blank|> as placeholder for audio features
audio_placeholder = "<|im_kimia_text_blank|>" * audio_count
# Default prompt for transcription
if not question:
question = "Please transcribe the audio"
prompt = f"{audio_placeholder}{question}"
# Stop at EOS token (151644) to prevent repetition
return ModelRequestData(
engine_args=engine_args,
prompt=prompt,
stop_token_ids=[151644],
)
# MiDashengLM
def run_midashenglm(question: str, audio_count: int):
model_name = "mispeech/midashenglm-7b"
@@ -485,6 +513,7 @@ model_example_map = {
"glmasr": run_glmasr,
"funaudiochat": run_funaudiochat,
"granite_speech": run_granite_speech,
"kimi_audio": run_kimi_audio,
"midashenglm": run_midashenglm,
"minicpmo": run_minicpmo,
"phi4_mm": run_phi4mm,
examples/offline_inference/mistral-small.py
+6 -6
View File
@@ -62,9 +62,9 @@ def run_simple_demo(args: argparse.Namespace):
llm = LLM(
model=model_name,
tokenizer_mode="mistral" if args.format == "mistral" else "auto",
config_format="mistral" if args.format == "mistral" else "auto",
load_format="mistral" if args.format == "mistral" else "auto",
tokenizer_mode="mistral" if args.format == "mistral" else "hf",
config_format="mistral" if args.format == "mistral" else "hf",
load_format="mistral" if args.format == "mistral" else "hf",
limit_mm_per_prompt={"image": 1},
max_model_len=4096,
max_num_seqs=2,
@@ -102,9 +102,9 @@ def run_advanced_demo(args: argparse.Namespace):
sampling_params = SamplingParams(max_tokens=8192, temperature=0.7)
llm = LLM(
model=model_name,
tokenizer_mode="mistral" if args.format == "mistral" else "auto",
config_format="mistral" if args.format == "mistral" else "auto",
load_format="mistral" if args.format == "mistral" else "auto",
tokenizer_mode="mistral" if args.format == "mistral" else "hf",
config_format="mistral" if args.format == "mistral" else "hf",
load_format="mistral" if args.format == "mistral" else "hf",
limit_mm_per_prompt={"image": max_img_per_msg},
max_model_len=max_img_per_msg * max_tokens_per_img,
tensor_parallel_size=2,
examples/offline_inference/prefix_caching_flexkv.py
+221
View File
@@ -0,0 +1,221 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
This example shows how to use FlexKV with vLLM for prefix caching.
FlexKV is a distributed KV Store and multi-level cache management system for
ultra-large-scale LLM inference.
Requirements:
- Install FlexKV (https://github.com/taco-project/FlexKV):
1. git clone git@github.com:taco-project/FlexKV.git
2. cd FlexKV && bash build.sh
- Ensure FlexKV is compatible with your vLLM version.
Usage:
1. Run this script:
python examples/offline_inference/prefix_caching_flexkv.py \
--model /path/to/your/model
2. Arguments:
--model Path or name of the model (required)
--tp-size Tensor parallel size (default: 1)
--gpu-memory-util GPU memory utilization (default: 0.4)
3. The script will:
- Create a FlexKV configuration file.
- Set the FLEXKV_CONFIG_PATH environment variable.
- Run vLLM with FlexKVConnectorV1 enabled.
- Compare results between regular execution, vLLM's default prefix
caching, and FlexKV.
"""
import argparse
import json
import os
import time
from vllm import LLM, SamplingParams
from vllm.distributed import cleanup_dist_env_and_memory
# NOTE: This is just a running example. For benchmarking purpose,
# please see benchmarks/benchmark_prefix_caching.py
def parse_args():
parser = argparse.ArgumentParser(
description="Example of using FlexKV with vLLM for prefix caching."
)
parser.add_argument(
"--model",
type=str,
required=True,
help="Path or name of the model to use.",
)
parser.add_argument(
"--tp-size",
type=int,
default=1,
help="Tensor parallel size (default: 1).",
)
parser.add_argument(
"--gpu-memory-util",
type=float,
default=0.4,
help="GPU memory utilization fraction (default: 0.4).",
)
return parser.parse_args()
def main():
args = parse_args()
flexkv_config = {
"server_recv_port": f"ipc:///tmp/flexkv_test_{os.getpid()}",
"cache_config": {
"enable_cpu": True,
"num_cpu_blocks": 10240,
},
"num_log_interval_requests": 200,
}
flexkv_config_path = f"./flexkv_config_{os.getpid()}.json"
with open(flexkv_config_path, "w") as f:
json.dump(flexkv_config, f)
os.environ["FLEXKV_CONFIG_PATH"] = flexkv_config_path
try:
_run(args)
finally:
if os.path.exists(flexkv_config_path):
os.remove(flexkv_config_path)
def _run(args):
# Common prefix.
prefix = (
"You are an expert school principal, skilled in effectively managing "
"faculty and staff. Draft 10-15 questions for a potential first grade "
"Head Teacher for my K-12, all-girls', independent school that emphasizes "
"community, joyful discovery, and life-long learning. The candidate is "
"coming in for a first-round panel interview for a 8th grade Math "
"teaching role. They have 5 years of previous teaching experience "
"as an assistant teacher at a co-ed, public school with experience "
"in middle school math teaching. Based on these information, fulfill "
"the following paragraph: "
)
# Sample prompts.
prompts = [
"Hello, my name is",
"The president of the United States is",
"The capital of France is",
"The future of AI is",
]
generating_prompts = [prefix + prompt for prompt in prompts]
# Create a sampling params object.
sampling_params = SamplingParams(temperature=0.0)
kv_transfer_config = {
"kv_connector": "FlexKVConnectorV1",
"kv_role": "kv_both",
}
# Create an LLM without prefix caching as a baseline.
regular_llm = LLM(
model=args.model,
enable_prefix_caching=False,
gpu_memory_utilization=args.gpu_memory_util,
tensor_parallel_size=args.tp_size,
)
print("Results without `enable_prefix_caching`")
# ruff: noqa: E501
# Generate texts from the prompts. The output is a list of RequestOutput
# objects that contain the prompt, generated text, and other information.
outputs = regular_llm.generate(generating_prompts, sampling_params)
regular_generated_texts = []
# Print the outputs.
print("-" * 50)
for output in outputs:
prompt = output.prompt
generated_text = output.outputs[0].text
regular_generated_texts.append(generated_text)
print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
print("-" * 50)
# Destroy the LLM object and free up the GPU memory.
del regular_llm
cleanup_dist_env_and_memory()
# Create an LLM with prefix caching enabled.
prefix_cached_llm = LLM(
model=args.model,
enable_prefix_caching=True,
gpu_memory_utilization=args.gpu_memory_util,
tensor_parallel_size=args.tp_size,
kv_transfer_config=kv_transfer_config,
)
# Warmup so that the shared prompt's KV cache is computed.
prefix_cached_llm.generate(generating_prompts[0], sampling_params)
# wait for offload kv task finished.
time.sleep(2)
# Generate with prefix caching.
outputs = prefix_cached_llm.generate(generating_prompts, sampling_params)
print("Results with `enable_prefix_caching`")
cached_generated_texts = []
# Print the outputs. You should see the same outputs as before.
print("-" * 50)
for output in outputs:
prompt = output.prompt
generated_text = output.outputs[0].text
cached_generated_texts.append(generated_text)
print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
print("-" * 50)
# Compare the results and display the speedup
generated_same = all(
regular_generated_texts[i] == cached_generated_texts[i]
for i in range(len(prompts))
)
print(f"Generated answers are the same: {generated_same}")
# wait for offload kv task finished.
time.sleep(2)
# reset prefix cache to use flexkv
prefix_cached_llm.reset_prefix_cache()
# Generate with prefix caching.
outputs = prefix_cached_llm.generate(generating_prompts, sampling_params)
print("Results with `flexkv`")
flexkv_generated_texts = []
# Print the outputs. You should see the same outputs as before.
print("-" * 50)
for output in outputs:
prompt = output.prompt
generated_text = output.outputs[0].text
flexkv_generated_texts.append(generated_text)
print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
print("-" * 50)
# Compare the results and display the speedup
generated_same = all(
regular_generated_texts[i] == flexkv_generated_texts[i]
for i in range(len(prompts))
)
print(f"Generated answers are the same: {generated_same}")
if __name__ == "__main__":
main()
examples/offline_inference/routed_experts_e2e.py
+384
View File
@@ -0,0 +1,384 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
End-to-end example for routed experts capture with hybrid models.
Validates that:
1. routed_experts is returned in CompletionOutput for MoE models.
2. Expert IDs are within valid range.
3. Results are deterministic across runs (baseline vs reference).
Usage:
python examples/offline_inference/routed_experts_e2e.py \
--model Qwen/Qwen3-30B-A3B \
--tp 4 \
--max-model-len 4096 \
--num-prompts 20 \
--max-new-tokens 50
"""
from __future__ import annotations
import argparse
import asyncio
import logging
import os
import uuid
from dataclasses import dataclass, field
import numpy as np
from vllm.engine.arg_utils import AsyncEngineArgs
logger = logging.getLogger(__name__)
DEFAULT_MODEL = "Qwen/Qwen3-30B-A3B"
TEST_PROMPTS = [
"Hello, my name is",
"The capital of France is",
"Explain quantum computing in simple terms:",
"Write a Python function that sorts a list:",
"The meaning of life is",
"In a distant galaxy, there was a",
"The best way to learn programming is",
"Once upon a time in a land far away,",
"The theory of relativity states that",
"How does photosynthesis work?",
"Describe the process of machine learning:",
"What are the benefits of exercise?",
"The history of artificial intelligence began",
"Translate the following to French: Hello world",
"Summarize the plot of Romeo and Juliet:",
"What is the difference between TCP and UDP?",
"The water cycle consists of",
"Explain how a neural network learns:",
"The periodic table organizes elements by",
"Write a haiku about the ocean:",
]
@dataclass
class InferenceResult:
"""Result from a single inference run."""
experts_list: list[np.ndarray] = field(default_factory=list)
token_ids_list: list[list[int]] = field(default_factory=list)
num_experts: int = 0
# ---------------------------------------------------------------------------
# Inference helpers
# ---------------------------------------------------------------------------
async def _run_async_inference(
engine_args: AsyncEngineArgs,
prompts: list[str],
max_new_tokens: int,
) -> InferenceResult:
"""Run inference using AsyncLLM."""
from vllm.sampling_params import SamplingParams
from vllm.v1.engine.async_llm import AsyncLLM
engine = AsyncLLM.from_engine_args(engine_args)
hf_config = engine.model_config.hf_text_config
num_experts: int = getattr(hf_config, "num_experts", 0) or getattr(
hf_config, "num_local_experts", 0
)
assert num_experts > 0, "Could not determine num_experts from model config"
sampling_params = SamplingParams(
temperature=0,
max_tokens=max_new_tokens,
)
async def _generate_one(prompt: str, idx: int):
request_id = str(uuid.uuid4())
final_output = None
async for output in engine.generate(prompt, sampling_params, request_id):
final_output = output
assert final_output is not None
completion = final_output.outputs[0]
routed = completion.routed_experts
num_prompt_tokens = len(final_output.prompt_token_ids)
num_generated_tokens = len(completion.token_ids)
expected_len = num_prompt_tokens + num_generated_tokens - 1
assert routed is not None, f"Prompt {idx}: routed_experts is None"
assert routed.shape[0] == expected_len, (
f"Prompt {idx}: routed_experts length {routed.shape[0]} != "
f"prompt ({num_prompt_tokens}) + generated ({num_generated_tokens})"
f" - 1 = {expected_len}"
)
return idx, routed, list(completion.token_ids)
tasks = [_generate_one(p, i) for i, p in enumerate(prompts)]
outputs = await asyncio.gather(*tasks)
# Sort by original index to maintain prompt order
outputs.sort(key=lambda x: x[0])
result = InferenceResult(num_experts=num_experts)
for _, routed, token_ids in outputs:
result.experts_list.append(routed)
result.token_ids_list.append(token_ids)
engine.shutdown()
return result
def run_inference(
model: str,
prompts: list[str],
max_new_tokens: int = 50,
tp: int = 1,
max_model_len: int = 4096,
) -> InferenceResult:
"""Run inference with routed experts capture enabled via AsyncLLM."""
engine_args = AsyncEngineArgs(
model=model,
enable_return_routed_experts=True,
tensor_parallel_size=tp,
max_model_len=max_model_len,
disable_log_stats=True,
attention_backend="FLASH_ATTN",
)
result = asyncio.run(_run_async_inference(engine_args, prompts, max_new_tokens))
from vllm.platforms import current_platform
if current_platform.is_cuda_alike():
current_platform.empty_cache()
return result
# ---------------------------------------------------------------------------
# Validation helpers
# ---------------------------------------------------------------------------
def validate_expert_ids(
experts_list: list[np.ndarray],
num_experts: int,
) -> None:
"""Check that all expert IDs are within valid range [0, num_experts)."""
for i, experts in enumerate(experts_list):
assert np.all(experts >= 0), (
f"Prompt {i}: negative expert IDs found, min={experts.min()}"
)
assert np.all(experts < num_experts), (
f"Prompt {i}: expert ID out of range [0, {num_experts}), "
f"max={experts.max()}"
)
def validate_shapes(experts_list: list[np.ndarray]) -> None:
"""Check that all routed_experts arrays have at least 2 dimensions."""
for i, experts in enumerate(experts_list):
assert experts.ndim >= 2, (
f"Prompt {i}: expected at least 2D array, got shape {experts.shape}"
)
logger.info("Prompt %d: routed_experts shape = %s", i, experts.shape)
# ---------------------------------------------------------------------------
# Comparison helpers
# ---------------------------------------------------------------------------
def compare_token_ids(
baseline: list[list[int]],
reference: list[list[int]],
) -> float:
"""Compare token IDs from two runs. Returns mismatch ratio."""
assert len(baseline) == len(reference), (
f"Length mismatch: {len(baseline)} vs {len(reference)}"
)
total_tokens = 0
total_mismatches = 0
for i, (base, ref) in enumerate(zip(baseline, reference)):
min_len = min(len(base), len(ref))
max_len = max(len(base), len(ref))
matches = 0
for a, b in zip(base[:min_len], ref[:min_len]):
if a != b:
break
matches += 1
total_mismatches += max_len - matches
total_tokens += max_len
if matches < min_len or len(base) != len(ref):
print(
f" Prompt {i}: token_ids len={len(base)} vs {len(ref)}, "
f"mismatches={max_len - matches}/{max_len}"
)
if total_tokens == 0:
raise ValueError("No tokens to compare")
mismatch_ratio = total_mismatches / total_tokens
print(
f"Token ID mismatches: {total_mismatches}/{total_tokens} ({mismatch_ratio:.4%})"
)
return mismatch_ratio
def compare_routed_experts(
baseline: list[np.ndarray],
reference: list[np.ndarray],
threshold: float = 0.05,
) -> float:
"""Compare two runs of routed experts. Returns mismatch ratio.
Raises AssertionError if ratio exceeds threshold.
"""
assert len(baseline) == len(reference), (
f"Length mismatch: {len(baseline)} vs {len(reference)}"
)
total_elements = 0
total_mismatches = 0
for i, (base, ref) in enumerate(zip(baseline, reference)):
min_len = min(len(base), len(ref))
max_len = max(len(base), len(ref))
if min_len == 0:
continue
base_trimmed = base[:min_len]
ref_trimmed = ref[:min_len]
matches = 0
for a, b in zip(base_trimmed, ref_trimmed):
if a.sum() != b.sum():
break
matches += 1
total_mismatches += max_len - matches
total_elements += max_len
if matches < min_len or len(base) != len(ref):
print(
f" Prompt {i}: routed_experts len={len(base)} vs {len(ref)}, "
f"mismatches={max_len - matches}/{max_len}"
)
if total_elements == 0:
raise ValueError("No elements to compare")
mismatch_ratio = total_mismatches / total_elements
print(
f"Routed experts mismatches: {total_mismatches}/{total_elements} "
f"({mismatch_ratio:.4%})"
)
assert mismatch_ratio < threshold, (
f"Too many mismatches: {total_mismatches}/{total_elements} "
f"({mismatch_ratio:.4%}) exceeds threshold {threshold:.4%}"
)
return mismatch_ratio
# ---------------------------------------------------------------------------
# CLI entry point
# ---------------------------------------------------------------------------
def main():
os.environ.setdefault("VLLM_BATCH_INVARIANT", "1")
parser = argparse.ArgumentParser(
description="Test routed experts capture for MoE models"
)
parser.add_argument("--model", type=str, default=DEFAULT_MODEL)
parser.add_argument("--tp", type=int, default=1)
parser.add_argument("--max-model-len", type=int, default=4096)
parser.add_argument("--num-prompts", type=int, default=20)
parser.add_argument("--max-new-tokens", type=int, default=50)
parser.add_argument(
"--deterministic",
action="store_true",
help="Run twice and compare results for determinism check",
)
parser.add_argument(
"--threshold",
type=float,
default=0.05,
help="Maximum allowed mismatch ratio for determinism check",
)
args = parser.parse_args()
logging.basicConfig(level=logging.INFO)
prompts = TEST_PROMPTS[: args.num_prompts]
print(f"Model: {args.model}")
print(f"TP: {args.tp}")
print(f"Prompts: {len(prompts)}")
print(f"Max new tokens: {args.max_new_tokens}")
print()
print("=== Run 1 (baseline) ===")
baseline = run_inference(
model=args.model,
prompts=prompts,
max_new_tokens=args.max_new_tokens,
tp=args.tp,
max_model_len=args.max_model_len,
)
print(f"num_experts (from model config): {baseline.num_experts}")
print("\n=== Validation ===")
validate_shapes(baseline.experts_list)
validate_expert_ids(baseline.experts_list, num_experts=baseline.num_experts)
print(f"All {len(baseline.experts_list)} results passed validation.")
for i, experts in enumerate(baseline.experts_list):
print(
f" Prompt {i}: shape={experts.shape}, "
f"min={experts.min()}, max={experts.max()}"
)
if args.deterministic:
print("\n=== Run 2 (reference) ===")
reference = run_inference(
model=args.model,
prompts=prompts,
max_new_tokens=args.max_new_tokens,
tp=args.tp,
max_model_len=args.max_model_len,
)
print("\n=== Determinism Check ===")
validate_expert_ids(reference.experts_list, num_experts=baseline.num_experts)
print("\n--- Token IDs ---")
token_mismatch = compare_token_ids(
baseline.token_ids_list, reference.token_ids_list
)
print("\n--- Routed Experts ---")
expert_mismatch = compare_routed_experts(
baseline.experts_list,
reference.experts_list,
threshold=args.threshold,
)
print(
f"\nDeterminism check passed. "
f"Token mismatch: {token_mismatch:.4%}, "
f"Expert mismatch: {expert_mismatch:.4%}"
)
print("\nAll tests passed!")
if __name__ == "__main__":
main()
examples/pooling/classify/vision_classification_online.py
+1 -1
View File
@@ -8,7 +8,7 @@ NOTE:
--runner pooling \
--max-model-len 5000 \
--limit-mm-per-prompt.video 1 \
--hf-overrides '{"text_config": {"architectures": ["Qwen2_5_VLForSequenceClassification"]}}'
--hf-overrides '{"architectures": ["Qwen2_5_VLForSequenceClassification"]}'
"""
import argparse
pyproject.toml
-5
View File
@@ -9,7 +9,6 @@ requires = [
"torch == 2.10.0",
"wheel",
"jinja2",
"grpcio-tools==1.78.0",
]
build-backend = "setuptools.build_meta"
@@ -57,10 +56,6 @@ include = ["vllm*"]
"vllm/third_party/**" = ["ALL"]
"vllm/version.py" = ["F401"]
"vllm/_version.py" = ["ALL"]
# Exclude generated protobuf files
"vllm/grpc/*_pb2.py" = ["ALL"]
"vllm/grpc/*_pb2_grpc.py" = ["ALL"]
"vllm/grpc/*_pb2.pyi" = ["ALL"]
[tool.ruff.lint]
select = [
requirements/build.txt
-1
View File
@@ -10,4 +10,3 @@ jinja2>=3.1.6
regex
build
protobuf >= 5.29.6, !=6.30.*, !=6.31.*, !=6.32.*, !=6.33.0.*, !=6.33.1.*, !=6.33.2.*, !=6.33.3.*, !=6.33.4.*
grpcio-tools==1.78.0 # Required for grpc entrypoints
requirements/common.txt
-2
View File
@@ -51,8 +51,6 @@ openai-harmony >= 0.0.3 # Required for gpt-oss
anthropic >= 0.71.0
model-hosting-container-standards >= 0.1.13, < 1.0.0
mcp
grpcio
grpcio-reflection
opentelemetry-sdk >= 1.27.0
opentelemetry-api >= 1.27.0
opentelemetry-exporter-otlp >= 1.27.0
requirements/cpu.txt
+3 -3
View File
@@ -7,13 +7,13 @@ numba == 0.61.2; platform_machine != "s390x" # Required for N-gram speculative d
# Dependencies for CPUs
torch==2.10.0+cpu; platform_machine == "x86_64" or platform_machine == "s390x"
torch==2.10.0; platform_machine == "aarch64" or platform_system == "Darwin" or platform_machine == "ppc64le"
torch==2.10.0; platform_machine == "aarch64" or platform_system == "Darwin" or platform_machine == "ppc64le" or platform_machine == "riscv64"
# required for the image processor of minicpm-o-2_6, this must be updated alongside torch
torchaudio; platform_machine != "s390x"
torchaudio; platform_machine != "s390x" and platform_machine != "riscv64"
# required for the image processor of phi3v, this must be updated alongside torch
torchvision; platform_machine != "s390x"
torchvision; platform_machine != "s390x" and platform_machine != "riscv64"
# Intel Extension for PyTorch, only for x86_64 CPUs
intel-openmp==2024.2.1; platform_machine == "x86_64"
requirements/cuda.txt
+3
View File
@@ -10,6 +10,9 @@ torchaudio==2.10.0
torchvision==0.25.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version
# FlashInfer should be updated together with the Dockerfile
flashinfer-python==0.6.4
# Cap nvidia-cudnn-frontend (transitive dep of flashinfer) due to
# breaking changes in 1.19.0
nvidia-cudnn-frontend>=1.13.0,<1.19.0
# QuACK and Cutlass DSL for FA4 (cute-DSL implementation)
nvidia-cutlass-dsl>=4.4.0.dev1
requirements/rocm.txt
-1
View File
@@ -4,7 +4,6 @@
# The version of gRPC libraries should be consistent with each other
grpcio==1.78.0
grpcio-reflection==1.78.0
grpcio-tools==1.78.0
numba == 0.61.2 # Required for N-gram speculative decoding
requirements/test.in
+4 -2
View File
@@ -51,7 +51,6 @@ tritonclient>=2.51.0
# The version of gRPC libraries should be consistent with each other
grpcio==1.78.0
grpcio-reflection==1.78.0
grpcio-tools==1.78.0
arctic-inference == 0.1.1 # Required for suffix decoding test
numba == 0.61.2 # Required for N-gram speculative decoding
@@ -71,4 +70,7 @@ kaldi-native-fbank >= 1.18.7 # required for fireredasr2 test
# Newer versions of datasets require torchcoded, that makes the tests fail in CI because of a missing library.
# Older versions are in conflict with teerratorch requirements.
datasets>=3.3.0,<=3.6.0
datasets>=3.3.0,<=3.6.0
openpyxl # required for perf comparison excel report
plotly # required for perf comparison html report
requirements/test.txt
+7 -6
View File
@@ -202,6 +202,8 @@ email-validator==2.2.0
# via pydantic
encodec==0.1.1
# via vocos
et-xmlfile==2.0.0
# via openpyxl
evaluate==0.4.3
# via lm-eval
fastapi==0.128.0
@@ -289,13 +291,10 @@ grpcio==1.78.0
# via
# -r requirements/test.in
# grpcio-reflection
# grpcio-tools
# ray
# tensorboard
grpcio-reflection==1.78.0
# via -r requirements/test.in
grpcio-tools==1.78.0
# via -r requirements/test.in
h11==0.14.0
# via
# httpcore
@@ -637,6 +636,8 @@ opencv-python-headless==4.13.0.90
# albucore
# albumentations
# mistral-common
openpyxl==3.1.5
# via -r requirements/test.in
opentelemetry-api==1.35.0
# via
# opentelemetry-exporter-prometheus
@@ -737,7 +738,9 @@ platformdirs==4.3.6
# virtualenv
# wandb
plotly==5.24.1
# via genai-perf
# via
# -r requirements/test.in
# genai-perf
pluggy==1.5.0
# via
# pytest
@@ -765,7 +768,6 @@ protobuf==6.33.2
# google-api-core
# googleapis-common-protos
# grpcio-reflection
# grpcio-tools
# opentelemetry-proto
# proto-plus
# ray
@@ -1045,7 +1047,6 @@ sentry-sdk==2.52.0
# via wandb
setuptools==77.0.3
# via
# grpcio-tools
# lightning-utilities
# pytablewriter
# tensorboard
scripts/autotune_helion_kernels.py
+12 -7
View File
@@ -27,6 +27,7 @@ import time
from dataclasses import dataclass
import torch
from torch._subclasses.fake_tensor import FakeTensorMode
try:
import helion
@@ -109,7 +110,8 @@ def autotune_kernel(
)
try:
inputs_dict = kernel_wrapper.get_inputs()
with FakeTensorMode():
all_config_keys = list(kernel_wrapper.get_inputs().keys())
except NotImplementedError:
error_msg = f"Kernel '{kernel_name}' has no input generator registered"
logger.error(error_msg)
@@ -126,15 +128,15 @@ def autotune_kernel(
"Autotuning kernel '%s' for platform '%s' with %d configs",
kernel_name,
platform,
len(inputs_dict),
len(all_config_keys),
)
configs_to_autotune = {}
if not force:
existing_configs = config_manager.get_platform_configs(
kernel_name, platform
)
for config_key, inputs in inputs_dict.items():
keys_to_autotune = []
for config_key in all_config_keys:
if config_key in existing_configs:
logger.debug(
"Config '%s' already exists for platform '%s', skipping",
@@ -142,12 +144,12 @@ def autotune_kernel(
platform,
)
else:
configs_to_autotune[config_key] = inputs
keys_to_autotune.append(config_key)
else:
logger.debug("Force mode enabled, will re-autotune all configs")
configs_to_autotune = inputs_dict
keys_to_autotune = all_config_keys
if not configs_to_autotune:
if not keys_to_autotune:
logger.info(
"All configs already exist for kernel '%s' on platform '%s'. "
"Use --force to re-autotune.",
@@ -162,6 +164,9 @@ def autotune_kernel(
configs={},
)
inputs_dict = kernel_wrapper.get_inputs()
configs_to_autotune = {k: inputs_dict[k] for k in keys_to_autotune}
total_start_time = time.time()
autotuned_configs = {}
failed_configs = []
setup.py
+3 -83
View File
@@ -18,8 +18,6 @@ import torch
from packaging.version import Version, parse
from setuptools import Extension, setup
from setuptools.command.build_ext import build_ext
from setuptools.command.build_py import build_py
from setuptools.command.develop import develop
from setuptools_scm import get_version
from torch.utils.cpp_extension import CUDA_HOME, ROCM_HOME
@@ -81,81 +79,6 @@ def is_freethreaded():
return bool(sysconfig.get_config_var("Py_GIL_DISABLED"))
def compile_grpc_protos():
"""Compile gRPC protobuf definitions during build.
This generates *_pb2.py, *_pb2_grpc.py, and *_pb2.pyi files from
the vllm_engine.proto definition.
"""
try:
from grpc_tools import protoc
except ImportError:
logger.warning(
"grpcio-tools not installed, skipping gRPC proto compilation. "
"gRPC server functionality will not be available."
)
return False
proto_file = ROOT_DIR / "vllm" / "grpc" / "vllm_engine.proto"
if not proto_file.exists():
logger.warning("Proto file not found at %s, skipping compilation", proto_file)
return False
logger.info("Compiling gRPC protobuf: %s", proto_file)
result = protoc.main(
[
"grpc_tools.protoc",
f"--proto_path={ROOT_DIR}",
f"--python_out={ROOT_DIR}",
f"--grpc_python_out={ROOT_DIR}",
f"--pyi_out={ROOT_DIR}",
str(proto_file),
]
)
if result != 0:
logger.error("protoc failed with exit code %s", result)
return False
# Add SPDX headers and mypy ignore to generated files
spdx_header = (
"# SPDX-License-Identifier: Apache-2.0\n"
"# SPDX-FileCopyrightText: Copyright contributors to the vLLM project\n"
"# mypy: ignore-errors\n"
)
grpc_dir = ROOT_DIR / "vllm" / "grpc"
for generated_file in [
grpc_dir / "vllm_engine_pb2.py",
grpc_dir / "vllm_engine_pb2_grpc.py",
grpc_dir / "vllm_engine_pb2.pyi",
]:
if generated_file.exists():
content = generated_file.read_text()
if not content.startswith("# SPDX-License-Identifier"):
generated_file.write_text(spdx_header + content)
logger.info("gRPC protobuf compilation successful")
return True
class BuildPyAndGenerateGrpc(build_py):
"""Build Python modules and generate gRPC stubs from proto files."""
def run(self):
compile_grpc_protos()
super().run()
class DevelopAndGenerateGrpc(develop):
"""Develop mode that also generates gRPC stubs from proto files."""
def run(self):
compile_grpc_protos()
super().run()
class CMakeExtension(Extension):
def __init__(self, name: str, cmake_lists_dir: str = ".", **kwa) -> None:
super().__init__(name, sources=[], py_limited_api=not is_freethreaded(), **kwa)
@@ -1028,17 +951,12 @@ if _no_device():
ext_modules = []
if not ext_modules:
cmdclass = {
"build_py": BuildPyAndGenerateGrpc,
"develop": DevelopAndGenerateGrpc,
}
cmdclass = {}
else:
cmdclass = {
"build_ext": precompiled_build_ext
if envs.VLLM_USE_PRECOMPILED
else cmake_build_ext,
"build_py": BuildPyAndGenerateGrpc,
"develop": DevelopAndGenerateGrpc,
}
setup(
@@ -1064,6 +982,8 @@ setup(
"petit-kernel": ["petit-kernel"],
# Optional deps for Helion kernel development
"helion": ["helion"],
# Optional deps for gRPC server (vllm serve --grpc)
"grpc": ["smg-grpc-servicer >= 0.4.2"],
# Optional deps for OpenTelemetry tracing
"otel": [
"opentelemetry-sdk>=1.26.0",
tests/basic_correctness/test_basic_correctness.py
+11
View File
@@ -11,6 +11,8 @@ from unittest.mock import Mock
import pytest
import torch
from packaging.version import Version
from transformers import __version__ as TRANSFORMERS_VERSION
from vllm import LLM
from vllm.platforms import current_platform
@@ -91,6 +93,15 @@ def test_models(
if enable_prompt_embeds:
with torch.no_grad():
prompt_embeds = hf_model.get_prompt_embeddings(example_prompts)
if model == "hmellor/tiny-random-Gemma2ForCausalLM" and (
Version(TRANSFORMERS_VERSION) < Version("5.3.0.dev0")
):
# For Gemma 1/2 models with Transformers 5.4.0+, the prompt embeddings
# are normalised in `get_prompt_embeddings`, like Gemma 3.
# For older versions, we need to manually normalise.
embed_scale = hf_model.config.hidden_size**0.5
normalizer = torch.tensor(embed_scale, dtype=prompt_embeds[0].dtype)
prompt_embeds = [p_e * normalizer for p_e in prompt_embeds]
with VllmRunner(
model,
tests/compile/fusions_e2e/conftest.py
+10
View File
@@ -72,6 +72,16 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
rocm_aiter_ops.refresh_env_variables()
# Filter here to reduce code duplication
requires_mla = "deepseek" in model_name.lower()
is_mla = "mla" in attn_backend.backend.name.lower()
if requires_mla != is_mla:
pytest.skip(
f"Incompatible model '{model_name}' and "
f"attention backend '{attn_backend.backend.name}'"
)
# Disable, compile cache to make sure custom passes run.
# Otherwise, we can't verify fusion happened through the logs.
monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1")
tests/compile/fusions_e2e/models.py
+36
View File
@@ -44,6 +44,20 @@ ROCM_AITER_UNIFIED_ATTN = pytest.param(
),
)
FLASHINFER_MLA_ATTN = pytest.param(
AttentionBackendCase(backend=AttentionBackendEnum.FLASHINFER_MLA),
id="FLASHINFER_MLA",
marks=pytest.mark.skipif(
not is_blackwell() or not has_flashinfer(),
reason="FI backend requires Blackwell and FlashInfer",
),
)
TRITON_MLA_ATTN = pytest.param(
AttentionBackendCase(backend=AttentionBackendEnum.TRITON_MLA),
id="TRITON_MLA",
)
# Models
llama3_8b = ModelFusionInfo(
model_name="meta-llama/Llama-3.1-8B-Instruct",
@@ -126,3 +140,25 @@ qwen3_a3b_fp8 = ModelFusionInfo(
async_tp=n_layers * 2,
),
)
deepseek_v3_fp8 = ModelFusionInfo(
model_name="deepseek-ai/DeepSeek-V3",
matches=lambda n_layers: Matches(
# 3 per dense layer (first 3):
# - input_rms + qkv_proj
# - q_a_layernorm + q_b_proj (inside MLA wrapper)
# - post_attn_layernorm + MLP
# 2 per MoE layer (remaining) due to MoE wrapping
rms_quant_fusion=n_layers * 2 + min(3, n_layers), # add for 3 dense layers
# TODO silu+block quant
# act_quant_fusion=min(3, n_layers), # dense layers only
act_quant_fusion=0,
# MLA attn + quant not supported yet:
# https://github.com/vllm-project/vllm/issues/35792
attn_quant_fusion=0,
ar_rms_fusion=n_layers * 2 + 1,
# TODO
# sequence_parallel= n_layers * 2 + 1,
# async_tp=n_layers * 2,
),
)
tests/compile/fusions_e2e/test_tp1_quant.py
+13 -8
View File
@@ -17,9 +17,12 @@ from .common import (
)
from .models import (
FLASHINFER_ATTN,
FLASHINFER_MLA_ATTN,
ROCM_AITER_UNIFIED_ATTN,
ROCM_ATTN,
TRITON_ATTN,
TRITON_MLA_ATTN,
deepseek_v3_fp8,
llama3_8b_fp4,
llama3_8b_fp8,
llama4_scout_fp4,
@@ -33,6 +36,9 @@ from .models import (
[
(*llama3_8b_fp8, False),
(*qwen3_a3b_fp8, False),
(*qwen3_a3b_fp8, True),
(*deepseek_v3_fp8, False),
(*deepseek_v3_fp8, True),
pytest.param(
*llama4_scout_fp8,
False,
@@ -41,13 +47,6 @@ from .models import (
reason="Llama4 Scout FP8 only supported on CUDA",
),
),
pytest.param(
*qwen3_a3b_fp8,
True,
marks=pytest.mark.skipif(
not current_platform.is_cuda(), reason="DeepGemm only supported on CUDA"
),
),
],
)
@pytest.mark.parametrize(
@@ -57,6 +56,8 @@ from .models import (
FLASHINFER_ATTN,
ROCM_ATTN,
ROCM_AITER_UNIFIED_ATTN,
FLASHINFER_MLA_ATTN,
TRITON_MLA_ATTN,
],
)
@pytest.mark.parametrize("n_layers", [6])
@@ -75,6 +76,9 @@ def test_tp1_fp8_fusions(
run_e2e_fusion_test,
monkeypatch,
):
if use_deepgemm and not current_platform.is_cuda():
pytest.skip("DeepGemm only supported on CUDA")
if use_deepgemm and is_flashinfer_fp8_blockscale_gemm_supported():
# Flashinfer block FP8 GEMM has internal quantization, so it can't
# be fused with other ops.
@@ -86,7 +90,8 @@ def test_tp1_fp8_fusions(
matches = matches_fn(n_layers)
if "qwen" in model_name.lower() and "-quant_fp8" in custom_ops:
block_fp8 = "qwen" in model_name.lower() or "deepseek" in model_name.lower()
if block_fp8 and "-quant_fp8" in custom_ops:
# This is why config forces +quant_fp8 by default
pytest.skip("native QuantFP8 matching not supported for group quant")
tests/compile/fusions_e2e/test_tp2_ar_rms.py
+9 -4
View File
@@ -17,7 +17,9 @@ from .common import (
)
from .models import (
FLASHINFER_ATTN,
FLASHINFER_MLA_ATTN,
TRITON_ATTN,
deepseek_v3_fp8,
llama3_8b,
llama3_8b_fp4,
llama3_8b_fp8,
@@ -33,10 +35,12 @@ pytestmark = pytest.mark.skipif(not current_platform.is_cuda(), reason="Only tes
@multi_gpu_test(num_gpus=2)
@pytest.mark.parametrize(
"model_name, matches_fn, model_kwargs, hf_overrides",
# qwen3-fp8 should still fuse AR+rms even though group quant is not yet supported
[llama3_8b_fp8, llama4_scout_fp8, qwen3_a3b_fp8],
# qwen3 & dsv3 should still fuse AR+rms even though group quant is not yet supported
[llama3_8b_fp8, llama4_scout_fp8, qwen3_a3b_fp8, deepseek_v3_fp8],
)
@pytest.mark.parametrize(
"attn_backend", [TRITON_ATTN, FLASHINFER_ATTN, FLASHINFER_MLA_ATTN]
)
@pytest.mark.parametrize("attn_backend", [TRITON_ATTN, FLASHINFER_ATTN])
@pytest.mark.parametrize("n_layers", [4])
@pytest.mark.parametrize("custom_ops", custom_ops_combos("quant_fp8", "rms_norm"))
@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION)
@@ -54,7 +58,8 @@ def test_tp2_ar_rms_fp8_fusions(
):
matches = matches_fn(n_layers)
if "qwen" in model_name.lower() and "-quant_fp8" in custom_ops:
block_fp8 = "qwen" in model_name.lower() or "deepseek" in model_name.lower()
if block_fp8 and "-quant_fp8" in custom_ops:
# This is why config forces +quant_fp8 by default
pytest.skip("native QuantFP8 matching not supported for group quant")
tests/compile/test_aot_compile.py
+114
View File
@@ -4,6 +4,7 @@
import functools
import hashlib
import multiprocessing
import os
import pickle
import tempfile
from contextlib import contextmanager
@@ -19,6 +20,7 @@ from vllm.compilation.caching import (
StandaloneCompiledArtifacts,
VllmSerializableFunction,
)
from vllm.compilation.counter import compilation_counter
from vllm.compilation.decorators import support_torch_compile
from vllm.config import (
CompilationConfig,
@@ -763,3 +765,115 @@ class TestStandaloneCompiledArtifactsIntegration:
assert isinstance(config, dict)
assert "bundled_autograd_cache" in config
assert config["bundled_autograd_cache"] is True
@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
def test_disable_compile_cache_skips_aot_save(
monkeypatch: pytest.MonkeyPatch, fresh_vllm_cache: str
):
"""When VLLM_DISABLE_COMPILE_CACHE=1, AOT artifacts must not be saved."""
monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1")
monkeypatch.setenv("VLLM_USE_AOT_COMPILE", "1")
disable_envs_cache()
args = (torch.randn(10, 10),)
expected = reference_fn(*args)
vllm_config = make_vllm_config()
with (
use_vllm_config(vllm_config),
compilation_counter.expect(
num_aot_compiles=1,
num_aot_artifacts_saved=0,
num_aot_artifacts_loaded=0,
),
):
mod = CompiledMod(vllm_config=vllm_config)
actual = mod(*args)
assert torch.allclose(actual, expected)
# No cached artifact should exist on disk
aot_dir = os.path.join(fresh_vllm_cache, "torch_compile_cache", "torch_aot_compile")
if os.path.isdir(aot_dir):
for root, _dirs, files in os.walk(aot_dir):
for f in files:
assert f != "model", (
f"AOT artifact unexpectedly saved at {os.path.join(root, f)}"
)
@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
def test_disable_compile_cache_skips_aot_load(
monkeypatch: pytest.MonkeyPatch, fresh_vllm_cache: str
):
"""When VLLM_DISABLE_COMPILE_CACHE=1, AOT artifacts must not be loaded."""
# Phase 1: compile and save with cache enabled
monkeypatch.setenv("VLLM_USE_AOT_COMPILE", "1")
disable_envs_cache()
args = (torch.randn(10, 10),)
vllm_config = make_vllm_config()
with (
use_vllm_config(vllm_config),
compilation_counter.expect(num_aot_artifacts_saved=1),
):
CompiledMod(vllm_config=vllm_config)(*args)
# Phase 2: disable cache, compile again — should NOT load from disk
monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1")
disable_envs_cache()
torch._dynamo.reset()
vllm_config = make_vllm_config()
with (
use_vllm_config(vllm_config),
compilation_counter.expect(
num_aot_compiles=1,
num_aot_artifacts_saved=0,
num_aot_artifacts_loaded=0,
),
):
mod = CompiledMod(vllm_config=vllm_config)
mod(*args)
assert not mod.was_aot_compile_fn_loaded_from_disk
@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
def test_aot_counters_on_save_and_load(
monkeypatch: pytest.MonkeyPatch, fresh_vllm_cache: str
):
"""Verify AOT counters are incremented correctly on save and load."""
monkeypatch.setenv("VLLM_USE_AOT_COMPILE", "1")
disable_envs_cache()
args = (torch.randn(10, 10),)
# Phase 1: fresh compile + save
vllm_config = make_vllm_config()
with (
use_vllm_config(vllm_config),
compilation_counter.expect(
num_aot_compiles=1,
num_aot_artifacts_saved=1,
num_aot_artifacts_loaded=0,
),
):
CompiledMod(vllm_config=vllm_config)(*args)
# Phase 2: load from cache
monkeypatch.setenv("VLLM_FORCE_AOT_LOAD", "1")
disable_envs_cache()
vllm_config = make_vllm_config()
with (
use_vllm_config(vllm_config),
compilation_counter.expect(
num_aot_compiles=0,
num_aot_artifacts_saved=0,
num_aot_artifacts_loaded=1,
),
):
CompiledMod(vllm_config=vllm_config)(*args)
tests/compile/test_compile_ranges.py
+82
View File
@@ -127,6 +127,88 @@ def test_compile_config_get_compile_ranges():
]
class PostGradStaticShapeChecker(InductorPass):
"""Asserts that compile_sizes entries produce graphs with fully concrete
(non-symbolic) shapes, and compile_ranges entries have symbolic shapes."""
def __init__(self):
self.num_static_calls = 0
self.num_dynamic_calls = 0
def __call__(self, graph: fx.Graph):
from torch.fx.experimental.symbolic_shapes import is_symbolic
compile_range = get_pass_context().compile_range
is_single = compile_range.is_single_size()
for node in graph.nodes:
val = node.meta.get("val")
if val is None:
val = node.meta.get("example_value")
if isinstance(val, torch.Tensor):
has_symbolic = any(is_symbolic(d) for d in val.shape)
if is_single:
assert not has_symbolic, (
f"compile_sizes entry {compile_range}: "
f"node '{node.name}' has symbolic shape "
f"{val.shape}"
)
else:
# compile_ranges should have at least some
# symbolic shapes (the batch dimension)
if has_symbolic:
self.num_dynamic_calls += 1
return
if is_single:
self.num_static_calls += 1
def uuid(self) -> str:
state: dict[str, Any] = {}
return InductorPass.hash_dict(state)
def test_compile_sizes_produce_static_shapes(use_fresh_inductor_cache):
"""Verify that compile_sizes entries are compiled with fully concrete
shapes (no SymInts), while compile_ranges entries retain dynamic shapes."""
checker = PostGradStaticShapeChecker()
torch.set_default_device("cuda")
vllm_config = VllmConfig(
scheduler_config=SchedulerConfig(
max_num_batched_tokens=8192,
max_model_len=8192,
is_encoder_decoder=False,
),
compilation_config=CompilationConfig(
mode=CompilationMode.VLLM_COMPILE,
compile_ranges_endpoints=[8],
compile_sizes=[16],
inductor_compile_config={
"post_grad_custom_post_pass": checker,
},
),
)
with set_current_vllm_config(vllm_config):
model = TestModel(vllm_config=vllm_config, prefix="").eval()
# 3 compilations: Range(1,8), Range(9,8192), single-size 16
with compilation_counter.expect(
num_graphs_seen=1,
num_piecewise_graphs_seen=1,
num_backend_compilations=3,
):
run_model(vllm_config, model, [1, 16, 64])
# compile_sizes=16 should produce static shapes
assert checker.num_static_calls == 1, (
f"Expected 1 static compilation, got {checker.num_static_calls}"
)
# compile_ranges should produce dynamic shapes
assert checker.num_dynamic_calls == 2, (
f"Expected 2 dynamic compilations, got {checker.num_dynamic_calls}"
)
def test_inductor_cache_compile_ranges(monkeypatch, use_fresh_inductor_cache):
# To force multiple compilations, we disable the compile cache
monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1")
tests/compile/test_graph_partition.py
+105 -15
View File
@@ -7,7 +7,7 @@ import pytest
import torch
from torch.fx.experimental.proxy_tensor import make_fx
from vllm.compilation.backends import split_graph
from vllm.compilation.backends import _is_empty_allocation_node, split_graph
from vllm.compilation.passes.fx_utils import find_op_nodes
# This import automatically registers `torch.ops.silly.attention`
@@ -186,10 +186,25 @@ def test_consecutive_ops_in_split():
] + ["output"]
def test_empty_only_partition_is_merged():
def _get_empty_nodes(split_item):
return [
node for node in split_item.graph.graph.nodes if _is_empty_allocation_node(node)
]
def _subgraphs_with_empty_nodes(split_items, *, is_splitting_graph):
return [
split_item
for split_item in split_items
if split_item.is_splitting_graph == is_splitting_graph
and _get_empty_nodes(split_item)
]
def test_empty_only_partition_stays_separate_after_splitting_predecessor():
"""
Test that an empty-allocation-only partition is merged into its previous
partition during Dynamo FX splitting.
Empty-only subgraphs should not be merged when the only predecessor is
a splitting-op subgraph.
"""
def model_fn(x: torch.Tensor) -> torch.Tensor:
@@ -204,9 +219,65 @@ def test_empty_only_partition_is_merged():
split_ops = ["aten::sin", "aten::cos.out"]
split_gm, split_items = split_graph(gm, split_ops)
# Without the merge, this graph is split into 3 partitions where the
# middle partition contains only aten::empty_like.
assert len(split_items) == 2, "Empty-only partition should be merged"
# Graph partitioning for this pattern is:
# [sin], [empty_like], [cos.out].
assert len(split_items) == 3, (
"Empty-only partition should not merge into splitting-op subgraph"
)
splitting_with_empty = _subgraphs_with_empty_nodes(
split_items, is_splitting_graph=True
)
assert len(splitting_with_empty) == 0, (
"Splitting-op subgraphs should not contain empty allocation nodes: "
f"{[item.submod_name for item in splitting_with_empty]}"
)
output_original = gm(x)
output_split = split_gm(x)
assert torch.allclose(output_original, output_split), "Output mismatch after split"
def test_empty_only_partition_is_merged():
"""
Empty-only subgraphs should still be merged when a non-splitting predecessor
exists. The merged empty node must remain outside splitting-op subgraphs.
"""
def model_fn(x: torch.Tensor) -> torch.Tensor:
base = x + 1
y = torch.sin(base)
out = torch.empty_like(base)
torch.ops.aten.cos.out(base, out=out)
return out + y
x = torch.randn(4, 3)
gm = make_fx(model_fn)(x)
split_gm, split_items = split_graph(gm, ["aten::sin", "aten::cos.out"])
# Partitioning should be:
# [add, empty_like], [sin], [cos.out], [add].
assert len(split_items) == 4, (
"Empty-only partition should be merged into non-splitting predecessor"
)
splitting_with_empty = _subgraphs_with_empty_nodes(
split_items, is_splitting_graph=True
)
assert len(splitting_with_empty) == 0, (
"Splitting-op subgraphs should not contain empty allocation nodes: "
f"{[item.submod_name for item in splitting_with_empty]}"
)
non_splitting_with_empty = _subgraphs_with_empty_nodes(
split_items, is_splitting_graph=False
)
assert len(non_splitting_with_empty) == 1, (
"Exactly one non-splitting subgraph should contain the merged empty node"
)
assert len(_get_empty_nodes(non_splitting_with_empty[0])) == 1, (
"Expected exactly one empty allocation node in merged subgraph"
)
output_original = gm(x)
output_split = split_gm(x)
@@ -220,18 +291,37 @@ def test_builtin_empty_only_partition_is_merged():
"""
def model_fn(x: torch.Tensor) -> torch.Tensor:
out1 = torch.empty_like(x)
torch.ops.silly.attention(x, x, x, out1)
out2 = torch.empty_like(x)
torch.ops.silly.attention(out1, out1, out1, out2)
return out2
hidden = x + 1
out1 = torch.empty_like(hidden)
torch.ops.silly.attention(hidden, hidden, hidden, out1)
out2 = torch.empty_like(hidden)
torch.ops.silly.attention(out1, out1, hidden, out2)
return out2 + hidden
gm = torch.fx.symbolic_trace(model_fn)
split_gm, split_items = split_graph(gm, ["silly::attention"])
# Without the empty-only merge, this graph creates 4 partitions:
# [empty_like], [attention], [empty_like], [attention].
assert len(split_items) == 3, "Builtin empty-only partition should be merged"
# Without empty-only merge, this graph would split into:
# [add, empty_like], [attention], [empty_like], [attention], [add].
assert len(split_items) == 4, "Builtin empty-only partition should be merged"
splitting_with_empty = _subgraphs_with_empty_nodes(
split_items, is_splitting_graph=True
)
assert len(splitting_with_empty) == 0, (
"Splitting-op subgraphs should not contain empty allocation nodes: "
f"{[item.submod_name for item in splitting_with_empty]}"
)
non_splitting_with_empty = _subgraphs_with_empty_nodes(
split_items, is_splitting_graph=False
)
assert len(non_splitting_with_empty) == 1, (
"Exactly one non-splitting subgraph should contain merged empty nodes"
)
assert len(_get_empty_nodes(non_splitting_with_empty[0])) == 2, (
"Expected two builtin empty_like nodes in merged non-splitting subgraph"
)
x = torch.randn(2, 3, device="cuda")
output_original = gm(x)
tests/distributed/test_pipeline_parallel.py
+3
View File
@@ -247,6 +247,7 @@ def _compare_tp(
hf_config = get_config(model_id, trust_remote_code)
require_embed_inputs = model_info.require_embed_inputs
max_num_seqs = model_info.max_num_seqs
enable_prefix_caching = model_info.enable_prefix_caching
dtype = "float16"
if hf_config.model_type in _FLOAT16_NOT_SUPPORTED_MODELS:
@@ -300,6 +301,8 @@ def _compare_tp(
common_args.extend(["--load-format", load_format])
if hf_overrides:
common_args.extend(["--hf-overrides", json.dumps(hf_overrides)])
if not enable_prefix_caching:
common_args.append("--no-enable-prefix-caching")
if require_embed_inputs:
common_args.extend(
[
tests/entrypoints/openai/test_anthropic_messages_conversion.py
+49
View File
@@ -324,3 +324,52 @@ class TestToolResultContent:
if m["role"] == "user" and isinstance(m.get("content"), list)
]
assert len(user_follow_ups) == 0
# ======================================================================
# Attribution header stripping
# ======================================================================
class TestAttributionHeaderStripping:
def test_billing_header_stripped_from_system(self):
"""Claude Code's x-anthropic-billing-header block should be
stripped to preserve prefix caching."""
request = _make_request(
[{"role": "user", "content": "Hello"}],
system=[
{"type": "text", "text": "You are a helpful assistant."},
{
"type": "text",
"text": "x-anthropic-billing-header: "
"cc_version=2.1.37.abc; cc_entrypoint=cli;",
},
],
)
result = _convert(request)
system_msg = result.messages[0]
assert system_msg["role"] == "system"
assert system_msg["content"] == "You are a helpful assistant."
def test_system_without_billing_header_unchanged(self):
"""Normal system blocks should pass through unchanged."""
request = _make_request(
[{"role": "user", "content": "Hello"}],
system=[
{"type": "text", "text": "You are a helpful assistant."},
{"type": "text", "text": " Be concise."},
],
)
result = _convert(request)
system_msg = result.messages[0]
assert system_msg["content"] == "You are a helpful assistant. Be concise."
def test_system_string_unchanged(self):
"""String system prompts should pass through unchanged."""
request = _make_request(
[{"role": "user", "content": "Hello"}],
system="You are a helpful assistant.",
)
result = _convert(request)
system_msg = result.messages[0]
assert system_msg["content"] == "You are a helpful assistant."
tests/entrypoints/openai/test_lora_adapters.py
+3 -3
View File
@@ -196,7 +196,7 @@ async def test_dynamic_lora_invalid_files(client: openai.AsyncOpenAI, tmp_path):
invalid_files.mkdir()
(invalid_files / "adapter_config.json").write_text("this is not json")
with pytest.raises(openai.BadRequestError):
with pytest.raises(openai.InternalServerError):
await client.post(
"load_lora_adapter",
cast_to=str,
@@ -232,7 +232,7 @@ async def test_dynamic_lora_badrequests(
json.dump(adapter_config, f)
# Test loading the adapter
with pytest.raises(openai.BadRequestError, match=expected_error):
with pytest.raises(openai.InternalServerError, match=expected_error):
await client.post(
"load_lora_adapter",
cast_to=str,
@@ -312,7 +312,7 @@ async def test_loading_invalid_adapters_does_not_break_others(
body={"lora_name": "notfound", "lora_path": "/not/an/adapter"},
)
for _ in range(25):
with suppress(openai.BadRequestError):
with suppress(openai.InternalServerError):
await client.post(
"load_lora_adapter",
cast_to=str,
tests/entrypoints/openai/test_realtime_validation.py
+1 -1
View File
@@ -118,7 +118,7 @@ async def test_multi_chunk_streaming(
# JIT compilation
warmup_done = False
while not warmup_done:
event = await receive_event(ws, timeout=360.0)
event = await receive_event(ws, timeout=600.0)
if event["type"] in ("transcription.done", "error"):
warmup_done = True
tests/entrypoints/openai/test_serving_responses.py
+6 -1
View File
@@ -659,9 +659,10 @@ class TestStreamingReasoningToContentTransition:
# Mock the reasoning parser on the serving instance
mock_parser = MagicMock()
mock_parser.extract_reasoning_streaming = mock_extract_reasoning_streaming
mock_parser.extract_tool_calls_streaming = mock_extract_reasoning_streaming
serving.parser = MagicMock()
serving.parser.reasoning_parser_cls = MagicMock(return_value=mock_parser)
serving.parser.tool_parser_cls = MagicMock(return_value=mock_parser)
# Create contexts for each streaming chunk
contexts = [
_make_simple_context_with_output("chunk1", [10]),
@@ -739,8 +740,10 @@ class TestStreamingReasoningToContentTransition:
mock_parser = MagicMock()
mock_parser.extract_reasoning_streaming = mock_extract_reasoning_streaming
mock_parser.extract_tool_calls_streaming = mock_extract_reasoning_streaming
serving.parser = MagicMock()
serving.parser.reasoning_parser_cls = MagicMock(return_value=mock_parser)
serving.parser.tool_parser_cls = MagicMock(return_value=mock_parser)
contexts = [
_make_simple_context_with_output("chunk1", [10]),
@@ -812,8 +815,10 @@ class TestStreamingReasoningToContentTransition:
mock_parser = MagicMock()
mock_parser.extract_reasoning_streaming = mock_extract_reasoning_streaming
mock_parser.extract_tool_calls_streaming = mock_extract_reasoning_streaming
serving.parser = MagicMock()
serving.parser.reasoning_parser_cls = MagicMock(return_value=mock_parser)
serving.parser.tool_parser_cls = MagicMock(return_value=mock_parser)
contexts = [
_make_simple_context_with_output("chunk1", [10]),
tests/entrypoints/openai/test_shutdown.py
-459
View File
@@ -1,20 +1,14 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Integration tests for shutdown behavior, timeout, and signal handling."""
import asyncio
import signal
import subprocess
import sys
import time
from dataclasses import dataclass, field
import httpx
import openai
import psutil
import pytest
from tests.utils import RemoteOpenAIServer
from vllm.platforms import current_platform
from vllm.utils.network_utils import get_open_port
@@ -24,101 +18,6 @@ MODEL_NAME = "hmellor/tiny-random-LlamaForCausalLM"
_IS_ROCM = current_platform.is_rocm()
_SERVER_STARTUP_TIMEOUT = 120
_PROCESS_EXIT_TIMEOUT = 15
_SHUTDOWN_DETECTION_TIMEOUT = 10
_CHILD_CLEANUP_TIMEOUT = 10
def _get_child_pids(parent_pid: int) -> list[int]:
try:
parent = psutil.Process(parent_pid)
return [c.pid for c in parent.children(recursive=True)]
except psutil.NoSuchProcess:
return []
async def _assert_children_cleaned_up(
child_pids: list[int],
timeout: float = _CHILD_CLEANUP_TIMEOUT,
):
"""Wait for child processes to exit and fail if any remain."""
if not child_pids:
return
deadline = time.time() + timeout
while time.time() < deadline:
still_alive = []
for pid in child_pids:
try:
p = psutil.Process(pid)
if p.is_running() and p.status() != psutil.STATUS_ZOMBIE:
still_alive.append(pid)
except psutil.NoSuchProcess:
pass
if not still_alive:
return
await asyncio.sleep(0.5)
pytest.fail(
f"Child processes {still_alive} still alive after {timeout}s. "
f"Process cleanup may not be working correctly."
)
@dataclass
class ShutdownState:
got_503: bool = False
got_500: bool = False
requests_after_sigterm: int = 0
aborted_requests: int = 0
connection_errors: int = 0
stop_requesting: bool = False
errors: list[str] = field(default_factory=list)
async def _concurrent_request_loop(
client: openai.AsyncOpenAI,
state: ShutdownState,
sigterm_sent: asyncio.Event | None = None,
concurrency: int = 10,
):
"""Run multiple concurrent requests to keep the server busy."""
async def single_request():
while not state.stop_requesting:
try:
response = await client.completions.create(
model=MODEL_NAME,
prompt="Write a story: ",
max_tokens=200,
)
if sigterm_sent is not None and sigterm_sent.is_set():
state.requests_after_sigterm += 1
# Check if any choice has finish_reason='abort'
if any(choice.finish_reason == "abort" for choice in response.choices):
state.aborted_requests += 1
except openai.APIStatusError as e:
if e.status_code == 503:
state.got_503 = True
elif e.status_code == 500:
state.got_500 = True
else:
state.errors.append(f"API error: {e}")
except (openai.APIConnectionError, httpx.RemoteProtocolError):
state.connection_errors += 1
if sigterm_sent is not None and sigterm_sent.is_set():
break
except Exception as e:
state.errors.append(f"Unexpected error: {e}")
break
await asyncio.sleep(0.01)
tasks = [asyncio.create_task(single_request()) for _ in range(concurrency)]
try:
await asyncio.gather(*tasks, return_exceptions=True)
finally:
for t in tasks:
if not t.done():
t.cancel()
@pytest.mark.asyncio
@@ -204,361 +103,3 @@ async def test_shutdown_on_engine_failure():
return_code = proc.wait(timeout=_PROCESS_EXIT_TIMEOUT)
assert return_code is not None
@pytest.mark.asyncio
async def test_wait_timeout_completes_requests():
"""Verify wait timeout: new requests rejected, in-flight requests complete."""
server_args = [
"--dtype",
"bfloat16",
"--max-model-len",
"256",
"--enforce-eager",
"--gpu-memory-utilization",
"0.05",
"--max-num-seqs",
"4",
"--shutdown-timeout",
"30",
]
with RemoteOpenAIServer(MODEL_NAME, server_args) as remote_server:
client = remote_server.get_async_client()
proc = remote_server.proc
child_pids = _get_child_pids(proc.pid)
state = ShutdownState()
sigterm_sent = asyncio.Event()
request_task = asyncio.create_task(
_concurrent_request_loop(client, state, sigterm_sent, concurrency=10)
)
await asyncio.sleep(0.5)
proc.send_signal(signal.SIGTERM)
sigterm_sent.set()
try:
await asyncio.wait_for(request_task, timeout=_SHUTDOWN_DETECTION_TIMEOUT)
except asyncio.TimeoutError:
pass
finally:
state.stop_requesting = True
if not request_task.done():
request_task.cancel()
await asyncio.gather(request_task, return_exceptions=True)
# wait timeout should complete in-flight requests
assert state.requests_after_sigterm > 0, (
f"Wait timeout should complete in-flight requests. "
f"503: {state.got_503}, 500: {state.got_500}, "
f"conn_errors: {state.connection_errors}, errors: {state.errors}"
)
# server must stop accepting new requests (503, 500, or connection close)
assert state.got_503 or state.got_500 or state.connection_errors > 0, (
f"Server should stop accepting requests. "
f"completed: {state.requests_after_sigterm}, errors: {state.errors}"
)
await _assert_children_cleaned_up(child_pids)
@pytest.mark.asyncio
@pytest.mark.parametrize("wait_for_engine_idle", [0.0, 2.0])
async def test_abort_timeout_exits_quickly(wait_for_engine_idle: float):
server_args = [
"--dtype",
"bfloat16",
"--max-model-len",
"256",
"--enforce-eager",
"--gpu-memory-utilization",
"0.05",
"--max-num-seqs",
"4",
"--shutdown-timeout",
"0",
]
with RemoteOpenAIServer(MODEL_NAME, server_args) as remote_server:
proc = remote_server.proc
child_pids = _get_child_pids(proc.pid)
if wait_for_engine_idle > 0:
client = remote_server.get_async_client()
# Send requests to ensure engine is fully initialized
for _ in range(2):
await client.completions.create(
model=MODEL_NAME,
prompt="Test request: ",
max_tokens=10,
)
# Wait for engine to become idle
await asyncio.sleep(wait_for_engine_idle)
start_time = time.time()
proc.send_signal(signal.SIGTERM)
# abort timeout (0) should exit promptly
for _ in range(20):
if proc.poll() is not None:
break
time.sleep(0.1)
if proc.poll() is None:
proc.kill()
proc.wait(timeout=5)
pytest.fail("Process did not exit after SIGTERM with abort timeout")
exit_time = time.time() - start_time
assert exit_time < 2, f"Default shutdown took too long: {exit_time:.1f}s"
assert proc.returncode in (0, -15, None), f"Unexpected: {proc.returncode}"
await _assert_children_cleaned_up(child_pids)
@pytest.mark.asyncio
async def test_wait_timeout_with_short_duration():
"""Verify server exits cleanly with a short wait timeout."""
wait_timeout = 3
server_args = [
"--dtype",
"bfloat16",
"--max-model-len",
"256",
"--enforce-eager",
"--gpu-memory-utilization",
"0.05",
"--max-num-seqs",
"4",
"--shutdown-timeout",
str(wait_timeout),
]
with RemoteOpenAIServer(MODEL_NAME, server_args) as remote_server:
client = remote_server.get_async_client()
proc = remote_server.proc
child_pids = _get_child_pids(proc.pid)
state = ShutdownState()
request_task = asyncio.create_task(
_concurrent_request_loop(client, state, concurrency=3)
)
await asyncio.sleep(0.5)
start_time = time.time()
proc.send_signal(signal.SIGTERM)
# server should exit within wait_timeout + buffer
max_wait = wait_timeout + 15
for _ in range(int(max_wait * 10)):
if proc.poll() is not None:
break
time.sleep(0.1)
exit_time = time.time() - start_time
state.stop_requesting = True
if not request_task.done():
request_task.cancel()
await asyncio.gather(request_task, return_exceptions=True)
if proc.poll() is None:
proc.kill()
proc.wait(timeout=5)
pytest.fail(f"Process did not exit within {max_wait}s after SIGTERM")
assert exit_time < wait_timeout + 10, (
f"Took too long to exit ({exit_time:.1f}s), expected <{wait_timeout + 10}s"
)
assert proc.returncode in (0, -15, None), f"Unexpected: {proc.returncode}"
await _assert_children_cleaned_up(child_pids)
@pytest.mark.asyncio
async def test_abort_timeout_fails_inflight_requests():
"""Verify abort timeout (0) immediately aborts in-flight requests."""
server_args = [
"--dtype",
"bfloat16",
"--max-model-len",
"256",
"--enforce-eager",
"--gpu-memory-utilization",
"0.05",
"--max-num-seqs",
"4",
"--shutdown-timeout",
"0",
]
with RemoteOpenAIServer(MODEL_NAME, server_args) as remote_server:
client = remote_server.get_async_client()
proc = remote_server.proc
child_pids = _get_child_pids(proc.pid)
state = ShutdownState()
sigterm_sent = asyncio.Event()
request_task = asyncio.create_task(
_concurrent_request_loop(client, state, sigterm_sent, concurrency=10)
)
await asyncio.sleep(0.5)
proc.send_signal(signal.SIGTERM)
sigterm_sent.set()
try:
await asyncio.wait_for(request_task, timeout=5)
except asyncio.TimeoutError:
pass
finally:
state.stop_requesting = True
if not request_task.done():
request_task.cancel()
await asyncio.gather(request_task, return_exceptions=True)
# With abort timeout (0), requests should be aborted (finish_reason='abort')
# or rejected (connection errors or API errors)
assert (
state.aborted_requests > 0
or state.connection_errors > 0
or state.got_500
or state.got_503
), (
f"Abort timeout should cause request aborts or failures. "
f"aborted: {state.aborted_requests}, "
f"503: {state.got_503}, 500: {state.got_500}, "
f"conn_errors: {state.connection_errors}, "
f"completed: {state.requests_after_sigterm}"
)
# Verify fast shutdown
start_time = time.time()
for _ in range(100):
if proc.poll() is not None:
break
time.sleep(0.1)
exit_time = time.time() - start_time
assert exit_time < 10, f"Abort timeout shutdown took too long: {exit_time:.1f}s"
await _assert_children_cleaned_up(child_pids)
@pytest.mark.asyncio
async def test_request_rejection_during_shutdown():
"""Verify new requests are rejected with error during shutdown."""
server_args = [
"--dtype",
"bfloat16",
"--max-model-len",
"256",
"--enforce-eager",
"--gpu-memory-utilization",
"0.05",
"--max-num-seqs",
"4",
"--shutdown-timeout",
"30",
]
with RemoteOpenAIServer(MODEL_NAME, server_args) as remote_server:
client = remote_server.get_async_client()
proc = remote_server.proc
child_pids = _get_child_pids(proc.pid)
proc.send_signal(signal.SIGTERM)
await asyncio.sleep(1.0)
# Try to send new requests - they should be rejected
rejected_count = 0
for _ in range(10):
try:
await client.completions.create(
model=MODEL_NAME, prompt="Hello", max_tokens=10
)
except (
openai.APIStatusError,
openai.APIConnectionError,
httpx.RemoteProtocolError,
):
rejected_count += 1
await asyncio.sleep(0.1)
assert rejected_count > 0, (
f"Expected requests to be rejected during shutdown, "
f"but {rejected_count} were rejected out of 10"
)
await _assert_children_cleaned_up(child_pids)
@pytest.mark.asyncio
async def test_multi_api_server_shutdown():
"""Verify shutdown works with multiple API servers."""
server_args = [
"--dtype",
"bfloat16",
"--max-model-len",
"256",
"--enforce-eager",
"--gpu-memory-utilization",
"0.05",
"--max-num-seqs",
"4",
"--shutdown-timeout",
"30",
"--api-server-count",
"2",
]
with RemoteOpenAIServer(MODEL_NAME, server_args, auto_port=True) as remote_server:
client = remote_server.get_async_client()
proc = remote_server.proc
child_pids = _get_child_pids(proc.pid)
assert len(child_pids) >= 2, (
f"Expected at least 2 child processes, got {len(child_pids)}"
)
state = ShutdownState()
sigterm_sent = asyncio.Event()
# Start concurrent requests across both API servers
request_task = asyncio.create_task(
_concurrent_request_loop(client, state, sigterm_sent, concurrency=8)
)
await asyncio.sleep(0.5)
# Send SIGTERM to parent - should propagate to all children
proc.send_signal(signal.SIGTERM)
sigterm_sent.set()
try:
await asyncio.wait_for(request_task, timeout=_SHUTDOWN_DETECTION_TIMEOUT)
except asyncio.TimeoutError:
pass
finally:
state.stop_requesting = True
if not request_task.done():
request_task.cancel()
await asyncio.gather(request_task, return_exceptions=True)
for _ in range(300): # up to 30 seconds
if proc.poll() is not None:
break
time.sleep(0.1)
if proc.poll() is None:
proc.kill()
proc.wait(timeout=5)
pytest.fail("Process did not exit after SIGTERM")
await _assert_children_cleaned_up(child_pids)
tests/entrypoints/pooling/classify/test_online_vision.py
+1 -5
View File
@@ -12,11 +12,7 @@ from vllm.multimodal.utils import encode_image_url, fetch_image
MODEL_NAME = "muziyongshixin/Qwen2.5-VL-7B-for-VideoCls"
MAXIMUM_VIDEOS = 1
HF_OVERRIDES = {
"text_config": {
"architectures": ["Qwen2_5_VLForSequenceClassification"],
},
}
HF_OVERRIDES = {"architectures": ["Qwen2_5_VLForSequenceClassification"]}
input_text = "This product was excellent and exceeded my expectations"
image_url = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/cat_snow.jpg"
image_base64 = {"url": encode_image_url(fetch_image(image_url))}
tests/entrypoints/pooling/score/test_utils.py
-36
View File
@@ -4,13 +4,10 @@
from unittest.mock import patch
import pytest
import torch
from vllm.config import ModelConfig
from vllm.entrypoints.chat_utils import ChatTemplateResolutionError
from vllm.entrypoints.pooling.score.utils import (
compute_maxsim_score,
compute_maxsim_scores,
get_score_prompt,
)
from vllm.inputs import TokensPrompt
@@ -354,36 +351,3 @@ class TestGetScorePrompt:
assert_prompt_tokenization_consistent(
cross_encoder_tokenizer, full_prompt, engine_prompt
)
def test_compute_maxsim_scores_matches_reference_per_pair() -> None:
generator = torch.Generator()
generator.manual_seed(7)
shared_query = torch.randn(5, 8, generator=generator)
q_embs = [
shared_query, # 1:N style shared query
shared_query,
torch.randn(2, 8, generator=generator),
torch.randn(4, 8, generator=generator),
]
d_embs = [
torch.randn(6, 8, generator=generator),
torch.randn(3, 8, generator=generator),
torch.randn(5, 8, generator=generator),
torch.randn(7, 8, generator=generator),
]
batched_scores = compute_maxsim_scores(
q_embs,
d_embs,
max_batch_size=4,
max_score_matrix_elements=40, # batch shrinking path.
)
reference_scores = [
compute_maxsim_score(q, d).to("cpu") for q, d in zip(q_embs, d_embs)
]
assert len(batched_scores) == len(reference_scores)
for batched, reference in zip(batched_scores, reference_scores):
torch.testing.assert_close(batched, reference, rtol=1e-4, atol=1e-4)
tests/entrypoints/sagemaker/test_sagemaker_lora_adapters.py
+1 -1
View File
@@ -88,7 +88,7 @@ async def test_sagemaker_load_adapter_invalid_files(
basic_server_with_lora.url_for("adapters"),
json={"name": "invalid-adapter", "src": str(invalid_files)},
)
assert load_response.status_code == 400
assert load_response.status_code == 500
@pytest.mark.asyncio
tests/entrypoints/test_api_server_process_manager.py
+7 -15
View File
@@ -79,7 +79,7 @@ def test_api_server_process_manager_init(api_server_args, with_stats_update):
finally:
# Always clean up the processes
print("Cleaning up processes...")
manager.shutdown()
manager.close()
# Give processes time to terminate
time.sleep(0.2)
@@ -111,8 +111,6 @@ def test_wait_for_completion_or_failure(api_server_args):
wait_for_completion_or_failure(api_server_manager=manager)
except Exception as e:
result["exception"] = e
finally:
manager.shutdown()
# Start a thread to run wait_for_completion_or_failure
wait_thread = threading.Thread(target=run_with_exception_capture, daemon=True)
@@ -145,7 +143,7 @@ def test_wait_for_completion_or_failure(api_server_args):
assert not proc.is_alive(), f"Process {i} should not be alive"
finally:
manager.shutdown()
manager.close()
time.sleep(0.2)
@@ -176,14 +174,11 @@ def test_normal_completion(api_server_args):
# since all processes have already
# terminated, it should return immediately
# with no error
try:
wait_for_completion_or_failure(api_server_manager=manager)
finally:
manager.shutdown()
wait_for_completion_or_failure(api_server_manager=manager)
finally:
# Clean up just in case
manager.shutdown()
manager.close()
time.sleep(0.2)
@@ -206,7 +201,7 @@ def test_external_process_monitoring(api_server_args):
def __init__(self, proc):
self.proc = proc
def shutdown(self):
def close(self):
if self.proc.is_alive():
self.proc.terminate()
self.proc.join(timeout=0.5)
@@ -231,9 +226,6 @@ def test_external_process_monitoring(api_server_args):
)
except Exception as e:
result["exception"] = e
finally:
manager.shutdown()
mock_coordinator.shutdown()
# Start a thread to run wait_for_completion_or_failure
wait_thread = threading.Thread(target=run_with_exception_capture, daemon=True)
@@ -267,6 +259,6 @@ def test_external_process_monitoring(api_server_args):
finally:
# Clean up
manager.shutdown()
mock_coordinator.shutdown()
manager.close()
mock_coordinator.close()
time.sleep(0.2)
tests/entrypoints/test_chat_utils.py
+32
View File
@@ -1458,6 +1458,38 @@ def test_parse_chat_messages_context_text_format(
assert mm_uuids is None
def test_parse_chat_messages_openai_format_image_url(
phi3v_model_config,
image_url,
):
content = [
{"type": "image_url", "image_url": {"url": image_url}},
{"type": "text", "text": "What's in the image?"},
]
conversation, mm_data, mm_uuids = parse_chat_messages(
[
{
"role": "user",
"content": content,
}
],
phi3v_model_config,
content_format="openai",
)
assert conversation == [
{
"role": "user",
"content": [
{"type": "image"},
{"type": "text", "text": "What's in the image?"},
],
}
]
_assert_mm_data_is_image_input(mm_data, 1)
_assert_mm_uuids(mm_uuids, 1, expected_uuids=[None])
def test_parse_chat_messages_rejects_too_many_images_in_one_message(
phi3v_model_config,
image_url,
tests/entrypoints/test_grpc_server.py
-428
View File
@@ -1,428 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
End-to-end tests for the vLLM gRPC server.
"""
import asyncio
import socket
import subprocess
import sys
import time
import grpc
import pytest
import pytest_asyncio
from vllm.grpc import vllm_engine_pb2, vllm_engine_pb2_grpc
# Use a small model for fast testing
MODEL_NAME = "hmellor/tiny-random-LlamaForCausalLM"
def find_free_port() -> int:
"""Find a free port on localhost."""
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.bind(("", 0))
s.listen(1)
port = s.getsockname()[1]
return port
async def wait_for_server(port: int, timeout: float = 60.0) -> bool:
"""Wait for the gRPC server to be ready by trying health checks."""
start_time = time.time()
print("waiting for server to start...")
while time.time() - start_time < timeout:
try:
channel = grpc.aio.insecure_channel(f"localhost:{port}")
stub = vllm_engine_pb2_grpc.VllmEngineStub(channel)
request = vllm_engine_pb2.HealthCheckRequest()
response = await stub.HealthCheck(request, timeout=5.0)
await channel.close()
if response.healthy:
print("server returned healthy=True")
return True
except Exception:
await asyncio.sleep(0.5)
return False
class GrpcServerProcess:
"""Manages a gRPC server running in a subprocess."""
def __init__(self):
self.process: subprocess.Popen | None = None
self.port: int | None = None
async def start(self):
"""Start the gRPC server process."""
self.port = find_free_port()
# Start the server as a subprocess
self.process = subprocess.Popen(
[
sys.executable,
"-m",
"vllm.entrypoints.grpc_server",
"--model",
MODEL_NAME,
"--host",
"localhost",
"--port",
str(self.port),
"--max-num-batched-tokens",
"512",
"--disable-log-stats-server",
],
)
# Wait for server to be ready
if not await wait_for_server(self.port):
self.stop()
raise RuntimeError("gRPC server failed to start within timeout")
def stop(self):
"""Stop the gRPC server process."""
if self.process:
self.process.terminate()
try:
self.process.wait(timeout=10)
except subprocess.TimeoutExpired:
self.process.kill()
self.process.wait()
@pytest_asyncio.fixture(scope="module")
async def grpc_server():
"""Fixture providing a running gRPC server in a subprocess."""
server = GrpcServerProcess()
await server.start()
yield server
server.stop()
@pytest_asyncio.fixture
async def grpc_client(grpc_server):
"""Fixture providing a gRPC client connected to the server."""
channel = grpc.aio.insecure_channel(f"localhost:{grpc_server.port}")
stub = vllm_engine_pb2_grpc.VllmEngineStub(channel)
yield stub
await channel.close()
@pytest.mark.asyncio
async def test_health_check(grpc_client):
"""Test the HealthCheck RPC."""
request = vllm_engine_pb2.HealthCheckRequest()
response = await grpc_client.HealthCheck(request)
assert response.healthy is True
assert response.message == "Health"
@pytest.mark.asyncio
async def test_get_model_info(grpc_client):
"""Test the GetModelInfo RPC."""
request = vllm_engine_pb2.GetModelInfoRequest()
response = await grpc_client.GetModelInfo(request)
assert response.model_path == MODEL_NAME
assert response.is_generation is True
assert response.max_context_length > 0
assert response.vocab_size > 0
assert response.supports_vision is False
@pytest.mark.asyncio
async def test_get_server_info(grpc_client):
"""Test the GetServerInfo RPC."""
request = vllm_engine_pb2.GetServerInfoRequest()
response = await grpc_client.GetServerInfo(request)
assert response.active_requests >= 0
assert response.is_paused is False
assert response.uptime_seconds >= 0
assert response.server_type == "vllm-grpc"
assert response.last_receive_timestamp > 0
@pytest.mark.asyncio
async def test_generate_non_streaming(grpc_client):
"""Test the Generate RPC in non-streaming mode."""
# Create a simple request
request = vllm_engine_pb2.GenerateRequest(
request_id="test-non-streaming-1",
tokenized=vllm_engine_pb2.TokenizedInput(
original_text="Hello, my name is",
input_ids=[15496, 11, 616, 1438, 318], # GPT-2 tokens for the prompt
),
sampling_params=vllm_engine_pb2.SamplingParams(
temperature=0.0,
max_tokens=10,
n=1,
),
stream=False,
)
# Collect all responses
responses = []
async for response in grpc_client.Generate(request):
responses.append(response)
# Should have exactly one response (complete)
assert len(responses) == 1
# Check the response
final_response = responses[0]
assert final_response.HasField("complete")
complete = final_response.complete
assert len(complete.output_ids) > 0
assert complete.finish_reason in ["stop", "length"]
assert complete.prompt_tokens > 0
assert complete.completion_tokens > 0
@pytest.mark.asyncio
async def test_generate_streaming(grpc_client):
"""Test the Generate RPC in streaming mode."""
request = vllm_engine_pb2.GenerateRequest(
request_id="test-streaming-1",
tokenized=vllm_engine_pb2.TokenizedInput(
original_text="The capital of France is",
input_ids=[464, 3139, 286, 4881, 318], # GPT-2 tokens
),
sampling_params=vllm_engine_pb2.SamplingParams(
temperature=0.0, max_tokens=10, n=1
),
stream=True,
)
# Collect all responses
chunks = []
complete_response = None
async for response in grpc_client.Generate(request):
if response.HasField("chunk"):
chunks.append(response.chunk)
elif response.HasField("complete"):
complete_response = response.complete
# Should have received some chunks
assert len(chunks) >= 0 # May have 0 chunks if generation is very fast
# Should have a final complete response
assert complete_response is not None
assert complete_response.finish_reason in ["stop", "length"]
assert complete_response.prompt_tokens > 0
# Verify chunk structure
for chunk in chunks:
assert chunk.prompt_tokens > 0
assert chunk.completion_tokens >= 0
@pytest.mark.asyncio
async def test_generate_with_different_sampling_params(grpc_client):
"""Test Generate with various sampling parameters."""
# Test with temperature
request = vllm_engine_pb2.GenerateRequest(
request_id="test-sampling-temp",
tokenized=vllm_engine_pb2.TokenizedInput(
original_text="Hello",
input_ids=[15496],
),
sampling_params=vllm_engine_pb2.SamplingParams(
temperature=0.8, top_p=0.95, max_tokens=5
),
stream=False,
)
responses = [r async for r in grpc_client.Generate(request)]
assert len(responses) == 1
assert responses[0].HasField("complete")
# Test with top_k
request = vllm_engine_pb2.GenerateRequest(
request_id="test-sampling-topk",
tokenized=vllm_engine_pb2.TokenizedInput(
original_text="Hello",
input_ids=[15496],
),
sampling_params=vllm_engine_pb2.SamplingParams(
temperature=1.0, top_k=50, max_tokens=5
),
stream=False,
)
responses = [r async for r in grpc_client.Generate(request)]
assert len(responses) == 1
assert responses[0].HasField("complete")
@pytest.mark.asyncio
async def test_generate_with_stop_strings(grpc_client):
"""Test Generate with stop strings."""
request = vllm_engine_pb2.GenerateRequest(
request_id="test-stop-strings",
tokenized=vllm_engine_pb2.TokenizedInput(
original_text="Hello",
input_ids=[15496],
),
sampling_params=vllm_engine_pb2.SamplingParams(
temperature=0.0,
max_tokens=20,
stop=["\n", "END"],
),
stream=False,
)
responses = [r async for r in grpc_client.Generate(request)]
assert len(responses) == 1
assert responses[0].HasField("complete")
complete = responses[0].complete
assert complete.finish_reason in ["stop", "length"]
@pytest.mark.asyncio
async def test_generate_multiple_requests(grpc_client):
"""Test handling multiple concurrent Generate requests."""
async def make_request(request_id: str):
request = vllm_engine_pb2.GenerateRequest(
request_id=request_id,
tokenized=vllm_engine_pb2.TokenizedInput(
original_text="Hello",
input_ids=[15496],
),
sampling_params=vllm_engine_pb2.SamplingParams(
temperature=0.0, max_tokens=5
),
stream=False,
)
responses = [r async for r in grpc_client.Generate(request)]
return responses[0]
# Send multiple requests concurrently
tasks = [make_request(f"test-concurrent-{i}") for i in range(3)]
responses = await asyncio.gather(*tasks)
# Verify all requests completed successfully
assert len(responses) == 3
for i, response in enumerate(responses):
assert response.HasField("complete")
@pytest.mark.asyncio
async def test_generate_with_seed(grpc_client):
"""Test Generate with a fixed seed for reproducibility."""
def make_request(request_id: str, seed: int):
return vllm_engine_pb2.GenerateRequest(
request_id=request_id,
tokenized=vllm_engine_pb2.TokenizedInput(
original_text="The future of AI is",
input_ids=[464, 2003, 286, 9552, 318],
),
sampling_params=vllm_engine_pb2.SamplingParams(
temperature=1.0, max_tokens=10, seed=seed
),
stream=False,
)
# Make two requests with the same seed
request1 = make_request("test-seed-1", 42)
request2 = make_request("test-seed-2", 42)
response_list1 = [r async for r in grpc_client.Generate(request1)]
response_list2 = [r async for r in grpc_client.Generate(request2)]
# Both should complete successfully
assert len(response_list1) == 1
assert len(response_list2) == 1
assert response_list1[0].HasField("complete")
assert response_list2[0].HasField("complete")
# With the same seed, outputs should be identical
output_ids1 = list(response_list1[0].complete.output_ids)
output_ids2 = list(response_list2[0].complete.output_ids)
assert output_ids1 == output_ids2
@pytest.mark.asyncio
async def test_generate_error_handling(grpc_client):
"""Test error handling in Generate RPC."""
# Request with invalid top_p value (-33)
request = vllm_engine_pb2.GenerateRequest(
request_id="test-error-invalid-topp",
sampling_params=vllm_engine_pb2.SamplingParams(
temperature=0.0, max_tokens=10, top_p=-33
),
stream=False,
)
# Should raise an error response
with pytest.raises(grpc.RpcError) as exc_info:
_ = [r async for r in grpc_client.Generate(request)]
assert exc_info.value.code() == grpc.StatusCode.INVALID_ARGUMENT
assert "top_p must be in (0, 1], got -33.0" in exc_info.value.details()
@pytest.mark.asyncio
async def test_abort_request(grpc_client):
"""Test the out-of-band Abort RPC."""
request_id = "test-abort-1"
# Start a long-running streaming generate request
generate_request = vllm_engine_pb2.GenerateRequest(
request_id=request_id,
tokenized=vllm_engine_pb2.TokenizedInput(
original_text="Hello",
input_ids=[15496],
),
sampling_params=vllm_engine_pb2.SamplingParams(
temperature=0.0,
min_tokens=500,
max_tokens=500, # Request many tokens to ensure it runs long enough
),
stream=True,
)
# Track whether we were aborted
was_aborted = False
received_chunks = 0
async def run_generate():
nonlocal was_aborted, received_chunks
async for response in grpc_client.Generate(generate_request):
if response.HasField("chunk"):
received_chunks += 1
if response.HasField("complete"):
complete = response.complete
was_aborted = complete.finish_reason == "abort"
else:
was_aborted = False
async def abort_after_delay():
# Small delay to ensure generate has started
await asyncio.sleep(0.1)
abort_request = vllm_engine_pb2.AbortRequest(request_ids=[request_id])
await grpc_client.Abort(abort_request)
# Run generate and abort concurrently
await asyncio.gather(run_generate(), abort_after_delay())
# The request should have been aborted (received final chunk with
# "abort" finish reason) and finished early due to the abort.
assert was_aborted and received_chunks < 500, (
"Request should have been aborted before generating all 500 tokens"
)
tests/entrypoints/test_utils.py
+14
View File
@@ -1,6 +1,8 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
from vllm.entrypoints.utils import get_max_tokens, sanitize_message
@@ -80,3 +82,15 @@ class TestGetMaxTokens:
default_sampling_params={"max_tokens": 2048},
)
assert result == 512
def test_input_length_exceeds_max_model_len(self):
with pytest.raises(
ValueError,
match="Input length .* exceeds model's maximum context length .*",
):
get_max_tokens(
max_model_len=100,
max_tokens=50,
input_length=150,
default_sampling_params={"max_tokens": 2048},
)
tests/kernels/attention/test_cache.py
+5 -2
View File
@@ -23,7 +23,7 @@ CACHE_LAYOUTS = ["NHD", "HND"]
KV_SCALE_TYPES = ["tensor", "attn_head"]
# Parameters for MLA tests.
KV_LORA_RANKS = [512]
KV_LORA_RANKS = [256, 512]
QK_ROPE_HEAD_DIMS = [64]
NUM_TOKENS_MLA = [42]
BLOCK_SIZES_MLA = [16]
@@ -627,6 +627,8 @@ def test_concat_and_cache_ds_mla(
pytest.skip("concat_and_cache_mla doesn't support fp8_ds_mla on ROCm")
if dtype.itemsize != 2:
pytest.skip("ds_mla only supports 16-bit input")
if kv_lora_rank != 512:
pytest.skip("fp8_ds_mla requires kv_lora_rank == 512")
kv_cache_dtype = "fp8_ds_mla"
set_random_seed(seed)
torch.set_default_device(device)
@@ -663,7 +665,8 @@ def test_concat_and_cache_ds_mla(
ref_cache_32bit = ref_cache_slice.view(torch.float32)
kv_c_data = kv_c[i]
for tile_idx in range(4):
num_tiles = kv_lora_rank // 128
for tile_idx in range(num_tiles):
tile_start = tile_idx * 128
tile_end = (tile_idx + 1) * 128
tile_data[:] = kv_c_data[tile_start:tile_end]
tests/kernels/attention/test_xpu_mla_sparse.py
+118
View File
@@ -0,0 +1,118 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import torch
from vllm.v1.attention.ops.xpu_mla_sparse import triton_bf16_mla_sparse_interface
# https://github.com/deepseek-ai/FlashMLA/blob/main/tests/ref.py#L7
def _merge_two_lse(
lse0: torch.Tensor, lse1: torch.Tensor | None, s_q: int, h_q: int
) -> torch.Tensor:
if lse1 is None:
return lse0
else:
return torch.logsumexp(
torch.stack([lse0.view(s_q, h_q), lse1.broadcast_to(s_q, h_q)], dim=0),
dim=0,
)
# Adapted from https://github.com/deepseek-ai/FlashMLA/blob/main/tests/ref.py#L19
def reference_mla_sparse_prefill(
q: torch.Tensor,
kv: torch.Tensor,
indices: torch.Tensor,
sm_scale: float,
d_v: int,
topk_length: torch.Tensor | None = None,
attn_sink: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
"""
Returns:
- o: [s_q, h_q, dv]
- o_fp32: [s_q, h_q, dv]
- max_logits: [s_q, h_q]
- lse: [s_q, h_q]
"""
s_q, h_q, d_qk = q.shape
s_kv, _, _ = kv.shape
_, _, topk = indices.shape
indices = indices.clone().squeeze(1)
if topk_length is not None:
mask = torch.arange(topk, device=topk_length.device).unsqueeze(0).broadcast_to(
s_q, topk
) >= topk_length.unsqueeze(1) # [s_q, topk]
indices[mask] = -1
invalid_mask = (indices < 0) | (indices >= s_kv) # [s_q, topk]
indices[invalid_mask] = 0
q = q.float()
gathered_kv = (
kv.index_select(dim=0, index=indices.flatten()).reshape(s_q, topk, d_qk).float()
) # [s_q, topk, d_qk]
P = q @ gathered_kv.transpose(1, 2) # [s_q, h_q, topk]
P *= sm_scale
P[invalid_mask.unsqueeze(1).broadcast_to(P.shape)] = float("-inf")
orig_lse = torch.logsumexp(P, dim=-1) # [s_q, h_q]
max_logits = P.max(dim=-1).values # [s_q, h_q]
lse_for_o = _merge_two_lse(orig_lse, attn_sink, s_q, h_q)
if not torch.is_inference_mode_enabled():
lse_for_o = lse_for_o.clone()
lse_for_o[lse_for_o == float("-inf")] = float(
"+inf"
) # So that corresponding O will be 0
s_for_o = torch.exp(P - lse_for_o.unsqueeze(-1))
out = s_for_o @ gathered_kv[..., :d_v] # [s_q, h_q, dv]
lonely_q_mask = orig_lse == float("-inf") # [s_q, h_q]
orig_lse[lonely_q_mask] = float("+inf")
return (out.to(kv.dtype), out, max_logits, orig_lse)
@pytest.mark.parametrize("device_str", ["xpu"])
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
@pytest.mark.skipif(
not torch.xpu.is_available(),
reason="XPU is required",
)
def test_bf16_triton_sparse_mla(device_str, dtype):
device = torch.device(device_str)
s_q = 1
s_kv = 256
h_q = 64 # kernel expects multiple of 64
h_kv = 1
d_qk = 576
d_v = 512
topk = 128
torch.random.manual_seed(1234)
q = torch.randn((s_q, h_q, d_qk), dtype=dtype, device=device)
kv = torch.randn((s_kv, h_kv, d_qk), dtype=dtype, device=device)
indices = torch.full((s_q, h_kv, topk), -1, dtype=torch.int32, device=device)
for t in range(s_q):
for h in range(h_kv):
i_i = torch.randperm(max(1, t))[:topk]
indices[t, h, : len(i_i)] = i_i
sm_scale = d_qk**-0.5
out, max_logits, lse = triton_bf16_mla_sparse_interface(
q, kv, indices, sm_scale, d_v
)
assert out.shape == (s_q, h_q, d_v)
assert max_logits.shape == (s_q, h_q)
assert lse.shape == (s_q, h_q)
ref_out, ref_out_fp32, ref_max_logits, ref_lse = reference_mla_sparse_prefill(
q, kv, indices, sm_scale, d_v
)
assert torch.allclose(out, ref_out, atol=1e-2, rtol=1e-2)
assert torch.allclose(max_logits, ref_max_logits, atol=1e-3, rtol=1e-3)
assert torch.allclose(lse, ref_lse, atol=1e-3, rtol=1e-3)
tests/kernels/core/test_fused_quant_layernorm.py
+51 -13
View File
@@ -162,6 +162,7 @@ def ops_impl(
)
@pytest.mark.parametrize("seed", SEEDS)
@pytest.mark.parametrize("device", CUDA_DEVICES)
@pytest.mark.parametrize("strided_input", [False, True])
@torch.inference_mode()
def test_rms_norm(
default_vllm_config,
@@ -175,6 +176,7 @@ def test_rms_norm(
tma_alignment: int,
seed: int,
device: str,
strided_input: bool,
) -> None:
torch.random.manual_seed(seed)
if torch.cuda.is_available():
@@ -184,17 +186,17 @@ def test_rms_norm(
if group_size is not None and hidden_size % group_size[1] != 0:
# skip
return
pytest.skip("Skip non-divisible group sizes")
if group_size is not None and has_scale_ub:
# blockwise baseline doesn't support scale_ub
return
pytest.skip("scale_ub not supported for blockwise/group quantization")
if (
group_size is None or quant_dtype != current_platform.fp8_dtype()
) and tma_alignment != 0:
# TMA alignment is only supported for groupwise fp8 kernels
return
pytest.skip("tma alignment not supported for per-token or int8 quantization")
if (
group_size is not None
@@ -202,21 +204,36 @@ def test_rms_norm(
and hidden_size // group_size[1] % tma_alignment == 0
):
# Skip tests where TMA alignment doesn't create extra padding to save time
return
pytest.skip("Skip TMA alignment cases where no extra padding is added")
if has_scale_ub and quant_dtype != current_platform.fp8_dtype():
# skip
return
pytest.skip("scale_ub only supported for fp8 quantization")
layer = RMSNorm(hidden_size, EPS).to(dtype=dtype)
# Make weights
layer.weight.data.normal_(mean=1.0, std=0.1)
# Make inputs
# Make inputs: use a wider tensor and slice to create a non-contiguous
# (strided) input when strided_input=True. The last dimension stride
# remains 1, which the kernel requires.
scale = 1 / (hidden_size)
x = torch.randn(num_tokens, hidden_size, dtype=dtype) * scale
residual = torch.randn_like(x) * scale if add_residual else None
last_dim = 2 * hidden_size if strided_input else hidden_size
x = torch.randn(num_tokens, last_dim, dtype=dtype) * scale
x = x[:, :hidden_size]
# dim 1 gets special-cased
x_is_strided = strided_input and num_tokens != 1
# check that the input is strided iff we expect it to be
assert x.is_contiguous() != x_is_strided
# Residual must still be contiguous
residual = (
torch.randn(num_tokens, hidden_size, dtype=dtype) * scale
if add_residual
else None
)
if has_scale_ub:
rms_x, _ = ref_rms_norm(layer, x, residual)
scale_ub = torch.mean(rms_x).to(dtype=torch.float32, device="cuda")
@@ -260,12 +277,33 @@ def test_rms_norm(
if add_residual:
assert torch.allclose(ref_residual, ops_residual)
output = torch.empty_like(x, dtype=quant_dtype)
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
)
opcheck(
torch.ops._C.rms_norm_dynamic_per_token_quant,
(output, x, layer.weight, scales, 1e-5, scale_ub, residual),
)
if group_size is None:
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
opcheck(
torch.ops._C.rms_norm_per_block_quant,
(
output,
x,
layer.weight,
scales,
1e-5,
scale_ub,
residual,
group_size[1],
True, # is_scale_transposed
),
)
tests/kernels/helion/test_config_manager.py
+38 -33
View File
@@ -160,10 +160,11 @@ class TestConfigManager:
"""Test getting config file path for a kernel."""
manager = ConfigManager(base_dir="/tmp")
file_path = manager.get_config_file_path("silu_mul_fp8")
dir_path = manager.get_config_file_path("silu_mul_fp8")
assert dir_path == Path("/tmp/silu_mul_fp8")
expected_path = Path("/tmp/silu_mul_fp8.json")
assert file_path == expected_path
file_path = manager.get_config_file_path("silu_mul_fp8", "nvidia_h100")
assert file_path == Path("/tmp/silu_mul_fp8/nvidia_h100.json")
def test_ensure_base_dir_exists(self):
"""Test ensuring base directory exists."""
@@ -189,19 +190,19 @@ class TestConfigManager:
assert config_set.get_platforms() == []
def test_load_config_set_valid_file(self):
"""Test loading config set from valid file."""
"""Test loading config set from per-platform files."""
with tempfile.TemporaryDirectory() as temp_dir:
# Use realistic config data
kernel_config = {
"block_sizes": [128, 64],
"num_warps": 8,
"num_stages": 6,
"pid_type": "persistent_interleaved",
}
config_data = {"h100": {"batch_32_hidden_4096": kernel_config}}
config_file = Path(temp_dir) / "test_kernel.json"
with open(config_file, "w") as f:
json.dump(config_data, f)
kernel_dir = Path(temp_dir) / "test_kernel"
kernel_dir.mkdir()
platform_file = kernel_dir / "h100.json"
with open(platform_file, "w") as f:
json.dump({"batch_32_hidden_4096": kernel_config}, f)
manager = ConfigManager(base_dir=temp_dir)
config_set = manager.load_config_set("test_kernel")
@@ -210,7 +211,6 @@ class TestConfigManager:
assert config_set.kernel_name == "test_kernel"
assert config_set.get_platforms() == ["h100"]
# Verify the config was loaded correctly
config = config_set.get_config("h100", "batch_32_hidden_4096")
assert isinstance(config, helion.Config)
assert config.block_sizes == [128, 64]
@@ -219,7 +219,9 @@ class TestConfigManager:
def test_load_config_set_invalid_json(self):
"""Test loading config set from file with invalid JSON."""
with tempfile.TemporaryDirectory() as temp_dir:
config_file = Path(temp_dir) / "test_kernel.json"
kernel_dir = Path(temp_dir) / "test_kernel"
kernel_dir.mkdir()
config_file = kernel_dir / "h100.json"
with open(config_file, "w") as f:
f.write("invalid json content {")
@@ -231,9 +233,8 @@ class TestConfigManager:
assert config_set.get_platforms() == []
def test_save_config_set(self):
"""Test saving ConfigSet to file."""
"""Test saving ConfigSet to per-platform files."""
with tempfile.TemporaryDirectory() as temp_dir:
# Use realistic config data
kernel_config = {
"block_sizes": [256, 128],
"num_warps": 16,
@@ -246,31 +247,34 @@ class TestConfigManager:
manager = ConfigManager(base_dir=temp_dir)
saved_path = manager.save_config_set(config_set)
expected_path = Path(temp_dir) / "test_kernel.json"
assert saved_path == expected_path
assert saved_path.exists()
expected_dir = Path(temp_dir) / "test_kernel"
assert saved_path == expected_dir
assert saved_path.is_dir()
with open(saved_path) as f:
platform_file = expected_dir / "h100.json"
assert platform_file.exists()
with open(platform_file) as f:
loaded_data = json.load(f)
assert loaded_data == data
assert loaded_data == data["h100"]
def test_save_config_set_creates_directory(self):
"""Test that save_config_set creates parent directories if needed."""
with tempfile.TemporaryDirectory() as temp_dir:
nested_dir = Path(temp_dir) / "nested" / "configs"
config_set = ConfigSet("test_kernel")
data = {"h100": {"default": {"num_warps": 4}}}
config_set = ConfigSet.from_dict("test_kernel", data)
manager = ConfigManager(base_dir=nested_dir)
saved_path = manager.save_config_set(config_set)
assert nested_dir.exists()
assert nested_dir.is_dir()
assert saved_path.exists()
assert saved_path.is_dir()
assert (saved_path / "h100.json").exists()
def test_get_platform_configs(self):
"""Test getting all configs for a specific platform."""
with tempfile.TemporaryDirectory() as temp_dir:
# Use realistic config data
config_1 = {"num_warps": 4, "num_stages": 3, "block_sizes": [64, 32]}
config_2 = {"num_warps": 8, "num_stages": 5, "block_sizes": [128, 64]}
default_config = {
@@ -280,17 +284,19 @@ class TestConfigManager:
}
config_3 = {"num_warps": 2, "num_stages": 2, "block_sizes": [32, 16]}
config_data = {
"h100": {
"batch_32_hidden_4096": config_1,
"batch_64_hidden_2048": config_2,
"default": default_config,
},
"a100": {"batch_16_hidden_1024": config_3},
}
config_file = Path(temp_dir) / "test_kernel.json"
with open(config_file, "w") as f:
json.dump(config_data, f)
kernel_dir = Path(temp_dir) / "test_kernel"
kernel_dir.mkdir()
with open(kernel_dir / "h100.json", "w") as f:
json.dump(
{
"batch_32_hidden_4096": config_1,
"batch_64_hidden_2048": config_2,
"default": default_config,
},
f,
)
with open(kernel_dir / "a100.json", "w") as f:
json.dump({"batch_16_hidden_1024": config_3}, f)
manager = ConfigManager(base_dir=temp_dir)
@@ -302,7 +308,6 @@ class TestConfigManager:
for config in h100_configs.values():
assert isinstance(config, helion.Config)
# Verify specific config details
assert h100_configs["batch_32_hidden_4096"].num_warps == 4
assert h100_configs["default"].num_stages == 7
tests/kernels/helion/test_register.py
+4 -57
View File
@@ -134,14 +134,14 @@ class TestValidateHelionSettings:
validate_helion_settings(settings, "test_kernel")
def test_warns_on_static_shapes_true(self):
"""Test that static_shapes=True emits a warning."""
"""Test that static_shapes=True emits a warning about being overridden."""
settings = helion.Settings()
settings.static_shapes = True
with patch("vllm.kernels.helion.register.logger") as mock_logger:
validate_helion_settings(settings, "test_kernel")
mock_logger.warning.assert_called_once()
assert "static_shapes=True" in mock_logger.warning.call_args[0][0]
assert "overridden to False" in mock_logger.warning.call_args[0][0]
def create_configured_kernel_with_configs(
@@ -259,7 +259,6 @@ class TestConfiguredHelionKernel:
settings = helion.Settings()
settings.print_output_code = True
# Note: helion.Settings() defaults static_shapes to True
mock_config_manager = Mock(spec=ConfigManager)
mock_config_manager.get_platform_configs = Mock(return_value=sample_configs)
@@ -288,46 +287,8 @@ class TestConfiguredHelionKernel:
call_kwargs = mock_kernel.call_args[1]
assert "print_output_code" in call_kwargs
assert call_kwargs["print_output_code"] is True
# helion.Settings() defaults to static_shapes=True, so it should remain True
assert call_kwargs["static_shapes"] is True
def test_create_decorated_kernel_preserves_static_shapes_true(
self, sample_kernel, sample_configs
):
"""Test that explicit static_shapes=True is preserved."""
def default_picker(args, config_keys):
return "default"
settings = helion.Settings()
settings.static_shapes = True
mock_config_manager = Mock(spec=ConfigManager)
mock_config_manager.get_platform_configs = Mock(return_value=sample_configs)
with (
patch("vllm.kernels.helion.register.helion.kernel") as mock_kernel,
patch(
"vllm.kernels.helion.config_manager.ConfigManager.get_instance",
return_value=mock_config_manager,
),
patch(
"vllm.kernels.helion.utils.get_canonical_gpu_name",
return_value="nvidia_h200",
),
):
mock_decorated = Mock()
mock_kernel.return_value = Mock(return_value=mock_decorated)
ConfiguredHelionKernel(
op_name="test_kernel",
config_picker=default_picker,
raw_kernel_func=sample_kernel,
helion_settings=settings,
)
call_kwargs = mock_kernel.call_args[1]
assert call_kwargs["static_shapes"] is True
# static_shapes is always forced to False by vLLM
assert call_kwargs["static_shapes"] is False
def test_key_and_config_selector_use_same_logic(
self, sample_kernel, sample_configs
@@ -761,20 +722,6 @@ class TestKernelRegistry:
def test_kernel(x):
return x
def test_register_kernel_warns_with_static_shapes_true(self):
"""Test register_kernel warns when static_shapes=True."""
mock_settings = Mock()
mock_settings.to_dict.return_value = {"static_shapes": True}
with patch("vllm.kernels.helion.register.logger") as mock_logger:
@register_kernel("test", helion_settings=mock_settings)
def test_kernel(x):
return x
mock_logger.warning.assert_called_once()
assert "static_shapes=True" in mock_logger.warning.call_args[0][0]
def test_register_kernel_no_warning_with_static_shapes_false(self):
"""Test register_kernel doesn't warn with static_shapes=False."""
mock_settings = Mock()
tests/kernels/quantization/test_rocm_skinny_gemms.py
+7 -4
View File
@@ -70,7 +70,6 @@ N_FACTORS_WVSPLITKRC = [
117,
128,
]
K_FACTORS_WVSPLITKRC = [2880, 2880 + 8, 3072, 3072 + 8]
M_FACTORS_WVSPLITKRC = [128, 128 + 16, 256, 256 + 16, 640, 640 + 16]
@@ -123,10 +122,11 @@ def pad_fp8(weight):
@pytest.mark.parametrize("m", M_FACTORS_WVSPLITKRC)
@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize("seed", SEEDS)
@pytest.mark.parametrize("padded_a", [False, True])
@pytest.mark.parametrize("bias_mode", BIAS_MODES)
@pytest.mark.skipif(not current_platform.is_rocm(), reason="only test for rocm")
@pytest.mark.skipif(not on_gfx950(), reason="only meant for gfx950")
def test_rocm_wvsplitkrc_kernel(xnorm, n, k, m, dtype, seed, bias_mode):
def test_rocm_wvsplitkrc_kernel(xnorm, n, k, m, dtype, seed, padded_a, bias_mode):
torch.manual_seed(seed)
cu_count = num_compute_units()
@@ -141,7 +141,8 @@ def test_rocm_wvsplitkrc_kernel(xnorm, n, k, m, dtype, seed, bias_mode):
# Given the above, how many CUs would we need?
CuNeeded = rndup_cus * GrpsShrB
# candidate for atomic reduce count splitk?
fits_wvsplitkrc = CuNeeded <= cu_count
fits_wvsplitkrc = (N_p2 * m * ((k + 512 - 1) // 512)) <= 128 * 1024 * 12
fits_wvsplitkrc &= CuNeeded <= cu_count
if not fits_wvsplitkrc:
pytest.skip("Too large for wvSplitKrc")
@@ -151,6 +152,8 @@ def test_rocm_wvsplitkrc_kernel(xnorm, n, k, m, dtype, seed, bias_mode):
) # normalize to avoid large output-bias deltas
A = (torch.rand(n, k, dtype=dtype, device="cuda") * 2 - 1) * xavier
B = (torch.rand(m, k, dtype=dtype, device="cuda") * 2 - 1) * xavier
if padded_a:
A = pad_fp8(A)
BIAS = None
if bias_mode == 1:
@@ -159,7 +162,7 @@ def test_rocm_wvsplitkrc_kernel(xnorm, n, k, m, dtype, seed, bias_mode):
BIAS = torch.rand(n, m, dtype=dtype, device="cuda") * 2 - 1
ref_out = torch.nn.functional.linear(A, B, BIAS)
out = ops.wvSplitKrc(B, A.view(-1, A.size(-1)), cu_count, BIAS)
out = ops.wvSplitKrc(A, B, cu_count, BIAS)
if xnorm:
torch.testing.assert_close(out, ref_out, atol=1e-3, rtol=1e-8)
tests/models/language/generation/test_common.py
+12
View File
@@ -3,6 +3,8 @@
import pytest
import torch
from packaging.version import Version
from transformers import __version__ as TRANSFORMERS_VERSION
from vllm.platforms import current_platform
@@ -151,6 +153,16 @@ def test_models(
if prompt_embeds is not None:
embed = hf_model.model.get_input_embeddings()(token_ids)
if "gemma" in model.lower() and (
Version(TRANSFORMERS_VERSION) < Version("5.3.0.dev0")
):
# For Gemma 1/2 models with Transformers 5.4.0+, the prompt
# embeddings are normalised in `get_prompt_embeddings`,
# like Gemma 3. For older versions, we need to manually normalise.
embed_scale = hf_model.config.hidden_size**0.5
normalizer = torch.tensor(embed_scale, dtype=embed.dtype)
embed *= normalizer
# MiniCPM models apply scale_emb to embeddings internally.
# vLLM expects pre-scaled embeddings when using inputs_embeds.
if model in EMBED_SCALING_MODELS:
tests/models/language/pooling/test_classification.py
+3 -1
View File
@@ -45,5 +45,7 @@ def test_models(
# half datatype tests in
# tests/models/language/pooling/test_embedding.py
assert torch.allclose(
hf_output, vllm_output, 1e-3 if dtype == "float" else 1e-2
hf_output,
vllm_output,
rtol=2e-3 if dtype == "float" else 1e-2,
)
tests/models/language/pooling/test_mm_classifier_conversion.py
+2 -1
View File
@@ -32,7 +32,8 @@ def test_idefics_multimodal(
def update_config(config):
config.text_config.update(
text_config = config.get_text_config()
text_config.update(
{
"architectures": ["Gemma3ForSequenceClassification"],
"classifier_from_token": ["A", "B", "C", "D", "E"],
tests/models/multimodal/generation/vlm_utils/core.py
+2
View File
@@ -74,6 +74,8 @@ def run_test(
if model_info.require_embed_inputs:
for k in ("skip_tokenizer_init", "enable_prompt_embeds", "enable_mm_embeds"):
vllm_runner_kwargs_[k] = model_info.require_embed_inputs
if not model_info.enable_prefix_caching:
vllm_runner_kwargs_["enable_prefix_caching"] = False
if vllm_runner_kwargs:
vllm_runner_kwargs_.update(vllm_runner_kwargs)
tests/models/multimodal/processing/test_common.py
+20 -59
View File
@@ -6,9 +6,6 @@ from functools import partial
import numpy as np
import pytest
from mistral_common.protocol.instruct.chunk import ImageChunk, TextChunk
from mistral_common.protocol.instruct.messages import UserMessage
from mistral_common.protocol.instruct.request import ChatCompletionRequest
from PIL import Image
from vllm.config import ModelConfig
@@ -21,7 +18,10 @@ from vllm.config.multimodal import (
from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalDataDict
from vllm.multimodal.cache import MultiModalProcessorOnlyCache
from vllm.multimodal.inputs import MultiModalInputs, batched_tensors_equal
from vllm.multimodal.processing import BaseMultiModalProcessor, InputProcessingContext
from vllm.multimodal.processing import (
BaseMultiModalProcessor,
InputProcessingContext,
)
from vllm.tokenizers import TokenizerLike, cached_tokenizer_from_config
from vllm.utils.mistral import is_mistral_tokenizer
@@ -74,20 +74,6 @@ def glmasr_patch_mm_data(mm_data: MultiModalDataDict) -> MultiModalDataDict:
return mm_data
# For some multimodal models, tokenizer will always add bos_token
# at the beginning of prompt by default, causing hf_processor outputs
# incorrect token ids. So we need use `add_special_tokens=False` here
# to leave bos_token to be added by the processor.
_ADD_SPECIAL_TOKENS_OVERRIDES = {
"lfm2_vl": False,
"nemotron_parse": False,
"ovis": False,
"ovis2_5": False,
"paligemma": False,
"ultravox": False,
"whisper": False,
}
_IGNORE_MM_KEYS = {
# In Ultravox, the audio_features can be different depending on padding
# The slight difference should not be a problem though, since
@@ -152,59 +138,34 @@ def get_text_token_prompts(
parsed_data = processor.info.parse_mm_data(mm_data)
mm_counts = {k: len(vs) for k, vs in parsed_data.items()}
text_prompt: str | None
token_prompt: list[int]
if is_mistral_tokenizer(tokenizer):
# ChatCompletionRequest only supports ImageChunk natively;
# for other modalities (e.g. audio), fall back to the model's
# own dummy inputs builder which knows the right placeholders.
has_non_image = any(
k != "image" and count > 0 for k, count in mm_counts.items()
inputs = dummy_inputs.get_dummy_processor_inputs(
model_config.max_model_len,
mm_counts,
mm_options={},
# Assume all Mistral models define this extra argument
mm_data=mm_data, # type: ignore[call-arg]
)
if has_non_image:
inputs = dummy_inputs.get_dummy_processor_inputs(
model_config.max_model_len,
mm_counts,
mm_options={},
)
text_prompt = None
token_prompt = (
inputs.prompt
if isinstance(inputs.prompt, list)
else tokenizer.encode(inputs.prompt, add_special_tokens=False)
)
else:
images = parsed_data.get("image", [])
request = ChatCompletionRequest(
messages=[
UserMessage(
content=[
TextChunk(text=""),
*(ImageChunk(image=image) for image in images),
]
),
]
)
res = tokenizer.mistral.encode_chat_completion(request)
# Mistral does not support decode_tokens with
# skip_special_tokens=False
text_prompt = None
token_prompt = res.tokens
else:
inputs = dummy_inputs.get_dummy_processor_inputs(
model_config.max_model_len,
mm_counts,
mm_options={},
)
assert isinstance(inputs.prompt, str)
text_prompt: str | None
token_prompt: list[int]
if isinstance(inputs.prompt, list):
text_prompt = None
token_prompt = inputs.prompt
elif isinstance(inputs.prompt, str):
text_prompt = inputs.prompt
token_prompt = tokenizer.encode(
text_prompt,
add_special_tokens=_ADD_SPECIAL_TOKENS_OVERRIDES.get(model_type, True),
**processor.info.get_default_tok_params().get_encode_kwargs(),
)
else:
raise TypeError(type(inputs.prompt))
return text_prompt, token_prompt
@@ -444,7 +405,7 @@ def test_processing_correctness(
)
if model_id == "mistralai/Voxtral-Mini-4B-Realtime-2602":
pytest.skip(
"Voxtral Realtime doesn't make use of any place-holder"
"Voxtral Realtime doesn't make use of any place-holder "
"tokens and hence cannot pass the processing "
"correctness test as is. Let's revisit adapting this "
"test once more realtime models exist."
tests/models/multimodal/processing/test_deepseek_ocr.py
+134
View File
@@ -0,0 +1,134 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Regression test for DeepSeek-OCR TensorSchema validation with empty images_crop.
When using the Gundam preset (BASE_SIZE=1024, IMAGE_SIZE=640, CROP_MODE=True),
images that are small enough to not require cropping produce an empty
images_crop tensor with shape (0, 3, 640, 640). The _parse_and_validate_image_input
method must correctly read image_size from this tensor's shape rather than
falling back to base_size, which would cause a TensorSchema mismatch.
Run with:
pytest tests/models/multimodal/processing/test_deepseek_ocr.py -v
"""
import pytest
from PIL import Image
from transformers import AutoTokenizer
from vllm.model_executor.models.deepseek_ocr import DeepseekOCRImagePixelInputs
from vllm.transformers_utils.processors.deepseek_ocr import DeepseekOCRProcessor
MODEL_ID = "deepseek-ai/DeepSeek-OCR"
@pytest.fixture(scope="module")
def processor():
"""Load the DeepseekOCRProcessor with tokenizer from HuggingFace."""
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
return DeepseekOCRProcessor(tokenizer=tokenizer)
class TestDeepseekOCREmptyImagesCrop:
"""Verify TensorSchema validation handles empty images_crop correctly."""
def test_empty_images_crop_small_image(self, processor):
"""A small image (<=640px) produces empty images_crop and should
not crash the TensorSchema validation.
Previously, the code used ``numel() > 0`` to decide whether to read
image_size from the tensor shape. When numel()==0, it fell back to
base_size=1024, mismatching the actual tensor dim of 640.
"""
# Small image: both dims <= IMAGE_SIZE (640) → no crops
small_image = Image.new("RGB", (100, 100), color="red")
result = processor(
prompt="<image>\nDescribe this image.",
images=[small_image],
)
pixel_values = result["pixel_values"]
images_crop = result["images_crop"]
images_spatial_crop = result["images_spatial_crop"]
# Processor must produce an empty crop tensor for a small image
assert images_crop.shape[0] == 0
base_size = pixel_values.shape[-1]
image_size = images_crop.shape[-1] if images_crop is not None else base_size
# This should NOT raise ValueError
schema = DeepseekOCRImagePixelInputs(
type="pixel_values",
data=pixel_values,
images_crop=images_crop,
images_spatial_crop=images_spatial_crop,
resolve_bindings={
"base_size": base_size,
"image_size": image_size,
},
)
assert schema.data.shape == (1, 3, 1024, 1024)
assert schema.images_crop.shape == (0, 3, 640, 640)
def test_populated_images_crop_large_image(self, processor):
"""A large image (>640px) produces populated images_crop."""
# Large image: exceeds IMAGE_SIZE (640) → dynamic crop tiles
large_image = Image.new("RGB", (1200, 800), color="blue")
result = processor(
prompt="<image>\nDescribe this image.",
images=[large_image],
)
pixel_values = result["pixel_values"]
images_crop = result["images_crop"]
images_spatial_crop = result["images_spatial_crop"]
assert images_crop.shape[0] > 0
base_size = pixel_values.shape[-1]
image_size = images_crop.shape[-1]
schema = DeepseekOCRImagePixelInputs(
type="pixel_values",
data=pixel_values,
images_crop=images_crop,
images_spatial_crop=images_spatial_crop,
resolve_bindings={
"base_size": base_size,
"image_size": image_size,
},
)
assert schema.data.shape == (1, 3, 1024, 1024)
assert schema.images_crop.shape[-1] == 640
def test_mismatched_image_size_raises(self, processor):
"""Deliberately wrong image_size binding should still be caught
by TensorSchema validation."""
small_image = Image.new("RGB", (100, 100), color="green")
result = processor(
prompt="<image>\nDescribe this image.",
images=[small_image],
)
pixel_values = result["pixel_values"]
images_crop = result["images_crop"]
images_spatial_crop = result["images_spatial_crop"]
with pytest.raises(ValueError, match="images_crop"):
DeepseekOCRImagePixelInputs(
type="pixel_values",
data=pixel_values,
images_crop=images_crop,
images_spatial_crop=images_spatial_crop,
resolve_bindings={
"base_size": 1024,
"image_size": 1024, # Wrong! Tensor has 640
},
)
tests/models/multimodal/processing/test_qwen3_vl.py
+94
View File
@@ -0,0 +1,94 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Regression tests for Qwen3-VL processor.
Covers the fix for num_frames-based timestamp calculation
(issue vllm-project/vllm#35909).
"""
from typing import Any
import numpy as np
import pytest
from vllm.multimodal import MULTIMODAL_REGISTRY
from ...utils import build_model_context
MODEL_ID = "Qwen/Qwen3-VL-4B-Instruct"
def _build_video_mm_data(
num_frames: int,
width: int = 128,
height: int = 128,
original_fps: float = 30.0,
) -> dict[str, Any]:
"""Create synthetic video data with metadata indicating that
HF processor should re-sample frames (do_sample_frames=True).
``total_num_frames`` is set equal to the ndarray frame count so
that HF's ``sample_frames`` indices stay within bounds of the
actual tensor that is passed."""
video = np.zeros((num_frames, height, width, 3), dtype=np.uint8)
metadata = {
"fps": original_fps,
"duration": num_frames / original_fps,
"total_num_frames": num_frames,
"frames_indices": list(range(num_frames)),
"video_backend": "opencv",
"do_sample_frames": True,
}
return {"video": [(video, metadata)]}
@pytest.mark.parametrize("model_id", [MODEL_ID])
@pytest.mark.parametrize(
"num_frames",
[8, 16],
)
def test_processor_num_frames_timestamp(
model_id: str,
num_frames: int,
) -> None:
"""Regression test: using ``num_frames`` (without ``fps``) must not
cause a timestamp / token-count mismatch.
Before the fix, ``_get_video_second_idx`` ignored the explicit
``num_frames`` and fell back to an fps-based calculation, which
produced a different number of timestamp entries and ultimately led
to shape mismatches in downstream token construction.
We deliberately choose ``num_frames`` values (8, 16) that differ
from what the default fps-based path would compute (which clamps
to ``min_frames=4`` for a short video at 30 fps), so this test
would fail without the fix.
"""
ctx = build_model_context(
model_id,
limit_mm_per_prompt={"image": 0, "video": 1},
)
processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
prompt = "<|vision_start|><|video_pad|><|vision_end|>"
mm_data = _build_video_mm_data(num_frames=num_frames)
# Process with explicit num_frames (no fps) -- this is the path
# that was broken before the fix.
hf_mm_kwargs: dict[str, Any] = {"num_frames": num_frames}
processed = processor(
prompt,
mm_items=processor.info.parse_mm_data(mm_data),
hf_processor_mm_kwargs=hf_mm_kwargs,
)
# Basic sanity: the processor must produce video tokens.
token_ids = processed["prompt_token_ids"]
assert len(token_ids) > 0, "Processor produced empty token list"
# Verify that video placeholders were actually inserted.
assert "mm_placeholders" in processed
video_phs = processed["mm_placeholders"].get("video", [])
assert len(video_phs) == 1, (
f"Expected exactly 1 video placeholder, got {len(video_phs)}"
)
tests/models/multimodal/test_mapping.py
+9 -6
View File
@@ -31,12 +31,6 @@ def create_dummy_model(repo: str, model_arch: str) -> PreTrainedModel:
config = AutoConfig.from_pretrained(repo)
with torch.device("meta"):
model = model_cls._from_config(config)
# TODO(hmellor): Remove this once Transformers has fixed tied weights on meta device
# https://github.com/huggingface/transformers/issues/43522
if getattr(config.get_text_config(), "tie_word_embeddings", False) or getattr(
config, "tie_word_embeddings", False
):
model.tie_weights()
return model
@@ -103,6 +97,15 @@ def test_hf_model_weights_mapper(model_arch: str):
# Some checkpoints may have buffers, we ignore them for this test
ref_weight_names -= buffer_names
# Some checkpoints include tied weights (e.g. lm_head tied to embed_tokens) in the
# safetensors file. In Transformers v5, named_parameters() will not include them
# after they are tied in the model, so the mapper will not be able to map them.
# We exclude them from the reference weight names for this test.
if isinstance(tied := getattr(hf_dummy_model, "_tied_weights_keys", None), dict):
mapped_tied_weights = mapper.apply((k, None) for k in tied)
tied_weight_names = set(map(lambda x: x[0], mapped_tied_weights))
ref_weight_names -= tied_weight_names
weights_missing = ref_weight_names - weight_names
weights_unmapped = weight_names - ref_weight_names
assert not weights_missing and not weights_unmapped, (
tests/models/registry.py
+108 -63
View File
@@ -72,6 +72,12 @@ class _HfExamplesInfo:
If False, we will use CUDA graph and eager execution in hybrid.
"""
enable_prefix_caching: bool = True
"""
Whether to enable prefix caching for the model. If True, we will test the model with
prefix caching enabled. If False, we will test the model without prefix caching.
"""
is_available_online: bool = True
"""
Set this to `False` if the name of this architecture no longer exists on
@@ -313,6 +319,10 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
"HunYuanMoEV1ForCausalLM": _HfExamplesInfo(
"tencent/Hunyuan-A13B-Instruct", trust_remote_code=True
),
"HyperCLOVAXForCausalLM": _HfExamplesInfo(
"naver-hyperclovax/HyperCLOVAX-SEED-Think-32B",
trust_remote_code=True,
),
"InternLMForCausalLM": _HfExamplesInfo(
"internlm/internlm-chat-7b", trust_remote_code=True
),
@@ -347,7 +357,11 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
),
"Lfm2ForCausalLM": _HfExamplesInfo("LiquidAI/LFM2-1.2B"),
"Lfm2MoeForCausalLM": _HfExamplesInfo(
"LiquidAI/LFM2-8B-A1B", min_transformers_version="4.58"
"LiquidAI/LFM2-8B-A1B",
min_transformers_version="5.0.0",
use_original_num_layers=True,
# Initialize at least one MoE layer
hf_overrides={"num_hidden_layers": 4},
),
"LlamaForCausalLM": _HfExamplesInfo(
"meta-llama/Llama-3.2-1B-Instruct",
@@ -507,9 +521,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
"stepfun-ai/Step-3.5-Flash",
use_original_num_layers=True,
# Initialize at least one MoE layer
hf_overrides={
"num_hidden_layers": 4,
},
hf_overrides={"num_hidden_layers": 4},
),
"Step3TextForCausalLM": _HfExamplesInfo("stepfun-ai/step3", trust_remote_code=True),
"SolarForCausalLM": _HfExamplesInfo(
@@ -540,15 +552,9 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
_EMBEDDING_EXAMPLE_MODELS = {
# [Text-only]
"BertModel": _HfExamplesInfo("BAAI/bge-base-en-v1.5"),
"HF_ColBERT": _HfExamplesInfo("answerdotai/answerai-colbert-small-v1"),
"ColBERTModernBertModel": _HfExamplesInfo(
"lightonai/GTE-ModernColBERT-v1",
hf_overrides={"architectures": ["ColBERTModernBertModel"]},
),
"ColBERTJinaRobertaModel": _HfExamplesInfo(
"jinaai/jina-colbert-v2",
trust_remote_code=True,
hf_overrides={"architectures": ["ColBERTJinaRobertaModel"]},
"BertSpladeSparseEmbeddingModel": _HfExamplesInfo(
"naver/splade-v3",
hf_overrides={"architectures": ["BertSpladeSparseEmbeddingModel"]},
),
"BgeM3EmbeddingModel": _HfExamplesInfo("BAAI/bge-m3"),
"Gemma2Model": _HfExamplesInfo("BAAI/bge-multilingual-gemma2"),
@@ -562,10 +568,6 @@ _EMBEDDING_EXAMPLE_MODELS = {
trust_remote_code=True,
hf_overrides={"architectures": ["GteNewModel"]},
),
"InternLM2ForRewardModel": _HfExamplesInfo(
"internlm/internlm2-1_8b-reward", trust_remote_code=True
),
"JambaForSequenceClassification": _HfExamplesInfo("ai21labs/Jamba-tiny-reward-dev"),
"LlamaModel": _HfExamplesInfo("llama", is_available_online=False),
"LlamaBidirectionalModel": _HfExamplesInfo(
"nvidia/llama-nemotron-embed-1b-v2", trust_remote_code=True
@@ -578,35 +580,14 @@ _EMBEDDING_EXAMPLE_MODELS = {
"nomic-ai/nomic-embed-text-v2-moe", trust_remote_code=True
),
"Qwen2Model": _HfExamplesInfo("ssmits/Qwen2-7B-Instruct-embed-base"),
"Qwen2ForRewardModel": _HfExamplesInfo(
"Qwen/Qwen2.5-Math-RM-72B",
max_transformers_version="4.53",
transformers_version_reason={
"hf": "HF model uses remote code that is not compatible with latest Transformers" # noqa: E501
},
),
"Qwen2ForProcessRewardModel": _HfExamplesInfo(
"Qwen/Qwen2.5-Math-PRM-7B",
max_transformers_version="4.53",
transformers_version_reason={
"hf": "HF model uses remote code that is not compatible with latest Transformers" # noqa: E501
},
),
"RobertaModel": _HfExamplesInfo("sentence-transformers/stsb-roberta-base-v2"),
"RobertaForMaskedLM": _HfExamplesInfo("sentence-transformers/all-roberta-large-v1"),
"VoyageQwen3BidirectionalEmbedModel": _HfExamplesInfo(
"voyageai/voyage-4-nano", trust_remote_code=True
),
"XLMRobertaModel": _HfExamplesInfo("intfloat/multilingual-e5-small"),
"BertSpladeSparseEmbeddingModel": _HfExamplesInfo(
"naver/splade-v3",
hf_overrides={"architectures": ["BertSpladeSparseEmbeddingModel"]},
),
# [Multimodal]
"CLIPModel": _HfExamplesInfo("openai/clip-vit-base-patch32"),
"ColModernVBertForRetrieval": _HfExamplesInfo(
"ModernVBERT/colmodernvbert-merged",
),
"LlamaNemotronVLModel": _HfExamplesInfo(
"nvidia/llama-nemotron-embed-vl-1b-v2", trust_remote_code=True
),
@@ -615,15 +596,6 @@ _EMBEDDING_EXAMPLE_MODELS = {
"TIGER-Lab/VLM2Vec-Full", trust_remote_code=True
),
"Qwen2VLForConditionalGeneration": _HfExamplesInfo("MrLight/dse-qwen2-2b-mrl-v1"),
"ColQwen3": _HfExamplesInfo(
"TomoroAI/tomoro-colqwen3-embed-4b", trust_remote_code=True
),
"OpsColQwen3Model": _HfExamplesInfo(
"OpenSearch-AI/Ops-Colqwen3-4B", trust_remote_code=True
),
"Qwen3VLNemotronEmbedModel": _HfExamplesInfo(
"nvidia/nemotron-colembed-vl-4b-v2",
),
"SiglipModel": _HfExamplesInfo("google/siglip-base-patch16-224"),
"PrithviGeoSpatialMAE": _HfExamplesInfo(
"ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11",
@@ -643,21 +615,74 @@ _EMBEDDING_EXAMPLE_MODELS = {
),
}
_SEQUENCE_CLASSIFICATION_EXAMPLE_MODELS = {
# [Decoder-only]
"GPT2ForSequenceClassification": _HfExamplesInfo(
"nie3e/sentiment-polish-gpt2-small"
_LATE_INTERACTION_EXAMPLE_MODELS = {
# [Text-only]
"HF_ColBERT": _HfExamplesInfo("answerdotai/answerai-colbert-small-v1"),
"ColBERTModernBertModel": _HfExamplesInfo(
"lightonai/GTE-ModernColBERT-v1",
hf_overrides={"architectures": ["ColBERTModernBertModel"]},
),
# [Cross-encoder]
"ColBERTJinaRobertaModel": _HfExamplesInfo(
"jinaai/jina-colbert-v2",
trust_remote_code=True,
hf_overrides={"architectures": ["ColBERTJinaRobertaModel"]},
),
# [Multimodal]
"ColModernVBertForRetrieval": _HfExamplesInfo(
"ModernVBERT/colmodernvbert-merged",
),
"ColQwen3": _HfExamplesInfo(
"TomoroAI/tomoro-colqwen3-embed-4b", trust_remote_code=True
),
"OpsColQwen3Model": _HfExamplesInfo(
"OpenSearch-AI/Ops-Colqwen3-4B", trust_remote_code=True
),
"Qwen3VLNemotronEmbedModel": _HfExamplesInfo(
"nvidia/nemotron-colembed-vl-4b-v2",
),
}
_REWARD_EXAMPLE_MODELS = {
"InternLM2ForRewardModel": _HfExamplesInfo(
"internlm/internlm2-1_8b-reward", trust_remote_code=True
),
"Qwen2ForRewardModel": _HfExamplesInfo(
"Qwen/Qwen2.5-Math-RM-72B",
max_transformers_version="4.53",
transformers_version_reason={
"hf": "HF model uses remote code that is not compatible with latest Transformers" # noqa: E501
},
),
"Qwen2ForProcessRewardModel": _HfExamplesInfo(
"Qwen/Qwen2.5-Math-PRM-7B",
max_transformers_version="4.53",
transformers_version_reason={
"hf": "HF model uses remote code that is not compatible with latest Transformers" # noqa: E501
},
),
}
_TOKEN_CLASSIFICATION_EXAMPLE_MODELS = {
"BertForTokenClassification": _HfExamplesInfo("boltuix/NeuroBERT-NER"),
"ModernBertForTokenClassification": _HfExamplesInfo(
"disham993/electrical-ner-ModernBERT-base"
),
}
_SEQUENCE_CLASSIFICATION_EXAMPLE_MODELS = {
"BertForSequenceClassification": _HfExamplesInfo(
"cross-encoder/ms-marco-MiniLM-L-6-v2"
),
"BertForTokenClassification": _HfExamplesInfo("boltuix/NeuroBERT-NER"),
"GPT2ForSequenceClassification": _HfExamplesInfo(
"nie3e/sentiment-polish-gpt2-small"
),
"GteNewForSequenceClassification": _HfExamplesInfo(
"Alibaba-NLP/gte-multilingual-reranker-base",
trust_remote_code=True,
hf_overrides={"architectures": ["GteNewForSequenceClassification"]},
),
"JambaForSequenceClassification": _HfExamplesInfo("ai21labs/Jamba-tiny-reward-dev"),
"LlamaBidirectionalForSequenceClassification": _HfExamplesInfo(
"nvidia/llama-nemotron-rerank-1b-v2", trust_remote_code=True
),
@@ -667,9 +692,6 @@ _SEQUENCE_CLASSIFICATION_EXAMPLE_MODELS = {
"ModernBertForSequenceClassification": _HfExamplesInfo(
"Alibaba-NLP/gte-reranker-modernbert-base"
),
"ModernBertForTokenClassification": _HfExamplesInfo(
"disham993/electrical-ner-ModernBERT-base"
),
"RobertaForSequenceClassification": _HfExamplesInfo(
"cross-encoder/quora-roberta-base"
),
@@ -793,6 +815,10 @@ _MULTIMODAL_EXAMPLE_MODELS = {
"naver-hyperclovax/HyperCLOVAX-SEED-Vision-Instruct-3B",
trust_remote_code=True,
),
"HCXVisionV2ForCausalLM": _HfExamplesInfo(
"naver-hyperclovax/HyperCLOVAX-SEED-Think-32B",
trust_remote_code=True,
),
"HunYuanVLForConditionalGeneration": _HfExamplesInfo(
"tencent/HunyuanOCR",
hf_overrides={"num_experts": 0},
@@ -837,6 +863,15 @@ _MULTIMODAL_EXAMPLE_MODELS = {
"Kwai-Keye/Keye-VL-1_5-8B",
trust_remote_code=True,
),
"MoonshotKimiaForCausalLM": _HfExamplesInfo(
"moonshotai/Kimi-Audio-7B-Instruct",
tokenizer_mode="kimi_audio",
trust_remote_code=True,
),
"KimiK25ForConditionalGeneration": _HfExamplesInfo(
"moonshotai/Kimi-K2.5",
trust_remote_code=True,
),
"KimiVLForConditionalGeneration": _HfExamplesInfo(
"moonshotai/Kimi-VL-A3B-Instruct",
extras={"thinking": "moonshotai/Kimi-VL-A3B-Thinking"},
@@ -850,10 +885,6 @@ _MULTIMODAL_EXAMPLE_MODELS = {
)
},
),
"KimiK25ForConditionalGeneration": _HfExamplesInfo(
"moonshotai/Kimi-K2.5",
trust_remote_code=True,
),
"LightOnOCRForConditionalGeneration": _HfExamplesInfo(
"lightonai/LightOnOCR-1B-1025"
),
@@ -1112,6 +1143,18 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
speculative_model="yuhuili/EAGLE-LLaMA3-Instruct-8B",
tokenizer="meta-llama/Meta-Llama-3-8B-Instruct",
),
"Eagle3DeepseekV2ForCausalLM": _HfExamplesInfo(
"moonshotai/Kimi-K2.5",
trust_remote_code=True,
speculative_model="AQ-MedAI/Kimi-K25-eagle3",
tokenizer="moonshotai/Kimi-K2.5",
),
"Eagle3DeepseekV3ForCausalLM": _HfExamplesInfo(
"moonshotai/Kimi-K2.5",
trust_remote_code=True,
speculative_model="AQ-MedAI/Kimi-K25-eagle3",
tokenizer="moonshotai/Kimi-K2.5",
),
"Eagle3LlamaForCausalLM": _HfExamplesInfo(
"meta-llama/Llama-3.1-8B-Instruct",
trust_remote_code=True,
@@ -1169,6 +1212,7 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
"LGAI-EXAONE/K-EXAONE-236B-A23B",
speculative_model="LGAI-EXAONE/K-EXAONE-236B-A23B",
min_transformers_version="5.1.0",
enable_prefix_caching=False,
),
"ExtractHiddenStatesModel": _HfExamplesInfo(
"Qwen/Qwen3-8B",
@@ -1225,9 +1269,7 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
speculative_model="stepfun-ai/Step-3.5-Flash",
use_original_num_layers=True,
# Initialize at least one MoE layer
hf_overrides={
"num_hidden_layers": 4,
},
hf_overrides={"num_hidden_layers": 4},
is_available_online=False,
),
}
@@ -1265,6 +1307,9 @@ _TRANSFORMERS_BACKEND_MODELS = {
_EXAMPLE_MODELS = {
**_TEXT_GENERATION_EXAMPLE_MODELS,
**_EMBEDDING_EXAMPLE_MODELS,
**_LATE_INTERACTION_EXAMPLE_MODELS,
**_REWARD_EXAMPLE_MODELS,
**_TOKEN_CLASSIFICATION_EXAMPLE_MODELS,
**_SEQUENCE_CLASSIFICATION_EXAMPLE_MODELS,
**_MULTIMODAL_EXAMPLE_MODELS,
**_SPECULATIVE_DECODING_EXAMPLE_MODELS,
tests/models/test_initialization.py
+19 -2
View File
@@ -88,15 +88,27 @@ def can_initialize(
[10 * GiB_bytes],
)
scheduler_kv_cache_config = generate_scheduler_kv_cache_config(kv_cache_configs)
vllm_config.cache_config.num_gpu_blocks = scheduler_kv_cache_config.num_blocks
kv_cache_groups = scheduler_kv_cache_config.kv_cache_groups
if kv_cache_groups:
vllm_config.cache_config.block_size = min(
g.kv_cache_spec.block_size for g in kv_cache_groups
)
# gpu_blocks (> 0), cpu_blocks, scheduler_kv_cache_config
return 1, 0, scheduler_kv_cache_config
vllm_config.validate_block_size()
return scheduler_kv_cache_config
if model_arch == "MiniMaxVL01ForConditionalGeneration":
pytest.skip(
"pickle error when loading `transformers.models.auto.CONFIG_MAPPING`"
)
if model_arch == "MoonshotKimiaForCausalLM":
pytest.skip(
"Kimi-Audio requires SpeechToTextConfig "
"which is not configured in test environment"
)
if model_arch in ["DeepseekV32ForCausalLM", "GlmMoeDsaForCausalLM"]:
from vllm.platforms import current_platform
@@ -124,6 +136,10 @@ def can_initialize(
if model_arch == "WhisperForConditionalGeneration":
m.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
kwargs = {}
if not model_info.enable_prefix_caching:
kwargs["enable_prefix_caching"] = False
LLM(
model_info.default,
tokenizer=model_info.tokenizer,
@@ -153,6 +169,7 @@ def can_initialize(
hf_overrides=hf_overrides_fn,
max_num_seqs=model_info.max_num_seqs,
attention_config=attention_config,
**kwargs,
)
tests/models/test_registry.py
+11 -8
View File
@@ -56,21 +56,24 @@ def test_registry_imports(model_arch):
@create_new_process_for_each_test()
@pytest.mark.parametrize(
"model_arch,is_mm,init_cuda,is_ce",
"model_arch,is_mm,init_cuda,score_type",
[
("LlamaForCausalLM", False, False, False),
("LlavaForConditionalGeneration", True, True, False),
("BertForSequenceClassification", False, False, True),
("RobertaForSequenceClassification", False, False, True),
("XLMRobertaForSequenceClassification", False, False, True),
("LlamaForCausalLM", False, False, "bi-encoder"),
("LlavaForConditionalGeneration", True, True, "bi-encoder"),
("BertForSequenceClassification", False, False, "cross-encoder"),
("RobertaForSequenceClassification", False, False, "cross-encoder"),
("XLMRobertaForSequenceClassification", False, False, "cross-encoder"),
("GteNewModel", False, False, "bi-encoder"),
("GteNewForSequenceClassification", False, False, "cross-encoder"),
("HF_ColBERT", False, False, "late-interaction"),
],
)
def test_registry_model_property(model_arch, is_mm, init_cuda, is_ce):
def test_registry_model_property(model_arch, is_mm, init_cuda, score_type):
model_info = ModelRegistry._try_inspect_model_cls(model_arch)
assert model_info is not None
assert model_info.supports_multimodal is is_mm
assert model_info.supports_cross_encoding is is_ce
assert model_info.score_type == score_type
if init_cuda and current_platform.is_cuda_alike():
assert not torch.cuda.is_initialized()
tests/quantization/test_mixed_precision.py
+1
View File
@@ -8,6 +8,7 @@ Run `pytest tests/quantization/test_mixed_precision.py`.
import importlib
import importlib.metadata
import importlib.util
from dataclasses import dataclass
import lm_eval
tests/reasoning/test_nemotron_v3_reasoning_parser.py
+22
View File
@@ -128,6 +128,28 @@ def test_nemotron_v3_without_thinking_returns_content(
assert content == "This is plain content"
def test_nemotron_v3_force_nonempty_content_returns_content(
tokenizer: FakeNemotronTokenizer,
):
parser_cls = ReasoningParserManager.get_reasoning_parser(parser_name)
parser = parser_cls(tokenizer)
request = ChatCompletionRequest(
model="test-model",
messages=[],
chat_template_kwargs={"force_nonempty_content": True},
)
reasoning, content = run_reasoning_extraction(
parser,
["<think>This is plain content"],
request=request,
streaming=False,
)
assert reasoning is None
assert content == "This is plain content"
def test_nemotron_v3_with_thinking_keeps_truncated_reasoning(
tokenizer: FakeNemotronTokenizer,
):
tests/tool_parsers/test_minimax_m2_tool_parser.py
+444
View File
@@ -0,0 +1,444 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import json
import pytest
from vllm.tool_parsers.minimax_m2_tool_parser import (
MinimaxM2ToolParser,
)
pytestmark = pytest.mark.cpu_test
# Token IDs matching FakeTokenizer.vocab
TC_START_ID = 1
TC_END_ID = 2
EOS_ID = 99
class FakeTokenizer:
"""Minimal fake tokenizer for unit tests."""
def __init__(self):
self.model_tokenizer = True
self.vocab = {
"<minimax:tool_call>": TC_START_ID,
"</minimax:tool_call>": TC_END_ID,
}
def get_vocab(self):
return self.vocab
@pytest.fixture
def parser():
return MinimaxM2ToolParser(FakeTokenizer())
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _feed(parser, chunks, request=None):
"""Feed chunks through the streaming parser and collect results.
Each element in *chunks* is either:
- a ``str``: used as delta_text (current_text accumulates automatically)
- a ``(delta_text, delta_token_ids)`` tuple for special-token scenarios
Returns a list of non-None DeltaMessage objects.
"""
previous = ""
results = []
for chunk in chunks:
if isinstance(chunk, tuple):
delta, delta_ids = chunk
else:
delta = chunk
delta_ids = []
current = previous + delta
result = parser.extract_tool_calls_streaming(
previous_text=previous,
current_text=current,
delta_text=delta,
previous_token_ids=[],
current_token_ids=[],
delta_token_ids=delta_ids,
request=request,
)
if result is not None:
results.append(result)
previous = current
return results
def _collect_content(results):
"""Join all content strings from a list of DeltaMessages."""
return "".join(r.content for r in results if r.content)
def _collect_tool_calls(results):
"""Aggregate tool calls by index from a list of DeltaMessages.
Returns a dict: index -> {"id": ..., "name": ..., "arguments": ...}
"""
tool_calls = {}
for r in results:
for tc in r.tool_calls or []:
if tc.index not in tool_calls:
tool_calls[tc.index] = {
"id": None,
"name": "",
"arguments": "",
}
if tc.id:
tool_calls[tc.index]["id"] = tc.id
if tc.function:
if tc.function.name:
tool_calls[tc.index]["name"] += tc.function.name
if tc.function.arguments:
tool_calls[tc.index]["arguments"] += tc.function.arguments
return tool_calls
# ---------------------------------------------------------------------------
# Phase 1: content before tool calls
# ---------------------------------------------------------------------------
class TestContentStreaming:
"""Tests for plain content (no tool calls)."""
def test_plain_content(self, parser):
"""No tool call tokens — all text is streamed as content."""
results = _feed(parser, ["Hello ", "world"])
assert _collect_content(results) == "Hello world"
assert not parser.prev_tool_call_arr
def test_content_before_tool_call(self, parser):
"""Text before <minimax:tool_call> is streamed as content."""
results = _feed(
parser,
[
"Let me check. ",
'<minimax:tool_call><invoke name="get_weather">'
'<parameter name="city">Seattle</parameter>'
"</invoke></minimax:tool_call>",
],
)
assert _collect_content(results) == "Let me check. "
assert len(parser.prev_tool_call_arr) == 1
def test_empty_delta_no_crash(self, parser):
"""Empty delta_text with no token IDs returns None."""
results = _feed(parser, [("", [])])
assert results == []
# ---------------------------------------------------------------------------
# Phase 2: tool call parsing
# ---------------------------------------------------------------------------
class TestSingleInvoke:
"""Tests for a single <invoke> block."""
def test_incremental_chunks(self, parser):
"""Each XML element arrives in a separate chunk."""
results = _feed(
parser,
[
"<minimax:tool_call>",
'<invoke name="get_weather">',
'<parameter name="city">Seattle</parameter>',
"</invoke></minimax:tool_call>",
],
)
tc = _collect_tool_calls(results)
assert len(tc) == 1
assert tc[0]["name"] == "get_weather"
assert json.loads(tc[0]["arguments"]) == {"city": "Seattle"}
assert tc[0]["id"] is not None
def test_single_chunk_complete(self, parser):
"""Entire tool call arrives in one delta."""
results = _feed(
parser,
[
'<minimax:tool_call><invoke name="get_weather">'
'<parameter name="city">Seattle</parameter>'
"</invoke></minimax:tool_call>",
],
)
tc = _collect_tool_calls(results)
assert len(tc) == 1
assert json.loads(tc[0]["arguments"]) == {"city": "Seattle"}
def test_multiple_params(self, parser):
"""Multiple parameters in one invoke."""
results = _feed(
parser,
[
"<minimax:tool_call>",
'<invoke name="get_weather">',
'<parameter name="city">Seattle</parameter>',
'<parameter name="days">5</parameter>',
"</invoke></minimax:tool_call>",
],
)
tc = _collect_tool_calls(results)
assert json.loads(tc[0]["arguments"]) == {
"city": "Seattle",
"days": "5",
}
class TestMultipleInvokes:
"""Tests for multiple <invoke> blocks in one tool call."""
def test_two_invokes_incremental(self, parser):
"""Two invokes arriving one chunk at a time."""
results = _feed(
parser,
[
"<minimax:tool_call>",
'<invoke name="search_web">'
'<parameter name="query">OpenAI</parameter>'
"</invoke>",
'<invoke name="search_web">'
'<parameter name="query">Gemini</parameter>'
"</invoke>",
"</minimax:tool_call>",
],
)
tc = _collect_tool_calls(results)
assert len(tc) == 2
assert tc[0]["name"] == "search_web"
assert tc[1]["name"] == "search_web"
assert json.loads(tc[0]["arguments"]) == {"query": "OpenAI"}
assert json.loads(tc[1]["arguments"]) == {"query": "Gemini"}
def test_two_invokes_in_single_delta(self, parser):
"""Both invokes close in the same delta — loop must emit both."""
results = _feed(
parser,
[
"<minimax:tool_call>",
'<invoke name="fn_a"><parameter name="x">1</parameter></invoke>'
'<invoke name="fn_b"><parameter name="y">2</parameter></invoke>',
"</minimax:tool_call>",
],
)
tc = _collect_tool_calls(results)
assert len(tc) == 2
assert tc[0]["name"] == "fn_a"
assert tc[1]["name"] == "fn_b"
def test_different_functions(self, parser):
"""Parallel calls to different functions."""
results = _feed(
parser,
[
"<minimax:tool_call>",
'<invoke name="get_weather">'
'<parameter name="city">NYC</parameter>'
"</invoke>",
'<invoke name="get_stock">'
'<parameter name="ticker">AAPL</parameter>'
"</invoke>",
"</minimax:tool_call>",
],
)
tc = _collect_tool_calls(results)
assert tc[0]["name"] == "get_weather"
assert tc[1]["name"] == "get_stock"
# ---------------------------------------------------------------------------
# Internal state: prev_tool_call_arr
# ---------------------------------------------------------------------------
class TestInternalState:
"""Verify prev_tool_call_arr is correct."""
def test_prev_tool_call_arr_single(self, parser):
_feed(
parser,
[
'<minimax:tool_call><invoke name="fn">'
'<parameter name="a">1</parameter>'
"</invoke></minimax:tool_call>",
],
)
assert len(parser.prev_tool_call_arr) == 1
assert parser.prev_tool_call_arr[0]["name"] == "fn"
assert parser.prev_tool_call_arr[0]["arguments"] == {"a": "1"}
def test_prev_tool_call_arr_multiple(self, parser):
"""prev_tool_call_arr records each invoke with correct arguments."""
_feed(
parser,
[
"<minimax:tool_call>",
'<invoke name="search"><parameter name="q">hello</parameter></invoke>',
'<invoke name="search"><parameter name="q">world</parameter></invoke>',
"</minimax:tool_call>",
],
)
assert len(parser.prev_tool_call_arr) == 2
assert parser.prev_tool_call_arr[0]["name"] == "search"
assert parser.prev_tool_call_arr[0]["arguments"] == {"q": "hello"}
assert parser.prev_tool_call_arr[1]["name"] == "search"
assert parser.prev_tool_call_arr[1]["arguments"] == {"q": "world"}
# ---------------------------------------------------------------------------
# DeltaMessage structure
# ---------------------------------------------------------------------------
class TestDeltaMessageFormat:
"""Verify the shape of emitted DeltaMessage / DeltaToolCall."""
def test_tool_call_fields(self, parser):
"""Each emitted tool call has id, name, arguments, type, index."""
results = _feed(
parser,
[
'<minimax:tool_call><invoke name="fn">'
'<parameter name="k">v</parameter>'
"</invoke></minimax:tool_call>",
],
)
tc_deltas = [tc for r in results for tc in (r.tool_calls or [])]
assert len(tc_deltas) == 1
tc = tc_deltas[0]
assert tc.index == 0
assert tc.type == "function"
assert tc.id is not None and tc.id.startswith("call_")
assert tc.function.name == "fn"
assert json.loads(tc.function.arguments) == {"k": "v"}
def test_multi_invoke_indices(self, parser):
"""Multiple invokes get sequential indices."""
results = _feed(
parser,
[
"<minimax:tool_call>",
'<invoke name="a"><parameter name="x">1</parameter></invoke>',
'<invoke name="b"><parameter name="x">2</parameter></invoke>',
"</minimax:tool_call>",
],
)
tc_deltas = [tc for r in results for tc in (r.tool_calls or [])]
indices = [tc.index for tc in tc_deltas]
assert indices == [0, 1]
# ---------------------------------------------------------------------------
# Phase 3: EOS handling
# ---------------------------------------------------------------------------
class TestEOSHandling:
"""Tests for the end-of-stream phase."""
def test_eos_after_tool_calls(self, parser):
"""EOS token (empty delta, non-special token id) returns content=''."""
results = _feed(
parser,
[
"<minimax:tool_call>",
'<invoke name="fn"><parameter name="k">v</parameter></invoke>',
"</minimax:tool_call>",
# EOS: empty delta_text, non-special token id
("", [EOS_ID]),
],
)
# Last result should be the EOS empty-content signal
assert results[-1].content == ""
def test_end_token_ignored(self, parser):
"""</minimax:tool_call> special token should NOT trigger EOS."""
results = _feed(
parser,
[
"<minimax:tool_call>",
'<invoke name="fn"><parameter name="k">v</parameter></invoke>',
# </minimax:tool_call> arrives as special token
("", [TC_END_ID]),
],
)
# The tool call delta should be emitted, but no EOS signal
assert not any(r.content == "" and r.tool_calls is None for r in results)
# ---------------------------------------------------------------------------
# Start token detection via token IDs
# ---------------------------------------------------------------------------
class TestSpecialTokenDetection:
"""Start token arrives as a special token (not in delta_text)."""
def test_start_token_via_id(self, parser):
"""<minimax:tool_call> detected via delta_token_ids, not text."""
results = _feed(parser, ["Hello "])
assert _collect_content(results) == "Hello "
# Start token as special token (empty delta_text)
previous = "Hello "
result = parser.extract_tool_calls_streaming(
previous_text=previous,
current_text=previous,
delta_text="",
previous_token_ids=[],
current_token_ids=[],
delta_token_ids=[TC_START_ID],
request=None,
)
assert result is None # no content to emit
assert parser.is_tool_call_started is True
# ---------------------------------------------------------------------------
# Large chunks (stream_interval > 1)
# ---------------------------------------------------------------------------
class TestLargeChunks:
"""Simulate stream_interval > 1 where many tokens arrive at once."""
def test_header_and_params_in_separate_chunks(self, parser):
"""Header in chunk 1, all params + close in chunk 2, then EOS."""
chunk1 = '<minimax:tool_call><invoke name="get_weather">'
chunk2 = (
'<parameter name="city">Seattle</parameter>'
'<parameter name="days">5</parameter>'
"</invoke></minimax:tool_call>"
)
results = _feed(
parser,
[
chunk1,
chunk2,
("", [EOS_ID]),
],
)
tc = _collect_tool_calls(results)
assert len(tc) == 1
parsed = json.loads(tc[0]["arguments"])
assert parsed == {"city": "Seattle", "days": "5"}
assert len(parser.prev_tool_call_arr) == 1
assert parser.prev_tool_call_arr[0]["arguments"] == {
"city": "Seattle",
"days": "5",
}
tests/tool_use/test_minimax_m2_tool_parser.py
-119
View File
@@ -1,119 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import json
import pytest
from vllm.tool_parsers.minimax_m2_tool_parser import (
MinimaxM2ToolParser,
)
pytestmark = pytest.mark.cpu_test
class FakeTokenizer:
"""Minimal fake tokenizer that exposes the attributes used by the
parser: a truthy model_tokenizer marker and a vocab mapping for the
special tokens.
"""
def __init__(self):
self.model_tokenizer = True
# The parser will look up start/end tokens by their literal strings
self.vocab = {
"<minimax:tool_call>": 1,
"</minimax:tool_call>": 2,
}
def get_vocab(self):
return self.vocab
@pytest.fixture
def minimax_m2_tool_parser():
return MinimaxM2ToolParser(FakeTokenizer())
def test_extract_tool_calls_streaming_incremental(minimax_m2_tool_parser):
parser = minimax_m2_tool_parser
parser._reset_streaming_state()
chunks = [
"<minimax:tool_call>",
'<invoke name="get_weather">',
'<parameter name="city">',
"Seattle</parameter>",
"</invoke></minimax:tool_call>",
]
previous = ""
for chunk in chunks:
current = previous + chunk
delta = chunk
parser.extract_tool_calls_streaming(
previous_text=previous,
current_text=current,
delta_text=delta,
previous_token_ids=[],
current_token_ids=[],
delta_token_ids=[],
request=None,
)
previous = current
assert len(parser.prev_tool_call_arr) == 1
entry = parser.prev_tool_call_arr[0]
assert entry["name"] == "get_weather"
args = entry["arguments"]
assert args["city"] == "Seattle"
def test_streaming_minimax_m2_multiple_invokes(minimax_m2_tool_parser):
parser = minimax_m2_tool_parser
parser._reset_streaming_state()
chunks = [
"<minimax:tool_call>",
'<invoke name="search_web">',
'<parameter name="query_tag">',
'["technology", "events"]</parameter>',
'<parameter name="query_list">',
'["OpenAI", "latest", "release"]</parameter>',
"</invoke>",
'<invoke name="search_web">',
'<parameter name="query_tag">',
'["technology", "events"]</parameter>',
'<parameter name="query_list">',
'["Gemini", "latest", "release"]</parameter>',
"</invoke>",
"</minimax:tool_call>",
]
previous = ""
for chunk in chunks:
current = previous + chunk
delta = chunk
parser.extract_tool_calls_streaming(
previous_text=previous,
current_text=current,
delta_text=delta,
previous_token_ids=[],
current_token_ids=[],
delta_token_ids=[],
request=None,
)
previous = current
assert len(parser.prev_tool_call_arr) == 2
for entry, expect_model in zip(parser.prev_tool_call_arr, ["OpenAI", "Gemini"]):
assert entry["name"] == "search_web"
args = json.dumps(entry["arguments"])
assert "technology" in args and "events" in args
assert expect_model in args
# check streamed_args_for_tool for serving_chat.py
for index in range(2):
expected_call = parser.prev_tool_call_arr[index].get("arguments", {})
expected_call = json.dumps(expected_call)
actual_call = parser.streamed_args_for_tool[index]
assert expected_call == actual_call
tests/utils.py
+25 -9
View File
@@ -144,6 +144,17 @@ class RemoteVLLMServer:
"""Subclasses override this method to customize server process launch"""
raise NotImplementedError
def _pre_download_model(self, model: str, args) -> None:
"""Download model weights before starting the server to avoid timeout."""
is_local = os.path.isdir(model)
if not is_local:
engine_args = AsyncEngineArgs.from_cli_args(args)
model_config = engine_args.create_model_config()
load_config = engine_args.create_load_config()
model_loader = get_model_loader(load_config)
model_loader.download_model(model_config)
def __init__(
self,
model: str,
@@ -195,15 +206,7 @@ class RemoteVLLMServer:
getattr(args, "show_hidden_metrics_for_version", None) is not None
)
# download the model before starting the server to avoid timeout
is_local = os.path.isdir(model)
if not is_local:
engine_args = AsyncEngineArgs.from_cli_args(args)
model_config = engine_args.create_model_config()
load_config = engine_args.create_load_config()
model_loader = get_model_loader(load_config)
model_loader.download_model(model_config)
self._pre_download_model(model, args)
# Record GPU memory before server start so we know what
# "released" looks like.
@@ -515,6 +518,19 @@ class RemoteLaunchRenderServer(RemoteVLLMServer):
start_new_session=True,
)
def _pre_download_model(self, model: str, args) -> None:
"""Download only the tokenizer files (no model weights needed)."""
is_local = os.path.isdir(model)
if not is_local:
engine_args = AsyncEngineArgs.from_cli_args(args)
model_config = engine_args.create_model_config()
get_tokenizer(
model_config.tokenizer,
tokenizer_mode=model_config.tokenizer_mode,
trust_remote_code=model_config.trust_remote_code,
revision=model_config.tokenizer_revision,
)
def _wait_for_gpu_memory_release(self, timeout: float = 30.0):
pass # No GPU used
tests/v1/core/test_scheduler.py
+98 -6
View File
@@ -1115,12 +1115,16 @@ def _step_until_done(
all_finished = all_done
def _num_waiting_requests(scheduler: Scheduler) -> int:
return len(scheduler.waiting) + len(scheduler.skipped_waiting)
def _step_until_kv_transfer_finished(scheduler: Scheduler, req_ids: list[str]):
"""Cycle requests through a KV transfer cycle."""
# Requests should first transition to WAITING_FOR_REMOTE_KVS
output = scheduler.schedule()
assert len(scheduler.waiting) == len(req_ids)
assert _num_waiting_requests(scheduler) == len(req_ids)
assert len(scheduler.running) == 0
assert len(output.scheduled_new_reqs) == 0
for req in scheduler.requests.values():
@@ -1139,7 +1143,7 @@ def _step_until_kv_transfer_finished(scheduler: Scheduler, req_ids: list[str]):
# Simulate KV transfer completion using KVConnectorOutput.finished_recving
output = scheduler.schedule()
assert len(scheduler.waiting) == len(req_ids)
assert _num_waiting_requests(scheduler) == len(req_ids)
assert len(scheduler.running) == 0
MODEL_RUNNER_OUTPUT = ModelRunnerOutput(
@@ -1546,7 +1550,7 @@ def test_kv_connector_handles_preemption(is_async, use_ec_connector, ec_role):
# All can be scheduled - 1st token.
output = scheduler.schedule()
if is_async:
assert len(scheduler.waiting) == 2
assert _num_waiting_requests(scheduler) == 2
assert scheduler.running == []
_step_until_kv_transfer_finished(scheduler, req_ids)
output = scheduler.schedule()
@@ -1604,7 +1608,11 @@ def test_kv_connector_handles_preemption(is_async, use_ec_connector, ec_role):
# This will have a local and remote cache hit.
output = scheduler.schedule()
if is_async:
waiting_req_ids = [req.request_id for req in scheduler.waiting]
waiting_req_ids = [
req.request_id
for req in scheduler.skipped_waiting
if req.status == RequestStatus.WAITING_FOR_REMOTE_KVS
]
assert len(waiting_req_ids) == 1
_step_until_kv_transfer_finished(scheduler, waiting_req_ids)
output = scheduler.schedule()
@@ -2439,7 +2447,8 @@ def test_schedule_skip_tokenizer_init_structured_output_request():
output = scheduler.schedule()
assert len(output.scheduled_new_reqs) == 0
assert len(scheduler.running) == 0
assert len(scheduler.waiting) == 1
assert len(scheduler.waiting) == 0
assert len(scheduler.skipped_waiting) == 1
@pytest.mark.parametrize(
@@ -3626,6 +3635,9 @@ def test_prepend_skipped_requests_order():
# simulate first 2 waiting requests are waiting for remote KVs
for req in expected_waiting_reqs[:2]:
req.status = RequestStatus.WAITING_FOR_REMOTE_KVS
scheduler.waiting.remove_requests(expected_waiting_reqs[:2])
for req in expected_waiting_reqs[:2]:
scheduler.skipped_waiting.add_request(req)
# schedule step
# expect the first 2 waiting to be skipped, the third running,
@@ -3636,7 +3648,87 @@ def test_prepend_skipped_requests_order():
expected_waiting_reqs.pop(2)
# verify waiting order is preserved
assert list(scheduler.waiting) == expected_waiting_reqs
waiting_reqs = list(scheduler.skipped_waiting) + list(scheduler.waiting)
assert waiting_reqs == expected_waiting_reqs
def test_remote_kv_promotion_keeps_fcfs_with_fsm_prefix():
scheduler = create_scheduler(max_num_seqs=1)
scheduler.connector = Mock()
scheduler.connector.get_num_new_matched_tokens.return_value = (0, False)
requests = create_requests(num_requests=4)
for request in requests:
scheduler.add_request(request)
req_fsm_1, req_fsm_2, req_remote, req_tail = list(scheduler.waiting)
# simulate two FSM requests at the waiting head that become ready now.
req_fsm_1.status = RequestStatus.WAITING_FOR_FSM
req_fsm_1.structured_output_request = Mock(grammar=object())
req_fsm_2.status = RequestStatus.WAITING_FOR_FSM
req_fsm_2.structured_output_request = Mock(grammar=object())
# simulate a remote-KV request that is ready to be promoted now.
req_remote.status = RequestStatus.WAITING_FOR_REMOTE_KVS
scheduler.waiting.remove_requests([req_fsm_1, req_fsm_2, req_remote])
scheduler.skipped_waiting.add_request(req_fsm_1)
scheduler.skipped_waiting.add_request(req_fsm_2)
scheduler.skipped_waiting.add_request(req_remote)
scheduler.finished_recving_kv_req_ids.add(req_remote.request_id)
scheduler._update_waiting_for_remote_kv = Mock()
output = scheduler.schedule()
assert output.scheduled_new_reqs
assert output.scheduled_new_reqs[0].req_id == req_fsm_1.request_id
waiting_req_ids = [
req.request_id
for req in list(scheduler.skipped_waiting) + list(scheduler.waiting)
]
assert waiting_req_ids == [
req_fsm_2.request_id,
req_remote.request_id,
req_tail.request_id,
]
def test_fcfs_mixed_skipped_waiting_types_keep_order():
scheduler = create_scheduler(max_num_batched_tokens=20)
scheduler._update_waiting_for_remote_kv = Mock()
mk_req = lambda req_id, num_tokens=1: create_requests( # noqa: E731
num_requests=1, num_tokens=num_tokens, req_ids=[req_id]
)[0]
req_fsm, req_remote, req_stream = mk_req("fsm"), mk_req("remote"), mk_req("stream")
req_regular, req_tail = mk_req("regular", 20), mk_req("tail")
req_fsm.status = RequestStatus.WAITING_FOR_FSM
req_fsm.structured_output_request = Mock(grammar=None)
req_remote.status = RequestStatus.WAITING_FOR_REMOTE_KVS
req_stream.status = RequestStatus.WAITING_FOR_STREAMING_REQ
for req in (req_fsm, req_remote, req_stream, req_regular, req_tail):
scheduler.add_request(req)
scheduler.schedule()
assert list(scheduler.skipped_waiting) == [req_fsm, req_remote, req_stream]
scheduler.finish_requests(req_regular.request_id, RequestStatus.FINISHED_ABORTED)
assert not scheduler.running
req_fsm.structured_output_request = Mock(grammar=object())
scheduler.finished_recving_kv_req_ids.add(req_remote.request_id)
req_stream.status = RequestStatus.WAITING
second_output = scheduler.schedule()
expected_order = [
req_fsm.request_id,
req_remote.request_id,
req_stream.request_id,
req_tail.request_id,
]
assert [req.req_id for req in second_output.scheduled_new_reqs] == expected_order
assert [req.request_id for req in scheduler.running] == expected_order
scheduler._update_waiting_for_remote_kv.assert_called_once_with(req_remote)
def test_abort_request_waiting_for_remote_kvs():
tests/v1/e2e/test_async_scheduling.py
+55 -46
View File
@@ -1,5 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import os
from itertools import repeat
from typing import Any
@@ -19,6 +20,8 @@ from ...models.utils import check_outputs_equal
MODEL = "Qwen/Qwen3-0.6B"
MTP_MODEL = "meta-llama/Llama-3.2-1B-Instruct"
# Need to enforce eager for MRV2 while we sort out cudagraph issues.
ENFORCE_EAGER = os.getenv("ENFORCE_EAGER", "0") == "1"
first_prompt = (
"The following numbers of the sequence "
@@ -47,10 +50,10 @@ def test_without_spec_decoding(
test_sampling_params: list[dict[str, Any]] = [
dict(),
# dict(min_tokens=20),
dict(presence_penalty=-1.0),
dict(frequency_penalty=-1.0),
dict(bad_words=["the", " the"]),
dict(logprobs=2),
dict(logprobs=2, presence_penalty=-1.0),
dict(logprobs=2, frequency_penalty=-1.0),
dict(structured_outputs=struct_outputs),
dict(
structured_outputs=struct_outputs,
@@ -58,12 +61,12 @@ def test_without_spec_decoding(
),
dict(
structured_outputs=struct_outputs,
presence_penalty=-1.0,
frequency_penalty=-1.0,
),
dict(
structured_outputs=struct_outputs,
logprobs=2,
presence_penalty=-1.0,
frequency_penalty=-1.0,
),
]
@@ -116,15 +119,15 @@ def test_with_eagle3_spec_decoding(sample_json_schema, monkeypatch: pytest.Monke
test_sampling_params = [
dict(),
dict(presence_penalty=-1.0),
dict(frequency_penalty=-1.0),
dict(bad_words=["the", " the"]),
dict(logprobs=2),
dict(logprobs=2, presence_penalty=-1.0),
dict(logprobs=2, frequency_penalty=-1.0),
dict(structured_outputs=struct_outputs),
dict(
structured_outputs=struct_outputs,
logprobs=2,
presence_penalty=-1.0,
frequency_penalty=-1.0,
),
]
@@ -144,14 +147,7 @@ def test_with_eagle3_spec_decoding(sample_json_schema, monkeypatch: pytest.Monke
(True, "uni", True, spec_config_short, True),
]
# On ROCm, use TRITON_ATTN + float32 for better numerical consistency
run_tests(
monkeypatch,
MTP_MODEL,
test_configs,
test_sampling_params,
is_testing_with_spec_decoding=True,
)
run_tests(monkeypatch, MTP_MODEL, test_configs, test_sampling_params)
def test_with_ngram_gpu_spec_decoding(monkeypatch: pytest.MonkeyPatch):
@@ -196,12 +192,11 @@ def run_tests(
model: str,
test_configs: list[tuple],
test_sampling_params: list[dict[str, Any]],
is_testing_with_spec_decoding: bool = False,
):
"""Test consistency of combos of async scheduling, preemption,
uni/multiproc executor with spec decoding."""
# Determine attention config based on platform
# Flex attention supports float32.
attention_config = {"backend": "FLEX_ATTENTION"}
with monkeypatch.context() as m:
@@ -226,7 +221,6 @@ def run_tests(
async_scheduling,
spec_config,
test_prefill_chunking=test_prefill_chunking,
is_testing_with_spec_decoding=is_testing_with_spec_decoding,
attention_config=attention_config,
)
outputs.append(test_results)
@@ -250,6 +244,7 @@ def run_tests(
test_acceptance_rates or repeat(None),
test_sampling_params,
):
reason = None
try:
check_outputs_equal(
outputs_0_lst=base_outs,
@@ -257,42 +252,57 @@ def run_tests(
name_0=f"baseline=[{baseline_config}], params={params}",
name_1=f"config=[{test_config}], params={params}",
)
except AssertionError as e:
reason = "outputs ", e
assert _all_logprobs_match(base_logprobs, test_logprobs)
if reason is None:
try:
assert _all_logprobs_match(base_logprobs, test_logprobs)
except AssertionError as e:
reason = "logprobs", e
if (
base_acceptance_rate is not None
and test_acceptance_rate is not None
):
if "spec_mml=None" in test_config:
# Preemption causes more variance in acceptance rates
if (
current_platform.is_rocm()
and "preemption=True" in test_config
):
tolerance = 0.10
if reason is None:
try:
if (
base_acceptance_rate is not None
and test_acceptance_rate is not None
):
if "spec_mml=None" in test_config:
# Preemption causes more variance in acceptance rates
if (
current_platform.is_rocm()
and "preemption=True" in test_config
):
tolerance = 0.10
else:
tolerance = 0.05
assert (
test_acceptance_rate > base_acceptance_rate
or test_acceptance_rate
== pytest.approx(base_acceptance_rate, rel=tolerance)
)
else:
tolerance = 0.05
assert (
test_acceptance_rate > base_acceptance_rate
or test_acceptance_rate
== pytest.approx(base_acceptance_rate, rel=tolerance)
)
else:
# Currently the reported acceptance rate is expected to be
# lower when we sometimes skip drafting altogether.
assert test_acceptance_rate > 0.1
# Currently the reported acceptance rate is expected to be
# lower when we sometimes skip drafting altogether.
assert test_acceptance_rate > 0.1
except AssertionError as e:
reason = "accept ", e
if reason is None:
print(
f"PASSED: config=[{test_config}], params={params}"
f"\033[32mPASSED\033[0m: "
f"config=[{test_config}], params={params}"
f" accept_rate={test_acceptance_rate}"
)
except AssertionError as e:
else:
reason_str, _ = reason
print(
f"FAILED: config=[{test_config}], params={params}"
f"\033[31mFAILED\033[0m({reason_str}): "
f"config=[{test_config}], params={params}"
f" accept_rate={test_acceptance_rate}"
)
if failure is None:
failure = e
_, failure = reason
if failure is not None:
raise failure
@@ -307,7 +317,6 @@ def run_test(
async_scheduling: bool,
spec_config: dict[str, Any] | None,
test_prefill_chunking: bool,
is_testing_with_spec_decoding: bool = False,
attention_config: dict[str, Any] | None = None,
):
spec_decoding = spec_config is not None
@@ -335,7 +344,7 @@ def run_test(
enable_chunked_prefill=test_prefill_chunking,
# Force prefill chunking
max_num_batched_tokens=48 if test_prefill_chunking else None,
# enforce_eager=True,
enforce_eager=ENFORCE_EAGER,
async_scheduling=async_scheduling,
distributed_executor_backend=executor,
dtype="float32",
tests/v1/engine/test_engine_core_client.py
+71
View File
@@ -24,17 +24,23 @@ from vllm import SamplingParams
from vllm.distributed.kv_events import BlockStored, KVEventBatch, ZmqEventPublisher
from vllm.engine.arg_utils import EngineArgs
from vllm.platforms import current_platform
from vllm.pooling_params import LateInteractionParams, PoolingParams
from vllm.usage.usage_lib import UsageContext
from vllm.utils.torch_utils import set_default_torch_num_threads
from vllm.v1.engine import EngineCoreRequest
from vllm.v1.engine.core import EngineCore
from vllm.v1.engine.core_client import (
AsyncMPClient,
DPLBAsyncMPClient,
EngineCoreClient,
SyncMPClient,
)
from vllm.v1.engine.utils import CoreEngineProcManager
from vllm.v1.executor.abstract import Executor
from vllm.v1.pool.late_interaction import (
LATE_INTERACTION_MODE_CACHE_QUERY,
LATE_INTERACTION_MODE_SCORE_DOC,
)
from ...distributed.conftest import MockSubscriber
from ...utils import create_new_process_for_each_test
@@ -164,6 +170,71 @@ def test_mp_client_uses_env_timeout(monkeypatch: pytest.MonkeyPatch):
client.shutdown()
def _make_pooling_request(
request_id: str, *, mode: str | None = None, query_key: str | None = None
) -> EngineCoreRequest:
late_interaction_params = None
if mode is not None and query_key is not None:
late_interaction_params = LateInteractionParams(
mode=mode,
query_key=query_key,
)
return EngineCoreRequest(
request_id=request_id,
prompt_token_ids=[1, 2, 3],
mm_features=None,
sampling_params=None,
pooling_params=PoolingParams(
task="token_embed",
late_interaction_params=late_interaction_params,
),
arrival_time=time.time(),
lora_request=None,
cache_salt=None,
data_parallel_rank=None,
)
def test_dplb_late_interaction_sticky_routing():
client = object.__new__(DPLBAsyncMPClient)
client.client_count = 1
client.reqs_in_flight = {}
client.core_engines = [b"\x00\x00", b"\x01\x00", b"\x02\x00"]
client.lb_engines = [[0, 0], [0, 0], [0, 0]]
client.eng_start_index = 0
query_key = "rerank-abc-query-0"
query_request = _make_pooling_request(
"query-req", mode=LATE_INTERACTION_MODE_CACHE_QUERY, query_key=query_key
)
doc_request = _make_pooling_request(
"doc-req", mode=LATE_INTERACTION_MODE_SCORE_DOC, query_key=query_key
)
query_engine = client.get_core_engine_for_request(query_request)
doc_engine = client.get_core_engine_for_request(doc_request)
assert query_engine == doc_engine
assert client.reqs_in_flight["query-req"] == query_engine
assert client.reqs_in_flight["doc-req"] == doc_engine
def test_dplb_non_late_interaction_still_uses_lb():
client = object.__new__(DPLBAsyncMPClient)
client.client_count = 1
client.reqs_in_flight = {}
client.core_engines = [b"\x00\x00", b"\x01\x00", b"\x02\x00"]
client.lb_engines = [[2, 1], [0, 0], [1, 0]]
client.eng_start_index = 0
request = make_request(SamplingParams(max_tokens=1))
chosen_engine = client.get_core_engine_for_request(request)
assert chosen_engine == client.core_engines[1]
assert client.lb_engines[1][0] == 1
def loop_until_done(client: EngineCoreClient, outputs: dict):
while True:
engine_core_outputs = client.get_output().outputs
tests/v1/entrypoints/openai/serving_responses/test_function_call.py
+105
View File
@@ -197,3 +197,108 @@ async def test_named_tool_use(client: openai.AsyncOpenAI):
response_2 = await client.responses.create(model=MODEL_NAME, input=input_messages)
# check the output
assert len(response_2.output_text) > 0
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
async def test_function_calling_with_streaming_expected_arguments(
client: openai.AsyncOpenAI, model_name: str
):
tools = [
{
"type": "function",
"name": "get_weather",
"description": "Get current temperature for provided location in celsius.",
"parameters": {
"type": "object",
"properties": {
"location": {"type": "string"},
},
"required": ["location"],
"additionalProperties": False,
},
"strict": True,
}
]
stream_response = await client.responses.create(
model=model_name,
input="Can you tell me what the current weather is in Berlin?",
tools=tools,
stream=True,
)
tool_call_item = None
completed_event = None
async for event in stream_response:
if (
event.type == "response.output_item.added"
and event.item.type == "function_call"
):
tool_call_item = event.item
elif event.type == "response.function_call_arguments.delta" and tool_call_item:
tool_call_item.arguments += event.delta
elif (
event.type == "response.output_item.done"
and event.item.type == "function_call"
):
completed_event = event
assert tool_call_item is not None
assert tool_call_item.type == "function_call"
assert tool_call_item.name == "get_weather"
assert completed_event is not None
assert tool_call_item.arguments == completed_event.item.arguments
assert tool_call_item.name == completed_event.item.name
args = json.loads(tool_call_item.arguments)
assert "location" in args
assert args["location"] is not None
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
async def test_function_calling_with_streaming_types(
client: openai.AsyncOpenAI, model_name: str
):
# this links the "done" type with the "start" type
# so every "done" type should have a corresponding "start" type
# and every open block should be closed by the end of the stream
pairs_of_event_types = {
"response.completed": "response.created",
"response.output_item.done": "response.output_item.added",
"response.output_text.done": "response.output_text.delta",
"response.content_part.done": "response.content_part.added",
"response.reasoning_text.done": "response.reasoning_text.delta",
"response.reasoning_part.done": "response.reasoning_part.added",
"response.function_call_arguments.done": "response.function_call_arguments.delta", # noqa
}
input_list = [
{
"role": "user",
"content": "Can you tell me what the current weather is in Berlin?",
}
]
stream_response = await client.responses.create(
model=model_name,
input=input_list,
tools=tools,
stream=True,
)
stack_of_event_types = []
async for event in stream_response:
if event.type == "response.created":
stack_of_event_types.append(event.type)
elif event.type == "response.completed":
assert stack_of_event_types[-1] == pairs_of_event_types[event.type]
stack_of_event_types.pop()
if event.type.endswith("added"):
stack_of_event_types.append(event.type)
elif event.type.endswith("delta"):
if stack_of_event_types[-1] == event.type:
continue
stack_of_event_types.append(event.type)
elif event.type.endswith("done"):
assert stack_of_event_types[-1] == pairs_of_event_types[event.type]
stack_of_event_types.pop()
assert len(stack_of_event_types) == 0
tests/v1/kv_connector/unit/test_error_propagation.py
+2 -1
View File
@@ -119,7 +119,7 @@ def test_error_propagation_async_load(fail_scheduler: Scheduler):
scheduler_output = fail_scheduler.schedule()
assert len(fail_scheduler.waiting) == 1
assert len(fail_scheduler.skipped_waiting) == 1
assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS
assert request.num_computed_tokens == num_external_computed_tokens
@@ -145,3 +145,4 @@ def test_error_propagation_async_load(fail_scheduler: Scheduler):
assert output.finish_reason == FinishReason.ERROR
assert len(fail_scheduler.waiting) == 0
assert len(fail_scheduler.skipped_waiting) == 0
tests/v1/kv_connector/unit/test_flexkv_connector.py
+232
View File
@@ -0,0 +1,232 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for FlexKVConnectorV1.
These tests mock the ``flexkv`` package so they can run without a real FlexKV
installation. They verify:
1. That ``FlexKVConnectorV1`` raises a helpful ``ImportError`` when FlexKV is
not installed.
2. That all public methods are correctly delegated to the underlying
``FlexKVConnectorV1Impl``.
"""
import sys
import types
from unittest.mock import MagicMock, patch
import pytest
import torch
from vllm.config import KVTransferConfig, VllmConfig
from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorRole
from vllm.v1.kv_cache_interface import KVCacheConfig
from .utils import create_vllm_config
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _make_vllm_config(
kv_connector: str = "FlexKVConnectorV1",
kv_role: str = "kv_both",
) -> VllmConfig:
"""Return a minimal VllmConfig with a KVTransferConfig attached."""
vllm_config = create_vllm_config(block_size=16, max_num_batched_tokens=512)
vllm_config.kv_transfer_config = KVTransferConfig(
kv_connector=kv_connector,
kv_role=kv_role,
)
return vllm_config
def _make_kv_cache_config() -> KVCacheConfig:
return MagicMock(spec=KVCacheConfig)
def _make_flexkv_module(
impl_mock: MagicMock,
) -> tuple[types.ModuleType, types.ModuleType]:
"""Build a fake ``flexkv`` package hierarchy that returns *impl_mock*
when ``FlexKVConnectorV1Impl`` is instantiated."""
flexkv_mod = types.ModuleType("flexkv")
integration_mod = types.ModuleType("flexkv.integration")
vllm_mod = types.ModuleType("flexkv.integration.vllm")
adapter_mod = types.ModuleType("flexkv.integration.vllm.vllm_v1_adapter")
# Make FlexKVConnectorV1Impl() return our mock instance.
# The "# type: ignore" markers below are needed because ModuleType does
# not declare these attributes statically; they are set dynamically.
FlexKVConnectorV1ImplCls = MagicMock(return_value=impl_mock)
adapter_mod.FlexKVConnectorV1Impl = FlexKVConnectorV1ImplCls # type: ignore
flexkv_mod.integration = integration_mod # type: ignore
integration_mod.vllm = vllm_mod # type: ignore
vllm_mod.vllm_v1_adapter = adapter_mod # type: ignore
return flexkv_mod, adapter_mod
def _install_flexkv_mock(impl_mock: MagicMock):
"""Insert fake flexkv modules into sys.modules and return a context that
cleans them up afterwards."""
flexkv_mod, adapter_mod = _make_flexkv_module(impl_mock)
mods = {
"flexkv": flexkv_mod,
"flexkv.integration": flexkv_mod.integration,
"flexkv.integration.vllm": flexkv_mod.integration.vllm,
"flexkv.integration.vllm.vllm_v1_adapter": adapter_mod,
}
return patch.dict(sys.modules, mods)
def _build_connector(vllm_config: VllmConfig, impl_mock: MagicMock):
"""Instantiate FlexKVConnectorV1 with faked flexkv modules."""
from vllm.distributed.kv_transfer.kv_connector.v1.flexkv_connector import (
FlexKVConnectorV1,
)
with _install_flexkv_mock(impl_mock):
connector = FlexKVConnectorV1(
vllm_config=vllm_config,
role=KVConnectorRole.WORKER,
kv_cache_config=_make_kv_cache_config(),
)
return connector
# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------
class TestFlexKVConnectorImportError:
"""FlexKVConnectorV1 should fail with a helpful message when flexkv is
absent."""
def test_import_error_message(self):
from vllm.distributed.kv_transfer.kv_connector.v1.flexkv_connector import (
FlexKVConnectorV1,
)
# Ensure flexkv is NOT in sys.modules
for key in list(sys.modules):
if key.startswith("flexkv"):
del sys.modules[key]
with pytest.raises(ImportError, match="(?i)flexkv") as exc_info:
FlexKVConnectorV1(
vllm_config=_make_vllm_config(),
role=KVConnectorRole.WORKER,
kv_cache_config=_make_kv_cache_config(),
)
assert "https://github.com/taco-project/FlexKV" in str(exc_info.value)
class TestFlexKVConnectorDelegation:
"""All public API methods should be forwarded to the impl."""
@pytest.fixture()
def connector_and_impl(self):
impl = MagicMock()
cfg = _make_vllm_config()
connector = _build_connector(cfg, impl)
return connector, impl
def test_shutdown(self, connector_and_impl):
connector, impl = connector_and_impl
connector.shutdown()
impl.shutdown.assert_called_once()
def test_start_load_kv(self, connector_and_impl):
connector, impl = connector_and_impl
ctx = MagicMock()
connector.start_load_kv(ctx, extra_arg="x")
impl.start_load_kv.assert_called_once_with(ctx, extra_arg="x")
def test_save_kv_layer(self, connector_and_impl):
connector, impl = connector_and_impl
kv_layer = torch.zeros(4, 4)
attn_meta = MagicMock()
connector.save_kv_layer("layer_0", kv_layer, attn_meta)
impl.save_kv_layer.assert_called_once_with("layer_0", kv_layer, attn_meta)
def test_wait_for_save(self, connector_and_impl):
connector, impl = connector_and_impl
connector.wait_for_save()
impl.wait_for_save.assert_called_once()
def test_get_finished(self, connector_and_impl):
connector, impl = connector_and_impl
impl.get_finished.return_value = ({"req1"}, None)
result = connector.get_finished({"req1"})
impl.get_finished.assert_called_once_with({"req1"})
assert result == ({"req1"}, None)
def test_register_kv_caches(self, connector_and_impl):
connector, impl = connector_and_impl
kv_caches = {"layer_0": torch.zeros(1)}
connector.register_kv_caches(kv_caches)
impl.register_kv_caches.assert_called_once_with(kv_caches)
def test_get_num_new_matched_tokens(self, connector_and_impl):
connector, impl = connector_and_impl
req = MagicMock()
impl.get_num_new_matched_tokens.return_value = (10, False)
result = connector.get_num_new_matched_tokens(req, 5)
impl.get_num_new_matched_tokens.assert_called_once_with(req, 5)
assert result == (10, False)
def test_update_state_after_alloc(self, connector_and_impl):
connector, impl = connector_and_impl
req = MagicMock()
blocks = MagicMock()
connector.update_state_after_alloc(req, blocks, 4)
impl.update_state_after_alloc.assert_called_once_with(req, blocks, 4)
def test_build_connector_meta(self, connector_and_impl):
connector, impl = connector_and_impl
sched_out = MagicMock()
connector.build_connector_meta(sched_out)
impl.build_connector_meta.assert_called_once_with(sched_out)
def test_update_connector_output(self, connector_and_impl):
connector, impl = connector_and_impl
out = MagicMock()
connector.update_connector_output(out)
impl.update_connector_output.assert_called_once_with(out)
def test_request_finished(self, connector_and_impl):
connector, impl = connector_and_impl
req = MagicMock()
impl.request_finished.return_value = (True, {"key": "val"})
result = connector.request_finished(req, [1, 2, 3])
impl.request_finished.assert_called_once_with(req, [1, 2, 3])
assert result == (True, {"key": "val"})
def test_take_events(self, connector_and_impl):
connector, impl = connector_and_impl
impl.take_events.return_value = iter([])
list(connector.take_events())
impl.take_events.assert_called_once()
def test_get_kv_connector_stats(self, connector_and_impl):
connector, impl = connector_and_impl
impl.get_kv_connector_stats.return_value = None
result = connector.get_kv_connector_stats()
impl.get_kv_connector_stats.assert_called_once()
assert result is None
def test_get_block_ids_with_load_errors(self, connector_and_impl):
connector, impl = connector_and_impl
impl.get_block_ids_with_load_errors.return_value = {7, 8}
result = connector.get_block_ids_with_load_errors()
assert result == {7, 8}
def test_wait_for_layer_load(self, connector_and_impl):
connector, impl = connector_and_impl
connector.wait_for_layer_load("layer_0")
impl.wait_for_layer_load.assert_called_once_with("layer_0")
tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py
+1 -1
View File
@@ -337,7 +337,7 @@ def test_async_recompute_blocks_not_cached_when_invalid(
scheduler_output = recompute_scheduler.schedule()
# request should be waiting for remote KVs
assert len(recompute_scheduler.waiting) == 1
assert len(recompute_scheduler.skipped_waiting) == 1
assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS
assert request.num_computed_tokens == num_external_computed_tokens
tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py
+12 -8
View File
@@ -76,8 +76,9 @@ def test_async_load_failure(
scheduler_output = scheduler.schedule()
assert len(scheduler.waiting) == 3
for request in scheduler.waiting:
assert len(scheduler.waiting) == 0
assert len(scheduler.skipped_waiting) == 3
for request in scheduler.skipped_waiting:
assert request.num_computed_tokens == num_external_computed_tokens
assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS
assert scheduler.connector.get_num_new_matched_tokens.call_count == 3
@@ -96,8 +97,9 @@ def test_async_load_failure(
min_invalid_block_idx = min(invalid_block_idxs)
assert len(scheduler.waiting) == 3
for request in scheduler.waiting:
assert len(scheduler.waiting) == 0
assert len(scheduler.skipped_waiting) == 3
for request in scheduler.skipped_waiting:
if request.request_id == request2.request_id:
assert request.num_computed_tokens == (
min_invalid_block_idx * scheduler.block_size
@@ -303,8 +305,9 @@ def test_async_progressive_load_failure(
scheduler_output = scheduler.schedule()
assert len(scheduler.waiting) == 1
assert scheduler.waiting.peek_request().request_id == request.request_id
assert len(scheduler.waiting) == 0
assert len(scheduler.skipped_waiting) == 1
assert scheduler.skipped_waiting.peek_request().request_id == request.request_id
assert request.num_computed_tokens == num_external_computed_tokens
assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS
assert scheduler.connector.get_num_new_matched_tokens.call_count == 1
@@ -325,8 +328,9 @@ def test_async_progressive_load_failure(
min_invalid_block_idx = min(min_invalid_block_idx, invalid_block_idx)
assert len(scheduler.waiting) == 1
assert scheduler.waiting.peek_request().request_id == request.request_id
assert len(scheduler.waiting) == 0
assert len(scheduler.skipped_waiting) == 1
assert scheduler.skipped_waiting.peek_request().request_id == request.request_id
assert request.num_computed_tokens == (
min_invalid_block_idx * scheduler.block_size
)
tests/v1/kv_connector/unit/test_multi_connector.py
+176 -25
View File
@@ -5,21 +5,27 @@ import shutil
import tempfile
from pathlib import Path
from typing import Any
from unittest.mock import MagicMock
import pytest
from tests.v1.kv_connector.unit.utils import create_vllm_config
from vllm import LLM, SamplingParams
from vllm.config import KVTransferConfig
from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory
from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorRole
from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorBase_V1
from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import (
MultiConnector,
MultiKVConnectorStats,
MultiKVConnectorWorkerMetadata,
)
from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (
NixlKVConnectorStats,
)
from vllm.v1.kv_cache_interface import KVCacheConfig
from vllm.v1.outputs import KVConnectorOutput, KVConnectorWorkerMetadata
MODEL_NAME = "meta-llama/Llama-3.2-1B-Instruct"
@@ -40,7 +46,14 @@ class MockConnectorStats(KVConnectorStats):
class MockConnector(KVConnectorBase_V1):
"""Mock connector that implements build_kv_connector_stats for testing."""
"""Mock connector for testing."""
def __new__(cls, *args, **kwargs):
# mock all KVConnectorBase_V1 functions
mock = MagicMock(spec_set=KVConnectorBase_V1)
# Override just build_kv_connector_stats
mock.build_kv_connector_stats = cls.build_kv_connector_stats
return mock
@classmethod
def build_kv_connector_stats(
@@ -70,16 +83,42 @@ class MockConnector(KVConnectorBase_V1):
pass
class MockCrossLayerConnector(MockConnector):
@property
def prefer_cross_layer_blocks(self) -> bool:
return True
# Register the mock connector
KVConnectorFactory.register_connector("MockConnector", __name__, MockConnector.__name__)
@pytest.fixture
def mc() -> MultiConnector:
"""MultiConnector using two mocked connectors"""
vllm_config = create_vllm_config()
mock_connector_config = {
"kv_connector": "MockConnector",
"kv_role": "kv_both",
"kv_connector_module_path": "tests.v1.kv_connector.unit.test_multi_connector",
}
vllm_config.kv_transfer_config = KVTransferConfig(
kv_connector="MultiConnector",
kv_role="kv_both",
kv_connector_extra_config={
"connectors": [mock_connector_config, mock_connector_config],
},
)
kv_cache_config = KVCacheConfig(
num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[]
)
mc = MultiConnector(
vllm_config=vllm_config,
role=KVConnectorRole.WORKER,
kv_cache_config=kv_cache_config,
)
return mc
# Helper function to compare directories recursively
def _compare_directories(dir1: Path, dir2: Path) -> bool:
"""Compares two directories recursively for identical content."""
@@ -715,24 +754,6 @@ class TestMultiConnectorStats:
assert not stats.is_empty()
class TestMultiConnectorPreferCrossLayerBlocks:
def test_all_connectors_prefer_cross_layer_blocks(self):
mc = MultiConnector.__new__(MultiConnector)
mc._connectors = [
MockCrossLayerConnector.__new__(MockCrossLayerConnector),
MockCrossLayerConnector.__new__(MockCrossLayerConnector),
]
assert mc.prefer_cross_layer_blocks is True
def test_mixed_connectors_do_not_prefer_cross_layer_blocks(self):
mc = MultiConnector.__new__(MultiConnector)
mc._connectors = [
MockCrossLayerConnector.__new__(MockCrossLayerConnector),
MockConnector.__new__(MockConnector), # default False
]
assert mc.prefer_cross_layer_blocks is False
def test_multi_connector_overrides_all_base_methods():
"""
Ensure MultiConnector overrides all public methods from KVConnectorBase_V1.
@@ -767,3 +788,133 @@ Options:
1. Add delegation in MultiConnector (preferred)
2. Add to INHERITED_OK if the base implementation works correctly
""")
def test_multi_connector_prefer_cross_layer_blocks(mc):
mc._connectors[0].prefer_cross_layer_blocks = False
mc._connectors[1].prefer_cross_layer_blocks = True
assert mc.prefer_cross_layer_blocks is False
mc._connectors[0].prefer_cross_layer_blocks = True
mc._connectors[1].prefer_cross_layer_blocks = True
assert mc.prefer_cross_layer_blocks is True
def test_multi_connector_worker_metadata(mc):
class MockConnectorWorkerMetadata(KVConnectorWorkerMetadata):
def __init__(self, data: set[str]):
self.data = data
class MockConnectorWorkerMetadata0(MockConnectorWorkerMetadata):
def aggregate(
self, other: KVConnectorWorkerMetadata
) -> KVConnectorWorkerMetadata:
assert isinstance(other, MockConnectorWorkerMetadata)
return MockConnectorWorkerMetadata0(data=self.data | other.data)
class MockConnectorWorkerMetadata1(MockConnectorWorkerMetadata):
def aggregate(
self, other: KVConnectorWorkerMetadata
) -> KVConnectorWorkerMetadata:
assert isinstance(other, MockConnectorWorkerMetadata)
return MockConnectorWorkerMetadata1(data=self.data | other.data)
# -------------------- test build_worker_connector_meta -------------------
# both connectors return None
mc._connectors[0].build_connector_worker_meta.return_value = None
mc._connectors[1].build_connector_worker_meta.return_value = None
assert mc.build_connector_worker_meta() is None
# only first connector returns None
worker_meta1a = MockConnectorWorkerMetadata1({"1a"})
mc._connectors[0].build_connector_worker_meta.return_value = None
mc._connectors[1].build_connector_worker_meta.return_value = worker_meta1a
mc_worker_meta_none_1a = mc.build_connector_worker_meta()
assert isinstance(mc_worker_meta_none_1a, MultiKVConnectorWorkerMetadata)
assert mc_worker_meta_none_1a.metadata == (None, worker_meta1a)
# only second connector returns None
worker_meta0a = MockConnectorWorkerMetadata0({"0a"})
mc._connectors[0].build_connector_worker_meta.return_value = worker_meta0a
mc._connectors[1].build_connector_worker_meta.return_value = None
mc_worker_meta_0a_none = mc.build_connector_worker_meta()
assert isinstance(mc_worker_meta_0a_none, MultiKVConnectorWorkerMetadata)
assert mc_worker_meta_0a_none.metadata == (worker_meta0a, None)
# both connectors do not return None
worker_meta0b = MockConnectorWorkerMetadata0({"0b"})
worker_meta1b = MockConnectorWorkerMetadata1({"1b"})
mc._connectors[0].build_connector_worker_meta.return_value = worker_meta0b
mc._connectors[1].build_connector_worker_meta.return_value = worker_meta1b
mc_worker_meta_0b_1b = mc.build_connector_worker_meta()
assert isinstance(mc_worker_meta_0b_1b, MultiKVConnectorWorkerMetadata)
assert mc_worker_meta_0b_1b.metadata == (worker_meta0b, worker_meta1b)
# ----------------------------- test aggregate ----------------------------
# aggregate ({"0a"}, None) and (None, {"1a"}) -> ({"0a"}, {"1a"})
mc_worker_meta_0a_1a = mc_worker_meta_0a_none.aggregate(mc_worker_meta_none_1a)
assert isinstance(mc_worker_meta_0a_1a, MultiKVConnectorWorkerMetadata)
assert mc_worker_meta_0a_1a.metadata == (worker_meta0a, worker_meta1a)
# aggregate ({"0a"}, None) and ({"0b"}, None) -> ({"0a", "0b"}, None)
mc._connectors[0].build_connector_worker_meta.return_value = worker_meta0b
mc._connectors[1].build_connector_worker_meta.return_value = None
mc_worker_meta_0b_none = mc.build_connector_worker_meta()
mc_worker_meta_0a_0b = mc_worker_meta_0a_none.aggregate(mc_worker_meta_0b_none)
assert isinstance(mc_worker_meta_0a_0b, MultiKVConnectorWorkerMetadata)
assert mc_worker_meta_0a_0b.metadata[1] is None
connector0_md = mc_worker_meta_0a_0b.metadata[0]
assert isinstance(connector0_md, MockConnectorWorkerMetadata0)
assert connector0_md.data == {"0a", "0b"}
# aggregate ({"0a"}, {"1a"}) and ({"0b"}, {"1b"}) -> ({"0a", "0b"}, {"1a", "1b"})
mc_worker_meta_01a_01b = mc_worker_meta_0a_1a.aggregate(mc_worker_meta_0b_1b)
assert isinstance(mc_worker_meta_01a_01b, MultiKVConnectorWorkerMetadata)
metadata = mc_worker_meta_01a_01b.metadata
assert len(metadata) == 2
connector0_md, connector1_md = metadata
assert isinstance(connector0_md, MockConnectorWorkerMetadata0)
assert isinstance(connector1_md, MockConnectorWorkerMetadata1)
assert connector0_md.data == {"0a", "0b"}
assert connector1_md.data == {"1a", "1b"}
# ---------------------- test update_connector_output ---------------------
def verify_worker_metadata(expected_metadata: MockConnectorWorkerMetadata | None):
def _verify_worker_metadata(connector_output: KVConnectorOutput):
worker_meta = connector_output.kv_connector_worker_meta
if expected_metadata is None:
assert worker_meta is None
return
assert isinstance(worker_meta, MockConnectorWorkerMetadata)
assert type(worker_meta) is type(expected_metadata)
assert expected_metadata.data == worker_meta.data
return _verify_worker_metadata
def assert_update_connector_output_called(mc: MultiConnector):
for c in mc._connectors:
c.update_connector_output.assert_called_once()
c.update_connector_output.reset_mock()
# no worker meta
kv_connector_output = KVConnectorOutput()
mc._connectors[0].update_connector_output.side_effect = verify_worker_metadata(None)
mc._connectors[1].update_connector_output.side_effect = verify_worker_metadata(None)
mc.update_connector_output(kv_connector_output)
assert_update_connector_output_called(mc)
# multi worker meta
kv_connector_output.kv_connector_worker_meta = mc_worker_meta_01a_01b
mc._connectors[0].update_connector_output.side_effect = verify_worker_metadata(
connector0_md
)
mc._connectors[1].update_connector_output.side_effect = verify_worker_metadata(
connector1_md
)
mc.update_connector_output(kv_connector_output)
assert_update_connector_output_called(mc)
assert kv_connector_output.kv_connector_worker_meta == mc_worker_meta_01a_01b
tests/v1/kv_connector/unit/test_nixl_connector.py
+59 -25
View File
@@ -9,7 +9,7 @@ import textwrap
import time
import uuid
from collections import defaultdict
from typing import Any
from typing import Any, cast
from unittest.mock import MagicMock, patch
import msgspec
@@ -332,14 +332,22 @@ def test_kv_transfer_handshake(dist_init):
# Prefill connector will register KV cache to populate proper handshake
# metadata.
# TODO this must match with values used in kv cache config
kv_cache_config = make_kv_cache_config(block_size=16, num_blocks=2)
prefill_connector = NixlConnector(
vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
vllm_config, KVConnectorRole.WORKER, kv_cache_config
)
kv_cache_spec = cast(
AttentionSpec, kv_cache_config.kv_cache_groups[0].kv_cache_spec
)
kv_cache_shape = FlashAttentionBackend.get_kv_cache_shape(
num_blocks=2, block_size=16, num_kv_heads=4, head_size=64
num_blocks=kv_cache_config.num_blocks,
block_size=kv_cache_spec.block_size,
num_kv_heads=kv_cache_spec.num_kv_heads,
head_size=kv_cache_spec.head_size,
)
shared_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
unique_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
shared_tensor = torch.zeros(*kv_cache_shape, dtype=kv_cache_spec.dtype)
unique_tensor = torch.zeros(*kv_cache_shape, dtype=kv_cache_spec.dtype)
kv_caches = {
"layer0": shared_tensor,
"layer1": unique_tensor,
@@ -383,7 +391,7 @@ def test_kv_transfer_handshake(dist_init):
# Decode connector will be able to create handshake with the prefill connector.
decode_connector = NixlConnector(
vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
vllm_config, KVConnectorRole.WORKER, kv_cache_config
)
decode_connector.register_kv_caches(kv_caches)
@@ -525,11 +533,13 @@ class TestNixlHandshake:
request_id = "req_id"
# Test worker role in decode server.
connector = NixlConnector(
vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
)
kv_cache_config = make_kv_cache_config(block_size=16, num_blocks=2)
connector = NixlConnector(vllm_config, KVConnectorRole.WORKER, kv_cache_config)
connector.connector_worker = FakeNixlConnectorWorker(
vllm_config, connector.engine_id, hand_shake_latency=0
vllm_config,
connector.engine_id,
hand_shake_latency=0,
kv_cache_config=kv_cache_config,
)
assert isinstance(connector.connector_worker.nixl_wrapper, FakeNixlWrapper)
worker = connector.connector_worker
@@ -1479,18 +1489,22 @@ def test_register_kv_caches(
patch(f"{nixl_module}.threading.Event"),
patch(f"{nixl_module}.threading.Thread") as mock_thread,
patch(f"{nixl_module}.get_current_attn_backend") as mock_get_attn_backend,
patch(f"{nixl_module}.get_current_attn_backends") as mock_get_attn_backends,
):
# Ensure get_attn_backend returns the correct value due to
# _cached_get_attn_backend returning the backend from previous
# test run if not mocking.
mock_get_attn_backend.return_value = backend_cls
mock_get_attn_backends.return_value = [backend_cls]
# Create connector
connector = NixlConnector(
vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
)
kv_cache_config = make_kv_cache_config(block_size=16, num_blocks=2)
connector = NixlConnector(vllm_config, KVConnectorRole.WORKER, kv_cache_config)
connector.connector_worker = FakeNixlConnectorWorker(
vllm_config, connector.engine_id, hand_shake_latency=0
vllm_config,
connector.engine_id,
hand_shake_latency=0,
kv_cache_config=kv_cache_config,
)
# Get the mock instance
@@ -1515,6 +1529,13 @@ def test_register_kv_caches(
num_layers = 32
block_size = 16
num_blocks = 8
# Keep the fake worker's expected num_blocks in sync with the
# cross-layer tensor we are about to register.
worker_kv_cache_config = make_kv_cache_config(
block_size=block_size, num_blocks=num_blocks
)
connector.connector_worker.kv_cache_config = worker_kv_cache_config
connector.connector_worker.num_blocks = worker_kv_cache_config.num_blocks
kv_cache_spec = AttentionSpec(
block_size=block_size,
num_kv_heads=4,
@@ -1568,11 +1589,17 @@ def test_register_kv_caches(
else:
# Create test kv cache tensors using proper backend shape
kv_cache_shape = backend_cls.get_kv_cache_shape(
num_blocks=2, block_size=16, num_kv_heads=4, head_size=64
kv_cache_spec = cast(
AttentionSpec, kv_cache_config.kv_cache_groups[0].kv_cache_spec
)
shared_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
unique_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
kv_cache_shape = backend_cls.get_kv_cache_shape(
num_blocks=kv_cache_config.num_blocks,
block_size=kv_cache_spec.block_size,
num_kv_heads=kv_cache_spec.num_kv_heads,
head_size=kv_cache_spec.head_size,
)
shared_tensor = torch.zeros(*kv_cache_shape, dtype=kv_cache_spec.dtype)
unique_tensor = torch.zeros(*kv_cache_shape, dtype=kv_cache_spec.dtype)
kv_caches = {
"layer0": shared_tensor,
"layer1": unique_tensor,
@@ -1606,7 +1633,7 @@ def test_register_kv_caches(
unique_tensor[1].data_ptr(),
]
expected_num_entries = 4
expected_blocks_count = 8
expected_blocks_count = kv_cache_config.num_blocks * 4
# Execute register_kv_caches
connector.register_kv_caches(kv_caches)
@@ -1639,7 +1666,7 @@ def test_register_kv_caches(
num_blocks = 8
expected_block_len = expected_tensor_size // num_blocks
else:
num_blocks = 2
num_blocks = kv_cache_config.num_blocks
if is_blocks_first:
expected_block_len = expected_tensor_size // num_blocks // 2
else:
@@ -2226,15 +2253,22 @@ def test_compatibility_hash_validation(
"enforce_handshake_compat": enforce_handshake_compat
},
)
kv_cache_config = make_kv_cache_config(block_size=16, num_blocks=2)
decode_connector = NixlConnector(
local_vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
local_vllm_config, KVConnectorRole.WORKER, kv_cache_config
)
decode_worker = decode_connector.connector_worker
kv_cache_shape = decode_worker.attn_backend.get_kv_cache_shape(
num_blocks=2, block_size=16, num_kv_heads=4, head_size=64
kv_cache_spec = cast(
AttentionSpec, kv_cache_config.kv_cache_groups[0].kv_cache_spec
)
shared_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
unique_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
kv_cache_shape = decode_worker.attn_backend.get_kv_cache_shape(
num_blocks=kv_cache_config.num_blocks,
block_size=kv_cache_spec.block_size,
num_kv_heads=kv_cache_spec.num_kv_heads,
head_size=kv_cache_spec.head_size,
)
shared_tensor = torch.zeros(*kv_cache_shape, dtype=kv_cache_spec.dtype)
unique_tensor = torch.zeros(*kv_cache_shape, dtype=kv_cache_spec.dtype)
kv_caches = {
"layer0": shared_tensor,
"layer1": unique_tensor,

Some files were not shown because too many files have changed in this diff Show More