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

Signed-off-by: AuYang <459461160@qq.com>

.buildkite/performance-benchmarks/scripts/compare-json-results.py

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

.readthedocs.yaml

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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 =

docs/benchmarking/dashboard.md

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -713,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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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()

requirements/cpu.txt

@@ -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

@@ -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/test.in

@@ -70,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

@@ -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
@@ -634,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
@@ -734,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

scripts/autotune_helion_kernels.py

@@ -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 = []

tests/basic_correctness/test_basic_correctness.py

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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/distributed/test_pipeline_parallel.py

@@ -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

@@ -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

@@ -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_serving_responses.py

@@ -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/pooling/score/test_utils.py

@@ -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

@@ -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_utils.py

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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/models/language/generation/test_common.py

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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
@@ -857,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"},
@@ -870,10 +885,6 @@ _MULTIMODAL_EXAMPLE_MODELS = {
             )
         },
     ),
-    "KimiK25ForConditionalGeneration": _HfExamplesInfo(
-        "moonshotai/Kimi-K2.5",
-        trust_remote_code=True,
-    ),
     "LightOnOCRForConditionalGeneration": _HfExamplesInfo(
         "lightonai/LightOnOCR-1B-1025"
     ),
@@ -1132,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,
@@ -1189,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",

tests/models/test_initialization.py

@@ -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/reasoning/test_nemotron_v3_reasoning_parser.py

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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/entrypoints/openai/serving_responses/test_function_call.py

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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,

tests/v1/kv_connector/unit/test_offloading_connector.py

@@ -148,17 +148,23 @@ class TransferSummary:
 
 class RequestRunner:
     def __init__(
-        self, offloaded_block_size: int, gpu_block_size: int, num_gpu_blocks: int
+        self,
+        offloaded_block_size: int,
+        gpu_block_size: int,
+        num_gpu_blocks: int,
+        async_scheduling: bool = True,
     ):
         self.offloaded_block_size: int = offloaded_block_size
         self.gpu_block_size: int = gpu_block_size
         self.num_gpu_blocks: int = num_gpu_blocks
+        self.async_scheduling: bool = async_scheduling
 
         self.req_id: int = -1
 
         vllm_config = create_vllm_config(
             block_size=gpu_block_size, max_num_batched_tokens=1000
         )
+        vllm_config.scheduler_config.async_scheduling = async_scheduling
         vllm_config.kv_transfer_config = KVTransferConfig(
             kv_connector="OffloadingConnector",
             kv_role="kv_both",
@@ -313,6 +319,8 @@ class RequestRunner:
 
         tokens_iter = iter(decoded_tokens)
         token_id = next(tokens_iter, None)
+        prev_scheduler_output = None
+        prev_model_runner_output = None
         while True:
             assert self.scheduler.requests
 
@@ -354,7 +362,16 @@ class RequestRunner:
             if self.scheduler.running:
                 token_id = next(tokens_iter, None)
 
-            self.scheduler.update_from_output(scheduler_output, model_runner_output)
+            if self.async_scheduling:
+                # in async scheduling we update the output of the previous step
+                if prev_model_runner_output is not None:
+                    self.scheduler.update_from_output(
+                        prev_scheduler_output, prev_model_runner_output
+                    )
+                prev_scheduler_output = scheduler_output
+                prev_model_runner_output = model_runner_output
+            else:
+                self.scheduler.update_from_output(scheduler_output, model_runner_output)
 
             if (
                 prev_token_id == EOS_TOKEN_ID
@@ -365,6 +382,11 @@ class RequestRunner:
                 continue
 
             if token_id is None:
+                if self.async_scheduling:
+                    # sample last token
+                    self.scheduler.update_from_output(
+                        prev_scheduler_output, prev_model_runner_output
+                    )
                 break
 
         self._parse_transfers()
@@ -445,11 +467,14 @@ class RequestRunner:
 def request_runner():
     runners = []
 
-    def runner_factory(offloaded_block_size, gpu_block_size, num_gpu_blocks):
+    def runner_factory(
+        offloaded_block_size, gpu_block_size, num_gpu_blocks, async_scheduling
+    ):
         runner = RequestRunner(
             offloaded_block_size=offloaded_block_size,
             gpu_block_size=gpu_block_size,
             num_gpu_blocks=num_gpu_blocks,
+            async_scheduling=async_scheduling,
         )
         runners.append(runner)
         return runner
@@ -466,7 +491,8 @@ def generate_store_output(block_hashes: Iterable[BlockHash]):
     )
 
 
-def test_offloading_connector(request_runner):
+@pytest.mark.parametrize("async_scheduling", [True, False])
+def test_offloading_connector(request_runner, async_scheduling: bool):
     offloaded_block_size = 12
     gpu_block_size = 4
     num_gpu_blocks = 100
@@ -476,6 +502,7 @@ def test_offloading_connector(request_runner):
         offloaded_block_size=offloaded_block_size,
         gpu_block_size=gpu_block_size,
         num_gpu_blocks=num_gpu_blocks,
+        async_scheduling=async_scheduling,
     )
 
     # 3 blocks, store just the middle block (skip first and last)
@@ -498,26 +525,28 @@ def test_offloading_connector(request_runner):
     runner.run(decoded_tokens=[0])
     runner.manager.prepare_store.assert_called()
 
-    # 1 more block, now set block_hashes_to_store = []
+    # 1 more block (+ token for async scheduling)
+    # now set block_hashes_to_store = []
     runner.manager.prepare_store.side_effect = (
         lambda block_hashes: generate_store_output([])
     )
-    runner.run(decoded_tokens=[0] * offloaded_block_size)
+    runner.run(decoded_tokens=[0] * (offloaded_block_size + 1))
 
-    # 1 more block, now check touch was called with all 6 blocks
+    # 1 more block (+ token for kicking off offloading)
+    # now check touch was called with all 6 blocks
     runner.manager.prepare_store.side_effect = (
         lambda block_hashes: generate_store_output(block_hashes)
     )
-    runner.run(decoded_tokens=[0] * offloaded_block_size)
+    runner.run(
+        decoded_tokens=[0] * (offloaded_block_size + 1),
+        expected_stored_gpu_block_indexes=(15, 16, 17),
+    )
     runner.manager.touch.assert_called()
     block_hashes1 = list(runner.manager.touch.call_args.args[0])
     assert len(block_hashes1) == 6
 
     # terminate request
-    runner.run(
-        decoded_tokens=[EOS_TOKEN_ID],
-        expected_stored_gpu_block_indexes=(15, 16, 17),
-    )
+    runner.run(decoded_tokens=[EOS_TOKEN_ID])
 
     # create a new request differing only on the last token
     runner.new_request(token_ids=[0] * (offloaded_block_size * 6 - 1) + [1])
@@ -608,7 +637,8 @@ def test_offloading_connector(request_runner):
     assert event.medium == "B"
 
 
-def test_request_preemption(request_runner):
+@pytest.mark.parametrize("async_scheduling", [True, False])
+def test_request_preemption(request_runner, async_scheduling: bool):
     offloaded_block_size = 12
     gpu_block_size = 4
     num_gpu_blocks = 100
@@ -617,6 +647,7 @@ def test_request_preemption(request_runner):
         offloaded_block_size=offloaded_block_size,
         gpu_block_size=gpu_block_size,
         num_gpu_blocks=num_gpu_blocks,
+        async_scheduling=async_scheduling,
     )
 
     free_block_queue = runner.scheduler.kv_cache_manager.block_pool.free_block_queue
@@ -674,7 +705,8 @@ def test_request_preemption(request_runner):
     )
 
 
-def test_concurrent_lookups_of_the_same_prefix(request_runner):
+@pytest.mark.parametrize("async_scheduling", [True, False])
+def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling: bool):
     offloaded_block_size = 12
     gpu_block_size = 4
     num_gpu_blocks = 100
@@ -683,6 +715,7 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner):
         offloaded_block_size=offloaded_block_size,
         gpu_block_size=gpu_block_size,
         num_gpu_blocks=num_gpu_blocks,
+        async_scheduling=async_scheduling,
     )
 
     # store 1 blocks
@@ -732,7 +765,8 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner):
     assert transfer_jobs == list(runner.offloading_spec.handler.transfer_specs)
 
 
-def test_abort_loading_requests(request_runner):
+@pytest.mark.parametrize("async_scheduling", [True, False])
+def test_abort_loading_requests(request_runner, async_scheduling: bool):
     offloaded_block_size = 12
     gpu_block_size = 4
     num_gpu_blocks = 100
@@ -741,6 +775,7 @@ def test_abort_loading_requests(request_runner):
         offloaded_block_size=offloaded_block_size,
         gpu_block_size=gpu_block_size,
         num_gpu_blocks=num_gpu_blocks,
+        async_scheduling=async_scheduling,
     )
 
     # store 1 blocks

tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py

@@ -18,6 +18,10 @@ from .utils import (
 pytestmark = pytest.mark.cpu_test
 
 
+def _num_waiting_requests(scheduler) -> int:
+    return len(scheduler.waiting) + len(scheduler.skipped_waiting)
+
+
 def test_basic_lifecycle():
     """Test lifecycle of a remote prefill."""
 
@@ -54,8 +58,8 @@ def test_basic_lifecycle():
     assert scheduler_output.total_num_scheduled_tokens == 0
 
     # Req waiting for KVs with no computed/scheduled toks ...
-    assert len(scheduler.waiting) == 1
-    assert request in scheduler.waiting
+    assert _num_waiting_requests(scheduler) == 1
+    assert request in scheduler.skipped_waiting
     assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS
     assert request.num_computed_tokens == NUM_TOKENS
 
@@ -81,7 +85,7 @@ def test_basic_lifecycle():
     # STEP (2):
     # (2a): schedule(): nothing happens!
     scheduler_output = scheduler.schedule()
-    assert len(scheduler.waiting) == 1
+    assert _num_waiting_requests(scheduler) == 1
     assert len(scheduler.running) == 0
 
     # (2b): forward(): request finishes recv.
@@ -94,7 +98,7 @@ def test_basic_lifecycle():
     engine_core_outputs = scheduler.update_from_output(
         scheduler_output, model_runner_output
     )
-    assert len(scheduler.waiting) == 1
+    assert _num_waiting_requests(scheduler) == 1
     assert request_id in scheduler.finished_recving_kv_req_ids
 
     # STEP (3):
@@ -180,7 +184,7 @@ def test_interleaved_lifecycle():
     scheduler.add_request(request_remote)
     scheduler_output = scheduler.schedule()
     assert len(scheduler.running) == 2
-    assert len(scheduler.waiting) == 1
+    assert _num_waiting_requests(scheduler) == 1
     assert len(scheduler_output.scheduled_new_reqs) == 1
     assert scheduler_output.scheduled_cached_reqs.num_reqs == 1
 
@@ -190,7 +194,7 @@ def test_interleaved_lifecycle():
     # STEP 3: continue running, KVs not arrived yet.
     scheduler_output = scheduler.schedule()
     assert len(scheduler.running) == 2
-    assert len(scheduler.waiting) == 1
+    assert _num_waiting_requests(scheduler) == 1
     assert len(scheduler_output.scheduled_new_reqs) == 0
     assert scheduler_output.scheduled_cached_reqs.num_reqs == 2
 
@@ -199,14 +203,14 @@ def test_interleaved_lifecycle():
     )
     scheduler.update_from_output(scheduler_output, model_runner_output)
     assert len(scheduler.running) == 2
-    assert len(scheduler.waiting) == 1
+    assert _num_waiting_requests(scheduler) == 1
     assert len(scheduler_output.scheduled_new_reqs) == 0
     assert scheduler_output.scheduled_cached_reqs.num_reqs == 2
 
     # STEP 4: KVs arrive.
     scheduler_output = scheduler.schedule()
     assert len(scheduler.running) == 2
-    assert len(scheduler.waiting) == 1
+    assert _num_waiting_requests(scheduler) == 1
     assert len(scheduler_output.scheduled_new_reqs) == 0
     assert scheduler_output.scheduled_cached_reqs.num_reqs == 2
 
@@ -218,7 +222,7 @@ def test_interleaved_lifecycle():
     # STEP 5: RECVed KVs are sent to ModelRunner.
     scheduler_output = scheduler.schedule()
     assert len(scheduler.running) == 3
-    assert len(scheduler.waiting) == 0
+    assert _num_waiting_requests(scheduler) == 0
     assert len(scheduler_output.scheduled_new_reqs) == 1
     assert scheduler_output.scheduled_cached_reqs.num_reqs == 2
 
@@ -279,14 +283,14 @@ def test_no_spurious_prefix_caching():
     scheduler.add_request(request_remote)
     scheduler_output = scheduler.schedule()
     scheduler.update_from_output(scheduler_output, EMPTY_MODEL_RUNNER_OUTPUT)
-    assert len(scheduler.waiting) == 1
+    assert _num_waiting_requests(scheduler) == 1
 
     # Schedule the local prefill request. This should
     # cause blocks to be cached, but separately from
     scheduler.add_request(request_local)
     scheduler_output = scheduler.schedule()
     assert len(scheduler.running) == 1
-    assert len(scheduler.waiting) == 1
+    assert _num_waiting_requests(scheduler) == 1
 
     local_blocks = scheduler.kv_cache_manager.coordinator.single_type_managers[
         0
@@ -348,7 +352,7 @@ def test_full_block_prompt():
         finished_recving={request_id}
     )
     scheduler.update_from_output(scheduler_output, model_runner_output)
-    assert len(scheduler.waiting) == 1
+    assert _num_waiting_requests(scheduler) == 1
     assert request_id in scheduler.finished_recving_kv_req_ids
 
     # # STEP (3): Run as usual.
@@ -418,7 +422,7 @@ def test_cannot_schedule_after_recv():
     model_runner_output = create_model_runner_output(reqs=[request_normal])
     scheduler.update_from_output(scheduler_output, model_runner_output)
     assert len(scheduler.running) == 1
-    assert len(scheduler.waiting) == 0
+    assert _num_waiting_requests(scheduler) == 0
 
     # Step 2: 5 blocks are in use (2 new for remote blocks).
     scheduler.add_request(request_remote)
@@ -426,7 +430,7 @@ def test_cannot_schedule_after_recv():
     model_runner_output = create_model_runner_output(reqs=[request_normal])
     scheduler.update_from_output(scheduler_output, model_runner_output)
     assert len(scheduler.running) == 1
-    assert len(scheduler.waiting) == 1
+    assert _num_waiting_requests(scheduler) == 1
 
     # Step 3: finish recving (5 blocks in use)
     scheduler_output = scheduler.schedule()
@@ -435,7 +439,7 @@ def test_cannot_schedule_after_recv():
     )
     scheduler.update_from_output(scheduler_output, model_runner_output)
     assert len(scheduler.running) == 1
-    assert len(scheduler.waiting) == 1
+    assert _num_waiting_requests(scheduler) == 1
 
     # Step 4: try to schedule, remote request is put to running list
     # because the transfer is completed.
@@ -445,7 +449,7 @@ def test_cannot_schedule_after_recv():
     )
     scheduler.update_from_output(scheduler_output, model_runner_output)
     assert len(scheduler.running) == 2
-    assert len(scheduler.waiting) == 0
+    assert _num_waiting_requests(scheduler) == 0
 
     # Step 5: Remote request will be put back to waiting list
     # because it needs new block to hold generated token.
@@ -453,7 +457,7 @@ def test_cannot_schedule_after_recv():
     model_runner_output = create_model_runner_output(reqs=[request_normal])
     scheduler.update_from_output(scheduler_output, model_runner_output)
     assert len(scheduler.running) == 1
-    assert len(scheduler.waiting) == 1
+    assert _num_waiting_requests(scheduler) == 1
 
     # Step 6: finish the request, free it.
     scheduler_output = scheduler.schedule()
@@ -462,7 +466,7 @@ def test_cannot_schedule_after_recv():
     )
     scheduler.update_from_output(scheduler_output, model_runner_output)
     assert len(scheduler.running) == 0
-    assert len(scheduler.waiting) == 1
+    assert _num_waiting_requests(scheduler) == 1
 
     # Step 7: now we can schedule (with 2 blocks computed),
     # request is retrieved from preempted list.
@@ -474,7 +478,7 @@ def test_cannot_schedule_after_recv():
     )
     scheduler.update_from_output(scheduler_output, model_runner_output)
     assert len(scheduler.running) == 1
-    assert len(scheduler.waiting) == 0
+    assert _num_waiting_requests(scheduler) == 0
 
     # Step 8: free everything.
     scheduler_output = scheduler.schedule()
@@ -521,7 +525,7 @@ def test_cannot_recv():
     model_runner_output = create_model_runner_output(reqs=[request_normal])
     scheduler.update_from_output(scheduler_output, model_runner_output)
     assert len(scheduler.running) == 1
-    assert len(scheduler.waiting) == 0
+    assert _num_waiting_requests(scheduler) == 0
 
     # Step 2: 3 blocks are in use,
     # need 3 new for remote blocks but only 2 are available.
@@ -530,7 +534,7 @@ def test_cannot_recv():
     model_runner_output = create_model_runner_output(reqs=[request_normal])
     scheduler.update_from_output(scheduler_output, model_runner_output)
     assert len(scheduler.running) == 1
-    assert len(scheduler.waiting) == 1
+    assert _num_waiting_requests(scheduler) == 1
     # Should not have KV transfer in progress.
     assert request_remote.status != RequestStatus.WAITING_FOR_REMOTE_KVS
 
@@ -541,14 +545,14 @@ def test_cannot_recv():
     )
     scheduler.update_from_output(scheduler_output, model_runner_output)
     assert len(scheduler.running) == 0
-    assert len(scheduler.waiting) == 1
+    assert _num_waiting_requests(scheduler) == 1
 
     # Step 4: now we can initiate KV transfer (with 2 blocks computed).
     scheduler_output = scheduler.schedule()
     model_runner_output = create_model_runner_output(reqs=[])
     scheduler.update_from_output(scheduler_output, model_runner_output)
     assert len(scheduler.running) == 0
-    assert len(scheduler.waiting) == 1
+    assert _num_waiting_requests(scheduler) == 1
     assert request_remote.status == RequestStatus.WAITING_FOR_REMOTE_KVS
 
     # Step 5: finish recving (5 blocks in use)
@@ -558,14 +562,14 @@ def test_cannot_recv():
     )
     scheduler.update_from_output(scheduler_output, model_runner_output)
     assert len(scheduler.running) == 0
-    assert len(scheduler.waiting) == 1
+    assert _num_waiting_requests(scheduler) == 1
 
     # Step 6: schedule remote request
     scheduler_output = scheduler.schedule()
     model_runner_output = create_model_runner_output(reqs=[request_remote])
     scheduler.update_from_output(scheduler_output, model_runner_output)
     assert len(scheduler.running) == 1
-    assert len(scheduler.waiting) == 0
+    assert _num_waiting_requests(scheduler) == 0
 
     # Step 7: free everything.
     scheduler_output = scheduler.schedule()

tests/v1/kv_connector/unit/utils.py

@@ -31,6 +31,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.example_connector import (  #
 from vllm.utils.hashing import sha256
 from vllm.v1.core.kv_cache_manager import KVCacheBlocks
 from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash
+from vllm.v1.core.sched.async_scheduler import AsyncScheduler
 from vllm.v1.core.sched.scheduler import Scheduler, SchedulerOutput
 from vllm.v1.kv_cache_interface import (
     FullAttentionSpec,
@@ -143,7 +144,7 @@ def create_scheduler(
     vllm_config: VllmConfig,
     num_blocks: int = 10000,
     kv_cache_config: KVCacheConfig | None = None,
-) -> Scheduler:
+) -> Scheduler | AsyncScheduler:
     """Initialize Scheduler For Testing."""
     block_size = vllm_config.cache_config.block_size
     if kv_cache_config is None:
@@ -163,7 +164,11 @@ def create_scheduler(
             ],
         )
     vllm_config.cache_config.num_gpu_blocks = num_blocks
-    return Scheduler(
+
+    scheduler_cls = (
+        AsyncScheduler if vllm_config.scheduler_config.async_scheduling else Scheduler
+    )
+    return scheduler_cls(
         vllm_config=vllm_config,
         kv_cache_config=kv_cache_config,
         log_stats=True,

tests/v1/spec_decode/test_eagle_step_kernel.py

@@ -0,0 +1,175 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for the fused EAGLE slot mapping kernel."""
+
+import pytest
+import torch
+
+from vllm.v1.spec_decode.utils import (
+    PADDING_SLOT_ID,
+    eagle_step_update_slot_mapping_and_metadata,
+)
+
+# Skip if no CUDA - Triton kernel requires GPU
+pytest.importorskip("triton")
+if not torch.cuda.is_available():
+    pytest.skip("CUDA required for EAGLE kernel tests", allow_module_level=True)
+
+
+def _reference_eagle_step_slot_mapping(
+    positions_1d: torch.Tensor,
+    block_table_tensor: torch.Tensor,
+    seq_lens: torch.Tensor,
+    block_size: int,
+    max_model_len: int,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Python reference for eagle_step_update_slot_mapping_and_metadata."""
+    new_positions = positions_1d + 1
+    exceeds_max = new_positions >= max_model_len
+    clamped_positions = torch.where(
+        exceeds_max, torch.zeros_like(positions_1d), new_positions
+    )
+    block_numbers = (clamped_positions // block_size).clamp(
+        max=block_table_tensor.shape[1] - 1
+    )
+    block_ids = block_table_tensor[
+        torch.arange(positions_1d.shape[0], device=positions_1d.device),
+        block_numbers.long(),
+    ].long()
+    slot_mapping = block_ids * block_size + (clamped_positions % block_size)
+    slot_mapping = torch.where(
+        exceeds_max, torch.full_like(slot_mapping, PADDING_SLOT_ID), slot_mapping
+    )
+    new_seq_lens = torch.where(exceeds_max, torch.ones_like(seq_lens), seq_lens + 1)
+    new_seq_lens = new_seq_lens.clamp(max=max_model_len)
+    return clamped_positions, slot_mapping, new_seq_lens
+
+
+def test_eagle_step_slot_mapping_kernel():
+    """Test fused kernel matches Python reference for slot mapping and metadata."""
+    device = torch.device("cuda")
+    batch_size = 32
+    block_size = 16
+    max_model_len = 4096
+    n_blocks_per_req = (max_model_len + block_size - 1) // block_size
+
+    positions_1d = torch.randint(
+        0, max_model_len - 10, (batch_size,), dtype=torch.int64, device=device
+    )
+    block_table_tensor = torch.randint(
+        0, 1000, (batch_size, n_blocks_per_req), dtype=torch.int32, device=device
+    )
+    seq_lens = torch.randint(1, 100, (batch_size,), dtype=torch.int32, device=device)
+
+    ref_clamped, ref_slot, ref_seq_lens = _reference_eagle_step_slot_mapping(
+        positions_1d.clone(),
+        block_table_tensor,
+        seq_lens.clone(),
+        block_size,
+        max_model_len,
+    )
+
+    out_clamped = torch.zeros(batch_size, dtype=torch.int64, device=device)
+    out_slot = torch.zeros(batch_size, dtype=torch.int64, device=device)
+    seq_lens_copy = seq_lens.clone()
+    eagle_step_update_slot_mapping_and_metadata(
+        positions_1d=positions_1d,
+        block_table_tensor=block_table_tensor,
+        seq_lens=seq_lens_copy,
+        block_size=block_size,
+        max_model_len=max_model_len,
+        out_clamped_positions=out_clamped,
+        out_slot_mapping=out_slot,
+    )
+
+    assert torch.equal(out_clamped, ref_clamped), (
+        f"clamped: {out_clamped} vs {ref_clamped}"
+    )
+    assert torch.equal(out_slot, ref_slot), f"slot: {out_slot} vs {ref_slot}"
+    assert torch.equal(seq_lens_copy, ref_seq_lens), (
+        f"seq_lens: {seq_lens_copy} vs {ref_seq_lens}"
+    )
+
+
+def test_eagle_step_slot_mapping_kernel_exceeds_max():
+    """Test fused kernel when position exceeds max_model_len."""
+    device = torch.device("cuda")
+    batch_size = 4
+    block_size = 16
+    max_model_len = 100
+    n_blocks_per_req = (max_model_len + block_size - 1) // block_size
+
+    positions_1d = torch.tensor([50, 98, 99, 100], dtype=torch.int64, device=device)
+    block_table_tensor = torch.randint(
+        0, 100, (batch_size, n_blocks_per_req), dtype=torch.int32, device=device
+    )
+    seq_lens = torch.tensor([51, 99, 100, 101], dtype=torch.int32, device=device)
+
+    out_clamped = torch.zeros(batch_size, dtype=torch.int64, device=device)
+    out_slot = torch.zeros(batch_size, dtype=torch.int64, device=device)
+    eagle_step_update_slot_mapping_and_metadata(
+        positions_1d=positions_1d,
+        block_table_tensor=block_table_tensor,
+        seq_lens=seq_lens,
+        block_size=block_size,
+        max_model_len=max_model_len,
+        out_clamped_positions=out_clamped,
+        out_slot_mapping=out_slot,
+    )
+
+    assert out_clamped[0].item() == 51
+    assert out_clamped[1].item() == 99
+    assert out_clamped[2].item() == 0
+    assert out_clamped[3].item() == 0
+    assert out_slot[2].item() == PADDING_SLOT_ID
+    assert out_slot[3].item() == PADDING_SLOT_ID
+    assert seq_lens[2].item() == 1
+    assert seq_lens[3].item() == 1
+
+
+def test_eagle_step_slot_mapping_kernel_cudagraph_padding():
+    """Test that padding threads write PADDING_SLOT_ID when
+    input_batch_size > batch_size (cudagraph padding)."""
+    device = torch.device("cuda")
+    batch_size = 4
+    input_batch_size = 8
+    block_size = 16
+    max_model_len = 4096
+    n_blocks_per_req = (max_model_len + block_size - 1) // block_size
+
+    positions_1d = torch.tensor([10, 20, 30, 40], dtype=torch.int64, device=device)
+    block_table_tensor = torch.randint(
+        0, 100, (batch_size, n_blocks_per_req), dtype=torch.int32, device=device
+    )
+    seq_lens = torch.tensor([11, 21, 31, 41], dtype=torch.int32, device=device)
+
+    ref_clamped, ref_slot, ref_seq_lens = _reference_eagle_step_slot_mapping(
+        positions_1d.clone(),
+        block_table_tensor,
+        seq_lens.clone(),
+        block_size,
+        max_model_len,
+    )
+
+    out_clamped = torch.zeros(batch_size, dtype=torch.int64, device=device)
+    out_slot = torch.full((input_batch_size,), -999, dtype=torch.int64, device=device)
+    seq_lens_copy = seq_lens.clone()
+    eagle_step_update_slot_mapping_and_metadata(
+        positions_1d=positions_1d,
+        block_table_tensor=block_table_tensor,
+        seq_lens=seq_lens_copy,
+        block_size=block_size,
+        max_model_len=max_model_len,
+        out_clamped_positions=out_clamped,
+        out_slot_mapping=out_slot,
+        input_batch_size=input_batch_size,
+    )
+
+    # Real slots should match the reference
+    assert torch.equal(out_clamped, ref_clamped)
+    assert torch.equal(out_slot[:batch_size], ref_slot)
+    assert torch.equal(seq_lens_copy, ref_seq_lens)
+
+    # Padding slots should be PADDING_SLOT_ID
+    for i in range(batch_size, input_batch_size):
+        assert out_slot[i].item() == PADDING_SLOT_ID

tests/v1/worker/test_gpu_model_runner.py

@@ -38,7 +38,7 @@ from vllm.v1.kv_cache_interface import (
 from vllm.v1.sample.metadata import SamplingMetadata
 from vllm.v1.worker.gpu_input_batch import InputBatch
 from vllm.v1.worker.gpu_model_runner import GPUModelRunner
-from vllm.v1.worker.utils import AttentionGroup, select_common_block_size
+from vllm.v1.worker.utils import select_common_block_size
 
 BLOCK_SIZE = 16
 NUM_BLOCKS = 10
@@ -203,37 +203,25 @@ def _make_kv_cache_spec() -> FullAttentionSpec:
 def test_select_common_block_size_prefers_manager_block_size():
     backend_a = _make_mock_backend_for_kernel_block_size([MultipleOf(32)])
     backend_b = _make_mock_backend_for_kernel_block_size([64, MultipleOf(16)])
-    attn_groups = [
-        AttentionGroup(backend_a, [], [], _make_kv_cache_spec(), 0),
-        AttentionGroup(backend_b, [], [], _make_kv_cache_spec(), 0),
-    ]
 
-    selected_size = select_common_block_size(128, attn_groups)
+    selected_size = select_common_block_size(128, [backend_a, backend_b])
     assert selected_size == 128
 
 
 def test_select_common_block_size_uses_largest_shared_int():
     backend_a = _make_mock_backend_for_kernel_block_size([128, 64])
     backend_b = _make_mock_backend_for_kernel_block_size([64, 32])
-    attn_groups = [
-        AttentionGroup(backend_a, [], [], _make_kv_cache_spec(), 0),
-        AttentionGroup(backend_b, [], [], _make_kv_cache_spec(), 0),
-    ]
 
-    selected_size = select_common_block_size(256, attn_groups)
+    selected_size = select_common_block_size(256, [backend_a, backend_b])
     assert selected_size == 64
 
 
 def test_select_common_block_size_no_valid_option():
     backend_a = _make_mock_backend_for_kernel_block_size([64])
     backend_b = _make_mock_backend_for_kernel_block_size([MultipleOf(16)])
-    attn_groups = [
-        AttentionGroup(backend_a, [], [], _make_kv_cache_spec(), 0),
-        AttentionGroup(backend_b, [], [], _make_kv_cache_spec(), 0),
-    ]
 
     with pytest.raises(ValueError):
-        select_common_block_size(48, attn_groups)
+        select_common_block_size(48, [backend_a, backend_b])
 
 
 def test_update_states_new_request(model_runner, dist_init):

tests/v1/worker/test_late_interaction_runner.py

@@ -64,6 +64,47 @@ def test_postprocess_scores_and_releases_query_cache():
         )
 
 
+def test_postprocess_scores_docs_in_batch():
+    runner = LateInteractionRunner()
+    query_key = "query-batch"
+    query_emb = torch.tensor([[1.0, 0.0], [0.0, 1.0]], dtype=torch.float32)
+    doc_emb_1 = torch.tensor([[1.0, 0.0], [0.5, 0.5]], dtype=torch.float32)
+    doc_emb_2 = torch.tensor([[0.0, 1.0], [0.3, 0.7], [1.0, 0.0]], dtype=torch.float32)
+
+    query_params = _make_pooling_params(
+        build_late_interaction_query_params(query_key=query_key, query_uses=2)
+    )
+    runner.postprocess_pooler_output(
+        raw_pooler_output=[query_emb],
+        pooling_params=[query_params],
+        req_ids=["query-req"],
+        finished_mask=[True],
+    )
+
+    doc_params = _make_pooling_params(
+        build_late_interaction_doc_params(query_key=query_key)
+    )
+    doc_output = runner.postprocess_pooler_output(
+        raw_pooler_output=[doc_emb_1, doc_emb_2],
+        pooling_params=[doc_params, doc_params],
+        req_ids=["doc-req-1", "doc-req-2"],
+        finished_mask=[True, True],
+    )
+    assert isinstance(doc_output, list)
+    assert doc_output[0] is not None
+    assert doc_output[1] is not None
+    assert torch.allclose(doc_output[0], compute_maxsim_score(query_emb, doc_emb_1))
+    assert torch.allclose(doc_output[1], compute_maxsim_score(query_emb, doc_emb_2))
+
+    with pytest.raises(ValueError, match="query cache miss"):
+        runner.postprocess_pooler_output(
+            raw_pooler_output=[doc_emb_1],
+            pooling_params=[doc_params],
+            req_ids=["doc-req-3"],
+            finished_mask=[True],
+        )
+
+
 def test_finished_request_releases_unscored_doc_use():
     runner = LateInteractionRunner()
     query_key = "query-cancel"

tools/pre_commit/check_torch_cuda.py

@@ -8,8 +8,8 @@ import regex as re
 # Regex: match `torch.cuda.xxx` but allow `torch.accelerator.xxx`
 # --------------------------------------------------------------------------- #
 _TORCH_CUDA_PATTERNS = [
-    r"\btorch\.cuda\.empty_cache\b",
-    r"\btorch\.cuda\.synchronize\b",
+    r"\btorch\.cuda\.(empty_cache|synchronize|device\()\b",
+    r"\bwith\btorch\.cuda\.device\b",
 ]
 
 ALLOWED_FILES = {"vllm/platforms/", "vllm/device_allocator/"}

vllm/_custom_ops.py

@@ -427,7 +427,7 @@ def rms_norm_dynamic_per_token_quant(
     scale_ub: torch.Tensor | None = None,
     residual: torch.Tensor | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    output = torch.empty_like(input, dtype=quant_dtype)
+    output = torch.empty(input.shape, dtype=quant_dtype, device=input.device)
     scales = torch.empty(
         (input.numel() // input.shape[-1], 1), device=input.device, dtype=torch.float32
     )
@@ -451,7 +451,7 @@ def rms_norm_per_block_quant(
     tma_alignment: int = 0,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     assert len(group_size) == 2
-    output = torch.empty_like(input, dtype=quant_dtype)
+    output = torch.empty(input.shape, dtype=quant_dtype, device=input.device)
     if is_scale_transposed:
         if tma_alignment == 0:
             scales = torch.empty(

vllm/_xpu_ops.py

@@ -7,6 +7,7 @@ import torch
 from vllm_xpu_kernels.flash_attn_interface import flash_attn_varlen_func
 
 from vllm.logger import init_logger
+from vllm.platforms import current_platform
 
 logger = init_logger(__name__)
 
@@ -157,3 +158,247 @@ class xpu_ops:
             "get_scheduler_metadata is not implemented for xpu_ops, returning None."
         )
         return None
+
+    @staticmethod
+    def indexer_k_quant_and_cache(
+        k: torch.Tensor,
+        kv_cache: torch.Tensor,
+        slot_mapping: torch.Tensor,
+        quant_block_size: int,
+        scale_fmt: str | None,
+    ) -> None:
+        head_dim = k.shape[-1]
+        k = k.view(-1, head_dim)  # [total_tokens, head_dim]
+
+        def group_quant_torch(
+            x: torch.Tensor,
+            group_size: int,
+            eps: float = 1e-10,
+            dtype: torch.dtype | None = None,
+            column_major_scales: bool = False,
+            out_q: torch.Tensor | None = None,
+            use_ue8m0: bool | None = None,
+        ) -> tuple[torch.Tensor, torch.Tensor]:
+            if use_ue8m0 is None:
+                # Default fallback - could import is_deep_gemm_e8m0_used if needed
+                use_ue8m0 = False
+
+            if dtype is None:
+                dtype = current_platform.fp8_dtype()
+
+            # Validate inputs
+            assert x.shape[-1] % group_size == 0, (
+                f"Last dimension {x.shape[-1]} must be divisible by "
+                f"group_size {group_size}"
+            )
+            assert x.stride(-1) == 1, "Input tensor groups must be contiguous"
+
+            # Prepare output tensor
+            if out_q is None:
+                x_q = torch.empty_like(x, dtype=dtype)
+            else:
+                assert out_q.shape == x.shape
+                x_q = out_q
+
+            # Reshape input for group processing
+            # Original shape: (..., last_dim)
+            # Target shape: (..., num_groups, group_size)
+            original_shape = x.shape
+            num_groups = original_shape[-1] // group_size
+
+            # Reshape to separate groups
+            group_shape = original_shape[:-1] + (num_groups, group_size)
+            x_grouped = x.view(group_shape)
+
+            # Compute per-group absolute maximum values
+            # Shape: (..., num_groups)
+            abs_max = torch.amax(torch.abs(x_grouped), dim=-1, keepdim=False)
+            abs_max = torch.maximum(
+                abs_max, torch.tensor(eps, device=x.device, dtype=x.dtype)
+            )
+
+            # Compute scales
+            FP8_MAX = torch.finfo(dtype).max
+            FP8_MIN = torch.finfo(dtype).min
+            scale_raw = abs_max / FP8_MAX
+
+            if use_ue8m0:
+                # For UE8M0 format, scales must be powers of 2
+                scales = torch.pow(2.0, torch.ceil(torch.log2(scale_raw)))
+            else:
+                scales = scale_raw
+
+            # Expand scales for broadcasting with grouped data
+            # Shape: (..., num_groups, 1)
+            scales_expanded = scales.unsqueeze(-1)
+
+            # Quantize the grouped data
+            x_scaled = x_grouped / scales_expanded
+            x_clamped = torch.clamp(x_scaled, FP8_MIN, FP8_MAX)
+            x_quantized = x_clamped.to(dtype)
+
+            # Reshape back to original shape
+            x_q.copy_(x_quantized.view(original_shape))
+
+            # Prepare scales tensor in requested format
+            if column_major_scales:
+                # Column-major: (num_groups,) + batch_dims
+                # Transpose the scales to put group dimension first
+                scales_shape = (num_groups,) + original_shape[:-1]
+                x_s = scales.permute(-1, *range(len(original_shape) - 1))
+                x_s = x_s.contiguous().view(scales_shape)
+            else:
+                # Row-major: batch_dims + (num_groups,)
+                x_s = scales.contiguous()
+
+            # Ensure scales are float32
+            return x_q, x_s.float()
+
+        k_fp8, k_scale = group_quant_torch(
+            k,
+            group_size=quant_block_size,
+            column_major_scales=False,
+            use_ue8m0=(scale_fmt == "ue8m0"),
+        )
+
+        k_fp8_bytes = k_fp8.view(-1, head_dim).view(torch.uint8)
+        scale_bytes = k_scale.view(torch.uint8).view(-1, 4)
+        k = torch.cat(
+            [k_fp8_bytes, scale_bytes], dim=-1
+        )  # [total_tokens, head_dim + 4]
+
+        slot_mapping = slot_mapping.flatten()
+        # kv_cache: [num_block, block_size, head_dim + 4]
+        kv_cache.view(-1, kv_cache.shape[-1]).index_copy_(0, slot_mapping, k)
+
+    @staticmethod
+    def cp_gather_indexer_k_quant_cache(
+        kv_cache: torch.Tensor,
+        dst_k: torch.Tensor,
+        dst_scale: torch.Tensor,
+        block_table: torch.Tensor,
+        cu_seq_lens: torch.Tensor,
+    ) -> None:
+        """
+        Args:
+            kv_cache: [num_blocks, block_size, cache_stride] - quantized KV cache
+                    Layout per block: [k_values, scale_values]
+                    - k_values: [block_size * head_dim]
+                    - scale_values: [block_size * head_dim * 4 / quant_block_size]
+            dst_k: [num_tokens, head_dim] - output tensor for K values
+            dst_scale: [num_tokens, head_dim / quant_block_size * 4]
+                - output tensor for scale values
+            block_table: [batch_size, num_blocks] - block table for indexing
+            cu_seq_lens: [batch_size + 1] - cumulative sequence lengths
+        """
+        batch_size = block_table.size(0)
+        num_tokens = dst_k.size(0)
+        head_dim = dst_k.size(1)
+        cache_block_size = kv_cache.size(1)
+        quant_block_size = head_dim * 4 // dst_scale.size(1)
+
+        # For each token, find which batch it belongs to using searchsorted
+        token_indices = torch.arange(num_tokens, device=dst_k.device) + 1
+        # cu_seq_lens is [batch_size + 1], we need to find which interval each
+        # token belongs to
+        batch_indices = torch.searchsorted(cu_seq_lens, token_indices) - 1
+        batch_indices = torch.clamp(batch_indices, 0, batch_size - 1)
+
+        # Calculate the in-batch sequence index for each token
+        inbatch_seq_indices = token_indices - cu_seq_lens[batch_indices]
+
+        # Find which block each token belongs to
+        block_indices_in_table = inbatch_seq_indices // cache_block_size
+        physical_block_indices = block_table[batch_indices, block_indices_in_table]
+
+        # Calculate the offset within each block
+        inblock_offsets = (inbatch_seq_indices - 1) % cache_block_size
+
+        # Calculate strides
+        block_stride = kv_cache.stride(0)  # stride for each block
+
+        # Flatten kv_cache for easier indexing
+        kv_cache_flat = kv_cache.view(-1)
+
+        # Calculate source offset for K values for all tokens (vectorized)
+        src_block_offsets = physical_block_indices * block_stride
+        src_k_offsets = src_block_offsets + inblock_offsets * head_dim
+
+        # Gather K values using advanced indexing
+        # Create indices for all elements we need to gather
+        k_indices = src_k_offsets.unsqueeze(1) + torch.arange(
+            head_dim, device=dst_k.device
+        )
+        dst_k[:] = kv_cache_flat[k_indices]
+
+        # Calculate source offset for scale values (vectorized)
+        # Scales are stored after all K values for each block
+        scale_size = head_dim * 4 // quant_block_size
+        src_scale_offsets = src_block_offsets + head_dim + inblock_offsets * scale_size
+
+        # Gather scale values
+        scale_indices = src_scale_offsets.unsqueeze(1) + torch.arange(
+            scale_size, device=dst_scale.device
+        )
+        dst_scale[:] = kv_cache_flat[scale_indices]
+
+    @staticmethod
+    def top_k_per_row_prefill(
+        logits: torch.Tensor,
+        cu_seqlen_ks: torch.Tensor,
+        cu_seqlen_ke: torch.Tensor,
+        raw_topk_indices: torch.Tensor,
+        num_rows: int,
+        stride0: int,
+        strdide1: int,
+        topk_tokens: int,
+    ) -> torch.Tensor:
+        real_topk = min(topk_tokens, logits.shape[-1])
+        topk_indices = logits.topk(real_topk, dim=-1)[1].to(torch.int32)
+        topk_indices -= cu_seqlen_ks[:, None]
+        mask_lo = topk_indices >= 0
+        mask_hi = topk_indices - (cu_seqlen_ke - cu_seqlen_ks)[:, None] < 0
+        mask = torch.full_like(
+            topk_indices, False, dtype=torch.bool, device=topk_indices.device
+        )
+        mask = mask_lo & mask_hi
+        topk_indices.masked_fill_(~mask, -1)
+        raw_topk_indices[: topk_indices.shape[0], : topk_indices.shape[1]] = (
+            topk_indices
+        )
+
+    @staticmethod
+    def top_k_per_row_decode(
+        logits: torch.Tensor,
+        next_n: int,
+        seq_lens: torch.Tensor,
+        raw_topk_indices: torch.Tensor,
+        num_rows: int,
+        stride0: int,
+        stride1: int,
+        topk_tokens: int,
+    ) -> torch.Tensor:
+        device = logits.device
+        batch_size = seq_lens.size(0)
+        # padded query len
+        padded_num_tokens = batch_size * next_n
+        positions = (
+            torch.arange(logits.shape[-1], device=device)
+            .unsqueeze(0)
+            .expand(batch_size * next_n, -1)
+        )
+        row_indices = torch.arange(padded_num_tokens, device=device) // next_n
+        next_n_offset = torch.arange(padded_num_tokens, device=device) % next_n
+        index_end_pos = (seq_lens[row_indices] - next_n + next_n_offset).unsqueeze(1)
+        # index_end_pos: [B * N, 1]
+        mask = positions <= index_end_pos
+        # mask: [B * N, L]
+        logits = logits.masked_fill(~mask, float("-inf"))
+        topk_indices = logits.topk(topk_tokens, dim=-1)[1].to(torch.int32)  # [B * N, K]
+        # ensure we don't set indices for the top k
+        # that is out of range(masked already)
+        # this will happen if context length is shorter than K
+        topk_indices[topk_indices > index_end_pos] = -1
+        raw_topk_indices[: topk_indices.shape[0], : topk_indices.shape[1]] = (
+            topk_indices
+        )

vllm/compilation/caching.py

@@ -369,8 +369,14 @@ class VllmSerializableFunction(SerializableCallable):  # type: ignore[misc]
 
         from vllm.compilation.backends import VllmBackend
 
+        saved_aot_autograd_config = self.aot_autograd_config
+        if saved_aot_autograd_config is not None:
+            functorch_ctx = torch._functorch.config.patch(saved_aot_autograd_config)
+        else:
+            functorch_ctx = contextlib.nullcontext()
+
         vllm_backend = VllmBackend(vllm_config, self.prefix, self.is_encoder)
-        with tracing(TracingContext(self._fake_mode)):
+        with tracing(TracingContext(self._fake_mode)), functorch_ctx:
             result = vllm_backend(self.graph_module, list(self.example_inputs))
             self.optimized_call = result.optimized_call
             self.vllm_backend = vllm_backend

vllm/compilation/compiler_interface.py

@@ -348,13 +348,39 @@ class InductorStandaloneAdaptor(CompilerInterface):
         # Can remove this after the following issue gets fixed
         # https://github.com/pytorch/pytorch/issues/174502
         if envs.VLLM_ENABLE_PREGRAD_PASSES:
-            ctx: Any = contextlib.nullcontext()
+            pregrad_ctx: Any = contextlib.nullcontext()
         else:
-            ctx = patch(
+            pregrad_ctx = patch(
                 "torch._inductor.compile_fx._recursive_pre_grad_passes",
                 lambda gm, _: gm,
             )
-        with ctx, _patch_constrain_to_fx_strides():
+
+        # When inputs are FakeTensors (from create_concrete_args),
+        # standalone_compile("from_example_inputs") would normally create
+        # a fresh FakeTensorMode, causing a mode mismatch assertion.
+        # Patch FakeTensorMode in standalone_compile so it reuses the
+        # mode already attached to our FakeTensors. This gives us both
+        # ignore_shape_env=True (from "from_example_inputs") and mode
+        # consistency (from reusing our mode).
+        # Can remove this after the following issue gets fixed:
+        # https://github.com/pytorch/pytorch/issues/176562
+        from torch._subclasses.fake_tensor import FakeTensor
+
+        input_fake_mode = None
+        for x in example_inputs:
+            if isinstance(x, FakeTensor):
+                input_fake_mode = x.fake_mode
+                break
+
+        if input_fake_mode is not None:
+            fake_mode_ctx: Any = patch(
+                "torch._inductor.standalone_compile.FakeTensorMode",
+                lambda *a, **kw: input_fake_mode,
+            )
+        else:
+            fake_mode_ctx = contextlib.nullcontext()
+
+        with pregrad_ctx, fake_mode_ctx, _patch_constrain_to_fx_strides():
             compiled_graph = standalone_compile(graph, example_inputs, **compile_kwargs)
 
         if use_aot:

vllm/compilation/counter.py

@@ -31,6 +31,12 @@ class CompilationCounter:
     num_compiled_artifacts_saved: int = 0
     # The number of standalone_compile compiled artifacts loaded from cache
     num_compiled_artifacts_loaded: int = 0
+    # The number of AOT compile invocations
+    num_aot_compiles: int = 0
+    # The number of AOT compiled artifacts saved to disk
+    num_aot_artifacts_saved: int = 0
+    # The number of AOT compiled artifacts loaded from disk
+    num_aot_artifacts_loaded: int = 0
     # Number of times a model was loaded with CompilationMode.STOCK_TORCH_COMPILE
     stock_torch_compile_count: int = 0
 

vllm/compilation/decorators.py

@@ -266,6 +266,51 @@ def _verify_source_unchanged(
         )
 
 
+def _try_load_aot_compiled_fn(
+    model: Any,
+    aot_compilation_path: str,
+) -> Any | None:
+    """Try to load an AOT-compiled function from disk.
+
+    Returns the loaded callable on success, or None on failure.
+    Re-raises on failure when ``VLLM_FORCE_AOT_LOAD`` is set.
+    """
+    try:
+        with monitor_torch_compile(model.vllm_config):
+            with (
+                set_current_vllm_config(model.vllm_config),
+                open(aot_compilation_path, "rb") as f,
+            ):
+                loaded_fn = torch.compiler.load_compiled_function(
+                    f, f_globals=model.forward.__globals__
+                )
+            _verify_source_unchanged(loaded_fn.source_info(), model.vllm_config)
+            ds_config = model.compilation_config.dynamic_shapes_config
+            if not ds_config.evaluate_guards:
+                loaded_fn.disable_guard_check()
+            # Eagerly load compiled artifacts now that traced_files
+            # is populated by _verify_source_unchanged.
+            with maybe_use_cudagraph_partition_wrapper(model.vllm_config):
+                loaded_fn._artifacts.compiled_fn.finalize_loading(model.vllm_config)
+        compilation_counter.num_aot_artifacts_loaded += 1
+        logger.info("Directly load AOT compilation from path %s", aot_compilation_path)
+        return loaded_fn
+    except Exception as e:
+        if os.path.exists(aot_compilation_path):
+            if isinstance(e, EOFError):
+                message = "Compile cache file corrupted."
+            else:
+                message = str(e)
+            logger.warning(
+                "Compiling model again due to a load failure from %s, reason: %s",
+                aot_compilation_path,
+                message,
+            )
+        if envs.VLLM_FORCE_AOT_LOAD:
+            raise e
+        return None
+
+
 def _support_torch_compile(
     cls: type[_T],
     dynamic_arg_dims: dict[str, int | list[int]],
@@ -438,51 +483,17 @@ def _support_torch_compile(
             dp_rank = self.vllm_config.parallel_config.data_parallel_index
             cache_dir = os.path.join(cache_dir, f"rank_{rank}_{dp_rank}")
             aot_compilation_path = os.path.join(cache_dir, "model")
-            try:
-                with monitor_torch_compile(self.vllm_config):
+            if not envs.VLLM_DISABLE_COMPILE_CACHE:
+                loaded_fn = _try_load_aot_compiled_fn(self, aot_compilation_path)
+                if loaded_fn is not None:
+                    self.aot_compiled_fn = loaded_fn
+                    self.was_aot_compile_fn_loaded_from_disk = True
                     with (
-                        set_current_vllm_config(self.vllm_config),
-                        open(aot_compilation_path, "rb") as f,
+                        monitor_profiling_run(),
+                        maybe_use_cudagraph_partition_wrapper(self.vllm_config),
                     ):
-                        loaded_fn = torch.compiler.load_compiled_function(
-                            f, f_globals=self.forward.__globals__
-                        )
-                    _verify_source_unchanged(loaded_fn.source_info(), self.vllm_config)
-                    ds_config = self.compilation_config.dynamic_shapes_config
-                    if not ds_config.evaluate_guards:
-                        loaded_fn.disable_guard_check()
-                    # Eagerly load compiled artifacts now that traced_files
-                    # is populated by _verify_source_unchanged.
-                    with maybe_use_cudagraph_partition_wrapper(self.vllm_config):
-                        loaded_fn._artifacts.compiled_fn.finalize_loading(
-                            self.vllm_config
-                        )
-                self.aot_compiled_fn = loaded_fn
-                self.was_aot_compile_fn_loaded_from_disk = True
-            except Exception as e:
-                if os.path.exists(aot_compilation_path):
-                    if isinstance(e, EOFError):
-                        message = "Compile cache file corrupted."
-                    else:
-                        message = str(e)
-                    logger.warning(
-                        "Compiling model again due to a load failure from %s, "
-                        "reason: %s",
-                        aot_compilation_path,
-                        message,
-                    )
-                if envs.VLLM_FORCE_AOT_LOAD:
-                    raise e
-            if getattr(self, "aot_compiled_fn", None) is not None:
-                logger.info(
-                    "Directly load AOT compilation from path %s", aot_compilation_path
-                )
-                with (
-                    monitor_profiling_run(),
-                    maybe_use_cudagraph_partition_wrapper(self.vllm_config),
-                ):
-                    output = self.aot_compiled_fn(self, *args, **kwargs)
-                return output
+                        output = self.aot_compiled_fn(self, *args, **kwargs)
+                    return output
 
         if self.compiled:
             assert (
@@ -570,6 +581,7 @@ def _support_torch_compile(
                 self._aot_cache_dir = cache_dir
                 with monitor_torch_compile(self.vllm_config):
                     self.aot_compiled_fn = self.aot_compile(*args, **kwargs)
+                    compilation_counter.num_aot_compiles += 1
                     # All compilation is done at this point, save the
                     # AOT artifact.
                     self.save_aot_compiled_function()
@@ -593,6 +605,9 @@ def _support_torch_compile(
 
     # triggers VllmSerializableFunction.serialize()
     def save_aot_compiled_function(self: type[_T]) -> None:
+        if envs.VLLM_DISABLE_COMPILE_CACHE:
+            return
+
         if self.was_aot_compile_fn_loaded_from_disk:
             logger.debug("AOT compiled function was loaded from cache, skipping save")
             return
@@ -608,6 +623,7 @@ def _support_torch_compile(
             tmp_file = f"{self._aot_compilation_path}.{os.getpid()}.tmp"
             self.aot_compiled_fn.save_compiled_function(tmp_file)
             os.replace(tmp_file, self._aot_compilation_path)
+            compilation_counter.num_aot_artifacts_saved += 1
             logger.info_once(
                 "saved AOT compiled function to %s",
                 self._aot_compilation_path,

vllm/compilation/passes/fusion/attn_quant_fusion.py

@@ -170,9 +170,8 @@ class AttentionFp8StaticQuantPattern(AttentionQuantPattern):
             kv_cache_dummy_dep: torch.Tensor,
         ) -> torch.Tensor:
             # attn output in quant_dtype
-            output_attn = torch.ops.aten.full.default(
+            output_attn = torch.empty(
                 [q.shape[0], self.num_heads, self.head_size],
-                0.0,
                 dtype=self.quant_dtype,
                 device=q.device,
             )
@@ -271,9 +270,8 @@ class AttentionNvfp4QuantPattern(AttentionQuantPattern):
             kv_cache_dummy_dep: torch.Tensor,
         ) -> tuple[torch.Tensor, torch.Tensor]:
             # attention output in quant_dtype
-            output_attn = torch.ops.aten.full.default(
+            output_attn = torch.empty(
                 [q.shape[0], self.num_heads, self.head_size // 2],
-                0.0,
                 dtype=self.quant_dtype,
                 device=q.device,
             )

vllm/compilation/piecewise_backend.py

@@ -34,13 +34,14 @@ def get_fake_args_from_graph(graph: fx.GraphModule) -> list[Any]:
 
 
 def create_concrete_args(graph: fx.GraphModule, size: int) -> list[Any]:
-    """Create example inputs with symbolic dims replaced by a concrete size.
+    """Create Fake example inputs with symbolic dims replaced by a concrete size.
 
-    Used for single-size eager compilation where we need concrete-shaped
-    inputs but don't have real runtime tensors yet.
+    Used for single-size compilation where we need concrete-shaped inputs.
+    The Dynamo-captured graph gives us example inputs with SymInts in them.
     """
     from torch._prims_common import compute_required_storage_length
-    from torch.fx.experimental.symbolic_shapes import is_symbolic
+    from torch._subclasses.fake_tensor import FakeTensorMode
+    from torch.fx.experimental.symbolic_shapes import ShapeEnv, is_symbolic
 
     def concretize(sym_val: Any) -> int:
         """Replace all symbolic variables in a SymInt expression with size."""
@@ -49,25 +50,28 @@ def create_concrete_args(graph: fx.GraphModule, size: int) -> list[Any]:
         expr = sym_val.node.expr
         return int(expr.subs({s: size for s in expr.free_symbols}))
 
+    fake_mode = FakeTensorMode(shape_env=ShapeEnv())
+
     args: list[Any] = []
-    for node in graph.graph.nodes:
-        if node.op != "placeholder":
-            break
-        val = node.meta["example_value"]
-        if isinstance(val, torch.SymInt):
-            args.append(concretize(val))
-        elif isinstance(val, torch.Tensor):
-            new_shape = tuple(concretize(d) for d in val.shape)
-            new_strides = tuple(concretize(s) for s in val.stride())
-            new_storage_offset = concretize(val.storage_offset())
-            needed_size = compute_required_storage_length(
-                new_shape, new_strides, new_storage_offset
-            )
-            t = torch.empty(needed_size, dtype=val.dtype, device=val.device)
-            t = t.as_strided(new_shape, new_strides, new_storage_offset)
-            args.append(t)
-        else:
-            args.append(val)
+    with fake_mode:
+        for node in graph.graph.nodes:
+            if node.op != "placeholder":
+                break
+            val = node.meta["example_value"]
+            if isinstance(val, torch.SymInt):
+                args.append(concretize(val))
+            elif isinstance(val, torch.Tensor):
+                new_shape = tuple(concretize(d) for d in val.shape)
+                new_strides = tuple(concretize(s) for s in val.stride())
+                new_storage_offset = concretize(val.storage_offset())
+                needed_size = compute_required_storage_length(
+                    new_shape, new_strides, new_storage_offset
+                )
+                t = torch.empty(needed_size, dtype=val.dtype, device=val.device)
+                t = t.as_strided(new_shape, new_strides, new_storage_offset)
+                args.append(t)
+            else:
+                args.append(val)
     return args
 
 
@@ -258,31 +262,15 @@ class PiecewiseBackend:
             else:
                 args_list = get_fake_args_from_graph(self.graph)
 
-            # TODO(https://github.com/vllm-project/vllm/issues/35766)
-            # Can we remove strict_autograd_cache and
-            # force_non_lazy_backward_lowering overrides?
-            # I added them explicitly because this is what they are
-            # set to before the refactor
-            # (https://github.com/vllm-project/vllm/pull/35472).
-            # They affect the aotautograd cache key computation
-            # but they shouldn't have any effect on the actual
-            # compilation.
-            config_patches = dict(
-                bundled_autograd_cache=True,
-                strict_autograd_cache=False,
+            range_entry.runnable = self.vllm_backend.compiler_manager.compile(
+                self.graph,
+                args_list,
+                self.vllm_backend.inductor_config,
+                self.compilation_config,
+                compile_range=range_entry.compile_range,
+                graph_index=self.piecewise_compile_index,
+                num_graphs=self.total_piecewise_compiles,
             )
-            if hasattr(torch._functorch.config, "force_non_lazy_backward_lowering"):
-                config_patches["force_non_lazy_backward_lowering"] = False
-            with torch._functorch.config.patch(**config_patches):
-                range_entry.runnable = self.vllm_backend.compiler_manager.compile(
-                    self.graph,
-                    args_list,
-                    self.vllm_backend.inductor_config,
-                    self.compilation_config,
-                    compile_range=range_entry.compile_range,
-                    graph_index=self.piecewise_compile_index,
-                    num_graphs=self.total_piecewise_compiles,
-                )
 
             range_entry.compiled = True
 

vllm/compilation/wrapper.py

@@ -349,6 +349,9 @@ def reset_compile_wrapper(model: torch.nn.Module) -> None:
     compilation_counter.num_cache_entries_updated = 0
     compilation_counter.num_compiled_artifacts_saved = 0
     compilation_counter.stock_torch_compile_count = 0
+    compilation_counter.num_aot_compiles = 0
+    compilation_counter.num_aot_artifacts_saved = 0
+    compilation_counter.num_aot_artifacts_loaded = 0
 
     # Clear the AOT compiled function so the model is forced to
     # recompile on the next call. Without this, decorators.py

vllm/config/kv_transfer.py

@@ -24,9 +24,9 @@ class KVTransferConfig:
     engine_id: str | None = None
     """The engine id for KV transfers."""
 
-    kv_buffer_device: str = "cuda"
-    """The device used by kv connector to buffer the KV cache. Choices are 
-    'cuda' and 'cpu'."""
+    kv_buffer_device: str | None = None
+    """The device used by kv connector to buffer the KV cache. Choices are
+    'cuda','cpu' and 'xpu'."""
 
     kv_buffer_size: float = 1e9
     """The buffer size for TorchDistributedConnector. Measured in number of
@@ -100,6 +100,11 @@ class KVTransferConfig:
                 f"is set, supported roles are {get_args(KVRole)}"
             )
 
+        if self.kv_buffer_device is None:
+            from vllm.platforms import current_platform
+
+            self.kv_buffer_device = current_platform.device_type
+
     @property
     def is_kv_transfer_instance(self) -> bool:
         return self.kv_connector is not None and self.kv_role in get_args(KVRole)

vllm/config/model.py

@@ -217,12 +217,13 @@ class ModelConfig:
     """Whether to disable sliding window. If True, we will disable the sliding
     window functionality of the model, capping to sliding window size. If the
     model does not support sliding window, this argument is ignored."""
-    disable_cascade_attn: bool = False
+    disable_cascade_attn: bool = True
     """Disable cascade attention for V1. While cascade attention does not
     change the mathematical correctness, disabling it could be useful for
-    preventing potential numerical issues. Note that even if this is set to
-    False, cascade attention will be only used when the heuristic tells that
-    it's beneficial."""
+    preventing potential numerical issues. This defaults to True, so users
+    must opt in to cascade attention by setting this to False. Even when this
+    is set to False, cascade attention will only be used when the heuristic
+    tells that it's beneficial."""
     skip_tokenizer_init: bool = False
     """Skip initialization of tokenizer and detokenizer. Expects valid
     `prompt_token_ids` and `None` for prompt from the input. The generated
@@ -531,6 +532,22 @@ class ModelConfig:
         self._architecture = arch
         logger.info("Resolved architecture: %s", arch)
 
+        # Set default tokenizer modes based on model architecture
+        if self.tokenizer_mode == "auto":
+            if arch == "Grok1ForCausalLM":
+                self.tokenizer_mode = "grok2"
+            elif arch == "MoonshotKimiaForCausalLM":
+                self.tokenizer_mode = "kimi_audio"
+            elif arch == "QwenVLForConditionalGeneration":
+                self.tokenizer_mode = "qwen_vl"
+
+            if self.tokenizer_mode != "auto":
+                logger.info(
+                    "Defaulting to tokenizer_mode=%r for %s",
+                    self.tokenizer_mode,
+                    arch,
+                )
+
         # Init pooler config if needed
         if self.runner_type == "pooling":
             if self.pooler_config is None:
@@ -1123,6 +1140,7 @@ class ModelConfig:
             return bool(self.hf_config.is_mm_prefix_lm)
         # fallback to list of known models
         MM_PREFIX_LM_MODELS = (
+            "bagel",
             "gemma3",
             "molmo2",
             "paligemma",

vllm/config/speculative.py

@@ -57,6 +57,10 @@ SpeculativeMethod = Literal[
     EagleModelTypes,
     NgramGPUTypes,
 ]
+RejectionSampleMethod = Literal[
+    "strict",
+    "probabilistic",
+]
 
 
 @config
@@ -171,6 +175,12 @@ class SpeculativeConfig:
     """Load config for the draft model. If not specified, will use the load
     config from the target model."""
 
+    rejection_sample_method: RejectionSampleMethod = "strict"
+    """Whether to use strict (target and draft sampled tokens match exactly)
+    or probabilistic rejection sampling. Both respect the target model
+    distribution, but the latter yields a higher acceptance rate at the cost
+    of more memory to cache draft logits."""
+
     def compute_hash(self) -> str:
         """
         WARNING: Whenever a new field is added to this config,
@@ -779,6 +789,10 @@ class SpeculativeConfig:
             "hunyuan_v1_dense",
             "afmoe",
             "nemotron_h",
+            "deepseek_v2",
+            "deepseek_v3",
+            "kimi_k2",
+            "kimi_k25",
         ]
         if (
             self.method in ("eagle3", "extract_hidden_states")