Merge ROCm:fused_aiter_rope_kvcache (#33443 )

Merge ROCm:rocm_ck_moe_a16w4 (#34301 )
small adjustment to wvSplitKrc (#34410 )
2026-02-23 17:23:40 +00:00 · 2026-02-23 17:22:16 +00:00 · 2026-02-23 17:15:20 +00:00 · 2026-02-23 17:15:20 +00:00 · 2026-02-23 17:15:20 +00:00 · 2026-02-23 17:15:20 +00:00
544 changed files with 9199 additions and 21151 deletions
@@ -1,7 +1,6 @@
-group: Hardware - AMD Build 
+group: Hardware
 steps:
  - label: "AMD: :docker: build image"
-    key: image-build-amd
    depends_on: []
    device: amd_cpu
    no_plugin: true
@@ -9,10 +9,8 @@ import json
 import os
 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
@@ -277,131 +275,6 @@ def _apply_two_decimals(
    return styler.format({c: "{:.2f}" for c in num_cols}, na_rep="")


-# -----------------------------
-# Export helpers (Excel + CSV)
-# -----------------------------
-def _sanitize_sheet_name(name: str) -> str:
-    """
-    Excel sheet constraints:
-      - max 31 chars
-      - cannot contain: : \ / ? * [ ]
-      - cannot be empty
-    """
-    name = "sheet" if name is None else str(name)
-    name = re.sub(r"[:\\/?*\[\]]", "_", name)
-    name = name.strip().strip("'")
-    name = re.sub(r"\s+", " ", name)
-    if not name:
-        name = "sheet"
-    return name[:31]
-
-
-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]
-    ilen = d.get("Input Len", "")
-    olen = d.get("Output Len", "")
-    lens = f"_{ilen}x{olen}" if ilen != "" and olen != "" else ""
-    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
-    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
-
-
-def _safe_filename(s: str) -> str:
-    s = re.sub(r"[^\w\-.]+", "_", str(s).strip())
-    return s[:180] if len(s) > 180 else s
-
-
-# -----------------------------
-# vLLM environment export helper
-# -----------------------------
-def _parse_vllm_env_txt(env_path: Path) -> pd.DataFrame:
-    """Parse vllm_env.txt into a flat table (Section, Key, Value).
-
-    Supports:
-      - section headers as standalone lines (no ':' or '=')
-      - key-value lines like 'OS: Ubuntu ...'
-      - env var lines like 'HF_HOME=/data/hf'
-    """
-    lines = env_path.read_text(encoding="utf-8", errors="replace").splitlines()
-    section = "General"
-    rows: list[dict] = []
-
-    def set_section(s: str):
-        nonlocal section
-        s = (s or "").strip()
-        if s:
-            section = s
-
-    for raw in lines:
-        stripped = raw.strip()
-        if not stripped:
-            continue
-        # divider lines like =====
-        if set(stripped) <= {"="}:
-            continue
-
-        # section header heuristic: short standalone line
-        if ":" not in stripped and "=" not in stripped and len(stripped) <= 64:
-            if stripped.lower().startswith("collecting environment information"):
-                continue
-            set_section(stripped)
-            continue
-
-        # env var style: KEY=VALUE (and not a URL with :)
-        if "=" in stripped and ":" not in stripped:
-            k, v = stripped.split("=", 1)
-            k = k.strip()
-            v = v.strip()
-            if k:
-                rows.append({"Section": section, "Key": k, "Value": v})
-            continue
-
-        # key: value
-        if ":" in stripped:
-            k, v = stripped.split(":", 1)
-            k = k.strip()
-            v = v.strip()
-            if k:
-                rows.append({"Section": section, "Key": k, "Value": v})
-            continue
-
-    return pd.DataFrame(rows, columns=["Section", "Key", "Value"])
-
-
-def _load_env_df_for_inputs(args, files: list[str]) -> pd.DataFrame | None:
-    """Load vllm_env.txt next to the *original* input JSON file.
-
-    Note: when only one -f is provided, the script may split JSON into ./splits/...,
-    but vllm_env.txt typically lives next to the original benchmark_results.json.
-    """
-    base_dir: Path | None = None
-    if getattr(args, "file", None):
-        base_dir = Path(args.file[0]).resolve().parent
-    elif files:
-        base_dir = Path(files[0]).resolve().parent
-    if base_dir is None:
-        return None
-
-    env_path = base_dir / "vllm_env.txt"
-    if not env_path.exists():
-        return None
-    df = _parse_vllm_env_txt(env_path)
-    return df
-
-
 # -----------------------------
 # Valid max concurrency summary helpers
 # -----------------------------
@@ -555,6 +428,7 @@ def build_valid_max_concurrency_summary_html(

    summary_df = pd.DataFrame(rows)

+    # --- Coerce numeric columns so Styler doesn't miss them due to object dtype ---
    for c in summary_df.columns:
        if c == "Configuration":
            continue
@@ -562,10 +436,12 @@ def build_valid_max_concurrency_summary_html(

    both_col = f"Max {conc_col} (Both)"

+    # --- Strict 2-decimal formatting for ALL non-Configuration columns ---
    formatters = {}
    for c in summary_df.columns:
        if c == "Configuration":
            continue
+        # default argument binds per-column formatter correctly
        formatters[c] = lambda v: "" if pd.isna(v) else f"{float(v):.2f}"

    styler = summary_df.style.format(formatters)
@@ -584,95 +460,6 @@ def build_valid_max_concurrency_summary_html(
    return title + styler.to_html(table_attributes='border="1" class="dataframe"')


-def build_valid_max_concurrency_summary_df(
-    tput_group_df: pd.DataFrame | None,
-    ttft_group_df: pd.DataFrame | None,
-    tpot_group_df: pd.DataFrame | None,
-    conc_col: str,
-    args,
-) -> pd.DataFrame | None:
-    if ttft_group_df is None and tpot_group_df is None:
-        return None
-
-    ttft_cols = (
-        _config_value_columns(ttft_group_df, conc_col)
-        if ttft_group_df is not None
-        else []
-    )
-    tpot_cols = (
-        _config_value_columns(tpot_group_df, conc_col)
-        if tpot_group_df is not None
-        else []
-    )
-    tput_cols = (
-        _config_value_columns(tput_group_df, conc_col)
-        if tput_group_df is not None
-        else []
-    )
-
-    if ttft_group_df is not None and tpot_group_df is not None:
-        cfg_cols = [c for c in ttft_cols if c in tpot_cols]
-        if tput_group_df is not None:
-            cfg_cols = [c for c in cfg_cols if c in tput_cols] or cfg_cols
-    else:
-        cfg_cols = ttft_cols or tpot_cols
-
-    if not cfg_cols:
-        cfg_cols = sorted(set(ttft_cols) | set(tpot_cols) | set(tput_cols), key=str)
-
-    rows = []
-    for cfg in cfg_cols:
-        ttft_max = (
-            _max_concurrency_ok(ttft_group_df, conc_col, cfg, args.ttft_max_ms)
-            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)
-            if tpot_group_df is not None
-            else pd.NA
-        )
-        both = (
-            pd.NA
-            if (pd.isna(ttft_max) or pd.isna(tpot_max))
-            else min(ttft_max, tpot_max)
-        )
-
-        tput_at_both = (
-            _value_at_concurrency(tput_group_df, conc_col, cfg, both)
-            if tput_group_df is not None
-            else pd.NA
-        )
-        ttft_at_both = (
-            _value_at_concurrency(ttft_group_df, conc_col, cfg, both)
-            if ttft_group_df is not None
-            else pd.NA
-        )
-        tpot_at_both = (
-            _value_at_concurrency(tpot_group_df, conc_col, cfg, both)
-            if tpot_group_df is not None
-            else pd.NA
-        )
-
-        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} (Both)": both,
-                "Output Tput @ Both (tok/s)": tput_at_both,
-                "TTFT @ Both (ms)": ttft_at_both,
-                "TPOT @ Both (ms)": tpot_at_both,
-            }
-        )
-
-    df = pd.DataFrame(rows)
-    for c in df.columns:
-        if c != "Configuration":
-            df[c] = pd.to_numeric(df[c], errors="coerce")
-    return df
-
-
 # -----------------------------
 # Plot helper
 # -----------------------------
@@ -750,21 +537,6 @@ def build_parser() -> argparse.ArgumentParser:
        default=100.0,
        help="Reference limit for TPOT plots (ms)",
    )
-
-    # ---- NEW: export options ----
-    parser.add_argument(
-        "--excel-out",
-        type=str,
-        default="perf_comparison.xlsx",
-        help="Write one sheet per (Model, Dataset, Input Len, Output Len).",
-    )
-    parser.add_argument(
-        "--csv-out-dir",
-        type=str,
-        default="",
-        help="If set, write per-group per-metric CSVs into this directory.",
-    )
-
    return parser


@@ -885,6 +657,7 @@ def maybe_write_plot(
        markers=True,
    )

+    # Ensure plot hover + y tick labels are also 2 decimals.
    fig.update_traces(hovertemplate="%{y:.2f}<extra></extra>")
    fig.update_yaxes(tickformat=".2f")

@@ -957,151 +730,87 @@ def write_report_group_first(
        for metric_label, (df, _) in metric_cache.items()
    }

-    csv_dir = Path(args.csv_out_dir) if args.csv_out_dir else None
-    if csv_dir:
-        csv_dir.mkdir(parents=True, exist_ok=True)
+    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:
+            gkey_tuple = normalize_group_key(gkey)
+            suffix = build_group_suffix(group_cols_canonical, gkey_tuple)
+            sub_path = group_filename(gkey_tuple)
+            group_header = (
+                '<div style="font-size: 1.4em; font-weight: 700; '
+                'margin: 18px 0 10px 0;">'
+                f"{_html.escape(suffix)}"
+                "</div>\n"
+            )

-    excel_path = args.excel_out or "perf_comparison.xlsx"
-    with pd.ExcelWriter(excel_path, engine="openpyxl") as xw:
-        # ---- 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)
-        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:
-                gkey_tuple = normalize_group_key(gkey)
-                suffix = build_group_suffix(group_cols_canonical, gkey_tuple)
-                sub_path = group_filename(gkey_tuple)
-                group_header = (
-                    '<div style="font-size: 1.4em; font-weight: 700; '
-                    'margin: 18px 0 10px 0;">'
-                    f"{_html.escape(suffix)}"
-                    "</div>\n"
+            main_fh.write(group_header)
+            with open(sub_path, "w", encoding="utf-8") as sub_fh:
+                sub_fh.write('<meta charset="utf-8">\n')
+                sub_fh.write(group_header)
+                tput_group_df = None
+                ttft_group_df = None
+                tpot_group_df = None
+                conc_col = args.xaxis
+
+                for metric_label in plan.data_cols:
+                    gb = metric_groupbys[metric_label]
+                    df_sorted, raw_data_cols = metric_cache[metric_label]
+
+                    try:
+                        group_df = gb.get_group(gkey)
+                    except KeyError:
+                        missing = (
+                            '<div style="font-size: 1.1em; font-weight: 600; '
+                            'margin: 10px 0;">'
+                            f"{_html.escape(metric_label)} — missing for this group"
+                            "</div>\n"
+                        )
+
+                        main_fh.write(missing)
+                        sub_fh.write(missing)
+                        continue
+
+                    if conc_col not in group_df.columns:
+                        conc_col = _find_concurrency_col(group_df)
+
+                    mn = metric_label.lower().strip()
+                    if "tok/s" in mn:
+                        tput_group_df = group_df
+                    elif "ttft" in mn:
+                        ttft_group_df = group_df
+                    elif mn in ("p99", "median") or "tpot" in mn:
+                        tpot_group_df = group_df
+
+                    display_group = group_df.drop(
+                        columns=group_cols_canonical, errors="ignore"
+                    )
+
+                    html = render_metric_table_html(
+                        display_group, metric_label, suffix, args
+                    )
+                    main_fh.write(html)
+                    sub_fh.write(html)
+
+                    maybe_write_plot(
+                        main_fh,
+                        sub_fh,
+                        group_df=group_df,
+                        raw_data_cols=raw_data_cols,
+                        metric_label=metric_label,
+                        y_axis_col=y_axis_col,
+                        args=args,
+                    )
+
+                summary_html = build_valid_max_concurrency_summary_html(
+                    tput_group_df=tput_group_df,
+                    ttft_group_df=ttft_group_df,
+                    tpot_group_df=tpot_group_df,
+                    conc_col=conc_col,
+                    args=args,
                )
-
-                main_fh.write(group_header)
-
-                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}")
-
-                excel_blocks: list[tuple[str, pd.DataFrame]] = []
-
-                with open(sub_path, "w", encoding="utf-8") as sub_fh:
-                    sub_fh.write('<meta charset="utf-8">\n')
-                    sub_fh.write(group_header)
-                    tput_group_df = None
-                    ttft_group_df = None
-                    tpot_group_df = None
-                    conc_col = args.xaxis
-
-                    for metric_label in plan.data_cols:
-                        gb = metric_groupbys[metric_label]
-                        df_sorted, raw_data_cols = metric_cache[metric_label]
-
-                        try:
-                            group_df = gb.get_group(gkey)
-                        except KeyError:
-                            missing = (
-                                '<div style="font-size: 1.1em; font-weight: 600; '
-                                'margin: 10px 0;">'
-                                f"{_html.escape(metric_label)} — missing for this group"
-                                "</div>\n"
-                            )
-                            main_fh.write(missing)
-                            sub_fh.write(missing)
-                            continue
-
-                        if conc_col not in group_df.columns:
-                            conc_col = _find_concurrency_col(group_df)
-
-                        mn = metric_label.lower().strip()
-                        if "tok/s" in mn:
-                            tput_group_df = group_df
-                        elif "ttft" in mn:
-                            ttft_group_df = group_df
-                        elif mn in ("p99", "median") or "tpot" in mn:
-                            tpot_group_df = group_df
-
-                        display_group = group_df.drop(
-                            columns=group_cols_canonical, errors="ignore"
-                        )
-
-                        html = render_metric_table_html(
-                            display_group, metric_label, suffix, args
-                        )
-                        main_fh.write(html)
-                        sub_fh.write(html)
-
-                        maybe_write_plot(
-                            main_fh,
-                            sub_fh,
-                            group_df=group_df,
-                            raw_data_cols=raw_data_cols,
-                            metric_label=metric_label,
-                            y_axis_col=y_axis_col,
-                            args=args,
-                        )
-
-                        excel_blocks.append(
-                            (metric_label, display_group.reset_index(drop=True))
-                        )
-                        if csv_dir:
-                            fn = _safe_filename(
-                                f"{sheet}__{metric_label}".replace(" ", "_").replace(
-                                    "/", "_"
-                                )
-                            )
-                            display_group.to_csv(csv_dir / f"{fn}.csv", index=False)
-
-                    summary_html = build_valid_max_concurrency_summary_html(
-                        tput_group_df=tput_group_df,
-                        ttft_group_df=ttft_group_df,
-                        tpot_group_df=tpot_group_df,
-                        conc_col=conc_col,
-                        args=args,
-                    )
-                    if summary_html:
-                        main_fh.write(summary_html)
-                        sub_fh.write(summary_html)
-
-                    summary_df = build_valid_max_concurrency_summary_df(
-                        tput_group_df=tput_group_df,
-                        ttft_group_df=ttft_group_df,
-                        tpot_group_df=tpot_group_df,
-                        conc_col=conc_col,
-                        args=args,
-                    )
-                    if summary_df is not None:
-                        excel_blocks.append(
-                            ("Valid Max Concurrency Summary", summary_df)
-                        )
-                        if csv_dir:
-                            fn = _safe_filename(
-                                f"{sheet}__Valid_Max_Concurrency_Summary"
-                            )
-                            summary_df.to_csv(csv_dir / f"{fn}.csv", index=False)
-
-                _write_tables_to_excel_sheet(xw, sheet, excel_blocks)
-
-    print(f"Wrote Excel: {excel_path}")
-    if csv_dir:
-        print(f"Wrote CSVs under: {csv_dir}")
+                if summary_html:
+                    main_fh.write(summary_html)
+                    sub_fh.write(summary_html)


 def main():
@@ -1,4 +1,6 @@
 #!/bin/bash
+
+# This script should be run inside the CI process
 # This script assumes that we are already inside the vllm/ directory
 # Benchmarking results will be available inside vllm/benchmarks/results/

@@ -7,11 +9,6 @@
 set -x
 set -o pipefail

-# Environment-driven debug controls (like ON_CPU=1)
-DRY_RUN="${DRY_RUN:-0}"
-MODEL_FILTER="${MODEL_FILTER:-}"
-DTYPE_FILTER="${DTYPE_FILTER:-}"
-
 check_gpus() {
  if command -v nvidia-smi; then
    # check the number of GPUs and GPU type.
@@ -115,12 +112,13 @@ json2envs() {
 }

 wait_for_server() {
+  # wait for vllm server to start
+  # return 1 if vllm server crashes
  local timeout_val="1200"
  timeout "$timeout_val" bash -c '
-    until curl -sf http://localhost:8000/v1/models >/dev/null; do
+    until curl -X POST localhost:8000/v1/completions; do
      sleep 1
-    done
-  '
+    done' && return 0 || return 1
 }

 kill_processes_launched_by_current_bash() {
@@ -254,16 +252,37 @@ run_benchmark_tests() {
  done
 }

-run_latency_tests() { run_benchmark_tests "latency" "$1"; }
-run_startup_tests() { run_benchmark_tests "startup" "$1"; }
-run_throughput_tests() { run_benchmark_tests "throughput" "$1"; }
+run_latency_tests() {
+  run_benchmark_tests "latency" "$1"
+}

-merge_serving_tests_stream() {
-  # Emit merged serving test objects, optionally filtered by MODEL_FILTER/DTYPE_FILTER in DRY_RUN mode.
-  # This helper does NOT modify JSON; it only filters the stream in dry-run mode.
-  local serving_test_file="$1"
-  # shellcheck disable=SC2016
-  local merged='
+run_startup_tests() {
+  run_benchmark_tests "startup" "$1"
+}
+
+run_throughput_tests() {
+  run_benchmark_tests "throughput" "$1"
+}
+
+run_serving_tests() {
+  # run serving tests using `vllm bench serve` command
+  # $1: a json file specifying serving test cases
+  #
+  # Supported JSON formats:
+  # 1) Plain format: top-level array
+  #    [ { "test_name": "...", "server_parameters": {...}, ... }, ... ]
+  #
+  # 2) Default parameters field + plain format tests
+  #    {
+  #      "defaults": { ... },
+  #      "tests": [ { "test_name": "...", "server_parameters": {...}, ... }, ... ]
+  #    }
+
+  local serving_test_file
+  serving_test_file=$1
+
+  # Iterate over serving tests
+  jq -c '
    if type == "array" then
      # Plain format: test cases array
      .[]
@@ -285,50 +304,7 @@ merge_serving_tests_stream() {
    else
      error("Unsupported serving test file format: must be array or object with .tests")
    end
-  '
-
-  jq -c "$merged" "$serving_test_file" | \
-  if [[ "${DRY_RUN:-0}" == "1" && ( "${MODEL_FILTER}${DTYPE_FILTER}" != "" ) ]]; then
-    jq -c --arg model "$MODEL_FILTER" --arg dtype "$DTYPE_FILTER" '
-      select((($model|length)==0)
-             or ((.server_parameters.model // "") == $model)
-             or ((.client_parameters.model // "") == $model))
-      | select((($dtype|length)==0) or ((.server_parameters.dtype // "") == $dtype))
-    '
-  else
-    cat
-  fi
-}
-
-run_serving_tests() {
-  # run serving tests using `vllm bench serve` command
-  # $1: a json file specifying serving test cases
-  #
-  # Supported JSON formats:
-  # 1) Plain format: top-level array
-  #    [ { "test_name": "...", "server_parameters": {...}, ... }, ... ]
-  #
-  # 2) Default parameters field + plain format tests
-  #    {
-  #      "defaults": { ... },
-  #      "tests": [ { "test_name": "...", "server_parameters": {...}, ... }, ... ]
-  #    }
-
-  local serving_test_file
-  serving_test_file=$1
-
-  # In dry-run mode, if filters are provided but no tests match, fail fast.
-  if [[ "${DRY_RUN:-0}" == "1" && ( "${MODEL_FILTER}${DTYPE_FILTER}" != "" ) ]]; then
-    local count
-    count=$(merge_serving_tests_stream "$serving_test_file" | wc -l | tr -d ' ')
-    if [[ "$count" -eq 0 ]]; then
-      echo "No matching serving tests found in $serving_test_file for model='$MODEL_FILTER' dtype='$DTYPE_FILTER'." >&2
-      return 0
-    fi
-  fi
-
-  # Iterate over serving tests (merged + optional filtered stream)
-  merge_serving_tests_stream "$serving_test_file" | while read -r params; do
+  ' "$serving_test_file" | while read -r params; do
    # get the test name, and append the GPU type back to it.
    test_name=$(echo "$params" | jq -r '.test_name')
    if [[ ! "$test_name" =~ ^serving_ ]]; then
@@ -397,7 +373,7 @@ run_serving_tests() {
    echo "Server command: $server_command"
    # support remote vllm server
    client_remote_args=""
-    if [[ -z "${REMOTE_HOST}" && "${DRY_RUN:-0}" != "1" ]]; then
+    if [[ -z "${REMOTE_HOST}" ]]; then
      bash -c "$server_command" &
      server_pid=$!
      # wait until the server is alive
@@ -408,9 +384,6 @@ run_serving_tests() {
        echo ""
        echo "vLLM failed to start within the timeout period."
      fi
-    elif [[ "${DRY_RUN:-0}" == "1" ]]; then
-        # dry-run: don't start server
-        echo "Dry Run."
    else
      server_command="Using Remote Server $REMOTE_HOST $REMOTE_PORT"
      if [[ ${REMOTE_PORT} ]]; then
@@ -429,7 +402,9 @@ run_serving_tests() {
    for qps in $qps_list; do
      # remove the surrounding single quote from qps
      if [[ "$qps" == *"inf"* ]]; then
+        echo "qps was $qps"
        qps="inf"
+        echo "now qps is $qps"
      fi

      # iterate over different max_concurrency
@@ -450,9 +425,7 @@ run_serving_tests() {
        echo "Running test case $test_name with qps $qps"
        echo "Client command: $client_command"

-        if [[ "${DRY_RUN:-0}" != "1" ]]; then
-          bash -c "$client_command"
-        fi
+        bash -c "$client_command"

        # record the benchmarking commands
        jq_output=$(jq -n \
@@ -470,15 +443,12 @@ run_serving_tests() {
    done

    # clean up
-    if [[ "${DRY_RUN:-0}" != "1" ]]; then
-      kill -9 $server_pid
-      kill_gpu_processes
-    fi
+    kill -9 $server_pid
+    kill_gpu_processes
  done
 }

 main() {
-
  local ARCH
  ARCH=''
  if [[ "$ON_CPU" == "1" ]]; then
@@ -488,13 +458,7 @@ main() {
     check_gpus
     ARCH="$arch_suffix"
  fi
-
-  # DRY_RUN does not execute vLLM; do not require HF_TOKEN.
-  if [[ "${DRY_RUN:-0}" != "1" ]]; then
-    check_hf_token
-  else
-    echo "DRY_RUN=1 -> skip HF_TOKEN validation"
-  fi
+  check_hf_token

  # dependencies
  (which wget && which curl) || (apt-get update && apt-get install -y wget curl)
@@ -515,16 +479,11 @@ main() {

  # dump vllm info via vllm collect-env
  env_output=$(vllm collect-env)
+
  echo "$env_output" >"$RESULTS_FOLDER/vllm_env.txt"

  # benchmarking
-  run_serving_tests $QUICK_BENCHMARK_ROOT/tests/"${SERVING_JSON:-serving-tests$ARCH.json}" || exit $?
-
-  if [[ "${DRY_RUN:-0}" == "1" ]]; then
-    echo "DRY_RUN=1 -> skip latency/startup/throughput suites"
-    exit 0
-  fi
-
+  run_serving_tests $QUICK_BENCHMARK_ROOT/tests/"${SERVING_JSON:-serving-tests$ARCH.json}"
  run_latency_tests $QUICK_BENCHMARK_ROOT/tests/"${LATENCY_JSON:-latency-tests$ARCH.json}"
  run_startup_tests $QUICK_BENCHMARK_ROOT/tests/"${STARTUP_JSON:-startup-tests$ARCH.json}"
  run_throughput_tests $QUICK_BENCHMARK_ROOT/tests/"${THROUGHPUT_JSON:-throughput-tests$ARCH.json}"
@@ -1,41 +0,0 @@
-{
-  "defaults": {
-    "qps_list": [
-      "inf"
-    ],
-    "max_concurrency_list": [
-      32,
-      64,
-      128
-    ],
-    "server_environment_variables": {
-      "VLLM_RPC_TIMEOUT": 100000,
-      "VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1,
-      "VLLM_ENGINE_ITERATION_TIMEOUT_S": 120,
-      "VLLM_CPU_SGL_KERNEL": 1,
-      "VLLM_CPU_KVCACHE_SPACE": 40
-    },
-    "server_parameters": {
-      "dtype": "bfloat16",
-      "model": "jinaai/jina-embeddings-v3",
-      "trust_remote_code": ""
-    },
-    "client_parameters": {
-      "model": "jinaai/jina-embeddings-v3",
-      "backend": "openai-embeddings",
-      "endpoint": "/v1/embeddings",
-      "dataset_name": "sharegpt",
-      "dataset_path": "ShareGPT_V3_unfiltered_cleaned_split.json",
-      "num_prompts": 200
-    }
-  },
-  "tests": [
-    {
-      "test_name": "serving_jina_embed_v3_tp1_sharegpt",
-      "server_parameters": {
-        "tensor_parallel_size": 1
-      },
-      "client_parameters": {}
-    }
-  ]
-}
@@ -1,283 +0,0 @@
-{
-  "defaults": {
-    "qps_list": [
-      "inf"
-    ],
-    "max_concurrency_list": [12, 16, 24, 32, 64, 128, 200],
-    "server_environment_variables": {
-      "VLLM_RPC_TIMEOUT": 100000,
-      "VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1,
-      "VLLM_ENGINE_ITERATION_TIMEOUT_S": 120,
-      "VLLM_CPU_SGL_KERNEL": 1,
-      "VLLM_CPU_KVCACHE_SPACE": 40
-    },
-    "server_parameters": {
-      "model": "meta-llama/Llama-3.1-8B-Instruct",
-      "tensor_parallel_size": 1,
-      "dtype": "bfloat16",
-      "distributed_executor_backend": "mp",
-      "block_size": 128,
-      "trust_remote_code": "",
-      "disable_log_stats": "",
-      "max_num_batched_tokens": 2048,
-      "max_num_seqs": 256
-    },
-    "client_parameters": {
-      "model": "meta-llama/Llama-3.1-8B-Instruct",
-      "backend": "vllm",
-      "ignore-eos": "",
-      "num_prompts": 200
-    }
-  },
-  "tests": [
-    {
-      "test_name": "serving_llama8B_tp1_sharegpt",
-      "server_parameters": {
-        "tensor_parallel_size": 1
-      },
-      "client_parameters": {
-        "dataset_name": "sharegpt",
-        "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json"
-      }
-    },
-    {
-      "test_name": "serving_llama8B_tp2_sharegpt",
-      "server_parameters": {
-        "tensor_parallel_size": 2
-      },
-      "client_parameters": {
-        "dataset_name": "sharegpt",
-        "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json"
-      }
-    },
-    {
-      "test_name": "serving_llama8B_tp1_random_128_128",
-      "server_parameters": {
-        "tensor_parallel_size": 1
-      },
-      "client_parameters": {
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "random-output-len": 128
-      }
-    },
-    {
-      "test_name": "serving_llama8B_tp2_random_128_128",
-      "server_parameters": {
-        "tensor_parallel_size": 2
-      },
-      "client_parameters": {
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "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": {
-        "tensor_parallel_size": 1
-      },
-      "client_parameters": {
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "random-output-len": 2048
-      }
-    },
-    {
-      "test_name": "serving_llama8B_tp2_random_128_2048",
-      "server_parameters": {
-        "tensor_parallel_size": 2
-      },
-      "client_parameters": {
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "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": {
-        "tensor_parallel_size": 1
-      },
-      "client_parameters": {
-        "dataset_name": "random",
-        "random-input-len": 2048,
-        "random-output-len": 128
-      }
-    },
-    {
-      "test_name": "serving_llama8B_tp2_random_2048_128",
-      "server_parameters": {
-        "tensor_parallel_size": 2
-      },
-      "client_parameters": {
-        "dataset_name": "random",
-        "random-input-len": 2048,
-        "random-output-len": 128
-      }
-    },
-    {
-      "test_name": "serving_llama8B_tp4_random_2048_128",
-      "server_parameters": {
-        "tensor_parallel_size": 4
-      },
-      "client_parameters": {
-        "dataset_name": "random",
-        "random-input-len": 2048,
-        "random-output-len": 128
-      }
-    },
-    {
-      "test_name": "serving_llama8B_int4_tp1_random_128_128",
-      "server_parameters": {
-        "model": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
-        "tensor_parallel_size": 1
-      },
-      "client_parameters": {
-        "model": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "random-output-len": 128
-      }
-    },
-    {
-      "test_name": "serving_llama8B_int4_tp2_random_128_128",
-      "server_parameters": {
-        "model": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
-        "tensor_parallel_size": 2
-      },
-      "client_parameters": {
-        "model": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "random-output-len": 128
-      }
-    },
-    {
-      "test_name": "serving_llama8B_int4_tp4_random_128_128",
-      "server_parameters": {
-        "model": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
-        "tensor_parallel_size": 4
-      },
-      "client_parameters": {
-        "model": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "random-output-len": 128
-      }
-    },
-    {
-      "test_name": "serving_llama3B_tp1_random_128_128",
-      "server_parameters": {
-        "model": "meta-llama/Llama-3.2-3B-Instruct",
-        "tensor_parallel_size": 1
-      },
-      "client_parameters": {
-        "model": "meta-llama/Llama-3.2-3B-Instruct",
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "random-output-len": 128
-      }
-    },
-    {
-      "test_name": "serving_granite2B_tp1_random_128_128",
-      "server_parameters": {
-        "model": "ibm-granite/granite-3.2-2b-instruct",
-        "tensor_parallel_size": 1
-      },
-      "client_parameters": {
-        "model": "ibm-granite/granite-3.2-2b-instruct",
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "random-output-len": 128
-      }
-    },
-    {
-      "test_name": "serving_qwen1.7B_tp1_random_128_128",
-      "server_parameters": {
-        "model": "Qwen/Qwen3-1.7B",
-        "tensor_parallel_size": 1
-      },
-      "client_parameters": {
-        "model": "Qwen/Qwen3-1.7B",
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "random-output-len": 128
-      }
-    },
-    {
-      "test_name": "serving_qwen4B_tp1_random_128_128",
-      "server_parameters": {
-        "model": "Qwen/Qwen3-4B",
-        "tensor_parallel_size": 1
-      },
-      "client_parameters": {
-        "model": "Qwen/Qwen3-4B",
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "random-output-len": 128
-      }
-    },
-    {
-      "test_name": "serving_qwen8B_tp1_random_128_128",
-      "server_parameters": {
-        "model": "Qwen/Qwen3-8B",
-        "tensor_parallel_size": 1
-      },
-      "client_parameters": {
-        "model": "Qwen/Qwen3-8B",
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "random-output-len": 128
-      }
-    },
-    {
-      "test_name": "serving_glm9B_tp1_random_128_128",
-      "server_parameters": {
-        "model": "zai-org/glm-4-9b-hf",
-        "tensor_parallel_size": 1
-      },
-      "client_parameters": {
-        "model": "zai-org/glm-4-9b-hf",
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "random-output-len": 128
-      }
-    },
-    {
-      "test_name": "serving_gemma7B_tp1_random_128_128",
-      "server_parameters": {
-        "model": "google/gemma-7b",
-        "tensor_parallel_size": 1
-      },
-      "client_parameters": {
-        "model": "google/gemma-7b",
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "random-output-len": 128
-      }
-    }
-  ]
-}
@@ -148,6 +148,136 @@
        "random-input-len": 2048,
        "random-output-len": 128
      }
+    },
+    {
+      "test_name": "serving_llama8B_int4_tp1_random_128_128",
+      "server_parameters": {
+        "model": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
+        "tensor_parallel_size": 1
+      },
+      "client_parameters": {
+        "model": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
+        "dataset_name": "random",
+        "random-input-len": 128,
+        "random-output-len": 128
+      }
+    },
+    {
+      "test_name": "serving_llama8B_int4_tp2_random_128_128",
+      "server_parameters": {
+        "model": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
+        "tensor_parallel_size": 2
+      },
+      "client_parameters": {
+        "model": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
+        "dataset_name": "random",
+        "random-input-len": 128,
+        "random-output-len": 128
+      }
+    },
+    {
+      "test_name": "serving_llama8B_int4_tp4_random_128_128",
+      "server_parameters": {
+        "model": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
+        "tensor_parallel_size": 4
+      },
+      "client_parameters": {
+        "model": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
+        "dataset_name": "random",
+        "random-input-len": 128,
+        "random-output-len": 128
+      }
+    },
+    {
+      "test_name": "serving_llama3B_tp1_random_128_128",
+      "server_parameters": {
+        "model": "meta-llama/Llama-3.2-3B-Instruct",
+        "tensor_parallel_size": 1
+      },
+      "client_parameters": {
+        "model": "meta-llama/Llama-3.2-3B-Instruct",
+        "dataset_name": "random",
+        "random-input-len": 128,
+        "random-output-len": 128
+      }
+    },
+    {
+      "test_name": "serving_granite2B_tp1_random_128_128",
+      "server_parameters": {
+        "model": "ibm-granite/granite-3.2-2b-instruct",
+        "tensor_parallel_size": 1
+      },
+      "client_parameters": {
+        "model": "ibm-granite/granite-3.2-2b-instruct",
+        "dataset_name": "random",
+        "random-input-len": 128,
+        "random-output-len": 128
+      }
+    },
+    {
+      "test_name": "serving_qwen1.7B_tp1_random_128_128",
+      "server_parameters": {
+        "model": "Qwen/Qwen3-1.7B",
+        "tensor_parallel_size": 1
+      },
+      "client_parameters": {
+        "model": "Qwen/Qwen3-1.7B",
+        "dataset_name": "random",
+        "random-input-len": 128,
+        "random-output-len": 128
+      }
+    },
+    {
+      "test_name": "serving_qwen4B_tp1_random_128_128",
+      "server_parameters": {
+        "model": "Qwen/Qwen3-4B",
+        "tensor_parallel_size": 1
+      },
+      "client_parameters": {
+        "model": "Qwen/Qwen3-4B",
+        "dataset_name": "random",
+        "random-input-len": 128,
+        "random-output-len": 128
+      }
+    },
+    {
+      "test_name": "serving_qwen8B_tp1_random_128_128",
+      "server_parameters": {
+        "model": "Qwen/Qwen3-8B",
+        "tensor_parallel_size": 1
+      },
+      "client_parameters": {
+        "model": "Qwen/Qwen3-8B",
+        "dataset_name": "random",
+        "random-input-len": 128,
+        "random-output-len": 128
+      }
+    },
+    {
+      "test_name": "serving_glm9B_tp1_random_128_128",
+      "server_parameters": {
+        "model": "zai-org/glm-4-9b-hf",
+        "tensor_parallel_size": 1
+      },
+      "client_parameters": {
+        "model": "zai-org/glm-4-9b-hf",
+        "dataset_name": "random",
+        "random-input-len": 128,
+        "random-output-len": 128
+      }
+    },
+    {
+      "test_name": "serving_gemma7B_tp1_random_128_128",
+      "server_parameters": {
+        "model": "google/gemma-7b",
+        "tensor_parallel_size": 1
+      },
+      "client_parameters": {
+        "model": "google/gemma-7b",
+        "dataset_name": "random",
+        "random-input-len": 128,
+        "random-output-len": 128
+      }
    }
  ]
 }
@@ -132,7 +132,7 @@ steps:
  - tests/entrypoints/
  commands:
  - pytest -v -s entrypoints/openai/tool_parsers
-  - pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/openai --ignore=entrypoints/rpc --ignore=entrypoints/instrumentator --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py  --ignore=entrypoints/pooling
+  - pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/openai --ignore=entrypoints/rpc --ignore=entrypoints/sleep --ignore=entrypoints/instrumentator --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py  --ignore=entrypoints/pooling

 - label: Entrypoints Integration Test (LLM) # 30min
  timeout_in_minutes: 40
@@ -179,14 +179,14 @@ steps:
  torch_nightly: true
  source_file_dependencies:
  - vllm/
+  - tests/entrypoints/sleep
  - tests/entrypoints/rpc
-  - tests/entrypoints/instrumentator
  - tests/tool_use
  commands:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/instrumentator
-  - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc
+  - pytest -v -s entrypoints/sleep
  - pytest -v -s tool_use
+  - PYTHONPATH=/vllm-workspace  pytest -v -s entrypoints/rpc

 - label: Entrypoints Integration Test (Pooling)
  timeout_in_minutes: 50
@@ -552,7 +552,7 @@ steps:
 - label: LoRA Test %N # 20min each
  timeout_in_minutes: 30
  mirror_hardwares: [amdexperimental]
-  agent_pool: mi325_1
+  agent_pool: mi325_8
  # grade: Blocking
  source_file_dependencies:
  - vllm/lora
@@ -648,7 +648,7 @@ steps:
 - label: Kernels Attention Test %N # 23min
  timeout_in_minutes: 35
  mirror_hardwares: [amdexperimental, amdproduction]
-  agent_pool: mi325_1
+  agent_pool: mi325_8
  # grade: Blocking
  source_file_dependencies:
  - csrc/attention/
@@ -663,7 +663,7 @@ steps:
 - label: Kernels Quantization Test %N # 64min
  timeout_in_minutes: 90
  mirror_hardwares: [amdexperimental]
-  agent_pool: mi325_1
+  agent_pool: mi325_8
  # grade: Blocking
  source_file_dependencies:
  - csrc/quantization/
@@ -676,7 +676,7 @@ steps:
 - label: Kernels MoE Test %N # 40min
  timeout_in_minutes: 60
  mirror_hardwares: [amdexperimental, amdproduction]
-  agent_pool: mi325_1
+  agent_pool: mi325_8
  # grade: Blocking
  source_file_dependencies:
  - csrc/quantization/cutlass_w8a8/moe/
@@ -839,7 +839,7 @@ steps:
 - label: Basic Models Tests (Extra Initialization) %N
  timeout_in_minutes: 45
  mirror_hardwares: [amdexperimental, amdproduction]
-  agent_pool: mi325_1
+  agent_pool: mi325_8
  # grade: Blocking
  torch_nightly: true
  source_file_dependencies:
@@ -901,7 +901,7 @@ steps:
 - label: Language Models Tests (Extra Standard) %N
  timeout_in_minutes: 45
  mirror_hardwares: [amdexperimental]
-  agent_pool: mi325_1
+  agent_pool: mi325_8
  # grade: Blocking
  torch_nightly: true
  source_file_dependencies:
@@ -922,7 +922,7 @@ steps:
 - label: Language Models Tests (Hybrid) %N
  timeout_in_minutes: 75
  mirror_hardwares: [amdexperimental]
-  agent_pool: mi325_1
+  agent_pool: mi325_8
  # grade: Blocking
  torch_nightly: true
  source_file_dependencies:
@@ -1334,7 +1334,7 @@ steps:
  - pytest -v -s ./compile/test_wrapper.py
  - VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
  - VLLM_TEST_SAME_HOST=1 VLLM_TEST_WITH_DEFAULT_DEVICE_SET=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
-  - pytest -v -s compile/correctness_e2e/test_sequence_parallel.py
+  - pytest -v -s distributed/test_sequence_parallel.py
  - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown
  - pytest -v -s v1/worker/test_worker_memory_snapshot.py

@@ -118,7 +118,7 @@ steps:
  - tests/entrypoints/
  commands:
  - pytest -v -s entrypoints/openai/tool_parsers
-  - pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/rpc --ignore=entrypoints/instrumentator --ignore=entrypoints/openai --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py  --ignore=entrypoints/pooling
+  - pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/rpc --ignore=entrypoints/sleep --ignore=entrypoints/instrumentator --ignore=entrypoints/openai --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py  --ignore=entrypoints/pooling

 - label: Entrypoints Integration Test (LLM) # 30min
  timeout_in_minutes: 40
@@ -148,7 +148,7 @@ steps:
  - tests/entrypoints/test_chat_utils
  commands:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/test_chat_with_tool_reasoning.py --ignore=entrypoints/instrumentator --ignore=entrypoints/openai/test_oot_registration.py --ignore=entrypoints/openai/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/  --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses
+  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/test_oot_registration.py --ignore=entrypoints/openai/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/  --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses
  - pytest -v -s entrypoints/test_chat_utils.py

 - label: Entrypoints Integration Test (API Server 2)
@@ -159,13 +159,13 @@ steps:
  torch_nightly: true
  source_file_dependencies:
  - vllm/
+  - tests/entrypoints/sleep
  - tests/entrypoints/rpc
-  - tests/entrypoints/instrumentator
  - tests/tool_use
  commands:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/instrumentator
-  - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc
+  - pytest -v -s entrypoints/sleep
+  - PYTHONPATH=/vllm-workspace  pytest -v -s entrypoints/rpc
  - pytest -v -s tool_use

 - label: Entrypoints Integration Test (Pooling)
@@ -862,7 +862,7 @@ steps:
  commands:
    # Install fast path packages for testing against transformers
    # Note: also needed to run plamo2 model in vLLM
-    - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
+    - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.2.5'
    - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.5.2'
    # Shard hybrid language model tests
    - pytest -v -s models/language/generation \
@@ -881,7 +881,7 @@ steps:
  commands:
    # Install fast path packages for testing against transformers
    # Note: also needed to run plamo2 model in vLLM
-    - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
+    - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.2.5'
    - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.5.2'
    - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'

@@ -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
@@ -17,15 +17,3 @@ steps:
  - tests/benchmarks/
  commands:
  - pytest -v -s benchmarks/
-
- label: Attention Benchmarks Smoke Test (B200)
-  device: b200
-  num_gpus: 2
-  optional: true
-  working_dir: "/vllm-workspace/"
-  timeout_in_minutes: 10
-  source_file_dependencies:
-  - benchmarks/attention_benchmarks/
-  - vllm/v1/attention/
-  commands:
-  - python3 benchmarks/attention_benchmarks/benchmark.py --backends flash flashinfer --batch-specs "8q1s1k" --repeats 1 --warmup-iters 1
@@ -121,10 +121,13 @@ steps:
  optional: true
  commands:
    - nvidia-smi
-    # Run all models but only FLASHINFER, Inductor partition and native custom ops
+    # Run all models and attn backends but only Inductor partition and native custom ops
+    # -k "inductor_partition and not +rms_norm and not +quant_fp8"
    # Qwen 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)"
+    # -k "inductor_partition and not +rms_norm and +quant_fp8 and qwen3"
+    # Run just llama3 (fp8 & fp4) for all config combinations
+    # -k "llama-3"
+    - pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "inductor_partition and not +rms_norm and not +quant_fp8" -k "inductor_partition and not +rms_norm and +quant_fp8 and qwen3" -k "llama-3"

 - label: Fusion E2E TP2 Quick (H100)
  timeout_in_minutes: 20
@@ -159,7 +162,7 @@ steps:
    - tests/compile/fusions_e2e/
  commands:
    - nvidia-smi
-    # Run just llama3 (fp8 & bf16) for all config combinations
+    # Run just llama3 (fp4 & fp8 & bf16) for all config combinations
    - pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k "llama-3"

 - label: Fusion E2E TP2 AsyncTP Config Sweep (H100)
@@ -194,8 +197,7 @@ steps:
    - tests/compile/fusions_e2e/
  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
+    # Run all models and attn backends but only Inductor partition and native custom ops
    # 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 "inductor_partition and not +rms_norm and not +quant_fp8" -k "Llama-3.1-8B-Instruct-FP4"
+    - pytest -v -s tests/compile/fusions_e2e/test_tp2_async_tp.py -k "inductor_partition and not +rms_norm and not +quant_fp8"
@@ -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
@@ -47,13 +42,15 @@ steps:
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
-  - tests/entrypoints/rpc
-  - tests/entrypoints/instrumentator
  - tests/tool_use
+  - tests/entrypoints/sleep
+  - tests/entrypoints/instrumentator
+  - tests/entrypoints/rpc
  commands:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/instrumentator
  - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc
+  - pytest -v -s entrypoints/instrumentator
+  - pytest -v -s entrypoints/sleep
  - pytest -v -s tool_use

 - label: Entrypoints Integration (Pooling)
@@ -4,6 +4,7 @@ depends_on:
 steps:
 - label: Basic Models Tests (Initialization)
  timeout_in_minutes: 45
+  mirror_hardwares: [amdexperimental]
  torch_nightly: true
  source_file_dependencies:
  - vllm/
@@ -15,6 +16,7 @@ steps:

 - label: Basic Models Tests (Extra Initialization) %N
  timeout_in_minutes: 45
+  mirror_hardwares: [amdexperimental]
  torch_nightly: true
  source_file_dependencies:
  - vllm/model_executor/models/
@@ -36,12 +38,6 @@ steps:
  - tests/models/test_registry.py
  commands:
    - pytest -v -s models/test_terratorch.py models/test_transformers.py models/test_registry.py
-  mirror:
-    amd:
-      device: mi325_1
-      depends_on:
-      - image-build-amd
-    

 - label: Basic Models Test (Other CPU) # 5min
  depends_on: 
@@ -4,6 +4,7 @@ depends_on:
 steps:
 - label: Language Models Tests (Standard)
  timeout_in_minutes: 25
+  mirror_hardwares: [amdexperimental]
  torch_nightly: true
  source_file_dependencies:
  - vllm/
@@ -15,6 +16,7 @@ steps:

 - label: Language Models Tests (Extra Standard) %N
  timeout_in_minutes: 45
+  mirror_hardwares: [amdexperimental]
  torch_nightly: true
  source_file_dependencies:
  - vllm/model_executor/models/
@@ -30,6 +32,7 @@ steps:

 - label: Language Models Tests (Hybrid) %N
  timeout_in_minutes: 75
+  mirror_hardwares: [amdexperimental]
  torch_nightly: true
  source_file_dependencies:
  - vllm/
@@ -37,7 +40,7 @@ steps:
  commands:
    # Install fast path packages for testing against transformers
    # Note: also needed to run plamo2 model in vLLM
-    - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
+    - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.2.5'
    - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.5.2'
    # Shard hybrid language model tests
    - pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
@@ -45,6 +48,7 @@ steps:

 - label: Language Models Test (Extended Generation) # 80min
  timeout_in_minutes: 110
+  mirror_hardwares: [amdexperimental]
  optional: true
  source_file_dependencies:
  - vllm/
@@ -52,12 +56,13 @@ steps:
  commands:
    # Install fast path packages for testing against transformers
    # Note: also needed to run plamo2 model in vLLM
-    - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
+    - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.2.5'
    - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.5.2'
    - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'

 - label: Language Models Test (PPL)
  timeout_in_minutes: 110
+  mirror_hardwares: [amdexperimental]
  optional: true
  source_file_dependencies:
  - vllm/
@@ -67,6 +72,7 @@ steps:

 - label: Language Models Test (Extended Pooling)  # 36min
  timeout_in_minutes: 50
+  mirror_hardwares: [amdexperimental]
  optional: true
  source_file_dependencies:
  - vllm/
@@ -76,6 +82,7 @@ steps:

 - label: Language Models Test (MTEB)
  timeout_in_minutes: 110
+  mirror_hardwares: [amdexperimental]
  optional: true
  source_file_dependencies:
  - vllm/
@@ -12,10 +12,3 @@ steps:
  commands:
    - pytest -v -s samplers
    - VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers
-  mirror:
-    amd:
-      device: mi325_1
-      depends_on:
-      - image-build-amd
-      commands:
-      - pytest -v -s -m 'not skip_v1' samplers
@@ -2,9 +2,7 @@
 # for more info about CODEOWNERS file

 # This lists cover the "core" components of vLLM that require careful review
-/vllm/compilation @zou3519 @youkaichao @ProExpertProg
-/vllm/distributed/kv_transfer @NickLucche @ApostaC @orozery
-/vllm/lora @jeejeelee
+/vllm/executor/executor_base.py @zhuohan123 @youkaichao @alexm-redhat @njhill @22quinn
 /vllm/model_executor/layers/attention @LucasWilkinson
 /vllm/model_executor/layers/fused_moe @mgoin @pavanimajety
 /vllm/model_executor/layers/quantization @mgoin @robertgshaw2-redhat @tlrmchlsmth @yewentao256 @pavanimajety
@@ -13,34 +11,18 @@
 /vllm/model_executor/layers/batch_invariant.py @yewentao256 
 /vllm/multimodal @DarkLight1337 @ywang96 @NickLucche @tjtanaa
 /vllm/vllm_flash_attn @LucasWilkinson
+/vllm/lora @jeejeelee
+/vllm/reasoning @aarnphm @chaunceyjiang
+/vllm/entrypoints @aarnphm @chaunceyjiang
+/vllm/tool_parsers @aarnphm @chaunceyjiang
+/vllm/compilation @zou3519 @youkaichao @ProExpertProg
+/vllm/distributed/kv_transfer @NickLucche @ApostaC @orozery
 CMakeLists.txt @tlrmchlsmth @LucasWilkinson

 # Any change to the VllmConfig changes can have a large user-facing impact,
 # so spam a lot of people
 /vllm/config @WoosukKwon @youkaichao @robertgshaw2-redhat @mgoin @tlrmchlsmth @houseroad @hmellor @yewentao256 @ProExpertProg
-/vllm/config/cache.py @heheda12345
-
-# Entrypoints
-/vllm/entrypoints/anthropic @mgoin @DarkLight1337
-/vllm/entrypoints/cli @hmellor @mgoin @DarkLight1337 @russellb
-/vllm/entrypoints/mcp @heheda12345
-/vllm/entrypoints/openai @aarnphm @chaunceyjiang @DarkLight1337 @russellb
-/vllm/entrypoints/openai/realtime @njhill
-/vllm/entrypoints/openai/speech_to_text @NickLucche
-/vllm/entrypoints/pooling @noooop
-/vllm/entrypoints/sagemaker @DarkLight1337
-/vllm/entrypoints/serve @njhill
-/vllm/entrypoints/*.py @njhill
-/vllm/entrypoints/chat_utils.py @DarkLight1337
-/vllm/entrypoints/llm.py @DarkLight1337
-
-# Input/Output Processing
-/vllm/sampling_params.py @njhill @NickLucche
-/vllm/pooling_params.py @noooop @DarkLight1337
-/vllm/tokenizers @DarkLight1337 @njhill
-/vllm/renderers @DarkLight1337 @njhill
-/vllm/reasoning @aarnphm @chaunceyjiang
-/vllm/tool_parsers @aarnphm @chaunceyjiang
+/vllm/config/cache.py @WoosukKwon @youkaichao @robertgshaw2-redhat @mgoin @tlrmchlsmth @houseroad @hmellor @yewentao256 @ProExpertProg @heheda12345

 # vLLM V1
 /vllm/v1/attention @LucasWilkinson
@@ -133,8 +115,8 @@ mkdocs.yaml @hmellor
 /vllm/model_executor/models/mixtral*.py @patrickvonplaten
 /vllm/model_executor/models/voxtral*.py @patrickvonplaten
 /vllm/model_executor/models/pixtral*.py @patrickvonplaten
-/vllm/tokenizers/mistral.py @patrickvonplaten
 /vllm/transformers_utils/configs/mistral.py @patrickvonplaten
+/vllm/transformers_utils/tokenizers/mistral.py @patrickvonplaten

 # Kernels
 /vllm/v1/attention/ops/chunked_prefill_paged_decode.py @tdoublep
@@ -170,7 +152,9 @@ mkdocs.yaml @hmellor
 /examples/pooling @noooop
 /tests/models/*/pooling* @noooop
 /tests/entrypoints/pooling @noooop
+/vllm/entrypoints/pooling @noooop
 /vllm/config/pooler.py @noooop
+/vllm/pooling_params.py @noooop
 /vllm/model_executor/layers/pooler @noooop

 # Security guide and policies
@@ -19,7 +19,6 @@ jobs:
        uses: actions/setup-python@83679a892e2d95755f2dac6acb0bfd1e9ac5d548 # v6.1.0
        with:
          python-version: '3.12'
-          cache: 'pip'

      - name: Install Python dependencies
        run: |
@@ -238,6 +238,3 @@ ep_kernels_workspace/
 vllm/grpc/vllm_engine_pb2.py
 vllm/grpc/vllm_engine_pb2_grpc.py
 vllm/grpc/vllm_engine_pb2.pyi
-
-# Ignore generated cpu headers 
-csrc/cpu/cpu_attn_dispatch_generated.h
@@ -143,11 +143,6 @@ repos:
    name: Check attention backend documentation is up to date
    entry: python tools/pre_commit/generate_attention_backend_docs.py --check
    language: python
-  - id: check-boolean-context-manager
-    name: Check for boolean ops in with-statements
-    entry: python tools/pre_commit/check_boolean_context_manager.py
-    language: python
-    types: [python]
  # Keep `suggestion` last
  - id: suggestion
    name: Suggestion
@@ -9,14 +9,13 @@ build:
    python: "3.12"
  jobs:
    post_checkout:
-      - git fetch origin main --unshallow --no-tags --filter=blob:none || true
-    pre_create_environment:
-      - pip install uv
-    create_environment:
-      - uv venv $READTHEDOCS_VIRTUALENV_PATH
-    install:
-      - uv pip install --python $READTHEDOCS_VIRTUALENV_PATH/bin/python --no-cache-dir -r requirements/docs.txt 
+      - git fetch --unshallow || true

 mkdocs:
  configuration: mkdocs.yaml
  fail_on_warning: true
+
+# Optionally declare the Python requirements required to build your docs
+python:
+  install:
+    - requirements: requirements/docs.txt
@@ -293,7 +293,6 @@ set(VLLM_EXT_SRC
  "csrc/fused_qknorm_rope_kernel.cu"
  "csrc/layernorm_quant_kernels.cu"
  "csrc/sampler.cu"
-  "csrc/topk.cu"
  "csrc/cuda_view.cu"
  "csrc/quantization/gptq/q_gemm.cu"
  "csrc/quantization/w8a8/int8/scaled_quant.cu"
@@ -229,40 +229,3 @@ def get_batch_stats(requests: list[BatchRequest]) -> dict:
            sum(r.kv_len for r in requests) / len(requests) if requests else 0
        ),
    }
-
-
-def get_batch_type(batch_spec: str, spec_decode_threshold: int = 8) -> str:
-    """
-    Classify a batch spec into a type string.
-
-    Args:
-        batch_spec: Batch specification string (e.g., "q2k", "8q1s1k", "2q2k_8q1s1k")
-        spec_decode_threshold: Max q_len to be considered spec-decode vs extend
-
-    Returns:
-        Type string: "prefill", "decode", "spec-decode", "extend", or "mixed (types...)"
-    """
-    requests = parse_batch_spec(batch_spec)
-
-    # Classify each request
-    types_present = set()
-    for req in requests:
-        if req.is_decode:
-            types_present.add("decode")
-        elif req.is_prefill:
-            types_present.add("prefill")
-        elif req.is_extend:
-            # Distinguish spec-decode (small q_len) from extend (chunked prefill)
-            if req.q_len <= spec_decode_threshold:
-                types_present.add("spec-decode")
-            else:
-                types_present.add("extend")
-
-    if len(types_present) == 1:
-        return types_present.pop()
-    elif len(types_present) > 1:
-        # Sort for consistent output
-        sorted_types = sorted(types_present)
-        return f"mixed ({'+'.join(sorted_types)})"
-    else:
-        return "unknown"
@@ -43,7 +43,6 @@ from common import (
    ModelParameterSweep,
    ParameterSweep,
    ResultsFormatter,
-    batch_spec_sort_key,
    is_mla_backend,
 )

@@ -219,13 +218,10 @@ def run_model_parameter_sweep(
                        by_param_and_spec[key].append(r)
                        break

-    # Sort by param value then spec (batch_size, q_len, kv_len)
+    # Sort by param value then spec
    sorted_keys = sorted(
        by_param_and_spec.keys(),
-        key=lambda x: (
-            int(x[0]) if x[0].isdigit() else x[0],
-            batch_spec_sort_key(x[1]),
-        ),
+        key=lambda x: (int(x[0]) if x[0].isdigit() else x[0], x[1]),
    )

    current_param_value = None
@@ -334,7 +330,7 @@ def run_parameter_sweep(
                by_spec[spec] = []
            by_spec[spec].append(r)

-    for spec in sorted(by_spec.keys(), key=batch_spec_sort_key):
+    for spec in sorted(by_spec.keys()):
        results = by_spec[spec]
        best = min(results, key=lambda r: r.mean_time)
        console.print(
@@ -500,18 +496,15 @@ def main():
        if "description" in yaml_config:
            console.print(f"[dim]{yaml_config['description']}[/]")

-        # Override args with YAML values, but CLI args take precedence
-        # Check if CLI provided backends (they would be non-None and not default)
-        cli_backends_provided = args.backends is not None or args.backend is not None
-
-        # Backend(s) - only use YAML if CLI didn't specify
-        if not cli_backends_provided:
-            if "backend" in yaml_config:
-                args.backend = yaml_config["backend"]
-                args.backends = None
-            elif "backends" in yaml_config:
-                args.backends = yaml_config["backends"]
-                args.backend = None
+        # Override args with YAML values
+        # (YAML takes precedence unless CLI arg was explicitly set)
+        # Backend(s)
+        if "backend" in yaml_config:
+            args.backend = yaml_config["backend"]
+            args.backends = None
+        elif "backends" in yaml_config:
+            args.backends = yaml_config["backends"]
+            args.backend = None

        # Check for special modes
        if "mode" in yaml_config:
@@ -551,15 +544,13 @@ def main():
            args.num_kv_heads = model.get("num_kv_heads", args.num_kv_heads)
            args.block_size = model.get("block_size", args.block_size)

-        # Benchmark settings (top-level keys)
-        if "device" in yaml_config:
-            args.device = yaml_config["device"]
-        if "repeats" in yaml_config:
-            args.repeats = yaml_config["repeats"]
-        if "warmup_iters" in yaml_config:
-            args.warmup_iters = yaml_config["warmup_iters"]
-        if "profile_memory" in yaml_config:
-            args.profile_memory = yaml_config["profile_memory"]
+        # Benchmark settings
+        if "benchmark" in yaml_config:
+            bench = yaml_config["benchmark"]
+            args.device = bench.get("device", args.device)
+            args.repeats = bench.get("repeats", args.repeats)
+            args.warmup_iters = bench.get("warmup_iters", args.warmup_iters)
+            args.profile_memory = bench.get("profile_memory", args.profile_memory)

        # Parameter sweep configuration
        if "parameter_sweep" in yaml_config:
@@ -12,36 +12,16 @@ from typing import Any

 import numpy as np
 import torch
-from batch_spec import get_batch_type, parse_batch_spec
 from rich.console import Console
 from rich.table import Table

-
-def batch_spec_sort_key(spec: str) -> tuple[int, int, int]:
-    """
-    Extract sorting key from batch spec: (batch_size, max_q_len, max_kv_len).
-
-    This ensures results are sorted by batch size first, then query length,
-    then sequence length, rather than alphabetically.
-    """
-    try:
-        requests = parse_batch_spec(spec)
-        batch_size = len(requests)
-        max_q_len = max(r.q_len for r in requests) if requests else 0
-        max_kv_len = max(r.kv_len for r in requests) if requests else 0
-        return (batch_size, max_q_len, max_kv_len)
-    except Exception:
-        # Fallback for unparseable specs
-        return (0, 0, 0)
-
-
 # Mock classes for vLLM attention infrastructure


 class MockHfConfig:
    """Mock HuggingFace config that satisfies vLLM's requirements."""

-    def __init__(self, mla_dims: dict, index_topk: int | None = None):
+    def __init__(self, mla_dims: dict):
        self.num_attention_heads = mla_dims["num_q_heads"]
        self.num_key_value_heads = mla_dims["num_kv_heads"]
        self.hidden_size = mla_dims["head_dim"] * mla_dims["num_q_heads"]
@@ -52,8 +32,6 @@ class MockHfConfig:
        self.qk_rope_head_dim = mla_dims["qk_rope_head_dim"]
        self.v_head_dim = mla_dims["v_head_dim"]
        self.qk_head_dim = mla_dims["qk_nope_head_dim"] + mla_dims["qk_rope_head_dim"]
-        if index_topk is not None:
-            self.index_topk = index_topk

    def get_text_config(self):
        return self
@@ -104,38 +82,6 @@ class MockKVBProj:
        return (result,)  # Return as tuple to match ColumnParallelLinear API


-class MockIndexer:
-    """Mock Indexer for sparse MLA backends.
-
-    Provides topk_indices_buffer that sparse MLA backends use to determine
-    which KV cache slots to attend to for each token.
-    """
-
-    def __init__(
-        self,
-        max_num_tokens: int,
-        topk_tokens: int,
-        device: torch.device,
-    ):
-        self.topk_tokens = topk_tokens
-        self.topk_indices_buffer = torch.zeros(
-            (max_num_tokens, topk_tokens),
-            dtype=torch.int32,
-            device=device,
-        )
-
-    def fill_random_indices(self, num_tokens: int, max_kv_len: int):
-        """Fill topk_indices_buffer with random valid indices for benchmarking."""
-        indices = torch.randint(
-            0,
-            max_kv_len,
-            (num_tokens, self.topk_tokens),
-            dtype=torch.int32,
-            device=self.topk_indices_buffer.device,
-        )
-        self.topk_indices_buffer[:num_tokens] = indices
-
-
 class MockLayer(AttentionLayerBase):
    """Mock attention layer with scale parameters and impl.

@@ -370,19 +316,14 @@ class ResultsFormatter:
            backends: List of backend names being compared
            compare_to_fastest: Show percentage comparison to fastest
        """
-        # Group by batch spec, preserving first-occurrence order
+        # Group by batch spec
        by_spec = {}
-        specs_order = []
        for r in results:
            spec = r.config.batch_spec
            if spec not in by_spec:
                by_spec[spec] = {}
-                specs_order.append(spec)
            by_spec[spec][r.config.backend] = r

-        # Sort specs by (batch_size, q_len, kv_len) instead of alphabetically
-        specs_order = sorted(by_spec.keys(), key=batch_spec_sort_key)
-
        # Create shortened backend names for display
        def shorten_backend_name(name: str) -> str:
            """Shorten long backend names for table display."""
@@ -396,8 +337,6 @@ class ResultsFormatter:

        table = Table(title="Attention Benchmark Results")
        table.add_column("Batch\nSpec", no_wrap=True)
-        table.add_column("Type", no_wrap=True)
-        table.add_column("Batch\nSize", justify="right", no_wrap=True)

        multi = len(backends) > 1
        for backend in backends:
@@ -411,14 +350,12 @@ class ResultsFormatter:
                table.add_column(col_rel, justify="right", no_wrap=False)

        # Add rows
-        for spec in specs_order:
+        for spec in sorted(by_spec.keys()):
            spec_results = by_spec[spec]
            times = {b: r.mean_time for b, r in spec_results.items() if r.success}
            best_time = min(times.values()) if times else 0.0

-            batch_type = get_batch_type(spec)
-            batch_size = len(parse_batch_spec(spec))
-            row = [spec, batch_type, str(batch_size)]
+            row = [spec]
            for backend in backends:
                if backend in spec_results:
                    r = spec_results[backend]
@@ -549,11 +486,10 @@ def get_attention_scale(head_dim: int) -> float:

 def is_mla_backend(backend: str) -> bool:
    """
-    Check if backend is an MLA backend using the AttentionBackendEnum.
+    Check if backend is an MLA backend using the backend's is_mla() property.

    Args:
-        backend: Backend name matching AttentionBackendEnum exactly
-        (e.g., "FLASHMLA_SPARSE")
+        backend: Backend name (e.g., "CUTLASS_MLA", "FLASHINFER_MLA")

    Returns:
        True if the backend is an MLA backend, False otherwise
@@ -561,8 +497,7 @@ def is_mla_backend(backend: str) -> bool:
    from vllm.v1.attention.backends.registry import AttentionBackendEnum

    try:
-        backend_enum = AttentionBackendEnum[backend]
-        backend_class = backend_enum.get_class()
+        backend_class = AttentionBackendEnum[backend.upper()].get_class()
        return backend_class.is_mla()
-    except (KeyError, ValueError, ImportError, AttributeError):
+    except (KeyError, ValueError, ImportError):
        return False
@@ -3,7 +3,7 @@
 model:
  name: "deepseek-v3"
  num_layers: 60
-  num_q_heads: 128  # Base value, can be swept for TP simulation
+  num_q_heads: 128
  num_kv_heads: 1  # MLA uses single latent KV
  head_dim: 576
  kv_lora_rank: 512
@@ -12,13 +12,6 @@ model:
  v_head_dim: 128
  block_size: 128  # CUTLASS MLA and FlashAttn MLA use 128

-# Model parameter sweep: simulate tensor parallelism by varying num_q_heads
-# TP=1: 128 heads, TP=2: 64 heads, TP=4: 32 heads, TP=8: 16 heads
-model_parameter_sweep:
-  param_name: "num_q_heads"
-  values: [128, 64, 32, 16]
-  label_format: "{backend}_{value}h"
-
 batch_specs:
  # Small batches, varying sequence lengths
  - "16q1s512"     # 16 requests, 512 KV cache
@@ -41,30 +34,28 @@ batch_specs:
  # Very large batches
  - "128q1s1k"     # 128 requests, 1k KV cache
  - "128q1s2k"     # 128 requests, 2k KV cache
-  - "128q1s4k"     # 128 requests, 4k KV cache
-  - "128q1s8k"     # 128 requests, 8k KV cache

  # Long context
  - "32q1s16k"     # 32 requests, 16k KV cache
  - "32q1s32k"     # 32 requests, 32k KV cache

 backends:
-  - CUTLASS_MLA
-  - FLASHINFER_MLA
-  - FLASH_ATTN_MLA  # Hopper only
-  - FLASHMLA        # Hopper only
+  - cutlass_mla
+  - flashinfer_mla
+  - flashattn_mla  # Hopper only
+  - flashmla        # Hopper only

 device: "cuda:0"
-repeats: 100
-warmup_iters: 10
+repeats: 5
+warmup_iters: 3
 profile_memory: true

 # Backend-specific tuning
-CUTLASS_MLA:
+cutlass_mla:
  num_kv_splits: auto  # or specific value like 4, 8, 16

-FLASH_ATTN_MLA:
+flashattn_mla:
  reorder_batch_threshold: 512

-FLASHMLA:
+flashmla:
  reorder_batch_threshold: 1
@@ -45,10 +45,10 @@ batch_specs:
  - "4q4k_60q1s4k"          # 4 prefill + 60 decode

 backends:
-  - CUTLASS_MLA
-  - FLASHINFER_MLA
-  - FLASH_ATTN_MLA   # Hopper only
-  - FLASHMLA         # Hopper only
+  - cutlass_mla
+  - flashinfer_mla
+  - flashattn_mla   # Hopper only
+  - flashmla        # Hopper only

 device: "cuda:0"
 repeats: 5
@@ -1,62 +0,0 @@
-# MLA prefill-only benchmark configuration for sparse backends
-
-model:
-  name: "deepseek-v3"
-  num_layers: 60
-  num_q_heads: 128
-  num_kv_heads: 1
-  head_dim: 576
-  kv_lora_rank: 512
-  qk_nope_head_dim: 128
-  qk_rope_head_dim: 64
-  v_head_dim: 128
-  block_size: 128
-
-# Model parameter sweep: simulate tensor parallelism by varying num_q_heads
-# TP=1: 128 heads, TP=2: 64 heads, TP=4: 32 heads, TP=8: 16 heads
-model_parameter_sweep:
-  param_name: "num_q_heads"
-  values: [128, 64, 32, 16]
-  label_format: "{backend}_{value}h"
-
-batch_specs:
-  # Pure prefill
-  - "1q512"
-  - "1q1k"
-  - "1q2k"
-  - "1q4k"
-  - "1q8k"
-
-  # Batched pure prefill
-  - "2q512"
-  - "2q1k"
-  - "2q2k"
-  - "2q4k"
-  - "2q8k"
-  - "4q512"
-  - "4q1k"
-  - "4q2k"
-  - "4q4k"
-  - "4q8k"
-  - "8q512"
-  - "8q1k"
-  - "8q2k"
-  - "8q4k"
-  - "8q8k"
-
-  # Extend
-  - "1q512s4k"
-  - "1q512s8k"
-  - "1q1ks8k"
-  - "1q2ks8k"
-  - "1q2ks16k"
-  - "1q4ks16k"
-
-backends:
-  - FLASHMLA_SPARSE
-  - FLASHINFER_MLA_SPARSE
-
-device: "cuda:0"
-repeats: 10
-warmup_iters: 3
-profile_memory: true
@@ -6,7 +6,7 @@
 description: "Decode vs Prefill pipeline crossover analysis"

 # Test FlashAttn MLA
-backend: FLASH_ATTN_MLA
+backend: flashattn_mla

 # Mode: decode_vs_prefill comparison (special sweep mode)
 # For each batch spec, we'll test both decode and prefill pipelines
@@ -62,10 +62,11 @@ model:
  block_size: 128

 # Benchmark settings
-device: "cuda:0"
-repeats: 15          # More repeats for spec decode variance
-warmup_iters: 5
-profile_memory: false
+benchmark:
+  device: "cuda:0"
+  repeats: 15          # More repeats for spec decode variance
+  warmup_iters: 5
+  profile_memory: false

 # Output
 output:
@@ -41,17 +41,18 @@ batch_specs:

 # Backends that support query length > 1
 backends:
-  - FLASH_ATTN_MLA    # reorder_batch_threshold = 512
-  - FLASHMLA          # reorder_batch_threshold = 1 (tunable)
+  - flashattn_mla    # reorder_batch_threshold = 512
+  - flashmla          # reorder_batch_threshold = 1 (tunable)

 # FlashInfer-MLA also supports uniform spec-as-decode but with different mechanism
-# - FLASHINFER_MLA
+# - flashinfer_mla

 # Benchmark settings
-device: "cuda:0"
-repeats: 10  # More repeats for statistical significance
-warmup_iters: 5
-profile_memory: false
+benchmark:
+  device: "cuda:0"
+  repeats: 10  # More repeats for statistical significance
+  warmup_iters: 5
+  profile_memory: false

 # Test these threshold values for optimization
 parameter_sweep:
@@ -25,22 +25,14 @@ batch_specs:
  - "4q1k_16q1s2k"     # 4 prefill + 16 decode
  - "2q4k_32q1s1k"     # 2 large prefill + 32 decode

-  # Speculative decode (q <= 8)
-  - "16q2s1k"         # 16 requests, 2 spec tokens, 1k KV cache
-  - "16q4s1k"         # 16 requests, 4 spec tokens, 1k KV cache
-  - "16q8s1k"         # 16 requests, 8 spec tokens, 1k KV cache
-  - "32q4s2k"         # 32 requests, 4 spec tokens, 2k KV cache
-  - "8q8s4k"          # 8 requests, 8 spec tokens, 4k KV cache
-
-  # Context extension (chunked prefill)
-  - "q1ks2k"          # 1k query, 2k sequence
+  # Context extension
+  - "q1ks2k"          # 1k query, 2k sequence (chunked prefill)
  - "2q1ks4k"         # 2 requests: 1k query, 4k sequence

-# Available backends: FLASH_ATTN, TRITON_ATTN, FLASHINFER
 backends:
-  - FLASH_ATTN
-  - TRITON_ATTN
-  - FLASHINFER
+  - flash
+  - triton
+  - flashinfer

 device: "cuda:0"
 repeats: 5
@@ -8,13 +8,14 @@ This module provides helpers for running MLA backends without
 needing full VllmConfig integration.
 """

+import importlib
+
 import numpy as np
 import torch
 from batch_spec import parse_batch_spec
 from common import (
    BenchmarkResult,
    MockHfConfig,
-    MockIndexer,
    MockKVBProj,
    MockLayer,
    setup_mla_dims,
@@ -61,7 +62,6 @@ def create_minimal_vllm_config(
    block_size: int = 128,
    max_num_seqs: int = 256,
    mla_dims: dict | None = None,
-    index_topk: int | None = None,
 ) -> VllmConfig:
    """
    Create minimal VllmConfig for MLA benchmarks.
@@ -73,8 +73,6 @@ def create_minimal_vllm_config(
        max_num_seqs: Maximum number of sequences
        mla_dims: Optional custom MLA dimensions dict. If not provided, uses
                  setup_mla_dims(model_name)
-        index_topk: Optional topk value for sparse MLA backends. If provided,
-                    the config will include index_topk for sparse attention.

    Returns:
        VllmConfig for benchmarking
@@ -84,7 +82,7 @@ def create_minimal_vllm_config(
        mla_dims = setup_mla_dims(model_name)

    # Create mock HF config first (avoids downloading from HuggingFace)
-    mock_hf_config = MockHfConfig(mla_dims, index_topk=index_topk)
+    mock_hf_config = MockHfConfig(mla_dims)

    # Create a temporary minimal config.json to avoid HF downloads
    # This ensures consistent ModelConfig construction without network access
@@ -122,12 +120,16 @@ def create_minimal_vllm_config(
            seed=0,
            max_model_len=32768,
            quantization=None,
+            quantization_param_path=None,
            enforce_eager=False,
+            max_context_len_to_capture=None,
+            max_seq_len_to_capture=8192,
            max_logprobs=20,
            disable_sliding_window=False,
            skip_tokenizer_init=True,
            served_model_name=None,
            limit_mm_per_prompt=None,
+            use_async_output_proc=True,
            config_format="auto",
        )
    finally:
@@ -178,65 +180,56 @@ def create_minimal_vllm_config(
 # ============================================================================


-# Backend-specific properties that can't be inferred from the backend class
-# Keys are AttentionBackendEnum names (uppercase)
+# Backend name to class name prefix mapping
+_BACKEND_NAME_MAP = {
+    "flashattn_mla": "FlashAttnMLA",
+    "flashmla": "FlashMLA",
+    "flashinfer_mla": "FlashInferMLA",
+    "cutlass_mla": "CutlassMLA",
+}
+
+# Special properties that differ from defaults
 _BACKEND_PROPERTIES = {
-    "FLASHMLA": {
+    "flashmla": {
        "query_format": "concat",  # Single concatenated tensor (vs tuple)
+        "block_size": 64,  # FlashMLA uses fixed block size
    },
-    "FLASHMLA_SPARSE": {
-        "query_format": "concat",  # Single concatenated tensor (vs tuple)
+    "flashinfer_mla": {
+        "block_size": 64,  # FlashInfer MLA only supports 32 or 64
    },
 }


 def _get_backend_config(backend: str) -> dict:
    """
-    Get backend configuration from AttentionBackendEnum.
+    Get backend configuration using naming conventions.

-    Uses the registry to get the backend class and extract configuration
-    from its methods (get_impl_cls, get_builder_cls, is_sparse, etc.).
-
-    Args:
-        backend: Backend name matching AttentionBackendEnum exactly
-        (e.g., "FLASHMLA_SPARSE")
-
-    Returns:
-        Dict with backend configuration
+    All MLA backends follow the pattern:
+    - Module: vllm.v1.attention.backends.mla.{backend}
+    - Impl: {Name}Impl
+    - Metadata: {Name}Metadata (or MLACommonMetadata)
+    - DecodeMetadata: {Name}DecodeMetadata (or MLACommonDecodeMetadata)
+    - MetadataBuilder: {Name}MetadataBuilder
    """
-    from vllm.v1.attention.backends.registry import AttentionBackendEnum
+    if backend not in _BACKEND_NAME_MAP:
+        raise ValueError(f"Unknown backend: {backend}")

-    try:
-        backend_enum = AttentionBackendEnum[backend]
-        backend_class = backend_enum.get_class()
-    except (KeyError, ValueError) as e:
-        valid_backends = [e.name for e in AttentionBackendEnum if e.name != "CUSTOM"]
-        raise ValueError(
-            f"Unknown backend: {backend}. "
-            f"Valid MLA backends: {[b for b in valid_backends if 'MLA' in b]}"
-        ) from e
-
-    # Get block size from backend class
-    block_sizes = backend_class.get_supported_kernel_block_sizes()
-    # Use first supported block size (backends typically support one for MLA)
-    block_size = block_sizes[0] if block_sizes else None
-    if hasattr(block_size, "value"):
-        # Handle MultipleOf enum
-        block_size = None
-
-    # Check if sparse via class method if available
-    is_sparse = getattr(backend_class, "is_sparse", lambda: False)()
-
-    # Get properties that can't be inferred
+    name = _BACKEND_NAME_MAP[backend]
    props = _BACKEND_PROPERTIES.get(backend, {})

+    # Check if backend uses common metadata (FlashInfer, CUTLASS)
+    uses_common = backend in ("flashinfer_mla", "cutlass_mla")
+
    return {
-        "backend_class": backend_class,
-        "impl_class": backend_class.get_impl_cls(),
-        "builder_class": backend_class.get_builder_cls(),
+        "module": f"vllm.v1.attention.backends.mla.{backend}",
+        "impl_class": f"{name}Impl",
+        "metadata_class": "MLACommonMetadata" if uses_common else f"{name}Metadata",
+        "decode_metadata_class": "MLACommonDecodeMetadata"
+        if uses_common
+        else f"{name}DecodeMetadata",
+        "builder_class": f"{name}MetadataBuilder",
        "query_format": props.get("query_format", "tuple"),
-        "block_size": block_size,
-        "is_sparse": is_sparse,
+        "block_size": props.get("block_size", None),
    }


@@ -454,26 +447,22 @@ def _create_backend_impl(
    mla_dims: dict,
    vllm_config: VllmConfig,
    device: torch.device,
-    max_num_tokens: int = 8192,
-    index_topk: int | None = None,
 ):
    """
    Create backend implementation instance.

    Args:
-        backend_cfg: Backend configuration dict from _get_backend_config()
+        backend_cfg: Backend configuration dict
        mla_dims: MLA dimension configuration
        vllm_config: VllmConfig instance
        device: Target device
-        max_num_tokens: Maximum number of tokens for sparse indexer buffer
-        index_topk: Topk value for sparse MLA backends

    Returns:
-        Tuple of (impl, layer, builder_instance, indexer)
+        Tuple of (impl, layer, builder_instance)
    """
-    # Get classes from backend config (already resolved by _get_backend_config)
-    impl_class = backend_cfg["impl_class"]
-    builder_class = backend_cfg["builder_class"]
+    # Import backend classes
+    backend_module = importlib.import_module(backend_cfg["module"])
+    impl_class = getattr(backend_module, backend_cfg["impl_class"])

    # Calculate scale
    scale = 1.0 / np.sqrt(mla_dims["qk_nope_head_dim"] + mla_dims["qk_rope_head_dim"])
@@ -485,44 +474,26 @@ def _create_backend_impl(
        v_head_dim=mla_dims["v_head_dim"],
    )

-    # Create indexer for sparse backends
-    indexer = None
-    if backend_cfg.get("is_sparse", False):
-        if index_topk is None:
-            index_topk = 2048  # Default topk for sparse MLA
-        indexer = MockIndexer(
-            max_num_tokens=max_num_tokens,
-            topk_tokens=index_topk,
-            device=device,
-        )
-
-    # Build impl kwargs
-    impl_kwargs = {
-        "num_heads": mla_dims["num_q_heads"],
-        "head_size": mla_dims["head_dim"],
-        "scale": scale,
-        "num_kv_heads": mla_dims["num_kv_heads"],
-        "alibi_slopes": None,
-        "sliding_window": None,
-        "kv_cache_dtype": "auto",
-        "logits_soft_cap": None,
-        "attn_type": "decoder",
-        "kv_sharing_target_layer_name": None,
-        "q_lora_rank": None,
-        "kv_lora_rank": mla_dims["kv_lora_rank"],
-        "qk_nope_head_dim": mla_dims["qk_nope_head_dim"],
-        "qk_rope_head_dim": mla_dims["qk_rope_head_dim"],
-        "qk_head_dim": mla_dims["qk_nope_head_dim"] + mla_dims["qk_rope_head_dim"],
-        "v_head_dim": mla_dims["v_head_dim"],
-        "kv_b_proj": mock_kv_b_proj,
-    }
-
-    # Add indexer for sparse backends
-    if indexer is not None:
-        impl_kwargs["indexer"] = indexer
-
    # Create impl
-    impl = impl_class(**impl_kwargs)
+    impl = impl_class(
+        num_heads=mla_dims["num_q_heads"],
+        head_size=mla_dims["head_dim"],
+        scale=scale,
+        num_kv_heads=mla_dims["num_kv_heads"],
+        alibi_slopes=None,
+        sliding_window=None,
+        kv_cache_dtype="auto",
+        logits_soft_cap=None,
+        attn_type="decoder",
+        kv_sharing_target_layer_name=None,
+        q_lora_rank=None,
+        kv_lora_rank=mla_dims["kv_lora_rank"],
+        qk_nope_head_dim=mla_dims["qk_nope_head_dim"],
+        qk_rope_head_dim=mla_dims["qk_rope_head_dim"],
+        qk_head_dim=mla_dims["qk_nope_head_dim"] + mla_dims["qk_rope_head_dim"],
+        v_head_dim=mla_dims["v_head_dim"],
+        kv_b_proj=mock_kv_b_proj,
+    )

    # Initialize DCP attributes
    if not hasattr(impl, "dcp_world_size") or impl.dcp_world_size in (None, -1):
@@ -544,7 +515,9 @@ def _create_backend_impl(

    # Create builder instance if needed
    builder_instance = None
-    if builder_class:
+    if backend_cfg["builder_class"]:
+        builder_class = getattr(backend_module, backend_cfg["builder_class"])
+
        # Populate static_forward_context so builder can find the layer
        # MockLayer inherits from AttentionLayerBase, so isinstance checks pass
        vllm_config.compilation_config.static_forward_context = {"placeholder": layer}
@@ -556,7 +529,7 @@ def _create_backend_impl(
            device=device,
        )

-    return impl, layer, builder_instance, indexer
+    return impl, layer, builder_instance


 # ============================================================================
@@ -621,7 +594,6 @@ def _run_single_benchmark(
    backend_cfg: dict,
    mla_dims: dict,
    device: torch.device,
-    indexer=None,
 ) -> BenchmarkResult:
    """
    Run a single benchmark iteration.
@@ -634,7 +606,6 @@ def _run_single_benchmark(
        backend_cfg: Backend configuration dict
        mla_dims: MLA dimension configuration
        device: Target device
-        indexer: Optional MockIndexer for sparse backends

    Returns:
        BenchmarkResult with timing statistics
@@ -642,9 +613,7 @@ def _run_single_benchmark(
    # Parse batch spec
    requests = parse_batch_spec(config.batch_spec)
    q_lens = [r.q_len for r in requests]
-    kv_lens = [r.kv_len for r in requests]
    total_q = sum(q_lens)
-    max_kv_len = max(kv_lens)

    # Determine block size
    block_size = backend_cfg["block_size"] or config.block_size
@@ -672,16 +641,8 @@ def _run_single_benchmark(
        torch.bfloat16,
    )

-    # Fill indexer with random indices for sparse backends
-    is_sparse = backend_cfg.get("is_sparse", False)
-    if is_sparse and indexer is not None:
-        indexer.fill_random_indices(total_q, max_kv_len)
-
-    # Determine which forward method to use
-    if is_sparse:
-        # Sparse backends use forward_mqa
-        forward_fn = lambda: impl.forward_mqa(decode_inputs, kv_cache, metadata, layer)
-    elif metadata.decode is not None:
+    # Determine which forward method to use based on metadata
+    if metadata.decode is not None:
        forward_fn = lambda: impl._forward_decode(
            decode_inputs, kv_cache, metadata, layer
        )
@@ -732,13 +693,11 @@ def _run_single_benchmark(
 def _run_mla_benchmark_batched(
    backend: str,
    configs_with_params: list[tuple],  # [(config, threshold, num_splits), ...]
-    index_topk: int = 2048,
 ) -> list[BenchmarkResult]:
    """
    Unified batched MLA benchmark runner for all backends.

-    Works for: flashattn_mla, flashmla, flashinfer_mla, cutlass_mla,
-               flashinfer_mla_sparse, flashmla_sparse
+    Works for: flashattn_mla, flashmla, flashinfer_mla, cutlass_mla

    This function reuses backend initialization across multiple benchmarks
    to avoid setup/teardown overhead.
@@ -748,7 +707,6 @@ def _run_mla_benchmark_batched(
        configs_with_params: List of (config, threshold, num_splits) tuples
            - threshold: reorder_batch_threshold (FlashAttn/FlashMLA only)
            - num_splits: num_kv_splits (CUTLASS only)
-        index_topk: Topk value for sparse MLA backends (default 2048)

    Returns:
        List of BenchmarkResult objects
@@ -772,27 +730,19 @@ def _run_mla_benchmark_batched(
    if mla_dims is None:
        mla_dims = setup_mla_dims("deepseek-v3")

-    # Determine if this is a sparse backend
-    is_sparse = backend_cfg.get("is_sparse", False)
-
    # Create and set vLLM config for MLA (reused across all benchmarks)
    vllm_config = create_minimal_vllm_config(
        model_name="deepseek-v3",  # Used only for model path
        block_size=block_size,
        mla_dims=mla_dims,  # Use custom dims from config or default
-        index_topk=index_topk if is_sparse else None,
    )

    results = []

    with set_current_vllm_config(vllm_config):
-        # Create backend impl, layer, builder, and indexer (reused across benchmarks)
-        impl, layer, builder_instance, indexer = _create_backend_impl(
-            backend_cfg,
-            mla_dims,
-            vllm_config,
-            device,
-            index_topk=index_topk if is_sparse else None,
+        # Create backend impl, layer, and builder (reused across benchmarks)
+        impl, layer, builder_instance = _create_backend_impl(
+            backend_cfg, mla_dims, vllm_config, device
        )

        # Run each benchmark with the shared impl
@@ -818,7 +768,6 @@ def _run_mla_benchmark_batched(
                    backend_cfg,
                    mla_dims,
                    device,
-                    indexer=indexer,
                )
                results.append(result)

@@ -844,24 +793,20 @@ def run_mla_benchmark(
    config,
    reorder_batch_threshold: int | None = None,
    num_kv_splits: int | None = None,
-    index_topk: int = 2048,
 ) -> BenchmarkResult | list[BenchmarkResult]:
    """
    Unified MLA benchmark runner for all backends.

-    Works for: flashattn_mla, flashmla, flashinfer_mla, cutlass_mla,
-               flashinfer_mla_sparse, flashmla_sparse
+    Works for: flashattn_mla, flashmla, flashinfer_mla, cutlass_mla

    Always uses batched execution internally for optimal performance.

    Args:
-        backend: Backend name (flashattn_mla, flashmla, flashinfer_mla, cutlass_mla,
-                 flashinfer_mla_sparse, flashmla_sparse)
+        backend: Backend name (flashattn_mla, flashmla, flashinfer_mla, cutlass_mla)
        config: BenchmarkConfig or list of (BenchmarkConfig, param) tuples
        reorder_batch_threshold: Threshold override for FlashAttn/FlashMLA
                                 (single config mode only)
        num_kv_splits: Number of KV splits for CUTLASS (single config mode only)
-        index_topk: Topk value for sparse MLA backends (default 2048)

    Returns:
        BenchmarkResult (single mode) or list of BenchmarkResult (batched mode)
@@ -871,9 +816,9 @@ def run_mla_benchmark(
        # Already in batched format
        if len(config) > 0 and isinstance(config[0], tuple):
            # Format: [(cfg, param), ...] where param is threshold or num_splits
-            if backend in ("flashattn_mla", "flashmla", "flashmla_sparse"):
+            if backend in ("flashattn_mla", "flashmla"):
                configs_with_params = [(cfg, param, None) for cfg, param in config]
-            else:  # cutlass_mla, flashinfer_mla, or sparse backends
+            else:  # cutlass_mla or flashinfer_mla
                configs_with_params = [(cfg, None, param) for cfg, param in config]
        else:
            # Format: [cfg, ...] - just configs
@@ -885,7 +830,7 @@ def run_mla_benchmark(
        return_single = True

    # Use unified batched execution
-    results = _run_mla_benchmark_batched(backend, configs_with_params, index_topk)
+    results = _run_mla_benchmark_batched(backend, configs_with_params)

    # Return single result or list based on input
    return results[0] if return_single else results
@@ -8,9 +8,7 @@ This module provides helpers for running standard attention backends
 (FlashAttention, Triton, FlashInfer) with real vLLM integration.
 """

-import logging
 import types
-from contextlib import contextmanager

 import numpy as np
 import torch
@@ -26,13 +24,8 @@ from vllm.config import (
    ParallelConfig,
    SchedulerConfig,
    VllmConfig,
-    set_current_vllm_config,
-)
-from vllm.v1.attention.backends.utils import (
-    CommonAttentionMetadata,
-    get_kv_cache_layout,
-    set_kv_cache_layout,
 )
+from vllm.v1.attention.backends.utils import CommonAttentionMetadata
 from vllm.v1.kv_cache_interface import FullAttentionSpec

 # ============================================================================
@@ -40,41 +33,37 @@ from vllm.v1.kv_cache_interface import FullAttentionSpec
 # ============================================================================


+_BACKEND_CONFIG = {
+    "flash": {
+        "module": "vllm.v1.attention.backends.flash_attn",
+        "backend_class": "FlashAttentionBackend",
+        "dtype": torch.float16,
+        "cache_layout": "standard",
+        # ^ [2, num_blocks, block_size, num_kv_heads, head_dim]
+    },
+    "triton": {
+        "module": "vllm.v1.attention.backends.triton_attn",
+        "backend_class": "TritonAttentionBackend",
+        "dtype": torch.float32,
+        "cache_layout": "standard",
+    },
+    "flashinfer": {
+        "module": "vllm.v1.attention.backends.flashinfer",
+        "backend_class": "FlashInferBackend",
+        "dtype": torch.float16,
+        "cache_layout": "flashinfer",
+        # ^ [num_blocks, 2, block_size, num_kv_heads, head_dim]
+    },
+}
+
+
 def _get_backend_config(backend: str) -> dict:
-    """
-    Get backend configuration from AttentionBackendEnum.
-
-    Args:
-        backend: Backend name matching AttentionBackendEnum exactly
-                 (e.g., "FLASH_ATTN", "TRITON_ATTN", "FLASHINFER")
-
-    Returns:
-        Dict with backend_class
-    """
-    from vllm.v1.attention.backends.registry import AttentionBackendEnum
-
-    try:
-        backend_enum = AttentionBackendEnum[backend]
-        backend_class = backend_enum.get_class()
-    except (KeyError, ValueError) as e:
-        valid_backends = [b.name for b in AttentionBackendEnum if b.name != "CUSTOM"]
+    if backend not in _BACKEND_CONFIG:
        raise ValueError(
-            f"Unknown backend: {backend}. Valid backends: {valid_backends}"
-        ) from e
-
-    return {"backend_class": backend_class}
-
-
-@contextmanager
-def log_warnings_and_errors_only():
-    """Temporarily set vLLM logger to WARNING level."""
-    logger = logging.getLogger("vllm")
-    old_level = logger.level
-    logger.setLevel(logging.WARNING)
-    try:
-        yield
-    finally:
-        logger.setLevel(old_level)
+            f"Unknown backend: {backend}. "
+            f"Available: {', '.join(_BACKEND_CONFIG.keys())}"
+        )
+    return _BACKEND_CONFIG[backend]


 # ============================================================================
@@ -99,7 +88,11 @@ def _build_common_attn_metadata(
    query_start_loc_cpu = query_start_loc.cpu()

    seq_lens = torch.tensor(kv_lens, dtype=torch.int32, device=device)
-    max_seq_len = int(seq_lens.max().item())
+    seq_lens_cpu = seq_lens.cpu()
+    max_seq_len = int(seq_lens_cpu.max())
+
+    context_lens = [kv - q for kv, q in zip(kv_lens, q_lens)]
+    num_computed_tokens_cpu = torch.tensor(context_lens, dtype=torch.int32)

    max_blocks = (max(kv_lens) + block_size - 1) // block_size
    num_blocks = batch_size * max_blocks
@@ -114,6 +107,8 @@ def _build_common_attn_metadata(
        query_start_loc=query_start_loc,
        query_start_loc_cpu=query_start_loc_cpu,
        seq_lens=seq_lens,
+        seq_lens_cpu=seq_lens_cpu,
+        num_computed_tokens_cpu=num_computed_tokens_cpu,
        num_reqs=batch_size,
        num_actual_tokens=total_tokens,
        max_query_len=max_query_len,
@@ -126,6 +121,7 @@ def _build_common_attn_metadata(

 def _create_vllm_config(
    config: BenchmarkConfig,
+    dtype: torch.dtype,
    max_num_blocks: int,
 ) -> VllmConfig:
    """Create a VllmConfig for benchmarking with mock model methods."""
@@ -133,7 +129,7 @@ def _create_vllm_config(
        model="meta-llama/Meta-Llama-3-8B",
        tokenizer="meta-llama/Meta-Llama-3-8B",
        trust_remote_code=False,
-        dtype="auto",  # Use model's native dtype
+        dtype=dtype,
        seed=0,
        max_model_len=1024,
    )
@@ -202,12 +198,15 @@ def _create_backend_impl(
    backend_cfg: dict,
    config: BenchmarkConfig,
    device: torch.device,
-    dtype: torch.dtype,
 ):
    """Create backend implementation instance."""
-    backend_class = backend_cfg["backend_class"]
+    import importlib
+
+    backend_module = importlib.import_module(backend_cfg["module"])
+    backend_class = getattr(backend_module, backend_cfg["backend_class"])

    scale = get_attention_scale(config.head_dim)
+    dtype = backend_cfg["dtype"]

    impl = backend_class.get_impl_cls()(
        num_heads=config.num_q_heads,
@@ -228,7 +227,7 @@ def _create_backend_impl(

    layer = MockLayer(device, kv_cache_spec=kv_cache_spec)

-    return backend_class, impl, layer
+    return backend_class, impl, layer, dtype


 def _create_metadata_builder(
@@ -236,44 +235,11 @@ def _create_metadata_builder(
    kv_cache_spec: FullAttentionSpec,
    vllm_config: VllmConfig,
    device: torch.device,
-    backend_name: str = "",
 ):
    """Create metadata builder instance."""
-    layer_names = ["layer_0"]
-    builder_cls = backend_class.get_builder_cls()
-
-    # Flashinfer needs get_per_layer_parameters mocked since we don't have
-    # real model layers registered
-    if backend_name == "FLASHINFER":
-        import unittest.mock
-
-        from vllm.v1.attention.backends.utils import PerLayerParameters
-
-        def mock_get_per_layer_parameters(vllm_config, layer_names, impl_cls):
-            head_size = vllm_config.model_config.get_head_size()
-            return {
-                layer_name: PerLayerParameters(
-                    window_left=-1,  # No sliding window
-                    logits_soft_cap=0.0,  # No soft cap
-                    sm_scale=1.0 / (head_size**0.5),  # Standard scale
-                )
-                for layer_name in layer_names
-            }
-
-        with unittest.mock.patch(
-            "vllm.v1.attention.backends.flashinfer.get_per_layer_parameters",
-            mock_get_per_layer_parameters,
-        ):
-            return builder_cls(
-                kv_cache_spec=kv_cache_spec,
-                layer_names=layer_names,
-                vllm_config=vllm_config,
-                device=device,
-            )
-
-    return builder_cls(
+    return backend_class.get_builder_cls()(
        kv_cache_spec=kv_cache_spec,
-        layer_names=layer_names,
+        layer_names=["layer_0"],
        vllm_config=vllm_config,
        device=device,
    )
@@ -315,44 +281,39 @@ def _create_input_tensors(
 def _create_kv_cache(
    config: BenchmarkConfig,
    max_num_blocks: int,
-    backend_class,
+    cache_layout: str,
    device: torch.device,
    dtype: torch.dtype,
 ) -> list:
-    """Create KV cache tensors for all layers using the backend's methods.
-
-    Uses the backend's get_kv_cache_shape() and get_kv_cache_stride_order()
-    to create the cache with the correct shape and memory layout.
-    """
-    # Get the logical shape from the backend
-    cache_shape = backend_class.get_kv_cache_shape(
-        num_blocks=max_num_blocks,
-        block_size=config.block_size,
-        num_kv_heads=config.num_kv_heads,
-        head_size=config.head_dim,
-    )
-
-    # Get the stride order for custom memory layout
-    try:
-        stride_order = backend_class.get_kv_cache_stride_order()
-        assert len(stride_order) == len(cache_shape)
-    except (AttributeError, NotImplementedError):
-        stride_order = tuple(range(len(cache_shape)))
-
-    # Permute shape to physical layout order
-    physical_shape = tuple(cache_shape[i] for i in stride_order)
-
-    # Compute inverse permutation to get back to logical view
-    inv_order = [stride_order.index(i) for i in range(len(stride_order))]
-
-    cache_list = []
-    for _ in range(config.num_layers):
-        # Allocate in physical layout order (contiguous in memory)
-        cache = torch.zeros(*physical_shape, device=device, dtype=dtype)
-        # Permute to logical view
-        cache = cache.permute(*inv_order)
-        cache_list.append(cache)
-
+    """Create KV cache tensors for all layers."""
+    if cache_layout == "flashinfer":
+        # FlashInfer layout: [num_blocks, 2, block_size, num_kv_heads, head_dim]
+        cache_list = [
+            torch.zeros(
+                max_num_blocks,
+                2,
+                config.block_size,
+                config.num_kv_heads,
+                config.head_dim,
+                device=device,
+                dtype=dtype,
+            )
+            for _ in range(config.num_layers)
+        ]
+    else:
+        # Standard layout: [2, num_blocks, block_size, num_kv_heads, head_dim]
+        cache_list = [
+            torch.zeros(
+                2,
+                max_num_blocks,
+                config.block_size,
+                config.num_kv_heads,
+                config.head_dim,
+                device=device,
+                dtype=dtype,
+            )
+            for _ in range(config.num_layers)
+        ]
    return cache_list


@@ -435,7 +396,7 @@ def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:
    """
    Run standard attention benchmark with real kernels.

-    Supports: FLASH_ATTN, TRITON_ATTN, FLASHINFER
+    Supports: flash, triton, flashinfer

    Args:
        config: Benchmark configuration
@@ -450,79 +411,60 @@ def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:

    requests = parse_batch_spec(config.batch_spec)

-    if config.backend == "FLASHINFER":
+    if config.backend == "flashinfer":
        requests = reorder_for_flashinfer(requests)

    q_lens = [r.q_len for r in requests]
    kv_lens = [r.kv_len for r in requests]
    total_q = sum(q_lens)
    max_kv = max(kv_lens)
-    batch_size = len(q_lens)

-    # Calculate total blocks needed: batch_size * max_blocks_per_request
-    max_blocks_per_request = (max_kv + config.block_size - 1) // config.block_size
-    max_num_blocks = batch_size * max_blocks_per_request
+    max_num_blocks = (max_kv + config.block_size - 1) // config.block_size

-    # Suppress vLLM logs during setup to reduce spam
-    with log_warnings_and_errors_only():
-        # Create vllm_config first - uses model's native dtype via "auto"
-        vllm_config = _create_vllm_config(config, max_num_blocks)
-        dtype = vllm_config.model_config.dtype
+    backend_class, impl, layer, dtype = _create_backend_impl(
+        backend_cfg, config, device
+    )

-        # Wrap everything in set_current_vllm_config context
-        # This is required for backends like flashinfer that need global config
-        with set_current_vllm_config(vllm_config):
-            backend_class, impl, layer = _create_backend_impl(
-                backend_cfg, config, device, dtype
-            )
+    common_metadata = _build_common_attn_metadata(
+        q_lens, kv_lens, config.block_size, device
+    )

-            # Set KV cache layout if the backend requires a specific one
-            # (e.g., FlashInfer requires HND on SM100/Blackwell for TRTLLM attention)
-            required_layout = backend_class.get_required_kv_cache_layout()
-            if required_layout is not None:
-                set_kv_cache_layout(required_layout)
-                get_kv_cache_layout.cache_clear()
+    kv_cache_spec = FullAttentionSpec(
+        block_size=config.block_size,
+        num_kv_heads=config.num_kv_heads,
+        head_size=config.head_dim,
+        dtype=dtype,
+    )

-            common_metadata = _build_common_attn_metadata(
-                q_lens, kv_lens, config.block_size, device
-            )
+    vllm_config = _create_vllm_config(config, dtype, max_num_blocks)

-            kv_cache_spec = FullAttentionSpec(
-                block_size=config.block_size,
-                num_kv_heads=config.num_kv_heads,
-                head_size=config.head_dim,
-                dtype=dtype,
-            )
+    builder = _create_metadata_builder(
+        backend_class, kv_cache_spec, vllm_config, device
+    )

-            builder = _create_metadata_builder(
-                backend_class, kv_cache_spec, vllm_config, device, config.backend
-            )
+    attn_metadata = builder.build(
+        common_prefix_len=0,
+        common_attn_metadata=common_metadata,
+    )

-            attn_metadata = builder.build(
-                common_prefix_len=0,
-                common_attn_metadata=common_metadata,
-            )
+    q_list, k_list, v_list = _create_input_tensors(config, total_q, device, dtype)

-            q_list, k_list, v_list = _create_input_tensors(
-                config, total_q, device, dtype
-            )
+    cache_list = _create_kv_cache(
+        config, max_num_blocks, backend_cfg["cache_layout"], device, dtype
+    )

-            cache_list = _create_kv_cache(
-                config, max_num_blocks, backend_class, device, dtype
-            )
-
-            times, mem_stats = _run_single_benchmark(
-                config,
-                impl,
-                layer,
-                q_list,
-                k_list,
-                v_list,
-                cache_list,
-                attn_metadata,
-                device,
-                dtype,
-            )
+    times, mem_stats = _run_single_benchmark(
+        config,
+        impl,
+        layer,
+        q_list,
+        k_list,
+        v_list,
+        cache_list,
+        attn_metadata,
+        device,
+        dtype,
+    )

    mean_time = np.mean(times)
    throughput = total_q / mean_time if mean_time > 0 else 0
@@ -11,7 +11,6 @@ import torch
 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
 from tests.kernels.moe.utils import make_dummy_moe_config
 from vllm import _custom_ops as ops
-from vllm.model_executor.layers.fused_moe.activation import MoEActivation
 from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_config
 from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassExpertsFp8
 from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts, fused_topk
@@ -162,7 +161,7 @@ def bench_run(
                w2_fp8q_cutlass,
                topk_weights,
                topk_ids,
-                activation=MoEActivation.SILU,
+                activation="silu",
                global_num_experts=num_experts,
            )
    torch.cuda.synchronize()
@@ -5,7 +5,7 @@
 Benchmark for FlashInfer fused collective operations vs standard operations.

 This benchmark compares:
-1. FlashInfer's allreduce_fusion (fused allreduce + rmsnorm + optional quant)
+1. FlashInfer's trtllm_allreduce_fusion (fused allreduce + rmsnorm + optional quant)
 2. Standard tensor_model_parallel_all_reduce + separate rmsnorm/quant operations

 Usage with torchrun:
@@ -24,6 +24,7 @@ import torch.distributed as dist  # type: ignore

 from vllm.config.vllm import CompilationConfig, VllmConfig, set_current_vllm_config
 from vllm.distributed import (
+    get_tp_group,
    tensor_model_parallel_all_reduce,
 )
 from vllm.distributed.parallel_state import (
@@ -51,12 +52,11 @@ logger = init_logger(__name__)
 try:
    import flashinfer.comm as flashinfer_comm  # type: ignore

-    if not (
-        hasattr(flashinfer_comm, "allreduce_fusion")
-        and hasattr(flashinfer_comm, "create_allreduce_fusion_workspace")
-    ):
+    if not hasattr(flashinfer_comm, "trtllm_allreduce_fusion"):
        flashinfer_comm = None
-        logger.warning("FlashInfer comm module found but missing allreduce_fusion API")
+        logger.warning(
+            "FlashInfer comm module found but missing trtllm_allreduce_fusion"
+        )
 except ImportError:
    flashinfer_comm = None
    logger.warning("FlashInfer not found, only benchmarking standard operations")
@@ -75,7 +75,7 @@ _FI_MAX_SIZES = {
 }

 # Global workspace tensor for FlashInfer
-_FI_WORKSPACE = None
+_FI_WORKSPACE_TENSOR = None


 def setup_flashinfer_workspace(
@@ -83,10 +83,10 @@ def setup_flashinfer_workspace(
    rank: int,
    hidden_dim: int,
    max_token_num: int,
-    dtype: torch.dtype,
+    use_fp32_lamport: bool = False,
 ):
    """Setup FlashInfer workspace for fused allreduce operations."""
-    global _FI_WORKSPACE
+    global _FI_WORKSPACE_TENSOR

    if flashinfer_comm is None:
        return None, None
@@ -96,29 +96,33 @@ def setup_flashinfer_workspace(
        return None, None

    try:
-        workspace = flashinfer_comm.create_allreduce_fusion_workspace(
-            backend="trtllm",
-            world_size=world_size,
-            rank=rank,
-            max_token_num=max_token_num,
-            hidden_dim=hidden_dim,
-            dtype=dtype,
+        # Create IPC workspace
+        ipc_handles, workspace_tensor = (
+            flashinfer_comm.trtllm_create_ipc_workspace_for_all_reduce_fusion(
+                tp_rank=rank,
+                tp_size=world_size,
+                max_token_num=max_token_num,
+                hidden_dim=hidden_dim,
+                group=get_tp_group().device_group,
+                use_fp32_lamport=use_fp32_lamport,
+            )
        )

-        _FI_WORKSPACE = workspace
-        return workspace
+        _FI_WORKSPACE_TENSOR = workspace_tensor
+        return ipc_handles, workspace_tensor
    except Exception as e:
        logger.error("Failed to setup FlashInfer workspace: %s", e)
-        return None
+        return None, None


-def cleanup_flashinfer_workspace(workspace):
+def cleanup_flashinfer_workspace(ipc_handles):
    """Cleanup FlashInfer workspace."""
-    if flashinfer_comm is None or workspace is None:
+    if flashinfer_comm is None or ipc_handles is None:
        return

    try:
-        workspace.destroy()
+        group = get_tp_group().device_group
+        flashinfer_comm.trtllm_destroy_ipc_workspace_for_all_reduce(ipc_handles, group)
    except Exception as e:
        logger.error("Failed to cleanup FlashInfer workspace: %s", e)

@@ -128,15 +132,25 @@ class FlashInferFusedAllReduceParams:

    def __init__(
        self,
+        rank: int,
+        world_size: int,
+        use_fp32_lamport: bool = False,
        max_token_num: int = 1024,
    ):
+        self.rank = rank
+        self.world_size = world_size
+        self.use_fp32_lamport = use_fp32_lamport
+        self.trigger_completion_at_end = True
        self.launch_with_pdl = True
        self.fp32_acc = True
        self.max_token_num = max_token_num

    def get_trtllm_fused_allreduce_kwargs(self):
        return {
+            "world_rank": self.rank,
+            "world_size": self.world_size,
            "launch_with_pdl": self.launch_with_pdl,
+            "trigger_completion_at_end": self.trigger_completion_at_end,
            "fp32_acc": self.fp32_acc,
        }

@@ -151,7 +165,7 @@ def flashinfer_fused_allreduce_rmsnorm(
    norm_out: torch.Tensor | None = None,
 ):
    """FlashInfer fused allreduce + rmsnorm operation."""
-    if flashinfer_comm is None or _FI_WORKSPACE is None:
+    if flashinfer_comm is None or _FI_WORKSPACE_TENSOR is None:
        raise RuntimeError("FlashInfer not available or workspace not initialized")

    if norm_out is None:
@@ -160,15 +174,18 @@ def flashinfer_fused_allreduce_rmsnorm(
    else:
        residual_out = input_tensor

-    flashinfer_comm.allreduce_fusion(
-        input=input_tensor,
-        workspace=_FI_WORKSPACE,
-        pattern=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNorm,
+    flashinfer_comm.trtllm_allreduce_fusion(
+        allreduce_in=input_tensor,
+        token_num=input_tensor.shape[0],
        residual_in=residual,
        residual_out=residual_out,
        norm_out=norm_out,
        rms_gamma=rms_gamma,
        rms_eps=rms_eps,
+        hidden_dim=input_tensor.shape[-1],
+        workspace_ptrs=_FI_WORKSPACE_TENSOR,
+        pattern_code=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNorm,
+        allreduce_out=None,
        quant_out=None,
        scale_out=None,
        layout_code=flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4,
@@ -190,7 +207,7 @@ def flashinfer_fused_allreduce_rmsnorm_fp8_quant(
    quant_out: torch.Tensor | None = None,
 ):
    """FlashInfer fused allreduce + rmsnorm + FP8 quantization."""
-    if flashinfer_comm is None or _FI_WORKSPACE is None:
+    if flashinfer_comm is None or _FI_WORKSPACE_TENSOR is None:
        raise RuntimeError("FlashInfer not available or workspace not initialized")

    if norm_out is None:
@@ -199,15 +216,18 @@ def flashinfer_fused_allreduce_rmsnorm_fp8_quant(
    else:
        residual_out = input_tensor

-    flashinfer_comm.allreduce_fusion(
-        input=input_tensor,
-        workspace=_FI_WORKSPACE,
-        pattern=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNormFP8Quant,
+    flashinfer_comm.trtllm_allreduce_fusion(
+        allreduce_in=input_tensor,
+        token_num=input_tensor.shape[0],
        residual_in=residual,
        residual_out=residual_out,
        norm_out=norm_out,
        rms_gamma=rms_gamma,
        rms_eps=rms_eps,
+        hidden_dim=input_tensor.shape[-1],
+        workspace_ptrs=_FI_WORKSPACE_TENSOR,
+        pattern_code=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNormFP8Quant,
+        allreduce_out=None,
        quant_out=quant_out,
        scale_out=None,
        layout_code=flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4,
@@ -230,7 +250,7 @@ def flashinfer_fused_allreduce_rmsnorm_fp4_quant(
    norm_out: torch.Tensor | None = None,
 ):
    """FlashInfer fused allreduce + rmsnorm + FP4 quantization."""
-    if flashinfer_comm is None or _FI_WORKSPACE is None:
+    if flashinfer_comm is None or _FI_WORKSPACE_TENSOR is None:
        raise RuntimeError("FlashInfer not available or workspace not initialized")

    if norm_out is None:
@@ -239,15 +259,18 @@ def flashinfer_fused_allreduce_rmsnorm_fp4_quant(
    else:
        residual_out = input_tensor

-    flashinfer_comm.allreduce_fusion(
-        input=input_tensor,
-        workspace=_FI_WORKSPACE,
-        pattern=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNormFP4Quant,
+    flashinfer_comm.trtllm_allreduce_fusion(
+        allreduce_in=input_tensor,
+        token_num=input_tensor.shape[0],
        residual_in=residual,
        residual_out=residual_out,
        norm_out=norm_out,
        rms_gamma=rms_gamma,
        rms_eps=rms_eps,
+        hidden_dim=input_tensor.shape[-1],
+        workspace_ptrs=_FI_WORKSPACE_TENSOR,
+        pattern_code=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNormFP4Quant,
+        allreduce_out=None,
        quant_out=quant_out,
        scale_out=output_scale,
        layout_code=flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4,
@@ -1017,31 +1040,23 @@ def main():
    configs = list(itertools.product(args.num_tokens, dtypes, residual_options))

    # Setup FlashInfer workspace if available
-    workspace = None
+    ipc_handles = None
    allreduce_params = None

    if flashinfer_comm is not None:
        # Use the largest hidden dimension for workspace setup
-        max_element_size = max(torch.finfo(dt).bits // 8 for dt in dtypes)
-        workspace_dtype = (
-            torch.float32
-            if max_element_size == 4
-            else (torch.bfloat16 if torch.bfloat16 in dtypes else torch.float16)
-        )
        max_num_token = _FI_MAX_SIZES.get(world_size) // (
-            args.hidden_dim * max_element_size
+            args.hidden_dim * world_size * 2
        )

-        workspace = setup_flashinfer_workspace(
-            world_size,
-            rank,
-            args.hidden_dim,
-            max_num_token,
-            dtype=workspace_dtype,
+        ipc_handles, workspace_tensor = setup_flashinfer_workspace(
+            world_size, rank, args.hidden_dim, max_num_token
        )

-        if workspace is not None:
+        if workspace_tensor is not None:
            allreduce_params = FlashInferFusedAllReduceParams(
+                rank=rank,
+                world_size=world_size,
                max_token_num=max_num_token,
            )

@@ -1104,8 +1119,8 @@ def main():

    finally:
        # Cleanup
-        if workspace is not None:
-            cleanup_flashinfer_workspace(workspace)
+        if ipc_handles is not None:
+            cleanup_flashinfer_workspace(ipc_handles)

        dist.barrier()

@@ -16,7 +16,6 @@ import torch
 from ray.experimental.tqdm_ray import tqdm

 from vllm.model_executor.layers.fused_moe import fused_topk
-from vllm.model_executor.layers.fused_moe.activation import MoEActivation
 from vllm.model_executor.layers.fused_moe.config import (
    FusedMoEConfig,
    FusedMoEParallelConfig,
@@ -100,38 +99,13 @@ def benchmark_config(
    dtype: torch.dtype,
    use_fp8_w8a8: bool,
    use_int8_w8a16: bool,
-    use_int4_w4a16: bool = False,
    num_iters: int = 100,
    block_quant_shape: list[int] = None,
    use_deep_gemm: bool = False,
 ) -> float:
    init_dtype = torch.float16 if use_fp8_w8a8 else dtype
    x = torch.randn(num_tokens, hidden_size, dtype=dtype)
-    if use_int4_w4a16:
-        # Int4 packed weights: 2 int4 values per uint8 byte
-        # K dimension is packed (halved)
-        intermediate_size = shard_intermediate_size // 2  # after silu_and_mul
-        w1 = torch.randint(
-            0,
-            255,
-            (
-                num_experts,
-                shard_intermediate_size,
-                hidden_size // 2,  # int4 packing
-            ),
-            dtype=torch.uint8,
-        )
-        w2 = torch.randint(
-            0,
-            255,
-            (
-                num_experts,
-                hidden_size,
-                intermediate_size // 2,  # int4 packing
-            ),
-            dtype=torch.uint8,
-        )
-    elif use_int8_w8a16:
+    if use_int8_w8a16:
        w1 = torch.randint(
            -127,
            127,
@@ -165,20 +139,7 @@ def benchmark_config(
    w2_scale = None
    a1_scale = None
    a2_scale = None
-    if use_int4_w4a16:
-        if block_quant_shape is None:
-            raise ValueError("block_quant_shape is required for int4_w4a16")
-        group_size = block_quant_shape[1]
-        # Scales shape: (E, N, K // group_size) in fp16
-        w1_scale = torch.rand(
-            (num_experts, shard_intermediate_size, hidden_size // group_size),
-            dtype=dtype,
-        )
-        w2_scale = torch.rand(
-            (num_experts, hidden_size, intermediate_size // group_size),
-            dtype=dtype,
-        )
-    elif use_int8_w8a16:
+    if use_int8_w8a16:
        w1_scale = torch.randn(
            (num_experts, 2 * shard_intermediate_size), dtype=torch.float32
        )
@@ -237,7 +198,6 @@ def benchmark_config(
            a1_scale=a1_scale,
            a2_scale=a2_scale,
            block_shape=block_quant_shape,
-            weight_dtype="int4" if use_int4_w4a16 else None,
        )

        deep_gemm_experts = None
@@ -251,8 +211,7 @@ def benchmark_config(
                        hidden_dim=hidden_size,
                        intermediate_size_per_partition=shard_intermediate_size,
                        num_local_experts=num_experts,
-                        num_logical_experts=num_experts,
-                        activation=MoEActivation.SILU,
+                        activation="silu",
                        moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
                        in_dtype=init_dtype,
                        routing_method=RoutingMethodType.TopK,
@@ -267,10 +226,9 @@ def benchmark_config(
                x, input_gating, topk, renormalize=not use_deep_gemm
            )

-            inplace = not disable_inplace()
            if use_deep_gemm:
                return deep_gemm_experts(
-                    x, w1, w2, topk_weights, topk_ids, inplace=inplace
+                    x, w1, w2, topk_weights, topk_ids, inplace=True
                )
            return fused_experts(
                x,
@@ -278,7 +236,7 @@ def benchmark_config(
                w2,
                topk_weights,
                topk_ids,
-                inplace=inplace,
+                inplace=True,
                quant_config=quant_config,
            )

@@ -520,7 +478,6 @@ class BenchmarkWorker:
        dtype: torch.dtype,
        use_fp8_w8a8: bool,
        use_int8_w8a16: bool,
-        use_int4_w4a16: bool = False,
        block_quant_shape: list[int] = None,
        use_deep_gemm: bool = False,
    ) -> tuple[dict[str, int], float]:
@@ -528,10 +485,7 @@ class BenchmarkWorker:

        set_random_seed(self.seed)
        dtype_str = _get_config_dtype_str(
-            dtype,
-            use_int8_w8a16=use_int8_w8a16,
-            use_fp8_w8a8=use_fp8_w8a8,
-            use_int4_w4a16=use_int4_w4a16,
+            dtype, use_int8_w8a16=use_int8_w8a16, use_fp8_w8a8=use_fp8_w8a8
        )
        # NOTE(woosuk): The current naming convention uses w2.shape[2], which
        # is the intermediate size after silu_and_mul.
@@ -562,7 +516,6 @@ class BenchmarkWorker:
            dtype,
            use_fp8_w8a8,
            use_int8_w8a16,
-            use_int4_w4a16=use_int4_w4a16,
            num_iters=100,
            block_quant_shape=block_quant_shape,
            use_deep_gemm=use_deep_gemm,
@@ -579,7 +532,6 @@ class BenchmarkWorker:
        dtype: torch.dtype,
        use_fp8_w8a8: bool,
        use_int8_w8a16: bool,
-        use_int4_w4a16: bool,
        search_space: list[dict[str, int]],
        block_quant_shape: list[int],
        use_deep_gemm: bool,
@@ -590,7 +542,7 @@ class BenchmarkWorker:
        best_config = None
        best_time = float("inf")
        if current_platform.is_rocm():
-            is_fp16 = not (use_fp8_w8a8 or use_int8_w8a16 or use_int4_w4a16)
+            is_fp16 = not (use_fp8_w8a8 or use_int8_w8a16)
            search_space = prune_rocm_search_space(
                num_tokens,
                shard_intermediate_size,
@@ -619,7 +571,6 @@ class BenchmarkWorker:
                        dtype,
                        use_fp8_w8a8,
                        use_int8_w8a16,
-                        use_int4_w4a16,
                        num_iters=20,
                        block_quant_shape=block_quant_shape,
                        use_deep_gemm=use_deep_gemm,
@@ -667,7 +618,6 @@ def sort_config(config: BenchmarkConfig) -> BenchmarkConfig:
            else {}
        ),
        **({"kpack": config["kpack"]} if "kpack" in config else {}),
-        **({"SPLIT_K": config["SPLIT_K"]} if "SPLIT_K" in config else {}),
    }


@@ -680,15 +630,11 @@ def save_configs(
    dtype: torch.dtype,
    use_fp8_w8a8: bool,
    use_int8_w8a16: bool,
-    use_int4_w4a16: bool,
    block_quant_shape: list[int],
    save_dir: str,
 ) -> None:
    dtype_str = _get_config_dtype_str(
-        dtype,
-        use_int8_w8a16=use_int8_w8a16,
-        use_fp8_w8a8=use_fp8_w8a8,
-        use_int4_w4a16=use_int4_w4a16,
+        dtype, use_int8_w8a16=use_int8_w8a16, use_fp8_w8a8=use_fp8_w8a8
    )

    # NOTE(woosuk): The current naming convention uses w2.shape[2], which
@@ -740,7 +686,6 @@ def get_model_params(config):
        "DeepseekV2ForCausalLM",
        "DeepseekV3ForCausalLM",
        "DeepseekV32ForCausalLM",
-        "GlmMoeDsaForCausalLM",
        "Glm4MoeForCausalLM",
        "Glm4MoeLiteForCausalLM",
        "NemotronHForCausalLM",
@@ -790,38 +735,6 @@ def get_model_params(config):
    return E, topk, intermediate_size, hidden_size


-def get_quantization_group_size(config) -> int | None:
-    """Extract the quantization group size from the HF model config.
-
-    This reads directly from the HuggingFace config object (as returned by
-    ``get_config()``), not from vLLM's quantization config classes.
-
-    Supports AWQ/GPTQ-style configs (direct 'group_size' key) and
-    compressed-tensors configs (nested inside 'config_groups').
-    """
-    quantization_config = getattr(config, "quantization_config", {})
-    if not isinstance(quantization_config, dict):
-        return None
-    # AWQ / GPTQ style: group_size is a top-level key
-    gs = quantization_config.get("group_size")
-    if gs is not None:
-        return gs
-    # compressed-tensors style: group_size is nested in config_groups
-    config_groups = quantization_config.get("config_groups", {})
-    if not isinstance(config_groups, dict):
-        return None
-    for group_cfg in config_groups.values():
-        if not isinstance(group_cfg, dict):
-            continue
-        weights = group_cfg.get("weights", {})
-        if not isinstance(weights, dict):
-            continue
-        gs = weights.get("group_size")
-        if gs is not None:
-            return gs
-    return None
-
-
 def main(args: argparse.Namespace):
    print(args)

@@ -840,20 +753,7 @@ def main(args: argparse.Namespace):
    dtype = torch.float16 if current_platform.is_rocm() else config.dtype
    use_fp8_w8a8 = args.dtype == "fp8_w8a8"
    use_int8_w8a16 = args.dtype == "int8_w8a16"
-    use_int4_w4a16 = args.dtype == "int4_w4a16"
    block_quant_shape = get_weight_block_size_safety(config)
-    if use_int4_w4a16:
-        group_size = get_quantization_group_size(config)
-        if group_size is None:
-            raise ValueError(
-                "Could not determine group_size from model config. "
-                "The model's quantization_config must contain a 'group_size' "
-                "field (AWQ/GPTQ) or 'config_groups.*.weights.group_size' "
-                "(compressed-tensors)."
-            )
-        # For int4_w4a16, block_shape = [0, group_size]
-        # block_shape[0]=0 means no block quantization on N dimension
-        block_quant_shape = [0, group_size]

    if args.batch_size is None:
        batch_sizes = [
@@ -907,20 +807,8 @@ def main(args: argparse.Namespace):
        return ray.get(outputs)

    if args.tune:
-        # int4_w4a16 weights are uint8-packed, not fp16; treat like fp8 for
-        # search space generation (no matrix_instr_nonkdim/kpack exploration).
-        is_fp16 = not (use_fp8_w8a8 or use_int8_w8a16 or use_int4_w4a16)
-        # For int4_w4a16, the group_size constraint on BLOCK_SIZE_K does not
-        # apply: the gptq_awq kernel handles arbitrary BLOCK_SIZE_K regardless
-        # of group_size. Skip block_quant_shape filtering to keep the full
-        # search space (e.g. BLOCK_SIZE_K=64 with group_size=128).
-        tune_block_quant_shape = None if use_int4_w4a16 else block_quant_shape
-        search_space = get_configs_compute_bound(is_fp16, tune_block_quant_shape)
-        if use_int4_w4a16:
-            # SPLIT_K is a required kernel constexpr for gptq_awq kernel;
-            # only SPLIT_K=1 is used at runtime, so fix it during tuning.
-            for cfg in search_space:
-                cfg["SPLIT_K"] = 1
+        is_fp16 = not (use_fp8_w8a8 or use_int8_w8a16)
+        search_space = get_configs_compute_bound(is_fp16, block_quant_shape)
        print(f"Start tuning over {len(search_space)} configurations...")
        if use_deep_gemm:
            raise ValueError(
@@ -940,7 +828,6 @@ def main(args: argparse.Namespace):
                    dtype,
                    use_fp8_w8a8,
                    use_int8_w8a16,
-                    use_int4_w4a16,
                    search_space,
                    block_quant_shape,
                    use_deep_gemm,
@@ -960,7 +847,6 @@ def main(args: argparse.Namespace):
            dtype,
            use_fp8_w8a8,
            use_int8_w8a16,
-            use_int4_w4a16,
            block_quant_shape,
            args.save_dir,
        )
@@ -979,7 +865,6 @@ def main(args: argparse.Namespace):
                    dtype,
                    use_fp8_w8a8,
                    use_int8_w8a16,
-                    use_int4_w4a16,
                    block_quant_shape,
                    use_deep_gemm,
                )
@@ -1002,10 +887,7 @@ if __name__ == "__main__":
    )
    parser.add_argument("--enable-expert-parallel", "-enable-ep", action="store_true")
    parser.add_argument(
-        "--dtype",
-        type=str,
-        choices=["auto", "fp8_w8a8", "int8_w8a16", "int4_w4a16"],
-        default="auto",
+        "--dtype", type=str, choices=["auto", "fp8_w8a8", "int8_w8a16"], default="auto"
    )
    parser.add_argument("--use-deep-gemm", action="store_true")
    parser.add_argument(
@@ -38,7 +38,7 @@ else()
  FetchContent_Declare(
          vllm-flash-attn
          GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
-          GIT_TAG 5824e6e2008271063c3229ab3e7032bd74abbbc6
+          GIT_TAG 188be16520ceefdc625fdf71365585d2ee348fe2
          GIT_PROGRESS TRUE
          # Don't share the vllm-flash-attn build between build types
          BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
@@ -9,111 +9,6 @@

 namespace vllm {

-struct alignas(32) u32x8_t {
-  uint32_t u0, u1, u2, u3, u4, u5, u6, u7;
-};
-
-__device__ __forceinline__ void ld256(u32x8_t& val, const u32x8_t* ptr) {
-#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000
-  asm volatile("ld.global.nc.v8.u32 {%0,%1,%2,%3,%4,%5,%6,%7}, [%8];\n"
-               : "=r"(val.u0), "=r"(val.u1), "=r"(val.u2), "=r"(val.u3),
-                 "=r"(val.u4), "=r"(val.u5), "=r"(val.u6), "=r"(val.u7)
-               : "l"(ptr));
-#else
-  const uint4* uint_ptr = reinterpret_cast<const uint4*>(ptr);
-  uint4 top_half = __ldg(&uint_ptr[0]);
-  uint4 bottom_half = __ldg(&uint_ptr[1]);
-  val.u0 = top_half.x;
-  val.u1 = top_half.y;
-  val.u2 = top_half.z;
-  val.u3 = top_half.w;
-  val.u4 = bottom_half.x;
-  val.u5 = bottom_half.y;
-  val.u6 = bottom_half.z;
-  val.u7 = bottom_half.w;
-#endif
-}
-
-__device__ __forceinline__ void st256(u32x8_t& val, u32x8_t* ptr) {
-#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000
-  asm volatile("st.global.v8.u32 [%0], {%1,%2,%3,%4,%5,%6,%7,%8};\n"
-               :
-               : "l"(ptr), "r"(val.u0), "r"(val.u1), "r"(val.u2), "r"(val.u3),
-                 "r"(val.u4), "r"(val.u5), "r"(val.u6), "r"(val.u7)
-               : "memory");
-#else
-  uint4* uint_ptr = reinterpret_cast<uint4*>(ptr);
-  uint_ptr[0] = make_uint4(val.u0, val.u1, val.u2, val.u3);
-  uint_ptr[1] = make_uint4(val.u4, val.u5, val.u6, val.u7);
-#endif
-}
-
-template <bool support_256>
-struct VecTraits;
-
-template <>
-struct VecTraits<true> {
-  static constexpr int ARCH_MAX_VEC_SIZE = 32;
-  using vec_t = u32x8_t;
-};
-
-template <>
-struct VecTraits<false> {
-  static constexpr int ARCH_MAX_VEC_SIZE = 16;
-  using vec_t = int4;
-};
-
-template <typename T>
-struct PackedTraits;
-
-template <>
-struct PackedTraits<c10::BFloat16> {
-  using packed_t = __nv_bfloat162;
-};
-
-template <>
-struct PackedTraits<c10::Half> {
-  using packed_t = __half2;
-};
-
-template <>
-struct PackedTraits<float> {
-  using packed_t = float2;
-};
-
-template <typename packed_t>
-__device__ __forceinline__ float2 cast_to_float2(const packed_t& val) {
-  if constexpr (std::is_same_v<packed_t, __nv_bfloat162>) {
-    return __bfloat1622float2(val);
-  } else if constexpr (std::is_same_v<packed_t, __half2>) {
-    return __half22float2(val);
-  } else if constexpr (std::is_same_v<packed_t, float2>) {
-    return float2(val);
-  }
-}
-
-template <typename packed_t>
-__device__ __forceinline__ packed_t cast_to_packed(const float2& val) {
-  if constexpr (std::is_same_v<packed_t, __nv_bfloat162>) {
-    return __float22bfloat162_rn(val);
-  } else if constexpr (std::is_same_v<packed_t, __half2>) {
-    return __float22half2_rn(val);
-  } else if constexpr (std::is_same_v<packed_t, float2>) {
-    return float2(val);
-  }
-}
-
-template <typename packed_t>
-__device__ __forceinline__ packed_t packed_mul(const packed_t& x,
-                                               const packed_t& y) {
-  if constexpr (std::is_same_v<packed_t, __nv_bfloat162> ||
-                std::is_same_v<packed_t, __half2>) {
-    return __hmul2(x, y);
-  } else if constexpr (std::is_same_v<packed_t, float2>) {
-    return make_float2(x.x * y.x, x.y * y.y);
-  }
-}
-
 template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&),
          bool act_first>
 __device__ __forceinline__ scalar_t compute(const scalar_t& x,
@@ -121,69 +16,52 @@ __device__ __forceinline__ scalar_t compute(const scalar_t& x,
  return act_first ? ACT_FN(x) * y : x * ACT_FN(y);
 }

-template <typename packed_t, packed_t (*PACKED_ACT_FN)(const packed_t&),
-          bool act_first>
-__device__ __forceinline__ packed_t packed_compute(const packed_t& x,
-                                                   const packed_t& y) {
-  return act_first ? packed_mul(PACKED_ACT_FN(x), y)
-                   : packed_mul(x, PACKED_ACT_FN(y));
-}
-
 // Check if all pointers are 16-byte aligned for int4 vectorized access
-__host__ __device__ __forceinline__ bool is_16byte_aligned(const void* ptr) {
+__device__ __forceinline__ bool is_16byte_aligned(const void* ptr) {
  return (reinterpret_cast<uintptr_t>(ptr) & 15) == 0;
 }

-// Check if all pointers are 16-byte aligned for longlong4_32a vectorized access
-__host__ __device__ __forceinline__ bool is_32byte_aligned(const void* ptr) {
-  return (reinterpret_cast<uintptr_t>(ptr) & 31) == 0;
-}
-
 // Activation and gating kernel template.
-template <typename scalar_t, typename packed_t,
-          scalar_t (*ACT_FN)(const scalar_t&),
-          packed_t (*PACKED_ACT_FN)(const packed_t&), bool act_first,
-          bool use_vec, bool use_256b = false>
+template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&),
+          bool act_first>
 __global__ void act_and_mul_kernel(
    scalar_t* __restrict__ out,          // [..., d]
    const scalar_t* __restrict__ input,  // [..., 2, d]
    const int d) {
-  const scalar_t* x_ptr = input + blockIdx.x * 2 * d;
+  constexpr int VEC_SIZE = 16 / sizeof(scalar_t);
+  const int64_t token_idx = blockIdx.x;
+  const scalar_t* x_ptr = input + token_idx * 2 * d;
  const scalar_t* y_ptr = x_ptr + d;
-  scalar_t* out_ptr = out + blockIdx.x * d;
+  scalar_t* out_ptr = out + token_idx * d;

-  if constexpr (use_vec) {
-    // Fast path: 128-bit/256-bit vectorized loop
-    using vec_t = typename VecTraits<use_256b>::vec_t;
-    constexpr int ARCH_MAX_VEC_SIZE = VecTraits<use_256b>::ARCH_MAX_VEC_SIZE;
-    constexpr int VEC_SIZE = ARCH_MAX_VEC_SIZE / sizeof(packed_t);
+  // Check alignment for 128-bit vectorized access.
+  // All three pointers must be 16-byte aligned for safe int4 operations.
+  const bool aligned = is_16byte_aligned(x_ptr) && is_16byte_aligned(y_ptr) &&
+                       is_16byte_aligned(out_ptr);

-    const vec_t* x_vec = reinterpret_cast<const vec_t*>(x_ptr);
-    const vec_t* y_vec = reinterpret_cast<const vec_t*>(y_ptr);
-    vec_t* out_vec = reinterpret_cast<vec_t*>(out_ptr);
-    const int num_vecs = d / 2 / VEC_SIZE;
+  if (aligned && d >= VEC_SIZE) {
+    // Fast path: 128-bit vectorized loop
+    const int4* x_vec = reinterpret_cast<const int4*>(x_ptr);
+    const int4* y_vec = reinterpret_cast<const int4*>(y_ptr);
+    int4* out_vec = reinterpret_cast<int4*>(out_ptr);
+    const int num_vecs = d / VEC_SIZE;
+    const int vec_end = num_vecs * VEC_SIZE;

    for (int i = threadIdx.x; i < num_vecs; i += blockDim.x) {
-      vec_t x, y;
-      if constexpr (use_256b) {
-        ld256(x, &x_vec[i]);
-        ld256(y, &y_vec[i]);
-      } else {
-        x = VLLM_LDG(&x_vec[i]);
-        y = VLLM_LDG(&y_vec[i]);
-      }
-      auto* xp = reinterpret_cast<packed_t*>(&x);
-      auto* yp = reinterpret_cast<packed_t*>(&y);
+      int4 x = VLLM_LDG(&x_vec[i]), y = VLLM_LDG(&y_vec[i]), r;
+      auto* xp = reinterpret_cast<scalar_t*>(&x);
+      auto* yp = reinterpret_cast<scalar_t*>(&y);
+      auto* rp = reinterpret_cast<scalar_t*>(&r);
 #pragma unroll
      for (int j = 0; j < VEC_SIZE; j++) {
-        xp[j] =
-            packed_compute<packed_t, PACKED_ACT_FN, act_first>(xp[j], yp[j]);
-      }
-      if constexpr (use_256b) {
-        st256(x, &out_vec[i]);
-      } else {
-        out_vec[i] = x;
+        rp[j] = compute<scalar_t, ACT_FN, act_first>(xp[j], yp[j]);
      }
+      out_vec[i] = r;
+    }
+    // Scalar cleanup for remaining elements
+    for (int i = vec_end + threadIdx.x; i < d; i += blockDim.x) {
+      out_ptr[i] = compute<scalar_t, ACT_FN, act_first>(VLLM_LDG(&x_ptr[i]),
+                                                        VLLM_LDG(&y_ptr[i]));
    }
  } else {
    // Scalar fallback for unaligned data or small d
@@ -201,15 +79,6 @@ __device__ __forceinline__ T silu_kernel(const T& x) {
  return (T)(((float)x) / (1.0f + expf((float)-x)));
 }

-template <typename packed_t>
-__device__ __forceinline__ packed_t packed_silu_kernel(const packed_t& val) {
-  // x * sigmoid(x)
-  float2 fval = cast_to_float2(val);
-  fval.x = fval.x / (1.0f + expf(-fval.x));
-  fval.y = fval.y / (1.0f + expf(-fval.y));
-  return cast_to_packed<packed_t>(fval);
-}
-
 template <typename T>
 __device__ __forceinline__ T gelu_kernel(const T& x) {
  // Equivalent to PyTorch GELU with 'none' approximation.
@@ -220,18 +89,6 @@ __device__ __forceinline__ T gelu_kernel(const T& x) {
  return (T)(f * 0.5f * (1.0f + ::erf(f * ALPHA)));
 }

-template <typename packed_t>
-__device__ __forceinline__ packed_t packed_gelu_kernel(const packed_t& val) {
-  // Equivalent to PyTorch GELU with 'none' approximation.
-  // Refer to:
-  // https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L36-L38
-  constexpr float ALPHA = M_SQRT1_2;
-  float2 fval = cast_to_float2(val);
-  fval.x = fval.x * 0.5f * (1.0f + ::erf(fval.x * ALPHA));
-  fval.y = fval.y * 0.5f * (1.0f + ::erf(fval.y * ALPHA));
-  return cast_to_packed<packed_t>(fval);
-}
-
 template <typename T>
 __device__ __forceinline__ T gelu_tanh_kernel(const T& x) {
  // Equivalent to PyTorch GELU with 'tanh' approximation.
@@ -245,83 +102,32 @@ __device__ __forceinline__ T gelu_tanh_kernel(const T& x) {
  return (T)(0.5f * f * (1.0f + ::tanhf(inner)));
 }

-template <typename packed_t>
-__device__ __forceinline__ packed_t
-packed_gelu_tanh_kernel(const packed_t& val) {
-  // Equivalent to PyTorch GELU with 'tanh' approximation.
-  // Refer to:
-  // https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L25-L30
-  float2 fval = cast_to_float2(val);
-  constexpr float BETA = M_SQRT2 * M_2_SQRTPI * 0.5f;
-  constexpr float KAPPA = 0.044715;
-
-  float x_cube = fval.x * fval.x * fval.x;
-  float inner = BETA * (fval.x + KAPPA * x_cube);
-  fval.x = 0.5f * fval.x * (1.0f + ::tanhf(inner));
-
-  x_cube = fval.y * fval.y * fval.y;
-  inner = BETA * (fval.y + KAPPA * x_cube);
-  fval.y = 0.5f * fval.y * (1.0f + ::tanhf(inner));
-  return cast_to_packed<packed_t>(fval);
-}
-
 }  // namespace vllm

 // Launch activation and gating kernel.
 // Use ACT_FIRST (bool) indicating whether to apply the activation function
 // first.
-#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, PACKED_KERNEL, ACT_FIRST)     \
-  auto dtype = input.scalar_type();                                         \
-  int d = input.size(-1) / 2;                                               \
-  int64_t num_tokens = input.numel() / input.size(-1);                      \
-  if (num_tokens == 0) {                                                    \
-    return;                                                                 \
-  }                                                                         \
-  dim3 grid(num_tokens);                                                    \
-  int cc_major = at::cuda::getCurrentDeviceProperties()->major;             \
-  int support_vec = (cc_major >= 10 && num_tokens > 128) ? 32 : 16;         \
-  int vec_size = support_vec / at::elementSize(dtype);                      \
-  const bool use_vec = (d % vec_size == 0);                                 \
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));         \
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();             \
-  if (use_vec) {                                                            \
-    dim3 block(std::min(d / vec_size, 1024));                               \
-    if (cc_major >= 10 && num_tokens > 128) {                               \
-      VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] {       \
-        vllm::act_and_mul_kernel<                                           \
-            scalar_t, typename vllm::PackedTraits<scalar_t>::packed_t,      \
-            KERNEL<scalar_t>,                                               \
-            PACKED_KERNEL<typename vllm::PackedTraits<scalar_t>::packed_t>, \
-            ACT_FIRST, true, true><<<grid, block, 0, stream>>>(             \
-            out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d);       \
-      });                                                                   \
-    } else {                                                                \
-      VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] {       \
-        vllm::act_and_mul_kernel<                                           \
-            scalar_t, typename vllm::PackedTraits<scalar_t>::packed_t,      \
-            KERNEL<scalar_t>,                                               \
-            PACKED_KERNEL<typename vllm::PackedTraits<scalar_t>::packed_t>, \
-            ACT_FIRST, true, false><<<grid, block, 0, stream>>>(            \
-            out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d);       \
-      });                                                                   \
-    }                                                                       \
-  } else {                                                                  \
-    dim3 block(std::min(d, 1024));                                          \
-    VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] {         \
-      vllm::act_and_mul_kernel<                                             \
-          scalar_t, typename vllm::PackedTraits<scalar_t>::packed_t,        \
-          KERNEL<scalar_t>,                                                 \
-          PACKED_KERNEL<typename vllm::PackedTraits<scalar_t>::packed_t>,   \
-          ACT_FIRST, false><<<grid, block, 0, stream>>>(                    \
-          out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d);         \
-    });                                                                     \
-  }
+#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, ACT_FIRST)                 \
+  int d = input.size(-1) / 2;                                            \
+  int64_t num_tokens = input.numel() / input.size(-1);                   \
+  dim3 grid(num_tokens);                                                 \
+  dim3 block(std::min(d, 1024));                                         \
+  if (num_tokens == 0) {                                                 \
+    return;                                                              \
+  }                                                                      \
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));      \
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();          \
+  VLLM_DISPATCH_FLOATING_TYPES(                                          \
+      input.scalar_type(), "act_and_mul_kernel", [&] {                   \
+        vllm::act_and_mul_kernel<scalar_t, KERNEL<scalar_t>, ACT_FIRST>  \
+            <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(),       \
+                                         input.data_ptr<scalar_t>(), d); \
+      });

 void silu_and_mul(torch::Tensor& out,    // [..., d]
                  torch::Tensor& input)  // [..., 2 * d]
 {
-  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
-                                true);
+  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, true);
 }

 void mul_and_silu(torch::Tensor& out,    // [..., d]
@@ -329,22 +135,19 @@ void mul_and_silu(torch::Tensor& out,    // [..., d]
 {
  // The difference between mul_and_silu and silu_and_mul is that mul_and_silu
  // applies the silu to the latter half of the input.
-  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
-                                false);
+  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, false);
 }

 void gelu_and_mul(torch::Tensor& out,    // [..., d]
                  torch::Tensor& input)  // [..., 2 * d]
 {
-  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_kernel, vllm::packed_gelu_kernel,
-                                true);
+  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_kernel, true);
 }

 void gelu_tanh_and_mul(torch::Tensor& out,    // [..., d]
                       torch::Tensor& input)  // [..., 2 * d]
 {
-  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_tanh_kernel,
-                                vllm::packed_gelu_tanh_kernel, true);
+  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_tanh_kernel, true);
 }

 namespace vllm {
@@ -355,57 +158,42 @@ __device__ __forceinline__ T fatrelu_kernel(const T& x, const float threshold) {
  return (T)(f > threshold ? f : 0.0f);
 }

-template <typename packed_t>
-__device__ __forceinline__ packed_t
-packed_fatrelu_kernel(const packed_t& val, const float threshold) {
-  float2 fval = cast_to_float2(val);
-  fval.x = fval.x > threshold ? fval.x : 0.0f;
-  fval.y = fval.y > threshold ? fval.y : 0.0f;
-  return cast_to_packed<packed_t>(fval);
-}
-
-template <typename scalar_t, typename packed_t,
-          scalar_t (*ACT_FN)(const scalar_t&, const float),
-          packed_t (*PACKED_ACT_FN)(const packed_t&, const float), bool use_vec,
-          bool use_256b = false>
+template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&, const float)>
 __global__ void act_and_mul_kernel_with_param(
    scalar_t* __restrict__ out, const scalar_t* __restrict__ input, const int d,
    const float param) {
-  const scalar_t* x_ptr = input + blockIdx.x * 2 * d;
+  constexpr int VEC_SIZE = 16 / sizeof(scalar_t);
+  const int64_t token_idx = blockIdx.x;
+  const scalar_t* x_ptr = input + token_idx * 2 * d;
  const scalar_t* y_ptr = x_ptr + d;
-  scalar_t* out_ptr = out + blockIdx.x * d;
+  scalar_t* out_ptr = out + token_idx * d;

-  if constexpr (use_vec) {
-    // Fast path: 128-bit/256-bit vectorized loop
-    using vec_t = typename VecTraits<use_256b>::vec_t;
-    constexpr int ARCH_MAX_VEC_SIZE = VecTraits<use_256b>::ARCH_MAX_VEC_SIZE;
-    constexpr int VEC_SIZE = ARCH_MAX_VEC_SIZE / sizeof(packed_t);
+  // Check alignment for 128-bit vectorized access
+  const bool aligned = is_16byte_aligned(x_ptr) && is_16byte_aligned(y_ptr) &&
+                       is_16byte_aligned(out_ptr);

-    const vec_t* x_vec = reinterpret_cast<const vec_t*>(x_ptr);
-    const vec_t* y_vec = reinterpret_cast<const vec_t*>(y_ptr);
-    vec_t* out_vec = reinterpret_cast<vec_t*>(out_ptr);
-    const int num_vecs = d / 2 / VEC_SIZE;
+  if (aligned && d >= VEC_SIZE) {
+    // Fast path: 128-bit vectorized loop
+    const int4* x_vec = reinterpret_cast<const int4*>(x_ptr);
+    const int4* y_vec = reinterpret_cast<const int4*>(y_ptr);
+    int4* out_vec = reinterpret_cast<int4*>(out_ptr);
+    const int num_vecs = d / VEC_SIZE;
+    const int vec_end = num_vecs * VEC_SIZE;

    for (int i = threadIdx.x; i < num_vecs; i += blockDim.x) {
-      vec_t x, y;
-      if constexpr (use_256b) {
-        ld256(x, &x_vec[i]);
-        ld256(y, &y_vec[i]);
-      } else {
-        x = VLLM_LDG(&x_vec[i]);
-        y = VLLM_LDG(&y_vec[i]);
-      }
-      auto* xp = reinterpret_cast<packed_t*>(&x);
-      auto* yp = reinterpret_cast<packed_t*>(&y);
+      int4 x = VLLM_LDG(&x_vec[i]), y = VLLM_LDG(&y_vec[i]), r;
+      auto* xp = reinterpret_cast<scalar_t*>(&x);
+      auto* yp = reinterpret_cast<scalar_t*>(&y);
+      auto* rp = reinterpret_cast<scalar_t*>(&r);
 #pragma unroll
      for (int j = 0; j < VEC_SIZE; j++) {
-        xp[j] = packed_mul(PACKED_ACT_FN(xp[j], param), yp[j]);
-      }
-      if constexpr (use_256b) {
-        st256(x, &out_vec[i]);
-      } else {
-        out_vec[i] = x;
+        rp[j] = ACT_FN(xp[j], param) * yp[j];
      }
+      out_vec[i] = r;
+    }
+    // Scalar cleanup for remaining elements
+    for (int i = vec_end + threadIdx.x; i < d; i += blockDim.x) {
+      out_ptr[i] = ACT_FN(VLLM_LDG(&x_ptr[i]), param) * VLLM_LDG(&y_ptr[i]);
    }
  } else {
    // Scalar fallback for unaligned data or small d
@@ -488,58 +276,20 @@ __global__ void swigluoai_and_mul_kernel(

 }  // namespace vllm

-#define LAUNCH_ACTIVATION_GATE_KERNEL_WITH_PARAM(KERNEL, PACKED_KERNEL, PARAM) \
-  auto dtype = input.scalar_type();                                            \
-  int d = input.size(-1) / 2;                                                  \
-  int64_t num_tokens = input.numel() / input.size(-1);                         \
-  if (num_tokens == 0) {                                                       \
-    return;                                                                    \
-  }                                                                            \
-  dim3 grid(num_tokens);                                                       \
-  int cc_major = at::cuda::getCurrentDeviceProperties()->major;                \
-  int support_vec = (cc_major >= 10 && num_tokens > 128) ? 32 : 16;            \
-  int vec_size = support_vec / at::elementSize(dtype);                         \
-  const bool use_vec = (d % vec_size == 0);                                    \
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));            \
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();                \
-  if (use_vec) {                                                               \
-    dim3 block(std::min(d / vec_size, 1024));                                  \
-    if (cc_major >= 10 && num_tokens > 128) {                                  \
-      VLLM_DISPATCH_FLOATING_TYPES(                                            \
-          dtype, "act_and_mul_kernel_with_param", [&] {                        \
-            vllm::act_and_mul_kernel_with_param<                               \
-                scalar_t, typename vllm::PackedTraits<scalar_t>::packed_t,     \
-                KERNEL<scalar_t>,                                              \
-                PACKED_KERNEL<                                                 \
-                    typename vllm::PackedTraits<scalar_t>::packed_t>,          \
-                true, true><<<grid, block, 0, stream>>>(                       \
-                out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d,       \
-                PARAM);                                                        \
-          });                                                                  \
-    } else {                                                                   \
-      VLLM_DISPATCH_FLOATING_TYPES(                                            \
-          dtype, "act_and_mul_kernel_with_param", [&] {                        \
-            vllm::act_and_mul_kernel_with_param<                               \
-                scalar_t, typename vllm::PackedTraits<scalar_t>::packed_t,     \
-                KERNEL<scalar_t>,                                              \
-                PACKED_KERNEL<                                                 \
-                    typename vllm::PackedTraits<scalar_t>::packed_t>,          \
-                true, false><<<grid, block, 0, stream>>>(                      \
-                out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d,       \
-                PARAM);                                                        \
-          });                                                                  \
-    }                                                                          \
-  } else {                                                                     \
-    dim3 block(std::min(d, 1024));                                             \
-    VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel_with_param", [&] { \
-      vllm::act_and_mul_kernel_with_param<                                     \
-          scalar_t, typename vllm::PackedTraits<scalar_t>::packed_t,           \
-          KERNEL<scalar_t>,                                                    \
-          PACKED_KERNEL<typename vllm::PackedTraits<scalar_t>::packed_t>,      \
-          false><<<grid, block, 0, stream>>>(                                  \
-          out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, PARAM);     \
-    });                                                                        \
-  }
+#define LAUNCH_ACTIVATION_GATE_KERNEL_WITH_PARAM(KERNEL, PARAM)         \
+  int d = input.size(-1) / 2;                                           \
+  int64_t num_tokens = input.numel() / input.size(-1);                  \
+  dim3 grid(num_tokens);                                                \
+  dim3 block(std::min(d, 1024));                                        \
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));     \
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();         \
+  VLLM_DISPATCH_FLOATING_TYPES(                                         \
+      input.scalar_type(), "act_and_mul_kernel_with_param", [&] {       \
+        vllm::act_and_mul_kernel_with_param<scalar_t, KERNEL<scalar_t>> \
+            <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(),      \
+                                         input.data_ptr<scalar_t>(), d, \
+                                         PARAM);                        \
+      });

 #define LAUNCH_SIGLUOAI_AND_MUL(KERNEL, ALPHA, LIMIT)                          \
  int d = input.size(-1) / 2;                                                  \
@@ -559,8 +309,7 @@ __global__ void swigluoai_and_mul_kernel(
 void fatrelu_and_mul(torch::Tensor& out,    // [..., d],
                     torch::Tensor& input,  // [..., 2 * d]
                     double threshold) {
-  LAUNCH_ACTIVATION_GATE_KERNEL_WITH_PARAM(
-      vllm::fatrelu_kernel, vllm::packed_fatrelu_kernel, threshold);
+  LAUNCH_ACTIVATION_GATE_KERNEL_WITH_PARAM(vllm::fatrelu_kernel, threshold);
 }
 void swigluoai_and_mul(torch::Tensor& out,    // [..., d]
                       torch::Tensor& input,  // [..., 2 * d]
@@ -570,41 +319,39 @@ void swigluoai_and_mul(torch::Tensor& out,    // [..., d]
 namespace vllm {

 // Element-wise activation kernel template.
-template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&), bool use_vec,
-          bool use_256b = false>
+template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&)>
 __global__ void activation_kernel(
    scalar_t* __restrict__ out,          // [..., d]
    const scalar_t* __restrict__ input,  // [..., d]
    const int d) {
-  const scalar_t* in_ptr = input + blockIdx.x * d;
-  scalar_t* out_ptr = out + blockIdx.x * d;
+  constexpr int VEC_SIZE = 16 / sizeof(scalar_t);
+  const int64_t token_idx = blockIdx.x;
+  const scalar_t* in_ptr = input + token_idx * d;
+  scalar_t* out_ptr = out + token_idx * d;

-  if constexpr (use_vec) {
-    // Fast path: 128-bit/256-bit vectorized loop
-    using vec_t = typename VecTraits<use_256b>::vec_t;
-    constexpr int ARCH_MAX_VEC_SIZE = VecTraits<use_256b>::ARCH_MAX_VEC_SIZE;
-    constexpr int VEC_SIZE = ARCH_MAX_VEC_SIZE / sizeof(scalar_t);
-    const vec_t* in_vec = reinterpret_cast<const vec_t*>(in_ptr);
-    vec_t* out_vec = reinterpret_cast<vec_t*>(out_ptr);
+  // Check alignment for 128-bit vectorized access
+  const bool aligned = is_16byte_aligned(in_ptr) && is_16byte_aligned(out_ptr);
+
+  if (aligned && d >= VEC_SIZE) {
+    // Fast path: 128-bit vectorized loop
+    const int4* in_vec = reinterpret_cast<const int4*>(in_ptr);
+    int4* out_vec = reinterpret_cast<int4*>(out_ptr);
    const int num_vecs = d / VEC_SIZE;
+    const int vec_end = num_vecs * VEC_SIZE;

    for (int i = threadIdx.x; i < num_vecs; i += blockDim.x) {
-      vec_t v;
-      if constexpr (use_256b) {
-        ld256(v, &in_vec[i]);
-      } else {
-        v = VLLM_LDG(&in_vec[i]);
-      }
+      int4 v = VLLM_LDG(&in_vec[i]), r;
      auto* vp = reinterpret_cast<scalar_t*>(&v);
+      auto* rp = reinterpret_cast<scalar_t*>(&r);
 #pragma unroll
      for (int j = 0; j < VEC_SIZE; j++) {
-        vp[j] = ACT_FN(vp[j]);
-      }
-      if constexpr (use_256b) {
-        st256(v, &out_vec[i]);
-      } else {
-        out_vec[i] = v;
+        rp[j] = ACT_FN(vp[j]);
      }
+      out_vec[i] = r;
+    }
+    // Scalar cleanup for remaining elements
+    for (int i = vec_end + threadIdx.x; i < d; i += blockDim.x) {
+      out_ptr[i] = ACT_FN(VLLM_LDG(&in_ptr[i]));
    }
  } else {
    // Scalar fallback for unaligned data or small d
@@ -618,43 +365,18 @@ __global__ void activation_kernel(
 }  // namespace vllm

 // Launch element-wise activation kernel.
-#define LAUNCH_ACTIVATION_KERNEL(KERNEL)                                 \
-  auto dtype = input.scalar_type();                                      \
-  int d = input.size(-1);                                                \
-  int64_t num_tokens = input.numel() / input.size(-1);                   \
-  if (num_tokens == 0) {                                                 \
-    return;                                                              \
-  }                                                                      \
-  dim3 grid(num_tokens);                                                 \
-  int cc_major = at::cuda::getCurrentDeviceProperties()->major;          \
-  int support_vec = (cc_major >= 10 && num_tokens > 128) ? 32 : 16;      \
-  int vec_size = support_vec / at::elementSize(dtype);                   \
-  const bool use_vec = (d % vec_size == 0);                              \
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));      \
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();          \
-  if (use_vec) {                                                         \
-    dim3 block(std::min(d / vec_size, 1024));                            \
-    if (cc_major >= 10 && num_tokens > 128) {                            \
-      VLLM_DISPATCH_FLOATING_TYPES(dtype, "activation_kernel", [&] {     \
-        vllm::activation_kernel<scalar_t, KERNEL<scalar_t>, true, true>  \
-            <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(),       \
-                                         input.data_ptr<scalar_t>(), d); \
-      });                                                                \
-    } else {                                                             \
-      VLLM_DISPATCH_FLOATING_TYPES(dtype, "activation_kernel", [&] {     \
-        vllm::activation_kernel<scalar_t, KERNEL<scalar_t>, true, false> \
-            <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(),       \
-                                         input.data_ptr<scalar_t>(), d); \
-      });                                                                \
-    }                                                                    \
-  } else {                                                               \
-    dim3 block(std::min(d, 1024));                                       \
-    VLLM_DISPATCH_FLOATING_TYPES(dtype, "activation_kernel", [&] {       \
-      vllm::activation_kernel<scalar_t, KERNEL<scalar_t>, false>         \
-          <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(),         \
-                                       input.data_ptr<scalar_t>(), d);   \
-    });                                                                  \
-  }
+#define LAUNCH_ACTIVATION_KERNEL(KERNEL)                                       \
+  int d = input.size(-1);                                                      \
+  int64_t num_tokens = input.numel() / d;                                      \
+  dim3 grid(num_tokens);                                                       \
+  dim3 block(std::min(d, 1024));                                               \
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));            \
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();                \
+  VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "activation_kernel", [&] { \
+    vllm::activation_kernel<scalar_t, KERNEL<scalar_t>>                        \
+        <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(),                 \
+                                     input.data_ptr<scalar_t>(), d);           \
+  });

 namespace vllm {

@@ -1234,13 +1234,8 @@ void cp_gather_and_upconvert_fp8_kv_cache(
              "src_cache and seq_lens must be on the same device");
  TORCH_CHECK(src_cache.device() == workspace_starts.device(),
              "src_cache and workspace_starts must be on the same device");
-  auto dtype = src_cache.scalar_type();
-  TORCH_CHECK(
-      dtype == at::ScalarType::Byte ||               // uint8
-          dtype == at::ScalarType::Float8_e4m3fn ||  // fp8 e4m3
-          dtype == at::ScalarType::Float8_e5m2,      // fp8 e5m2
-      "src_cache must be uint8, float8_e4m3fn, or float8_e5m2, but got ",
-      src_cache.dtype());
+
+  TORCH_CHECK(src_cache.dtype() == torch::kUInt8, "src_cache must be uint8");
  TORCH_CHECK(dst.dtype() == torch::kBFloat16, "dst must be bfloat16");
  TORCH_CHECK(head_dim == 576, "head_dim must be 576 for MLA");

@@ -1249,21 +1244,14 @@ void cp_gather_and_upconvert_fp8_kv_cache(
  int64_t cache_entry_stride = src_cache.stride(1);
  int64_t dst_entry_stride = dst.stride(0);

-  const uint8_t* src_ptr = nullptr;
-  if (dtype == at::ScalarType::Byte) {
-    src_ptr = src_cache.data_ptr<uint8_t>();
-  } else {
-    // float8_e4m3fn or float8_e5m2
-    src_ptr = reinterpret_cast<const uint8_t*>(src_cache.data_ptr());
-  }
-
  // Decide on the number of splits based on the batch size
  int num_splits = batch_size > 128 ? 2 : batch_size > 64 ? 4 : 16;
  dim3 grid(batch_size, num_splits);
  dim3 block(576);

  vllm::cp_gather_and_upconvert_fp8_kv_cache<<<grid, block, 0, stream>>>(
-      src_ptr, reinterpret_cast<__nv_bfloat16*>(dst.data_ptr()),
+      src_cache.data_ptr<uint8_t>(),
+      reinterpret_cast<__nv_bfloat16*>(dst.data_ptr()),
      block_table.data_ptr<int32_t>(), seq_lens.data_ptr<int32_t>(),
      workspace_starts.data_ptr<int32_t>(), block_size, head_dim,
      block_table_stride, cache_block_stride, cache_entry_stride,
@@ -821,7 +821,7 @@ struct VecTypeTrait<c10::BFloat16> {
  using vec_t = vec_op::BF16Vec16;
 };

-#if !defined(__powerpc__)
+#if !defined(__powerpc__) && !defined(__s390x__)
 template <>
 struct VecTypeTrait<c10::Half> {
  using vec_t = vec_op::FP16Vec16;
@@ -147,7 +147,7 @@ void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input,
                    const int32_t token_num, const int32_t expert_num,
                    const int32_t topk_num, const int32_t input_size_13,
                    const int32_t output_size_13, const int32_t input_size_2,
-                    const int32_t output_size_2, const bool skip_weighted) {
+                    const int32_t output_size_2) {
  using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
  constexpr int32_t gemm_n_tile_size = gemm_t::NSize;
  constexpr int32_t gemm_m_tile_size = gemm_t::MaxMSize;
@@ -582,11 +582,6 @@ void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input,
        scalar_t* __restrict__ curr_output_buffer =
            output + token_id * output_size_2;

-        if (skip_weighted) {
-          // Only for topk_num == 1
-          *curr_weight = 1.0f;
-        }
-
        if (topk_num > 1) {
          {
            int32_t w2_output_idx = curr_expand_token_id_index_buffer[0];
@@ -704,7 +699,7 @@ void cpu_fused_moe(
    const std::optional<torch::Tensor>& w2_bias,  // [expert_num, output_size_2]
    const torch::Tensor& topk_weights,            // [token_num, k], float32
    const torch::Tensor& topk_id,                 // [token_num, k], int32
-    const bool skip_weighted, const std::string& act, const std::string& isa) {
+    const std::string& act, const std::string& isa) {
  const int32_t token_num = input.size(0);
  const int32_t input_size_13 = input.size(1);
  const int64_t input_stride = input.stride(0);
@@ -716,8 +711,6 @@ void cpu_fused_moe(
  const int32_t topk_num = topk_id.size(1);
  const FusedMOEAct act_type = get_act_type(act);
  cpu_utils::ISA isa_type = cpu_utils::get_isa(isa);
-  TORCH_CHECK(!skip_weighted || topk_num == 1,
-              "skip_weighted is only supported for topk=1 on CPU");

  VLLM_DISPATCH_FLOATING_TYPES(w13.scalar_type(), "cpu_fused_moe", [&]() {
    CPU_ISA_DISPATCH_IMPL(isa_type, [&]() {
@@ -728,7 +721,7 @@ void cpu_fused_moe(
          w2_bias.has_value() ? w2_bias->data_ptr<scalar_t>() : nullptr,
          topk_weights.data_ptr<float>(), topk_id.data_ptr<int32_t>(), act_type,
          token_num, expert_num, topk_num, input_size_13, output_size_13,
-          input_size_2, output_size_2, skip_weighted);
+          input_size_2, output_size_2);
    });
  });
 }
@@ -16,12 +16,10 @@ namespace vec_op {
 #define vec_sr(a, b) ((a) >> (b))  // Vector Shift Right Algebraic
 #define vec_sl(a, b) ((a) << (b))  // Vector Shift Left

-// NOTE: FP16 (Half) is supported on s390x via custom bit-manipulation
-// conversion. PyTorch itself lacks native s390x FP16 support.
-#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...)            \
-  AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__)    \
-  AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) \
-  AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__)
+// FIXME: FP16 is not fully supported in Torch-CPU
+#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...)         \
+  AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
+  AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__)

 #define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
  AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
@@ -88,39 +86,6 @@ struct BF16Vec8 : public Vec<BF16Vec8> {
  }
 };

-struct FP16Vec8 : public Vec<FP16Vec8> {
-  constexpr static int VEC_ELEM_NUM = 8;
-
-  __vector signed short reg;
-
-  explicit FP16Vec8(const void* ptr) : reg(*(__vector signed short*)ptr) {}
-  explicit FP16Vec8(const FP32Vec8&);
-
-  void save(void* ptr) const {
-    *reinterpret_cast<__vector signed short*>(ptr) = reg;
-  }
-};
-
-struct FP16Vec16 : public Vec<FP16Vec16> {
-  constexpr static int VEC_ELEM_NUM = 16;
-
-  ss16x8x2_t reg;
-
-  explicit FP16Vec16(const void* ptr) {
-    // Load 256 bits (16 FP16 values) in two parts
-    reg.val[0] = (__vector signed short)vec_xl(0, (signed short*)ptr);
-    reg.val[1] = (__vector signed short)vec_xl(16, (signed short*)ptr);
-  }
-
-  explicit FP16Vec16(const FP32Vec16&);
-
-  void save(void* ptr) const {
-    // Save 256 bits in two parts
-    vec_xst(reg.val[0], 0, (signed short*)ptr);
-    vec_xst(reg.val[1], 16, (signed short*)ptr);
-  }
-};
-
 struct BF16Vec16 : public Vec<BF16Vec16> {
  constexpr static int VEC_ELEM_NUM = 16;

@@ -143,92 +108,6 @@ struct BF16Vec16 : public Vec<BF16Vec16> {

 const static __vector signed short zero = vec_splats((signed short)0);

-FORCE_INLINE __vector float fp16_to_fp32_bits(__vector unsigned int x) {
-  const __vector unsigned int mask_sign = {0x8000, 0x8000, 0x8000, 0x8000};
-  const __vector unsigned int mask_exp = {0x7C00, 0x7C00, 0x7C00, 0x7C00};
-  const __vector unsigned int mask_mant = {0x03FF, 0x03FF, 0x03FF, 0x03FF};
-  const __vector unsigned int bias_adj = {112, 112, 112, 112};
-  const __vector unsigned int exp_max_fp16 = {0x1F, 0x1F, 0x1F,
-                                              0x1F};  // FP16 NaN/Inf exponent
-  const __vector unsigned int exp_max_fp32 = {0xFF, 0xFF, 0xFF,
-                                              0xFF};  // FP32 NaN/Inf exponent
-
-  __vector unsigned int s = (x & mask_sign) << 16;
-  __vector unsigned int e = (x & mask_exp) >> 10;
-  __vector unsigned int m = (x & mask_mant) << 13;
-
-  // Check for NaN/Inf: exponent = 0x1F in FP16
-  __vector __bool int is_nan_inf = vec_cmpeq(e, exp_max_fp16);
-
-  // Normal: adjust bias; NaN/Inf: set to 0xFF
-  __vector unsigned int e_normal = e + bias_adj;
-  e = vec_sel(e_normal, exp_max_fp32, is_nan_inf);
-
-  return (__vector float)(s | (e << 23) | m);
-}
-
-FORCE_INLINE __vector unsigned int fp32_to_fp16_bits(__vector float f_in) {
-  __vector unsigned int in = (__vector unsigned int)f_in;
-
-  const __vector unsigned int mask_sign_32 = {0x80000000, 0x80000000,
-                                              0x80000000, 0x80000000};
-  const __vector unsigned int mask_exp_32 = {0x7F800000, 0x7F800000, 0x7F800000,
-                                             0x7F800000};
-  const __vector unsigned int mask_mant_32 = {0x007FFFFF, 0x007FFFFF,
-                                              0x007FFFFF, 0x007FFFFF};
-
-  // Use SIGNED integers for exponent math to handle underflow check
-  const __vector signed int bias_adj = {112, 112, 112, 112};
-  const __vector signed int zero = {0, 0, 0, 0};
-  const __vector signed int max_exp = {31, 31, 31, 31};  // Max FP16 exp
-  const __vector unsigned int exp_max_fp32 = {0xFF, 0xFF, 0xFF, 0xFF};
-  const __vector unsigned int exp_max_fp16 = {0x1F, 0x1F, 0x1F, 0x1F};
-
-  __vector unsigned int s = (in & mask_sign_32) >> 16;
-  __vector unsigned int e_u = (in & mask_exp_32) >> 23;
-
-  // Check for NaN/Inf: exponent = 0xFF in FP32
-  __vector __bool int is_nan_inf = vec_cmpeq(e_u, exp_max_fp32);
-
-  __vector signed int e_s = (__vector signed int)e_u;
-  e_s = vec_sub(e_s, bias_adj);
-  e_s = vec_max(e_s, zero);
-  e_s = vec_min(e_s, max_exp);
-  __vector unsigned int e_normal = (__vector unsigned int)e_s;
-
-  __vector unsigned int e_final = vec_sel(e_normal, exp_max_fp16, is_nan_inf);
-
-  const __vector unsigned int one_v = {1, 1, 1, 1};
-  const __vector unsigned int mask_sticky = {0xFFF, 0xFFF, 0xFFF, 0xFFF};
-
-  __vector unsigned int round_bit = (in >> 12) & one_v;
-  __vector unsigned int sticky = in & mask_sticky;
-  __vector unsigned int m = (in & mask_mant_32) >> 13;
-  __vector unsigned int lsb = m & one_v;  // LSB of mantissa for tie-breaking
-
-  // Round up if: round_bit && (sticky || lsb)
-  __vector __bool int sticky_nonzero =
-      vec_cmpgt(sticky, (__vector unsigned int){0, 0, 0, 0});
-  __vector __bool int lsb_set = vec_cmpeq(lsb, one_v);
-  __vector __bool int round_up =
-      vec_and(vec_cmpeq(round_bit, one_v), vec_or(sticky_nonzero, lsb_set));
-
-  m = vec_sel(m, m + one_v, round_up);
-
-  const __vector unsigned int mant_mask = {0x3FF, 0x3FF, 0x3FF, 0x3FF};
-  const __vector unsigned int max_normal_exp = {0x1E, 0x1E, 0x1E, 0x1E};
-  __vector __bool int mant_overflows = vec_cmpgt(m, mant_mask);
-  __vector __bool int would_overflow_to_inf =
-      vec_and(mant_overflows, vec_cmpeq(e_final, max_normal_exp));
-  __vector unsigned int e_inc = vec_min(e_final + one_v, exp_max_fp16);
-  e_final = vec_sel(e_final, e_inc, mant_overflows);
-  m = vec_and(m, mant_mask);
-  e_final = vec_sel(e_final, max_normal_exp, would_overflow_to_inf);
-  m = vec_sel(m, mant_mask, would_overflow_to_inf);
-
-  return s | (e_final << 10) | m;
-}
-
 struct BF16Vec32 : public Vec<BF16Vec32> {
  constexpr static int VEC_ELEM_NUM = 32;

@@ -301,18 +180,6 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
    reg.val[1] = (__vector float)vec_mergel(v.reg, zero);
  }

-  explicit FP32Vec8(const FP16Vec8& v) {
-    // Cast to UNSIGNED short vector to prevent sign-extension during unpack
-    __vector unsigned short raw_u = (__vector unsigned short)v.reg;
-
-    // Unpack 8x16-bit to two 4x32-bit vectors (Zero extended)
-    __vector unsigned int raw_hi = (__vector unsigned int)vec_unpackh(raw_u);
-    __vector unsigned int raw_lo = (__vector unsigned int)vec_unpackl(raw_u);
-
-    reg.val[0] = fp16_to_fp32_bits(raw_hi);
-    reg.val[1] = fp16_to_fp32_bits(raw_lo);
-  }
-
  float reduce_sum() const {
    AliasReg ar;
    ar.reg = reg;
@@ -664,22 +531,6 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
    reg.val[3] = (__vector float)vec_mergel(v.reg.val[1], zero);
  }

-  explicit FP32Vec16(const FP16Vec16& v) {
-    __vector unsigned int raw_hi_0 =
-        (__vector unsigned int)vec_unpackh(v.reg.val[0]);
-    __vector unsigned int raw_lo_0 =
-        (__vector unsigned int)vec_unpackl(v.reg.val[0]);
-    reg.val[0] = fp16_to_fp32_bits(raw_hi_0);
-    reg.val[1] = fp16_to_fp32_bits(raw_lo_0);
-
-    __vector unsigned int raw_hi_1 =
-        (__vector unsigned int)vec_unpackh(v.reg.val[1]);
-    __vector unsigned int raw_lo_1 =
-        (__vector unsigned int)vec_unpackl(v.reg.val[1]);
-    reg.val[2] = fp16_to_fp32_bits(raw_hi_1);
-    reg.val[3] = fp16_to_fp32_bits(raw_lo_1);
-  }
-
  explicit FP32Vec16(const BF16Vec8& v) : FP32Vec16(FP32Vec8(v)) {}

  FP32Vec16 operator*(const FP32Vec16& b) const {
@@ -777,10 +628,8 @@ struct VecType<c10::BFloat16> {
  using vec_type = BF16Vec8;
 };

-template <>
-struct VecType<c10::Half> {
-  using vec_type = FP16Vec8;
-};
+// On s390x, FP16 (Half) is not natively supported, use FP32 vectors instead
+using FP16Vec16 = FP32Vec16;

 template <typename T>
 void storeFP32(float v, T* ptr) {
@@ -801,52 +650,6 @@ inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
  *ptr = *(v_ptr + 1);
 }

-template <>
-inline void storeFP32<::c10::Half>(float v, ::c10::Half* ptr) {
-  // Use bit-manipulation for IEEE FP32 to FP16 conversion since vector
-  // intrinsics for FP32 to FP16 conversion does not use IEEE rounding and can
-  // produce incorrect results for some inputs. Process each of the 4 vectors
-  // separately.
-  uint32_t in;
-  std::memcpy(&in, &v, sizeof(in));
-
-  uint32_t s = (in & 0x80000000) >> 16;  // Sign
-  uint32_t e = (in & 0x7F800000) >> 23;  // Exponent
-  uint32_t round_bit = (in >> 12) & 1;
-  uint32_t sticky = (in & 0xFFF) != 0;  // Any bits in [11..0]
-  uint32_t m = (in & 0x007FFFFF) >> 13;
-  uint32_t lsb = m & 1;  // LSB of mantissa for tie-breaking
-
-  // Check for NaN/Inf before rounding
-  bool is_nan_inf = (e == 0xFF);
-
-  if (round_bit && (sticky || lsb)) {
-    m++;
-    // Handle mantissa overflow: if m overflows 10 bits, increment exponent
-    if (m > 0x3FF) {
-      m = 0;
-      e++;
-    }
-  }
-
-  if (is_nan_inf) {
-    // NaN/Inf: preserve it
-    e = 0x1F;
-  } else {
-    // Normal: adjust bias (127 - 15), flush subnormals to zero
-    e = (e >= 112) ? (e - 112) : 0;
-    // If exponent overflows to Inf range, saturate to max normal FP16 value
-    if (e > 0x1E) {
-      e = 0x1E;   // Max normal exponent
-      m = 0x3FF;  // Max mantissa
-    }
-  }
-
-  uint16_t fp16 = (uint16_t)(s | (e << 10) | m);
-
-  *reinterpret_cast<uint16_t*>(ptr) = fp16;
-}
-
 #ifndef __VEC_CLASS_FP_NAN
  #define __VEC_CLASS_FP_NAN (1 << 6)
 #endif
@@ -1000,44 +803,6 @@ inline BF16Vec16::BF16Vec16(const FP32Vec16& v) {
  reg.val[1] = (__vector signed short)vec_perm(inp2, inp3, omask);
 }

-inline FP16Vec8::FP16Vec8(const FP32Vec8& v) {
-  // Use bit-manipulation for IEEE FP32 to FP16 conversion since vector
-  // intrinsics for FP32 to FP16 conversion does not use IEEE rounding and can
-  // produce incorrect results for some inputs. Process each of the 4 vectors
-  // separately.
-  __vector unsigned int res_hi = fp32_to_fp16_bits(v.reg.val[0]);
-  __vector unsigned int res_lo = fp32_to_fp16_bits(v.reg.val[1]);
-
-  const __vector unsigned char perm_pack = {
-      2,  3,  6,  7,  10, 11, 14, 15,  // Select lower 2 bytes from res_hi
-      18, 19, 22, 23, 26, 27, 30, 31   // Select lower 2 bytes from res_lo
-  };
-
-  reg = vec_perm((__vector signed short)res_hi, (__vector signed short)res_lo,
-                 perm_pack);
-}
-
-inline FP16Vec16::FP16Vec16(const FP32Vec16& v) {
-  // Use bit-manipulation for IEEE FP32 to FP16 conversion since vector
-  // intrinsics for FP32 to FP16 conversion does not use IEEE rounding and can
-  // produce incorrect results for some inputs. Process each of the 4 vectors
-  // separately.
-  __vector unsigned int res_0 = fp32_to_fp16_bits(v.reg.val[0]);
-  __vector unsigned int res_1 = fp32_to_fp16_bits(v.reg.val[1]);
-  __vector unsigned int res_2 = fp32_to_fp16_bits(v.reg.val[2]);
-  __vector unsigned int res_3 = fp32_to_fp16_bits(v.reg.val[3]);
-
-  const __vector unsigned char perm_pack = {
-      2,  3,  6,  7,  10, 11, 14, 15,  // Lower 2 bytes from first vector
-      18, 19, 22, 23, 26, 27, 30, 31   // Lower 2 bytes from second vector
-  };
-
-  reg.val[0] = vec_perm((__vector signed short)res_0,
-                        (__vector signed short)res_1, perm_pack);
-  reg.val[1] = vec_perm((__vector signed short)res_2,
-                        (__vector signed short)res_3, perm_pack);
-}
-
 // 1D softmax over `n` elements in `input`, writes result to `output`.
 // Uses FP32Vec8 for main body, scalar tail handling.
 // Requirement: n > 0
@@ -237,20 +237,12 @@ W8A8MatMulPrimitiveHandler::W8A8MatMulPrimitiveHandler(const Args& args)
  };
  dnnl::memory::desc original_b_md({b_k_size_, b_n_size_}, b_type_,
                                   {b_k_stride_, b_n_stride_});
-#ifdef __aarch64__
-  // dummy M size for prepacking weights
-  // Prepacking weights improves performance and avoid runtime reorders
-  constexpr dnnl_dim_t kProbeM = 128;
-#else
-  constexpr dnnl_dim_t kProbeM = DNNL_RUNTIME_DIM_VAL;
-#endif
-
  prepack_weight(args.b_ptr, original_b_md,
                 create_primitive_desc(
-                     MSizeCacheKey{.a_m_size = kProbeM,
+                     MSizeCacheKey{.a_m_size = DNNL_RUNTIME_DIM_VAL,
                                   .use_bias = false,
                                   .bias_type = dnnl::memory::data_type::undef},
-                     /*first_time=*/true)
+                     true)
                     .weights_desc());
  init_runtime_memory_cache(args);
 }
@@ -18,8 +18,8 @@ struct KernelVecType<float> {

 template <>
 struct KernelVecType<c10::Half> {
-#if defined(__powerpc64__)
-  // Power specific vector types
+#if defined(__powerpc64__) || defined(__s390x__)
+  // Power and s390x architecture-specific vector types
  using qk_load_vec_type = vec_op::FP32Vec16;
  using qk_vec_type = vec_op::FP32Vec16;
  using v_load_vec_type = vec_op::FP32Vec16;
@@ -38,7 +38,16 @@ struct KernelVecType<c10::BFloat16> {
  using qk_vec_type = vec_op::BF16Vec32;
  using v_load_vec_type = vec_op::BF16Vec16;
 };
-#else
+
+#elif defined(__s390x__)
+template <>
+struct KernelVecType<c10::BFloat16> {
+  using qk_load_vec_type = vec_op::BF16Vec16;
+  using qk_vec_type = vec_op::FP32Vec16;
+  using v_load_vec_type = vec_op::BF16Vec16;
+};
+
+#elif defined(__aarch64__)
 template <>
 struct KernelVecType<c10::BFloat16> {
  using qk_load_vec_type = vec_op::BF16Vec16;
@@ -119,8 +119,8 @@ void cpu_fused_moe(torch::Tensor& output, const torch::Tensor& input,
                   const std::optional<torch::Tensor>& w13_bias,
                   const std::optional<torch::Tensor>& w2_bias,
                   const torch::Tensor& topk_weights,
-                   const torch::Tensor& topk_id, const bool skip_weighted,
-                   const std::string& act, const std::string& isa);
+                   const torch::Tensor& topk_id, const std::string& act,
+                   const std::string& isa);

 TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  // vLLM custom ops
@@ -320,7 +320,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  ops.def(
      "cpu_fused_moe(Tensor(a0!) output, Tensor input, Tensor w13, Tensor w2, "
      "Tensor? w13_bias, Tensor? w2_bias, Tensor topk_weights, Tensor topk_id, "
-      "bool skip_weighted, "
      "str act, str isa) -> ()");
  ops.impl("cpu_fused_moe", torch::kCPU, &cpu_fused_moe);
 #endif
@@ -2,58 +2,33 @@
 #include <torch/cuda.h>
 #include <cuda_runtime.h>

-// This function assumes that `cpu_tensor` is a CPU tensor,
-// and that UVA (Unified Virtual Addressing) is enabled.
+// This function assumes that `cpu_tensor` is a CPU tensor allocated with pinned
+// memory, and that UVA (Unified Virtual Addressing) is enabled.
 torch::Tensor get_cuda_view_from_cpu_tensor(torch::Tensor& cpu_tensor) {
  TORCH_CHECK(cpu_tensor.device().is_cpu(), "Input tensor must be on CPU");

-  // handle empty tensor
-  if (cpu_tensor.numel() == 0) {
-    return torch::empty(cpu_tensor.sizes(),
-                        cpu_tensor.options().device(torch::kCUDA));
-  }
-
-  if (cpu_tensor.is_pinned()) {
-    // If CPU tensor is pinned, directly get the device pointer.
-    void* host_ptr = const_cast<void*>(cpu_tensor.data_ptr());
-    void* device_ptr = nullptr;
-    cudaError_t err = cudaHostGetDevicePointer(&device_ptr, host_ptr, 0);
-    TORCH_CHECK(err == cudaSuccess,
-                "cudaHostGetDevicePointer failed: ", cudaGetErrorString(err));
-
-    return torch::from_blob(
-        device_ptr, cpu_tensor.sizes(), cpu_tensor.strides(),
-        [base = cpu_tensor](void*) {},  // keep cpu tensor alive
-        cpu_tensor.options().device(torch::kCUDA));
-  }
-
-  // If CPU tensor is not pinned, allocate a new pinned memory buffer.
-  torch::Tensor contiguous_cpu = cpu_tensor.contiguous();
-  size_t nbytes = contiguous_cpu.nbytes();
-
-  void* host_ptr = nullptr;
-  cudaError_t err = cudaHostAlloc(&host_ptr, nbytes, cudaHostAllocMapped);
-  if (err != cudaSuccess) {
-    AT_ERROR("cudaHostAlloc failed: ", cudaGetErrorString(err));
-  }
-
-  err = cudaMemcpy(host_ptr, contiguous_cpu.data_ptr(), nbytes,
-                   cudaMemcpyDefault);
-  if (err != cudaSuccess) {
-    cudaFreeHost(host_ptr);
-    AT_ERROR("cudaMemcpy failed: ", cudaGetErrorString(err));
-  }
+  // Get raw host pointer from CPU tensor
+  void* host_ptr = cpu_tensor.data_ptr();

+  // Get a device pointer corresponding to the pinned host memory
  void* device_ptr = nullptr;
-  err = cudaHostGetDevicePointer(&device_ptr, host_ptr, 0);
-  if (err != cudaSuccess) {
-    cudaFreeHost(host_ptr);
-    AT_ERROR("cudaHostGetDevicePointer failed: ", cudaGetErrorString(err));
-  }
+  cudaError_t err = cudaHostGetDevicePointer(&device_ptr, host_ptr, 0);
+  TORCH_CHECK(err == cudaSuccess,
+              "cudaHostGetDevicePointer failed: ", cudaGetErrorString(err));

-  auto deleter = [host_ptr](void*) { cudaFreeHost(host_ptr); };
+  // We'll use the same sizes, strides, and dtype as the CPU tensor.
+  // TODO: check if layout is respected.
+  auto sizes = cpu_tensor.sizes();
+  auto strides = cpu_tensor.strides();
+  auto options = cpu_tensor.options().device(torch::kCUDA);

-  return torch::from_blob(device_ptr, contiguous_cpu.sizes(),
-                          contiguous_cpu.strides(), deleter,
-                          contiguous_cpu.options().device(torch::kCUDA));
-}
+  // use default no-op deleter, since the memory is owned by the original CPU
+  // tensor
+  torch::Tensor cuda_tensor =
+      torch::from_blob(device_ptr, sizes, strides, options);
+
+  TORCH_CHECK(cuda_tensor.device().is_cuda(),
+              "Resulting tensor is not on CUDA device");
+
+  return cuda_tensor;
+}
@@ -114,10 +114,6 @@ void top_k_per_row_decode(const torch::Tensor& logits, int64_t next_n,
                          int64_t numRows, int64_t stride0, int64_t stride1,
                          int64_t topK);

-void large_context_topk(const torch::Tensor& score, torch::Tensor& indices,
-                        const torch::Tensor& lengths,
-                        std::optional<torch::Tensor> row_starts_opt);
-
 void rms_norm_static_fp8_quant(torch::Tensor& out, torch::Tensor& input,
                               torch::Tensor& weight, torch::Tensor& scale,
                               double epsilon);
@@ -725,4 +725,4 @@ void top_k_per_row_prefill(const torch::Tensor& logits,
                     static_cast<int>(stride0), static_cast<int>(stride1),
                     static_cast<int>(topK), kSortingAlgorithmThreshold);
  }
-}
+}
@@ -6,11 +6,11 @@
 #include "cutlass_extensions/common.hpp"

 bool cutlass_sparse_scaled_mm_supported(int64_t cuda_device_capability) {
-  // sparse CUTLASS kernels need exactly hopper and are not forward compatible
+  // sparse CUTLASS kernels need at least
  //   CUDA 12.2 and SM90 (Hopper)

 #if defined CUDA_VERSION
-  return CUDA_VERSION >= 12020 && cuda_device_capability == 90;
+  return CUDA_VERSION >= 12020 && cuda_device_capability >= 90;
 #endif

  return false;
@@ -98,7 +98,7 @@ std::vector<torch::Tensor> cutlass_sparse_compress(torch::Tensor const& a) {

  TORCH_CHECK_NOT_IMPLEMENTED(
      false,
-      "No compiled cutlass_sparse_compress for a compute capability equal to "
+      "No compiled cutlass_sparse_compress for a compute capability less than "
      "CUDA device capability: ",
      version_num);
 }
@@ -1,373 +0,0 @@
-// Portions of this file are adapted from SGLang PR:
-// https://github.com/sgl-project/sglang/pull/11194
-// and
-// https://github.com/sgl-project/sglang/pull/17747
-
-#include "cuda_compat.h"
-#include "dispatch_utils.h"
-
-#include <torch/cuda.h>
-#include <c10/cuda/CUDAGuard.h>
-
-#ifndef USE_ROCM
-  #include <cub/cub.cuh>
-#else
-  #include <hipcub/hipcub.hpp>
-#endif
-
-namespace vllm {
-
-constexpr int TopK = 2048;              // DeepSeek V3 sparse attention top-k
-constexpr int kThreadsPerBlock = 1024;  // Threads per block
-
-// Shared memory budget
-#if defined(USE_ROCM)
-constexpr size_t kSmem = 48 * 1024;  // ROCm default: 48KB
-#else
-// Reduced from 128KB to 32KB to improve occupancy.
-// Each radix pass needs at most ~TopK candidates in the threshold bin,
-// so 4K entries per round (2 rounds = 8K entries = 32KB) is sufficient.
-constexpr size_t kSmem = 8 * 1024 * sizeof(uint32_t);  // 32KB (bytes)
-#endif
-
-struct FastTopKParams {
-  const float* __restrict__ input;         // [batch, seq_len] Logits
-  const int32_t* __restrict__ row_starts;  // [batch] Offset into each row
-                                           // (optional)
-  int32_t* __restrict__ indices;           // [batch, TopK] Output top-k indices
-  int32_t* __restrict__ lengths;           // [batch] Sequence lengths per row
-  int64_t input_stride;                    // Stride between rows
-};
-
-__device__ __forceinline__ auto convert_to_uint32_v2(float x) -> uint32_t {
-  uint32_t bits = __float_as_uint(x);
-  return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u);
-}
-
-__device__ __forceinline__ auto convert_to_uint8(float x) -> uint8_t {
-  __half h = __float2half_rn(x);
-  uint16_t bits = __half_as_ushort(h);
-  uint16_t key = (bits & 0x8000) ? static_cast<uint16_t>(~bits)
-                                 : static_cast<uint16_t>(bits | 0x8000);
-  return static_cast<uint8_t>(key >> 8);
-}
-
-__device__ void naive_topk_cuda(const float* __restrict__ logits,
-                                int32_t* __restrict__ output_indices,
-                                int32_t seq_len) {
-  const int thread_id = threadIdx.x;
-  for (int i = thread_id; i < TopK; i += kThreadsPerBlock) {
-    output_indices[i] = (i < seq_len) ? i : -1;
-  }
-}
-
-// Adapted from:
-// https://github.com/sgl-project/sglang/blob/v0.5.8/sgl-kernel/csrc/elementwise/topk.cu#L87
-// by: DarkSharpness
-// which at the same time is an optimized topk kernel copied from tilelang
-// kernel
-__device__ void fast_topk_cuda_tl(
-    const float* __restrict__ logits,  // Input logits [seq_len]
-    int* __restrict__ output_indices,  // Output top-k indices [TopK]
-    int logits_offset,                 // Starting offset in logits array
-    int seq_len)                       // Number of valid logits to process
-{
-  constexpr int RADIX = 256;
-  constexpr int MAX_BUFFERED_ITEMS = kSmem / (2 * sizeof(int));
-
-  alignas(128) __shared__ int shared_histogram[2][RADIX + 128];
-  alignas(128) __shared__ int shared_output_count;
-  alignas(128) __shared__ int shared_threshold_bin;
-  alignas(128) __shared__ int shared_buffered_count[2];
-
-  extern __shared__ int buffered_indices[][MAX_BUFFERED_ITEMS];
-
-  const int thread_id = threadIdx.x;
-  int remaining_k = TopK;
-
-  // Pass 0: Build coarse 8-bit histogram using FP16 high bits
-  if (thread_id < RADIX + 1) {
-    shared_histogram[0][thread_id] = 0;
-  }
-  __syncthreads();
-
-  for (int idx = thread_id; idx < seq_len; idx += kThreadsPerBlock) {
-    const auto bin = convert_to_uint8(logits[idx + logits_offset]);
-    ::atomicAdd(&shared_histogram[0][bin], 1);
-  }
-  __syncthreads();
-
-  // Helper: Compute cumulative sum (suffix sum) over histogram using ping-pong
-  // buffers
-  auto compute_cumulative_sum = [&]() {
-    static_assert(1 << 8 == RADIX,
-                  "Radix must be 256 for 8 unrolled iterations");
-#pragma unroll 8
-    for (int i = 0; i < 8; ++i) {
-      if (C10_LIKELY(thread_id < RADIX)) {
-        const int stride = 1 << i;
-        const int src_buffer = i & 1;
-        const int dst_buffer = src_buffer ^ 1;
-
-        int value = shared_histogram[src_buffer][thread_id];
-        if (thread_id < RADIX - stride) {
-          value += shared_histogram[src_buffer][thread_id + stride];
-        }
-        shared_histogram[dst_buffer][thread_id] = value;
-      }
-      __syncthreads();
-    }
-  };
-
-  compute_cumulative_sum();
-
-  // Find threshold bin where cumsum crosses remaining_k
-  if (thread_id < RADIX && shared_histogram[0][thread_id] > remaining_k &&
-      shared_histogram[0][thread_id + 1] <= remaining_k) {
-    shared_threshold_bin = thread_id;
-    shared_buffered_count[0] = 0;
-    shared_output_count = 0;
-  }
-  __syncthreads();
-
-  const int threshold_bin = shared_threshold_bin;
-  remaining_k -= shared_histogram[0][threshold_bin + 1];
-
-  // Early exit if threshold bin perfectly matches remaining_k
-  if (remaining_k == 0) {
-    for (int idx = thread_id; idx < seq_len; idx += kThreadsPerBlock) {
-      const int bin = convert_to_uint8(logits[idx + logits_offset]);
-      if (bin > threshold_bin) {
-        const int output_pos = ::atomicAdd(&shared_output_count, 1);
-        output_indices[output_pos] = idx;
-      }
-    }
-    __syncthreads();
-    return;
-  }
-
-  // Prepare for refinement passes: Process threshold bin
-  __syncthreads();
-  if (thread_id < RADIX + 1) {
-    shared_histogram[0][thread_id] = 0;
-  }
-  __syncthreads();
-
-  // Scan all elements and:
-  // 1. Write indices > threshold_bin to output
-  // 2. Buffer indices == threshold_bin for refinement
-  // 3. Build histogram for next refinement pass (fused optimization)
-  for (int idx = thread_id; idx < seq_len; idx += kThreadsPerBlock) {
-    const float logit_value = logits[idx + logits_offset];
-    const int bin = convert_to_uint8(logit_value);
-
-    if (bin > threshold_bin) {
-      // in top-k, write to output
-      const int output_pos = ::atomicAdd(&shared_output_count, 1);
-      output_indices[output_pos] = idx;
-    } else if (bin == threshold_bin) {
-      // Candidate for top-k, needs refinement
-      const int buffer_pos = ::atomicAdd(&shared_buffered_count[0], 1);
-      if (C10_LIKELY(buffer_pos < MAX_BUFFERED_ITEMS)) {
-        buffered_indices[0][buffer_pos] = idx;
-        // Fused: Build histogram for next pass
-        const uint32_t fp32_bits = convert_to_uint32_v2(logit_value);
-        const int next_bin = (fp32_bits >> 24) & 0xFF;
-        ::atomicAdd(&shared_histogram[0][next_bin], 1);
-      }
-    }
-  }
-  __syncthreads();
-
-  // ============================================================================
-  // Passes 1-4: Refine using 8-bit passes over FP32 bits
-  // ============================================================================
-  // FP32 bits [31:0] split into 4 bytes processed MSB-first:
-  // Pass 1: bits [31:24], Pass 2: bits [23:16], Pass 3: bits [15:8], Pass 4:
-  // bits [7:0]
-#pragma unroll 4
-  for (int pass = 0; pass < 4; ++pass) {
-    __shared__ int shared_final_k;  // For final pass: remaining slots to fill
-    const int src_buffer = pass % 2;
-    const int dst_buffer = src_buffer ^ 1;
-
-    // Clamp buffered count to prevent overflow
-    const int raw_buffered = shared_buffered_count[src_buffer];
-    const int num_buffered =
-        (raw_buffered < MAX_BUFFERED_ITEMS) ? raw_buffered : MAX_BUFFERED_ITEMS;
-
-    compute_cumulative_sum();
-
-    // Find threshold bin for this pass
-    if (thread_id < RADIX && shared_histogram[0][thread_id] > remaining_k &&
-        shared_histogram[0][thread_id + 1] <= remaining_k) {
-      shared_threshold_bin = thread_id;
-      shared_buffered_count[dst_buffer] = 0;
-      shared_final_k = remaining_k - shared_histogram[0][thread_id + 1];
-    }
-    __syncthreads();
-
-    const int threshold_bin = shared_threshold_bin;
-    remaining_k -= shared_histogram[0][threshold_bin + 1];
-
-    // Bit offset for this pass: 24, 16, 8, 0
-    const int bit_offset = 24 - pass * 8;
-
-    // Early exit if threshold bin perfectly matches
-    if (remaining_k == 0) {
-      for (int i = thread_id; i < num_buffered; i += kThreadsPerBlock) {
-        const int idx = buffered_indices[src_buffer][i];
-        const uint32_t fp32_bits =
-            convert_to_uint32_v2(logits[idx + logits_offset]);
-        const int bin = (fp32_bits >> bit_offset) & 0xFF;
-        if (bin > threshold_bin) {
-          const int output_pos = ::atomicAdd(&shared_output_count, 1);
-          output_indices[output_pos] = idx;
-        }
-      }
-      __syncthreads();
-      break;
-    }
-
-    // Continue refinement
-    __syncthreads();
-    if (thread_id < RADIX + 1) {
-      shared_histogram[0][thread_id] = 0;
-    }
-    __syncthreads();
-
-    for (int i = thread_id; i < num_buffered; i += kThreadsPerBlock) {
-      const int idx = buffered_indices[src_buffer][i];
-      const float logit_value = logits[idx + logits_offset];
-      const uint32_t fp32_bits = convert_to_uint32_v2(logit_value);
-      const int bin = (fp32_bits >> bit_offset) & 0xFF;
-
-      if (bin > threshold_bin) {
-        // Definitely in top-k
-        const int output_pos = ::atomicAdd(&shared_output_count, 1);
-        output_indices[output_pos] = idx;
-      } else if (bin == threshold_bin) {
-        if (pass == 3) {
-          // Final pass (bits [7:0]): No more refinement possible
-          // Fill remaining slots in reverse order to maintain descending order
-          const int slot = ::atomicAdd(&shared_final_k, -1);
-          if (slot > 0) {
-            output_indices[TopK - slot] = idx;
-          }
-        } else {
-          // Buffer for next pass and build next histogram
-          const int buffer_pos =
-              ::atomicAdd(&shared_buffered_count[dst_buffer], 1);
-          if (C10_LIKELY(buffer_pos < MAX_BUFFERED_ITEMS)) {
-            buffered_indices[dst_buffer][buffer_pos] = idx;
-            // Fused: Build histogram for next pass
-            const int next_bit_offset = bit_offset - 8;
-            const int next_bin = (fp32_bits >> next_bit_offset) & 0xFF;
-            ::atomicAdd(&shared_histogram[0][next_bin], 1);
-          }
-        }
-      }
-    }
-    __syncthreads();
-  }
-}
-
-__global__ __launch_bounds__(kThreadsPerBlock) void topk_kernel(
-    const FastTopKParams params) {
-  const auto& [input, row_starts, indices, lengths, input_stride] = params;
-  const uint64_t batch_idx = blockIdx.x;
-  const int logits_offset = row_starts == nullptr ? 0 : row_starts[batch_idx];
-  const int seq_len = lengths[batch_idx];
-  int* output_indices = indices + batch_idx * TopK;
-  const float* logits = input + batch_idx * input_stride;
-
-  if (seq_len <= TopK) {
-    // Shortcut: All elements are in top-k
-    return naive_topk_cuda(logits, output_indices, seq_len);
-  } else {
-    return fast_topk_cuda_tl(logits, output_indices, logits_offset, seq_len);
-  }
-}
-
-FastTopKParams get_params(
-    const at::Tensor& score, const at::Tensor& lengths,
-    std::optional<at::Tensor> row_starts_opt = std::nullopt,
-    std::optional<at::Tensor> indices_opt = std::nullopt) {
-  const int64_t batch_size = score.size(0);
-
-  TORCH_CHECK(score.dim() == 2 && score.stride(1) == 1,
-              "score must be 2D with contiguous rows");
-  TORCH_CHECK(lengths.dim() == 1 && lengths.is_contiguous() &&
-                  lengths.size(0) == batch_size,
-              "lengths must be 1D contiguous with size matching batch");
-
-  const int32_t* row_starts_ptr = nullptr;
-  if (row_starts_opt.has_value()) {
-    const auto& row_starts = *row_starts_opt;
-    TORCH_CHECK(row_starts.dim() == 1 && row_starts.size(0) == batch_size,
-                "row_starts must be 1D with size matching batch");
-    row_starts_ptr = row_starts.data_ptr<int32_t>();
-  }
-
-  int32_t* indices_ptr = nullptr;
-  if (indices_opt.has_value()) {
-    const auto& indices = *indices_opt;
-    TORCH_CHECK(indices.dim() == 2 && indices.is_contiguous() &&
-                    indices.size(0) == batch_size && indices.size(1) == TopK,
-                "indices must be 2D contiguous [batch, TopK]");
-    indices_ptr = indices.data_ptr<int32_t>();
-  }
-
-  return FastTopKParams{
-      .input = score.data_ptr<float>(),
-      .row_starts = row_starts_ptr,
-      .indices = indices_ptr,
-      .lengths = lengths.data_ptr<int32_t>(),
-      .input_stride = score.stride(0),
-  };
-}
-
-template <auto* kernel_func, size_t smem_bytes>
-void setup_kernel_smem_once() {
-  static const cudaError_t result = []() -> cudaError_t {
-#ifdef USE_ROCM
-    auto func_ptr = reinterpret_cast<const void*>(kernel_func);
-#else
-    auto func_ptr = kernel_func;
-#endif
-    return cudaFuncSetAttribute(
-        func_ptr, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes);
-  }();
-
-  TORCH_CHECK(
-      result == cudaSuccess,
-      "Failed to set kernel shared memory limit: ", cudaGetErrorString(result));
-}
-
-}  // namespace vllm
-
-void large_context_topk(
-    const torch::Tensor& logits, torch::Tensor& indices,
-    const torch::Tensor& seq_lens,
-    std::optional<torch::Tensor> row_starts = std::nullopt) {
-  TORCH_CHECK(logits.is_cuda(), "logits must be a CUDA tensor");
-  TORCH_CHECK(indices.is_cuda(), "indices must be a CUDA tensor");
-  TORCH_CHECK(seq_lens.is_cuda(), "seq_lens must be a CUDA tensor");
-  if (row_starts.has_value()) {
-    TORCH_CHECK(row_starts->is_cuda(), "row_starts must be a CUDA tensor");
-  }
-
-  const auto params = vllm::get_params(logits, seq_lens, row_starts, indices);
-  const int64_t batch_size = logits.size(0);
-
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-  const dim3 grid(static_cast<uint32_t>(batch_size));
-  const dim3 block(vllm::kThreadsPerBlock);
-
-  vllm::setup_kernel_smem_once<vllm::topk_kernel, vllm::kSmem>();
-  vllm::topk_kernel<<<grid, block, vllm::kSmem, stream>>>(params);
-
-  const cudaError_t result = cudaGetLastError();
-  TORCH_CHECK(result == cudaSuccess,
-              "large_context_topk kernel failed: ", cudaGetErrorString(result));
-}
@@ -190,12 +190,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
      "int numRows, int stride0, int stride1, int topK) -> ()");
  ops.impl("top_k_per_row_decode", torch::kCUDA, &top_k_per_row_decode);

-  ops.def(
-      "large_context_topk(Tensor score, Tensor indices, Tensor lengths, "
-      "Tensor? "
-      "row_starts_opt) -> ()");
-  ops.impl("large_context_topk", torch::kCUDA, &large_context_topk);
-
  // Layernorm-quant
  // Apply Root Mean Square (RMS) Normalization to the input tensor.
  ops.def(
@@ -1,7 +1,7 @@
 # default base image
 ARG REMOTE_VLLM="0"
 ARG COMMON_WORKDIR=/app
-ARG BASE_IMAGE=rocm/vllm-dev:base
+ARG BASE_IMAGE=rocm/vllm-dev:base_custom_releases_rocm_v0.16.0_20260211

 # Sccache configuration (only used in release pipeline)
 ARG USE_SCCACHE
@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=rocm/dev-ubuntu-22.04:7.0-complete
+ARG BASE_IMAGE=rocm/dev-ubuntu-22.04:7.2-complete
 ARG TRITON_BRANCH="57c693b6"
 ARG TRITON_REPO="https://github.com/ROCm/triton.git"
 ARG PYTORCH_BRANCH="89075173"
@@ -142,6 +142,7 @@ ARG PYTORCH_VISION_REPO
 ARG PYTORCH_AUDIO_REPO
 ARG USE_SCCACHE

+RUN apt-get update && apt-get install -y pkg-config liblzma-dev
 RUN git clone ${PYTORCH_REPO} pytorch
 RUN cd pytorch && git checkout ${PYTORCH_BRANCH} \
    && pip install -r requirements.txt && git submodule update --init --recursive \
@@ -30,7 +30,6 @@ th {
 | HuggingFace-Other | ✅ | ✅ | `lmms-lab/LLaVA-OneVision-Data`, `Aeala/ShareGPT_Vicuna_unfiltered` |
 | HuggingFace-MTBench | ✅ | ✅ | `philschmid/mt-bench` |
 | HuggingFace-Blazedit | ✅ | ✅ | `vdaita/edit_5k_char`, `vdaita/edit_10k_char` |
-| HuggingFace-ASR | ✅ | ✅ | `openslr/librispeech_asr`, `facebook/voxpopuli`,  `LIUM/tedlium`, `edinburghcstr/ami`,        `speechcolab/gigaspeech`,        `kensho/spgispeech` |
 | Spec Bench | ✅ | ✅ | `wget https://raw.githubusercontent.com/hemingkx/Spec-Bench/refs/heads/main/data/spec_bench/question.jsonl` |
 | Custom | ✅ | ✅ | Local file: `data.jsonl` |
 | Custom MM | ✅ | ✅ | Local file: `mm_data.jsonl` |
@@ -300,22 +299,6 @@ vllm bench serve \
    --blazedit-max-distance 0.99
 ```

-`openslr/librispeech_asr`, `facebook/voxpopuli`, `LIUM/tedlium`, `edinburghcstr/ami`, `speechcolab/gigaspeech`, `kensho/spgispeech`
-
-```bash
-vllm bench serve \
-    --model openai/whisper-large-v3-turbo \
-    --backend openai-audio \
-    --dataset-name hf \
-    --dataset-path facebook/voxpopuli --hf-subset en --hf-split test --no-stream --trust-remote-code \
-    --num-prompts 99999999 \
-    --no-oversample \
-    --endpoint /v1/audio/transcriptions \
-    --ready-check-timeout-sec 600 \
-    --save-result \
-    --max-concurrency 512
-```
-
 #### Running With Sampling Parameters

 When using OpenAI-compatible backends such as `vllm`, optional sampling
@@ -128,7 +128,6 @@ Priority is **1 = highest** (tried first).
 | 4 | `FLASHMLA` |
 | 5 | `TRITON_MLA` |
 | 6 | `FLASHMLA_SPARSE` |
-| 7 | `FLASHINFER_MLA_SPARSE` |

 **Ampere/Hopper (SM 8.x-9.x):**

@@ -153,7 +152,6 @@ Priority is **1 = highest** (tried first).
 | **Sink** | Attention sink support (for StreamingLLM) |
 | **Sparse** | Sparse attention support (MLA only) |
 | **MM Prefix** | Multimodal prefix full attention support |
-| **DCP** | Decode Context Parallelism support (`--decode-context-parallel-size`) |
 | **Attention Types** | Supported attention patterns (Decoder, Encoder, Enc-Dec) |
 | **Compute Cap.** | Required CUDA compute capability (N/A for non-CUDA backends) |

@@ -161,20 +159,20 @@ Priority is **1 = highest** (tried first).

 ## Standard Attention (MHA, MQA, GQA) Backends

-| Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | MM Prefix | DCP | Attention Types | Compute Cap. |
-|---------|---------|--------|-----------|-------------|------------|------|-----------|-----|-----------------|--------------|
-| `CPU_ATTN` |  | fp16, bf16, fp32 | `auto` | Any | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | All | N/A |
-| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
-| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ✅ | Decoder | 10.x |
-| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥8.0 |
-| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ❌ | ✅ | All | 9.x |
-| `FLASH_ATTN_DIFFKV` |  | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ✅ | Decoder | Any |
-| `FLEX_ATTENTION` |  | fp16, bf16, fp32 | `auto`, `bfloat16` | Any | Any | ❌ | ✅ | ❌ | Decoder, Encoder Only | Any |
-| `ROCM_AITER_FA` |  | fp16, bf16 | `auto` | 16, 32 | 64, 128, 256 | ❌ | ❌ | ❌ | Decoder | N/A |
-| `ROCM_AITER_UNIFIED_ATTN` |  | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | Decoder | N/A |
-| `ROCM_ATTN` |  | fp16, bf16, fp32 | `auto` | 16, 32, 544 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | Decoder | 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 |
+| Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | MM Prefix | Attention Types | Compute Cap. |
+|---------|---------|--------|-----------|-------------|------------|------|-----------|-----------------|--------------|
+| `CPU_ATTN` |  | fp16, bf16, fp32 | `auto` | Any | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256 | ❌ | ❌ | All | N/A |
+| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ❌ | ❌ | Decoder | 7.x-9.x |
+| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | Decoder | 10.x |
+| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | All | ≥8.0 |
+| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ❌ | All | 9.x |
+| `FLASH_ATTN_DIFFKV` |  | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | Decoder | Any |
+| `FLEX_ATTENTION` |  | fp16, bf16, fp32 | `auto`, `bfloat16` | Any | Any | ❌ | ✅ | Decoder, Encoder Only | Any |
+| `ROCM_AITER_FA` |  | fp16, bf16 | `auto` | 16, 32 | 64, 128, 256 | ❌ | ❌ | Decoder | N/A |
+| `ROCM_AITER_UNIFIED_ATTN` |  | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | Decoder | N/A |
+| `ROCM_ATTN` |  | fp16, bf16, fp32 | `auto` | 16, 32, 544 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | Decoder | 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 |

 > **†** FlashInfer uses TRTLLM attention on Blackwell (SM100), which supports sinks. Disable via `--attention-config.use_trtllm_attention=0`.
 >
@@ -201,15 +199,14 @@ configuration.

 ### Decode Backends

-| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. |
-|---------|--------|-----------|-------------|------------|------|--------|-----------|-----|-----------------|--------------|
-| `CUTLASS_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
-| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
-| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `bfloat16` | 32, 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x |
-| `FLASHMLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |
-| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
-| `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x |
-| `ROCM_AITER_MLA` | fp16, bf16 | `auto` | 1 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
-| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto` | Any | 576 | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
-| `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
-| `TRITON_MLA` | fp16, bf16 | `auto`, `bfloat16` | Any | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
+| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Sparse | MM Prefix | Attention Types | Compute Cap. |
+|---------|--------|-----------|-------------|------------|------|--------|-----------|-----------------|--------------|
+| `CUTLASS_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | Decoder | 10.x |
+| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | Decoder | 10.x |
+| `FLASHMLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | Decoder | 9.x-10.x |
+| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ✅ | ❌ | Decoder | 9.x-10.x |
+| `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | Decoder | 9.x |
+| `ROCM_AITER_MLA` | fp16, bf16 | `auto` | 1 | Any | ❌ | ❌ | ❌ | Decoder | N/A |
+| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto` | Any | 576 | ❌ | ❌ | ❌ | Decoder | N/A |
+| `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | Decoder | N/A |
+| `TRITON_MLA` | fp16, bf16 | `auto`, `bfloat16` | Any | Any | ❌ | ❌ | ❌ | Decoder | Any |
@@ -14,26 +14,8 @@ IOProcessorOutput = TypeVar("IOProcessorOutput")

 class IOProcessor(ABC, Generic[IOProcessorInput, IOProcessorOutput]):
    def __init__(self, vllm_config: VllmConfig):
-        super().__init__()
-
        self.vllm_config = vllm_config

-    @abstractmethod
-    def parse_data(self, data: object) -> IOProcessorInput:
-        raise NotImplementedError
-
-    def merge_sampling_params(
-        self,
-        params: SamplingParams | None = None,
-    ) -> SamplingParams:
-        return params or SamplingParams()
-
-    def merge_pooling_params(
-        self,
-        params: PoolingParams | None = None,
-    ) -> PoolingParams:
-        return params or PoolingParams()
-
    @abstractmethod
    def pre_process(
        self,
@@ -73,13 +55,29 @@ class IOProcessor(ABC, Generic[IOProcessorInput, IOProcessorOutput]):
            [(i, item) async for i, item in model_output], key=lambda output: output[0]
        )
        collected_output = [output[1] for output in sorted_output]
-        return self.post_process(collected_output, request_id=request_id, **kwargs)
+        return self.post_process(collected_output, request_id, **kwargs)
+
+    @abstractmethod
+    def parse_request(self, request: Any) -> IOProcessorInput:
+        raise NotImplementedError
+
+    def validate_or_generate_params(
+        self, params: SamplingParams | PoolingParams | None = None
+    ) -> SamplingParams | PoolingParams:
+        return params or PoolingParams()
+
+    @abstractmethod
+    def output_to_response(
+        self, plugin_output: IOProcessorOutput
+    ) -> IOProcessorResponse:
+        raise NotImplementedError
 ```

-The `parse_data` method is used for validating the user data and converting it into the input expected by the `pre_process*` methods.
-The `merge_sampling_params` and `merge_pooling_params` methods merge input `SamplingParams` or `PoolingParams` (if any) with the default one.
+The `parse_request` method is used for validating the user prompt and converting it into the input expected by the `pre_process`/`pre_process_async` methods.
 The `pre_process*` methods take the validated plugin input to generate vLLM's model prompts for regular inference.
 The `post_process*` methods take `PoolingRequestOutput` objects as input and generate a custom plugin output.
+The `validate_or_generate_params` method is used for validating with the plugin any `SamplingParameters`/`PoolingParameters` received with the user request, or to generate new ones if none are specified. The function always returns the validated/generated parameters.
+The `output_to_response` method is used only for online serving and converts the plugin output to the `IOProcessorResponse` type that is then returned by the API Server. The implementation of the `/pooling` serving endpoint is available here [vllm/entrypoints/openai/serving_pooling.py](../../vllm/entrypoints/pooling/pooling/serving.py).

 An example implementation of a plugin that enables generating geotiff images with the PrithviGeospatialMAE model is available [here](https://github.com/IBM/terratorch/tree/main/terratorch/vllm/plugins/segmentation). Please, also refer to our online ([examples/pooling/plugin/prithvi_geospatial_mae_online.py](../../examples/pooling/plugin/prithvi_geospatial_mae_online.py)) and offline ([examples/pooling/plugin/prithvi_geospatial_mae_io_processor.py](../../examples/pooling/plugin/prithvi_geospatial_mae_io_processor.py)) inference examples.

@@ -32,7 +32,7 @@ th {

 | Backend | Output act. format | Quant. types | Quant. format | Async | Apply Weight On Input | Subclass |
 |---------|--------------------|--------------|---------------|-------|-----------------------|-----------|
-| naive | standard | all<sup>1</sup> | G,A,T | N | <sup>6</sup> | [layer.py][vllm.model_executor.layers.fused_moe.layer.FusedMoE |
+| naive | standard | all<sup>1</sup> | G,A,T | N | <sup>6</sup> | [layer.py][vllm.model_executor.layers.fused_moe.layer.FusedMoE.forward_impl] |
 | pplx | batched | fp8,int8 | G,A,T | Y | Y | [`PplxPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.pplx_prepare_finalize.PplxPrepareAndFinalize] |
 | deepep_high_throughput | standard | fp8 | G(128),A,T<sup>2</sup> | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ll_prepare_finalize.DeepEPLLPrepareAndFinalize] |
 | deepep_low_latency | batched | fp8 | G(128),A,T<sup>3</sup> | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize.DeepEPHTPrepareAndFinalize] |
@@ -521,7 +521,7 @@ First, launch the OpenAI-compatible server:

 ```bash
 vllm serve microsoft/Phi-3.5-vision-instruct --runner generate \
-  --trust-remote-code --max-model-len 4096 --limit-mm-per-prompt.image 2
+  --trust-remote-code --max-model-len 4096 --limit-mm-per-prompt '{"image":2}'
 ```

 Then, you can use the OpenAI client as follows:
@@ -48,7 +48,7 @@ th:not(:first-child) {
 |-----------------------|---------|----------|----------|-------|----------|-----------|-------------|-----------|
 | AWQ                   | ❌      | ✅︎       | ✅︎       | ✅︎    | ✅︎       | ❌         | ✅︎          | ✅︎        |
 | GPTQ                  | ✅︎      | ✅︎       | ✅︎       | ✅︎    | ✅︎       | ❌         | ✅︎          | ✅︎        |
-| Marlin (GPTQ/AWQ/FP8/FP4) | ❌      | ✅︎*       | ✅︎       | ✅︎    | ✅︎       | ❌         | ❌          | ❌        |
+| Marlin (GPTQ/AWQ/FP8) | ❌      | ❌       | ✅︎       | ✅︎    | ✅︎       | ❌         | ❌          | ❌        |
 | INT8 (W8A8)           | ❌      | ✅︎       | ✅︎       | ✅︎    | ✅︎       | ❌         | ❌          | ✅︎        |
 | FP8 (W8A8)            | ❌      | ❌       | ❌       | ✅︎    | ✅︎       | ✅︎         | ❌          | ❌        |
 | bitsandbytes          | ✅︎      | ✅︎       | ✅︎       | ✅︎    | ✅︎       | ❌         | ❌          | ❌        |
@@ -59,7 +59,6 @@ th:not(:first-child) {
 - ✅︎ indicates that the quantization method is supported on the specified hardware.
 - ❌ indicates that the quantization method is not supported on the specified hardware.
 - All Intel Gaudi quantization support has been migrated to [vLLM-Gaudi](https://github.com/vllm-project/vllm-gaudi).
- *Turing does not support Marlin MXFP4.

 !!! note
    For information on quantization support on Google TPU, please refer to the [TPU-Inference Recommended Models and Features](https://docs.vllm.ai/projects/tpu/en/latest/recommended_models_features/) documentation.
@@ -2,7 +2,7 @@

 vLLM supports FP8 (8-bit floating point) weight and activation quantization using hardware acceleration on GPUs such as Nvidia H100 and AMD MI300x.
 Currently, only Hopper and Ada Lovelace GPUs are officially supported for W8A8.
-Turing/Ampere GPUs are supported for W8A16 (weight-only FP8) utilizing Marlin kernels.
+Ampere GPUs are supported for W8A16 (weight-only FP8) utilizing Marlin kernels.
 Quantization of models with FP8 allows for a 2x reduction in model memory requirements and up to a 1.6x improvement in throughput with minimal impact on accuracy.

 Please visit the HF collection of [quantized FP8 checkpoints of popular LLMs ready to use with vLLM](https://huggingface.co/collections/neuralmagic/fp8-llms-for-vllm-666742ed2b78b7ac8df13127).
@@ -13,8 +13,8 @@ The FP8 types typically supported in hardware have two distinct representations,
 - **E5M2**: Consists of 1 sign bit, 5 exponent bits, and 2 bits of mantissa. It can store values up to +/-57344, +/- `inf`, and `nan`. The tradeoff for the increased dynamic range is lower precision of the stored values.

 !!! note
-    FP8 computation is supported on NVIDIA GPUs with compute capability >= 8.9 (Ada Lovelace, Hopper).
-    FP8 models will run on compute capability >= 7.5 (Turing) as weight-only W8A16, utilizing FP8 Marlin.
+    FP8 computation is supported on NVIDIA GPUs with compute capability > 8.9 (Ada Lovelace, Hopper).
+    FP8 models will run on compute capability > 8.0 (Ampere) as weight-only W8A16, utilizing FP8 Marlin.

 ## Installation

@@ -1,5 +1,10 @@
 # Speculative Decoding

+!!! warning
+    Please note that speculative decoding in vLLM is not yet optimized and does
+    not usually yield inter-token latency reductions for all prompt datasets or sampling parameters.
+    The work to optimize it is ongoing and can be followed here: <https://github.com/vllm-project/vllm/issues/4630>
+
 !!! warning
    Currently, speculative decoding in vLLM is not compatible with pipeline parallelism.

@@ -176,7 +176,7 @@ For the full and up-to-date list of models validated on CPU platforms, please se

 ### How to find benchmark configuration examples for supported CPU models?

-For any model listed under [Supported Models on CPU](../../models/hardware_supported_models/cpu.md), optimized runtime configurations are provided in the vLLM Benchmark Suite’s CPU test cases, defined in cpu test cases as serving-tests-cpu.json. Full test cases for Text-only models, Multi-Modal models and Embedded models are in cpu Text-Only test cases as serving-tests-cpu-text.json, cpu Multi-Modal test cases as serving-tests-cpu-multimodal.json and cpu Embedded test cases as serving-tests-cpu-embed.json.  
+For any model listed under [Supported Models on CPU](../../models/hardware_supported_models/cpu.md), optimized runtime configurations are provided in the vLLM Benchmark Suite’s CPU test cases, defined in [cpu test cases](../../../.buildkite/performance-benchmarks/tests/serving-tests-cpu.json)
 For details on how these optimized configurations are determined, see: [performance-benchmark-details](../../../.buildkite/performance-benchmarks/README.md#performance-benchmark-details).
 To benchmark the supported models using these optimized settings, follow the steps in [running vLLM Benchmark Suite manually](../../benchmarking/dashboard.md#manually-trigger-the-benchmark) and run the Benchmark Suite on a CPU environment.  

@@ -199,28 +199,6 @@ lscpu | grep "NUMA node(s):" | awk '{print $3}'
 For performance reference, users may also consult the [vLLM Performance Dashboard](https://hud.pytorch.org/benchmark/llms?repoName=vllm-project%2Fvllm&deviceName=cpu)
 , which publishes default-model CPU results produced using the same Benchmark Suite.

-#### Dry-Run
-
-For users only need to get the optimized runtime configurations without running benchmark, a Dry-Run mode is provided.
-By passing an environment variable DRY_RUN=1 with run-performance-benchmarks.sh,
-all commands will be generated under `./benchmark/results/`.
-
-```bash
-ON_CPU=1 DRY_RUN=1 bash .buildkite/performance-benchmarks/scripts/run-performance-benchmarks.sh
-```
-
-By providing different JSON file, users can get runtime configurations for different models such as Embedded Models.
-
-```bash
-ON_CPU=1 SERVING_JSON=serving-tests-cpu-embed.json DRY_RUN=1 bash .buildkite/performance-benchmarks/scripts/run-performance-benchmarks.sh
-```
-
-By providing MODEL_FILTER and DTYPE_FILTER, only commands for related model ID and Data Type will be generated.
-
-```bash
-ON_CPU=1 SERVING_JSON=serving-tests-cpu-text.json DRY_RUN=1 MODEL_FILTER=meta-llama/Llama-3.1-8B-Instruct DTYPE_FILTER=bfloat16  bash .buildkite/performance-benchmarks/scripts/run-performance-benchmarks.sh
-```
-
 ### How to decide `VLLM_CPU_OMP_THREADS_BIND`?

 - Default `auto` thread-binding is recommended for most cases. Ideally, each OpenMP thread will be bound to a dedicated physical core respectively, threads of each rank will be bound to the same NUMA node respectively, and 1 CPU per rank will be reserved for other vLLM components when `world_size > 1`. If you have any performance problems or unexpected binding behaviours, please try to bind threads as following.
@@ -311,31 +311,20 @@ An OpenAI client example can be found here: [examples/pooling/embed/openai_embed

 [ColBERT](https://arxiv.org/abs/2004.12832) (Contextualized Late Interaction over BERT) is a retrieval model that uses per-token embeddings and MaxSim scoring for document ranking. Unlike single-vector embedding models, ColBERT retains token-level representations and computes relevance scores through late interaction, providing better accuracy while being more efficient than cross-encoders.

-vLLM supports ColBERT models with multiple encoder backbones:
-
-| Architecture | Backbone | Example HF Models |
-|---|---|---|
-| `HF_ColBERT` | BERT | `answerdotai/answerai-colbert-small-v1`, `colbert-ir/colbertv2.0` |
-| `ColBERTModernBertModel` | ModernBERT | `lightonai/GTE-ModernColBERT-v1` |
-| `ColBERTJinaRobertaModel` | Jina XLM-RoBERTa | `jinaai/jina-colbert-v2` |
-
-**BERT-based ColBERT** models work out of the box:
+vLLM supports ColBERT models for reranking tasks, automatically applying MaxSim scoring for query-document relevance:

 ```shell
 vllm serve answerdotai/answerai-colbert-small-v1
 ```

-For **non-BERT backbones**, use `--hf-overrides` to set the correct architecture:
+Currently supports ColBERT models with standard BERT encoders (e.g., `answerdotai/answerai-colbert-small-v1`, `colbert-ir/colbertv2.0`).
+
+ColBERT models with modified encoder architectures are not yet supported, including BERT variants with rotary embeddings (e.g., `jinaai/jina-colbert-v2`) or other custom encoders (e.g., `LiquidAI/LFM2-ColBERT-350M`).
+
+If your standard BERT ColBERT model's config doesn't specify the architecture as `HF_ColBERT`, override it with:

 ```shell
-# ModernBERT backbone
-vllm serve lightonai/GTE-ModernColBERT-v1 \
-    --hf-overrides '{"architectures": ["ColBERTModernBertModel"]}'
-
-# Jina XLM-RoBERTa backbone
-vllm serve jinaai/jina-colbert-v2 \
-    --hf-overrides '{"architectures": ["ColBERTJinaRobertaModel"]}' \
-    --trust-remote-code
+vllm serve your-colbert-model --hf-overrides '{"architectures": ["HF_ColBERT"]}'
 ```

 Then you can use the rerank endpoint:
@@ -374,77 +363,6 @@ curl -s http://localhost:8000/pooling -H "Content-Type: application/json" -d '{

 An example can be found here: [examples/pooling/score/colbert_rerank_online.py](../../examples/pooling/score/colbert_rerank_online.py)

-### ColQwen3 Multi-Modal Late Interaction Models
-
-ColQwen3 is based on [ColPali](https://arxiv.org/abs/2407.01449), which extends ColBERT's late interaction approach to **multi-modal** inputs. While ColBERT operates on text-only token embeddings, ColPali/ColQwen3 can embed both **text and images** (e.g. PDF pages, screenshots, diagrams) into per-token L2-normalized vectors and compute relevance via MaxSim scoring. ColQwen3 specifically uses Qwen3-VL as its vision-language backbone.
-
-| Architecture | Backbone | Example HF Models |
-|---|---|---|
-| `ColQwen3` | Qwen3-VL | `TomoroAI/tomoro-colqwen3-embed-4b`, `TomoroAI/tomoro-colqwen3-embed-8b` |
-| `OpsColQwen3Model` | Qwen3-VL | `OpenSearch-AI/Ops-Colqwen3-4B`, `OpenSearch-AI/Ops-Colqwen3-8B` |
-
-Start the server:
-
-```shell
-vllm serve TomoroAI/tomoro-colqwen3-embed-4b --max-model-len 4096
-```
-
-Then you can use the rerank endpoint:
-
-```shell
-curl -s http://localhost:8000/rerank -H "Content-Type: application/json" -d '{
-    "model": "TomoroAI/tomoro-colqwen3-embed-4b",
-    "query": "What is machine learning?",
-    "documents": [
-        "Machine learning is a subset of artificial intelligence.",
-        "Python is a programming language.",
-        "Deep learning uses neural networks."
-    ]
-}'
-```
-
-Or the score endpoint:
-
-```shell
-curl -s http://localhost:8000/score -H "Content-Type: application/json" -d '{
-    "model": "TomoroAI/tomoro-colqwen3-embed-4b",
-    "text_1": "What is the capital of France?",
-    "text_2": ["The capital of France is Paris.", "Python is a programming language."]
-}'
-```
-
-You can also get the raw token embeddings using the pooling endpoint with `token_embed` task:
-
-```shell
-curl -s http://localhost:8000/pooling -H "Content-Type: application/json" -d '{
-    "model": "TomoroAI/tomoro-colqwen3-embed-4b",
-    "input": "What is machine learning?",
-    "task": "token_embed"
-}'
-```
-
-For **image inputs**, use the chat-style `messages` field so that the vLLM multimodal processor handles them correctly:
-
-```shell
-curl -s http://localhost:8000/pooling -H "Content-Type: application/json" -d '{
-    "model": "TomoroAI/tomoro-colqwen3-embed-4b",
-    "messages": [
-        {
-            "role": "user",
-            "content": [
-                {"type": "image_url", "image_url": {"url": "data:image/png;base64,<BASE64>"}},
-                {"type": "text", "text": "Describe the image."}
-            ]
-        }
-    ]
-}'
-```
-
-Examples can be found here:
-
- Multi-vector retrieval: [examples/pooling/token_embed/colqwen3_token_embed_online.py](../../examples/pooling/token_embed/colqwen3_token_embed_online.py)
- Reranking: [examples/pooling/score/colqwen3_rerank_online.py](../../examples/pooling/score/colqwen3_rerank_online.py)
-
 ### BAAI/bge-m3

 The `BAAI/bge-m3` model comes with extra weights for sparse and colbert embeddings but unfortunately in its `config.json`
@@ -658,7 +658,7 @@ On the other hand, modalities separated by `/` are mutually exclusive.
 See [this page](../features/multimodal_inputs.md) on how to pass multi-modal inputs to the model.

 !!! tip
-    For hybrid-only models such as Llama-4, Step3, Mistral-3 and Qwen-3.5, a text-only mode can be enabled by setting all supported multimodal modalities to 0 (`--language-model-only`) so that their multimodal modules will not be loaded to free up more GPU memory for KV cache.
+    For hybrid-only models such as Llama-4, Step3 and Mistral-3, a text-only mode can be enabled by setting all supported multimodal modalities to 0 (e.g, `--limit-mm-per-prompt '{"image":0}`) so that their multimodal modules will not be loaded to free up more GPU memory for KV cache.

 !!! note
    vLLM currently supports adding LoRA adapters to the language backbone for most multimodal models. Additionally, vLLM now experimentally supports adding LoRA to the tower and connector modules for some multimodal models. See [this page](../features/lora.md).
@@ -728,8 +728,6 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
 | `OpenPanguVLForConditionalGeneration` | openpangu-VL | T + I<sup>E+</sup> + V<sup>E+</sup> |`FreedomIntelligence/openPangu-VL-7B` | ✅︎ | ✅︎ |
 | `Ovis` | Ovis2, Ovis1.6 | T + I<sup>+</sup> | `AIDC-AI/Ovis2-1B`, `AIDC-AI/Ovis1.6-Llama3.2-3B`, etc. | | ✅︎ |
 | `Ovis2_5` | Ovis2.5 | T + I<sup>+</sup> + V | `AIDC-AI/Ovis2.5-9B`, etc. | | |
-| `Ovis2_6ForCausalLM` | Ovis2.6 | T + I<sup>+</sup> + V | `AIDC-AI/Ovis2.6-2B`, etc. | | |
-| `Ovis2_6_MoeForCausalLM` | Ovis2.6 | T + I<sup>+</sup> + V | `AIDC-AI/Ovis2.6-30B-A3B`, etc. | | |
 | `PaddleOCRVLForConditionalGeneration` | Paddle-OCR | T + I<sup>+</sup> | `PaddlePaddle/PaddleOCR-VL`, etc. | | |
 | `PaliGemmaForConditionalGeneration` | PaliGemma, PaliGemma 2 | T + I<sup>E</sup> | `google/paligemma-3b-pt-224`, `google/paligemma-3b-mix-224`, `google/paligemma2-3b-ft-docci-448`, etc. | ✅︎ | ✅︎ |
 | `Phi3VForCausalLM` | Phi-3-Vision, Phi-3.5-Vision | T + I<sup>E+</sup> | `microsoft/Phi-3-vision-128k-instruct`, `microsoft/Phi-3.5-vision-instruct`, etc. | | ✅︎ |
@@ -740,8 +738,6 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
 | `Qwen2VLForConditionalGeneration` | QVQ, Qwen2-VL | T + I<sup>E+</sup> + V<sup>E+</sup> | `Qwen/QVQ-72B-Preview`, `Qwen/Qwen2-VL-7B-Instruct`, `Qwen/Qwen2-VL-72B-Instruct`, etc. | ✅︎ | ✅︎ |
 | `Qwen2_5_VLForConditionalGeneration` | Qwen2.5-VL | T + I<sup>E+</sup> + V<sup>E+</sup> | `Qwen/Qwen2.5-VL-3B-Instruct`, `Qwen/Qwen2.5-VL-72B-Instruct`, etc. | ✅︎ | ✅︎ |
 | `Qwen2_5OmniThinkerForConditionalGeneration` | Qwen2.5-Omni | T + I<sup>E+</sup> + V<sup>E+</sup> + A<sup>+</sup> | `Qwen/Qwen2.5-Omni-3B`, `Qwen/Qwen2.5-Omni-7B` | ✅︎ | ✅︎ |
-| `Qwen3_5ForConditionalGeneration` | Qwen3.5 | T + I<sup>E+</sup> + V<sup>E+</sup> | `Qwen/Qwen3.5-9B-Instruct`, etc. | ✅︎ | ✅︎ |
-| `Qwen3_5MoeForConditionalGeneration` | Qwen3.5-MOE | T + I<sup>E+</sup> + V<sup>E+</sup> | `Qwen/Qwen3.5-35B-A3B-Instruct`, etc. | ✅︎ | ✅︎ |
 | `Qwen3VLForConditionalGeneration` | Qwen3-VL | T + I<sup>E+</sup> + V<sup>E+</sup> | `Qwen/Qwen3-VL-4B-Instruct`, etc. | ✅︎ | ✅︎ |
 | `Qwen3VLMoeForConditionalGeneration` | Qwen3-VL-MOE | T + I<sup>E+</sup> + V<sup>E+</sup> | `Qwen/Qwen3-VL-30B-A3B-Instruct`, etc. | ✅︎ | ✅︎ |
 | `Qwen3OmniMoeThinkerForConditionalGeneration` | Qwen3-Omni | T + I<sup>E+</sup> + V<sup>E+</sup> + A<sup>+</sup> | `Qwen/Qwen3-Omni-30B-A3B-Instruct`, `Qwen/Qwen3-Omni-30B-A3B-Thinking` | ✅︎ | ✅︎ |
@@ -792,7 +788,6 @@ Speech2Text models trained specifically for Automatic Speech Recognition.

 | Architecture | Models | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) |
 |--------------|--------|-------------------|----------------------|---------------------------|
-| `FunASRForConditionalGeneration` | FunASR | `allendou/Fun-ASR-Nano-2512-vllm`, etc. | | |
 | `Gemma3nForConditionalGeneration` | Gemma3n | `google/gemma-3n-E2B-it`, `google/gemma-3n-E4B-it`, etc. | | |
 | `GlmAsrForConditionalGeneration` | GLM-ASR | `zai-org/GLM-ASR-Nano-2512` | ✅︎ | ✅︎ |
 | `GraniteSpeechForConditionalGeneration` | Granite Speech | `ibm-granite/granite-speech-3.3-2b`, `ibm-granite/granite-speech-3.3-8b`, etc. | ✅︎ | ✅︎ |
@@ -28,4 +28,3 @@ It demonstrates vLLM's ability to recover from KV load failures in both synchron

 ```bash
 ./run.sh
-```
@@ -18,11 +18,11 @@ from vllm.assets.image import ImageAsset
 # # Mistral format
 # vllm serve mistralai/Mistral-Small-3.1-24B-Instruct-2503 \
 #   --tokenizer-mode mistral --config-format mistral --load-format mistral \
-#   --limit-mm-per-prompt.image 4 --max-model-len 16384
+#   --limit-mm-per-prompt '{"image":4}' --max-model-len 16384
 #
 # # HF format
 # vllm serve mistralai/Mistral-Small-3.1-24B-Instruct-2503 \
-#   --limit-mm-per-prompt.image 4 --max-model-len 16384
+#   --limit-mm-per-prompt '{"image":4}' --max-model-len 16384
 # ```
 #
 # - Client:
@@ -1,112 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from __future__ import annotations
-
-from vllm import LLM, EngineArgs
-from vllm.config import ProfilerConfig
-from vllm.utils.argparse_utils import FlexibleArgumentParser
-
-DEFAULT_MAX_TOKENS = 16
-
-
-def create_parser() -> FlexibleArgumentParser:
-    parser = FlexibleArgumentParser()
-    EngineArgs.add_cli_args(parser)
-    parser.set_defaults(model="meta-llama/Llama-3.2-1B-Instruct")
-
-    batch_group = parser.add_argument_group("Batch parameters")
-    batch_group.add_argument("--batch-size", type=int, default=1)
-    batch_group.add_argument("--prompt-size", type=int, default=128)
-    batch_group.add_argument("--prompt-prefix", type=str, default="Hello, my name is")
-
-    profile_group = parser.add_argument_group("Profiling parameters")
-    profile_group.add_argument(
-        "--profile",
-        choices=["none", "prefill", "decode", "both"],
-        default="none",
-    )
-    profile_group.add_argument(
-        "--profile-dir",
-        type=str,
-        default="",
-        help="Required when --profile is not 'none'.",
-    )
-
-    return parser
-
-
-def _build_prompt(prefix: str, prompt_size: int) -> str:
-    if prompt_size <= 0:
-        return ""
-    if not prefix:
-        prefix = " "
-    if len(prefix) >= prompt_size:
-        return prefix[:prompt_size]
-    repeat_count = (prompt_size + len(prefix) - 1) // len(prefix)
-    return (prefix * repeat_count)[:prompt_size]
-
-
-def _build_profiler_config(
-    profile: str, profile_dir: str, max_tokens: int
-) -> ProfilerConfig | None:
-    if profile == "none":
-        return None
-    if not profile_dir:
-        raise ValueError("--profile-dir must be set when profiling is enabled.")
-    if profile == "prefill":
-        delay_iterations = 0
-        max_iterations = 1
-    elif profile == "decode":
-        delay_iterations = 1
-        max_iterations = max(1, max_tokens)
-    else:
-        delay_iterations = 0
-        max_iterations = 0
-
-    return ProfilerConfig(
-        profiler="torch",
-        torch_profiler_dir=profile_dir,
-        delay_iterations=delay_iterations,
-        max_iterations=max_iterations,
-    )
-
-
-def main(args: dict) -> None:
-    max_tokens = DEFAULT_MAX_TOKENS
-    batch_size = args.pop("batch_size")
-    prompt_size = args.pop("prompt_size")
-    prompt_prefix = args.pop("prompt_prefix")
-    profile = args.pop("profile")
-    profile_dir = args.pop("profile_dir")
-
-    profiler_config = _build_profiler_config(profile, profile_dir, max_tokens)
-    if profiler_config is not None:
-        args["profiler_config"] = profiler_config
-
-    llm = LLM(**args)
-
-    sampling_params = llm.get_default_sampling_params()
-    sampling_params.max_tokens = max_tokens
-    sampling_params.min_tokens = max_tokens
-    sampling_params.ignore_eos = True
-
-    prompt = _build_prompt(prompt_prefix, prompt_size)
-    prompts = [prompt] * batch_size
-
-    if profile != "none":
-        llm.start_profile()
-    outputs = llm.generate(prompts, sampling_params)
-    if profile != "none":
-        llm.stop_profile()
-
-    print("-" * 50)
-    for output in outputs:
-        generated_text = output.outputs[0].text
-        print(f"Prompt: {output.prompt!r}\nGenerated text: {generated_text!r}")
-        print("-" * 50)
-
-
-if __name__ == "__main__":
-    parser = create_parser()
-    main(vars(parser.parse_args()))
@@ -5,9 +5,14 @@ from transformers import AutoTokenizer

 from vllm import LLM, SamplingParams
 from vllm.benchmarks.datasets import add_dataset_parser, get_samples
-from vllm.utils.argparse_utils import FlexibleArgumentParser
 from vllm.v1.metrics.reader import Counter, Vector

+try:
+    from vllm.utils.argparse_utils import FlexibleArgumentParser
+except ImportError:
+    from argparse import ArgumentParser as FlexibleArgumentParser
+
+
 QUESTION = "What is the content of each image?"
 IMAGE_URLS = [
    "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/duck.jpg",
@@ -1,135 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Test pause/resume with Data Parallel (DP) via HTTP API.
-
-This example demonstrates coordinated pause/resume across multiple DP ranks.
-The pause synchronizes across all DP engines via all-reduce.
-
-Prerequisites:
-    Start a vLLM server with data parallelism:
-
-    $ VLLM_SERVER_DEV_MODE=1 vllm serve facebook/opt-125m \
-        --enforce-eager \
-        --data-parallel-size 4 \
-        --tensor-parallel-size 1
-
-    Then run this script:
-
-    $ python data_parallel_pause_resume.py
-
-The test verifies pause works by:
-1. Starting a streaming generation request
-2. Pausing the server mid-generation
-3. Sleeping for PAUSE_DURATION seconds
-4. Resuming the server
-5. Verifying there was a gap in token generation matching the pause duration
-"""
-
-import argparse
-import threading
-import time
-
-import requests
-from openai import OpenAI
-
-BASE_URL = "http://localhost:8000"
-MODEL_NAME = "facebook/opt-125m"
-PAUSE_DURATION = 3.0
-
-
-def pause_generation(base_url: str, mode: str = "keep") -> None:
-    """Pause generation via HTTP endpoint."""
-    url = f"{base_url}/pause"
-    response = requests.post(url, params={"mode": mode}, timeout=60)
-    response.raise_for_status()
-    print("Server paused")
-
-
-def resume_generation(base_url: str) -> None:
-    """Resume generation via HTTP endpoint."""
-    url = f"{base_url}/resume"
-    response = requests.post(url, timeout=60)
-    response.raise_for_status()
-    print("Server resumed")
-
-
-def main():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--base-url", default=BASE_URL)
-    parser.add_argument("--model", default=MODEL_NAME)
-    args = parser.parse_args()
-
-    client = OpenAI(
-        base_url=f"{args.base_url}/v1",
-        api_key="EMPTY",
-    )
-
-    prompt = "Write a long story about a dragon. Once upon a time"
-    token_times: list[float] = []
-    pause_token_idx = 0
-    pause_triggered = threading.Event()
-
-    def generator_thread():
-        """Stream tokens and record timestamps."""
-        stream = client.completions.create(
-            model=args.model,
-            prompt=prompt,
-            max_tokens=50,
-            stream=True,
-        )
-        for chunk in stream:
-            if chunk.choices[0].text:
-                token_times.append(time.monotonic())
-                token_count = len(token_times)
-                print(f"Token {token_count}: {chunk.choices[0].text!r}")
-
-                # Signal controller after some tokens
-                if token_count >= 5 and not pause_triggered.is_set():
-                    pause_triggered.set()
-
-    def controller_thread():
-        """Pause and resume the server."""
-        nonlocal pause_token_idx
-
-        # Wait for some tokens
-        pause_triggered.wait()
-
-        print(f"\nPausing server (keep mode) at token {len(token_times)}...")
-        pause_generation(args.base_url, mode="keep")
-        pause_token_idx = len(token_times)
-        print(f"Sleeping for {PAUSE_DURATION}s...")
-
-        time.sleep(PAUSE_DURATION)
-
-        print("Resuming server...")
-        resume_generation(args.base_url)
-        print("Resumed!\n")
-
-    # Run both threads
-    gen_thread = threading.Thread(target=generator_thread)
-    ctrl_thread = threading.Thread(target=controller_thread)
-
-    gen_thread.start()
-    ctrl_thread.start()
-
-    gen_thread.join()
-    ctrl_thread.join()
-
-    # Check gap at the pause point
-    if pause_token_idx < len(token_times):
-        pause_gap = token_times[pause_token_idx] - token_times[pause_token_idx - 1]
-        print(
-            f"\nGap after pause (token {pause_token_idx} -> "
-            f"{pause_token_idx + 1}): {pause_gap:.3f}s"
-        )
-        if pause_gap >= PAUSE_DURATION * 0.9:
-            print("Test passed! Pause synchronized across DP ranks.")
-        else:
-            print(f"Test failed! Expected ~{PAUSE_DURATION}s gap, got {pause_gap:.3f}s")
-    else:
-        print("Test failed! No tokens were generated after resuming.")
-
-
-if __name__ == "__main__":
-    main()
@@ -10,7 +10,7 @@ vllm serve llava-hf/llava-1.5-7b-hf

 (multi-image inference with Phi-3.5-vision-instruct)
 vllm serve microsoft/Phi-3.5-vision-instruct --runner generate \
-    --trust-remote-code --max-model-len 4096 --limit-mm-per-prompt.image 2
+    --trust-remote-code --max-model-len 4096 --limit-mm-per-prompt '{"image":2}'

 (audio inference with Ultravox)
 vllm serve fixie-ai/ultravox-v0_5-llama-3_2-1b \
@@ -26,9 +26,7 @@ from openai import AsyncOpenAI, OpenAI
 from vllm.assets.audio import AudioAsset


-def sync_openai(
-    audio_path: str, client: OpenAI, model: str, *, repetition_penalty: float = 1.3
-):
+def sync_openai(audio_path: str, client: OpenAI, model: str):
    """
    Perform synchronous transcription using OpenAI-compatible API.
    """
@@ -42,7 +40,7 @@ def sync_openai(
            # Additional sampling params not provided by OpenAI API.
            extra_body=dict(
                seed=4419,
-                repetition_penalty=repetition_penalty,
+                repetition_penalty=1.3,
            ),
        )
        print("transcription result [sync]:", transcription.text)
@@ -131,12 +129,7 @@ def main(args):
    print(f"Using model: {model}")

    # Run the synchronous function
-    sync_openai(
-        audio_path=args.audio_path if args.audio_path else mary_had_lamb,
-        client=client,
-        model=model,
-        repetition_penalty=args.repetition_penalty,
-    )
+    sync_openai(args.audio_path if args.audio_path else mary_had_lamb, client, model)

    # Run the asynchronous function
    if "openai" in model:
@@ -168,11 +161,5 @@ if __name__ == "__main__":
        default=None,
        help="The path to the audio file to transcribe.",
    )
-    parser.add_argument(
-        "--repetition_penalty",
-        type=float,
-        default=1.3,
-        help="repetition penalty",
-    )
    args = parser.parse_args()
    main(args)
@@ -7,7 +7,7 @@ NOTE:
    vllm serve muziyongshixin/Qwen2.5-VL-7B-for-VideoCls \
         --runner pooling \
         --max-model-len 5000 \
-         --limit-mm-per-prompt.video 1 \
+         --limit-mm-per-prompt '{"video": 1}' \
         --hf-overrides '{"text_config": {"architectures": ["Qwen2_5_VLForSequenceClassification"]}}'
 """

@@ -28,7 +28,7 @@ def main():
    )

    llm = LLM(
-        model="ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11",
+        model="christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM",
        skip_tokenizer_init=True,
        trust_remote_code=True,
        enforce_eager=True,
@@ -391,7 +391,7 @@ if __name__ == "__main__":
    parser.add_argument(
        "--model",
        type=str,
-        default="ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11",
+        default="christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM",
        help="Path to a checkpoint file to load from.",
    )
    parser.add_argument(
@@ -14,7 +14,9 @@ import requests
 # - install TerraTorch v1.1 (or later):
 #   pip install terratorch>=v1.1
 # - start vllm in serving mode with the below args
-#   --model='ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11'
+#   --model='christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM'
+#   --model-impl terratorch
+#   --trust-remote-code
 #   --skip-tokenizer-init --enforce-eager
 #   --io-processor-plugin terratorch_segmentation
 #   --enable-mm-embeds
@@ -32,7 +34,7 @@ def main():
            "out_data_format": "b64_json",
        },
        "priority": 0,
-        "model": "ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11",
+        "model": "christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM",
    }

    ret = requests.post(server_endpoint, json=request_payload_url)
@@ -1,27 +1,15 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """
-Example of using ColBERT late interaction models for reranking and scoring.
+Example of using ColBERT late interaction model for reranking.

 ColBERT (Contextualized Late Interaction over BERT) uses per-token embeddings
 and MaxSim scoring for document reranking, providing better accuracy than
 single-vector models while being more efficient than cross-encoders.

-vLLM supports ColBERT with multiple encoder backbones. Start the server
-with one of the following:
-
-    # BERT backbone (works out of the box)
+Start the server with:
    vllm serve answerdotai/answerai-colbert-small-v1

-    # ModernBERT backbone
-    vllm serve lightonai/GTE-ModernColBERT-v1 \
-        --hf-overrides '{"architectures": ["ColBERTModernBertModel"]}'
-
-    # Jina XLM-RoBERTa backbone
-    vllm serve jinaai/jina-colbert-v2 \
-        --hf-overrides '{"architectures": ["ColBERTJinaRobertaModel"]}' \
-        --trust-remote-code
-
 Then run this script:
    python colbert_rerank_online.py
 """
@@ -30,62 +18,39 @@ import json

 import requests

-# Change this to match the model you started the server with
-MODEL = "answerdotai/answerai-colbert-small-v1"
-BASE_URL = "http://127.0.0.1:8000"
+url = "http://127.0.0.1:8000/rerank"

 headers = {"accept": "application/json", "Content-Type": "application/json"}

-documents = [
-    "Machine learning is a subset of artificial intelligence.",
-    "Python is a programming language.",
-    "Deep learning uses neural networks for complex tasks.",
-    "The weather today is sunny.",
-]
-
-
-def rerank_example():
-    """Use the /rerank endpoint to rank documents by query relevance."""
-    print("=== Rerank Example ===")
-
-    data = {
-        "model": MODEL,
-        "query": "What is machine learning?",
-        "documents": documents,
-    }
-
-    response = requests.post(f"{BASE_URL}/rerank", headers=headers, json=data)
-    result = response.json()
-    print(json.dumps(result, indent=2))
-
-    print("\nRanked documents (most relevant first):")
-    for item in result["results"]:
-        doc_idx = item["index"]
-        score = item["relevance_score"]
-        print(f"  Score {score:.4f}: {documents[doc_idx]}")
-
-
-def score_example():
-    """Use the /score endpoint for pairwise query-document scoring."""
-    print("\n=== Score Example ===")
-
-    data = {
-        "model": MODEL,
-        "text_1": "What is machine learning?",
-        "text_2": [
-            "Machine learning is a subset of AI.",
-            "The weather is sunny.",
-        ],
-    }
-
-    response = requests.post(f"{BASE_URL}/score", headers=headers, json=data)
-    result = response.json()
-    print(json.dumps(result, indent=2))
+data = {
+    "model": "answerdotai/answerai-colbert-small-v1",
+    "query": "What is machine learning?",
+    "documents": [
+        "Machine learning is a subset of artificial intelligence.",
+        "Python is a programming language.",
+        "Deep learning uses neural networks for complex tasks.",
+        "The weather today is sunny.",
+    ],
+}


 def main():
-    rerank_example()
-    score_example()
+    response = requests.post(url, headers=headers, json=data)
+
+    if response.status_code == 200:
+        print("ColBERT Rerank Request successful!")
+        result = response.json()
+        print(json.dumps(result, indent=2))
+
+        # Show ranked results
+        print("\nRanked documents (most relevant first):")
+        for item in result["results"]:
+            doc_idx = item["index"]
+            score = item["relevance_score"]
+            print(f"  Score {score:.4f}: {data['documents'][doc_idx]}")
+    else:
+        print(f"Request failed with status code: {response.status_code}")
+        print(response.text)


 if __name__ == "__main__":
@@ -1,130 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Example of using ColQwen3 late interaction model for reranking.
-
-ColQwen3 is a multi-modal ColBERT-style model based on Qwen3-VL.
-It produces per-token embeddings and uses MaxSim scoring for retrieval
-and reranking. Supports both text and image inputs.
-
-Start the server with:
-    vllm serve TomoroAI/tomoro-colqwen3-embed-4b --max-model-len 50000
-
-Then run this script:
-    python colqwen3_rerank_online.py
-"""
-
-import requests
-
-MODEL = "TomoroAI/tomoro-colqwen3-embed-4b"
-BASE_URL = "http://127.0.0.1:8000"
-
-headers = {"accept": "application/json", "Content-Type": "application/json"}
-
-
-def rerank_text():
-    """Text-only reranking via /rerank endpoint."""
-    print("=" * 60)
-    print("1. Text reranking (/rerank)")
-    print("=" * 60)
-
-    data = {
-        "model": MODEL,
-        "query": "What is machine learning?",
-        "documents": [
-            "Machine learning is a subset of artificial intelligence.",
-            "Python is a programming language.",
-            "Deep learning uses neural networks for complex tasks.",
-            "The weather today is sunny.",
-        ],
-    }
-
-    response = requests.post(f"{BASE_URL}/rerank", headers=headers, json=data)
-
-    if response.status_code == 200:
-        result = response.json()
-        print("\n  Ranked documents (most relevant first):")
-        for item in result["results"]:
-            doc_idx = item["index"]
-            score = item["relevance_score"]
-            print(f"    [{score:.4f}] {data['documents'][doc_idx]}")
-    else:
-        print(f"  Request failed: {response.status_code}")
-        print(f"  {response.text[:300]}")
-
-
-def score_text():
-    """Text-only scoring via /score endpoint."""
-    print()
-    print("=" * 60)
-    print("2. Text scoring (/score)")
-    print("=" * 60)
-
-    query = "What is the capital of France?"
-    documents = [
-        "The capital of France is Paris.",
-        "Berlin is the capital of Germany.",
-        "Python is a programming language.",
-    ]
-
-    data = {
-        "model": MODEL,
-        "text_1": query,
-        "text_2": documents,
-    }
-
-    response = requests.post(f"{BASE_URL}/score", headers=headers, json=data)
-
-    if response.status_code == 200:
-        result = response.json()
-        print(f"\n  Query: {query}\n")
-        for item in result["data"]:
-            idx = item["index"]
-            score = item["score"]
-            print(f"    Doc {idx} (score={score:.4f}): {documents[idx]}")
-    else:
-        print(f"  Request failed: {response.status_code}")
-        print(f"  {response.text[:300]}")
-
-
-def score_text_top_n():
-    """Text reranking with top_n filtering via /rerank endpoint."""
-    print()
-    print("=" * 60)
-    print("3. Text reranking with top_n=2 (/rerank)")
-    print("=" * 60)
-
-    data = {
-        "model": MODEL,
-        "query": "What is the capital of France?",
-        "documents": [
-            "The capital of France is Paris.",
-            "Berlin is the capital of Germany.",
-            "Python is a programming language.",
-            "The Eiffel Tower is in Paris.",
-        ],
-        "top_n": 2,
-    }
-
-    response = requests.post(f"{BASE_URL}/rerank", headers=headers, json=data)
-
-    if response.status_code == 200:
-        result = response.json()
-        print(f"\n  Top {data['top_n']} results:")
-        for item in result["results"]:
-            doc_idx = item["index"]
-            score = item["relevance_score"]
-            print(f"    [{score:.4f}] {data['documents'][doc_idx]}")
-    else:
-        print(f"  Request failed: {response.status_code}")
-        print(f"  {response.text[:300]}")
-
-
-def main():
-    rerank_text()
-    score_text()
-    score_text_top_n()
-
-
-if __name__ == "__main__":
-    main()
@@ -1,198 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-# ruff: noqa: E501
-
-"""
-Example online usage of Pooling API for ColQwen3 multi-vector retrieval.
-
-ColQwen3 is a multi-modal late interaction model based on Qwen3-VL that
-produces per-token embeddings (320-dim, L2-normalized) for both text and
-image inputs. Similarity is computed via MaxSim scoring.
-
-This example mirrors the official TomoroAI inference code
-(https://huggingface.co/TomoroAI/tomoro-colqwen3-embed-4b) but uses the
-vLLM serving API instead of local HuggingFace model loading.
-
-Start the server with:
-    vllm serve TomoroAI/tomoro-colqwen3-embed-4b --max-model-len 4096
-
-Then run this script:
-    python colqwen3_token_embed_online.py
-"""
-
-import argparse
-import base64
-from io import BytesIO
-
-import numpy as np
-import requests
-from PIL import Image
-
-# ── Helpers ─────────────────────────────────────────────────
-
-
-def post_http_request(payload: dict, api_url: str) -> requests.Response:
-    headers = {"User-Agent": "Test Client"}
-    return requests.post(api_url, headers=headers, json=payload)
-
-
-def load_image(url: str) -> Image.Image:
-    """Download an image from URL (handles Wikimedia 403)."""
-    for hdrs in ({}, {"User-Agent": "Mozilla/5.0 (compatible; ColQwen3-demo/1.0)"}):
-        resp = requests.get(url, headers=hdrs, timeout=10)
-        if resp.status_code == 403:
-            continue
-        resp.raise_for_status()
-        return Image.open(BytesIO(resp.content)).convert("RGB")
-    raise RuntimeError(f"Could not fetch image from {url}")
-
-
-def encode_image_base64(image: Image.Image) -> str:
-    """Encode a PIL image to a base64 data URI."""
-    buf = BytesIO()
-    image.save(buf, format="PNG")
-    return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()
-
-
-def compute_maxsim(q_emb: np.ndarray, d_emb: np.ndarray) -> float:
-    """Compute ColBERT-style MaxSim score between query and document."""
-    sim = q_emb @ d_emb.T
-    return float(sim.max(axis=-1).sum())
-
-
-# ── Encode functions ────────────────────────────────────────
-
-
-def encode_queries(texts: list[str], model: str, api_url: str) -> list[np.ndarray]:
-    """Encode text queries → list of multi-vector embeddings."""
-    resp = post_http_request({"model": model, "input": texts}, api_url)
-    return [np.array(item["data"]) for item in resp.json()["data"]]
-
-
-def encode_images(image_urls: list[str], model: str, api_url: str) -> list[np.ndarray]:
-    """Encode image documents → list of multi-vector embeddings.
-
-    Images are sent via the chat-style `messages` field so that the
-    vLLM multimodal processor handles them correctly.
-    """
-    embeddings = []
-    for url in image_urls:
-        print(f"  Loading: {url.split('/')[-1]}...")
-        image = load_image(url)
-        image_uri = encode_image_base64(image)
-        resp = post_http_request(
-            {
-                "model": model,
-                "messages": [
-                    {
-                        "role": "user",
-                        "content": [
-                            {"type": "image_url", "image_url": {"url": image_uri}},
-                            {"type": "text", "text": "Describe the image."},
-                        ],
-                    }
-                ],
-            },
-            api_url,
-        )
-        result = resp.json()
-        if resp.status_code != 200 or "data" not in result:
-            print(f"    Error ({resp.status_code}): {str(result)[:200]}")
-            continue
-        embeddings.append(np.array(result["data"][0]["data"]))
-    return embeddings
-
-
-# ── Main ────────────────────────────────────────────────────
-
-
-def parse_args():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--host", type=str, default="localhost")
-    parser.add_argument("--port", type=int, default=8000)
-    parser.add_argument(
-        "--model",
-        type=str,
-        default="TomoroAI/tomoro-colqwen3-embed-4b",
-    )
-    return parser.parse_args()
-
-
-def main(args):
-    pooling_url = f"http://{args.host}:{args.port}/pooling"
-    score_url = f"http://{args.host}:{args.port}/score"
-    model = args.model
-
-    # Same sample data as the official TomoroAI example
-    queries = [
-        "Retrieve the city of Singapore",
-        "Retrieve the city of Beijing",
-        "Retrieve the city of London",
-    ]
-    image_urls = [
-        "https://upload.wikimedia.org/wikipedia/commons/2/27/Singapore_skyline_2022.jpg",
-        "https://upload.wikimedia.org/wikipedia/commons/6/61/Beijing_skyline_at_night.JPG",
-        "https://upload.wikimedia.org/wikipedia/commons/4/49/London_skyline.jpg",
-    ]
-
-    # ── 1) Text query embeddings ────────────────────────────
-    print("=" * 60)
-    print("1. Encode text queries (multi-vector)")
-    print("=" * 60)
-    query_embeddings = encode_queries(queries, model, pooling_url)
-    for i, emb in enumerate(query_embeddings):
-        norm = float(np.linalg.norm(emb[0]))
-        print(f'  Query {i}: {emb.shape}  (L2 norm: {norm:.4f})  "{queries[i]}"')
-
-    # ── 2) Image document embeddings ────────────────────────
-    print()
-    print("=" * 60)
-    print("2. Encode image documents (multi-vector)")
-    print("=" * 60)
-    doc_embeddings = encode_images(image_urls, model, pooling_url)
-    for i, emb in enumerate(doc_embeddings):
-        print(f"  Doc {i}:   {emb.shape}  {image_urls[i].split('/')[-1]}")
-
-    # ── 3) Cross-modal MaxSim scoring ───────────────────────
-    if doc_embeddings:
-        print()
-        print("=" * 60)
-        print("3. Cross-modal MaxSim scores (text queries × image docs)")
-        print("=" * 60)
-        # Header
-        print(f"{'':>35s}", end="")
-        for j in range(len(doc_embeddings)):
-            print(f"  Doc {j:>2d}", end="")
-        print()
-        # Score matrix
-        for i, q_emb in enumerate(query_embeddings):
-            print(f"  {queries[i]:<33s}", end="")
-            for j, d_emb in enumerate(doc_embeddings):
-                score = compute_maxsim(q_emb, d_emb)
-                print(f"  {score:6.2f}", end="")
-            print()
-
-    # ── 4) Text-only /score endpoint ────────────────────────
-    print()
-    print("=" * 60)
-    print("4. Text-only late interaction scoring (/score endpoint)")
-    print("=" * 60)
-    text_query = "What is the capital of France?"
-    text_docs = [
-        "The capital of France is Paris.",
-        "Berlin is the capital of Germany.",
-        "Python is a programming language.",
-    ]
-    resp = post_http_request(
-        {"model": model, "text_1": text_query, "text_2": text_docs},
-        score_url,
-    )
-    print(f'  Query: "{text_query}"\n')
-    for item in resp.json()["data"]:
-        idx = item["index"]
-        print(f"  Doc {idx} (score={item['score']:.4f}): {text_docs[idx]}")
-
-
-if __name__ == "__main__":
-    args = parse_args()
-    main(args)
@@ -63,9 +63,8 @@ plugins:
  - git-revision-date-localized:
      # exclude autogenerated files
      exclude:
-        - api/*
+        - argparse/*
        - examples/*
-        - generated/*
  - minify:
      minify_html: true
      minify_js: true
@@ -93,6 +92,7 @@ plugins:
              - "!.*_pb2_grpc"  # Exclude auto-generated gRPC stubs
            summary:
              modules: true
+            show_if_no_docstring: true
            show_signature_annotations: true
            separate_signature: true
            show_overloads: true
@@ -31,7 +31,7 @@ partial-json-parser # used for parsing partial JSON outputs
 pyzmq >= 25.0.0
 msgspec
 gguf >= 0.17.0
-mistral_common[image] >= 1.9.1
+mistral_common[image] >= 1.9.0
 opencv-python-headless >= 4.13.0    # required for video IO
 pyyaml
 six>=1.16.0; python_version > '3.11' # transitive dependency of pandas that needs to be the latest version for python 3.12
@@ -52,4 +52,4 @@ anthropic >= 0.71.0
 model-hosting-container-standards >= 0.1.13, < 1.0.0
 mcp
 grpcio
-grpcio-reflection
+grpcio-reflection
@@ -23,7 +23,7 @@ jiwer # required for audio tests
 timm # required for internvl test
 transformers_stream_generator # required for qwen-vl test
 matplotlib # required for qwen-vl test
-mistral_common[image,audio] >= 1.9.1 # required for voxtral test
+mistral_common[image,audio] >= 1.9.0 # required for voxtral test
 num2words # required for smolvlm test
 opencv-python-headless >= 4.13.0 # required for video test
 datamodel_code_generator # required for minicpm3 test
@@ -43,5 +43,5 @@ tritonclient>=2.51.0
 numba == 0.61.2 # Required for N-gram speculative decoding
 numpy
 runai-model-streamer[s3,gcs]==0.15.3
-fastsafetensors>=0.2.2
+fastsafetensors>=0.1.10
 pydantic>=2.12 # 2.11 leads to error on python 3.13
@@ -96,5 +96,3 @@ albumentations==1.4.6
 transformers==4.57.3
 # Pin HF Hub version
 huggingface-hub==0.36.2
-# Pin Mistral Common
-mistral-common[image,audio]==1.9.1
@@ -30,7 +30,7 @@ torchaudio==2.10.0
 torchvision==0.25.0
 transformers_stream_generator # required for qwen-vl test
 matplotlib # required for qwen-vl test
-mistral_common[image,audio] >= 1.9.1 # required for voxtral test
+mistral_common[image,audio] >= 1.9.0 # required for voxtral test
 num2words # required for smolvlm test
 open_clip_torch==2.32.0 # Required for nemotron_vl test, Nemotron Parse in test_common.py
 opencv-python-headless >= 4.13.0 # required for video test
@@ -53,17 +53,10 @@ arctic-inference == 0.1.1 # Required for suffix decoding test
 numba == 0.61.2 # Required for N-gram speculative decoding
 numpy
 runai-model-streamer[s3,gcs]==0.15.3
-fastsafetensors>=0.2.2 # 0.2.2 contains important fixes for multi-GPU mem usage
+fastsafetensors>=0.1.10
 pydantic>=2.12 # 2.11 leads to error on python 3.13
 decord==0.6.0
-terratorch >= 1.2.2 # Required for Prithvi tests
-imagehash # Required for Prithvi tests
-segmentation-models-pytorch > 0.4.0 # Required for Prithvi tests
-
+terratorch @ git+https://github.com/IBM/terratorch.git@1.1.rc3 # required for PrithviMAE test
 gpt-oss >= 0.0.7; python_version > '3.11'

 perceptron # required for isaac 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
@@ -1,9 +1,7 @@
 # This file was autogenerated by uv via the following command:
 #    uv pip compile requirements/test.in -o requirements/test.txt --index-strategy unsafe-best-match --torch-backend cu129 --python-platform x86_64-manylinux_2_28 --python-version 3.12
 absl-py==2.1.0
-    # via
-    #   rouge-score
-    #   tensorboard
+    # via rouge-score
 accelerate==1.0.1
    # via
    #   lm-eval
@@ -33,7 +31,9 @@ albumentations==1.4.6
    #   -r requirements/test.in
    #   terratorch
 alembic==1.16.4
-    # via optuna
+    # via
+    #   mlflow
+    #   optuna
 annotated-doc==0.0.4
    # via fastapi
 annotated-types==0.7.0
@@ -74,6 +74,8 @@ bitsandbytes==0.46.1
    #   lightning
 black==24.10.0
    # via datamodel-code-generator
+blinker==1.9.0
+    # via flask
 blobfile==3.0.0
    # via -r requirements/test.in
 bm25s==0.2.13
@@ -91,7 +93,9 @@ bounded-pool-executor==0.0.3
 buildkite-test-collector==0.1.9
    # via -r requirements/test.in
 cachetools==5.5.2
-    # via google-auth
+    # via
+    #   google-auth
+    #   mlflow-skinny
 certifi==2024.8.30
    # via
    #   fiona
@@ -102,7 +106,6 @@ certifi==2024.8.30
    #   pyproj
    #   rasterio
    #   requests
-    #   sentry-sdk
 cffi==1.17.1
    # via soundfile
 chardet==5.2.0
@@ -117,14 +120,15 @@ click==8.1.7
    #   click-plugins
    #   cligj
    #   fiona
+    #   flask
    #   jiwer
+    #   mlflow-skinny
    #   nltk
    #   rasterio
    #   ray
    #   schemathesis
    #   typer
    #   uvicorn
-    #   wandb
 click-plugins==1.1.1.2
    # via
    #   fiona
@@ -133,6 +137,8 @@ cligj==0.7.2
    # via
    #   fiona
    #   rasterio
+cloudpickle==3.1.1
+    # via mlflow-skinny
 colorama==0.4.6
    # via
    #   perceptron
@@ -157,15 +163,16 @@ cupy-cuda12x==13.6.0
    # via ray
 cycler==0.12.1
    # via matplotlib
+databricks-sdk==0.59.0
+    # via mlflow-skinny
 datamodel-code-generator==0.26.3
    # via -r requirements/test.in
 dataproperty==1.0.1
    # via
    #   pytablewriter
    #   tabledata
-datasets==3.3.0
+datasets==3.0.2
    # via
-    #   -r requirements/test.in
    #   evaluate
    #   lm-eval
    #   mteb
@@ -173,8 +180,6 @@ decorator==5.1.1
    # via librosa
 decord==0.6.0
    # via -r requirements/test.in
-diffusers==0.36.0
-    # via terratorch
 dill==0.3.8
    # via
    #   datasets
@@ -186,11 +191,15 @@ distlib==0.3.9
 dnspython==2.7.0
    # via email-validator
 docker==7.1.0
-    # via gpt-oss
+    # via
+    #   gpt-oss
+    #   mlflow
 docopt==0.6.2
    # via num2words
 docstring-parser==0.17.0
    # via jsonargparse
+efficientnet-pytorch==0.7.1
+    # via segmentation-models-pytorch
 einops==0.8.1
    # via
    #   -r requirements/test.in
@@ -208,18 +217,19 @@ encodec==0.1.1
 evaluate==0.4.3
    # via lm-eval
 fastapi==0.128.0
-    # via gpt-oss
+    # via
+    #   gpt-oss
+    #   mlflow-skinny
 fastparquet==2024.11.0
    # via genai-perf
 fastrlock==0.8.2
    # via cupy-cuda12x
-fastsafetensors==0.2.2
+fastsafetensors==0.1.10
    # via -r requirements/test.in
 filelock==3.16.1
    # via
    #   blobfile
    #   datasets
-    #   diffusers
    #   huggingface-hub
    #   ray
    #   torch
@@ -227,6 +237,8 @@ filelock==3.16.1
    #   virtualenv
 fiona==1.10.1
    # via torchgeo
+flask==3.1.1
+    # via mlflow
 fonttools==4.55.0
    # via matplotlib
 fqdn==1.5.1
@@ -237,7 +249,7 @@ frozenlist==1.5.0
    # via
    #   aiohttp
    #   aiosignal
-fsspec==2024.12.0
+fsspec==2024.9.0
    # via
    #   datasets
    #   evaluate
@@ -245,7 +257,6 @@ fsspec==2024.12.0
    #   huggingface-hub
    #   lightning
    #   pytorch-lightning
-    #   tacoreader
    #   torch
 ftfy==6.3.1
    # via open-clip-torch
@@ -258,7 +269,7 @@ geopandas==1.0.1
 gitdb==4.0.12
    # via gitpython
 gitpython==3.1.44
-    # via wandb
+    # via mlflow-skinny
 google-api-core==2.24.2
    # via
    #   google-cloud-core
@@ -266,6 +277,7 @@ google-api-core==2.24.2
    #   opencensus
 google-auth==2.40.2
    # via
+    #   databricks-sdk
    #   google-api-core
    #   google-cloud-core
    #   google-cloud-storage
@@ -284,17 +296,25 @@ googleapis-common-protos==1.70.0
    # via google-api-core
 gpt-oss==0.0.8
    # via -r requirements/test.in
+graphene==3.4.3
+    # via mlflow
 graphql-core==3.2.6
-    # via hypothesis-graphql
+    # via
+    #   graphene
+    #   graphql-relay
+    #   hypothesis-graphql
+graphql-relay==3.2.0
+    # via graphene
 greenlet==3.2.3
    # via sqlalchemy
 grpcio==1.78.0
    # via
    #   grpcio-tools
    #   ray
-    #   tensorboard
 grpcio-tools==1.78.0
    # via -r requirements/test.in
+gunicorn==23.0.0
+    # via mlflow
 h11==0.14.0
    # via
    #   httpcore
@@ -318,14 +338,12 @@ httpcore==1.0.6
 httpx==0.27.2
    # via
    #   -r requirements/test.in
-    #   diffusers
    #   perceptron
    #   schemathesis
 huggingface-hub==0.36.2
    # via
    #   accelerate
    #   datasets
-    #   diffusers
    #   evaluate
    #   open-clip-torch
    #   peft
@@ -361,13 +379,11 @@ idna==3.10
    #   jsonschema
    #   requests
    #   yarl
-imagehash==4.3.2
-    # via -r requirements/test.in
 imageio==2.37.0
    # via scikit-image
 importlib-metadata==8.7.0
    # via
-    #   diffusers
+    #   mlflow-skinny
    #   opentelemetry-api
 importlib-resources==6.5.2
    # via typeshed-client
@@ -379,10 +395,14 @@ isoduration==20.11.0
    # via jsonschema
 isort==5.13.2
    # via datamodel-code-generator
+itsdangerous==2.2.0
+    # via flask
 jinja2==3.1.6
    # via
    #   datamodel-code-generator
+    #   flask
    #   genai-perf
+    #   mlflow
    #   torch
 jiwer==3.0.5
    # via -r requirements/test.in
@@ -395,14 +415,12 @@ joblib==1.4.2
    #   librosa
    #   nltk
    #   scikit-learn
-jsonargparse==4.46.0
+jsonargparse==4.35.0
    # via
    #   lightning
    #   terratorch
 jsonlines==4.0.0
    # via lm-eval
-jsonnet==0.21.0
-    # via jsonargparse
 jsonpointer==3.0.0
    # via jsonschema
 jsonschema==4.23.0
@@ -431,13 +449,13 @@ libnacl==2.1.0
    # via tensorizer
 librosa==0.10.2.post1
    # via -r requirements/test.in
-lightly==1.5.22
+lightly==1.5.20
    # via
    #   terratorch
    #   torchgeo
 lightly-utils==0.0.2
    # via lightly
-lightning==2.6.1
+lightning==2.5.1.post0
    # via
    #   terratorch
    #   torchgeo
@@ -458,11 +476,12 @@ lxml==5.3.0
 mako==1.3.10
    # via alembic
 markdown==3.8.2
-    # via tensorboard
+    # via mlflow
 markdown-it-py==3.0.0
    # via rich
 markupsafe==3.0.1
    # via
+    #   flask
    #   jinja2
    #   mako
    #   werkzeug
@@ -470,6 +489,7 @@ matplotlib==3.9.2
    # via
    #   -r requirements/test.in
    #   lightning
+    #   mlflow
    #   pycocotools
    #   torchgeo
 mbstrdecoder==1.1.3
@@ -479,8 +499,12 @@ mbstrdecoder==1.1.3
    #   typepy
 mdurl==0.1.2
    # via markdown-it-py
-mistral-common==1.9.1
+mistral-common==1.9.0
    # via -r requirements/test.in
+mlflow==2.22.0
+    # via terratorch
+mlflow-skinny==2.22.0
+    # via mlflow
 more-itertools==10.5.0
    # via lm-eval
 mpmath==1.3.0
@@ -499,6 +523,8 @@ multiprocess==0.70.16
    # via
    #   datasets
    #   evaluate
+munch==4.0.0
+    # via pretrainedmodels
 mypy-extensions==1.0.0
    # via black
 networkx==3.2.1
@@ -527,7 +553,6 @@ numpy==2.2.6
    #   cupy-cuda12x
    #   datasets
    #   decord
-    #   diffusers
    #   einx
    #   encodec
    #   evaluate
@@ -535,13 +560,13 @@ numpy==2.2.6
    #   genai-perf
    #   geopandas
    #   h5py
-    #   imagehash
    #   imageio
    #   librosa
    #   lightly
    #   lightly-utils
    #   matplotlib
    #   mistral-common
+    #   mlflow
    #   mteb
    #   numba
    #   numexpr
@@ -553,7 +578,6 @@ numpy==2.2.6
    #   perceptron
    #   pycocotools
    #   pyogrio
-    #   pywavelets
    #   rasterio
    #   rioxarray
    #   rouge-score
@@ -566,10 +590,8 @@ numpy==2.2.6
    #   shapely
    #   soxr
    #   statsmodels
-    #   tensorboard
    #   tensorboardx
    #   tensorizer
-    #   terratorch
    #   tifffile
    #   torchgeo
    #   torchmetrics
@@ -637,6 +659,7 @@ opencv-python-headless==4.13.0.90
    #   mistral-common
 opentelemetry-api==1.35.0
    # via
+    #   mlflow-skinny
    #   opentelemetry-exporter-prometheus
    #   opentelemetry-sdk
    #   opentelemetry-semantic-conventions
@@ -646,6 +669,7 @@ opentelemetry-proto==1.36.0
    # via ray
 opentelemetry-sdk==1.35.0
    # via
+    #   mlflow-skinny
    #   opentelemetry-exporter-prometheus
    #   ray
 opentelemetry-semantic-conventions==0.56b0
@@ -663,6 +687,7 @@ packaging==24.2
    #   evaluate
    #   fastparquet
    #   geopandas
+    #   gunicorn
    #   huggingface-hub
    #   hydra-core
    #   kornia
@@ -670,6 +695,7 @@ packaging==24.2
    #   lightning
    #   lightning-utilities
    #   matplotlib
+    #   mlflow-skinny
    #   optuna
    #   peft
    #   plotly
@@ -682,12 +708,10 @@ packaging==24.2
    #   rioxarray
    #   scikit-image
    #   statsmodels
-    #   tensorboard
    #   tensorboardx
    #   torchmetrics
    #   transformers
    #   typepy
-    #   wandb
    #   xarray
 pandas==2.2.3
    # via
@@ -696,8 +720,8 @@ pandas==2.2.3
    #   fastparquet
    #   genai-perf
    #   geopandas
+    #   mlflow
    #   statsmodels
-    #   tacoreader
    #   torchgeo
    #   xarray
 pathspec==0.12.1
@@ -716,9 +740,7 @@ perf-analyzer==0.1.0
    # via genai-perf
 pillow==10.4.0
    # via
-    #   diffusers
    #   genai-perf
-    #   imagehash
    #   imageio
    #   lightly-utils
    #   matplotlib
@@ -726,7 +748,6 @@ pillow==10.4.0
    #   perceptron
    #   scikit-image
    #   segmentation-models-pytorch
-    #   tensorboard
    #   torchgeo
    #   torchvision
 platformdirs==4.3.6
@@ -734,7 +755,6 @@ platformdirs==4.3.6
    #   black
    #   pooch
    #   virtualenv
-    #   wandb
 plotly==5.24.1
    # via genai-perf
 pluggy==1.5.0
@@ -749,6 +769,8 @@ portalocker==2.10.1
    # via sacrebleu
 pqdm==0.2.0
    # via -r requirements/test.in
+pretrainedmodels==0.7.4
+    # via segmentation-models-pytorch
 prometheus-client==0.22.0
    # via
    #   opentelemetry-exporter-prometheus
@@ -764,13 +786,12 @@ protobuf==6.33.2
    #   google-api-core
    #   googleapis-common-protos
    #   grpcio-tools
+    #   mlflow-skinny
    #   opentelemetry-proto
    #   proto-plus
    #   ray
-    #   tensorboard
    #   tensorboardx
    #   tensorizer
-    #   wandb
 psutil==6.1.0
    # via
    #   accelerate
@@ -780,12 +801,11 @@ py==1.11.0
    # via pytest-forked
 py-spy==0.4.0
    # via ray
-pyarrow==23.0.0
+pyarrow==18.0.0
    # via
    #   datasets
    #   genai-perf
-    #   tacoreader
-    #   terratorch
+    #   mlflow
 pyasn1==0.6.1
    # via
    #   pyasn1-modules
@@ -811,11 +831,11 @@ pydantic==2.12.0
    #   gpt-oss
    #   lightly
    #   mistral-common
+    #   mlflow-skinny
    #   mteb
    #   openai-harmony
    #   pydantic-extra-types
    #   ray
-    #   wandb
 pydantic-core==2.41.1
    # via pydantic
 pydantic-extra-types==2.10.5
@@ -853,6 +873,7 @@ pytest==8.3.5
    #   pytest-subtests
    #   pytest-timeout
    #   schemathesis
+    #   terratorch
 pytest-asyncio==0.24.0
    # via -r requirements/test.in
 pytest-cov==6.3.0
@@ -875,6 +896,7 @@ python-dateutil==2.9.0.post0
    # via
    #   arrow
    #   botocore
+    #   graphene
    #   lightly
    #   matplotlib
    #   pandas
@@ -891,8 +913,6 @@ pytz==2024.2
    # via
    #   pandas
    #   typepy
-pywavelets==1.9.0
-    # via imagehash
 pyyaml==6.0.2
    # via
    #   accelerate
@@ -903,6 +923,7 @@ pyyaml==6.0.2
    #   huggingface-hub
    #   jsonargparse
    #   lightning
+    #   mlflow-skinny
    #   omegaconf
    #   optuna
    #   peft
@@ -913,7 +934,6 @@ pyyaml==6.0.2
    #   timm
    #   transformers
    #   vocos
-    #   wandb
 rapidfuzz==3.12.1
    # via jiwer
 rasterio==1.4.3
@@ -931,7 +951,6 @@ referencing==0.35.1
    #   jsonschema-specifications
 regex==2024.9.11
    # via
-    #   diffusers
    #   nltk
    #   open-clip-torch
    #   sacrebleu
@@ -940,8 +959,8 @@ regex==2024.9.11
 requests==2.32.3
    # via
    #   buildkite-test-collector
+    #   databricks-sdk
    #   datasets
-    #   diffusers
    #   docker
    #   evaluate
    #   google-api-core
@@ -951,16 +970,15 @@ requests==2.32.3
    #   lightly
    #   lm-eval
    #   mistral-common
+    #   mlflow-skinny
    #   mteb
    #   pooch
    #   ray
    #   responses
    #   schemathesis
    #   starlette-testclient
-    #   tacoreader
    #   tiktoken
    #   transformers
-    #   wandb
 responses==0.25.3
    # via genai-perf
 rfc3339-validator==0.1.4
@@ -973,7 +991,6 @@ rich==13.9.4
    #   lightning
    #   mteb
    #   perceptron
-    #   terratorch
    #   typer
 rioxarray==0.19.0
    # via terratorch
@@ -1000,55 +1017,47 @@ sacrebleu==2.4.3
 safetensors==0.4.5
    # via
    #   accelerate
-    #   diffusers
    #   open-clip-torch
    #   peft
-    #   segmentation-models-pytorch
    #   timm
    #   transformers
 schemathesis==3.39.15
    # via -r requirements/test.in
 scikit-image==0.25.2
-    # via
-    #   albumentations
-    #   terratorch
+    # via albumentations
 scikit-learn==1.5.2
    # via
    #   albumentations
    #   librosa
    #   lm-eval
+    #   mlflow
    #   mteb
    #   sentence-transformers
-    #   terratorch
 scipy==1.13.1
    # via
    #   albumentations
    #   bm25s
-    #   imagehash
    #   librosa
+    #   mlflow
    #   mteb
    #   scikit-image
    #   scikit-learn
    #   sentence-transformers
    #   statsmodels
    #   vocos
-segmentation-models-pytorch==0.5.0
+segmentation-models-pytorch==0.4.0
    # via
-    #   -r requirements/test.in
    #   terratorch
    #   torchgeo
 sentence-transformers==5.2.0
    # via
    #   -r requirements/test.in
    #   mteb
-sentry-sdk==2.52.0
-    # via wandb
 setuptools==77.0.3
    # via
    #   grpcio-tools
    #   lightning-utilities
    #   pytablewriter
-    #   tensorboard
    #   torch
 shapely==2.1.1
    # via
@@ -1066,6 +1075,7 @@ six==1.16.0
    #   python-dateutil
    #   rfc3339-validator
    #   rouge-score
+    #   segmentation-models-pytorch
 smart-open==7.1.0
    # via ray
 smmap==5.0.2
@@ -1089,9 +1099,12 @@ soxr==0.5.0.post1
 sqlalchemy==2.0.41
    # via
    #   alembic
+    #   mlflow
    #   optuna
 sqlitedict==2.1.0
    # via lm-eval
+sqlparse==0.5.3
+    # via mlflow-skinny
 starlette==0.50.0
    # via
    #   fastapi
@@ -1111,8 +1124,6 @@ tabledata==1.3.3
    # via pytablewriter
 tabulate==0.9.0
    # via sacrebleu
-tacoreader==0.5.6
-    # via terratorch
 tblib==3.1.0
    # via -r requirements/test.in
 tcolorpy==0.1.6
@@ -1122,19 +1133,13 @@ tenacity==9.1.2
    #   gpt-oss
    #   lm-eval
    #   plotly
-tensorboard==2.20.0
-    # via terratorch
-tensorboard-data-server==0.7.2
-    # via tensorboard
 tensorboardx==2.6.4
    # via lightning
 tensorizer==2.10.1
    # via -r requirements/test.in
 termcolor==3.1.0
-    # via
-    #   gpt-oss
-    #   terratorch
-terratorch==1.2.2
+    # via gpt-oss
+terratorch @ git+https://github.com/IBM/terratorch.git@07184fcf91a1324f831ff521dd238d97fe350e3e
    # via -r requirements/test.in
 threadpoolctl==3.5.0
    # via scikit-learn
@@ -1167,7 +1172,9 @@ torch==2.10.0+cu129
    #   -r requirements/test.in
    #   accelerate
    #   bitsandbytes
+    #   efficientnet-pytorch
    #   encodec
+    #   fastsafetensors
    #   kornia
    #   lightly
    #   lightning
@@ -1175,6 +1182,7 @@ torch==2.10.0+cu129
    #   mteb
    #   open-clip-torch
    #   peft
+    #   pretrainedmodels
    #   pytorch-lightning
    #   runai-model-streamer
    #   segmentation-models-pytorch
@@ -1206,11 +1214,12 @@ torchvision==0.25.0+cu129
    #   -r requirements/test.in
    #   lightly
    #   open-clip-torch
+    #   pretrainedmodels
    #   segmentation-models-pytorch
    #   terratorch
    #   timm
    #   torchgeo
-tqdm==4.67.3
+tqdm==4.66.6
    # via
    #   datasets
    #   evaluate
@@ -1224,11 +1233,10 @@ tqdm==4.67.3
    #   optuna
    #   peft
    #   pqdm
+    #   pretrainedmodels
    #   pytorch-lightning
    #   segmentation-models-pytorch
    #   sentence-transformers
-    #   tacoreader
-    #   terratorch
    #   tqdm-multiprocess
    #   transformers
 tqdm-multiprocess==0.0.11
@@ -1267,12 +1275,14 @@ typing-extensions==4.15.0
    #   alembic
    #   chz
    #   fastapi
+    #   graphene
    #   grpcio
    #   huggingface-hub
    #   librosa
    #   lightning
    #   lightning-utilities
    #   mistral-common
+    #   mlflow-skinny
    #   mteb
    #   opentelemetry-api
    #   opentelemetry-sdk
@@ -1290,7 +1300,6 @@ typing-extensions==4.15.0
    #   typer
    #   typeshed-client
    #   typing-inspection
-    #   wandb
 typing-inspection==0.4.2
    # via pydantic
 tzdata==2024.2
@@ -1305,26 +1314,25 @@ urllib3==2.2.3
    #   lightly
    #   requests
    #   responses
-    #   sentry-sdk
    #   tritonclient
 uvicorn==0.35.0
-    # via gpt-oss
+    # via
+    #   gpt-oss
+    #   mlflow-skinny
 vector-quantize-pytorch==1.21.2
    # via -r requirements/test.in
 virtualenv==20.31.2
    # via ray
 vocos==0.1.0
    # via -r requirements/test.in
-wandb==0.24.2
-    # via terratorch
 wcwidth==0.2.13
    # via ftfy
 webcolors==24.11.1
    # via jsonschema
 werkzeug==3.1.3
    # via
+    #   flask
    #   schemathesis
-    #   tensorboard
 word2number==1.1
    # via lm-eval
 wrapt==1.17.2
@@ -1,430 +0,0 @@
-#!/usr/bin/env python3
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Autotune registered Helion kernels for optimal configurations.
-
-Usage:
-    # Autotune all registered kernels
-    python scripts/autotune_helion_kernels.py
-
-    # Autotune specific kernel
-    python scripts/autotune_helion_kernels.py --kernels silu_mul_fp8
-
-    # Autotune multiple kernels
-    python scripts/autotune_helion_kernels.py --kernels silu_mul_fp8 rms_norm_fp8
-
-    # Force re-autotuning
-    python scripts/autotune_helion_kernels.py --force
-
-    # List available kernels
-    python scripts/autotune_helion_kernels.py --list
-"""
-
-import argparse
-import sys
-import time
-from dataclasses import dataclass
-
-import torch
-
-try:
-    import helion
-
-    from vllm.kernels.helion import (
-        ConfigManager,
-        get_kernel_by_name,
-        get_registered_kernels,
-    )
-    from vllm.kernels.helion.utils import get_canonical_gpu_name
-    from vllm.logger import init_logger
-    from vllm.utils.import_utils import has_helion
-except ImportError as e:
-    print(f"Error importing vLLM: {e}")
-    print("Please ensure vLLM is installed and in your Python path")
-    sys.exit(1)
-
-logger = init_logger("vllm.scripts.autotune_helion_kernels")
-
-
-@dataclass
-class AutotuneResult:
-    status: str  # "success" | "partial" | "error" | "skipped"
-    successful: int
-    failed: int
-    configs: dict[str, "helion.Config"]
-    message: str = ""
-
-
-def list_kernels() -> None:
-    kernels = get_registered_kernels()
-
-    if not kernels:
-        print("No Helion kernels found in registry.")
-        return
-
-    print("Available Helion kernels:")
-    print("=" * 50)
-
-    for name in sorted(kernels.keys()):
-        print(f"  {name}")
-
-    print(f"\nTotal: {len(kernels)} kernels")
-
-
-def check_requirements() -> bool:
-    if not torch.cuda.is_available():
-        logger.error("CUDA is not available. Helion autotuning requires GPU.")
-        return False
-
-    if not has_helion():
-        logger.error("Helion is not installed. Please install Helion package.")
-        return False
-
-    return True
-
-
-def autotune_kernel(
-    kernel_name: str,
-    platform: str,
-    config_manager: ConfigManager,
-    force: bool = False,
-    autotune_effort: str = "quick",
-) -> AutotuneResult:
-    logger.debug(
-        "Starting autotune for kernel '%s' with effort='%s'",
-        kernel_name,
-        autotune_effort,
-    )
-    kernel_wrapper = get_kernel_by_name(kernel_name)
-    if kernel_wrapper is None:
-        error_msg = f"Kernel '{kernel_name}' not found in registry"
-        logger.error(error_msg)
-        return AutotuneResult(
-            status="error",
-            message=error_msg,
-            successful=0,
-            failed=0,
-            configs={},
-        )
-
-    try:
-        inputs_dict = kernel_wrapper.get_inputs()
-    except NotImplementedError:
-        error_msg = f"Kernel '{kernel_name}' has no input generator registered"
-        logger.error(error_msg)
-        return AutotuneResult(
-            status="error",
-            message=error_msg,
-            successful=0,
-            failed=0,
-            configs={},
-        )
-
-    try:
-        logger.info(
-            "Autotuning kernel '%s' for platform '%s' with %d configs",
-            kernel_name,
-            platform,
-            len(inputs_dict),
-        )
-
-        configs_to_autotune = {}
-        if not force:
-            existing_configs = config_manager.get_platform_configs(
-                kernel_name, platform
-            )
-            for config_key, inputs in inputs_dict.items():
-                if config_key in existing_configs:
-                    logger.debug(
-                        "Config '%s' already exists for platform '%s', skipping",
-                        config_key,
-                        platform,
-                    )
-                else:
-                    configs_to_autotune[config_key] = inputs
-        else:
-            logger.debug("Force mode enabled, will re-autotune all configs")
-            configs_to_autotune = inputs_dict
-
-        if not configs_to_autotune:
-            logger.info(
-                "All configs already exist for kernel '%s' on platform '%s'. "
-                "Use --force to re-autotune.",
-                kernel_name,
-                platform,
-            )
-            return AutotuneResult(
-                status="skipped",
-                message="All configs already exist",
-                successful=0,
-                failed=0,
-                configs={},
-            )
-
-        total_start_time = time.time()
-        autotuned_configs = {}
-        failed_configs = []
-
-        for config_key, inputs in configs_to_autotune.items():
-            logger.info("Autotuning config: %s", config_key)
-            logger.debug(
-                "Input shapes: %s",
-                [getattr(inp, "shape", type(inp).__name__) for inp in inputs],
-            )
-
-            try:
-                config_start_time = time.time()
-                config = kernel_wrapper.run_autotune(inputs, autotune_effort)
-                config_duration = time.time() - config_start_time
-
-                # Save immediately for checkpointing
-                config_manager.save_configs(kernel_name, platform, {config_key: config})
-
-                autotuned_configs[config_key] = config
-                logger.debug("Config details: %s", config)
-
-                logger.info(
-                    "✓ Autotuned and saved config '%s' (%.2fs)",
-                    config_key,
-                    config_duration,
-                )
-
-            except (RuntimeError, ValueError, OSError) as e:
-                logger.exception(
-                    "Failed to autotune config '%s': %s",
-                    config_key,
-                    e,
-                )
-                failed_configs.append(config_key)
-
-        total_duration = time.time() - total_start_time
-        successful = len(autotuned_configs)
-        failed = len(failed_configs)
-
-        logger.info(
-            "Completed autotuning for kernel '%s': %d successful, %d failed (%.2fs)",
-            kernel_name,
-            successful,
-            failed,
-            total_duration,
-        )
-
-        status = "success" if failed == 0 else "partial"
-        return AutotuneResult(
-            status=status,
-            successful=successful,
-            failed=failed,
-            configs=autotuned_configs,
-        )
-
-    except (KeyError, RuntimeError, ValueError, OSError) as e:
-        error_msg = f"Unexpected error: {e}"
-        logger.exception("Failed to autotune kernel '%s': %s", kernel_name, e)
-        return AutotuneResult(
-            status="error",
-            message=error_msg,
-            successful=0,
-            failed=0,
-            configs={},
-        )
-
-
-def summarize_results(results: dict[str, AutotuneResult]) -> bool:
-    logger.info("=" * 50)
-    logger.info("Autotuning Results Summary")
-    logger.info("=" * 50)
-
-    total_successful = 0
-    total_failed = 0
-    success_kernels = []
-    partial_kernels = []
-    error_kernels = []
-    skipped_kernels = []
-
-    for kernel_name, result in results.items():
-        total_successful += result.successful
-        total_failed += result.failed
-
-        if result.status == "success":
-            success_kernels.append(f"{kernel_name} ({result.successful} configs)")
-            logger.info("✓ %s: %d configs successful", kernel_name, result.successful)
-        elif result.status == "partial":
-            partial_kernels.append(
-                f"{kernel_name} ({result.successful} ok, {result.failed} failed)"
-            )
-            logger.warning(
-                "⚠ %s: %d successful, %d failed",
-                kernel_name,
-                result.successful,
-                result.failed,
-            )
-        elif result.status == "error":
-            error_kernels.append(f"{kernel_name}: {result.message or 'Unknown error'}")
-            logger.error("✗ %s: %s", kernel_name, result.message or "Unknown error")
-        elif result.status == "skipped":
-            skipped_kernels.append(f"{kernel_name}: {result.message or 'Skipped'}")
-            logger.info("- %s: %s", kernel_name, result.message or "Skipped")
-
-    logger.info("=" * 50)
-    logger.info(
-        "Summary: %d total configs (%d successful, %d failed)",
-        total_successful + total_failed,
-        total_successful,
-        total_failed,
-    )
-    logger.info(
-        "Kernels: %d success, %d partial, %d error, %d skipped",
-        len(success_kernels),
-        len(partial_kernels),
-        len(error_kernels),
-        len(skipped_kernels),
-    )
-
-    has_failures = bool(error_kernels or partial_kernels)
-
-    if not has_failures:
-        if total_successful > 0:
-            logger.info("All configs autotuned successfully!")
-        else:
-            logger.info("No new configs were generated (all may already exist)")
-
-    return not has_failures
-
-
-def get_kernels_to_autotune(requested_kernels: list[str] | None) -> list[str]:
-    all_kernels = get_registered_kernels()
-    if not all_kernels:
-        logger.error("No Helion kernels found in registry")
-        sys.exit(1)
-
-    if not requested_kernels:
-        return list(all_kernels.keys())
-
-    if len(requested_kernels) != len(set(requested_kernels)):
-        duplicates = [
-            k for k in set(requested_kernels) if requested_kernels.count(k) > 1
-        ]
-        logger.error("Duplicate kernel names in --kernels flag: %s", duplicates)
-        sys.exit(1)
-
-    kernels_to_autotune = []
-    missing_kernels = []
-
-    for kernel_name in requested_kernels:
-        if kernel_name in all_kernels:
-            kernels_to_autotune.append(kernel_name)
-        else:
-            missing_kernels.append(kernel_name)
-
-    if missing_kernels:
-        logger.error("Kernel(s) not found: %s", missing_kernels)
-        logger.error("Available kernels: %s", list(all_kernels.keys()))
-        sys.exit(1)
-
-    return kernels_to_autotune
-
-
-def main():
-    parser = argparse.ArgumentParser(
-        description="Autotune Helion kernels",
-        formatter_class=argparse.RawDescriptionHelpFormatter,
-        epilog=__doc__.split("Usage:")[1] if "Usage:" in __doc__ else "",
-    )
-
-    parser.add_argument(
-        "--kernels",
-        nargs="+",
-        help="Kernel(s) to autotune (default: all kernels)",
-    )
-
-    parser.add_argument(
-        "--config-dir",
-        type=str,
-        help="Config directory for config files (default: vLLM helion configs dir)",
-    )
-
-    parser.add_argument(
-        "--list",
-        action="store_true",
-        help="List available Helion kernels and exit",
-    )
-
-    parser.add_argument(
-        "--force",
-        action="store_true",
-        help=(
-            "Force re-autotuning even if configs already exist for the "
-            "platform and config keys"
-        ),
-    )
-
-    parser.add_argument(
-        "--autotune-effort",
-        type=str,
-        default="quick",
-        help=(
-            "Helion autotune effort level: 'quick' (smaller search) or "
-            "'full' (full search budget) (default: quick)"
-        ),
-    )
-
-    parser.add_argument(
-        "--verbose",
-        action="store_true",
-        help="Enable verbose logging",
-    )
-
-    args = parser.parse_args()
-
-    import logging
-
-    if args.verbose:
-        logging.getLogger("vllm").setLevel(logging.DEBUG)
-        logger.debug("Verbose mode enabled")
-        logger.debug("Arguments: %s", vars(args))
-    else:
-        logging.getLogger("vllm").setLevel(logging.INFO)
-
-    if args.list:
-        list_kernels()
-        return
-
-    if not check_requirements():
-        sys.exit(1)
-
-    platform = get_canonical_gpu_name()
-    logger.info("Detected GPU platform: %s", platform)
-
-    config_manager = (
-        ConfigManager(args.config_dir) if args.config_dir else ConfigManager()
-    )
-
-    try:
-        config_manager.ensure_base_dir_writable()
-    except OSError as e:
-        logger.error("Failed to access config directory: %s", e)
-        sys.exit(1)
-
-    kernels_to_autotune = get_kernels_to_autotune(args.kernels)
-
-    logger.info(
-        "Will autotune %d kernel(s) for platform '%s': %s",
-        len(kernels_to_autotune),
-        platform,
-        kernels_to_autotune,
-    )
-
-    results = {}
-    for kernel_name in kernels_to_autotune:
-        result = autotune_kernel(
-            kernel_name, platform, config_manager, args.force, args.autotune_effort
-        )
-        results[kernel_name] = result
-
-    success = summarize_results(results)
-    sys.exit(0 if success else 1)
-
-
-if __name__ == "__main__":
-    main()
@@ -1035,7 +1035,7 @@ setup(
    extras_require={
        "bench": ["pandas", "matplotlib", "seaborn", "datasets", "scipy"],
        "tensorizer": ["tensorizer==2.10.1"],
-        "fastsafetensors": ["fastsafetensors >= 0.2.2"],
+        "fastsafetensors": ["fastsafetensors >= 0.1.10"],
        "runai": ["runai-model-streamer[s3,gcs] >= 0.15.3"],
        "audio": [
            "librosa",
@@ -1,29 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

-import pytest
-
 from ..utils import compare_two_settings


-@pytest.mark.parametrize("disable_pin_memory", [False, True])
-@pytest.mark.parametrize("disable_uva", [False, True])
-def test_cpu_offload(disable_pin_memory, disable_uva):
-    env_vars = {
-        "VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY": str(int(disable_pin_memory)),
-        "VLLM_WEIGHT_OFFLOADING_DISABLE_UVA": str(int(disable_uva)),
-    }
-
-    args = ["--cpu-offload-gb", "1"]
-
-    # cuda graph only works with UVA offloading
-    if disable_uva:
-        args.append("--enforce-eager")
-
+def test_cpu_offload():
    compare_two_settings(
-        model="hmellor/tiny-random-LlamaForCausalLM",
-        arg1=[],
-        arg2=args,
-        env1=None,
-        env2=env_vars,
+        "hmellor/tiny-random-LlamaForCausalLM", [], ["--cpu-offload-gb", "1"]
    )
@@ -1,76 +1,15 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import subprocess
-import tempfile
-import time
-from pathlib import Path

 import pytest
-import requests
-import urllib3

 from ..utils import RemoteOpenAIServer

 MODEL_NAME = "meta-llama/Llama-3.2-1B-Instruct"


-def generate_self_signed_cert(cert_dir: Path) -> tuple[Path, Path]:
-    """Generate a self-signed certificate for testing."""
-    cert_file = cert_dir / "cert.pem"
-    key_file = cert_dir / "key.pem"
-
-    # Generate self-signed certificate using openssl
-    subprocess.run(
-        [
-            "openssl",
-            "req",
-            "-x509",
-            "-newkey",
-            "rsa:2048",
-            "-keyout",
-            str(key_file),
-            "-out",
-            str(cert_file),
-            "-days",
-            "1",
-            "-nodes",
-            "-subj",
-            "/CN=localhost",
-        ],
-        check=True,
-        capture_output=True,
-    )
-    return cert_file, key_file
-
-
-class RemoteOpenAIServerSSL(RemoteOpenAIServer):
-    """RemoteOpenAIServer subclass that supports SSL with self-signed certs."""
-
-    @property
-    def url_root(self) -> str:
-        return f"https://{self.host}:{self.port}"
-
-    def _wait_for_server(self, *, url: str, timeout: float):
-        """Override to use HTTPS with SSL verification disabled."""
-        # Suppress InsecureRequestWarning for self-signed certs
-        urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
-
-        start = time.time()
-        while True:
-            try:
-                if requests.get(url, verify=False).status_code == 200:
-                    break
-            except Exception:
-                result = self._poll()
-                if result is not None and result != 0:
-                    raise RuntimeError("Server exited unexpectedly.") from None
-
-                time.sleep(0.5)
-                if time.time() - start > timeout:
-                    raise RuntimeError("Server failed to start in time.") from None
-
-
-@pytest.fixture(scope="function")
+@pytest.fixture(scope="module")
 def server():
    args = ["--max-model-len", "1024", "--enforce-eager", "--load-format", "dummy"]

@@ -78,27 +17,6 @@ def server():
        yield remote_server


-@pytest.fixture(scope="function")
-def ssl_server():
-    """Start a vLLM server with SSL enabled using a self-signed certificate."""
-    with tempfile.TemporaryDirectory() as cert_dir:
-        cert_file, key_file = generate_self_signed_cert(Path(cert_dir))
-        args = [
-            "--max-model-len",
-            "1024",
-            "--enforce-eager",
-            "--load-format",
-            "dummy",
-            "--ssl-certfile",
-            str(cert_file),
-            "--ssl-keyfile",
-            str(key_file),
-        ]
-
-        with RemoteOpenAIServerSSL(MODEL_NAME, args) as remote_server:
-            yield remote_server
-
-
@pytest.mark.benchmark
 def test_bench_serve(server):
    # Test default model detection and input/output len
@@ -124,31 +42,6 @@ def test_bench_serve(server):
    assert result.returncode == 0, f"Benchmark failed: {result.stderr}"


-@pytest.mark.benchmark
-def test_bench_serve_insecure(ssl_server):
-    """Test --insecure flag with an HTTPS server using a self-signed certificate."""
-    base_url = f"https://{ssl_server.host}:{ssl_server.port}"
-    command = [
-        "vllm",
-        "bench",
-        "serve",
-        "--base-url",
-        base_url,
-        "--input-len",
-        "32",
-        "--output-len",
-        "4",
-        "--num-prompts",
-        "5",
-        "--insecure",
-    ]
-    result = subprocess.run(command, capture_output=True, text=True)
-    print(result.stdout)
-    print(result.stderr)
-
-    assert result.returncode == 0, f"Benchmark failed: {result.stderr}"
-
-
@pytest.mark.benchmark
 def test_bench_serve_chat(server):
    command = [
@@ -27,29 +27,10 @@ from ...utils import create_new_process_for_each_test
 from ..silly_attention import get_global_counter, reset_global_counter


-# Custom op that returns an unbacked symint during graph capture
-@torch.library.custom_op("mylib::foo", mutates_args=())
-def foo(x: torch.Tensor) -> int:
-    return 3
-
-
-@foo.register_fake
-def _(x):
-    return torch.library.get_ctx().new_dynamic_size()
-
-
@support_torch_compile
 class SillyModel(nn.Module):
-    def __init__(
-        self,
-        *,
-        vllm_config: VllmConfig,
-        prefix: str = "",
-        intermediate_unbacked=False,
-        **kwargs,
-    ) -> None:
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "", **kwargs) -> None:
        super().__init__()
-        self.intermediate_unbacked = intermediate_unbacked

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
@@ -63,13 +44,6 @@ class SillyModel(nn.Module):
        torch.ops.silly.attention(x, x, x, out)
        x = out
        x = x - 2
-
-        if self.intermediate_unbacked:
-            # Test for unbacked symints: the following is a fancy way to multiply by 1
-            u0 = foo(x)
-            ones = x.new_ones(x.shape[0], u0).sum(-1) / 3
-            x = x * ones
-
        x = x - 1
        out = torch.empty_like(x)
        torch.ops.silly.attention(x, x, x, out)
@@ -78,7 +52,6 @@ class SillyModel(nn.Module):
        return x


-@torch._dynamo.config.patch(capture_dynamic_output_shape_ops=True)
 def _run_simple_model(
    splitting_ops,
    use_inductor_graph_partition,
@@ -87,8 +60,6 @@ def _run_simple_model(
    expected_num_piecewise_capturable_graphs_seen,
    expected_num_backend_compilations,
    expected_num_cudagraph_captured,
-    *,
-    intermediate_unbacked=False,
 ):
    vllm_config = VllmConfig(
        compilation_config=CompilationConfig(
@@ -101,11 +72,7 @@ def _run_simple_model(
        )
    )
    with set_current_vllm_config(vllm_config):
-        model = SillyModel(
-            vllm_config=vllm_config,
-            prefix="",
-            intermediate_unbacked=intermediate_unbacked,
-        )
+        model = SillyModel(vllm_config=vllm_config, prefix="")

    inputs = torch.randn(100).cuda()

@@ -158,10 +125,9 @@ def _run_simple_model(


@pytest.mark.parametrize("backend", ["inductor", "eager"])
-@pytest.mark.parametrize("intermediate_unbacked", [True, False])
@torch.inference_mode()
@create_new_process_for_each_test("spawn")
-def test_simple_piecewise_compile(backend, intermediate_unbacked):
+def test_simple_piecewise_compile(backend):
    _run_simple_model(
        splitting_ops=["silly::attention"],
        use_inductor_graph_partition=False,
@@ -174,7 +140,6 @@ def test_simple_piecewise_compile(backend, intermediate_unbacked):
        expected_num_backend_compilations=3,
        # num_cudagraph_sizes * num_piecewise_capturable_graphs_seen
        expected_num_cudagraph_captured=6,
-        intermediate_unbacked=intermediate_unbacked,
    )


@@ -202,10 +202,9 @@ class TestAllReduceFusedAddRMSNormStaticQuantFP4Model(torch.nn.Module):
@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"], reason="Only test on CUDA")
@pytest.mark.skipif(
    not find_spec("flashinfer")
-    or not has_module_attribute("flashinfer.comm", "allreduce_fusion")
-    or not has_module_attribute("flashinfer.comm", "create_allreduce_fusion_workspace"),
+    or not has_module_attribute("flashinfer.comm", "trtllm_allreduce_fusion"),
    reason="flashinfer is not found or flashinfer "
-    "is not compiled with allreduce_fusion",
+    "is not compiled with trtllm_allreduce_fusion",
 )
 def test_all_reduce_fusion_pass_replace(
    test_model: torch.nn.Module,
@@ -1,10 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

+import copy
+
 import pytest
 import torch

-import vllm.envs as envs
 from tests.compile.backend import TestBackend
 from tests.utils import TestFP8Layer
 from vllm.compilation.passes.fusion.act_quant_fusion import (
@@ -31,6 +32,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
 )
 from vllm.model_executor.layers.rotary_embedding import get_rope
 from vllm.platforms import current_platform
+from vllm.utils.torch_utils import direct_register_custom_op

 TEST_FP8 = current_platform.supports_fp8()
 FP8_DTYPE = current_platform.fp8_dtype()
@@ -198,23 +200,82 @@ class TestRotaryEmbeddingSliceScatter(torch.nn.Module):
        return [torch.ops.aten.slice_scatter.default]


-MODELS = [
-    TestSiluMul,
-    TestFusedAddRMSNorm,
-    TestRotaryEmbedding,
-    TestRotaryEmbeddingSliceScatter,
-]
+class TestFunctionWithMutatedArgsAndReturn(torch.nn.Module):
+    OP_REGISTERED = False
+
+    def __init__(self):
+        super().__init__()
+        self.register_test_custom_op()
+
+    @classmethod
+    def register_test_custom_op(cls):
+        if not cls.OP_REGISTERED:
+
+            def function_with_mutated_args_and_return_impl(
+                x: torch.Tensor,
+            ) -> torch.Tensor:
+                ret = x + 1
+                x.add_(2)
+                return ret
+
+            def function_with_mutated_args_and_return_fake(
+                x: torch.Tensor,
+            ) -> torch.Tensor:
+                return torch.empty_like(x)
+
+            direct_register_custom_op(
+                op_name="function_with_mutated_args_and_return",
+                op_func=function_with_mutated_args_and_return_impl,
+                mutates_args=["x"],
+                fake_impl=function_with_mutated_args_and_return_fake,
+            )
+
+            cls.OP_REGISTERED = True
+
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        # Clone x to avoid mutating the original tensor
+        ret = torch.ops.vllm.function_with_mutated_args_and_return(x)
+        return x, ret
+
+    def example_inputs(self, num_tokens=32):
+        hidden_states = torch.randn(num_tokens)
+        return (hidden_states,)
+
+    def ops_in_model(self, do_fusion):
+        return [torch.ops.vllm.function_with_mutated_args_and_return.default]
+
+    def ops_not_in_model(self):
+        return []
+
+
+MODELS_AND_DO_FUSION = {
+    TestSiluMul: [True, False],
+    TestFusedAddRMSNorm: [True, False],
+    TestRotaryEmbedding: [False],
+    TestRotaryEmbeddingSliceScatter: [False],
+    TestFunctionWithMutatedArgsAndReturn: [False],
+}


@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
-@pytest.mark.parametrize("model_class", MODELS)
-@pytest.mark.parametrize("do_fusion", [True, False])
-@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE != "cuda", reason="Only test on CUDA")
+@pytest.mark.parametrize(
+    "model_class, do_fusion",
+    [
+        (model_class, do_fusion)
+        for model_class, fusions in MODELS_AND_DO_FUSION.items()
+        for do_fusion in fusions
+    ],
+)
+@pytest.mark.skipif(
+    not current_platform.is_cuda_alike(),
+    reason="Only test on cuda and rocm platform",
+)
 def test_fix_functionalization(
    model_class: torch.nn.Module, do_fusion: bool, dtype: torch.dtype
 ):
    torch.set_default_device("cuda")
    torch.set_default_dtype(dtype)
+    torch.manual_seed(0)

    vllm_config = VllmConfig(
        model_config=ModelConfig(dtype=dtype),
@@ -246,8 +307,14 @@ def test_fix_functionalization(
        backend_no_func = TestBackend(*passes)

        model = model_class()
-        torch.compile(model, backend=backend_func)(*model.example_inputs())
-        torch.compile(model, backend=backend_no_func)(*model.example_inputs())
+        inputs_func = model.example_inputs()
+        inputs_no_func = copy.deepcopy(inputs_func)
+        model_func = model_class()
+        model_no_func = copy.deepcopy(model_func)
+        model_func = torch.compile(model_func, backend=backend_func)
+        model_no_func = torch.compile(model_no_func, backend=backend_no_func)
+        model_func(*inputs_func)
+        model_no_func(*inputs_no_func)

        # check if the functionalization pass is applied
        for op in model.ops_in_model(do_fusion):
@@ -265,3 +332,8 @@ def test_fix_functionalization(
                    found[op] = True
        assert all(found[op] for op in model.ops_in_model(do_fusion))
        assert all(not found.get(op) for op in model.ops_not_in_model())
+
+        # TODO (Rohan138): compare the outputs from model_func and model_no_func
+        # currently runs into errors while comparing `TestFusedAddRMSNorm`
+        # Linked issue: https://github.com/vllm-project/vllm/issues/34996
+        # torch.testing.assert_close(outputs_func, outputs_no_func)
@@ -92,8 +92,6 @@ class AttentionQuantPatternModel(torch.nn.Module):
    def build_attn_metadata(self, batch_size: int) -> AttentionMetadata:
        """Initialize attention metadata."""

-        # TODO (Rohan138) reuse utils from vllm/v1/worker/gpu/attn_utils.py
-
        # Create common attn metadata
        batch_spec = BatchSpec(seq_lens=[1] * batch_size, query_lens=[1] * batch_size)
        common_attn_metadata = create_common_attn_metadata(
@@ -102,31 +100,58 @@ class AttentionQuantPatternModel(torch.nn.Module):

        max_blocks = (max(batch_spec.seq_lens) + self.block_size - 1) // self.block_size
        num_blocks = batch_size * max_blocks
+        backend = self.attn.backend

-        # Fetch the attention backend and kv cache shape and stride order
-        attn_backend = self.attn.attn_backend
-        kv_cache_shape = attn_backend.get_kv_cache_shape(
-            num_blocks, self.block_size, self.num_kv_heads, self.head_size
-        )
-        try:
-            kv_cache_stride_order = attn_backend.get_kv_cache_stride_order()
-        except (AttributeError, NotImplementedError):
-            kv_cache_stride_order = tuple(range(len(kv_cache_shape)))
-
-        kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order)
-        inv_order = [
-            kv_cache_stride_order.index(i) for i in range(len(kv_cache_stride_order))
-        ]
-
-        # Create dummy KV cache
-        raw_tensor = torch.zeros(
-            2 * num_blocks * self.block_size * self.num_kv_heads * self.head_size,
-            dtype=self.kv_cache_dtype,
-            device=self.device,
-        )
-        raw_tensor = raw_tensor.view(kv_cache_shape)
-        kv_cache = raw_tensor.permute(*inv_order)
-
+        # TODO(luka) use get_kv_cache_stride_order
+        # Create dummy KV cache for the selected backend
+        if backend == AttentionBackendEnum.ROCM_ATTN:
+            # k/v as 1st dimention
+            # HND: [num_blocks, num_kv_heads, block_size, head_size]
+            kv_cache = torch.zeros(
+                2,
+                num_blocks,
+                self.num_kv_heads,
+                self.block_size,
+                self.head_size,
+                dtype=self.kv_cache_dtype,
+                device=self.device,
+            )
+        elif backend == AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN:
+            # k/v as 1st dimention
+            # NHD: [num_blocks, block_size, num_kv_heads, head_size]
+            kv_cache = torch.zeros(
+                2,
+                num_blocks,
+                self.block_size,
+                self.num_kv_heads,
+                self.head_size,
+                dtype=self.kv_cache_dtype,
+                device=self.device,
+            )
+        elif backend == AttentionBackendEnum.TRITON_ATTN:
+            # k/v as 2nd dimention
+            # NHD: [num_blocks, block_size, num_kv_heads, head_size]
+            kv_cache = torch.zeros(
+                num_blocks,
+                2,
+                self.num_kv_heads,
+                self.block_size,
+                self.head_size,
+                dtype=self.kv_cache_dtype,
+                device=self.device,
+            )
+        elif backend == AttentionBackendEnum.FLASHINFER:
+            kv_cache = torch.zeros(
+                num_blocks,
+                2,
+                self.num_kv_heads,
+                self.block_size,
+                self.head_size,
+                dtype=self.kv_cache_dtype,
+                device=self.device,
+            ).permute(0, 1, 3, 2, 4)
+        else:
+            raise ValueError(f"Unsupported backend: {backend}")
        self.attn.kv_cache = [kv_cache]

        # Build attn metadata
@@ -0,0 +1,325 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+import vllm.config
+from tests.compile.backend import TestBackend
+from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata
+from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
+from vllm.compilation.passes.fusion.matcher_utils import ROTARY_OP
+from vllm.compilation.passes.fusion.rope_kvcache_fusion import RopeKVCacheFusionPass
+from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass
+from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass
+from vllm.compilation.passes.utility.scatter_split_replace import (
+    ScatterSplitReplacementPass,
+)
+from vllm.compilation.passes.utility.split_coalescing import SplitCoalescingPass
+from vllm.config import (
+    CacheConfig,
+    CompilationConfig,
+    CompilationMode,
+    ModelConfig,
+    PassConfig,
+    VllmConfig,
+)
+from vllm.forward_context import get_forward_context, set_forward_context
+from vllm.model_executor.layers.attention import Attention
+from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
+from vllm.platforms import current_platform
+from vllm.v1.attention.backend import (
+    AttentionBackend,
+    CommonAttentionMetadata,
+)
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
+from vllm.v1.kv_cache_interface import AttentionSpec
+
+INDEX_SELECT_OP = torch.ops.aten.index.Tensor
+VLLM_UNIFIED_KV_CACHE_UPDATE_OP = torch.ops.vllm.unified_kv_cache_update
+FP8_DTYPE = current_platform.fp8_dtype()
+
+
+class QKRoPEKVCacheTestModel(torch.nn.Module):
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        attn_backend: AttentionBackendEnum,
+        num_heads: int,
+        num_kv_heads: int,
+        head_size: int,
+        is_neox: bool,
+        dtype: torch.dtype,
+        device: torch.device,
+        prefix: str = "model.layers.0.self_attn.attn",
+    ):
+        super().__init__()
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads
+        self.head_size = head_size
+        self.block_size = vllm_config.cache_config.block_size
+        self.q_size = num_heads * head_size
+        self.kv_size = num_kv_heads * head_size
+        self.is_neox = is_neox
+        self.dtype = dtype
+        self.device = device
+        self.layer_name = prefix
+
+        self.rotary_emb = RotaryEmbedding(
+            head_size,
+            rotary_dim=head_size,
+            max_position_embeddings=4096,
+            base=10000,
+            is_neox_style=is_neox,
+            dtype=self.dtype,
+        )
+
+        # Whether to check for the RoPE custom op or component index_select
+        self.enable_rope_custom_op = self.rotary_emb.enabled()
+
+        # Register layer metadata for the fusion pass via Attention.
+        self.attn = Attention(
+            num_heads=num_heads,
+            head_size=head_size,
+            scale=1.0 / head_size**0.5,
+            num_kv_heads=num_kv_heads,
+            cache_config=vllm_config.cache_config,
+            quant_config=vllm_config.quant_config,
+            prefix=prefix,
+            attn_backend=attn_backend.get_class(),
+        )
+        self.attn_backend: type[AttentionBackend] = self.attn.get_attn_backend()
+        assert not self.attn_backend.forward_includes_kv_cache_update, (
+            f"Attention backend {self.attn_backend} does not support fuse_rope_kvcache."
+        )
+        self.attn._k_scale = self.attn._k_scale.to(device)
+        self.attn._v_scale = self.attn._v_scale.to(device)
+
+        kv_cache_dtype_str = vllm_config.cache_config.cache_dtype
+        self.kv_cache_dtype = (
+            FP8_DTYPE if kv_cache_dtype_str.startswith("fp8") else self.dtype
+        )
+
+        # Initialize attn MetadataBuilder
+        self.builder = self.attn.attn_backend.get_builder_cls()(
+            kv_cache_spec=AttentionSpec(
+                block_size=self.block_size,
+                num_kv_heads=self.num_kv_heads,
+                head_size=head_size,
+                dtype=self.kv_cache_dtype,
+            ),
+            layer_names=[self.attn.layer_name],
+            vllm_config=vllm_config,
+            device=device,
+        )
+
+    def build_attn_metadata(self, batch_size: int) -> CommonAttentionMetadata:
+        """Initialize attention metadata."""
+        # Create common attn metadata
+        batch_spec = BatchSpec(seq_lens=[1] * batch_size, query_lens=[1] * batch_size)
+        common_attn_metadata = create_common_attn_metadata(
+            batch_spec, self.block_size, self.device, arange_block_indices=True
+        )
+
+        max_blocks = (max(batch_spec.seq_lens) + self.block_size - 1) // self.block_size
+        num_blocks = batch_size * max_blocks
+
+        # Fetch the attention backend and kv cache shape and stride order
+        attn_backend = self.attn.attn_backend
+        kv_cache_shape = attn_backend.get_kv_cache_shape(
+            num_blocks, self.block_size, self.num_kv_heads, self.head_size
+        )
+        try:
+            kv_cache_stride_order = attn_backend.get_kv_cache_stride_order()
+        except (AttributeError, NotImplementedError):
+            kv_cache_stride_order = tuple(range(len(kv_cache_shape)))
+
+        kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order)
+        inv_order = [
+            kv_cache_stride_order.index(i) for i in range(len(kv_cache_stride_order))
+        ]
+
+        # Create dummy KV cache
+        raw_tensor = torch.zeros(
+            2 * num_blocks * self.block_size * self.num_kv_heads * self.head_size,
+            dtype=self.kv_cache_dtype,
+            device=self.device,
+        )
+        raw_tensor = raw_tensor.view(kv_cache_shape)
+        kv_cache = raw_tensor.permute(*inv_order)
+
+        self.attn.kv_cache = [kv_cache]
+
+        # Build attn metadata
+        attn_metadata = self.builder.build(
+            common_prefix_len=0, common_attn_metadata=common_attn_metadata
+        )
+
+        return attn_metadata
+
+    def forward(
+        self, qkv: torch.Tensor, positions: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        # Create copy so inplace ops do not modify the original tensors
+        qkv = qkv.clone()
+        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+        q, k = self.rotary_emb(positions, q, k)
+
+        # Instead of a full forward pass, match only the KV cache update op here
+        q = q.view(-1, self.num_heads, self.head_size)
+        k = k.view(-1, self.num_kv_heads, self.head_size)
+        v = v.view(-1, self.num_kv_heads, self.head_size)
+        kv_cache_dummy_dep = torch.ops.vllm.unified_kv_cache_update(
+            k, v, self.layer_name
+        )
+        return q, k, v, kv_cache_dummy_dep
+
+    def ops_in_model_before(self) -> list[torch._ops.OpOverload]:
+        ops = []
+        if self.enable_rope_custom_op:
+            ops.append(ROTARY_OP)
+        else:
+            ops.append(INDEX_SELECT_OP)
+        ops.append(torch.ops.vllm.unified_kv_cache_update.default)
+        return ops
+
+    def ops_in_model_after(self) -> list[torch._ops.OpOverload]:
+        return [torch.ops.vllm.fused_rope_and_unified_kv_cache_update.default]
+
+
+@pytest.mark.parametrize(
+    "attn_backend",
+    [
+        AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN,
+        AttentionBackendEnum.TRITON_ATTN,
+        AttentionBackendEnum.ROCM_ATTN,
+    ],
+)
+@pytest.mark.parametrize("enable_rope_custom_op", [True])  # [True, False])
+@pytest.mark.parametrize("num_heads", [64])
+@pytest.mark.parametrize("num_kv_heads", [8])
+@pytest.mark.parametrize("head_size", [64])
+@pytest.mark.parametrize("block_size", [16])
+@pytest.mark.parametrize("is_neox", [True, False])
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
+@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
+@pytest.mark.skipif(
+    not is_aiter_found_and_supported(),
+    reason="Only test on ROCm with AITER installed and supported",
+)
+def test_rope_kvcache_fusion(
+    attn_backend: AttentionBackendEnum,
+    enable_rope_custom_op: bool,
+    num_heads: int,
+    num_kv_heads: int,
+    head_size: int,
+    block_size: int,
+    is_neox: bool,
+    dtype: torch.dtype,
+    kv_cache_dtype: str,
+    monkeypatch: pytest.MonkeyPatch,
+):
+    torch.set_default_device("cuda")
+    torch.set_default_dtype(dtype)
+    torch.manual_seed(0)
+
+    custom_ops: list[str] = []
+    if enable_rope_custom_op:
+        custom_ops.append("+rotary_embedding")
+
+    vllm_config = VllmConfig(
+        model_config=ModelConfig(dtype=dtype),
+        cache_config=CacheConfig(
+            block_size=block_size,
+            cache_dtype=kv_cache_dtype,
+        ),
+        compilation_config=CompilationConfig(
+            mode=CompilationMode.VLLM_COMPILE,
+            custom_ops=custom_ops,
+            pass_config=PassConfig(
+                fuse_rope_kvcache=True,
+                eliminate_noops=True,
+            ),
+        ),
+    )
+
+    with vllm.config.set_current_vllm_config(vllm_config), monkeypatch.context() as m:
+        m.setenv("VLLM_ROCM_USE_AITER", "1")
+        rocm_aiter_ops.refresh_env_variables()
+
+        model = QKRoPEKVCacheTestModel(
+            vllm_config=vllm_config,
+            attn_backend=attn_backend,
+            num_heads=num_heads,
+            num_kv_heads=num_kv_heads,
+            head_size=head_size,
+            is_neox=is_neox,
+            dtype=dtype,
+            device=torch.get_default_device(),
+        )
+
+        fusion_pass = RopeKVCacheFusionPass(vllm_config)
+        passes = [
+            NoOpEliminationPass(vllm_config),
+            SplitCoalescingPass(vllm_config),
+            ScatterSplitReplacementPass(vllm_config),
+            fusion_pass,
+            PostCleanupPass(vllm_config),
+        ]
+        backend = TestBackend(*passes)
+
+        T = 5
+
+        qkv = torch.randn(
+            T, num_heads * head_size + 2 * num_kv_heads * head_size, dtype=dtype
+        )
+        pos = torch.arange(T, dtype=torch.long)
+
+        qkv_unfused = qkv.clone()
+        pos_unfused = pos.clone()
+
+        with set_forward_context(None, vllm_config):
+            forward_context = get_forward_context()
+            attn_metadata = model.build_attn_metadata(T)
+            forward_context.slot_mapping = {
+                model.layer_name: attn_metadata.slot_mapping
+            }
+            q_unfused, k_unfused, v_unfused, dummy = model(qkv_unfused, pos_unfused)
+            attn_layer = forward_context.no_compile_layers[model.layer_name]
+            kv_cache_unfused = attn_layer.kv_cache[forward_context.virtual_engine]
+        del dummy
+
+        torch._dynamo.mark_dynamic(qkv, 0)
+        torch._dynamo.mark_dynamic(pos, 0)
+        with set_forward_context(None, vllm_config):
+            model_fused = torch.compile(model, backend=backend)
+            forward_context = get_forward_context()
+            attn_metadata = model_fused.build_attn_metadata(T)
+            forward_context.slot_mapping = {
+                model.layer_name: attn_metadata.slot_mapping
+            }
+            q_fused, k_fused, v_fused, dummy = model_fused(qkv, pos)
+            attn_layer = forward_context.no_compile_layers[model.layer_name]
+            kv_cache_fused = attn_layer.kv_cache[forward_context.virtual_engine]
+        del dummy
+
+        assert fusion_pass.matched_count == 1
+
+        backend.check_before_ops(model.ops_in_model_before())
+        backend.check_after_ops(model.ops_in_model_after())
+
+        if dtype == torch.float16:
+            ATOL, RTOL = (2e-3, 2e-3)
+        else:
+            ATOL, RTOL = (1e-2, 1e-2)
+
+        torch.testing.assert_close(q_unfused, q_fused, atol=ATOL, rtol=RTOL)
+        torch.testing.assert_close(k_unfused, k_fused, atol=ATOL, rtol=RTOL)
+        torch.testing.assert_close(v_unfused, v_fused, atol=ATOL, rtol=RTOL)
+        # Cannot compare fp8_* directly here, cast to model dtype instead
+        torch.testing.assert_close(
+            kv_cache_unfused.view(dtype),
+            kv_cache_fused.view(dtype),
+            atol=ATOL,
+            rtol=RTOL,
+        )
@@ -0,0 +1,107 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+import torch.nn as nn
+
+import vllm
+from tests.compile.backend import TestBackend
+from vllm.compilation.passes.utility.scatter_split_replace import (
+    ScatterSplitReplacementPass,
+)
+from vllm.compilation.passes.utility.split_coalescing import SplitCoalescingPass
+from vllm.config import CompilationConfig, CompilationMode, VllmConfig
+from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
+
+
+class ScatterSplitReplacementModel(nn.Module):
+    """Model with a rope+getitem+slice_scatter+split_with_sizes sequence."""
+
+    def __init__(
+        self,
+        num_heads: int,
+        num_kv_heads: int,
+        head_size: int,
+        dtype: torch.dtype,
+    ):
+        super().__init__()
+        self.q_size = num_heads * head_size
+        self.kv_size = num_kv_heads * head_size
+
+        self.rotary_emb = RotaryEmbedding(
+            head_size,
+            rotary_dim=head_size,
+            max_position_embeddings=4096,
+            base=10000,
+            is_neox_style=True,
+            dtype=dtype,
+        )
+
+    def forward(self, qkv: torch.Tensor, positions: torch.Tensor):
+        # Create copy so inplace ops do not modify the original tensors
+        qkv = qkv.clone()
+        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+        q, k = self.rotary_emb(positions, q, k)
+        q = q + 1
+        k = k + 2
+        v = v + 3
+        return q, k, v
+
+    def ops_in_model_before(self) -> list[torch._ops.OpOverload]:
+        return [
+            torch.ops.aten.slice_scatter.default,
+            torch.ops.aten.split_with_sizes.default,
+            torch.ops.aten.getitem.default,
+        ]
+
+    def ops_in_model_after(self) -> list[torch._ops.OpOverload]:
+        return [torch.ops.aten.getitem.default]
+
+
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+def test_scatter_split_replace(dtype):
+    torch.set_default_device("cuda")
+    torch.set_default_dtype(dtype)
+    torch.manual_seed(0)
+
+    num_heads = 8
+    num_kv_heads = 4
+    head_size = 64
+
+    vllm_config = VllmConfig(
+        compilation_config=CompilationConfig(
+            mode=CompilationMode.VLLM_COMPILE,
+            custom_ops=["+rotary_embedding"],
+        ),
+    )
+    with vllm.config.set_current_vllm_config(vllm_config):
+        # ScatterSplitReplacementPass requires SplitCoalescingPass to be run before it
+        coalesce_pass = SplitCoalescingPass(vllm_config)
+        replace_pass = ScatterSplitReplacementPass(vllm_config)
+        passes = [coalesce_pass, replace_pass]
+        backend = TestBackend(*passes)
+
+        model = ScatterSplitReplacementModel(num_heads, num_kv_heads, head_size, dtype)
+
+        T = 5
+        qkv = torch.randn(
+            T, num_heads * head_size + 2 * num_kv_heads * head_size, dtype=dtype
+        )
+        pos = torch.arange(T, dtype=torch.long)
+
+        qkv_eager = qkv.clone()
+        pos_eager = pos.clone()
+        result_eager = model(qkv_eager, pos_eager)
+
+        torch._dynamo.mark_dynamic(qkv, 0)
+        torch._dynamo.mark_dynamic(pos, 0)
+
+        model_compiled = torch.compile(model, backend=backend)
+        result_compiled = model_compiled(qkv, pos)
+
+        for eager, compiled in zip(result_eager, result_compiled):
+            torch.testing.assert_close(eager, compiled)
+
+        assert backend.op_count(torch.ops.aten.slice_scatter.default) == 0
+        assert backend.op_count(torch.ops.aten.split_with_sizes.default) == 1
@@ -80,19 +80,12 @@ def test_ray_runtime_env(monkeypatch: pytest.MonkeyPatch):
    ray.shutdown()


-def test_unrecognized_env(monkeypatch):
+def test_unrecognized_env():
    import os

-    from vllm.envs import environment_variables
-
-    # Remove any existing unrecognized VLLM env vars that might interfere
-    for env in list(os.environ):
-        if env.startswith("VLLM_") and env not in environment_variables:
-            monkeypatch.delenv(env, raising=False)
-
    # Test that if fail_on_environ_validation is True, then an error
    # is raised when an unrecognized vLLM environment variable is set
-    monkeypatch.setenv("VLLM_UNRECOGNIZED_ENV_VAR", "some_value")
+    os.environ["VLLM_UNRECOGNIZED_ENV_VAR"] = "some_value"
    engine_args = EngineArgs(
        fail_on_environ_validation=True,
    )
@@ -104,7 +97,7 @@ def test_unrecognized_env(monkeypatch):
    engine_args.create_engine_config()

    # Test that when the unrecognized env var is removed, no error is raised
-    monkeypatch.delenv("VLLM_UNRECOGNIZED_ENV_VAR")
+    os.environ.pop("VLLM_UNRECOGNIZED_ENV_VAR", None)
    engine_args = EngineArgs(
        fail_on_environ_validation=True,
    )
@@ -3,7 +3,6 @@

 import pytest

-from vllm.config.model import ModelConfig
 from vllm.config.multimodal import MultiModalConfig
 from vllm.v1.attention.backends.registry import AttentionBackendEnum

@@ -24,20 +23,3 @@ def test_mm_encoder_attn_backend_hash_updates():
        mm_encoder_attn_backend=AttentionBackendEnum.FLASH_ATTN
    ).compute_hash()
    assert base_hash != overridden_hash
-
-
-def test_language_model_only_does_not_affect_mm_hash():
-    """language_model_only does not affect the ViT computation graph,
-    so it should not change the multimodal config hash."""
-    base_hash = MultiModalConfig().compute_hash()
-    lm_only_hash = MultiModalConfig(language_model_only=True).compute_hash()
-    assert base_hash == lm_only_hash
-
-
-def test_language_model_only_affects_model_hash():
-    """language_model_only affects the LM computation graph,
-    so it should change the model config hash."""
-    model = "llava-hf/llava-1.5-7b-hf"
-    base_hash = ModelConfig(model).compute_hash()
-    lm_only_hash = ModelConfig(model, language_model_only=True).compute_hash()
-    assert base_hash != lm_only_hash
@@ -419,6 +419,7 @@ class HfRunner:
            self.tokenizer: "PreTrainedTokenizer | PreTrainedTokenizerFast" = (
                AutoTokenizer.from_pretrained(
                    model_name,
+                    dtype=dtype,
                    trust_remote_code=trust_remote_code,
                )
            )
@@ -429,6 +430,7 @@ class HfRunner:

        self.processor = AutoProcessor.from_pretrained(
            model_name,
+            dtype=dtype,
            trust_remote_code=trust_remote_code,
        )
        if skip_tokenizer_init:
@@ -39,6 +39,7 @@ def test_min_tokens_with_stop(min_tokens: int, stop: str, truth: str):
        mm_features=None,
        sampling_params=params,
        pooling_params=None,
+        eos_token_id=None,
        arrival_time=0.0,
        lora_request=None,
        cache_salt=None,
@@ -35,6 +35,7 @@ def _make_request(stop, include_stop_str_in_output: bool, min_tokens: int = 0):
        mm_features=None,
        sampling_params=params,
        pooling_params=None,
+        eos_token_id=None,
        arrival_time=0.0,
        lora_request=None,
        cache_salt=None,
@@ -19,8 +19,6 @@ from vllm.distributed import (
    tensor_model_parallel_all_reduce,
    tensor_model_parallel_reduce_scatter,
 )
-from vllm.distributed.parallel_state import GroupCoordinator, TensorMetadata
-from vllm.v1.worker.gpu_worker import AsyncIntermediateTensors

 from ..utils import (
    init_test_distributed_environment,
@@ -202,111 +200,6 @@ def send_recv_tensor_dict_test_worker(
        torch.testing.assert_close(recv_dict["f"], test_dict["f"])


-class _DummyWork:
-    def __init__(self) -> None:
-        self.wait_calls = 0
-
-    def wait(self) -> None:
-        self.wait_calls += 1
-
-
-class _DummyAllGatherGroup:
-    def __init__(self, world_size: int, rank_in_group: int) -> None:
-        self.world_size = world_size
-        self.rank_in_group = rank_in_group
-
-    def all_gather(self, t: torch.Tensor, dim: int = 0) -> torch.Tensor:
-        # duplicate local slice across ranks.
-        assert dim == 0
-        return torch.cat([t for _ in range(self.world_size)], dim=0)
-
-
-def _make_group_for_unit_test(
-    rank_in_group: int = 0, world_size: int = 2
-) -> GroupCoordinator:
-    # avoid running GroupCoordinator.__init__ (it wires up real process groups).
-    g = GroupCoordinator.__new__(GroupCoordinator)
-    g.world_size = world_size
-    g.rank_in_group = rank_in_group
-    g.ranks = list(range(world_size))
-    g.use_cpu_custom_send_recv = False
-    g.device_group = None
-    g.cpu_group = None
-    return g
-
-
-def test_irecv_tensor_dict_send_allgather_postprocess_binds_keys(
-    monkeypatch: pytest.MonkeyPatch,
-) -> None:
-    def fake_irecv(t: torch.Tensor, *args: Any, **kwargs: Any) -> _DummyWork:
-        t.fill_(1)
-        return _DummyWork()
-
-    monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True)
-    monkeypatch.setattr(torch.distributed, "irecv", fake_irecv)
-
-    g = _make_group_for_unit_test(rank_in_group=0, world_size=2)
-    # 2 tensors so we can catch late-binding bugs in postprocess closures.
-    metadata_list = [
-        ("a", TensorMetadata("cpu", torch.int32, torch.Size([4]))),
-        ("b", TensorMetadata("cpu", torch.int32, torch.Size([4]))),
-    ]
-    g.recv_object = lambda src=None: metadata_list  # type: ignore[method-assign]
-
-    ag = _DummyAllGatherGroup(world_size=2, rank_in_group=0)
-    td, handles, postprocess = g.irecv_tensor_dict(all_gather_group=ag)
-
-    assert td is not None
-    assert len(handles) == 2
-    assert len(postprocess) == 2
-
-    # before postprocess, dict holds the TP slice (shape 2).
-    assert td["a"].shape == torch.Size([2])
-    assert td["b"].shape == torch.Size([2])
-
-    # simulate worker-side "defer wait": wait + postprocess later.
-    for handle in handles:
-        handle.wait()
-    for fn in postprocess:
-        fn()
-
-    # after postprocess, dict values are reconstructed to full shape (shape 4),
-    # and each key should be updated independently
-    assert td["a"].shape == torch.Size([4])
-    assert td["b"].shape == torch.Size([4])
-    torch.testing.assert_close(td["a"], torch.ones(4, dtype=torch.int32))
-    torch.testing.assert_close(td["b"], torch.ones(4, dtype=torch.int32))
-
-
-def test_async_intermediate_tensors_lazy_wait() -> None:
-    work = _DummyWork()
-    post_calls = {"n": 0}
-
-    def post() -> None:
-        post_calls["n"] += 1
-
-    it = AsyncIntermediateTensors(
-        {"x": torch.tensor([1])},
-        comm_handles=[work],
-        comm_postprocess=[post],
-    )
-
-    # accessing non-tensor attributes should not trigger wait.
-    assert it.kv_connector_output is None
-    assert work.wait_calls == 0
-    assert post_calls["n"] == 0
-
-    # first access of `.tensors` triggers wait + postprocess.
-    _ = it.tensors
-    assert work.wait_calls == 1
-    assert post_calls["n"] == 1
-
-    # subsequent access should not re-wait.
-    _ = it.tensors
-    assert work.wait_calls == 1
-    assert post_calls["n"] == 1
-
-
@ray.remote(num_gpus=1, max_calls=1)
 def send_recv_test_worker(
    monkeypatch: pytest.MonkeyPatch,
@@ -295,11 +295,12 @@ def _test_async_transfer_layer_without_mtp_worker(
    for layer_idx in range(num_layers):
        is_unchanged, is_received_locally, recv_metadata = asyncio.run(
            transfer_layer(
-                old_layer_indices=old_indices_cpu[layer_idx],
-                new_layer_indices=new_indices_cpu[layer_idx],
-                expert_weights=expert_weights[layer_idx],
+                old_global_expert_indices=old_indices_cpu,
+                new_global_expert_indices=new_indices_cpu,
+                expert_weights=expert_weights,
                expert_weights_buffer=expert_buffer,
                ep_group=ep_group,
+                layer=layer_idx,
                cuda_stream=cuda_stream,
            )
        )
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
Gregory Shtrasberg	dbab37f36b	Merge ROCm:fused_aiter_rope_kvcache (#33443 )	2026-02-23 17:23:40 +00:00
Gregory Shtrasberg	ea9721a447	Merge ROCm:rocm_ck_moe_a16w4 (#34301 )	2026-02-23 17:22:16 +00:00
Hashem HashemiandGregory Shtrasberg	956da72c21	small adjustment to wvSplitKrc (#34410 ) Signed-off-by: Hashem Hashemi <hashem.hashemi@amd.com>	2026-02-23 17:15:20 +00:00
kliuaeandGregory Shtrasberg	eb65a75111	[ROCm] [aiter] Split KV cache update for AiterFlashAttention (#33681 ) Signed-off-by: kliuae <kuanfu.liu@embeddedllm.com>	2026-02-23 17:15:20 +00:00
Gregory Shtrasberg	a3a1428d03	Merge ROCm:rocm_dsfp4_dyn_quant (#34157 ) Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com>	2026-02-23 17:15:20 +00:00
Gregory Shtrasberg	f867dfcb29	ROCm 7.2 and disabled prebuild kernels for aiter. Added ROCm7.2 required libs to build aotriton; Use the new base image for this release Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com>	2026-02-23 17:15:20 +00:00
Gregory Shtrasberg	59916839c9	Merge ROCm:upstream_fp8_with_static_scales_gpt_oss (#30357 ) Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com>	2026-02-23 17:15:20 +00:00
Gregory Shtrasberg	45219818d0	Merge ROCm:rocm_dsfp4_moe_shuffle (#34192 ) Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com>	2026-02-20 16:10:16 +00:00
xuebwang-amd Gregory Shtrasberg Robert Shaw Robert Shaw	9e9432ac96	[ROCm][Quantization] GPT_OSS in amd-quark format model loading and emulations (#29008 ) Signed-off-by: xuebwang-amd <xuebwang@amd.com> Signed-off-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>	2026-02-20 16:10:10 +00:00
Gregory Shtrasberg	dadbceb1d9	Merge ROCm:rocm_defaults (#33271 ) Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com>	2026-02-20 16:09:17 +00:00
Gregory ShtrasbergandGregory Shtrasberg	830de08dc4	[Feature] Warn about unrecognized environment variables (#33581 ) Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com>	2026-02-11 16:20:24 +00:00
Seiji Eicherandkhluu	c44d0c6d66	Patch protobuf for CVE-2026-0994 (#34253 ) Signed-off-by: Seiji Eicher <seiji@anyscale.com> Co-authored-by: Kevin H. Luu <khluu000@gmail.com> (cherry picked from commit `5045d5c983`)	2026-02-11 02:33:40 -08:00
Kunshang Jiandkhluu	83db96d8cd	[XPU][9/N] clean up existing ipex code/doc (#34111 ) Signed-off-by: Kunshang Ji <kunshang.ji@intel.com> (cherry picked from commit `cb9574eb85`)	2026-02-11 02:33:27 -08:00
zofiaandkhluu	dbfb79fe45	[XPU][7/N] enable xpu fp8 moe (#34202 ) Signed-off-by: Zhu, Zufang <zufang.zhu@intel.com> (cherry picked from commit `b482f71e9f`)	2026-02-11 02:33:15 -08:00
Roger Wangandkhluu	b2e1fc3589	[Bugfix][Core] Fix CPU memory leak from Request reference cycle in prefix caching (#34183 ) Signed-off-by: Roger Wang <hey@rogerw.io> (cherry picked from commit `8a5e0e2b2b`)	2026-02-11 02:33:04 -08:00
Gregory Shtrasbergandkhluu	55a1baebc5	[Bugfix][ROCm][GPT-OSS] Use old triton_kernels implementation on ROCm if the new API is not available (#34153 ) Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com> (cherry picked from commit `c60f8e3b49`)	2026-02-11 02:32:52 -08:00
Charlie Fuandkhluu	e1e9841631	[torch.compile][Fusion] Fix attention fusion pass removing kv_udpate op. (#33945 ) Signed-off-by: charlifu <charlifu@amd.com> (cherry picked from commit `bb9f97308d`)	2026-02-11 02:32:41 -08:00
zofiaandkhluu	5bd63387c3	[XPU][6/N] add xpu scaled_mm kernel (#34117 ) Signed-off-by: Zhu, Zufang <zufang.zhu@intel.com> (cherry picked from commit `9bdb06b436`)	2026-02-11 02:32:27 -08:00