[ci] fix wheel names for arm wheels (#24898)

Signed-off-by: simon-mo <simon.mo@hey.com>

.buildkite/nightly-benchmarks/nightly-descriptions.md

@@ -8,7 +8,7 @@ This benchmark aims to:
 
 Latest results: [results link](https://blog.vllm.ai/2024/09/05/perf-update.html), scroll to the end.
 
-Latest reproduction guide: [github issue link](https://github.com/vllm-project/vllm/issues/8176)
+Latest reproduction guilde: [github issue link](https://github.com/vllm-project/vllm/issues/8176)
 
 ## Setup
 

.buildkite/scripts/hardware_ci/run-amd-test.sh

@@ -86,6 +86,10 @@ if [[ $commands == *"pytest -v -s models/test_registry.py"* ]]; then
   commands=${commands//"pytest -v -s models/test_registry.py"/"pytest -v -s models/test_registry.py -k 'not BambaForCausalLM and not GritLM and not Mamba2ForCausalLM and not Zamba2ForCausalLM'"}
 fi
 
+if [[ $commands == *"VLLM_USE_V1=0 pytest -v -s models/test_initialization.py -k 'not llama4 and not plamo2'"* ]]; then
+  commands=${commands//"VLLM_USE_V1=0 pytest -v -s models/test_initialization.py -k 'not llama4 and not plamo2'"/"VLLM_USE_V1=0 pytest -v -s models/test_initialization.py -k 'not llama4 and not plamo2 and not BambaForCausalLM and not Gemma2ForCausalLM and not Grok1ModelForCausalLM and not Zamba2ForCausalLM and not Gemma2Model and not GritLM'"}
+fi
+
 if [[ $commands == *"pytest -v -s compile/test_basic_correctness.py"* ]]; then
   commands=${commands//"pytest -v -s compile/test_basic_correctness.py"/"VLLM_USE_TRITON_FLASH_ATTN=0 pytest -v -s compile/test_basic_correctness.py"}
 fi
@@ -163,6 +167,12 @@ if [[ $commands == *" entrypoints/llm "* ]]; then
   --ignore=entrypoints/llm/test_prompt_validation.py "}
 fi
 
+#Obsolete currently
+##ignore certain Entrypoints/llm tests
+#if [[ $commands == *" && pytest -v -s entrypoints/llm/test_guided_generate.py"* ]]; then
+#  commands=${commands//" && pytest -v -s entrypoints/llm/test_guided_generate.py"/" "}
+#fi
+
 # --ignore=entrypoints/openai/test_encoder_decoder.py \
 # --ignore=entrypoints/openai/test_embedding.py \
 # --ignore=entrypoints/openai/test_oot_registration.py

.buildkite/scripts/hardware_ci/run-cpu-test.sh

@@ -58,11 +58,15 @@ function cpu_tests() {
     # pytest -x -v -s tests/kernels/attention/test_cache.py -m cpu_model
     # pytest -x -v -s tests/kernels/attention/test_mla_decode_cpu.py -m cpu_model
 
-    pytest -x -v -s tests/models/language/generation -m cpu_model
-    VLLM_CPU_SGL_KERNEL=1 pytest -x -v -s tests/models/language/generation -m cpu_model
+    # Note: disable Bart until supports V1
+    pytest -x -v -s tests/models/language/generation -m cpu_model \
+                --ignore=tests/models/language/generation/test_bart.py
+    VLLM_CPU_SGL_KERNEL=1 pytest -x -v -s tests/models/language/generation -m cpu_model \
+                --ignore=tests/models/language/generation/test_bart.py
 
     pytest -x -v -s tests/models/language/pooling -m cpu_model
     pytest -x -v -s tests/models/multimodal/generation \
+                --ignore=tests/models/multimodal/generation/test_mllama.py \
                 --ignore=tests/models/multimodal/generation/test_pixtral.py \
                 -m cpu_model"
 

.buildkite/scripts/hardware_ci/run-tpu-v1-test-part2.sh

@@ -62,7 +62,7 @@ echo "--- Installing Python dependencies ---"
 python3 -m pip install --progress-bar off git+https://github.com/thuml/depyf.git \
     && python3 -m pip install --progress-bar off pytest pytest-asyncio tpu-info \
     && python3 -m pip install --progress-bar off "lm-eval @ git+https://github.com/EleutherAI/lm-evaluation-harness.git@206b7722158f58c35b7ffcd53b035fdbdda5126d" \
-    && python3 -m pip install --progress-bar off hf-transfer tblib==3.1.0
+    && python3 -m pip install --progress-bar off hf-transfer
 echo "--- Python dependencies installed ---"
 export VLLM_USE_V1=1
 export VLLM_XLA_CHECK_RECOMPILATION=1

.buildkite/scripts/hardware_ci/run-tpu-v1-test.sh

@@ -62,7 +62,7 @@ echo "--- Installing Python dependencies ---"
 python3 -m pip install --progress-bar off git+https://github.com/thuml/depyf.git \
     && python3 -m pip install --progress-bar off pytest pytest-asyncio tpu-info \
     && python3 -m pip install --progress-bar off "lm-eval @ git+https://github.com/EleutherAI/lm-evaluation-harness.git@206b7722158f58c35b7ffcd53b035fdbdda5126d" \
-    && python3 -m pip install --progress-bar off hf-transfer tblib==3.1.0
+    && python3 -m pip install --progress-bar off hf-transfer
 echo "--- Python dependencies installed ---"
 export VLLM_USE_V1=1
 export VLLM_XLA_CHECK_RECOMPILATION=1

.buildkite/scripts/hardware_ci/run-xpu-test.sh

@@ -35,7 +35,7 @@ docker run \
     python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 -O3 -O.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
-    VLLM_ATTENTION_BACKEND=TRITON_ATTN python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager
+    VLLM_ATTENTION_BACKEND=TRITON_ATTN_VLLM_V1 python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager
     cd tests
     pytest -v -s v1/core
     pytest -v -s v1/engine

.buildkite/scripts/run-prime-rl-test.sh

@@ -1,59 +0,0 @@
-#!/bin/bash
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-# Setup script for Prime-RL integration tests
-# This script prepares the environment for running Prime-RL tests with nightly vLLM
-
-set -euo pipefail
-
-SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
-REPO_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)"
-PRIME_RL_REPO="https://github.com/PrimeIntellect-ai/prime-rl.git"
-PRIME_RL_DIR="${REPO_ROOT}/prime-rl"
-
-echo "Setting up Prime-RL integration test environment..."
-
-# Clean up any existing Prime-RL directory
-if [ -d "${PRIME_RL_DIR}" ]; then
-    echo "Removing existing Prime-RL directory..."
-    rm -rf "${PRIME_RL_DIR}"
-fi
-
-# Install UV if not available
-if ! command -v uv &> /dev/null; then
-    echo "Installing UV package manager..."
-    curl -LsSf https://astral.sh/uv/install.sh | sh
-    source $HOME/.local/bin/env
-fi
-
-# Clone Prime-RL repository at specific branch for reproducible tests
-PRIME_RL_BRANCH="integ-vllm-main"
-echo "Cloning Prime-RL repository at branch: ${PRIME_RL_BRANCH}..."
-git clone --branch "${PRIME_RL_BRANCH}" --single-branch "${PRIME_RL_REPO}" "${PRIME_RL_DIR}"
-cd "${PRIME_RL_DIR}"
-
-echo "Setting up UV project environment..."
-export UV_PROJECT_ENVIRONMENT=/usr/local
-ln -s /usr/bin/python3 /usr/local/bin/python
-
-# Remove vllm pin from pyproject.toml
-echo "Removing vllm pin from pyproject.toml..."
-sed -i '/vllm==/d' pyproject.toml
-
-# Sync Prime-RL dependencies
-echo "Installing Prime-RL dependencies..."
-uv sync --inexact && uv sync --inexact --all-extras
-
-# Verify installation
-echo "Verifying installations..."
-uv run python -c "import vllm; print(f'vLLM version: {vllm.__version__}')"
-uv run python -c "import prime_rl; print('Prime-RL imported successfully')"
-
-echo "Prime-RL integration test environment setup complete!"
-
-echo "Running Prime-RL integration tests..."
-export WANDB_MODE=offline # this makes this test not require a WANDB_API_KEY
-uv run pytest -vs tests/integration/test_rl.py -m gpu
-
-echo "Prime-RL integration tests completed!"

.buildkite/test-pipeline.yaml

@@ -6,28 +6,24 @@
 # to generate the final pipeline yaml file.
 
 # Documentation
-# label(str): the name of the test. emojis allowed.
-# fast_check(bool): whether to run this on each commit on the fastcheck pipeline.
-# torch_nightly(bool): whether to run this on vllm against the torch nightly pipeline.
-# fast_check_only(bool): run this test on the fastcheck pipeline only
-# optional(bool): never run this test by default (i.e. need to unblock manually) unless it's a scheduled nightly run.
-# soft_fail(bool): allow this step to fail without failing the entire pipeline (useful for flaky or experimental tests).
+# label(str): the name of the test. emoji allowed.
+# fast_check(bool): whether to run this on each commit on fastcheck pipeline.
+# torch_nightly(bool): whether to run this on vllm against torch nightly pipeline.
+# fast_check_only(bool): run this test on fastcheck pipeline only
+# optional(bool): never run this test by default (i.e. need to unblock manually) unless it's scheduled nightly run.
 # command(str): the single command to run for tests. incompatible with commands.
-# commands(list): the list of commands to run for the test. incompatible with command.
-# mirror_hardwares(list): the list of hardware to run the test on as well. currently only supports [amdexperimental]
-# gpu(str): override the GPU selection for the test. default is L4 GPUs. supports a100, b200, h200
-# num_gpus(int): override the number of GPUs for the test. defaults to 1 GPU. currently supports 2,4.
-# num_nodes(int): whether to simulate multi-node setup by launching multiple containers on one host,
-#     in this case, commands must be specified. the first command runs on the first host, the second
+# commands(list): the list of commands to run for test. incompatbile with command.
+# mirror_hardwares(list): the list of hardwares to run the test on as well. currently only supports [amd]
+# gpu(str): override the GPU selection for the test. default is on L4 GPUs. currently only supports a100
+# num_gpus(int): override the number of GPUs for the test. default to 1 GPU. currently support 2,4.
+# num_nodes(int): whether to simulate multi-node setup by launch multiple containers on one host,
+#     in this case, commands must be specified. the first command runs on first host, the second
 #     command runs on the second host.
-# timeout_in_minutes(int): sets a timeout for the step in minutes. if not specified, uses the default timeout.
-# parallelism(int): number of parallel jobs to run for this step. enables test sharding using $$BUILDKITE_PARALLEL_JOB
-#     and $$BUILDKITE_PARALLEL_JOB_COUNT environment variables.
-# working_dir(str): specify the place where the command should execute, default to /vllm-workspace/tests
-# source_file_dependencies(list): the list of prefixes to opt-in the test for, if empty, the test will always run.
+# working_dir(str): specify the place where command should execute, default to /vllm-workspace/tests
+# source_file_dependencies(list): the list of prefix to opt-in the test for, if empty, the test will always run.
 
 # When adding a test
-# - If the test belongs to an existing group, add it there
+# - If the test belong to an existing group, add it there
 # - If the test is short, add to any existing step
 # - If the test takes more than 10min, then it is okay to create a new step.
 #   Note that all steps execute in parallel.
@@ -50,18 +46,24 @@ steps:
   mirror_hardwares: [amdexperimental]
   source_file_dependencies:
   - vllm/
+  - tests/mq_llm_engine
+  - tests/async_engine
   - tests/test_inputs.py
   - tests/test_outputs.py
   - tests/multimodal
   - tests/utils_
+  - tests/worker
   - tests/standalone_tests/lazy_imports.py
   - tests/transformers_utils
   commands:
   - python3 standalone_tests/lazy_imports.py
+  - pytest -v -s mq_llm_engine # MQLLMEngine
+  - pytest -v -s async_engine # AsyncLLMEngine
   - pytest -v -s test_inputs.py
   - pytest -v -s test_outputs.py
   - pytest -v -s multimodal
   - pytest -v -s utils_ # Utils
+  - pytest -v -s worker # Worker
   - pytest -v -s transformers_utils # transformers_utils
 
 - label: Python-only Installation Test # 10min
@@ -82,12 +84,25 @@ steps:
   - vllm/
   - tests/basic_correctness/test_basic_correctness
   - tests/basic_correctness/test_cpu_offload
+  - tests/basic_correctness/test_preemption
   - tests/basic_correctness/test_cumem.py
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
   - pytest -v -s basic_correctness/test_cumem.py
   - pytest -v -s basic_correctness/test_basic_correctness.py
   - pytest -v -s basic_correctness/test_cpu_offload.py
+  - VLLM_TEST_ENABLE_ARTIFICIAL_PREEMPT=1 pytest -v -s basic_correctness/test_preemption.py
+
+- label: Core Test # 22min
+  timeout_in_minutes: 35
+  mirror_hardwares: [amdexperimental]
+  fast_check: true
+  source_file_dependencies:
+  - vllm/core
+  - vllm/distributed
+  - tests/core
+  commands:
+  - pytest -v -s core
 
 - label: Entrypoints Unit Tests # 5min
   timeout_in_minutes: 10
@@ -112,9 +127,10 @@ steps:
   - tests/entrypoints/offline_mode
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/llm --ignore=entrypoints/llm/test_generate.py --ignore=entrypoints/llm/test_collective_rpc.py
+  - pytest -v -s entrypoints/llm --ignore=entrypoints/llm/test_lazy_outlines.py --ignore=entrypoints/llm/test_generate.py --ignore=entrypoints/llm/test_collective_rpc.py
+  - pytest -v -s entrypoints/llm/test_lazy_outlines.py # it needs a clean process
   - pytest -v -s entrypoints/llm/test_generate.py # it needs a clean process
-  - pytest -v -s entrypoints/offline_mode # Needs to avoid interference with other tests
+  - VLLM_USE_V1=0 pytest -v -s entrypoints/offline_mode # Needs to avoid interference with other tests
 
 - label: Entrypoints Integration Test (API Server) # 100min
   timeout_in_minutes: 130
@@ -152,6 +168,7 @@ steps:
   num_gpus: 4
   source_file_dependencies:
   - vllm/distributed/
+  - vllm/core/
   - tests/distributed/test_utils
   - tests/distributed/test_pynccl
   - tests/distributed/test_events
@@ -164,20 +181,12 @@ steps:
   - tests/v1/test_internal_lb_dp.py
   - tests/v1/test_hybrid_lb_dp.py
   - tests/v1/engine/test_engine_core_client.py
-  - tests/distributed/test_symm_mem_allreduce.py
   commands:
-  # test with torchrun tp=2 and external_dp=2
+  # test with tp=2 and external_dp=2
+  - VLLM_USE_V1=0 torchrun --nproc-per-node=4 distributed/test_torchrun_example.py
   - torchrun --nproc-per-node=4 distributed/test_torchrun_example.py
-  # test with torchrun tp=2 and pp=2
+  # test with tp=2 and pp=2
   - PP_SIZE=2 torchrun --nproc-per-node=4 distributed/test_torchrun_example.py
-  # test with torchrun tp=4 and dp=1
-  - TP_SIZE=4 torchrun --nproc-per-node=4 distributed/test_torchrun_example_moe.py
-  # test with torchrun tp=2, pp=2 and dp=1
-  - PP_SIZE=2 TP_SIZE=2 torchrun --nproc-per-node=4 distributed/test_torchrun_example_moe.py
-  # test with torchrun tp=1 and dp=4 with ep
-  - DP_SIZE=4 ENABLE_EP=1 torchrun --nproc-per-node=4 distributed/test_torchrun_example_moe.py
-  # test with torchrun tp=2 and dp=2 with ep
-  - TP_SIZE=2 DP_SIZE=2 ENABLE_EP=1 torchrun --nproc-per-node=4 distributed/test_torchrun_example_moe.py
   # test with internal dp
   - python3 ../examples/offline_inference/data_parallel.py --enforce-eager
   - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/test_async_llm_dp.py
@@ -189,7 +198,6 @@ steps:
   - pytest -v -s compile/test_basic_correctness.py
   - pytest -v -s distributed/test_pynccl.py
   - pytest -v -s distributed/test_events.py
-  - pytest -v -s distributed/test_symm_mem_allreduce.py
   # TODO: create a dedicated test section for multi-GPU example tests
   # when we have multiple distributed example tests
   - pushd ../examples/offline_inference
@@ -222,14 +230,16 @@ steps:
   num_gpus: 2
   source_file_dependencies:
   - vllm/
+  - tests/metrics
   - tests/v1/tracing
   commands:
+  - pytest -v -s metrics
   - "pip install \
       'opentelemetry-sdk>=1.26.0' \
       'opentelemetry-api>=1.26.0' \
       'opentelemetry-exporter-otlp>=1.26.0' \
       'opentelemetry-semantic-conventions-ai>=0.4.1'"
-  - pytest -v -s v1/tracing
+  - pytest -v -s tracing
 
 ##### fast check tests  #####
 #####  1 GPU test  #####
@@ -292,7 +302,6 @@ steps:
     # split the test to avoid interference
     - pytest -v -s v1/core
     - pytest -v -s v1/executor
-    - pytest -v -s v1/kv_offload
     - pytest -v -s v1/sample
     - pytest -v -s v1/logits_processors
     - pytest -v -s v1/worker
@@ -300,12 +309,10 @@ steps:
     - pytest -v -s v1/spec_decode
     - pytest -v -s v1/kv_connector/unit
     - pytest -v -s v1/metrics
-    - pytest -v -s v1/test_kv_sharing.py
-    - pytest -v -s v1/test_metrics_reader.py
-    - pytest -v -s v1/test_oracle.py
-    - pytest -v -s v1/test_request.py
     - pytest -v -s v1/test_serial_utils.py
     - pytest -v -s v1/test_utils.py
+    - pytest -v -s v1/test_oracle.py
+    - pytest -v -s v1/test_metrics_reader.py
     # Integration test for streaming correctness (requires special branch).
     - 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
@@ -328,13 +335,12 @@ steps:
     - python3 offline_inference/vision_language.py --seed 0
     - python3 offline_inference/vision_language_pooling.py --seed 0
     - python3 offline_inference/vision_language_multi_image.py --seed 0
-    - python3 others/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 others/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
+    - VLLM_USE_V1=0 python3 others/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 others/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
     - python3 offline_inference/encoder_decoder_multimodal.py --model-type whisper --seed 0
     - python3 offline_inference/basic/classify.py
     - python3 offline_inference/basic/embed.py
     - python3 offline_inference/basic/score.py
-    - 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
-    - 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 2048
+    - VLLM_USE_V1=0 python3 offline_inference/profiling.py --model facebook/opt-125m run_num_steps --num-steps 2
 
 - label: Platform Tests (CUDA) # 4min
   timeout_in_minutes: 15
@@ -388,7 +394,6 @@ steps:
     - pytest -v -s compile/test_async_tp.py
     - pytest -v -s compile/test_fusion_all_reduce.py
     - pytest -v -s compile/test_decorator.py
-    - pytest -v -s compile/test_noop_elimination.py
 
 - label: PyTorch Fullgraph Smoke Test # 15min
   timeout_in_minutes: 30
@@ -543,6 +548,15 @@ steps:
   commands: # LMEval+Transcription WER check
   - pytest -s entrypoints/openai/correctness/
 
+- label: Encoder Decoder tests # 12min
+  timeout_in_minutes: 20
+  mirror_hardwares: [amdexperimental]
+  source_file_dependencies:
+  - vllm/
+  - tests/encoder_decoder
+  commands:
+    - pytest -v -s encoder_decoder
+
 - label: OpenAI-Compatible Tool Use # 23 min
   timeout_in_minutes: 35
   mirror_hardwares: [amdexperimental]
@@ -557,85 +571,36 @@ steps:
 
 #####  models test  #####
 
-- label: Basic Models Tests (Initialization)
-  timeout_in_minutes: 45
+- label: Basic Models Test # 57min
+  timeout_in_minutes: 75
   mirror_hardwares: [amdexperimental]
   torch_nightly: true
   source_file_dependencies:
   - vllm/
-  - tests/models/test_initialization.py
+  - tests/models
   commands:
-    # Run a subset of model initialization tests
-    - pytest -v -s models/test_initialization.py::test_can_initialize_small_subset
+    - pytest -v -s models/test_transformers.py
+    - pytest -v -s models/test_registry.py
+    - pytest -v -s models/test_utils.py
+    - pytest -v -s models/test_vision.py
+    - pytest -v -s models/test_initialization.py
 
-- label: Basic Models Tests (Extra Initialization) %N
+- label: Language Models Test (Standard) # 35min
   timeout_in_minutes: 45
   mirror_hardwares: [amdexperimental]
   torch_nightly: true
   source_file_dependencies:
-  - vllm/model_executor/models/
-  - tests/models/test_initialization.py
-  commands:
-    # Only when vLLM model source is modified - test initialization of a large
-    # subset of supported models (the complement of the small subset in the above
-    # test.) Also run if model initialization test file is modified
-    - pytest -v -s models/test_initialization.py \
-             -k 'not test_can_initialize_small_subset' \
-             --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT \
-             --shard-id=$$BUILDKITE_PARALLEL_JOB
-  parallelism: 2
-
-- label: Basic Models Tests (Other)
-  timeout_in_minutes: 45
-  mirror_hardwares: [amdexperimental]
-  torch_nightly: true
-  source_file_dependencies:
-  - vllm/
-  - tests/models/test_transformers.py
-  - tests/models/test_registry.py
-  - tests/models/test_utils.py
-  - tests/models/test_vision.py
-  commands:
-    - pytest -v -s models/test_transformers.py \
-                   models/test_registry.py \
-                   models/test_utils.py \
-                   models/test_vision.py
-
-- label: Language Models Tests (Standard)
-  timeout_in_minutes: 25
-  mirror_hardwares: [amdexperimental]
-  torch_nightly: true
-  source_file_dependencies:
   - vllm/
   - tests/models/language
   commands:
-    # Test standard language models, excluding a subset of slow tests
     - pip freeze | grep -E 'torch'
-    - pytest -v -s models/language -m 'core_model and (not slow_test)'
+    - pytest -v -s models/language -m core_model
 
-- label: Language Models Tests (Extra Standard) %N
+- label: Language Models Test (Hybrid) # 35 min
   timeout_in_minutes: 45
   mirror_hardwares: [amdexperimental]
   torch_nightly: true
   source_file_dependencies:
-  - vllm/model_executor/models/
-  - tests/models/language/pooling/test_embedding.py
-  - tests/models/language/generation/test_common.py
-  - tests/models/language/pooling/test_classification.py
-  commands:
-    # Shard slow subset of standard language models tests. Only run when model
-    # source is modified, or when specified test files are modified
-    - pip freeze | grep -E 'torch'
-    - pytest -v -s models/language -m 'core_model and slow_test' \
-             --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT \
-             --shard-id=$$BUILDKITE_PARALLEL_JOB
-  parallelism: 2
-
-- label: Language Models Tests (Hybrid) %N
-  timeout_in_minutes: 75
-  mirror_hardwares: [amdexperimental]
-  torch_nightly: true
-  source_file_dependencies:
   - vllm/
   - tests/models/language/generation
   commands:
@@ -643,12 +608,7 @@ steps:
     # Note: also needed to run plamo2 model in vLLM
     - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.2.5'
     - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.5.2'
-    # Shard hybrid language model tests
-    - pytest -v -s models/language/generation \
-                   -m hybrid_model \
-                   --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT \
-                   --shard-id=$$BUILDKITE_PARALLEL_JOB
-  parallelism: 2
+    - pytest -v -s models/language/generation -m hybrid_model
 
 - label: Language Models Test (Extended Generation) # 80min
   timeout_in_minutes: 110
@@ -772,9 +732,8 @@ steps:
     - pytest -v -s tests/models/multimodal/processing/
     - pytest -v -s tests/models/multimodal/test_mapping.py
     - python3 examples/offline_inference/basic/chat.py
+    - python3 examples/offline_inference/audio_language.py --model-type whisper
     - 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 # 38 min
   timeout_in_minutes: 60
@@ -804,7 +763,7 @@ steps:
     # 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_silu_nvfp4_quant_fusion.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
@@ -816,20 +775,6 @@ steps:
     - pytest -v -s tests/kernels/moe/test_flashinfer.py
     - pytest -v -s tests/compile/test_silu_mul_quant_fusion.py
 
-- label: GPT-OSS Eval (Blackwell)
-  timeout_in_minutes: 60
-  working_dir: "/vllm-workspace/"
-  gpu: b200
-  optional: true # disable while debugging
-  source_file_dependencies:
-  - tests/evals/gpt_oss
-  - vllm/model_executor/models/gpt_oss.py
-  - 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 tests/evals/gpt_oss/test_gpqa_correctness.py --model openai/gpt-oss-20b --metric 0.58 --server-args '--tensor-parallel-size 2'
-
 #####  1 GPU test  #####
 #####  multi gpus test  #####
 
@@ -844,8 +789,6 @@ steps:
   commands:
   - pytest -v -s distributed/test_comm_ops.py
   - pytest -v -s distributed/test_shm_broadcast.py
-  - pytest -v -s distributed/test_shm_buffer.py
-  - pytest -v -s distributed/test_shm_storage.py
 
 - label: 2 Node Tests (4 GPUs in total) # 16min
   timeout_in_minutes: 30
@@ -872,28 +815,26 @@ steps:
     - NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed'
     - python3 ../examples/offline_inference/data_parallel.py --dp-size=2 --tp-size=1 --node-size=2 --node-rank=1 --master-addr=192.168.10.10 --master-port=12345 --enforce-eager --trust-remote-code
 
-- label: Distributed Tests (2 GPUs) # 68min
-  timeout_in_minutes: 90
+- label: Distributed Tests (2 GPUs) # 110min
+  timeout_in_minutes: 150
   mirror_hardwares: [amdexperimental]
   working_dir: "/vllm-workspace/tests"
   num_gpus: 2
   source_file_dependencies:
-  - vllm/compilation/
   - vllm/distributed/
   - vllm/engine/
   - vllm/executor/
-  - vllm/worker/worker_base.py
-  - vllm/v1/engine/
-  - vllm/v1/worker/
-  - tests/compile/test_basic_correctness.py
-  - tests/compile/test_wrapper.py
+  - vllm/model_executor/models/
   - tests/distributed/
-  - tests/entrypoints/llm/test_collective_rpc.py
+  - vllm/compilation
+  - vllm/worker/worker_base.py
+  - vllm/worker/worker.py
+  - vllm/worker/model_runner.py
+  - entrypoints/llm/test_collective_rpc.py
   - tests/v1/test_async_llm_dp.py
   - tests/v1/test_external_lb_dp.py
   - tests/v1/entrypoints/openai/test_multi_api_servers.py
-  - tests/v1/shutdown
-  - tests/v1/worker/test_worker_memory_snapshot.py
+  - vllm/v1/engine/
   commands:
   - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/test_async_llm_dp.py
   - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/test_external_lb_dp.py
@@ -902,29 +843,19 @@ steps:
   - pytest -v -s ./compile/test_basic_correctness.py
   - 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'
-  - pytest -v -s distributed/test_sequence_parallel.py
-  - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown
-  - pytest -v -s v1/worker/test_worker_memory_snapshot.py
-
-- label: Distributed Model Tests (2 GPUs) # 37min
-  timeout_in_minutes: 50
-  mirror_hardwares: [amdexperimental]
-  working_dir: "/vllm-workspace/tests"
-  num_gpus: 2
-  source_file_dependencies:
-  - vllm/model_executor/model_loader/sharded_state_loader.py
-  - vllm/model_executor/models/
-  - tests/basic_correctness/
-  - tests/model_executor/model_loader/test_sharded_state_loader.py
-  - tests/models/
-  commands:
   - TARGET_TEST_SUITE=L4 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
-  - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py
   # Avoid importing model tests that cause CUDA reinitialization error
   - pytest models/test_transformers.py -v -s -m 'distributed(num_gpus=2)'
   - pytest models/language -v -s -m 'distributed(num_gpus=2)'
   - pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py
   - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest models/multimodal/generation/test_whisper.py -v -s -m 'distributed(num_gpus=2)'
+  # test sequence parallel
+  - pytest -v -s distributed/test_sequence_parallel.py
+  # this test fails consistently.
+  # TODO: investigate and fix
+  - VLLM_USE_V1=0 CUDA_VISIBLE_DEVICES=0,1 pytest -v -s test_sharded_state_loader.py
+  - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown
+  - pytest -v -s models/multimodal/generation/test_maverick.py
 
 - label: Plugin Tests (2 GPUs) # 40min
   timeout_in_minutes: 60
@@ -967,6 +898,7 @@ steps:
   commands:
   - pytest -v -s distributed/test_pp_cudagraph.py
   - pytest -v -s distributed/test_pipeline_parallel.py
+  # - pytest -v -s distributed/test_context_parallel.py # TODO: enable it on Hopper runners or add triton MLA support
 
 - label: LoRA TP Test (Distributed) # 17 min
   timeout_in_minutes: 30
@@ -1040,34 +972,9 @@ 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
 
-##### H200 test #####
-- label: Distrubted Tests (H200) # optional
+- label: Qwen MoE EP Test # optional
   gpu: h200
   optional: true
-  working_dir: "/vllm-workspace/"
   num_gpus: 2
   commands:
-    - pytest -v -s tests/distributed/test_context_parallel.py
-    - CUDA_VISIBLE_DEVICES=1,2 VLLM_ALL2ALL_BACKEND=deepep_high_throughput VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model Qwen/Qwen1.5-MoE-A2.7B --tp-size=1  --dp-size=2 --max-model-len 2048
-
-##### B200 test #####
-- label: Distributed Tests (B200) # optional
-  gpu: b200
-  optional: true
-  working_dir: "/vllm-workspace/"
-  num_gpus: 2
-  commands:
-    - pytest -v -s tests/distributed/test_context_parallel.py
-    - pytest -v -s tests/distributed/test_nccl_symm_mem_allreduce.py
-
-##### RL Integration Tests #####
-- label: Prime-RL Integration Test # 15min
-  timeout_in_minutes: 30
-  optional: true
-  num_gpus: 2
-  working_dir: "/vllm-workspace"
-  source_file_dependencies:
-  - vllm/
-  - .buildkite/scripts/run-prime-rl-test.sh
-  commands:
-    - bash .buildkite/scripts/run-prime-rl-test.sh
+    - CUDA_VISIBLE_DEVICES=1,2 VLLM_ALL2ALL_BACKEND=deepep_high_throughput VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 /vllm-workspace/examples/offline_inference/data_parallel.py --model Qwen/Qwen1.5-MoE-A2.7B --tp-size=1  --dp-size=2 --max-model-len 2048

.coveragerc

@@ -1,32 +0,0 @@
-[run]
-source = vllm
-omit =
-    */tests/*
-    */test_*
-    */__pycache__/*
-    */build/*
-    */dist/*
-    */vllm.egg-info/*
-    */third_party/*
-    */examples/*
-    */benchmarks/*
-    */docs/*
-
-[report]
-exclude_lines =
-    pragma: no cover
-    def __repr__
-    if self.debug:
-    if settings.DEBUG
-    raise AssertionError
-    raise NotImplementedError
-    if 0:
-    if __name__ == .__main__.:
-    class .*\bProtocol\):
-    @(abc\.)?abstractmethod
-
-[html]
-directory = htmlcov
-
-[xml]
-output = coverage.xml

.github/CODEOWNERS

@@ -2,24 +2,24 @@
 # for more info about CODEOWNERS file
 
 # This lists cover the "core" components of vLLM that require careful review
-/vllm/attention @LucasWilkinson
 /vllm/attention/backends/abstract.py @WoosukKwon @zhuohan123 @youkaichao @alexm-redhat @comaniac @njhill
+/vllm/core @zhuohan123 @youkaichao @alexm-redhat @comaniac @njhill
+/vllm/engine/llm_engine.py @zhuohan123 @youkaichao @alexm-redhat @comaniac @njhill
 /vllm/executor/executor_base.py @zhuohan123 @youkaichao @alexm-redhat @comaniac @njhill @22quinn
 /vllm/worker/worker_base.py @zhuohan123 @youkaichao @alexm-redhat @comaniac @njhill @22quinn
-/vllm/model_executor/layers/fused_moe @mgoin
+/vllm/worker/worker.py @zhuohan123 @youkaichao @alexm-redhat @comaniac @njhill
 /vllm/model_executor/layers/sampler.py @zhuohan123 @youkaichao @alexm-redhat @comaniac @njhill @NickLucche
 /vllm/model_executor/layers/quantization @mgoin @robertgshaw2-redhat @tlrmchlsmth @yewentao256
 /vllm/model_executor/layers/mamba @tdoublep
 /vllm/model_executor/model_loader @22quinn
 /vllm/multimodal @DarkLight1337 @ywang96 @NickLucche
-/vllm/v1/attention @LucasWilkinson
 /vllm/v1/sample @22quinn @houseroad
 /vllm/vllm_flash_attn @LucasWilkinson
 /vllm/lora @jeejeelee
 /vllm/reasoning @aarnphm @chaunceyjiang
 /vllm/entrypoints @aarnphm @chaunceyjiang
 /vllm/compilation @zou3519 @youkaichao @ProExpertProg
-/vllm/distributed/kv_transfer @NickLucche @ApostaC
+/vllm/distributed/kv_transfer @NickLucche
 CMakeLists.txt @tlrmchlsmth @LucasWilkinson
 
 # Any change to the VllmConfig changes can have a large user-facing impact,
@@ -30,59 +30,44 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson
 /vllm/v1 @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @comaniac @alexm-redhat
 /vllm/v1/structured_output @mgoin @russellb @aarnphm @benchislett
 /vllm/v1/spec_decode @benchislett @luccafong
-/vllm/v1/attention/backends/flashinfer.py @mgoin
 /vllm/v1/attention/backends/triton_attn.py @tdoublep
-/vllm/v1/core @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @comaniac @alexm-redhat @heheda12345 @ApostaC
+/vllm/v1/core @heheda12345
 /vllm/v1/kv_cache_interface.py @heheda12345
-/vllm/v1/offloading @ApostaC
 
 # Test ownership
 /.buildkite/lm-eval-harness @mgoin @simon-mo
+/tests/async_engine @njhill @robertgshaw2-redhat @simon-mo
 /tests/distributed/test_multi_node_assignment.py @youkaichao
 /tests/distributed/test_pipeline_parallel.py @youkaichao
 /tests/distributed/test_same_node.py @youkaichao
 /tests/entrypoints @DarkLight1337 @robertgshaw2-redhat @simon-mo @aarnphm @NickLucche
-/tests/evals @mgoin
-/tests/kernels @mgoin @tlrmchlsmth @WoosukKwon @yewentao256
+/tests/kernels @tlrmchlsmth @WoosukKwon @yewentao256
 /tests/models @DarkLight1337 @ywang96
 /tests/multimodal @DarkLight1337 @ywang96 @NickLucche
+/tests/prefix_caching @comaniac @KuntaiDu
 /tests/quantization @mgoin @robertgshaw2-redhat @yewentao256
 /tests/test_inputs.py @DarkLight1337 @ywang96
 /tests/v1/entrypoints/llm/test_struct_output_generate.py @mgoin @russellb @aarnphm
 /tests/v1/structured_output @mgoin @russellb @aarnphm
-/tests/v1/core @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @comaniac @alexm-redhat @heheda12345 @ApostaC
+/tests/v1/core @heheda12345
 /tests/weight_loading @mgoin @youkaichao @yewentao256
 /tests/lora @jeejeelee
 /tests/models/language/generation/test_hybrid.py @tdoublep
-/tests/v1/kv_connector/nixl_integration @NickLucche 
-/tests/v1/kv_connector @ApostaC
-/tests/v1/offloading @ApostaC
-
-# Transformers backend
-/vllm/model_executor/models/transformers.py @hmellor
-/tests/models/test_transformers.py @hmellor
+/tests/v1/kv_connector/nixl_integration @NickLucche
 
 # Docs
-/docs/mkdocs @hmellor
-/docs/**/*.yml @hmellor
-/requirements/docs.txt @hmellor
-.readthedocs.yaml @hmellor
+/docs @hmellor
 mkdocs.yaml @hmellor
 
-# Linting
-.markdownlint.yaml @hmellor
-.pre-commit-config.yaml @hmellor
-/tools/pre_commit @hmellor
-
 # CPU
-/vllm/v1/worker/cpu* @bigPYJ1151
+/vllm/v1/worker/^cpu @bigPYJ1151
 /csrc/cpu @bigPYJ1151
 /vllm/platforms/cpu.py @bigPYJ1151
 /cmake/cpu_extension.cmake @bigPYJ1151
 /docker/Dockerfile.cpu @bigPYJ1151
 
 # Intel GPU
-/vllm/v1/worker/xpu* @jikunshang
+/vllm/v1/worker/^xpu @jikunshang
 /vllm/platforms/xpu.py @jikunshang
 /docker/Dockerfile.xpu @jikunshang
 
@@ -116,7 +101,4 @@ mkdocs.yaml @hmellor
 /vllm/v1/worker/tpu* @NickLucche
 /vllm/platforms/tpu.py @NickLucche
 /vllm/v1/sample/tpu @NickLucche
-/vllm/tests/v1/tpu @NickLucche
-
-# KVConnector installation files
-/requirements/kv_connectors.txt @NickLucche
+/vllm/tests/v1/tpu @NickLucche

.github/ISSUE_TEMPLATE/750-RFC.yml

@@ -43,6 +43,10 @@ body:
       Any other things you would like to mention.
   validations:
     required: false
+- type: markdown
+  attributes:
+    value: >
+      Thanks for contributing 🎉! The vLLM core team hosts a biweekly RFC review session at 9:30AM Pacific Time, while most RFCs can be discussed online, you can optionally sign up for a slot to discuss your RFC online [here](https://docs.google.com/document/d/1CiLVBZeIVfR7_PNAKVSusxpceywkoOOB78qoWqHvSZc/edit).
 - type: checkboxes
   id: askllm
   attributes:

.github/mergify.yml

@@ -124,16 +124,9 @@ pull_request_rules:
     - or:
       - files~=^examples/.*gpt[-_]?oss.*\.py
       - files~=^tests/.*gpt[-_]?oss.*\.py
-      - files~=^tests/entrypoints/openai/test_response_api_with_harmony.py
-      - files~=^tests/entrypoints/test_context.py
       - files~=^vllm/model_executor/models/.*gpt[-_]?oss.*\.py
       - files~=^vllm/model_executor/layers/.*gpt[-_]?oss.*\.py
-      - files~=^vllm/entrypoints/harmony_utils.py
-      - files~=^vllm/entrypoints/tool_server.py
-      - files~=^vllm/entrypoints/tool.py
-      - files~=^vllm/entrypoints/context.py
       - title~=(?i)gpt[-_]?oss
-      - title~=(?i)harmony
   actions:
     label:
       add:
@@ -171,7 +164,7 @@ pull_request_rules:
       - files=examples/online_serving/openai_chat_completion_structured_outputs.py
       - files=examples/online_serving/openai_chat_completion_structured_outputs_with_reasoning.py
       - files~=^tests/v1/structured_output/
-      - files=tests/v1/entrypoints/llm/test_struct_output_generate.py
+      - files=tests/v1/entrypoints/llm/test_guided_generate.py
       - files~=^vllm/v1/structured_output/
   actions:
     label:
@@ -302,20 +295,3 @@ pull_request_rules:
     label:
       remove:
         - needs-rebase
-
-- name: label-kv-connector
-  description: Automatically apply kv-connector label
-  conditions:
-    - or:
-      - files~=^examples/online_serving/disaggregated[^/]*/.*
-      - files~=^examples/offline_inference/disaggregated[^/]*/.*
-      - files~=^examples/others/lmcache/
-      - files~=^tests/v1/kv_connector/
-      - files~=^vllm/distributed/kv_transfer/
-      - title~=(?i)\bP/?D\b
-      - title~=(?i)NIXL
-      - title~=(?i)LMCache
-  actions:
-    label:
-      add:
-        - kv-connector

.github/workflows/bc-lint.yml

@@ -6,8 +6,6 @@ on:
       - opened
       - synchronize
       - reopened
-      - labeled
-      - unlabeled
 
 jobs:
   bc_lint:

.pre-commit-config.yaml

@@ -49,7 +49,7 @@ repos:
   rev: 0.6.17
   hooks:
     - id: pip-compile
-      args: [requirements/test.in, -o, requirements/test.txt, --index-strategy, unsafe-best-match, --torch-backend, cu128, --python-platform, x86_64-manylinux_2_28]
+      args: [requirements/test.in, -o, requirements/test.txt, --index-strategy, unsafe-best-match, --torch-backend, cu128]
       files: ^requirements/test\.(in|txt)$
 - repo: local
   hooks:
@@ -60,32 +60,38 @@ repos:
     files: ^requirements/test\.(in|txt)$
   - id: mypy-local
     name: Run mypy for local Python installation
-    entry: python tools/pre_commit/mypy.py 0 "local"
+    entry: tools/mypy.sh 0 "local"
+    language: python
+    types: [python]
+    additional_dependencies: &mypy_deps [mypy==1.11.1, types-cachetools, types-setuptools, types-PyYAML, types-requests, pydantic]
     stages: [pre-commit] # Don't run in CI
-    <<: &mypy_common
-      language: python
-      types_or: [python, pyi]
-      require_serial: true
-      additional_dependencies: [mypy==1.11.1, regex, types-cachetools, types-setuptools, types-PyYAML, types-requests, types-torch, pydantic]
   - id: mypy-3.9 # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup is less awkward
     name: Run mypy for Python 3.9
-    entry: python tools/pre_commit/mypy.py 1 "3.9"
-    <<: *mypy_common
+    entry: tools/mypy.sh 1 "3.9"
+    language: python
+    types: [python]
+    additional_dependencies: *mypy_deps
     stages: [manual] # Only run in CI
   - 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"
-    <<: *mypy_common
+    entry: tools/mypy.sh 1 "3.10"
+    language: python
+    types: [python]
+    additional_dependencies: *mypy_deps
     stages: [manual] # Only run in CI
   - id: mypy-3.11 # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup is less awkward
     name: Run mypy for Python 3.11
-    entry: python tools/pre_commit/mypy.py 1 "3.11"
-    <<: *mypy_common
+    entry: tools/mypy.sh 1 "3.11"
+    language: python
+    types: [python]
+    additional_dependencies: *mypy_deps
     stages: [manual] # Only run in CI
   - id: mypy-3.12 # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup is less awkward
     name: Run mypy for Python 3.12
-    entry: python tools/pre_commit/mypy.py 1 "3.12"
-    <<: *mypy_common
+    entry: tools/mypy.sh 1 "3.12"
+    language: python
+    types: [python]
+    additional_dependencies: *mypy_deps
     stages: [manual] # Only run in CI
   - id: shellcheck
     name: Lint shell scripts
@@ -149,15 +155,18 @@ repos:
     additional_dependencies: [regex]
   - id: check-pickle-imports
     name: Prevent new pickle/cloudpickle imports
-    entry: python tools/pre_commit/check_pickle_imports.py
+    entry: python tools/check_pickle_imports.py
     language: python
     types: [python]
-    additional_dependencies: [regex]
+    pass_filenames: false
+    additional_dependencies: [pathspec, regex]
   - id: validate-config
     name: Validate configuration has default values and that each field has a docstring
     entry: python tools/validate_config.py
     language: python
-    additional_dependencies: [regex]
+    types: [python]
+    pass_filenames: true
+    files: vllm/config.py|tests/test_config.py|vllm/entrypoints/openai/cli_args.py
   # Keep `suggestion` last
   - id: suggestion
     name: Suggestion

.readthedocs.yaml

@@ -13,7 +13,6 @@ build:
 
 mkdocs:
   configuration: mkdocs.yaml
-  fail_on_warning: true
 
 # Optionally declare the Python requirements required to build your docs
 python:

CMakeLists.txt

@@ -13,10 +13,6 @@ cmake_minimum_required(VERSION 3.26)
 # cmake --install . --component _C
 project(vllm_extensions LANGUAGES CXX)
 
-set(CMAKE_CXX_STANDARD 17)
-set(CMAKE_CXX_STANDARD_REQUIRED ON)
-
-
 # CUDA by default, can be overridden by using -DVLLM_TARGET_DEVICE=... (used by setup.py)
 set(VLLM_TARGET_DEVICE "cuda" CACHE STRING "Target device backend for vLLM")
 message(STATUS "Build type: ${CMAKE_BUILD_TYPE}")
@@ -175,16 +171,6 @@ if(NVCC_THREADS AND VLLM_GPU_LANG STREQUAL "CUDA")
   list(APPEND VLLM_GPU_FLAGS "--threads=${NVCC_THREADS}")
 endif()
 
-#
-# Set CUDA include flags for CXX compiler.
-#
-if(VLLM_GPU_LANG STREQUAL "CUDA")
-  set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -I${CUDA_TOOLKIT_ROOT_DIR}/include")
-  if(CUDA_VERSION VERSION_GREATER_EQUAL 13.0)
-    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -I${CUDA_TOOLKIT_ROOT_DIR}/include/cccl")
-  endif()
-endif()
-
 #
 # Use FetchContent for C++ dependencies that are compiled as part of vLLM's build process.
 # setup.py will override FETCHCONTENT_BASE_DIR to play nicely with sccache.
@@ -308,6 +294,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
     "csrc/quantization/fp4/nvfp4_blockwise_moe_kernel.cu"
     "csrc/sparse/cutlass/sparse_scaled_mm_entry.cu"
     "csrc/cutlass_extensions/common.cpp"
+    "csrc/attention/mla/cutlass_mla_entry.cu"
     "csrc/quantization/fp8/per_token_group_quant.cu")
 
   set_gencode_flags_for_srcs(
@@ -594,6 +581,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
   cuda_archs_loose_intersection(MLA_ARCHS "10.0a" "${CUDA_ARCHS}")
   if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND MLA_ARCHS)
     set(SRCS
+      "csrc/attention/mla/cutlass_mla_kernels.cu"
       "csrc/attention/mla/sm100_cutlass_mla_kernel.cu")
     set_gencode_flags_for_srcs(
       SRCS "${SRCS}"
@@ -791,17 +779,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
     endif()
   endif()
 
-  # Hadacore kernels
-  cuda_archs_loose_intersection(HADACORE_ARCHS "8.0;8.9;9.0" "${CUDA_ARCHS}")
-  if(HADACORE_ARCHS)
-    set(SRCS "csrc/quantization/hadamard/hadacore/hadamard_transform_cuda.cu")
-    set_gencode_flags_for_srcs(
-      SRCS "${SRCS}"
-      CUDA_ARCHS "${HADACORE_ARCHS}")
-    list(APPEND VLLM_EXT_SRC "${SRCS}")
-    message(STATUS "Building hadacore")
-  endif()
-
 # if CUDA endif
 endif()
 

README.md

@@ -81,7 +81,7 @@ vLLM is flexible and easy to use with:
 - Tensor, pipeline, data and expert parallelism support for distributed inference
 - Streaming outputs
 - OpenAI-compatible API server
-- Support for NVIDIA GPUs, AMD CPUs and GPUs, Intel CPUs and GPUs, PowerPC CPUs, and TPU. Additionally, support for diverse hardware plugins such as Intel Gaudi, IBM Spyre and Huawei Ascend.
+- Support NVIDIA GPUs, AMD CPUs and GPUs, Intel CPUs and GPUs, PowerPC CPUs, TPU, and AWS Neuron
 - Prefix caching support
 - Multi-LoRA support
 

benchmarks/README.md

@@ -1,20 +1,874 @@
-# Benchmarks
+# Benchmarking vLLM
 
-This directory used to contain vLLM's benchmark scripts and utilities for performance testing and evaluation.
+This README guides you through running benchmark tests with the extensive
+datasets supported on vLLM. It’s a living document, updated as new features and datasets
+become available.
 
-## Contents
+## Dataset Overview
 
-- **Serving benchmarks**: Scripts for testing online inference performance (latency, throughput)
-- **Throughput benchmarks**: Scripts for testing offline batch inference performance
-- **Specialized benchmarks**: Tools for testing specific features like structured output, prefix caching, long document QA, request prioritization, and multi-modal inference
-- **Dataset utilities**: Framework for loading and sampling from various benchmark datasets (ShareGPT, HuggingFace datasets, synthetic data, etc.)
+<table style="width:100%; border-collapse: collapse;">
+  <thead>
+    <tr>
+      <th style="width:15%; text-align: left;">Dataset</th>
+      <th style="width:10%; text-align: center;">Online</th>
+      <th style="width:10%; text-align: center;">Offline</th>
+      <th style="width:65%; text-align: left;">Data Path</th>
+    </tr>
+  </thead>
+  <tbody>
+    <tr>
+      <td><strong>ShareGPT</strong></td>
+      <td style="text-align: center;">✅</td>
+      <td style="text-align: center;">✅</td>
+      <td><code>wget https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json</code></td>
+    </tr>
+    <tr>
+      <td><strong>ShareGPT4V (Image)</strong></td>
+      <td style="text-align: center;">✅</td>
+      <td style="text-align: center;">✅</td>
+      <td>
+        <code>wget https://huggingface.co/datasets/Lin-Chen/ShareGPT4V/blob/main/sharegpt4v_instruct_gpt4-vision_cap100k.json</code>
+        <br>
+        <div>Note that the images need to be downloaded separately. For example, to download COCO's 2017 Train images:</div>
+        <code>wget http://images.cocodataset.org/zips/train2017.zip</code>
+      </td>
+    </tr>
+        <tr>
+      <td><strong>ShareGPT4Video (Video)</strong></td>
+      <td style="text-align: center;">✅</td>
+      <td style="text-align: center;">✅</td>
+      <td>
+        <code>git clone https://huggingface.co/datasets/ShareGPT4Video/ShareGPT4Video</code>
+      </td>
+    </tr>
+    <tr>
+      <td><strong>BurstGPT</strong></td>
+      <td style="text-align: center;">✅</td>
+      <td style="text-align: center;">✅</td>
+      <td><code>wget https://github.com/HPMLL/BurstGPT/releases/download/v1.1/BurstGPT_without_fails_2.csv</code></td>
+    </tr>
+    <tr>
+      <td><strong>Sonnet (deprecated)</strong></td>
+      <td style="text-align: center;">✅</td>
+      <td style="text-align: center;">✅</td>
+      <td>Local file: <code>benchmarks/sonnet.txt</code></td>
+    </tr>
+    <tr>
+      <td><strong>Random</strong></td>
+      <td style="text-align: center;">✅</td>
+      <td style="text-align: center;">✅</td>
+      <td><code>synthetic</code></td>
+    </tr>
+    <tr>
+      <td><strong>RandomMultiModal (Image/Video)</strong></td>
+      <td style="text-align: center;">🟡</td>
+      <td style="text-align: center;">🚧</td>
+      <td><code>synthetic</code> </td>
+    </tr>
+    <tr>
+      <td><strong>Prefix Repetition</strong></td>
+      <td style="text-align: center;">✅</td>
+      <td style="text-align: center;">✅</td>
+      <td><code>synthetic</code></td>
+    </tr>
+    <tr>
+      <td><strong>HuggingFace-VisionArena</strong></td>
+      <td style="text-align: center;">✅</td>
+      <td style="text-align: center;">✅</td>
+      <td><code>lmarena-ai/VisionArena-Chat</code></td>
+    </tr>
+    <tr>
+      <td><strong>HuggingFace-InstructCoder</strong></td>
+      <td style="text-align: center;">✅</td>
+      <td style="text-align: center;">✅</td>
+      <td><code>likaixin/InstructCoder</code></td>
+    </tr>
+      <tr>
+      <td><strong>HuggingFace-AIMO</strong></td>
+      <td style="text-align: center;">✅</td>
+      <td style="text-align: center;">✅</td>
+      <td><code>AI-MO/aimo-validation-aime</code> , <code>AI-MO/NuminaMath-1.5</code>, <code>AI-MO/NuminaMath-CoT</code></td>
+    </tr>
+    <tr>
+      <td><strong>HuggingFace-Other</strong></td>
+      <td style="text-align: center;">✅</td>
+      <td style="text-align: center;">✅</td>
+      <td><code>lmms-lab/LLaVA-OneVision-Data</code>, <code>Aeala/ShareGPT_Vicuna_unfiltered</code></td>
+    </tr>
+    <tr>
+      <td><strong>HuggingFace-MTBench</strong></td>
+      <td style="text-align: center;">✅</td>
+      <td style="text-align: center;">✅</td>
+      <td><code>philschmid/mt-bench</code></td>
+    </tr>
+    <tr>
+      <td><strong>HuggingFace-Blazedit</strong></td>
+      <td style="text-align: center;">✅</td>
+      <td style="text-align: center;">✅</td>
+      <td><code>vdaita/edit_5k_char</code>, <code>vdaita/edit_10k_char</code></td>
+    </tr>
+    <tr>
+      <td><strong>Spec Bench</strong></td>
+      <td style="text-align: center;">✅</td>
+      <td style="text-align: center;">✅</td>
+      <td><code>wget https://raw.githubusercontent.com/hemingkx/Spec-Bench/refs/heads/main/data/spec_bench/question.jsonl</code></td>
+    </tr>
+    <tr>
+      <td><strong>Custom</strong></td>
+      <td style="text-align: center;">✅</td>
+      <td style="text-align: center;">✅</td>
+      <td>Local file: <code>data.jsonl</code></td>
+    </tr>
+  </tbody>
+</table>
 
-## Usage
+✅: supported
 
-For detailed usage instructions, examples, and dataset information, see the [Benchmark CLI documentation](https://docs.vllm.ai/en/latest/contributing/benchmarks.html#benchmark-cli).
+🟡: Partial support
 
-For full CLI reference see:
+🚧: to be supported
 
-- <https://docs.vllm.ai/en/latest/cli/bench/latency.html>
-- <https://docs.vllm.ai/en/latest/cli/bench/serve.html>
-- <https://docs.vllm.ai/en/latest/cli/bench/throughput.html>
+**Note**: HuggingFace dataset's `dataset-name` should be set to `hf`.
+For local `dataset-path`, please set `hf-name` to its Hugging Face ID like
+
+```bash
+--dataset-path /datasets/VisionArena-Chat/ --hf-name lmarena-ai/VisionArena-Chat
+```
+
+## 🚀 Example - Online Benchmark
+
+<details>
+<summary>Show more</summary>
+
+<br/>
+
+First start serving your model
+
+```bash
+vllm serve NousResearch/Hermes-3-Llama-3.1-8B
+```
+
+Then run the benchmarking script
+
+```bash
+# download dataset
+# wget https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json
+vllm bench serve \
+  --backend vllm \
+  --model NousResearch/Hermes-3-Llama-3.1-8B \
+  --endpoint /v1/completions \
+  --dataset-name sharegpt \
+  --dataset-path <your data path>/ShareGPT_V3_unfiltered_cleaned_split.json \
+  --num-prompts 10
+```
+
+If successful, you will see the following output
+
+```text
+============ Serving Benchmark Result ============
+Successful requests:                     10
+Benchmark duration (s):                  5.78
+Total input tokens:                      1369
+Total generated tokens:                  2212
+Request throughput (req/s):              1.73
+Output token throughput (tok/s):         382.89
+Total Token throughput (tok/s):          619.85
+---------------Time to First Token----------------
+Mean TTFT (ms):                          71.54
+Median TTFT (ms):                        73.88
+P99 TTFT (ms):                           79.49
+-----Time per Output Token (excl. 1st token)------
+Mean TPOT (ms):                          7.91
+Median TPOT (ms):                        7.96
+P99 TPOT (ms):                           8.03
+---------------Inter-token Latency----------------
+Mean ITL (ms):                           7.74
+Median ITL (ms):                         7.70
+P99 ITL (ms):                            8.39
+==================================================
+```
+
+### Custom Dataset
+
+If the dataset you want to benchmark is not supported yet in vLLM, even then you can benchmark on it using `CustomDataset`. Your data needs to be in `.jsonl` format and needs to have "prompt" field per entry, e.g., data.jsonl
+
+```json
+{"prompt": "What is the capital of India?"}
+{"prompt": "What is the capital of Iran?"}
+{"prompt": "What is the capital of China?"}
+```
+
+```bash
+# start server
+VLLM_USE_V1=1 vllm serve meta-llama/Llama-3.1-8B-Instruct
+```
+
+```bash
+# run benchmarking script
+vllm bench serve --port 9001 --save-result --save-detailed \
+  --backend vllm \
+  --model meta-llama/Llama-3.1-8B-Instruct \
+  --endpoint /v1/completions \
+  --dataset-name custom \
+  --dataset-path <path-to-your-data-jsonl> \
+  --custom-skip-chat-template \
+  --num-prompts 80 \
+  --max-concurrency 1 \
+  --temperature=0.3 \
+  --top-p=0.75 \
+  --result-dir "./log/"
+```
+
+You can skip applying chat template if your data already has it by using `--custom-skip-chat-template`.
+
+### VisionArena Benchmark for Vision Language Models
+
+```bash
+# need a model with vision capability here
+vllm serve Qwen/Qwen2-VL-7B-Instruct
+```
+
+```bash
+vllm bench serve \
+  --backend openai-chat \
+  --endpoint-type openai-chat \
+  --model Qwen/Qwen2-VL-7B-Instruct \
+  --endpoint /v1/chat/completions \
+  --dataset-name hf \
+  --dataset-path lmarena-ai/VisionArena-Chat \
+  --hf-split train \
+  --num-prompts 1000
+```
+
+### InstructCoder Benchmark with Speculative Decoding
+
+``` bash
+VLLM_USE_V1=1 vllm serve meta-llama/Meta-Llama-3-8B-Instruct \
+    --speculative-config $'{"method": "ngram",
+    "num_speculative_tokens": 5, "prompt_lookup_max": 5,
+    "prompt_lookup_min": 2}'
+```
+
+``` bash
+vllm bench serve \
+    --model meta-llama/Meta-Llama-3-8B-Instruct \
+    --dataset-name hf \
+    --dataset-path likaixin/InstructCoder \
+    --num-prompts 2048
+```
+
+### Spec Bench Benchmark with Speculative Decoding
+
+``` bash
+VLLM_USE_V1=1 vllm serve meta-llama/Meta-Llama-3-8B-Instruct \
+    --speculative-config $'{"method": "ngram",
+    "num_speculative_tokens": 5, "prompt_lookup_max": 5,
+    "prompt_lookup_min": 2}'
+```
+
+[SpecBench dataset](https://github.com/hemingkx/Spec-Bench)
+
+Run all categories:
+
+``` bash
+# Download the dataset using:
+# wget https://raw.githubusercontent.com/hemingkx/Spec-Bench/refs/heads/main/data/spec_bench/question.jsonl
+
+vllm bench serve \
+    --model meta-llama/Meta-Llama-3-8B-Instruct \
+    --dataset-name spec_bench \ 
+    --dataset-path "<YOUR_DOWNLOADED_PATH>/data/spec_bench/question.jsonl" \
+    --num-prompts -1
+```
+
+Available categories include `[writing, roleplay, reasoning, math, coding, extraction, stem, humanities, translation, summarization, qa, math_reasoning, rag]`.
+
+Run only a specific category like "summarization":
+
+``` bash
+vllm bench serve \
+    --model meta-llama/Meta-Llama-3-8B-Instruct \
+    --dataset-name spec_bench \ 
+    --dataset-path "<YOUR_DOWNLOADED_PATH>/data/spec_bench/question.jsonl" \
+    --num-prompts -1
+    --spec-bench-category "summarization"
+```
+
+### Other HuggingFaceDataset Examples
+
+```bash
+vllm serve Qwen/Qwen2-VL-7B-Instruct
+```
+
+`lmms-lab/LLaVA-OneVision-Data`:
+
+```bash
+vllm bench serve \
+  --backend openai-chat \
+  --endpoint-type openai-chat \
+  --model Qwen/Qwen2-VL-7B-Instruct \
+  --endpoint /v1/chat/completions \
+  --dataset-name hf \
+  --dataset-path lmms-lab/LLaVA-OneVision-Data \
+  --hf-split train \
+  --hf-subset "chart2text(cauldron)" \
+  --num-prompts 10
+```
+
+`Aeala/ShareGPT_Vicuna_unfiltered`:
+
+```bash
+vllm bench serve \
+  --backend openai-chat \
+  --endpoint-type openai-chat \
+  --model Qwen/Qwen2-VL-7B-Instruct \
+  --endpoint /v1/chat/completions \
+  --dataset-name hf \
+  --dataset-path Aeala/ShareGPT_Vicuna_unfiltered \
+  --hf-split train \
+  --num-prompts 10
+```
+
+`AI-MO/aimo-validation-aime`:
+
+``` bash
+vllm bench serve \
+    --model Qwen/QwQ-32B \
+    --dataset-name hf \
+    --dataset-path AI-MO/aimo-validation-aime \
+    --num-prompts 10 \
+    --seed 42
+```
+
+`philschmid/mt-bench`:
+
+``` bash
+vllm bench serve \
+    --model Qwen/QwQ-32B \
+    --dataset-name hf \
+    --dataset-path philschmid/mt-bench \
+    --num-prompts 80
+```
+
+`vdaita/edit_5k_char` or `vdaita/edit_10k_char`:
+
+``` bash
+vllm bench serve \
+    --model Qwen/QwQ-32B \
+    --dataset-name hf \
+    --dataset-path vdaita/edit_5k_char \
+    --num-prompts 90 \
+    --blazedit-min-distance 0.01 \
+    --blazedit-max-distance 0.99
+```
+
+### Running With Sampling Parameters
+
+When using OpenAI-compatible backends such as `vllm`, optional sampling
+parameters can be specified. Example client command:
+
+```bash
+vllm bench serve \
+  --backend vllm \
+  --model NousResearch/Hermes-3-Llama-3.1-8B \
+  --endpoint /v1/completions \
+  --dataset-name sharegpt \
+  --dataset-path <your data path>/ShareGPT_V3_unfiltered_cleaned_split.json \
+  --top-k 10 \
+  --top-p 0.9 \
+  --temperature 0.5 \
+  --num-prompts 10
+```
+
+### Running With Ramp-Up Request Rate
+
+The benchmark tool also supports ramping up the request rate over the
+duration of the benchmark run. This can be useful for stress testing the
+server or finding the maximum throughput that it can handle, given some latency budget.
+
+Two ramp-up strategies are supported:
+
+- `linear`: Increases the request rate linearly from a start value to an end value.
+- `exponential`: Increases the request rate exponentially.
+
+The following arguments can be used to control the ramp-up:
+
+- `--ramp-up-strategy`: The ramp-up strategy to use (`linear` or `exponential`).
+- `--ramp-up-start-rps`: The request rate at the beginning of the benchmark.
+- `--ramp-up-end-rps`: The request rate at the end of the benchmark.
+
+</details>
+
+## 📈 Example - Offline Throughput Benchmark
+
+<details>
+<summary>Show more</summary>
+
+<br/>
+
+```bash
+vllm bench throughput \
+  --model NousResearch/Hermes-3-Llama-3.1-8B \
+  --dataset-name sonnet \
+  --dataset-path vllm/benchmarks/sonnet.txt \
+  --num-prompts 10
+```
+
+If successful, you will see the following output
+
+```text
+Throughput: 7.15 requests/s, 4656.00 total tokens/s, 1072.15 output tokens/s
+Total num prompt tokens:  5014
+Total num output tokens:  1500
+```
+
+### VisionArena Benchmark for Vision Language Models
+
+```bash
+vllm bench throughput \
+  --model Qwen/Qwen2-VL-7B-Instruct \
+  --backend vllm-chat \
+  --dataset-name hf \
+  --dataset-path lmarena-ai/VisionArena-Chat \
+  --num-prompts 1000 \
+  --hf-split train
+```
+
+The `num prompt tokens` now includes image token counts
+
+```text
+Throughput: 2.55 requests/s, 4036.92 total tokens/s, 326.90 output tokens/s
+Total num prompt tokens:  14527
+Total num output tokens:  1280
+```
+
+### InstructCoder Benchmark with Speculative Decoding
+
+``` bash
+VLLM_WORKER_MULTIPROC_METHOD=spawn \
+VLLM_USE_V1=1 \
+vllm bench throughput \
+    --dataset-name=hf \
+    --dataset-path=likaixin/InstructCoder \
+    --model=meta-llama/Meta-Llama-3-8B-Instruct \
+    --input-len=1000 \
+    --output-len=100 \
+    --num-prompts=2048 \
+    --async-engine \
+    --speculative-config $'{"method": "ngram",
+    "num_speculative_tokens": 5, "prompt_lookup_max": 5,
+    "prompt_lookup_min": 2}'
+```
+
+```text
+Throughput: 104.77 requests/s, 23836.22 total tokens/s, 10477.10 output tokens/s
+Total num prompt tokens:  261136
+Total num output tokens:  204800
+```
+
+### Other HuggingFaceDataset Examples
+
+`lmms-lab/LLaVA-OneVision-Data`:
+
+```bash
+vllm bench throughput \
+  --model Qwen/Qwen2-VL-7B-Instruct \
+  --backend vllm-chat \
+  --dataset-name hf \
+  --dataset-path lmms-lab/LLaVA-OneVision-Data \
+  --hf-split train \
+  --hf-subset "chart2text(cauldron)" \
+  --num-prompts 10
+```
+
+`Aeala/ShareGPT_Vicuna_unfiltered`:
+
+```bash
+vllm bench throughput \
+  --model Qwen/Qwen2-VL-7B-Instruct \
+  --backend vllm-chat \
+  --dataset-name hf \
+  --dataset-path Aeala/ShareGPT_Vicuna_unfiltered \
+  --hf-split train \
+  --num-prompts 10
+```
+
+`AI-MO/aimo-validation-aime`:
+
+```bash
+vllm bench throughput \
+  --model Qwen/QwQ-32B \
+  --backend vllm \
+  --dataset-name hf \
+  --dataset-path AI-MO/aimo-validation-aime \
+  --hf-split train \
+  --num-prompts 10
+```
+
+Benchmark with LoRA adapters:
+
+``` bash
+# download dataset
+# wget https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json
+vllm bench throughput \
+  --model meta-llama/Llama-2-7b-hf \
+  --backend vllm \
+  --dataset_path <your data path>/ShareGPT_V3_unfiltered_cleaned_split.json \
+  --dataset_name sharegpt \
+  --num-prompts 10 \
+  --max-loras 2 \
+  --max-lora-rank 8 \
+  --enable-lora \
+  --lora-path yard1/llama-2-7b-sql-lora-test
+  ```
+
+</details>
+
+## 🛠️ Example - Structured Output Benchmark
+
+<details>
+<summary>Show more</summary>
+
+<br/>
+
+Benchmark the performance of structured output generation (JSON, grammar, regex).
+
+### Server Setup
+
+```bash
+vllm serve NousResearch/Hermes-3-Llama-3.1-8B
+```
+
+### JSON Schema Benchmark
+
+```bash
+python3 benchmarks/benchmark_serving_structured_output.py \
+  --backend vllm \
+  --model NousResearch/Hermes-3-Llama-3.1-8B \
+  --dataset json \
+  --structured-output-ratio 1.0 \
+  --request-rate 10 \
+  --num-prompts 1000
+```
+
+### Grammar-based Generation Benchmark
+
+```bash
+python3 benchmarks/benchmark_serving_structured_output.py \
+  --backend vllm \
+  --model NousResearch/Hermes-3-Llama-3.1-8B \
+  --dataset grammar \
+  --structure-type grammar \
+  --request-rate 10 \
+  --num-prompts 1000
+```
+
+### Regex-based Generation Benchmark
+
+```bash
+python3 benchmarks/benchmark_serving_structured_output.py \
+  --backend vllm \
+  --model NousResearch/Hermes-3-Llama-3.1-8B \
+  --dataset regex \
+  --request-rate 10 \
+  --num-prompts 1000
+```
+
+### Choice-based Generation Benchmark
+
+```bash
+python3 benchmarks/benchmark_serving_structured_output.py \
+  --backend vllm \
+  --model NousResearch/Hermes-3-Llama-3.1-8B \
+  --dataset choice \
+  --request-rate 10 \
+  --num-prompts 1000
+```
+
+### XGrammar Benchmark Dataset
+
+```bash
+python3 benchmarks/benchmark_serving_structured_output.py \
+  --backend vllm \
+  --model NousResearch/Hermes-3-Llama-3.1-8B \
+  --dataset xgrammar_bench \
+  --request-rate 10 \
+  --num-prompts 1000
+```
+
+</details>
+
+## 📚 Example - Long Document QA Benchmark
+
+<details>
+<summary>Show more</summary>
+
+<br/>
+
+Benchmark the performance of long document question-answering with prefix caching.
+
+### Basic Long Document QA Test
+
+```bash
+python3 benchmarks/benchmark_long_document_qa_throughput.py \
+  --model meta-llama/Llama-2-7b-chat-hf \
+  --enable-prefix-caching \
+  --num-documents 16 \
+  --document-length 2000 \
+  --output-len 50 \
+  --repeat-count 5
+```
+
+### Different Repeat Modes
+
+```bash
+# Random mode (default) - shuffle prompts randomly
+python3 benchmarks/benchmark_long_document_qa_throughput.py \
+  --model meta-llama/Llama-2-7b-chat-hf \
+  --enable-prefix-caching \
+  --num-documents 8 \
+  --document-length 3000 \
+  --repeat-count 3 \
+  --repeat-mode random
+
+# Tile mode - repeat entire prompt list in sequence
+python3 benchmarks/benchmark_long_document_qa_throughput.py \
+  --model meta-llama/Llama-2-7b-chat-hf \
+  --enable-prefix-caching \
+  --num-documents 8 \
+  --document-length 3000 \
+  --repeat-count 3 \
+  --repeat-mode tile
+
+# Interleave mode - repeat each prompt consecutively
+python3 benchmarks/benchmark_long_document_qa_throughput.py \
+  --model meta-llama/Llama-2-7b-chat-hf \
+  --enable-prefix-caching \
+  --num-documents 8 \
+  --document-length 3000 \
+  --repeat-count 3 \
+  --repeat-mode interleave
+```
+
+</details>
+
+## 🗂️ Example - Prefix Caching Benchmark
+
+<details>
+<summary>Show more</summary>
+
+<br/>
+
+Benchmark the efficiency of automatic prefix caching.
+
+### Fixed Prompt with Prefix Caching
+
+```bash
+python3 benchmarks/benchmark_prefix_caching.py \
+  --model meta-llama/Llama-2-7b-chat-hf \
+  --enable-prefix-caching \
+  --num-prompts 1 \
+  --repeat-count 100 \
+  --input-length-range 128:256
+```
+
+### ShareGPT Dataset with Prefix Caching
+
+```bash
+# download dataset
+# wget https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json
+
+python3 benchmarks/benchmark_prefix_caching.py \
+  --model meta-llama/Llama-2-7b-chat-hf \
+  --dataset-path /path/ShareGPT_V3_unfiltered_cleaned_split.json \
+  --enable-prefix-caching \
+  --num-prompts 20 \
+  --repeat-count 5 \
+  --input-length-range 128:256
+```
+
+### Prefix Repetition Dataset
+
+```bash
+vllm bench serve \
+  --backend openai \
+  --model meta-llama/Llama-2-7b-chat-hf \
+  --dataset-name prefix_repetition \
+  --num-prompts 100 \
+  --prefix-repetition-prefix-len 512 \
+  --prefix-repetition-suffix-len 128 \
+  --prefix-repetition-num-prefixes 5 \
+  --prefix-repetition-output-len 128
+```
+
+</details>
+
+## ⚡ Example - Request Prioritization Benchmark
+
+<details>
+<summary>Show more</summary>
+
+<br/>
+
+Benchmark the performance of request prioritization in vLLM.
+
+### Basic Prioritization Test
+
+```bash
+python3 benchmarks/benchmark_prioritization.py \
+  --model meta-llama/Llama-2-7b-chat-hf \
+  --input-len 128 \
+  --output-len 64 \
+  --num-prompts 100 \
+  --scheduling-policy priority
+```
+
+### Multiple Sequences per Prompt
+
+```bash
+python3 benchmarks/benchmark_prioritization.py \
+  --model meta-llama/Llama-2-7b-chat-hf \
+  --input-len 128 \
+  --output-len 64 \
+  --num-prompts 100 \
+  --scheduling-policy priority \
+  --n 2
+```
+
+</details>
+
+## 👁️ Example - Multi-Modal Benchmark
+
+<details>
+<summary>Show more</summary>
+
+<br/>
+
+Benchmark the performance of multi-modal requests in vLLM.
+
+### Images (ShareGPT4V)
+
+Start vLLM:
+
+```bash
+python -m vllm.entrypoints.openai.api_server \
+  --model Qwen/Qwen2.5-VL-7B-Instruct \
+  --dtype bfloat16 \
+  --limit-mm-per-prompt '{"image": 1}' \
+  --allowed-local-media-path /path/to/sharegpt4v/images
+```
+
+Send requests with images:
+
+```bash
+vllm bench serve \
+  --backend openai-chat \
+  --model Qwen/Qwen2.5-VL-7B-Instruct \
+  --dataset-name sharegpt \
+  --dataset-path /path/to/ShareGPT4V/sharegpt4v_instruct_gpt4-vision_cap100k.json \
+  --num-prompts 100 \
+  --save-result \
+  --result-dir ~/vllm_benchmark_results \
+  --save-detailed \
+  --endpoint /v1/chat/completion
+```
+
+### Videos (ShareGPT4Video)
+
+Start vLLM:
+
+```bash
+python -m vllm.entrypoints.openai.api_server \
+  --model Qwen/Qwen2.5-VL-7B-Instruct \
+  --dtype bfloat16 \
+  --limit-mm-per-prompt '{"video": 1}' \
+  --allowed-local-media-path /path/to/sharegpt4video/videos
+```
+
+Send requests with videos:
+
+```bash
+vllm bench serve \
+  --backend openai-chat \
+  --model Qwen/Qwen2.5-VL-7B-Instruct \
+  --dataset-name sharegpt \
+  --dataset-path /path/to/ShareGPT4Video/llava_v1_5_mix665k_with_video_chatgpt72k_share4video28k.json \
+  --num-prompts 100 \
+  --save-result \
+  --result-dir ~/vllm_benchmark_results \
+  --save-detailed \
+  --endpoint /v1/chat/completion
+```
+
+### Synthetic Random Images (random-mm)
+
+Generate synthetic image inputs alongside random text prompts to stress-test vision models without external datasets.
+
+Notes:
+
+- 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):
+
+```bash
+vllm serve Qwen/Qwen2.5-VL-3B-Instruct \
+  --dtype bfloat16 \
+  --max-model-len 16384 \
+  --limit-mm-per-prompt '{"image": 3, "video": 0}' \
+  --mm-processor-kwargs max_pixels=1003520
+```
+
+Benchmark. It is recommended to use the flag `--ignore-eos` to simulate real responses. You can set the size of the output via the arg `random-output-len`.
+
+Ex.1: Fixed number of items and a single image resolution, enforcing generation of approx 40 tokens:
+
+```bash
+vllm bench serve \
+  --backend openai-chat \
+  --model Qwen/Qwen2.5-VL-3B-Instruct \
+  --endpoint /v1/chat/completions \
+  --dataset-name random-mm \
+  --num-prompts 100 \
+  --max-concurrency 10 \
+  --random-prefix-len 25 \
+  --random-input-len 300 \
+  --random-output-len 40 \
+  --random-range-ratio 0.2 \
+  --random-mm-base-items-per-request 2 \
+  --random-mm-limit-mm-per-prompt '{"image": 3, "video": 0}' \
+  --random-mm-bucket-config '{(224, 224, 1): 1.0}' \
+  --request-rate inf \
+  --ignore-eos \
+  --seed 42
+```
+
+The number of items per request can be controlled by passing multiple image buckets:
+
+```bash
+  --random-mm-base-items-per-request 2 \
+  --random-mm-num-mm-items-range-ratio 0.5 \
+  --random-mm-limit-mm-per-prompt '{"image": 4, "video": 0}' \
+  --random-mm-bucket-config '{(256, 256, 1): 0.7, (720, 1280, 1): 0.3}' \
+```
+
+Flags specific to `random-mm`:
+
+- `--random-mm-base-items-per-request`: base number of multimodal items per request.
+- `--random-mm-num-mm-items-range-ratio`: vary item count uniformly in the closed integer range [floor(n·(1−r)), ceil(n·(1+r))]. Set r=0 to keep it fixed; r=1 allows 0 items.
+- `--random-mm-limit-mm-per-prompt`: per-modality hard caps, e.g. '{"image": 3, "video": 0}'.
+- `--random-mm-bucket-config`: dict mapping (H, W, T) → probability. Entries with probability 0 are removed; remaining probabilities are renormalized to sum to 1. Use T=1 for images. Set any T>1 for videos (video sampling not yet supported).
+
+Behavioral notes:
+
+- If the requested base item count cannot be satisfied under the provided per-prompt limits, the tool raises an error rather than silently clamping.
+
+How sampling works:
+
+- Determine per-request item count k by sampling uniformly from the integer range defined by `--random-mm-base-items-per-request` and `--random-mm-num-mm-items-range-ratio`, then clamp k to at most the sum of per-modality limits.
+- For each of the k items, sample a bucket (H, W, T) according to the normalized probabilities in `--random-mm-bucket-config`, while tracking how many items of each modality have been added.
+- If a modality (e.g., image) reaches its limit from `--random-mm-limit-mm-per-prompt`, all buckets of that modality are excluded and the remaining bucket probabilities are renormalized before continuing.
+This should be seen as an edge case, and if this behavior can be avoided by setting `--random-mm-limit-mm-per-prompt` to a large number. Note that this might result in errors due to engine config `--limit-mm-per-prompt`.
+- The resulting request contains synthetic image data in `multi_modal_data` (OpenAI Chat format). When `random-mm` is used with the OpenAI Chat backend, prompts remain text and MM content is attached via `multi_modal_data`.
+
+</details>

benchmarks/auto_tune/README.md

@@ -149,70 +149,3 @@ The script follows a systematic process to find the optimal parameters:
 4. **Track Best Result**: Throughout the process, the script tracks the parameter combination that has yielded the highest valid throughput so far.
 
 5. **Profile Collection**: For the best-performing run, the script saves the vLLM profiler output, which can be used for deep-dive performance analysis with tools like TensorBoard.
-
-## Batched `auto_tune`
-
-The `batch_auto_tune.sh` script allows you to run multiple `auto_tune.sh` experiments sequentially from a single configuration file. It iterates through a list of parameter sets, executes `auto_tune.sh` for each, and records the results back into the input file.
-
-### Prerequisites
-
-- **jq**: This script requires `jq` to parse the JSON configuration file.
-- **gcloud**: If you plan to upload results to Google Cloud Storage, the `gcloud` CLI must be installed and authenticated.
-
-### How to Run
-
-1. **Create a JSON configuration file**: Create a file (e.g., `runs_config.json`) containing an array of JSON objects. Each object defines the parameters for a single `auto_tune.sh` run.
-
-2. **Execute the script**:
-
-    ```bash
-    bash batch_auto_tune.sh <path_to_json_file> [gcs_upload_path]
-    ```
-
-    - `<path_to_json_file>`: **Required.** Path to your JSON configuration file.
-    - `[gcs_upload_path]`: **Optional.** A GCS path (e.g., `gs://my-bucket/benchmark-results`) where the detailed results and profiles for each run will be uploaded. If this is empty, the results will be available on the local filesystem (see the log for `RESULT_FILE=/path/to/results/file.txt`).
-
-### Configuration File
-
-The JSON configuration file should contain an array of objects. Each object's keys correspond to the configuration variables for `auto_tune.sh` (see the [Configuration table above](#configuration)). These keys will be converted to uppercase environment variables for each run.
-
-Here is an example `runs_config.json` with two benchmark configurations:
-
-```json
-[
-  {
-    "base": "/home/user",
-    "model": "meta-llama/Llama-3.1-8B-Instruct",
-    "system": "TPU", # OR GPU
-    "tp": 8,
-    "input_len": 128,
-    "output_len": 2048,
-    "max_model_len": 2300,
-    "num_seqs_list": "128 256",
-    "num_batched_tokens_list": "8192 16384"
-  },
-  {
-    "base": "/home/user",
-    "model": "meta-llama/Llama-3.1-70B-Instruct",
-    "system": "TPU", # OR GPU
-    "tp": 8,
-    "input_len": 4000,
-    "output_len": 16,
-    "max_model_len": 4096,
-    "num_seqs_list": "64 128",
-    "num_batched_tokens_list": "4096 8192",
-    "max_latency_allowed_ms": 500
-  }
-]
-```
-
-### Output
-
-The script modifies the input JSON file in place, adding the results of each run to the corresponding object. The following fields are added:
-
-- `run_id`: A unique identifier for the run, derived from the timestamp.
-- `status`: The outcome of the run (`SUCCESS`, `FAILURE`, or `WARNING_NO_RESULT_FILE`).
-- `results`: The content of the `result.txt` file from the `auto_tune.sh` run.
-- `gcs_results`: The GCS URL where the run's artifacts are stored (if a GCS path was provided).
-
-A summary of successful and failed runs is also printed to the console upon completion.

benchmarks/auto_tune/auto_tune.sh

@@ -103,15 +103,10 @@ start_server() {
         VLLM_USE_V1=1 VLLM_SERVER_DEV_MODE=1 \
             vllm serve "${common_args_array[@]}" > "$vllm_log" 2>&1 &
     fi
-    local server_pid=$!
 
     # wait for 10 minutes...
     server_started=0
     for i in {1..60}; do
-        # This line checks whether the server is still alive or not,
-        # since that we should always have permission to send signal to the server process.
-        kill -0 $server_pid 2> /dev/null || break
-
         RESPONSE=$(curl -s -X GET "http://0.0.0.0:8004/health" -w "%{http_code}" -o /dev/stdout)
         STATUS_CODE=$(echo "$RESPONSE" | tail -n 1)
         if [[ "$STATUS_CODE" -eq 200 ]]; then
@@ -123,7 +118,7 @@ start_server() {
     done
 
     if (( ! server_started )); then
-        echo "server did not start within 10 minutes or crashed. Please check server log at $vllm_log".
+        echo "server did not start within 10 minutes. Please check server log at $vllm_log".
         return 1
     else
         return 0

benchmarks/auto_tune/batch_auto_tune.sh

@@ -1,128 +0,0 @@
-#!/bin/bash
-
-INPUT_JSON="$1"
-GCS_PATH="$2" # Optional GCS path for uploading results for each run
-
-SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)
-AUTOTUNE_SCRIPT="$SCRIPT_DIR/auto_tune.sh"
-
-if [[ -z "$INPUT_JSON" ]]; then
-  echo "Error: Input JSON file not provided."
-  echo "Usage: $0 <path_to_json_file> [gcs_upload_path]"
-  exit 1
-fi
-
-if [[ ! -f "$INPUT_JSON" ]]; then
-  echo "Error: File not found at '$INPUT_JSON'"
-  exit 1
-fi
-
-if ! command -v jq &> /dev/null; then
-    echo "Error: 'jq' command not found. Please install jq to process the JSON input."
-    exit 1
-fi
-
-if [[ -n "$GCS_PATH" ]] && ! command -v gcloud &> /dev/null; then
-    echo "Error: 'gcloud' command not found, but a GCS_PATH was provided."
-    exit 1
-fi
-
-SUCCESS_COUNT=0
-FAILURE_COUNT=0
-FAILED_RUNS=()
-SCRIPT_START_TIME=$(date +%s)
-
-json_content=$(cat "$INPUT_JSON")
-if ! num_runs=$(echo "$json_content" | jq 'length'); then
-  echo "Error: Invalid JSON in $INPUT_JSON. 'jq' failed to get array length." >&2
-  exit 1
-fi
-
-echo "Found $num_runs benchmark configurations in $INPUT_JSON."
-echo "Starting benchmark runs..."
-echo "--------------------------------------------------"
-
-for i in $(seq 0 $(($num_runs - 1))); do
-  run_object=$(echo "$json_content" | jq ".[$i]")
-
-  RUN_START_TIME=$(date +%s)
-  ENV_VARS_ARRAY=()
-  # Dynamically create env vars from the JSON object's keys
-  for key in $(echo "$run_object" | jq -r 'keys_unsorted[]'); do
-    value=$(echo "$run_object" | jq -r ".$key")
-    var_name=$(echo "$key" | tr '[:lower:]' '[:upper:]' | tr -cd 'A-Z0-9_')
-    ENV_VARS_ARRAY+=("${var_name}=${value}")
-  done
-
-  echo "Executing run #$((i+1))/$num_runs with parameters: ${ENV_VARS_ARRAY[*]}"
-
-  # Execute auto_tune.sh and capture output
-  RUN_OUTPUT_FILE=$(mktemp)
-  if env "${ENV_VARS_ARRAY[@]}" bash "$AUTOTUNE_SCRIPT" > >(tee -a "$RUN_OUTPUT_FILE") 2>&1; then
-    STATUS="SUCCESS"
-    ((SUCCESS_COUNT++))
-  else
-    STATUS="FAILURE"
-    ((FAILURE_COUNT++))
-    FAILED_RUNS+=("Run #$((i+1)): $(echo $run_object | jq -c .)")
-  fi
-
-  RUN_OUTPUT=$(<"$RUN_OUTPUT_FILE")
-  rm "$RUN_OUTPUT_FILE"
-
-  # Parse results and optionally upload them to GCS
-  RUN_ID=""
-  RESULTS=""
-  GCS_RESULTS_URL=""
-  if [[ "$STATUS" == "SUCCESS" ]]; then
-    RESULT_FILE_PATH=$(echo "$RUN_OUTPUT" | grep 'RESULT_FILE=' | tail -n 1 | cut -d'=' -f2 | tr -s '/' || true)
-
-    if [[ -n "$RESULT_FILE_PATH" && -f "$RESULT_FILE_PATH" ]]; then
-      RUN_ID=$(basename "$(dirname "$RESULT_FILE_PATH")")
-      RESULT_DIR=$(dirname "$RESULT_FILE_PATH")
-      RESULTS=$(cat "$RESULT_FILE_PATH")
-
-      if [[ -n "$GCS_PATH" ]]; then
-        GCS_RESULTS_URL="${GCS_PATH}/${RUN_ID}"
-        echo "Uploading results to GCS..."
-        if gcloud storage rsync --recursive "$RESULT_DIR/" "$GCS_RESULTS_URL"; then
-          echo "GCS upload successful."
-        else
-          echo "Warning: GCS upload failed for RUN_ID $RUN_ID."
-        fi
-      fi
-    else
-      echo "Warning: Could not find result file for a successful run."
-      STATUS="WARNING_NO_RESULT_FILE"
-    fi
-  fi
-
-  # Add the results back into the JSON object for this run
-  json_content=$(echo "$json_content" | jq --argjson i "$i" --arg run_id "$RUN_ID" --arg status "$STATUS" --arg results "$RESULTS" --arg gcs_results "$GCS_RESULTS_URL" \
-    '.[$i] += {run_id: $run_id, status: $status, results: $results, gcs_results: $gcs_results}')
-
-  RUN_END_TIME=$(date +%s)
-  echo "Run finished in $((RUN_END_TIME - RUN_START_TIME)) seconds. Status: $STATUS"
-  echo "--------------------------------------------------"
-
-  # Save intermediate progress back to the file
-  echo "$json_content" > "$INPUT_JSON.tmp" && mv "$INPUT_JSON.tmp" "$INPUT_JSON"
-
-done
-
-SCRIPT_END_TIME=$(date +%s)
-echo "All benchmark runs completed in $((SCRIPT_END_TIME - SCRIPT_START_TIME)) seconds."
-echo
-echo "====================== SUMMARY ======================"
-echo "Successful runs: $SUCCESS_COUNT"
-echo "Failed runs:     $FAILURE_COUNT"
-echo "==================================================="
-
-if [[ $FAILURE_COUNT -gt 0 ]]; then
-  echo "Details of failed runs (see JSON file for full parameters):"
-  for failed in "${FAILED_RUNS[@]}"; do
-    echo "  - $failed"
-  done
-fi
-
-echo "Updated results have been saved to '$INPUT_JSON'."

benchmarks/benchmark_ngram_proposer.py

@@ -1,31 +1,17 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import gc
-import time
-from unittest import mock
 
 import numpy as np
 from tabulate import tabulate
 
 from benchmark_utils import TimeCollector
-from vllm.config import (
-    CacheConfig,
-    DeviceConfig,
-    LoadConfig,
-    ModelConfig,
-    ParallelConfig,
-    SchedulerConfig,
-    SpeculativeConfig,
-    VllmConfig,
-)
-from vllm.platforms import current_platform
+from vllm.config import ModelConfig, SpeculativeConfig, VllmConfig
 from vllm.utils import FlexibleArgumentParser
 from vllm.v1.spec_decode.ngram_proposer import NgramProposer
-from vllm.v1.worker.gpu_input_batch import InputBatch
-from vllm.v1.worker.gpu_model_runner import GPUModelRunner
 
 
-def benchmark_propose(args):
+def main(args):
     rows = []
     for max_ngram in args.max_ngram:
         collector = TimeCollector(TimeCollector.US)
@@ -83,88 +69,10 @@ def benchmark_propose(args):
     )
 
 
-def benchmark_batched_propose(args):
-    NUM_SPECULATIVE_TOKENS_NGRAM = 10
-    PROMPT_LOOKUP_MIN = 5
-    PROMPT_LOOKUP_MAX = 15
-    MAX_MODEL_LEN = int(1e7)
-    DEVICE = current_platform.device_type
-
-    model_config = ModelConfig(model="facebook/opt-125m", runner="generate")
-
-    speculative_config = SpeculativeConfig(
-        target_model_config=model_config,
-        target_parallel_config=ParallelConfig(),
-        method="ngram",
-        num_speculative_tokens=NUM_SPECULATIVE_TOKENS_NGRAM,
-        prompt_lookup_max=PROMPT_LOOKUP_MAX,
-        prompt_lookup_min=PROMPT_LOOKUP_MIN,
-    )
-
-    vllm_config = VllmConfig(
-        model_config=model_config,
-        cache_config=CacheConfig(),
-        speculative_config=speculative_config,
-        device_config=DeviceConfig(device=current_platform.device_type),
-        parallel_config=ParallelConfig(),
-        load_config=LoadConfig(),
-        scheduler_config=SchedulerConfig(),
-    )
-
-    # monkey patch vllm.v1.worker.gpu_model_runner.get_pp_group
-    mock_pp_group = mock.MagicMock()
-    mock_pp_group.world_size = 1
-    with mock.patch(
-        "vllm.v1.worker.gpu_model_runner.get_pp_group", return_value=mock_pp_group
-    ):
-        runner = GPUModelRunner(vllm_config, DEVICE)
-
-        # hack max model len
-        runner.max_model_len = MAX_MODEL_LEN
-        runner.drafter.max_model_len = MAX_MODEL_LEN
-
-        dummy_input_batch = InputBatch(
-            max_num_reqs=args.num_req,
-            max_model_len=MAX_MODEL_LEN,
-            max_num_batched_tokens=args.num_req * args.num_token,
-            device=DEVICE,
-            pin_memory=False,
-            vocab_size=256000,
-            block_sizes=[16],
-        )
-        dummy_input_batch._req_ids = list(str(id) for id in range(args.num_req))
-        dummy_input_batch.spec_decode_unsupported_reqs = ()
-        dummy_input_batch.num_tokens_no_spec = [args.num_token] * args.num_req
-        dummy_input_batch.token_ids_cpu = np.random.randint(
-            0, 20, (args.num_req, args.num_token)
-        )
-
-        runner.input_batch = dummy_input_batch
-
-        sampled_token_ids = [[0]] * args.num_req
-
-        print("Starting benchmark")
-        # first run is warmup so ignore it
-        for _ in range(args.num_iteration):
-            start = time.time()
-            runner.drafter.propose(
-                sampled_token_ids,
-                dummy_input_batch.req_ids,
-                dummy_input_batch.num_tokens_no_spec,
-                dummy_input_batch.token_ids_cpu,
-                dummy_input_batch.spec_decode_unsupported_reqs,
-            )
-            end = time.time()
-            print(f"Iteration time (s): {end - start}")
-
-
 def invoke_main() -> None:
     parser = FlexibleArgumentParser(
         description="Benchmark the performance of N-gram speculative decode drafting"
     )
-    parser.add_argument(
-        "--batched", action="store_true", help="consider time to prepare batch"
-    )  # noqa: E501
     parser.add_argument(
         "--num-iteration",
         type=int,
@@ -197,17 +105,8 @@ def invoke_main() -> None:
         help="Number of speculative tokens to generate",
     )
     args = parser.parse_args()
-
-    if not args.batched:
-        benchmark_propose(args)
-    else:
-        benchmark_batched_propose(args)
+    main(args)
 
 
-"""
-# Example command lines:
-# time python3 benchmarks/benchmark_ngram_proposer.py
-# time python3 benchmarks/benchmark_ngram_proposer.py --batched --num-iteration 4 --num-token 1000000 --num-req 128
-"""  # noqa: E501
 if __name__ == "__main__":
     invoke_main()  # pragma: no cover

benchmarks/benchmark_serving_structured_output.py

@@ -449,8 +449,7 @@ async def benchmark(
     def prepare_extra_body(request) -> dict:
         extra_body = {}
         # Add the schema to the extra_body
-        extra_body["structured_outputs"] = {}
-        extra_body["structured_outputs"][request.structure_type] = request.schema
+        extra_body[request.structure_type] = request.schema
         return extra_body
 
     print("Starting initial single prompt test run...")
@@ -697,11 +696,11 @@ def evaluate(ret, args):
         return re.match(args.regex, actual) is not None
 
     def _eval_correctness(expected, actual):
-        if args.structure_type == "json":
+        if args.structure_type == "guided_json":
             return _eval_correctness_json(expected, actual)
-        elif args.structure_type == "regex":
+        elif args.structure_type == "guided_regex":
             return _eval_correctness_regex(expected, actual)
-        elif args.structure_type == "choice":
+        elif args.structure_type == "guided_choice":
             return _eval_correctness_choice(expected, actual)
         else:
             return None
@@ -781,18 +780,18 @@ def main(args: argparse.Namespace):
     )
 
     if args.dataset == "grammar":
-        args.structure_type = "grammar"
+        args.structure_type = "guided_grammar"
     elif args.dataset == "regex":
-        args.structure_type = "regex"
+        args.structure_type = "guided_regex"
     elif args.dataset == "choice":
-        args.structure_type = "choice"
+        args.structure_type = "guided_choice"
     else:
-        args.structure_type = "json"
+        args.structure_type = "guided_json"
 
     if args.no_structured_output:
         args.structured_output_ratio = 0
     if args.save_results:
-        result_file_name = f"{args.structured_output_ratio}so"
+        result_file_name = f"{args.structured_output_ratio}guided"
         result_file_name += f"_{backend}"
         result_file_name += f"_{args.request_rate}qps"
         result_file_name += f"_{args.model.split('/')[-1]}"

benchmarks/kernels/bench_nvfp4_gemm.py

@@ -3,7 +3,6 @@
 import argparse
 import copy
 import itertools
-import os
 
 import torch
 from weight_shapes import WEIGHT_SHAPES
@@ -24,45 +23,21 @@ PROVIDER_CFGS = {
     "torch-bf16": dict(enabled=True),
     "nvfp4": dict(no_a_quant=False, enabled=True),
     "nvfp4-noquant": dict(no_a_quant=True, enabled=True),
-    "fbgemm-nvfp4": dict(fbgemm=True, no_a_quant=False, enabled=True),
-    "fbgemm-nvfp4-noquant": dict(fbgemm=True, no_a_quant=True, enabled=True),
 }
 
-_needs_fbgemm = any(
-    v.get("fbgemm", False) for v in PROVIDER_CFGS.values() if v.get("enabled", False)
-)
-if _needs_fbgemm:
-    try:
-        from fbgemm_gpu.experimental.gemm.triton_gemm.fp4_quantize import (
-            triton_scale_nvfp4_quant,
-        )
-    except ImportError:
-        print(
-            "WARNING: FBGEMM providers are enabled but fbgemm_gpu is not installed. "
-            "These providers will be skipped. Please install fbgemm_gpu with: "
-            "'pip install fbgemm-gpu-genai' to run them."
-        )
-        # Disable FBGEMM providers so the benchmark can run.
-        for cfg in PROVIDER_CFGS.values():
-            if cfg.get("fbgemm"):
-                cfg["enabled"] = False
-
 _enabled = [k for k, v in PROVIDER_CFGS.items() if v["enabled"]]
 
 
-def _quant_weight_nvfp4(b: torch.Tensor, device: str, cfg):
+def _quant_weight_nvfp4(b: torch.Tensor, device: str):
     # Compute global scale for weight
     b_amax = torch.abs(b).max().to(torch.float32)
     b_global_scale = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / b_amax
-    if "fbgemm" in cfg and cfg["fbgemm"]:
-        b_fp4, scale_b_fp4 = triton_scale_nvfp4_quant(b, b_global_scale)
-    else:
-        b_fp4, scale_b_fp4 = ops.scaled_fp4_quant(b, b_global_scale)
+    b_fp4, scale_b_fp4 = ops.scaled_fp4_quant(b, b_global_scale)
     return b_fp4, scale_b_fp4, b_global_scale
 
 
 def build_nvfp4_runner(cfg, a, b, dtype, device):
-    b_fp4, scale_b_fp4, b_global_scale = _quant_weight_nvfp4(b, device, cfg)
+    b_fp4, scale_b_fp4, b_global_scale = _quant_weight_nvfp4(b, device)
 
     # Compute global scale for activation
     # NOTE: This is generally provided ahead-of-time by the model checkpoint.
@@ -71,35 +46,6 @@ def build_nvfp4_runner(cfg, a, b, dtype, device):
 
     # Alpha for the GEMM operation
     alpha = 1.0 / (a_global_scale * b_global_scale)
-    if "fbgemm" in cfg and cfg["fbgemm"]:
-        if cfg["no_a_quant"]:
-            a_fp4, scale_a_fp4 = triton_scale_nvfp4_quant(a, a_global_scale)
-
-            def run():
-                return torch.ops.fbgemm.f4f4bf16(
-                    a_fp4,
-                    b_fp4,
-                    scale_a_fp4,
-                    scale_b_fp4,
-                    global_scale=alpha,
-                    use_mx=False,
-                )
-
-            return run
-        else:
-
-            def run():
-                a_fp4, scale_a_fp4 = triton_scale_nvfp4_quant(a, a_global_scale)
-                return torch.ops.fbgemm.f4f4bf16(
-                    a_fp4,
-                    b_fp4,
-                    scale_a_fp4,
-                    scale_b_fp4,
-                    global_scale=alpha,
-                    use_mx=False,
-                )
-
-            return run
 
     if cfg["no_a_quant"]:
         # Pre-quantize activation
@@ -184,13 +130,10 @@ if __name__ == "__main__":
 
     for K, N, model in prepare_shapes(args):
         print(f"{model}, N={N} K={K}, BF16 vs NVFP4 GEMMs TFLOP/s:")
-        save_dir = f"bench_nvfp4_res_n{N}_k{K}"
-        os.makedirs(save_dir, exist_ok=True)
-
         benchmark.run(
             print_data=True,
             show_plots=True,
-            save_path=save_dir,
+            save_path=f"bench_nvfp4_res_n{N}_k{K}",
             N=N,
             K=K,
         )

benchmarks/kernels/bench_per_token_quant_fp8.py

@@ -2,25 +2,14 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import itertools
 from typing import Callable
-from unittest.mock import patch
 
-import pandas as pd
 import torch
 
+from vllm import _custom_ops as ops
+from vllm.config import CompilationConfig, VllmConfig, set_current_vllm_config
 from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
 from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
 from vllm.triton_utils import triton
-from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE, FlexibleArgumentParser
-
-
-def with_triton_mode(fn):
-    """Temporarily force the Triton fallback path"""
-
-    def wrapped(*args, **kwargs):
-        with patch("vllm.platforms.current_platform.is_cuda", return_value=False):
-            return fn(*args, **kwargs)
-
-    return wrapped
 
 
 # TODO(luka): use standalone_compile utility
@@ -32,238 +21,78 @@ def with_dyn_arg(fn: Callable, arg_index: int, dim_index: int):
     return inner
 
 
-def bench_compile(fn: Callable):
-    # recompile for different shapes
-    fwd = torch.compile(fn, fullgraph=True, dynamic=False)
+torch._dynamo.config.recompile_limit = 8888
+compilation_config = CompilationConfig(custom_ops=["none"])
+with set_current_vllm_config(VllmConfig(compilation_config=compilation_config)):
+    torch_per_token_quant_fp8 = torch.compile(
+        QuantFP8(False, GroupShape.PER_TOKEN),
+        fullgraph=True,
+        dynamic=False,  # recompile for different shapes
+    )
 
     # First dim is explicitly dynamic to simulate vLLM usage
-    return with_dyn_arg(fwd, 0, 0)
+    torch_per_token_quant_fp8 = with_dyn_arg(torch_per_token_quant_fp8, 0, 0)
 
 
-torch._dynamo.config.recompile_limit = 8888
+def cuda_per_token_quant_fp8(
+    input: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    return ops.scaled_fp8_quant(input)
 
 
-def calculate_diff(
-    batch_size: int,
-    hidden_size: int,
-    group_shape: GroupShape,
-    dtype: torch.dtype,
-):
-    """Calculate the difference between Inductor and CUDA implementations."""
+def calculate_diff(batch_size: int, seq_len: int):
+    """Calculate difference between Triton and CUDA implementations."""
     device = torch.device("cuda")
-    x = torch.randn((batch_size, hidden_size), dtype=dtype, device=device)
+    x = torch.rand((batch_size * seq_len, 4096), dtype=torch.float16, device=device)
 
-    quant_fp8 = QuantFP8(False, group_shape, column_major_scales=False)
+    torch_out, torch_scale = torch_per_token_quant_fp8(x)
+    cuda_out, cuda_scale = cuda_per_token_quant_fp8(x)
 
-    torch_out, torch_scale = bench_compile(quant_fp8.forward_native)(x)
-    torch_eager_out, torch_eager_scale = quant_fp8.forward_native(x)
-    cuda_out, cuda_scale = quant_fp8.forward_cuda(x)
-
-    try:
-        torch.testing.assert_close(
-            cuda_out.to(torch.float32),
-            torch_out.to(torch.float32),
-            rtol=1e-3,
-            atol=1e-5,
-        )
-        torch.testing.assert_close(cuda_scale, torch_scale, rtol=1e-3, atol=1e-5)
-        torch.testing.assert_close(
-            cuda_out.to(torch.float32),
-            torch_eager_out.to(torch.float32),
-            rtol=1e-3,
-            atol=1e-5,
-        )
-        torch.testing.assert_close(cuda_scale, torch_eager_scale, rtol=1e-3, atol=1e-5)
+    if torch.allclose(
+        cuda_out.to(torch.float32), torch_out.to(torch.float32), rtol=1e-3, atol=1e-5
+    ) and torch.allclose(cuda_scale, torch_scale, rtol=1e-3, atol=1e-5):
         print("✅ All implementations match")
-    except AssertionError as e:
+    else:
         print("❌ Implementations differ")
-        print(e)
 
 
-configs = []
+batch_size_range = [1, 16, 32, 64, 128]
+seq_len_range = [1, 16, 64, 128, 256, 512, 1024, 2048, 4096]
+
+configs = list(itertools.product(batch_size_range, seq_len_range))
 
 
-def benchmark_quantization(
-    batch_size,
-    hidden_size,
-    provider,
-    group_shape: GroupShape,
-    col_major: bool,
-    dtype: torch.dtype,
-):
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["batch_size", "seq_len"],
+        x_vals=configs,
+        line_arg="provider",
+        line_vals=["torch", "cuda"],
+        line_names=["Torch", "CUDA"],
+        styles=[("blue", "-"), ("green", "-")],
+        ylabel="us",
+        plot_name="per-token-dynamic-quant-fp8-performance",
+        args={},
+    )
+)
+def benchmark_quantization(batch_size, seq_len, provider):
+    dtype = torch.float16
     device = torch.device("cuda")
 
-    x = torch.randn(batch_size, hidden_size, device=device, dtype=dtype)
+    x = torch.randn(batch_size * seq_len, 4096, device=device, dtype=dtype)
 
     quantiles = [0.5, 0.2, 0.8]
-    quant_fp8 = QuantFP8(False, group_shape, column_major_scales=col_major)
 
     if provider == "torch":
-        fn = lambda: bench_compile(quant_fp8.forward_native)(x.clone())
+        fn = lambda: torch_per_token_quant_fp8(x.clone())
     elif provider == "cuda":
-        fn = lambda: quant_fp8.forward_cuda(x.clone())
-    elif provider == "triton":
-        if not group_shape.is_per_group():
-            # Triton only supported for per-group
-            return 0, 0, 0
-
-        fn = lambda: with_triton_mode(quant_fp8.forward_cuda)(x.clone())
+        fn = lambda: cuda_per_token_quant_fp8(x.clone())
 
     ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)
 
     return 1000 * ms, 1000 * max_ms, 1000 * min_ms
 
 
-# TODO(luka) extract to utils
-def compute_geomean_speedups(
-    df: pd.DataFrame,
-    baseline_col: str,
-    speedup_cols: list[str],
-    groupby_cols: list[str] | None = None,
-) -> pd.DataFrame:
-    """
-    Compute geometric mean speedups over a baseline column.
-
-    Args:
-        df: Input dataframe
-        baseline_col: Column to use as baseline
-        speedup_cols: Columns to compute speedups for
-        groupby_cols: Columns to group by. If None, compute over entire df.
-
-    Returns:
-        pd.DataFrame with geometric mean speedups
-    """
-    from scipy.stats import gmean
-
-    def geo_speedup(group: pd.DataFrame) -> pd.Series:
-        ratios = {
-            col: (group[baseline_col] / group[col]).values for col in speedup_cols
-        }
-        return pd.Series({col: gmean(vals) for col, vals in ratios.items()})
-
-    if groupby_cols is None:
-        result = geo_speedup(df).to_frame().T
-    else:
-        result = (
-            df.groupby(groupby_cols)
-            .apply(geo_speedup, include_groups=False)
-            .reset_index()
-        )
-
-    return result
-
-
 if __name__ == "__main__":
-    parser = FlexibleArgumentParser(
-        description="Benchmark the various implementations of QuantFP8 (dynamic-only)"
-    )
-    parser.add_argument("-c", "--check", action="store_true")
-    parser.add_argument(
-        "--dtype", type=str, choices=["half", "bfloat16", "float"], default="bfloat16"
-    )
-    parser.add_argument(
-        "--hidden-sizes",
-        type=int,
-        nargs="+",
-        default=[896, 1024, 2048, 4096, 7168],
-        help="Hidden sizes to benchmark",
-    )
-    parser.add_argument(
-        "--batch-sizes",
-        type=int,
-        nargs="+",
-        default=[1, 16, 128, 512, 1024],
-        help="Batch sizes to benchmark",
-    )
-    parser.add_argument(
-        "--group-sizes",
-        type=int,
-        nargs="+",
-        default=None,
-        help="Group sizes for GroupShape(1,N) to benchmark. "
-        "Use 0 for PER_TENSOR, -1 for PER_TOKEN (default: 0,-1,64,128)",
-    )
-    parser.add_argument(
-        "--no-column-major",
-        action="store_true",
-        help="Disable column-major scales testing",
-    )
-
-    args = parser.parse_args()
-    assert args
-
-    dtype = STR_DTYPE_TO_TORCH_DTYPE[args.dtype]
-
-    hidden_sizes = args.hidden_sizes
-    batch_sizes = args.batch_sizes
-
-    if args.group_sizes is not None:
-        group_shapes = []
-        for size in args.group_sizes:
-            if size == 0:
-                group_shapes.append(GroupShape.PER_TENSOR)
-            elif size == -1:
-                group_shapes.append(GroupShape.PER_TOKEN)
-            else:
-                group_shapes.append(GroupShape(1, size))
-    else:
-        group_shapes = [
-            GroupShape.PER_TENSOR,
-            GroupShape.PER_TOKEN,
-            GroupShape(1, 64),
-            GroupShape(1, 128),
-        ]
-
-    column_major_scales = [False] if args.no_column_major else [True, False]
-
-    config_gen = itertools.product(
-        group_shapes,
-        column_major_scales,
-        batch_sizes,
-        hidden_sizes,
-    )
-
-    # filter out column-major scales for non-group, reverse order
-    configs.extend(c[::-1] for c in config_gen if (c[0].is_per_group() or not c[1]))
-
-    print(f"Running {len(configs)} configurations:")
-    print(f"  Hidden sizes: {hidden_sizes}")
-    print(f"  Batch sizes: {batch_sizes}")
-    print(f"  Group shapes: {[str(g) for g in group_shapes]}")
-    print(f"  Column major scales: {column_major_scales}")
-    print()
-
-    if args.check:
-        for group_shape in group_shapes:
-            group_size = group_shape[1]
-            print(f"{group_size=}")
-            calculate_diff(
-                batch_size=4, hidden_size=4096, group_shape=group_shape, dtype=dtype
-            )
-
-    benchmark = triton.testing.perf_report(
-        triton.testing.Benchmark(
-            x_names=["hidden_size", "batch_size", "col_major", "group_shape"],
-            x_vals=configs,
-            line_arg="provider",
-            line_vals=["torch", "cuda", "triton"],
-            line_names=["Torch (Compiled)", "CUDA", "Triton"],
-            styles=[("blue", "-"), ("green", "-"), ("black", "-")],
-            ylabel="us",
-            plot_name="QuantFP8 performance",
-            args={},
-        )
-    )(benchmark_quantization)
-
-    df = benchmark.run(print_data=True, dtype=dtype, return_df=True)
-
-    # Print geomean speedups
-    geo_table_grouped = compute_geomean_speedups(
-        df,
-        baseline_col="Torch (Compiled)",
-        speedup_cols=["CUDA", "Triton"],
-        groupby_cols=["col_major", "group_shape"],
-    )
-
-    print("Speedup over Torch (Compiled)")
-    print(geo_table_grouped.to_string(index=False))
+    calculate_diff(batch_size=4, seq_len=4096)
+    benchmark_quantization.run(print_data=True)

benchmarks/kernels/benchmark_cutlass_fp4_moe.py

@@ -13,10 +13,6 @@ import torch.utils.benchmark as benchmark
 
 from vllm import _custom_ops as ops
 from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
-from vllm.model_executor.layers.fused_moe.config import (
-    fp8_w8a8_moe_quant_config,
-    nvfp4_moe_quant_config,
-)
 from vllm.model_executor.layers.fused_moe.cutlass_moe import cutlass_moe_fp4
 from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts, fused_topk
 from vllm.scalar_type import scalar_types
@@ -144,12 +140,6 @@ def bench_run(
         a_fp8_scale: torch.Tensor,
         num_repeats: int,
     ):
-        quant_config = fp8_w8a8_moe_quant_config(
-            w1_scale=w1_scale,
-            w2_scale=w2_scale,
-            a1_scale=a_fp8_scale,
-        )
-
         for _ in range(num_repeats):
             fused_experts(
                 a,
@@ -157,7 +147,10 @@ def bench_run(
                 w2,
                 topk_weights,
                 topk_ids,
-                quant_config=quant_config,
+                use_fp8_w8a8=True,
+                w1_scale=w1_scale,
+                w2_scale=w2_scale,
+                a1_scale=a_fp8_scale,
             )
 
     def run_cutlass_moe_fp4(
@@ -179,27 +172,25 @@ def bench_run(
         device: torch.device,
         num_repeats: int,
     ):
-        quant_config = nvfp4_moe_quant_config(
-            a1_gscale=a1_gs,
-            a2_gscale=a2_gs,
-            w1_scale=w1_blockscale,
-            w2_scale=w2_blockscale,
-            g1_alphas=w1_gs,
-            g2_alphas=w2_gs,
-        )
         for _ in range(num_repeats):
             with nvtx.annotate("cutlass_moe_fp4", color="green"):
                 cutlass_moe_fp4(
                     a=a,
+                    a1_gscale=a1_gs,
+                    a2_gscale=a2_gs,
                     w1_fp4=w1_fp4,
+                    w1_blockscale=w1_blockscale,
+                    w1_alphas=w1_gs,
                     w2_fp4=w2_fp4,
+                    w2_blockscale=w2_blockscale,
+                    w2_alphas=w2_gs,
                     topk_weights=topk_weights,
                     topk_ids=topk_ids,
                     m=m,
                     n=n,
                     k=k,
                     e=num_experts,
-                    quant_config=quant_config,
+                    device=device,
                 )
 
     def run_cutlass_from_graph(
@@ -220,29 +211,26 @@ def bench_run(
         e: int,
         device: torch.device,
     ):
-        quant_config = nvfp4_moe_quant_config(
-            a1_gscale=a1_gs,
-            a2_gscale=a2_gs,
-            w1_scale=w1_blockscale,
-            w2_scale=w2_blockscale,
-            g1_alphas=w1_gs,
-            g2_alphas=w2_gs,
-        )
-
         with set_current_vllm_config(
             VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1))
         ):
             return cutlass_moe_fp4(
                 a=a,
+                a1_gscale=a1_gs,
                 w1_fp4=w1_fp4,
+                w1_blockscale=w1_blockscale,
+                w1_alphas=w1_alphas,
+                a2_gscale=a2_gs,
                 w2_fp4=w2_fp4,
+                w2_blockscale=w2_blockscale,
+                w2_alphas=w2_alphas,
                 topk_weights=topk_weights,
                 topk_ids=topk_ids,
                 m=m,
                 n=n,
                 k=k,
                 e=num_experts,
-                quant_config=quant_config,
+                device=device,
             )
 
     def run_triton_from_graph(
@@ -258,18 +246,16 @@ def bench_run(
         with set_current_vllm_config(
             VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1))
         ):
-            quant_config = fp8_w8a8_moe_quant_config(
-                w1_scale=w1_scale,
-                w2_scale=w2_scale,
-                a1_scale=a_fp8_scale,
-            )
             return fused_experts(
                 a,
                 w1,
                 w2,
                 topk_weights,
                 topk_ids,
-                quant_config=quant_config,
+                use_fp8_w8a8=True,
+                w1_scale=w1_scale,
+                w2_scale=w2_scale,
+                a1_scale=a_fp8_scale,
             )
 
     def replay_graph(graph, num_repeats):

benchmarks/kernels/benchmark_cutlass_moe_fp8.py

@@ -1,406 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Benchmark the performance of the cutlass_moe_fp8 kernel vs the triton_moe
-kernel. Both kernels take in fp8 quantized weights and 16-bit activations,
-but use different quantization strategies and backends.
-"""
-
-import nvtx
-import torch
-
-from vllm import _custom_ops as ops
-from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_config
-from vllm.model_executor.layers.fused_moe.cutlass_moe import cutlass_moe_fp8
-from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts, fused_topk
-from vllm.platforms import current_platform
-from vllm.utils import FlexibleArgumentParser
-
-# Weight shapes for different models: [num_experts, topk, hidden_size,
-# intermediate_size]
-WEIGHT_SHAPES_MOE = {
-    "mixtral-8x7b": [
-        [8, 2, 4096, 14336],
-    ],
-    "deepseek-v2": [
-        [160, 6, 5120, 12288],
-    ],
-    "custom-small": [
-        [8, 2, 2048, 7168],
-    ],
-    "glm45-fp8": [
-        [128, 8, 4096, 1408],
-    ],
-    "Llama-4-Maverick-17B-128E-Instruct-FP8": [
-        [128, 1, 5120, 8192],
-    ],
-}
-
-DEFAULT_MODELS = [
-    "mixtral-8x7b",
-]
-
-DEFAULT_BATCH_SIZES = [4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048]
-DEFAULT_TP_SIZES = [1]
-
-PER_ACT_TOKEN_OPTS = [False, True]
-PER_OUT_CH_OPTS = [False, True]
-
-FP8_DTYPE = current_platform.fp8_dtype()
-
-
-def bench_run(
-    results: list,
-    model: str,
-    num_experts: int,
-    topk: int,
-    per_act_token: bool,
-    per_out_ch: bool,
-    mkn: tuple[int, int, int],
-):
-    (m, k, n) = mkn
-
-    dtype = torch.half
-    device = "cuda"
-
-    # Create input activations
-    a = torch.randn((m, k), device=device, dtype=dtype) / 10
-
-    # Create weights
-    w1 = torch.randn((num_experts, 2 * n, k), device=device, dtype=dtype) / 10
-    w2 = torch.randn((num_experts, k, n), device=device, dtype=dtype) / 10
-
-    # Create FP8 quantized weights and scales for both kernels
-    w1_fp8q = torch.empty((num_experts, 2 * n, k), device=device, dtype=FP8_DTYPE)
-    w2_fp8q = torch.empty((num_experts, k, n), device=device, dtype=FP8_DTYPE)
-
-    # Create scales based on quantization strategy
-    if per_out_ch:
-        # Per-channel quantization
-        w1_scale = torch.empty(
-            (num_experts, 2 * n, 1), device=device, dtype=torch.float32
-        )
-        w2_scale = torch.empty((num_experts, k, 1), device=device, dtype=torch.float32)
-    else:
-        # Per-tensor quantization
-        w1_scale = torch.empty((num_experts, 1, 1), device=device, dtype=torch.float32)
-        w2_scale = torch.empty((num_experts, 1, 1), device=device, dtype=torch.float32)
-
-    # Quantize weights
-    for expert in range(num_experts):
-        if per_out_ch:
-            # Per-channel quantization - not yet implemented properly
-            # For now, fall back to per-tensor quantization
-            w1_fp8q[expert], w1_scale_temp = ops.scaled_fp8_quant(w1[expert])
-            w2_fp8q[expert], w2_scale_temp = ops.scaled_fp8_quant(w2[expert])
-            # Expand scalar scales to the expected per-channel shape
-            w1_scale[expert] = w1_scale_temp.expand(2 * n, 1)
-            w2_scale[expert] = w2_scale_temp.expand(k, 1)
-        else:
-            # Per-tensor quantization
-            w1_fp8q[expert], w1_scale_temp = ops.scaled_fp8_quant(w1[expert])
-            w2_fp8q[expert], w2_scale_temp = ops.scaled_fp8_quant(w2[expert])
-            # Store scalar scales in [1, 1] tensors
-            w1_scale[expert, 0, 0] = w1_scale_temp
-            w2_scale[expert, 0, 0] = w2_scale_temp
-
-    # Prepare weights for CUTLASS (no transpose needed)
-    w1_fp8q_cutlass = w1_fp8q  # Keep original [E, 2N, K]
-    w2_fp8q_cutlass = w2_fp8q  # Keep original [E, K, N]
-
-    # Create router scores and get topk
-    score = torch.randn((m, num_experts), device=device, dtype=dtype)
-    topk_weights, topk_ids, _ = fused_topk(a, score, topk, renormalize=False)
-
-    # WORKAROUND: CUTLASS MoE FP8 has issues with per-token quantization
-    # Force per-tensor quantization for all cases to match working e2e setup
-    a1_scale = torch.full((), 1e-2, device=device, dtype=torch.float32)
-    a2_scale = torch.full((), 1e-2, device=device, dtype=torch.float32)
-
-    # Force per-tensor quantization for all cases
-    per_act_token = False
-
-    # Create stride tensors for CUTLASS
-    ab_strides1 = torch.full((num_experts,), k, dtype=torch.int64, device=device)
-    ab_strides2 = torch.full((num_experts,), n, dtype=torch.int64, device=device)
-    c_strides1 = torch.full((num_experts,), 2 * n, dtype=torch.int64, device=device)
-    c_strides2 = torch.full((num_experts,), k, dtype=torch.int64, device=device)
-
-    def run_triton_moe(
-        a: torch.Tensor,
-        w1: torch.Tensor,
-        w2: torch.Tensor,
-        topk_weights: torch.Tensor,
-        topk_ids: torch.Tensor,
-        w1_scale: torch.Tensor,
-        w2_scale: torch.Tensor,
-        a1_scale: torch.Tensor,
-        a2_scale: torch.Tensor,
-        num_repeats: int,
-    ):
-        quant_config = fp8_w8a8_moe_quant_config(
-            w1_scale=w1_scale,
-            w2_scale=w2_scale,
-            a1_scale=a1_scale,
-            a2_scale=a2_scale,
-            per_act_token_quant=per_act_token,
-            per_out_ch_quant=per_out_ch,
-        )
-
-        for _ in range(num_repeats):
-            fused_experts(
-                a,
-                w1,
-                w2,
-                topk_weights,
-                topk_ids,
-                quant_config=quant_config,
-            )
-
-    def run_cutlass_moe_fp8(
-        a: torch.Tensor,
-        w1: torch.Tensor,
-        w2: torch.Tensor,
-        topk_weights: torch.Tensor,
-        topk_ids: torch.Tensor,
-        ab_strides1: torch.Tensor,
-        ab_strides2: torch.Tensor,
-        c_strides1: torch.Tensor,
-        c_strides2: torch.Tensor,
-        w1_scale: torch.Tensor,
-        w2_scale: torch.Tensor,
-        a1_scale: torch.Tensor,
-        a2_scale: torch.Tensor,
-        num_repeats: int,
-    ):
-        quant_config = fp8_w8a8_moe_quant_config(
-            w1_scale=w1_scale,
-            w2_scale=w2_scale,
-            a1_scale=a1_scale,
-            a2_scale=a2_scale,
-            per_act_token_quant=per_act_token,
-            per_out_ch_quant=per_out_ch,
-        )
-
-        for _ in range(num_repeats):
-            with nvtx.annotate("cutlass_moe_fp8", color="blue"):
-                cutlass_moe_fp8(
-                    a=a,
-                    w1_q=w1,
-                    w2_q=w2,
-                    topk_weights=topk_weights,
-                    topk_ids=topk_ids,
-                    ab_strides1=ab_strides1,
-                    ab_strides2=ab_strides2,
-                    c_strides1=c_strides1,
-                    c_strides2=c_strides2,
-                    quant_config=quant_config,
-                    activation="silu",
-                    global_num_experts=num_experts,
-                )
-
-    # Pre-create quantization config to avoid creating it inside CUDA graph
-    quant_config = fp8_w8a8_moe_quant_config(
-        w1_scale=w1_scale,
-        w2_scale=w2_scale,
-        a1_scale=a1_scale,
-        a2_scale=a2_scale,
-        per_act_token_quant=per_act_token,
-        per_out_ch_quant=per_out_ch,
-    )
-
-    # Create CUDA graphs for CUTLASS (match benchmark_moe.py pattern exactly)
-    cutlass_stream = torch.cuda.Stream()
-    cutlass_graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(cutlass_graph, stream=cutlass_stream):
-        # Capture 10 invocations like benchmark_moe.py
-        for _ in range(10):
-            cutlass_moe_fp8(
-                a=a,
-                w1_q=w1_fp8q_cutlass,
-                w2_q=w2_fp8q_cutlass,
-                topk_weights=topk_weights,
-                topk_ids=topk_ids,
-                ab_strides1=ab_strides1,
-                ab_strides2=ab_strides2,
-                c_strides1=c_strides1,
-                c_strides2=c_strides2,
-                quant_config=quant_config,
-                activation="silu",
-                global_num_experts=num_experts,
-            )
-    torch.cuda.synchronize()
-
-    # Create CUDA graphs for Triton (match benchmark_moe.py pattern exactly)
-    triton_stream = torch.cuda.Stream()
-    triton_graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(triton_graph, stream=triton_stream):
-        # Capture 10 invocations like benchmark_moe.py
-        for _ in range(10):
-            fused_experts(
-                a,
-                w1_fp8q,
-                w2_fp8q,
-                topk_weights,
-                topk_ids,
-                quant_config=quant_config,
-            )
-    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.cuda.synchronize()
-
-        # Timing
-        start_event = torch.cuda.Event(enable_timing=True)
-        end_event = torch.cuda.Event(enable_timing=True)
-
-        latencies = []
-        for _ in range(num_iters):
-            torch.cuda.synchronize()
-            start_event.record()
-            graph.replay()
-            end_event.record()
-            end_event.synchronize()
-            latencies.append(start_event.elapsed_time(end_event))
-
-        # Divide by 10 since graph contains 10 calls
-        return sum(latencies) / (num_iters * 10)
-
-    # Benchmark parameters
-    num_warmup = 5
-    num_iters = 100
-
-    # Benchmark only CUDA graphs (more reliable and faster)
-    # Benchmark Triton MoE with CUDA graphs
-    triton_graph_time = bench_cuda_graph(
-        triton_graph, num_warmup=num_warmup, num_iters=num_iters
-    )
-
-    # Benchmark CUTLASS MoE with CUDA graphs
-    cutlass_graph_time = bench_cuda_graph(
-        cutlass_graph, num_warmup=num_warmup, num_iters=num_iters
-    )
-
-    # Convert ms to us and return results
-    triton_time_us = triton_graph_time * 1000
-    cutlass_time_us = cutlass_graph_time * 1000
-
-    return {
-        "batch_size": m,
-        "triton_time_us": triton_time_us,
-        "cutlass_time_us": cutlass_time_us,
-    }
-
-
-def main(args):
-    print("Benchmarking models:")
-    for i, model in enumerate(args.models):
-        print(f"[{i}]  {model}")
-
-    all_results = []
-
-    for model in args.models:
-        for tp in args.tp_sizes:
-            for layer in WEIGHT_SHAPES_MOE[model]:
-                num_experts = layer[0]
-                topk = layer[1]
-                size_k = layer[2]
-                size_n = layer[3] // tp
-
-                if len(args.limit_k) > 0 and size_k not in args.limit_k:
-                    continue
-
-                if len(args.limit_n) > 0 and size_n not in args.limit_n:
-                    continue
-
-                for per_act_token in args.per_act_token_opts:
-                    for per_out_ch in args.per_out_ch_opts:
-                        print(
-                            f"\n=== {model}, experts={num_experts}, topk={topk},"
-                            f"per_act={per_act_token}, per_out_ch={per_out_ch} ==="
-                        )
-
-                        config_results = []
-                        for size_m in args.batch_sizes:
-                            mkn = (size_m, size_k, size_n)
-                            result = bench_run(
-                                [],  # Not used anymore
-                                model,
-                                num_experts,
-                                topk,
-                                per_act_token,
-                                per_out_ch,
-                                mkn,
-                            )
-                            if result:
-                                config_results.append(result)
-
-                        # Print results table for this configuration
-                        if config_results:
-                            print(
-                                f"\n{'Batch Size':<12}"
-                                f"{'Triton (us)':<15}"
-                                f"{'CUTLASS (us)':<15}"
-                            )
-                            print("-" * 45)
-                            for result in config_results:
-                                print(
-                                    f"{result['batch_size']:<12}"
-                                    f"{result['triton_time_us']:<15.2f}"
-                                    f"{result['cutlass_time_us']:<15.2f}"
-                                )
-
-                            all_results.extend(config_results)
-
-    print(f"\nTotal benchmarks completed: {len(all_results)}")
-
-
-if __name__ == "__main__":
-    parser = FlexibleArgumentParser(
-        description="""Benchmark CUTLASS FP8 MOE vs Triton FP8 FUSED MOE
-         across specified models/shapes/batches
-
-        Example usage:
-        python benchmark_cutlass_moe_fp8.py  \
-            --model "Llama-4-Maverick-17B-128E-Instruct-FP8"  \
-            --tp-sizes 8 \
-            --batch-size 2 4 8  \
-            --per-act-token-opts false \
-            --per-out-ch-opts false
-
-        """
-    )
-    parser.add_argument(
-        "--models",
-        nargs="+",
-        type=str,
-        default=DEFAULT_MODELS,
-        choices=WEIGHT_SHAPES_MOE.keys(),
-    )
-    parser.add_argument("--tp-sizes", nargs="+", type=int, default=DEFAULT_TP_SIZES)
-    parser.add_argument(
-        "--batch-sizes", nargs="+", type=int, default=DEFAULT_BATCH_SIZES
-    )
-    parser.add_argument("--limit-k", nargs="+", type=int, default=[])
-    parser.add_argument("--limit-n", nargs="+", type=int, default=[])
-    parser.add_argument(
-        "--per-act-token-opts",
-        nargs="+",
-        type=lambda x: x.lower() == "true",
-        default=[False, True],
-        help="Per-activation token quantization options (true/false)",
-    )
-    parser.add_argument(
-        "--per-out-ch-opts",
-        nargs="+",
-        type=lambda x: x.lower() == "true",
-        default=[False, True],
-        help="Per-output channel quantization options (true/false)",
-    )
-
-    args = parser.parse_args()
-    main(args)

benchmarks/kernels/benchmark_device_communicators.py

@@ -7,10 +7,6 @@ Benchmark script for device communicators:
 CustomAllreduce (oneshot, twoshot), PyNcclCommunicator,
 and SymmMemCommunicator (multimem, two-shot).
 
-for NCCL symmetric memory you need to set the environment variables
-NCCL_NVLS_ENABLE=1 NCCL_CUMEM_ENABLE=1 VLLM_USE_NCCL_SYMM_MEM=1, otherwise NCCL does
-not use fast NVLS implementation for all reduce.
-
 Usage:
     torchrun --nproc_per_node=<N> benchmark_device_communicators.py [options]
 
@@ -30,13 +26,7 @@ import torch.distributed as dist
 from torch.distributed import ProcessGroup
 
 from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce
-from vllm.distributed.device_communicators.pynccl import (
-    PyNcclCommunicator,
-    register_nccl_symmetric_ops,
-)
-from vllm.distributed.device_communicators.pynccl_allocator import (
-    set_graph_pool_id,
-)
+from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
 from vllm.distributed.device_communicators.symm_mem import SymmMemCommunicator
 from vllm.logger import init_logger
 from vllm.utils import FlexibleArgumentParser
@@ -108,7 +98,6 @@ class CommunicatorBenchmark:
             )
             if not self.pynccl_comm.disabled:
                 logger.info("Rank %s: PyNcclCommunicator initialized", self.rank)
-                register_nccl_symmetric_ops(self.pynccl_comm)
             else:
                 logger.info("Rank %s: PyNcclCommunicator disabled", self.rank)
                 self.pynccl_comm = None
@@ -205,15 +194,6 @@ class CommunicatorBenchmark:
                     None,  # no env variable needed
                 )
             )
-            communicators.append(
-                (
-                    "pynccl-symm",
-                    lambda t: torch.ops.vllm.all_reduce_symmetric_with_copy(t),
-                    lambda t: True,  # Always available if initialized
-                    nullcontext(),
-                    None,  # no env variable needed
-                )
-            )
 
         if self.symm_mem_comm_multimem is not None:
             comm = self.symm_mem_comm_multimem
@@ -291,9 +271,7 @@ class CommunicatorBenchmark:
                 # Capture the graph using context manager
                 with context:
                     graph = torch.cuda.CUDAGraph()
-                    graph_pool = torch.cuda.graph_pool_handle()
-                    set_graph_pool_id(graph_pool)
-                    with torch.cuda.graph(graph, pool=graph_pool):
+                    with torch.cuda.graph(graph):
                         for _ in range(CUDA_GRAPH_CAPTURE_CYCLES):
                             allreduce_fn(graph_input)
 

benchmarks/kernels/benchmark_grouped_gemm_cutlass.py

@@ -7,7 +7,6 @@ from benchmark_shapes import WEIGHT_SHAPES_MOE
 
 from vllm import _custom_ops as ops
 from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
-from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_config
 from vllm.model_executor.layers.fused_moe.cutlass_moe import cutlass_moe_fp8
 from vllm.model_executor.layers.fused_moe.fused_moe import (
     fused_experts,
@@ -97,11 +96,6 @@ def bench_run(
         a_scale: torch.Tensor,
         num_repeats: int,
     ):
-        quant_config = fp8_w8a8_moe_quant_config(
-            w1_scale=w1_scale,
-            w2_scale=w2_scale,
-            a1_scale=a_scale,
-        )
         for _ in range(num_repeats):
             fused_experts(
                 a,
@@ -109,7 +103,10 @@ def bench_run(
                 w2,
                 topk_weights,
                 topk_ids,
-                quant_config=quant_config,
+                use_fp8_w8a8=True,
+                w1_scale=w1_scale,
+                w2_scale=w2_scale,
+                a1_scale=a_scale,
             )
 
     def run_cutlass_moe(
@@ -128,12 +125,6 @@ def bench_run(
         per_act_token: bool,
         num_repeats: int,
     ):
-        quant_config = fp8_w8a8_moe_quant_config(
-            w1_scale=w1_scale,
-            w2_scale=w2_scale,
-            per_act_token_quant=per_act_token,
-        )
-
         for _ in range(num_repeats):
             cutlass_moe_fp8(
                 a,
@@ -141,11 +132,14 @@ def bench_run(
                 w2,
                 topk_weights,
                 topk_ids,
+                w1_scale,
+                w2_scale,
                 ab_strides1,
                 ab_strides2,
                 c_strides1,
                 c_strides2,
-                quant_config=quant_config,
+                per_act_token,
+                a1_scale=None,
             )
 
     def run_cutlass_from_graph(
@@ -162,12 +156,6 @@ def bench_run(
         topk_weights: torch.Tensor,
         topk_ids: torch.Tensor,
     ):
-        quant_config = fp8_w8a8_moe_quant_config(
-            w1_scale=w1_scale,
-            w2_scale=w2_scale,
-            per_act_token_quant=per_act_token,
-        )
-
         with set_current_vllm_config(
             VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1))
         ):
@@ -177,11 +165,14 @@ def bench_run(
                 w2_q,
                 topk_weights,
                 topk_ids,
+                w1_scale,
+                w2_scale,
                 ab_strides1,
                 ab_strides2,
                 c_strides1,
                 c_strides2,
-                quant_config=quant_config,
+                per_act_token,
+                a1_scale=None,
             )
 
     def run_triton_from_graph(
@@ -194,11 +185,6 @@ def bench_run(
         w2_scale: torch.Tensor,
         a_scale: torch.Tensor,
     ):
-        quant_config = fp8_w8a8_moe_quant_config(
-            w1_scale=w1_scale,
-            w2_scale=w2_scale,
-            a1_scale=a_scale,
-        )
         with set_current_vllm_config(
             VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1))
         ):
@@ -208,7 +194,10 @@ def bench_run(
                 w2,
                 topk_weights,
                 topk_ids,
-                quant_config=quant_config,
+                use_fp8_w8a8=True,
+                w1_scale=w1_scale,
+                w2_scale=w2_scale,
+                a1_scale=a_scale,
             )
 
     def replay_graph(graph, num_repeats):

benchmarks/kernels/benchmark_lora.py

@@ -79,9 +79,9 @@ def make_rand_lora_weight_tensor(
 
 
 def make_rand_tensors(
-    a_shape: tuple[int, ...],
-    b_shape: tuple[int, ...],
-    c_shape: tuple[int, ...],
+    a_shape: tuple[int],
+    b_shape: tuple[int],
+    c_shape: tuple[int],
     a_dtype: torch.dtype,
     b_dtype: torch.dtype,
     c_dtype: torch.dtype,
@@ -243,7 +243,7 @@ class OpType(Enum):
         lora_rank: int,
         num_loras: int,
         num_slices: int,
-    ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
+    ) -> tuple[tuple[int], tuple[int], tuple[int]]:
         """
         Given num_slices, return the shapes of the A, B, and C matrices
         in A x B = C, for the op_type
@@ -464,11 +464,7 @@ class BenchmarkTensors:
         for field_name in LoRAKernelMeta.__dataclass_fields__:
             field = getattr(self.lora_kernel_meta, field_name)
             assert isinstance(field, torch.Tensor)
-            setattr(
-                self.lora_kernel_meta,
-                field_name,
-                to_device(field) if field_name != "no_lora_flag_cpu" else field,
-            )
+            setattr(self.lora_kernel_meta, field_name, to_device(field))
 
     def metadata(self) -> tuple[int, int, int]:
         """
@@ -516,7 +512,6 @@ class BenchmarkTensors:
             "lora_token_start_loc": self.lora_kernel_meta.lora_token_start_loc,
             "lora_ids": self.lora_kernel_meta.active_lora_ids,
             "scaling": 1.0,
-            "no_lora_flag_cpu": self.lora_kernel_meta.no_lora_flag_cpu,
         }
 
     def as_lora_expand_kwargs(self, add_inputs: bool) -> dict[str, Any]:
@@ -557,7 +552,6 @@ class BenchmarkTensors:
             "lora_ids": self.lora_kernel_meta.active_lora_ids,
             "offset_start": 0,
             "add_inputs": add_inputs,
-            "no_lora_flag_cpu": self.lora_kernel_meta.no_lora_flag_cpu,
         }
 
     def bench_fn_kwargs(

benchmarks/kernels/benchmark_moe.py

@@ -14,10 +14,6 @@ import ray
 import torch
 from ray.experimental.tqdm_ray import tqdm
 
-from vllm.model_executor.layers.fused_moe.config import (
-    FusedMoEQuantConfig,
-    _get_config_dtype_str,
-)
 from vllm.model_executor.layers.fused_moe.fused_moe import *
 from vllm.platforms import current_platform
 from vllm.transformers_utils.config import get_config
@@ -138,36 +134,43 @@ def benchmark_config(
     def run():
         from vllm.model_executor.layers.fused_moe import override_config
 
-        if use_fp8_w8a8:
-            quant_dtype = torch.float8_e4m3fn
-        elif use_int8_w8a16:
-            quant_dtype = torch.int8
-        else:
-            quant_dtype = None
-
-        quant_config = FusedMoEQuantConfig.make(
-            quant_dtype=quant_dtype,
-            w1_scale=w1_scale,
-            w2_scale=w2_scale,
-            a1_scale=a1_scale,
-            a2_scale=a2_scale,
-            block_shape=block_quant_shape,
-        )
-
         with override_config(config):
-            topk_weights, topk_ids, token_expert_indices = fused_topk(
-                x, input_gating, topk, renormalize=not use_deep_gemm
-            )
-            return fused_experts(
-                x,
-                w1,
-                w2,
-                topk_weights,
-                topk_ids,
-                inplace=True,
-                quant_config=quant_config,
-                allow_deep_gemm=use_deep_gemm,
-            )
+            if use_deep_gemm:
+                topk_weights, topk_ids, token_expert_indices = fused_topk(
+                    x, input_gating, topk, False
+                )
+                return fused_experts(
+                    x,
+                    w1,
+                    w2,
+                    topk_weights,
+                    topk_ids,
+                    inplace=True,
+                    use_fp8_w8a8=use_fp8_w8a8,
+                    w1_scale=w1_scale,
+                    w2_scale=w2_scale,
+                    a1_scale=a1_scale,
+                    a2_scale=a2_scale,
+                    block_shape=block_quant_shape,
+                    allow_deep_gemm=True,
+                )
+            else:
+                fused_moe(
+                    x,
+                    w1,
+                    w2,
+                    input_gating,
+                    topk,
+                    renormalize=True,
+                    inplace=True,
+                    use_fp8_w8a8=use_fp8_w8a8,
+                    use_int8_w8a16=use_int8_w8a16,
+                    w1_scale=w1_scale,
+                    w2_scale=w2_scale,
+                    a1_scale=a1_scale,
+                    a2_scale=a2_scale,
+                    block_shape=block_quant_shape,
+                )
 
     # JIT compilation & warmup
     run()
@@ -411,7 +414,7 @@ class BenchmarkWorker:
         use_deep_gemm: bool = False,
     ) -> tuple[dict[str, int], float]:
         current_platform.seed_everything(self.seed)
-        dtype_str = _get_config_dtype_str(
+        dtype_str = get_config_dtype_str(
             dtype, use_int8_w8a16=use_int8_w8a16, use_fp8_w8a8=use_fp8_w8a8
         )
         # NOTE(woosuk): The current naming convention uses w2.shape[2], which
@@ -544,7 +547,7 @@ def save_configs(
     block_quant_shape: list[int],
     save_dir: str,
 ) -> None:
-    dtype_str = _get_config_dtype_str(
+    dtype_str = get_config_dtype_str(
         dtype, use_int8_w8a16=use_int8_w8a16, use_fp8_w8a8=use_fp8_w8a8
     )
 
@@ -557,7 +560,7 @@ def save_configs(
     filename = os.path.join(save_dir, filename)
     print(f"Writing best config to {filename}...")
     with open(filename, "w") as f:
-        json.dump({"triton_version": triton.__version__, **configs}, f, indent=4)
+        json.dump(configs, f, indent=4)
         f.write("\n")
 
 

benchmarks/kernels/benchmark_reshape_and_cache_flash.py

@@ -9,9 +9,6 @@ import torch
 from tabulate import tabulate
 
 from vllm import _custom_ops as ops
-from vllm.attention.ops.triton_reshape_and_cache_flash import (
-    triton_reshape_and_cache_flash,
-)
 from vllm.logger import init_logger
 from vllm.platforms import current_platform
 from vllm.utils import (
@@ -34,8 +31,6 @@ def run_benchmark(
     kv_cache_dtype: str,
     kv_cache_layout: str,
     num_iters: int,
-    implementation: str,
-    benchmark_mode: str,
     device: str = "cuda",
 ) -> float:
     """Return latency (seconds) for given num_tokens."""
@@ -43,14 +38,6 @@ def run_benchmark(
     if kv_cache_dtype == "fp8" and head_size % 16:
         raise ValueError("fp8 kv-cache requires head_size to be a multiple of 16.")
 
-    if implementation not in ("cuda", "triton"):
-        raise ValueError(
-            f"Unsupported implementation: {implementation}. "
-            "Only 'cuda' and 'triton' are supported."
-        )
-    if implementation == "triton" and kv_cache_layout == "HND":
-        return float("nan")  # Triton does not support HND layout yet.
-
     current_platform.seed_everything(42)
     torch.set_default_device(device)
 
@@ -78,49 +65,27 @@ def run_benchmark(
         cache_layout=kv_cache_layout,
     )
     key_cache, value_cache = key_caches[0], value_caches[0]
-    # to free unused memory
-    del key_caches, value_caches
 
     # compute per-kernel scaling factors for fp8 conversion (if used).
     k_scale = (key.amax() / 64.0).to(torch.float32)
     v_scale = (value.amax() / 64.0).to(torch.float32)
 
-    if implementation == "cuda":
-        function_under_test = lambda: ops.reshape_and_cache_flash(
-            key,  # noqa: F821
-            value,  # noqa: F821
-            key_cache,  # noqa: F821
-            value_cache,  # noqa: F821
-            slot_mapping,  # noqa: F821
-            kv_cache_dtype,
-            k_scale,
-            v_scale,
-        )
-    else:
-        function_under_test = lambda: triton_reshape_and_cache_flash(
-            key,  # noqa: F821
-            value,  # noqa: F821
-            key_cache,  # noqa: F821
-            value_cache,  # noqa: F821
-            slot_mapping,  # noqa: F821
-            kv_cache_dtype,
-            k_scale,
-            v_scale,
-        )
-    if benchmark_mode == "cudagraph":
-        g = torch.cuda.CUDAGraph()
-        with torch.cuda.graph(g):
-            function_under_test()
-        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.cuda.synchronize()
         start = time.perf_counter()
         for _ in range(n_iters):
-            function_under_test()
-            torch.cuda.synchronize()
+            ops.reshape_and_cache_flash(
+                key,
+                value,
+                key_cache,
+                value_cache,
+                slot_mapping,
+                kv_cache_dtype,
+                k_scale,
+                v_scale,
+            )
+        torch.cuda.synchronize()
         end = time.perf_counter()
         return (end - start) / n_iters
 
@@ -151,16 +116,10 @@ def main(args):
                 kv_cache_dtype=args.kv_cache_dtype,
                 kv_cache_layout=layout,
                 num_iters=args.iters,
-                implementation=args.implementation,
-                benchmark_mode=args.mode,
                 device="cuda",
             )
             rows.append([n_tok, layout, f"{lat * 1e6:.3f}"])
 
-    print(
-        f"Benchmark results for implementation {args.implementation}"
-        f" (measuring with {args.mode}):"
-    )
     print(tabulate(rows, headers=["num_tokens", "layout", "latency (µs)"]))
 
 
@@ -192,21 +151,6 @@ if __name__ == "__main__":
     )
 
     parser.add_argument("--iters", type=int, default=100)
-
-    parser.add_argument(
-        "--implementation",
-        type=str,
-        choices=["cuda", "triton"],
-        default="cuda",
-    )
-
-    parser.add_argument(
-        "--mode",
-        type=str,
-        choices=["cudagraph", "no_graph"],
-        default="cudagraph",
-    )
-
     args = parser.parse_args()
 
     main(args)

benchmarks/kernels/benchmark_silu_mul_fp8_quant.py

@@ -1,675 +1,77 @@
+#!/usr/bin/env python3
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from collections.abc import Callable
+import time
 
-import matplotlib.pyplot as plt
-import numpy as np
 import torch
 
 from vllm.model_executor.layers.fused_moe.batched_deep_gemm_moe import (
-    silu_mul_fp8_quant_deep_gemm_cuda,
+    silu_mul_fp8_quant_deep_gemm,
 )
 from vllm.platforms import current_platform
-from vllm.triton_utils import tl, triton
-from vllm.utils.deep_gemm import is_deep_gemm_e8m0_used
 
 
-@triton.jit
-def _silu_mul_fp8_quant_deep_gemm(
-    # Pointers ------------------------------------------------------------
-    input_ptr,  # 16-bit activations (E, T, 2*H)
-    y_q_ptr,  # fp8 quantized activations (E, T, H)
-    y_s_ptr,  # 16-bit scales (E, T, G)
-    counts_ptr,  # int32 num tokens per expert (E)
-    # Sizes ---------------------------------------------------------------
-    H: tl.constexpr,  # hidden dimension (per output)
-    GROUP_SIZE: tl.constexpr,  # elements per group (usually 128)
-    # Strides for input (elements) ---------------------------------------
-    stride_i_e,
-    stride_i_t,
-    stride_i_h,
-    # Strides for y_q (elements) -----------------------------------------
-    stride_yq_e,
-    stride_yq_t,
-    stride_yq_h,
-    # Strides for y_s (elements) -----------------------------------------
-    stride_ys_e,
-    stride_ys_t,
-    stride_ys_g,
-    # Stride for counts (elements)
-    stride_counts_e,
-    # Numeric params ------------------------------------------------------
-    eps: tl.constexpr,
-    fp8_min: tl.constexpr,
-    fp8_max: tl.constexpr,
-    use_ue8m0: tl.constexpr,
-    # Meta ---------------------------------------------------------------
-    BLOCK: tl.constexpr,
-    NUM_STAGES: tl.constexpr,
-):
-    G = H // GROUP_SIZE
-
-    # map program id -> (e, g)
-    pid = tl.program_id(0)
-    e = pid // G
-    g = pid % G
-
-    e = e.to(tl.int64)
-    g = g.to(tl.int64)
-
-    # number of valid tokens for this expert
-    n_tokens = tl.load(counts_ptr + e * stride_counts_e).to(tl.int64)
-
-    cols = tl.arange(0, BLOCK).to(tl.int64)
-    mask = cols < BLOCK
-
-    base_input_offset = e * stride_i_e + g * GROUP_SIZE * stride_i_h
-    base_gate_offset = base_input_offset + cols * stride_i_h
-    base_up_offset = base_input_offset + H * stride_i_h + cols * stride_i_h
-    base_yq_offset = e * stride_yq_e + g * GROUP_SIZE * stride_yq_h + cols * stride_yq_h
-    base_ys_offset = e * stride_ys_e + g * stride_ys_g
-
-    for t in tl.range(0, n_tokens, num_stages=NUM_STAGES):
-        gate = tl.load(
-            input_ptr + base_gate_offset + t * stride_i_t, mask=mask, other=0.0
-        ).to(tl.float32)
-        up = tl.load(input_ptr + base_up_offset + t * stride_i_t, mask=mask, other=0.0)
-
-        gate = gate * (1.0 / (1.0 + tl.exp(-gate)))
-        y = gate * up
-
-        y_s = tl.maximum(tl.max(tl.abs(y)), eps) / fp8_max
-        if use_ue8m0:
-            y_s = tl.exp2(tl.ceil(tl.log2(y_s)))
-
-        y_q = tl.clamp(y / y_s, fp8_min, fp8_max).to(y_q_ptr.dtype.element_ty)
-
-        tl.store(y_q_ptr + base_yq_offset + t * stride_yq_t, y_q, mask=mask)
-        tl.store(y_s_ptr + base_ys_offset + t * stride_ys_t, y_s)
-
-
-def silu_mul_fp8_quant_deep_gemm_triton(
-    y: torch.Tensor,  # (E, T, 2*H)
-    tokens_per_expert: torch.Tensor,  # (E,) number of valid tokens per expert
-    num_parallel_tokens,
-    group_size: int = 128,
-    eps: float = 1e-10,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """Quantize silu(y[..., :H]) * y[..., H:] to FP8 with group per-token scales
-
-    y has shape (E, T, 2*H). The first half of the last dimension is
-    silu-activated, multiplied by the second half, then quantized into FP8.
-
-    Returns `(y_q, y_s)` where
-    * `y_q`: FP8 tensor, shape (E, T, H), same layout as y[..., :H]
-    * `y_s`: FP32 tensor, shape (E, T, H // group_size), strides (T*G, 1, T)
-    """
-    assert y.ndim == 3, "y must be (E, T, 2*H)"
-    E, T, H2 = y.shape
-    assert H2 % 2 == 0, "last dim of y must be even (2*H)"
-    H = H2 // 2
-    G = (H + group_size - 1) // group_size
-    assert H % group_size == 0, "H must be divisible by group_size"
-    assert tokens_per_expert.ndim == 1 and tokens_per_expert.shape[0] == E, (
-        "tokens_per_expert must be shape (E,)"
+def benchmark(E, T, H, G=128, runs=50):
+    current_platform.seed_everything(42)
+    y = torch.randn((E, T, 2 * H), dtype=torch.bfloat16, device="cuda")
+    tokens_per_expert = torch.randint(
+        T // 2, T, size=(E,), dtype=torch.int32, device="cuda"
     )
-    tokens_per_expert = tokens_per_expert.to(device=y.device, dtype=torch.int32)
-
-    # allocate outputs
-    fp8_dtype = torch.float8_e4m3fn
-    y_q = torch.empty((E, T, H), dtype=fp8_dtype, device=y.device)
-
-    # strides (elements)
-    stride_i_e, stride_i_t, stride_i_h = y.stride()
-    stride_yq_e, stride_yq_t, stride_yq_h = y_q.stride()
-
-    # desired scale strides (elements): (T*G, 1, T)
-    stride_ys_e = T * G
-    stride_ys_t = 1
-    stride_ys_g = T
-    y_s = torch.empty_strided(
-        (E, T, G),
-        (stride_ys_e, stride_ys_t, stride_ys_g),
-        dtype=torch.float32,
-        device=y.device,
-    )
-
-    stride_cnt_e = tokens_per_expert.stride()[0]
-
-    # Static grid over experts and H-groups.
-    # A loop inside the kernel handles the token dim
-    grid = (E * G,)
-
-    f_info = torch.finfo(fp8_dtype)
-    fp8_max = f_info.max
-    fp8_min = f_info.min
-
-    _silu_mul_fp8_quant_deep_gemm[grid](
-        y,
-        y_q,
-        y_s,
-        tokens_per_expert,
-        H,
-        group_size,
-        stride_i_e,
-        stride_i_t,
-        stride_i_h,
-        stride_yq_e,
-        stride_yq_t,
-        stride_yq_h,
-        stride_ys_e,
-        stride_ys_t,
-        stride_ys_g,
-        stride_cnt_e,
-        eps,
-        fp8_min,
-        fp8_max,
-        is_deep_gemm_e8m0_used(),
-        BLOCK=group_size,
-        NUM_STAGES=4,
-        num_warps=1,
-    )
-
-    return y_q, y_s
-
-
-# Parse generation strategies
-strategies = ["uniform", "max_t", "first_t"]
-
-
-def benchmark(
-    kernel: Callable,
-    E: int,
-    T: int,
-    H: int,
-    total_tokens: int,
-    num_parallel_tokens: int = 64,
-    G: int = 128,
-    runs: int = 200,
-    num_warmups: int = 20,
-    gen_strategy: str = "default",
-    iterations_per_run: int = 20,
-):
-    def generate_data(seed_offset=0):
-        """Generate input data with given seed offset"""
-        current_platform.seed_everything(42 + seed_offset)
-        y = torch.rand((E, T, 2 * H), dtype=torch.bfloat16, device="cuda").contiguous()
-
-        if gen_strategy == "uniform":
-            r = torch.rand(size=(E,), device="cuda")
-            r /= r.sum()
-            r *= total_tokens
-            tokens_per_expert = r.int()
-            tokens_per_expert = torch.minimum(
-                tokens_per_expert,
-                torch.ones((E,), device=r.device, dtype=torch.int) * T,
-            )
-        elif gen_strategy == "max_t":
-            tokens_per_expert = torch.empty(size=(E,), dtype=torch.int32, device="cuda")
-            tokens_per_expert.fill_(total_tokens / E)
-        elif gen_strategy == "first_t":
-            tokens_per_expert = torch.zeros(size=(E,), dtype=torch.int32, device="cuda")
-            tokens_per_expert[0] = min(T, total_tokens)
-        else:
-            raise ValueError(f"Unknown generation strategy: {gen_strategy}")
-        return y, tokens_per_expert
-
-    dataset_count = 4
-    # Pre-generate different input matrices for each iteration to avoid cache effects
-    data_sets = [generate_data(i) for i in range(dataset_count)]
 
     # Warmup
-    y, tokens_per_expert = data_sets[0]
-    for _ in range(num_warmups):
-        kernel(
-            y, tokens_per_expert, num_parallel_tokens=num_parallel_tokens, group_size=G
-        )
-    torch.cuda.synchronize()
-
-    start_event = torch.cuda.Event(enable_timing=True)
-    end_event = torch.cuda.Event(enable_timing=True)
-
-    # Benchmark
-    latencies: list[float] = []
-    for _ in range(runs):
+    for _ in range(10):
+        silu_mul_fp8_quant_deep_gemm(y, tokens_per_expert, group_size=G)
         torch.cuda.synchronize()
 
-        start_event.record()
-        for i in range(iterations_per_run):
-            y, tokens_per_expert = data_sets[i % dataset_count]
-            kernel(
-                y,
-                tokens_per_expert,
-                num_parallel_tokens=num_parallel_tokens,
-                group_size=G,
-            )
-        end_event.record()
-        end_event.synchronize()
+    # Benchmark
+    torch.cuda.synchronize()
+    start = time.perf_counter()
+    for _ in range(runs):
+        silu_mul_fp8_quant_deep_gemm(y, tokens_per_expert, group_size=G)
+    torch.cuda.synchronize()
 
-        total_time_ms = start_event.elapsed_time(end_event)
-        per_iter_time_ms = total_time_ms / iterations_per_run
-        latencies.append(per_iter_time_ms)
+    avg_time = (time.perf_counter() - start) / runs * 1000
 
-    # Use median instead of average for better outlier handling
-    median_time_ms = np.median(latencies)
-    median_time_s = median_time_ms / 1000
-
-    # Calculate actual work done (using first dataset for consistency)
-    _, tokens_per_expert = data_sets[0]
+    # Calculate actual work done (only count valid tokens)
     actual_tokens = tokens_per_expert.sum().item()
     actual_elements = actual_tokens * H
 
     # GFLOPS: operations per element = exp + 3 muls + 1 div + quantization ops ≈ 8 ops
     ops_per_element = 8
     total_ops = actual_elements * ops_per_element
-    gflops = total_ops / median_time_s / 1e9
+    gflops = total_ops / (avg_time / 1000) / 1e9
 
     # Memory bandwidth: bfloat16 inputs (2 bytes), fp8 output (1 byte), scales (4 bytes)
     input_bytes = actual_tokens * 2 * H * 2  # 2*H bfloat16 inputs
     output_bytes = actual_tokens * H * 1  # H fp8 outputs
     scale_bytes = actual_tokens * (H // G) * 4  # scales in float32
     total_bytes = input_bytes + output_bytes + scale_bytes
-    memory_bw = total_bytes / median_time_s / 1e9
+    memory_bw = total_bytes / (avg_time / 1000) / 1e9
 
-    HOPPER_BANDWIDTH_TBPS = 3.35
-    return (
-        median_time_ms,
-        gflops,
-        memory_bw,
-        (memory_bw / (HOPPER_BANDWIDTH_TBPS * 1024)) * 100,
-    )
+    return avg_time, gflops, memory_bw
 
 
-def create_comparison_plot(
-    ratio, cuda_times, baseline_times, config_labels, strategy_name, id
-):
-    """Create a comparison plot for a specific generation strategy"""
-    fig, ax = plt.subplots(1, 1, figsize=(16, 6))
-
-    # Configure x-axis positions
-    x = np.arange(len(config_labels))
-    width = 0.35
-
-    # Execution Time plot (lower is better)
-    ax.bar(
-        x - width / 2, cuda_times, width, label="CUDA Kernel", alpha=0.8, color="blue"
-    )
-    ax.bar(
-        x + width / 2,
-        baseline_times,
-        width,
-        label="Baseline",
-        alpha=0.8,
-        color="orange",
-    )
-
-    # Add speedup labels over each bar pair
-    for i in range(len(x)):
-        speedup = ratio[i]
-        max_height = max(cuda_times[i], baseline_times[i])
-        ax.text(
-            x[i],
-            max_height + max_height * 0.02,
-            f"{speedup:.2f}x",
-            ha="center",
-            va="bottom",
-            fontweight="bold",
-            fontsize=9,
-        )
-
-    ax.set_xlabel("Configuration")
-    ax.set_ylabel("% Utilization")
-    ax.set_title(
-        f"Memory Bandwidth Utilization (%) - {strategy_name}\n(Higher is Better)"
-    )
-    ax.set_xticks(x)
-    ax.set_xticklabels(config_labels, rotation=45, ha="right")
-    ax.legend()
-    ax.grid(True, alpha=0.3)
-
-    plt.tight_layout()
-    return fig, ax
-
-
-def create_combined_plot(all_results):
-    """Create a combined plot with all strategies in one PNG"""
-    num_strategies = len(all_results)
-    fig, axes = plt.subplots(num_strategies, 1, figsize=(20, 6 * num_strategies))
-
-    if num_strategies == 1:
-        axes = [axes]
-
-    for idx, (
-        strategy_name,
-        ratio,
-        cuda_times,
-        baseline_times,
-        config_labels,
-    ) in enumerate(all_results):
-        ax = axes[idx]
-
-        # Configure x-axis positions
-        x = np.arange(len(config_labels))
-        width = 0.35
-
-        # Execution Time plot (lower is better)
-        ax.bar(
-            x - width / 2,
-            cuda_times,
-            width,
-            label="CUDA Kernel",
-            alpha=0.8,
-            color="blue",
-        )
-        ax.bar(
-            x + width / 2,
-            baseline_times,
-            width,
-            label="Baseline",
-            alpha=0.8,
-            color="orange",
-        )
-
-        # Add speedup labels over each bar pair
-        for i in range(len(x)):
-            speedup = ratio[i]
-            max_height = max(cuda_times[i], baseline_times[i])
-            ax.text(
-                x[i],
-                max_height + max_height * 0.02,
-                f"{speedup:.2f}x",
-                ha="center",
-                va="bottom",
-                fontweight="bold",
-                fontsize=9,
-            )
-
-        ax.set_xlabel("Configuration")
-        ax.set_ylabel("% Utilization")
-        ax.set_title(
-            f"Memory Bandwidth Utilization (%) - {strategy_name}\n(Higher is Better)"
-        )
-        ax.set_xticks(x)
-        ax.set_xticklabels(config_labels, rotation=45, ha="right")
-        ax.legend()
-        ax.grid(True, alpha=0.3)
-
-    plt.tight_layout()
-    filename = "../../silu_bench/silu_benchmark_combined.png"
-    plt.savefig(filename, dpi=300, bbox_inches="tight")
-    plt.show()
-
-    return filename
-
-
-outer_dim = 7168
 configs = [
+    (8, 32, 1024),
+    (16, 64, 2048),
+    (32, 128, 4096),
     # DeepSeekV3 Configs
-    (8, 1024, 7168),
-    # DeepSeekV3 Configs
-    (32, 1024, 7168),
-    # DeepSeekV3 Configs
+    (256, 16, 7168),
+    (256, 32, 7168),
+    (256, 64, 7168),
+    (256, 128, 7168),
+    (256, 256, 7168),
+    (256, 512, 7168),
     (256, 1024, 7168),
 ]
 
-runs = 100
-num_warmups = 20
-
-strategy_descriptions = {
-    "uniform": "Uniform Random",
-    "max_t": "Even Assignment",
-    "first_t": "experts[0] = T, experts[1:] = 0",
-}
-
 print(f"GPU: {torch.cuda.get_device_name()}")
-print(f"Testing strategies: {', '.join(strategies)}")
-print(f"Configurations: {len(configs)} configs")
+print(f"{'Config':<20} {'Time(ms)':<10} {'GFLOPS':<10} {'GB/s':<10}")
+print("-" * 50)
 
-all_results = []
-
-# Run benchmarks for each strategy
-for id, strategy in enumerate(strategies):
-    print(f"\n{'=' * 60}")
-    print(f"Testing strategy: {strategy_descriptions[strategy]}")
-    print(f"{'=' * 60}")
-
-    # Collect benchmark data for both algorithms
-    config_labels = []
-    config_x_axis = []
-    all_cuda_results = []
-    all_baseline_results = []
-    all_ratios = []
-
-    for E, T, H in configs:
-        total_tokens_config = [8 * E, 16 * E, 32 * E, 64 * E, 128 * E, 256 * E]
-        config_x_axis.append(total_tokens_config)
-
-        cuda_results = []
-        baseline_results = []
-        ratios = []
-
-        for total_tokens in total_tokens_config:
-            config_label = f"E={E},T={T},H={H},TT={total_tokens}"
-            config_labels.append(config_label)
-
-            # CUDA kernel results
-            time_ms_cuda, gflops, gbps, perc = benchmark(
-                silu_mul_fp8_quant_deep_gemm_cuda,
-                E,
-                T,
-                H,
-                total_tokens,
-                runs=runs,
-                num_warmups=num_warmups,
-                gen_strategy=strategy,
-            )
-            cuda_results.append((time_ms_cuda, gflops, gbps, perc))
-
-            # Baseline results
-            time_ms_triton, gflops, gbps, perc = benchmark(
-                silu_mul_fp8_quant_deep_gemm_triton,
-                E,
-                T,
-                H,
-                total_tokens,
-                runs=runs,
-                num_warmups=num_warmups,
-                gen_strategy=strategy,
-            )
-            baseline_results.append((time_ms_triton, gflops, gbps, perc))
-            ratios.append(time_ms_triton / time_ms_cuda)
-
-            print(f"Completed: {config_label}")
-        all_cuda_results.append(cuda_results)
-        all_baseline_results.append(baseline_results)
-        all_ratios.append(ratios)
-
-    # Store results for combined plotting
-    all_results.append(
-        (
-            strategy_descriptions[strategy],
-            all_ratios,
-            all_cuda_results,
-            all_baseline_results,
-            config_labels,
-            config_x_axis,
-        )
-    )
-
-    # Print summary table for this strategy
-    print(f"\nSummary Table - {strategy_descriptions[strategy]}:")
-    print(f"{'Config':<20} {'CUDA Time(ms)':<12} {'Base Time(ms)':<12} {'Speedup':<8}")
-    print("-" * 60)
-
-    for i, (E, T, H) in enumerate(configs):
-        speedup = baseline_results[i][0] / cuda_results[i][0]
-        config_label = f"E={E:3d},T={T:4d},H={H:4d}"
-        print(
-            f"{config_label:<20} {cuda_results[i][0]:8.5f} "
-            f"{baseline_results[i][0]:8.5f} {speedup:6.2f}x"
-        )
-
-
-def create_total_tokens_plot(all_results):
-    num_strategies = len(all_results)
-    num_configs = len(configs)
-
-    # Create side-by-side subplots: 2 columns for speedup and bandwidth percentage
-    fig, axs = plt.subplots(
-        num_strategies, num_configs * 2, figsize=(28, 6 * num_strategies)
-    )
-
-    # Add main title to the entire figure
-    fig.suptitle(
-        "Performance Analysis: Speedup vs Bandwidth Utilization (Triton & CUDA)",
-        fontsize=16,
-        fontweight="bold",
-        y=0.98,
-    )
-
-    # Handle single strategy case
-    if num_strategies == 1:
-        axs = axs.reshape(1, -1)
-
-    # Handle single config case
-    if num_configs == 1:
-        axs = axs.reshape(-1, 2)
-
-    for strategy_idx, result in enumerate(all_results):
-        (
-            strategy_name,
-            all_ratios,
-            all_cuda_results,
-            all_baseline_results,
-            config_labels,
-            config_x_axis,
-        ) = result
-
-        for config_idx in range(num_configs):
-            # Speedup plot (left column)
-            ax_speedup = axs[strategy_idx, config_idx * 2]
-            # Bandwidth plot (right column)
-            ax_bandwidth = axs[strategy_idx, config_idx * 2 + 1]
-
-            E, T, H = configs[config_idx]
-            ratios = all_ratios[config_idx]
-            total_tokens_values = config_x_axis[config_idx]
-
-            # Extract CUDA and Triton bandwidth percentages
-            cuda_bandwidth_percentages = [
-                result[3] for result in all_cuda_results[config_idx]
-            ]
-            triton_bandwidth_percentages = [
-                result[3] for result in all_baseline_results[config_idx]
-            ]
-
-            # Plot speedup ratios vs total tokens (left plot)
-            ax_speedup.plot(
-                total_tokens_values, ratios, "bo-", linewidth=3, markersize=8
-            )
-            ax_speedup.set_title(
-                f"{strategy_name}\nSpeedup (CUDA/Triton)\nE={E}, T={T}, H={H}",
-                fontsize=12,
-                fontweight="bold",
-            )
-            ax_speedup.set_xlabel("Total Tokens", fontweight="bold", fontsize=11)
-            ax_speedup.set_ylabel("Speedup Ratio", fontweight="bold", fontsize=11)
-            ax_speedup.grid(True, alpha=0.3)
-
-            ax_bandwidth.plot(
-                total_tokens_values,
-                cuda_bandwidth_percentages,
-                "ro-",
-                linewidth=3,
-                markersize=8,
-                label="CUDA",
-            )
-            ax_bandwidth.plot(
-                total_tokens_values,
-                triton_bandwidth_percentages,
-                "go-",
-                linewidth=3,
-                markersize=8,
-                label="Triton",
-            )
-            ax_bandwidth.set_title(
-                f"{strategy_name}\nBandwidth Utilization (Hopper)\nE={E}, T={T}, H={H}",
-                fontsize=12,
-                fontweight="bold",
-            )
-            ax_bandwidth.set_xlabel("Total Tokens", fontweight="bold", fontsize=11)
-            ax_bandwidth.set_ylabel(
-                "% of Peak Bandwidth", fontweight="bold", fontsize=11
-            )
-            ax_bandwidth.legend(prop={"weight": "bold"})
-            ax_bandwidth.grid(True, alpha=0.3)
-
-            # Format x-axis labels for both plots
-            for ax in [ax_speedup, ax_bandwidth]:
-                ax.set_xticks(total_tokens_values)
-                ax.set_xticklabels(
-                    [
-                        f"{tt // 1000}K" if tt >= 1000 else str(tt)
-                        for tt in total_tokens_values
-                    ],
-                    fontweight="bold",
-                )
-                # Make tick labels bold
-                for label in ax.get_xticklabels() + ax.get_yticklabels():
-                    label.set_fontweight("bold")
-
-            # Add value labels on speedup points
-            for x, y in zip(total_tokens_values, ratios):
-                ax_speedup.annotate(
-                    f"{y:.2f}x",
-                    (x, y),
-                    textcoords="offset points",
-                    xytext=(0, 12),
-                    ha="center",
-                    fontsize=10,
-                    fontweight="bold",
-                    bbox=dict(boxstyle="round,pad=0.3", facecolor="white", alpha=0.7),
-                )
-
-            # Add value labels on CUDA bandwidth points
-            for x, y in zip(total_tokens_values, cuda_bandwidth_percentages):
-                ax_bandwidth.annotate(
-                    f"{y:.1f}%",
-                    (x, y),
-                    textcoords="offset points",
-                    xytext=(0, 12),
-                    ha="center",
-                    fontsize=9,
-                    fontweight="bold",
-                    bbox=dict(boxstyle="round,pad=0.2", facecolor="red", alpha=0.3),
-                )
-
-            # Add value labels on Triton bandwidth points
-            for x, y in zip(total_tokens_values, triton_bandwidth_percentages):
-                ax_bandwidth.annotate(
-                    f"{y:.1f}%",
-                    (x, y),
-                    textcoords="offset points",
-                    xytext=(0, -15),
-                    ha="center",
-                    fontsize=9,
-                    fontweight="bold",
-                    bbox=dict(boxstyle="round,pad=0.2", facecolor="green", alpha=0.3),
-                )
-
-    plt.tight_layout()
-    plt.subplots_adjust(top=0.93)  # Make room for main title
-    filename = "silu_benchmark_total_tokens.png"
-    plt.savefig(filename, dpi=300, bbox_inches="tight")
-    plt.show()
-
-    return filename
-
-
-# Create combined plot with all strategies
-combined_plot_filename = create_total_tokens_plot(all_results)
-
-print(f"\n{'=' * 60}")
-print("Benchmark Complete!")
-print(f"Generated combined plot: {combined_plot_filename}")
-print(f"{'=' * 60}")
+for E, T, H in configs:
+    try:
+        time_ms, gflops, gbps = benchmark(E, T, H)
+        print(f"E={E:3d},T={T:4d},H={H:4d} {time_ms:8.3f} {gflops:8.1f} {gbps:8.1f}")
+    except Exception:
+        print(f"E={E:3d},T={T:4d},H={H:4d} FAILED")

benchmarks/kernels/benchmark_w8a8_block_fp8.py

@@ -11,13 +11,13 @@ from datetime import datetime
 from typing import Any
 
 import torch
+import triton
 from tqdm import tqdm
 
 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
     _w8a8_block_fp8_matmul,
 )
 from vllm.platforms import current_platform
-from vllm.triton_utils import triton
 from vllm.utils import FlexibleArgumentParser
 
 mp.set_start_method("spawn", force=True)
@@ -56,7 +56,7 @@ def w8a8_block_matmul(
         Bs: The per-block quantization scale for `B`.
         block_size: The block size for per-block quantization.
                     It should be 2-dim, e.g., [128, 128].
-        output_dtype: The dtype of the returned tensor.
+        output_dytpe: The dtype of the returned tensor.
 
     Returns:
         torch.Tensor: The result of matmul.

benchmarks/kernels/deepgemm/benchmark_fp8_block_dense_gemm.py

@@ -8,16 +8,12 @@ import torch
 
 from vllm import _custom_ops as ops
 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
+    get_col_major_tma_aligned_tensor,
     per_token_group_quant_fp8,
     w8a8_block_fp8_matmul,
 )
 from vllm.triton_utils import triton
-from vllm.utils.deep_gemm import (
-    calc_diff,
-    fp8_gemm_nt,
-    get_col_major_tma_aligned_tensor,
-    per_block_cast_to_fp8,
-)
+from vllm.utils.deep_gemm import calc_diff, fp8_gemm_nt, per_block_cast_to_fp8
 
 
 def benchmark_shape(m: int,

benchmarks/multi_turn/README.md

@@ -55,107 +55,6 @@ output_num_chunks  166.0    99.01   11.80    79.00    90.00    98.00   108.75
 ----------------------------------------------------------------------------------------------------
 ```
 
-### JSON configuration file for synthetic conversations generation
-
-The input flag `--input-file` is used to determine the input conversations for the benchmark.<br/>
-When the input is a JSON file with the field `"filetype": "generate_conversations"` the tool will generate synthetic multi-turn (questions and answers) conversations.
-
-The file `generate_multi_turn.json` is an example file.
-
-The file must contain the sections `prompt_input` and `prompt_output`.
-
-The `prompt_input` section must contain `num_turns`, `prefix_num_tokens` and `num_tokens`:
-
-* `num_turns` - Number of total turns in the conversation (both user & assistant).<br/>
-The final value will always be rounded to an even number so each user turn has a reply.
-* `prefix_num_tokens` - Tokens added at the start of only the **first user turn** in a conversation (unique per conversation).
-* `num_tokens` - Total token length of each **user** message (one turn).
-
-The `prompt_output` section must contain `num_tokens`:
-
-* `num_tokens` - Total token length of each **assistant** message (one turn).
-
-### Random distributions for synthetic conversations generation
-
-When creating an input JSON file (such as `generate_multi_turn.json`),<br/>
-every numeric field (such as `num_turns` or `num_tokens`) requires a distribution.<br/>
-The distribution determines how to randomly sample values for the field.
-
-The available distributions are listed below.
-
-**Note:** The optional `max` field (for lognormal, zipf, and poisson) can be used to cap sampled values at an upper bound.</br>
-Can be used to make sure that the total number of tokens in every request does not exceed `--max-model-len`.
-
-#### constant
-
-```json
-{
-    "distribution": "constant",
-    "value": 500
-}
-```
-
-* `value` - the fixed integer value (always returns the same number).
-
-#### uniform
-
-```json
-{
-    "distribution": "uniform",
-    "min": 12,
-    "max": 18
-}
-```
-
-* `min` - minimum value (inclusive).
-* `max` - maximum value (inclusive), should be equal or larger than min.
-
-#### lognormal
-
-```json
-{
-    "distribution": "lognormal",
-    "average": 1000,
-    "max": 5000
-}
-```
-
-You can parameterize the lognormal distribution in one of two ways:
-
-Using the average and optional median ratio:
-
-* `average` - target average value of the distribution.
-* `median_ratio` - the ratio of the median to the average; controls the skewness. Must be in the range (0, 1).
-
-Using the parameters of the underlying normal distribution:
-
-* `mean` - mean of the underlying normal distribution.
-* `sigma` - standard deviation of the underlying normal distribution.
-
-#### zipf
-
-```json
-{
-    "distribution": "zipf",
-    "alpha": 1.2,
-    "max": 100
-}
-```
-
-* `alpha` - skew parameter (> 1). Larger values produce stronger skew toward smaller integers.
-
-#### poisson
-
-```json
-{
-    "distribution": "poisson",
-    "alpha": 10,
-    "max": 50
-}
-```
-
-* `alpha` - expected value (λ). Also the variance of the distribution.
-
 ## ShareGPT Conversations
 
 To run with the ShareGPT data, download the following ShareGPT dataset:

benchmarks/multi_turn/bench_dataset.py

@@ -99,105 +99,21 @@ class PoissonDistribution(Distribution):
 
 class LognormalDistribution(Distribution):
     def __init__(
-        self,
-        mean: Optional[float] = None,
-        sigma: Optional[float] = None,
-        average: Optional[int] = None,
-        median_ratio: Optional[float] = None,
-        max_val: Optional[int] = None,
+        self, mean: float, sigma: float, max_val: Optional[int] = None
     ) -> None:
-        self.average = average
-        self.median_ratio = median_ratio
-        self.max_val = max_val
-
-        if average is not None:
-            if average < 1:
-                raise ValueError("Lognormal average must be positive")
-
-            if mean or sigma:
-                raise ValueError(
-                    "When using lognormal average, you can't provide mean/sigma"
-                )
-
-            if self.median_ratio is None:
-                # Default value that provides relatively wide range of values
-                self.median_ratio = 0.85
-
-            # Calculate mean/sigma of np.random.lognormal based on the average
-            mean, sigma = self._generate_lognormal_by_median(
-                target_average=self.average, median_ratio=self.median_ratio
-            )
-        else:
-            if mean is None or sigma is None:
-                raise ValueError(
-                    "Must provide both mean and sigma if average is not used"
-                )
-
-            if mean <= 0 or sigma < 0:
-                raise ValueError(
-                    "Lognormal mean must be positive and sigma must be non-negative"
-                )
-
-        # Mean and standard deviation of the underlying normal distribution
-        # Based on numpy.random.lognormal
         self.mean = mean
         self.sigma = sigma
-
-    @staticmethod
-    def _generate_lognormal_by_median(
-        target_average: int, median_ratio: float
-    ) -> tuple[float, float]:
-        """
-        Compute (mu, sigma) for a lognormal distribution given:
-        - a target average (mean of the distribution)
-        - a ratio of median / mean (controls skewness), assume mean > median
-
-        Background:
-        If Z ~ Normal(mu, sigma^2), then X = exp(Z) ~ LogNormal(mu, sigma).
-        * mean(X)   = exp(mu + sigma^2 / 2)
-        * median(X) = exp(mu)
-
-        So:
-        median / mean = exp(mu) / exp(mu + sigma^2 / 2)
-                      = exp(-sigma^2 / 2)
-
-        Rearranging:
-        sigma^2 = 2 * ln(mean / median)
-        mu      = ln(median)
-
-        This gives a unique (mu, sigma) for any valid mean and median.
-        """
-        # Check input validity: median must be smaller than mean
-        if median_ratio <= 0 or median_ratio >= 1:
-            raise ValueError("median_ratio must be in range (0, 1)")
-
-        target_median = target_average * median_ratio
-
-        # Solve sigma^2 = 2 * ln(mean / median)
-        sigma = np.sqrt(2 * np.log(target_average / target_median))
-        mu = np.log(target_median)
-
-        return mu, sigma
+        self.max_val = max_val
 
     def sample(self, size: int = 1) -> np.ndarray:
         samples = np.random.lognormal(mean=self.mean, sigma=self.sigma, size=size)
-
-        if self.average is not None:
-            # Scale to average
-            samples *= self.average / samples.mean()
-
         if self.max_val:
             samples = np.minimum(samples, self.max_val)
 
         return np.round(samples).astype(int)
 
     def __repr__(self) -> str:
-        if self.average:
-            return (
-                f"LognormalDistribution[{self.average}, "
-                f"{self.median_ratio}, {self.max_val}]"
-            )
-        return f"LognormalDistribution[{self.mean}, {self.sigma}, {self.max_val}]"
+        return f"LognormalDistribution[{self.mean}, {self.sigma}]"
 
 
 class GenConvArgs(NamedTuple):
@@ -257,21 +173,10 @@ def get_random_distribution(
         return PoissonDistribution(conf["alpha"], max_val=max_val)
 
     elif distribution == "lognormal":
-        max_val = conf.get("max", None)
-
-        if "average" in conf:
-            # Infer lognormal mean/sigma (numpy) from input average
-            median_ratio = conf.get("median_ratio", None)
-            return LognormalDistribution(
-                average=conf["average"], median_ratio=median_ratio, max_val=max_val
-            )
-
-        # Use mean/sigma directly (for full control over the distribution)
         verify_field_exists(conf, "mean", section, subsection)
         verify_field_exists(conf, "sigma", section, subsection)
-        return LognormalDistribution(
-            mean=conf["mean"], sigma=conf["sigma"], max_val=max_val
-        )
+        max_val = conf.get("max", None)
+        return LognormalDistribution(conf["mean"], conf["sigma"], max_val=max_val)
 
     elif distribution == "uniform":
         verify_field_exists(conf, "min", section, subsection)

benchmarks/multi_turn/generate_multi_turn.json

@@ -15,8 +15,9 @@
         },
         "prefix_num_tokens": {
             "distribution": "lognormal",
-            "average": 1000,
-            "max": 5000
+            "mean": 6,
+            "sigma": 4,
+            "max": 1500
         },
         "num_tokens": {
             "distribution": "uniform",

cmake/cpu_extension.cmake

@@ -101,7 +101,6 @@ else()
     find_isa(${CPUINFO} "asimd" ASIMD_FOUND) # Check for ARM NEON support
     find_isa(${CPUINFO} "bf16" ARM_BF16_FOUND) # Check for ARM BF16 support
     find_isa(${CPUINFO} "S390" S390_FOUND)
-    find_isa(${CPUINFO} "v" RVV_FOUND) # Check for RISC-V RVV support
 endif()
 
 if (AVX512_FOUND AND NOT AVX512_DISABLED)
@@ -178,14 +177,8 @@ elseif (S390_FOUND)
         "-mzvector"
         "-march=native"
         "-mtune=native")
-elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64")
-    if(RVV_FOUND)
-	    message(FAIL_ERROR "Can't support rvv now.")
-    else()
-        list(APPEND CXX_COMPILE_FLAGS "-march=rv64gc")
-    endif()
 else()
-    message(FATAL_ERROR "vLLM CPU backend requires AVX512, AVX2, Power9+ ISA, S390X ISA, ARMv8 or RISC-V support.")
+    message(FATAL_ERROR "vLLM CPU backend requires AVX512, AVX2, Power9+ ISA, S390X ISA or ARMv8 support.")
 endif()
 
 #
@@ -265,8 +258,7 @@ set(VLLM_EXT_SRC
     "csrc/cpu/layernorm.cpp"
     "csrc/cpu/mla_decode.cpp"
     "csrc/cpu/pos_encoding.cpp"
-    "csrc/cpu/torch_bindings.cpp"
-    "csrc/moe/dynamic_4bit_int_moe_cpu.cpp")
+    "csrc/cpu/torch_bindings.cpp")
 
 if (AVX512_FOUND AND NOT AVX512_DISABLED)
     set(VLLM_EXT_SRC

cmake/utils.cmake

@@ -480,6 +480,7 @@ function (define_gpu_extension_target GPU_MOD_NAME)
       ${GPU_LANGUAGE}_ARCHITECTURES "${GPU_ARCHITECTURES}")
   endif()
 
+  set_property(TARGET ${GPU_MOD_NAME} PROPERTY CXX_STANDARD 17)
 
   target_compile_options(${GPU_MOD_NAME} PRIVATE
     $<$<COMPILE_LANGUAGE:${GPU_LANGUAGE}>:${GPU_COMPILE_FLAGS}>)

csrc/attention/mla/cutlass_mla_entry.cu

@@ -0,0 +1,38 @@
+/*
+ * Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <torch/all.h>
+
+#if defined ENABLE_CUTLASS_MLA && ENABLE_CUTLASS_MLA
+void cutlass_mla_decode_sm100a(torch::Tensor const& out,
+                               torch::Tensor const& q_nope,
+                               torch::Tensor const& q_pe,
+                               torch::Tensor const& kv_c_and_k_pe_cache,
+                               torch::Tensor const& seq_lens,
+                               torch::Tensor const& page_table, double scale);
+#endif
+
+void cutlass_mla_decode(torch::Tensor const& out, torch::Tensor const& q_nope,
+                        torch::Tensor const& q_pe,
+                        torch::Tensor const& kv_c_and_k_pe_cache,
+                        torch::Tensor const& seq_lens,
+                        torch::Tensor const& page_table, double scale) {
+#if defined ENABLE_CUTLASS_MLA && ENABLE_CUTLASS_MLA
+  return cutlass_mla_decode_sm100a(out, q_nope, q_pe, kv_c_and_k_pe_cache,
+                                   seq_lens, page_table, scale);
+#endif
+  TORCH_CHECK_NOT_IMPLEMENTED(false, "No compiled cutlass MLA");
+}

csrc/attention/mla/cutlass_mla_kernels.cu

@@ -0,0 +1,225 @@
+/*
+ * Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <torch/all.h>
+
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+
+#include "cute/tensor.hpp"
+
+#include "cutlass/cutlass.h"
+#include "cutlass/kernel_hardware_info.h"
+
+#include "cutlass_extensions/common.hpp"
+
+#include "device/sm100_mla.hpp"
+#include "kernel/sm100_mla_tile_scheduler.hpp"
+
+using namespace cute;
+using namespace cutlass::fmha::kernel;
+
+template <typename T, bool PersistenceOption = true>
+struct MlaSm100 {
+  using Element = T;
+  using ElementAcc = float;
+  using ElementOut = T;
+
+  using TileShape = Shape<_128, _128, Shape<_512, _64>>;
+  using TileShapeH = cute::tuple_element_t<0, TileShape>;
+  using TileShapeD = cute::tuple_element_t<2, TileShape>;
+
+  // H K (D_latent D_rope) B
+  using ProblemShape = cute::tuple<TileShapeH, int, TileShapeD, int>;
+
+  using StrideQ = cute::tuple<int64_t, _1, int64_t>;  // H D B
+  using StrideK = cute::tuple<int64_t, _1, int64_t>;  // K D B
+  using StrideO = StrideK;                            // H D B
+  using StrideLSE = cute::tuple<_1, int>;             // H B
+
+  using TileScheduler =
+      std::conditional_t<PersistenceOption, Sm100MlaPersistentTileScheduler,
+                         Sm100MlaIndividualTileScheduler>;
+
+  using FmhaKernel =
+      cutlass::fmha::kernel::Sm100FmhaMlaKernelTmaWarpspecialized<
+          TileShape, Element, ElementAcc, ElementOut, ElementAcc, TileScheduler,
+          /*kIsCpAsync=*/true>;
+  using Fmha = cutlass::fmha::device::MLA<FmhaKernel>;
+};
+
+template <typename T>
+typename T::Fmha::Arguments args_from_options(
+    at::Tensor const& out, at::Tensor const& q_nope, at::Tensor const& q_pe,
+    at::Tensor const& kv_c_and_k_pe_cache, at::Tensor const& seq_lens,
+    at::Tensor const& page_table, double scale) {
+  cutlass::KernelHardwareInfo hw_info;
+  hw_info.device_id = q_nope.device().index();
+  hw_info.sm_count =
+      cutlass::KernelHardwareInfo::query_device_multiprocessor_count(
+          hw_info.device_id);
+
+  int batches = q_nope.sizes()[0];
+  int page_count_per_seq = page_table.sizes()[1];
+  int page_count_total = kv_c_and_k_pe_cache.sizes()[0];
+  int page_size = kv_c_and_k_pe_cache.sizes()[1];
+  int max_seq_len = page_size * page_count_per_seq;
+  using TileShapeH = typename T::TileShapeH;
+  using TileShapeD = typename T::TileShapeD;
+  auto problem_shape =
+      cute::make_tuple(TileShapeH{}, max_seq_len, TileShapeD{}, batches);
+
+  auto [H, K, D, B] = problem_shape;
+  auto [D_latent, D_rope] = D;
+
+  using StrideQ = typename T::StrideQ;
+  using StrideK = typename T::StrideK;
+  using StrideO = typename T::StrideO;
+  using StrideLSE = typename T::StrideLSE;
+
+  StrideQ stride_Q_latent = cute::make_tuple(
+      static_cast<int64_t>(D_latent), _1{}, static_cast<int64_t>(H * D_latent));
+  StrideQ stride_Q_rope = cute::make_tuple(static_cast<int64_t>(D_rope), _1{},
+                                           static_cast<int64_t>(H * D_rope));
+  StrideK stride_C =
+      cute::make_tuple(static_cast<int64_t>(D_latent + D_rope), _1{},
+                       static_cast<int64_t>(page_size * (D_latent + D_rope)));
+  StrideLSE stride_PT = cute::make_stride(_1{}, page_count_per_seq);
+  StrideLSE stride_LSE = cute::make_tuple(_1{}, static_cast<int>(H));
+  StrideO stride_O = cute::make_tuple(static_cast<int64_t>(D_latent), _1{},
+                                      static_cast<int64_t>(H * D_latent));
+
+  using Element = typename T::Element;
+  using ElementOut = typename T::ElementOut;
+  using ElementAcc = typename T::ElementAcc;
+  auto Q_latent_ptr = static_cast<Element*>(q_nope.data_ptr());
+  auto Q_rope_ptr = static_cast<Element*>(q_pe.data_ptr());
+  auto C_ptr = static_cast<Element*>(kv_c_and_k_pe_cache.data_ptr());
+  auto scale_f = static_cast<float>(scale);
+  typename T::Fmha::Arguments arguments{
+      problem_shape,
+      {scale_f, Q_latent_ptr, stride_Q_latent, Q_rope_ptr, stride_Q_rope, C_ptr,
+       stride_C, C_ptr + D_latent, stride_C,
+       static_cast<int*>(seq_lens.data_ptr()),
+       static_cast<int*>(page_table.data_ptr()), stride_PT, page_count_total,
+       page_size},
+      {static_cast<ElementOut*>(out.data_ptr()), stride_O,
+       static_cast<ElementAcc*>(nullptr), stride_LSE},
+      hw_info,
+      1,        // split_kv
+      nullptr,  // is_var_split_kv
+  };
+  // TODO(kaixih@nvidia): When split_kv=-1 and is_var_split_kv=false, we compute
+  // split_kv automatically based on batch size and sequence length to balance
+  // workload across available SMs. Consider using var_split_kv for manual
+  // control if needed.
+  T::Fmha::set_split_kv(arguments);
+  return arguments;
+}
+
+template <typename Element>
+void runMla(at::Tensor const& out, at::Tensor const& q_nope,
+            at::Tensor const& q_pe, at::Tensor const& kv_c_and_k_pe_cache,
+            at::Tensor const& seq_lens, at::Tensor const& page_table,
+            float scale, cudaStream_t stream) {
+  using MlaSm100Type = MlaSm100<Element>;
+  typename MlaSm100Type::Fmha fmha;
+  auto arguments = args_from_options<MlaSm100Type>(
+      out, q_nope, q_pe, kv_c_and_k_pe_cache, seq_lens, page_table, scale);
+  size_t workspace_size = MlaSm100Type::Fmha::get_workspace_size(arguments);
+  auto const workspace_options =
+      torch::TensorOptions().dtype(torch::kUInt8).device(q_nope.device());
+  auto workspace = torch::empty(workspace_size, workspace_options);
+
+  CUTLASS_CHECK(fmha.can_implement(arguments));
+
+  CUTLASS_CHECK(fmha.initialize(arguments, workspace.data_ptr(), stream));
+
+  CUTLASS_CHECK(fmha.run(arguments, workspace.data_ptr(), stream));
+}
+
+void cutlass_mla_decode_sm100a(torch::Tensor const& out,
+                               torch::Tensor const& q_nope,
+                               torch::Tensor const& q_pe,
+                               torch::Tensor const& kv_c_and_k_pe_cache,
+                               torch::Tensor const& seq_lens,
+                               torch::Tensor const& page_table, double scale) {
+  TORCH_CHECK(q_nope.device().is_cuda(), "q_nope must be on CUDA");
+  TORCH_CHECK(q_nope.dim() == 3, "q_nope must be a 3D tensor");
+  TORCH_CHECK(q_pe.dim() == 3, "q_pe must be a 3D tensor");
+  TORCH_CHECK(kv_c_and_k_pe_cache.dim() == 3,
+              "kv_c_and_k_pe_cache must be a 3D tensor");
+  TORCH_CHECK(seq_lens.dim() == 1, "seq_lens must be a 1D tensor");
+  TORCH_CHECK(page_table.dim() == 2, "page_table must be a 2D tensor");
+  TORCH_CHECK(out.dim() == 3, "out must be a 3D tensor");
+
+  auto B_q_nope = q_nope.size(0);
+  auto H_q_nope = q_nope.size(1);
+  auto D_q_nope = q_nope.size(2);
+  auto B_q_pe = q_pe.size(0);
+  auto H_q_pe = q_pe.size(1);
+  auto D_q_pe = q_pe.size(2);
+  auto B_pt = page_table.size(0);
+  auto PAGE_NUM = page_table.size(1);
+  auto PAGE_SIZE = kv_c_and_k_pe_cache.size(1);
+  auto D_ckv = kv_c_and_k_pe_cache.size(2);
+  auto B_o = out.size(0);
+  auto H_o = out.size(1);
+  auto D_o = out.size(2);
+
+  TORCH_CHECK(D_q_nope == 512, "D_q_nope must be equal to 512");
+  TORCH_CHECK(D_q_pe == 64, "D_q_pe must be equal to 64");
+  TORCH_CHECK(D_ckv == 576, "D_ckv must be equal to 576");
+  TORCH_CHECK(H_q_nope == H_q_pe && H_q_nope == H_o && H_o == 128,
+              "H_q_nope, H_q_pe, and H_o must be equal to 128");
+  TORCH_CHECK(PAGE_SIZE > 0 && (PAGE_SIZE & (PAGE_SIZE - 1)) == 0,
+              "PAGE_SIZE must be a power of 2");
+  TORCH_CHECK(
+      B_q_nope == B_q_pe && B_q_nope == B_pt && B_q_nope == B_o,
+      "Batch dims must be same for page_table, q_nope and q_pe, and out");
+  TORCH_CHECK(PAGE_NUM % (128 / PAGE_SIZE) == 0,
+              "PAGE_NUM must be divisible by 128 / PAGE_SIZE");
+  TORCH_CHECK(D_o == 512, "D_o must be equal to 512");
+
+  TORCH_CHECK(q_nope.dtype() == at::ScalarType::Half ||
+                  q_nope.dtype() == at::ScalarType::BFloat16 ||
+                  q_nope.dtype() == at::ScalarType::Float8_e4m3fn,
+              "q_nope must be a half, bfloat16, or float8_e4m3fn tensor");
+  TORCH_CHECK(kv_c_and_k_pe_cache.dtype() == q_nope.dtype() &&
+                  q_nope.dtype() == q_pe.dtype(),
+              "kv_c_and_k_pe_cache, q_nope, and q_pe must be the same type");
+  TORCH_CHECK(seq_lens.dtype() == torch::kInt32,
+              "seq_lens must be a 32-bit integer tensor");
+  TORCH_CHECK(page_table.dtype() == torch::kInt32,
+              "page_table must be a 32-bit integer tensor");
+
+  auto in_dtype = q_nope.dtype();
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(q_nope));
+  const cudaStream_t stream =
+      at::cuda::getCurrentCUDAStream(q_nope.get_device());
+  if (in_dtype == at::ScalarType::Half) {
+    runMla<cutlass::half_t>(out, q_nope, q_pe, kv_c_and_k_pe_cache, seq_lens,
+                            page_table, scale, stream);
+  } else if (in_dtype == at::ScalarType::BFloat16) {
+    runMla<cutlass::bfloat16_t>(out, q_nope, q_pe, kv_c_and_k_pe_cache,
+                                seq_lens, page_table, scale, stream);
+  } else if (in_dtype == at::ScalarType::Float8_e4m3fn) {
+    runMla<cutlass::float_e4m3_t>(out, q_nope, q_pe, kv_c_and_k_pe_cache,
+                                  seq_lens, page_table, scale, stream);
+  } else {
+    TORCH_CHECK(false, "Unsupported input data type of MLA");
+  }
+}

csrc/attention/mla/cutlass_sm100_mla/device/sm100_mla.hpp

@@ -133,14 +133,6 @@ public:
     // printf("    sm_count = %d\n", sm_count);
     int max_splits = ceil_div(K, 128);
     max_splits = min(16, max_splits);
-
-    // TODO: This avoids a hang when the batch size larger than 1 and 
-    // there is more than 1 kv_splits. 
-    // Discuss with NVIDIA how this can be fixed.
-    if (B > 1) {
-      max_splits = min(1, max_splits);
-    }
-    
     // printf("    max_splits = %d\n", max_splits);
     int sms_per_batch = max(1, sm_count / B);
     // printf("    sms_per_batch = %d\n", sms_per_batch);

csrc/cpu/cpu_types.hpp

@@ -14,12 +14,7 @@
   // arm implementation
   #include "cpu_types_arm.hpp"
 #else
-  #warning "unsupported vLLM cpu implementation, vLLM will compile with scalar"
-  #include "cpu_types_scalar.hpp"
-#endif
-
-#ifdef _OPENMP
-  #include <omp.h>
+  #warning "unsupported vLLM cpu implementation"
 #endif
 
 #endif

csrc/cpu/cpu_types_scalar.hpp

@@ -1,513 +0,0 @@
-#include <cmath>
-#include <cstdint>
-#include <cstring>
-#include <torch/all.h>
-#include "float_convert.hpp"
-
-namespace vec_op {
-
-#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...)            \
-  AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__)    \
-  AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) \
-  AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__)
-
-#define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
-  AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
-
-#ifndef CPU_OP_GUARD
-  #define CPU_KERNEL_GUARD_IN(NAME)
-  #define CPU_KERNEL_GUARD_OUT(NAME)
-#else
-  #define CPU_KERNEL_GUARD_IN(NAME) \
-    std::cout << #NAME << " invoked." << std::endl;
-  #define CPU_KERNEL_GUARD_OUT(NAME) \
-    std::cout << #NAME << " exit." << std::endl;
-#endif
-
-#define FORCE_INLINE __attribute__((always_inline)) inline
-
-#define __max(a, b) ((a) > (b) ? (a) : (b))
-#define __min(a, b) ((a) < (b) ? (a) : (b))
-#define __abs(a) ((a) < (0) ? (0 - a) : (a))
-
-typedef struct f16x8_t {
-  uint16_t val[8];
-} f16x8_t;
-
-typedef struct f16x16_t {
-  uint16_t val[16];
-} f16x16_t;
-
-typedef struct f16x32_t {
-  uint16_t val[32];
-} f16x32_t;
-
-typedef struct f32x4_t {
-  float val[4];
-} f32x4_t;
-
-typedef struct f32x8_t {
-  float val[8];
-} f32x8_t;
-
-typedef struct f32x16_t {
-  float val[16];
-} f32x16_t;
-
-namespace {
-template <typename T, T... indexes, typename F>
-constexpr void unroll_loop_item(std::integer_sequence<T, indexes...>, F&& f) {
-  (f(std::integral_constant<T, indexes>{}), ...);
-};
-};  // namespace
-
-template <typename T, T count, typename F,
-          typename = std::enable_if_t<std::is_invocable_v<F, T> > >
-constexpr void unroll_loop(F&& f) {
-  unroll_loop_item(std::make_integer_sequence<T, count>{}, std::forward<F>(f));
-}
-
-template <typename T>
-struct Vec {
-  constexpr static int get_elem_num() { return T::VEC_ELEM_NUM; }
-};
-
-struct FP32Vec8;
-struct FP32Vec16;
-
-struct FP16Vec8 : public Vec<FP16Vec8> {
-  constexpr static int VEC_ELEM_NUM = 8;
-  f16x8_t reg;
-
-  explicit FP16Vec8(const void* ptr)
-      : reg(*reinterpret_cast<const f16x8_t*>(ptr)) {};
-
-  explicit FP16Vec8(const FP32Vec8&);
-
-  void save(void* ptr) const { *reinterpret_cast<f16x8_t*>(ptr) = reg; }
-};
-
-struct FP16Vec16 : public Vec<FP16Vec16> {
-  constexpr static int VEC_ELEM_NUM = 16;
-  f16x16_t reg;
-
-  explicit FP16Vec16(const void* ptr)
-      : reg(*reinterpret_cast<const f16x16_t*>(ptr)) {};
-
-  explicit FP16Vec16(const FP32Vec16&);
-
-  void save(void* ptr) const { *reinterpret_cast<f16x16_t*>(ptr) = reg; }
-
-  void save(void* ptr, const int elem_num) const {
-    int num = __min(elem_num, VEC_ELEM_NUM);
-    std::memcpy(ptr, &(reg.val[0]), num * sizeof(uint16_t));
-  }
-};
-
-struct BF16Vec8 : public Vec<BF16Vec8> {
-  constexpr static int VEC_ELEM_NUM = 8;
-  f16x8_t reg;
-
-  explicit BF16Vec8(const void* ptr)
-      : reg(*reinterpret_cast<const f16x8_t*>(ptr)) {};
-
-  explicit BF16Vec8(const FP32Vec8&);
-
-  void save(void* ptr) const { *reinterpret_cast<f16x8_t*>(ptr) = reg; }
-};
-
-struct BF16Vec16 : public Vec<BF16Vec16> {
-  constexpr static int VEC_ELEM_NUM = 16;
-  f16x16_t reg;
-
-  explicit BF16Vec16(const void* ptr)
-      : reg(*reinterpret_cast<const f16x16_t*>(ptr)) {};
-
-  explicit BF16Vec16(const FP32Vec16&);
-
-  void save(void* ptr) const { *reinterpret_cast<f16x16_t*>(ptr) = reg; }
-
-  void save(void* ptr, const int elem_num) const {
-    int num = __min(elem_num, VEC_ELEM_NUM);
-    std::memcpy(ptr, &(reg.val[0]), num * sizeof(uint16_t));
-  }
-};
-
-struct BF16Vec32 : public Vec<BF16Vec32> {
-  constexpr static int VEC_ELEM_NUM = 32;
-  f16x32_t reg;
-
-  explicit BF16Vec32(const void* ptr)
-      : reg(*reinterpret_cast<const f16x32_t*>(ptr)) {};
-
-  explicit BF16Vec32(f16x32_t data) : reg(data) {};
-
-  explicit BF16Vec32(BF16Vec8& vec8_data) {
-    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
-      reg.val[i] = vec8_data.reg.val[i % BF16Vec8::VEC_ELEM_NUM];
-    }
-  }
-
-  void save(void* ptr) const { *reinterpret_cast<f16x32_t*>(ptr) = reg; }
-};
-
-struct FP32Vec4 : public Vec<FP32Vec4> {
-  constexpr static int VEC_ELEM_NUM = 4;
-
-  f32x4_t reg;
-
-  explicit FP32Vec4(float v) {
-    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
-      reg.val[i] = v;
-    }
-  }
-
-  explicit FP32Vec4() {
-    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
-      reg.val[i] = 0.0f;
-    }
-  }
-
-  explicit FP32Vec4(const float* ptr)
-      : reg(*reinterpret_cast<const f32x4_t*>(ptr)) {};
-
-  explicit FP32Vec4(f32x4_t data) : reg(data) {};
-
-  explicit FP32Vec4(const FP32Vec4& data) : reg(data.reg) {};
-};
-
-struct FP32Vec8 : public Vec<FP32Vec8> {
-  constexpr static int VEC_ELEM_NUM = 8;
-
-  f32x8_t reg;
-
-  explicit FP32Vec8(float v) {
-    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
-      reg.val[i] = v;
-    }
-  }
-
-  explicit FP32Vec8() {
-    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
-      reg.val[i] = 0.0f;
-    }
-  }
-
-  explicit FP32Vec8(const float* ptr)
-      : reg(*reinterpret_cast<const f32x8_t*>(ptr)) {};
-
-  explicit FP32Vec8(f32x8_t data) : reg(data) {};
-
-  explicit FP32Vec8(const FP32Vec8& data) : reg(data.reg) {};
-
-  explicit FP32Vec8(const FP16Vec8& v) {
-    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
-      reg.val[i] = fp16_to_float(v.reg.val[i]);
-    }
-  }
-
-  FP32Vec8(const BF16Vec8& v) {
-    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
-      reg.val[i] = bf16_to_float(v.reg.val[i]);
-    }
-  }
-
-  float reduce_sum() const {
-    float result = 0;
-    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
-      result += reg.val[i];
-    }
-    return result;
-  }
-
-  FP32Vec8 exp() const {
-    f32x8_t ret;
-    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
-      ret.val[i] = expf(reg.val[i]);
-    }
-    return FP32Vec8(ret);
-  }
-
-  FP32Vec8 tanh() const {
-    f32x8_t ret;
-    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
-      ret.val[i] = tanhf(reg.val[i]);
-    }
-    return FP32Vec8(ret);
-  }
-
-  FP32Vec8 er() const {
-    f32x8_t ret;
-    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
-      ret.val[i] = erf(reg.val[i]);
-    }
-    return FP32Vec8(ret);
-  }
-
-  FP32Vec8 operator*(const FP32Vec8& b) const {
-    f32x8_t ret;
-    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
-      ret.val[i] = reg.val[i] * b.reg.val[i];
-    }
-    return FP32Vec8(ret);
-  }
-
-  FP32Vec8 operator+(const FP32Vec8& b) const {
-    f32x8_t ret;
-    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
-      ret.val[i] = reg.val[i] + b.reg.val[i];
-    }
-    return FP32Vec8(ret);
-  }
-
-  FP32Vec8 operator-(const FP32Vec8& b) const {
-    f32x8_t ret;
-    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
-      ret.val[i] = reg.val[i] - b.reg.val[i];
-    }
-    return FP32Vec8(ret);
-  }
-
-  FP32Vec8 operator/(const FP32Vec8& b) const {
-    f32x8_t ret;
-    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
-      ret.val[i] = reg.val[i] / b.reg.val[i];
-    }
-    return FP32Vec8(ret);
-  }
-
-  void save(void* ptr) const { *reinterpret_cast<f32x8_t*>(ptr) = reg; }
-};
-
-struct FP32Vec16 : public Vec<FP32Vec16> {
-  constexpr static int VEC_ELEM_NUM = 16;
-  f32x16_t reg;
-
-  explicit FP32Vec16(float v) {
-    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
-      reg.val[i] = v;
-    }
-  }
-
-  explicit FP32Vec16() {
-    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
-      reg.val[i] = 0.0f;
-    }
-  }
-
-  explicit FP32Vec16(const float* ptr)
-      : reg(*reinterpret_cast<const f32x16_t*>(ptr)) {};
-
-  explicit FP32Vec16(f32x16_t data) : reg(data) {};
-
-  FP32Vec16(const FP32Vec4& data) {
-    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
-      reg.val[i] = data.reg.val[i % FP32Vec4::VEC_ELEM_NUM];
-    }
-  }
-
-  FP32Vec16(const FP32Vec8& data) {
-    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
-      reg.val[i] = data.reg.val[i % FP32Vec8::VEC_ELEM_NUM];
-    }
-  }
-
-  FP32Vec16(const FP32Vec16& data) : reg(data.reg) {};
-
-  explicit FP32Vec16(const FP16Vec16& v) {
-    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
-      reg.val[i] = fp16_to_float(v.reg.val[i]);
-    }
-  }
-
-  explicit FP32Vec16(const BF16Vec16& v) {
-    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
-      reg.val[i] = bf16_to_float(v.reg.val[i]);
-    }
-  }
-
-  explicit FP32Vec16(const FP16Vec8& v) : FP32Vec16(FP32Vec8(v)) {};
-
-  FP32Vec16(const BF16Vec8& v) : FP32Vec16(FP32Vec8(v)) {};
-
-  FP32Vec16 operator*(const FP32Vec16& b) const {
-    FP32Vec16 result(0.0f);
-    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
-      result.reg.val[i] = reg.val[i] * b.reg.val[i];
-    }
-    return result;
-  }
-
-  FP32Vec16 operator+(const FP32Vec16& b) const {
-    FP32Vec16 result(0.0f);
-    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
-      result.reg.val[i] = reg.val[i] + b.reg.val[i];
-    }
-    return result;
-  }
-
-  FP32Vec16 operator-(const FP32Vec16& b) const {
-    FP32Vec16 result(0.0f);
-    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
-      result.reg.val[i] = reg.val[i] - b.reg.val[i];
-    }
-    return result;
-  }
-
-  FP32Vec16 operator/(const FP32Vec16& b) const {
-    FP32Vec16 result(0.0f);
-    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
-      result.reg.val[i] = reg.val[i] / b.reg.val[i];
-    }
-    return result;
-  }
-
-  FP32Vec16 max(const FP32Vec16& b) const {
-    FP32Vec16 result(0.0f);
-    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
-      result.reg.val[i] = __max(reg.val[i], b.reg.val[i]);
-    }
-    return result;
-  }
-
-  FP32Vec16 min(const FP32Vec16& b) const {
-    FP32Vec16 result(0.0f);
-    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
-      result.reg.val[i] = __min(reg.val[i], b.reg.val[i]);
-    }
-    return result;
-  }
-
-  FP32Vec16 abs() const {
-    FP32Vec16 result(0.0f);
-    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
-      result.reg.val[i] = __abs(reg.val[i]);
-    }
-    return result;
-  }
-
-  float reduce_sum() const {
-    float result = 0.0f;
-    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
-      result += reg.val[i];
-    }
-    return result;
-  }
-
-  float reduce_max() const {
-    float result = reg.val[0];
-    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
-      result = __max(reg.val[i], result);
-    }
-    return result;
-  }
-
-  float reduce_min() const {
-    float result = reg.val[0];
-    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
-      result = __min(reg.val[i], result);
-    }
-    return result;
-  }
-
-  template <int group_size>
-  float reduce_sub_sum(int idx) {
-    static_assert(VEC_ELEM_NUM % group_size == 0);
-    float sum = 0.0;
-    int start = idx * group_size;
-    int end = (idx + 1) * group_size;
-
-    for (; (start < VEC_ELEM_NUM) && (start < end); ++start) {
-      sum += reg.val[start];
-    }
-
-    return sum;
-  }
-
-  void save(void* ptr) const { *reinterpret_cast<f32x16_t*>(ptr) = reg; }
-};
-
-template <typename T>
-struct VecType {
-  using vec_type = void;
-};
-
-template <typename T>
-using vec_t = typename VecType<T>::vec_type;
-
-template <>
-struct VecType<float> {
-  using vec_type = FP32Vec8;
-};
-
-template <>
-struct VecType<c10::Half> {
-  using vec_type = FP16Vec8;
-};
-
-template <>
-struct VecType<c10::BFloat16> {
-  using vec_type = BF16Vec8;
-};
-
-template <typename T>
-void storeFP32(float v, T* ptr) {
-  *ptr = v;
-}
-
-/*
-template <> inline void storeFP32<c10::Half>(float v, c10::Half *ptr) {
-  c10::Half __attribute__((__may_alias__)) *v_ptr =
-      reinterpret_cast<c10::Half *>(&v);
-  *ptr = *(v_ptr + 1);
-}
-*/
-
-template <>
-inline void storeFP32<c10::Half>(float v, c10::Half* ptr) {
-  uint16_t fp16 = float_to_fp16(v);
-  *reinterpret_cast<uint16_t*>(ptr) = fp16;
-}
-
-template <>
-inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
-  c10::BFloat16 __attribute__((__may_alias__))* v_ptr =
-      reinterpret_cast<c10::BFloat16*>(&v);
-  *ptr = *(v_ptr + 1);
-}
-
-inline FP16Vec16::FP16Vec16(const FP32Vec16& v) {
-  int i = 0;
-  for (i = 0; i < FP16Vec16::VEC_ELEM_NUM; ++i) {
-    reg.val[i] = float_to_fp16(v.reg.val[i]);
-  }
-}
-
-inline FP16Vec8 ::FP16Vec8(const FP32Vec8& v) {
-  int i = 0;
-  for (i = 0; i < FP16Vec8::VEC_ELEM_NUM; ++i) {
-    reg.val[i] = float_to_fp16(v.reg.val[i]);
-  }
-}
-
-inline void fma(FP32Vec16& acc, FP32Vec16& a, FP32Vec16& b) {
-  acc = acc + a * b;
-}
-
-inline BF16Vec8::BF16Vec8(const FP32Vec8& v) {
-  int i = 0;
-  for (i = 0; i < BF16Vec8::VEC_ELEM_NUM; ++i) {
-    reg.val[i] = float_to_bf16(v.reg.val[i]);
-  }
-}
-
-inline BF16Vec16::BF16Vec16(const FP32Vec16& v) {
-  int i = 0;
-  for (i = 0; i < BF16Vec16::VEC_ELEM_NUM; ++i) {
-    reg.val[i] = float_to_bf16(v.reg.val[i]);
-  }
-}
-
-inline void prefetch(const void* addr) { __builtin_prefetch(addr, 0, 3); }
-
-};  // namespace vec_op

csrc/cpu/cpu_types_vxe.hpp

@@ -12,7 +12,7 @@ namespace vec_op {
 #define vec_sub(a, b) ((a) - (b))
 #define vec_mul(a, b) ((a) * (b))
 #define vec_div(a, b) ((a) / (b))
-#define vec_sr(a, b) ((a) >> (b))  // Vector Shift Right Algebraic
+#define vec_sr(a, b) ((a) >> (b))  // Vector Shift Right Algebaic
 #define vec_sl(a, b) ((a) << (b))  // Vector Shift Left
 
 // FIXME: FP16 is not fully supported in Torch-CPU

csrc/cpu/dnnl_kernels.cpp

@@ -523,7 +523,7 @@ void onednn_mm(torch::Tensor& c,        // [M, OC], row-major
   CPU_KERNEL_GUARD_IN(onednn_mm)
   TORCH_CHECK(a.dim() == 2);
   TORCH_CHECK(a.stride(-1) == 1);
-  TORCH_CHECK(c.stride(-1) == 1);
+  TORCH_CHECK(c.is_contiguous());
   MatMulPrimitiveHandler* ptr =
       reinterpret_cast<MatMulPrimitiveHandler*>(handler);
 

csrc/cpu/float_convert.hpp

@@ -1,106 +0,0 @@
-
-static float bf16_to_float(uint16_t bf16) {
-  uint32_t bits = static_cast<uint32_t>(bf16) << 16;
-  float fp32;
-  std::memcpy(&fp32, &bits, sizeof(fp32));
-  return fp32;
-}
-
-static uint16_t float_to_bf16(float fp32) {
-  uint32_t bits;
-  std::memcpy(&bits, &fp32, sizeof(fp32));
-  return static_cast<uint16_t>(bits >> 16);
-}
-
-/************************************************
- * Copyright (c) 2015 Princeton Vision Group
- * Licensed under the MIT license.
- * Codes below copied from
- * https://github.com/PrincetonVision/marvin/tree/master/tools/tensorIO_matlab
- *************************************************/
-static uint16_t float_to_fp16(float fp32) {
-  uint16_t fp16;
-
-  unsigned x;
-  unsigned u, remainder, shift, lsb, lsb_s1, lsb_m1;
-  unsigned sign, exponent, mantissa;
-
-  std::memcpy(&x, &fp32, sizeof(fp32));
-  u = (x & 0x7fffffff);
-
-  // Get rid of +NaN/-NaN case first.
-  if (u > 0x7f800000) {
-    fp16 = 0x7fffU;
-    return fp16;
-  }
-
-  sign = ((x >> 16) & 0x8000);
-
-  // Get rid of +Inf/-Inf, +0/-0.
-  if (u > 0x477fefff) {
-    fp16 = sign | 0x7c00U;
-    return fp16;
-  }
-  if (u < 0x33000001) {
-    fp16 = (sign | 0x0000);
-    return fp16;
-  }
-
-  exponent = ((u >> 23) & 0xff);
-  mantissa = (u & 0x7fffff);
-
-  if (exponent > 0x70) {
-    shift = 13;
-    exponent -= 0x70;
-  } else {
-    shift = 0x7e - exponent;
-    exponent = 0;
-    mantissa |= 0x800000;
-  }
-  lsb = (1 << shift);
-  lsb_s1 = (lsb >> 1);
-  lsb_m1 = (lsb - 1);
-
-  // Round to nearest even.
-  remainder = (mantissa & lsb_m1);
-  mantissa >>= shift;
-  if (remainder > lsb_s1 || (remainder == lsb_s1 && (mantissa & 0x1))) {
-    ++mantissa;
-    if (!(mantissa & 0x3ff)) {
-      ++exponent;
-      mantissa = 0;
-    }
-  }
-
-  fp16 = (sign | (exponent << 10) | mantissa);
-
-  return fp16;
-}
-
-static float fp16_to_float(uint16_t fp16) {
-  unsigned sign = ((fp16 >> 15) & 1);
-  unsigned exponent = ((fp16 >> 10) & 0x1f);
-  unsigned mantissa = ((fp16 & 0x3ff) << 13);
-  int temp;
-  float fp32;
-  if (exponent == 0x1f) { /* NaN or Inf */
-    mantissa = (mantissa ? (sign = 0, 0x7fffff) : 0);
-    exponent = 0xff;
-  } else if (!exponent) { /* Denorm or Zero */
-    if (mantissa) {
-      unsigned int msb;
-      exponent = 0x71;
-      do {
-        msb = (mantissa & 0x400000);
-        mantissa <<= 1; /* normalize */
-        --exponent;
-      } while (!msb);
-      mantissa &= 0x7fffff; /* 1.mantissa is implicit */
-    }
-  } else {
-    exponent += 0x70;
-  }
-  temp = ((sign << 31) | (exponent << 23) | mantissa);
-  std::memcpy(&fp32, &temp, sizeof(temp));
-  return fp32;
-}

csrc/cpu/sgl-kernels/moe.cpp

@@ -215,7 +215,7 @@ int moe_align_block_size(
       offsets[mb + 1] = sorted_id_size(sorted_ids + mb * BLOCK_M);
     }
   });
-  // TODO: do we need to vectorize this ?
+  // TODO: do we need to vecterize this ?
   for (int mb = 0; mb < num_token_blocks; ++mb) {
     offsets[mb + 1] += offsets[mb];
   }

csrc/cpu/torch_bindings.cpp

@@ -88,18 +88,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       "    int tp_rank, int blocksparse_local_blocks,"
       "    int blocksparse_vert_stride, int blocksparse_block_size,"
       "    int blocksparse_head_sliding_step) -> ()");
-
   ops.impl("paged_attention_v1", torch::kCPU, &paged_attention_v1);
 
-  ops.def(
-      "dynamic_4bit_int_moe("
-      "Tensor x, Tensor topk_ids, Tensor topk_weights,"
-      "Tensor w13_packed, Tensor w2_packed, int H, int I, int I2,"
-      "int group_size, bool apply_router_weight_on_input, int activation_kind"
-      ") -> Tensor");
-
-  ops.impl("dynamic_4bit_int_moe", torch::kCPU, &dynamic_4bit_int_moe_cpu);
-
   // PagedAttention V2.
   ops.def(
       "paged_attention_v2("

csrc/cub_helpers.h

@@ -1,17 +0,0 @@
-#pragma once
-
-#ifndef USE_ROCM
-  #include <cub/cub.cuh>
-  #if CUB_VERSION >= 200800
-    #include <cuda/std/functional>
-using CubAddOp = cuda::std::plus<>;
-using CubMaxOp = cuda::maximum<>;
-  #else   // if CUB_VERSION < 200800
-using CubAddOp = cub::Sum;
-using CubMaxOp = cub::Max;
-  #endif  // CUB_VERSION
-#else
-  #include <hipcub/hipcub.hpp>
-using CubAddOp = cub::Sum;
-using CubMaxOp = cub::Max;
-#endif  // USE_ROCM

csrc/launch_bounds_utils.h

@@ -1,38 +0,0 @@
-#pragma once
-
-#include <cuda_runtime_api.h>
-#include <algorithm>
-
-// maximum blocks per SM cap
-#ifndef VLLM_LAUNCH_BLOCKS_CAP
-  #define VLLM_LAUNCH_BLOCKS_CAP 4
-#endif
-
-// compile-time estimate of max threads per SM for launch bounds.
-#ifndef VLLM_MAX_THREADS_PER_SM
-  #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 300
-    #define VLLM_MAX_THREADS_PER_SM 1536
-  #else
-    #define VLLM_MAX_THREADS_PER_SM 2048
-  #endif
-#endif
-
-// compute the number of blocks per SM to request in __launch_bounds__
-#define VLLM_BLOCKS_DIV(VAL) (VLLM_MAX_THREADS_PER_SM / (VAL))
-#define VLLM_CLAMP_BLOCKS_PER_SM(VAL) \
-  (((VAL) <= 0)                       \
-       ? 1                            \
-       : (((VAL) < VLLM_LAUNCH_BLOCKS_CAP) ? (VAL) : VLLM_LAUNCH_BLOCKS_CAP))
-#define VLLM_BLOCKS_PER_SM(BLOCK_THREADS) \
-  VLLM_CLAMP_BLOCKS_PER_SM(VLLM_BLOCKS_DIV(BLOCK_THREADS))
-
-// runtime-time helper to compute blocks/SM
-static inline int vllm_runtime_blocks_per_sm(int block_threads) {
-  int device = -1;
-  cudaGetDevice(&device);
-  int max_threads_per_sm = VLLM_MAX_THREADS_PER_SM;
-  cudaDeviceGetAttribute(&max_threads_per_sm,
-                         cudaDevAttrMaxThreadsPerMultiProcessor, device);
-  int blocks = (block_threads > 0) ? (max_threads_per_sm / block_threads) : 1;
-  return VLLM_CLAMP_BLOCKS_PER_SM(blocks);
-}

csrc/layernorm_kernels.cu

@@ -1,10 +1,15 @@
 #include "type_convert.cuh"
 #include "dispatch_utils.h"
-#include "cub_helpers.h"
 
 #include <torch/cuda.h>
 #include <c10/cuda/CUDAGuard.h>
 
+#ifndef USE_ROCM
+  #include <cub/cub.cuh>
+#else
+  #include <hipcub/hipcub.hpp>
+#endif
+
 namespace vllm {
 
 // TODO(woosuk): Further optimize this kernel.
@@ -25,7 +30,7 @@ __global__ void rms_norm_kernel(
 
   using BlockReduce = cub::BlockReduce<float, 1024>;
   __shared__ typename BlockReduce::TempStorage reduceStore;
-  variance = BlockReduce(reduceStore).Reduce(variance, CubAddOp{}, blockDim.x);
+  variance = BlockReduce(reduceStore).Reduce(variance, cub::Sum{}, blockDim.x);
 
   if (threadIdx.x == 0) {
     s_variance = rsqrtf(variance / hidden_size + epsilon);
@@ -80,7 +85,7 @@ fused_add_rms_norm_kernel(
 
   using BlockReduce = cub::BlockReduce<float, 1024>;
   __shared__ typename BlockReduce::TempStorage reduceStore;
-  variance = BlockReduce(reduceStore).Reduce(variance, CubAddOp{}, blockDim.x);
+  variance = BlockReduce(reduceStore).Reduce(variance, cub::Sum{}, blockDim.x);
 
   if (threadIdx.x == 0) {
     s_variance = rsqrtf(variance / hidden_size + epsilon);
@@ -121,7 +126,7 @@ fused_add_rms_norm_kernel(
 
   using BlockReduce = cub::BlockReduce<float, 1024>;
   __shared__ typename BlockReduce::TempStorage reduceStore;
-  variance = BlockReduce(reduceStore).Reduce(variance, CubAddOp{}, blockDim.x);
+  variance = BlockReduce(reduceStore).Reduce(variance, cub::Sum{}, blockDim.x);
 
   if (threadIdx.x == 0) {
     s_variance = rsqrtf(variance / hidden_size + epsilon);

csrc/layernorm_quant_kernels.cu

@@ -8,11 +8,16 @@
 #include "type_convert.cuh"
 #include "quantization/fp8/common.cuh"
 #include "dispatch_utils.h"
-#include "cub_helpers.h"
 
 #include <torch/cuda.h>
 #include <c10/cuda/CUDAGuard.h>
 
+#ifndef USE_ROCM
+  #include <cub/cub.cuh>
+#else
+  #include <hipcub/hipcub.hpp>
+#endif
+
 namespace vllm {
 
 // TODO(woosuk): Further optimize this kernel.
@@ -34,7 +39,7 @@ __global__ void rms_norm_static_fp8_quant_kernel(
 
   using BlockReduce = cub::BlockReduce<float, 1024>;
   __shared__ typename BlockReduce::TempStorage reduceStore;
-  variance = BlockReduce(reduceStore).Reduce(variance, CubAddOp{}, blockDim.x);
+  variance = BlockReduce(reduceStore).Reduce(variance, cub::Sum{}, blockDim.x);
 
   if (threadIdx.x == 0) {
     s_variance = rsqrtf(variance / hidden_size + epsilon);
@@ -95,7 +100,7 @@ fused_add_rms_norm_static_fp8_quant_kernel(
 
   using BlockReduce = cub::BlockReduce<float, 1024>;
   __shared__ typename BlockReduce::TempStorage reduceStore;
-  variance = BlockReduce(reduceStore).Reduce(variance, CubAddOp{}, blockDim.x);
+  variance = BlockReduce(reduceStore).Reduce(variance, cub::Sum{}, blockDim.x);
 
   if (threadIdx.x == 0) {
     s_variance = rsqrtf(variance / hidden_size + epsilon);
@@ -144,7 +149,7 @@ fused_add_rms_norm_static_fp8_quant_kernel(
 
   using BlockReduce = cub::BlockReduce<float, 1024>;
   __shared__ typename BlockReduce::TempStorage reduceStore;
-  variance = BlockReduce(reduceStore).Reduce(variance, CubAddOp{}, blockDim.x);
+  variance = BlockReduce(reduceStore).Reduce(variance, cub::Sum{}, blockDim.x);
 
   if (threadIdx.x == 0) {
     s_variance = rsqrtf(variance / hidden_size + epsilon);

csrc/moe/dynamic_4bit_int_moe_cpu.cpp

@@ -1,156 +0,0 @@
-#include <ATen/ATen.h>
-#include <ATen/Parallel.h>
-#include <torch/all.h>
-
-// _dyn_quant_matmul_4bit is only available on AArch64.
-#if defined(__aarch64__)
-  #include <ATen/ops/_dyn_quant_matmul_4bit.h>
-#endif
-
-inline torch::Tensor mm(const torch::Tensor& a, const torch::Tensor& packed_w,
-                        int64_t group_size_eff, int64_t in_features,
-                        int64_t out_features) {
-#if defined(__aarch64__)
-  return at::_ops::_dyn_quant_matmul_4bit::call(a, packed_w, group_size_eff,
-                                                in_features, out_features);
-#else
-  TORCH_CHECK(false,
-              "dynamic 4-bit int MoE path requires AArch64 (ARM64); "
-              "_dyn_quant_matmul_4bit is unavailable on this architecture");
-  return {};
-#endif
-}
-
-enum ActivationKind : int64_t {
-  SwiGLU_Gu = 0,  // act = SiLU(g) * u
-  SwiGLUOAI = 1,  // act = SiLU(u) * g
-  SiLU = 2        // SiLU
-};
-
-torch::Tensor dynamic_4bit_int_moe_cpu(
-    torch::Tensor x, torch::Tensor topk_ids, torch::Tensor topk_weights,
-    torch::Tensor w13_packed, torch::Tensor w2_packed, int64_t H, int64_t I,
-    int64_t I2, int64_t group_size, bool apply_router_weight_on_input,
-    int64_t activation_kind) {
-  TORCH_CHECK(x.dim() == 2, "x must be 2D");
-  TORCH_CHECK(topk_ids.dim() == 2 && topk_weights.dim() == 2,
-              "topk tensors must be [T, K]");
-  TORCH_CHECK(
-      w13_packed.size(0) == w2_packed.size(0),
-      "w13_packed and w2_packed must have same number of experts in dim 0");
-  TORCH_CHECK(I2 == 2 * I, "I2 must equal 2*I");
-
-  const int64_t T = x.size(0);
-  const int64_t K = topk_ids.size(1);
-  const int64_t E = w13_packed.size(0);
-  const int64_t N = T * K;
-
-  auto x_c = x.contiguous();
-  auto ids_c = topk_ids.contiguous();
-  auto gates_c = topk_weights.to(at::kFloat).contiguous();
-
-  // bucketing tokens -> experts
-  c10::SmallVector<int64_t, 64> counts(
-      E, 0);  // Small vector uses stack allocation
-  {
-    const auto* ids_ptr = ids_c.data_ptr<int64_t>();
-    for (int64_t i = 0; i < N; ++i) {
-      const int64_t e_id = ids_ptr[i];
-      TORCH_CHECK(0 <= e_id && e_id < E, "expert id out of range");
-      counts[e_id]++;
-    }
-  }
-  c10::SmallVector<int64_t, 65> offsets(E + 1, 0);  // ( E +1 )
-  for (int64_t e = 0; e < E; ++e) offsets[e + 1] = offsets[e] + counts[e];
-
-  auto expert_tokens = at::empty({offsets[E]}, ids_c.options());
-  auto expert_gates = at::empty({offsets[E]}, gates_c.options());
-  {
-    c10::SmallVector<int64_t, 64> cursor(E, 0);
-    const auto* ids_ptr = ids_c.data_ptr<int64_t>();
-    const auto* gts_ptr = gates_c.data_ptr<float>();
-    auto* tok_ptr = expert_tokens.data_ptr<int64_t>();
-    auto* gate_ptr = expert_gates.data_ptr<float>();
-
-    for (int64_t t = 0; t < T; ++t) {
-      const int64_t base = t * K;
-      for (int64_t k = 0; k < K; ++k) {
-        const int64_t idx = base + k;
-        const int64_t e = ids_ptr[idx];
-        const int64_t p = offsets[e] + (cursor[e]++);
-        tok_ptr[p] = t;
-        gate_ptr[p] = gts_ptr[idx];
-      }
-    }
-  }
-
-  const int64_t g_eff_13 = (group_size != -1) ? group_size : H;
-  const int64_t g_eff_2 = (group_size != -1) ? group_size : I;
-
-  // Per-expert outputs filled in parallel
-  std::vector<torch::Tensor> y_list(E);
-  y_list.resize(E);
-
-  at::parallel_for(0, E, 1, [&](int64_t e_begin, int64_t e_end) {
-    for (int64_t e = e_begin; e < e_end; ++e) {
-      const int64_t te = counts[e];
-      if (te == 0) {
-        y_list[e] = at::empty({0, H}, x_c.options());
-        continue;
-      }
-
-      const int64_t start = offsets[e];
-
-      auto sel_tokens =
-          expert_tokens.narrow(/*dim=*/0, /*start=*/start, /*length=*/te);
-      auto gates_e =
-          expert_gates.narrow(/*dim=*/0, /*start=*/start, /*length=*/te);
-
-      auto x_e = x_c.index_select(/*dim=*/0, sel_tokens);
-
-      if (apply_router_weight_on_input) {
-        x_e = x_e.mul(gates_e.unsqueeze(1));
-      }
-
-      auto w13_e = w13_packed.select(/*dim=*/0, e);
-      auto w2_e = w2_packed.select(/*dim=*/0, e);
-
-      // W13
-      auto y13 =
-          mm(x_e, w13_e, g_eff_13, /*in_features=*/H, /*out_features=*/I2);
-
-      auto g_part = y13.narrow(/*dim=*/1, /*start=*/0, /*length=*/I);
-      auto u_part = y13.narrow(/*dim=*/1, /*start=*/I, /*length=*/I);
-
-      torch::Tensor act;
-      if (activation_kind == ActivationKind::SwiGLUOAI) {  // SwiGLUOAI
-        constexpr double kAlpha = 1.702;                   // GPT-OSS default
-        constexpr double kLimit = 7.0;                     // GPT-OSS default
-        auto gate_c = at::clamp_max(g_part, kLimit);
-        auto up_c = at::clamp(u_part, -kLimit, kLimit);
-        auto glu = gate_c.mul(at::sigmoid(gate_c.mul(kAlpha)));
-        act = up_c.add(1.0).mul(glu);
-      } else {  // SiLU , SwiGLU_GU, vLLM maps silu to SiluAndMul()
-        act = at::silu(g_part).mul(u_part);
-      }
-
-      // W2
-      auto y = mm(act, w2_e, g_eff_2, /*in_features=*/I, /*out_features=*/H);
-
-      if (!apply_router_weight_on_input) {
-        y = y.mul(gates_e.unsqueeze(1));
-      }
-
-      // Store per-expert result
-      y_list[e] = y;
-    }
-  });
-
-  // Concatenate all expert outputs to match expert_tokens order
-  auto Y_all = at::cat(y_list, /*dim=*/0);
-  auto out = at::zeros({T, H}, x.options());
-  out =
-      at::index_add(out, /*dim=*/0, /*index=*/expert_tokens, /*source=*/Y_all);
-
-  return out;
-}

csrc/moe/grouped_topk_kernels.cu

@@ -21,7 +21,6 @@
 #include <torch/all.h>
 #include <cuda_fp16.h>
 #include <cuda_bf16.h>
-#include <cuda/std/limits>
 #include <cooperative_groups.h>
 #include <cooperative_groups/reduce.h>
 namespace cg = cooperative_groups;
@@ -29,6 +28,7 @@ namespace cg = cooperative_groups;
 namespace vllm {
 namespace moe {
 
+constexpr float kNegInfinity = INFINITY * -1;
 constexpr unsigned FULL_WARP_MASK = 0xffffffff;
 constexpr int32_t WARP_SIZE = 32;
 constexpr int32_t BLOCK_SIZE = 512;
@@ -411,30 +411,14 @@ __device__ inline float cuda_cast<float, __nv_bfloat16>(__nv_bfloat16 val) {
   return __bfloat162float(val);
 }
 
-template <typename T>
-__device__ inline T neg_inf() {
-  // cuda::std::numeric_limits<T>::infinity() returns `0` for [T=bf16 or fp16]
-  // so we need to cast from fp32
-  return cuda_cast<T, float>(-cuda::std::numeric_limits<float>::infinity());
-}
-
-template <typename T>
-__device__ inline bool is_finite(const T val) {
-#if (__CUDACC_VER_MAJOR__ * 10000 + __CUDACC_VER_MINOR__ * 100 >= 120800)
-  return cuda::std::isfinite(val);
-#else
-  return isfinite(cuda_cast<float, T>(val));
-#endif
-}
-
 template <typename T>
 __device__ void topk_with_k2(T* output, T const* input,
                              cg::thread_block_tile<32> const& tile,
                              int32_t const lane_id,
                              int const num_experts_per_group) {
   // Get the top2 per thread
-  T largest = neg_inf<T>();
-  T second_largest = neg_inf<T>();
+  T largest = -INFINITY;
+  T second_largest = -INFINITY;
 
   if (num_experts_per_group > WARP_SIZE) {
     for (int i = lane_id; i < num_experts_per_group; i += WARP_SIZE) {
@@ -529,8 +513,8 @@ __global__ void group_idx_and_topk_idx_kernel(
       warp_id * topk;
   s_topk_idx += warp_id * topk;
 
-  T value = neg_inf<T>();
-  T topk_group_value = neg_inf<T>();
+  T value = kNegInfinity;
+  T topk_group_value = kNegInfinity;
   int32_t num_equalto_topkth_group;
 
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
@@ -541,8 +525,11 @@ __global__ void group_idx_and_topk_idx_kernel(
   if (case_id < num_tokens) {
     // calculate group_idx
     int32_t target_num_min = WARP_SIZE - n_group + topk_group;
-    // The check is necessary to avoid abnormal input
-    if (lane_id < n_group && is_finite(group_scores[lane_id])) {
+    if (lane_id < n_group &&
+        (isfinite(cuda_cast<float, T>(
+            group_scores[lane_id]))))  // The check is necessary to avoid
+                                       // abnormal input
+    {
       value = group_scores[lane_id];
     }
 
@@ -553,11 +540,11 @@ __global__ void group_idx_and_topk_idx_kernel(
       __syncwarp();  // Ensure all threads have valid data before reduction
       topk_group_value = cg::reduce(tile, value, cg::greater<T>());
       if (value == topk_group_value) {
-        value = neg_inf<T>();
+        value = kNegInfinity;
       }
       pre_count_equal_to_top_value = count_equal_to_top_value;
-      count_equal_to_top_value =
-          __popc(__ballot_sync(FULL_WARP_MASK, (value == neg_inf<T>())));
+      count_equal_to_top_value = __popc(__ballot_sync(
+          FULL_WARP_MASK, (value == cuda_cast<T, float>(kNegInfinity))));
     }
     num_equalto_topkth_group = target_num_min - pre_count_equal_to_top_value;
   }
@@ -565,10 +552,11 @@ __global__ void group_idx_and_topk_idx_kernel(
 
   warp_topk::WarpSelect</*capability*/ WARP_SIZE, /*greater*/ true, T, int32_t,
                         /* is_stable */ true>
-      queue((int32_t)topk, neg_inf<T>());
+      queue((int32_t)topk, -INFINITY);
 
   int count_equalto_topkth_group = 0;
-  bool if_proceed_next_topk = topk_group_value != neg_inf<T>();
+  bool if_proceed_next_topk =
+      (topk_group_value != cuda_cast<T, float>(kNegInfinity));
   if (case_id < num_tokens && if_proceed_next_topk) {
     for (int i_group = 0; i_group < n_group; i_group++) {
       if ((group_scores[i_group] > topk_group_value) ||
@@ -577,10 +565,11 @@ __global__ void group_idx_and_topk_idx_kernel(
         int32_t offset = i_group * num_experts_per_group;
         for (int32_t i = lane_id; i < align_num_experts_per_group;
              i += WARP_SIZE) {
-          T candidates = (i < num_experts_per_group) &&
-                                 is_finite(scores_with_bias[offset + i])
-                             ? scores_with_bias[offset + i]
-                             : neg_inf<T>();
+          T candidates =
+              (i < num_experts_per_group) && isfinite(cuda_cast<float, T>(
+                                                 scores_with_bias[offset + i]))
+                  ? scores_with_bias[offset + i]
+                  : cuda_cast<T, float>(kNegInfinity);
           queue.add(candidates, offset + i);
         }
         if (group_scores[i_group] == topk_group_value) {
@@ -609,8 +598,7 @@ __global__ void group_idx_and_topk_idx_kernel(
       if (i < topk) {
         s_topk_value[i] = value;
       }
-      topk_sum +=
-          cg::reduce(tile, cuda_cast<float, T>(value), cg::plus<float>());
+      topk_sum += reduce(tile, cuda_cast<float, T>(value), cg::plus<float>());
     }
   }
 

csrc/moe/moe_align_sum_kernels.cu

@@ -44,9 +44,6 @@ __global__ void moe_align_block_size_kernel(
 
   for (size_t i = tid; i < numel; i += stride) {
     int expert_id = topk_ids[i];
-    if (expert_id >= num_experts) {
-      continue;
-    }
     int warp_idx = expert_id / experts_per_warp;
     int expert_offset = expert_id % experts_per_warp;
     atomicAdd(&shared_counts[warp_idx * experts_per_warp + expert_offset], 1);
@@ -98,15 +95,12 @@ template <typename scalar_t>
 __global__ void count_and_sort_expert_tokens_kernel(
     const scalar_t* __restrict__ topk_ids,
     int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ cumsum_buffer,
-    size_t numel, int32_t num_experts) {
+    size_t numel) {
   const size_t tid = blockIdx.x * blockDim.x + threadIdx.x;
   const size_t stride = blockDim.x * gridDim.x;
 
   for (size_t i = tid; i < numel; i += stride) {
     int32_t expert_id = topk_ids[i];
-    if (expert_id >= num_experts) {
-      continue;
-    }
     int32_t rank_post_pad = atomicAdd(&cumsum_buffer[expert_id], 1);
     sorted_token_ids[rank_post_pad] = i;
   }
@@ -275,7 +269,7 @@ void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
           sort_kernel<<<actual_blocks, block_threads, 0, stream>>>(
               topk_ids.data_ptr<scalar_t>(),
               sorted_token_ids.data_ptr<int32_t>(),
-              cumsum_buffer.data_ptr<int32_t>(), topk_ids.numel(), num_experts);
+              cumsum_buffer.data_ptr<int32_t>(), topk_ids.numel());
         }
       });
 }

csrc/moe/topk_softmax_kernels.cu

@@ -20,7 +20,17 @@
 #include <ATen/cuda/CUDAContext.h>
 #include <c10/cuda/CUDAGuard.h>
 #include "../cuda_compat.h"
-#include "../cub_helpers.h"
+
+#ifndef USE_ROCM
+    #include <cub/util_type.cuh>
+    #include <cub/cub.cuh>
+    #include <cuda/std/functional>
+    using AddOp = cuda::std::plus<float>;
+#else
+    #include <hipcub/util_type.hpp>
+    #include <hipcub/hipcub.hpp>
+    using AddOp = cub::Sum; 
+#endif
 
 #define MAX(a, b) ((a) > (b) ? (a) : (b))
 #define MIN(a, b) ((a) < (b) ? (a) : (b))
@@ -69,7 +79,7 @@ __launch_bounds__(TPB) __global__
         threadData = max(static_cast<float>(input[idx]), threadData);
     }
 
-    const float maxElem = BlockReduce(tmpStorage).Reduce(threadData, CubMaxOp());
+    const float maxElem = BlockReduce(tmpStorage).Reduce(threadData, cub::Max());
     if (threadIdx.x == 0)
     {
         float_max = maxElem;
@@ -84,7 +94,7 @@ __launch_bounds__(TPB) __global__
         threadData += exp((static_cast<float>(input[idx]) - float_max));
     }
 
-    const auto Z = BlockReduce(tmpStorage).Reduce(threadData, CubAddOp());
+    const auto Z = BlockReduce(tmpStorage).Reduce(threadData, AddOp());
 
     if (threadIdx.x == 0)
     {

csrc/ops.h

@@ -122,6 +122,12 @@ void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
                       std::optional<torch::Tensor> key, int64_t head_size,
                       torch::Tensor& cos_sin_cache, bool is_neox);
 
+void batched_rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
+                              std::optional<torch::Tensor> key,
+                              int64_t head_size, torch::Tensor& cos_sin_cache,
+                              bool is_neox, int64_t rot_dim,
+                              torch::Tensor& cos_sin_cache_offsets);
+
 void silu_and_mul(torch::Tensor& out, torch::Tensor& input);
 
 void silu_and_mul_quant(torch::Tensor& out, torch::Tensor& input,
@@ -133,12 +139,6 @@ void silu_and_mul_nvfp4_quant(torch::Tensor& out,
                               torch::Tensor& input,
                               torch::Tensor& input_global_scale);
 #endif
-void silu_mul_fp8_quant_deep_gemm_cuda(
-    const at::Tensor& input,   // (E, T, 2*H)
-    const at::Tensor& counts,  // (E)
-    at::Tensor& y_q,           // (E, T, H) [OUT]
-    at::Tensor& y_s,           // (E, T, H//group_size) [OUT]
-    int64_t group_size, bool use_ue8m0, int64_t num_parallel_tokens);
 
 void mul_and_silu(torch::Tensor& out, torch::Tensor& input);
 
@@ -328,12 +328,6 @@ void selective_scan_fwd(const torch::Tensor& u, const torch::Tensor& delta,
                         const std::optional<torch::Tensor>& has_initial_state,
                         const torch::Tensor& ssm_states, int64_t pad_slot_id);
 
-torch::Tensor dynamic_4bit_int_moe_cpu(
-    torch::Tensor x, torch::Tensor topk_ids, torch::Tensor topk_weights,
-    torch::Tensor w13_packed, torch::Tensor w2_packed, int64_t H, int64_t I,
-    int64_t I2, int64_t group_size, bool apply_router_weight_on_input,
-    int64_t activation_kind);
-
 using fptr_t = int64_t;
 fptr_t init_custom_ar(const std::vector<int64_t>& fake_ipc_ptrs,
                       torch::Tensor& rank_data, int64_t rank,
@@ -353,8 +347,6 @@ std::tuple<int64_t, torch::Tensor> allocate_shared_buffer_and_handle(
 int64_t open_mem_handle(torch::Tensor& mem_handle);
 void free_shared_buffer(int64_t buffer);
 
-torch::Tensor hadacore_transform(torch::Tensor& x, bool inplace);
-
 #ifdef USE_ROCM
 fptr_t init_custom_qr(int64_t rank, int64_t world_size,
                       std::optional<int64_t> qr_max_size = std::nullopt);

csrc/pos_encoding_kernels.cu

@@ -99,6 +99,35 @@ __global__ void rotary_embedding_kernel(
       token_idx, query_stride, key_stride, head_stride);
 }
 
+template <typename scalar_t, bool IS_NEOX>
+__global__ void batched_rotary_embedding_kernel(
+    const int64_t* __restrict__ positions,  // [batch_size, seq_len] or
+                                            // [num_tokens]
+    scalar_t* __restrict__ query,           // [batch_size, seq_len, num_heads,
+                                   // head_size] or [num_tokens, num_heads,
+                                   // head_size]
+    scalar_t* __restrict__ key,  // nullptr or
+                                 // [batch_size, seq_len, num_kv_heads,
+                                 // head_size] or [num_tokens, num_kv_heads,
+                                 // head_size]
+    const scalar_t* __restrict__ cos_sin_cache,  // [max_position, 2, rot_dim //
+                                                 // 2]
+    const int64_t* __restrict__ cos_sin_cache_offsets,  // [batch_size, seq_len]
+    const int rot_dim, const int64_t query_stride, const int64_t key_stride,
+    const int64_t head_stride, const int num_heads, const int num_kv_heads,
+    const int head_size) {
+  // Each thread block is responsible for one token.
+  const int token_idx = blockIdx.x;
+  int64_t pos = positions[token_idx];
+  int64_t cos_sin_cache_offset = cos_sin_cache_offsets[token_idx];
+  const scalar_t* cache_ptr =
+      cos_sin_cache + (cos_sin_cache_offset + pos) * rot_dim;
+
+  apply_rotary_embedding<scalar_t, IS_NEOX>(
+      query, key, cache_ptr, head_size, num_heads, num_kv_heads, rot_dim,
+      token_idx, query_stride, key_stride, head_stride);
+}
+
 }  // namespace vllm
 
 void rotary_embedding(
@@ -182,3 +211,96 @@ void rotary_embedding(
     }
   });
 }
+
+/*
+Batched version of rotary embedding, pack multiple LoRAs together
+and process in batched manner.
+*/
+void batched_rotary_embedding(
+    torch::Tensor& positions,  // [batch_size, seq_len] or [num_tokens]
+    torch::Tensor& query,  // [batch_size, seq_len, num_heads * head_size] or
+                           // [num_tokens, num_heads * head_size] or
+                           // [batch_size, seq_len, num_heads, head_size] or
+                           // [num_tokens, num_heads, head_size]
+    std::optional<torch::Tensor>
+        key,  // null or
+              // [batch_size, seq_len, num_kv_heads * head_size] or
+              // [num_tokens, num_kv_heads * head_size] or
+              // [batch_size, seq_len, num_heads, head_size] or
+              // [num_tokens, num_heads, head_size]
+    int64_t head_size,
+    torch::Tensor& cos_sin_cache,  // [max_position, rot_dim]
+    bool is_neox, int64_t rot_dim,
+    torch::Tensor& cos_sin_cache_offsets  // [num_tokens] or [batch_size]
+) {
+  // num_tokens = batch_size * seq_len
+  int64_t num_tokens = cos_sin_cache_offsets.size(0);
+  TORCH_CHECK(
+      positions.size(0) == num_tokens || positions.numel() == num_tokens,
+      "positions must have the same num_tokens or batch_size as "
+      "cos_sin_cache_offsets");
+
+  int positions_ndim = positions.dim();
+  // Make sure num_tokens dim is consistent across positions, query, and key
+  TORCH_CHECK(
+      positions_ndim == 1 || positions_ndim == 2,
+      "positions must have shape [num_tokens] or [batch_size, seq_len]");
+  if (positions_ndim == 1) {
+    TORCH_CHECK(query.size(0) == positions.size(0) &&
+                    (!key.has_value() || key->size(0) == positions.size(0)),
+                "query, key and positions must have the same number of tokens");
+  }
+  if (positions_ndim == 2) {
+    TORCH_CHECK(
+        query.size(0) == positions.size(0) &&
+            (!key.has_value() || key->size(0) == positions.size(0)) &&
+            query.size(1) == positions.size(1) &&
+            (!key.has_value() || key->size(1) == positions.size(1)),
+        "query, key and positions must have the same batch_size and seq_len");
+  }
+
+  // Make sure head_size is valid for query and key
+  int query_hidden_size = query.numel() / num_tokens;
+  int key_hidden_size = key.has_value() ? key->numel() / num_tokens : 0;
+  TORCH_CHECK(query_hidden_size % head_size == 0);
+  TORCH_CHECK(key_hidden_size % head_size == 0);
+
+  // Make sure query and key have concistent number of heads
+  int num_heads = query_hidden_size / head_size;
+  int num_kv_heads = key.has_value() ? key_hidden_size / head_size : num_heads;
+  TORCH_CHECK(num_heads % num_kv_heads == 0);
+
+  int seq_dim_idx = positions_ndim - 1;
+  int64_t query_stride = query.stride(seq_dim_idx);
+  int64_t key_stride = key.has_value() ? key->stride(seq_dim_idx) : 0;
+  // Determine head stride: for [*, heads, head_size] use stride of last dim;
+  // for flat [*, heads*head_size], heads blocks are contiguous of size
+  // head_size
+  int query_ndim = query.dim();
+  int64_t head_stride =
+      (query_ndim == positions_ndim + 2) ? query.stride(-2) : head_size;
+
+  dim3 grid(num_tokens);
+  dim3 block(std::min<int64_t>(num_heads * rot_dim / 2, 512));
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  VLLM_DISPATCH_FLOATING_TYPES(query.scalar_type(), "rotary_embedding", [&] {
+    if (is_neox) {
+      vllm::batched_rotary_embedding_kernel<scalar_t, true>
+          <<<grid, block, 0, stream>>>(
+              positions.data_ptr<int64_t>(), query.data_ptr<scalar_t>(),
+              key.has_value() ? key->data_ptr<scalar_t>() : nullptr,
+              cos_sin_cache.data_ptr<scalar_t>(),
+              cos_sin_cache_offsets.data_ptr<int64_t>(), rot_dim, query_stride,
+              key_stride, head_stride, num_heads, num_kv_heads, head_size);
+    } else {
+      vllm::batched_rotary_embedding_kernel<scalar_t, false>
+          <<<grid, block, 0, stream>>>(
+              positions.data_ptr<int64_t>(), query.data_ptr<scalar_t>(),
+              key.has_value() ? key->data_ptr<scalar_t>() : nullptr,
+              cos_sin_cache.data_ptr<scalar_t>(),
+              cos_sin_cache_offsets.data_ptr<int64_t>(), rot_dim, query_stride,
+              key_stride, head_stride, num_heads, num_kv_heads, head_size);
+    }
+  });
+}

csrc/quantization/activation_kernels.cu

@@ -9,31 +9,6 @@
 
 #include "quantization/fp8/common.cuh"
 
-#include <c10/util/Float8_e4m3fn.h>
-
-#ifndef USE_ROCM
-  #include <cuda_bf16.h>
-  #include <cuda_fp16.h>
-  #include <cuda_fp8.h>
-#else
-  #include <hip/hip_bf16.h>
-  #include <hip/hip_fp16.h>
-  #include <hip/hip_fp8.h>
-
-typedef __hip_bfloat162 __nv_bfloat162;
-typedef __hip_bfloat16 __nv_bfloat16;
-typedef __hip_bfloat16_raw __nv_bfloat16_raw;
-  #if defined(HIP_FP8_TYPE_OCP)
-typedef __hip_fp8_e4m3 __nv_fp8_e4m3;
-typedef __hip_fp8x4_e4m3 __nv_fp8x4_e4m3;
-  #else
-// ROCm 6.2 fallback: only *_fnuz types exist
-typedef __hip_fp8_e4m3_fnuz __nv_fp8_e4m3;
-typedef __hip_fp8x4_e4m3_fnuz __nv_fp8x4_e4m3;
-  #endif
-#endif
-
-#include "core/registration.h"
 namespace vllm {
 
 template <typename T>
@@ -112,336 +87,6 @@ __global__ void act_and_mul_quant_kernel(
     }
   }
 }
-
-__device__ __forceinline__ float silu(float x) {
-  return (__fdividef(x, (1.f + expf(-x))));
-}
-
-__device__ __forceinline__ float2 silu2(float2 x) {
-  return make_float2(silu(x.x), silu(x.y));
-}
-
-#ifndef USE_ROCM
-__device__ __forceinline__ float warp_max(float v) {
-  static constexpr unsigned FULL_MASK = 0xffffffffu;
-  for (int offset = 1; offset < WARP_SIZE; offset *= 2) {
-    v = fmaxf(v, __shfl_xor_sync(FULL_MASK, v, offset));
-  }
-  return v;
-}
-
-__device__ __forceinline__ __nv_bfloat16 warp_max(__nv_bfloat16 v) {
-  static constexpr unsigned FULL_MASK = 0xffffffffu;
-  for (int offset = 1; offset < WARP_SIZE; offset *= 2) {
-    v = __hmax(v, __shfl_xor_sync(FULL_MASK, v, offset));
-  }
-  return v;
-}
-#endif
-
-template <typename T, typename U>
-__device__ __forceinline__ void cp_async4(T* _smem_ptr, const U* _glob_ptr) {
-#if __CUDACC_VER_MAJOR__ >= 11 && __CUDA_ARCH__ >= 800
-  auto smem_ptr = reinterpret_cast<void*>(_smem_ptr);
-  auto glob_ptr = reinterpret_cast<const void*>(_glob_ptr);
-  const int BYTES = 16;
-  uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
-  asm volatile(
-      "{\n"
-      "   cp.async.cg.shared.global [%0], [%1], %2;\n"
-      "}\n" ::"r"(smem),
-      "l"(glob_ptr), "n"(BYTES));
-#else
-  _smem_ptr[0] = _glob_ptr[0];
-#endif
-}
-
-__device__ __forceinline__ void cp_async_fence() {
-#if __CUDACC_VER_MAJOR__ >= 11 && __CUDA_ARCH__ >= 800
-  asm volatile("cp.async.commit_group;\n" ::);
-#else
-#endif
-}
-
-template <int N>
-__device__ __forceinline__ void cp_async_wait() {
-#if __CUDACC_VER_MAJOR__ >= 11 && __CUDA_ARCH__ >= 800
-  asm volatile("cp.async.wait_group %0;\n" ::"n"(N));
-#else
-#endif
-}
-
-template <>
-__device__ __forceinline__ void cp_async_wait<0>() {
-#if __CUDACC_VER_MAJOR__ >= 11 && __CUDA_ARCH__ >= 800
-  asm volatile("cp.async.wait_all;\n" ::);
-#else
-#endif
-}
-
-__device__ __forceinline__ float clip(float v, float mmin, float mmax) {
-#if __CUDACC_VER_MAJOR__ >= 11 && __CUDA_ARCH__ >= 800
-  return fminf(mmax, fmaxf(v, mmin));
-#else
-#endif
-}
-
-__device__ __forceinline__ __nv_bfloat16 clip(__nv_bfloat16 v,
-                                              __nv_bfloat16 mmin,
-                                              __nv_bfloat16 mmax) {
-  return __hmin(mmax, __hmax(v, mmin));
-}
-
-__device__ __forceinline__ __nv_bfloat162 clip(__nv_bfloat162 v,
-                                               __nv_bfloat162 mmin,
-                                               __nv_bfloat162 mmax) {
-  return __hmin2(mmax, __hmax2(v, mmin));
-}
-
-// We use the following values for fp8 min/max:
-//  __nv_fp8_e4m3 = (-448, +448)
-//  __nv_fp8_e4m3uz = (-240.0, +240.0)
-// It is currently assumed that only
-template <class T>
-constexpr __nv_bfloat16 get_fp8_max() {
-  static_assert(std::is_same_v<T, c10::Float8_e4m3fn> ||
-                std::is_same_v<T, c10::Float8_e4m3fnuz>);
-  if constexpr (std::is_same_v<T, c10::Float8_e4m3fn>) {
-    return __nv_bfloat16(__nv_bfloat16_raw{.x = 17376});
-  } else {
-    return __nv_bfloat16(__nv_bfloat16_raw{.x = 17264});
-  }
-}
-
-template <class T>
-constexpr __nv_bfloat16 get_fp8_min() {
-  static_assert(std::is_same_v<T, c10::Float8_e4m3fn> ||
-                std::is_same_v<T, c10::Float8_e4m3fnuz>);
-  if constexpr (std::is_same_v<T, c10::Float8_e4m3fn>) {
-    return __nv_bfloat16(__nv_bfloat16_raw{.x = 50144});
-  } else {
-    return __nv_bfloat16(__nv_bfloat16_raw{.x = 50032});
-  }
-}
-#ifndef USE_ROCM
-template <typename fp8_type, int32_t NUM_WARPS, typename Idx_t,
-          int NUM_PARALLEL_TOKENS, bool USE_UE8M0, int GROUP_SIZE = 128,
-          int NUM_STAGES = 3>
-__global__ void silu_mul_fp8_quant_deep_gemm_kernel(
-    const __nv_bfloat16* __restrict__ _input, fp8_type* __restrict__ _y_q,
-    float* __restrict__ _y_s, const int32_t* __restrict__ counts,
-
-    // sizes
-    int H, int G,
-
-    // strides (in elements)
-    Idx_t stride_i_e, Idx_t stride_i_t, Idx_t stride_i_h, Idx_t stride_yq_e,
-    Idx_t stride_yq_t, Idx_t stride_yq_h, Idx_t stride_ys_e, Idx_t stride_ys_t,
-    Idx_t stride_ys_g, Idx_t stride_counts_e) {
-  static constexpr __nv_bfloat16 fp8_min = get_fp8_min<fp8_type>();
-  static constexpr __nv_bfloat16 fp8_max = get_fp8_max<fp8_type>();
-  // We assign EPS with its 16-bit unsigned counterpart to allow constexpr.
-  static constexpr __nv_bfloat16 EPS = (__nv_bfloat16_raw{.x = 11996});
-
-  // We pack 8 16-bit bfloat16 values into a 128-bit __int128_t.
-  static constexpr int32_t BFLOAT16_PER_GROUP = 8;
-
-  // We split the shared memory in half, corresponding to gate and up matrices:
-  // [...gate_i, ...up_i]  where 0 <= i < stages.
-  static constexpr int32_t S_NUM_128 =
-      2u * (GROUP_SIZE / BFLOAT16_PER_GROUP) * NUM_WARPS * NUM_STAGES;
-  static constexpr auto THREAD_COUNT = NUM_WARPS * WARP_SIZE;
-  static constexpr int HALF_THREAD_COUNT = THREAD_COUNT / 2;
-  static constexpr int32_t S_NUM_64 = S_NUM_128 * 2;
-  __shared__ __int128_t __align__(16) s_buff_128[S_NUM_128];
-
-  const int32_t tid = threadIdx.x;
-  const int32_t warp_id = tid / WARP_SIZE;
-  const int32_t lane_id = tid % WARP_SIZE;
-
-  auto s_buff_compute_32 = reinterpret_cast<__nv_bfloat162*>(s_buff_128);
-
-  // block handles one (expert e, group g)
-  int32_t pid = blockIdx.x;
-  int32_t e = pid / G;
-  int32_t g = pid % G;
-
-  const int32_t n_tokens = counts[e * stride_counts_e];
-
-  if (!n_tokens) {
-    return;  // Exit ASAP.
-  }
-
-  const Idx_t stride_i_t_128 = stride_i_t / 8u;
-
-  int32_t n_tokens_lower, n_tokens_upper;
-
-  // Each block i iterates over tokens of a slice of n_tokens =
-  // expert_counts[i], with the size of chunk being
-  // (n_tokens / NUM_PARALLEL_TOKENS) + residual, instead of
-  // updiv(n_tokens, NUM_PARALLEL_TOKENS) for better scheduling.
-  if (n_tokens < NUM_PARALLEL_TOKENS && blockIdx.y < n_tokens) {
-    // Specialize this, but can be likely fused.
-    if (blockIdx.y >= NUM_PARALLEL_TOKENS) {
-      return;
-    }
-    n_tokens_lower = blockIdx.y;
-    n_tokens_upper = blockIdx.y + 1;
-  } else {
-    auto chunk_size = n_tokens / NUM_PARALLEL_TOKENS;
-    auto residual = n_tokens - chunk_size * NUM_PARALLEL_TOKENS;
-    auto calc_id = [&](int32_t id) {
-      if (id < residual) {
-        return min(n_tokens, id * (chunk_size + 1));
-      } else {
-        return min(n_tokens, id * chunk_size + residual);
-      }
-    };
-    n_tokens_lower = calc_id(blockIdx.y);
-    n_tokens_upper = calc_id(blockIdx.y + 1);
-  }
-
-  if (n_tokens_lower >= n_tokens_upper) {
-    return;
-  }
-
-  // We do calculations here, using constexpr wherever possible.
-  const Idx_t base_i = e * stride_i_e + NUM_WARPS * g * GROUP_SIZE * stride_i_h;
-  const Idx_t base_ys = e * stride_ys_e + NUM_WARPS * g * stride_ys_g;
-  const Idx_t base_yq =
-      e * stride_yq_e + NUM_WARPS * g * GROUP_SIZE * stride_yq_h;
-  Idx_t gate_off_128 = (base_i / static_cast<Idx_t>(8u));
-  auto input_128_ptr = reinterpret_cast<const __int128_t*>(_input);
-  auto gate_128_ptr = input_128_ptr + gate_off_128 + (tid % HALF_THREAD_COUNT) +
-                      stride_i_t_128 * n_tokens_lower;
-  auto up_128_ptr = gate_128_ptr + (H * stride_i_h) / 8u;
-  auto y_s_ptr =
-      _y_s + base_ys + warp_id * stride_ys_g + n_tokens_lower * stride_ys_t;
-  auto y_q_ptr = _y_q + base_yq + warp_id * GROUP_SIZE +
-                 stride_yq_t * n_tokens_lower + 4 * lane_id;
-  int32_t t_load = n_tokens_lower, load_stage_id = 0;
-  auto s_buff_gate_load_128 = s_buff_128 + (tid % HALF_THREAD_COUNT);
-  auto s_buff_up_load_128 = s_buff_gate_load_128 + S_NUM_128 / 2u;
-  int32_t stage_offset{};
-
-  static constexpr int32_t LOAD_STAGE_SIZE = (NUM_WARPS * WARP_SIZE / 2);
-  static constexpr int32_t LOAD_STAGE_MOD =
-      NUM_STAGES * (NUM_WARPS * WARP_SIZE / 2);
-
-  // Two halves of all threads in a block conduct global loads for gate and up,
-  // repsectively.
-  auto load_and_advance_y_pred = [&] {
-    if (t_load < n_tokens_upper) {
-      auto s_gate_stage_128_staged_ptr = s_buff_gate_load_128 + stage_offset;
-      auto s_up_stage_128_staged_ptr = s_buff_up_load_128 + stage_offset;
-
-      // It is very important that LOAD_STAGE_SIZE is constexpr to avoid
-      // unnecessary ALU ops.
-      stage_offset += LOAD_STAGE_SIZE;
-      stage_offset %= LOAD_STAGE_MOD;
-
-      if (tid < HALF_THREAD_COUNT) {
-        cp_async4(s_gate_stage_128_staged_ptr, gate_128_ptr);
-        gate_128_ptr += stride_i_t_128;
-      } else {
-        cp_async4(s_up_stage_128_staged_ptr, up_128_ptr);
-        up_128_ptr += stride_i_t_128;
-      }
-      ++t_load;
-      ++load_stage_id;
-    }
-    // We fence even if there is nothing to load to simplify pipelining.
-    cp_async_fence();
-  };
-
-  #pragma unroll
-  for (int i = 0; i < NUM_STAGES - 1; i++) {
-    load_and_advance_y_pred();
-  }
-
-  __int64_t* s_gate_ptr = reinterpret_cast<__int64_t*>(
-                              s_buff_compute_32 + warp_id * (GROUP_SIZE / 2)) +
-                          lane_id;
-  __int64_t* s_up_ptr = s_gate_ptr + S_NUM_64 / 2;
-
-  static constexpr int32_t STAGE_SIZE = (GROUP_SIZE * NUM_WARPS) / 4u;
-  static constexpr int32_t STAGE_MOD = STAGE_SIZE * NUM_STAGES;
-
-  int32_t compute_pipeline_offset_64 = 0;
-
-  for (int32_t t = n_tokens_lower; t < n_tokens_upper; ++t) {
-    __nv_bfloat162 results_bf162[2];
-
-    cp_async_wait<NUM_STAGES - 2>();
-    __syncthreads();
-
-    // We double-buffer pipelined loads so that the next load will
-    // concurrently run with compute without overwrites.
-    load_and_advance_y_pred();
-
-    auto s_gate_compute_64 = s_gate_ptr + compute_pipeline_offset_64;
-    auto s_up_compute_64 = s_up_ptr + compute_pipeline_offset_64;
-
-    // STAGE_SIZE must also be constexpr!
-    compute_pipeline_offset_64 += STAGE_SIZE;
-    compute_pipeline_offset_64 %= STAGE_MOD;
-
-    // Each thread loads (gate/up) 2X 4X bfloat16 values into registers.
-    __int64_t gate64 = *s_gate_compute_64;
-    __nv_bfloat162* s_gate_compute_32 =
-        reinterpret_cast<__nv_bfloat162*>(&gate64);
-
-    __int64_t up64 = *s_up_compute_64;
-    __nv_bfloat162* s_up_compute_32 = reinterpret_cast<__nv_bfloat162*>(&up64);
-
-  #pragma unroll
-    for (int i = 0; i < 2; i++) {
-      // For silu, we make sure that div is emitted.
-      float2 gate = silu2(__bfloat1622float2(s_gate_compute_32[i]));
-      results_bf162[i] = __float22bfloat162_rn(gate);
-    }
-
-  #pragma unroll
-    for (int i = 0; i < 2; i++) {
-      results_bf162[i] = __hmul2(results_bf162[i], s_up_compute_32[i]);
-    }
-
-    auto _y_max2 =
-        __hmax2(__habs2(results_bf162[0]), __habs2(results_bf162[1]));
-
-    __nv_bfloat16 y_max_bf16 = __hmax(EPS, __hmax(_y_max2.x, _y_max2.y));
-
-    // An entire group is assigned to a single warp, so a simple warp reduce
-    // is used.
-    __nv_bfloat16 y_s = warp_max(y_max_bf16) / fp8_max;
-
-    if constexpr (USE_UE8M0) {
-      y_s = hexp2(hceil(hlog2(y_s)));
-    }
-
-    auto inv_y = __float2bfloat16_rn(1.f) / y_s;
-
-    auto y_s2 = make_bfloat162(inv_y, inv_y);
-
-  #pragma unroll
-    for (int32_t i = 0; i < 2; ++i) {
-      results_bf162[i] =
-          clip(__hmul2(results_bf162[i], y_s2), __bfloat162bfloat162(fp8_min),
-               __bfloat162bfloat162(fp8_max));
-    }
-
-    auto fp8x4 = __nv_fp8x4_e4m3(results_bf162[0], results_bf162[1]);
-    *reinterpret_cast<__nv_fp8x4_e4m3*>(y_q_ptr) = fp8x4;
-    y_q_ptr += stride_yq_t;
-
-    if (lane_id == 0) {
-      *y_s_ptr = y_s;
-      y_s_ptr += stride_ys_t;
-    }
-  }
-}
-#endif
-
 }  // namespace vllm
 
 // Launch activation, gating, and quantize kernel.
@@ -474,117 +119,3 @@ void silu_and_mul_quant(torch::Tensor& out,    // [..., d]
   TORCH_CHECK(input.size(-1) % 2 == 0);
   LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel);
 }
-
-void silu_mul_fp8_quant_deep_gemm_cuda(
-    const at::Tensor& input,   // (E, T, 2*H)
-    const at::Tensor& counts,  // (E)
-    at::Tensor& y_q,           // (E, T, H) [OUT]
-    at::Tensor& y_s,           // (E, T, H//group_size) [OUT]
-    int64_t group_size, bool use_ue8m0, int64_t num_parallel_tokens) {
-#ifndef USE_ROCM
-  // This kernel relies heavily on cp.async and fp8 support.
-  // This kernel currently only supports H % 128 == 0 and assumes a
-  // fixed GROUP_SIZE of 128.
-  TORCH_CHECK(input.dtype() == torch::kBFloat16);
-  TORCH_CHECK(y_q.dtype() == torch::kFloat8_e4m3fn ||
-              y_q.dtype() == torch::kFloat8_e4m3fnuz);
-  TORCH_CHECK(y_s.dtype() == torch::kFloat32);
-  TORCH_CHECK(input.size(-1) % 256 == 0);
-
-  // Check that num_parallel_tokens is of power of 2 and between 1 and 64.
-  TORCH_CHECK(1 <= num_parallel_tokens && num_parallel_tokens <= 64);
-  TORCH_CHECK(!(num_parallel_tokens & (num_parallel_tokens - 1)));
-
-  using Idx_t = int64_t;
-
-  Idx_t E = input.size(0);
-  Idx_t T = input.size(1);
-  Idx_t H = input.size(2) / 2;
-  Idx_t stride_i_e = input.stride(0);
-  Idx_t stride_i_t = input.stride(1);
-  Idx_t stride_i_h = input.stride(2);
-  Idx_t stride_yq_e = y_q.stride(0);
-  Idx_t stride_yq_t = y_q.stride(1);
-  Idx_t stride_yq_h = y_q.stride(2);
-  Idx_t stride_ys_e = y_s.stride(0);
-  Idx_t stride_ys_t = y_s.stride(1);
-  Idx_t stride_ys_g = y_s.stride(2);
-
-  Idx_t stride_counts_e = counts.stride(0);
-
-  static constexpr int GROUP_SIZE = 128;
-
-  #define KERNEL_FN                                                         \
-    if (use_ue8m0) {                                                        \
-      vllm::silu_mul_fp8_quant_deep_gemm_kernel<fp8_t, NUM_WARPS, Idx_t,    \
-                                                NUM_PARALLEL_TOKENS, true>  \
-          <<<grid, block, 0, stream>>>(                                     \
-              reinterpret_cast<__nv_bfloat16*>(input.data_ptr()),           \
-              (fp8_t*)y_q.data_ptr(), y_s.data_ptr<float>(),                \
-              reinterpret_cast<int32_t*>(counts.data_ptr<int>()), H, G,     \
-              stride_i_e, stride_i_t, stride_i_h, stride_yq_e, stride_yq_t, \
-              stride_yq_h, stride_ys_e, stride_ys_t, stride_ys_g,           \
-              stride_counts_e);                                             \
-    } else {                                                                \
-      vllm::silu_mul_fp8_quant_deep_gemm_kernel<fp8_t, NUM_WARPS, Idx_t,    \
-                                                NUM_PARALLEL_TOKENS, false> \
-          <<<grid, block, 0, stream>>>(                                     \
-              reinterpret_cast<__nv_bfloat16*>(input.data_ptr()),           \
-              (fp8_t*)y_q.data_ptr(), y_s.data_ptr<float>(),                \
-              reinterpret_cast<int32_t*>(counts.data_ptr<int>()), H, G,     \
-              stride_i_e, stride_i_t, stride_i_h, stride_yq_e, stride_yq_t, \
-              stride_yq_h, stride_ys_e, stride_ys_t, stride_ys_g,           \
-              stride_counts_e);                                             \
-    }
-
-  #define KERNEL_CALL_H                                       \
-    if (H % (4 * GROUP_SIZE) == 0) {                          \
-      static constexpr int NUM_WARPS = 4;                     \
-      populate_launch_params(NUM_WARPS, NUM_PARALLEL_TOKENS); \
-      KERNEL_FN                                               \
-    } else {                                                  \
-      static constexpr int NUM_WARPS = 1;                     \
-      populate_launch_params(NUM_WARPS, NUM_PARALLEL_TOKENS); \
-      KERNEL_FN                                               \
-    }
-
-  #define KERNEL_CALL_TOP_LEVEL                      \
-    if (num_parallel_tokens == 1) {                  \
-      static constexpr int NUM_PARALLEL_TOKENS = 1;  \
-      KERNEL_CALL_H                                  \
-    } else if (num_parallel_tokens == 2) {           \
-      static constexpr int NUM_PARALLEL_TOKENS = 2;  \
-      KERNEL_CALL_H                                  \
-    } else if (num_parallel_tokens == 4) {           \
-      static constexpr int NUM_PARALLEL_TOKENS = 4;  \
-      KERNEL_CALL_H                                  \
-    } else if (num_parallel_tokens == 8) {           \
-      static constexpr int NUM_PARALLEL_TOKENS = 8;  \
-      KERNEL_CALL_H                                  \
-    } else if (num_parallel_tokens == 16) {          \
-      static constexpr int NUM_PARALLEL_TOKENS = 16; \
-      KERNEL_CALL_H                                  \
-    } else if (num_parallel_tokens == 32) {          \
-      static constexpr int NUM_PARALLEL_TOKENS = 32; \
-      KERNEL_CALL_H                                  \
-    } else if (num_parallel_tokens == 64) {          \
-      static constexpr int NUM_PARALLEL_TOKENS = 64; \
-      KERNEL_CALL_H                                  \
-    }
-
-  Idx_t G;
-  dim3 block, grid;
-  auto populate_launch_params = [&](int num_warps, int _num_parallel_tokens) {
-    G = H / Idx_t(group_size * num_warps);
-    grid = dim3(E * G, _num_parallel_tokens);
-    block = dim3(num_warps * WARP_SIZE);
-  };
-
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
-  VLLM_DISPATCH_FP8_TYPES(y_q.scalar_type(),
-                          "silu_mul_fp8_quant_deep_gemm_kernel",
-                          [&] { KERNEL_CALL_TOP_LEVEL });
-
-#endif
-}

csrc/quantization/compressed_tensors/int8_quant_kernels.cu

@@ -7,10 +7,17 @@
 
 #include <cmath>
 
-#include "../../cub_helpers.h"
 #include "../../dispatch_utils.h"
 #include "../vectorization_utils.cuh"
 
+#ifndef USE_ROCM
+  #include <cub/cub.cuh>
+  #include <cub/util_type.cuh>
+#else
+  #include <hipcub/hipcub.hpp>
+  #include <hipcub/util_type.hpp>
+#endif
+
 static inline __device__ int8_t float_to_int8_rn(float x) {
 #ifdef USE_ROCM
   static constexpr auto i8_min =
@@ -166,7 +173,7 @@ __global__ void dynamic_scaled_int8_quant_kernel(
       });
   using BlockReduce = cub::BlockReduce<float, 256>;
   __shared__ typename BlockReduce::TempStorage tmp;
-  float block_max = BlockReduce(tmp).Reduce(thread_max, CubMaxOp{}, blockDim.x);
+  float block_max = BlockReduce(tmp).Reduce(thread_max, cub::Max{}, blockDim.x);
   __shared__ float absmax;
   if (tid == 0) {
     absmax = block_max;

csrc/quantization/cutlass_w4a8/w4a8_mm_entry.cu

@@ -25,8 +25,6 @@
 #include "cutlass_extensions/common.hpp"
 #include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
 
-#include <cuda_runtime.h>
-
 namespace vllm::cutlass_w4a8 {
 
 using namespace cute;
@@ -395,71 +393,6 @@ torch::Tensor pack_scale_fp8(torch::Tensor const& scales) {
   return packed_scales;
 }
 
-/*
-  GPU-accelerated implementation of cutlass::unified_encode_int4b.
-  Constructs a lookup table in constant memory to map 8 bits
-  (two 4-bit values) at a time. Assumes memory is contiguous
-  and pointers are 16-byte aligned.
-*/
-__constant__ uint8_t kNibbleLUT[256];
-
-__global__ void unified_encode_int4b_device(const uint8_t* in, uint8_t* out,
-                                            size_t nbytes) {
-  constexpr size_t V = sizeof(uint4);  // 16 bytes
-  const size_t tid = blockIdx.x * blockDim.x + threadIdx.x;
-  const size_t nthreads = size_t(gridDim.x) * blockDim.x;
-  const size_t nvec = nbytes / V;
-
-  // 1-D grid-stride loop over 16-byte chunks
-  for (size_t vec = tid; vec < nvec; vec += nthreads) {
-    uint4 v = reinterpret_cast<const uint4*>(in)[vec];
-    uint8_t* b = reinterpret_cast<uint8_t*>(&v);
-#pragma unroll
-    for (int i = 0; i < int(V); ++i) b[i] = kNibbleLUT[b[i]];
-    reinterpret_cast<uint4*>(out)[vec] = v;
-  }
-}
-
-static bool upload_lut() {
-  std::array<uint8_t, 256> lut{};
-  auto map_nib = [](uint8_t v) -> uint8_t {
-    // 1..7 -> (8 - v); keep 0 and 8..15
-    return (v == 0 || (v & 0x8)) ? v : uint8_t(8 - v);
-  };
-  for (int b = 0; b < 256; ++b) {
-    uint8_t lo = b & 0xF;
-    uint8_t hi = (b >> 4) & 0xF;
-    lut[b] = uint8_t((map_nib(hi) << 4) | map_nib(lo));
-  }
-  cudaError_t e = cudaMemcpyToSymbol(kNibbleLUT, lut.data(), lut.size(),
-                                     /*offset=*/0, cudaMemcpyHostToDevice);
-
-  return (e == cudaSuccess);
-}
-
-static bool unified_encode_int4b(cutlass::int4b_t const* in,
-                                 cutlass::int4b_t* out, size_t num_int4_elems) {
-  // Build/upload LUT
-  if (!upload_lut()) return false;
-
-  static_assert(sizeof(typename cutlass::int4b_t::Storage) == 1,
-                "int4 storage must be 1 byte");
-  const size_t nbytes = num_int4_elems >> 1;
-
-  auto* in_bytes = reinterpret_cast<uint8_t const*>(in);
-  auto* out_bytes = reinterpret_cast<uint8_t*>(out);
-
-  // kernel launch params
-  constexpr int block = 256;
-  const size_t nvec = nbytes / sizeof(uint4);  // # of 16B vectors
-  int grid = int((nvec + block - 1) / block);
-  if (grid == 0) grid = 1;  // ensure we still cover the tail in the kernel
-
-  unified_encode_int4b_device<<<grid, block>>>(in_bytes, out_bytes, nbytes);
-  cudaError_t err = cudaGetLastError();
-  return (err == cudaSuccess);
-}
-
 torch::Tensor encode_and_reorder_int4b(torch::Tensor const& B) {
   TORCH_CHECK(B.dtype() == torch::kInt32);
   TORCH_CHECK(B.dim() == 2);
@@ -468,7 +401,6 @@ torch::Tensor encode_and_reorder_int4b(torch::Tensor const& B) {
 
   int k = B.size(0) * PackFactor;  // logical k
   int n = B.size(1);
-  TORCH_CHECK((n * k) % 32 == 0, "need multiples of 32 int4s for 16B chunks");
 
   auto B_ptr = static_cast<QuantType const*>(B.const_data_ptr());
   auto B_packed_ptr = static_cast<QuantType*>(B_packed.data_ptr());
@@ -477,9 +409,7 @@ torch::Tensor encode_and_reorder_int4b(torch::Tensor const& B) {
   LayoutB_Reordered layout_B_reordered =
       cute::tile_to_shape(LayoutAtomQuant{}, shape_B);
 
-  bool ok =
-      vllm::cutlass_w4a8::unified_encode_int4b(B_ptr, B_packed_ptr, n * k);
-  TORCH_CHECK(ok, "unified_encode_int4b failed");
+  cutlass::unified_encode_int4b(B_ptr, B_packed_ptr, n * k);
   cutlass::reorder_tensor(B_packed_ptr, layout_B, layout_B_reordered);
 
   return B_packed;

csrc/quantization/cutlass_w8a8/c3x/scaled_mm_blockwise_sm100_fp8_dispatch.cuh

@@ -146,7 +146,6 @@ void cutlass_gemm_caller_blockwise(torch::Tensor& out, torch::Tensor const& a,
 
   using ElementAB = typename Gemm::ElementAB;
   using ElementD = typename Gemm::ElementD;
-  using ElementBlockScale = typename Gemm::ElementBlockScale;
 
   int32_t m = a.size(0), n = b.size(1), k = a.size(1);
 
@@ -167,29 +166,26 @@ void cutlass_gemm_caller_blockwise(torch::Tensor& out, torch::Tensor const& a,
       ScaleConfig::tile_atom_to_shape_SFB(make_shape(n, m, k, 1)) :
       ScaleConfig::tile_atom_to_shape_SFB(make_shape(m, n, k, 1));
 
-  auto a_ptr = static_cast<ElementAB const*>(a.data_ptr());
-  auto b_ptr = static_cast<ElementAB const*>(b.data_ptr());
-  auto a_scales_ptr = static_cast<ElementBlockScale const*>(a_scales.data_ptr());
-  auto b_scales_ptr = static_cast<ElementBlockScale const*>(b_scales.data_ptr());
+  auto a_ptr = static_cast<ElementAB*>(a.data_ptr());
+  auto b_ptr = static_cast<ElementAB*>(b.data_ptr());
+  auto a_scales_ptr = static_cast<float*>(a_scales.data_ptr());
+  auto b_scales_ptr = static_cast<float*>(b_scales.data_ptr());
 
-  typename GemmKernel::MainloopArguments mainloop_args{};
-  mainloop_args.layout_SFA = layout_SFA;
-  mainloop_args.layout_SFB = layout_SFB;
-  if (swap_ab) {
-    mainloop_args.ptr_A = b_ptr;
-    mainloop_args.dA = b_stride;
-    mainloop_args.ptr_B = a_ptr;
-    mainloop_args.dB = a_stride;
-    mainloop_args.ptr_SFA = b_scales_ptr;
-    mainloop_args.ptr_SFB = a_scales_ptr;
-  } else {
-    mainloop_args.ptr_A = a_ptr;
-    mainloop_args.dA = a_stride;
-    mainloop_args.ptr_B = b_ptr;
-    mainloop_args.dB = b_stride;
-    mainloop_args.ptr_SFA = a_scales_ptr;
-    mainloop_args.ptr_SFB = b_scales_ptr;
-  }
+  auto mainloop_args = [&](){
+    // layout_SFA and layout_SFB cannot be swapped since they are deduced.
+    if (swap_ab) {
+      return typename GemmKernel::MainloopArguments{
+          b_ptr,        b_stride,   a_ptr,        a_stride,
+          b_scales_ptr, layout_SFA, a_scales_ptr, layout_SFB
+      };
+    }
+    else {
+      return typename GemmKernel::MainloopArguments{
+          a_ptr,        a_stride,   b_ptr,        b_stride,
+          a_scales_ptr, layout_SFA, b_scales_ptr, layout_SFB
+      };
+    }
+  }();
   auto prob_shape = swap_ab ? cute::make_shape(n, m, k, 1) : cute::make_shape(m, n, k, 1);
 
   auto c_ptr = static_cast<ElementD*>(out.data_ptr());

csrc/quantization/cutlass_w8a8/c3x/scaled_mm_blockwise_sm120_fp8_dispatch.cuh

@@ -125,7 +125,6 @@ void cutlass_gemm_caller_blockwise(torch::Tensor& out, torch::Tensor const& a,
 
   using ElementAB = typename Gemm::ElementAB;
   using ElementD = typename Gemm::ElementD;
-  using ElementBlockScale = typename Gemm::ElementBlockScale;
 
   int32_t m = a.size(0), n = b.size(1), k = a.size(1);
 
@@ -144,20 +143,17 @@ void cutlass_gemm_caller_blockwise(torch::Tensor& out, torch::Tensor const& a,
   LayoutSFB layout_SFB = 
       ScaleConfig::tile_atom_to_shape_SFB(make_shape(m, n, k, 1));
 
-  auto a_ptr = static_cast<ElementAB const*>(a.data_ptr());
-  auto b_ptr = static_cast<ElementAB const*>(b.data_ptr());
-  auto a_scales_ptr = static_cast<ElementBlockScale const*>(a_scales.data_ptr());
-  auto b_scales_ptr = static_cast<ElementBlockScale const*>(b_scales.data_ptr());
+  auto a_ptr = static_cast<ElementAB*>(a.data_ptr());
+  auto b_ptr = static_cast<ElementAB*>(b.data_ptr());
+  auto a_scales_ptr = static_cast<float*>(a_scales.data_ptr());
+  auto b_scales_ptr = static_cast<float*>(b_scales.data_ptr());
 
-  typename GemmKernel::MainloopArguments mainloop_args{};
-  mainloop_args.ptr_A = a_ptr;
-  mainloop_args.dA = a_stride;
-  mainloop_args.ptr_B = b_ptr;
-  mainloop_args.dB = b_stride;
-  mainloop_args.ptr_SFA = a_scales_ptr;
-  mainloop_args.layout_SFA = layout_SFA;
-  mainloop_args.ptr_SFB = b_scales_ptr;
-  mainloop_args.layout_SFB = layout_SFB;
+  auto mainloop_args = [&](){
+    return typename GemmKernel::MainloopArguments{
+        a_ptr,        a_stride,   b_ptr,        b_stride,
+        a_scales_ptr, layout_SFA, b_scales_ptr, layout_SFB
+    };
+  }();
   auto prob_shape = cute::make_shape(m, n, k, 1);
 
   auto c_ptr = static_cast<ElementD*>(out.data_ptr());

csrc/quantization/cutlass_w8a8/c3x/scaled_mm_blockwise_sm90_fp8_dispatch.cuh

@@ -115,7 +115,6 @@ void cutlass_gemm_caller_blockwise(torch::Tensor& out, torch::Tensor const& a,
 
   using ElementAB = typename Gemm::ElementAB;
   using ElementD = typename Gemm::ElementD;
-  using ElementBlockScale = typename Gemm::ElementBlockScale;
 
   int32_t m = a.size(0), n = b.size(1), k = a.size(1);
 
@@ -136,20 +135,17 @@ void cutlass_gemm_caller_blockwise(torch::Tensor& out, torch::Tensor const& a,
   LayoutSFB layout_SFB = 
       ScaleConfig::tile_atom_to_shape_SFB(make_shape(m, n, k, 1));
 
-  auto a_ptr = static_cast<ElementAB const*>(a.data_ptr());
-  auto b_ptr = static_cast<ElementAB const*>(b.data_ptr());
-  auto a_scales_ptr = static_cast<ElementBlockScale const*>(a_scales.data_ptr());
-  auto b_scales_ptr = static_cast<ElementBlockScale const*>(b_scales.data_ptr());
+  auto a_ptr = static_cast<ElementAB*>(a.data_ptr());
+  auto b_ptr = static_cast<ElementAB*>(b.data_ptr());
+  auto a_scales_ptr = static_cast<float*>(a_scales.data_ptr());
+  auto b_scales_ptr = static_cast<float*>(b_scales.data_ptr());
 
-  typename GemmKernel::MainloopArguments mainloop_args{};
-  mainloop_args.ptr_A = a_ptr;
-  mainloop_args.dA = a_stride;
-  mainloop_args.ptr_B = b_ptr;
-  mainloop_args.dB = b_stride;
-  mainloop_args.ptr_SFA = a_scales_ptr;
-  mainloop_args.layout_SFA = layout_SFA;
-  mainloop_args.ptr_SFB = b_scales_ptr;
-  mainloop_args.layout_SFB = layout_SFB;
+  auto mainloop_args = [&](){
+    return typename GemmKernel::MainloopArguments{
+        a_ptr,        a_stride,   b_ptr,        b_stride,
+        a_scales_ptr, layout_SFA, b_scales_ptr, layout_SFB
+    };
+  }();
   auto prob_shape = cute::make_shape(m, n, k, 1);
 
   auto c_ptr = static_cast<ElementD*>(out.data_ptr());

csrc/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu

@@ -26,46 +26,113 @@
 #include "dispatch_utils.h"
 
 #include "cuda_utils.h"
-#include "launch_bounds_utils.h"
 #include "nvfp4_utils.cuh"
 
 namespace vllm {
 
-// silu in float32
-__device__ __forceinline__ float silu(float x) {
-  return __fdividef(x, (1.f + __expf(-x)));
-}
-
-__device__ __forceinline__ float2 silu2(float2 x) {
-  return make_float2(silu(x.x), silu(x.y));
-}
-
 template <class Type>
-__inline__ __device__ PackedVec<Type> compute_silu_mul(PackedVec<Type>& vec,
-                                                       PackedVec<Type>& vec2) {
+__inline__ __device__ PackedVec<Type> compute_silu(PackedVec<Type>& vec,
+                                                   PackedVec<Type>& vec2) {
   PackedVec<Type> result;
-  using packed_type = typename TypeConverter<Type>::Type;
-
 #pragma unroll
   for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; ++i) {
-    // silu_mul in float32
     if constexpr (std::is_same_v<Type, half>) {
-      float2 silu_vec = silu2(__half22float2(vec.elts[i]));
-      result.elts[i] =
-          __float22half2_rn(__fmul2_rn(silu_vec, __half22float2(vec2.elts[i])));
+      half2 val(0.5f, 0.5f);
+      half2 t0 = __hmul2(vec.elts[i], val);
+      half2 t1 = __hfma2(h2tanh(t0), val, val);
+      half2 t2 = __hmul2(vec.elts[i], t1);
+      result.elts[i] = __hmul2(t2, vec2.elts[i]);
     } else {
-      float2 silu_vec = silu2(__bfloat1622float2(vec.elts[i]));
-      result.elts[i] = __float22bfloat162_rn(
-          __fmul2_rn(silu_vec, __bfloat1622float2(vec2.elts[i])));
+      __nv_bfloat162 val(0.5f, 0.5f);
+      __nv_bfloat162 t0 = __hmul2(vec.elts[i], val);
+      __nv_bfloat162 t1 = __hfma2(h2tanh(t0), val, val);
+      __nv_bfloat162 t2 = __hmul2(vec.elts[i], t1);
+      result.elts[i] = __hmul2(t2, vec2.elts[i]);
     }
   }
   return result;
 }
 
+// Quantizes the provided PackedVec into the uint32_t output
+template <class Type, bool UE8M0_SF = false>
+__device__ uint32_t silu_and_cvt_warp_fp16_to_fp4(PackedVec<Type>& vec,
+                                                  PackedVec<Type>& vec2,
+                                                  float SFScaleVal,
+                                                  uint8_t* SFout) {
+  PackedVec<Type> out_silu = compute_silu(vec, vec2);
+  // Get absolute maximum values among the local 8 values.
+  auto localMax = __habs2(out_silu.elts[0]);
+
+// Local maximum value.
+#pragma unroll
+  for (int i = 1; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) {
+    localMax = __hmax2(localMax, __habs2(out_silu.elts[i]));
+  }
+
+  // Get the absolute maximum among all 16 values (two threads).
+  localMax = __hmax2(__shfl_xor_sync(uint32_t(-1), localMax, 1), localMax);
+  // Get the final absolute maximum values.
+  float vecMax = float(__hmax(localMax.x, localMax.y));
+
+  // Get the SF (max value of the vector / max value of e2m1).
+  // maximum value of e2m1 = 6.0.
+  // TODO: use half as compute data type.
+  float SFValue = SFScaleVal * (vecMax * reciprocal_approximate_ftz(6.0f));
+  // 8 bits representation of the SF.
+  uint8_t fp8SFVal;
+  // Write the SF to global memory (STG.8).
+  if constexpr (UE8M0_SF) {
+    // Extract the 8 exponent bits from float32.
+    // float 32bits = 1 sign bit + 8 exponent bits + 23 mantissa bits.
+    uint32_t tmp = reinterpret_cast<uint32_t&>(SFValue) >> 23;
+    fp8SFVal = tmp & 0xff;
+    // Convert back to fp32.
+    reinterpret_cast<uint32_t&>(SFValue) = tmp << 23;
+  } else {
+    // Here SFValue is always positive, so E4M3 is the same as UE4M3.
+    __nv_fp8_e4m3 tmp = __nv_fp8_e4m3(SFValue);
+    reinterpret_cast<__nv_fp8_e4m3&>(fp8SFVal) = tmp;
+    // Convert back to fp32.
+    SFValue = float(tmp);
+  }
+  // Get the output scale.
+  // Recipe: final_scale = reciprocal(fp32(fp8(SFValue * SFScaleVal))) *
+  //                       reciprocal(SFScaleVal))
+  float outputScale =
+      SFValue != 0 ? reciprocal_approximate_ftz(
+                         SFValue * reciprocal_approximate_ftz(SFScaleVal))
+                   : 0.0f;
+
+  if (SFout) {
+    // Write the SF to global memory (STG.8).
+    *SFout = fp8SFVal;
+  }
+
+  // Convert the input to float.
+  float2 fp2Vals[CVT_FP4_ELTS_PER_THREAD / 2];
+
+#pragma unroll
+  for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) {
+    if constexpr (std::is_same_v<Type, half>) {
+      fp2Vals[i] = __half22float2(out_silu.elts[i]);
+    } else {
+      fp2Vals[i] = __bfloat1622float2(out_silu.elts[i]);
+    }
+    fp2Vals[i].x *= outputScale;
+    fp2Vals[i].y *= outputScale;
+  }
+
+  // Convert to e2m1 values.
+  uint32_t e2m1Vec = fp32_vec_to_e2m1(fp2Vals);
+
+  // Write the e2m1 values to global memory.
+  return e2m1Vec;
+}
+
 // Use UE4M3 by default.
 template <class Type, bool UE8M0_SF = false>
-__global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024))
-    silu_mul_cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, Type const* in,
+__global__ void __launch_bounds__(1024, 4)
+    silu_and_cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, Type const* in,
                              float const* SFScale, uint32_t* out,
                              uint32_t* SFout) {
   using PackedVec = PackedVec<Type>;
@@ -93,18 +160,16 @@ __global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024))
       // Get the output tensor offset.
       // Same as inOffset because 8 elements are packed into one uint32_t.
       int64_t outOffset = rowIdx * (numCols / CVT_FP4_ELTS_PER_THREAD) + colIdx;
+      ;
       auto& out_pos = out[outOffset];
 
-      // Compute silu and mul
-      PackedVec out_silu_mul = compute_silu_mul(in_vec, in_vec2);
-
       auto sf_out =
           cvt_quant_to_fp4_get_sf_out_offset<uint32_t,
                                              CVT_FP4_NUM_THREADS_PER_SF>(
               rowIdx, colIdx, numCols, SFout);
 
-      out_pos = cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(out_silu_mul, SFScaleVal,
-                                                     sf_out);
+      out_pos = silu_and_cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(
+          in_vec, in_vec2, SFScaleVal, sf_out);
     }
   }
 }
@@ -132,15 +197,14 @@ void silu_and_mul_nvfp4_quant_sm1xxa(torch::Tensor& output,  // [..., d]
   const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
   auto stream = at::cuda::getCurrentCUDAStream(input.get_device());
   dim3 block(std::min(int(n / ELTS_PER_THREAD), 1024));
-  int const numBlocksPerSM =
-      vllm_runtime_blocks_per_sm(static_cast<int>(block.x));
+  int const numBlocksPerSM = 2048 / block.x;
   dim3 grid(std::min(int(m), multiProcessorCount * numBlocksPerSM));
 
   VLLM_DISPATCH_HALF_TYPES(
       input.scalar_type(), "silu_and_mul_nvfp4_quant_kernel", [&] {
         using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
         auto input_ptr = static_cast<cuda_type const*>(input.data_ptr());
-        vllm::silu_mul_cvt_fp16_to_fp4<cuda_type><<<grid, block, 0, stream>>>(
+        vllm::silu_and_cvt_fp16_to_fp4<cuda_type><<<grid, block, 0, stream>>>(
             m, n, input_ptr, input_sf_ptr,
             reinterpret_cast<uint32_t*>(output_ptr),
             reinterpret_cast<uint32_t*>(sf_out));

csrc/quantization/fp4/nvfp4_experts_quant.cu

@@ -26,13 +26,12 @@
 #include "dispatch_utils.h"
 
 #include "nvfp4_utils.cuh"
-#include "launch_bounds_utils.h"
 
 namespace vllm {
 
 // Use UE4M3 by default.
 template <class Type, bool UE8M0_SF = false, bool SMALL_NUM_EXPERTS = false>
-__global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
+__global__ void __launch_bounds__(512, 4)
     cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, Type const* in,
                     float const* SFScale, uint32_t* out, uint32_t* SFout,
                     uint32_t* input_offset_by_experts,
@@ -130,7 +129,7 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
 
 // Kernel for LARGE_M_TOPK = true (large m_topk optimized version)
 template <class Type, bool UE8M0_SF = false, bool SMALL_NUM_EXPERTS = false>
-__global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024))
+__global__ void __launch_bounds__(1024, 4)
     cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, Type const* in,
                     float const* SFScale, uint32_t* out, uint32_t* SFout,
                     uint32_t* input_offset_by_experts,
@@ -234,9 +233,8 @@ void quant_impl(void* output, void* output_scale, void* input,
   int const workSizePerRow = k / ELTS_PER_THREAD;
   int const totalWorkSize = m_topk * workSizePerRow;
   dim3 block(std::min(workSizePerRow, 512));
-  // Get number of blocks per SM
-  int const numBlocksPerSM =
-      vllm_runtime_blocks_per_sm(static_cast<int>(block.x));
+  // Get number of blocks per SM (assume we can fully utilize the SM).
+  int const numBlocksPerSM = 2048 / block.x;
   dim3 grid(std::min(static_cast<int>((totalWorkSize + block.x - 1) / block.x),
                      multiProcessorCount * numBlocksPerSM));
   while (grid.x <= multiProcessorCount && block.x > 64) {

csrc/quantization/fp4/nvfp4_quant_kernels.cu

@@ -26,14 +26,13 @@
 #include "dispatch_utils.h"
 
 #include "cuda_utils.h"
-#include "launch_bounds_utils.h"
 #include "nvfp4_utils.cuh"
 
 namespace vllm {
 
 // Use UE4M3 by default.
 template <class Type, bool UE8M0_SF = false>
-__global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
+__global__ void __launch_bounds__(512, 4)
     cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, Type const* in,
                     float const* SFScale, uint32_t* out, uint32_t* SFout) {
   using PackedVec = PackedVec<Type>;
@@ -76,9 +75,8 @@ void invokeFP4Quantization(int m, int n, T const* input, float const* SFScale,
   // Grid, Block size.
   // Each thread converts 8 values.
   dim3 block(std::min(int(n / ELTS_PER_THREAD), 512));
-  // Get number of blocks per SM
-  int const numBlocksPerSM =
-      vllm_runtime_blocks_per_sm(static_cast<int>(block.x));
+  // Get number of blocks per SM (assume we can fully utilize the SM).
+  int const numBlocksPerSM = 2048 / block.x;
   dim3 grid(std::min(int(m), multiProcessorCount * numBlocksPerSM));
 
   // Launch the cvt kernel.

csrc/quantization/fp8/common.cu

@@ -1,10 +1,15 @@
 #include "common.cuh"
 #include "dispatch_utils.h"
-#include "../../cub_helpers.h"
 #include "../vectorization_utils.cuh"
 #include <c10/cuda/CUDAGuard.h>
 #include <ATen/cuda/Exceptions.h>
 
+#ifndef USE_ROCM
+  #include <cub/cub.cuh>
+#else
+  #include <hipcub/hipcub.hpp>
+#endif
+
 namespace vllm {
 
 template <typename scalar_t, typename fp8_type>
@@ -111,7 +116,7 @@ __global__ void dynamic_per_token_scaled_fp8_quant_kernel_strided(
   using BlockReduce = cub::BlockReduce<float, 256>;
   __shared__ typename BlockReduce::TempStorage tmp;
   const float block_max =
-      BlockReduce(tmp).Reduce(absmax_val, CubMaxOp{}, blockDim.x);
+      BlockReduce(tmp).Reduce(absmax_val, cub::Max{}, blockDim.x);
 
   __shared__ float token_scale;
   if (tid == 0) {

csrc/quantization/fp8/common.cuh

@@ -5,9 +5,7 @@
 
 #include <cmath>
 
-#ifndef USE_ROCM
-  #include "nvidia/quant_utils.cuh"
-#else
+#ifdef USE_ROCM
   #include "amd/quant_utils.cuh"
 #endif
 
@@ -50,9 +48,7 @@ __device__ __forceinline__ fp8_type scaled_fp8_conversion(float const val,
   float r =
       fmaxf(-quant_type_max_v<fp8_type>, fminf(x, quant_type_max_v<fp8_type>));
 #ifndef USE_ROCM
-  // Use hardware cvt instruction for fp8 on nvidia
-  // Currently only support fp8_type = c10::Float8_e4m3fn
-  return fp8::vec_conversion<fp8_type, float>(r);
+  return static_cast<fp8_type>(r);
 #else
   // Use hardware cvt instruction for fp8 on rocm
   return fp8::cvt_c10<fp8_type>(r);

csrc/quantization/fp8/nvidia/quant_utils.cuh

@@ -12,26 +12,13 @@ namespace vllm {
 namespace fp8 {
   #ifdef ENABLE_FP8
 
+    #if 0  // Disable the following code to reduce the binary size.
 template <typename Tout, typename Tin>
-__inline__ __device__ Tout vec_conversion(
-    const Tin& x, const __nv_fp8_interpretation_t fp8_type = __NV_E4M3) {
+__inline__ __device__ Tout
+vec_conversion(const Tin &x, const __nv_fp8_interpretation_t fp8_type) {
   return x;
 }
 
-// float -> c10::Float8_e4m3fn
-template <>
-__inline__ __device__ c10::Float8_e4m3fn
-vec_conversion<c10::Float8_e4m3fn, float>(
-    const float& a, const __nv_fp8_interpretation_t fp8_type) {
-    #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
-  return static_cast<c10::Float8_e4m3fn>(a);
-    #else
-  return c10::Float8_e4m3fn(__nv_cvt_float_to_fp8(a, __NV_SATFINITE, fp8_type),
-                            c10::Float8_e4m3fn::from_bits());
-    #endif
-}
-
-    #if 0  // Disable the following code to reduce the binary size.
 // fp8 -> half
 template <>
 __inline__ __device__ uint16_t vec_conversion<uint16_t, uint8_t>(

csrc/quantization/fp8/per_token_group_quant.cu

@@ -12,8 +12,8 @@
 #include "../vectorization_utils.cuh"
 #include "../../dispatch_utils.h"
 
-__device__ __forceinline__ float GroupReduceMax(float val) {
-  unsigned mask = threadIdx.x % 32 >= 16 ? 0xffff0000 : 0x0000ffff;
+__device__ __forceinline__ float GroupReduceMax(float val, const int tid) {
+  unsigned mask = 0xffff;
 
   val = fmaxf(val, __shfl_xor_sync(mask, val, 8));
   val = fmaxf(val, __shfl_xor_sync(mask, val, 4));
@@ -86,7 +86,7 @@ __global__ void per_token_group_quant_8bit_kernel(
       threads_per_group,  // stride in group
       scalar_op_cache);   // scalar handler
 
-  local_absmax = GroupReduceMax(local_absmax);
+  local_absmax = GroupReduceMax(local_absmax, lane_id);
 
   float y_s = local_absmax / max_8bit;
   if constexpr (SCALE_UE8M0) {

csrc/quantization/fused_kernels/layernorm_utils.cuh

@@ -8,7 +8,11 @@
 #include "quantization/utils.cuh"
 #include "quant_conversions.cuh"
 
-#include "../../cub_helpers.h"
+#ifndef USE_ROCM
+  #include <cub/cub.cuh>
+#else
+  #include <hipcub/hipcub.hpp>
+#endif
 
 namespace vllm {
 
@@ -32,7 +36,7 @@ __device__ void compute_rms(float* rms, scalar_t const* __restrict__ input,
 
   using BlockReduce = cub::BlockReduce<float, 1024>;
   __shared__ typename BlockReduce::TempStorage reduceStore;
-  ss = BlockReduce(reduceStore).Reduce(ss, CubAddOp{}, blockDim.x);
+  ss = BlockReduce(reduceStore).Reduce(ss, cub::Sum{}, blockDim.x);
 
   __shared__ float s_rms;
   if (threadIdx.x == 0) {
@@ -69,7 +73,7 @@ __device__ void compute_dynamic_per_token_scales(
   __shared__ typename BlockReduce::TempStorage reduceStore;
   block_absmax_val_maybe =
       BlockReduce(reduceStore)
-          .Reduce(block_absmax_val_maybe, CubMaxOp{}, blockDim.x);
+          .Reduce(block_absmax_val_maybe, cub::Max{}, blockDim.x);
 
   __shared__ float s_token_scale;
   if (threadIdx.x == 0) {
@@ -165,7 +169,7 @@ __device__ void compute_rms(float* rms, scalar_t const* __restrict__ input,
 
   using BlockReduce = cub::BlockReduce<float, 1024>;
   __shared__ typename BlockReduce::TempStorage reduceStore;
-  ss = BlockReduce(reduceStore).Reduce(ss, CubAddOp{}, blockDim.x);
+  ss = BlockReduce(reduceStore).Reduce(ss, cub::Sum{}, blockDim.x);
 
   __shared__ float s_rms;
   if (threadIdx.x == 0) {
@@ -236,7 +240,7 @@ __device__ void compute_dynamic_per_token_scales(
   __shared__ typename BlockReduce::TempStorage reduceStore;
   block_absmax_val_maybe =
       BlockReduce(reduceStore)
-          .Reduce(block_absmax_val_maybe, CubMaxOp{}, blockDim.x);
+          .Reduce(block_absmax_val_maybe, cub::Max{}, blockDim.x);
 
   __shared__ float s_token_scale;
   if (threadIdx.x == 0) {

csrc/quantization/hadamard/hadacore/hadamard_transform_cuda.cu

@@ -1,817 +0,0 @@
-// clang-format off
-// Adapted from: https://github.com/meta-pytorch/applied-ai/blob/main/kernels/cuda/inference/hadamard_transform/hadamard_transform_cuda.cu
-
-/***********
-Copyright 2024 Meta
-
-Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
-1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
-2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.
-3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.
-THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS “AS IS” AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
-***********/
-
-#include <torch/all.h>
-#include <stdint.h>
-#include <cuda_runtime.h>
-#include <mma.h>
-#include <cuda/annotated_ptr>
-#include <c10/cuda/CUDAException.h>
-
-#include <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAGuard.h>
-
-#include "core/registration.h"
-#include "dispatch_utils.h"
-
-namespace hadacore {
-
-#ifndef __CUDACC__
-#define __launch_bounds__(x,y)
-#endif
-
-#define MAX_WARPS_PER_SM 48
-
-#define MIN(a, b) ((a) < (b) ? (a) : (b))
-
-using b16 = uint16_t;
-using b32 = uint32_t;
-
-constexpr int launch_configs_big[7][3] = {
-    // default
-    {2, 1, 24},
-    {2, 2, 16}, 
-    {2, 4, 8}, 
-    {2, 8, 4}, 
-    {2, 16, 3},
-    {4, 16, 2},
-    {8, 16, 1}
-    // // extra coalescing
-    // {2, 1, 24},
-    // {2, 2, 16}, 
-    // {2, 4, 8}, 
-    // {2, 8, 4}, 
-    // {4, 8, 3},
-    // {8, 8, 2},
-    // {16, 8, 1}
-    // // less coalescing
-    // {2, 1, 24},
-    // {2, 2, 16}, 
-    // {2, 4, 8}, 
-    // {2, 8, 4}, 
-    // {1, 32, 1},
-    // {2, 32, 1},
-    // {4, 32, 1}
-};
-
-// a 4x2, b 2x2, c 2x2
-template <torch::ScalarType dtype>
-__device__ __forceinline__ void mma_m16_n8_k16_b16_b16_b16_noacc(b32 a0, b32 a1, b32 a2, b32 a3, b32 b0, b32 b1, b32& c0, b32& c1){
-    static_assert(dtype == torch::ScalarType::Half || dtype == torch::ScalarType::BFloat16);
-    // d, a, b, c
-    b32 zero = 0;
-    if constexpr(dtype == torch::ScalarType::Half) {
-        asm (
-            "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 "
-            "{%0, %1}, {%2, %3, %4, %5}, {%6, %7}, {%8, %9};\n\t"
-            : "=r"(c0), "=r"(c1) : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1), "r"(zero), "r"(zero)
-        );
-    } else {
-        b32 temp0, temp1, temp2, temp3;
-        asm (
-            "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
-            "{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%10, %11, %12, %13};\n\t"
-            : "=r"(temp0), "=r"(temp1), "=r"(temp2), "=r"(temp3) : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1), "r"(zero), "r"(zero), "r"(zero), "r"(zero)
-        );
-        asm ("cvt.rn.bf16x2.f32 %0, %1, %2;\n\t" : "=r"(c0) : "r"(temp1), "r"(temp0));
-        asm ("cvt.rn.bf16x2.f32 %0, %1, %2;\n\t" : "=r"(c1) : "r"(temp3), "r"(temp2));
-    }
-}
-
-// a 4x2, b 4x2, c 4x2
-template <torch::ScalarType dtype>
-__device__ __forceinline__ void mma_m16_n16_k16_b16_b16_b16_noacc(b32 a0, b32 a1, b32 a2, b32 a3, b32 b0, b32 b1, b32 b2, b32 b3, b32& c0, b32& c1, b32& c2, b32& c3){
-    mma_m16_n8_k16_b16_b16_b16_noacc<dtype>(a0, a1, a2, a3, b0, b1, c0, c1);
-    mma_m16_n8_k16_b16_b16_b16_noacc<dtype>(a0, a1, a2, a3, b2, b3, c2, c3);
-}
-
-__device__ __forceinline__ void matrix_transpose_m8_n8_b16_inplace(b32& a0) {
-    asm (
-        "movmatrix.sync.aligned.m8n8.trans.b16 "
-        "%0, %1;\n\t"
-        : "=r"(a0) : "r"(a0)
-    );
-}
-
-#define p_p(i) ((val_1p[i] & 0x0000FFFF) | val_1p[i] << 16)
-#define p_n(i) ((val_1p[i] & 0x0000FFFF) | val_1n[i] << 16)
-#define n_p(i) ((val_1n[i] & 0x0000FFFF) | val_1p[i] << 16)
-#define n_n(i) ((val_1n[i] & 0x0000FFFF) | val_1n[i] << 16)
-
-template<int64_t num_chunks, int64_t warps_per_block, int64_t log_had_size, int64_t blocks_per_sm, bool enable_mask, torch::ScalarType dtype>
-__global__ void __launch_bounds__(32 * warps_per_block, blocks_per_sm)
-// a is column major, b is row major
-hadamard_transform_kernel(b16* a, b16* out, int total_num_chunks) {
-    static_assert(dtype == torch::ScalarType::Half || dtype == torch::ScalarType::BFloat16, "Only fp16 and bf16 supported currently");
-
-    b32 b_frag_all[num_chunks][4]; // for all chunks, holds matrix fragment (which takes 4 regs of b16x2 * 32 threads)
-
-    int64_t blockid = blockIdx.x * warps_per_block + threadIdx.x / 32;
-    int64_t threadid = threadIdx.x % 32;
-    extern __shared__ b32 bfrag_arr[]; // num_chunks * warps_per_block * 128
-    int64_t real_num_chunks = ((blockid + 1) * num_chunks) > total_num_chunks ? (total_num_chunks - (blockid * num_chunks)) : num_chunks;
-    int64_t diff_num_chunks = real_num_chunks - num_chunks;
-
-    b32* a_start_ptr = (b32*) (a + blockid * num_chunks * 256); // offset a to where this warp starts
-    b32* out_start_ptr = (b32*) (out + blockid * num_chunks * 256);
-    b32* a_ptr = a_start_ptr + threadid * 4;
-    b32* b_frag_ptr = bfrag_arr + (blockid % warps_per_block) * num_chunks * 128 + threadid * 4;
-
-    #if (__CUDA_ARCH__ < 900) // SM80, SM89
-    uint64_t cache_policy;
-    asm volatile(
-        "createpolicy.fractional.L2::evict_first.b64 %0, 1.0;\n"
-        : "=l"(cache_policy)
-    );
-    #endif
-
-    #pragma unroll
-    for (int64_t k = 0; k < num_chunks; k++) {
-        size_t shared_ptr = __cvta_generic_to_shared(b_frag_ptr);
-        #if (__CUDA_ARCH__ >= 900) // SM90
-            asm volatile(
-                "cp.async.cg.shared.global [%0], [%1], 16;\n"
-                "cp.async.commit_group;\n"
-                :: "l"(shared_ptr), "l"(a_ptr)
-            );
-        #else // SM80, SM89
-            asm volatile(
-                "cp.async.cg.shared.global.L2::cache_hint.L2::256B [%0], [%1], 16, %2;\n"
-                "cp.async.commit_group;\n"
-                :: "l"(shared_ptr), "l"(a_ptr), "l"(cache_policy)
-            );
-        #endif
-
-        a_ptr += 128;
-        b_frag_ptr += 128;
-    }
-
-    // generate hadamard 16x16 (up to 2 of them)
-    constexpr b16 fp16_1p[4] = {0b0011100110101000, 0b0011100000000000, 0b0011010110101000, 0b0011010000000000};
-    constexpr b16 fp16_1n[4] = {0b1011100110101000, 0b1011100000000000, 0b1011010110101000, 0b1011010000000000};
-    constexpr b16 bf16_1p[4] = {0b0011111100110101, 0b0011111100000000, 0b0011111010110101, 0b0011111010000000};
-    constexpr b16 bf16_1n[4] = {0b1011111100110101, 0b1011111100000000, 0b1011111010110101, 0b1011111010000000};
-
-    #define val_type_1p(i) (((dtype) == torch::ScalarType::Half) ? (fp16_1p[i]) : (bf16_1p[i]))
-    #define val_type_1n(i) (((dtype) == torch::ScalarType::Half) ? (fp16_1n[i]) : (bf16_1n[i]))
-    constexpr b16 val_1p[4] = {val_type_1p(0), val_type_1p(1), val_type_1p(2), val_type_1p(3)};
-    constexpr b16 val_1n[4] = {val_type_1n(0), val_type_1n(1), val_type_1n(2), val_type_1n(3)};
-
-    constexpr b32 p_p[4] = {p_p(0), p_p(1), p_p(2), p_p(3)};
-    constexpr b32 p_n[4] = {p_n(0), p_n(1), p_n(2), p_n(3)};
-    constexpr b32 n_p[4] = {n_p(0), n_p(1), n_p(2), n_p(3)};
-    constexpr b32 n_n[4] = {n_n(0), n_n(1), n_n(2), n_n(3)};
-    const b32 had_16_p1[4][4] = {
-        {
-            0b10001000010001000010001000010001,
-            0b00000000000000000000000000000000,
-            0b00000000000000000000000000000000,
-            0b10001000010001000010001000010001
-        },
-        {
-            0b11001100100010000011001100100010,
-            0b00000000000000000000000000000000,
-            0b00000000000000000000000000000000,
-            0b11001100100010000011001100100010
-        },
-        {
-            0b11111111101010101100110010011001,
-            0b00000000000000000000000000000000,
-            0b00000000000000000000000000000000,
-            0b11111111101010101100110010011001
-        },
-        {
-            0b11111111101010101100110010011001,
-            0b11111111101010101100110010011001,
-            0b11111111101010101100110010011001,
-            0b00000000010101010011001101100110
-        }
-    };
-    const b32 had_16_p2[4][4] = {
-        {
-            0b10000000010000000010000000010000,
-            0b00000000000000000000000000000000,
-            0b00000000000000000000000000000000,
-            0b10000000010000000010000000010000
-        },
-        {
-            0b11000000100001000011000000100001,
-            0b00000000000000000000000000000000,
-            0b00000000000000000000000000000000,
-            0b11000000100001000011000000100001
-        },
-        {
-            0b11110000101001011100001110010110,
-            0b00000000000000000000000000000000,
-            0b00000000000000000000000000000000,
-            0b11110000101001011100001110010110
-        },
-        {
-            0b11110000101001011100001110010110,
-            0b11110000101001011100001110010110,
-            0b11110000101001011100001110010110,
-            0b00001111010110100011110001101001
-        }
-    };
-    const b32 had_16_mask[3][4] = {
-        {
-            0b10001000010001000010001000010001,
-            0b00000000000000000000000000000000,
-            0b00000000000000000000000000000000,
-            0b10001000010001000010001000010001
-        },
-        {
-            0b11001100110011000011001100110011,
-            0b00000000000000000000000000000000,
-            0b00000000000000000000000000000000,
-            0b11001100110011000011001100110011
-        },
-        {
-            0b11111111111111111111111111111111,
-            0b00000000000000000000000000000000,
-            0b00000000000000000000000000000000,
-            0b11111111111111111111111111111111
-        }
-    };
-    b32 had_frag[8];
-    #pragma unroll
-    for (int64_t i = 0; i < 2; i++) {
-        int64_t c_log_h = (i == 0) ? MIN(4, log_had_size) : log_had_size % 4;
-        #pragma unroll
-        for (int64_t j = 0; j < 4; j++) {
-            if (c_log_h < 4) {
-                bool mask = had_16_mask[c_log_h - 1][j] & (1 << (31 - threadid));
-                if (!mask) {
-                    had_frag[i * 4 + j] = 0;
-                    continue;
-                }
-            }
-            bool pred1 = had_16_p1[c_log_h - 1][j] & (1 << (31 - threadid));
-            bool pred2 = had_16_p2[c_log_h - 1][j] & (1 << (31 - threadid));
-            b32 val = pred1 ? (pred2 ? p_p[c_log_h - 1] : p_n[c_log_h - 1]) : (pred2 ? n_p[c_log_h - 1] : n_n[c_log_h - 1]);
-            had_frag[i * 4 + j] = val;
-        }
-        if constexpr(log_had_size <= 4 || log_had_size % 4 == 0) break;
-    }
-
-    // log had size above 8, only used for above 2^8 = 256 size
-    constexpr int64_t part8_log_had_size = log_had_size - 8;
-
-    b32* a_chunk_ptr = a_start_ptr; // first chunk starts at this warp's data starts
-    b32* out_chunk_ptr = out_start_ptr;
-
-    #pragma unroll
-    for (int64_t l = 0; l < 2; l++) {
-        if constexpr(log_had_size <= 8) { // l == 0 guaranteed, redundant simplified version of else body, to help compiler warnings
-            b_frag_ptr = bfrag_arr + (blockid % warps_per_block) * num_chunks * 128;
-        } else {
-            b_frag_ptr = bfrag_arr + (blockid % warps_per_block) * num_chunks * (l == 0 ? 128 : (128 >> part8_log_had_size));
-        }
-
-        if (l == 1) {
-            if constexpr(log_had_size > 8) {
-                __syncthreads(); // sync between first and second iterations if above size 256
-
-                if constexpr(log_had_size >= 12) {
-                    // sizes 4k and above
-
-                    // a + threadblock offset + warp offset
-                    // can then index into all chunks owned by this warp
-                    b32* store = bfrag_arr + (128 >> part8_log_had_size) * (num_chunks * (blockid % warps_per_block));
-
-                    #pragma unroll
-                    for (int64_t j = 0; j < 4; j++) {
-                        #pragma unroll
-                        for (int64_t k = 0; k < num_chunks; k++) {
-                            // here, j represents register, and k represents 8-offset/chunk
-                            uint64_t real_chunk_num = (num_chunks - (threadid % num_chunks) + k) % num_chunks; // chunk at which you have target thread #'s data
-                            
-                            int64_t real_thread_id = (threadid / num_chunks) * num_chunks + k; // target thread #
-                            int64_t chunk_idx = 128 * real_chunk_num; // index due to fetching from another chunk (chunk in which this thread has the target thread's original data)
-                            int64_t thread_group_idx = (real_thread_id / 4) * 16; // index due to fetching from another group of num_chunk threads (since shuffle is between num_chunk threads)
-                            int64_t thread_idx = (real_thread_id % 4) * 2; // index due to original thread's position within the group of num_chunk threads
-                            int64_t reg_idx = (j / 2) * 8 + (j % 2); // index due to target register
-                            int64_t idx = chunk_idx + thread_group_idx + thread_idx + reg_idx; // final index
-
-                            // fix idx for majorness
-                            int64_t rowidx = idx % (1 << part8_log_had_size);
-                            int64_t colidx = idx >> part8_log_had_size;
-
-                            // store[rowidx * 128 + colidx] = data;
-                            b32 data = store[rowidx * 128 + colidx];
-
-                            // compiler generates excessive instructions, so we manually do the if statement
-                            #pragma unroll
-                            for (uint64_t i = 0; i < num_chunks; i++) {
-                                asm volatile (
-                                    "{\n\t"
-                                    "  .reg .pred p0;\n\t"
-                                    "  setp.eq.s64 p0, %1, %2;\n\t"
-                                    "  @p0 mov.b32 %0, %3;\n\t"
-                                    "}\n\t"
-                                    : "+r"(b_frag_all[i][j]) // Output operand %0
-                                    : "l"(real_chunk_num), "l"(i), "r"(data) // Input operands %1, %2, %3
-                                );
-                            }
-                        }
-                    }
-
-                    #pragma unroll
-                    for (int64_t j = 0; j < 4; j++) {
-                        #pragma unroll
-                        for (int64_t k = 1; k < num_chunks; k++) {
-                            int64_t threadid_contig = threadid % num_chunks;
-                            int64_t threadid_mul = threadid / num_chunks;
-                            int64_t threadid2 = (threadid_contig + num_chunks - k) % num_chunks + threadid_mul * num_chunks; // thread to give your data to
-                            b_frag_all[k][j] = __shfl_sync(0xFFFFFFFF, b_frag_all[k][j], threadid2);
-                        }
-                    }
-                }
-            }
-        }
-
-        #pragma unroll
-        for (int64_t k = 0; k < num_chunks; k++) {
-            if constexpr(enable_mask) {
-                if (k >= real_num_chunks)
-                    break;
-            }
-            if (l == 0) {
-                // bad fix for k not being recognized as a constexpr by compiler
-                // asm("cp.async.wait_group %0;\n" :: "n"(num_chunks - k - 1));
-                #define SWITCH_WAIT_ASYNC_LOAD_GROUP(i) case i: asm volatile("cp.async.wait_group %0;\n" :: "n"(num_chunks - i - 1)); break;
-                if constexpr(enable_mask) {
-                    switch(k + diff_num_chunks) {
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(0)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(1)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(2)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(3)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(4)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(5)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(6)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(7)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(8)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(9)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(10)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(11)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(12)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(13)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(14)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(15)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(16)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(17)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(18)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(19)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(20)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(21)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(22)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(23)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(24)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(25)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(26)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(27)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(28)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(29)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(30)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(31)
-                    }
-                } else {
-                    switch(k) {
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(0)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(1)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(2)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(3)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(4)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(5)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(6)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(7)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(8)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(9)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(10)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(11)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(12)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(13)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(14)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(15)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(16)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(17)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(18)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(19)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(20)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(21)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(22)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(23)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(24)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(25)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(26)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(27)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(28)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(29)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(30)
-                        SWITCH_WAIT_ASYNC_LOAD_GROUP(31)
-                    }
-                }
-            }
-
-            if (l == 0) {
-                // loading for the first iteration
-
-                // thread 0 loads  [t0r0, t16r1, t0r2, t16r3]
-                // thread 16 loads [t0r1, t16r0, t0r3, t16r2]
-                // allows full coalescing, same for t1/t17, t2/t18, etc.
-                #pragma unroll
-                for (int64_t j = 0; j < 4; j++) {
-                    int64_t reg = ((threadid & 16) == 0) ? j : (j / 2 * 2 + (1 - j % 2));
-                    int64_t real_thread_id = (reg == 0 || reg == 2) ? threadid : (threadid ^ 16);
-                    int64_t real_row = real_thread_id % 4;
-                    int64_t real_col = real_thread_id / 4;
-                    b_frag_all[k][j] = b_frag_ptr[(real_row + (reg % 2) * 4) + (real_col + (j / 2) * 8) * 8];
-                }
-
-                // for t16 swap r0/r1 and r2/r3 to have [t16r0, t0r1, t16r2, t0r3]
-                // so registers are in right order, same for t17, t18, etc.
-                if ((threadid & 16) != 0) {
-                    b32 temp = b_frag_all[k][0];
-                    b_frag_all[k][0] = b_frag_all[k][1];
-                    b_frag_all[k][1] = temp;
-
-                    temp = b_frag_all[k][2];
-                    b_frag_all[k][2] = b_frag_all[k][3];
-                    b_frag_all[k][3] = temp;
-                }
-
-                // t0 and t16 swap r1 and r3 to have their own data,
-                // same for t1/t17, t2/18, etc.
-                #pragma unroll
-                for (int64_t j = 1; j < 4; j += 2) {
-                    b_frag_all[k][j] = __shfl_xor_sync(0xFFFFFFFF, b_frag_all[k][j], 16);
-                }
-            } else if constexpr(log_had_size > 8) { // condition is redundant to help compiler warnings
-                if constexpr(log_had_size < 12) {
-                    // sizes 512, 1k, and 2k
-
-                    // for 512:
-                    //     thread 0 loads  [t0r0, t0r1, t16r2, t16r3]
-                    //     thread 16 loads [t0r2, t0r3, t16r0, t16r1]
-                    //     same for t1/t17, t2/t18, etc.
-                    // for 1k and 2k:
-                    //     thread 0 loads [t0r0, t0r1, t1r2, t1r3]
-                    //     thread 1 loads [t0r2, t0r3, t1r0, t1r1]
-                    //     same for t2/t3, t4/t5, etc.
-                    // allows full coalescing for 512 and 1k, 16x coalescing for 2k
-                    constexpr int64_t xor_val = log_had_size == 9 ? 16 : 1;
-
-                    #pragma unroll
-                    for (int64_t j = 0; j < 4; j++) {
-                        int64_t reg = ((threadid & xor_val) == 0) ? j : (j + 2) % 4;
-                        int64_t real_thread_id = reg < 2 ? threadid : (threadid ^ xor_val);
-                        int64_t idx = (real_thread_id / 4 * 16) + (real_thread_id % 4 * 2) + (reg / 2 * 8) + (reg % 2);
-                        int64_t rowidx = idx % (1 << part8_log_had_size);
-                        int64_t colidx = idx >> part8_log_had_size;
-                        b_frag_all[k][j] = b_frag_ptr[rowidx * 128 + colidx];
-                    }
-
-                    if ((threadid & xor_val) != 0) {
-                        b32 temp = b_frag_all[k][0];
-                        b_frag_all[k][0] = b_frag_all[k][2];
-                        b_frag_all[k][2] = temp;
-
-                        temp = b_frag_all[k][1];
-                        b_frag_all[k][1] = b_frag_all[k][3];
-                        b_frag_all[k][3] = temp;
-                    }
-
-                    #pragma unroll
-                    for (int64_t j = 2; j < 4; j++) {
-                        b_frag_all[k][j] = __shfl_xor_sync(0xFFFFFFFF, b_frag_all[k][j], xor_val);
-                    }
-                }
-            }
-
-            if (l == 1) {
-                // for second iteration, we load 2 consecutive b16s (1 b32) per register,
-                // but tensor core register layout requires 2 b16s that are in the
-                // same column/consecutive rows to be in the same register, so do the swap
-                b32 f0 = ((b_frag_all[k][1] & 0xFFFF) << 16) | (b_frag_all[k][0] & 0xFFFF);
-                b32 f1 = ((b_frag_all[k][3] & 0xFFFF) << 16) | (b_frag_all[k][2] & 0xFFFF);
-                b32 f2 = (b_frag_all[k][1] & 0xFFFF0000) | (b_frag_all[k][0] >> 16);
-                b32 f3 = (b_frag_all[k][3] & 0xFFFF0000) | (b_frag_all[k][2] >> 16);
-                b_frag_all[k][0] = f0;
-                b_frag_all[k][1] = f1;
-                b_frag_all[k][2] = f2;
-                b_frag_all[k][3] = f3;
-            }
-
-            #pragma unroll
-            for(int64_t i = 0, remaining_log_had_size = log_had_size - l * 8; i < 2 && remaining_log_had_size > 0; i++) {
-                int64_t had_off = ((remaining_log_had_size < 4) && !(log_had_size <= 4 || log_had_size % 4 == 0)) ? 4 : 0;
-                mma_m16_n16_k16_b16_b16_b16_noacc<dtype>(had_frag[had_off + 0], had_frag[had_off + 1], had_frag[had_off + 2], had_frag[had_off + 3], b_frag_all[k][0], b_frag_all[k][1], b_frag_all[k][2], b_frag_all[k][3], b_frag_all[k][0], b_frag_all[k][1], b_frag_all[k][2], b_frag_all[k][3]);
-
-                remaining_log_had_size -= 4;
-                if (remaining_log_had_size <= 0 && i == 0) {
-                    // TODO: consider different storing so no need for transpose
-                    matrix_transpose_m8_n8_b16_inplace(b_frag_all[k][0]);
-                    matrix_transpose_m8_n8_b16_inplace(b_frag_all[k][1]);
-                    matrix_transpose_m8_n8_b16_inplace(b_frag_all[k][2]);
-                    matrix_transpose_m8_n8_b16_inplace(b_frag_all[k][3]);
-                } else {
-                    // swap and use output directly as b_frag for next iteration as an actually free transpose
-                    b32 temp = b_frag_all[k][1];
-                    b_frag_all[k][1] = b_frag_all[k][2];
-                    b_frag_all[k][2] = temp;
-                }
-            }
-
-            if (l == 1) {
-                // invert swap from above for second iteration
-                b32 f0 = ((b_frag_all[k][2] & 0xFFFF) << 16) | (b_frag_all[k][0] & 0xFFFF);
-                b32 f1 = (b_frag_all[k][2] & 0xFFFF0000) | (b_frag_all[k][0] >> 16);
-                b32 f2 = ((b_frag_all[k][3] & 0xFFFF) << 16) | (b_frag_all[k][1] & 0xFFFF);
-                b32 f3 = (b_frag_all[k][3] & 0xFFFF0000) | (b_frag_all[k][1] >> 16);
-                b_frag_all[k][0] = f0;
-                b_frag_all[k][1] = f1;
-                b_frag_all[k][2] = f2;
-                b_frag_all[k][3] = f3;
-            }
-
-            if (l == 0) {
-                // inverse of coalesced load for first iteration to store result
-                #pragma unroll
-                for (int64_t j = 1; j < 4; j += 2) {
-                    b_frag_all[k][j] = __shfl_xor_sync(0xFFFFFFFF, b_frag_all[k][j], 16);
-                }
-
-                if ((threadid & 16) != 0) {
-                    b32 temp = b_frag_all[k][0];
-                    b_frag_all[k][0] = b_frag_all[k][1];
-                    b_frag_all[k][1] = temp;
-
-                    temp = b_frag_all[k][2];
-                    b_frag_all[k][2] = b_frag_all[k][3];
-                    b_frag_all[k][3] = temp;
-                }
-
-                // if only going up to 256 size, store directly back to global memory,
-                // otherwise store back to shared memory for next iteration
-                b32* store = (log_had_size <= 8) ? out_chunk_ptr : b_frag_ptr;
-
-                #pragma unroll
-                for (int64_t j = 0; j < 4; j++) {
-                    int64_t reg = ((threadid & 16) == 0) ? j : (j / 2 * 2 + (1 - j % 2));
-                    int64_t real_thread_id = (reg == 0 || reg == 2) ? threadid : (threadid ^ 16);
-                    int64_t real_row = real_thread_id % 4;
-                    int64_t real_col = real_thread_id / 4;
-                    store[(real_row + (reg % 2) * 4) + (real_col + (reg / 2) * 8) * 8] = b_frag_all[k][j];
-                }
-            } else if constexpr(log_had_size > 8) { // condition is redundant to help compiler warnings
-                if (log_had_size < 12) {
-                    // inverse of coalesced load for sizes 512, 1k and 2k to store result
-                    constexpr int xor_val = log_had_size == 9 ? 16 : 1;
-                    #pragma unroll
-                    for (int64_t j = 2; j < 4; j++) {
-                        b_frag_all[k][j] = __shfl_xor_sync(0xFFFFFFFF, b_frag_all[k][j], xor_val);
-                    }
-
-                    if ((threadid & xor_val) != 0) {
-                        b32 temp = b_frag_all[k][0];
-                        b_frag_all[k][0] = b_frag_all[k][2];
-                        b_frag_all[k][2] = temp;
-
-                        temp = b_frag_all[k][1];
-                        b_frag_all[k][1] = b_frag_all[k][3];
-                        b_frag_all[k][3] = temp;
-                    }
-
-                    b32* store = (b32*)(out + (blockid / warps_per_block) * (num_chunks * warps_per_block) * 256 + (256 >> part8_log_had_size) * (num_chunks * (blockid % warps_per_block) + k));
-                    #pragma unroll
-                    for (int64_t j = 0; j < 4; j++) {
-                        int64_t reg = ((threadid & xor_val) == 0) ? j : (j + 2) % 4;
-                        b32 data = b_frag_all[k][j];
-                        int64_t real_thread_id = reg < 2 ? threadid : (threadid ^ xor_val);
-                        int64_t idx = (real_thread_id / 4 * 16) + (real_thread_id % 4 * 2) + (reg / 2 * 8) + (reg % 2);
-                        int64_t rowidx = idx % (1 << part8_log_had_size);
-                        int64_t colidx = idx >> part8_log_had_size;
-                        store[rowidx * 128 + colidx] = data;
-                    }
-                }
-                // for size 4k and above, wait to process all chunks so a final store can be performed coalesced
-            }
-
-            a_chunk_ptr += 128; // (only affects first 256 size) move on to next chunk by skipping 256 elements in b16 (= 128 in b32)
-            out_chunk_ptr += 128;
-            if constexpr(log_had_size > 8) {
-                b_frag_ptr += (l == 0 ? 128 : (128 >> part8_log_had_size));
-            } else { // else is redundant, simplified version of if body, to help compiler warnings
-                b_frag_ptr += 128;
-            }
-        }
-        if (log_had_size <= 8)
-            break;
-    }
-
-    if constexpr(log_had_size >= 12) {
-        // for sizes 4k and above, perform final coalesced store after processing all chunks
-        #pragma unroll
-        for (int64_t j = 0; j < 4; j++) {
-            #pragma unroll
-            for (int64_t k = 1; k < num_chunks; k++) {
-                int64_t threadid_contig = threadid % num_chunks;
-                int64_t threadid_mul = threadid / num_chunks;
-                int64_t threadid2 = (threadid_contig + k) % num_chunks + threadid_mul * num_chunks; // thread to give your data to
-                b_frag_all[k][j] = __shfl_sync(0xFFFFFFFF, b_frag_all[k][j], threadid2);
-            }
-        }
-
-        // a + threadblock offset + warp offset
-        // can then index into all chunks owned by this warp
-        b32* store = bfrag_arr + (128 >> part8_log_had_size) * (num_chunks * (blockid % warps_per_block));
-
-        #pragma unroll
-        for (int64_t j = 0; j < 4; j++) {
-            #pragma unroll
-            for (int64_t k = 0; k < num_chunks; k++) {
-                // here, j represents register, and k represents 8-offset/chunk
-                int64_t real_chunk_num = (num_chunks - (threadid % num_chunks) + k) % num_chunks; // chunk at which you have target thread #'s data
-
-                // b32 data = b_frag_all[real_chunk_num][j]; // target thread data
-                b32 data;
-                #pragma unroll
-                for (int64_t i = 0; i < num_chunks; i++) {
-                    if (real_chunk_num == i) data = b_frag_all[i][j];
-                }
-                
-                int64_t real_thread_id = (threadid / num_chunks) * num_chunks + k; // target thread #
-                int64_t chunk_idx = 128 * real_chunk_num; // index due to fetching from another chunk (chunk in which this thread has the target thread's original data)
-                int64_t thread_group_idx = (real_thread_id / 4) * 16; // index due to fetching from another group of num_chunk threads (since shuffle is between num_chunk threads)
-                int64_t thread_idx = (real_thread_id % 4) * 2; // index due to original thread's position within the group of num_chunk threads
-                int64_t reg_idx = (j / 2) * 8 + (j % 2); // index due to target register
-                int64_t idx = chunk_idx + thread_group_idx + thread_idx + reg_idx; // final index
-
-                // fix idx for majorness
-                int64_t rowidx = idx % (1 << part8_log_had_size);
-                int64_t colidx = idx >> part8_log_had_size;
-
-                store[rowidx * 128 + colidx] = data;
-            }
-        }
-
-        __syncthreads();
-        store = ((b32*) out) + (blockid / warps_per_block) * (num_chunks * warps_per_block) * 128;
-        int4* store4 = (int4*) store;
-        int4* bfrag_arr4 = (int4*) bfrag_arr;
-        // flush smem, simply linearly write to store
-        // always divisible by 128*32b, so (32*4)*32b is ok
-        #pragma unroll
-        for (int64_t warp_off = 0; warp_off < (num_chunks * warps_per_block * 128 / 4); warp_off += 32 * warps_per_block) {
-            int64_t total_off = warp_off + threadid + (blockid % warps_per_block) * 32;
-            store4[total_off] = bfrag_arr4[total_off];
-        }
-    }
-
-}
-
-constexpr int64_t ceil_div(int64_t a, int64_t b) {
-    return (a + b - 1) / b;
-}
-
-template <torch::ScalarType dtype, int64_t chunks_per_warp, int64_t warps_per_block, int64_t log_had_size, int64_t blocks_per_sm, bool check_masking = false>
-void __forceinline__ run_kernel(b16* a_mat, b16* out, int64_t num_chunks, cudaStream_t stream) {
-    int64_t shared_size = chunks_per_warp * warps_per_block * 128 * 4;
-    dim3 block_size = 32 * warps_per_block;
-
-    #define CHECK_SHARED_LIM() {                                                                              \
-        if (shared_size > 48 * 1024) {                                                                        \
-            C10_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, 65536)); \
-        }                                                                                                     \
-    }                                                                                                         \
-
-    if constexpr(check_masking) {
-        if (num_chunks % (chunks_per_warp * warps_per_block) != 0) {
-            dim3 grid_size = ceil_div(ceil_div(num_chunks, chunks_per_warp), warps_per_block);
-            auto kernel = hadamard_transform_kernel<chunks_per_warp, warps_per_block, log_had_size, blocks_per_sm, true, dtype>;
-            CHECK_SHARED_LIM();
-            kernel<<<dim3(grid_size), dim3(block_size), shared_size, stream>>>(a_mat, out, num_chunks);
-        } else {
-            dim3 grid_size = num_chunks / chunks_per_warp / warps_per_block;
-            auto kernel = hadamard_transform_kernel<chunks_per_warp, warps_per_block, log_had_size, blocks_per_sm, false, dtype>;
-            CHECK_SHARED_LIM();
-            kernel<<<dim3(grid_size), dim3(block_size), shared_size, stream>>>(a_mat, out, num_chunks);
-        }
-    } else {
-        dim3 grid_size = num_chunks / chunks_per_warp / warps_per_block;
-        auto kernel = hadamard_transform_kernel<chunks_per_warp, warps_per_block, log_had_size, blocks_per_sm, false, dtype>;
-        CHECK_SHARED_LIM();
-        kernel<<<dim3(grid_size), dim3(block_size), shared_size, stream>>>(a_mat, out, num_chunks);
-    }
-    
-    C10_CUDA_KERNEL_LAUNCH_CHECK();
-}
-
-template <torch::ScalarType dtype>
-void run_fht(void* a_mat_ptr, void* out_ptr, int64_t numel, int64_t had_size, cudaStream_t stream) {
-    int64_t num_chunks = numel / 256; // caller required to ensure divisible by 256
-    // for size 256, use (2, 1)
-    // for size 32k use (8, 16)
-    constexpr int64_t chunks_per_warp_small = 1;// 8;
-    constexpr int64_t warps_per_block_small = 1;//2;//16;
-    constexpr int64_t blocks_per_sm_small = 24;
-    constexpr int64_t chunks_per_warp_large = 2;
-    constexpr int64_t warps_per_block_large = 1;
-    constexpr int64_t blocks_per_sm_large = 24;
-
-    b16* a_mat = (b16*) a_mat_ptr;
-    b16* out = (b16*) out_ptr;
-
-    if (numel <= 256) {
-        switch (had_size) {
-            case (1<<1): run_kernel<dtype, chunks_per_warp_small, warps_per_block_small, 1, blocks_per_sm_small>(a_mat, out, num_chunks, stream); break;
-            case (1<<2): run_kernel<dtype, chunks_per_warp_small, warps_per_block_small, 2, blocks_per_sm_small>(a_mat, out, num_chunks, stream); break;
-            case (1<<3): run_kernel<dtype, chunks_per_warp_small, warps_per_block_small, 3, blocks_per_sm_small>(a_mat, out, num_chunks, stream); break;
-            case (1<<4): run_kernel<dtype, chunks_per_warp_small, warps_per_block_small, 4, blocks_per_sm_small>(a_mat, out, num_chunks, stream); break;
-            case (1<<5): run_kernel<dtype, chunks_per_warp_small, warps_per_block_small, 5, blocks_per_sm_small>(a_mat, out, num_chunks, stream); break;
-            case (1<<6): run_kernel<dtype, chunks_per_warp_small, warps_per_block_small, 6, blocks_per_sm_small>(a_mat, out, num_chunks, stream); break;
-            case (1<<7): run_kernel<dtype, chunks_per_warp_small, warps_per_block_small, 7, blocks_per_sm_small>(a_mat, out, num_chunks, stream); break;
-            case (1<<8): run_kernel<dtype, chunks_per_warp_small, warps_per_block_small, 8, blocks_per_sm_small>(a_mat, out, num_chunks, stream); break;
-        }
-    } else {
-        switch (had_size) {
-            case (1<<1):  run_kernel<dtype, chunks_per_warp_large, warps_per_block_large, 1, blocks_per_sm_large, true>(a_mat, out, num_chunks, stream); break;
-            case (1<<2):  run_kernel<dtype, chunks_per_warp_large, warps_per_block_large, 2, blocks_per_sm_large, true>(a_mat, out, num_chunks, stream); break;
-            case (1<<3):  run_kernel<dtype, chunks_per_warp_large, warps_per_block_large, 3, blocks_per_sm_large, true>(a_mat, out, num_chunks, stream); break;
-            case (1<<4):  run_kernel<dtype, chunks_per_warp_large, warps_per_block_large, 4, blocks_per_sm_large, true>(a_mat, out, num_chunks, stream); break;
-            case (1<<5):  run_kernel<dtype, chunks_per_warp_large, warps_per_block_large, 5, blocks_per_sm_large, true>(a_mat, out, num_chunks, stream); break;
-            case (1<<6):  run_kernel<dtype, chunks_per_warp_large, warps_per_block_large, 6, blocks_per_sm_large, true>(a_mat, out, num_chunks, stream); break;
-            case (1<<7):  run_kernel<dtype, chunks_per_warp_large, warps_per_block_large, 7, blocks_per_sm_large, true>(a_mat, out, num_chunks, stream); break;
-            case (1<<8):  run_kernel<dtype, chunks_per_warp_large, warps_per_block_large, 8, blocks_per_sm_large, true>(a_mat, out, num_chunks, stream); break;
-            case (1<<9):  run_kernel<dtype, launch_configs_big[0][0], launch_configs_big[0][1], 9 , launch_configs_big[0][2]>(a_mat, out, num_chunks, stream); break;
-            case (1<<10): run_kernel<dtype, launch_configs_big[1][0], launch_configs_big[1][1], 10, launch_configs_big[1][2]>(a_mat, out, num_chunks, stream); break;
-            case (1<<11): run_kernel<dtype, launch_configs_big[2][0], launch_configs_big[2][1], 11, launch_configs_big[2][2]>(a_mat, out, num_chunks, stream); break;
-            case (1<<12): run_kernel<dtype, launch_configs_big[3][0], launch_configs_big[3][1], 12, launch_configs_big[3][2]>(a_mat, out, num_chunks, stream); break;
-            case (1<<13): run_kernel<dtype, launch_configs_big[4][0], launch_configs_big[4][1], 13, launch_configs_big[4][2]>(a_mat, out, num_chunks, stream); break;
-            case (1<<14): run_kernel<dtype, launch_configs_big[5][0], launch_configs_big[5][1], 14, launch_configs_big[5][2]>(a_mat, out, num_chunks, stream); break;
-            case (1<<15): run_kernel<dtype, launch_configs_big[6][0], launch_configs_big[6][1], 15, launch_configs_big[6][2]>(a_mat, out, num_chunks, stream); break;
-        }
-    }
-}
-
-template void run_fht<torch::ScalarType::Half>(void* a_mat_ptr, void* out_ptr, int64_t numel, int64_t had_size, cudaStream_t stream);
-template void run_fht<torch::ScalarType::BFloat16>(void* a_mat_ptr, void* out_ptr, int64_t numel, int64_t had_size, cudaStream_t stream);
-
-}  // namespace hadacore
-
-constexpr bool is_power_of_two(int x) { return x && !(x & (x - 1)); }
-
-torch::Tensor hadacore_transform(torch::Tensor& x, bool inplace) {
-    auto dtype = x.scalar_type();
-    TORCH_CHECK(dtype == torch::ScalarType::Half || dtype == torch::ScalarType::BFloat16, "Only fp16 and bf16 supported currently");
-    TORCH_CHECK(x.is_cuda());
-    
-    const int had_size = x.size(-1);
-    TORCH_CHECK(is_power_of_two(had_size) && (had_size <= (1U << 15)),
-        "Only power of two Hadamard sizes up to 2^15 are supported, got ", had_size);
-    
-    const auto res_shape = x.sizes();
-    x = x.reshape({-1, had_size});
-    
-    auto numel = x.numel();
-    if (numel % 256 != 0) {
-        x = torch::nn::functional::pad(x, torch::nn::functional::PadFuncOptions({0, 0, 0, (256 - numel % 256) / had_size}));
-    }
-    
-    if (x.stride(-1) != 1) {
-        x = x.contiguous();
-    }
-    torch::Tensor out = inplace ? x : torch::empty_like(x);
-
-    at::cuda::CUDAGuard device_guard{(char)x.get_device()};
-    auto stream = at::cuda::getCurrentCUDAStream().stream();
-
-    VLLM_DISPATCH_HALF_TYPES(x.scalar_type(), "hadacore_transform_runfht", [&] {
-      auto constexpr SCALAR_TYPE = c10::CppTypeToScalarType<scalar_t>::value;
-      hadacore::run_fht<SCALAR_TYPE>(x.data_ptr(), x.data_ptr(), x.numel(), had_size, stream);
-    });
-
-    if (numel % 256 != 0) {
-        out = out.index({torch::indexing::Slice(0, numel / had_size)});
-    }
-
-    if (inplace && out.data_ptr() != x.data_ptr()) {
-        x.copy_(out.view(res_shape));
-        return x;
-    }
-    return out.reshape(res_shape);
-}
-
-TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
-    m.impl("hadacore_transform", &hadacore_transform);
-}

csrc/rocm/attention.cu

@@ -25,21 +25,11 @@
 #include "../attention/dtype_fp8.cuh"
 #include "../quantization/fp8/amd/quant_utils.cuh"
 
-// ROCm 6.2 compatibility: map OCP fp8 types to FNUZ variants if OCP is absent
-#if !defined(HIP_FP8_TYPE_OCP)
-using __hip_fp8_e4m3 = __hip_fp8_e4m3_fnuz;
-using __hip_fp8_e5m2 = __hip_fp8_e5m2_fnuz;
-#endif
-
 #if defined(__HIPCC__) && \
     (defined(__gfx90a__) || defined(__gfx942__) || defined(__gfx950__))
   #define __HIP__GFX9__
 #endif
 
-#if defined(__HIPCC__) && (defined(__gfx942__) || defined(__gfx950__))
-  #define __HIP__FP8MFMA__
-#endif
-
 #if defined(__HIPCC__) && (defined(__gfx1100__) || defined(__gfx1101__))
   #define __HIP__GFX11__
 #endif
@@ -61,12 +51,6 @@ using __hip_fp8_e5m2 = __hip_fp8_e5m2_fnuz;
 #define MIN(a, b) ((a) < (b) ? (a) : (b))
 #define DIVIDE_ROUND_UP(a, b) (((a) + (b) - 1) / (b))
 
-enum class MFMAType {
-  F16 = 0,
-  Fp8 = 1,
-  Fp4 = 2,
-};
-
 #if defined(__HIP__GFX9__)
 
   #define GCN_MFMA_INSTR1 __builtin_amdgcn_mfma_f32_16x16x4f32
@@ -128,21 +112,6 @@ __device__ __forceinline__ floatx4 gcn_mfma16x16x16_instr(const _B16x4& inpA,
   }
 }
 
-template <typename T, int absz, int cbid, int blgp>
-__device__ __forceinline__ floatx4 gcn_mfma16x16x32_instr(const long& inpA,
-                                                          const long& inpB,
-                                                          const floatx4& inpC) {
-  if constexpr (std::is_same<T, __hip_fp8_e4m3>::value) {
-    return __builtin_amdgcn_mfma_f32_16x16x32_fp8_fp8(inpA, inpB, inpC, absz,
-                                                      cbid, blgp);
-  } else if constexpr (std::is_same<T, __hip_fp8_e5m2>::value) {
-    return __builtin_amdgcn_mfma_f32_16x16x32_bf8_bf8(inpA, inpB, inpC, absz,
-                                                      cbid, blgp);
-  } else {
-    static_assert(false, "unsupported 8b dtype");
-  }
-}
-
 template <typename T>
 __device__ __forceinline__ float to_float(const T& inp) {
   if constexpr (std::is_same<T, _Float16>::value) {
@@ -287,44 +256,12 @@ __device__ __forceinline__ _B16x8 convert_b8x8_custom(const _B8x8 input) {
   return ret;
 }
 
-typedef union u64_cvt {
-  half f16x4[4];
-  int16_t b16x4[4];
-  _B8x8 b8x8;
-  _B16x4 b64;
-  int64_t i64;
-} _T8x8;
-
-__device__ __forceinline__ _B8x8 convert_b16x8(const _B16x8& input,
-                                               _T8x8& Mtemp) {
-  _T8x8 Qtmp8x8;
-
-  for (int i = 0; i < 2; i++) {
-    floatx4 q_out = {0, 0, 0, 0};
-    q_out = gcn_mfma16x16x16_instr<_Float16, 0, 0, 0>(Mtemp.b64, input.xy[i],
-                                                      q_out);
-    Qtmp8x8.b16x4[i * 2] =
-        __builtin_amdgcn_cvt_pk_fp8_f32(q_out[0], q_out[1], 0, false);
-    Qtmp8x8.b16x4[i * 2 + 1] =
-        __builtin_amdgcn_cvt_pk_fp8_f32(q_out[2], q_out[3], 0, false);
-  }
-  return Qtmp8x8.b8x8;
-}
-
-__device__ float warpReduceMax(float val) {
-  for (int offset = warpSize / 2; offset > 0; offset /= 2) {
-    val = max(
-        val, __shfl_down(val, offset, WARP_SIZE));  // Using max() for reduction
-  }
-  return val;
-}
-
 // grid (num_seqs, num_partitions,num_kv_heads)
 // block (256)
 // clang-format off
 template <typename scalar_t, typename cache_t,
           vllm::Fp8KVCacheDataType KV_DTYPE, typename OUTT, int BLOCK_SIZE,
-          int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED, int GQA_RATIO, MFMAType MFMA_TYPE>
+          int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED, int GQA_RATIO>
 __global__
 __launch_bounds__(NUM_THREADS, 5) void paged_attention_ll4mi_QKV_mfma16_kernel(
     const scalar_t* __restrict__ q,         // [num_seqs, num_heads, head_size]
@@ -430,10 +367,6 @@ __launch_bounds__(NUM_THREADS, 5) void paged_attention_ll4mi_QKV_mfma16_kernel(
   const int* block_table_seq = block_tables + seq_idx * max_num_blocks_per_seq;
 
   int kphysical_block_number[TLOOP];
-  #if defined(__HIP__FP8MFMA__)
-  float q_max = 0;
-  float q_scale = 1.0;
-  #endif
 
   // fetch k physical block numbers
   for (int token_depth = 0; token_depth < TLOOP; token_depth++) {
@@ -483,15 +416,6 @@ __launch_bounds__(NUM_THREADS, 5) void paged_attention_ll4mi_QKV_mfma16_kernel(
         Qlocal[qkhe_depth][qkratio].xy[i] =
             shared_logits[qkhe_depth][rowid][lane16id % GQA_RATIO]
                          [2 * qkratio + i];
-  #if defined(__HIP__FP8MFMA__)
-        if constexpr (KV_DTYPE != vllm::Fp8KVCacheDataType::kAuto &&
-                      MFMA_TYPE == MFMAType::Fp8) {
-          scalar_t* qptr =
-              reinterpret_cast<scalar_t*>(&Qlocal[qkhe_depth][qkratio].xy[i]);
-          for (int k = 0; k < 4; k++)
-            q_max = fmax(fabs(to_float<scalar_t>(qptr[k])), q_max);
-        }
-  #endif
       }
     }
   }
@@ -591,14 +515,6 @@ __launch_bounds__(NUM_THREADS, 5) void paged_attention_ll4mi_QKV_mfma16_kernel(
   if constexpr (KV_DTYPE != vllm::Fp8KVCacheDataType::kAuto) {
     // multiply by k_scale if fp8 kv cache
     scale2 *= *k_scale;
-  #if defined(__HIP__FP8MFMA__)
-    q_max = warpReduceMax(q_max);
-    constexpr float FP8_E4M3_SCALE_TARGET = 224.0f;
-    if constexpr (MFMA_TYPE == MFMAType::Fp8) {
-      q_scale = q_max > 0 ? FP8_E4M3_SCALE_TARGET / q_max : 1.0f;
-      scale2 /= q_scale;
-    }
-  #endif
   }
 
   floatx4 d_out[TLOOP];
@@ -618,41 +534,12 @@ __launch_bounds__(NUM_THREADS, 5) void paged_attention_ll4mi_QKV_mfma16_kernel(
         auto Ktmp = Klocal[token_depth][qkhe_depth];
         _B8x16 Ktmp8x16 = *reinterpret_cast<_B8x16*>(&Ktmp);
         for (int qkratio = 0; qkratio < QK_SIZE_RATIO; qkratio++) {
-          if constexpr (MFMA_TYPE == MFMAType::F16) {
-            _B8x8 Ktmp8x8 = Ktmp8x16.xy[qkratio];
-            _B16x8 Klocaltmp = convert_b8x8_custom<scalar_t>(Ktmp8x8);
-            for (int i = 0; i < 2; i++) {
-              d_out[token_depth] = gcn_mfma16x16x16_instr<scalar_t, 0, 0, 0>(
-                  Klocaltmp.xy[i], Qlocal[qkhe_depth][qkratio].xy[i],
-                  d_out[token_depth]);
-            }
-          } else {
-  #if defined(__HIP__FP8MFMA__)
-            _T8x8 Ktmp8x8, Qtmp8x8;
-            Ktmp8x8.b8x8 = Ktmp8x16.xy[qkratio];
-
-            for (int n = 0; n < 2; n++) {
-              scalar_t* qptr = reinterpret_cast<scalar_t*>(
-                  &Qlocal[qkhe_depth][qkratio].xy[n]);
-
-              Qtmp8x8.b16x4[n * 2] =
-                  vllm::fp8::scaled_vec_conversion<uint16_t, float2>(
-                      make_float2(to_float<scalar_t>(qptr[0]),
-                                  to_float<scalar_t>(qptr[1])),
-                      q_scale);
-              Qtmp8x8.b16x4[n * 2 + 1] =
-                  vllm::fp8::scaled_vec_conversion<uint16_t, float2>(
-                      make_float2(to_float<scalar_t>(qptr[2]),
-                                  to_float<scalar_t>(qptr[3])),
-                      q_scale);
-            }
-
-            d_out[token_depth] =
-                gcn_mfma16x16x32_instr<__hip_fp8_e4m3, 0, 0, 0>(
-                    Ktmp8x8.i64, Qtmp8x8.i64, d_out[token_depth]);
-  #else
-            UNREACHABLE_CODE
-  #endif
+          _B8x8 Ktmp8x8 = Ktmp8x16.xy[qkratio];
+          _B16x8 Klocaltmp = convert_b8x8_custom<scalar_t>(Ktmp8x8);
+          for (int i = 0; i < 2; i++) {
+            d_out[token_depth] = gcn_mfma16x16x16_instr<scalar_t, 0, 0, 0>(
+                Klocaltmp.xy[i], Qlocal[qkhe_depth][qkratio].xy[i],
+                d_out[token_depth]);
           }
         }
       }
@@ -742,36 +629,17 @@ __launch_bounds__(NUM_THREADS, 5) void paged_attention_ll4mi_QKV_mfma16_kernel(
   // disable rtz conversion due to its impact on accuracy.
   constexpr bool LOGITS_RTZ_CONVERSION = false;
 
-  #if defined(__HIP__FP8MFMA__)
-  int rowid_8x8 = rowid / 2;
-  int offset = rowid % 2;
-  #endif
-
   // write logits to shared mem
   for (int token_depth = 0; token_depth < TLOOP; token_depth++) {
     d_out[token_depth] *= inv_sum_scale;
-    if constexpr (MFMA_TYPE != MFMAType::Fp8) {
-      if constexpr (LOGITS_RTZ_CONVERSION) {
-        // use rtz conversion for better performance, with negligible impact on
-        // accuracy
-        shared_logits[warpid][token_depth][lane16id][rowid] =
-            from_floatx4_rtz<scalar_t>(d_out[token_depth]);
-      } else {
-        shared_logits[warpid][token_depth][lane16id][rowid] =
-            from_floatx4<scalar_t>(d_out[token_depth]);
-      }
+    if constexpr (LOGITS_RTZ_CONVERSION) {
+      // use rtz conversion for better performance, with negligible impact on
+      // accuracy
+      shared_logits[warpid][token_depth][lane16id][rowid] =
+          from_floatx4_rtz<scalar_t>(d_out[token_depth]);
     } else {
-  #if defined(__HIP__FP8MFMA__)
-      // cast _B16x4* to _B8x8*
-      _T8x8& logits_8x8 = *reinterpret_cast<_T8x8*>(
-          &shared_logits[warpid][token_depth][lane16id][rowid_8x8]);
-      logits_8x8.b16x4[offset * 2] = __builtin_amdgcn_cvt_pk_fp8_f32(
-          d_out[token_depth][0], d_out[token_depth][1], 0, false);
-      logits_8x8.b16x4[offset * 2 + 1] = __builtin_amdgcn_cvt_pk_fp8_f32(
-          d_out[token_depth][2], d_out[token_depth][3], 0, false);
-  #else
-      UNREACHABLE_CODE
-  #endif
+      shared_logits[warpid][token_depth][lane16id][rowid] =
+          from_floatx4<scalar_t>(d_out[token_depth]);
     }
   }
 
@@ -824,42 +692,19 @@ __launch_bounds__(NUM_THREADS, 5) void paged_attention_ll4mi_QKV_mfma16_kernel(
           _B8x16 Vtmp8x16 = *reinterpret_cast<_B8x16*>(&Vtmp);
           for (int j = 0; j < ELEMS16_ELEMS8_RATIO; j++) {
             _B8x8 Vtmp8x8 = Vtmp8x16.xy[j];
-            if constexpr (MFMA_TYPE == MFMAType::F16) {
-              _B16x8 Vlocaltmp = convert_b8x8_custom<scalar_t>(Vtmp8x8);
-              for (int i = 0; i < ELEMS8_ELEMS4_RATIO; i++) {
-                const int offset =
-                    rowid * ELEMS16_ELEMS8_RATIO * ELEMS8_ELEMS4_RATIO +
-                    j * ELEMS8_ELEMS4_RATIO + i;
-                const int offset1 = offset % ROWS_PER_WARP;
-                const int offset2 = offset / ROWS_PER_WARP;
-                // output format is 16 qheads across 16 lanes, 16 head elems
-                // spread across 4 rows
-                tmp_out = gcn_mfma16x16x16_instr<scalar_t, 0, 0, 0>(
-                    Vlocaltmp.xy[i],
-                    shared_logits[vtoken_depth][offset2][lane16id][offset1],
-                    tmp_out);
-              }
-            } else {
-  #if defined(__HIP__FP8MFMA__)
-              for (int i = 0; i < ELEMS8_ELEMS4_RATIO / 2; i++) {
-                const int offset =
-                    rowid * ELEMS16_ELEMS8_RATIO * ELEMS8_ELEMS4_RATIO +
-                    j * ELEMS8_ELEMS4_RATIO + i;
-                const int offset1 = (offset % ROWS_PER_WARP) / 2;
-                const int offset2 = offset / ROWS_PER_WARP;
-                // output format is 16 qheads across 16 lanes, 16 head elems
-                // spread across 4 rows
-                tmp_out = gcn_mfma16x16x32_instr<__hip_fp8_e4m3, 0, 0, 0>(
-                    reinterpret_cast<_T8x8*>(&Vtmp8x8)->i64,
-                    reinterpret_cast<_T8x8*>(
-                        &shared_logits[vtoken_depth][offset2][lane16id]
-                                      [offset1])
-                        ->i64,
-                    tmp_out);
-              }
-  #else
-              UNREACHABLE_CODE
-  #endif
+            _B16x8 Vlocaltmp = convert_b8x8_custom<scalar_t>(Vtmp8x8);
+            for (int i = 0; i < ELEMS8_ELEMS4_RATIO; i++) {
+              const int offset =
+                  rowid * ELEMS16_ELEMS8_RATIO * ELEMS8_ELEMS4_RATIO +
+                  j * ELEMS8_ELEMS4_RATIO + i;
+              const int offset1 = offset % ROWS_PER_WARP;
+              const int offset2 = offset / ROWS_PER_WARP;
+              // output format is 16 qheads across 16 lanes, 16 head elems
+              // spread across 4 rows
+              tmp_out = gcn_mfma16x16x16_instr<scalar_t, 0, 0, 0>(
+                  Vlocaltmp.xy[i],
+                  shared_logits[vtoken_depth][offset2][lane16id][offset1],
+                  tmp_out);
             }
           }
         }
@@ -1725,8 +1570,7 @@ __device__ __forceinline__ _B16x8 from_floatx8(const floatx8& inp) {
 // clang-format off
 template <typename scalar_t, typename cache_t,
           vllm::Fp8KVCacheDataType KV_DTYPE, typename OUTT, int BLOCK_SIZE,
-          int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED, int GQA_RATIO,
-          MFMAType MFMA_TYPE>
+          int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED, int GQA_RATIO>
 __global__
 __launch_bounds__(NUM_THREADS, 3) void paged_attention_ll4mi_QKV_mfma16_kernel(
     const scalar_t* __restrict__ q,       // [num_seqs, num_heads, head_size]
@@ -2493,8 +2337,7 @@ __device__ __forceinline__ _B16x8 from_floatx8(const floatx8& inp) {
 // clang-format off
 template <typename scalar_t, typename cache_t,
           vllm::Fp8KVCacheDataType KV_DTYPE, typename OUTT, int BLOCK_SIZE,
-          int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED, int GQA_RATIO,
-          MFMAType MFMA_TYPE>
+          int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED, int GQA_RATIO>
 __global__
 __launch_bounds__(NUM_THREADS, 3) void paged_attention_ll4mi_QKV_mfma16_kernel(
     const scalar_t* __restrict__ q,       // [num_seqs, num_heads, head_size]
@@ -3126,7 +2969,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
 template <typename scalar_t, typename cache_t,
           vllm::Fp8KVCacheDataType KV_DTYPE, typename OUTT, int BLOCK_SIZE,
           int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED,
-          int GQA_RATIO, MFMAType MFMA_TYPE>
+          int GQA_RATIO>
 __global__
 __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma16_kernel(
     const scalar_t* __restrict__ q,         // [num_seqs, num_heads, head_size]
@@ -3198,7 +3041,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
 #define LAUNCH_CUSTOM_ATTENTION_MFMA16(GQA_RATIO)                              \
   paged_attention_ll4mi_QKV_mfma16_kernel<T, KVT, KV_DTYPE, OUTT, BLOCK_SIZE,  \
                                           HEAD_SIZE, NTHR, ALIBI_ENABLED,      \
-                                          GQA_RATIO, MFMA_TYPE>                \
+                                          GQA_RATIO>                           \
       <<<grid, block, 0, stream>>>(                                            \
           query_ptr, key_cache_ptr, value_cache_ptr, num_kv_heads, scale,      \
           block_tables_ptr, seq_lens_ptr, query_start_loc_ptr,                 \
@@ -3226,7 +3069,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
 
 template <typename T, typename KVT, vllm::Fp8KVCacheDataType KV_DTYPE,
           int BLOCK_SIZE, int HEAD_SIZE, typename OUTT, int PARTITION_SIZE_OLD,
-          bool ALIBI_ENABLED, MFMAType MFMA_TYPE>
+          bool ALIBI_ENABLED>
 void paged_attention_custom_launcher(
     torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits,
     torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache,
@@ -3382,7 +3225,7 @@ void paged_attention_custom_launcher(
 
 template <typename T, typename KVT, vllm::Fp8KVCacheDataType KV_DTYPE,
           int BLOCK_SIZE, int HEAD_SIZE, typename OUTT, int PARTITION_SIZE_OLD,
-          bool ALIBI_ENABLED, MFMAType MFMA_TYPE>
+          bool ALIBI_ENABLED>
 void paged_attention_custom_launcher_navi(
     torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits,
     torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache,
@@ -3554,77 +3397,74 @@ void paged_attention_custom_launcher_navi(
 }
 
 #define CALL_CUSTOM_LAUNCHER(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, OUTT,   \
-                             PSIZE, ALIBI_ENABLED, MFMA_TYPE)               \
+                             PSIZE, ALIBI_ENABLED)                          \
   if (!is_navi) {                                                           \
     paged_attention_custom_launcher<T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE,  \
-                                    OUTT, PSIZE, ALIBI_ENABLED, MFMA_TYPE>( \
+                                    OUTT, PSIZE, ALIBI_ENABLED>(            \
         out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache,  \
         num_kv_heads, scale, block_tables, seq_lens, query_start_loc,       \
         max_seq_len, alibi_slopes, k_scale, v_scale, fp8_out_scale);        \
   } else {                                                                  \
-    paged_attention_custom_launcher_navi<T, KVT, KV_DTYPE, BLK_SIZE,        \
-                                         HEAD_SIZE, OUTT, PSIZE,            \
-                                         ALIBI_ENABLED, MFMA_TYPE>(         \
+    paged_attention_custom_launcher_navi<                                   \
+        T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, OUTT, PSIZE, ALIBI_ENABLED>( \
         out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache,  \
         num_kv_heads, scale, block_tables, seq_lens, query_start_loc,       \
         max_seq_len, alibi_slopes, k_scale, v_scale);                       \
   }
 
 #define CALL_CUSTOM_LAUNCHER_ALIBI(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE,    \
-                                   OUTT, PSIZE, MFMA_TYPE)                   \
+                                   OUTT, PSIZE)                              \
   if (alibi_slopes) {                                                        \
     CALL_CUSTOM_LAUNCHER(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, OUTT, PSIZE, \
-                         true, MFMA_TYPE);                                   \
+                         true);                                              \
   } else {                                                                   \
     CALL_CUSTOM_LAUNCHER(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, OUTT, PSIZE, \
-                         false, MFMA_TYPE);                                  \
+                         false);                                             \
   }
 
 #if defined(__HIPCC__) && defined(__gfx90a__)
-  #define CALL_CUSTOM_LAUNCHER_OUT(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE,  \
-                                   MFMA_TYPE)                              \
+  #define CALL_CUSTOM_LAUNCHER_OUT(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE)  \
     if (fp8_out_scale) {                                                   \
       TORCH_CHECK(false, "fp8 out scale unsupported for gfx90a");          \
     } else {                                                               \
       CALL_CUSTOM_LAUNCHER_ALIBI(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, T, \
-                                 256, MFMA_TYPE);                          \
+                                 256);                                     \
     }
 #else
-  #define CALL_CUSTOM_LAUNCHER_OUT(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE,  \
-                                   MFMA_TYPE)                              \
+  #define CALL_CUSTOM_LAUNCHER_OUT(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE)  \
     if (fp8_out_scale) {                                                   \
       CALL_CUSTOM_LAUNCHER_ALIBI(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE,    \
-                                 uint8_t, 256, MFMA_TYPE);                 \
+                                 uint8_t, 256);                            \
     } else {                                                               \
       CALL_CUSTOM_LAUNCHER_ALIBI(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, T, \
-                                 256, MFMA_TYPE);                          \
+                                 256);                                     \
     }
 #endif
 
-#define CALL_CUSTOM_LAUNCHER_BLK(T, KVT, KV_DTYPE, HEAD_SIZE, MFMA_TYPE)    \
-  switch (block_size) {                                                     \
-    case 16:                                                                \
-      CALL_CUSTOM_LAUNCHER_OUT(T, KVT, KV_DTYPE, 16, HEAD_SIZE, MFMA_TYPE); \
-      break;                                                                \
-    case 32:                                                                \
-      CALL_CUSTOM_LAUNCHER_OUT(T, KVT, KV_DTYPE, 32, HEAD_SIZE, MFMA_TYPE); \
-      break;                                                                \
-    default:                                                                \
-      TORCH_CHECK(false, "Unsupported block size: ", block_size);           \
-      break;                                                                \
+#define CALL_CUSTOM_LAUNCHER_BLK(T, KVT, KV_DTYPE, HEAD_SIZE)     \
+  switch (block_size) {                                           \
+    case 16:                                                      \
+      CALL_CUSTOM_LAUNCHER_OUT(T, KVT, KV_DTYPE, 16, HEAD_SIZE);  \
+      break;                                                      \
+    case 32:                                                      \
+      CALL_CUSTOM_LAUNCHER_OUT(T, KVT, KV_DTYPE, 32, HEAD_SIZE);  \
+      break;                                                      \
+    default:                                                      \
+      TORCH_CHECK(false, "Unsupported block size: ", block_size); \
+      break;                                                      \
   }
 
-#define CALL_CUSTOM_LAUNCHER_BLK_HEAD(T, KVT, KV_DTYPE, MFMA_TYPE) \
-  switch (head_size) {                                             \
-    case 64:                                                       \
-      CALL_CUSTOM_LAUNCHER_BLK(T, KVT, KV_DTYPE, 64, MFMA_TYPE);   \
-      break;                                                       \
-    case 128:                                                      \
-      CALL_CUSTOM_LAUNCHER_BLK(T, KVT, KV_DTYPE, 128, MFMA_TYPE);  \
-      break;                                                       \
-    default:                                                       \
-      TORCH_CHECK(false, "Unsupported head size: ", head_size);    \
-      break;                                                       \
+#define CALL_CUSTOM_LAUNCHER_BLK_HEAD(T, KVT, KV_DTYPE)         \
+  switch (head_size) {                                          \
+    case 64:                                                    \
+      CALL_CUSTOM_LAUNCHER_BLK(T, KVT, KV_DTYPE, 64);           \
+      break;                                                    \
+    case 128:                                                   \
+      CALL_CUSTOM_LAUNCHER_BLK(T, KVT, KV_DTYPE, 128);          \
+      break;                                                    \
+    default:                                                    \
+      TORCH_CHECK(false, "Unsupported head size: ", head_size); \
+      break;                                                    \
   }
 
 bool is_navi_gpu() {
@@ -3663,43 +3503,28 @@ void paged_attention(
     const std::optional<torch::Tensor>& alibi_slopes,
     const std::string& kv_cache_dtype, torch::Tensor& k_scale,
     torch::Tensor& v_scale,
-    const std::optional<torch::Tensor>& fp8_out_scale,
-    const std::string& mfma_type) {
+    const std::optional<torch::Tensor>& fp8_out_scale) {
   // clang-format on
   bool is_navi = is_navi_gpu();
+
   const int head_size = query.size(2);
   if (kv_cache_dtype == "auto") {
     if (query.dtype() == at::ScalarType::Half) {
-      CALL_CUSTOM_LAUNCHER_BLK_HEAD(
-          _Float16, _Float16, vllm::Fp8KVCacheDataType::kAuto, MFMAType::F16);
+      CALL_CUSTOM_LAUNCHER_BLK_HEAD(_Float16, _Float16,
+                                    vllm::Fp8KVCacheDataType::kAuto);
     } else if (query.dtype() == at::ScalarType::BFloat16) {
       CALL_CUSTOM_LAUNCHER_BLK_HEAD(__hip_bfloat16, __hip_bfloat16,
-                                    vllm::Fp8KVCacheDataType::kAuto,
-                                    MFMAType::F16);
+                                    vllm::Fp8KVCacheDataType::kAuto);
     } else {
       TORCH_CHECK(false, "Unsupported data type: ", query.dtype());
     }
   } else if (kv_cache_dtype == "fp8" || kv_cache_dtype == "fp8_e4m3") {
     if (query.dtype() == at::ScalarType::Half) {
-      if (mfma_type == "fp8") {
-        CALL_CUSTOM_LAUNCHER_BLK_HEAD(_Float16, uint8_t,
-                                      vllm::Fp8KVCacheDataType::kFp8E4M3,
-                                      MFMAType::Fp8);
-      } else {
-        CALL_CUSTOM_LAUNCHER_BLK_HEAD(_Float16, uint8_t,
-                                      vllm::Fp8KVCacheDataType::kFp8E4M3,
-                                      MFMAType::F16);
-      }
+      CALL_CUSTOM_LAUNCHER_BLK_HEAD(_Float16, uint8_t,
+                                    vllm::Fp8KVCacheDataType::kFp8E4M3);
     } else if (query.dtype() == at::ScalarType::BFloat16) {
-      if (mfma_type == "fp8") {
-        CALL_CUSTOM_LAUNCHER_BLK_HEAD(__hip_bfloat16, uint8_t,
-                                      vllm::Fp8KVCacheDataType::kFp8E4M3,
-                                      MFMAType::Fp8);
-      } else {
-        CALL_CUSTOM_LAUNCHER_BLK_HEAD(__hip_bfloat16, uint8_t,
-                                      vllm::Fp8KVCacheDataType::kFp8E4M3,
-                                      MFMAType::F16);
-      }
+      CALL_CUSTOM_LAUNCHER_BLK_HEAD(__hip_bfloat16, uint8_t,
+                                    vllm::Fp8KVCacheDataType::kFp8E4M3);
     } else {
       TORCH_CHECK(false, "Unsupported data type: ", query.dtype());
     }

csrc/rocm/ops.h

@@ -5,14 +5,11 @@
 torch::Tensor LLMM1(at::Tensor& in_a, at::Tensor& in_b,
                     const int64_t rows_per_block);
 
-torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b,
-                       const std::optional<at::Tensor>& in_bias,
+torch::Tensor wvSplitK(at::Tensor& in_a, at::Tensor& in_b,
                        const int64_t CuCount);
 
-void wvSplitKQ(const at::Tensor& in_a, const at::Tensor& in_b,
-               const std::optional<at::Tensor>& in_bias, at::Tensor& out_c,
-               const at::Tensor& scale_a, const at::Tensor& scale_b,
-               const int64_t CuCount);
+void wvSplitKQ(at::Tensor& in_a, at::Tensor& in_b, at::Tensor& out_c,
+               at::Tensor& scale_a, at::Tensor& scale_b, const int64_t CuCount);
 
 void paged_attention(
     torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits,
@@ -22,5 +19,4 @@ void paged_attention(
     const std::optional<torch::Tensor>& query_start_loc, int64_t block_size,
     int64_t max_seq_len, const std::optional<torch::Tensor>& alibi_slopes,
     const std::string& kv_cache_dtype, torch::Tensor& k_scale,
-    torch::Tensor& v_scale, const std::optional<torch::Tensor>& fp8_out_scale,
-    const std::string& mfma_type);
+    torch::Tensor& v_scale, const std::optional<torch::Tensor>& fp8_out_scale);

csrc/rocm/skinny_gemms.cu

@@ -292,9 +292,8 @@ torch::Tensor LLMM1(at::Tensor& in_a, at::Tensor& in_b,
 template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
           int UNRL, int N>
 __global__ void __launch_bounds__(WvPrGrp* THRDS)
-    wvSplitK_hf_sml_(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, scalar_t* C,
+    wvSplitK_hf_sml_(const int K, const int M, const scalar_t* B,
+                     const scalar_t* __restrict__ A, scalar_t* C,
                      const int _WvPrGrp, const int CuCount) {
   constexpr int max_lds_len = LDS_SIZE / 2;
   #if defined(__HIP__MI3XX__)
@@ -485,14 +484,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
       if (threadIdx.x == 63) {
         for (int n = 0; n < N; n++) {
           for (int i = 0; i < YTILE; i++) {
-            if constexpr (std::is_same_v<scalar_t, half>) {
-              if (BIAS)
-                sum[n][i] += __half2float(BIAS[(m + i) % Bx + (n % By) * M]);
-            } else if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) {
-              if (BIAS)
-                sum[n][i] +=
-                    __bfloat162float(BIAS[(m + i) % Bx + (n % By) * M]);
-            }
+            // if (commitColumn[i]) C[m + i + n * M] = __float2half(sum[n][i]);
             C[m + i + n * M] = __float2s<scalar_t>(sum[n][i]);
           }
         }
@@ -537,9 +529,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
       if (threadIdx.x == 63) {
         for (int n = 0; n < N; n++) {
           for (int i = 0; i < YTILE; i++) {
-            if (BIAS)
-              sum4[n][i][0] +=
-                  __bfloat162float(BIAS[(m + i) % Bx + (n % By) * M]);
+            // if (commitColumn[i]) C[n + i + m * N] = __float2half(sum[n][i]);
             C[m + i + n * M] = __float2bfloat16(sum4[n][i][0]);
           }
         }
@@ -551,10 +541,8 @@ __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>
-__global__ void wvSplitK_hf_sml_(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, scalar_t* C,
+__global__ void wvSplitK_hf_sml_(const int K, const int M, const scalar_t* B,
+                                 const scalar_t* __restrict__ A, scalar_t* C,
                                  const int _WvPrGrp, const int CuCount) {
   UNREACHABLE_CODE
 }
@@ -565,9 +553,8 @@ __global__ void wvSplitK_hf_sml_(const int K, const int M, const int Bx,
 template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
           int UNRL, int N>
 __global__ void __launch_bounds__(WvPrGrp* THRDS)
-    wvSplitK_hf_(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, scalar_t* C,
+    wvSplitK_hf_(const int K, const int M, const scalar_t* B,
+                 const scalar_t* __restrict__ A, scalar_t* C,
                  const int _WvPrGrp, const int CuCount) {
   constexpr int max_lds_len = LDS_SIZE / 2;
   #if defined(__HIP__MI3XX__)
@@ -785,17 +772,8 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
       if (threadIdx.x == 63) {
         for (int n = 0; n < N; n++) {
           for (int i = 0; i < YTILE; i++) {
-            if (commitColumn[i]) {
-              if constexpr (std::is_same_v<scalar_t, half>) {
-                if (BIAS)
-                  sum[n][i] += __half2float(BIAS[(m + i) % Bx + (n % By) * M]);
-              } else if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) {
-                if (BIAS)
-                  sum[n][i] +=
-                      __bfloat162float(BIAS[(m + i) % Bx + (n % By) * M]);
-              }
+            if (commitColumn[i])
               C[m + i + n * M] = __float2s<scalar_t>(sum[n][i]);
-            }
           }
         }
       }
@@ -840,12 +818,8 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
       if (threadIdx.x == 63) {
         for (int n = 0; n < N; n++) {
           for (int i = 0; i < YTILE; i++) {
-            if (commitColumn[i]) {
-              if (BIAS)
-                sum4[n][i][0] +=
-                    __bfloat162float(BIAS[(m + i) % Bx + (n % By) * M]);
-              C[m + i + n * M] = __float2bfloat16(sum4[n][i][0]);
-            }
+            // if (commitColumn[i]) C[n + i + m * N] = __float2half(sum[n][i]);
+            C[m + i + n * M] = __float2bfloat16(sum4[n][i][0]);
           }
         }
       }
@@ -868,10 +842,8 @@ __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>
-__global__ void wvSplitK_hf_(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, scalar_t* C,
+__global__ void wvSplitK_hf_(const int K, const int M, const scalar_t* B,
+                             const scalar_t* __restrict__ A, scalar_t* C,
                              const int _WvPrGrp, const int CuCount) {
   UNREACHABLE_CODE
 }
@@ -882,9 +854,8 @@ __global__ void wvSplitK_hf_(const int K, const int M, const int Bx,
 template <typename scalar_t, int THRDS, int YTILE, int WvPrGrp, int A_CHUNK,
           int UNRL, int N>
 __global__ void __launch_bounds__(WvPrGrp* THRDS)
-    wvSplitK_hf_big_(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, scalar_t* C,
+    wvSplitK_hf_big_(const int K, const int M, const scalar_t* B,
+                     const scalar_t* __restrict__ A, scalar_t* C,
                      const int _WvPrGrp, const int CuCount) {
   constexpr int max_lds_len = LDS_SIZE / 2;
   #if defined(__HIP__MI3XX__)
@@ -1153,17 +1124,8 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
       if (threadIdx.x == 63) {
         for (int n = 0; n < N; n++) {
           for (int i = 0; i < YTILE; i++) {
-            if (commitColumn[i]) {
-              if constexpr (std::is_same_v<scalar_t, half>) {
-                if (BIAS)
-                  sum[n][i] += __half2float(BIAS[(m + i) % Bx + (n % By) * M]);
-              } else if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) {
-                if (BIAS)
-                  sum[n][i] +=
-                      __bfloat162float(BIAS[(m + i) % Bx + (n % By) * M]);
-              }
+            if (commitColumn[i])
               C[m + i + n * M] = __float2s<scalar_t>(sum[n][i]);
-            }
           }
         }
       }
@@ -1204,12 +1166,8 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
       if (threadIdx.x == 63) {
         for (int n = 0; n < N; n++) {
           for (int i = 0; i < YTILE; i++) {
-            if (commitColumn[i]) {
-              if (BIAS)
-                sum4[n][i][0] +=
-                    __bfloat162float(BIAS[(m + i) % Bx + (n % By) * M]);
-              C[m + i + n * M] = __float2bfloat16(sum4[n][i][0]);
-            }
+            // if (commitColumn[i]) C[n + i + m * N] = __float2half(sum[n][i]);
+            C[m + i + n * M] = __float2bfloat16(sum4[n][i][0]);
           }
         }
       }
@@ -1232,10 +1190,8 @@ __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>
-__global__ void wvSplitK_hf_big_(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, scalar_t* C,
+__global__ void wvSplitK_hf_big_(const int K, const int M, const scalar_t* B,
+                                 const scalar_t* __restrict__ A, scalar_t* C,
                                  const int _WvPrGrp, const int CuCount) {
   UNREACHABLE_CODE
 }
@@ -1270,20 +1226,11 @@ int mindiv(int N, int div1, int div2) {
   return rtn;
 }
 
-torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b,
-                       const std::optional<at::Tensor>& in_bias,
+torch::Tensor wvSplitK(at::Tensor& in_a, at::Tensor& in_b,
                        const int64_t CuCount) {
   auto M_in = in_a.size(0);
   auto K_in = in_a.size(1);
   auto N_in = in_b.size(0);
-  auto Bx_in =
-      (in_bias.has_value() && in_bias->numel() > 0)
-          ? (in_bias->sizes().size() == 2) ? in_bias->size(1) : in_bias->size(0)
-          : 1;
-  auto By_in = (in_bias.has_value() && in_bias->numel() > 0 &&
-                in_bias->sizes().size() == 2)
-                   ? in_bias->size(0)
-                   : 1;
 
   TORCH_CHECK(in_a.dtype() == in_b.dtype());
   TORCH_CHECK(K_in % 8 == 0, "k % 8 == 0");
@@ -1307,18 +1254,18 @@ torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b,
     if ((K_in * N_in <= max_lds_len) && (M_in % _YTILEs == 0)) {              \
       int __wvPrGrp = mindiv(M_in, CuCount * _YTILEs, _WvPrGrp);              \
       wvSplitK_hf_sml_<fptype, 64, _YTILEs, _WvPrGrp, 8, _UNRLs, _N>          \
-          <<<grid, block, 0, stream>>>(K_in, M_in, Bx_in, By_in, af4, bf4,    \
-                                       biasf4, c, __wvPrGrp, CuCount);        \
+          <<<grid, block, 0, stream>>>(K_in, M_in, af4, bf4, c, __wvPrGrp,    \
+                                       CuCount);                              \
     } else if (K_in * N_in <= max_lds_len * 1.2) {                            \
       int __wvPrGrp = mindiv(M_in, CuCount * _YTILEm, _WvPrGrp);              \
       wvSplitK_hf_<fptype, 64, _YTILEm, _WvPrGrp, 8, _UNRLm, _N>              \
-          <<<grid, block, 0, stream>>>(K_in, M_in, Bx_in, By_in, af4, bf4,    \
-                                       biasf4, c, __wvPrGrp, CuCount);        \
+          <<<grid, block, 0, stream>>>(K_in, M_in, af4, bf4, c, __wvPrGrp,    \
+                                       CuCount);                              \
     } else {                                                                  \
       int __wvPrGrp = mindiv(M_in, CuCount * _YTILEb, _WvPrGrp);              \
       wvSplitK_hf_big_<fptype, 64, _YTILEb, _WvPrGrp, 8, _UNRLb, _N>          \
-          <<<grid, block, 0, stream>>>(K_in, M_in, Bx_in, By_in, af4, bf4,    \
-                                       biasf4, c, __wvPrGrp, CuCount);        \
+          <<<grid, block, 0, stream>>>(K_in, M_in, af4, bf4, c, __wvPrGrp,    \
+                                       CuCount);                              \
     }                                                                         \
   }
 
@@ -1326,10 +1273,6 @@ torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b,
     using fptype = typename scalar<scalar_t>::type;
     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());
     switch (N_in) {
       case 1:
@@ -1357,9 +1300,8 @@ torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b,
 template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
           int A_CHUNK, int UNRL, int N>
 __global__ void __launch_bounds__(WvPrGrp* THRDS)
-    wvSplitKQ_hf_sml_(const int K, const int Kp, const int M, const int Bx,
-                      const int By, const fp8_t* B, const fp8_t* __restrict__ A,
-                      const scalar_t* __restrict__ BIAS, scalar_t* C,
+    wvSplitKQ_hf_sml_(const int K, const int Kp, const int M, const fp8_t* B,
+                      const fp8_t* __restrict__ A, scalar_t* C,
                       const float* __restrict__ s_A,
                       const float* __restrict__ s_B, const int _WvPrGrp,
                       const int CuCount) {
@@ -1511,17 +1453,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
     if (threadIdx.x == 0) {
       for (int n = 0; n < N; n++) {
         for (int y = 0; y < YTILE; y++) {
-          if (y + m >= M) break;  // To avoid mem access fault.
-          sum[n][y][0] *= sA * sB;
-          if constexpr (std::is_same_v<scalar_t, half>) {
-            if (BIAS)
-              sum[n][y][0] += __half2float(BIAS[(m + y) % Bx + (n % By) * M]);
-          } else if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) {
-            if (BIAS)
-              sum[n][y][0] +=
-                  __bfloat162float(BIAS[(m + y) % Bx + (n % By) * M]);
-          }
-          C[m + y + n * M] = __float2s<scalar_t>(sum[n][y][0]);  // * sA * sB);
+          C[m + y + n * M] = __float2s<scalar_t>(sum[n][y][0] * sA * sB);
         }
       }
     }
@@ -1533,9 +1465,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
 template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
           int A_CHUNK, int UNRL, int N>
 __global__ void wvSplitKQ_hf_sml_(const int K, const int Kp, const int M,
-                                  const int Bx, const int By, const fp8_t* B,
-                                  const fp8_t* __restrict__ A,
-                                  const scalar_t* __restrict__ BIAS,
+                                  const fp8_t* B, const fp8_t* __restrict__ A,
                                   scalar_t* C, const float* __restrict__ s_A,
                                   const float* __restrict__ s_B,
                                   const int _WvPrGrp, const int CuCount) {
@@ -1547,9 +1477,8 @@ __global__ void wvSplitKQ_hf_sml_(const int K, const int Kp, const int M,
 template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
           int A_CHUNK, int UNRL, int N>
 __global__ void __launch_bounds__(WvPrGrp* THRDS)
-    wvSplitKQ_hf_(const int K, const int Kp, const int M, const int Bx,
-                  const int By, const fp8_t* B, const fp8_t* __restrict__ A,
-                  const scalar_t* __restrict__ BIAS, scalar_t* C,
+    wvSplitKQ_hf_(const int K, const int Kp, const int M, const fp8_t* B,
+                  const fp8_t* __restrict__ A, scalar_t* C,
                   const float* __restrict__ s_A, const float* __restrict__ s_B,
                   const int _WvPrGrp, const int CuCount) {
   constexpr int max_lds_len = LDS_SIZE;
@@ -1697,16 +1626,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
       for (int n = 0; n < N; n++) {
         for (int y = 0; y < YTILE; y++) {
           if (y + m >= M) break;  // To avoid mem access fault.
-          sum[n][y][0] *= sA * sB;
-          if constexpr (std::is_same_v<scalar_t, half>) {
-            if (BIAS)
-              sum[n][y][0] += __half2float(BIAS[(m + y) % Bx + (n % By) * M]);
-          } else if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) {
-            if (BIAS)
-              sum[n][y][0] +=
-                  __bfloat162float(BIAS[(m + y) % Bx + (n % By) * M]);
-          }
-          C[m + y + n * M] = __float2s<scalar_t>(sum[n][y][0]);
+          C[m + y + n * M] = __float2s<scalar_t>(sum[n][y][0] * sA * sB);
         }
       }
     }
@@ -1718,19 +1638,16 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
 template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
           int A_CHUNK, int UNRL, int N>
 __global__ void wvSplitKQ_hf_(const int K, const int Kp, const int M,
-                              const int Bx, const int By, const fp8_t* B,
-                              const fp8_t* __restrict__ A,
-                              const scalar_t* __restrict__ BIAS, scalar_t* C,
-                              const float* __restrict__ s_A,
+                              const fp8_t* B, const fp8_t* __restrict__ A,
+                              scalar_t* C, const float* __restrict__ s_A,
                               const float* __restrict__ s_B, const int _WvPrGrp,
                               const int CuCount) {
   UNREACHABLE_CODE
 }
 #endif  // defined(__HIP__MI3XX__) TODO: Add NAVI support
 
-void wvSplitKQ(const at::Tensor& in_a, const at::Tensor& in_b,
-               const std::optional<at::Tensor>& in_bias, at::Tensor& out_c,
-               const at::Tensor& scale_a, const at::Tensor& scale_b,
+void wvSplitKQ(at::Tensor& in_a, at::Tensor& in_b, at::Tensor& out_c,
+               at::Tensor& scale_a, at::Tensor& scale_b,
                const int64_t CuCount) {
   static c10::ScalarType kFp8Type = is_fp8_ocp()
                                         ? c10::ScalarType::Float8_e4m3fn
@@ -1739,15 +1656,6 @@ void wvSplitKQ(const at::Tensor& in_a, const at::Tensor& in_b,
   auto K_in = in_a.size(1);
   auto N_in = in_b.size(0);
   auto Kp_in = in_a.stride(0);
-  auto Bx_in =
-      (in_bias.has_value() && in_bias->numel() > 0)
-          ? (in_bias->sizes().size() == 2) ? in_bias->size(1) : in_bias->size(0)
-          : 1;
-  auto By_in = (in_bias.has_value() && in_bias->numel() > 0 &&
-                in_bias->sizes().size() == 2)
-                   ? in_bias->size(0)
-                   : 1;
-
   TORCH_CHECK(K_in % 16 == 0, "k % 16 == 0");
   TORCH_CHECK(in_a.dtype() == in_b.dtype() && in_a.dtype() == kFp8Type);
   TORCH_CHECK(out_c.dtype() == torch::kFloat16 ||
@@ -1765,15 +1673,13 @@ void wvSplitKQ(const at::Tensor& in_a, const at::Tensor& in_b,
     if ((K_in * N_in <= max_lds_len) && (M_in % _YTILEs == 0)) {               \
       int __wvPrGrp = mindiv(M_in, CuCount * _YTILEs, _WvPrGrp);               \
       wvSplitKQ_hf_sml_<fptype, fp8_t, 64, _YTILEs, _WvPrGrp, 16, _UNRLs, _N>  \
-          <<<grid, block, 0, stream>>>(K_in, Kp_in, M_in, Bx_in, By_in, a_ptr, \
-                                       b_ptr, bias_ptr, c_ptr, s_a, s_b,       \
-                                       __wvPrGrp, CuCount);                    \
+          <<<grid, block, 0, stream>>>(K_in, Kp_in, M_in, a_ptr, b_ptr, c_ptr, \
+                                       s_a, s_b, __wvPrGrp, CuCount);          \
     } else {                                                                   \
       int __wvPrGrp = mindiv(M_in, CuCount * _YTILEm, _WvPrGrp);               \
       wvSplitKQ_hf_<fptype, fp8_t, 64, _YTILEm, _WvPrGrp, 16, _UNRLm, _N>      \
-          <<<grid, block, 0, stream>>>(K_in, Kp_in, M_in, Bx_in, By_in, a_ptr, \
-                                       b_ptr, bias_ptr, c_ptr, s_a, s_b,       \
-                                       __wvPrGrp, CuCount);                    \
+          <<<grid, block, 0, stream>>>(K_in, Kp_in, M_in, a_ptr, b_ptr, c_ptr, \
+                                       s_a, s_b, __wvPrGrp, CuCount);          \
     }                                                                          \
   }
 
@@ -1785,9 +1691,6 @@ void wvSplitKQ(const at::Tensor& in_a, const at::Tensor& in_b,
     VLLM_DISPATCH_FP8_TYPES(in_a.scalar_type(), "wvSplitKQ", [&] {
       auto a_ptr = in_a.data_ptr<fp8_t>();
       auto b_ptr = in_b.data_ptr<fp8_t>();
-      auto bias_ptr = (in_bias.has_value() && in_bias->numel() > 0)
-                          ? reinterpret_cast<fptype*>(in_bias->data_ptr())
-                          : nullptr;
       switch (N_in) {
         case 1:
           WVSPLITKQ(16, 2, 2, 2, 2, 2, 2, 1)

csrc/rocm/torch_bindings.cpp

@@ -22,14 +22,13 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) {
 
   // Custom gemm op for skinny matrix-matrix multiplication
   rocm_ops.def(
-      "wvSplitK(Tensor in_a, Tensor in_b, Tensor? in_bias, int CuCount) -> "
+      "wvSplitK(Tensor in_a, Tensor in_b, int CuCount) -> "
       "Tensor");
   rocm_ops.impl("wvSplitK", torch::kCUDA, &wvSplitK);
 
   // wvSplitK for fp8
   rocm_ops.def(
-      "wvSplitKQ(Tensor in_a, Tensor in_b, Tensor? in_bias, Tensor! out_c, "
-      "Tensor scale_a, "
+      "wvSplitKQ(Tensor in_a, Tensor in_b, Tensor! out_c, Tensor scale_a, "
       "          Tensor scale_b, int CuCount) -> ()");
   rocm_ops.impl("wvSplitKQ", torch::kCUDA, &wvSplitKQ);
 
@@ -49,8 +48,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) {
       "                Tensor? alibi_slopes,"
       "                str kv_cache_dtype,"
       "                Tensor k_scale, Tensor v_scale,"
-      "                Tensor? fp8_out_scale,"
-      "                str mfma_type) -> ()");
+      "                Tensor? fp8_out_scale) -> ()");
   rocm_ops.impl("paged_attention", torch::kCUDA, &paged_attention);
 }
 

csrc/torch_bindings.cpp

@@ -32,13 +32,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   #define stride_tag
 #endif
 
-  ops.def(
-      "silu_mul_fp8_quant_deep_gemm_cuda(Tensor input, Tensor counts, Tensor! "
-      "y_q, Tensor! y_s, int group_size, "
-      "bool use_ue8m0, int num_parallel_tokens) -> ()");
-  ops.impl("silu_mul_fp8_quant_deep_gemm_cuda", torch::kCUDA,
-           &silu_mul_fp8_quant_deep_gemm_cuda);
-
   ops.def("weak_ref_tensor(Tensor input) -> Tensor");
   ops.impl("weak_ref_tensor", torch::kCUDA, &weak_ref_tensor);
 
@@ -221,6 +214,16 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       "                 Tensor cos_sin_cache, bool is_neox) -> ()");
   ops.impl("rotary_embedding", torch::kCUDA, &rotary_embedding);
 
+  // Apply GPT-NeoX or GPT-J style rotary embedding to query and key
+  // (supports multiple loras).
+  ops.def(
+      "batched_rotary_embedding(Tensor positions, Tensor! query,"
+      "                         Tensor!? key, int head_size,"
+      "                         Tensor cos_sin_cache, bool is_neox,"
+      "                         int rot_dim,"
+      "                         Tensor cos_sin_cache_offsets) -> ()");
+  ops.impl("batched_rotary_embedding", torch::kCUDA, &batched_rotary_embedding);
+
   // Quantization ops
 #ifndef USE_ROCM
   // Quantized GEMM for AWQ.
@@ -510,6 +513,13 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.def("cutlass_sparse_compress(Tensor a) -> Tensor[]");
   ops.impl("cutlass_sparse_compress", &cutlass_sparse_compress);
 
+  // CUTLASS MLA decode
+  ops.def(
+      "cutlass_mla_decode(Tensor! out, Tensor q_nope, Tensor q_pe,"
+      "                   Tensor kv_c_and_k_pe_cache, Tensor seq_lens,"
+      "                   Tensor page_table, float scale) -> ()");
+  ops.impl("cutlass_mla_decode", torch::kCUDA, &cutlass_mla_decode);
+
   // SM100 CUTLASS MLA decode
   ops.def(
       "sm100_cutlass_mla_decode(Tensor! out, Tensor! lse, Tensor q_nope,"
@@ -606,9 +616,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       "int pad_slot_id) -> ()");
   ops.impl("selective_scan_fwd", torch::kCUDA, &selective_scan_fwd);
 
-  // Hadamard transforms
-  ops.def("hadacore_transform(Tensor! x, bool inplace) -> Tensor");
-
 #ifndef USE_ROCM
   // Compute per-token-group FP8 quantized tensor and scaling factor.
   ops.def(

docker/Dockerfile

@@ -283,10 +283,6 @@ WORKDIR /vllm-workspace
 ENV DEBIAN_FRONTEND=noninteractive
 ARG TARGETPLATFORM
 
-ARG GDRCOPY_CUDA_VERSION=12.8
-# Keep in line with FINAL_BASE_IMAGE
-ARG GDRCOPY_OS_VERSION=Ubuntu22_04
-
 SHELL ["/bin/bash", "-c"]
 
 ARG DEADSNAKES_MIRROR_URL
@@ -381,7 +377,7 @@ RUN --mount=type=bind,from=build,src=/workspace/dist,target=/vllm-workspace/dist
 # Install FlashInfer from source
 ARG FLASHINFER_GIT_REPO="https://github.com/flashinfer-ai/flashinfer.git"
 # Keep this in sync with "flashinfer" extra in setup.py
-ARG FLASHINFER_GIT_REF="v0.3.1"
+ARG FLASHINFER_GIT_REF="v0.3.0"
 # Flag to control whether to compile FlashInfer AOT kernels
 # Set to "true" to enable AOT compilation:
 # docker build --build-arg FLASHINFER_AOT_COMPILE=true ...
@@ -445,21 +441,13 @@ COPY tools/install_deepgemm.sh /tmp/install_deepgemm.sh
 RUN --mount=type=cache,target=/root/.cache/uv \
     VLLM_DOCKER_BUILD_CONTEXT=1 /tmp/install_deepgemm.sh --cuda-version "${CUDA_VERSION}" ${DEEPGEMM_GIT_REF:+--ref "$DEEPGEMM_GIT_REF"} 
 
-COPY tools/install_gdrcopy.sh install_gdrcopy.sh
-RUN set -eux; \
-    case "${TARGETPLATFORM}" in \
-      linux/arm64) UUARCH="aarch64" ;; \
-      linux/amd64) UUARCH="x64" ;; \
-      *) echo "Unsupported TARGETPLATFORM: ${TARGETPLATFORM}" >&2; exit 1 ;; \
-    esac; \
-    ./install_gdrcopy.sh "${GDRCOPY_OS_VERSION}" "${GDRCOPY_CUDA_VERSION}" "${UUARCH}"; \
-    rm ./install_gdrcopy.sh
-
-# Install EP kernels(pplx-kernels and DeepEP)
+# Install EP kernels(pplx-kernels and DeepEP), NixL
 COPY tools/ep_kernels/install_python_libraries.sh install_python_libraries.sh
+COPY tools/install_nixl.sh install_nixl.sh
 ENV CUDA_HOME=/usr/local/cuda
 RUN export TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST:-9.0a+PTX}" \
-    && bash install_python_libraries.sh
+    && bash install_python_libraries.sh \
+    && bash install_nixl.sh --force
 
 #################### vLLM installation IMAGE ####################
 

docker/Dockerfile.cpu

@@ -114,6 +114,9 @@ WORKDIR /workspace/vllm
 RUN --mount=type=bind,src=requirements/test.in,target=requirements/test.in \
     cp requirements/test.in requirements/cpu-test.in && \
     sed -i '/mamba_ssm/d' requirements/cpu-test.in && \
+    sed -i 's/^torch==.*/torch==2.6.0/g' requirements/cpu-test.in && \
+    sed -i 's/torchaudio.*/torchaudio/g' requirements/cpu-test.in && \
+    sed -i 's/torchvision.*/torchvision/g' requirements/cpu-test.in && \
     uv pip compile requirements/cpu-test.in -o requirements/cpu-test.txt --index-strategy unsafe-best-match --torch-backend cpu
 
 RUN --mount=type=cache,target=/root/.cache/uv \

docker/Dockerfile.nightly_torch

@@ -246,7 +246,7 @@ RUN pip install setuptools==75.6.0 packaging==23.2 ninja==1.11.1.3 build==1.2.2.
 
 
 # build flashinfer for torch nightly from source around 10 mins
-# release version: v0.3.1
+# release version: v0.2.2.post1
 # todo(elainewy): cache flashinfer build result for faster build
 ENV CCACHE_DIR=/root/.cache/ccache
 RUN --mount=type=cache,target=/root/.cache/ccache \
@@ -254,7 +254,7 @@ RUN --mount=type=cache,target=/root/.cache/ccache \
     echo "git clone flashinfer..." \
     && git clone --recursive https://github.com/flashinfer-ai/flashinfer.git \
     && cd flashinfer \
-    && git checkout v0.3.1 \
+    && git checkout v0.2.2.post1 \
     && git submodule update --init --recursive \
     && echo "finish git clone flashinfer..." \
     && rm -rf build \

docker/Dockerfile.rocm

@@ -29,10 +29,7 @@ ARG VLLM_BRANCH="main"
 ONBUILD RUN git clone ${VLLM_REPO} \
 	    && cd vllm \
 	    && git fetch -v --prune -- origin ${VLLM_BRANCH} \
-	    && git checkout FETCH_HEAD \
-        && if [ ${VLLM_REPO} != "https://github.com/vllm-project/vllm.git" ] ; then \
-               git remote add upstream "https://github.com/vllm-project/vllm.git" \
-               && git fetch upstream ; fi
+	    && git checkout FETCH_HEAD
 FROM fetch_vllm_${REMOTE_VLLM} AS fetch_vllm
 
 # -----------------------

docker/Dockerfile.rocm_base

@@ -1,23 +1,25 @@
-ARG BASE_IMAGE=rocm/dev-ubuntu-22.04:7.0-complete
-ARG TRITON_BRANCH="f9e5bf54"
-ARG TRITON_REPO="https://github.com/ROCm/triton.git"
-ARG PYTORCH_BRANCH="b2fb6885"
-ARG PYTORCH_VISION_BRANCH="v0.23.0"
+ARG BASE_IMAGE=rocm/dev-ubuntu-22.04:6.4.1-complete
+ARG HIPBLASLT_BRANCH="aa0bda7b"
+ARG HIPBLAS_COMMON_BRANCH="9b80ba8e"
+ARG LEGACY_HIPBLASLT_OPTION=
+ARG TRITON_BRANCH="e5be006"
+ARG TRITON_REPO="https://github.com/triton-lang/triton.git"
+ARG PYTORCH_BRANCH="f717b2af"
+ARG PYTORCH_VISION_BRANCH="v0.21.0"
 ARG PYTORCH_REPO="https://github.com/ROCm/pytorch.git"
 ARG PYTORCH_VISION_REPO="https://github.com/pytorch/vision.git"
-ARG FA_BRANCH="0e60e394"
+ARG FA_BRANCH="1a7f4dfa"
 ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git"
-ARG AITER_BRANCH="2ab9f4cd"
+ARG AITER_BRANCH="4822e675"
 ARG AITER_REPO="https://github.com/ROCm/aiter.git"
 
 FROM ${BASE_IMAGE} AS base
 
-ENV PATH=/opt/rocm/llvm/bin:/opt/rocm/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin
+ENV PATH=/opt/rocm/llvm/bin:$PATH
 ENV ROCM_PATH=/opt/rocm
 ENV LD_LIBRARY_PATH=/opt/rocm/lib:/usr/local/lib:
-ARG PYTORCH_ROCM_ARCH=gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1200;gfx1201
+ARG PYTORCH_ROCM_ARCH=gfx90a;gfx942;gfx1100;gfx1101;gfx1200;gfx1201
 ENV PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH}
-ENV AITER_ROCM_ARCH=gfx942;gfx950
 
 ARG PYTHON_VERSION=3.12
 
@@ -43,6 +45,29 @@ RUN apt-get update -y \
 
 RUN pip install -U packaging 'cmake<4' ninja wheel 'setuptools<80' pybind11 Cython
 
+FROM base AS build_hipblaslt
+ARG HIPBLASLT_BRANCH
+ARG HIPBLAS_COMMON_BRANCH
+# Set to "--legacy_hipblas_direct" for ROCm<=6.2
+ARG LEGACY_HIPBLASLT_OPTION
+RUN git clone https://github.com/ROCm/hipBLAS-common.git
+RUN apt-get remove -y hipblaslt && apt-get autoremove -y && apt-get autoclean -y
+RUN cd hipBLAS-common \
+    && git checkout ${HIPBLAS_COMMON_BRANCH} \
+    && mkdir build \
+    && cd build \
+    && cmake .. \
+    && make package \
+    && dpkg -i ./*.deb
+RUN git clone https://github.com/ROCm/hipBLASLt
+RUN cd hipBLASLt \
+    && git checkout ${HIPBLASLT_BRANCH} \
+    && apt-get install -y llvm-dev \
+    && ./install.sh -dc --architecture ${PYTORCH_ROCM_ARCH} ${LEGACY_HIPBLASLT_OPTION} \
+    && cd build/release \
+    && make package
+RUN mkdir -p /app/install && cp /app/hipBLASLt/build/release/*.deb /app/hipBLAS-common/build/*.deb /app/install
+
 FROM base AS build_triton
 ARG TRITON_BRANCH
 ARG TRITON_REPO
@@ -65,6 +90,8 @@ ARG PYTORCH_BRANCH
 ARG PYTORCH_VISION_BRANCH
 ARG PYTORCH_REPO
 ARG PYTORCH_VISION_REPO
+ARG FA_BRANCH
+ARG FA_REPO
 RUN git clone ${PYTORCH_REPO} pytorch
 RUN cd pytorch && git checkout ${PYTORCH_BRANCH} && \
     pip install -r requirements.txt && git submodule update --init --recursive \
@@ -75,20 +102,14 @@ RUN git clone ${PYTORCH_VISION_REPO} vision
 RUN cd vision && git checkout ${PYTORCH_VISION_BRANCH} \
     && python3 setup.py bdist_wheel --dist-dir=dist \
     && pip install dist/*.whl
-RUN mkdir -p /app/install && cp /app/pytorch/dist/*.whl /app/install \
-    && cp /app/vision/dist/*.whl /app/install
-
-FROM base AS build_fa
-ARG FA_BRANCH
-ARG FA_REPO
-RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \
-    pip install /install/*.whl
 RUN git clone ${FA_REPO}
 RUN cd flash-attention \
     && git checkout ${FA_BRANCH} \
     && git submodule update --init \
     && GPU_ARCHS=$(echo ${PYTORCH_ROCM_ARCH} | sed -e 's/;gfx1[0-9]\{3\}//g') python3 setup.py bdist_wheel --dist-dir=dist
-RUN mkdir -p /app/install && cp /app/flash-attention/dist/*.whl /app/install
+RUN mkdir -p /app/install && cp /app/pytorch/dist/*.whl /app/install \
+    && cp /app/vision/dist/*.whl /app/install \
+    && cp /app/flash-attention/dist/*.whl /app/install
 
 FROM base AS build_aiter
 ARG AITER_BRANCH
@@ -100,15 +121,15 @@ RUN cd aiter \
     && git checkout ${AITER_BRANCH} \
     && git submodule update --init --recursive \
     && pip install -r requirements.txt
-RUN pip install pyyaml && cd aiter && PREBUILD_KERNELS=1 GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist && ls /app/aiter/dist/*.whl
+RUN pip install pyyaml && cd aiter && PREBUILD_KERNELS=1 GPU_ARCHS=gfx942 python3 setup.py bdist_wheel --dist-dir=dist && ls /app/aiter/dist/*.whl
 RUN mkdir -p /app/install && cp /app/aiter/dist/*.whl /app/install
 
 FROM base AS debs
 RUN mkdir /app/debs
+RUN --mount=type=bind,from=build_hipblaslt,src=/app/install/,target=/install \
+    cp /install/*.deb /app/debs
 RUN --mount=type=bind,from=build_triton,src=/app/install/,target=/install \
     cp /install/*.whl /app/debs
-RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \
-    cp /install/*.whl /app/debs
 RUN --mount=type=bind,from=build_amdsmi,src=/app/install/,target=/install \
     cp /install/*.whl /app/debs
 RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \
@@ -117,10 +138,24 @@ RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \
     cp /install/*.whl /app/debs
 
 FROM base AS final
-RUN --mount=type=bind,from=debs,src=/app/debs,target=/install \
+RUN --mount=type=bind,from=build_hipblaslt,src=/app/install/,target=/install \
+    dpkg -i /install/*deb \
+    && perl -p -i -e 's/, hipblas-common-dev \([^)]*?\), /, /g' /var/lib/dpkg/status \
+    && perl -p -i -e 's/, hipblaslt-dev \([^)]*?\), /, /g' /var/lib/dpkg/status \
+    && perl -p -i -e 's/, hipblaslt \([^)]*?\), /, /g' /var/lib/dpkg/status
+RUN --mount=type=bind,from=build_triton,src=/app/install/,target=/install \
+    pip install /install/*.whl
+RUN --mount=type=bind,from=build_amdsmi,src=/app/install/,target=/install \
+    pip install /install/*.whl
+RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \
+    pip install /install/*.whl
+RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \
     pip install /install/*.whl
 
 ARG BASE_IMAGE
+ARG HIPBLAS_COMMON_BRANCH
+ARG HIPBLASLT_BRANCH
+ARG LEGACY_HIPBLASLT_OPTION
 ARG TRITON_BRANCH
 ARG TRITON_REPO
 ARG PYTORCH_BRANCH
@@ -132,6 +167,9 @@ ARG FA_REPO
 ARG AITER_BRANCH
 ARG AITER_REPO
 RUN echo "BASE_IMAGE: ${BASE_IMAGE}" > /app/versions.txt \
+    && echo "HIPBLAS_COMMON_BRANCH: ${HIPBLAS_COMMON_BRANCH}" >> /app/versions.txt \
+    && echo "HIPBLASLT_BRANCH: ${HIPBLASLT_BRANCH}" >> /app/versions.txt \
+    && echo "LEGACY_HIPBLASLT_OPTION: ${LEGACY_HIPBLASLT_OPTION}" >> /app/versions.txt \
     && echo "TRITON_BRANCH: ${TRITON_BRANCH}" >> /app/versions.txt \
     && echo "TRITON_REPO: ${TRITON_REPO}" >> /app/versions.txt \
     && echo "PYTORCH_BRANCH: ${PYTORCH_BRANCH}" >> /app/versions.txt \
@@ -139,6 +177,5 @@ RUN echo "BASE_IMAGE: ${BASE_IMAGE}" > /app/versions.txt \
     && echo "PYTORCH_REPO: ${PYTORCH_REPO}" >> /app/versions.txt \
     && echo "PYTORCH_VISION_REPO: ${PYTORCH_VISION_REPO}" >> /app/versions.txt \
     && echo "FA_BRANCH: ${FA_BRANCH}" >> /app/versions.txt \
-    && echo "FA_REPO: ${FA_REPO}" >> /app/versions.txt \
     && echo "AITER_BRANCH: ${AITER_BRANCH}" >> /app/versions.txt \
     && echo "AITER_REPO: ${AITER_REPO}" >> /app/versions.txt

docs/README.md

@@ -56,7 +56,7 @@ vLLM is flexible and easy to use with:
 - Tensor, pipeline, data and expert parallelism support for distributed inference
 - Streaming outputs
 - OpenAI-compatible API server
-- Support for NVIDIA GPUs, AMD CPUs and GPUs, Intel CPUs and GPUs, PowerPC CPUs, and TPU. Additionally, support for diverse hardware plugins such as Intel Gaudi, IBM Spyre and Huawei Ascend.
+- Support NVIDIA GPUs, AMD CPUs and GPUs, Intel CPUs, Gaudi® accelerators and GPUs, IBM Power CPUs, TPU, and AWS Trainium and Inferentia Accelerators.
 - Prefix caching support
 - Multi-LoRA support
 

docs/api/README.md

@@ -14,7 +14,7 @@ API documentation for vLLM's configuration classes.
 - [vllm.config.LoRAConfig][]
 - [vllm.config.MultiModalConfig][]
 - [vllm.config.PoolerConfig][]
-- [vllm.config.StructuredOutputsConfig][]
+- [vllm.config.DecodingConfig][]
 - [vllm.config.ObservabilityConfig][]
 - [vllm.config.KVTransferConfig][]
 - [vllm.config.CompilationConfig][]
@@ -46,6 +46,7 @@ Engine classes for offline and online inference.
 Inference parameters for vLLM APIs.
 
 [](){ #sampling-params }
+[](){ #pooling-params }
 
 - [vllm.SamplingParams][]
 - [vllm.PoolingParams][]

docs/configuration/optimization.md

@@ -139,9 +139,9 @@ there is relatively little gain from TP. On the other hand, TP incurs significan
 overhead because of all-reduce being performed after every layer.
 
 Given this, it may be advantageous to instead shard the batched input data using TP, essentially
-performing batch-level DP. This has been shown to improve the throughput and TTFT by around 10% for
+performing batch-level DP. This has been shown to improve the throughput by around 10% for
 `tensor_parallel_size=8`. For vision encoders that use hardware-unoptimized Conv3D operations,
-batch-level DP can provide another 40% improvement compared to regular TP.
+batch-level DP can provide another 40% increase to throughput compared to regular TP.
 
 Nevertheless, since the weights of the multi-modal encoder are replicated across each TP rank,
 there will be a minor increase in memory consumption and may cause OOM if you can barely fit the model already.
@@ -172,15 +172,13 @@ Batch-level DP needs to be implemented on a per-model basis,
 and enabled by setting `supports_encoder_tp_data = True` in the model class.
 Regardless, you need to set `mm_encoder_tp_mode="data"` in engine arguments to use this feature.
 
-Known supported models (with corresponding benchmarks):
+Known supported models:
 
-- dots_ocr (<gh-pr:25466>)
-- GLM-4.1V or above (<gh-pr:23168>)
-- InternVL (<gh-pr:23909>)
+- GLM-4.5V GLM-4.1V (<gh-pr:23168>)
 - Kimi-VL (<gh-pr:23817>)
 - Llama4 (<gh-pr:18368>)
 - MiniCPM-V-2.5 or above (<gh-pr:23327>, <gh-pr:23948>)
-- Qwen2-VL or above (<gh-pr:22742>, <gh-pr:24955>, <gh-pr:25445>)
+- Qwen2.5-VL (<gh-pr:22742>)
 - Step3 (<gh-pr:22697>)
 
 ## Input Processing
@@ -232,20 +230,6 @@ Multi-modal IPC caching is automatically enabled when
 there is a one-to-one correspondence between API (`P0`) and engine core (`P1`) processes,
 to avoid repeatedly transferring the same multi-modal inputs between them.
 
-#### Key-Replicated Cache
-
-By default, IPC caching uses a **key-replicated cache**, where cache keys exist
-in both the API (`P0`) and engine core (`P1`) processes, but the actual cache
-data resides only in `P1`.
-
-#### Shared Memory Cache
-
-When multiple worker processes are involved (e.g., when TP > 1), a
-**shared-memory cache** is more efficient. This can be enabled by setting
-`mm_processor_cache_type="shm"`. In this mode, cache keys are stored
-on `P0`, while the cache data itself lives in shared memory accessible by all
-processes.
-
 ### Configuration
 
 You can adjust the size of the cache by setting the value of `mm_processor_cache_gb` (default 4 GiB).
@@ -260,12 +244,6 @@ Examples:
 llm = LLM(model="Qwen/Qwen2.5-VL-3B-Instruct",
           mm_processor_cache_gb=8)
 
-# Use a shared-memory based IPC cache
-llm = LLM(model="Qwen/Qwen2.5-VL-3B-Instruct",
-          tensor_parallel_size=2,
-          mm_processor_cache_type="shm",
-          mm_processor_cache_gb=8)
-
 # Disable the cache
 llm = LLM(model="Qwen/Qwen2.5-VL-3B-Instruct",
           mm_processor_cache_gb=0)
@@ -275,12 +253,11 @@ llm = LLM(model="Qwen/Qwen2.5-VL-3B-Instruct",
 
 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` |
-| N/A | Disabled | N/A | N/A | N/A | `0` |
+| Processor Caching | IPC Caching | `P0` Cache | `P1` Cache | Max. Memory |
+|-------------------|-------------|------------|------------|-------------|
+| ✅ | ✅ | K | K + V | `mm_processor_cache_gb * data_parallel_size` |
+| ✅ | ❌ | K + V | N/A | `mm_processor_cache_gb * api_server_count` |
+| ❌ | ❌ | N/A | N/A | `0` |
 
 K: Stores the hashes of multi-modal items  
 V: Stores the processed tensor data of multi-modal items

docs/contributing/README.md

@@ -26,123 +26,113 @@ See <gh-file:LICENSE>.
 
 ## Developing
 
-The first step of contributing to vLLM is to clone the GitHub repository:
+--8<-- "docs/getting_started/installation/python_env_setup.inc.md"
+
+Depending on the kind of development you'd like to do (e.g. Python, CUDA), you can choose to build vLLM with or without compilation.
+Check out the [building from source][build-from-source] documentation for details.
+
+For an optimized workflow when iterating on C++/CUDA kernels, see the [Incremental Compilation Workflow](./incremental_build.md) for recommendations.
+
+### Building the docs with MkDocs
+
+#### Introduction to MkDocs
+
+[MkDocs](https://github.com/mkdocs/mkdocs) is a fast, simple and downright gorgeous static site generator that's geared towards building project documentation. Documentation source files are written in Markdown, and configured with a single YAML configuration file.
+
+#### Install MkDocs and Plugins
+
+Install MkDocs along with the [plugins](https://github.com/vllm-project/vllm/blob/main/mkdocs.yaml) used in the vLLM documentation, as well as required dependencies:
+
+```bash
+uv pip install -r requirements/docs.txt
+```
+
+!!! note
+    Ensure that your Python version is compatible with the plugins (e.g., `mkdocs-awesome-nav` requires Python 3.10+)
+
+#### Verify Installation
+
+Confirm that MkDocs is correctly installed:
+
+```bash
+mkdocs --version
+```
+
+Example output:
+
+```console
+mkdocs, version 1.6.1 from /opt/miniconda3/envs/mkdoc/lib/python3.10/site-packages/mkdocs (Python 3.10)
+```
+
+#### Clone the `vLLM` repository
 
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
 ```
 
-Then, configure your Python virtual environment.
+#### Start the Development Server
 
---8<-- "docs/getting_started/installation/python_env_setup.inc.md"
-
-If you are only developing vLLM's Python code, install vLLM using:
+MkDocs comes with a built-in dev-server that lets you preview your documentation as you work on it. Make sure you're in the same directory as the `mkdocs.yml` configuration file, and then start the server by running the `mkdocs serve` command:
 
 ```bash
-VLLM_USE_PRECOMPILED=1 uv pip install -e .
+mkdocs serve
 ```
 
-If you are developing vLLM's Python and CUDA/C++ code, install vLLM using:
+Example output:
 
-```bash
-uv pip install -e .
+```console
+INFO    -  Documentation built in 106.83 seconds
+INFO    -  [22:02:02] Watching paths for changes: 'docs', 'mkdocs.yaml'
+INFO    -  [22:02:02] Serving on http://127.0.0.1:8000/
 ```
 
-For more details about installing from source and installing for other hardware, check out the [installation instructions](../getting_started/installation/README.md) for your hardware and head to the "Build wheel from source" section.
+#### View in Your Browser
 
-For an optimized workflow when iterating on C++/CUDA kernels, see the [Incremental Compilation Workflow](./incremental_build.md) for recommendations.
+Open up [http://127.0.0.1:8000/](http://127.0.0.1:8000/) in your browser to see a live preview:.
+
+#### Learn More
+
+For additional features and advanced configurations, refer to the official [MkDocs Documentation](https://www.mkdocs.org/).
+
+## Testing
+
+??? console "Commands"
+
+    ```bash
+    # These commands are only for Nvidia CUDA platforms.
+    uv pip install -r requirements/common.txt -r requirements/dev.txt --torch-backend=auto
+
+    # Linting, formatting and static type checking
+    pre-commit install
+
+    # You can manually run pre-commit with
+    pre-commit run --all-files --show-diff-on-failure
+
+    # To manually run something from CI that does not run
+    # locally by default, you can run:
+    pre-commit run mypy-3.9 --hook-stage manual --all-files
+
+    # Unit tests
+    pytest tests/
+
+    # Run tests for a single test file with detailed output
+    pytest -s -v tests/test_logger.py
+    ```
 
 !!! tip
-    vLLM is compatible with Python versions 3.9 to 3.12. However, vLLM's default [Dockerfile](gh-file:docker/Dockerfile) ships with Python 3.12 and tests in CI (except `mypy`) are run with Python 3.12.
+    Since the <gh-file:docker/Dockerfile> ships with Python 3.12, all tests in CI (except `mypy`) are run with Python 3.12.
 
     Therefore, we recommend developing with Python 3.12 to minimise the chance of your local environment clashing with our CI environment.
 
-### Linting
-
-vLLM uses `pre-commit` to lint and format the codebase. See <https://pre-commit.com/#usage> if `pre-commit` is new to you. Setting up `pre-commit` is as easy as:
-
-```bash
-uv pip install pre-commit
-pre-commit install
-```
-
-vLLM's `pre-commit` hooks will now run automatically every time you commit.
-
-!!! tip "Tips"
-    You can manually run the `pre-commit` hooks using:
-
-    ```bash
-    pre-commit run     # runs on staged files
-    pre-commit run -a  # runs on all files (short for --all-files)
-    ```
-
-    ---
-
-    Some `pre-commit` hooks only run in CI. If you need to, you can run them locally with:
-
-    ```bash
-    pre-commit run --hook-stage manual markdownlint
-    pre-commit run --hook-stage manual mypy-3.9
-    ```
-
-### Documentation
-
-MkDocs is a fast, simple and downright gorgeous static site generator that's geared towards building project documentation. Documentation source files are written in Markdown, and configured with a single YAML configuration file, <gh-file:mkdocs.yaml>.
-
-Get started with:
-
-```bash
-uv pip install -r requirements/docs.txt
-```
-
-!!! tip
-    Ensure that your Python version is compatible with the plugins
-    (e.g., `mkdocs-awesome-nav` requires Python 3.10+)
-
-MkDocs comes with a built-in dev-server that lets you preview your documentation as you work on it.
-From the root of the repository, run:
-
-```bash
-mkdocs serve                           # with API ref (~10 minutes)
-API_AUTONAV_EXCLUDE=vllm mkdocs serve  # API ref off (~15 seconds)
-```
-
-Once you see `Serving on http://127.0.0.1:8000/` in the logs, the live preview is ready!
-Open <http://127.0.0.1:8000/> in your browser to see it.
-
-For additional features and advanced configurations, refer to the:
-
-- [MkDocs documentation](https://www.mkdocs.org/)
-- [Material for MkDocs documentation](https://squidfunk.github.io/mkdocs-material/) (the MkDocs theme we use)
-
-### Testing
-
-vLLM uses `pytest` to test the codebase.
-
-```bash
-# Install the test dependencies used in CI (CUDA only)
-uv pip install -r requirements/common.txt -r requirements/dev.txt --torch-backend=auto
-
-# Install some common test dependencies (hardware agnostic)
-uv pip install pytest pytest-asyncio
-
-# Run all tests
-pytest tests/
-
-# Run tests for a single test file with detailed output
-pytest -s -v tests/test_logger.py
-```
-
-!!! tip "Install python3-dev if Python.h is missing"
+!!! note "Install python3-dev if Python.h is missing"
     If any of the above commands fails with `Python.h: No such file or directory`, install
     `python3-dev` with `sudo apt install python3-dev`.
 
-!!! warning "Warnings"
+!!! note
     Currently, the repository is not fully checked by `mypy`.
 
-    ---
-
+!!! note
     Currently, not all unit tests pass when run on CPU platforms. If you don't have access to a GPU
     platform to run unit tests locally, rely on the continuous integration system to run the tests for
     now.
@@ -204,7 +194,8 @@ appropriately to indicate the type of change. Please use one of the following:
 The PR needs to meet the following code quality standards:
 
 - We adhere to [Google Python style guide](https://google.github.io/styleguide/pyguide.html) and [Google C++ style guide](https://google.github.io/styleguide/cppguide.html).
-- Pass all linter checks.
+- Pass all linter checks. Please use `pre-commit` to format your code. See
+  <https://pre-commit.com/#usage> if `pre-commit` is new to you.
 - The code needs to be well-documented to ensure future contributors can easily
   understand the code.
 - Include sufficient tests to ensure the project stays correct and robust. This

docs/contributing/benchmarks.md

@@ -1,787 +1,9 @@
----
-toc_depth: 4
----
-
 # Benchmark Suites
 
-vLLM provides comprehensive benchmarking tools for performance testing and evaluation:
+vLLM contains two sets of benchmarks:
 
-- **[Benchmark CLI]**: `vllm bench` CLI tools and specialized benchmark scripts for interactive performance testing
-- **[Performance benchmarks][performance-benchmarks]**: Automated CI benchmarks for development
-- **[Nightly benchmarks][nightly-benchmarks]**: Comparative benchmarks against alternatives
-
-[Benchmark CLI]: #benchmark-cli
-
-## Benchmark CLI
-
-This section guides you through running benchmark tests with the extensive
-datasets supported on vLLM. It's a living document, updated as new features and datasets
-become available.
-
-### Dataset Overview
-
-<style>
-th {
-  min-width: 0 !important;
-}
-</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/blob/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` |
-| Prefix Repetition | ✅ | ✅ | `synthetic` |
-| HuggingFace-VisionArena | ✅ | ✅ | `lmarena-ai/VisionArena-Chat` |
-| HuggingFace-MMVU | ✅ | ✅ | `yale-nlp/MMVU` |
-| HuggingFace-InstructCoder | ✅ | ✅ | `likaixin/InstructCoder` |
-| HuggingFace-AIMO | ✅ | ✅ | `AI-MO/aimo-validation-aime`, `AI-MO/NuminaMath-1.5`, `AI-MO/NuminaMath-CoT` |
-| HuggingFace-Other | ✅ | ✅ | `lmms-lab/LLaVA-OneVision-Data`, `Aeala/ShareGPT_Vicuna_unfiltered` |
-| HuggingFace-MTBench | ✅ | ✅ | `philschmid/mt-bench` |
-| HuggingFace-Blazedit | ✅ | ✅ | `vdaita/edit_5k_char`, `vdaita/edit_10k_char` |
-| Spec Bench | ✅ | ✅ | `wget https://raw.githubusercontent.com/hemingkx/Spec-Bench/refs/heads/main/data/spec_bench/question.jsonl` |
-| Custom | ✅ | ✅ | Local file: `data.jsonl` |
-
-Legend:
-
-- ✅ - supported
-- 🟡 - Partial support
-- 🚧 - to be supported
-
-!!! note
-    HuggingFace dataset's `dataset-name` should be set to `hf`.
-    For local `dataset-path`, please set `hf-name` to its Hugging Face ID like
-
-    ```bash
-    --dataset-path /datasets/VisionArena-Chat/ --hf-name lmarena-ai/VisionArena-Chat
-    ```
-
-### Examples
-
-#### 🚀 Online Benchmark
-
-<details class="admonition abstract" markdown="1">
-<summary>Show more</summary>
-
-First start serving your model
-
-```bash
-vllm serve NousResearch/Hermes-3-Llama-3.1-8B
-```
-
-Then run the benchmarking script
-
-```bash
-# download dataset
-# wget https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json
-vllm bench serve \
-  --backend vllm \
-  --model NousResearch/Hermes-3-Llama-3.1-8B \
-  --endpoint /v1/completions \
-  --dataset-name sharegpt \
-  --dataset-path <your data path>/ShareGPT_V3_unfiltered_cleaned_split.json \
-  --num-prompts 10
-```
-
-If successful, you will see the following output
-
-```text
-============ Serving Benchmark Result ============
-Successful requests:                     10
-Benchmark duration (s):                  5.78
-Total input tokens:                      1369
-Total generated tokens:                  2212
-Request throughput (req/s):              1.73
-Output token throughput (tok/s):         382.89
-Total Token throughput (tok/s):          619.85
----------------Time to First Token----------------
-Mean TTFT (ms):                          71.54
-Median TTFT (ms):                        73.88
-P99 TTFT (ms):                           79.49
------Time per Output Token (excl. 1st token)------
-Mean TPOT (ms):                          7.91
-Median TPOT (ms):                        7.96
-P99 TPOT (ms):                           8.03
----------------Inter-token Latency----------------
-Mean ITL (ms):                           7.74
-Median ITL (ms):                         7.70
-P99 ITL (ms):                            8.39
-==================================================
-```
-
-##### Custom Dataset
-
-If the dataset you want to benchmark is not supported yet in vLLM, even then you can benchmark on it using `CustomDataset`. Your data needs to be in `.jsonl` format and needs to have "prompt" field per entry, e.g., data.jsonl
-
-```json
-{"prompt": "What is the capital of India?"}
-{"prompt": "What is the capital of Iran?"}
-{"prompt": "What is the capital of China?"}
-```
-
-```bash
-# start server
-VLLM_USE_V1=1 vllm serve meta-llama/Llama-3.1-8B-Instruct
-```
-
-```bash
-# run benchmarking script
-vllm bench serve --port 9001 --save-result --save-detailed \
-  --backend vllm \
-  --model meta-llama/Llama-3.1-8B-Instruct \
-  --endpoint /v1/completions \
-  --dataset-name custom \
-  --dataset-path <path-to-your-data-jsonl> \
-  --custom-skip-chat-template \
-  --num-prompts 80 \
-  --max-concurrency 1 \
-  --temperature=0.3 \
-  --top-p=0.75 \
-  --result-dir "./log/"
-```
-
-You can skip applying chat template if your data already has it by using `--custom-skip-chat-template`.
-
-##### VisionArena Benchmark for Vision Language Models
-
-```bash
-# need a model with vision capability here
-vllm serve Qwen/Qwen2-VL-7B-Instruct
-```
-
-```bash
-vllm bench serve \
-  --backend openai-chat \
-  --model Qwen/Qwen2-VL-7B-Instruct \
-  --endpoint /v1/chat/completions \
-  --dataset-name hf \
-  --dataset-path lmarena-ai/VisionArena-Chat \
-  --hf-split train \
-  --num-prompts 1000
-```
-
-##### InstructCoder Benchmark with Speculative Decoding
-
-``` bash
-VLLM_USE_V1=1 vllm serve meta-llama/Meta-Llama-3-8B-Instruct \
-    --speculative-config $'{"method": "ngram",
-    "num_speculative_tokens": 5, "prompt_lookup_max": 5,
-    "prompt_lookup_min": 2}'
-```
-
-``` bash
-vllm bench serve \
-    --model meta-llama/Meta-Llama-3-8B-Instruct \
-    --dataset-name hf \
-    --dataset-path likaixin/InstructCoder \
-    --num-prompts 2048
-```
-
-##### Spec Bench Benchmark with Speculative Decoding
-
-``` bash
-VLLM_USE_V1=1 vllm serve meta-llama/Meta-Llama-3-8B-Instruct \
-    --speculative-config $'{"method": "ngram",
-    "num_speculative_tokens": 5, "prompt_lookup_max": 5,
-    "prompt_lookup_min": 2}'
-```
-
-[SpecBench dataset](https://github.com/hemingkx/Spec-Bench)
-
-Run all categories:
-
-``` bash
-# Download the dataset using:
-# wget https://raw.githubusercontent.com/hemingkx/Spec-Bench/refs/heads/main/data/spec_bench/question.jsonl
-
-vllm bench serve \
-    --model meta-llama/Meta-Llama-3-8B-Instruct \
-    --dataset-name spec_bench \
-    --dataset-path "<YOUR_DOWNLOADED_PATH>/data/spec_bench/question.jsonl" \
-    --num-prompts -1
-```
-
-Available categories include `[writing, roleplay, reasoning, math, coding, extraction, stem, humanities, translation, summarization, qa, math_reasoning, rag]`.
-
-Run only a specific category like "summarization":
-
-``` bash
-vllm bench serve \
-    --model meta-llama/Meta-Llama-3-8B-Instruct \
-    --dataset-name spec_bench \
-    --dataset-path "<YOUR_DOWNLOADED_PATH>/data/spec_bench/question.jsonl" \
-    --num-prompts -1
-    --spec-bench-category "summarization"
-```
-
-##### Other HuggingFaceDataset Examples
-
-```bash
-vllm serve Qwen/Qwen2-VL-7B-Instruct
-```
-
-`lmms-lab/LLaVA-OneVision-Data`:
-
-```bash
-vllm bench serve \
-  --backend openai-chat \
-  --model Qwen/Qwen2-VL-7B-Instruct \
-  --endpoint /v1/chat/completions \
-  --dataset-name hf \
-  --dataset-path lmms-lab/LLaVA-OneVision-Data \
-  --hf-split train \
-  --hf-subset "chart2text(cauldron)" \
-  --num-prompts 10
-```
-
-`Aeala/ShareGPT_Vicuna_unfiltered`:
-
-```bash
-vllm bench serve \
-  --backend openai-chat \
-  --model Qwen/Qwen2-VL-7B-Instruct \
-  --endpoint /v1/chat/completions \
-  --dataset-name hf \
-  --dataset-path Aeala/ShareGPT_Vicuna_unfiltered \
-  --hf-split train \
-  --num-prompts 10
-```
-
-`AI-MO/aimo-validation-aime`:
-
-``` bash
-vllm bench serve \
-    --model Qwen/QwQ-32B \
-    --dataset-name hf \
-    --dataset-path AI-MO/aimo-validation-aime \
-    --num-prompts 10 \
-    --seed 42
-```
-
-`philschmid/mt-bench`:
-
-``` bash
-vllm bench serve \
-    --model Qwen/QwQ-32B \
-    --dataset-name hf \
-    --dataset-path philschmid/mt-bench \
-    --num-prompts 80
-```
-
-`vdaita/edit_5k_char` or `vdaita/edit_10k_char`:
-
-``` bash
-vllm bench serve \
-    --model Qwen/QwQ-32B \
-    --dataset-name hf \
-    --dataset-path vdaita/edit_5k_char \
-    --num-prompts 90 \
-    --blazedit-min-distance 0.01 \
-    --blazedit-max-distance 0.99
-```
-
-##### Running With Sampling Parameters
-
-When using OpenAI-compatible backends such as `vllm`, optional sampling
-parameters can be specified. Example client command:
-
-```bash
-vllm bench serve \
-  --backend vllm \
-  --model NousResearch/Hermes-3-Llama-3.1-8B \
-  --endpoint /v1/completions \
-  --dataset-name sharegpt \
-  --dataset-path <your data path>/ShareGPT_V3_unfiltered_cleaned_split.json \
-  --top-k 10 \
-  --top-p 0.9 \
-  --temperature 0.5 \
-  --num-prompts 10
-```
-
-##### Running With Ramp-Up Request Rate
-
-The benchmark tool also supports ramping up the request rate over the
-duration of the benchmark run. This can be useful for stress testing the
-server or finding the maximum throughput that it can handle, given some latency budget.
-
-Two ramp-up strategies are supported:
-
-- `linear`: Increases the request rate linearly from a start value to an end value.
-- `exponential`: Increases the request rate exponentially.
-
-The following arguments can be used to control the ramp-up:
-
-- `--ramp-up-strategy`: The ramp-up strategy to use (`linear` or `exponential`).
-- `--ramp-up-start-rps`: The request rate at the beginning of the benchmark.
-- `--ramp-up-end-rps`: The request rate at the end of the benchmark.
-
-</details>
-
-#### 📈 Offline Throughput Benchmark
-
-<details class="admonition abstract" markdown="1">
-<summary>Show more</summary>
-
-```bash
-vllm bench throughput \
-  --model NousResearch/Hermes-3-Llama-3.1-8B \
-  --dataset-name sonnet \
-  --dataset-path vllm/benchmarks/sonnet.txt \
-  --num-prompts 10
-```
-
-If successful, you will see the following output
-
-```text
-Throughput: 7.15 requests/s, 4656.00 total tokens/s, 1072.15 output tokens/s
-Total num prompt tokens:  5014
-Total num output tokens:  1500
-```
-
-##### VisionArena Benchmark for Vision Language Models
-
-```bash
-vllm bench throughput \
-  --model Qwen/Qwen2-VL-7B-Instruct \
-  --backend vllm-chat \
-  --dataset-name hf \
-  --dataset-path lmarena-ai/VisionArena-Chat \
-  --num-prompts 1000 \
-  --hf-split train
-```
-
-The `num prompt tokens` now includes image token counts
-
-```text
-Throughput: 2.55 requests/s, 4036.92 total tokens/s, 326.90 output tokens/s
-Total num prompt tokens:  14527
-Total num output tokens:  1280
-```
-
-##### InstructCoder Benchmark with Speculative Decoding
-
-``` bash
-VLLM_WORKER_MULTIPROC_METHOD=spawn \
-VLLM_USE_V1=1 \
-vllm bench throughput \
-    --dataset-name=hf \
-    --dataset-path=likaixin/InstructCoder \
-    --model=meta-llama/Meta-Llama-3-8B-Instruct \
-    --input-len=1000 \
-    --output-len=100 \
-    --num-prompts=2048 \
-    --async-engine \
-    --speculative-config $'{"method": "ngram",
-    "num_speculative_tokens": 5, "prompt_lookup_max": 5,
-    "prompt_lookup_min": 2}'
-```
-
-```text
-Throughput: 104.77 requests/s, 23836.22 total tokens/s, 10477.10 output tokens/s
-Total num prompt tokens:  261136
-Total num output tokens:  204800
-```
-
-##### Other HuggingFaceDataset Examples
-
-`lmms-lab/LLaVA-OneVision-Data`:
-
-```bash
-vllm bench throughput \
-  --model Qwen/Qwen2-VL-7B-Instruct \
-  --backend vllm-chat \
-  --dataset-name hf \
-  --dataset-path lmms-lab/LLaVA-OneVision-Data \
-  --hf-split train \
-  --hf-subset "chart2text(cauldron)" \
-  --num-prompts 10
-```
-
-`Aeala/ShareGPT_Vicuna_unfiltered`:
-
-```bash
-vllm bench throughput \
-  --model Qwen/Qwen2-VL-7B-Instruct \
-  --backend vllm-chat \
-  --dataset-name hf \
-  --dataset-path Aeala/ShareGPT_Vicuna_unfiltered \
-  --hf-split train \
-  --num-prompts 10
-```
-
-`AI-MO/aimo-validation-aime`:
-
-```bash
-vllm bench throughput \
-  --model Qwen/QwQ-32B \
-  --backend vllm \
-  --dataset-name hf \
-  --dataset-path AI-MO/aimo-validation-aime \
-  --hf-split train \
-  --num-prompts 10
-```
-
-Benchmark with LoRA adapters:
-
-``` bash
-# download dataset
-# wget https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json
-vllm bench throughput \
-  --model meta-llama/Llama-2-7b-hf \
-  --backend vllm \
-  --dataset_path <your data path>/ShareGPT_V3_unfiltered_cleaned_split.json \
-  --dataset_name sharegpt \
-  --num-prompts 10 \
-  --max-loras 2 \
-  --max-lora-rank 8 \
-  --enable-lora \
-  --lora-path yard1/llama-2-7b-sql-lora-test
-```
-
-</details>
-
-#### 🛠️ Structured Output Benchmark
-
-<details class="admonition abstract" markdown="1">
-<summary>Show more</summary>
-
-Benchmark the performance of structured output generation (JSON, grammar, regex).
-
-##### Server Setup
-
-```bash
-vllm serve NousResearch/Hermes-3-Llama-3.1-8B
-```
-
-##### JSON Schema Benchmark
-
-```bash
-python3 benchmarks/benchmark_serving_structured_output.py \
-  --backend vllm \
-  --model NousResearch/Hermes-3-Llama-3.1-8B \
-  --dataset json \
-  --structured-output-ratio 1.0 \
-  --request-rate 10 \
-  --num-prompts 1000
-```
-
-##### Grammar-based Generation Benchmark
-
-```bash
-python3 benchmarks/benchmark_serving_structured_output.py \
-  --backend vllm \
-  --model NousResearch/Hermes-3-Llama-3.1-8B \
-  --dataset grammar \
-  --structure-type grammar \
-  --request-rate 10 \
-  --num-prompts 1000
-```
-
-##### Regex-based Generation Benchmark
-
-```bash
-python3 benchmarks/benchmark_serving_structured_output.py \
-  --backend vllm \
-  --model NousResearch/Hermes-3-Llama-3.1-8B \
-  --dataset regex \
-  --request-rate 10 \
-  --num-prompts 1000
-```
-
-##### Choice-based Generation Benchmark
-
-```bash
-python3 benchmarks/benchmark_serving_structured_output.py \
-  --backend vllm \
-  --model NousResearch/Hermes-3-Llama-3.1-8B \
-  --dataset choice \
-  --request-rate 10 \
-  --num-prompts 1000
-```
-
-##### XGrammar Benchmark Dataset
-
-```bash
-python3 benchmarks/benchmark_serving_structured_output.py \
-  --backend vllm \
-  --model NousResearch/Hermes-3-Llama-3.1-8B \
-  --dataset xgrammar_bench \
-  --request-rate 10 \
-  --num-prompts 1000
-```
-
-</details>
-
-#### 📚 Long Document QA Benchmark
-
-<details class="admonition abstract" markdown="1">
-<summary>Show more</summary>
-
-Benchmark the performance of long document question-answering with prefix caching.
-
-##### Basic Long Document QA Test
-
-```bash
-python3 benchmarks/benchmark_long_document_qa_throughput.py \
-  --model meta-llama/Llama-2-7b-chat-hf \
-  --enable-prefix-caching \
-  --num-documents 16 \
-  --document-length 2000 \
-  --output-len 50 \
-  --repeat-count 5
-```
-
-##### Different Repeat Modes
-
-```bash
-# Random mode (default) - shuffle prompts randomly
-python3 benchmarks/benchmark_long_document_qa_throughput.py \
-  --model meta-llama/Llama-2-7b-chat-hf \
-  --enable-prefix-caching \
-  --num-documents 8 \
-  --document-length 3000 \
-  --repeat-count 3 \
-  --repeat-mode random
-
-# Tile mode - repeat entire prompt list in sequence
-python3 benchmarks/benchmark_long_document_qa_throughput.py \
-  --model meta-llama/Llama-2-7b-chat-hf \
-  --enable-prefix-caching \
-  --num-documents 8 \
-  --document-length 3000 \
-  --repeat-count 3 \
-  --repeat-mode tile
-
-# Interleave mode - repeat each prompt consecutively
-python3 benchmarks/benchmark_long_document_qa_throughput.py \
-  --model meta-llama/Llama-2-7b-chat-hf \
-  --enable-prefix-caching \
-  --num-documents 8 \
-  --document-length 3000 \
-  --repeat-count 3 \
-  --repeat-mode interleave
-```
-
-</details>
-
-#### 🗂️ Prefix Caching Benchmark
-
-<details class="admonition abstract" markdown="1">
-<summary>Show more</summary>
-
-Benchmark the efficiency of automatic prefix caching.
-
-##### Fixed Prompt with Prefix Caching
-
-```bash
-python3 benchmarks/benchmark_prefix_caching.py \
-  --model meta-llama/Llama-2-7b-chat-hf \
-  --enable-prefix-caching \
-  --num-prompts 1 \
-  --repeat-count 100 \
-  --input-length-range 128:256
-```
-
-##### ShareGPT Dataset with Prefix Caching
-
-```bash
-# download dataset
-# wget https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json
-
-python3 benchmarks/benchmark_prefix_caching.py \
-  --model meta-llama/Llama-2-7b-chat-hf \
-  --dataset-path /path/ShareGPT_V3_unfiltered_cleaned_split.json \
-  --enable-prefix-caching \
-  --num-prompts 20 \
-  --repeat-count 5 \
-  --input-length-range 128:256
-```
-
-##### Prefix Repetition Dataset
-
-```bash
-vllm bench serve \
-  --backend openai \
-  --model meta-llama/Llama-2-7b-chat-hf \
-  --dataset-name prefix_repetition \
-  --num-prompts 100 \
-  --prefix-repetition-prefix-len 512 \
-  --prefix-repetition-suffix-len 128 \
-  --prefix-repetition-num-prefixes 5 \
-  --prefix-repetition-output-len 128
-```
-
-</details>
-
-#### ⚡ Request Prioritization Benchmark
-
-<details class="admonition abstract" markdown="1">
-<summary>Show more</summary>
-
-Benchmark the performance of request prioritization in vLLM.
-
-##### Basic Prioritization Test
-
-```bash
-python3 benchmarks/benchmark_prioritization.py \
-  --model meta-llama/Llama-2-7b-chat-hf \
-  --input-len 128 \
-  --output-len 64 \
-  --num-prompts 100 \
-  --scheduling-policy priority
-```
-
-##### Multiple Sequences per Prompt
-
-```bash
-python3 benchmarks/benchmark_prioritization.py \
-  --model meta-llama/Llama-2-7b-chat-hf \
-  --input-len 128 \
-  --output-len 64 \
-  --num-prompts 100 \
-  --scheduling-policy priority \
-  --n 2
-```
-
-</details>
-
-#### 👁️ Multi-Modal Benchmark
-
-<details class="admonition abstract" markdown="1">
-<summary>Show more</summary>
-
-Benchmark the performance of multi-modal requests in vLLM.
-
-##### Images (ShareGPT4V)
-
-Start vLLM:
-
-```bash
-python -m vllm.entrypoints.openai.api_server \
-  --model Qwen/Qwen2.5-VL-7B-Instruct \
-  --dtype bfloat16 \
-  --limit-mm-per-prompt '{"image": 1}' \
-  --allowed-local-media-path /path/to/sharegpt4v/images
-```
-
-Send requests with images:
-
-```bash
-vllm bench serve \
-  --backend openai-chat \
-  --model Qwen/Qwen2.5-VL-7B-Instruct \
-  --dataset-name sharegpt \
-  --dataset-path /path/to/ShareGPT4V/sharegpt4v_instruct_gpt4-vision_cap100k.json \
-  --num-prompts 100 \
-  --save-result \
-  --result-dir ~/vllm_benchmark_results \
-  --save-detailed \
-  --endpoint /v1/chat/completions
-```
-
-##### Videos (ShareGPT4Video)
-
-Start vLLM:
-
-```bash
-python -m vllm.entrypoints.openai.api_server \
-  --model Qwen/Qwen2.5-VL-7B-Instruct \
-  --dtype bfloat16 \
-  --limit-mm-per-prompt '{"video": 1}' \
-  --allowed-local-media-path /path/to/sharegpt4video/videos
-```
-
-Send requests with videos:
-
-```bash
-vllm bench serve \
-  --backend openai-chat \
-  --model Qwen/Qwen2.5-VL-7B-Instruct \
-  --dataset-name sharegpt \
-  --dataset-path /path/to/ShareGPT4Video/llava_v1_5_mix665k_with_video_chatgpt72k_share4video28k.json \
-  --num-prompts 100 \
-  --save-result \
-  --result-dir ~/vllm_benchmark_results \
-  --save-detailed \
-  --endpoint /v1/chat/completions
-```
-
-##### Synthetic Random Images (random-mm)
-
-Generate synthetic image inputs alongside random text prompts to stress-test vision models without external datasets.
-
-Notes:
-
-- 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):
-
-```bash
-vllm serve Qwen/Qwen2.5-VL-3B-Instruct \
-  --dtype bfloat16 \
-  --max-model-len 16384 \
-  --limit-mm-per-prompt '{"image": 3, "video": 0}' \
-  --mm-processor-kwargs max_pixels=1003520
-```
-
-Benchmark. It is recommended to use the flag `--ignore-eos` to simulate real responses. You can set the size of the output via the arg `random-output-len`.
-
-Ex.1: Fixed number of items and a single image resolution, enforcing generation of approx 40 tokens:
-
-```bash
-vllm bench serve \
-  --backend openai-chat \
-  --model Qwen/Qwen2.5-VL-3B-Instruct \
-  --endpoint /v1/chat/completions \
-  --dataset-name random-mm \
-  --num-prompts 100 \
-  --max-concurrency 10 \
-  --random-prefix-len 25 \
-  --random-input-len 300 \
-  --random-output-len 40 \
-  --random-range-ratio 0.2 \
-  --random-mm-base-items-per-request 2 \
-  --random-mm-limit-mm-per-prompt '{"image": 3, "video": 0}' \
-  --random-mm-bucket-config '{(224, 224, 1): 1.0}' \
-  --request-rate inf \
-  --ignore-eos \
-  --seed 42
-```
-
-The number of items per request can be controlled by passing multiple image buckets:
-
-```bash
-  --random-mm-base-items-per-request 2 \
-  --random-mm-num-mm-items-range-ratio 0.5 \
-  --random-mm-limit-mm-per-prompt '{"image": 4, "video": 0}' \
-  --random-mm-bucket-config '{(256, 256, 1): 0.7, (720, 1280, 1): 0.3}' \
-```
-
-Flags specific to `random-mm`:
-
-- `--random-mm-base-items-per-request`: base number of multimodal items per request.
-- `--random-mm-num-mm-items-range-ratio`: vary item count uniformly in the closed integer range [floor(n·(1−r)), ceil(n·(1+r))]. Set r=0 to keep it fixed; r=1 allows 0 items.
-- `--random-mm-limit-mm-per-prompt`: per-modality hard caps, e.g. '{"image": 3, "video": 0}'.
-- `--random-mm-bucket-config`: dict mapping (H, W, T) → probability. Entries with probability 0 are removed; remaining probabilities are renormalized to sum to 1. Use T=1 for images. Set any T>1 for videos (video sampling not yet supported).
-
-Behavioral notes:
-
-- If the requested base item count cannot be satisfied under the provided per-prompt limits, the tool raises an error rather than silently clamping.
-
-How sampling works:
-
-- Determine per-request item count k by sampling uniformly from the integer range defined by `--random-mm-base-items-per-request` and `--random-mm-num-mm-items-range-ratio`, then clamp k to at most the sum of per-modality limits.
-- For each of the k items, sample a bucket (H, W, T) according to the normalized probabilities in `--random-mm-bucket-config`, while tracking how many items of each modality have been added.
-- If a modality (e.g., image) reaches its limit from `--random-mm-limit-mm-per-prompt`, all buckets of that modality are excluded and the remaining bucket probabilities are renormalized before continuing.
-This should be seen as an edge case, and if this behavior can be avoided by setting `--random-mm-limit-mm-per-prompt` to a large number. Note that this might result in errors due to engine config `--limit-mm-per-prompt`.
-- The resulting request contains synthetic image data in `multi_modal_data` (OpenAI Chat format). When `random-mm` is used with the OpenAI Chat backend, prompts remain text and MM content is attached via `multi_modal_data`.
-
-</details>
+- [Performance benchmarks][performance-benchmarks]
+- [Nightly benchmarks][nightly-benchmarks]
 
 [](){ #performance-benchmarks }
 
@@ -791,22 +13,22 @@ The performance benchmarks are used for development to confirm whether new chang
 
 ### Manually Trigger the benchmark
 
-Use [vllm-ci-test-repo images](https://gallery.ecr.aws/q9t5s3a7/vllm-ci-test-repo) with vLLM benchmark suite.
+Use [vllm-ci-test-repo images](https://gallery.ecr.aws/q9t5s3a7/vllm-ci-test-repo) with vLLM benchmark suite.  
 For CPU environment, please use the image with "-cpu" postfix.
 
-Here is an example for docker run command for CPU.
+Here is an example for docker run command for CPU.  
 
 ```bash
 docker run -it --entrypoint /bin/bash -v /data/huggingface:/root/.cache/huggingface  -e HF_TOKEN=''  --shm-size=16g --name vllm-cpu-ci  public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:1da94e673c257373280026f75ceb4effac80e892-cpu
 ```
 
-Then, run below command inside the docker instance.
+Then, run below command inside the docker instance.  
 
 ```bash
 bash .buildkite/nightly-benchmarks/scripts/run-performance-benchmarks.sh
 ```
 
-When run, benchmark script generates results under **benchmark/results** folder, along with the benchmark_results.md and benchmark_results.json.
+When run, benchmark script generates results under **benchmark/results** folder, along with the benchmark_results.md and benchmark_results.json.  
 
 #### Runtime environment variables
 

docs/contributing/incremental_build.md

@@ -40,16 +40,6 @@ python tools/generate_cmake_presets.py
 
 The script will prompt you if it cannot automatically determine certain paths (e.g., `nvcc` or a specific Python executable for your vLLM development environment). Follow the on-screen prompts. If an existing `CMakeUserPresets.json` is found, the script will ask for confirmation before overwriting it.
 
-**Force overwrite existing file:**
-
-To automatically overwrite an existing `CMakeUserPresets.json` without prompting, use the `--force-overwrite` flag:
-
-```console
-python tools/generate_cmake_presets.py --force-overwrite
-```
-
-This is particularly useful in automated scripts or CI/CD environments where interactive prompts are not desired.
-
 After running the script, a `CMakeUserPresets.json` file will be created in the root of your vLLM repository.
 
 ### Example `CMakeUserPresets.json`

docs/contributing/model/README.md

@@ -3,7 +3,7 @@
 !!! important
     Many decoder language models can now be automatically loaded using the [Transformers backend][transformers-backend] without having to implement them in vLLM. See if `vllm serve <model>` works first!
 
-vLLM models are specialized [PyTorch](https://pytorch.org/) models that take advantage of various [features](../../features/README.md#compatibility-matrix) to optimize their performance.
+vLLM models are specialized [PyTorch](https://pytorch.org/) models that take advantage of various [features](../../features/compatibility_matrix.md) to optimize their performance.
 
 The complexity of integrating a model into vLLM depends heavily on the model's architecture.
 The process is considerably straightforward if the model shares a similar architecture with an existing model in vLLM.

docs/contributing/model/multimodal.md

@@ -840,6 +840,7 @@ Some HF processors directly insert feature tokens without replacing anything in
 Examples:
 
 - BLIP-2 (insert at start of prompt): <gh-file:vllm/model_executor/models/blip2.py>
+- Florence2 (insert at start of prompt): <gh-file:vllm/model_executor/models/florence2.py>
 - Molmo (insert after `<|endoftext|>` token): <gh-file:vllm/model_executor/models/molmo.py>
 
 ### Handling prompt updates unrelated to multi-modal data

docs/deployment/frameworks/streamlit.md

@@ -6,33 +6,35 @@ It can be quickly integrated with vLLM as a backend API server, enabling powerfu
 
 ## Prerequisites
 
-Set up the vLLM environment by installing all required packages:
-
-```bash
-pip install vllm streamlit openai
-```
+- Setup vLLM environment
 
 ## Deploy
 
-1. Start the vLLM server with a supported chat completion model, e.g.
+- Start the vLLM server with the supported chat completion model, e.g.
 
-    ```bash
-    vllm serve Qwen/Qwen1.5-0.5B-Chat
-    ```
+```bash
+vllm serve qwen/Qwen1.5-0.5B-Chat
+```
 
-1. Use the script: <gh-file:examples/online_serving/streamlit_openai_chatbot_webserver.py>
+- Install streamlit and openai:
 
-1. Start the streamlit web UI and start to chat:
+```bash
+pip install streamlit openai
+```
 
-    ```bash
+- Use the script: <gh-file:examples/online_serving/streamlit_openai_chatbot_webserver.py>
+
+- Start the streamlit web UI and start to chat:
+
+```bash
+streamlit run streamlit_openai_chatbot_webserver.py
+
+# or specify the VLLM_API_BASE or VLLM_API_KEY
+VLLM_API_BASE="http://vllm-server-host:vllm-server-port/v1" \
     streamlit run streamlit_openai_chatbot_webserver.py
 
-    # or specify the VLLM_API_BASE or VLLM_API_KEY
-    VLLM_API_BASE="http://vllm-server-host:vllm-server-port/v1" \
-        streamlit run streamlit_openai_chatbot_webserver.py
+# start with debug mode to view more details
+streamlit run streamlit_openai_chatbot_webserver.py --logger.level=debug
+```
 
-    # start with debug mode to view more details
-    streamlit run streamlit_openai_chatbot_webserver.py --logger.level=debug
-    ```
-
-    ![Chat with vLLM assistant in Streamlit](../../assets/deployment/streamlit-chat.png)
+![](../../assets/deployment/streamlit-chat.png)

docs/design/huggingface_integration.md

@@ -1,31 +1,31 @@
 # Integration with Hugging Face
 
-This document describes how vLLM integrates with Hugging Face libraries. We will explain step by step what happens under the hood when we run `vllm serve`.
+This document describes how vLLM integrates with HuggingFace libraries. We will explain step by step what happens under the hood when we run `vllm serve`.
 
-Let's say we want to serve the popular Qwen model by running `vllm serve Qwen/Qwen2-7B`.
+Let's say we want to serve the popular QWen model by running `vllm serve Qwen/Qwen2-7B`.
 
 1. The `model` argument is `Qwen/Qwen2-7B`. vLLM determines whether this model exists by checking for the corresponding config file `config.json`. See this [code snippet](https://github.com/vllm-project/vllm/blob/10b67d865d92e376956345becafc249d4c3c0ab7/vllm/transformers_utils/config.py#L162-L182) for the implementation. Within this process:
     - If the `model` argument corresponds to an existing local path, vLLM will load the config file directly from this path.
-    - If the `model` argument is a Hugging Face model ID consisting of a username and model name, vLLM will first try to use the config file from the Hugging Face local cache, using the `model` argument as the model name and the `--revision` argument as the revision. See [their website](https://huggingface.co/docs/huggingface_hub/en/package_reference/environment_variables#hfhome) for more information on how the Hugging Face cache works.
-    - If the `model` argument is a Hugging Face model ID but it is not found in the cache, vLLM will download the config file from the Hugging Face model hub. Refer to [this function](https://github.com/vllm-project/vllm/blob/10b67d865d92e376956345becafc249d4c3c0ab7/vllm/transformers_utils/config.py#L91) for the implementation. The input arguments include the `model` argument as the model name, the `--revision` argument as the revision, and the environment variable `HF_TOKEN` as the token to access the model hub. In our case, vLLM will download the [config.json](https://huggingface.co/Qwen/Qwen2-7B/blob/main/config.json) file.
+    - If the `model` argument is a HuggingFace model ID consisting of a username and model name, vLLM will first try to use the config file from the HuggingFace local cache, using the `model` argument as the model name and the `--revision` argument as the revision. See [their website](https://huggingface.co/docs/huggingface_hub/en/package_reference/environment_variables#hfhome) for more information on how the HuggingFace cache works.
+    - If the `model` argument is a HuggingFace model ID but it is not found in the cache, vLLM will download the config file from the HuggingFace model hub. Refer to [this function](https://github.com/vllm-project/vllm/blob/10b67d865d92e376956345becafc249d4c3c0ab7/vllm/transformers_utils/config.py#L91) for the implementation. The input arguments include the `model` argument as the model name, the `--revision` argument as the revision, and the environment variable `HF_TOKEN` as the token to access the model hub. In our case, vLLM will download the [config.json](https://huggingface.co/Qwen/Qwen2-7B/blob/main/config.json) file.
 
 2. After confirming the existence of the model, vLLM loads its config file and converts it into a dictionary. See this [code snippet](https://github.com/vllm-project/vllm/blob/10b67d865d92e376956345becafc249d4c3c0ab7/vllm/transformers_utils/config.py#L185-L186) for the implementation.
 
 3. Next, vLLM [inspects](https://github.com/vllm-project/vllm/blob/10b67d865d92e376956345becafc249d4c3c0ab7/vllm/transformers_utils/config.py#L189) the `model_type` field in the config dictionary to [generate](https://github.com/vllm-project/vllm/blob/10b67d865d92e376956345becafc249d4c3c0ab7/vllm/transformers_utils/config.py#L190-L216) the config object to use. There are some `model_type` values that vLLM directly supports; see [here](https://github.com/vllm-project/vllm/blob/10b67d865d92e376956345becafc249d4c3c0ab7/vllm/transformers_utils/config.py#L48) for the list. If the `model_type` is not in the list, vLLM will use [AutoConfig.from_pretrained](https://huggingface.co/docs/transformers/en/model_doc/auto#transformers.AutoConfig.from_pretrained) to load the config class, with `model`, `--revision`, and `--trust_remote_code` as the arguments. Please note that:
-    - Hugging Face also has its own logic to determine the config class to use. It will again use the `model_type` field to search for the class name in the transformers library; see [here](https://github.com/huggingface/transformers/tree/main/src/transformers/models) for the list of supported models. If the `model_type` is not found, Hugging Face will use the `auto_map` field from the config JSON file to determine the class name. Specifically, it is the `AutoConfig` field under `auto_map`. See [DeepSeek](https://huggingface.co/deepseek-ai/DeepSeek-V2.5/blob/main/config.json) for an example.
-    - The `AutoConfig` field under `auto_map` points to a module path in the model's repository. To create the config class, Hugging Face will import the module and use the `from_pretrained` method to load the config class. This can generally cause arbitrary code execution, so it is only executed when `--trust_remote_code` is enabled.
+    - HuggingFace also has its own logic to determine the config class to use. It will again use the `model_type` field to search for the class name in the transformers library; see [here](https://github.com/huggingface/transformers/tree/main/src/transformers/models) for the list of supported models. If the `model_type` is not found, HuggingFace will use the `auto_map` field from the config JSON file to determine the class name. Specifically, it is the `AutoConfig` field under `auto_map`. See [DeepSeek](https://huggingface.co/deepseek-ai/DeepSeek-V2.5/blob/main/config.json) for an example.
+    - The `AutoConfig` field under `auto_map` points to a module path in the model's repository. To create the config class, HuggingFace will import the module and use the `from_pretrained` method to load the config class. This can generally cause arbitrary code execution, so it is only executed when `--trust_remote_code` is enabled.
 
 4. Subsequently, vLLM applies some historical patches to the config object. These are mostly related to RoPE configuration; see [here](https://github.com/vllm-project/vllm/blob/127c07480ecea15e4c2990820c457807ff78a057/vllm/transformers_utils/config.py#L244) for the implementation.
 
 5. Finally, vLLM can reach the model class we want to initialize. vLLM uses the `architectures` field in the config object to determine the model class to initialize, as it maintains the mapping from architecture name to model class in [its registry](https://github.com/vllm-project/vllm/blob/127c07480ecea15e4c2990820c457807ff78a057/vllm/model_executor/models/registry.py#L80). If the architecture name is not found in the registry, it means this model architecture is not supported by vLLM. For `Qwen/Qwen2-7B`, the `architectures` field is `["Qwen2ForCausalLM"]`, which corresponds to the `Qwen2ForCausalLM` class in [vLLM's code](https://github.com/vllm-project/vllm/blob/127c07480ecea15e4c2990820c457807ff78a057/vllm/model_executor/models/qwen2.py#L364). This class will initialize itself depending on various configs.
 
-Beyond that, there are two more things vLLM depends on Hugging Face for.
+Beyond that, there are two more things vLLM depends on HuggingFace for.
 
-1. **Tokenizer**: vLLM uses the tokenizer from Hugging Face to tokenize the input text. The tokenizer is loaded using [AutoTokenizer.from_pretrained](https://huggingface.co/docs/transformers/en/model_doc/auto#transformers.AutoTokenizer.from_pretrained) with the `model` argument as the model name and the `--revision` argument as the revision. It is also possible to use a tokenizer from another model by specifying the `--tokenizer` argument in the `vllm serve` command. Other relevant arguments are `--tokenizer-revision` and `--tokenizer-mode`. Please check Hugging Face's documentation for the meaning of these arguments. This part of the logic can be found in the [get_tokenizer](https://github.com/vllm-project/vllm/blob/127c07480ecea15e4c2990820c457807ff78a057/vllm/transformers_utils/tokenizer.py#L87) function. After obtaining the tokenizer, notably, vLLM will cache some expensive attributes of the tokenizer in [get_cached_tokenizer](https://github.com/vllm-project/vllm/blob/127c07480ecea15e4c2990820c457807ff78a057/vllm/transformers_utils/tokenizer.py#L24).
+1. **Tokenizer**: vLLM uses the tokenizer from HuggingFace to tokenize the input text. The tokenizer is loaded using [AutoTokenizer.from_pretrained](https://huggingface.co/docs/transformers/en/model_doc/auto#transformers.AutoTokenizer.from_pretrained) with the `model` argument as the model name and the `--revision` argument as the revision. It is also possible to use a tokenizer from another model by specifying the `--tokenizer` argument in the `vllm serve` command. Other relevant arguments are `--tokenizer-revision` and `--tokenizer-mode`. Please check HuggingFace's documentation for the meaning of these arguments. This part of the logic can be found in the [get_tokenizer](https://github.com/vllm-project/vllm/blob/127c07480ecea15e4c2990820c457807ff78a057/vllm/transformers_utils/tokenizer.py#L87) function. After obtaining the tokenizer, notably, vLLM will cache some expensive attributes of the tokenizer in [get_cached_tokenizer](https://github.com/vllm-project/vllm/blob/127c07480ecea15e4c2990820c457807ff78a057/vllm/transformers_utils/tokenizer.py#L24).
 
-2. **Model weight**: vLLM downloads the model weight from the Hugging Face model hub using the `model` argument as the model name and the `--revision` argument as the revision. vLLM provides the argument `--load-format` to control what files to download from the model hub. By default, it will try to load the weights in the safetensors format and fall back to the PyTorch bin format if the safetensors format is not available. We can also pass `--load-format dummy` to skip downloading the weights.
+2. **Model weight**: vLLM downloads the model weight from the HuggingFace model hub using the `model` argument as the model name and the `--revision` argument as the revision. vLLM provides the argument `--load-format` to control what files to download from the model hub. By default, it will try to load the weights in the safetensors format and fall back to the PyTorch bin format if the safetensors format is not available. We can also pass `--load-format dummy` to skip downloading the weights.
     - It is recommended to use the safetensors format, as it is efficient for loading in distributed inference and also safe from arbitrary code execution. See the [documentation](https://huggingface.co/docs/safetensors/en/index) for more information on the safetensors format. This part of the logic can be found [here](https://github.com/vllm-project/vllm/blob/10b67d865d92e376956345becafc249d4c3c0ab7/vllm/model_executor/model_loader/loader.py#L385). Please note that:
 
-This completes the integration between vLLM and Hugging Face.
+This completes the integration between vLLM and HuggingFace.
 
-In summary, vLLM reads the config file `config.json`, tokenizer, and model weight from the Hugging Face model hub or a local directory. It uses the config class from either vLLM, Hugging Face transformers, or loads the config class from the model's repository.
+In summary, vLLM reads the config file `config.json`, tokenizer, and model weight from the HuggingFace model hub or a local directory. It uses the config class from either vLLM, HuggingFace transformers, or loads the config class from the model's repository.

docs/design/logits_processors.md

@@ -1,559 +0,0 @@
-# Logits Processors
-
-!!! important
-    Some logits processors design changes are still in progress and the API may
-    change in the near future. We hope to stabilize this part of the API soon
-
-This document describes how the vLLM engine interacts with logits processors, and the programming model which vLLM supports for implementing logits processors.
-
-## Logits Processors Background
-
-A logits processor adjusts the next-token probability distribution, usually with the intention of steering the model towards a desired type of behavior.
-
-In vLLM, logits processors operate at batch granularity. During a given engine step, the logits processor consumes a `(num_requests) x (vocab_size)` tensor of raw logits output by the model. For all requests which enable the logits processor, the logits processor applies a transformation to the corresponding row of the logits tensor, while leaving other rows unmodified. The transformed logits tensor is then passed to softmax.  
-
-## Logits Processors in the vLLM engine
-
-The vLLM engine's persistent batch data structure maintains a list of loaded logits processors.
-
-In order to operate on the entire batch at once, each logits processor may maintain metadata about the requests in the batch (i.e. each request's logits-processor-specific configuration settings). Therefore, logits processors are stateful.
-
-In each engine step, the vLLM engine will (1) update each logits processor's internal state and (2) apply logits processors to the model output logits.
-
-### Updating Logits Processor Internal State
-
-At the beginning of each engine step, the persistent batch may add, discard and/or reorder requests in response to the scheduler output. After the persistent batch has reorganized, the vLLM engine invokes each logits processor's `update_state()` method. This is necessary to ensure that logits processors' internal states are reorganized to match the new persistent batch state at the beginning of the engine step.
-
-The pseudocode below shows the process by which the vLLM persistent batch notifies each logits processor of changes in batch state:
-
-??? code "Model Runner Updates Logits Processor States"
-
-    ``` python
-    # gpu_model_runner.py
-
-    class GPUModelRunner(...):
-
-        ...
-
-        def execute_model(self, scheduler_output, ...):
-            self._update_states(scheduler_output)
-
-            ...
-
-        def _update_states(...):
-
-            ...
-
-            # ...update persistent batch to reflect new/finished requests & reordering
-            # of requests within batch...
-
-            ...
-
-            self.input_batch.refresh_metadata()
-
-
-    # gpu_input_batch.py
-
-    class InputBatch:
-
-        ...
-
-        def refresh_metadata(self):
-
-            ...
-
-            # Update each logits processor's state to reflect persistent batch state
-            batch_update = self.batch_update_builder.get_and_reset(self.num_reqs)
-            for logit_proc in self.logitsprocs.all:
-                logit_proc.update_state(batch_update)
-
-            ...
-
-
-    # vllm/v1/sample/logits_processor/interface.py
-
-    @dataclass(frozen=True)
-    class BatchUpdate:
-        # Batch state-change data structure which is passed to logits processors'
-        # update_state() methods
-
-        batch_size: int
-
-        removed: Sequence[RemovedRequest]
-        added: Sequence[AddedRequest]
-        moved: Sequence[MovedRequest]
-    
-    ```
-
-### Applying Logits Processors to the Model Output Logits
-
-After updating persistent batch state, the vLLM model runner performs model inference to obtain logits. Then, the model runner invokes the sampler against the logits. In turn, part of the sampler's operation is to invoke the logits processors' `apply()` methods against the model output logit processors, yielding transformed logits (the `apply()` methods may modify the logits in-place or out-of-place, although in-place is more memory-efficient). This process is shown in the pseudocode below.
-
-Note that the sampler will access the logits processors via `SamplingMetadata.logitsprocs`. When the vLLM engine constructs `SamplingMetadata` (not shown in the code below), the reference to the list of logits processors is passed from the persistent batch data structure to `SamplingMetadata`.
-
-??? code "Apply logits processors to model output logits"
-
-    ``` python
-    # gpu_model_runner.py
-
-    class GPUModelRunner(...):
-
-        ...
-
-        def execute_model(self, scheduler_output, ...):
-            # (discussed in previous section)
-            self._update_states(scheduler_output)
-
-            ...
-
-            # ...run model inference to obtain logits...
-
-            ...
-
-            # Invoke sampler, which applies logits processors
-            sampler_output = self.sampler(logits=logits,
-                                          sampling_metadata=sampling_metadata)
-
-            ...
-
-
-    # sampler.py
-
-    class Sampler(nn.Module):
-
-        ...
-
-        def forward(self, logits, sampling_metadata):
-
-            ...
-
-            # Apply non-argmax-invariant logits processors to model output logits
-            for processor in (sampling_metadata.logitsprocs.non_argmax_invariant):
-                logits = processor.apply(logits)
-
-            sampled = self.sample(logits, sampling_metadata)
-
-            ...
-
-            # ...return sampler output data structure...
-
-
-        def sample(self, logits, sampling_metadta)
-
-            ...
-
-            # ...exit early if all requests are greedy-sampling...
-
-            ...
-
-            # Apply argmax-invariant logits processors
-            for processor in sampling_metadata.logitsprocs.argmax_invariant:
-                logits = processor.apply(logits)
-
-            ...
-
-            # ...perform sampling and return sampling result...
-    ``` 
-
-At sampling time, the sampler checks whether all requests in the persistent batch employ greedy sampling. If that is the case, the sampler saves compute by skipping "argmax-invariant" logits processors. Here, "argmax" is shorthand for the token ID with the highest logit value in a given row of the logits tensor (i.e. the token which the model weighted the highest for a given request).
-
-* An **argmax-invariant logits processor** is a logits processor (such as Min-P) which does not modify the argmax. For example, a logits processor which masks out the lowest-probability tokens will not change which token ID has the max logit. Greedy sampling always picks the highest-logit-value token ID, and so conceptually an argmax-invariant logits processor can be skipped for greedy sampling requests.
-
-* A **non-argmax-invariant logits processor** is a logits processor which may modify the argmax. For example, a logits processor which masks all tokens except for EOS after a certain number of steps in order to force decoding to terminate might end up masking the max-logit-value token and therefore change the argmax. Conceptually, these logits processors cannot be skipped for greedy sampling requests.
-
-The vLLM logits processor abstraction requires the engine to apply logits processors at batch granularity; therefore in practice the argmax-invariant logits processors can only be skipped when the entire batch uses greedy sampling.
-
-## Logits Processor Programming Model
-
-The previous sections alluded to the interfaces which vLLM logits processors must support. This section introduces in full the programming model for implementing logits processors that are compatible with the vLLM engine, including the `LogitsProcessor` base class and its interface methods as well as the `BatchUpdate` data structure for representing persistent batch state changes, both of which are shown in the code below:
-
-??? code "`LogitsProcessor` base class and `BatchUpdate` data structure"
-
-    ``` python
-    from abc import ABC, abstractmethod
-    from collections.abc import Sequence
-    from dataclasses import dataclass
-    from enum import Enum, auto
-    from typing import TYPE_CHECKING, Optional
-
-    import torch
-
-    from vllm import SamplingParams
-
-    if TYPE_CHECKING:
-        from vllm.config import VllmConfig
-
-
-    class MoveDirectionality(Enum):
-        # One-way i1->i2 req move within batch
-        UNIDIRECTIONAL = auto()
-        # Two-way i1<->i2 req swap within batch
-        SWAP = auto()
-
-
-    # (index, params, prompt_tok_ids, output_tok_ids) tuples for new
-    # requests added to the batch.
-    AddedRequest = tuple[int, SamplingParams, list[int], list[int]]
-
-    # (index 1, index 2, directionality) tuples representing
-    # one-way moves or two-way swaps of requests in batch
-    MovedRequest = tuple[int, int, MoveDirectionality]
-
-    # Batch indices of any removed requests.
-    RemovedRequest = int
-
-
-    @dataclass(frozen=True)
-    class BatchUpdate:
-        """Persistent batch state change info for logitsprocs"""
-        batch_size: int  # Current num reqs in batch
-
-        # Metadata for requests added to, removed from, and moved
-        # within the persistent batch.
-        #
-        # Key assumption: the `output_tok_ids` list (which is an element of each
-        # tuple in `added`) is a reference to the request's running output tokens
-        # list; via this reference, the logits processors always see the latest
-        # list of generated output tokens
-        removed: Sequence[RemovedRequest]
-        moved: Sequence[MovedRequest]
-        added: Sequence[AddedRequest]
-
-
-    class LogitsProcessor(ABC):
-
-        @abstractmethod
-        def __init__(self, vllm_config: "VllmConfig", device: torch.device,
-                    is_pin_memory: bool) -> None:
-            raise NotImplementedError
-
-        @abstractmethod
-        def apply(self, logits: torch.Tensor) -> torch.Tensor:
-            raise NotImplementedError
-
-        @abstractmethod
-        def is_argmax_invariant(self) -> bool:
-            """True if logits processor has no impact on the
-            argmax computation in greedy sampling.
-            NOTE: may or may not have the same value for all
-            instances of a given LogitsProcessor subclass,
-            depending on subclass implementation.
-            """
-            raise NotImplementedError
-
-        @abstractmethod
-        def update_state(
-            self,
-            batch_update: Optional["BatchUpdate"],
-        ) -> None:
-            """Called when there are new output tokens, prior
-            to each forward pass.
-
-            Args:
-                batch_update is non-None iff there have been
-                changes to the batch makeup.
-            """
-            raise NotImplementedError
-            
-    ```
-
-A vLLM logits processor must subclass `LogitsProcessor` and define (at minimum) the following methods:
-
-* `__init__(self, vllm_config: VllmConfig, device: torch.device, is_pin_memory: bool)`
-    * `vllm_config`: engine configuration data structure
-    * `device`: hardware accelerator device info
-    * `is_pin_memory`: flag indicating whether pin memory is available to support logits processor implementation
-
-* `apply(self, logits: torch.Tensor) -> torch.Tensor`:
-    * Consume a `(num_requests) x (vocab_size)` logits tensor (`logits`)
-    * Apply logits processor transformation at batch granularity
-    * Return a transformed `(num_requests) x (vocab_size)` logits tensor
-    * You can modify the input logits processors in-place or out-of-place; in-place is more memory-efficient
-
-* `is_argmax_invariant(self) -> bool`:
-    * Return `True` if the logits processor is argmax invariant (never changes what is the highest-logit-value token ID for a given request), `False` if the logits processor may modify argmax
-    * `is_argmax_invariant()` is evaluated once at startup; if `True`, vLLM will skip applying this logits processor in a given step when all requests use greedy sampling
-
-* `update_state(self, batch_update: Optional["BatchUpdate"]) -> None`:
-    * Consume a `BatchUpdate` data structure representing persistent batch state changes at the beginning of the current engine step
-    * Use the `BatchUpdate` members to update logits processor internal state
-    * **Note:** batch update data structure may be `None`, signaling no change to the batch constituents. In this case, the LogitsProcessor might still want to update its state based on the updated `output_token_ids` lists that it could have retained when they were added.
-
-### `BatchUpdate` data structure
-
-The `BatchUpdate` abstraction models the persistent batch as a list of requests, supporting the following operations to change batch state (note that the order in which the operations are mentioned below reflects the order in which they should be processed in `update_state()`):
-
-* **Remove:** remove (without replacement) request at index `i`
-
-    * A Remove is represented in `Batchupdate.removed` by an `int` (representing `i`)
-
-    * Effect of remove-at-index on batch:
-
-        ``` text
-        Batch: [A,B,C]
-        Remove @ i:  1
-
-        =>
-
-        New Batch: [A,x,C] # Discard B and leave an empty slot
-        ```
-
-* **Add:** add (or replace existing request with) a new request at index `i`. If a request is replaced, its associated state should be discarded.
-
-    * An Add is represented in `Batchupdate.added` as a tuple of
-
-        ``` text
-        (index, new request SamplingParams, prompt token ids, output token ids)
-        ```
-
-    * `prompt token ids` and `output token ids` are references to the request's prompt token ids and output token ids lists, respectively. Note that the output token ids list grows with each engine step, and this growth is visible to the logits processor because output token ids are passed by reference. **This is important for LogitsProcessors that take into account the tokens generated so far**.
-
-    * The implementation of the particular logits processor subclass determines whether or how the fields in the added request tuple are digested into an internal representation. For example, a logits processor that does not utilize prompt or output token ids may only need to utilize `index` and `SamplingParams` and discard the other tuple fields
-
-    * If index `i` currently holds a request, a replacement occurs:
-
-        ``` text
-        Batch: [A,B,C]
-        New request to be added @ i: D @ 1
-
-        =>
-
-        New Batch: [A,D,C] # Add D, discard B
-        ```
-
-    * If index `i` does not currently hold a request (because `i` is out of bounds of the current batch size):
-
-        ``` text
-        Batch: [A,B,C]
-        New request to be added @ i: D @ 3
-
-        =>
-
-        New Batch: [A,B,C,D] # Add D, extending batch
-        ```
-
-* **Move:** move request at index `s` to index `d` OR swap requests at indices `s` and `d`
-
-    * A Move is represented in `Batchupdate.moved` as a tuple of
-
-        ``` text
-        (s, d, UNIDIRECTIONAL or SWAP)
-        ```
-
-    * If the Move specifies `UNIDRECTIONAL`:
-
-        * The request at index `s` is moved to index `d`; index `s` becomes an empty slot
-
-            ``` text
-            Batch: [A,x,C,D]
-            Unidirectionally Move s -> d:  3 -> 1
-
-            =>
-
-            New Batch: [A,D,C,x] # Move D to 1, leaving empty slot at 3
-            ```
-
-        * If another request already resided at index `d`, it is replaced and discarded
-
-            ``` text
-            Batch: [A,B,C,D]
-            Unidirectionally Move s -> d:  3 -> 1
-
-            =>
-
-            New Batch: [A,D,C,x] # Move D to 1, discarding B and leaving empty slot at 3
-            ```
-
-    * If the Move specifies `SWAP`, the requests at `s` and `d` exchange indices
-
-        ``` text
-        Batch: [A,B,C,D]
-        Swap Move s <-> d:  3 <-> 1
-
-        =>
-
-        New Batch: [A,D,C,B] # Swap B and D
-        ```
-
-Additionally, the `BatchUpdate` data structure includes a representation (`batch_size`) of the size of the persistent batch at the beginning of the engine step.
-
-### How the vLLM engine builds the `BatchUpdate` data structure
-
-Logits processor `update_state()` implementations should assume the following model for how the model runner updates persistent batch state (expressed here in terms of the `BatchUpdate` abstraction):
-
-1. Identify indices of requests which finished in the current engine step
-
-2. Identify new requests introduced in the current step
-
-3. Use Add operations to replace as many finished requests with new requests, in order of increasing index of the replaced request starting with the lowest index
-
-4. Based on the relative number of new and finished requests:
-
-    1. If the numbers of new and finished requests are the same, proceed to next step
-
-    2. *If there are more new requests than finished requests:* apply Add operations to extend the batch with the remaining new requests which did not replace finished requests. Assign consecutive indices to these new requests, starting with `current_max_batch_index + 1`
-
-    3. *If there are fewer new requests than finished requests:*
-
-        * Apply Remove operations to finished requests which were not replaced with new requests. These removed request indices will necessarily be greater than the greatest index of the finished requests which were replaced in the previous step. The Removes may leave the batch in a non-contiguous state
-
-        * **"Condense" the batch to be contiguous:** starting with the lowest-index empty slot (which was caused by a Remove), apply a Unidirectional Move from the current highest non-empty slot in the batch to fill the empty slot. Proceed with additional Unidirectional Move operations in order of increasing empty slot destination index and decreasing non-empty slot source index until the batch is contiguous
-
-        * **Shrink the batch:** a side-effect of condensing the batch is that empty slots resulting from Remove operations are grouped in a contiguous block at the end of the batch array. Thus, after condensing, update `BatchUpdate.batch_size` to reflect the number of non-empty slots
-
-5. Reorder the batch for improved efficiency. Depending on the attention backend implementation and the current characteristics of the batch, zero or more Swap Move operations may be applied to reorder the batch
-
-Notes:
-
-* A logits processor `update_state()` method must process batch update operations in the following order: removes, adds, moves
-
-* The index argument for Add operations refers to the index *at the time the Add occurred*, i.e. before any Move operations
-    * Example: if a request is Added at index 5 and then swapped with index 3, the Add operation in `BatchUpdate.added` will be associated with index 5 not 3
-    * In other words Move operations can be assumed to be applied after Adds and Removes
-
-* Move operations can be assumed to be applied in the order in which they appear in `BatchUpdate.moved`
-
-* If there are no new/finished requests and there is no batch reordering, then the batch update for the logits processors will be `None`
-
-#### Example: Batch Update with Fewer New Requests Than Finished Requests
-
-The following example models an engine step where 1 new request is introduced and 2 finished requests are eliminated, additionally the attention backend performs a swap to optimize the batch ordering.
-
-``` text
-Batch state (beginning of engine step): [A,B,C,D]
-Batch size: 4
-
-New requests: E
-
-Finished requests: A, C
-
-Processing steps (using BatchUpdate abstraction):
-
-1. Add E at index 0
-
-[E,B,C,D] # Discard A
-Batch size: 4
-
-2. Remove at index 2
-
-[E,B,x,D] # Discard C, empty slot at index 2
-Batch size: 4
-
-3. Condense batch with a Unidirectional Move 3 -> 2 operation and shrink batch
-
-[E,B,D] x # Empty slot is now outside batch
-Batch size: 3
-
-4. Attention backend optimization: reorder batch with Swap 0 <-> 1
-
-[B,E,D]
-Batch size: 3
-
-```
-
-The resulting `BatchUpdate` data structure will look like
-
-``` text
-BatchUpdate instance
-* added: [(0,E's SamplingParams,E's prompt tokens ref,E's output tokens ref)]
-* removed: [2] # request C was removed without replacement
-* moved: [(3,2,UNIDIRECTIONAL),(0,1,SWAP)]
-```
-
-#### Example: Batch Update with More New Requests Than Finished Requests
-
-The following example models an engine step where 2 new requests are introduced and 1 finished request is eliminated, additionally the attention backend performs a swap to optimize the batch ordering.
-
-``` text
-Batch state (beginning of engine step): [A,B,C,D]
-Batch size: 4
-
-New requests: E,F
-
-Finished requests: C
-
-Processing steps (using BatchUpdate abstraction):
-
-1. Add E at index 2
-
-[A,B,E,D] # Discard C
-Batch size: 4
-
-2. Add F at index 4 (current max batch index + 1)
-
-[A,B,E,D,F] # Extend batch by 1
-Batch size: 5
-
-4. Attention backend optimization: reorder batch with Swap 0 <-> 1
-
-[B,A,E,D,F]
-Batch size: 5
-
-```
-
-Note that batch condensation is skipped because there are no empty slots left behind by Remove operations.
-
-The resulting `BatchUpdate` data structure will look like
-
-``` text
-BatchUpdate instance
-* added: [(2,E's SamplingParams,E's prompt tokens ref,E's output tokens ref),(4,F's SamplingParams,F's prompt tokens ref,F's output tokens ref)]
-* removed: [] # no requests were removed without replacement
-* moved: [(0,1,SWAP)]
-```
-
-## How to Introduce a New Logits Processor to vLLM
-
-### Best Practices for Writing Built-In Logits Processors
-
-* Write efficient `apply()` and `update_state()` implementations in light of the fact that logits processors operate at batch granularity
-    * For example, you may be able to use efficient vectorized operations to implement `apply()` or update internal state vectors in `update_state()`
-    * However, if you think that a logits processor may be used infrequently, it may be appropriate to use a "sparse" representation of request state i.e. the class can represent request configuration using a dictionary which only stores metadata about requests that enable the logits processor
-
-* It is up to the logits processor author to determine:
-
-    1. **The per-request attributes which configure the logits processor's behavior against that request.** For example, if you are writing a new built-in logits processor for vLLM, you may or may not need to add additional fields to `SamplingParams` and the vLLM REST API
-
-    2. **The conditions under which the logits processor is or is not enabled on a per-request basis.** Unless your intention is for the built-in logits processor to act on all requests all the time, you should write your logits processor in such a way that it is possible to disable the logits processor for a given request, i.e. by defaulting an argument to `None` or by passing in a specific do-nothing argument value i.e. `0.0`. Try to save compute and memory for requests which disable the logits processor
-
-    3. **The conditions under which the logits processor is short-circuited at the batch level.** Even if you have defined a way to disable the built-in logits processor at the request level, it may be difficult to translate this into compute savings i.e. if your `update_state()` and `apply()` implementations use efficient vectorized implementations that operate on the whole persistent batch in a single command. For example, you cannot skip an entire vectorized operation in `apply()` just because one request disabled the logits processor. To save compute in the edge-case where no running requests utilize the built-in logits processor, we recommend designing `apply()` to return the unmodified input tensor if all requests have the logits processor disabled. Similarly, consider whether steps can be skipped in `update_state()` if no requests enable the logits processor
-
-        * Additionally, an easy way to save compute in `update_state()` is to exit early when the batch_update is `None`
-
-* Ensure that the logits processor `update_state` method discards information about finished requests (i.e. requests which are replaced by an Add or which are subject to a Remove)
-
-* `is_argmax_invariant()` can be hard-coded to `True` or `False` if the logits processor has consistent behavior. However the argmax invariance may also be determined programmatically (i.e. if your logits processor is user-customizable in some way that impacts whether the logits processor is argmax invariant). For this reason, `is_argmax_invariant()` is not a class method
-
-### Built-In Logits Processors
-
-Built-in logits processors are always loaded when the vLLM engine starts. See the existing vLLM built-in logits processors in `vllm/v1/sample/logits_processor/builtin.py` for examples of how to write a new built-in vLLM logits processor. It makes sense to write a PR to introduce a new logits processor as a built-in if it is likely to be useful to a wide audience. vLLM currently employs the following built-in logits processors based on the programming model described above:
-
-* Min-P
-
-* Logit bias
-
-* Min-tokens
-
-Review these logits processor implementations for guidance on writing built-in logits processors.
-
-Additionally, the following logits-processor-like functionalities are hard-coded into the sampler and do not yet utilize the programming model described above. Most of them will be refactored to use the aforemented logits processor programming model.
-
-* Allowed token IDs
-
-* Bad words
-
-* Repetition penalty
-
-* Frequency penalty
-
-* Presence penalty
-
-* Temperature
-
-* Top-K
-
-* Top-P
-
-### Custom Logits Processors
-
-vLLM can be augmented with [user-provided custom logits processors](../features/custom_logitsprocs.md).

docs/design/multiprocessing.md

@@ -8,7 +8,7 @@ page for information on known issues and how to solve them.
 ## Introduction
 
 !!! important
-    The source code references are to the state of the code at the time of writing in December 2024.
+    The source code references are to the state of the code at the time of writing in December, 2024.
 
 The use of Python multiprocessing in vLLM is complicated by:
 

docs/examples/README.md

@@ -2,6 +2,6 @@
 
 vLLM's examples are split into three categories:
 
-- If you are using vLLM from within Python code, see the *Offline Inference* section.
-- If you are using vLLM from an HTTP application or client, see the *Online Serving* section.
-- For examples of using some of vLLM's advanced features (e.g. LMCache or Tensorizer) which are not specific to either of the above use cases, see the *Others* section.
+- If you are using vLLM from within Python code, see [Offline Inference](./offline_inference)
+- If you are using vLLM from an HTTP application or client, see [Online Serving](./online_serving)
+- For examples of using some of vLLM's advanced features (e.g. LMCache or Tensorizer) which are not specific to either of the above use cases, see [Others](./others)

docs/features/README.md

@@ -36,23 +36,22 @@ th:not(:first-child) {
 }
 </style>
 
-| Feature | [CP][chunked-prefill] | [APC](automatic_prefix_caching.md) | [LoRA](lora.md) | [SD](spec_decode.md) | CUDA graph | [pooling](../models/pooling_models.md) | <abbr title="Encoder-Decoder Models">enc-dec</abbr> | <abbr title="Logprobs">logP</abbr> | <abbr title="Prompt Logprobs">prmpt logP</abbr> | <abbr title="Async Output Processing">async output</abbr> | multi-step | <abbr title="Multimodal Inputs">mm</abbr> | best-of | beam-search | [prompt-embeds](prompt_embeds.md) |
-|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
-| [CP][chunked-prefill] | ✅ | | | | | | | | | | | | | | |
-| [APC](automatic_prefix_caching.md) | ✅ | ✅ | | | | | | | | | | | | | |
-| [LoRA](lora.md) | ✅ | ✅ | ✅ | | | | | | | | | | | | |
-| [SD](spec_decode.md) | ✅ | ✅ | ❌ | ✅ | | | | | | | | | | | |
-| CUDA graph | ✅ | ✅ | ✅ | ✅ | ✅ | | | | | | | | | | |
-| [pooling](../models/pooling_models.md) | 🟠\* | 🟠\* | ✅ | ❌ | ✅ | ✅ | | | | | | | | | |
-| <abbr title="Encoder-Decoder Models">enc-dec</abbr> | ❌ | [❌](gh-issue:7366) | ❌ | [❌](gh-issue:7366) | ✅ | ✅ | ✅ | | | | | | | | |
-| <abbr title="Logprobs">logP</abbr> | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | | | | | | | |
-| <abbr title="Prompt Logprobs">prmpt logP</abbr> | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | | | | | | |
-| <abbr title="Async Output Processing">async output</abbr> | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ✅ | | | | | |
-| multi-step | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | | | | |
-| [mm](multimodal_inputs.md) | ✅ | ✅ | [🟠](gh-pr:4194)<sup>^</sup> | ❔ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❔ | ✅ | | | |
-| best-of | ✅ | ✅ | ✅ | [❌](gh-issue:6137) | ✅ | ❌ | ✅ | ✅ | ✅ | ❔ | [❌](gh-issue:7968) | ✅ | ✅ | | |
-| beam-search | ✅ | ✅ | ✅ | [❌](gh-issue:6137) | ✅ | ❌ | ✅ | ✅ | ✅ | ❔ | [❌](gh-issue:7968) | ❔ | ✅ | ✅ | |
-| [prompt-embeds](prompt_embeds.md) | ✅ | [❌](gh-issue:25096) | ? | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ? | ? | ❌ | ? | ? | ✅ |
+| Feature | [CP][chunked-prefill] | [APC](automatic_prefix_caching.md) | [LoRA](lora.md) | [SD](spec_decode.md) | CUDA graph | [pooling](../models/pooling_models.md) | <abbr title="Encoder-Decoder Models">enc-dec</abbr> | <abbr title="Logprobs">logP</abbr> | <abbr title="Prompt Logprobs">prmpt logP</abbr> | <abbr title="Async Output Processing">async output</abbr> | multi-step | <abbr title="Multimodal Inputs">mm</abbr> | best-of | beam-search |
+|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
+| [CP][chunked-prefill] | ✅ | | | | | | | | | | | | | |
+| [APC](automatic_prefix_caching.md) | ✅ | ✅ | | | | | | | | | | | | |
+| [LoRA](lora.md) | ✅ | ✅ | ✅ | | | | | | | | | | | |
+| [SD](spec_decode.md) | ✅ | ✅ | ❌ | ✅ | | | | | | | | | | |
+| CUDA graph | ✅ | ✅ | ✅ | ✅ | ✅ | | | | | | | | | |
+| [pooling](../models/pooling_models.md) | 🟠\* | 🟠\* | ✅ | ❌ | ✅ | ✅ | | | | | | | | |
+| <abbr title="Encoder-Decoder Models">enc-dec</abbr> | ❌ | [❌](gh-issue:7366) | ❌ | [❌](gh-issue:7366) | ✅ | ✅ | ✅ | | | | | | | |
+| <abbr title="Logprobs">logP</abbr> | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | | | | | | |
+| <abbr title="Prompt Logprobs">prmpt logP</abbr> | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | | | | | |
+| <abbr title="Async Output Processing">async output</abbr> | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ✅ | | | | |
+| multi-step | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | | | |
+| [mm](multimodal_inputs.md) | ✅ | ✅ | [🟠](gh-pr:4194)<sup>^</sup> | ❔ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❔ | ✅ | | |
+| best-of | ✅ | ✅ | ✅ | [❌](gh-issue:6137) | ✅ | ❌ | ✅ | ✅ | ✅ | ❔ | [❌](gh-issue:7968) | ✅ | ✅ | |
+| beam-search | ✅ | ✅ | ✅ | [❌](gh-issue:6137) | ✅ | ❌ | ✅ | ✅ | ✅ | ❔ | [❌](gh-issue:7968) | ❔ | ✅ | ✅ |
 
 \* Chunked prefill and prefix caching are only applicable to last-token pooling.  
 <sup>^</sup> LoRA is only applicable to the language backbone of multimodal models.
@@ -77,4 +76,6 @@ th:not(:first-child) {
 | multi-step                                                | ✅                  | ✅        | ✅        | ✅     | ✅        | [❌](gh-issue:8477) | ✅     | ❌ |
 | best-of                                                   | ✅                  | ✅        | ✅        | ✅     | ✅        | ✅                  | ✅     | ❌ |
 | beam-search                                               | ✅                  | ✅        | ✅        | ✅     | ✅        | ✅                  | ✅     | ❌ |
-| [prompt-embeds](prompt_embeds.md)                         | ✅                  | ✅        | ✅        | ✅     | ✅        | ✅                  | ?     | [❌](gh-issue:25097) |
+
+!!! note
+    Please refer to [Feature support through NxD Inference backend][feature-support-through-nxd-inference-backend] for features supported on AWS Neuron hardware