Merge branch 'main' into wye-refactor-quant-folder

Signed-off-by: yewentao256 <zhyanwentao@126.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/release-pipeline.yaml

@@ -1,22 +1,24 @@
 steps:
   # aarch64 + CUDA builds. PyTorch 2.8 aarch64 + CUDA wheel is only available on CUDA 12.9
   - label: "Build arm64 wheel - CUDA 12.9"
-    depends_on: ~
     id: build-wheel-arm64-cuda-12-9
     agents:
       queue: arm64_cpu_queue_postmerge
     commands:
       # #NOTE: torch_cuda_arch_list is derived from upstream PyTorch build files here:
       # https://github.com/pytorch/pytorch/blob/main/.ci/aarch64_linux/aarch64_ci_build.sh#L7
-      - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg VLLM_MAIN_CUDA_VERSION=12.9 --build-arg torch_cuda_arch_list='8.7 9.0 10.0+PTX 12.0' --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
+      - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list='8.7 9.0 10.0+PTX 12.0' --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
       - "mkdir artifacts"
       - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
       - "bash .buildkite/scripts/upload-wheels.sh"
     env:
       DOCKER_BUILDKIT: "1"
 
+  - block: "Build CUDA 12.8 wheel"
+    key: block-build-cu128-wheel
+
   - label: "Build wheel - CUDA 12.8"
-    depends_on: ~
+    depends_on: block-build-cu128-wheel
     id: build-wheel-cuda-12-8
     agents:
       queue: cpu_queue_postmerge
@@ -28,8 +30,12 @@ steps:
     env:
       DOCKER_BUILDKIT: "1"
 
-  - label: "Build wheel - CUDA 12.6"
+  - block: "Build CUDA 12.6 wheel"
+    key: block-build-cu126-wheel
     depends_on: ~
+
+  - label: "Build wheel - CUDA 12.6"
+    depends_on: block-build-cu126-wheel
     id: build-wheel-cuda-12-6
     agents:
       queue: cpu_queue_postmerge
@@ -96,6 +102,8 @@ steps:
     depends_on:
       - create-multi-arch-manifest
       - build-wheel-cuda-12-8
+      - build-wheel-cuda-12-6
+      - build-wheel-cuda-12-9
     id: annotate-release-workflow
     agents:
       queue: cpu_queue_postmerge

.buildkite/scripts/annotate-release.sh

@@ -14,33 +14,18 @@ buildkite-agent annotate --style 'info' --context 'release-workflow' << EOF
 To download the wheel:
 \`\`\`
 aws s3 cp s3://vllm-wheels/${RELEASE_VERSION}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux1_x86_64.whl .
-aws s3 cp s3://vllm-wheels/${RELEASE_VERSION}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux2014_aarch64.whl .
-
 aws s3 cp s3://vllm-wheels/${RELEASE_VERSION}+cu126/vllm-${RELEASE_VERSION}+cu126-cp38-abi3-manylinux1_x86_64.whl .
-aws s3 cp s3://vllm-wheels/${RELEASE_VERSION}+cu129/vllm-${RELEASE_VERSION}+cu129-cp38-abi3-manylinux1_x86_64.whl .
+aws s3 cp s3://vllm-wheels/${RELEASE_VERSION}+cu118/vllm-${RELEASE_VERSION}+cu118-cp38-abi3-manylinux1_x86_64.whl . 
 \`\`\`
 
 To download and upload the image:
 
 \`\`\`
-docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64
-docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64
-
-docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64 vllm/vllm-openai:x86_64
-docker tag vllm/vllm-openai:x86_64 vllm/vllm-openai:latest-x86_64
-docker tag vllm/vllm-openai:x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64
-docker push vllm/vllm-openai:latest-x86_64
-docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64
-
-docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64 vllm/vllm-openai:aarch64
-docker tag vllm/vllm-openai:aarch64 vllm/vllm-openai:latest-aarch64
-docker tag vllm/vllm-openai:aarch64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
-docker push vllm/vllm-openai:latest-aarch64
-docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
-
-docker manifest create vllm/vllm-openai:latest vllm/vllm-openai:latest-x86_64 vllm/vllm-openai:latest-aarch64 --amend
-docker manifest create vllm/vllm-openai:v${RELEASE_VERSION} vllm/vllm-openai:v${RELEASE_VERSION}-x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64 --amend
-docker manifest push vllm/vllm-openai:latest
-docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}
+docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}
+docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT} vllm/vllm-openai
+docker tag vllm/vllm-openai vllm/vllm-openai:latest
+docker tag vllm/vllm-openai vllm/vllm-openai:v${RELEASE_VERSION}
+docker push vllm/vllm-openai:latest
+docker push vllm/vllm-openai:v${RELEASE_VERSION}
 \`\`\`
 EOF 

.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

@@ -66,6 +66,7 @@ function cpu_tests() {
 
     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/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
@@ -326,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
@@ -386,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
@@ -541,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]
@@ -801,7 +817,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
@@ -813,20 +829,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  #####
 
@@ -882,13 +884,13 @@ steps:
   - tests/distributed/
   - 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
   - vllm/v1/engine/
-  - vllm/v1/worker/
-  - tests/v1/worker/test_worker_memory_snapshot.py
   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
@@ -907,10 +909,9 @@ steps:
   - pytest -v -s distributed/test_sequence_parallel.py
   # this test fails consistently.
   # TODO: investigate and fix
-  - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s test_sharded_state_loader.py
+  - 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
-  - pytest -v -s v1/worker/test_worker_memory_snapshot.py
 
 - label: Plugin Tests (2 GPUs) # 40min
   timeout_in_minutes: 60
@@ -953,6 +954,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
@@ -1026,34 +1028,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

@@ -171,7 +171,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 +302,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

.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

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.
@@ -257,8 +243,8 @@ set(VLLM_EXT_SRC
   "csrc/sampler.cu"
   "csrc/cuda_view.cu"
   "csrc/quantization/gptq/q_gemm.cu"
-  "csrc/quantization/compressed_tensors/int8_quant_kernels.cu"
-  "csrc/quantization/fp8/common.cu"
+  "csrc/quantization/w8a8/int8/scaled_quant.cu"
+  "csrc/quantization/w8a8/fp8/common.cu"
   "csrc/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu"
   "csrc/quantization/gguf/gguf_kernel.cu"
   "csrc/quantization/activation_kernels.cu"
@@ -302,13 +288,15 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
   list(APPEND VLLM_EXT_SRC
     "csrc/quantization/awq/gemm_kernels.cu"
     "csrc/permute_cols.cu"
-    "csrc/quantization/cutlass_w8a8/scaled_mm_entry.cu"
+    "csrc/quantization/w8a8/cutlass/scaled_mm_entry.cu"
     "csrc/quantization/fp4/nvfp4_quant_entry.cu"
     "csrc/quantization/fp4/nvfp4_scaled_mm_entry.cu"
     "csrc/quantization/fp4/nvfp4_blockwise_moe_kernel.cu"
     "csrc/sparse/cutlass/sparse_scaled_mm_entry.cu"
     "csrc/cutlass_extensions/common.cpp"
-    "csrc/quantization/fp8/per_token_group_quant.cu")
+    "csrc/attention/mla/cutlass_mla_entry.cu"
+    "csrc/quantization/w8a8/fp8/per_token_group_quant.cu"
+    "csrc/quantization/w8a8/int8/per_token_group_quant.cu")
 
   set_gencode_flags_for_srcs(
     SRCS "${VLLM_EXT_SRC}"
@@ -412,11 +400,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
   cuda_archs_loose_intersection(SCALED_MM_ARCHS "9.0a;" "${CUDA_ARCHS}")
   if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND SCALED_MM_ARCHS)
     set(SRCS
-       "csrc/quantization/cutlass_w8a8/scaled_mm_c3x_sm90.cu"
-       "csrc/quantization/cutlass_w8a8/c3x/scaled_mm_sm90_fp8.cu"
-       "csrc/quantization/cutlass_w8a8/c3x/scaled_mm_sm90_int8.cu"
-       "csrc/quantization/cutlass_w8a8/c3x/scaled_mm_azp_sm90_int8.cu"
-       "csrc/quantization/cutlass_w8a8/c3x/scaled_mm_blockwise_sm90_fp8.cu")
+       "csrc/quantization/w8a8/cutlass/scaled_mm_c3x_sm90.cu"
+       "csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8.cu"
+       "csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_int8.cu"
+       "csrc/quantization/w8a8/cutlass/c3x/scaled_mm_azp_sm90_int8.cu"
+       "csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm90_fp8.cu")
     set_gencode_flags_for_srcs(
       SRCS "${SRCS}"
       CUDA_ARCHS "${SCALED_MM_ARCHS}")
@@ -443,9 +431,9 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
   cuda_archs_loose_intersection(SCALED_MM_ARCHS "12.0;12.0a" "${CUDA_ARCHS}")
   if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS)
     set(SRCS
-      "csrc/quantization/cutlass_w8a8/scaled_mm_c3x_sm120.cu"
-      "csrc/quantization/cutlass_w8a8/c3x/scaled_mm_sm120_fp8.cu"
-      "csrc/quantization/cutlass_w8a8/c3x/scaled_mm_blockwise_sm120_fp8.cu"
+      "csrc/quantization/w8a8/cutlass/scaled_mm_c3x_sm120.cu"
+      "csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8.cu"
+      "csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8.cu"
     )
     set_gencode_flags_for_srcs(
       SRCS "${SRCS}"
@@ -473,9 +461,9 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
   cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.1a" "${CUDA_ARCHS}")
   if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS)
     set(SRCS
-      "csrc/quantization/cutlass_w8a8/scaled_mm_c3x_sm100.cu"
-      "csrc/quantization/cutlass_w8a8/c3x/scaled_mm_sm100_fp8.cu"
-      "csrc/quantization/cutlass_w8a8/c3x/scaled_mm_blockwise_sm100_fp8.cu"
+      "csrc/quantization/w8a8/cutlass/scaled_mm_c3x_sm100.cu"
+      "csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8.cu"
+      "csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm100_fp8.cu"
     )
     set_gencode_flags_for_srcs(
       SRCS "${SRCS}"
@@ -506,7 +494,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
   # subtract out the archs that are already built for 3x
   list(REMOVE_ITEM SCALED_MM_2X_ARCHS ${SCALED_MM_3X_ARCHS})
   if (SCALED_MM_2X_ARCHS)
-    set(SRCS "csrc/quantization/cutlass_w8a8/scaled_mm_c2x.cu")
+    set(SRCS "csrc/quantization/w8a8/cutlass/scaled_mm_c2x.cu")
     set_gencode_flags_for_srcs(
       SRCS "${SRCS}"
       CUDA_ARCHS "${SCALED_MM_2X_ARCHS}")
@@ -594,6 +582,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}"
@@ -617,7 +606,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
   # if it's possible to compile MoE kernels that use its output.
   cuda_archs_loose_intersection(SCALED_MM_ARCHS "9.0a" "${CUDA_ARCHS}")
   if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3 AND SCALED_MM_ARCHS)
-    set(SRCS "csrc/quantization/cutlass_w8a8/moe/grouped_mm_c3x_sm90.cu")
+    set(SRCS "csrc/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm90.cu")
     set_gencode_flags_for_srcs(
       SRCS "${SRCS}"
       CUDA_ARCHS "${SCALED_MM_ARCHS}")
@@ -637,7 +626,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
 
   cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a" "${CUDA_ARCHS}")
   if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS)
-    set(SRCS "csrc/quantization/cutlass_w8a8/moe/grouped_mm_c3x_sm100.cu")
+    set(SRCS "csrc/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm100.cu")
     set_gencode_flags_for_srcs(
       SRCS "${SRCS}"
       CUDA_ARCHS "${SCALED_MM_ARCHS}")
@@ -658,7 +647,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
   # moe_data.cu is used by all CUTLASS MoE kernels.
   cuda_archs_loose_intersection(CUTLASS_MOE_DATA_ARCHS "9.0a;10.0a" "${CUDA_ARCHS}")
   if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3 AND CUTLASS_MOE_DATA_ARCHS)
-    set(SRCS "csrc/quantization/cutlass_w8a8/moe/moe_data.cu")
+    set(SRCS "csrc/quantization/w8a8/cutlass/moe/moe_data.cu")
     set_gencode_flags_for_srcs(
       SRCS "${SRCS}"
       CUDA_ARCHS "${CUTLASS_MOE_DATA_ARCHS}")
@@ -677,7 +666,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
 
   cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a" "${CUDA_ARCHS}")
   if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS)
-    set(SRCS "csrc/quantization/cutlass_w8a8/moe/blockwise_scaled_group_mm_sm100.cu")
+    set(SRCS "csrc/quantization/w8a8/cutlass/moe/blockwise_scaled_group_mm_sm100.cu")
     set_gencode_flags_for_srcs(
       SRCS "${SRCS}"
       CUDA_ARCHS "${SCALED_MM_ARCHS}")
@@ -791,17 +780,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()
 

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_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/cutlass_benchmarks/w8a8_benchmarks.py

@@ -17,7 +17,7 @@ from weight_shapes import WEIGHT_SHAPES
 
 from vllm import _custom_ops as ops
 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
-    w8a8_triton_block_scaled_mm,
+    w8a8_block_fp8_matmul,
 )
 from vllm.utils import FlexibleArgumentParser, cdiv
 
@@ -158,7 +158,7 @@ def bench_fp8(
         "cutlass_fp8_fp8_fp16_scaled_mm_bias": lambda: ops.cutlass_scaled_mm(
             a, b, scale_a, scale_b, torch.float16, bias.to(dtype=torch.float16)
         ),
-        "triton_fp8_fp8_fp16_scaled_mm_blockwise": lambda: w8a8_triton_block_scaled_mm(
+        "triton_fp8_fp8_fp16_scaled_mm_blockwise": lambda: w8a8_block_fp8_matmul(
             a_cont, b.t(), block_scale_a, block_scale_b.t(), (128, 128)
         ),
         "cutlass_fp8_fp8_fp16_scaled_mm_blockwise": lambda: ops.cutlass_scaled_mm(

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

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

benchmarks/kernels/deepgemm/benchmark_fp8_block_dense_gemm.py

@@ -10,7 +10,7 @@ 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_triton_block_scaled_mm,
+    w8a8_block_fp8_matmul,
 )
 from vllm.triton_utils import triton
 from vllm.utils.deep_gemm import calc_diff, fp8_gemm_nt, per_block_cast_to_fp8
@@ -59,7 +59,7 @@ def benchmark_shape(m: int,
 
     # === vLLM Triton Implementation ===
     def vllm_triton_gemm():
-        return w8a8_triton_block_scaled_mm(A_vllm,
+        return w8a8_block_fp8_matmul(A_vllm,
                                      B_vllm,
                                      A_scale_vllm,
                                      B_scale_vllm,

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

@@ -258,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/attention_kernels.cuh

@@ -28,10 +28,10 @@
 
 #ifdef USE_ROCM
   #include <hip/hip_bf16.h>
-  #include "../quantization/fp8/amd/quant_utils.cuh"
+  #include "../quantization/w8a8/fp8/amd/quant_utils.cuh"
 typedef __hip_bfloat16 __nv_bfloat16;
 #else
-  #include "../quantization/fp8/nvidia/quant_utils.cuh"
+  #include "../quantization/w8a8/fp8/nvidia/quant_utils.cuh"
 #endif
 
 #define MAX(a, b) ((a) > (b) ? (a) : (b))

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/cache_kernels.cu

@@ -9,9 +9,9 @@
 #include "quantization/vectorization_utils.cuh"
 
 #ifdef USE_ROCM
-  #include "quantization/fp8/amd/quant_utils.cuh"
+  #include "quantization/w8a8/fp8/amd/quant_utils.cuh"
 #else
-  #include "quantization/fp8/nvidia/quant_utils.cuh"
+  #include "quantization/w8a8/fp8/nvidia/quant_utils.cuh"
 #endif
 
 #include <algorithm>

csrc/cpu/cpu_types.hpp

@@ -17,8 +17,4 @@
   #warning "unsupported vLLM cpu implementation"
 #endif
 
-#ifdef _OPENMP
-  #include <omp.h>
-#endif
-
 #endif

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

@@ -6,13 +6,18 @@
  */
 
 #include "type_convert.cuh"
-#include "quantization/fp8/common.cuh"
+#include "quantization/w8a8/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/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

@@ -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/quantization/activation_kernels.cu

@@ -7,7 +7,7 @@
 #include "../cuda_compat.h"
 #include "dispatch_utils.h"
 
-#include "quantization/fp8/common.cuh"
+#include "quantization/w8a8/fp8/common.cuh"
 
 #include <c10/util/Float8_e4m3fn.h>
 
@@ -23,14 +23,9 @@
 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"
@@ -370,6 +365,7 @@ __global__ void silu_mul_fp8_quant_deep_gemm_kernel(
   int32_t compute_pipeline_offset_64 = 0;
 
   for (int32_t t = n_tokens_lower; t < n_tokens_upper; ++t) {
+    __nv_bfloat16 y_max_bf16 = EPS;
     __nv_bfloat162 results_bf162[2];
 
     cp_async_wait<NUM_STAGES - 2>();
@@ -409,7 +405,7 @@ __global__ void silu_mul_fp8_quant_deep_gemm_kernel(
     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));
+    y_max_bf16 = __hmax(_y_max2.x, _y_max2.y);
 
     // An entire group is assigned to a single warp, so a simple warp reduce
     // is used.

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/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/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/fused_kernels/quant_conversions.cuh

@@ -6,7 +6,7 @@
 
 #include "quantization/vectorization.cuh"
 // TODO(luka/varun):refactor common.cuh to use this file instead
-#include "quantization/fp8/common.cuh"
+#include "quantization/w8a8/fp8/common.cuh"
 
 namespace vllm {
 

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/quantization/w8a8/cutlass/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/w8a8/cutlass/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/w8a8/cutlass/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/w8a8/fp8/amd/quant_utils.cuh

@@ -5,7 +5,7 @@
 #include <hip/hip_bf16.h>
 #include <hip/hip_bfloat16.h>
 
-#include "../../../attention/attention_dtypes.h"
+#include "../../../../attention/attention_dtypes.h"
 
 namespace vllm {
 #ifdef USE_ROCM

csrc/quantization/w8a8/fp8/common.cu

@@ -1,10 +1,15 @@
 #include "common.cuh"
 #include "dispatch_utils.h"
-#include "../../cub_helpers.h"
-#include "../vectorization_utils.cuh"
+#include "quantization/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) {