[Core] Add Prefill Context Parallelism (PCP) and All-to-All DCP communication

This PR adds Prefill Context Parallelism (PCP) support for splitting prefill tokens across ranks using a DualChunkSwap pattern, and integrates an All-to-All communication backend for Decode Context Parallelism (DCP). Key changes: - Add PCP with DualChunkSwap token partitioning for balanced prefill computation - Add All-to-All DCP communication backend reducing NCCL calls from 3 to 2 - Restrict DCP+PCP to two clean configurations: - Case 1: DCP = PCP (same TP position, all-reduce only) - Case 2: DCP = TP × PCP (full TP all-gather, all-reduce + slice) - Add PCPManager for buffer management and input partitioning - Update attention backends (FlashAttention, FlashInfer, MLA) for PCP support - Add comprehensive tests for DCP operations Co-Authored-By: QiuChunshuo <qiuchunshuo@huawei.com> Co-Authored-By: zhenwenqi2024 <zhenwenqi_2022@qq.com> Co-Authored-By: FENP <yuanyongjie.yyj@antgroup.com> Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
2026-03-02 05:40:48 +00:00
888 changed files with 34170 additions and 54272 deletions
@@ -13,10 +13,9 @@ import os
 from contextlib import contextmanager

 import lm_eval
+import numpy as np
 import yaml

-from vllm.platforms import current_platform
-
 DEFAULT_RTOL = 0.08


@@ -64,9 +63,6 @@ def launch_lm_eval(eval_config, tp_size):
        "allow_deprecated_quantization=True,"
    )

-    if current_platform.is_rocm() and "Nemotron-3" in eval_config["model_name"]:
-        model_args += "attention_backend=TRITON_ATTN"
-
    env_vars = eval_config.get("env_vars", None)
    with scoped_env_vars(env_vars):
        results = lm_eval.simple_evaluate(
@@ -106,8 +102,6 @@ def test_lm_eval_correctness_param(config_filename, tp_size):
                f"ground_truth={ground_truth:.3f} | "
                f"measured={measured_value:.3f} | rtol={rtol}"
            )
-
-            min_acceptable = ground_truth * (1 - rtol)
-            success = success and measured_value >= min_acceptable
+            success = success and np.isclose(ground_truth, measured_value, rtol=rtol)

    assert success
@@ -83,6 +83,7 @@ We test the throughput by using `vllm bench serve` with request rate = inf to co
        "server_parameters": {
            "model": "meta-llama/Meta-Llama-3-8B",
            "tensor_parallel_size": 1,
+            "swap_space": 16,
            "disable_log_stats": "",
            "load_format": "dummy"
        },
@@ -51,56 +51,5 @@
            "max-model-len": 256,
            "async-scheduling": ""
        }
-    },
-    {
-        "test_name": "latency_deepseek_r1",
-        "environment_variables": {
-            "PT_HPU_LAZY_MODE": 1,
-            "PT_HPU_ENABLE_LAZY_COLLECTIVES": 1,
-            "VLLM_CONTIGUOUS_PA": 1,
-            "VLLM_DEFRAG": 1
-        },
-        "parameters": {
-            "model": "deepseek-ai/DeepSeek-R1",
-            "tensor_parallel_size": 8,
-            "load_format": "dummy",
-            "max-model-len": 2048,
-            "dtype": "bfloat16"
-        }
-    },
-    {
-        "test_name": "latency_llama4_maverick_17b128e_instruct_fp8",
-        "environment_variables": {
-            "PT_HPU_LAZY_MODE": 1,
-            "PT_HPU_ENABLE_LAZY_COLLECTIVES": 1,
-            "VLLM_CONTIGUOUS_PA": 1,
-            "VLLM_DEFRAG": 1
-        },
-        "parameters": {
-            "model": "meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8",
-            "tensor_parallel_size": 8,
-            "max-model-len": 512,
-            "max-num-seqs": 128,
-            "async-scheduling": "",
-            "gpu-memory-utilization": 0.95,
-            "enable_expert_parallel": ""
-        }
-    },
-    {
-        "test_name": "latency_qwen3_8b",
-        "environment_variables": {
-            "PT_HPU_LAZY_MODE": 1,
-            "PT_HPU_ENABLE_LAZY_COLLECTIVES": 1,
-            "VLLM_CONTIGUOUS_PA": 1,
-            "VLLM_DEFRAG": 1
-        },
-        "parameters": {
-            "model": "Qwen/Qwen3-8B",
-            "tensor_parallel_size": 1,
-            "max-model-len": 2048,
-            "max-num-seqs": 128,
-            "dtype": "bfloat16",
-            "async-scheduling": ""
-        }
    }
 ]
@@ -10,6 +10,7 @@
        "server_parameters": {
            "model": "meta-llama/Meta-Llama-3.1-8B-Instruct",
            "tensor_parallel_size": 1,
+            "swap_space": 16,
            "disable_log_stats": "",
            "load_format": "dummy",
            "max-model-len": 2048,
@@ -36,6 +37,7 @@
        "server_parameters": {
            "model": "meta-llama/Meta-Llama-3.1-70B-Instruct",
            "tensor_parallel_size": 4,
+            "swap_space": 16,
            "disable_log_stats": "",
            "load_format": "dummy",
            "max-model-len": 2048,
@@ -62,6 +64,7 @@
        "server_parameters": {
            "model": "mistralai/Mixtral-8x7B-Instruct-v0.1",
            "tensor_parallel_size": 2,
+            "swap_space": 16,
            "disable_log_stats": "",
            "load_format": "dummy",
            "max-model-len": 2048,
@@ -75,83 +78,5 @@
            "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json",
            "num_prompts": 200
        }
-    },
-    {
-        "test_name": "serving_deepseek_r1",
-        "qps_list": [1, 4, 16, "inf"],
-        "server_environment_variables": {
-            "PT_HPU_LAZY_MODE": 1,
-            "PT_HPU_ENABLE_LAZY_COLLECTIVES": 1,
-            "VLLM_CONTIGUOUS_PA": 1,
-            "VLLM_DEFRAG": 1
-        },
-        "server_parameters": {
-            "model": "deepseek-ai/DeepSeek-R1",
-            "tensor_parallel_size": 8,
-            "disable_log_stats": "",
-            "load_format": "dummy",
-            "max-model-len": 2048,
-            "max-num-seqs": 200,
-            "async-scheduling": "",
-            "dtype": "bfloat16"
-        },
-        "client_parameters": {
-            "model": "deepseek-ai/DeepSeek-R1",
-            "backend": "vllm",
-            "dataset_name": "sharegpt",
-            "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json",
-            "num_prompts": 200
-        }
-    },
-    {
-        "test_name": "serving_llama4_maverick_17b128e_instruct_fp8",
-        "qps_list": [1, 4, 16, "inf"],
-        "server_environment_variables": {
-            "PT_HPU_LAZY_MODE": 1,
-            "PT_HPU_ENABLE_LAZY_COLLECTIVES": 1,
-            "VLLM_CONTIGUOUS_PA": 1,
-            "VLLM_DEFRAG": 1
-        },
-        "server_parameters": {
-            "model": "meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8",
-            "tensor_parallel_size": 8,
-            "disable_log_stats": "",
-            "max-model-len": 2048,
-            "max-num-seqs": 128,
-            "async-scheduling": "",
-            "enable_expert_parallel": "",
-            "max-num-batched-tokens": 4096
-        },
-        "client_parameters": {
-            "model": "meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8",
-            "backend": "vllm",
-            "dataset_name": "sharegpt",
-            "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json",
-            "num_prompts": 200
-        }
-    },
-    {
-        "test_name": "serving_qwen3_8b",
-        "qps_list": [1, 4, 10, "inf"],
-        "server_environment_variables": {
-            "PT_HPU_LAZY_MODE": 1,
-            "PT_HPU_ENABLE_LAZY_COLLECTIVES": 1,
-            "VLLM_CONTIGUOUS_PA": 1,
-            "VLLM_DEFRAG": 1
-        },
-        "server_parameters": {
-            "model": "Qwen/Qwen-3-8B",
-            "tensor_parallel_size": 1,
-            "dtype": "bfloat16",
-            "disable_log_stats": "",
-            "async-scheduling": ""
-        },
-        "client_parameters": {
-            "model": "Qwen/Qwen-3-8B",
-            "backend": "vllm",
-            "dataset_name": "sharegpt",
-            "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json",
-            "num_prompts": 200
-        }
    }
 ]
@@ -5,6 +5,7 @@
        "server_parameters": {
            "model": "meta-llama/Meta-Llama-3.1-8B-Instruct",
            "tensor_parallel_size": 1,
+            "swap_space": 16,
            "disable_log_stats": "",
            "load_format": "dummy"
        },
@@ -22,6 +23,7 @@
        "server_parameters": {
            "model": "meta-llama/Meta-Llama-3.1-70B-Instruct",
            "tensor_parallel_size": 4,
+            "swap_space": 16,
            "disable_log_stats": "",
            "load_format": "dummy"
        },
@@ -39,6 +41,7 @@
        "server_parameters": {
            "model": "mistralai/Mixtral-8x7B-Instruct-v0.1",
            "tensor_parallel_size": 2,
+            "swap_space": 16,
            "disable_log_stats": "",
            "load_format": "dummy"
        },
@@ -56,6 +59,7 @@
        "server_parameters": {
            "model": "meta-llama/Meta-Llama-3.1-70B-Instruct", 
            "tensor_parallel_size": 4,
+            "swap_space": 16,
            "speculative_config": {
                "model": "turboderp/Qwama-0.5B-Instruct",
                "num_speculative_tokens": 4,
@@ -57,67 +57,5 @@
            "max-num-seqs": 512,
            "async-scheduling": ""
        }
-    },
-    {
-        "test_name": "throughput_deepseek_r1",
-        "environment_variables": {
-            "PT_HPU_LAZY_MODE": 1,
-            "PT_HPU_ENABLE_LAZY_COLLECTIVES": 1,
-            "VLLM_CONTIGUOUS_PA": 1,
-            "VLLM_DEFRAG": 1
-        },
-        "parameters": {
-            "model": "deepseek-ai/DeepSeek-R1",
-            "tensor_parallel_size": 8,
-            "load_format": "dummy",
-            "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json",
-            "dataset_name": "sharegpt",
-            "num_prompts": 1000,
-            "backend": "vllm",
-            "max-model-len": 2048,
-            "max-num-seqs": 384,
-            "async-scheduling": ""
-        }
-    },
-    {
-        "test_name": "throughput_llama4_maverick_17b128e_instruct_fp8",
-        "environment_variables": {
-            "PT_HPU_LAZY_MODE": 1,
-            "PT_HPU_ENABLE_LAZY_COLLECTIVES": 1,
-            "VLLM_CONTIGUOUS_PA": 1,
-            "VLLM_DEFRAG": 1
-        },
-        "parameters": {
-            "model": "meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8",
-            "tensor_parallel_size": 8,
-            "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json",
-            "dataset_name": "sharegpt",
-            "num_prompts": 1000,
-            "backend": "vllm",
-            "max-model-len": 2048,
-            "max-num-seqs": 512,
-            "async-scheduling": "",
-            "enable_expert_parallel": ""
-        }
-    },
-    {
-        "test_name": "throughput_qwen3_8b",
-        "environment_variables": {
-            "PT_HPU_LAZY_MODE": 1,
-            "PT_HPU_ENABLE_LAZY_COLLECTIVES": 1,
-            "VLLM_CONTIGUOUS_PA": 1,
-            "VLLM_DEFRAG": 1
-        },
-        "parameters": {
-            "model": "Qwen/Qwen-3-8B",
-            "tensor_parallel_size": 1,
-            "load_format": "dummy",
-            "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json",
-            "dataset_name": "sharegpt",
-            "num_prompts": 1000,
-            "max-num-seqs": 512,
-            "backend": "vllm",
-            "async-scheduling": ""
-        }
    }
 ]
@@ -68,7 +68,7 @@ aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/triton
 aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/torchvision-*.whl .
 aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/torchaudio-*.whl .
 aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/amdsmi-*.whl .
-aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/amd_aiter-*.whl .
+aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/aiter-*.whl .
 aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/flash-attn-*.whl .
 \`\`\`

@@ -80,7 +80,7 @@ aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/flash-
 - **torchvision**: TorchVision for ROCm PyTorch
 - **torchaudio**: Torchaudio for ROCm PyTorch
 - **amdsmi**: AMD SMI Python bindings
- **amd_aiter**: Aiter for ROCm
+- **aiter**: Aiter for ROCm
 - **flash-attn**: Flash Attention for ROCm

 ### :warning: Notes
@@ -1,213 +0,0 @@
-#!/bin/bash
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-#
-# Check if Ray LLM can generate lock files that are compatible with this
-# version of vllm. Downloads Ray's requirement files and runs a full
-# dependency resolution with the installed vllm's constraints to see if
-# a valid lock file can be produced.
-#
-# See: https://github.com/vllm-project/vllm/issues/33599
-
-set -eo pipefail
-
-RAY_BASE_URL="https://raw.githubusercontent.com/ray-project/ray/master/python"
-
-WORK_DIR=$(mktemp -d)
-trap 'rm -rf "$WORK_DIR"' EXIT
-
-# Fetch all Ray requirement files used in the LLM depset pipeline
-echo ">>> Fetching Ray requirement files"
-RAY_FILES=(
-    "requirements.txt"
-    "requirements/cloud-requirements.txt"
-    "requirements/base-test-requirements.txt"
-    "requirements/llm/llm-requirements.txt"
-    "requirements/llm/llm-test-requirements.txt"
-)
-for FILE in "${RAY_FILES[@]}"; do
-    LOCAL_PATH="${WORK_DIR}/$(basename "$FILE")"
-    echo "    ${FILE}"
-    curl -fsSL -o "$LOCAL_PATH" "${RAY_BASE_URL}/${FILE}"
-done
-
-# Extract installed vllm deps
-echo ">>> Extracting installed vllm dependency constraints"
-python3 - "${WORK_DIR}/vllm-constraints.txt" <<'PYEOF'
-"""Write out the installed vllm's dependencies as pip constraint lines.
-
-Ray uses vllm[audio], so audio-extra deps are included with their extra
-markers stripped. The resolver cannot evaluate extra markers for a
-package that is not itself being resolved from an index, so we activate
-them manually here.
-"""
-import importlib.metadata
-import re
-import sys
-
-out_path = sys.argv[1]
-raw_reqs = importlib.metadata.requires("vllm") or []
-
-# Ray uses vllm[audio] – activate that extra.
-ACTIVE_EXTRAS = {"audio"}
-EXTRA_RE = re.compile(r"""extra\s*==\s*['"]([^'"]+)['"]""")
-
-lines = []
-for r in raw_reqs:
-    if ";" not in r:
-        # Unconditional dep — always include.
-        lines.append(r.strip())
-        continue
-
-    req_part, _, marker_part = r.partition(";")
-    marker_part = marker_part.strip()
-
-    extra_matches = EXTRA_RE.findall(marker_part)
-    if not extra_matches:
-        # Non-extra marker (python_version, etc.) — keep as-is.
-        lines.append(r.strip())
-        continue
-
-    if not ACTIVE_EXTRAS.intersection(extra_matches):
-        continue  # Skip inactive extras (tensorizer, bench, …).
-
-    # Strip the extra== conditions but keep any remaining markers
-    # (e.g. python_version).
-    cleaned = EXTRA_RE.sub("", marker_part)
-    cleaned = re.sub(r"\band\b\s*\band\b", "and", cleaned)
-    cleaned = re.sub(r"^\s*and\s+|\s+and\s*$", "", cleaned).strip()
-
-    if cleaned:
-        lines.append(f"{req_part.strip()} ; {cleaned}")
-    else:
-        lines.append(req_part.strip())
-
-with open(out_path, "w") as f:
-    for line in lines:
-        f.write(line + "\n")
-
-print(f"Wrote {len(lines)} constraints to {out_path}")
-PYEOF
-
-echo ">>> Installed vllm deps (first 20 lines):"
-head -20 "${WORK_DIR}/vllm-constraints.txt"
-
-# Remove Ray's vllm pin — the installed vllm's transitive deps
-# (written above) replace it in the resolution. vllm itself cannot
-# be resolved from PyPI for in-development versions, so we test
-# whether Ray's requirements can coexist with vllm's dependency
-# constraints instead.
-sed -i '/^vllm/d' "${WORK_DIR}/llm-requirements.txt"
-
-# Install uv if needed
-if ! command -v uv &>/dev/null; then
-    echo ">>> Installing uv"
-    pip install uv -q
-fi
-
-# Resolve: given vllm's constraints, can Ray compile a lock file?
-#
-# vllm's dependency constraints are the fixed side — Ray is flexible and
-# can regenerate its lock files. We pass vllm's constraints via -c so
-# the resolver treats them as non-negotiable bounds, then check whether
-# Ray's own requirements can still be satisfied within those bounds.
-echo ""
-echo "============================================================"
-echo ">>> Resolving: Can Ray generate compatible lock files?"
-echo "============================================================"
-
-set +e
-uv pip compile \
-    "${WORK_DIR}/requirements.txt" \
-    "${WORK_DIR}/cloud-requirements.txt" \
-    "${WORK_DIR}/base-test-requirements.txt" \
-    "${WORK_DIR}/llm-requirements.txt" \
-    "${WORK_DIR}/llm-test-requirements.txt" \
-    -c "${WORK_DIR}/vllm-constraints.txt" \
-    --python-version 3.12 \
-    --python-platform x86_64-manylinux_2_31 \
-    --extra-index-url https://download.pytorch.org/whl/cu129 \
-    --index-strategy unsafe-best-match \
-    --unsafe-package setuptools \
-    --unsafe-package ray \
-    --no-header \
-    -o "${WORK_DIR}/resolved.txt" \
-    2>&1
-EXIT_CODE=$?
-set -e
-
-echo ""
-echo "=========================================="
-if [ $EXIT_CODE -eq 0 ]; then
-    echo "SUCCESS: Ray can generate lock files compatible with this vllm."
-    echo ""
-    echo "Key resolved versions:"
-    grep -E '^(protobuf|torch|numpy|transformers)==' \
-        "${WORK_DIR}/resolved.txt" | sort || true
-    echo "=========================================="
-    exit 0
-fi
-
-echo "FAILURE: Ray cannot generate lock files compatible with this vllm."
-echo "This means a fundamental dependency conflict exists that Ray"
-echo "cannot resolve by regenerating its lock files."
-echo "See: https://github.com/vllm-project/vllm/issues/33599"
-echo "=========================================="
-
-# Buildkite annotation
-if [ -f /usr/bin/buildkite-agent ]; then
-    buildkite-agent annotate --style 'warning' --context 'ray-compat' << EOF
-### :warning: Ray Dependency Compatibility Warning
-This PR introduces dependencies that **cannot** be resolved with Ray's requirements.
-Ray would not be able to regenerate its lock files to accommodate this vllm version.
-
-Please check the **Ray Dependency Compatibility Check** step logs for details.
-See [issue #33599](https://github.com/vllm-project/vllm/issues/33599) for context.
-EOF
-fi
-
-# Notify Slack if webhook is configured and PR/branch are valid.
-if [ -n "$RAY_COMPAT_SLACK_WEBHOOK_URL" ]; then
-    PR="${BUILDKITE_PULL_REQUEST:-}"
-    BRANCH="${BUILDKITE_BRANCH:-}"
-
-    # Skip notification if PR is invalid or branch is empty
-    if [[ "$PR" = "false" || -z "$PR" || -z "$BRANCH" ]]; then
-        echo ">>> Skipping Slack notification (invalid PR or empty branch: PR=$PR, branch=$BRANCH)"
-    else
-        echo ">>> Sending Slack notification"
-        # Single quotes are intentional: the f-string expressions are Python, not shell.
-        # shellcheck disable=SC2016
-        PAYLOAD=$(python3 -c '
-import json, os, sys
-pr = os.getenv("BUILDKITE_PULL_REQUEST", "N/A")
-branch = os.getenv("BUILDKITE_BRANCH", "unknown")
-url = os.getenv("BUILDKITE_BUILD_URL", "#")
-data = {
-    "text": ":warning: Ray Dependency Compatibility Check Failed",
-    "blocks": [{
-        "type": "section",
-        "text": {
-            "type": "mrkdwn",
-            "text": (
-                "*:warning: Ray Dependency Compatibility Check Failed*\n"
-                f"PR #{pr} on branch `{branch}` introduces dependencies "
-                f"that cannot be resolved with Ray'\''s requirements.\n"
-                f"<{url}|View Build>"
-            ),
-        },
-    }],
-}
-print(json.dumps(data))
-')
-
-        HTTP_CODE=$(curl -s -o /dev/null -w "%{http_code}" -X POST "$RAY_COMPAT_SLACK_WEBHOOK_URL" \
-            -H 'Content-type: application/json' \
-            -d "$PAYLOAD")
-        echo "    Slack webhook response: $HTTP_CODE"
-    fi
-else
-    echo ">>> Skipping Slack notification (RAY_COMPAT_SLACK_WEBHOOK_URL not set)"
-fi
-
-exit 1
@@ -99,15 +99,6 @@ is_multi_node() {
  return 1
 }

-handle_pytest_exit() {
-  local exit_code=$1
-  if [ "$exit_code" -eq 5 ]; then
-    echo "Pytest exit code 5 (no tests collected) - treating as success."
-    exit 0
-  fi
-  exit "$exit_code"
-}
-
 ###############################################################################
 # Pytest marker/keyword re-quoting
 #
@@ -144,9 +135,8 @@ re_quote_pytest_markers() {
  local collecting=false
  local marker_buf=""

-  # Strip backslash-newline continuations, then flatten remaining newlines
-  local flat="${input//$'\\\n'/ }"
-  flat="${flat//$'\n'/ }"
+  # Flatten newlines for consistent tokenization
+  local flat="${input//$'\n'/ }"

  # Disable globbing to prevent *.py etc. from expanding during read -ra
  local restore_glob
@@ -174,9 +164,6 @@ re_quote_pytest_markers() {

      local is_boundary=false
      case "$word" in
-        # Line-continuation artifact
-        "\\")
-          is_boundary=true ;;
        # Command separators
        "&&"|"||"|";"|"|")
          is_boundary=true ;;
@@ -217,9 +204,6 @@ re_quote_pytest_markers() {
        if [[ "$word" == "-m" || "$word" == "-k" ]]; then
          output+="${word} "
          collecting=true
-        # Drop stray backslash tokens silently
-        elif [[ "$word" == "\\" ]]; then
-          :
        else
          output+="${word} "
        fi
@@ -469,9 +453,7 @@ if is_multi_node "$commands"; then
    done

    /bin/bash -c "${composite_command}"
-    exit_code=$?
    cleanup_network
-    handle_pytest_exit "$exit_code"
  else
    echo "Multi-node job detected but failed to parse bracket command syntax."
    echo "Expected format: prefix ; [node0_cmd1, node0_cmd2] && [node1_cmd1, node1_cmd2]"
@@ -498,7 +480,4 @@ else
    --name "${container_name}" \
    "${image_name}" \
    /bin/bash -c "${commands}"
-
-  exit_code=$?
-  handle_pytest_exit "$exit_code"
 fi
@@ -1,43 +1,26 @@
 #!/bin/bash
 set -euox pipefail
-export VLLM_CPU_CI_ENV=0

 echo "--- PP+TP"
 vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -pp=2 &
 server_pid=$!
-timeout 600 bash -c "until curl localhost:8000/v1/models > /dev/null 2>&1; do sleep 1; done" || exit 1
+timeout 600 bash -c "until curl localhost:8000/v1/models; do sleep 1; done" || exit 1
 vllm bench serve \
    --backend vllm \
    --dataset-name random \
    --model meta-llama/Llama-3.2-3B-Instruct \
    --num-prompts 20 \
-    --result-dir ./test_results \
-    --result-filename tp_pp.json \
-    --save-result \
    --endpoint /v1/completions
-kill -s SIGTERM $server_pid; wait $server_pid || true
-failed_req=$(jq '.failed' ./test_results/tp_pp.json)
-if [ "$failed_req" -ne 0 ]; then
-  echo "Some requests were failed!"
-  exit 1
-fi
+kill -s SIGTERM $server_pid &

 echo "--- DP+TP"
 vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -dp=2 &
 server_pid=$!
-timeout 600 bash -c "until curl localhost:8000/v1/models > /dev/null 2>&1; do sleep 1; done" || exit 1
+timeout 600 bash -c "until curl localhost:8000/v1/models; do sleep 1; done" || exit 1
 vllm bench serve \
    --backend vllm \
    --dataset-name random \
    --model meta-llama/Llama-3.2-3B-Instruct \
    --num-prompts 20 \
-    --result-dir ./test_results \
-    --result-filename dp_pp.json \
-    --save-result \
    --endpoint /v1/completions
-kill -s SIGTERM $server_pid; wait $server_pid || true
-failed_req=$(jq '.failed' ./test_results/dp_pp.json)
-if [ "$failed_req" -ne 0 ]; then
-  echo "Some requests were failed!"
-  exit 1
-fi
+kill -s SIGTERM $server_pid &
@@ -34,7 +34,7 @@ function cpu_tests() {
  # offline inference
  docker exec cpu-test bash -c "
    set -e
-    python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m"
+    python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m"

  # Run model tests
  docker exec cpu-test bash -c "
@@ -27,7 +27,7 @@ function cpu_tests() {
  podman exec -it "$container_id" bash -c "
    export TORCH_COMPILE_DISABLE=1
    set -xve
-    python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m" >> "$HOME"/test_basic.log
+    python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m" >> "$HOME"/test_basic.log

  # Run basic model test
  podman exec -it "$container_id" bash -c "
@@ -25,5 +25,5 @@ remove_docker_container

 # Run the image and test offline inference
 docker run -e HF_TOKEN -e VLLM_WORKER_MULTIPROC_METHOD=spawn -v /root/.cache/huggingface:/root/.cache/huggingface --name gh200-test --gpus=all --entrypoint="" gh200-test bash -c '
-    python3 examples/basic/offline_inference/generate.py --model meta-llama/Llama-3.2-1B
+    python3 examples/offline_inference/basic/generate.py --model meta-llama/Llama-3.2-1B
 '
@@ -1,27 +1,9 @@
 #!/bin/bash

-# This script builds the HPU docker image and runs the offline inference inside the container.
+# This script build the CPU docker image and run the offline inference inside the container.
 # It serves a sanity check for compilation and basic model usage.
-#
-# vllm-gaudi compatibility pinning:
-#   The vllm-gaudi plugin is installed on top of the vllm upstream checkout used by this CI job.
-#   When upstream vllm changes its API, the plugin may break before it has been updated.
-#   To handle this, the vllm-gaudi repository maintains a file:
-#     vllm/last-good-commit-for-vllm-gaudi/VLLM_COMMUNITY_COMMIT
-#   The first line of that file controls what version of vllm is used inside the Docker image:
-#     - "latest"        : no checkout override; the current Buildkite CI commit is used as-is.
-#     - "<commit SHA>"  : vllm is checked out to that specific commit before building, pinning
-#                         the test to a known-compatible baseline.
-#   To unpin (resume testing against the live vllm tip), set the file content back to "latest".
 set -exuo pipefail

-# Fetch the vllm community commit reference from vllm-gaudi (first line only).
-VLLM_COMMUNITY_COMMIT=$(curl -s \
-  https://raw.githubusercontent.com/vllm-project/vllm-gaudi/vllm/last-good-commit-for-vllm-gaudi/VLLM_COMMUNITY_COMMIT \
-  | head -1 | tr -d '\n')
-
-echo "Using vllm community commit: ${VLLM_COMMUNITY_COMMIT}"
-
 # Try building the docker image
 image_name="hpu/upstream-vllm-ci:${BUILDKITE_COMMIT}"
 container_name="hpu-upstream-vllm-ci-${BUILDKITE_COMMIT}-container"
@@ -30,13 +12,6 @@ FROM gaudi-base-image:latest

 COPY ./ /workspace/vllm

-# If VLLM_COMMUNITY_COMMIT is a specific commit (not "latest"), check it out to pin vllm
-# to the version known to be compatible with vllm-gaudi. When the value is "latest",
-# the current checkout (the Buildkite CI commit) is used unchanged.
-RUN if [ "${VLLM_COMMUNITY_COMMIT}" != "latest" ]; then \
-      cd /workspace/vllm && git fetch --unshallow 2>/dev/null || true && git checkout ${VLLM_COMMUNITY_COMMIT}; \
-    fi
-
 WORKDIR /workspace/vllm

 ENV no_proxy=localhost,127.0.0.1
@@ -76,7 +51,7 @@ docker run --rm --runtime=habana --name="${container_name}" --network=host \
  -e PT_HPU_LAZY_MODE=1 \
  "${image_name}" \
  /bin/bash -c '
-  cd vllm; timeout 120s python -u examples/basic/offline_inference/generate.py --model facebook/opt-125m
+  cd vllm; timeout 120s python -u examples/offline_inference/basic/generate.py --model facebook/opt-125m
 '

 EXITCODE=$?
@@ -34,17 +34,17 @@ docker run \
    set -e
    echo $ZE_AFFINITY_MASK
    pip install tblib==3.1.0
-    python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager
-    python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 -O3 -cc.cudagraph_mode=NONE
-    python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager -tp 2 --distributed-executor-backend ray
-    python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager -tp 2 --distributed-executor-backend mp
-    python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --attention-backend=TRITON_ATTN
-    python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --quantization fp8
-    python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4  --block-size 64 --enforce-eager
-    python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b  --block-size 64 --enforce-eager -tp 2
-    python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b  --block-size 64 --enforce-eager -tp 2 --enable-expert-parallel
+    python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager
+    python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 -O3 -cc.cudagraph_mode=NONE
+    python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager -tp 2 --distributed-executor-backend ray
+    python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager -tp 2 --distributed-executor-backend mp
+    python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --attention-backend=TRITON_ATTN
+    python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --quantization fp8
+    python3 examples/offline_inference/basic/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4  --block-size 64 --enforce-eager
+    python3 examples/offline_inference/basic/generate.py --model ibm-research/PowerMoE-3b  --block-size 64 --enforce-eager -tp 2
+    python3 examples/offline_inference/basic/generate.py --model ibm-research/PowerMoE-3b  --block-size 64 --enforce-eager -tp 2 --enable-expert-parallel
    cd tests
-    pytest -v -s v1/core --ignore=v1/core/test_reset_prefix_cache_e2e.py --ignore=v1/core/test_scheduler_e2e.py
+    pytest -v -s v1/core --ignore=v1/core/test_reset_prefix_cache_e2e.py
    pytest -v -s v1/engine
    pytest -v -s v1/sample --ignore=v1/sample/test_logprobs.py --ignore=v1/sample/test_logprobs_e2e.py
    pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py
@@ -24,7 +24,7 @@ if command -v rocm-smi &> /dev/null || [[ -d /opt/rocm ]] || [[ -n "${ROCM_PATH:
  BACKENDS=("allgather_reducescatter")
  # Disable MOE padding for ROCm since it is causing eplb to fail
  export VLLM_ROCM_MOE_PADDING=0
-  PLATFORM_ARGS=("--no-async-scheduling" "--attention-backend=TRITON_ATTN")
+  PLATFORM_ARGS=("--no-async-scheduling")
  echo "Disabled async scheduling for ROCm platform due to issues with spec decode."
 else
  # Non-ROCm platform (CUDA/other)
@@ -72,7 +72,7 @@ obj_json="objects.json"
 aws s3api list-objects-v2 --bucket "$BUCKET" --prefix "$SUBPATH/" --delimiter / --output json > "$obj_json"
 mkdir -p "$INDICES_OUTPUT_DIR"

-# call script to generate indices for all existing wheels
+# call script to generate indicies for all existing wheels
 # this indices have relative paths that could work as long as it is next to the wheel directory in s3
 # i.e., the wheels are always in s3://vllm-wheels/<commit>/
 # and indices can be placed in /<commit>/, or /nightly/, or /<version>/
@@ -54,13 +54,10 @@ mkdir -p $DIST_DIR
 # include only wheels for the release version, ignore all files with "dev" or "rc" in the name (without excluding 'aarch64')
 aws s3 cp --recursive --exclude "*" --include "vllm-${PURE_VERSION}*.whl" --exclude "*dev*" --exclude "*rc[0-9]*" "$S3_COMMIT_PREFIX" $DIST_DIR
 echo "Wheels copied to local directory"
-# generate source distribution using setup.py
-python setup.py sdist --dist-dir=$DIST_DIR
+# generate source tarball
+git archive --format=tar.gz --output="$DIST_DIR/vllm-${PURE_VERSION}.tar.gz" "$BUILDKITE_COMMIT"
 ls -la $DIST_DIR

-SDIST_FILE=$(find $DIST_DIR -name "vllm*.tar.gz")
-echo "Found sdist: $SDIST_FILE"
-
 # upload wheels to PyPI (only default variant, i.e. files without '+' in the name)
 PYPI_WHEEL_FILES=$(find $DIST_DIR -name "vllm-${PURE_VERSION}*.whl" -not -name "*+*")
 if [[ -z "$PYPI_WHEEL_FILES" ]]; then
@@ -68,6 +65,6 @@ if [[ -z "$PYPI_WHEEL_FILES" ]]; then
  exit 1
 fi

-python3 -m twine check "$PYPI_WHEEL_FILES" "$SDIST_FILE"
-python3 -m twine upload --non-interactive --verbose "$PYPI_WHEEL_FILES" "$SDIST_FILE"
-echo "Wheels and source distribution uploaded to PyPI"
+python3 -m twine check "$PYPI_WHEEL_FILES"
+python3 -m twine upload --non-interactive --verbose "$PYPI_WHEEL_FILES"
+echo "Wheels uploaded to PyPI"
@@ -388,7 +388,9 @@ steps:
 - label: V1 Test e2e + engine # 65min
  timeout_in_minutes: 90
  mirror_hardwares: [amdexperimental, amdproduction]
-  agent_pool: mi325_1
+  # The test uses 4 GPUs, but we schedule it on 8-GPU machines for stability.
+  # See discussion here: https://github.com/vllm-project/vllm/pull/31040
+  agent_pool: mi325_8
  optional: true
  # grade: Blocking
  source_file_dependencies:
@@ -400,34 +402,6 @@ steps:
    - pytest -v -s v1/e2e
    - pytest -v -s v1/engine

- label: V1 Test e2e (2 GPUs) # 65min
-  timeout_in_minutes: 90
-  mirror_hardwares: [amdexperimental, amdproduction]
-  agent_pool: mi325_2
-  optional: true
-  # grade: Blocking
-  source_file_dependencies:
-    - vllm/
-    - tests/v1
-  commands:
-    # Only run tests that need exactly 2 GPUs
-    - pytest -v -s v1/e2e/test_spec_decode.py -k "tensor_parallelism"
-
- label: V1 Test e2e (4 GPUs) # 65min
-  timeout_in_minutes: 90
-  mirror_hardwares: [amdexperimental, amdproduction]
-  # The test uses 4 GPUs, but we schedule it on 8-GPU machines for stability.
-  # See discussion here: https://github.com/vllm-project/vllm/pull/31040
-  agent_pool: mi325_4
-  optional: true
-  # grade: Blocking
-  source_file_dependencies:
-    - vllm/
-    - tests/v1
-  commands:
-    # Only run tests that need 4 GPUs
-    - pytest -v -s v1/e2e/test_spec_decode.py -k "eagle_correctness_heavy"
-
 - label: V1 Test entrypoints # 35min
  timeout_in_minutes: 50
  mirror_hardwares: [amdexperimental, amdproduction, amdtentative]
@@ -467,7 +441,7 @@ steps:
    - pip install -U git+https://github.com/robertgshaw2-redhat/lm-evaluation-harness.git@streaming-api
    - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine

-# TODO: Add the "V1 Test attention (MI300)" test group
+# TODO: Add the "V1 Test attetion (MI300)" test group

 - label: V1 Test attention (H100) # 10min
  mirror_hardwares: [amdexperimental, amdproduction]
@@ -499,6 +473,17 @@ steps:
    - pytest -v -s v1/determinism/test_batch_invariance.py
    - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py

+- label: V1 Test attention (B200) # 10min
+  timeout_in_minutes: 30
+  gpu: b200
+  source_file_dependencies:
+    - vllm/config/attention.py
+    - vllm/model_executor/layers/attention
+    - vllm/v1/attention
+    - tests/v1/attention
+  commands:
+    - pytest -v -s v1/attention
+
 - label: V1 Test others (CPU) # 5 mins
  mirror_hardwares: [amdexperimental, amdproduction, amdtentative]
  agent_pool: mi325_1
@@ -529,12 +514,12 @@ steps:
  commands:
    - pip install tensorizer # for tensorizer test
    # for basic
-    - python3 basic/offline_inference/chat.py --attention-backend TRITON_ATTN
-    - python3 basic/offline_inference/generate.py --model facebook/opt-125m
-    - python3 basic/offline_inference/generate.py --model meta-llama/Llama-2-13b-chat-hf --cpu-offload-gb 10
-    - python3 basic/offline_inference/classify.py
-    - python3 basic/offline_inference/embed.py
-    - python3 basic/offline_inference/score.py
+    - python3 offline_inference/basic/chat.py
+    - python3 offline_inference/basic/generate.py --model facebook/opt-125m
+    - python3 offline_inference/basic/generate.py --model meta-llama/Llama-2-13b-chat-hf --cpu-offload-gb 10
+    - python3 offline_inference/basic/classify.py
+    - python3 offline_inference/basic/embed.py
+    - python3 offline_inference/basic/score.py
    # for multi-modal models
    - python3 offline_inference/audio_language.py --seed 0
    - python3 offline_inference/vision_language.py --seed 0
@@ -599,8 +584,6 @@ steps:
      --ignore=lora/test_qwen3moe_tp.py
  parallelism: 4

-##### .buildkite/test_areas/pytorch.yaml #####
-# corresponds to .buildkite/test_areas/pytorch.yaml
 - label: PyTorch Compilation Unit Tests # 15min
  timeout_in_minutes: 30
  mirror_hardwares: [amdexperimental, amdproduction]
@@ -618,20 +601,6 @@ steps:
  # they do not suffer from https://github.com/vllm-project/vllm/issues/28965
  - "find compile/ -maxdepth 1 -name 'test_*.py' -exec pytest -s -v {} \\\\;"

-# corresponds to .buildkite/test_areas/pytorch.yaml
- label: PyTorch Compilation Passes Unit Tests
-  timeout_in_minutes: 20
-  mirror_hardwares: [amdexperimental, amdproduction]
-  agent_pool: mi325_1
-  source_file_dependencies:
-    - vllm/
-    - tests/compile/passes
-  commands:
-  # TODO: clean up this comment if not needed. It is used to 
-  # keep track of the tests changes during vLLM IR Ops refactoring.
-  # Use `find` to launch multiple instances of pytest.
-  - "find compile/passes -maxdepth 1 -name 'test_*.py' -exec pytest -s -v {} \\\\;"
-
 - label: PyTorch Fullgraph Smoke Test # 15min
  timeout_in_minutes: 30
  mirror_hardwares: [amdexperimental, amdproduction]
@@ -1049,7 +1018,6 @@ steps:
  source_file_dependencies:
  - vllm/
  - tests/models/multimodal
-  - tests/models/registry.py
  no_gpu: true
  commands:
    - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
@@ -1063,7 +1031,6 @@ steps:
  source_file_dependencies:
  - vllm/
  - tests/models/multimodal
-  - tests/models/registry.py
  commands:
    - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
    - pytest -v -s models/multimodal/processing
@@ -1169,11 +1136,53 @@ steps:
    - pytest -v -s tests/models/test_transformers.py
    # - pytest -v -s tests/models/multimodal/processing/
    - pytest -v -s tests/models/multimodal/test_mapping.py -k 'not (Gemma3 or Qwen2VL or Qwen2_5_VL)'
-    - python3 examples/basic/offline_inference/chat.py
+    - python3 examples/offline_inference/basic/chat.py
    # - python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl
    # Whisper needs spawn method to avoid deadlock
    - VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper

+- label: Blackwell Test # 21 min
+  timeout_in_minutes: 30
+  working_dir: "/vllm-workspace/"
+  gpu: b200
+  # optional: true
+  source_file_dependencies:
+  - csrc/quantization/fp4/
+  - csrc/attention/mla/
+  - csrc/quantization/cutlass_w8a8/moe/
+  - vllm/model_executor/layers/fused_moe/cutlass_moe.py
+  - vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py
+  - vllm/model_executor/layers/fused_moe/flashinfer_a2a_prepare_finalize.py
+  - vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
+  - vllm/v1/attention/backends/flashinfer.py
+  - vllm/v1/attention/backends/mla/cutlass_mla.py
+  - vllm/v1/attention/backends/mla/flashinfer_mla.py
+  - vllm/v1/attention/selector.py
+  - vllm/platforms/cuda.py
+  commands:
+    - nvidia-smi
+    - python3 examples/offline_inference/basic/chat.py
+    # Attention
+    # num_heads2 broken by https://github.com/flashinfer-ai/flashinfer/issues/1353
+    - pytest -v -s tests/kernels/attention/test_attention_selector.py
+    - pytest -v -s tests/kernels/attention/test_flashinfer.py -k 'not num_heads2'
+    - pytest -v -s tests/kernels/attention/test_flashinfer_trtllm_attention.py
+    - pytest -v -s tests/kernels/attention/test_cutlass_mla_decode.py
+    - pytest -v -s tests/kernels/attention/test_flashinfer_mla_decode.py
+    # Quantization
+    - pytest -v -s tests/kernels/quantization/test_cutlass_scaled_mm.py -k 'fp8'
+    - pytest -v -s tests/kernels/quantization/test_nvfp4_quant.py
+    - pytest -v -s tests/kernels/quantization/test_silu_mul_nvfp4_quant.py
+    - pytest -v -s tests/kernels/quantization/test_nvfp4_scaled_mm.py
+    - pytest -v -s tests/kernels/quantization/test_flashinfer_scaled_mm.py
+    - pytest -v -s tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py
+    - pytest -v -s tests/kernels/quantization/test_nvfp4_qutlass.py
+    - pytest -v -s tests/kernels/quantization/test_mxfp4_qutlass.py
+    - pytest -v -s tests/kernels/moe/test_nvfp4_moe.py
+    - pytest -v -s tests/kernels/moe/test_ocp_mx_moe.py
+    - pytest -v -s tests/kernels/moe/test_flashinfer.py
+    - pytest -v -s tests/kernels/moe/test_cutedsl_moe.py
+
 - label: Blackwell Fusion and Compile Tests # 30 min
  timeout_in_minutes: 40
  working_dir: "/vllm-workspace/"
@@ -1240,6 +1249,16 @@ steps:
  commands:
    - pytest -s -v tests/quantization/test_blackwell_moe.py

+- label: Blackwell LM Eval Small Models
+  timeout_in_minutes: 120
+  gpu: b200
+  optional: true # run on nightlies
+  source_file_dependencies:
+  - csrc/
+  - vllm/model_executor/layers/quantization
+  commands:
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-blackwell.txt
+
 #####  1 GPU test  #####
 #####  multi gpus test  #####

@@ -1311,7 +1330,6 @@ steps:
  - tests/v1/entrypoints/openai/test_multi_api_servers.py
  - tests/v1/shutdown
  - tests/v1/worker/test_worker_memory_snapshot.py
-  - examples/offline_inference/new_weight_syncing/
  commands:
  # Work around HIP bug tracked here: https://github.com/ROCm/hip/issues/3876
  # TODO: Remove when the bug is fixed in a future ROCm release
@@ -1325,6 +1343,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
  - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown
  - pytest -v -s v1/worker/test_worker_memory_snapshot.py

@@ -1486,20 +1505,6 @@ steps:
    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
    - DP_EP=1 ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh

- label: CrossLayer KV layout Distributed NixlConnector PD accuracy tests (4 GPUs)
-  mirror_hardwares: [amdexperimental, amdproduction]
-  agent_pool: mi325_4
-  # grade: Blocking
-  timeout_in_minutes: 30
-  working_dir: "/vllm-workspace/tests"
-  num_devices: 4
-  source_file_dependencies:
-    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
-    - tests/v1/kv_connector/nixl_integration/
-  commands:
-    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
-    - CROSS_LAYERS_BLOCKS=1 ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
-
 ##### multi gpus test #####
 ##### A100 test #####

@@ -1539,8 +1544,8 @@ steps:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
  - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large.txt --tp-size=4

-##### FP8 test #####
- label: LM Eval Large Models (H100) # optional, still use H100 for consistency
+##### H100 test #####
+- label: LM Eval Large Models (H100) # optional
  gpu: h100
  optional: true
  mirror_hardwares: [amdexperimental, amdproduction]
@@ -1552,8 +1557,8 @@ steps:
  - csrc/
  - vllm/model_executor/layers/quantization
  commands:
-    - export VLLM_USE_DEEP_GEMM=0 
-    - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-rocm.txt --tp-size=4
+    - export VLLM_USE_DEEP_GEMM=0  # We found Triton is faster than DeepGEMM for H100
+    - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-hopper.txt --tp-size=4


 ##### H200 test #####
@@ -1568,16 +1573,16 @@ steps:
  commands:
    - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/passes/distributed/test_async_tp.py
    - pytest -v -s tests/compile/passes/distributed/test_sequence_parallelism.py
-    # TODO: this test is not supported on ROCm, there are aiter kernels for this.
-    # - pytest -v -s tests/compile/passes/distributed/test_fusion_all_reduce.py
+    - pytest -v -s tests/compile/passes/distributed/test_fusion_all_reduce.py
    #- pytest -v -s tests/compile/distributed/test_fusions_e2e.py::test_tp2_attn_quant_allreduce_rmsnorm
    # - "VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/distributed/test_fusions_e2e.py -k 'not Llama-4'"
    # Old E2E tests were removed in https://github.com/vllm-project/vllm/pull/33293
    # in favor of new tests in fusions_e2e. We avoid replicating the new jobs in this file as it's deprecated.
+
+    - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py
    - pytest -v -s tests/distributed/test_context_parallel.py
    - HIP_VISIBLE_DEVICES=0,1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=allgather_reducescatter --disable-nccl-for-dp-synchronization
-    # this test is not supported on ROCm
-    # - pytest -v -s tests/v1/distributed/test_dbo.py
+    - pytest -v -s tests/v1/distributed/test_dbo.py

 ##### B200 test #####
 - label: Distributed Tests (B200) # optional
@@ -1639,8 +1644,8 @@ steps:
  - vllm/model_executor/layers/quantization/mxfp4.py
  - vllm/v1/attention/backends/flashinfer.py
  commands:
-  - uv pip install --system 'gpt-oss[eval]==0.0.5'
-  - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-gfx942.txt
+    - uv pip install --system 'gpt-oss[eval]==0.0.5'
+    - VLLM_ROCM_USE_AITER_MHA=0 VLLM_ROCM_USE_AITER=1 VLLM_USE_AITER_UNIFIED_ATTENTION=1 pytest -s -v tests/evals/gpt_oss/test_gpqa_correctness.py --model openai/gpt-oss-20b --metric 0.58

 ##### EPLB Accuracy Tests #####
 - label: DeepSeek V2-Lite Accuracy
@@ -1667,6 +1672,16 @@ steps:
  commands:
  - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020

+- label: Qwen3-30B-A3B-FP8-block Accuracy (B200)
+  timeout_in_minutes: 60
+  gpu: b200
+  optional: true
+  num_gpus: 2
+  working_dir: "/vllm-workspace"
+  commands:
+  - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020 2 1
+
+
 - label: Qwen3-Next-80B-A3B-Instruct MTP Async EPLB Accuracy
  timeout_in_minutes: 60
  mirror_hardwares: [amdexperimental, amdproduction]
@@ -1678,93 +1693,6 @@ steps:
  commands:
  - bash .buildkite/scripts/scheduled_integration_test/qwen3_next_mtp_async_eplb.sh 0.8 1319 8040

-##### .buildkite/test_areas/compile.yaml #####
-# Slowly setting up the tests so that it is also easier for the 
-# CI team to review and upstream to the pipelinev2.
-# The following tests are important for vLLM IR Ops refactoring,
-# which affects fusion passes on ROCm. So we have to 
-# enable them as as soon as possible.
-
-## TODO: Enable the test in this group
-# # corresponds to .buildkite/test_areas/compile.yaml
-# - label: Fusion and Compile Unit Tests (2xMI325 GPUs)
-#   timeout_in_minutes: 20
-#   working_dir: "/vllm-workspace/"
-#   mirror_hardwares: [amdexperimental, amdproduction, tj]
-#   agent_pool: mi325_1 # changed to 1 GPU until the fusion all reduce is enabled then only revert back to 2 GPUs
-#   source_file_dependencies:
-#   - csrc/quantization/fp4/
-#   - vllm/model_executor/layers/quantization/
-#   - vllm/model_executor/layers/layernorm.py
-#   - vllm/model_executor/layers/activation.py
-#   - vllm/model_executor/layers/attention/attention.py
-#   - vllm/v1/attention/backends/flashinfer.py
-#   - vllm/compilation/ # TODO(luka) limit to vllm/compilation/passes
-#   - tests/compile/test_fusion_attn.py
-#   - tests/compile/test_silu_mul_quant_fusion.py
-#   - tests/compile/distributed/test_fusion_all_reduce.py
-#   - tests/compile/fullgraph/test_full_graph.py
-#   commands:
-#     - rocm-smi
-#     # we run all backend tests on ROCm
-#     # These two tests are covered in "PyTorch Compilation Passes Unit Tests"
-#     # - "pytest -v -s tests/compile/passes/test_fusion_attn.py"
-#     # - "pytest -v -s tests/compile/passes/test_silu_mul_quant_fusion.py"
-#     # TODO: this test is not supported on ROCm, there are aiter kernels for this.
-#     # - pytest -v -s tests/compile/passes/distributed/test_fusion_all_reduce.py
-#     # TODO: find out more details
-#     # - pytest -v -s tests/compile/fullgraph/test_full_graph.py::test_fp8_kv_scale_compile
-
-# corresponds to .buildkite/test_areas/compile.yaml
- label: Fusion E2E Quick (MI325)
-  timeout_in_minutes: 15
-  working_dir: "/vllm-workspace/"
-  mirror_hardwares: [amdexperimental, amdproduction]
-  agent_pool: mi325_1
-  num_devices: 1
-  source_file_dependencies:
-    - csrc/quantization/
-    - vllm/model_executor/
-    - vllm/v1/attention/
-    - vllm/compilation/
-    - tests/compile/fusions_e2e/
-  commands:
-    - rocm-smi
-    # Run all models and attn backends but only Inductor partition and native custom ops
-    - "pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k 'inductor_partition and not +rms_norm and not +quant_fp8'"
-    # Different from CUDA, Qwen requires +rms_norm and +quant_fp8 as rms+quant fusion is only supported on AITER
-    - "pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k 'inductor_partition and +rms_norm and +quant_fp8 and qwen3'"
-
-# corresponds to .buildkite/test_areas/compile.yaml
- label: Fusion E2E Config Sweep (MI325)
-  timeout_in_minutes: 30
-  working_dir: "/vllm-workspace/"
-  mirror_hardwares: [amdexperimental, amdproduction]
-  agent_pool: mi325_1
-  num_devices: 1
-  source_file_dependencies:
-    - csrc/quantization/
-    - vllm/compilation/
-    # can affect pattern matching
-    - vllm/model_executor/layers/layernorm.py
-    - vllm/model_executor/layers/activation.py
-    - vllm/model_executor/layers/attention/attention.py
-    - vllm/model_executor/layers/quantization/input_quant_fp8.py
-    - tests/compile/fusions_e2e/
-  commands:
-    - rocm-smi
-    # Run just llama3 (fp8) for all config combinations
-    - pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "llama-3"
-
-## There are no ops on ROCm for these tests.
-## The test still passes but the logs are not useful.
-## fused ops just call torch.ops.symm_mem which 
-## exists in ROCm even though they don't work
-# - label: AsyncTP Correctness Tests  (2xMI325 GPUs)
-# - label: Fusion E2E TP2 Quick (MI325)
-# - label: Fusion E2E TP2 AsyncTP Config Sweep (MI325)
-# - label: Fusion E2E TP2 (MI325)
-# - label: Sequence Parallel Correctness Tests (2xMI325 GPUs)


 #####################################################################################################################################
@@ -1947,10 +1875,8 @@ steps:

 - label: Distributed Tests (4 GPUs) # 35min
  timeout_in_minutes: 50
-  mirror_hardwares: [amdexperimental, amdproduction]
+  mirror_hardwares: [amdexperimental]
  agent_pool: mi355_4
-  optional: true
-  # grade: Blocking
  working_dir: "/vllm-workspace/tests"
  num_gpus: 4
  source_file_dependencies:
@@ -2004,8 +1930,7 @@ steps:
  - popd
  # NEW rlhf examples
  - pushd ../examples/offline_inference/new_weight_syncing
-  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_nccl.py
-  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_ipc.py
+  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf.py
  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_async_new_apis.py
  - popd

@@ -2150,7 +2075,20 @@ steps:
    - pip install -U git+https://github.com/robertgshaw2-redhat/lm-evaluation-harness.git@streaming-api
    - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine

-# TODO: Add the "V1 Test attention (MI300)" test group
+# TODO: Add the "V1 Test attetion (MI300)" test group
+
+- label: V1 Test attention (H100) # 10min
+  mirror_hardwares: [amdexperimental]
+  agent_pool: mi355_1
+  timeout_in_minutes: 30
+  gpu: h100
+  source_file_dependencies:
+    - vllm/config/attention.py
+    - vllm/model_executor/layers/attention
+    - vllm/v1/attention
+    - tests/v1/attention
+  commands:
+    - pytest -v -s v1/attention

 - label: Batch Invariance Tests (H100) # 10min
  mirror_hardwares: [amdexperimental]
@@ -2168,8 +2106,6 @@ steps:
    - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py

 - label: V1 Test attention (B200) # 10min
-  mirror_hardwares: [amdexperimental, amdmi355]
-  agent_pool: mi355_1
  timeout_in_minutes: 30
  gpu: b200
  source_file_dependencies:
@@ -2208,12 +2144,12 @@ steps:
  commands:
    - pip install tensorizer # for tensorizer test
    # for basic
-    - python3 basic/offline_inference/chat.py --attention-backend TRITON_ATTN
-    - python3 basic/offline_inference/generate.py --model facebook/opt-125m
-    - python3 basic/offline_inference/generate.py --model meta-llama/Llama-2-13b-chat-hf --cpu-offload-gb 10
-    - python3 basic/offline_inference/classify.py
-    - python3 basic/offline_inference/embed.py
-    - python3 basic/offline_inference/score.py
+    - python3 offline_inference/basic/chat.py
+    - python3 offline_inference/basic/generate.py --model facebook/opt-125m
+    - python3 offline_inference/basic/generate.py --model meta-llama/Llama-2-13b-chat-hf --cpu-offload-gb 10
+    - python3 offline_inference/basic/classify.py
+    - python3 offline_inference/basic/embed.py
+    - python3 offline_inference/basic/score.py
    # for multi-modal models
    - python3 offline_inference/audio_language.py --seed 0
    - python3 offline_inference/vision_language.py --seed 0
@@ -2789,14 +2725,12 @@ steps:
    - pytest -v -s tests/models/test_transformers.py
    # - pytest -v -s tests/models/multimodal/processing/
    - pytest -v -s tests/models/multimodal/test_mapping.py -k 'not (Gemma3 or Qwen2VL or Qwen2_5_VL)'
-    - python3 examples/basic/offline_inference/chat.py
+    - python3 examples/offline_inference/basic/chat.py
    # - python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl
    # Whisper needs spawn method to avoid deadlock
    - VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper

- label: Blackwell Test (MI355) # 21 min
-  mirror_hardwares: [amdexperimental, amdmi355]
-  agent_pool: mi355_1
+- label: Blackwell Test # 21 min
  timeout_in_minutes: 30
  working_dir: "/vllm-workspace/"
  gpu: b200
@@ -2815,28 +2749,28 @@ steps:
  - vllm/v1/attention/selector.py
  - vllm/platforms/cuda.py
  commands:
-    - rocm-smi
-    - python3 examples/basic/offline_inference/chat.py
+    - nvidia-smi
+    - python3 examples/offline_inference/basic/chat.py
    # Attention
    # num_heads2 broken by https://github.com/flashinfer-ai/flashinfer/issues/1353
-    - pytest -v -s tests/kernels/attention/test_attention_selector.py 
-    #- pytest -v -s tests/kernels/attention/test_flashinfer.py -k 'not num_heads2'
-    #- pytest -v -s tests/kernels/attention/test_flashinfer_trtllm_attention.py
-    #- pytest -v -s tests/kernels/attention/test_cutlass_mla_decode.py
-    #- pytest -v -s tests/kernels/attention/test_flashinfer_mla_decode.py
-    ## Quantization
-    #- pytest -v -s tests/kernels/quantization/test_cutlass_scaled_mm.py -k 'fp8'
-    #- pytest -v -s tests/kernels/quantization/test_nvfp4_quant.py
-    #- pytest -v -s tests/kernels/quantization/test_silu_mul_nvfp4_quant.py
-    #- pytest -v -s tests/kernels/quantization/test_nvfp4_scaled_mm.py
-    #- pytest -v -s tests/kernels/quantization/test_flashinfer_scaled_mm.py
-    #- pytest -v -s tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py
-    #- pytest -v -s tests/kernels/quantization/test_nvfp4_qutlass.py
-    #- pytest -v -s tests/kernels/quantization/test_mxfp4_qutlass.py
-    #- pytest -v -s tests/kernels/moe/test_nvfp4_moe.py
-    #- pytest -v -s tests/kernels/moe/test_ocp_mx_moe.py
-    #- pytest -v -s tests/kernels/moe/test_flashinfer.py
-    #- pytest -v -s tests/kernels/moe/test_cutedsl_moe.py
+    - pytest -v -s tests/kernels/attention/test_attention_selector.py
+    - pytest -v -s tests/kernels/attention/test_flashinfer.py -k 'not num_heads2'
+    - pytest -v -s tests/kernels/attention/test_flashinfer_trtllm_attention.py
+    - pytest -v -s tests/kernels/attention/test_cutlass_mla_decode.py
+    - pytest -v -s tests/kernels/attention/test_flashinfer_mla_decode.py
+    # Quantization
+    - pytest -v -s tests/kernels/quantization/test_cutlass_scaled_mm.py -k 'fp8'
+    - pytest -v -s tests/kernels/quantization/test_nvfp4_quant.py
+    - pytest -v -s tests/kernels/quantization/test_silu_mul_nvfp4_quant.py
+    - pytest -v -s tests/kernels/quantization/test_nvfp4_scaled_mm.py
+    - pytest -v -s tests/kernels/quantization/test_flashinfer_scaled_mm.py
+    - pytest -v -s tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py
+    - pytest -v -s tests/kernels/quantization/test_nvfp4_qutlass.py
+    - pytest -v -s tests/kernels/quantization/test_mxfp4_qutlass.py
+    - pytest -v -s tests/kernels/moe/test_nvfp4_moe.py
+    - pytest -v -s tests/kernels/moe/test_ocp_mx_moe.py
+    - pytest -v -s tests/kernels/moe/test_flashinfer.py
+    - pytest -v -s tests/kernels/moe/test_cutedsl_moe.py

 - label: Blackwell Fusion and Compile Tests # 30 min
  timeout_in_minutes: 40
@@ -2906,15 +2840,13 @@ steps:

 - label: Blackwell LM Eval Small Models
  timeout_in_minutes: 120
-  mirror_hardwares: [amdexperimental, amdproduction, amdmi355]
-  agent_pool: mi355_2
  gpu: b200
  optional: true # run on nightlies
  source_file_dependencies:
  - csrc/
  - vllm/model_executor/layers/quantization
  commands:
-  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi355.txt
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-blackwell.txt

 #####  1 GPU test  #####
 #####  multi gpus test  #####
@@ -2962,10 +2894,8 @@ steps:

 - label: Distributed Tests (2 GPUs) # 68min
  timeout_in_minutes: 90
-  mirror_hardwares: [amdexperimental, amdproduction]
+  mirror_hardwares: [amdexperimental]
  agent_pool: mi355_2
-  optional: true
-  # grade: Blocking
  working_dir: "/vllm-workspace/tests"
  num_gpus: 2
  source_file_dependencies:
@@ -3150,20 +3080,6 @@ steps:
    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
    - DP_EP=1 ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh

- label: CrossLayer KV layout Distributed NixlConnector PD accuracy tests (4 GPUs)
-  mirror_hardwares: [amdexperimental, amdproduction]
-  agent_pool: mi355_4
-  # grade: Blocking
-  timeout_in_minutes: 30
-  working_dir: "/vllm-workspace/tests"
-  num_devices: 4
-  source_file_dependencies:
-    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
-    - tests/v1/kv_connector/nixl_integration/
-  commands:
-    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
-    - CROSS_LAYERS_BLOCKS=1 ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
-
 ##### multi gpus test #####
 ##### A100 test #####

@@ -3296,8 +3212,8 @@ steps:
  - vllm/model_executor/layers/quantization/mxfp4.py
  - vllm/v1/attention/backends/flashinfer.py
  commands:
-  - uv pip install --system 'gpt-oss[eval]==0.0.5'
-  - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-gfx950.txt
+    - uv pip install --system 'gpt-oss[eval]==0.0.5'
+    - VLLM_ROCM_USE_AITER_MHA=0 VLLM_ROCM_USE_AITER=1 VLLM_USE_AITER_UNIFIED_ATTENTION=1 pytest -s -v tests/evals/gpt_oss/test_gpqa_correctness.py --model openai/gpt-oss-20b --metric 0.58

 ##### EPLB Accuracy Tests #####
 - label: DeepSeek V2-Lite Accuracy
@@ -3311,9 +3227,18 @@ steps:
  commands:
  - bash .buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_ep_eplb.sh 0.25 200 8010

- label: Qwen3-30B-A3B-FP8-block Accuracy (B200-MI355)
-  mirror_hardwares: [amdexperimental, amdproduction, amdmi355]
-  agent_pool: mi355_2
+- label: Qwen3-30B-A3B-FP8-block Accuracy (H100)
+  mirror_hardwares: [amdexperimental, amdproduction]
+  agent_pool: mi355_4
+  timeout_in_minutes: 60
+  gpu: h100
+  optional: true
+  num_gpus: 4
+  working_dir: "/vllm-workspace"
+  commands:
+  - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020
+
+- label: Qwen3-30B-A3B-FP8-block Accuracy (B200)
  timeout_in_minutes: 60
  gpu: b200
  optional: true
@@ -3332,18 +3257,3 @@ steps:
  working_dir: "/vllm-workspace"
  commands:
  - bash .buildkite/scripts/scheduled_integration_test/qwen3_next_mtp_async_eplb.sh 0.8 1319 8040
-
- label: Attention Benchmarks Smoke Test (B200-MI355)
-  device: b200
-  mirror_hardwares: [amdexperimental, amdmi355]
-  agent_pool: mi355_2
-  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 ROCM_ATTN ROCM_AITER_FA ROCM_AITER_UNIFIED_ATTN --batch-specs "8q1s1k" --repeats 1 --warmup-iters 1
-
@@ -36,16 +36,6 @@ steps:
  - export VLLM_TEST_CLEAN_GPU_MEMORY=1
  - pytest -v -s tests/compile/correctness_e2e/test_async_tp.py

- label: AsyncTP Correctness Tests (B200)
-  timeout_in_minutes: 50
-  working_dir: "/vllm-workspace/"
-  device: b200
-  optional: true
-  num_devices: 2
-  commands:
-  - export VLLM_TEST_CLEAN_GPU_MEMORY=1
-  - pytest -v -s tests/compile/correctness_e2e/test_async_tp.py
-
 - label: Distributed Compile Unit Tests (2xH100)
  timeout_in_minutes: 20
  working_dir: "/vllm-workspace/"
@@ -67,7 +67,6 @@ steps:
  - tests/v1/distributed
  - tests/v1/engine/test_engine_core_client.py
  - tests/distributed/test_symm_mem_allreduce.py
-  - tests/distributed/test_multiproc_executor.py
  commands:
  # https://github.com/NVIDIA/nccl/issues/1838
  - export NCCL_CUMEM_HOST_ENABLE=0
@@ -96,8 +95,6 @@ steps:
  - pytest -v -s distributed/test_pynccl.py
  - pytest -v -s distributed/test_events.py
  - pytest -v -s distributed/test_symm_mem_allreduce.py
-  # test multi-node TP with multiproc executor (simulated on single node)
-  - pytest -v -s distributed/test_multiproc_executor.py::test_multiproc_executor_multi_node
  # TODO: create a dedicated test section for multi-GPU example tests
  # when we have multiple distributed example tests
  # OLD rlhf examples
@@ -149,7 +146,7 @@ steps:
  num_devices: 2
  commands:
    - pytest -v -s tests/distributed/test_context_parallel.py
-    # - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/offline_inference/new_weight_syncing/rlhf_async_new_apis.py --- failing, need to re-enable
+    - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/offline_inference/new_weight_syncing/rlhf_async_new_apis.py
    - VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
    - pytest -v -s tests/v1/distributed/test_dbo.py

@@ -213,19 +210,6 @@ steps:
    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
    - CROSS_LAYERS_BLOCKS=True bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh

- label: NixlConnector PD + Spec Decode acceptance (2 GPUs)
-  timeout_in_minutes: 30
-  device: a100
-  working_dir: "/vllm-workspace/tests"
-  num_devices: 2
-  source_file_dependencies:
-    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
-    - vllm/v1/worker/kv_connector_model_runner_mixin.py
-    - tests/v1/kv_connector/nixl_integration/
-  commands:
-    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
-    - bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh
-
 - label: Pipeline + Context Parallelism (4 GPUs)
  timeout_in_minutes: 60
  working_dir: "/vllm-workspace/tests"
@@ -14,7 +14,7 @@ steps:
  commands:
  - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py

- label: V1 e2e + engine (1 GPU)
+- label: V1 e2e + engine
  timeout_in_minutes: 45
  source_file_dependencies:
    - vllm/
@@ -36,35 +36,3 @@ steps:
      commands:
      - pytest -v -s v1/e2e
      - pytest -v -s v1/engine
-
- label: V1 e2e (2 GPUs)
-  timeout_in_minutes: 60 # TODO: Fix timeout after we have more confidence in the test stability
-  optional: true
-  num_devices: 2
-  source_file_dependencies:
-    - vllm/
-    - tests/v1/e2e
-  commands:
-    # Only run tests that need exactly 2 GPUs
-    - pytest -v -s v1/e2e/test_spec_decode.py -k "tensor_parallelism"
-  mirror:
-    amd:
-      device: mi325_2
-      depends_on:
-      - image-build-amd
-
- label: V1 e2e (4 GPUs)
-  timeout_in_minutes: 60 # TODO: Fix timeout after we have more confidence in the test stability
-  optional: true
-  num_devices: 4
-  source_file_dependencies:
-    - vllm/
-    - tests/v1/e2e
-  commands:
-    # Only run tests that need 4 GPUs
-    - pytest -v -s v1/e2e/test_spec_decode.py -k "eagle_correctness_heavy"
-  mirror:
-    amd:
-      device: mi325_4
-      depends_on:
-      - image-build-amd
@@ -41,11 +41,6 @@ steps:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/test_oot_registration.py --ignore=entrypoints/openai/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/  --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses
  - pytest -v -s entrypoints/test_chat_utils.py
-  mirror:
-    amd:
-      device: mi325_1
-      depends_on:
-      - image-build-amd

 - label: Entrypoints Integration (API Server 2)
  timeout_in_minutes: 130
@@ -60,11 +55,6 @@ steps:
  - pytest -v -s entrypoints/instrumentator
  - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc
  - pytest -v -s tool_use
-  mirror:
-    amd:
-      device: mi325_1
-      depends_on:
-      - image-build-amd

 - label: Entrypoints Integration (Pooling)
  timeout_in_minutes: 50
@@ -97,11 +87,6 @@ steps:
    - tests/v1
  commands:
    - pytest -v -s v1/entrypoints
-  mirror:
-    amd:
-      device: mi325_1
-      depends_on:
-      - image-build-amd

 - label: OpenAI API Correctness
  timeout_in_minutes: 30
@@ -8,9 +8,8 @@ steps:
  - csrc/
  - tests/kernels/core
  - tests/kernels/test_top_k_per_row.py
-  - tests/kernels/test_concat_mla_q.py
  commands:
-    - pytest -v -s kernels/core kernels/test_top_k_per_row.py kernels/test_concat_mla_q.py
+    - pytest -v -s kernels/core kernels/test_top_k_per_row.py

 - label: Kernels Attention Test %N
  timeout_in_minutes: 35
@@ -45,8 +44,7 @@ steps:
  - vllm/envs.py
  - vllm/config
  commands:
-    - pytest -v -s kernels/moe --ignore=kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
-    - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
+    - pytest -v -s kernels/moe --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
  parallelism: 2

 - label: Kernels Mamba Test
@@ -97,7 +95,7 @@ steps:
  - vllm/platforms/cuda.py
  commands:
    - nvidia-smi
-    - python3 examples/basic/offline_inference/chat.py
+    - python3 examples/offline_inference/basic/chat.py
    # Attention
    # num_heads2 broken by https://github.com/flashinfer-ai/flashinfer/issues/1353
    - pytest -v -s tests/kernels/attention/test_attention_selector.py
@@ -11,17 +11,17 @@ steps:
  commands:
  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.txt

-# - label: LM Eval Large Models (4 GPUs)(A100)
-#   device: a100
-#   optional: true
-#   num_devices: 4
-#   working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
-#   source_file_dependencies:
-#   - csrc/
-#   - vllm/model_executor/layers/quantization
-#   commands:
-#   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-#   - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large.txt --tp-size=4
+- label: LM Eval Large Models (4 GPUs)(A100)
+  device: a100
+  optional: true
+  num_devices: 4
+  working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
+  source_file_dependencies:
+  - csrc/
+  - vllm/model_executor/layers/quantization
+  commands:
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large.txt --tp-size=4

 - label: LM Eval Large Models (4 GPUs)(H100)
  device: h100
@@ -67,13 +67,12 @@ steps:
  - examples/
  commands:
    - pip install tensorizer # for tensorizer test
-     # for basic
-    - python3 basic/offline_inference/chat.py
-    - python3 basic/offline_inference/generate.py --model facebook/opt-125m
-    - python3 basic/offline_inference/generate.py --model meta-llama/Llama-2-13b-chat-hf --cpu-offload-gb 10
-    - python3 basic/offline_inference/classify.py
-    - python3 basic/offline_inference/embed.py
-    - python3 basic/offline_inference/score.py
+    - python3 offline_inference/basic/chat.py # for basic
+    - python3 offline_inference/basic/generate.py --model facebook/opt-125m
+    - python3 offline_inference/basic/generate.py --model meta-llama/Llama-2-13b-chat-hf --cpu-offload-gb 10
+    - python3 offline_inference/basic/classify.py
+    - python3 offline_inference/basic/embed.py
+    - python3 offline_inference/basic/score.py
    # for multi-modal models
    - python3 offline_inference/audio_language.py --seed 0
    - python3 offline_inference/vision_language.py --seed 0
@@ -88,11 +87,6 @@ steps:
    - python3 offline_inference/spec_decode.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
    # https://github.com/vllm-project/vllm/pull/26682 uses slightly more memory in PyTorch 2.9+ causing this test to OOM in 1xL4 GPU
    - python3 offline_inference/spec_decode.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536
-  mirror:
-    amd:
-      device: mi325_1
-      depends_on:
-      - image-build-amd

 - label: Metrics, Tracing (2 GPUs)
  timeout_in_minutes: 20
@@ -65,7 +65,7 @@ steps:
    - pytest -v -s tests/models/test_transformers.py
    - pytest -v -s tests/models/multimodal/processing/
    - pytest -v -s tests/models/multimodal/test_mapping.py
-    - python3 examples/basic/offline_inference/chat.py
+    - python3 examples/offline_inference/basic/chat.py
    - python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl
    # Whisper needs spawn method to avoid deadlock
    - VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper
@@ -12,11 +12,6 @@ steps:
    - pip freeze | grep -E 'torch'
    - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/processing
    - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model  # Otherwise, mp_method="spawn" doesn't work
-  mirror:
-    amd:
-      device: mi325_1
-      depends_on:
-      - image-build-amd

 - label: Multi-Modal Processor Test (CPU)
  depends_on: 
@@ -25,7 +20,6 @@ steps:
  source_file_dependencies:
  - vllm/
  - tests/models/multimodal
-  - tests/models/registry.py
  device: cpu
  commands:
    - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
@@ -36,7 +30,6 @@ steps:
  source_file_dependencies:
  - vllm/
  - tests/models/multimodal
-  - tests/models/registry.py
  commands:
    - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
    - pytest -v -s models/multimodal/processing/test_tensor_schema.py
@@ -59,11 +52,6 @@ steps:
  commands:
    - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
    - pytest -v -s models/multimodal -m 'not core_model' --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/processing
-  mirror:
-    amd:
-      device: mi325_1
-      depends_on:
-      - image-build-amd

 - label: Multi-Modal Models (Extended) 2
  optional: true
@@ -82,3 +70,12 @@ steps:
  commands:
    - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
    - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=1) and not core_model'
+
+# This test is used only in PR development phase to test individual models and should never run on main
+- label: Custom Models
+  optional: true
+  commands:
+    - echo 'Testing custom models...'
+    # PR authors can temporarily add commands below to test individual models
+    # e.g. pytest -v -s models/encoder_decoder/vision_language/test_mllama.py
+    # *To avoid merge conflicts, remember to REMOVE (not just comment out) them before merging the PR*
@@ -15,12 +15,9 @@ steps:
  - pytest -v -s plugins_tests/test_platform_plugins.py
  - pip uninstall vllm_add_dummy_platform -y
  # end platform plugin tests
-  # begin io_processor plugins test
-  # test generic io_processor plugins functions
-  - pytest -v -s ./plugins_tests/test_io_processor_plugins.py
-  # test Terratorch io_processor plugins
+  # begin io_processor plugins test, all the code in between uses the prithvi_io_processor plugin
  - pip install -e ./plugins/prithvi_io_processor_plugin
-  - pytest -v -s plugins_tests/test_terratorch_io_processor_plugins.py
+  - pytest -v -s plugins_tests/test_io_processor_plugins.py
  - pip uninstall prithvi_io_processor_plugin -y
  # test bge_m3_sparse io_processor plugin
  - pip install -e ./plugins/bge_m3_sparse_plugin
@@ -39,8 +36,3 @@ steps:
  - pytest -v -s entrypoints/openai/test_oot_registration.py # it needs a clean process
  - pytest -v -s models/test_oot_registration.py # it needs a clean process
  - pytest -v -s plugins/lora_resolvers # unit tests for in-tree lora resolver plugins
-  mirror:
-    amd:
-      device: mi325_2
-      depends_on:
-      - image-build-amd
@@ -1,16 +0,0 @@
-group: Ray Compatibility
-depends_on:
-  - image-build
-steps:
- label: Ray Dependency Compatibility Check
-  # Informational only — does not block the pipeline.
-  # If this fails, it means the PR introduces a dependency that
-  # conflicts with Ray's dependency constraints.
-  # See https://github.com/vllm-project/vllm/issues/33599
-  soft_fail: true
-  timeout_in_minutes: 10
-  source_file_dependencies:
-  - requirements/
-  - setup.py
-  commands:
-  - bash /vllm-workspace/.buildkite/scripts/check-ray-compatibility.sh
@@ -13,13 +13,13 @@ steps:
  commands:
    - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models.txt

-# - label: Weight Loading Multiple GPU - Large Models # optional
-#   working_dir: "/vllm-workspace/tests"
-#   num_devices: 2
-#   device: a100
-#   optional: true
-#   source_file_dependencies:
-#   - vllm/
-#   - tests/weight_loading
-#   commands:
-#     - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-large.txt
+- label: Weight Loading Multiple GPU - Large Models # optional
+  working_dir: "/vllm-workspace/tests"
+  num_devices: 2
+  device: a100
+  optional: true
+  source_file_dependencies:
+  - vllm/
+  - tests/weight_loading
+  commands:
+    - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-large.txt
@@ -3,7 +3,6 @@ pull_request_rules:
  description: Automatically apply documentation label
  conditions:
    - label != stale
-    - -closed
    - or:
      - files~=^[^/]+\.md$
      - files~=^docs/
@@ -38,13 +37,15 @@ pull_request_rules:

        > [!TIP]
        > <details>
-        > <summary>Is <code>mypy</code> failing?</summary>
+        > <summary>Is <code>mypy</code> or <code>markdownlint</code> failing?</summary>
        > <br/>
-        > <code>mypy</code> is run differently in CI. If the failure is related to this check, please use the following command to run it locally:
+        > <code>mypy</code> and <code>markdownlint</code> are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:
        >
        > ```bash
        > # For mypy (substitute "3.10" with the failing version if needed)
        > pre-commit run --hook-stage manual mypy-3.10
+        > # For markdownlint
+        > pre-commit run --hook-stage manual markdownlint
        > ```
        > </details>

@@ -6,9 +6,6 @@ on:
      - main
  workflow_dispatch:  # Manual trigger

-permissions:
-  contents: read
-
 jobs:
  macos-m1-smoke-test:
    runs-on: macos-latest
@@ -13,7 +13,7 @@ repos:
    args: [--output-format, github, --fix]
  - id: ruff-format
 - repo: https://github.com/crate-ci/typos
-  rev: v1.43.5
+  rev: v1.38.1
  hooks:
  - id: typos
    args: [--force-exclude]
@@ -24,12 +24,12 @@ repos:
    exclude: 'csrc/(moe/topk_softmax_kernels.cu|quantization/gguf/(ggml-common.h|dequantize.cuh|vecdotq.cuh|mmq.cuh|mmvq.cuh))|vllm/third_party/.*'
    types_or: [c++, cuda]
    args: [--style=file, --verbose]
- repo: https://github.com/DavidAnson/markdownlint-cli2
-  rev: v0.21.0
+- repo: https://github.com/igorshubovych/markdownlint-cli
+  rev: v0.45.0
  hooks:
-  - id: markdownlint-cli2
-    language_version: lts
-    args: [--fix]
+  - id: markdownlint
+    exclude: '.*\.inc\.md'
+    stages: [manual] # Only run in CI
 - repo: https://github.com/rhysd/actionlint
  rev: v1.7.7
  hooks:
@@ -55,7 +55,7 @@ repos:
      language: python
      types_or: [python, pyi]
      require_serial: true
-      additional_dependencies: ["mypy[faster-cache]==1.19.1", regex, types-cachetools, types-setuptools, types-PyYAML, types-requests, types-torch, pydantic]
+      additional_dependencies: [mypy==1.11.1, regex, types-cachetools, types-setuptools, types-PyYAML, types-requests, types-torch, pydantic]
  - id: mypy-3.10 # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup is less awkward
    name: Run mypy for Python 3.10
    entry: python tools/pre_commit/mypy.py 1 "3.10"
@@ -127,13 +127,6 @@ repos:
    language: python
    types: [python]
    additional_dependencies: [regex]
-  # prevent use torch.cuda APIs
-  - id: check-torch-cuda-call
-    name: "Prevent new 'torch.cuda' APIs call"
-    entry: python tools/pre_commit/check_torch_cuda.py
-    language: python
-    types: [python]
-    additional_dependencies: [regex]
  - id: validate-config
    name: Validate configuration has default values and that each field has a docstring
    entry: python tools/pre_commit/validate_config.py
@@ -9,7 +9,6 @@ build:
    python: "3.12"
  jobs:
    post_checkout:
-      - bash docs/maybe_skip_pr_build.sh
      - git fetch origin main --unshallow --no-tags --filter=blob:none || true
    pre_create_environment:
      - pip install uv
@@ -771,33 +771,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
    endif()
  endif()

-  # Expert-specialization MXFP8 blockscaled grouped kernels (SM100+).
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(ES_MXFP8_GROUPED_MM_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
-  else()
-    cuda_archs_loose_intersection(ES_MXFP8_GROUPED_MM_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
-  endif()
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND ES_MXFP8_GROUPED_MM_ARCHS)
-    set(SRCS
-      "csrc/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm.cu"
-      "csrc/moe/mxfp8_moe/mxfp8_experts_quant.cu")
-    set_gencode_flags_for_srcs(
-      SRCS "${SRCS}"
-      CUDA_ARCHS "${ES_MXFP8_GROUPED_MM_ARCHS}")
-    list(APPEND VLLM_EXT_SRC "${SRCS}")
-    list(APPEND VLLM_GPU_FLAGS "-DENABLE_ES_MXFP8_GROUPED_MM_SM100=1")
-    message(STATUS "Building ES MXFP8 grouped kernels for archs: ${ES_MXFP8_GROUPED_MM_ARCHS}")
-  else()
-    if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8
-        AND ES_MXFP8_GROUPED_MM_ARCHS)
-      message(STATUS "Not building ES MXFP8 grouped kernels as CUDA Compiler version is "
-                     "not >= 12.8.")
-    else()
-      message(STATUS "Not building ES MXFP8 grouped kernels as no compatible archs found "
-                     "in CUDA target architectures.")
-    endif()
-  endif()
-
  # DeepSeek V3 fused A GEMM kernel (requires SM 9.0+, Hopper and later)
  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
    cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0f;11.0f" "${CUDA_ARCHS}")
@@ -187,7 +187,7 @@ python benchmark.py \
 ## Hardware Requirements

 | Backend | Hardware |
-| ------- | -------- |
+|---------|----------|
 | Flash/Triton/FlashInfer | Any CUDA GPU |
 | CUTLASS MLA | Blackwell (SM100+) |
 | FlashAttn MLA | Hopper (SM90+) |
@@ -30,7 +30,7 @@ def batch_spec_sort_key(spec: str) -> tuple[int, int, int]:
        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 unparsable specs
+        # Fallback for unparseable specs
        return (0, 0, 0)


@@ -145,6 +145,7 @@ def create_minimal_vllm_config(
    cache_config = CacheConfig(
        block_size=block_size,
        gpu_memory_utilization=0.9,
+        swap_space=0,
        cache_dtype="auto",
        enable_prefix_caching=False,
    )
@@ -700,7 +701,7 @@ def _run_single_benchmark(
    # Warmup
    for _ in range(config.warmup_iters):
        forward_fn()
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    # Benchmark
    times = []
@@ -713,7 +714,7 @@ def _run_single_benchmark(
            forward_fn()
        end.record()

-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()
        elapsed_ms = start.elapsed_time(end)
        times.append(elapsed_ms / 1000.0 / config.num_layers)

@@ -141,6 +141,7 @@ def _create_vllm_config(
    cache_config = CacheConfig(
        block_size=config.block_size,
        cache_dtype="auto",
+        swap_space=0,
    )
    cache_config.num_gpu_blocks = max_num_blocks
    cache_config.num_cpu_blocks = 0
@@ -390,7 +391,7 @@ def _run_single_benchmark(
                attn_metadata,
                output=out,
            )
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    # Benchmark
    times = []
@@ -411,7 +412,7 @@ def _run_single_benchmark(
            )
        end.record()

-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()
        elapsed_ms = start.elapsed_time(end)
        times.append(elapsed_ms / 1000.0 / config.num_layers)  # seconds per layer

@@ -41,7 +41,7 @@ MODEL=meta-llama/Llama-3.3-70B-Instruct SYSTEM=TPU TP=8 DOWNLOAD_DIR='' INPUT_LE
 | --- | --- | --- |
 | `BASE` | **Required.** The absolute path to the parent directory of your vLLM repository directory. | `"$HOME"` |
 | `MODEL` | **Required.** The Hugging Face model identifier to be served by vllm. | `"meta-llama/Llama-3.1-8B-Instruct"` |
-| `SYSTEM` | **Required.** The hardware you are running on. Choices: `TPU` or `GPU`. (For other systems, it might not support saving profiles) | `"TPU"` |
+| `SYSTEM`| **Required.** The hardware you are running on. Choices: `TPU` or `GPU`. (For other systems, it might not support saving profiles) | `"TPU"` |
 | `TP` | **Required.** The tensor-parallelism size. | `1` |
 | `DOWNLOAD_DIR` | **Required.** Directory to download and load model weights from. | `""` (default download path) |
 | `INPUT_LEN` | **Required.** Request input length. | `4000` |
@@ -85,6 +85,7 @@ start_server() {
    # Each argument and its value are separate elements.
    local common_args_array=(
        "$MODEL"
+        "--disable-log-requests"
        "--port" "8004"
        "--host" "$HOSTNAME"
        "--gpu-memory-utilization" "$gpu_memory_utilization"
@@ -94,7 +94,7 @@ def create_logits(

 def measure_memory() -> tuple[int, int]:
    """Return (allocated, reserved) memory in bytes."""
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()
    return torch.cuda.memory_allocated(), torch.cuda.max_memory_allocated()


@@ -102,7 +102,7 @@ def reset_memory_stats():
    """Reset peak memory statistics."""
    reset_buffer_cache()
    torch.cuda.reset_peak_memory_stats()
-    torch.accelerator.empty_cache()
+    torch.cuda.empty_cache()
    gc.collect()


@@ -123,7 +123,7 @@ def benchmark_function(
    for _ in range(warmup_iters):
        logits_copy = logits.clone()
        func(logits_copy, k, p)
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    # Reset memory stats before benchmark
    reset_memory_stats()
@@ -140,7 +140,7 @@ def benchmark_function(
        func(logits_copy, k, p)
        end_events[i].record()

-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    # Calculate timing
    times = [
@@ -1,98 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import argparse
-
-import torch
-
-from vllm import _custom_ops as ops
-from vllm.triton_utils import triton
-
-# DeepSeek V3 dimensions
-NOPE_DIM = 512
-ROPE_DIM = 64
-NUM_HEADS = 128
-
-NUM_TOKENS = [8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192]
-
-
-def get_configs():
-    return NUM_TOKENS
-
-
-def make_inputs(num_tokens, dtype):
-    """Create inputs matching the real code path.
-
-    Args:
-        contiguous_nope: If False, simulate the transposed BMM output
-                         (non-contiguous nope with stride pattern from
-                         [N,B,L].transpose(0,1)).
-    """
-    # Simulate: bmm output [N, B, L].transpose(0, 1) -> [B, N, L]
-    raw = torch.randn(NUM_HEADS, num_tokens, NOPE_DIM, dtype=dtype, device="cuda")
-    ql_nope = raw.transpose(0, 1)
-
-    q_pe = torch.randn(num_tokens, NUM_HEADS, ROPE_DIM, dtype=dtype, device="cuda")
-    return ql_nope, q_pe
-
-
-# ---- Non-contiguous nope benchmark (real code path) ----
-@triton.testing.perf_report(
-    triton.testing.Benchmark(
-        x_names=["num_tokens"],
-        x_vals=get_configs(),
-        line_arg="provider",
-        line_vals=["torch_cat", "concat_mla_q"],
-        line_names=["torch.cat", "concat_mla_q (v8)"],
-        styles=[("blue", "--"), ("green", "-")],
-        ylabel="Latency (us)",
-        plot_name="concat_mla_q-transposed",
-        args={},
-    )
-)
-def bench_transposed(num_tokens, provider):
-    dtype = torch.bfloat16
-    ql_nope, q_pe = make_inputs(num_tokens, dtype)
-
-    q_out = torch.empty(
-        num_tokens, NUM_HEADS, NOPE_DIM + ROPE_DIM, dtype=dtype, device="cuda"
-    )
-
-    quantiles = [0.5, 0.2, 0.8]
-
-    if provider == "torch_cat":
-        ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
-            lambda: torch.cat((ql_nope, q_pe), dim=-1), quantiles=quantiles, rep=500
-        )
-    else:
-        ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
-            lambda: ops.concat_mla_q(ql_nope, q_pe, q_out), quantiles=quantiles, rep=500
-        )
-
-    return ms * 1000, max_ms * 1000, min_ms * 1000  # us
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Benchmark concat_mla_q vs torch.cat")
-    parser.add_argument(
-        "--save-path", type=str, default=None, help="Path to save benchmark results"
-    )
-    args = parser.parse_args()
-
-    print("\n" + "=" * 70)
-    print("CONCAT MLA Q KERNEL BENCHMARKS")
-    print("=" * 70)
-    print(f"Dimensions: nope={NOPE_DIM}, rope={ROPE_DIM}, heads={NUM_HEADS}")
-    print(
-        f"Per-head output: {NOPE_DIM + ROPE_DIM} bf16 = "
-        f"{(NOPE_DIM + ROPE_DIM) * 2} bytes"
-    )
-    print(f"num_tokens (decode=batch_size, prefill=chunk_size): {NUM_TOKENS}")
-    print("=" * 70)
-
-    print("\n--- Non-contiguous nope inputs (transposed BMM output) ---")
-    bench_transposed.run(print_data=True, save_path=args.save_path)
-
-    print("\n" + "=" * 70)
-    print("Benchmarking complete!")
-    print("=" * 70)
@@ -1,153 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import argparse
-import math
-
-import torch
-
-from vllm import _custom_ops as ops
-from vllm.triton_utils import triton
-
-# DeepSeek V3 MLA dimensions
-NOPE_DIM = 512
-ROPE_DIM = 64
-HEAD_DIM = NOPE_DIM + ROPE_DIM  # 576 BF16 output elements per token
-ENTRY_BYTES = 656  # 512 FP8 + 16 scales + 128 BF16 RoPE
-BLOCK_SIZE = 64  # tokens per physical cache block - get_supported_kernel_block_sizes
-
-# Realistic prefill scenarios:
-#   - 1 long prefill: single request, 16K-96K tokens
-#   - 4 medium prefills: 4 requests, 4K-24K tokens each
-#   - 16 shorter prefills: 16 requests, 1K-6K tokens each
-SCENARIOS = [
-    # (label, num_reqs, total_tokens_list)
-    ("1-req", 1, [8192, 16384, 32768, 65536, 98304]),
-    ("4-reqs", 4, [8192, 16384, 32768, 65536, 98304]),
-    ("16-reqs", 16, [8192, 16384, 32768, 65536, 98304]),
-]
-
-
-def make_inputs(total_tokens, num_reqs, block_size):
-    """Create synthetic FP8 cache, block table, and output buffer.
-
-    Fills the cache with random bytes (we only measure throughput,
-    not correctness). Block table maps each request to contiguous
-    physical blocks.
-    """
-    # Divide tokens evenly across requests
-    base_len = total_tokens // num_reqs
-    remainder = total_tokens % num_reqs
-    seq_lens = [base_len + (1 if r < remainder else 0) for r in range(num_reqs)]
-
-    # workspace_starts: cumulative sum of seq_lens
-    workspace_starts = [0] * num_reqs
-    for r in range(1, num_reqs):
-        workspace_starts[r] = workspace_starts[r - 1] + seq_lens[r - 1]
-
-    # Physical blocks needed per request
-    blocks_per_req = [math.ceil(s / block_size) for s in seq_lens]
-    total_blocks = sum(blocks_per_req)
-    max_blocks = max(blocks_per_req)
-
-    # Allocate cache with random data (content doesn't matter for perf)
-    cache = torch.randint(
-        0,
-        256,
-        (total_blocks, block_size, ENTRY_BYTES),
-        dtype=torch.uint8,
-        device="cuda",
-    )
-
-    # Block table: contiguous block assignments
-    block_table = torch.zeros(num_reqs, max_blocks, dtype=torch.int32, device="cuda")
-    block_idx = 0
-    for r in range(num_reqs):
-        for b in range(blocks_per_req[r]):
-            block_table[r, b] = block_idx
-            block_idx += 1
-
-    # Output workspace
-    dst = torch.zeros(total_tokens, HEAD_DIM, dtype=torch.bfloat16, device="cuda")
-
-    seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device="cuda")
-    workspace_starts_t = torch.tensor(
-        workspace_starts, dtype=torch.int32, device="cuda"
-    )
-
-    return cache, dst, block_table, seq_lens_t, workspace_starts_t
-
-
-def bench_scenario(label, num_reqs, total_tokens_list, save_path):
-    """Run benchmark for a specific (num_reqs, total_tokens) scenario."""
-
-    @triton.testing.perf_report(
-        triton.testing.Benchmark(
-            x_names=["total_tokens"],
-            x_vals=total_tokens_list,
-            line_arg="provider",
-            line_vals=["cuda_kernel"],
-            line_names=["cp_gather_fp8 (CUDA)"],
-            styles=[("green", "-")],
-            ylabel="Latency (us)",
-            plot_name=f"cp_gather_fp8-{label}-bs{BLOCK_SIZE}",
-            args={"num_reqs": num_reqs},
-        )
-    )
-    def bench_fn(total_tokens, provider, num_reqs):
-        cache, dst, block_table, seq_lens_t, ws_starts = make_inputs(
-            total_tokens, num_reqs, BLOCK_SIZE
-        )
-
-        quantiles = [0.5, 0.2, 0.8]
-
-        ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
-            lambda: ops.cp_gather_and_upconvert_fp8_kv_cache(
-                cache, dst, block_table, seq_lens_t, ws_starts, num_reqs
-            ),
-            quantiles=quantiles,
-            rep=500,
-        )
-
-        return ms * 1000, max_ms * 1000, min_ms * 1000  # us
-
-    seq_len_per_req = total_tokens_list[0] // num_reqs
-    seq_len_per_req_max = total_tokens_list[-1] // num_reqs
-    print(
-        f"\n--- {label}: {num_reqs} request(s), "
-        f"~{seq_len_per_req}-{seq_len_per_req_max} tokens/req ---"
-    )
-    bench_fn.run(print_data=True, save_path=save_path)
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(
-        description="Benchmark cp_gather_and_upconvert_fp8_kv_cache"
-    )
-    parser.add_argument(
-        "--save-path",
-        type=str,
-        default=None,
-        help="Path to save benchmark results as CSV",
-    )
-    args = parser.parse_args()
-
-    # Print data volume info for bandwidth analysis
-    read_per_token = ENTRY_BYTES  # 656 bytes from cache
-    write_per_token = HEAD_DIM * 2  # 576 * 2 = 1152 bytes to workspace
-    total_per_token = read_per_token + write_per_token  # 1808 bytes
-
-    print("\n" + "=" * 70)
-    print("CP_GATHER_AND_UPCONVERT_FP8_KV_CACHE BENCHMARKS")
-    print("=" * 70)
-    print(f"Cache entry: {ENTRY_BYTES} bytes (512 FP8 + 16 scales + 128 RoPE)")
-    print(f"Output row:  {HEAD_DIM} BF16 = {HEAD_DIM * 2} bytes")
-    print(f"Per token:   {total_per_token} bytes (read + write)")
-    print(f"Block size:  {BLOCK_SIZE} tokens/block")
-    print("=" * 70)
-
-    for label, num_reqs, total_tokens_list in SCENARIOS:
-        bench_scenario(label, num_reqs, total_tokens_list, args.save_path)
-
-    print("\n" + "=" * 70)
-    print("Benchmarking complete!")
-    print("=" * 70)
@@ -168,7 +168,7 @@ def bench_impl(
    # warmup
    for kwargs in kwargs_list:
        impl_type.get_impl()(**kwargs)
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    # Merge into a single kwargs and qualify arguments as ArgPool
    kwargs = {k: ArgPool([]) for k in kwargs_list[0]}
@@ -202,7 +202,7 @@ def test_correctness(T: int, N: int):
    # reference output
    ref_out_q, ref_out_s = output_from_impl(ImplType.REFERENCE)

-    # test output
+    # test ouptut
    out_q, out_s = output_from_impl(
        ImplType.SILU_MUL_PER_TOKEN_GROUP_QUANT_FP8_COLMAJOR
    )
@@ -12,12 +12,12 @@ 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.all2all_utils import (
-    maybe_make_prepare_finalize,
-)
 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
+from vllm.model_executor.layers.fused_moe.prepare_finalize import (
+    MoEPrepareAndFinalizeNoEP,
+)
 from vllm.platforms import current_platform
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 from vllm.v1.worker.workspace import init_workspace_manager
@@ -137,21 +137,15 @@ def bench_run(
        per_out_ch_quant=per_out_ch,
    )

-    moe_config = make_dummy_moe_config(
-        num_experts=num_experts,
-        hidden_dim=k,
-        intermediate_size_per_partition=n,
-        in_dtype=a.dtype,
-    )
-    fn = mk.FusedMoEKernel(
-        maybe_make_prepare_finalize(
-            moe=moe_config,
-            quant_config=quant_config,
-            allow_new_interface=True,
-            use_monolithic=False,
-        ),
+    fn = mk.FusedMoEModularKernel(
+        MoEPrepareAndFinalizeNoEP(),
        CutlassExpertsFp8(
-            moe_config=moe_config,
+            moe_config=make_dummy_moe_config(
+                num_experts=num_experts,
+                hidden_dim=k,
+                intermediate_size_per_partition=n,
+                in_dtype=a.dtype,
+            ),
            quant_config=quant_config,
        ),
    )
@@ -171,7 +165,7 @@ def bench_run(
                activation=MoEActivation.SILU,
                global_num_experts=num_experts,
            )
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    # Create CUDA graphs for Triton (match benchmark_moe.py pattern exactly)
    triton_stream = torch.cuda.Stream()
@@ -187,14 +181,14 @@ def bench_run(
                topk_ids,
                quant_config=quant_config,
            )
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    def bench_cuda_graph(graph, num_warmup=5, num_iters=100):
        """Benchmark CUDA graph using events like benchmark_moe.py"""
        # Warmup
        for _ in range(num_warmup):
            graph.replay()
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()

        # Timing
        start_event = torch.Event(enable_timing=True)
@@ -202,7 +196,7 @@ def bench_run(

        latencies = []
        for _ in range(num_iters):
-            torch.accelerator.synchronize()
+            torch.cuda.synchronize()
            start_event.record()
            graph.replay()
            end_event.record()
@@ -15,9 +15,6 @@ 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.config import ParallelConfig, VllmConfig, set_current_vllm_config
-from vllm.model_executor.layers.fused_moe.all2all_utils import (
-    maybe_make_prepare_finalize,
-)
 from vllm.model_executor.layers.fused_moe.config import (
    fp8_w8a8_moe_quant_config,
    nvfp4_moe_quant_config,
@@ -26,6 +23,9 @@ from vllm.model_executor.layers.fused_moe.cutlass_moe import (
    CutlassExpertsFp4,
 )
 from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts, fused_topk
+from vllm.model_executor.layers.fused_moe.prepare_finalize import (
+    MoEPrepareAndFinalizeNoEP,
+)
 from vllm.scalar_type import scalar_types
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 from vllm.v1.worker.workspace import init_workspace_manager
@@ -196,21 +196,10 @@ def bench_run(
            g2_alphas=w2_gs,
        )

-        moe_config = make_dummy_moe_config(
-            num_experts=num_experts,
-            hidden_dim=k,
-            intermediate_size_per_partition=n,
-            in_dtype=a.dtype,
-        )
-        kernel = mk.FusedMoEKernel(
-            maybe_make_prepare_finalize(
-                moe=moe_config,
-                quant_config=quant_config,
-                allow_new_interface=True,
-                use_monolithic=False,
-            ),
+        kernel = mk.FusedMoEModularKernel(
+            MoEPrepareAndFinalizeNoEP(),
            CutlassExpertsFp4(
-                moe_config=moe_config,
+                make_dummy_moe_config(),
                quant_config=quant_config,
            ),
        )
@@ -251,17 +240,11 @@ def bench_run(
            g1_alphas=w1_gs,
            g2_alphas=w2_gs,
        )
-        moe_config = make_dummy_moe_config()

-        kernel = mk.FusedMoEKernel(
-            maybe_make_prepare_finalize(
-                moe=moe_config,
-                quant_config=quant_config,
-                allow_new_interface=True,
-                use_monolithic=False,
-            ),
+        kernel = mk.FusedMoEModularKernel(
+            MoEPrepareAndFinalizeNoEP(),
            CutlassExpertsFp4(
-                moe_config=moe_config,
+                make_dummy_moe_config(),
                quant_config=quant_config,
            ),
        )
@@ -307,7 +290,7 @@ def bench_run(
    def replay_graph(graph, num_repeats):
        for _ in range(num_repeats):
            graph.replay()
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()

    cutlass_stream = torch.cuda.Stream()
    cutlass_graph = torch.cuda.CUDAGraph()
@@ -330,7 +313,7 @@ def bench_run(
            e=num_experts,
            device=device,
        )
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    triton_stream = torch.cuda.Stream()
    triton_graph = torch.cuda.CUDAGraph()
@@ -345,7 +328,7 @@ def bench_run(
            w2_fp8scale,
            a_fp8_scale,
        )
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    min_run_time = 5
    num_warmup = 5
@@ -342,7 +342,7 @@ class CommunicatorBenchmark:
            if not should_use_fn(tensor):
                return None

-            torch.accelerator.synchronize()
+            torch.cuda.synchronize()
            stream = torch.cuda.Stream()
            with torch.cuda.stream(stream):
                graph_input = tensor.clone()
@@ -360,17 +360,17 @@ class CommunicatorBenchmark:
                        for _ in range(CUDA_GRAPH_CAPTURE_CYCLES):
                            allreduce_fn(graph_input)

-            torch.accelerator.synchronize()
+            torch.cuda.synchronize()
            for _ in range(num_warmup):
                graph.replay()
-            torch.accelerator.synchronize()
+            torch.cuda.synchronize()

-            torch.accelerator.synchronize()
+            torch.cuda.synchronize()
            start_time = time.perf_counter()

            for _ in range(num_trials):
                graph.replay()
-            torch.accelerator.synchronize()
+            torch.cuda.synchronize()

            end_time = time.perf_counter()

@@ -385,7 +385,7 @@ def benchmark_operation(
    # Warmup before graph capture
    for _ in range(warmup):
        operation_func(*args, **kwargs)
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    # Create CUDA graph
    graph = torch.cuda.CUDAGraph()
@@ -398,19 +398,19 @@ def benchmark_operation(
            operation_func(*args, **kwargs)

    # Graph warmup
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()
    for _ in range(warmup):
        graph.replay()

    # Benchmark with CUDA graph
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()
    start_time = time.perf_counter()

    for _ in range(trials // num_op_per_cudagraph):
        # operation_func(*args, **kwargs)
        graph.replay()

-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()
    end_time = time.perf_counter()

    avg_time_ms = ((end_time - start_time) / trials) * 1000
@@ -9,15 +9,15 @@ 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.config import ParallelConfig, VllmConfig, set_current_vllm_config
-from vllm.model_executor.layers.fused_moe.all2all_utils import (
-    maybe_make_prepare_finalize,
-)
 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,
 )
+from vllm.model_executor.layers.fused_moe.prepare_finalize import (
+    MoEPrepareAndFinalizeNoEP,
+)
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 from vllm.v1.worker.workspace import init_workspace_manager

@@ -131,22 +131,16 @@ def bench_run(
            w2_scale=w2_scale,
            per_act_token_quant=per_act_token,
        )
-        moe_config = make_dummy_moe_config(
-            num_experts=w2.shape[0],
-            hidden_dim=w2.shape[1],
-            intermediate_size_per_partition=w2.shape[2],
-            in_dtype=a.dtype,
-        )

-        fn = mk.FusedMoEKernel(
-            maybe_make_prepare_finalize(
-                moe=moe_config,
-                quant_config=quant_config,
-                allow_new_interface=True,
-                use_monolithic=False,
-            ),
+        fn = mk.FusedMoEModularKernel(
+            MoEPrepareAndFinalizeNoEP(),
            CutlassExpertsFp8(
-                moe_config=moe_config,
+                moe_config=make_dummy_moe_config(
+                    num_experts=w2.shape[0],
+                    hidden_dim=w2.shape[1],
+                    intermediate_size_per_partition=w2.shape[2],
+                    in_dtype=a.dtype,
+                ),
                quant_config=quant_config,
            ),
        )
@@ -169,22 +163,16 @@ def bench_run(
            w2_scale=w2_scale,
            per_act_token_quant=per_act_token,
        )
-        moe_config = make_dummy_moe_config(
-            num_experts=w2.shape[0],
-            hidden_dim=w2.shape[1],
-            intermediate_size_per_partition=w2.shape[2],
-            in_dtype=a.dtype,
-        )

-        fn = mk.FusedMoEKernel(
-            maybe_make_prepare_finalize(
-                moe=moe_config,
-                quant_config=quant_config,
-                allow_new_interface=True,
-                use_monolithic=False,
-            ),
+        fn = mk.FusedMoEModularKernel(
+            MoEPrepareAndFinalizeNoEP(),
            CutlassExpertsFp8(
-                moe_config=moe_config,
+                moe_config=make_dummy_moe_config(
+                    num_experts=w2.shape[0],
+                    hidden_dim=w2.shape[1],
+                    intermediate_size_per_partition=w2.shape[2],
+                    in_dtype=a.dtype,
+                ),
                quant_config=quant_config,
            ),
        )
@@ -224,7 +212,7 @@ def bench_run(
    def replay_graph(graph, num_repeats):
        for _ in range(num_repeats):
            graph.replay()
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()

    cutlass_stream = torch.cuda.Stream()
    cutlass_graph = torch.cuda.CUDAGraph()
@@ -239,7 +227,7 @@ def bench_run(
            topk_weights,
            topk_ids,
        )
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    triton_stream = torch.cuda.Stream()
    triton_graph = torch.cuda.CUDAGraph()
@@ -254,7 +242,7 @@ def bench_run(
            w2_scale,
            a_scale,
        )
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    min_run_time = 5
    num_warmup = 5
@@ -34,14 +34,14 @@ def main(
    residual = torch.randn_like(x) * scale if add_residual else None

    def run_cuda_benchmark(num_iters: int, profile: bool = False) -> float:
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()
        if profile:
            torch.cuda.cudart().cudaProfilerStart()
        start_time = time.perf_counter()

        for _ in range(num_iters):
            layer(x, residual)
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()

        end_time = time.perf_counter()
        if profile:
@@ -1035,7 +1035,7 @@ def bench_optype(
    # Run bench function so that _LORA_A_PTR_DICT and _LORA_B_PTR_DICT are set up
    for kwargs in kwargs_list:
        op_type.bench_fn()(**kwargs)
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    # Merge into a single kwargs and qualify arguments as ArgPool
    kwargs = {k: ArgPool([]) for k in kwargs_list[0]}
@@ -47,13 +47,13 @@ def benchmark_method(
    # Warmup
    for _ in range(num_warmup):
        _ = method(k_nope, k_pe)
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    # Benchmark
    start = time.perf_counter()
    for _ in range(num_iters):
        _ = method(k_nope, k_pe)
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()
    end = time.perf_counter()

    return (end - start) / num_iters * 1000  # Convert to ms
@@ -17,9 +17,6 @@ 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.all2all_utils import (
-    maybe_make_prepare_finalize,
-)
 from vllm.model_executor.layers.fused_moe.config import (
    FusedMoEConfig,
    FusedMoEParallelConfig,
@@ -54,7 +51,7 @@ def clear_triton_cache():

    # Clear CUDA memory cache
    if torch.cuda.is_available():
-        torch.accelerator.empty_cache()
+        torch.cuda.empty_cache()

    # Try to clear Triton's runtime cache
    try:
@@ -245,33 +242,24 @@ def benchmark_config(

        deep_gemm_experts = None
        if use_deep_gemm:
-            moe_config = (
-                FusedMoEConfig(
-                    num_experts=num_experts,
-                    experts_per_token=topk,
-                    hidden_dim=hidden_size,
-                    intermediate_size_per_partition=shard_intermediate_size,
-                    num_local_experts=num_experts,
-                    num_logical_experts=num_experts,
-                    activation=MoEActivation.SILU,
-                    moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
-                    in_dtype=init_dtype,
-                    routing_method=RoutingMethodType.TopK,
-                    device="cuda",
-                ),
-            )
-            deep_gemm_experts = mk.FusedMoEKernel(
-                prepare_finalize=maybe_make_prepare_finalize(
-                    moe=moe_config,
-                    quant_config=quant_config,
-                    allow_new_interface=True,
-                    use_monolithic=False,
-                ),
+            deep_gemm_experts = mk.FusedMoEModularKernel(
+                prepare_finalize=MoEPrepareAndFinalizeNoEP(),
                fused_experts=TritonOrDeepGemmExperts(
-                    moe_config=moe_config,
+                    moe_config=FusedMoEConfig(
+                        num_experts=num_experts,
+                        experts_per_token=topk,
+                        hidden_dim=hidden_size,
+                        intermediate_size_per_partition=shard_intermediate_size,
+                        num_local_experts=num_experts,
+                        num_logical_experts=num_experts,
+                        activation=MoEActivation.SILU,
+                        moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
+                        in_dtype=init_dtype,
+                        routing_method=RoutingMethodType.TopK,
+                        device="cuda",
+                    ),
                    quant_config=quant_config,
                ),
-                inplace=not disable_inplace(),
            )

        with override_config(config):
@@ -281,16 +269,8 @@ def benchmark_config(

            inplace = not disable_inplace()
            if use_deep_gemm:
-                return deep_gemm_experts.apply(
-                    x,
-                    w1,
-                    w2,
-                    topk_weights,
-                    topk_ids,
-                    activation=MoEActivation.SILU,
-                    global_num_experts=num_experts,
-                    apply_router_weight_on_input=False,
-                    expert_map=False,
+                return deep_gemm_experts(
+                    x, w1, w2, topk_weights, topk_ids, inplace=inplace
                )
            return fused_experts(
                x,
@@ -304,19 +284,19 @@ def benchmark_config(

    # JIT compilation & warmup
    run()
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    # Capture 10 invocations with CUDA graph
    graph = torch.cuda.CUDAGraph()
    with torch.cuda.graph(graph):
        for _ in range(10):
            run()
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    # Warmup
    for _ in range(5):
        graph.replay()
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    start_event = torch.Event(enable_timing=True)
    end_event = torch.Event(enable_timing=True)
@@ -324,7 +304,7 @@ def benchmark_config(
    latencies: list[float] = []
    for i in range(num_iters):
        prepare(i)
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()

        start_event.record()
        graph.replay()
@@ -131,7 +131,7 @@ def benchmark_config(
                topk_ids,
                quant_config=quant_config,
            )
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    # Benchmark
    start = torch.cuda.Event(enable_timing=True)
@@ -149,7 +149,7 @@ def benchmark_config(
                quant_config=quant_config,
            )
    end.record()
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()
    return start.elapsed_time(end) / num_iters * 1000  # ms -> us


@@ -69,19 +69,19 @@ def benchmark_permute(

    # JIT compilation & warmup
    run()
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    # Capture 10 invocations with CUDA graph
    graph = torch.cuda.CUDAGraph()
    with torch.cuda.graph(graph):
        for _ in range(10):
            run()
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    # Warmup
    for _ in range(5):
        graph.replay()
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    start_event = torch.Event(enable_timing=True)
    end_event = torch.Event(enable_timing=True)
@@ -89,7 +89,7 @@ def benchmark_permute(
    latencies: list[float] = []
    for i in range(num_iters):
        prepare(i)
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()

        start_event.record()
        graph.replay()
@@ -159,26 +159,26 @@ def benchmark_unpermute(
    # JIT compilation & warmup
    input = prepare()
    run(input)
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    # Capture 10 invocations with CUDA graph
    graph = torch.cuda.CUDAGraph()
    with torch.cuda.graph(graph):
        for _ in range(10):
            run(input)
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    # Warmup
    for _ in range(5):
        graph.replay()
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    start_event = torch.Event(enable_timing=True)
    end_event = torch.Event(enable_timing=True)

    latencies: list[float] = []
    for i in range(num_iters):
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()
        start_event.record()
        graph.replay()
        end_event.record()
@@ -135,14 +135,14 @@ def benchmark_mrope(
            key.clone(),
        )

-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    # Time reference implementation
    torch_times = []
    for _ in range(benchmark_iter):
        query_clone = query.clone()
        key_clone = key.clone()
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()
        start_time = time.time()

        mrope_helper_class.forward_native(
@@ -151,7 +151,7 @@ def benchmark_mrope(
            key_clone,
        )

-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()
        torch_times.append(time.time() - start_time)

    # Time triton kernel implementation
@@ -159,14 +159,14 @@ def benchmark_mrope(
    for _ in range(benchmark_iter):
        query_clone = query.clone()
        key_clone = key.clone()
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()
        start_time = time.time()
        mrope_helper_class.forward_cuda(
            positions,
            query_clone,
            key_clone,
        )
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()
        triton_times.append(time.time() - start_time)

    # Calculate statistics
@@ -103,7 +103,7 @@ def main(
        max_logits = torch.empty_like(exp_sums)

    def run_cuda_benchmark(num_iters: int, profile: bool = False) -> float:
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()
        if profile:
            torch.cuda.cudart().cudaProfilerStart()
        start_time = time.perf_counter()
@@ -173,7 +173,7 @@ def main(
                    )
            else:
                raise ValueError(f"Invalid version: {version}")
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()

        end_time = time.perf_counter()
        if profile:
@@ -28,7 +28,7 @@ def _time_cuda(
    # warmup
    for _ in range(warmup_iters):
        fn()
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    start = torch.Event(enable_timing=True)
    end = torch.Event(enable_timing=True)
@@ -37,7 +37,7 @@ def _time_cuda(
    for _ in range(bench_iters):
        fn()
    end.record()
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    return start.elapsed_time(end) / bench_iters  # ms/iter

@@ -29,7 +29,7 @@ def main(
    scale = torch.randn(1, 1, dtype=torch.float32) if static_scale else None

    def run_cuda_benchmark(num_iters: int, profile: bool = False) -> float:
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()
        if profile:
            torch.cuda.cudart().cudaProfilerStart()
        start_time = time.perf_counter()
@@ -39,7 +39,7 @@ def main(
                ops.scaled_int8_quant(x, scale)
            else:
                ops.scaled_fp8_quant(x, scale)
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()

        end_time = time.perf_counter()
        if profile:
@@ -84,16 +84,16 @@ def run_benchmark(
        g = torch.cuda.CUDAGraph()
        with torch.cuda.graph(g):
            function_under_test()
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()
        function_under_test = lambda: g.replay()

    def run_cuda_benchmark(n_iters: int) -> float:
        nonlocal key, value, key_cache, value_cache, slot_mapping
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()
        start = time.perf_counter()
        for _ in range(n_iters):
            function_under_test()
-            torch.accelerator.synchronize()
+            torch.cuda.synchronize()
        end = time.perf_counter()
        return (end - start) / n_iters

@@ -104,7 +104,7 @@ def run_benchmark(

    # free tensors to mitigate OOM when sweeping
    del key, value, key_cache, value_cache, slot_mapping
-    torch.accelerator.empty_cache()
+    torch.cuda.empty_cache()

    return lat

@@ -109,16 +109,16 @@ def run_benchmark(
        g = torch.cuda.CUDAGraph()
        with torch.cuda.graph(g):
            function_under_test()
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()
        function_under_test = lambda: g.replay()

    def run_cuda_benchmark(n_iters: int) -> float:
        nonlocal key, value, key_cache, value_cache, slot_mapping
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()
        start = time.perf_counter()
        for _ in range(n_iters):
            function_under_test()
-            torch.accelerator.synchronize()
+            torch.cuda.synchronize()
        end = time.perf_counter()
        return (end - start) / n_iters

@@ -129,7 +129,7 @@ def run_benchmark(

    # free tensors to mitigate OOM when sweeping
    del key, value, key_cache, value_cache, slot_mapping
-    torch.accelerator.empty_cache()
+    torch.cuda.empty_cache()

    return lat

@@ -251,7 +251,7 @@ def benchmark(
        kernel(
            y, tokens_per_expert, num_parallel_tokens=num_parallel_tokens, group_size=G
        )
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    start_event = torch.Event(enable_timing=True)
    end_event = torch.Event(enable_timing=True)
@@ -259,7 +259,7 @@ def benchmark(
    # Benchmark
    latencies: list[float] = []
    for _ in range(runs):
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()

        start_event.record()
        for i in range(iterations_per_run):
@@ -126,7 +126,7 @@ def benchmark_decode(
    )

    def time_fn(fn, warmup=10, trials=20):
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()
        start = torch.Event(enable_timing=True)
        end = torch.Event(enable_timing=True)
        times = []
@@ -136,7 +136,7 @@ def benchmark_decode(
            start.record()
            fn()
            end.record()
-            torch.accelerator.synchronize()
+            torch.cuda.synchronize()
            times.append(start.elapsed_time(end))  # ms
        return sum(times) / len(times), torch.std(torch.tensor(times))

@@ -138,7 +138,7 @@ def benchmark_prefill(
    )

    def time_fn(fn, warmup=10, trials=20):
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()
        start = torch.Event(enable_timing=True)
        end = torch.Event(enable_timing=True)
        times = []
@@ -148,7 +148,7 @@ def benchmark_prefill(
            start.record()
            fn()
            end.record()
-            torch.accelerator.synchronize()
+            torch.cuda.synchronize()
            times.append(start.elapsed_time(end))  # ms
        return sum(times) / len(times), torch.std(torch.tensor(times))

@@ -177,18 +177,18 @@ def benchmark_config(
    def run():
        w8a8_block_matmul(A, B, As, Bs, block_size, config, out_dtype)

-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()
    # JIT complication & warmup
    for _ in range(5):
        run()
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()

    start_event = torch.Event(enable_timing=True)
    end_event = torch.Event(enable_timing=True)

    latencies: list[float] = []
    for i in range(num_iters):
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()
        start_event.record()
        run()
        end_event.record()
@@ -35,7 +35,7 @@ def benchmark_shape(
    B = torch.randn((n, k), device="cuda", dtype=torch.bfloat16)

    # Reference result in BF16
-    torch.accelerator.synchronize()
+    torch.cuda.synchronize()
    C_ref = A @ B.t()

    # Pre-quantize B for all implementations
@@ -121,14 +121,14 @@ def benchmark_shape(
        # Warmup
        for _ in range(warmup):
            func()
-            torch.accelerator.synchronize()
+            torch.cuda.synchronize()

        # Timing loop
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()
        start = time.time()
        for _ in range(repeat):
            func()
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()
        end = time.time()

        # Calculate timing and TFLOPS
@@ -7,7 +7,7 @@ First start serving your model
 ```bash
 export MODEL_PATH=/models/meta-llama/Meta-Llama-3.1-8B-Instruct/

-vllm serve $MODEL_PATH --served-model-name Llama
+vllm serve $MODEL_PATH --served-model-name Llama --disable-log-requests
 ```

 The variable `MODEL_PATH` should be a path to the model files (e.g. downloaded from huggingface).
@@ -242,24 +242,13 @@ if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND
        )
    else()
        message(STATUS "Downloading oneDNN from GitHub")
-        if(ASIMD_FOUND AND NOT APPLE_SILICON_FOUND)
-            message(STATUS "aarch64 detected: using pinned oneDNN commit 9c5be1cc59e368aebf0909e6cf20f981ea61462a")
-            FetchContent_Declare(
-                oneDNN
-                GIT_REPOSITORY https://github.com/oneapi-src/oneDNN.git
-                GIT_TAG        9c5be1cc59e368aebf0909e6cf20f981ea61462a
-                GIT_PROGRESS   TRUE
-                GIT_SHALLOW    FALSE
-            )
-        else()
-            FetchContent_Declare(
-                oneDNN
-                GIT_REPOSITORY https://github.com/oneapi-src/oneDNN.git
-                GIT_TAG        v3.10
-                GIT_PROGRESS   TRUE
-                GIT_SHALLOW    TRUE
-            )
-        endif()
+        FetchContent_Declare(
+            oneDNN
+            GIT_REPOSITORY https://github.com/oneapi-src/oneDNN.git
+            GIT_TAG v3.10
+            GIT_PROGRESS TRUE
+            GIT_SHALLOW TRUE
+        )
    endif()

    set(ONEDNN_LIBRARY_TYPE "STATIC")
@@ -46,20 +46,24 @@ else()
  )
 endif()

-# Make sure vllm-flash-attn install rules are nested under vllm/
-# ALL_COMPONENTS ensures the save/modify/restore runs exactly once regardless
-# of how many components are being installed, avoiding double-append of /vllm/.
-install(CODE "set(CMAKE_INSTALL_LOCAL_ONLY FALSE)" ALL_COMPONENTS)
-install(CODE "set(OLD_CMAKE_INSTALL_PREFIX \"\${CMAKE_INSTALL_PREFIX}\")" ALL_COMPONENTS)
-install(CODE "set(CMAKE_INSTALL_PREFIX \"\${CMAKE_INSTALL_PREFIX}/vllm/\")" ALL_COMPONENTS)
+
+# Install rules for FA components need the install prefix nested under vllm/
+# These run at install time, before the FA library's own install rules
+foreach(_FA_COMPONENT _vllm_fa2_C _vllm_fa3_C)
+  install(CODE "set(CMAKE_INSTALL_LOCAL_ONLY FALSE)" COMPONENT ${_FA_COMPONENT})
+  install(CODE "set(OLD_CMAKE_INSTALL_PREFIX \"\${CMAKE_INSTALL_PREFIX}\")" COMPONENT ${_FA_COMPONENT})
+  install(CODE "set(CMAKE_INSTALL_PREFIX \"\${CMAKE_INSTALL_PREFIX}/vllm/\")" COMPONENT ${_FA_COMPONENT})
+endforeach()

 # Fetch the vllm-flash-attn library
 FetchContent_MakeAvailable(vllm-flash-attn)
 message(STATUS "vllm-flash-attn is available at ${vllm-flash-attn_SOURCE_DIR}")

 # Restore the install prefix after FA's install rules
-install(CODE "set(CMAKE_INSTALL_PREFIX \"\${OLD_CMAKE_INSTALL_PREFIX}\")" ALL_COMPONENTS)
-install(CODE "set(CMAKE_INSTALL_LOCAL_ONLY TRUE)" ALL_COMPONENTS)
+foreach(_FA_COMPONENT _vllm_fa2_C _vllm_fa3_C)
+  install(CODE "set(CMAKE_INSTALL_PREFIX \"\${OLD_CMAKE_INSTALL_PREFIX}\")" COMPONENT ${_FA_COMPONENT})
+  install(CODE "set(CMAKE_INSTALL_LOCAL_ONLY TRUE)" COMPONENT ${_FA_COMPONENT})
+endforeach()

 # Install shared Python files for both FA2 and FA3 components
 foreach(_FA_COMPONENT _vllm_fa2_C _vllm_fa3_C)
@@ -74,12 +74,6 @@ void indexer_k_quant_and_cache(
    int64_t quant_block_size,     // quantization block size
    const std::string& scale_fmt);

-// Concatenate query nope and rope for MLA/DSA attention
-void concat_mla_q(
-    torch::Tensor& ql_nope,  // [num_tokens, num_heads, nope_dim]
-    torch::Tensor& q_pe,     // [num_tokens, num_heads, rope_dim]
-    torch::Tensor& q_out);   // [num_tokens, num_heads, nope_dim + rope_dim]
-
 // Extract function to gather quantized K cache
 void cp_gather_indexer_k_quant_cache(
    const torch::Tensor& kv_cache,  // [num_blocks, block_size, cache_stride]
@@ -8,7 +8,6 @@
 #include "cuda_compat.h"
 #include "dispatch_utils.h"
 #include "quantization/vectorization_utils.cuh"
-#include "concat_mla_q.cuh"

 #ifdef USE_ROCM
  #include "quantization/w8a8/fp8/amd/quant_utils.cuh"
@@ -996,67 +995,75 @@ namespace vllm {
 // Similar to cp_gather_cache but specifically for FP8->BF16 conversion
 __global__ void cp_gather_and_upconvert_fp8_kv_cache(
    const uint8_t* __restrict__ src_cache,    // [NUM_BLOCKS, BLOCK_SIZE, 656]
-    __nv_bfloat16* __restrict__ dst,          // [total_tokens, 576]
-    const int32_t* __restrict__ block_table,  // [num_reqs, BLOCK_INDICES]
-    const int32_t* __restrict__ workspace_starts,  // [num_reqs]
-    const int32_t num_reqs, const int32_t block_size,
-    const int32_t total_tokens, const int64_t block_table_stride,
-    const int64_t cache_block_stride, const int64_t cache_entry_stride,
-    const int64_t dst_entry_stride) {
-  const int flat_warp_id = (blockIdx.x * blockDim.x + threadIdx.x) >> 5;
-  if (flat_warp_id >= total_tokens) return;
-  const int lane_id = threadIdx.x & 31;
+    __nv_bfloat16* __restrict__ dst,          // [TOT_TOKENS, 576]
+    const int32_t* __restrict__ block_table,  // [BATCH, BLOCK_INDICES]
+    const int32_t* __restrict__ seq_lens,     // [BATCH]
+    const int32_t* __restrict__ workspace_starts,  // [BATCH]
+    const int32_t block_size, const int32_t head_dim,
+    const int64_t block_table_stride, const int64_t cache_block_stride,
+    const int64_t cache_entry_stride, const int64_t dst_entry_stride) {
+  const int64_t bid = blockIdx.x;  // Batch ID
+  const int32_t num_splits = gridDim.y;
+  const int32_t split = blockIdx.y;
+  const int32_t seq_start = workspace_starts[bid];
+  const int32_t seq_len = seq_lens[bid];
+  const int32_t tot_slots = seq_len;
+  const int32_t split_slots = cuda_utils::ceil_div(tot_slots, num_splits);

-  // Binary search to find which request owns this output token
-  int lo = 0, hi = num_reqs - 1;
-  while (lo < hi) {
-    int mid = (lo + hi + 1) >> 1;
-    if (workspace_starts[mid] <= flat_warp_id)
-      lo = mid;
-    else
-      hi = mid - 1;
+  const int32_t split_start = split * split_slots;
+  const int32_t split_end = min((split + 1) * split_slots, tot_slots);
+
+  const bool is_active_split = (split_start < tot_slots);
+
+  if (!is_active_split) return;
+
+  // Adjust the pointer for the block_table for this batch
+  const int32_t batch_offset = bid * block_table_stride;
+  int32_t offset = split_start;
+  int32_t offset_div = offset / block_size;
+  offset = offset % block_size;
+  const int32_t* batch_block_table = block_table + batch_offset;
+
+  // Adjust dst pointer based on the cumulative sequence lengths
+  dst += seq_start * dst_entry_stride;
+
+  const int tid = threadIdx.x;
+
+  // Process each token in this split
+  for (int pid = split_start; pid < split_end; ++pid) {
+    auto block_id = batch_block_table[offset_div];
+    const uint8_t* token_ptr =
+        src_cache + block_id * cache_block_stride + offset * cache_entry_stride;
+    __nv_bfloat16* dst_ptr = dst + pid * dst_entry_stride;
+
+    // FP8 format: 512 bytes fp8 + 16 bytes scales + 128 bytes rope (64 bf16)
+    const uint8_t* no_pe_ptr = token_ptr;
+    const float* scales_ptr = reinterpret_cast<const float*>(token_ptr + 512);
+    const __nv_bfloat16* rope_ptr =
+        reinterpret_cast<const __nv_bfloat16*>(token_ptr + 512 + 16);
+
+    // Parallelize fp8 dequant (512 elements) and rope copy (64 elements)
+    if (tid < 512) {
+      // FP8 dequantization
+      const int tile = tid >> 7;  // each tile is 128 elements
+      const float scale = scales_ptr[tile];
+      const uint8_t val = no_pe_ptr[tid];
+      dst_ptr[tid] =
+          fp8::scaled_convert<__nv_bfloat16, uint8_t,
+                              vllm::Fp8KVCacheDataType::kFp8E4M3>(val, scale);
+    } else if (tid < 576) {
+      // Rope copy (64 bf16 elements)
+      const int rope_idx = tid - 512;
+      dst_ptr[512 + rope_idx] = rope_ptr[rope_idx];
+    }
+
+    // Move to next token
+    offset += 1;
+    if (offset == block_size) {
+      offset_div += 1;
+      offset = 0;
+    }
  }
-  const int req_id = lo;
-
-  // Compute physical token address via block table
-  const int out_token_id = flat_warp_id;
-  const int token_offset = out_token_id - workspace_starts[req_id];
-  const int cache_block_idx = token_offset / block_size;
-  const int offset_in_block = token_offset % block_size;
-  const int physical_block =
-      block_table[req_id * block_table_stride + cache_block_idx];
-
-  const uint8_t* token_ptr = src_cache + physical_block * cache_block_stride +
-                             offset_in_block * cache_entry_stride;
-
-  const int4* nope_src = reinterpret_cast<const int4*>(token_ptr);
-  const int4 fp8_data = nope_src[lane_id];
-
-  const float* scales_ptr = reinterpret_cast<const float*>(token_ptr + 512);
-  const float scale = scales_ptr[lane_id >> 3];
-
-  const uint2 fp8_lo = make_uint2(fp8_data.x, fp8_data.y);
-  const uint2 fp8_hi = make_uint2(fp8_data.z, fp8_data.w);
-#ifdef USE_ROCM
-  const bf16_8_t bf16_lo =
-      fp8::scaled_vec_conversion<bf16_8_t, uint2>(fp8_lo, scale);
-  const bf16_8_t bf16_hi =
-      fp8::scaled_vec_conversion<bf16_8_t, uint2>(fp8_hi, scale);
-#else
-  const bf16_8_t bf16_lo =
-      fp8::scaled_vec_conversion<bf16_8_t, uint2>(fp8_lo, scale, __NV_E4M3);
-  const bf16_8_t bf16_hi =
-      fp8::scaled_vec_conversion<bf16_8_t, uint2>(fp8_hi, scale, __NV_E4M3);
-#endif
-
-  __nv_bfloat16* dst_ptr = dst + out_token_id * dst_entry_stride;
-  int4* nope_dst = reinterpret_cast<int4*>(dst_ptr) + lane_id * 2;
-  nope_dst[0] = *reinterpret_cast<const int4*>(&bf16_lo);
-  nope_dst[1] = *reinterpret_cast<const int4*>(&bf16_hi);
-
-  const int* rope_src = reinterpret_cast<const int*>(token_ptr + 528);
-  int* rope_dst = reinterpret_cast<int*>(dst_ptr + 512);
-  rope_dst[lane_id] = rope_src[lane_id];
 }

 template <typename scalar_t>
@@ -1250,16 +1257,15 @@ void cp_gather_and_upconvert_fp8_kv_cache(
    src_ptr = reinterpret_cast<const uint8_t*>(src_cache.data_ptr());
  }

-  const int total_tokens = dst.size(0);
-  constexpr int warps_per_block = 8;
-  const int grid_size = (total_tokens + warps_per_block - 1) / warps_per_block;
-  const int block_size_threads = warps_per_block * 32;  // 256 threads
+  // 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_size, block_size_threads, 0,
-                                               stream>>>(
+  vllm::cp_gather_and_upconvert_fp8_kv_cache<<<grid, block, 0, stream>>>(
      src_ptr, reinterpret_cast<__nv_bfloat16*>(dst.data_ptr()),
-      block_table.data_ptr<int32_t>(), workspace_starts.data_ptr<int32_t>(),
-      static_cast<int32_t>(batch_size), block_size, total_tokens,
+      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,
      dst_entry_stride);
 }
@@ -1359,43 +1365,3 @@ void cp_gather_indexer_k_quant_cache(
    CALL_CP_GATHER_INDEXER_K_QUANT_CACHE(32);
  }
 }
-
-// Concatenate ql_nope and q_pe into a contiguous q_out tensor for MLA/DSA.
-// Replaces torch.cat((ql_nope, q_pe), dim=-1).
-void concat_mla_q(torch::Tensor& ql_nope,  // [num_tokens, num_heads, nope_dim]
-                  torch::Tensor& q_pe,     // [num_tokens, num_heads, rope_dim]
-                  torch::Tensor& q_out     // [num_tokens, num_heads, nope_dim +
-                                           // rope_dim]
-) {
-  const int num_tokens = ql_nope.size(0);
-  const int num_heads = ql_nope.size(1);
-  const int nope_dim = ql_nope.size(2);
-  const int rope_dim = q_pe.size(2);
-
-  TORCH_CHECK(nope_dim % 512 == 0, "nope_dim must be a multiple of 512, got ",
-              nope_dim);
-  TORCH_CHECK(rope_dim == 64, "rope_dim must be 64, got ", rope_dim);
-  TORCH_CHECK(q_out.size(2) == nope_dim + rope_dim);
-
-  TORCH_CHECK(ql_nope.stride(2) == 1, "ql_nope must have stride 1 in dim 2");
-  TORCH_CHECK(q_pe.stride(2) == 1, "q_pe must have stride 1 in dim 2");
-  TORCH_CHECK(q_out.stride(2) == 1, "q_out must have stride 1 in dim 2");
-
-  if (num_tokens == 0) return;
-
-  constexpr int warps_per_block = 8;
-  const int total_warps = num_tokens * num_heads;
-  const int grid_size = (total_warps + warps_per_block - 1) / warps_per_block;
-  const int block_size = warps_per_block * 32;
-
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(ql_nope));
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-
-  VLLM_DISPATCH_FLOATING_TYPES(ql_nope.scalar_type(), "concat_mla_q", [&] {
-    vllm::ConcatMLAQKernel<scalar_t, 512><<<grid_size, block_size, 0, stream>>>(
-        q_out.data_ptr<scalar_t>(), ql_nope.data_ptr<scalar_t>(),
-        q_pe.data_ptr<scalar_t>(), num_tokens, num_heads, q_out.stride(0),
-        q_out.stride(1), ql_nope.stride(0), ql_nope.stride(1), q_pe.stride(0),
-        q_pe.stride(1));
-  });
-}
@@ -1,60 +0,0 @@
-#ifndef CONCAT_MLA_Q_CUH_
-#define CONCAT_MLA_Q_CUH_
-
-#include <cuda_bf16.h>
-#include <cuda_fp16.h>
-
-#include "cuda_vec_utils.cuh"
-
-namespace vllm {
-
-// Concatenates ql_nope [num_tokens, num_heads, NOPE_DIM] and
-// q_pe [num_tokens, num_heads, 64]
-// into q_out [num_tokens, num_heads, NOPE_DIM+64].
-// Currently instantiated only for NOPE_DIM=512.
-// Rope dim is hardcoded to 64 (DeepSeek V3.2 MLA)
-template <typename DType, int NOPE_DIM>
-__global__ void ConcatMLAQKernel(
-    DType* __restrict__ q_out, const DType* __restrict__ ql_nope,
-    const DType* __restrict__ q_pe, const int num_tokens, const int num_heads,
-    const int64_t out_stride_0, const int64_t out_stride_1,
-    const int64_t nope_stride_0, const int64_t nope_stride_1,
-    const int64_t pe_stride_0, const int64_t pe_stride_1) {
-  const int flat_warp_id = (blockIdx.x * blockDim.x + threadIdx.x) >> 5;
-  if (flat_warp_id >= num_tokens * num_heads) return;
-
-  const int token_id = flat_warp_id / num_heads;
-  const int head_id = flat_warp_id % num_heads;
-  const int lane_id = threadIdx.x & 31;
-
-  constexpr bool use_256b = VLLM_256B_PTX_ENABLED;
-  constexpr int nope_vec_loads =
-      NOPE_DIM * sizeof(DType) / (VecTraits<use_256b>::ARCH_MAX_VEC_SIZE * 32);
-
-  const DType* nope_src =
-      ql_nope + token_id * nope_stride_0 + head_id * nope_stride_1;
-  DType* nope_dst = q_out + token_id * out_stride_0 + head_id * out_stride_1;
-
-#pragma unroll
-  for (int i = 0; i < nope_vec_loads; i++) {
-    const int offset = i * 32 + lane_id;
-    if constexpr (use_256b) {
-      st256_cs(reinterpret_cast<u32x8_t*>(nope_dst) + offset,
-               ld256_cs(reinterpret_cast<const u32x8_t*>(nope_src) + offset));
-    } else {
-      st128_cs(reinterpret_cast<int4*>(nope_dst) + offset,
-               ld128_cs(reinterpret_cast<const int4*>(nope_src) + offset));
-    }
-  }
-
-  const int* rope_src = reinterpret_cast<const int*>(
-      q_pe + token_id * pe_stride_0 + head_id * pe_stride_1);
-  int* rope_dst = reinterpret_cast<int*>(q_out + token_id * out_stride_0 +
-                                         head_id * out_stride_1 + NOPE_DIM);
-
-  st32_cs(rope_dst + lane_id, ld32_cs(rope_src + lane_id));
-}
-
-}  // namespace vllm
-
-#endif  // CONCAT_MLA_Q_CUH_
@@ -420,7 +420,7 @@ class AttentionImpl<ISA::AMX, scalar_t, head_dim> {
      const int64_t block_size, const int64_t block_size_stride) {
    // For AMX 2D tiles, size of each line is 64 bytes
    constexpr int64_t amx_tile_row_size = AMX_TILE_ROW_BYTES;
-    // For AMX B matrix, N always is 16
+    // For AMX B martix, N always is 16
    constexpr int64_t amx_b_tile_n_size = AMX_TILE_ROW_BYTES / 4;
    constexpr int64_t amx_b_tile_k_size = amx_tile_row_size / sizeof(scalar_t);
    // For now suppose block_size is divisible by amx_tile_column_num
@@ -237,10 +237,13 @@ 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(
@@ -408,19 +411,21 @@ MatMulPrimitiveHandler::MatMulPrimitiveHandler(const Args& args)
  dnnl::memory::desc original_b_md({b_k_size_, b_n_size_}, b_type_,
                                   {b_k_stride_, b_n_stride_});

-  // dummy M size for prepacking weights
-  // Prepacking weights improves performance and avoid runtime reorders
-  constexpr dnnl_dim_t kProbeM = 128;
-
  prepack_weight(args.b_ptr, original_b_md,
                 create_primitive_desc(
-                     MSizeCacheKey{// Use a concrete M so oneDNN's kernel
-                                   // selector can choose an optimally blocked
-                                   // weight layout.
-                                   .a_m_size = kProbeM,
-                                   .a_m_stride = b_k_size_,
-                                   .use_bias = false,
-                                   .bias_type = dnnl::memory::data_type::undef},
+                     MSizeCacheKey{
+#ifdef VLLM_USE_ACL
+                         // Arm Compute Library (ACL) backend for oneDNN does
+                         // not support runtime
+                         // dimensions, so we set M to a default value
+                         .a_m_size = 128,
+                         .a_m_stride = b_k_size_,
+#else
+                         .a_m_size = DNNL_RUNTIME_DIM_VAL,
+                         .a_m_stride = DNNL_RUNTIME_DIM_VAL,
+#endif
+                         .use_bias = false,
+                         .bias_type = dnnl::memory::data_type::undef},
                     true)
                     .weights_desc());
  init_runtime_memory_cache(args);
@@ -4,7 +4,7 @@

 #include <torch/library.h>

-// Note: overwrite the external definition for sharing same name between
+// Note: overwrite the external defination for sharing same name between
 // libraries use different ISAs.
 #define TORCH_EXTENSION_NAME _C

@@ -196,6 +196,7 @@ __forceinline__ __device__ u32x8_t ld256_cs(const u32x8_t* addr) {
  return val;
 #else
  assert(false && "ld256_cs requires SM100+ with CUDA 12.9+");
+  return {};
 #endif
 }

@@ -210,51 +211,23 @@ __forceinline__ __device__ void st256_cs(u32x8_t* addr, u32x8_t val) {
 #endif
 }

-// 32-bit load / store.
-__device__ __forceinline__ int ld32(const int* addr) { return __ldg(addr); }
-
-__device__ __forceinline__ void st32(int* addr, int val) { *addr = val; }
-
-// 32-bit cache-streaming (.cs) load / store.
-// Falls back to ld32/st32 on ROCm (no .cs hint).
+// 32-bit cache-streaming (.cs) load / store  — SM100+ only.
 __forceinline__ __device__ int ld32_cs(const int* addr) {
+#if VLLM_256B_PTX_ENABLED
  int val;
-#ifndef USE_ROCM
  asm volatile("ld.global.cs.b32 %0, [%1];" : "=r"(val) : "l"(addr));
-#else
-  val = ld32(addr);
-#endif
  return val;
+#else
+  assert(false && "ld32_cs requires SM100+ with CUDA 12.9+");
+  return 0;
+#endif
 }

 __forceinline__ __device__ void st32_cs(int* addr, int val) {
-#ifndef USE_ROCM
+#if VLLM_256B_PTX_ENABLED
  asm volatile("st.global.cs.b32 [%0], %1;" ::"l"(addr), "r"(val));
 #else
-  st32(addr, val);
-#endif
-}
-
-// 128-bit cache-streaming (.cs) load / store.
-// Falls back to ld128/st128 on ROCm (no .cs hint).
-__forceinline__ __device__ int4 ld128_cs(const int4* addr) {
-  int4 val;
-#ifndef USE_ROCM
-  asm volatile("ld.global.cs.v4.u32 {%0,%1,%2,%3}, [%4];"
-               : "=r"(val.x), "=r"(val.y), "=r"(val.z), "=r"(val.w)
-               : "l"(addr));
-#else
-  ld128(val, addr);
-#endif
-  return val;
-}
-
-__forceinline__ __device__ void st128_cs(int4* addr, int4 val) {
-#ifndef USE_ROCM
-  asm volatile("st.global.cs.v4.u32 [%0], {%1,%2,%3,%4};" ::"l"(addr),
-               "r"(val.x), "r"(val.y), "r"(val.z), "r"(val.w));
-#else
-  st128(val, addr);
+  assert(false && "st32_cs requires SM100+ with CUDA 12.9+");
 #endif
 }

@@ -287,7 +260,7 @@ __device__ __forceinline__ void ld256_cg_or_zero(u32x8_t& val, const void* ptr,

 __device__ __forceinline__ void ld128_cg_or_zero(uint4& val, const void* ptr,
                                                 bool pred) {
-#ifndef USE_ROCM
+#if VLLM_256B_PTX_ENABLED
  uint32_t r0, r1, r2, r3;

  asm volatile(
@@ -305,7 +278,7 @@ __device__ __forceinline__ void ld128_cg_or_zero(uint4& val, const void* ptr,

  val = uint4{r0, r1, r2, r3};
 #else
-  assert(false && "ld128_cg_or_zero is not supported on ROCm");
+  assert(false && "ld128_cg_or_zero requires SM100+ with CUDA 12.9+");
 #endif
 }

@@ -35,11 +35,11 @@ __global__ void batched_moe_align_block_size_kernel(
  int32_t const block_ids_size = sorted_ids_size / block_size;
  int32_t const SENTINEL =
      num_batches * max_tokens_per_batch;  // To denote invalid entries.
-  // Initialize sorted_ids
+  // Intialize sorted_ids
  for (size_t i = threadIdx.x; i < sorted_ids_size; i += stride) {
    sorted_ids[i] = SENTINEL;
  }
-  // Initialize expert_ids with -1
+  // Intialize expert_ids with -1
  for (size_t i = threadIdx.x; i < block_ids_size; i += stride) {
    block_ids[i] = -1;
  }
@@ -1,60 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-// Adapted from SGLang:
-// https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled.cu
-
-#include <torch/all.h>
-
-#include "cutlass_mxfp8_grouped_mm_launcher.cuh"
-
-void cutlass_mxfp8_grouped_mm(const torch::Tensor& a, const torch::Tensor& b,
-                              const torch::Tensor& sfa,
-                              const torch::Tensor& sfb, torch::Tensor& d,
-                              const torch::Tensor& problem_sizes,
-                              const torch::Tensor& expert_offsets,
-                              const torch::Tensor& blockscale_offsets) {
-#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
-  TORCH_CHECK(problem_sizes.dim() == 2, "problem_sizes must be 2D tensor");
-  TORCH_CHECK(problem_sizes.size(1) == 3,
-              "problem_sizes must have shape (num_experts, 3)");
-  TORCH_CHECK(problem_sizes.size(0) == expert_offsets.size(0),
-              "Number of experts in problem_sizes must match expert_offsets");
-  TORCH_CHECK(problem_sizes.dtype() == torch::kInt32,
-              "problem_sizes must be int32");
-  TORCH_CHECK(expert_offsets.dtype() == torch::kInt32,
-              "expert_offsets must be int32");
-  TORCH_CHECK(blockscale_offsets.dtype() == torch::kInt32,
-              "blockscale_offsets must be int32");
-  TORCH_CHECK(a.dim() == 2, "a must be a 2D tensor of shape (num_tokens, k)");
-  TORCH_CHECK(b.dim() == 3,
-              "b must be a 3D tensor of shape (num_experts, k, n)");
-  TORCH_CHECK(a.size(1) == b.size(1) && a.size(1) % 128 == 0,
-              "k should align 128");
-  TORCH_CHECK(b.size(2) % 128 == 0, "n should align 128");
-  TORCH_CHECK(a.strides()[1] == 1, "a must be row major");
-  TORCH_CHECK(b.strides()[1] == 1, "b must be column major");
-
-  auto stream = at::cuda::getCurrentCUDAStream();
-  if (d.dtype() == torch::kBFloat16) {
-    expert_specialization::cutlass_mxfp8_grouped_mm_dispatch_out_dtype<
-        cutlass::bfloat16_t>(a, b, sfa, sfb, d, problem_sizes, expert_offsets,
-                             blockscale_offsets, stream);
-  } else if (d.dtype() == torch::kFloat16) {
-    expert_specialization::cutlass_mxfp8_grouped_mm_dispatch_out_dtype<
-        cutlass::half_t>(a, b, sfa, sfb, d, problem_sizes, expert_offsets,
-                         blockscale_offsets, stream);
-  } else {
-    TORCH_CHECK(false, "dtype must be kFloat16 or kBFloat16");
-  }
-#else
-  TORCH_CHECK(false,
-              "No implemented cutlass_mxfp8_grouped_mm for "
-              "current device");
-#endif
-}
-
-#include "core/registration.h"
-
-TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
-  m.impl("cutlass_mxfp8_grouped_mm", cutlass_mxfp8_grouped_mm);
-}
@@ -1,141 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-// Adapted from SGLang:
-// https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled_functor.cuh
-
-#pragma once
-#include <cuda.h>
-
-#include "cute/tensor.hpp"
-#include "cutlass/util/packed_stride.hpp"
-#include "cutlass_mxfp8_grouped_mm_traits.cuh"
-
-namespace expert_specialization {
-
-using namespace cute;
-
-template <typename GemmTraits>
-struct CutlassMxfp8GroupedMmOffsetFunctor {
-  using Gemm = typename GemmTraits::Gemm;
-  using ElementA = typename Gemm::ElementA;
-  using ElementB = typename Gemm::ElementB;
-  using ElementSF = typename GemmTraits::ElementSF;
-  using ElementD = typename GemmTraits::ElementOutput;
-  // Input
-  int* expert_offsets{nullptr};
-  int* blockscale_offsets{nullptr};
-  // Output
-  ElementA* a_base{nullptr};
-  ElementB* b_base{nullptr};
-  ElementSF* sfa_base{nullptr};
-  ElementSF* sfb_base{nullptr};
-  ElementD* d_base{nullptr};
-  ElementA** a_offsets{nullptr};
-  ElementB** b_offsets{nullptr};
-  ElementSF** sfa_offsets{nullptr};
-  ElementSF** sfb_offsets{nullptr};
-  ElementD** d_offsets{nullptr};
-
-  CutlassMxfp8GroupedMmOffsetFunctor() = default;
-  CutlassMxfp8GroupedMmOffsetFunctor(
-      int* _expert_offsets, int* _blockscale_offsets, ElementA* _a_base,
-      ElementB* _b_base, ElementSF* _sfa_base, ElementSF* _sfb_base,
-      ElementD* _d_base, ElementA** _a_offsets, ElementB** _b_offsets,
-      ElementSF** _sfa_offsets, ElementSF** _sfb_offsets, ElementD** _d_offsets)
-      : expert_offsets{_expert_offsets},
-        blockscale_offsets{_blockscale_offsets},
-        a_base(_a_base),
-        b_base(_b_base),
-        sfa_base(_sfa_base),
-        sfb_base(_sfb_base),
-        d_base(_d_base),
-        a_offsets(_a_offsets),
-        b_offsets(_b_offsets),
-        sfa_offsets(_sfa_offsets),
-        sfb_offsets(_sfb_offsets),
-        d_offsets(_d_offsets) {}
-
-  void CUTE_DEVICE operator()(int64_t expert_id, int m, int n, int k) {
-    int64_t expert_offset = static_cast<int64_t>(expert_offsets[expert_id]);
-    int64_t blockscale_offset =
-        static_cast<int64_t>(blockscale_offsets[expert_id]);
-    int64_t a_stride = expert_offset * k;
-    int64_t b_stride = expert_id * k * n;
-    int64_t d_stride = expert_offset * n;
-    int64_t sfa_stride = blockscale_offset * (k / 32);
-    int64_t sfb_stride = expert_id * n * (k / 32);
-
-    a_offsets[expert_id] = a_base + a_stride;
-    b_offsets[expert_id] = b_base + b_stride;
-    sfa_offsets[expert_id] = sfa_base + sfa_stride;
-    sfb_offsets[expert_id] = sfb_base + sfb_stride;
-    d_offsets[expert_id] = d_base + d_stride;
-  }
-};
-
-template <typename GemmTraits>
-struct CutlassMxfp8GroupedMmLayoutFunctor {
-  using Sm1xxBlkScaledConfig = typename GemmTraits::Sm1xxBlkScaledConfig;
-  using LayoutSFA = typename GemmTraits::LayoutSFA;
-  using LayoutSFB = typename GemmTraits::LayoutSFB;
-  LayoutSFA* layout_sfa_base{nullptr};
-  LayoutSFB* layout_sfb_base{nullptr};
-
-  CutlassMxfp8GroupedMmLayoutFunctor() = default;
-  CutlassMxfp8GroupedMmLayoutFunctor(LayoutSFA* _layout_sfa_base,
-                                     LayoutSFB* _layout_sfb_base)
-      : layout_sfa_base(_layout_sfa_base), layout_sfb_base(_layout_sfb_base) {}
-
-  void CUTE_DEVICE operator()(int64_t expert_id, int m, int n, int k) {
-    LayoutSFA* layout_sfa_ptr = layout_sfa_base + expert_id;
-    LayoutSFB* layout_sfb_ptr = layout_sfb_base + expert_id;
-    *layout_sfa_ptr = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(
-        cute::make_shape(m, n, k, 1));
-    *layout_sfb_ptr = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(
-        cute::make_shape(m, n, k, 1));
-  }
-};
-
-template <typename GemmTraits>
-struct CutlassMxfp8GroupedMmStrideFunctor {
-  using StrideA = typename GemmTraits::StrideA;
-  using StrideB = typename GemmTraits::StrideB;
-  using StrideD = typename GemmTraits::StrideD;
-  StrideA* stride_A_base{nullptr};
-  StrideB* stride_B_base{nullptr};
-  StrideD* stride_D_base{nullptr};
-
-  CutlassMxfp8GroupedMmStrideFunctor() = default;
-  CutlassMxfp8GroupedMmStrideFunctor(StrideA* _stride_A_base,
-                                     StrideB* _stride_B_base,
-                                     StrideD* _stride_D_base)
-      : stride_A_base(_stride_A_base),
-        stride_B_base(_stride_B_base),
-        stride_D_base(_stride_D_base) {}
-
-  void CUTE_DEVICE operator()(int64_t expert_id, int m, int n, int k) {
-    StrideA* stride_A = stride_A_base + expert_id;
-    StrideB* stride_B = stride_B_base + expert_id;
-    StrideD* stride_D = stride_D_base + expert_id;
-    *stride_A = cutlass::make_cute_packed_stride(StrideA{}, {m, k, 1});
-    *stride_B = cutlass::make_cute_packed_stride(StrideB{}, {n, k, 1});
-    *stride_D = cutlass::make_cute_packed_stride(StrideD{}, {m, n, 1});
-  }
-};
-
-template <typename OffsetFunctor, typename LayoutFunctor,
-          typename StrideFunctor>
-__global__ void cutlassMxfp8GroupedMmPreComputeKernel(
-    int* problem_sizes, OffsetFunctor offset_functor,
-    LayoutFunctor layout_functor, StrideFunctor stride_functor) {
-  int64_t expert_id = static_cast<int64_t>(threadIdx.x);
-  int m = problem_sizes[expert_id * 3 + 0];
-  int n = problem_sizes[expert_id * 3 + 1];
-  int k = problem_sizes[expert_id * 3 + 2];
-
-  offset_functor(expert_id, m, n, k);
-  layout_functor(expert_id, m, n, k);
-  stride_functor(expert_id, m, n, k);
-}
-
-}  // namespace expert_specialization
@@ -1,179 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-// Adapted from SGLang:
-// https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled_launcher.cuh
-
-#pragma once
-#include <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAGuard.h>
-#include <torch/all.h>
-
-#include <cassert>
-#include <iostream>
-#include <string>
-
-#include "cute/tensor.hpp"
-#include "cutlass_mxfp8_grouped_mm_functor.cuh"
-#include "cutlass_mxfp8_grouped_mm_traits.cuh"
-
-namespace expert_specialization {
-
-template <typename GemmTraits>
-void cutlass_mxfp8_grouped_mm_pre_compute(
-    torch::Tensor& a_ptrs, torch::Tensor& b_ptrs, torch::Tensor& sfa_ptrs,
-    torch::Tensor& sfb_ptrs, torch::Tensor& d_ptrs, torch::Tensor& stride_a,
-    torch::Tensor& stride_b, torch::Tensor& stride_d, torch::Tensor& layout_sfa,
-    torch::Tensor& layout_sfb, const torch::Tensor& a, const torch::Tensor& b,
-    const torch::Tensor& sfa, const torch::Tensor& sfb, const torch::Tensor& d,
-    const torch::Tensor& problem_sizes, const torch::Tensor& expert_offsets,
-    const torch::Tensor& blockscale_offsets, cudaStream_t stream) {
-  using OffsetFunctor = CutlassMxfp8GroupedMmOffsetFunctor<GemmTraits>;
-  using ElementA = typename OffsetFunctor::ElementA;
-  using ElementB = typename OffsetFunctor::ElementB;
-  using ElementSF = typename OffsetFunctor::ElementSF;
-  using ElementD = typename OffsetFunctor::ElementD;
-
-  using LayoutFunctor = CutlassMxfp8GroupedMmLayoutFunctor<GemmTraits>;
-  using LayoutSFA = typename LayoutFunctor::LayoutSFA;
-  using LayoutSFB = typename LayoutFunctor::LayoutSFB;
-
-  using StrideFunctor = CutlassMxfp8GroupedMmStrideFunctor<GemmTraits>;
-  using StrideA = typename StrideFunctor::StrideA;
-  using StrideB = typename StrideFunctor::StrideB;
-  using StrideD = typename StrideFunctor::StrideD;
-
-  int num_experts = (int)expert_offsets.size(0);
-  TORCH_CHECK(num_experts <= 1024,
-              "Number of experts cannot exceed 1024, the maximum number of "
-              "threads per block.");
-
-  OffsetFunctor offset_functor(
-      reinterpret_cast<int*>(expert_offsets.data_ptr()),
-      reinterpret_cast<int*>(blockscale_offsets.data_ptr()),
-      reinterpret_cast<ElementA*>(a.data_ptr()),
-      reinterpret_cast<ElementB*>(b.data_ptr()),
-      reinterpret_cast<ElementSF*>(sfa.data_ptr()),
-      reinterpret_cast<ElementSF*>(sfb.data_ptr()),
-      reinterpret_cast<ElementD*>(d.data_ptr()),
-      reinterpret_cast<ElementA**>(a_ptrs.data_ptr()),
-      reinterpret_cast<ElementB**>(b_ptrs.data_ptr()),
-      reinterpret_cast<ElementSF**>(sfa_ptrs.data_ptr()),
-      reinterpret_cast<ElementSF**>(sfb_ptrs.data_ptr()),
-      reinterpret_cast<ElementD**>(d_ptrs.data_ptr()));
-  LayoutFunctor layout_functor(
-      reinterpret_cast<LayoutSFA*>(layout_sfa.data_ptr()),
-      reinterpret_cast<LayoutSFB*>(layout_sfb.data_ptr()));
-  StrideFunctor stride_functor(reinterpret_cast<StrideA*>(stride_a.data_ptr()),
-                               reinterpret_cast<StrideB*>(stride_b.data_ptr()),
-                               reinterpret_cast<StrideD*>(stride_d.data_ptr()));
-  cutlassMxfp8GroupedMmPreComputeKernel<<<1, num_experts, 0, stream>>>(
-      static_cast<int*>(problem_sizes.data_ptr()), offset_functor,
-      layout_functor, stride_functor);
-}
-
-template <typename GemmTraits>
-void cutlass_mxfp8_grouped_mm(
-    const torch::Tensor& a_ptrs, const torch::Tensor& b_ptrs,
-    const torch::Tensor& sfa_ptrs, const torch::Tensor& sfb_ptrs,
-    const torch::Tensor& d_ptrs, const torch::Tensor& stride_a,
-    const torch::Tensor& stride_b, const torch::Tensor& stride_d,
-    const torch::Tensor& layout_sfa, const torch::Tensor& layout_sfb,
-    const torch::Tensor& problem_sizes, cudaStream_t stream) {
-  using Gemm = typename GemmTraits::Gemm;
-  using ElementA = typename Gemm::ElementA;
-  using ElementB = typename Gemm::ElementB;
-  using ElementSF = typename GemmTraits::ElementSF;
-  using ElementD = typename GemmTraits::ElementOutput;
-  using StrideA = typename GemmTraits::StrideA;
-  using StrideB = typename GemmTraits::StrideB;
-  using StrideD = typename GemmTraits::StrideD;
-  using LayoutSFA = typename GemmTraits::LayoutSFA;
-  using LayoutSFB = typename GemmTraits::LayoutSFB;
-  using UnderlyingProblemShape =
-      typename GemmTraits::ProblemShape::UnderlyingProblemShape;
-
-  cutlass::KernelHardwareInfo hw_info;
-  hw_info.device_id = c10::cuda::current_device();
-  hw_info.sm_count =
-      at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
-  hw_info.cluster_shape = GemmTraits::MMAConfig::preferred_cluster;
-  hw_info.cluster_shape_fallback = GemmTraits::MMAConfig::fallback_cluster;
-
-  int num_experts = (int)problem_sizes.size(0);
-
-  UnderlyingProblemShape* underlying_problem_shape =
-      reinterpret_cast<UnderlyingProblemShape*>(problem_sizes.data_ptr());
-
-  typename Gemm::Arguments arguments = {
-      cutlass::gemm::GemmUniversalMode::kGrouped,
-      {num_experts, underlying_problem_shape, nullptr},
-      {reinterpret_cast<const ElementA**>(a_ptrs.data_ptr()),
-       reinterpret_cast<StrideA*>(stride_a.data_ptr()),
-       reinterpret_cast<const ElementB**>(b_ptrs.data_ptr()),
-       reinterpret_cast<StrideB*>(stride_b.data_ptr()),
-       reinterpret_cast<const ElementSF**>(sfa_ptrs.data_ptr()),
-       reinterpret_cast<LayoutSFA*>(layout_sfa.data_ptr()),
-       reinterpret_cast<const ElementSF**>(sfb_ptrs.data_ptr()),
-       reinterpret_cast<LayoutSFB*>(layout_sfb.data_ptr())},
-      {{},
-       nullptr,
-       nullptr,
-       reinterpret_cast<ElementD**>(d_ptrs.data_ptr()),
-       reinterpret_cast<StrideD*>(stride_d.data_ptr())},
-      hw_info,
-      {}  // Scheduler
-  };
-
-  Gemm gemm;
-
-  auto can_implement_status = gemm.can_implement(arguments);
-  TORCH_CHECK(can_implement_status == cutlass::Status::kSuccess,
-              "Failed to implement GEMM");
-
-  torch::TensorOptions options_uint8 =
-      torch::TensorOptions().dtype(torch::kUInt8).device(d_ptrs.device());
-  size_t workspace_size = gemm.get_workspace_size(arguments);
-  torch::Tensor workspace = torch::empty(workspace_size, options_uint8);
-
-  auto status = gemm.initialize(arguments, workspace.data_ptr(), stream);
-  TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to initialize GEMM");
-
-  status = gemm.run(stream, nullptr, true);  // Enable PDL
-  TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to run GEMM");
-}
-
-template <typename OutType>
-void cutlass_mxfp8_grouped_mm_dispatch_out_dtype(
-    const torch::Tensor& a, const torch::Tensor& b, const torch::Tensor& sfa,
-    const torch::Tensor& sfb, torch::Tensor& d,
-    const torch::Tensor& problem_sizes, const torch::Tensor& expert_offsets,
-    const torch::Tensor& blockscale_offsets, cudaStream_t stream) {
-  int num_experts = (int)problem_sizes.size(0);
-  torch::TensorOptions options_int64 =
-      torch::TensorOptions().dtype(torch::kInt64).device(a.device());
-  torch::TensorOptions options_int32 =
-      torch::TensorOptions().dtype(torch::kInt32).device(a.device());
-
-  torch::Tensor a_ptrs = torch::empty(num_experts, options_int64);
-  torch::Tensor b_ptrs = torch::empty(num_experts, options_int64);
-  torch::Tensor sfa_ptrs = torch::empty(num_experts, options_int64);
-  torch::Tensor sfb_ptrs = torch::empty(num_experts, options_int64);
-  torch::Tensor d_ptrs = torch::empty(num_experts, options_int64);
-
-  torch::Tensor stride_a = torch::empty(num_experts, options_int64);
-  torch::Tensor stride_b = torch::empty(num_experts, options_int64);
-  torch::Tensor stride_d = torch::empty(num_experts, options_int64);
-  torch::Tensor layout_sfa = torch::empty({num_experts, 5}, options_int32);
-  torch::Tensor layout_sfb = torch::empty({num_experts, 5}, options_int32);
-
-  using GemmTraits = CutlassMxfp8GroupedMmGemmTraits<MMA1SMConfig, OutType>;
-  cutlass_mxfp8_grouped_mm_pre_compute<GemmTraits>(
-      a_ptrs, b_ptrs, sfa_ptrs, sfb_ptrs, d_ptrs, stride_a, stride_b, stride_d,
-      layout_sfa, layout_sfb, a, b, sfa, sfb, d, problem_sizes, expert_offsets,
-      blockscale_offsets, stream);
-  cutlass_mxfp8_grouped_mm<GemmTraits>(
-      a_ptrs, b_ptrs, sfa_ptrs, sfb_ptrs, d_ptrs, stride_a, stride_b, stride_d,
-      layout_sfa, layout_sfb, problem_sizes, stream);
-}
-
-}  // namespace expert_specialization
@@ -1,127 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-// Adapted from SGLang:
-// https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled_traits.cuh
-
-#pragma once
-
-// Misc
-#include "cute/tensor.hpp"
-#include "cutlass/arch/arch.h"
-#include "cutlass/arch/mma.h"
-#include "cutlass/cutlass.h"
-#include "cutlass/detail/sm100_blockscaled_layout.hpp"
-#include "cutlass/epilogue/dispatch_policy.hpp"
-#include "cutlass/gemm/dispatch_policy.hpp"
-#include "cutlass/gemm/group_array_problem_shape.hpp"
-#include "cutlass/layout/layout.h"
-#include "cutlass/numeric_conversion.h"
-#include "cutlass/numeric_size.h"
-
-// Collective Builder
-#include "cutlass/epilogue/collective/collective_builder.hpp"
-#include "cutlass/epilogue/fusion/sm90_callbacks_tma_warpspecialized.hpp"
-#include "cutlass/epilogue/thread/activation.h"
-#include "cutlass/gemm/collective/collective_builder.hpp"
-
-// Integration
-#include "cutlass/gemm/device/gemm_universal_adapter.h"
-#include "cutlass/gemm/kernel/gemm_universal.hpp"
-
-namespace expert_specialization {
-
-using namespace cute;
-
-// Different configs for 1SM and 2SM MMA kernel
-struct MMA1SMConfig {
-  using MmaTileShape = Shape<_128, _128, _128>;
-  using KernelSchedule =
-      cutlass::gemm::KernelPtrArrayTmaWarpSpecialized1SmMxf8f6f4Sm100;
-  using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecialized1Sm;
-  const static dim3 preferred_cluster;
-  const static dim3 fallback_cluster;
-};
-const dim3 MMA1SMConfig::preferred_cluster(1, 4, 1);
-const dim3 MMA1SMConfig::fallback_cluster(1, 2, 1);
-
-template <typename _MMAConfig, typename OutputDtype>
-struct CutlassMxfp8GroupedMmGemmTraits {
-  using MMAConfig = _MMAConfig;
-  using ElementInput = cutlass::float_e4m3_t;
-  using ElementOutput = OutputDtype;
-  using ProblemShape = cutlass::gemm::GroupProblemShape<Shape<int, int, int>>;
-
-  // A matrix configuration
-  using ElementA = cutlass::mx_float8_t<ElementInput>;
-  using LayoutA = cutlass::layout::RowMajor;
-  constexpr static int AlignmentA = 32;
-
-  // B matrix configuration
-  using ElementB = cutlass::mx_float8_t<ElementInput>;
-  using LayoutB = cutlass::layout::ColumnMajor;
-  constexpr static int AlignmentB = 32;
-
-  // C/D matrix configuration
-  using ElementC = void;
-  using ElementD = ElementOutput;
-  using LayoutC = cutlass::layout::RowMajor;
-  using LayoutD = cutlass::layout::RowMajor;
-  constexpr static int AlignmentC = 128 / cutlass::sizeof_bits<ElementD>::value;
-  constexpr static int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
-  using ElementAccumulator = float;
-
-  static constexpr auto RoundStyle = cutlass::FloatRoundStyle::round_to_nearest;
-  using CustomEVTIdentity =  // acc
-      cutlass::epilogue::fusion::Sm90EVT<
-          cutlass::epilogue::fusion::Sm90Compute<
-              cutlass::epilogue::thread::Identity, ElementD, ElementAccumulator,
-              RoundStyle>,
-          cutlass::epilogue::fusion::Sm90AccFetch>;
-
-  // Core kernel configurations
-  using ArchTag = cutlass::arch::Sm100;
-  using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp;
-  using StageCountType = cutlass::gemm::collective::StageCountAuto;
-
-  // Runtime Cluster Shape
-  using ClusterShape = Shape<int32_t, int32_t, _1>;
-
-  // Define Epilogue
-  using CollectiveEpilogue =
-      typename cutlass::epilogue::collective::CollectiveBuilder<
-          ArchTag, OperatorClass, typename MMAConfig::MmaTileShape,
-          ClusterShape, Shape<_64, _64>, ElementAccumulator, ElementAccumulator,
-          ElementC, LayoutC*, AlignmentC, ElementD, LayoutD*, AlignmentD,
-          typename MMAConfig::EpilogueSchedule,
-          CustomEVTIdentity>::CollectiveOp;
-
-  // Define Mainloop
-  using CollectiveMainloop =
-      typename cutlass::gemm::collective::CollectiveBuilder<
-          ArchTag, OperatorClass, ElementA, LayoutA*, AlignmentA, ElementB,
-          LayoutB*, AlignmentB, ElementAccumulator,
-          typename MMAConfig::MmaTileShape, ClusterShape,
-          cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
-              sizeof(typename CollectiveEpilogue::SharedStorage))>,
-          typename MMAConfig::KernelSchedule>::CollectiveOp;
-
-  // Define GemmKernel
-  using GemmKernel =
-      cutlass::gemm::kernel::GemmUniversal<ProblemShape, CollectiveMainloop,
-                                           CollectiveEpilogue>;
-  using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
-
-  using ElementSF = typename Gemm::GemmKernel::ElementSF;
-  using StrideA = typename Gemm::GemmKernel::InternalStrideA;
-  using StrideB = typename Gemm::GemmKernel::InternalStrideB;
-  using StrideC = typename Gemm::GemmKernel::InternalStrideC;
-  using StrideD = typename Gemm::GemmKernel::InternalStrideD;
-  using LayoutSFA =
-      typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFA;
-  using LayoutSFB =
-      typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFB;
-  using Sm1xxBlkScaledConfig =
-      typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig;
-};
-
-}  // namespace expert_specialization
@@ -1,60 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-// Adapted from SGLang:
-// https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled_group_quant.cu
-
-#include <torch/all.h>
-
-#include "mxfp8_experts_quant.cuh"
-
-void mxfp8_experts_quant(const torch::Tensor& input,
-                         const torch::Tensor& problem_sizes,
-                         const torch::Tensor& expert_offsets,
-                         const torch::Tensor& blockscale_offsets,
-                         torch::Tensor& quant_output,
-                         torch::Tensor& scale_factor) {
-#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
-  TORCH_CHECK(input.dim() == 2, "input must be 2D tensor");
-  TORCH_CHECK(input.size(1) % 128 == 0, "k must align to 128");
-  TORCH_CHECK(input.strides()[1] == 1, "input must be row major");
-  TORCH_CHECK(problem_sizes.dim() == 2, "problem_sizes must be 2D tensor");
-  TORCH_CHECK(problem_sizes.dtype() == torch::kInt32,
-              "problem_sizes must be int32");
-  TORCH_CHECK(expert_offsets.dtype() == torch::kInt32,
-              "expert_offsets must be int32");
-  TORCH_CHECK(blockscale_offsets.dtype() == torch::kInt32,
-              "blockscale_offsets must be int32");
-
-  auto groups = problem_sizes.size(0);
-  TORCH_CHECK(
-      expert_offsets.dim() == 1 && expert_offsets.size(0) == groups,
-      "expert_offsets must be 1D and have size equal to the number of groups");
-  TORCH_CHECK(
-      blockscale_offsets.dim() == 1 && blockscale_offsets.size(0) == groups,
-      "blockscale_offsets must be 1D and have size equal to the number of "
-      "groups");
-
-  auto stream = at::cuda::getCurrentCUDAStream();
-  if (input.dtype() == torch::kBFloat16) {
-    expert_specialization::launch_mxfp8_experts_quant<__nv_bfloat16>(
-        input, problem_sizes, expert_offsets, blockscale_offsets, quant_output,
-        scale_factor);
-  } else if (input.dtype() == torch::kFloat16) {
-    expert_specialization::launch_mxfp8_experts_quant<__half>(
-        input, problem_sizes, expert_offsets, blockscale_offsets, quant_output,
-        scale_factor);
-  } else {
-    TORCH_CHECK(false, "dtype must be kFloat16 or kBFloat16");
-  }
-#else
-  TORCH_CHECK(false,
-              "No implemented mxfp8_experts_quant for "
-              "current device");
-#endif
-}
-
-#include "core/registration.h"
-
-TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
-  m.impl("mxfp8_experts_quant", mxfp8_experts_quant);
-}
@@ -1,414 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-// Adapted from SGLang:
-// https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled_group_quant.cuh
-
-#pragma once
-#include <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAGuard.h>
-#include <cuda.h>
-#include <cuda_bf16.h>
-#include <cuda_fp16.h>
-#include <torch/all.h>
-
-#include <cuda/ptx>
-
-#include "cute/tensor.hpp"
-
-namespace expert_specialization {
-
-using namespace cute;
-
-constexpr uint32_t THREAD_BLOCK_SIZE = 128;
-constexpr uint32_t WARP_SIZE = 32;
-constexpr int BLOCK_M = 128;
-constexpr int BLOCK_K = 128;
-using ThrLayout = Layout<Shape<_16, _8>, Stride<_8, _1>>;
-using ValLayout = Layout<Shape<_1, _16>>;
-using SfR2SThrLayout = Layout<Shape<_16, _4>, Stride<_4, _1>>;
-using SfR2SValLayout = Layout<Shape<_1, _1>>;
-using ScaleFactorTileLayout =
-    Layout<Shape<Shape<_32, _4>, _4>, Stride<Stride<_16, _4>, _1>>;
-
-// Fast reciprocal.
-inline __device__ float reciprocal_approximate_ftz(float a) {
-  float b;
-  asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a));
-  return b;
-}
-
-// Some code references TRT-LLM:
-// https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/quantization.cuh
-template <typename FragmentS, typename FragmentD>
-__inline__ __device__ uint8_t cvt_warp_fp16_to_mxfp8(FragmentS& fragment_s,
-                                                     FragmentD& fragment_d) {
-  using FragmentSLayout = typename FragmentS::layout_type;
-  using FragmentDLayout = typename FragmentD::layout_type;
-  FragmentSLayout fragment_s_layout;
-  FragmentDLayout fragment_d_layout;
-  static_assert(is_static<FragmentSLayout>::value &&
-                size(fragment_s_layout) == 16);
-  static_assert(is_static<FragmentDLayout>::value &&
-                size(fragment_d_layout) == 16);
-
-  constexpr int eles_per_thr = 16;
-  using ValType = typename FragmentS::element_type;
-  using VecType = std::conditional_t<std::is_same_v<ValType, __nv_bfloat16>,
-                                     __nv_bfloat162, __half2>;
-  VecType vec[8];
-  // Assign vals
-  vec[0].x = fragment_s(Int<0>{});
-  vec[0].y = fragment_s(Int<1>{});
-  vec[1].x = fragment_s(Int<2>{});
-  vec[1].y = fragment_s(Int<3>{});
-  vec[2].x = fragment_s(Int<4>{});
-  vec[2].y = fragment_s(Int<5>{});
-  vec[3].x = fragment_s(Int<6>{});
-  vec[3].y = fragment_s(Int<7>{});
-  vec[4].x = fragment_s(Int<8>{});
-  vec[4].y = fragment_s(Int<9>{});
-  vec[5].x = fragment_s(Int<10>{});
-  vec[5].y = fragment_s(Int<11>{});
-  vec[6].x = fragment_s(Int<12>{});
-  vec[6].y = fragment_s(Int<13>{});
-  vec[7].x = fragment_s(Int<14>{});
-  vec[7].y = fragment_s(Int<15>{});
-
-  auto local_max = __habs2(vec[0]);
-  for (int i = 1; i < eles_per_thr / 2; i++) {
-    local_max = __hmax2(__habs2(vec[i]), local_max);
-  }
-  local_max = __hmax2(__shfl_xor_sync(uint32_t(-1), local_max, 1), local_max);
-
-  // Get the final absolute maximum values.
-  float block_max(0.0f);
-  if constexpr (std::is_same_v<ValType, __nv_bfloat16>) {
-    block_max = __bfloat162float(__hmax(local_max.x, local_max.y));
-  } else {
-    block_max = __half2float(__hmax(local_max.x, local_max.y));
-  }
-  // Get the SF (max value of the vector / max value of mxfp8).
-  float sf_val = block_max * reciprocal_approximate_ftz(448.0f);
-  // 8 bits representation of the SF.
-  uint8_t fp8_sf_val;
-
-  __nv_fp8_e8m0 tmp_sf_val;
-  tmp_sf_val.__x =
-      __nv_cvt_float_to_e8m0(sf_val, __NV_SATFINITE, cudaRoundPosInf);
-  sf_val = static_cast<float>(tmp_sf_val);
-  fp8_sf_val = tmp_sf_val.__x;
-  // Get the output scale (reciprocal of the SFValue).
-  float output_scale =
-      block_max != 0.f ? reciprocal_approximate_ftz(sf_val) : 0.0f;
-
-  // Convert the input to float.
-  float2 fp2_vals[eles_per_thr / 2];
-
-#pragma unroll
-  for (int i = 0; i < eles_per_thr / 2; i++) {
-    if constexpr (std::is_same_v<ValType, __half>) {
-      fp2_vals[i] = __half22float2(vec[i]);
-    } else {
-      fp2_vals[i] = __bfloat1622float2(vec[i]);
-    }
-    fp2_vals[i].x *= output_scale;
-    fp2_vals[i].y *= output_scale;
-  }
-  union {
-    uint8_t bytes[16];
-    __nv_fp8x2_e4m3 elts[8];
-  } u;
-  u.elts[0] = __nv_fp8x2_e4m3(fp2_vals[0]);
-  u.elts[1] = __nv_fp8x2_e4m3(fp2_vals[1]);
-  u.elts[2] = __nv_fp8x2_e4m3(fp2_vals[2]);
-  u.elts[3] = __nv_fp8x2_e4m3(fp2_vals[3]);
-  u.elts[4] = __nv_fp8x2_e4m3(fp2_vals[4]);
-  u.elts[5] = __nv_fp8x2_e4m3(fp2_vals[5]);
-  u.elts[6] = __nv_fp8x2_e4m3(fp2_vals[6]);
-  u.elts[7] = __nv_fp8x2_e4m3(fp2_vals[7]);
-  fragment_d(Int<0>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[0]);
-  fragment_d(Int<1>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[1]);
-  fragment_d(Int<2>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[2]);
-  fragment_d(Int<3>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[3]);
-  fragment_d(Int<4>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[4]);
-  fragment_d(Int<5>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[5]);
-  fragment_d(Int<6>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[6]);
-  fragment_d(Int<7>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[7]);
-  fragment_d(Int<8>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[8]);
-  fragment_d(Int<9>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[9]);
-  fragment_d(Int<10>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[10]);
-  fragment_d(Int<11>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[11]);
-  fragment_d(Int<12>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[12]);
-  fragment_d(Int<13>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[13]);
-  fragment_d(Int<14>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[14]);
-  fragment_d(Int<15>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[15]);
-  return fp8_sf_val;
-}
-
-template <typename TensorS, typename TensorP, typename TensorD,
-          typename TensorSharedSF, typename TensorSF, typename TiledCopyG2R,
-          typename TiledCopyR2G, typename TiledCopyR2S>
-__inline__ __device__ void mxfp8_experts_quant_tile(
-    TensorS& tensor_s, TensorP& tensor_p, TensorD& tensor_d,
-    TensorSharedSF& tensor_shared_sf, TensorSF& tensor_sf, int m,
-    TiledCopyG2R& tiled_copy_g2r, TiledCopyR2G& tiled_copy_r2g,
-    TiledCopyR2S& tiled_copy_r2s) {
-  static_assert(size(get<0>(typename TensorS::layout_type{})) == 128 &&
-                size(get<1>(typename TensorS::layout_type{})) == 128 &&
-                stride(get<1>(typename TensorS::layout_type{})) == 1);
-  static_assert(size(get<0>(typename TensorD::layout_type{})) == 128 &&
-                size(get<1>(typename TensorD::layout_type{})) == 128 &&
-                stride(get<1>(typename TensorD::layout_type{})) == 1);
-  static_assert(size(get<0>(typename TensorP::layout_type{})) == 128 &&
-                size(get<1>(typename TensorP::layout_type{})) == 128);
-  static_assert(size(get<0>(typename TensorSharedSF::layout_type{})) == 128 &&
-                size(get<1>(typename TensorSharedSF::layout_type{})) == 4);
-  static_assert(size(get<0>(typename TensorSF::layout_type{})) == 128 &&
-                size(get<1>(typename TensorSF::layout_type{})) == 4);
-
-  using Tiler_MN = typename TiledCopyG2R::Tiler_MN;
-  auto tiler_mn = Tiler_MN{};
-  static_assert(size<0>(tiler_mn) == 16 && size<1>(tiler_mn) == 128);
-
-  auto tiled_tensor_s = tiled_divide(tensor_s, tiler_mn);
-  auto tiled_tensor_p = tiled_divide(tensor_p, tiler_mn);
-  auto tiled_tensor_d = tiled_divide(tensor_d, tiler_mn);
-  static_assert(size<2>(tiled_tensor_s) == 1);
-  static_assert(size<2>(tiled_tensor_p) == 1);
-  static_assert(size<2>(tiled_tensor_d) == 1);
-  auto squeeze_tiled_tensor_s = take<0, 2>(tiled_tensor_s);
-  auto squeeze_tiled_tensor_p = take<0, 2>(tiled_tensor_p);
-  auto squeeze_tiled_tensor_d = take<0, 2>(tiled_tensor_d);
-
-  using SF_Tiler_MN = typename TiledCopyR2S::Tiler_MN;
-  auto sf_tiler_mn = SF_Tiler_MN{};
-  static_assert(size<0>(sf_tiler_mn) == 16 && size<1>(sf_tiler_mn) == 4);
-
-  auto tiled_tensor_sf = tiled_divide(tensor_sf, sf_tiler_mn);
-  auto tiled_tensor_shared_sf = tiled_divide(tensor_shared_sf, sf_tiler_mn);
-  auto squeeze_tiled_tensor_sf = take<0, 2>(tiled_tensor_sf);
-  auto squeeze_tiled_tensor_shared_sf = take<0, 2>(tiled_tensor_shared_sf);
-
-  constexpr int tile_loop_count = size<1>(tiled_tensor_s);
-  constexpr int rows_in_tile = 16;
-  // We don't need to clear shared memory
-  // clear(squeeze_tiled_tensor_shared_sf);
-#pragma unroll 4
-  for (int t = 0; t < tile_loop_count; t++) {
-    if (t * rows_in_tile >= m) {
-      break;
-    }
-    auto current_copy_tile_s = tensor<0>(squeeze_tiled_tensor_s(_, t));
-    auto current_copy_tile_p = tensor<0>(squeeze_tiled_tensor_p(_, t));
-    auto current_copy_tile_d = tensor<0>(squeeze_tiled_tensor_d(_, t));
-    auto current_copy_tile_sf = tensor<0>(squeeze_tiled_tensor_sf(_, t));
-    auto current_copy_tile_shared_sf =
-        tensor<0>(squeeze_tiled_tensor_shared_sf(_, t));
-
-    // Global to Register copy
-    auto thr_copy_g2r = tiled_copy_g2r.get_thread_slice(threadIdx.x);
-    auto thr_tile_g2r_s = thr_copy_g2r.partition_S(current_copy_tile_s);
-    auto thr_tile_g2r_p = thr_copy_g2r.partition_S(current_copy_tile_p);
-    auto input_fragment = make_fragment_like(thr_tile_g2r_s);
-
-    // Register to Global copy
-    auto thr_copy_r2g = tiled_copy_r2g.get_thread_slice(threadIdx.x);
-    auto thr_tile_r2g_d = thr_copy_r2g.partition_D(current_copy_tile_d);
-    auto thr_tile_r2g_p = thr_copy_r2g.partition_D(current_copy_tile_p);
-    auto output_fragment = make_fragment_like(thr_tile_r2g_d);
-
-    // Register to Shared copy
-    auto thr_copy_r2s = tiled_copy_r2s.get_thread_slice(threadIdx.x / 2);
-    auto thr_tile_r2s_shared_sf =
-        thr_copy_r2s.partition_D(current_copy_tile_shared_sf);
-    auto shared_sf_fragment = make_fragment_like(thr_tile_r2s_shared_sf);
-
-    // CopyG2R & convert & CopyR2G
-    copy_if(tiled_copy_g2r, thr_tile_g2r_p, thr_tile_g2r_s, input_fragment);
-    uint8_t fp8_sf_val =
-        cvt_warp_fp16_to_mxfp8(input_fragment, output_fragment);
-    copy_if(tiled_copy_r2g, thr_tile_r2g_p, output_fragment, thr_tile_r2g_d);
-    shared_sf_fragment[0] = fp8_sf_val;
-
-    // Before first copy r2s, clear shared memory and wait previous group
-    if (t == 0 && threadIdx.x == 0) {
-      // Wait for the group to have completed reading from shared memory.
-      cuda::ptx::cp_async_bulk_wait_group_read(cuda::ptx::n32_t<0>());
-    }
-    __syncthreads();
-
-    if (threadIdx.x % 2 == 0) {
-      copy(tiled_copy_r2s, shared_sf_fragment, thr_tile_r2s_shared_sf);
-    }
-    __syncthreads();
-  }
-
-  // Wait for shared memory writes to be visible to TMA engine.
-  cuda::ptx::fence_proxy_async(cuda::ptx::space_shared);  // b)
-  __syncthreads();
-
-  if (threadIdx.x == 0) {
-    cuda::ptx::cp_async_bulk(cuda::ptx::space_global, cuda::ptx::space_shared,
-                             squeeze_tiled_tensor_sf.data().get(),
-                             squeeze_tiled_tensor_shared_sf.data().get(), 512);
-    // Wait for TMA transfer to have finished reading shared memory.
-    // Create a "bulk async-group" out of the previous bulk copy operation.
-    cuda::ptx::cp_async_bulk_commit_group();
-  }
-  __syncthreads();
-}
-
-template <typename T_IN, typename TiledCopyG2R, typename TiledCopyR2G,
-          typename TiledCopyR2S>
-__global__ void mxfp8_experts_quant_kernel(
-    const T_IN* input, const int* problem_sizes, const int* expert_offsets,
-    const int* blockscale_offsets, cutlass::float_e4m3_t* quant_output,
-    uint8_t* scale_factor, int groups, TiledCopyG2R tiled_copy_g2r,
-    TiledCopyR2G tiled_copy_r2g, TiledCopyR2S tiled_copy_r2s) {
-#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000
-  __shared__ __align__(512) uint8_t shared_memory[512];
-  ScaleFactorTileLayout scale_factor_tile_layout{};
-  auto scale_factor_shared =
-      make_tensor(make_smem_ptr(shared_memory),
-                  scale_factor_tile_layout);  // ((_32,_4), _4):((_16,_4), _1)
-  // TODO: Transform Groupwise Schedule into a more efficient Schedule
-  for (int g = 0; g < groups; g++) {
-    int m = problem_sizes[g * 3 + 0];
-    int k = problem_sizes[g * 3 + 2];
-    int64_t expert_offset = static_cast<int64_t>(expert_offsets[g]);
-    int64_t blockscale_offset = static_cast<int64_t>(blockscale_offsets[g]);
-
-    auto input_tensor = make_tensor(
-        make_gmem_ptr(input + expert_offset * k),
-        make_layout(make_shape(m, k),
-                    LayoutRight{}));  // (M, K):(K, 1) half_t/bfloat16_t
-
-    auto quant_output_tensor = make_tensor(
-        make_gmem_ptr(quant_output + expert_offset * k),
-        make_layout(make_shape(m, k),
-                    LayoutRight{}));  // (M, K):(K, 1) cutlass::float_e4m3_t
-
-    auto scale_factor_shape = make_shape(ceil_div(m, 128) * 128, k / 32);
-    auto scale_factor_layout = tile_to_shape(scale_factor_tile_layout,
-                                             scale_factor_shape, LayoutRight{});
-    // layout<0>(layout<0>(scale_factor_layout))  (_32,_4):(_16,_4) -- static
-    // layout<1>(layout<0>(scale_factor_layout))  M_align_128 / 128 -- dynamic
-    // shape dynamic stride layout<0>(layout<1>(scale_factor_layout))  _4:_1 --
-    // static layout<1>(layout<1>(scale_factor_layout))  (K / 32) / 4 : _512 --
-    // dynamic shape static stride
-
-    // Reshape to zipped layout for 1D indexing
-    auto zipped_scale_factor_layout = make_layout(
-        make_layout(layout<0>(layout<0>(scale_factor_layout)),
-                    layout<0>(layout<1>(scale_factor_layout))),
-        make_layout(
-            layout<1>(layout<0>(scale_factor_layout)),
-            layout<1>(layout<1>(
-                scale_factor_layout))));  // (((_32,_4),_4),(M_align_128 /
-                                          // 128,(K / 32) /
-                                          // 4)):(((_16,_4),_1),(?,_512))
-
-    auto scale_factor_tensor =
-        make_tensor(make_gmem_ptr(scale_factor + blockscale_offset * (k / 32)),
-                    zipped_scale_factor_layout);
-
-    // Used for cases where M is not divisible by 128 (most scenarios).
-    auto input_shape = shape(input_tensor);  // (M, K):(K, 1)
-    auto identity_tensor = make_identity_tensor(input_shape);
-    auto predict_tensor = cute::lazy::transform(
-        identity_tensor, [&](auto c) { return elem_less(c, input_shape); });
-
-    // (_128, _128)
-    auto tiler = make_shape(Int<BLOCK_M>{}, Int<BLOCK_K>{});
-
-    auto tiled_input_tensor = zipped_divide(
-        input_tensor, tiler);  // ((128, 128), (cdiv(M, 128), cdiv(K, 128)))
-    auto tiled_quant_output_tensor =
-        zipped_divide(quant_output_tensor,
-                      tiler);  // ((128, 128), (cdiv(M, 128), cdiv(K, 128)))
-    auto tiled_predict_tensor = zipped_divide(
-        predict_tensor, tiler);  // ((128, 128), (cdiv(M, 128), cdiv(K, 128)))
-
-    auto total_tiles =
-        size<1>(tiled_input_tensor);  // cdiv(M, 128) * cdiv(K, 128)
-    decltype(total_tiles) blk_offset = blockIdx.x;
-    while (blk_offset < total_tiles) {
-      auto current_input_tile = tensor<0>(tiled_input_tensor(_, blk_offset));
-      auto current_quant_output_tile =
-          tensor<0>(tiled_quant_output_tensor(_, blk_offset));
-      auto current_predict_tile =
-          tensor<0>(tiled_predict_tensor(_, blk_offset));
-      auto current_scale_factor_tile =
-          tensor<0>(scale_factor_tensor(_, blk_offset));
-
-      mxfp8_experts_quant_tile<
-          decltype(current_input_tile), decltype(current_predict_tile),
-          decltype(current_quant_output_tile), decltype(scale_factor_shared),
-          decltype(current_scale_factor_tile), TiledCopyG2R, TiledCopyR2G,
-          TiledCopyR2S>(current_input_tile, current_predict_tile,
-                        current_quant_output_tile, scale_factor_shared,
-                        current_scale_factor_tile, m, tiled_copy_g2r,
-                        tiled_copy_r2g, tiled_copy_r2s);
-      blk_offset += gridDim.x;
-    }
-  }
-#endif
-}
-
-template <typename T_IN>
-void launch_mxfp8_experts_quant(const torch::Tensor& input,
-                                const torch::Tensor& problem_sizes,
-                                const torch::Tensor& expert_offsets,
-                                const torch::Tensor& blockscale_offsets,
-                                torch::Tensor& quant_output,
-                                torch::Tensor& scale_factor) {
-  ThrLayout thr_layout{};
-  ValLayout val_layout{};
-  SfR2SThrLayout r2s_thr_layout{};
-  SfR2SValLayout r2s_val_layout{};
-
-  using CopyOpG2R =
-      UniversalCopy<cutlass::AlignedArray<T_IN, size(val_layout)>>;
-  using CopyAtomG2R = cute::Copy_Atom<CopyOpG2R, T_IN>;
-  auto tiled_copy_g2r = cute::make_tiled_copy(
-      CopyAtomG2R{}, thr_layout, val_layout);  // Tiler_MN: (16, 128)
-
-  using CopyOpR2G = UniversalCopy<
-      cutlass::AlignedArray<cutlass::float_e4m3_t, size(val_layout)>>;
-  using CopyAtomR2G = cute::Copy_Atom<CopyOpR2G, cutlass::float_e4m3_t>;
-  auto tiled_copy_r2g = cute::make_tiled_copy(
-      CopyAtomR2G{}, thr_layout, val_layout);  // Tiler_MN: (16, 128)
-
-  using CopyOpR2S =
-      UniversalCopy<cutlass::AlignedArray<uint8_t, size(r2s_val_layout)>>;
-  using CopyAtomR2S = cute::Copy_Atom<CopyOpR2S, uint8_t>;
-  auto tiled_copy_r2s = cute::make_tiled_copy(
-      CopyAtomR2S{}, r2s_thr_layout, r2s_val_layout);  // Tiler_MN: (16, 4)
-
-  int max_active_blocks_per_sm = -1;
-  AT_CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(
-      &max_active_blocks_per_sm,
-      mxfp8_experts_quant_kernel<T_IN, decltype(tiled_copy_g2r),
-                                 decltype(tiled_copy_r2g),
-                                 decltype(tiled_copy_r2s)>,
-      THREAD_BLOCK_SIZE, 0));
-
-  dim3 grid(at::cuda::getCurrentDeviceProperties()->multiProcessorCount *
-                max_active_blocks_per_sm,
-            1, 1);
-  dim3 block(THREAD_BLOCK_SIZE, 1, 1);
-  int num_experts = (int)problem_sizes.size(0);
-  auto stream = at::cuda::getCurrentCUDAStream();
-  mxfp8_experts_quant_kernel<T_IN, decltype(tiled_copy_g2r),
-                             decltype(tiled_copy_r2g), decltype(tiled_copy_r2s)>
-      <<<grid, block, 0, stream>>>(
-          reinterpret_cast<const T_IN*>(input.data_ptr()),
-          reinterpret_cast<const int*>(problem_sizes.data_ptr()),
-          reinterpret_cast<const int*>(expert_offsets.data_ptr()),
-          reinterpret_cast<const int*>(blockscale_offsets.data_ptr()),
-          reinterpret_cast<cutlass::float_e4m3_t*>(quant_output.data_ptr()),
-          reinterpret_cast<uint8_t*>(scale_factor.data_ptr()), num_experts,
-          tiled_copy_g2r, tiled_copy_r2g, tiled_copy_r2s);
-}
-
-}  // namespace expert_specialization
@@ -542,7 +542,7 @@ __global__ void silu_mul_fp8_quant_deep_gemm_kernel(
      if (!lane_id) {
        // Store scales.
        if constexpr (std::is_same<scale_t, uint8_t>::value) {
-          // Packed UE8M0 format. Remove Mantissa.
+          // Packed UE8MO format. Remove Mantissa.
          *y_s_ptr = reinterpret_cast<int16_t&>(y_s) >> 7;

          bool const jump_pack = (current_group_id + 1) % 4 == 0;
@@ -12,7 +12,6 @@
 #include "../cuda_compat.h"
 #include "dispatch_utils.h"
 #include "quantization/w8a8/fp8/common.cuh"
-#include "core/batch_invariant.hpp"

 // TODO(rasmith): The kernels in this file are susceptible to integer overflow
 // issues, do not take strides, and are unable to handle PyTorch tensors that
@@ -1225,14 +1224,17 @@ torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b,
 #if defined(__gfx950__)
  #define WVSPLITKRC_1KPASS
 template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
-          int UNRL, int N, int GrpsShrB, int CHUNKK, int DTRMNSTC>
+          int UNRL, int N, int GrpsShrB, int CHUNKK>
 __global__ void __launch_bounds__(WvPrGrp* THRDS)
    __attribute__((amdgpu_waves_per_eu(1, 1)))
-    wvSplitKrc_(const int actlN, const int K, const int Kap, const int M,
-                const int Bx, const int By, const scalar_t* __restrict__ A,
-                const scalar_t* __restrict__ B,
-                const scalar_t* __restrict__ BIAS, float* glbl, int* cntr,
-                scalar_t* C, const int CuCount) {
+    wvSplitKrc_(const int actlN, const int K, const int M, const int Bx,
+                const int By, const scalar_t* __restrict__ B,
+                const scalar_t* __restrict__ A,
+                const scalar_t* __restrict__ BIAS, float* glbl, scalar_t* C,
+                const int CuCount) {
+  // Use upper half of glbl buffer for atomic reduce counting
+  int* cntr = (int*)(&glbl[M * N]);
+
  constexpr int NTILE = 16;
  constexpr int APAD = 1;
  constexpr int ASTRD = 64;
@@ -1423,11 +1425,11 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
          unsigned int kOffcp = min__(K - A_CHUNK, k_str + kOff);
          for (unsigned int n = 0; n < N; n += CHUNKK * sprdN) {
            __builtin_amdgcn_global_load_lds(
-                (int*)(&A[min__(Kap * actlN - A_CHUNK,
-                                kOffcp + Kap * (n / CHUNKK +
-                                                (N / CHUNKK) * (threadIdx.x /
-                                                                (64 / CHUNKK)) +
-                                                (threadIdx.y % sprdN)))]),
+                (int*)(&A[min__(
+                    K * actlN - A_CHUNK,
+                    kOffcp + K * (n / CHUNKK +
+                                  (N / CHUNKK) * (threadIdx.x / (64 / CHUNKK)) +
+                                  (threadIdx.y % sprdN)))]),
                (int*)(&s[(k +
                           kFitPdd * ((n / CHUNKK) + (threadIdx.y % sprdN)))]),
                16, 0, 0);
@@ -1474,7 +1476,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
  #endif

    // B[] staging is cooperative across GrpsShrB, so sync here before reading
-    // back. This wait is currently inserted by compiler, but not guaranteed.
+    // back. This wait is currently inserted by compiler, but not gauranteed.
    asm volatile("s_waitcnt 0");
    __syncthreads();

@@ -1531,98 +1533,45 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
    }
  }

-  union flt4 {
-    scalar8 s8;
-    float2 f2[2];
-    float4 f4;
-  };
  if (m + (threadIdx.x % 16) < M) {
    int my_cntr;
    int mindx = m + (threadIdx.x % 16);
    int g_mindx = m * 4 + (threadIdx.x % 64);  // coalesced atomic reduction
    scalar_t biases[N / NTILE / GrpsShrB][4] = {};
    // Atomic add the output, read biases
-    for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
-      int g_nindx =
-          (nt * NTILE + (N / GrpsShrB) * (threadIdx.y % GrpsShrB)) / 4;
-      int g_adr = g_mindx * 4 + 0 + M * g_nindx * 4;
-      if (DTRMNSTC) {
-        flt4 flt4_ = {.s8 = sum4[nt][0]};
-        __hip_atomic_store((float2*)&glbl[g_adr + M * N * (m0 / Mmod)],
-                           flt4_.f2[0], __ATOMIC_RELAXED,
-                           __HIP_MEMORY_SCOPE_AGENT);
-        __hip_atomic_store((float2*)&glbl[g_adr + 2 + M * N * (m0 / Mmod)],
-                           flt4_.f2[1], __ATOMIC_RELAXED,
-                           __HIP_MEMORY_SCOPE_AGENT);
-      } else {
-        for (uint32_t j = 0; j < 4; j++)
-          atomicAdd((&glbl[g_adr + j]), sum4[nt][0][j]);
+    for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++)
+      for (uint32_t j = 0; j < 4; j++) {
+        // int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE +
+        //             (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
+        // int adr = mindx + M * nindx;
+        int g_nindx =
+            j + (nt * NTILE + (N / GrpsShrB) * (threadIdx.y % GrpsShrB)) / 4;
+        int g_adr = g_mindx + M * g_nindx * 4;
+        atomicAdd(&glbl[g_adr], sum4[nt][0][j]);
      }
-    }
-
-    __atomic_signal_fence(__ATOMIC_SEQ_CST);
-    asm volatile("s_waitcnt vmcnt(0)" ::: "memory");
-    __atomic_signal_fence(__ATOMIC_SEQ_CST);
-
    int nindx_ = (0 + (threadIdx.x / 16) * 4) + 0 * NTILE +
                 (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
    int adr_ = mindx + M * nindx_ / 4;
-    my_cntr = atomicAdd(&cntr[adr_], 1);
-
-    // make sure LDS is free for write out staging
-    if (DTRMNSTC) __syncthreads();
-
    // Update the complete counter
-    flt4 vals[N / NTILE / GrpsShrB] = {};
+    my_cntr = atomicAdd(&cntr[adr_], 1);
+    float vals[N / NTILE / GrpsShrB][4] = {};
    // If we're the last k-shard, read back the value and convert...
    if (my_cntr + 1 == k_rnd) {
-      cntr[adr_] = 0;  // clear for next round
-      if constexpr (DTRMNSTC) {
-  #pragma unroll
-        for (int ks = 0; ks < k_rnd; ks++) {
-          for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
-            int g_nindx =
-                (nt * NTILE + (N / GrpsShrB) * (threadIdx.y % GrpsShrB)) / 4;
-            int g_adr = g_mindx * 4 + 0 + M * g_nindx * 4;
-            __builtin_amdgcn_global_load_lds(
-                (float4*)(&glbl[g_adr + M * N * ks]),
-                &(((float4*)s)[(threadIdx.y * THRDS) + ks * THRDS * 4 +
-                               nt * THRDS * 4 * k_rnd]),
-                16, 0, 0);
-          }
-        }
-        if (BIAS)
-          for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
-            for (uint32_t j = 0; j < 4; j++) {
-              int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE +
-                          (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
-              biases[nt][j] = BIAS[(mindx % Bx) + (nindx % By) * Bx];
-            }
-          }
-        asm volatile("s_waitcnt 0");
-        for (int ks = 0; ks < k_rnd; ks++) {
-          for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
-            float4 eval = ((float4*)s)[(threadIdx.x + threadIdx.y * THRDS) +
-                                       ks * THRDS * 4 + nt * THRDS * 4 * k_rnd];
-            vals[nt].f4 += eval;
-          }
-        }
-      } else {
+      if (BIAS)
        for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
-          int g_nindx =
-              (nt * NTILE + (N / GrpsShrB) * (threadIdx.y % GrpsShrB)) / 4;
-          int g_adr = g_mindx * 4 + 0 + M * g_nindx * 4;
-          vals[nt].f4 = *(float4*)(&glbl[g_adr]);
-          *(float4*)(&glbl[g_adr]) = {};  // clear out for next round
-        }
-        if (BIAS)
-          for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
-            for (uint32_t j = 0; j < 4; j++) {
-              int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE +
-                          (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
-              biases[nt][j] = BIAS[(mindx % Bx) + (nindx % By) * Bx];
-            }
+          for (uint32_t j = 0; j < 4; j++) {
+            int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE +
+                        (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
+            biases[nt][j] = BIAS[(mindx % Bx) + (nindx % By) * Bx];
          }
+        }
+      for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
+        for (uint32_t j = 0; j < 4; j++) {
+          int g_nindx =
+              j + (nt * NTILE + (N / GrpsShrB) * (threadIdx.y % GrpsShrB)) / 4;
+          int g_adr = g_mindx + M * g_nindx * 4;
+          vals[nt][j] = glbl[g_adr];
+        }
      }
      __builtin_amdgcn_sched_barrier(0);
      for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
@@ -1632,11 +1581,11 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
          if (nindx < actlN) {
            int adr = mindx + M * nindx;
            if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) {
-              vals[nt].s8[j] += __bfloat162float(biases[nt][j]);
-              C[adr] = __float2bfloat16(vals[nt].s8[j]);
+              vals[nt][j] += __bfloat162float(biases[nt][j]);
+              C[adr] = __float2bfloat16(vals[nt][j]);
            } else {
-              vals[nt].s8[j] += __half2float(biases[nt][j]);
-              C[adr] = __float2half(vals[nt].s8[j]);
+              vals[nt][j] += __half2float(biases[nt][j]);
+              C[adr] = __float2half(vals[nt][j]);
            }
          }
        }
@@ -1655,25 +1604,21 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
 }
 #else   // !defined(__HIP__GFX9__) TODO: Add NAVI support
 template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
-          int UNRL, int N, int GrpsShrB, int CHUNKK, int DTRMNSTC>
-__global__ void wvSplitKrc_(const int actlN, const int K, const int Kap,
-                            const int M, const int Bx, const int By,
-                            const scalar_t* B, const scalar_t* __restrict__ A,
+          int UNRL, int N, int GrpsShrB, int CHUNKK>
+__global__ void wvSplitKrc_(const int actlN, const int K, const int M,
+                            const int Bx, const int By, const scalar_t* B,
+                            const scalar_t* __restrict__ A,
                            const scalar_t* __restrict__ BIAS, float* glbl,
-                            int* cntr, scalar_t* C,
-                            const int CuCount){UNREACHABLE_CODE}
+                            // int* cntr,
+                            scalar_t* C, const int CuCount){UNREACHABLE_CODE}
 #endif  // defined(__HIP__GFX9__) TODO: Add NAVI support

 torch::Tensor wvSplitKrc(const at::Tensor& in_a, const at::Tensor& in_b,
                         const std::optional<at::Tensor>& in_bias,
                         const int64_t CuCount) {
-  int _DTRMNSTC = 1;  // vllm::vllm_is_batch_invariant();
-
-  auto M_in = in_b.size(0);
-  auto N_in = in_a.size(0);
-  auto K_in = in_b.size(1);
-  auto Kap_in = in_a.stride(0);
-
+  auto M_in = in_a.size(0);
+  auto N_in = in_b.size(0);
+  auto K_in = in_a.size(1);
  auto Bx_in =
      (in_bias.has_value() && in_bias->numel() > 0)
          ? (in_bias->sizes().size() == 2) ? in_bias->size(1) : in_bias->size(0)
@@ -1690,9 +1635,13 @@ torch::Tensor wvSplitKrc(const at::Tensor& in_a, const at::Tensor& in_b,

  auto out_c = torch::empty(
      {N_in, M_in},
-      torch::TensorOptions().dtype(in_a.dtype()).device(in_a.device()));
+      torch::TensorOptions().dtype(in_b.dtype()).device(in_b.device()));

  auto N_p2 = 1U << (32 - __builtin_clz(N_in - 1));
+  auto axl_glbl = torch::empty(
+      {N_p2 + N_p2 / 4, M_in + M_in / 4},
+      torch::TensorOptions().dtype(torch::kFloat32).device(in_b.device()));
+  axl_glbl.zero_();  // disable for FAST_UNSAFE_RDC_INIT

  dim3 grid(CuCount);

@@ -1700,70 +1649,55 @@ torch::Tensor wvSplitKrc(const at::Tensor& in_a, const at::Tensor& in_b,
  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
  // const int max_lds_len = get_lds_size() / 2;

-  // With 64 Ms per CU (each of 4 SIMDs working on a 16x16 tile),
-  // and each working on a 512-shard of K, how many CUs would we need?
-  int rndup_cus = ((M_in + 64 - 1) / 64) * ((K_in + 512 - 1) / 512);
-
-  // How many of 4 waves in a group can work on same 16 Ms at same time? First
-  // try to maximize this. This reduces the Ms each group works on, i.e.
-  // increasing the number of CUs needed.
-  int GrpsShrB = min(N_p2 / 16, 4);
-
-  // Given the above, how many CUs would we need?
-  int CuNeeded = rndup_cus * GrpsShrB;
-
-  if (CuNeeded > CuCount) throw std::runtime_error("Invalid wvSplitKrc size");
-
-  // Can we increase SplitK by shrinking the K-shared to 256?
-  int chunkk = (CuNeeded * 2 <= CuCount) ? 2 : 1;
-
-  static torch::Tensor axl_glbl =
-      torch::zeros(
-          128 * 1024 * (_DTRMNSTC ? 12 : 1),
-          torch::TensorOptions().dtype(torch::kFloat32).device(in_a.device()))
-          .detach();
-  static torch::Tensor axl_cntr =
-      torch::zeros(
-          128 * 1024 * (_DTRMNSTC ? 12 : 1) / 4,
-          torch::TensorOptions().dtype(torch::kInt).device(in_a.device()))
-          .detach();
-  auto glbl = axl_glbl.data_ptr<float>();
-  auto cntr = axl_cntr.data_ptr<int>();
-
 #define WVSPLITKrc(_N, _GrpsShrB, _CHUNKK)                                     \
  {                                                                            \
    dim3 block(64, 4);                                                         \
-    if (_DTRMNSTC)                                                             \
-      wvSplitKrc_<fptype, 64, 16, 4, 8, 1, _N, _GrpsShrB, _CHUNKK, 1>          \
-          <<<grid, block, 0, stream>>>(N_in, K_in, Kap_in, M_in, Bx_in, By_in, \
-                                       af4, bf4, biasf4, glbl, cntr, c,        \
-                                       CuCount);                               \
-    else                                                                       \
-      wvSplitKrc_<fptype, 64, 16, 4, 8, 1, _N, _GrpsShrB, _CHUNKK, 0>          \
-          <<<grid, block, 0, stream>>>(N_in, K_in, Kap_in, M_in, Bx_in, By_in, \
-                                       af4, bf4, biasf4, glbl, cntr, c,        \
-                                       CuCount);                               \
+    wvSplitKrc_<fptype, 64, 16, 4, 8, 1, _N, _GrpsShrB, _CHUNKK>               \
+        <<<grid, block, 0, stream>>>(N_in, K_in, M_in, Bx_in, By_in, af4, bf4, \
+                                     biasf4, glbl, c, CuCount);                \
  }

-  AT_DISPATCH_REDUCED_FLOATING_TYPES(in_a.scalar_type(), "wvSplitKrc", [&] {
+  AT_DISPATCH_REDUCED_FLOATING_TYPES(in_b.scalar_type(), "wvSplitKrc", [&] {
    using fptype = typename scalar<scalar_t>::type;
-    const fptype* af4 = reinterpret_cast<const fptype*>(in_a.data_ptr());
+    fptype* af4 = reinterpret_cast<fptype*>(in_a.data_ptr());
    const fptype* bf4 = reinterpret_cast<const fptype*>(in_b.data_ptr());
    const fptype* biasf4 =
        (in_bias.has_value() && in_bias->numel() > 0)
            ? reinterpret_cast<const fptype*>(in_bias->data_ptr())
            : nullptr;
    fptype* c = reinterpret_cast<fptype*>(out_c.data_ptr());
+    auto glbl = axl_glbl.data_ptr<float>();
+
+    // With 64 Ms per CU (each of 4 SIMDs working on a 16x16 tile),
+    // and each working on a 512-shard of K, how many CUs would we need?
+    int rndup_cus = ((M_in + 64 - 1) / 64) * ((K_in + 512 - 1) / 512);
+
+    // How many of 4 waves in a group can work on same 16 Ms at same time? First
+    // try to maximize this. This reduces the Ms each group works on, i.e.
+    // increasing the number of CUs needed.
+    int GrpsShrB = min(N_p2 / 16, 4);
+
+    // Given the above, how many CUs would we need?
+    int CuNeeded = rndup_cus * GrpsShrB;
+
+    if (CuNeeded > CuCount) std::runtime_error("Invalid wvSplitKrc size");
+
+    // Can we increase SplitK by shrinking the K-shared to 256?
+    int chunkk = (CuNeeded * 2 <= CuCount) ? 2 : 1;

    switch (N_p2) {
      case 16:
        WVSPLITKrc(16, 1, 1) break;
      case 32:
-        if (chunkk == 2) WVSPLITKrc(32, 2, 2) else WVSPLITKrc(32, 2, 1) break;
+        if (chunkk == 2)
+          WVSPLITKrc(32, 2, 2) else if (chunkk == 1) WVSPLITKrc(32, 2, 1) break;
      case 64:
-        if (chunkk == 2) WVSPLITKrc(64, 4, 2) else WVSPLITKrc(64, 4, 1) break;
+        if (chunkk == 2)
+          WVSPLITKrc(64, 4, 2) else if (chunkk == 1) WVSPLITKrc(64, 4, 1) break;
      case 128:
-        if (chunkk == 2) WVSPLITKrc(128, 4, 2) else WVSPLITKrc(128, 4, 1) break;
+        if (chunkk == 2)
+          WVSPLITKrc(128, 4, 2) else if (chunkk == 1)
+              WVSPLITKrc(128, 4, 1) break;
      default:
        throw std::runtime_error(
            "Unsupported N value: " + std::to_string(M_in) + "," +
@@ -426,22 +426,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
      " Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()");
  // conditionally compiled so impl registration is in source file

-  // Expert-specialization mxfp8 blockscaled grouped quantization (SM100+).
-  ops.def(
-      "mxfp8_experts_quant("
-      " Tensor input, Tensor problem_sizes, Tensor expert_offsets,"
-      " Tensor blockscale_offsets, Tensor! quant_output, Tensor! scale_factor)"
-      " -> ()");
-  // conditionally compiled so impl registration is in source file
-
-  // Expert-specialization mxfp8 blockscaled grouped GEMM (SM100+).
-  ops.def(
-      "cutlass_mxfp8_grouped_mm("
-      " Tensor a, Tensor b, Tensor sfa, Tensor sfb, Tensor! out,"
-      " Tensor problem_sizes, Tensor expert_offsets, Tensor blockscale_offsets)"
-      " -> ()");
-  // conditionally compiled so impl registration is in source file
-
  // CUTLASS w8a8 GEMM, supporting symmetric per-tensor or per-row/column
  // quantization, as well as bias
  ops.def(
@@ -802,10 +786,6 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cache_ops), cache_ops) {
  cache_ops.impl("indexer_k_quant_and_cache", torch::kCUDA,
                 &indexer_k_quant_and_cache);

-  cache_ops.def(
-      "concat_mla_q(Tensor ql_nope, Tensor q_pe, Tensor! q_out) -> ()");
-  cache_ops.impl("concat_mla_q", torch::kCUDA, &concat_mla_q);
-
  cache_ops.def(
      "cp_gather_indexer_k_quant_cache(Tensor kv_cache, Tensor! dst_k, Tensor! "
      "dst_scale, Tensor block_table, Tensor cu_seq_lens) -> ()");
@@ -262,9 +262,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \

 # Build the vLLM wheel
 # if USE_SCCACHE is set, use sccache to speed up compilation
-# AWS credentials mounted at ~/.aws/credentials for sccache S3 auth (optional)
 RUN --mount=type=cache,target=/root/.cache/uv \
-    --mount=type=secret,id=aws-credentials,target=/root/.aws/credentials,required=false \
    if [ "$USE_SCCACHE" = "1" ]; then \
        echo "Installing sccache..." \
        && case "${TARGETPLATFORM}" in \
@@ -115,57 +115,9 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 # install development dependencies (for testing)
 RUN uv pip install -e tests/vllm_test_utils

-# install NIXL and UCX from source code
-ARG UCX_VERSION=e5d98879705239d254ede40b4a52891850cb5349
-ARG NIXL_VERSION=0.7.0
-
-RUN apt-get update && apt-get install -y \
-    pciutils \
-    net-tools \
-    iproute2 \
-    hwloc \
-    numactl \
-    wget \
-    curl \
-    git \
-    build-essential \
-    autoconf \
-    automake \
-    libtool \
-    pkg-config \
-    rdma-core \
-    libibverbs-dev \
-    ibverbs-utils \
-    libibverbs1 \
-    librdmacm-dev \
-    librdmacm1 \
-    libibumad-dev \
-    libibumad3 \
-    libibmad-dev \
-    libibmad5 \
-    infiniband-diags \
-    perftest \
-    ibutils \
-    libmlx5-1 \
-    libmlx4-1 \
-    ibverbs-providers \
-    librdmacm1t64
-
-ENV PKG_CONFIG_PATH=/tmp/ucx_install/lib/pkgconfig:${PKG_CONFIG_PATH}
-ENV LD_LIBRARY_PATH=/tmp/ucx_install/lib:${LD_LIBRARY_PATH}
-RUN --mount=type=cache,target=/root/.cache/uv \
-    git clone https://github.com/openucx/ucx /tmp/ucx_source && \
-    cd /tmp/ucx_source && git checkout "${UCX_VERSION}" && \
-    bash autogen.sh && \
-    ./configure --prefix=/tmp/ucx_install --with-ze=yes --enable-examples --enable-mt && \
-    make CFLAGS="-Wno-error=incompatible-pointer-types" -j8 && make install && \
-    git clone https://github.com/ai-dynamo/nixl /tmp/nixl_source && \
-    cd /tmp/nixl_source && git checkout "${NIXL_VERSION}" && \
-    cd /tmp/nixl_source && \
-    uv pip install --upgrade meson pybind11 patchelf && \
-    uv pip install -r requirements.txt && \
-    uv pip install . && \
-    rm -rf /tmp/ucx_source /tmp/nixl_source
+# install nixl from source code
+ENV NIXL_VERSION=0.7.0
+RUN python /workspace/vllm/tools/install_nixl_from_source_ubuntu.py

 # FIX triton
 RUN --mount=type=cache,target=/root/.cache/uv \
@@ -18,14 +18,14 @@ th {
 </style>

 | Dataset | Online | Offline | Data Path |
-| ------- | ------ | ------- | --------- |
+|---------|--------|---------|-----------|
 | ShareGPT | ✅ | ✅ | `wget https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json` |
 | ShareGPT4V (Image) | ✅ | ✅ | `wget https://huggingface.co/datasets/Lin-Chen/ShareGPT4V/resolve/main/sharegpt4v_instruct_gpt4-vision_cap100k.json`<br>Note that the images need to be downloaded separately. For example, to download COCO's 2017 Train images:<br>`wget http://images.cocodataset.org/zips/train2017.zip` |
 | ShareGPT4Video (Video) | ✅ | ✅ | `git clone https://huggingface.co/datasets/ShareGPT4Video/ShareGPT4Video` |
 | BurstGPT | ✅ | ✅ | `wget https://github.com/HPMLL/BurstGPT/releases/download/v1.1/BurstGPT_without_fails_2.csv` |
 | Sonnet (deprecated) | ✅ | ✅ | Local file: `benchmarks/sonnet.txt` |
 | Random | ✅ | ✅ | `synthetic` |
-| RandomMultiModal (Image/Video) | ✅ | ✅ | `synthetic` |
+| RandomMultiModal (Image/Video) | 🟡 | 🚧 | `synthetic` |
 | RandomForReranking | ✅ | ✅ | `synthetic` |
 | Prefix Repetition | ✅ | ✅ | `synthetic` |
 | HuggingFace-VisionArena | ✅ | ✅ | `lmarena-ai/VisionArena-Chat` |
@@ -383,14 +383,14 @@ The `--burstiness` parameter mathematically controls request arrival patterns us

 Load Pattern Recommendations by Use Case:

-| Use Case           | Burstiness   | Request Rate    | Max Concurrency | Description                                                                        |
-| ---                | ---          | ---             | ---             | ---                                                                                |
+| Use Case           | Burstiness   | Request Rate    | Max Concurrency | Description                                               |
+| ---                | ---          | ---             | ---             | ---                                                       |
 | Maximum Throughput | N/A          | Infinite        | Limited         | **Most common**: Simulates load balancer/gateway limits with unlimited user demand |
-| Realistic Testing  | 1.0          | Moderate (5-20) | Infinite        | Natural Poisson traffic patterns for baseline performance                          |
-| Stress Testing     | 0.1-0.5      | High (20-100)   | Infinite        | Challenging burst patterns to test resilience                                      |
-| Latency Profiling  | 2.0-5.0      | Low (1-10)      | Infinite        | Uniform load for consistent timing analysis                                        |
-| Capacity Planning  | 1.0          | Variable        | Limited         | Test resource limits with realistic constraints                                    |
-| SLA Validation     | 1.0          | Target rate     | SLA limit       | Production-like constraints for compliance testing                                 |
+| Realistic Testing  | 1.0          | Moderate (5-20) | Infinite        | Natural Poisson traffic patterns for baseline performance |
+| Stress Testing     | 0.1-0.5      | High (20-100)   | Infinite        | Challenging burst patterns to test resilience             |
+| Latency Profiling  | 2.0-5.0      | Low (1-10)      | Infinite        | Uniform load for consistent timing analysis               |
+| Capacity Planning  | 1.0          | Variable        | Limited         | Test resource limits with realistic constraints           |
+| SLA Validation     | 1.0          | Target rate     | SLA limit       | Production-like constraints for compliance testing        |

 These load patterns help evaluate different aspects of your vLLM deployment, from basic performance characteristics to resilience under challenging traffic conditions.

@@ -545,24 +545,6 @@ vllm bench throughput \
  --lora-path yard1/llama-2-7b-sql-lora-test
 ```

-#### Synthetic Random Multimodal (random-mm)
-
-Generate synthetic multimodal inputs for offline throughput testing without external datasets.
-Use `--backend vllm-chat` so that image tokens are counted correctly.
-
-```bash
-vllm bench throughput \
-  --model Qwen/Qwen2-VL-7B-Instruct \
-  --backend vllm-chat \
-  --dataset-name random-mm \
-  --num-prompts 100 \
-  --random-input-len 300 \
-  --random-output-len 40 \
-  --random-mm-base-items-per-request 2 \
-  --random-mm-limit-mm-per-prompt '{"image": 3, "video": 0}' \
-  --random-mm-bucket-config '{(256, 256, 1): 0.7, (720, 1280, 1): 0.3}'
-```
-
 </details>

 ### 🛠️ Structured Output Benchmark
@@ -864,8 +846,8 @@ Generate synthetic image inputs alongside random text prompts to stress-test vis

 Notes:

- For online benchmarks, use `--backend openai-chat` with endpoint `/v1/chat/completions`.
- For offline benchmarks, use `--backend vllm-chat` (see [Offline Throughput Benchmark](#-offline-throughput-benchmark) for an example).
+- Works only with online benchmark via the OpenAI backend (`--backend openai-chat`) and endpoint `/v1/chat/completions`.
+- Video sampling is not yet implemented.

 Start the server (example):

@@ -931,74 +913,6 @@ This should be seen as an edge case, and if this behavior can be avoided by sett

 </details>

-### 🔬 Multimodal Processor Benchmark
-
-Benchmark per-stage latency of the multimodal (MM) input processor pipeline, including the encoder forward pass. This is useful for profiling preprocessing bottlenecks in vision-language models.
-
-<details class="admonition abstract" markdown="1">
-<summary>Show more</summary>
-
-The benchmark measures the following stages for each request:
-
-| Stage | Description |
-| ----- | ----------- |
-| `get_mm_hashes_secs` | Time spent hashing multimodal inputs |
-| `get_cache_missing_items_secs` | Time spent looking up the processor cache |
-| `apply_hf_processor_secs` | Time spent in the HuggingFace processor |
-| `merge_mm_kwargs_secs` | Time spent merging multimodal kwargs |
-| `apply_prompt_updates_secs` | Time spent updating prompt tokens |
-| `preprocessor_total_secs` | Total preprocessing time |
-| `encoder_forward_secs` | Time spent in the encoder model forward pass |
-| `num_encoder_calls` | Number of encoder invocations per request |
-
-The benchmark also reports end-to-end latency (TTFT + decode time) per
-request. Use `--metric-percentiles` to select which percentiles to report
-(default: p99) and `--output-json` to save results.
-
-#### Basic Example with Synthetic Data (random-mm)
-
-```bash
-vllm bench mm-processor \
-  --model Qwen/Qwen2-VL-7B-Instruct \
-  --dataset-name random-mm \
-  --num-prompts 50 \
-  --random-input-len 300 \
-  --random-output-len 40 \
-  --random-mm-base-items-per-request 2 \
-  --random-mm-limit-mm-per-prompt '{"image": 3, "video": 0}' \
-  --random-mm-bucket-config '{(256, 256, 1): 0.7, (720, 1280, 1): 0.3}'
-```
-
-#### Using a HuggingFace Dataset
-
-```bash
-vllm bench mm-processor \
-  --model Qwen/Qwen2-VL-7B-Instruct \
-  --dataset-name hf \
-  --dataset-path lmarena-ai/VisionArena-Chat \
-  --hf-split train \
-  --num-prompts 100
-```
-
-#### Warmup, Custom Percentiles, and JSON Output
-
-```bash
-vllm bench mm-processor \
-  --model Qwen/Qwen2-VL-7B-Instruct \
-  --dataset-name random-mm \
-  --num-prompts 200 \
-  --num-warmups 5 \
-  --random-input-len 300 \
-  --random-output-len 40 \
-  --random-mm-base-items-per-request 1 \
-  --metric-percentiles 50,90,95,99 \
-  --output-json results.json
-```
-
-See [`vllm bench mm-processor`](../cli/bench/mm_processor.md) for the full argument reference.
-
-</details>
-
 ### Embedding Benchmark

 Benchmark the performance of embedding requests in vLLM.
@@ -60,12 +60,12 @@ Here is an example using the script to compare result_a and result_b with max co

 ***Output Tput (tok/s) — Model : [ meta-llama/Llama-3.1-8B-Instruct ] , Dataset Name : [ random ] , Input Len : [ 2048.0 ] , Output Len : [ 2048.0 ]***

-| | # of max concurrency | qps | results_a/benchmark_results.json | results_b/benchmark_results.json | perf_ratio |
-| | -------------------- | --- | -------------------------------- | -------------------------------- | ---------- |
-| 0 | 12 | inf | 24.98 | 186.03 |  7.45 |
-| 1 | 16 | inf |  25.49 | 246.92 | 9.69 |
-| 2 | 24 | inf | 27.74 | 293.34 |  10.57 |
-| 3 | 32 | inf | 28.61 |306.69 | 10.72 |
+|    | # of max concurrency | qps  | results_a/benchmark_results.json | results_b/benchmark_results.json | perf_ratio        |
+|----|------|-----|-----------|----------|----------|
+| 0  | 12 | inf | 24.98   | 186.03 |  7.45 |
+| 1  | 16 | inf|  25.49  | 246.92 | 9.69 |
+| 2  | 24 | inf| 27.74  | 293.34 |  10.57 |
+| 3  | 32 | inf| 28.61  |306.69 | 10.72 |

 ***compare-json-results.py – Command-Line Parameters***  

@@ -1,51 +1,5 @@
 # vllm bench mm-processor

-## Overview
-
-`vllm bench mm-processor` profiles the multimodal input processor pipeline of
-vision-language models. It measures per-stage latency from the HuggingFace
-processor through to the encoder forward pass, helping you identify
-preprocessing bottlenecks and understand how different image resolutions or
-item counts affect end-to-end request time.
-
-The benchmark supports two data sources: synthetic random multimodal inputs
-(`random-mm`) and HuggingFace datasets (`hf`). Warmup requests are run before
-measurement to ensure stable results.
-
-## Quick Start
-
-```bash
-vllm bench mm-processor \
-  --model Qwen/Qwen2-VL-7B-Instruct \
-  --dataset-name random-mm \
-  --num-prompts 50 \
-  --random-input-len 300 \
-  --random-output-len 40 \
-  --random-mm-base-items-per-request 2 \
-  --random-mm-limit-mm-per-prompt '{"image": 3, "video": 0}' \
-  --random-mm-bucket-config '{(256, 256, 1): 0.7, (720, 1280, 1): 0.3}'
-```
-
-## Measured Stages
-
-| Stage | Description |
-| ----- | ----------- |
-| `get_mm_hashes_secs` | Time spent hashing multimodal inputs |
-| `get_cache_missing_items_secs` | Time spent looking up the processor cache |
-| `apply_hf_processor_secs` | Time spent in the HuggingFace processor |
-| `merge_mm_kwargs_secs` | Time spent merging multimodal kwargs |
-| `apply_prompt_updates_secs` | Time spent updating prompt tokens |
-| `preprocessor_total_secs` | Total preprocessing time |
-| `encoder_forward_secs` | Time spent in the encoder model forward pass |
-| `num_encoder_calls` | Number of encoder invocations per request |
-
-The benchmark also reports end-to-end latency (TTFT + decode time) per
-request. Use `--metric-percentiles` to select which percentiles to report
-(default: p99) and `--output-json` to save results.
-
-For more examples (HF datasets, warmup, JSON output), see
-[Benchmarking CLI — Multimodal Processor Benchmark](../../benchmarking/cli.md#multimodal-processor-benchmark).
-
 ## JSON CLI Arguments

 --8<-- "docs/cli/json_tip.inc.md"
@@ -1,4 +1,3 @@
-<!-- markdownlint-disable MD041 -->
 When passing JSON CLI arguments, the following sets of arguments are equivalent:

 - `--json-arg '{"key1": "value1", "key2": {"key3": "value2"}}'`
@@ -7,4 +6,4 @@ When passing JSON CLI arguments, the following sets of arguments are equivalent:
 Additionally, list elements can be passed individually using `+`:

 - `--json-arg '{"key4": ["value3", "value4", "value5"]}'`
- `--json-arg.key4+ value3 --json-arg.key4+='value4,value5'`
+- `--json-arg.key4+ value3 --json-arg.key4+='value4,value5'`
@@ -5,17 +5,6 @@ This guide covers optimization strategies and performance tuning for vLLM V1.
 !!! tip
    Running out of memory? Consult [this guide](./conserving_memory.md) on how to conserve memory.

-## Optimization Levels
-
-vLLM provides 4 optimization levels (`-O0`, `-O1`, `-O2`, `-O3`) that allow users to trade off startup time for performance:
-
- `-O0`: No optimizations. Fastest startup time, but lowest performance.
- `-O1`: Fast optimization. Simple compilation and fast fusions, and PIECEWISE cudagraphs.
- `-O2`: Default optimization. Additional compilation ranges, additional fusions, FULL_AND_PIECEWISE cudagraphs.
- `-O3`: Aggressive optimization. Currently equal to `-O2`, but may include additional time-consuming or experimental optimizations in the future.
-
-For more information, see the [optimization level documentation](../design/optimization_levels.md).
-
 ## Preemption

 Due to the autoregressive nature of transformer architecture, there are times when KV cache space is insufficient to handle all batched requests.
@@ -293,7 +282,7 @@ llm = LLM(
 Based on the configuration, the content of the multi-modal caches on `P0` and `P1` are as follows:

 | mm_processor_cache_type | Cache Type | `P0` Cache | `P1` Engine Cache | `P1` Worker Cache | Max. Memory |
-| ----------------- | ----------- | ---------- | ---------- | ----------- | ----------- |
+|-------------------|-------------|------------|------------|-------------|-------------|
 | lru | Processor Caching | K + V | N/A | N/A | `mm_processor_cache_gb * data_parallel_size` |
 | lru | Key-Replicated Caching | K | K + V | N/A | `mm_processor_cache_gb * api_server_count` |
 | shm | Shared Memory Caching | K | N/A | V | `mm_processor_cache_gb * api_server_count` |
--- a/Show More
+++ b/Show More