Sage Moore fixes for full cuda graph support for DeepEP+DeepGEMM LL

Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com>

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

@@ -51,7 +51,6 @@ function cpu_tests() {
     pytest -v -s tests/kernels/attention/test_cache.py -m cpu_model
     pytest -v -s tests/kernels/attention/test_mla_decode_cpu.py -m cpu_model
     pytest -v -s tests/models/language/generation -m cpu_model
-    VLLM_CPU_SGL_KERNEL=1 pytest -v -s tests/models/language/generation -m cpu_model
     pytest -v -s tests/models/language/pooling -m cpu_model
     pytest -v -s tests/models/multimodal/generation \
                 --ignore=tests/models/multimodal/generation/test_mllama.py \
@@ -99,4 +98,4 @@ function cpu_tests() {
 
 # All of CPU tests are expected to be finished less than 40 mins.
 export -f cpu_tests
-timeout 1.5h bash -c "cpu_tests $CORE_RANGE $NUMA_NODE"
+timeout 1h bash -c "cpu_tests $CORE_RANGE $NUMA_NODE"

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

@@ -2,34 +2,10 @@
 
 # This script build the CPU docker image and run the offline inference inside the container.
 # It serves a sanity check for compilation and basic model usage.
-set -exuo pipefail
+set -ex
 
 # Try building the docker image
-cat <<EOF | docker build -t hpu-plugin-v1-test-env -f - .
-FROM 1.22-413-pt2.7.1:latest
-
-COPY ./ /workspace/vllm
-
-WORKDIR /workspace/vllm
-
-RUN pip install -v -r requirements/hpu.txt
-RUN pip install git+https://github.com/vllm-project/vllm-gaudi.git
-
-ENV no_proxy=localhost,127.0.0.1
-ENV PT_HPU_ENABLE_LAZY_COLLECTIVES=true
-
-RUN VLLM_TARGET_DEVICE=hpu python3 setup.py install
-
-# install development dependencies (for testing)
-RUN python3 -m pip install -e tests/vllm_test_utils
-
-WORKDIR /workspace/
-
-RUN git clone https://github.com/vllm-project/vllm-gaudi.git
-
-RUN ln -s /workspace/vllm/tests && ln -s /workspace/vllm/examples && ln -s /workspace/vllm/benchmarks
-
-EOF
+docker build -t hpu-test-env -f docker/Dockerfile.hpu .
 
 # Setup cleanup
 # certain versions of HPU software stack have a bug that can
@@ -38,21 +14,13 @@ EOF
 # functions, while other platforms only need one remove_docker_container
 # function.
 EXITCODE=1
-remove_docker_containers() { docker rm -f hpu-plugin-v1-test || true; }
-trap 'remove_docker_containers; exit $EXITCODE;' EXIT
+remove_docker_containers() { docker rm -f hpu-test || true; docker rm -f hpu-test-tp2 || true; }
+remove_docker_containers_and_exit() { remove_docker_containers; exit $EXITCODE; }
+trap remove_docker_containers_and_exit EXIT
 remove_docker_containers
 
-echo "Running HPU plugin v1 test"
-docker run --rm --runtime=habana --name=hpu-plugin-v1-test --network=host \
-  -e HABANA_VISIBLE_DEVICES=all \
-  hpu-plugin-v1-test-env \
-  /bin/bash "/workspace/vllm-gaudi/tests/upstream_tests/ci_tests.sh"
+# Run the image and launch offline inference
+docker run --runtime=habana --name=hpu-test --network=host -e HABANA_VISIBLE_DEVICES=all -e VLLM_SKIP_WARMUP=true --entrypoint="" hpu-test-env python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m
+docker run --runtime=habana --name=hpu-test-tp2 --network=host -e HABANA_VISIBLE_DEVICES=all -e VLLM_SKIP_WARMUP=true --entrypoint="" hpu-test-env python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --tensor-parallel-size 2
 
 EXITCODE=$?
-if [ $EXITCODE -eq 0 ]; then
-  echo "Test with basic model passed"
-else
-  echo "Test with basic model FAILED with exit code: $EXITCODE" >&2
-fi
-
-# The trap will handle the container removal and final exit.

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

@@ -159,8 +159,6 @@ run_and_track_test 14 "test_tpu_qkv_linear.py" \
     "python3 -m pytest -s -v /workspace/vllm/tests/v1/tpu/test_tpu_qkv_linear.py"
 run_and_track_test 15 "test_spmd_model_weight_loading.py" \
     "python3 -m pytest -s -v /workspace/vllm/tests/v1/tpu/test_spmd_model_weight_loading.py"
-run_and_track_test 16 "test_kv_cache_update_kernel.py" \
-    "python3 -m pytest -s -v /workspace/vllm/tests/v1/tpu/test_kv_cache_update_kernel.py"
 
 # After all tests have been attempted, exit with the overall status.
 if [ "$overall_script_exit_code" -ne 0 ]; then

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

@@ -28,5 +28,4 @@ docker run \
     sh -c '
     VLLM_USE_V1=0 python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m
     VLLM_USE_V1=0 python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m -tp 2
-    VLLM_USE_V1=1 python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager
 '

.buildkite/scripts/tpu/docker_run_bm.sh

@@ -68,7 +68,7 @@ docker run \
 
 echo "run script..."
 echo
-docker exec "$CONTAINER_NAME" /bin/bash -c ".buildkite/scripts/tpu/run_bm.sh"
+docker exec "$CONTAINER_NAME" /bin/bash -c ".buildkite/scripts/hardware_ci/run_bm.sh"
 
 echo "copy result back..."
 VLLM_LOG="$LOG_ROOT/$TEST_NAME"_vllm_log.txt

.buildkite/test-pipeline.yaml

@@ -41,16 +41,6 @@ steps:
   # TODO: add `--strict` once warnings in docstrings are fixed
   - mkdocs build
 
-- label: Pytorch Nightly Dependency Override Check # 2min
-  # if this test fails, it means the nightly torch version is not compatible with some
-  # of the dependencies. Please check the error message and add the package to whitelist
-  # in /vllm/tools/generate_nightly_torch_test.py
-  soft_fail: true
-  source_file_dependencies:
-  - requirements/nightly_torch_test.txt
-  commands:
-  - bash standalone_tests/pytorch_nightly_dependency.sh
-
 - label: Async Engine, Inputs, Utils, Worker Test # 24min
   mirror_hardwares: [amdexperimental]
   source_file_dependencies:
@@ -178,23 +168,6 @@ steps:
   - VLLM_ALLOW_INSECURE_SERIALIZATION=1 RAY_DEDUP_LOGS=0 python3 rlhf_colocate.py
   - popd
 
-- label: EPLB Algorithm Test
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/distributed/eplb
-  - tests/distributed/test_eplb_algo.py
-  commands:
-  - pytest -v -s distributed/test_eplb_algo.py
-
-- label: EPLB Execution Test # 5min
-  working_dir: "/vllm-workspace/tests"
-  num_gpus: 4
-  source_file_dependencies:
-  - vllm/distributed/eplb
-  - tests/distributed/test_eplb_execute.py
-  commands:
-  - pytest -v -s distributed/test_eplb_execute.py
-
 - label: Metrics, Tracing Test # 10min
   mirror_hardwares: [amdexperimental, amdproduction]
   num_gpus: 2
@@ -536,17 +509,6 @@ steps:
     - pip freeze | grep -E 'torch'
     - pytest -v -s models/language -m core_model
 
-- label: Language Models Test (Hybrid) # 35 min
-  mirror_hardwares: [amdexperimental]
-  torch_nightly: true
-  source_file_dependencies:
-  - vllm/
-  - tests/models/language/generation
-  commands:
-    # Install causal-conv1d for plamo2 models here, as it is not compatible with pip-compile.
-    - pip install 'git+https://github.com/Dao-AILab/causal-conv1d@v1.5.0.post8'
-    - pytest -v -s models/language/generation -m hybrid_model
-
 - label: Language Models Test (Extended Generation) # 1hr20min
   mirror_hardwares: [amdexperimental]
   optional: true
@@ -556,7 +518,7 @@ steps:
   commands:
     # Install causal-conv1d for plamo2 models here, as it is not compatible with pip-compile.
     - pip install 'git+https://github.com/Dao-AILab/causal-conv1d@v1.5.0.post8'
-    - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'
+    - pytest -v -s models/language/generation -m 'not core_model'
 
 - label: Language Models Test (Extended Pooling)  # 36min
   mirror_hardwares: [amdexperimental]
@@ -653,18 +615,13 @@ steps:
   - vllm/executor/
   - vllm/model_executor/models/
   - tests/distributed/
-  - tests/examples/offline_inference/data_parallel.py
   commands:
   - # the following commands are for the first node, with ip 192.168.10.10 (ray environment already set up)
     - VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed'
-    - NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed'
-    - python3 ../examples/offline_inference/data_parallel.py --dp-size=2 --tp-size=1 --node-size=2 --node-rank=0 --master-addr=192.168.10.10 --master-port=12345 --enforce-eager --trust-remote-code
     - VLLM_MULTI_NODE=1 pytest -v -s distributed/test_multi_node_assignment.py
     - VLLM_MULTI_NODE=1 pytest -v -s distributed/test_pipeline_parallel.py
   - # the following commands are for the second node, with ip 192.168.10.11 (ray environment already set up)
     - VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed'
-    - NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed'
-    - python3 ../examples/offline_inference/data_parallel.py --dp-size=2 --tp-size=1 --node-size=2 --node-rank=1 --master-addr=192.168.10.10 --master-port=12345 --enforce-eager --trust-remote-code
 
 - label: Distributed Tests (2 GPUs) # 40min
   mirror_hardwares: [amdexperimental]
@@ -788,7 +745,7 @@ steps:
     - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models.txt
 
 - label: Weight Loading Multiple GPU Test - Large Models # optional
-  mirror_hardwares: [amdexperimental]
+  mirror_hardwares: [amdexperimental] 
   working_dir: "/vllm-workspace/tests"
   num_gpus: 2
   gpu: a100

.github/CODEOWNERS

@@ -18,10 +18,6 @@
 /vllm/entrypoints @aarnphm
 CMakeLists.txt @tlrmchlsmth
 
-# Any change to the VllmConfig changes can have a large user-facing impact,
-# so spam a lot of people
-/vllm/config.py @simon-mo @WoosukKwon @youkaichao @robertgshaw2-redhat @mgoin @tlrmchlsmth @houseroad @hmellor
-
 # vLLM V1
 /vllm/v1 @WoosukKwon @robertgshaw2-redhat @njhill @ywang96 @comaniac @alexm-redhat
 /vllm/v1/structured_output @mgoin @russellb @aarnphm

.github/mergify.yml

@@ -27,22 +27,6 @@ pull_request_rules:
       add:
         - ci/build
 
-- name: label-deepseek
-  description: Automatically apply deepseek label
-  conditions:
-    - or:
-      - files~=^examples/.*deepseek.*\.py
-      - files~=^tests/.*deepseek.*\.py
-      - files~=^vllm/entrypoints/openai/tool_parsers/.*deepseek.*\.py
-      - files~=^vllm/model_executor/models/.*deepseek.*\.py
-      - files~=^vllm/reasoning/.*deepseek.*\.py
-      - files~=^vllm/transformers_utils/.*deepseek.*\.py
-      - title~=(?i)DeepSeek
-  actions:
-    label:
-      add:
-        - deepseek
-
 - name: label-frontend
   description: Automatically apply frontend label
   conditions:

.pre-commit-config.yaml

@@ -53,11 +53,6 @@ repos:
       files: ^requirements/test\.(in|txt)$
 - repo: local
   hooks:
-  - id: format-torch-nightly-test
-    name: reformat nightly_torch_test.txt to be in sync with test.in
-    language: python
-    entry: python tools/generate_nightly_torch_test.py
-    files: ^requirements/test\.(in|txt)$
   - id: mypy-local
     name: Run mypy for local Python installation
     entry: tools/mypy.sh 0 "local"
@@ -160,13 +155,6 @@ repos:
     types: [python]
     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
-    types: [python]
-    pass_filenames: true
-    files: vllm/config.py|tests/test_config.py
   # Keep `suggestion` last
   - id: suggestion
     name: Suggestion

CMakeLists.txt

@@ -420,36 +420,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
     endif()
   endif()
 
-
-  # The cutlass_scaled_mm kernels for Geforce Blackwell SM120 (c3x, i.e. CUTLASS 3.x) require
-  # CUDA 12.8 or later
-  cuda_archs_loose_intersection(SCALED_MM_ARCHS "12.0;12.0a" "${CUDA_ARCHS}")
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER 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"
-    )
-    set_gencode_flags_for_srcs(
-      SRCS "${SRCS}"
-      CUDA_ARCHS "${SCALED_MM_ARCHS}")
-    list(APPEND VLLM_EXT_SRC "${SRCS}")
-    list(APPEND VLLM_GPU_FLAGS "-DENABLE_SCALED_MM_SM120=1")
-    # Let scaled_mm_c2x know it doesn't need to build these arches
-    list(APPEND SCALED_MM_3X_ARCHS "${SCALED_MM_ARCHS}")
-    message(STATUS "Building scaled_mm_c3x_sm120 for archs: ${SCALED_MM_ARCHS}")
-  else()
-    if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER 12.8 AND SCALED_MM_ARCHS)
-      message(STATUS "Not building scaled_mm_c3x_sm120 as CUDA Compiler version is "
-                     "not >= 12.8, we recommend upgrading to CUDA 12.8 or "
-                     "later if you intend on running FP8 quantized models on "
-                     "Blackwell.")
-    else()
-      message(STATUS "Not building scaled_mm_c3x_120 as no compatible archs found "
-                     "in CUDA target architectures")
-    endif()
-  endif()
-
-
   # The cutlass_scaled_mm kernels for Blackwell SM100 (c3x, i.e. CUTLASS 3.x)
   # require CUDA 12.8 or later
   cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.1a" "${CUDA_ARCHS}")
@@ -543,7 +513,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
       CUDA_ARCHS "${FP4_ARCHS}")
     list(APPEND VLLM_EXT_SRC "${SRCS}")
     list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4=1")
-    list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
     message(STATUS "Building NVFP4 for archs: ${FP4_ARCHS}")
   else()
     message(STATUS "Not building NVFP4 as no compatible archs were found.")
@@ -578,7 +547,8 @@ 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.cu")
+    set(SRCS "csrc/quantization/cutlass_w8a8/moe/grouped_mm_c3x.cu"
+             "csrc/quantization/cutlass_w8a8/moe/moe_data.cu")
     set_gencode_flags_for_srcs(
       SRCS "${SRCS}"
       CUDA_ARCHS "${SCALED_MM_ARCHS}")
@@ -592,27 +562,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
                      "if you intend on running FP8 quantized MoE models on Hopper.")
     else()
       message(STATUS "Not building grouped_mm_c3x as no compatible archs found "
-                     "in CUDA target architectures.")
-    endif()
-  endif()
-
-  # 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_gencode_flags_for_srcs(
-      SRCS "${SRCS}"
-      CUDA_ARCHS "${CUTLASS_MOE_DATA_ARCHS}")
-    list(APPEND VLLM_EXT_SRC "${SRCS}")
-    message(STATUS "Building moe_data for archs: ${CUTLASS_MOE_DATA_ARCHS}")
-  else()
-    if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3 AND CUTLASS_MOE_DATA_ARCHS)
-      message(STATUS "Not building moe_data as CUDA Compiler version is "
-                     "not >= 12.3, we recommend upgrading to CUDA 12.3 or later "
-                     "if you intend on running FP8 quantized MoE models on Hopper or Blackwell.")
-    else()
-      message(STATUS "Not building moe_data as no compatible archs found "
-                     "in CUDA target architectures.")
+                     "in CUDA target architectures")
     endif()
   endif()
 
@@ -688,14 +638,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
 # if CUDA endif
 endif()
 
-if (VLLM_GPU_LANG STREQUAL "HIP")
-  # Add QuickReduce kernels
-  list(APPEND VLLM_EXT_SRC
-    "csrc/custom_quickreduce.cu"
-  )
-# if ROCM endif
-endif()
-
 message(STATUS "Enabling C extension.")
 define_gpu_extension_target(
   _C

benchmarks/README.md

@@ -4,7 +4,7 @@ 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.
 
-**Dataset Overview**
+## Dataset Overview
 
 <table style="width:100%; border-collapse: collapse;">
   <thead>
@@ -82,10 +82,7 @@ become available.
 **Note**: HuggingFace dataset's `dataset-name` should be set to `hf`
 
 ---
-<details>
-<summary><b>🚀 Example - Online Benchmark</b></summary>
-
-<br/>
+## Example - Online Benchmark
 
 First start serving your model
 
@@ -133,8 +130,7 @@ P99 ITL (ms):                            8.39
 ==================================================
 ```
 
-**Custom Dataset**
-
+### 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
 
 ```
@@ -166,7 +162,7 @@ python3 benchmarks/benchmark_serving.py --port 9001 --save-result --save-detaile
 
 You can skip applying chat template if your data already has it by using `--custom-skip-chat-template`.
 
-**VisionArena Benchmark for Vision Language Models**
+### VisionArena Benchmark for Vision Language Models
 
 ```bash
 # need a model with vision capability here
@@ -184,7 +180,7 @@ python3 vllm/benchmarks/benchmark_serving.py \
   --num-prompts 1000
 ```
 
-**InstructCoder Benchmark with Speculative Decoding**
+### InstructCoder Benchmark with Speculative Decoding
 
 ``` bash
 VLLM_USE_V1=1 vllm serve meta-llama/Meta-Llama-3-8B-Instruct \
@@ -201,7 +197,7 @@ python3 benchmarks/benchmark_serving.py \
     --num-prompts 2048
 ```
 
-**Other HuggingFaceDataset Examples**
+### Other HuggingFaceDataset Examples
 
 ```bash
 vllm serve Qwen/Qwen2-VL-7B-Instruct --disable-log-requests
@@ -255,7 +251,7 @@ python3 vllm/benchmarks/benchmark_serving.py \
     --num-prompts 80
 ```
 
-**Running With Sampling Parameters**
+### Running With Sampling Parameters
 
 When using OpenAI-compatible backends such as `vllm`, optional sampling
 parameters can be specified. Example client command:
@@ -273,27 +269,8 @@ python3 vllm/benchmarks/benchmark_serving.py \
   --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>
-
-<details>
-<summary><b>📈 Example - Offline Throughput Benchmark</b></summary>
-
-<br/>
+---
+## Example - Offline Throughput Benchmark
 
 ```bash
 python3 vllm/benchmarks/benchmark_throughput.py \
@@ -311,7 +288,7 @@ Total num prompt tokens:  5014
 Total num output tokens:  1500
 ```
 
-**VisionArena Benchmark for Vision Language Models**
+### VisionArena Benchmark for Vision Language Models
 
 ``` bash
 python3 vllm/benchmarks/benchmark_throughput.py \
@@ -331,7 +308,7 @@ Total num prompt tokens:  14527
 Total num output tokens:  1280
 ```
 
-**InstructCoder Benchmark with Speculative Decoding**
+### InstructCoder Benchmark with Speculative Decoding
 
 ``` bash
 VLLM_WORKER_MULTIPROC_METHOD=spawn \
@@ -355,7 +332,7 @@ Total num prompt tokens:  261136
 Total num output tokens:  204800
 ```
 
-**Other HuggingFaceDataset Examples**
+### Other HuggingFaceDataset Examples
 
 **`lmms-lab/LLaVA-OneVision-Data`**
 
@@ -394,7 +371,7 @@ python3 benchmarks/benchmark_throughput.py \
   --num-prompts 10
 ```
 
-**Benchmark with LoRA Adapters**
+### Benchmark with LoRA Adapters
 
 ``` bash
 # download dataset
@@ -411,22 +388,18 @@ python3 vllm/benchmarks/benchmark_throughput.py \
   --lora-path yard1/llama-2-7b-sql-lora-test
   ```
 
-</details>
-
-<details>
-<summary><b>🛠️ Example - Structured Output Benchmark</b></summary>
-
-<br/>
+---
+## Example - Structured Output Benchmark
 
 Benchmark the performance of structured output generation (JSON, grammar, regex).
 
-**Server Setup**
+### Server Setup
 
 ```bash
 vllm serve NousResearch/Hermes-3-Llama-3.1-8B --disable-log-requests
 ```
 
-**JSON Schema Benchmark**
+### JSON Schema Benchmark
 
 ```bash
 python3 benchmarks/benchmark_serving_structured_output.py \
@@ -438,7 +411,7 @@ python3 benchmarks/benchmark_serving_structured_output.py \
   --num-prompts 1000
 ```
 
-**Grammar-based Generation Benchmark**
+### Grammar-based Generation Benchmark
 
 ```bash
 python3 benchmarks/benchmark_serving_structured_output.py \
@@ -450,7 +423,7 @@ python3 benchmarks/benchmark_serving_structured_output.py \
   --num-prompts 1000
 ```
 
-**Regex-based Generation Benchmark**
+### Regex-based Generation Benchmark
 
 ```bash
 python3 benchmarks/benchmark_serving_structured_output.py \
@@ -461,7 +434,7 @@ python3 benchmarks/benchmark_serving_structured_output.py \
   --num-prompts 1000
 ```
 
-**Choice-based Generation Benchmark**
+### Choice-based Generation Benchmark
 
 ```bash
 python3 benchmarks/benchmark_serving_structured_output.py \
@@ -472,7 +445,7 @@ python3 benchmarks/benchmark_serving_structured_output.py \
   --num-prompts 1000
 ```
 
-**XGrammar Benchmark Dataset**
+### XGrammar Benchmark Dataset
 
 ```bash
 python3 benchmarks/benchmark_serving_structured_output.py \
@@ -483,16 +456,12 @@ python3 benchmarks/benchmark_serving_structured_output.py \
   --num-prompts 1000
 ```
 
-</details>
-
-<details>
-<summary><b>📚 Example - Long Document QA Benchmark</b></summary>
-
-<br/>
+---
+## Example - Long Document QA Throughput Benchmark
 
 Benchmark the performance of long document question-answering with prefix caching.
 
-**Basic Long Document QA Test**
+### Basic Long Document QA Test
 
 ```bash
 python3 benchmarks/benchmark_long_document_qa_throughput.py \
@@ -504,7 +473,7 @@ python3 benchmarks/benchmark_long_document_qa_throughput.py \
   --repeat-count 5
 ```
 
-**Different Repeat Modes**
+### Different Repeat Modes
 
 ```bash
 # Random mode (default) - shuffle prompts randomly
@@ -535,16 +504,12 @@ python3 benchmarks/benchmark_long_document_qa_throughput.py \
   --repeat-mode interleave
 ```
 
-</details>
-
-<details>
-<summary><b>🗂️ Example - Prefix Caching Benchmark</b></summary>
-
-<br/>
+---
+## Example - Prefix Caching Benchmark
 
 Benchmark the efficiency of automatic prefix caching.
 
-**Fixed Prompt with Prefix Caching**
+### Fixed Prompt with Prefix Caching
 
 ```bash
 python3 benchmarks/benchmark_prefix_caching.py \
@@ -555,7 +520,7 @@ python3 benchmarks/benchmark_prefix_caching.py \
   --input-length-range 128:256
 ```
 
-**ShareGPT Dataset with Prefix Caching**
+### ShareGPT Dataset with Prefix Caching
 
 ```bash
 # download dataset
@@ -570,16 +535,12 @@ python3 benchmarks/benchmark_prefix_caching.py \
   --input-length-range 128:256
 ```
 
-</details>
-
-<details>
-<summary><b>⚡ Example - Request Prioritization Benchmark</b></summary>
-
-<br/>
+---
+## Example - Request Prioritization Benchmark
 
 Benchmark the performance of request prioritization in vLLM.
 
-**Basic Prioritization Test**
+### Basic Prioritization Test
 
 ```bash
 python3 benchmarks/benchmark_prioritization.py \
@@ -590,7 +551,7 @@ python3 benchmarks/benchmark_prioritization.py \
   --scheduling-policy priority
 ```
 
-**Multiple Sequences per Prompt**
+### Multiple Sequences per Prompt
 
 ```bash
 python3 benchmarks/benchmark_prioritization.py \
@@ -601,5 +562,3 @@ python3 benchmarks/benchmark_prioritization.py \
   --scheduling-policy priority \
   --n 2
 ```
-
-</details>

benchmarks/benchmark_dataset.py

@@ -349,9 +349,8 @@ class RandomDataset(BenchmarkDataset):
             # [1650, 939, 486] -> ['Ġcall', 'sh', 'ere']
             # To avoid uncontrolled change of the prompt length,
             # the encoded sequence is truncated before being decode again.
-            total_input_len = prefix_len + int(input_lens[i])
             re_encoded_sequence = tokenizer.encode(prompt, add_special_tokens=False)[
-                :total_input_len
+                : input_lens[i]
             ]
             prompt = tokenizer.decode(re_encoded_sequence)
             total_input_len = len(re_encoded_sequence)

benchmarks/benchmark_serving.py

@@ -33,7 +33,7 @@ import warnings
 from collections.abc import AsyncGenerator, Iterable
 from dataclasses import dataclass
 from datetime import datetime
-from typing import Any, Literal, Optional
+from typing import Any, Optional
 
 import numpy as np
 from tqdm.asyncio import tqdm
@@ -107,42 +107,14 @@ class BenchmarkMetrics:
     percentiles_e2el_ms: list[tuple[float, float]]
 
 
-def _get_current_request_rate(
-    ramp_up_strategy: Optional[Literal["linear", "exponential"]],
-    ramp_up_start_rps: Optional[int],
-    ramp_up_end_rps: Optional[int],
-    request_index: int,
-    total_requests: int,
-    request_rate: float,
-) -> float:
-    if (
-        ramp_up_strategy
-        and ramp_up_start_rps is not None
-        and ramp_up_end_rps is not None
-    ):
-        progress = request_index / max(total_requests - 1, 1)
-        if ramp_up_strategy == "linear":
-            increase = (ramp_up_end_rps - ramp_up_start_rps) * progress
-            return ramp_up_start_rps + increase
-        elif ramp_up_strategy == "exponential":
-            ratio = ramp_up_end_rps / ramp_up_start_rps
-            return ramp_up_start_rps * (ratio**progress)
-        else:
-            raise ValueError(f"Unknown ramp-up strategy: {ramp_up_strategy}")
-    return request_rate
-
-
 async def get_request(
     input_requests: list[SampleRequest],
     request_rate: float,
     burstiness: float = 1.0,
-    ramp_up_strategy: Optional[Literal["linear", "exponential"]] = None,
-    ramp_up_start_rps: Optional[int] = None,
-    ramp_up_end_rps: Optional[int] = None,
-) -> AsyncGenerator[tuple[SampleRequest, float], None]:
+) -> AsyncGenerator[SampleRequest, None]:
     """
     Asynchronously generates requests at a specified rate
-    with OPTIONAL burstiness and OPTIONAL ramp-up strategy.
+    with OPTIONAL burstiness.
 
     Args:
         input_requests:
@@ -157,44 +129,22 @@ async def get_request(
             A lower burstiness value (0 < burstiness < 1) results
             in more bursty requests, while a higher burstiness value
             (burstiness > 1) results in a more uniform arrival of requests.
-         ramp_up_strategy (optional):
-            The ramp-up strategy. Can be "linear" or "exponential".
-            If None, uses constant request rate (specified by request_rate).
-        ramp_up_start_rps (optional):
-            The starting request rate for ramp-up.
-        ramp_up_end_rps (optional):
-            The ending request rate for ramp-up.
     """
+    input_requests: Iterable[SampleRequest] = iter(input_requests)
+
+    # Calculate scale parameter theta to maintain the desired request_rate.
     assert burstiness > 0, (
         f"A positive burstiness factor is expected, but given {burstiness}."
     )
-    # Convert to list to get length for ramp-up calculations
-    if isinstance(input_requests, Iterable) and not isinstance(input_requests, list):
-        input_requests = list(input_requests)
-
-    total_requests = len(input_requests)
-    request_index = 0
+    theta = 1.0 / (request_rate * burstiness)
 
     for request in input_requests:
-        current_request_rate = _get_current_request_rate(
-            ramp_up_strategy,
-            ramp_up_start_rps,
-            ramp_up_end_rps,
-            request_index,
-            total_requests,
-            request_rate,
-        )
+        yield request
 
-        yield request, current_request_rate
-
-        request_index += 1
-
-        if current_request_rate == float("inf"):
+        if request_rate == float("inf"):
             # If the request rate is infinity, then we don't need to wait.
             continue
 
-        theta = 1.0 / (current_request_rate * burstiness)
-
         # Sample the request interval from the gamma distribution.
         # If burstiness is 1, it follows exponential distribution.
         interval = np.random.gamma(shape=burstiness, scale=theta)
@@ -340,9 +290,6 @@ async def benchmark(
     max_concurrency: Optional[int],
     lora_modules: Optional[Iterable[str]],
     extra_body: Optional[dict],
-    ramp_up_strategy: Optional[Literal["linear", "exponential"]] = None,
-    ramp_up_start_rps: Optional[int] = None,
-    ramp_up_end_rps: Optional[int] = None,
 ):
     if backend in ASYNC_REQUEST_FUNCS:
         request_func = ASYNC_REQUEST_FUNCS[backend]
@@ -406,15 +353,7 @@ async def benchmark(
 
     distribution = "Poisson process" if burstiness == 1.0 else "Gamma distribution"
 
-    if ramp_up_strategy is not None:
-        print(
-            f"Traffic ramp-up strategy: {ramp_up_strategy}. Will increase "
-            f"RPS from {ramp_up_start_rps} to {ramp_up_end_rps} RPS over "
-            "the duration of the benchmark."
-        )
-    else:
-        print(f"Traffic request rate: {request_rate} RPS.")
-
+    print(f"Traffic request rate: {request_rate}")
     print(f"Burstiness factor: {burstiness} ({distribution})")
     print(f"Maximum request concurrency: {max_concurrency}")
 
@@ -434,34 +373,7 @@ async def benchmark(
 
     benchmark_start_time = time.perf_counter()
     tasks: list[asyncio.Task] = []
-
-    rps_change_events = []
-    last_int_rps = -1
-    if ramp_up_strategy is not None and ramp_up_start_rps is not None:
-        last_int_rps = ramp_up_start_rps
-        rps_change_events.append(
-            {
-                "rps": last_int_rps,
-                "timestamp": datetime.now().isoformat(),
-            }
-        )
-
-    async for request, current_request_rate in get_request(
-        input_requests,
-        request_rate,
-        burstiness,
-        ramp_up_strategy,
-        ramp_up_start_rps,
-        ramp_up_end_rps,
-    ):
-        if ramp_up_strategy is not None:
-            current_int_rps = int(current_request_rate)
-            if current_int_rps > last_int_rps:
-                timestamp = datetime.now().isoformat()
-                for rps_val in range(last_int_rps + 1, current_int_rps + 1):
-                    rps_change_events.append({"rps": rps_val, "timestamp": timestamp})
-                last_int_rps = current_int_rps
-
+    async for request in get_request(input_requests, request_rate, burstiness):
         prompt, prompt_len, output_len, mm_content = (
             request.prompt,
             request.prompt_len,
@@ -485,8 +397,11 @@ async def benchmark(
             ignore_eos=ignore_eos,
             extra_body=extra_body,
         )
-        task = limited_request_func(request_func_input=request_func_input, pbar=pbar)
-        tasks.append(asyncio.create_task(task))
+        tasks.append(
+            asyncio.create_task(
+                limited_request_func(request_func_input=request_func_input, pbar=pbar)
+            )
+        )
     outputs: list[RequestFuncOutput] = await asyncio.gather(*tasks)
 
     if profile:
@@ -551,7 +466,7 @@ async def benchmark(
         "total_input_tokens": metrics.total_input,
         "total_output_tokens": metrics.total_output,
         "request_throughput": metrics.request_throughput,
-        "request_goodput": metrics.request_goodput if goodput_config_dict else None,
+        "request_goodput:": metrics.request_goodput if goodput_config_dict else None,
         "output_throughput": metrics.output_throughput,
         "total_token_throughput": metrics.total_token_throughput,
         "input_lens": [output.prompt_len for output in outputs],
@@ -562,9 +477,6 @@ async def benchmark(
         "errors": [output.error for output in outputs],
     }
 
-    if rps_change_events:
-        result["rps_change_events"] = rps_change_events
-
     def process_one_metric(
         # E.g., "ttft"
         metric_attribute_name: str,
@@ -698,26 +610,6 @@ def main(args: argparse.Namespace):
     tokenizer_id = args.tokenizer if args.tokenizer is not None else args.model
     tokenizer_mode = args.tokenizer_mode
 
-    # Validate ramp-up arguments
-    if args.ramp_up_strategy is not None:
-        if args.request_rate != float("inf"):
-            raise ValueError(
-                "When using ramp-up, do not specify --request-rate. "
-                "The request rate will be controlled by ramp-up parameters. "
-                "Please remove the --request-rate argument."
-            )
-        if args.ramp_up_start_rps is None or args.ramp_up_end_rps is None:
-            raise ValueError(
-                "When using --ramp-up-strategy, both --ramp-up-start-rps and "
-                "--ramp-up-end-rps must be specified"
-            )
-        if args.ramp_up_start_rps < 0 or args.ramp_up_end_rps < 0:
-            raise ValueError("Ramp-up start and end RPS must be non-negative")
-        if args.ramp_up_start_rps > args.ramp_up_end_rps:
-            raise ValueError("Ramp-up start RPS must be less than end RPS")
-        if args.ramp_up_strategy == "exponential" and args.ramp_up_start_rps == 0:
-            raise ValueError("For exponential ramp-up, the start RPS cannot be 0.")
-
     if args.base_url is not None:
         api_url = f"{args.base_url}{args.endpoint}"
         base_url = f"{args.base_url}"
@@ -910,9 +802,6 @@ def main(args: argparse.Namespace):
             max_concurrency=args.max_concurrency,
             lora_modules=args.lora_modules,
             extra_body=sampling_params,
-            ramp_up_strategy=args.ramp_up_strategy,
-            ramp_up_start_rps=args.ramp_up_start_rps,
-            ramp_up_end_rps=args.ramp_up_end_rps,
         )
     )
 
@@ -945,11 +834,6 @@ def main(args: argparse.Namespace):
         result_json["burstiness"] = args.burstiness
         result_json["max_concurrency"] = args.max_concurrency
 
-        if args.ramp_up_strategy is not None:
-            result_json["ramp_up_strategy"] = args.ramp_up_strategy
-            result_json["ramp_up_start_rps"] = args.ramp_up_start_rps
-            result_json["ramp_up_end_rps"] = args.ramp_up_end_rps
-
         # Merge with benchmark result
         result_json = {**result_json, **benchmark_result}
 
@@ -975,10 +859,7 @@ def main(args: argparse.Namespace):
             if args.max_concurrency is not None
             else ""
         )
-        if args.ramp_up_strategy is not None:
-            file_name = f"{backend}-ramp-up-{args.ramp_up_strategy}-{args.ramp_up_start_rps}qps-{args.ramp_up_end_rps}qps{max_concurrency_str}-{base_model_id}-{current_dt}.json"  # noqa
-        else:
-            file_name = f"{backend}-{args.request_rate}qps{max_concurrency_str}-{base_model_id}-{current_dt}.json"  # noqa
+        file_name = f"{backend}-{args.request_rate}qps{max_concurrency_str}-{base_model_id}-{current_dt}.json"  # noqa
         if args.result_filename:
             file_name = args.result_filename
         if args.result_dir:
@@ -1344,31 +1225,6 @@ def create_argument_parser():
         "script chooses a LoRA module at random.",
     )
 
-    parser.add_argument(
-        "--ramp-up-strategy",
-        type=str,
-        default=None,
-        choices=["linear", "exponential"],
-        help="The ramp-up strategy. This would be used to "
-        "ramp up the request rate from initial RPS to final "
-        "RPS rate (specified by --ramp-up-start-rps and --ramp-up-end-rps). "
-        "over the duration of the benchmark.",
-    )
-    parser.add_argument(
-        "--ramp-up-start-rps",
-        type=int,
-        default=None,
-        help="The starting request rate for ramp-up (RPS). "
-        "Needs to be specified when --ramp-up-strategy is used.",
-    )
-    parser.add_argument(
-        "--ramp-up-end-rps",
-        type=int,
-        default=None,
-        help="The ending request rate for ramp-up (RPS). "
-        "Needs to be specified when --ramp-up-strategy is used.",
-    )
-
     return parser
 
 

benchmarks/cutlass_benchmarks/w8a8_benchmarks.py

@@ -19,7 +19,7 @@ from vllm import _custom_ops as ops
 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
     w8a8_block_fp8_matmul,
 )
-from vllm.utils import FlexibleArgumentParser, cdiv
+from vllm.utils import FlexibleArgumentParser
 
 DEFAULT_MODELS = list(WEIGHT_SHAPES.keys())
 DEFAULT_BATCH_SIZES = [1, 16, 32, 64, 128, 256, 512]
@@ -117,9 +117,14 @@ def bench_fp8(
     scale_a = torch.tensor(1.0, device="cuda", dtype=torch.float32)
     scale_b = torch.tensor(1.0, device="cuda", dtype=torch.float32)
 
-    block_scale_a = torch.rand((m, cdiv(k, 128)), device="cuda", dtype=torch.float32)
+    def ceil_div(x: int, y: int) -> int:
+        return (x + y - 1) // y
+
+    block_scale_a = torch.rand(
+        (m, ceil_div(k, 128)), device="cuda", dtype=torch.float32
+    )
     block_scale_b = torch.rand(
-        cdiv(k, 128), cdiv(n, 128), device="cuda", dtype=torch.float32
+        ceil_div(k, 128), ceil_div(n, 128), device="cuda", dtype=torch.float32
     )
     block_scale_a_M_major = block_scale_a.t().contiguous().t()
     block_scale_b_K_major = block_scale_b.t().contiguous().t()

benchmarks/kernels/benchmark_grouped_gemm_cutlass.py

@@ -113,7 +113,6 @@ def bench_run(
         w2_scale: torch.Tensor,
         topk_weights: torch.Tensor,
         topk_ids: torch.Tensor,
-        per_act_token: bool,
         num_repeats: int,
     ):
         for _ in range(num_repeats):
@@ -125,8 +124,7 @@ def bench_run(
                 topk_ids,
                 w1_scale,
                 w2_scale,
-                per_act_token,
-                a1_scale=None,
+                a1_scale=a_scale,
             )
 
     def run_cutlass_from_graph(
@@ -150,8 +148,7 @@ def bench_run(
                 topk_ids,
                 w1_scale,
                 w2_scale,
-                per_act_token,
-                a1_scale=None,
+                a1_scale=a_scale,
             )
 
     def run_triton_from_graph(
@@ -230,7 +227,6 @@ def bench_run(
         "w2_q": w2_q,
         "w1_scale": w1_scale,
         "w2_scale": w2_scale,
-        "per_act_token": per_act_token,
         # cuda graph params
         "cutlass_graph": cutlass_graph,
         "triton_graph": triton_graph,
@@ -291,13 +287,12 @@ def bench_run(
         w2_scale,
         topk_weights,
         topk_ids,
-        per_act_token,
         num_warmup,
     )
 
     results.append(
         benchmark.Timer(
-            stmt="run_cutlass_moe(a, a_scale, w1_q, w2_q, w1_scale, w2_scale, topk_weights, topk_ids, per_act_token, num_runs)",  # noqa: E501
+            stmt="run_cutlass_moe(a, a_scale, w1_q, w2_q, w1_scale, w2_scale, topk_weights, topk_ids, num_runs)",  # noqa: E501
             globals=globals,
             label=label,
             sub_label=sub_label,

benchmarks/kernels/benchmark_machete.py

@@ -234,10 +234,8 @@ def marlin_create_bench_fn(bt: BenchmarkTensors) -> Callable:
 
         fn = lambda: ops.gptq_marlin_gemm(
             a=bt.a,
-            c=None,
             b_q_weight=w_q,
             b_scales=w_s,
-            global_scale=None,
             b_zeros=w_zp,
             g_idx=g_idx,
             perm=sort_indices,

benchmarks/kernels/deepgemm/benchmark_fp8_block_dense_gemm.py

@@ -85,6 +85,12 @@ def benchmark_shape(m: int,
 
     # === DeepGEMM Implementation ===
     def deepgemm_gemm():
+        # A quantization is inside the loop as it depends on activations
+        # A_deepgemm, A_scale_deepgemm = per_token_cast_to_fp8(A)
+        # A_deepgemm, A_scale_deepgemm = per_token_group_quant_fp8(
+        #     A, block_size[1])
+        # A_scale_aligned = get_col_major_tma_aligned_tensor(A_scale_deepgemm)
+        # C_deepgemm = torch.empty((m, n), device='cuda', dtype=torch.bfloat16)
         deep_gemm.gemm_fp8_fp8_bf16_nt((A_deepgemm, A_scale_deepgemm),
                                        (B_deepgemm, B_scale_deepgemm),
                                        C_deepgemm)
@@ -92,6 +98,8 @@ def benchmark_shape(m: int,
 
     # === vLLM Triton Implementation ===
     def vllm_triton_gemm():
+        # A quantization is inside the loop as it depends on activations
+        # A_vllm, A_scale_vllm = per_token_group_quant_fp8(A, block_size[1])
         return w8a8_block_fp8_matmul(A_vllm,
                                      B_vllm,
                                      A_scale_vllm,
@@ -101,6 +109,9 @@ def benchmark_shape(m: int,
 
     # === vLLM CUTLASS Implementation ===
     def vllm_cutlass_gemm():
+        # A quantization is inside the loop as it depends on activations
+        # A_vllm_cutlass, A_scale_vllm_cutlass = per_token_group_quant_fp8(
+        #     A, block_size[1], column_major_scales=True)
         return ops.cutlass_scaled_mm(A_vllm_cutlass,
                                      B_vllm.T,
                                      scale_a=A_scale_vllm_cutlass,

cmake/cpu_extension.cmake

@@ -96,21 +96,12 @@ if (AVX512_FOUND AND NOT AVX512_DISABLED)
         if (CMAKE_CXX_COMPILER_ID STREQUAL "GNU" AND
             CMAKE_CXX_COMPILER_VERSION VERSION_GREATER_EQUAL 12.3)
             list(APPEND CXX_COMPILE_FLAGS "-mavx512bf16")
-            set(ENABLE_AVX512BF16 ON)
         else()
-            set(ENABLE_AVX512BF16 OFF)
             message(WARNING "Disable AVX512-BF16 ISA support, requires gcc/g++ >= 12.3")
         endif()
     else()
-        set(ENABLE_AVX512BF16 OFF)
         message(WARNING "Disable AVX512-BF16 ISA support, no avx512_bf16 found in local CPU flags." " If cross-compilation is required, please set env VLLM_CPU_AVX512BF16=1.")
     endif()
-
-    find_isa(${CPUINFO} "avx512_vnni" AVX512VNNI_FOUND)
-    if (AVX512VNNI_FOUND)
-        list(APPEND CXX_COMPILE_FLAGS "-mavx512vnni")
-        set(ENABLE_AVX512VNNI ON) 
-    endif() 
     
 elseif (AVX2_FOUND)
     list(APPEND CXX_COMPILE_FLAGS "-mavx2")
@@ -240,17 +231,6 @@ if (AVX512_FOUND AND NOT AVX512_DISABLED)
         "csrc/cpu/quant.cpp"
         "csrc/cpu/shm.cpp"
         ${VLLM_EXT_SRC})
-    if (ENABLE_AVX512BF16 AND ENABLE_AVX512VNNI)
-        set(VLLM_EXT_SRC
-            "csrc/cpu/sgl-kernels/gemm.cpp"
-            "csrc/cpu/sgl-kernels/gemm_int8.cpp"
-            "csrc/cpu/sgl-kernels/gemm_fp8.cpp"
-            "csrc/cpu/sgl-kernels/moe.cpp"
-            "csrc/cpu/sgl-kernels/moe_int8.cpp"
-            "csrc/cpu/sgl-kernels/moe_fp8.cpp"
-            ${VLLM_EXT_SRC})
-        add_compile_definitions(-DCPU_CAPABILITY_AVX512)
-    endif()
 elseif(POWER10_FOUND)
     set(VLLM_EXT_SRC
         "csrc/cpu/quant.cpp"

cmake/external_projects/vllm_flash_attn.cmake

@@ -38,7 +38,7 @@ else()
   FetchContent_Declare(
           vllm-flash-attn
           GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
-          GIT_TAG 1c2624e53c078854e0637ee566c72fe2107e75f4
+          GIT_TAG 763ad155a1c826f71ff318f41edb1e4e5e376ddb
           GIT_PROGRESS TRUE
           # Don't share the vllm-flash-attn build between build types
           BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn

cmake/utils.cmake

@@ -265,8 +265,8 @@ macro(set_gencode_flags_for_srcs)
 endmacro()
 
 #
-# For the given `SRC_CUDA_ARCHS` list of gencode versions in the form
-#  `<major>.<minor>[letter]` compute the "loose intersection" with the
+# For the given `SRC_CUDA_ARCHS` list of gencode versions in the form 
+#  `<major>.<minor>[letter]` compute the "loose intersection" with the 
 #  `TGT_CUDA_ARCHS` list of gencodes. We also support the `+PTX` suffix in
 #  `SRC_CUDA_ARCHS` which indicates that the PTX code should be built when there
 #  is a CUDA_ARCH in `TGT_CUDA_ARCHS` that is equal to or larger than the
@@ -278,7 +278,7 @@ endmacro()
 #  in `SRC_CUDA_ARCHS` that is less or equal to the version in `TGT_CUDA_ARCHS`.
 # We have special handling for x.0a, if x.0a is in `SRC_CUDA_ARCHS` and x.0 is
 #  in `TGT_CUDA_ARCHS` then we should remove x.0a from `SRC_CUDA_ARCHS` and add
-#  x.0a to the result (and remove x.0 from TGT_CUDA_ARCHS).
+#  x.0a to the result (and remove x.0 from TGT_CUDA_ARCHS). 
 # The result is stored in `OUT_CUDA_ARCHS`.
 #
 # Example:
@@ -313,16 +313,21 @@ function(cuda_archs_loose_intersection OUT_CUDA_ARCHS SRC_CUDA_ARCHS TGT_CUDA_AR
   # if x.0a is in SRC_CUDA_ARCHS and x.0 is in CUDA_ARCHS then we should
   # remove x.0a from SRC_CUDA_ARCHS and add x.0a to _CUDA_ARCHS
   set(_CUDA_ARCHS)
-  foreach(_arch ${_SRC_CUDA_ARCHS})
-    if(_arch MATCHES "\\a$")
-      list(REMOVE_ITEM _SRC_CUDA_ARCHS "${_arch}")
-      string(REPLACE "a" "" _base "${_arch}")
-      if ("${_base}" IN_LIST TGT_CUDA_ARCHS)
-        list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_base}")
-        list(APPEND _CUDA_ARCHS "${_arch}")
-      endif()
+  if ("9.0a" IN_LIST _SRC_CUDA_ARCHS)
+    list(REMOVE_ITEM _SRC_CUDA_ARCHS "9.0a")
+    if ("9.0" IN_LIST TGT_CUDA_ARCHS)
+      list(REMOVE_ITEM _TGT_CUDA_ARCHS "9.0")
+      set(_CUDA_ARCHS "9.0a")
     endif()
-  endforeach()
+  endif()
+
+  if ("10.0a" IN_LIST _SRC_CUDA_ARCHS)
+    list(REMOVE_ITEM _SRC_CUDA_ARCHS "10.0a")
+    if ("10.0" IN_LIST TGT_CUDA_ARCHS)
+      list(REMOVE_ITEM _TGT_CUDA_ARCHS "10.0")
+      set(_CUDA_ARCHS "10.0a")
+    endif()
+  endif()
 
   list(SORT _SRC_CUDA_ARCHS COMPARE NATURAL ORDER ASCENDING)
 
@@ -354,7 +359,7 @@ function(cuda_archs_loose_intersection OUT_CUDA_ARCHS SRC_CUDA_ARCHS TGT_CUDA_AR
   endforeach()
 
   list(REMOVE_DUPLICATES _CUDA_ARCHS)
-
+  
   # reapply +PTX suffix to architectures that requested PTX
   set(_FINAL_ARCHS)
   foreach(_arch ${_CUDA_ARCHS})
@@ -365,7 +370,7 @@ function(cuda_archs_loose_intersection OUT_CUDA_ARCHS SRC_CUDA_ARCHS TGT_CUDA_AR
     endif()
   endforeach()
   set(_CUDA_ARCHS ${_FINAL_ARCHS})
-
+  
   set(${OUT_CUDA_ARCHS} ${_CUDA_ARCHS} PARENT_SCOPE)
 endfunction()
 

csrc/attention/mla/cutlass_mla_kernels.cu

@@ -207,7 +207,7 @@ void cutlass_mla_decode_sm100a(torch::Tensor const& out,
               "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));
+  at::cuda::CUDAGuard device_guard{(char)q_nope.get_device()};
   const cudaStream_t stream =
       at::cuda::getCurrentCUDAStream(q_nope.get_device());
   if (in_dtype == at::ScalarType::Half) {

csrc/cpu/sgl-kernels/common.h

@@ -1,238 +0,0 @@
-// Adapted from
-// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu
-
-#pragma once
-
-#include <ATen/ATen.h>
-#include <ATen/Parallel.h>
-#include <ATen/record_function.h>
-
-// clang-format off
-
-#if defined(_OPENMP)
-#include <omp.h>
-#endif
-
-namespace {
-
-// dispatch bool
-#define AT_DISPATCH_BOOL(BOOL_V, BOOL_NAME, ...)                                 \
-  [&] {                                                                          \
-    if (BOOL_V) {                                                                \
-      constexpr bool BOOL_NAME = true;                                           \
-      return __VA_ARGS__();                                                      \
-    } else {                                                                     \
-      constexpr bool BOOL_NAME = false;                                          \
-      return __VA_ARGS__();                                                      \
-    }                                                                            \
-  }()
-
-// dispatch: bfloat16, float16, int8_t, fp8_e4m3
-#define CPU_DISPATCH_PACKED_TYPES(TYPE, ...)                                    \
-  [&] {                                                                         \
-    switch (TYPE) {                                                             \
-      case at::ScalarType::BFloat16 : {                                         \
-        using packed_t = at::BFloat16;                                          \
-        return __VA_ARGS__();                                                   \
-      }                                                                         \
-      case at::ScalarType::Half: {                                              \
-        using packed_t = at::Half;                                              \
-        return __VA_ARGS__();                                                   \
-      }                                                                         \
-      case at::ScalarType::Char : {                                             \
-        using packed_t = int8_t;                                                \
-        return __VA_ARGS__();                                                   \
-      }                                                                         \
-      case at::ScalarType::Float8_e4m3fn : {                                    \
-        using packed_t = at::Float8_e4m3fn;                                     \
-        return __VA_ARGS__();                                                   \
-      }                                                                         \
-      default:                                                                  \
-        TORCH_CHECK(false, "Unsupported floating data type.\n");                \
-    }                                                                           \
-  }()
-
-#define UNUSED(x) (void)(x)
-
-#define CHECK_CPU(x) TORCH_CHECK(x.device().type() == at::kCPU, #x " must be a CPU tensor")
-
-#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
-#define CHECK_LAST_DIM_CONTIGUOUS(x) \
-  TORCH_CHECK(x.strides()[x.strides().size() - 1] == 1, #x "must be contiguous at last dimention")
-
-#define CHECK_INPUT(x) \
-  CHECK_CPU(x);        \
-  CHECK_CONTIGUOUS(x)
-#define CHECK_LAST_DIM_CONTIGUOUS_INPUT(x) \
-  CHECK_CPU(x);                            \
-  CHECK_LAST_DIM_CONTIGUOUS(x)
-
-#define CHECK_DIM(d, x) TORCH_CHECK(x.dim() == d, #x " must be a " #d "D tensor")
-
-#define CHECK_EQ(a, b) TORCH_CHECK((a) == (b), "CHECK_EQ(" #a ", " #b ") failed. ", a, " vs ", b)
-
-// parallel routines
-constexpr int GRAIN_SIZE = 1024;
-
-template <typename T, typename std::enable_if<std::is_integral<T>::value, int>::type = 0>
-inline T div_up(T x, T y) { return (x + y - 1) / y; }
-
-template <typename T>
-inline void balance211(T n, T nth, T ith, T& n_start, T& n_end) {
-#if 0
-    // onednn partition pattern
-    T& n_my = n_end;
-    if (nth <= 1 || n == 0) {
-        n_start = 0;
-        n_my = n;
-    } else {
-        T n1 = div_up(n, nth);
-        T n2 = n1 - 1;
-        T T1 = n - n2 * nth;
-        n_my = ith < T1 ? n1 : n2;
-        n_start = ith <= T1 ? ith*n1 : T1 * n1 + (ith - T1) * n2;
-    }
-    n_end += n_start;
-#else
-    // pytorch aten partition pattern
-    T n_my = div_up(n, nth);
-    n_start = ith * n_my;
-    n_end = std::min(n_start + n_my, n);
-#endif
-}
-
-template <typename func_t>
-inline void parallel_for(int n, const func_t& f) {
-#if defined(_OPENMP)
-#pragma omp parallel
-{
-    int nth = omp_get_num_threads();
-    int ith = omp_get_thread_num();
-    int tbegin, tend;
-    balance211(n, nth, ith, tbegin, tend);
-    f(tbegin, tend);
-}
-#else
-    f(0, n);
-#endif
-}
-
-// for 1d parallel, use `actual_nth`
-// for 2d parallel, use even nths, e.g. 43->42
-int inline adjust_num_threads(int m) {
-  int actual_nth = at::get_num_threads();
-  if (m == 1) {
-    return actual_nth;
-  }
-  return std::max(1, (actual_nth >> 1) * 2);
-}
-
-template <typename func_t>
-inline void parallel_2d(int m, int n, const func_t& f) {
-
-  // make sure we have even num_threads
-  int nth = adjust_num_threads(m);
-
-  // [NOTE] thread blocking:
-  //
-  //   1) prefer square block per thread
-  //   2) use even number of CPU cores
-  //   3) use all `num_threads` cores
-  //
-  //   we have:
-  //     TM * TN = T
-  //     BM / TM = BN / TN
-  //   then:
-  //     TM = ((BM / BN) * T) ^ 0.5
-  //
-  float r = float(m) / n;
-  int nth_m = std::ceil(std::sqrt(r * nth));
-  int nth_n = 1;
-  for (; nth_m > 0; --nth_m) {
-    nth_n = nth / nth_m;
-    if (nth_m * nth_n == nth) {
-      break;
-    }
-  }
-
-#if defined(_OPENMP)
-#pragma omp parallel num_threads(nth)
-{
-  int ith = omp_get_thread_num();
-  int ith_m = ith / nth_n;
-  int ith_n = ith % nth_n;
-
-  int thread_block_m = div_up(m, nth_m);
-  int thread_block_n = div_up(n, nth_n);
-
-  int begin_m = ith_m * thread_block_m;
-  int end_m = std::min(m, begin_m + thread_block_m);
-  int begin_n = ith_n * thread_block_n;
-  int end_n = std::min(n, begin_n + thread_block_n);
-
-  f(begin_m, end_m, begin_n, end_n);
-}
-#else
-  f(0, m, 0, n);
-#endif
-}
-
-template <typename T>
-int get_cache_blocks(int BLOCK_SIZE, int K) {
-  // L2 2MB and ratio of 50%
-  const int L2_size = 2048 * 1024 >> 1;
-  return std::max(1, int(L2_size / (BLOCK_SIZE * K * sizeof(T))));
-}
-
-// data indexing for dimension collapse
-template <typename T>
-inline T data_index_init(T offset) {
-  return offset;
-}
-
-template <typename T, typename... Args>
-inline T data_index_init(T offset, T& x, const T& X, Args&&... args) {
-  offset = data_index_init(offset, std::forward<Args>(args)...);
-  x = offset % X;
-  return offset / X;
-}
-
-inline bool data_index_step() {
-  return true;
-}
-
-template <typename T, typename... Args>
-inline bool data_index_step(T& x, const T& X, Args&&... args) {
-  if (data_index_step(std::forward<Args>(args)...)) {
-    x = ((x + 1) == X) ? 0 : (x + 1);
-    return x == 0;
-  }
-  return false;
-}
-
-// forced unroll for perf critical path
-
-#if __has_attribute(always_inline)
-#define ALWAYS_INLINE __attribute__((__always_inline__)) inline
-#else
-#define ALWAYS_INLINE inline
-#endif
-
-template <int n>
-struct Unroll {
-  template <typename Func, typename... Args>
-  ALWAYS_INLINE void operator()(const Func& f, Args... args) const {
-    Unroll<n - 1>{}(f, args...);
-    f(std::integral_constant<int, n - 1>{}, args...);
-  }
-};
-
-template <>
-struct Unroll<1> {
-  template <typename Func, typename... Args>
-  ALWAYS_INLINE void operator()(const Func& f, Args... args) const {
-    f(std::integral_constant<int, 0>{}, args...);
-  }
-};
-
-} // anonymous namespace

csrc/cpu/sgl-kernels/gemm.cpp

@@ -1,464 +0,0 @@
-// Adapted from
-// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu
-
-#include "common.h"
-#include "vec.h"
-#include "gemm.h"
-
-// clang-format off
-
-namespace {
-
-// packed   layout:
-//   quants {N, K}  int8_t
-//   comp   {N}     int32_t
-template <int BLOCK_N>
-inline void s8s8_compensation(int8_t* __restrict__ packed, int K) {
-#if defined(CPU_CAPABILITY_AVX512)
-  constexpr int COLS = BLOCK_N / 16;
-  __m512i vcomp[COLS];
-
-  for (int col = 0; col < COLS; ++col) {
-    vcomp[col] = _mm512_setzero_si512();
-  }
-
-  const int64_t offset = BLOCK_N * K;
-  const __m512i off = _mm512_set1_epi8(static_cast<char>(0x80));
-  for (int k = 0; k < K / 4; ++k) {
-    for (int col = 0; col < COLS; ++col) {
-      __m512i vb = _mm512_loadu_si512((const __m512i *)(packed + k * BLOCK_N * 4 + col * 64));
-      vcomp[col] = _mm512_dpbusd_epi32(vcomp[col], off, vb);
-    }
-  }
-
-  for (int col = 0; col < COLS; ++col) {
-    _mm512_storeu_si512((__m512i *)(packed + offset + col * 64), vcomp[col]);
-  }
-#else
-  TORCH_CHECK(false, "s8s8_compensation not implemented!");
-#endif
-}
-
-// convert to vnni format
-// from [N, K] to [K/2, N, 2] for bfloat16 and float16
-template <typename packed_t>
-inline void pack_vnni(packed_t* __restrict__ packed, const packed_t* __restrict__ weight, int N, int K) {
-  const int VNNI_BLK = 2;
-  for (int n = 0; n < N; ++n) {
-    for (int k = 0; k < K / VNNI_BLK; ++k) {
-      for (int d = 0; d < VNNI_BLK; ++d) {
-        packed[k * N * VNNI_BLK + n * VNNI_BLK + d] = weight[n * K + k * VNNI_BLK + d];
-      }
-    }
-  }
-}
-
-template <>
-inline void pack_vnni<int8_t>(int8_t* __restrict__ packed, const int8_t* __restrict__ weight, int N, int K) {
-  constexpr int BLOCK_N = block_size_n();
-  TORCH_CHECK(N == BLOCK_N);
-
-  const int VNNI_BLK = 4;
-  for (int n = 0; n < N; ++n) {
-    for (int k = 0; k < K / VNNI_BLK; ++k) {
-      for (int d = 0; d < VNNI_BLK; ++d) {
-        packed[k * N * VNNI_BLK + n * VNNI_BLK + d] = weight[n * K + k * VNNI_BLK + d];
-      }
-    }
-  }
-  s8s8_compensation<BLOCK_N>(packed, K);
-}
-
-template <typename scalar_t>
-inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ input, int64_t size) {
-  using bVec = at::vec::Vectorized<scalar_t>;
-  using fVec = at::vec::Vectorized<float>;
-  constexpr int kVecSize = bVec::size();
-
-  int64_t d;
-  #pragma GCC unroll 4
-  for (d = 0; d <= size - kVecSize; d += kVecSize) {
-    fVec data0 = fVec::loadu(input + d);
-    fVec data1 = fVec::loadu(input + d + fVec::size());
-    bVec out_vec = convert_from_float_ext<scalar_t>(data0, data1);
-    out_vec.store(out + d);
-  }
-  for (; d < size; ++d) {
-    out[d] = static_cast<scalar_t>(input[d]);
-  }
-}
-
-template <typename scalar_t>
-inline void copy_add_stub(scalar_t* __restrict__ out, const float* __restrict__ input, const float* __restrict__ bias, int64_t size) {
-  using bVec = at::vec::Vectorized<scalar_t>;
-  using fVec = at::vec::Vectorized<float>;
-  constexpr int kVecSize = bVec::size();
-
-  int64_t d;
-  #pragma GCC unroll 4
-  for (d = 0; d <= size - kVecSize; d += kVecSize) {
-    fVec data0 = fVec::loadu(input + d) + fVec::loadu(bias + d);
-    fVec data1 = fVec::loadu(input + d + fVec::size()) + fVec::loadu(bias + d + fVec::size());
-    bVec out_vec = convert_from_float_ext<scalar_t>(data0, data1);
-    out_vec.store(out + d);
-  }
-  for (; d < size; ++d) {
-    out[d] = static_cast<scalar_t>(input[d] + bias[d]);
-  }
-}
-
-template <typename scalar_t, bool has_bias, int BLOCK_M, int BLOCK_N>
-struct tinygemm_kernel_nn {
-  static inline void apply(
-      const scalar_t* __restrict__ A, const scalar_t* __restrict__ B, scalar_t* __restrict__ C,
-      const float* __restrict__ bias, int64_t K, int64_t lda, int64_t ldb, int64_t ldc) {
-    TORCH_CHECK(false, "tinygemm_kernel_nn: scalar path not implemented!");
-  }
-};
-
-#if defined(CPU_CAPABILITY_AVX512)
-template <bool has_bias, int BLOCK_M, int BLOCK_N>
-struct tinygemm_kernel_nn<at::BFloat16, has_bias, BLOCK_M, BLOCK_N> {
-  static inline void apply(
-      const at::BFloat16* __restrict__ A, const at::BFloat16* __restrict__ B, at::BFloat16* __restrict__ C,
-      const float* __restrict__ bias, int64_t K, int64_t lda, int64_t ldb, int64_t ldc) {
-
-    constexpr int ROWS = BLOCK_M;
-    constexpr int COLS = BLOCK_N / 16;
-
-    // prefetch distance
-    constexpr int PREFETCH_SIZE_K = 0;
-
-    __m512bh va;
-    __m512bh vb[COLS];
-    __m512 vc[ROWS * COLS];
-
-    auto loadc = [&](auto i) {
-      constexpr int col = i % COLS;
-      if constexpr (has_bias) {
-        vc[i] = _mm512_loadu_ps(bias + col * 16);
-      } else {
-        vc[i] = _mm512_set1_ps(0.f);
-      }
-    };
-    Unroll<ROWS * COLS>{}(loadc);
-
-    const int64_t K2 = K >> 1;
-    const int64_t lda2 = lda >> 1;
-    const int64_t ldb2 = ldb; // ldb * 2 >> 1;
-    const float* a_ptr = reinterpret_cast<const float*>(A);
-    const float* b_ptr = reinterpret_cast<const float*>(B);
-
-    auto compute = [&](auto i, int64_t k) {
-      constexpr int row = i / COLS;
-      constexpr int col = i % COLS;
-
-      if constexpr (col == 0) {
-        va = (__m512bh)(_mm512_set1_ps(a_ptr[row * lda2 + k]));
-      }
-      if constexpr (row == 0) {
-        vb[col] = (__m512bh)(_mm512_loadu_si512(b_ptr + k * ldb2 + col * 16));
-        if constexpr (PREFETCH_SIZE_K > 0) {
-          _mm_prefetch(b_ptr + (k + PREFETCH_SIZE_K) * ldb2 + col * 16, _MM_HINT_T0);
-        }
-      }
-      vc[i] = _mm512_dpbf16_ps(vc[i], va, vb[col]);
-    };
-    for (int64_t k = 0; k < K2; ++k) {
-      Unroll<ROWS * COLS>{}(compute, k);
-    }
-
-    auto storec = [&](auto i) {
-      constexpr int row = i / COLS;
-      constexpr int col = i % COLS;
-      // for COLS = 2, 4 use 512bit store
-      // for COLS = 1, 3 use 256bit store
-      if constexpr (COLS % 2 == 0) {
-        if constexpr (col % 2 == 0) {
-          _mm512_storeu_si512(
-              reinterpret_cast<__m512i*>((C + row * ldc + col * 16)),
-              (__m512i)(_mm512_cvtne2ps_pbh(vc[row * COLS + col + 1], vc[row * COLS + col])));
-        }
-      } else {
-        _mm256_storeu_si256(
-            reinterpret_cast<__m256i*>(C + row * ldc + col * 16),
-            (__m256i)(_mm512_cvtneps_pbh(vc[i])));
-      }
-    };
-    Unroll<ROWS * COLS>{}(storec);
-  }
-};
-#endif
-
-#define LAUNCH_TINYGEMM_KERNEL_NN(MB_SIZE, NB_SIZE)                          \
-    tinygemm_kernel_nn<scalar_t, has_bias, MB_SIZE, NB_SIZE>::apply(         \
-        A + mb_start * lda, B + nb_start * 2, C + mb_start * ldc + nb_start, \
-        has_bias ? bias + nb_start : nullptr, K, lda, ldb, ldc);
-
-template <typename scalar_t, bool has_bias>
-struct brgemm {
-  static inline void apply(
-      const scalar_t* __restrict__ A, const scalar_t* __restrict__ B, scalar_t* __restrict__ C,
-      float* __restrict__ Ctmp, const float* __restrict__ bias,
-      int64_t M, int64_t N, int64_t K, int64_t lda, int64_t ldb, int64_t ldc) {
-
-    constexpr int BLOCK_N = block_size_n();
-    at::native::cpublas::brgemm(
-        M, N, K, lda, ldb, BLOCK_N, /* add_C */false,
-        A, B, Ctmp);
-
-    // copy from Ctmp to C
-    for (int64_t m = 0; m < M; ++m) {
-      if constexpr (has_bias) {
-        copy_add_stub(C + m * ldc, Ctmp + m * BLOCK_N, bias, N);
-      } else {
-        copy_stub(C + m * ldc, Ctmp + m * BLOCK_N, N);
-      }
-    }
-  }
-};
-
-template <typename scalar_t, bool has_bias>
-void tinygemm_kernel(
-    const scalar_t* __restrict__ A,
-    const scalar_t* __restrict__ B,
-    scalar_t* __restrict__ C,
-    float* __restrict__ Ctmp,
-    const float* __restrict__ bias,
-    int64_t M,
-    int64_t N,
-    int64_t K,
-    int64_t lda,
-    int64_t ldb,
-    int64_t ldc,
-    bool brg) {
-
-  if (brg) {
-    brgemm<scalar_t, has_bias>::apply(
-        A, B, C, Ctmp, bias,
-        M, N, K, lda, ldb, ldc);
-    return;
-  }
-
-  // pattern: 1-4-16
-  constexpr int64_t BLOCK_M = 4;
-  constexpr int64_t BLOCK_N = 64;
-  const int64_t MB = div_up(M, BLOCK_M);
-  const int64_t NB = div_up(N, BLOCK_N);
-  for (int mb = 0; mb < MB; ++mb) {
-    int64_t mb_start = mb * BLOCK_M;
-    int64_t mb_size = std::min(BLOCK_M, M - mb_start);
-    for (int64_t nb = 0; nb < NB; ++nb) {
-      int64_t nb_start = nb * BLOCK_N;
-      int64_t nb_size = std::min(BLOCK_N, N - nb_start);
-
-      switch(mb_size << 4 | nb_size >> 4) {
-        // mb_size = 1
-        case 0x12: LAUNCH_TINYGEMM_KERNEL_NN(1, 32); break;
-        case 0x14: LAUNCH_TINYGEMM_KERNEL_NN(1, 64); break;
-        // mb_size = 2
-        case 0x22: LAUNCH_TINYGEMM_KERNEL_NN(2, 32); break;
-        case 0x24: LAUNCH_TINYGEMM_KERNEL_NN(2, 64); break;
-        // mb_size = 3
-        case 0x32: LAUNCH_TINYGEMM_KERNEL_NN(3, 32); break;
-        case 0x34: LAUNCH_TINYGEMM_KERNEL_NN(3, 64); break;
-        // mb_size = 4
-        case 0x42: LAUNCH_TINYGEMM_KERNEL_NN(4, 32); break;
-        case 0x44: LAUNCH_TINYGEMM_KERNEL_NN(4, 64); break;
-        default: TORCH_CHECK(false, "Unexpected block size, ", mb_size, "x", "nb_size");
-      }
-    }
-  }
-}
-
-template <typename scalar_t>
-void weight_packed_linear_kernel_impl(
-    scalar_t* __restrict__ out,
-    const scalar_t* __restrict__ mat1,
-    const scalar_t* __restrict__ mat2,
-    const float* __restrict__ bias,
-    int64_t M,
-    int64_t N,
-    int64_t K,
-    int64_t mat1_strideM,
-    int64_t out_strideM) {
-
-  constexpr int64_t BLOCK_M = block_size_m();
-  constexpr int64_t BLOCK_N = block_size_n();
-  const int64_t MB = div_up(M, BLOCK_M);
-  const int64_t NB = div_up(N, BLOCK_N);
-
-  // use avx512-bf16 when a) M is small; b) dtype is bfloat16, otherwise use amx
-  const bool use_brgemm = (M > 4) || (!std::is_same_v<scalar_t, at::BFloat16>);
-
-  // l2 cache block for n
-  int64_t cache_blocks_nb = get_cache_blocks<scalar_t>(BLOCK_N, K);
-
-  // parallel on [MB, NB]
-  AT_DISPATCH_BOOL(bias != nullptr, has_bias, [&] {
-    parallel_2d(MB, NB, [&](int64_t begin_mb, int64_t end_mb, int64_t begin_nb, int64_t end_nb) {
-
-      // for brgemm, use float32 for accumulate
-      alignas(64) float Ctmp[BLOCK_M * BLOCK_N];
-
-      for (int64_t nbb = begin_nb; nbb < end_nb; nbb += cache_blocks_nb) {
-      for (int64_t mb = begin_mb; mb < end_mb; ++mb) {
-      for (int64_t nb = nbb; nb < std::min(nbb + cache_blocks_nb, end_nb); ++nb) {
-
-        int64_t mb_start = mb * BLOCK_M;
-        int64_t mb_size = std::min(M - mb_start, BLOCK_M);
-        int64_t nb_start = nb * BLOCK_N;
-        int64_t nb_size = std::min(N - nb_start, BLOCK_N);
-
-        tinygemm_kernel<scalar_t, has_bias>(
-            /*   A */ mat1 + mb_start * mat1_strideM,
-            /*   B */ mat2 + nb_start * K /* nb * BLOCK_N * K */,
-            /*   C */ out + mb_start * out_strideM + nb_start,
-            /* Ctmp*/ Ctmp,
-            /* bias*/ bias + nb_start,
-            /*   M */ mb_size,
-            /*   N */ nb_size,
-            /*   K */ K,
-            /* lda */ mat1_strideM,
-            /* ldb */ nb_size,
-            /* ldc */ out_strideM,
-            /* brg */ use_brgemm);
-      }}}
-
-      if (use_brgemm) {
-        at::native::cpublas::brgemm_release();
-      }
-    });
-  });
-}
-
-} // anonymous namespace
-
-// tinygemm interface
-template <typename scalar_t>
-void tinygemm_kernel(const scalar_t* __restrict__ A, const scalar_t* __restrict__ B, scalar_t* __restrict__ C,
-    float* __restrict__ Ctmp, int64_t M, int64_t N, int64_t K, int64_t lda, int64_t ldb, int64_t ldc, bool brg) {
-  tinygemm_kernel<scalar_t, false>(A, B, C, Ctmp, nullptr, M, N, K, lda, ldb, ldc, brg);
-}
-
-#define INSTANTIATE_TINYGEMM_TEMPLATE(TYPE)                                             \
-    template void tinygemm_kernel<TYPE>(                                                \
-        const TYPE* __restrict__ A, const TYPE* __restrict__ B, TYPE* __restrict__ C,   \
-        float* __restrict__ Ctmp, int64_t M, int64_t N, int64_t K, int64_t lda,         \
-        int64_t ldb, int64_t ldc, bool brg)
-
-INSTANTIATE_TINYGEMM_TEMPLATE(at::BFloat16);
-INSTANTIATE_TINYGEMM_TEMPLATE(at::Half);
-
-at::Tensor convert_weight_packed(at::Tensor& weight) {
-  // for 3d moe weights
-  // weight : [E, OC, IC]
-  //     w1 : [E, 2N,  K]
-  //     w2 : [E,  K,  N]
-  CHECK_INPUT(weight);
-
-  const int64_t ndim = weight.ndimension();
-  TORCH_CHECK(ndim == 2 || ndim == 3, "expect weight to be 2d or 3d, got ", ndim, "d tensor.");
-  const auto st = weight.scalar_type();
-  const int64_t E = ndim == 3 ? weight.size(0) : 1;
-  const int64_t OC = ndim == 3 ? weight.size(1) : weight.size(0);
-  const int64_t IC = ndim == 3 ? weight.size(2) : weight.size(1);
-
-  // we handle 2 TILE_N at a time.
-  TORCH_CHECK(OC % TILE_N == 0, "invalid weight out features ", OC);
-  TORCH_CHECK(IC % TILE_K == 0, "invalid weight input features ", IC);
-
-  constexpr int64_t BLOCK_N = block_size_n();
-  const int64_t NB = div_up(OC, BLOCK_N);
-
-  // use phony sizes here [E, OC, IC], for each [E], [OC, IC] -> [IC / 2, OC, 2]
-  auto packed_weight = at::empty({}, weight.options());
-  const int64_t stride = OC * IC;
-
-  TORCH_CHECK(st == at::kBFloat16 || st == at::kHalf || st == at::kChar || st == at::kFloat8_e4m3fn,
-      "expect weight to be bfloat16, float16, int8 or fp8_e4m3.");
-
-  CPU_DISPATCH_PACKED_TYPES(st, [&] {
-    // adjust most inner dimension size
-    const int packed_row_size = get_row_size<packed_t>(IC);
-    auto sizes = weight.sizes().vec();
-    sizes[ndim - 1] = packed_row_size;
-    packed_weight.resize_(sizes);
-
-    const packed_t* w_data = weight.data_ptr<packed_t>();
-    packed_t* packed_data = packed_weight.data_ptr<packed_t>();
-
-    // parallel on {E, NB}
-    at::parallel_for(0, E * NB, 0, [&](int64_t begin, int64_t end) {
-      int64_t e{0}, nb{0};
-      data_index_init(begin, e, E, nb, NB);
-
-      for (int64_t i = begin; i < end; ++i) {
-        UNUSED(i);
-
-        int64_t n = nb * BLOCK_N;
-        int64_t n_size = std::min(BLOCK_N, OC - n);
-        pack_vnni<packed_t>(
-            packed_data + e * OC * packed_row_size + n * packed_row_size,
-            w_data + e * stride + n * IC,
-            n_size,
-            IC);
-
-        // move to the next index
-        data_index_step(e, E, nb, NB);
-      }
-    });
-  });
-  return packed_weight;
-}
-
-// mat1 : [M, K]
-// mat2 : [N, K]
-// bias : [N]
-// out  : [M, N]
-//
-at::Tensor weight_packed_linear(at::Tensor& mat1, at::Tensor& mat2,
-    const std::optional<at::Tensor>& bias, bool is_vnni) {
-  RECORD_FUNCTION(
-    "sgl-kernel::weight_packed_linear", std::vector<c10::IValue>({mat1, mat2, bias}));
-
-  auto packed_w = is_vnni ? mat2 : convert_weight_packed(mat2);
-
-  CHECK_LAST_DIM_CONTIGUOUS_INPUT(mat1);
-  CHECK_INPUT(mat2);
-
-  int64_t M = mat1.size(0);
-  int64_t N = mat2.size(0);
-  int64_t K = mat2.size(1);
-  CHECK_EQ(mat1.size(1), K);
-  CHECK_DIM(2, mat1);
-  CHECK_DIM(2, mat2);
-
-  auto out = at::empty({M, N}, mat1.options());
-
-  // strides
-  int64_t mat1_strideM = mat1.stride(0);
-  int64_t out_strideM = out.stride(0);
-
-  const bool has_bias = bias.has_value();
-  const float* bias_data = nullptr;
-  if (has_bias) {
-    CHECK_EQ(bias.value().size(0), N);
-    bias_data = bias.value().data_ptr<float>();
-  }
-
-  AT_DISPATCH_REDUCED_FLOATING_TYPES(mat1.scalar_type(), "weight_packed_linear_kernel_impl", [&] {
-    weight_packed_linear_kernel_impl<scalar_t>(
-        out.data_ptr<scalar_t>(),
-        mat1.data_ptr<scalar_t>(),
-        packed_w.data_ptr<scalar_t>(),
-        bias_data,
-        M,
-        N,
-        K,
-        mat1_strideM,
-        out_strideM);
-  });
-
-  return out;
-}

csrc/cpu/sgl-kernels/gemm.h

@@ -1,266 +0,0 @@
-#pragma once
-
-#include <ATen/native/CPUBlas.h>
-
-// clang-format off
-
-// amx-bf16
-#define TILE_M 16
-#define TILE_N 16
-#define TILE_K 32
-
-// block size for AMX gemm
-constexpr int block_size_m() { return 2 * TILE_M; }
-constexpr int block_size_n() { return 2 * TILE_N; }
-
-// define threshold using brgemm (intel AMX)
-template <typename T> inline bool can_use_brgemm(int M);
-template <> inline bool can_use_brgemm<at::BFloat16>(int M) { return M > 4; }
-template <> inline bool can_use_brgemm<at::Half>(int M) { return true; }
-// TODO: add u8s8 brgemm, this requires PyTorch 2.7
-template <> inline bool can_use_brgemm<int8_t>(int M) { return false; }
-template <> inline bool can_use_brgemm<at::Float8_e4m3fn>(int M) { return M > 4; }
-template <> inline bool can_use_brgemm<at::quint4x2>(int M) { return M > 4; }
-
-// work around compiler internal error
-#define BLOCK_K 128 // 4 * TILE_K
-
-// adjust leading dimension size for K
-template <typename T>
-inline int64_t get_row_size(int64_t K) {
-  return K;
-}
-
-template <>
-inline int64_t get_row_size<int8_t>(int64_t K) {
-  return K + sizeof(int32_t);
-}
-
-inline int64_t get_row_size(int64_t K, bool use_int8_w8a8) {
-  return use_int8_w8a8 ? K + sizeof(int32_t) : K;
-}
-
-// pack weight to vnni format
-at::Tensor convert_weight_packed(at::Tensor& weight);
-
-// moe implementations for int8 w8a8
-template <typename scalar_t>
-void fused_experts_int8_kernel_impl(
-    scalar_t* __restrict__ output,
-    scalar_t* __restrict__ ic1,
-    scalar_t* __restrict__ ic2,
-    uint8_t* __restrict__ A_tmp,
-    float* __restrict__ C_tmp,
-    uint8_t* __restrict__ Aq_tmp,
-    float* __restrict__ As_tmp,
-    const scalar_t* __restrict__ input,
-    const int8_t* __restrict__ packed_w1,
-    const int8_t* __restrict__ packed_w2,
-    const float* __restrict__ w1s,
-    const float* __restrict__ w2s,
-    const float* __restrict__ topk_weights,
-    const int32_t* __restrict__ sorted_ids,
-    const int32_t* __restrict__ expert_ids,
-    const int32_t* __restrict__ offsets,
-    int64_t M,
-    int64_t N,
-    int64_t K,
-    int64_t E,
-    int64_t topk,
-    int64_t num_tokens_post_pad);
-
-// moe implementations for fp8 w8a16
-template <typename scalar_t>
-void fused_experts_fp8_kernel_impl(
-    scalar_t* __restrict__ output,
-    scalar_t* __restrict__ ic0,
-    scalar_t* __restrict__ ic1,
-    scalar_t* __restrict__ ic2,
-    scalar_t* __restrict__ A_tmp,
-    scalar_t* __restrict__ B_tmp,
-    float* __restrict__ C_tmp,
-    const scalar_t* __restrict__ input,
-    const at::Float8_e4m3fn* __restrict__ packed_w1,
-    const at::Float8_e4m3fn* __restrict__ packed_w2,
-    const float* __restrict__ w1s,
-    const float* __restrict__ w2s,
-    int64_t block_size_N,
-    int64_t block_size_K,
-    const float* __restrict__ topk_weights,
-    const int32_t* __restrict__ sorted_ids,
-    const int32_t* __restrict__ expert_ids,
-    const int32_t* __restrict__ offsets,
-    int64_t M,
-    int64_t N,
-    int64_t K,
-    int64_t E,
-    int64_t topk,
-    int64_t num_tokens_post_pad);
-
-// moe implementations for int4 w4a16
-template <typename scalar_t>
-void fused_experts_int4_w4a16_kernel_impl(
-    scalar_t* __restrict__ output,
-    scalar_t* __restrict__ ic0,
-    scalar_t* __restrict__ ic1,
-    scalar_t* __restrict__ ic2,
-    scalar_t* __restrict__ A_tmp,
-    scalar_t* __restrict__ B_tmp,
-    float* __restrict__ C_tmp,
-    const scalar_t* __restrict__ input,
-    const at::quint4x2* __restrict__ packed_w1,
-    const at::quint4x2* __restrict__ packed_w2,
-    const uint8_t* __restrict__ w1z,
-    const uint8_t* __restrict__ w2z,
-    const scalar_t* __restrict__ w1s,
-    const scalar_t* __restrict__ w2s,
-    int group_size,
-    const float* __restrict__ topk_weights,
-    const int32_t* __restrict__ sorted_ids,
-    const int32_t* __restrict__ expert_ids,
-    const int32_t* __restrict__ offsets,
-    int64_t M,
-    int64_t N,
-    int64_t K,
-    int64_t E,
-    int64_t topk,
-    int64_t num_tokens_post_pad);
-
-// shared expert implememntation for int8 w8a8
-template <typename scalar_t>
-void shared_expert_int8_kernel_impl(
-    scalar_t* __restrict__ output,
-    scalar_t* __restrict__ ic1,
-    float* __restrict__ C_tmp,
-    uint8_t* __restrict__ Aq_tmp,
-    float* __restrict__ As_tmp,
-    const scalar_t* __restrict__ input,
-    const int8_t* __restrict__ packed_w1,
-    const int8_t* __restrict__ packed_w2,
-    const float* __restrict__ w1s,
-    const float* __restrict__ w2s,
-    const scalar_t* __restrict__ fused_experts_out,
-    float routed_scaling_factor,
-    int64_t M,
-    int64_t N,
-    int64_t K);
-
-template <typename scalar_t>
-void shared_expert_fp8_kernel_impl(
-    scalar_t* __restrict__ output,
-    scalar_t* __restrict__ ic0,
-    scalar_t* __restrict__ ic1,
-    scalar_t* __restrict__ B_tmp,
-    float* __restrict__ C_tmp,
-    const scalar_t* __restrict__ input,
-    const at::Float8_e4m3fn* __restrict__ packed_w1,
-    const at::Float8_e4m3fn* __restrict__ packed_w2,
-    const float* __restrict__ w1s,
-    const float* __restrict__ w2s,
-    int64_t block_size_N,
-    int64_t block_size_K,
-    const scalar_t* __restrict__ fused_experts_out,
-    float routed_scaling_factor,
-    int64_t M,
-    int64_t N,
-    int64_t K);
-
-// tinygemm interface
-template <typename scalar_t>
-void tinygemm_kernel(
-    const scalar_t* __restrict__ A,
-    const scalar_t* __restrict__ B,
-    scalar_t* __restrict__ C,
-    float* __restrict__ Ctmp,
-    int64_t M,
-    int64_t N,
-    int64_t K,
-    int64_t lda,
-    int64_t ldb,
-    int64_t ldc,
-    bool brg);
-
-template <typename scalar_t>
-void tinygemm_kernel(
-    const uint8_t* __restrict__ A,
-    const int8_t* __restrict__ B,
-    scalar_t* __restrict__ C,
-    int32_t* __restrict__ Ctmp,
-    const float* __restrict__ As,
-    const float* __restrict__ Bs,
-    int64_t M,
-    int64_t N,
-    int64_t K,
-    int64_t lda,
-    int64_t ldb,
-    int64_t ldc,
-    bool brg);
-
-template <typename scalar_t>
-void tinygemm_kernel(
-    const scalar_t* __restrict__ A,
-    const at::Float8_e4m3fn* __restrict__ B,
-    scalar_t* __restrict__ C,
-    scalar_t* __restrict__ Btmp,
-    float* __restrict__ Ctmp,
-    const float* __restrict__ scale,
-    int64_t M,
-    int64_t N,
-    int64_t K,
-    int64_t lda,
-    int64_t ldb,
-    int64_t ldc,
-    bool brg,
-    int64_t block_size_K);
-
-template <typename scalar_t>
-void tinygemm_kernel(
-    const scalar_t* __restrict__ A,
-    const at::quint4x2* __restrict__ B,
-    scalar_t* __restrict__ C,
-    const uint8_t* __restrict__ Bz,
-    const scalar_t* __restrict__ Bs,
-    scalar_t* __restrict__ Btmp,
-    float* __restrict__ Ctmp,
-    int64_t M,
-    int64_t N,
-    int64_t K,
-    int group_size,
-    int64_t lda,
-    int64_t ldb,
-    int64_t ldc,
-    int64_t strideBz,
-    int64_t strideBs,
-    bool brg);
-
-// TODO: debug print, remove me later
-inline void print_16x32i(const __m512i x) {
-  int32_t a[16];
-  _mm512_storeu_si512((__m512i *)a, x);
-
-  for (int i = 0; i < 16; i++){
-    std::cout << a[i] << " ";
-  }
-  std::cout << std::endl;
-}
-
-inline void print_16x32(const __m512 x) {
-  float a[16];
-  _mm512_storeu_ps((__m512 *)a, x);
-
-  for (int i = 0; i < 16; i++){
-    std::cout << a[i] << " ";
-  }
-  std::cout << std::endl;
-}
-
-
-inline void print_32x8u(const __m256i x) {
-  uint8_t a[32];
-  _mm256_storeu_si256((__m256i *)a, x);
-
-  for (int i = 0; i < 32; ++i) {
-    std::cout << int32_t(a[i]) << " ";
-  }
-  std::cout << std::endl;
-}

csrc/cpu/sgl-kernels/gemm_fp8.cpp

@@ -1,530 +0,0 @@
-// Adapted from
-// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu
-
-#include "common.h"
-#include "vec.h"
-#include "gemm.h"
-
-// clang-format off
-
-// we use 4x32 for BLOCK_M
-#define BLOCK_SIZE_M_SCALE 4
-
-namespace {
-
-template <typename scalar_t>
-inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ input, int64_t size) {
-  using bVec = at::vec::Vectorized<scalar_t>;
-  using fVec = at::vec::Vectorized<float>;
-  constexpr int kVecSize = bVec::size();
-
-  int64_t d;
-  #pragma GCC unroll 4
-  for (d = 0; d <= size - kVecSize; d += kVecSize) {
-    fVec data0 = fVec::loadu(input + d);
-    fVec data1 = fVec::loadu(input + d + fVec::size());
-    bVec out_vec = convert_from_float_ext<scalar_t>(data0, data1);
-    out_vec.store(out + d);
-  }
-  for (; d < size; ++d) {
-    out[d] = static_cast<scalar_t>(input[d]);
-  }
-}
-
-template <typename scalar_t>
-inline void copy_add_stub(scalar_t* __restrict__ out, const float* __restrict__ input, const float* __restrict__ bias, int64_t size) {
-  using bVec = at::vec::Vectorized<scalar_t>;
-  using fVec = at::vec::Vectorized<float>;
-  constexpr int kVecSize = bVec::size();
-
-  int64_t d;
-  #pragma GCC unroll 4
-  for (d = 0; d <= size - kVecSize; d += kVecSize) {
-    fVec data0 = fVec::loadu(input + d) + fVec::loadu(bias + d);
-    fVec data1 = fVec::loadu(input + d + fVec::size()) + fVec::loadu(bias + d + fVec::size());
-    bVec out_vec = convert_from_float_ext<scalar_t>(data0, data1);
-    out_vec.store(out + d);
-  }
-  for (; d < size; ++d) {
-    out[d] = static_cast<scalar_t>(input[d] + bias[d]);
-  }
-}
-
-inline void unpack_B(
-    at::BFloat16* __restrict__ Btmp,
-    const at::Float8_e4m3fn* __restrict__ packed_B,
-    int N,
-    int K,
-    int ldb,
-    int ldb_tmp,
-    float scale) {
-#if defined(CPU_CAPABILITY_AVX512)
-  // [K/2, N, 2]
-  const int K2 = K >> 1;
-  const int ldb2 = ldb; // ldb * 2 >> 1;
-  const uint16_t* b_ptr = reinterpret_cast<const uint16_t*>(packed_B);
-  const __m512 vd = _mm512_set1_ps(scale);
-
-  constexpr int BLOCK_N = block_size_n();
-  static_assert(BLOCK_N == 32);
-
-  // prefetch distance
-  constexpr int PREFETCH_SIZE_K = 64;
-
-#pragma GCC unroll 4
-  for (int k = 0; k < K2; ++k) {
-    __m512i b8 = _mm512_loadu_si512(b_ptr + k * ldb2);
-    if constexpr (PREFETCH_SIZE_K > 0) {
-      _mm_prefetch(b_ptr + (k + PREFETCH_SIZE_K) * ldb2, _MM_HINT_T0);
-    }
-
-    __m256i b8_0 = _mm512_extracti32x8_epi32(b8, 0);
-    __m256i b8_1 = _mm512_extracti32x8_epi32(b8, 1);
-
-    __m512bh bf16_0 = CVT_FP8_TO_BF16(b8_0);
-    __m512bh bf16_1 = CVT_FP8_TO_BF16(b8_1);
-
-    // Apply scale
-    __m512 f0_lo = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32((__m512i)bf16_0, 0));
-    __m512 f0_hi = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32((__m512i)bf16_0, 1));
-    __m512 f1_lo = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32((__m512i)bf16_1, 0));
-    __m512 f1_hi = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32((__m512i)bf16_1, 1));
-
-    f0_lo = _mm512_mul_ps(f0_lo, vd);
-    f0_hi = _mm512_mul_ps(f0_hi, vd);
-    f1_lo = _mm512_mul_ps(f1_lo, vd);
-    f1_hi = _mm512_mul_ps(f1_hi, vd);
-
-    bf16_0 = _mm512_cvtne2ps_pbh(f0_hi, f0_lo);
-    bf16_1 = _mm512_cvtne2ps_pbh(f1_hi, f1_lo);
-
-    _mm512_storeu_si512(Btmp + k * ldb_tmp * 2 + 0, (__m512i)bf16_0);
-    _mm512_storeu_si512(Btmp + k * ldb_tmp * 2 + 32, (__m512i)bf16_1);
-  }
-#else
-  TORCH_CHECK(false, "unpack_B: scalar path not implemented!");
-#endif
-}
-
-template <typename scalar_t, typename packed_t, bool has_bias, int BLOCK_M, int BLOCK_N>
-struct tinygemm_kernel_nn {
-  static inline void apply(
-      const scalar_t* __restrict__ A, const packed_t* __restrict__ B, scalar_t* __restrict__ C,
-      const float* __restrict__ bias, const float* __restrict__ scale, int K, int lda, int ldb, int ldc, int64_t block_size_K) {
-    TORCH_CHECK(false, "tinygemm_kernel_nn: scalar path not implemented!");
-  }
-};
-
-#if defined(CPU_CAPABILITY_AVX512)
-template <bool has_bias, int BLOCK_M, int BLOCK_N>
-struct tinygemm_kernel_nn<at::BFloat16, at::Float8_e4m3fn, has_bias, BLOCK_M, BLOCK_N> {
-  static inline void apply(
-      const at::BFloat16* __restrict__ A, const at::Float8_e4m3fn* __restrict__ B, at::BFloat16* __restrict__ C,
-      const float* __restrict__ bias, const float* __restrict__ scale, int K, int lda, int ldb, int ldc, int64_t block_size_K) {
-
-    constexpr int ROWS = BLOCK_M;
-    constexpr int COLS = BLOCK_N / 16;
-
-    const int KB = div_up(K, BLOCK_K);
-
-    // prefetch distance
-    constexpr int PREFETCH_SIZE_K = 64;
-    constexpr int PREFETCH_SIZE_KB = 1;
-
-    __m512bh va;
-    __m512bh vb[COLS];
-    __m512 vc[ROWS * COLS];
-    __m512 vsum[ROWS * COLS];
-
-    // block quant scale
-    __m512 vscale;
-
-    auto loadc = [&](auto i) {
-      constexpr int col = i % COLS;
-      if constexpr (has_bias) {
-        vc[i] = _mm512_loadu_ps(bias + col * 16);
-      } else {
-        vc[i] = _mm512_setzero_ps();
-      }
-    };
-    Unroll<ROWS * COLS>{}(loadc);
-
-    const int lda2 = lda >> 1;
-    const int ldb2 = ldb; // ldb * 2 >> 1;
-    const float* a_ptr = reinterpret_cast<const float*>(A);
-    const uint16_t* b_ptr = reinterpret_cast<const uint16_t*>(B);
-
-    auto compute = [&](auto i, int k) {
-      constexpr int row = i / COLS;
-      constexpr int col = i % COLS;
-
-      if constexpr (col == 0) {
-        va = (__m512bh)(_mm512_set1_ps(a_ptr[row * lda2 + k]));
-        if constexpr (PREFETCH_SIZE_K > 0) {
-          _mm_prefetch(a_ptr + row * lda2 + k + PREFETCH_SIZE_K, _MM_HINT_T0);
-        }
-      }
-      if constexpr (row == 0) {
-        if constexpr (col % 2 == 0) {
-          __m512i b8 = _mm512_loadu_si512(b_ptr + k * ldb2 + col * 16);
-          if constexpr (PREFETCH_SIZE_K > 0) {
-            _mm_prefetch(b_ptr + (k + PREFETCH_SIZE_K) * ldb2 + col * 16, _MM_HINT_T0);
-          }
-          vb[col + 0] = CVT_FP8_TO_BF16(_mm512_extracti32x8_epi32(b8, 0));
-          vb[col + 1] = CVT_FP8_TO_BF16(_mm512_extracti32x8_epi32(b8, 1));
-        }
-      }
-      vsum[i] = _mm512_dpbf16_ps(vsum[i], va, vb[col]);
-    };
-
-    constexpr int BLOCK_K2 = BLOCK_K >> 1;
-    for (int kb = 0; kb < KB; ++kb) {
-      int kb_start = kb * BLOCK_K2;
-      int kb_end = std::min(K, kb_start + BLOCK_K2);
-      // 1. load scale vector
-      vscale = _mm512_set1_ps(scale[kb]);
-      if constexpr (PREFETCH_SIZE_KB > 0) {
-        _mm_prefetch(scale + kb + PREFETCH_SIZE_KB, _MM_HINT_T0);
-      }
-      // 2. zero vsum for each block
-      Unroll<ROWS * COLS>{}([&](auto i) {
-        vsum[i] = _mm512_setzero_ps();
-      });
-      // 3. accumulate across each block
-      for (int k = kb_start; k < kb_end; ++k) {
-        Unroll<ROWS * COLS>{}(compute, k);
-      }
-      // 4. apply scale
-      Unroll<ROWS * COLS>{}([&](auto i) {
-        vc[i] = _mm512_fmadd_ps(vsum[i], vscale, vc[i]);
-      });
-    }
-
-    auto storec = [&](auto i) {
-      constexpr int row = i / COLS;
-      constexpr int col = i % COLS;
-      // for COLS = 2,4 use 512bit store
-      if constexpr (col % 2 == 0) {
-        _mm512_storeu_si512(
-            reinterpret_cast<__m512i*>((C + row * ldc + col * 16)),
-            (__m512i)(_mm512_cvtne2ps_pbh(vc[row * COLS + col + 1], vc[row * COLS + col])));
-      }
-    };
-    Unroll<ROWS * COLS>{}(storec);
-  }
-};
-#endif
-
-#define LAUNCH_TINYGEMM_KERNEL_NN(MB_SIZE, NB_SIZE)                          \
-    tinygemm_kernel_nn<scalar_t, at::Float8_e4m3fn, has_bias, MB_SIZE, NB_SIZE>::apply(         \
-        A + mb_start * lda, B + nb_start * 2, C + mb_start * ldc + nb_start, \
-        has_bias ? bias + nb_start : nullptr, scale, K, lda, ldb, ldc, block_size_K);
-
-template <typename scalar_t, typename packed_t, bool has_bias>
-struct brgemm {
-  static inline void apply(
-      const scalar_t* __restrict__ A,
-      const packed_t* __restrict__ B,
-      scalar_t* __restrict__ C,
-      scalar_t* __restrict__ Btmp,
-      float* __restrict__ Ctmp,
-      const float* __restrict__ bias,
-      const float* __restrict__ scale,
-      int M,
-      int N,
-      int K,
-      int lda,
-      int ldb,
-      int ldc) {
-    TORCH_CHECK(false, "struct brgemm: primary template not implemented!");
-  }
-};
-
-template <bool has_bias>
-struct brgemm<at::BFloat16, at::Float8_e4m3fn, has_bias> {
-  static inline void apply(
-      const at::BFloat16* __restrict__ A,
-      const at::Float8_e4m3fn* __restrict__ B,
-      at::BFloat16* __restrict__ C,
-      at::BFloat16* __restrict__ Btmp,
-      float* __restrict__ Ctmp,
-      const float* __restrict__ bias,
-      const float* __restrict__ scale,
-      int M,
-      int N,
-      int K,
-      int lda,
-      int ldb,
-      int ldc) {
-
-    constexpr int BLOCK_N = block_size_n();
-
-    // [K, BLOCK_N] -> [K / 2, BLOCK_N * 2]
-    const int ldb_tmp = BLOCK_N;
-
-    for (int k = 0; k < K; k += BLOCK_K) {
-      int kb_size = std::min(BLOCK_K, K - k);
-
-      int idx = k >> 7; // k / BLOCK_K where BLOCK_K = 128
-      unpack_B(Btmp + k * ldb_tmp, B + k * ldb, N, kb_size, ldb, ldb_tmp, scale[idx]);
-    }
-
-    at::native::cpublas::brgemm(
-        M, N, K, lda, ldb_tmp, BLOCK_N, /* add_C */ false, A, Btmp, Ctmp);
-
-    // copy from Ctmp to C
-    for (int m = 0; m < M; ++m) {
-      if constexpr (has_bias) {
-        copy_add_stub(C + m * ldc, Ctmp + m * BLOCK_N, bias, N);
-      } else {
-        copy_stub(C + m * ldc, Ctmp + m * BLOCK_N, N);
-      }
-    }
-  }
-};
-
-template <typename scalar_t, bool has_bias>
-void tinygemm_kernel(
-    const scalar_t* __restrict__ A,
-    const at::Float8_e4m3fn* __restrict__ B,
-    scalar_t* __restrict__ C,
-    scalar_t* __restrict__ Btmp,
-    float* __restrict__ Ctmp,
-    const float* __restrict__ scale,
-    const float* __restrict__ bias,
-    int64_t M,
-    int64_t N,
-    int64_t K,
-    int64_t lda,
-    int64_t ldb,
-    int64_t ldc,
-    bool brg,
-    int64_t block_size_K) {
-
-  if (brg) {
-    brgemm<scalar_t, at::Float8_e4m3fn, has_bias>::apply(
-        A, B, C, Btmp, Ctmp, bias, scale, M, N, K, lda, ldb, ldc);
-    return;
-  }
-
-  // pattern: 1-4-16
-  constexpr int64_t BLOCK_M = 4;
-  constexpr int64_t BLOCK_N = 64;
-  const int64_t MB = div_up(M, BLOCK_M);
-  const int64_t NB = div_up(N, BLOCK_N);
-  for (int mb = 0; mb < MB; ++mb) {
-    int64_t mb_start = mb * BLOCK_M;
-    int64_t mb_size = std::min(BLOCK_M, M - mb_start);
-    for (int64_t nb = 0; nb < NB; ++nb) {
-      int64_t nb_start = nb * BLOCK_N;
-      int64_t nb_size = std::min(BLOCK_N, N - nb_start);
-
-      switch(mb_size << 4 | nb_size >> 4) {
-        case 0x12: LAUNCH_TINYGEMM_KERNEL_NN(1, 32); break;
-        case 0x22: LAUNCH_TINYGEMM_KERNEL_NN(2, 32); break;
-        case 0x32: LAUNCH_TINYGEMM_KERNEL_NN(3, 32); break;
-        case 0x42: LAUNCH_TINYGEMM_KERNEL_NN(4, 32); break;
-        default: TORCH_CHECK(false, "Unexpected block size, ", mb_size, "x", "nb_size");
-      }
-    }
-  }
-}
-
-template <typename scalar_t>
-void fp8_scaled_mm_kernel_impl(
-    scalar_t* __restrict__ out,
-    const scalar_t* __restrict__ mat1,
-    const at::Float8_e4m3fn* __restrict__ mat2,
-    const float* __restrict__ scales2,
-    const float* __restrict__ bias,
-    scalar_t* __restrict__ buffer,
-    int64_t M,
-    int64_t N,
-    int64_t K,
-    int64_t mat1_strideM,
-    int64_t out_strideM,
-    int64_t block_size_N,
-    int64_t block_size_K,
-    int64_t buffer_size_per_thread) {
-
-  constexpr int64_t BLOCK_M = block_size_m() * BLOCK_SIZE_M_SCALE;
-  constexpr int64_t BLOCK_N = block_size_n();
-  const int64_t MB = div_up(M, BLOCK_M);
-  const int64_t NB = div_up(N, BLOCK_N);
-
-  const int64_t scale_size_K = div_up(K, block_size_K);
-  const int64_t blocks_n_per_group = block_size_N / BLOCK_N;
-
-  const bool use_brgemm = can_use_brgemm<at::Float8_e4m3fn>(M);
-
-  // parallel on [MB, NB]
-  AT_DISPATCH_BOOL(bias != nullptr, has_bias, [&] {
-    at::parallel_for(0, MB * NB, 0, [&](int64_t begin, int64_t end) {
-      int64_t mb{0}, nb{0};
-      data_index_init(begin, mb, MB, nb, NB);
-
-      int tid = at::get_thread_num();
-      scalar_t* __restrict__ Btmp = buffer + tid * buffer_size_per_thread;
-      float* __restrict__ Ctmp = (float*)((void*)(Btmp + BLOCK_N * K));
-
-      for (int64_t i = begin; i < end; ++i) {
-        UNUSED(i);
-        const float* scale_ptr = scales2 + (nb / blocks_n_per_group) * scale_size_K;
-
-        int64_t mb_start = mb * BLOCK_M;
-        int64_t mb_size = std::min(M - mb_start, BLOCK_M);
-        int64_t nb_start = nb * BLOCK_N;
-        int64_t nb_size = std::min(N - nb_start, BLOCK_N);
-
-        tinygemm_kernel<scalar_t, has_bias>(
-            /*   A            */ mat1 + mb_start * mat1_strideM,
-            /*   B            */ mat2 + nb_start * K, // nb * BLOCK_N * K
-            /*   C            */ out + mb_start * out_strideM + nb_start,
-            /*   Btmp         */ Btmp,
-            /*   Ctmp         */ Ctmp,
-            /*   scale        */ scale_ptr,
-            /*   bias         */ bias + nb_start,
-            /*   M            */ mb_size,
-            /*   N            */ nb_size,
-            /*   K            */ K,
-            /*   lda          */ mat1_strideM,
-            /*   ldb          */ nb_size,
-            /*   ldc          */ out_strideM,
-            /*   brg          */ use_brgemm,
-            /*   block_size_K */ block_size_K);
-
-        // move to the next index
-        data_index_step(mb, MB, nb, NB);
-      }
-
-      if (use_brgemm) {
-        at::native::cpublas::brgemm_release();
-      }
-    });
-  });
-}
-
-} // anonymous namespace
-
-// tinygemm interface
-template <typename scalar_t>
-void tinygemm_kernel(
-    const scalar_t* __restrict__ A,
-    const at::Float8_e4m3fn* __restrict__ B,
-    scalar_t* __restrict__ C,
-    scalar_t* __restrict__ Btmp,
-    float* __restrict__ Ctmp,
-    const float* __restrict__ scale,
-    int64_t M,
-    int64_t N,
-    int64_t K,
-    int64_t lda,
-    int64_t ldb,
-    int64_t ldc,
-    bool brg,
-    int64_t block_size_K) {
-  tinygemm_kernel<scalar_t, false>(A, B, C, Btmp, Ctmp, scale, nullptr, M, N, K, lda, ldb, ldc, brg, block_size_K);
-}
-
-#define INSTANTIATE_TINYGEMM_TEMPLATE(TYPE)    \
-  template void tinygemm_kernel<TYPE>(         \
-      const TYPE* __restrict__ A,              \
-      const at::Float8_e4m3fn* __restrict__ B, \
-      TYPE* __restrict__ C,                    \
-      TYPE* __restrict__ Btmp,                 \
-      float* __restrict__ Ctmp,                \
-      const float* __restrict__ scale,         \
-      int64_t M,                               \
-      int64_t N,                               \
-      int64_t K,                               \
-      int64_t lda,                             \
-      int64_t ldb,                             \
-      int64_t ldc,                             \
-      bool brg,                                \
-      int64_t block_size_K)
-
-INSTANTIATE_TINYGEMM_TEMPLATE(at::BFloat16);
-INSTANTIATE_TINYGEMM_TEMPLATE(at::Half);
-
-at::Tensor fp8_scaled_mm_cpu(at::Tensor& mat1, at::Tensor& mat2, at::Tensor& scales2,
-    std::vector<int64_t> block_size, std::optional<at::Tensor>& bias,
-    at::ScalarType out_dtype, bool is_vnni) {
-  RECORD_FUNCTION("sgl-kernel::fp8_scaled_mm_cpu", std::vector<c10::IValue>({mat1, mat2, scales2, block_size, bias}));
-
-  auto packed_w = is_vnni ? mat2 : convert_weight_packed(mat2);
-
-  CHECK_LAST_DIM_CONTIGUOUS_INPUT(mat1);
-  CHECK_INPUT(mat2);
-  CHECK_INPUT(scales2);
-  TORCH_CHECK(scales2.scalar_type() == at::kFloat,
-      "fp8_scaled_mm_cpu: expect scales2 to be float32.");
-
-  int64_t M = mat1.size(0);
-  int64_t N = mat2.size(0);
-  int64_t K = mat2.size(1);
-
-  CHECK_EQ(mat1.size(1), K);
-  CHECK_DIM(2, mat1);
-  CHECK_DIM(2, mat2);
-
-  TORCH_CHECK(block_size.size() == 2,
-      "fp8_scaled_mm_cpu: expect block_size.size() to be 2.");
-
-  int64_t block_size_N = block_size[0];
-  int64_t block_size_K = block_size[1];
-
-  constexpr int64_t BLOCK_M = block_size_m() * BLOCK_SIZE_M_SCALE;
-  constexpr int64_t BLOCK_N = block_size_n();
-  TORCH_CHECK(block_size_N % BLOCK_N == 0, "fp8_scaled_mm_cpu: expect block_size_N to be multiples of BLOCK_N");
-  TORCH_CHECK(block_size_K == BLOCK_K, "fp8_scaled_mm_cpu: expect block_size_K equals to BLOCK_K");
-  CHECK_EQ(scales2.size(0), div_up(N, block_size_N));
-  CHECK_EQ(scales2.size(1), div_up(K, block_size_K));
-
-  const auto st = mat1.scalar_type();
-  TORCH_CHECK(st == at::kBFloat16 || st == at::kHalf,
-      "fp8_scaled_mm_cpu: expect A to be bfloat16 or half.");
-  TORCH_CHECK(st == out_dtype,
-      "fp8_scaled_mm_cpu: expect A has same dtype with out_dtype.");
-  TORCH_CHECK(mat2.scalar_type() == at::kFloat8_e4m3fn,
-      "fp8_scaled_mm_cpu: expect mat2 to be fp8_e4m3.");
-  TORCH_CHECK(scales2.scalar_type() == at::kFloat,
-      "fp8_scaled_mm_cpu: expect scales to be float32.");
-  auto out = at::empty({M, N}, mat1.options().dtype(out_dtype));
-
-  // strides
-  int64_t mat1_strideM = mat1.stride(0);
-  int64_t out_strideM = out.stride(0);
-
-  const bool has_bias = bias.has_value();
-  const float* bias_data = nullptr;
-  if (has_bias) {
-    CHECK_EQ(bias.value().size(0), N);
-    bias_data = bias.value().data_ptr<float>();
-  }
-
-  // Btmp : [T, BLOCK_N * K]
-  // Ctmp : [T, BLOCK_M * BLOCK_N]
-  int num_threads = at::get_num_threads();
-  int64_t size_per_thread = BLOCK_N * K + BLOCK_M * BLOCK_N * 2;
-  auto buffer = at::empty({num_threads, size_per_thread}, mat1.options());
-
-  AT_DISPATCH_REDUCED_FLOATING_TYPES(out_dtype, "fp8_scaled_mm_kernel_impl", [&] {
-    fp8_scaled_mm_kernel_impl<scalar_t>(
-        out.data_ptr<scalar_t>(),
-        mat1.data_ptr<scalar_t>(),
-        packed_w.data_ptr<at::Float8_e4m3fn>(),
-        scales2.data_ptr<float>(),
-        bias_data,
-        buffer.data_ptr<scalar_t>(),
-        M,
-        N,
-        K,
-        mat1_strideM,
-        out_strideM,
-        block_size_N,
-        block_size_K,
-        size_per_thread);
-  });
-
-  return out;
-}

csrc/cpu/sgl-kernels/gemm_int8.cpp

@@ -1,440 +0,0 @@
-// Adapted from
-// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu
-
-#include "common.h"
-#include "vec.h"
-#include "gemm.h"
-
-// clang-format off
-
-namespace {
-
-template <typename scalar_t, bool has_bias, int BLOCK_M, int BLOCK_N>
-struct tinygemm_kernel_nn {
-  static inline void apply(
-      const uint8_t* __restrict__ A, const int8_t* __restrict__ B, scalar_t* __restrict__ C,
-      const float* __restrict__ As, const float* __restrict__ Bs, const int32_t* __restrict__ Bcomp,
-      const float* __restrict__ bias, int64_t K, int64_t lda, int64_t ldb, int64_t ldc) {
-    TORCH_CHECK(false, "tinygemm_kernel_nn: scalar path not implemented!");
-  }
-};
-
-#if defined(CPU_CAPABILITY_AVX512)
-template <bool has_bias, int BLOCK_M, int BLOCK_N>
-struct tinygemm_kernel_nn<at::BFloat16, has_bias, BLOCK_M, BLOCK_N> {
-  static inline void apply(
-      const uint8_t* __restrict__ A, const int8_t* __restrict__ B, at::BFloat16* __restrict__ C,
-      const float* __restrict__ As, const float* __restrict__ Bs, const int32_t* __restrict__ Bcomp,
-      const float* __restrict__ bias, int64_t K, int64_t lda, int64_t ldb, int64_t ldc) {
-
-    constexpr int ROWS = BLOCK_M;
-    constexpr int COLS = BLOCK_N / 16;
-    static_assert(COLS % 2 == 0);
-
-    // prefetch distance
-    constexpr int PREFETCH_SIZE_K = 0;
-
-    __m512i va;
-    __m512i vb[COLS];
-    __m512i vc[ROWS * COLS];
-    __m512i vcomp[COLS];
-    __m512  vd0;
-    __m512  vd1[COLS];
-
-    // oops! 4x4 spills but luckly we use 4x2
-    __m512 vbias[COLS];
-
-    // [NOTE]: s8s8 igemm compensation in avx512-vnni
-    //
-    // avx512-vnni has no s8s8, so we need to change s8s8 to u8s8 with compensate:
-    //
-    //   a * b = (a + 128) * b - 128 * b
-    //   s   s       u       s    u    s
-    //
-    // 1) 128 * b is pre-computed when packing B to vnni formats
-    // 2) a + 128 is fused when dynamically quantize A
-    //
-    auto loadc = [&](auto i) {
-      vc[i] = _mm512_set1_epi32(0);
-    };
-    Unroll<ROWS * COLS>{}(loadc);
-
-    const int64_t K4 = K >> 2;
-    const int64_t lda4 = lda >> 2;
-    const int64_t ldb4 = ldb; // ldb * 4 >> 2;
-    const int32_t* a_ptr = reinterpret_cast<const int32_t*>(A);
-    const int32_t* b_ptr = reinterpret_cast<const int32_t*>(B);
-
-    auto compute = [&](auto i, int64_t k) {
-      constexpr int row = i / COLS;
-      constexpr int col = i % COLS;
-
-      if constexpr (col == 0) {
-        va = _mm512_set1_epi32(a_ptr[row * lda4 + k]);
-      }
-      if constexpr (row == 0) {
-        vb[col] = _mm512_loadu_si512(b_ptr + k * ldb4 + col * 16);
-        if constexpr (PREFETCH_SIZE_K > 0) {
-          _mm_prefetch(b_ptr + (k + PREFETCH_SIZE_K) * ldb4 + col * 16, _MM_HINT_T0);
-        }
-      }
-      vc[i] = _mm512_dpbusd_epi32(vc[i], va, vb[col]);
-    };
-    for (int64_t k = 0; k < K4; ++k) {
-      Unroll<ROWS * COLS>{}(compute, k);
-    }
-
-    auto storec = [&](auto i) {
-      constexpr int row = i / COLS;
-      constexpr int col = i % COLS;
-
-      // load a scale
-      if constexpr(col == 0) {
-        vd0 = _mm512_set1_ps(As[row]);
-      }
-      // load b scale and vcomp per 2 vectors
-      // also load bias if any
-      if constexpr (row == 0) {
-        if constexpr (col % 2 == 0) {
-          vd1[col + 0] = _mm512_loadu_ps(Bs + col * 16);
-          vd1[col + 1] = _mm512_loadu_ps(Bs + col * 16 + 16);
-          vcomp[col + 0] = _mm512_loadu_si512(Bcomp + col * 16);
-          vcomp[col + 1] = _mm512_loadu_si512(Bcomp + col * 16 + 16);
-          if constexpr (has_bias) {
-            vbias[col + 0] = _mm512_loadu_ps(bias + col * 16);
-            vbias[col + 1] = _mm512_loadu_ps(bias + col * 16 + 16);
-          }
-        }
-      }
-
-      // for COLS = 2, 4 use 512bit store
-      if constexpr (col % 2 == 0) {
-        __m512 vc0 = _mm512_cvtepi32_ps(_mm512_sub_epi32(vc[row * COLS + col + 0], vcomp[col + 0]));
-        __m512 vc1 = _mm512_cvtepi32_ps(_mm512_sub_epi32(vc[row * COLS + col + 1], vcomp[col + 1]));
-        if constexpr (has_bias) {
-          vc0 = _mm512_fmadd_ps(_mm512_mul_ps(vc0, vd0), vd1[col + 0], vbias[col + 0]);
-          vc1 = _mm512_fmadd_ps(_mm512_mul_ps(vc1, vd0), vd1[col + 1], vbias[col + 1]);
-        } else {
-          vc0 = _mm512_mul_ps(_mm512_mul_ps(vc0, vd0), vd1[col + 0]);
-          vc1 = _mm512_mul_ps(_mm512_mul_ps(vc1, vd0), vd1[col + 1]);
-        }
-
-        _mm512_storeu_si512(
-            reinterpret_cast<__m512i*>((C + row * ldc + col * 16)),
-            (__m512i)(_mm512_cvtne2ps_pbh(vc1, vc0)));
-      }
-    };
-    Unroll<ROWS * COLS>{}(storec);
-  }
-};
-#endif
-
-#define LAUNCH_TINYGEMM_KERNEL_NN(MB_SIZE, NB_SIZE)                          \
-    tinygemm_kernel_nn<scalar_t, has_bias, MB_SIZE, NB_SIZE>::apply(         \
-        A + mb_start * lda, B + nb_start * 4, C + mb_start * ldc + nb_start, \
-        As + mb_start, Bs + nb_start, Bcomp + nb_start,                      \
-        has_bias ? bias + nb_start : nullptr, K, lda, ldb, ldc);
-
-template <typename scalar_t, bool has_bias>
-void tinygemm_kernel(
-    const uint8_t* __restrict__ A,
-    const int8_t* __restrict__ B,
-    scalar_t* __restrict__ C,
-    int32_t* __restrict__ Ctmp,
-    const float* __restrict__ As,
-    const float* __restrict__ Bs,
-    const float* __restrict__ bias,
-    int64_t M,
-    int64_t N,
-    int64_t K,
-    int64_t lda,
-    int64_t ldb,
-    int64_t ldc,
-    bool brg) {
-
-  // B compensation
-  const int32_t* Bcomp = reinterpret_cast<const int32_t*>(B + block_size_n() * K);
-
-  // pattern: 1-4-16
-  constexpr int64_t BLOCK_M = 4;
-  constexpr int64_t BLOCK_N = 64;
-  const int64_t MB = div_up(M, BLOCK_M);
-  const int64_t NB = div_up(N, BLOCK_N);
-  for (int64_t mb = 0; mb < MB; ++mb) {
-    int64_t mb_start = mb * BLOCK_M;
-    int64_t mb_size = std::min(BLOCK_M, M - mb_start);
-    for (int64_t nb = 0; nb < NB; ++nb) {
-      int64_t nb_start = nb * BLOCK_N;
-      int64_t nb_size = std::min(BLOCK_N, N - nb_start);
-
-      switch(mb_size << 4 | nb_size >> 4) {
-        // mb_size = 1
-        case 0x12: LAUNCH_TINYGEMM_KERNEL_NN(1, 32); break;
-        case 0x14: LAUNCH_TINYGEMM_KERNEL_NN(1, 64); break;
-        // mb_size = 2
-        case 0x22: LAUNCH_TINYGEMM_KERNEL_NN(2, 32); break;
-        case 0x24: LAUNCH_TINYGEMM_KERNEL_NN(2, 64); break;
-        // mb_size = 3
-        case 0x32: LAUNCH_TINYGEMM_KERNEL_NN(3, 32); break;
-        case 0x34: LAUNCH_TINYGEMM_KERNEL_NN(3, 64); break;
-        // mb_size = 4
-        case 0x42: LAUNCH_TINYGEMM_KERNEL_NN(4, 32); break;
-        case 0x44: LAUNCH_TINYGEMM_KERNEL_NN(4, 64); break;
-        default: TORCH_CHECK(false, "Unexpected block size, ", mb_size, "x", "nb_size");
-      }
-    }
-  }
-}
-
-template<typename scalar_t>
-void int8_scaled_mm_kernel_impl(
-    scalar_t* __restrict__ out,
-    const uint8_t* __restrict__ mat1,
-    const int8_t* __restrict__ mat2,
-    const float* __restrict__ scales1,
-    const float* __restrict__ scales2,
-    const float* __restrict__ bias,
-    int64_t M,
-    int64_t N,
-    int64_t K) {
-
-  constexpr int64_t BLOCK_M = block_size_m();
-  constexpr int64_t BLOCK_N = block_size_n();
-  const int64_t MB = div_up(M, BLOCK_M);
-  const int64_t NB = div_up(N, BLOCK_N);
-
-  // TODO: brgemm u8s8 depends on PyTorch 2.7 release.
-  const bool use_brgemm = false;
-
-  // K + 4 after compensation
-  const int64_t packed_row_size = get_row_size<int8_t>(K);
-
-  AT_DISPATCH_BOOL(bias != nullptr, has_bias, [&] {
-    at::parallel_for(0, MB * NB, 0, [&](int64_t begin, int64_t end) {
-      int64_t mb{0}, nb{0};
-      data_index_init(begin, mb, MB, nb, NB);
-
-      // for brgemm, use int32_t for accumulate
-      alignas(64) int32_t Ctmp[BLOCK_M * BLOCK_N];
-
-      for (int i = begin; i < end; ++i) {
-        UNUSED(i);
-        int mb_start = mb * BLOCK_M;
-        int mb_size = std::min(M - mb_start, BLOCK_M);
-        int nb_start = nb * BLOCK_N;
-        int nb_size = std::min(N - nb_start, BLOCK_N);
-
-        tinygemm_kernel<scalar_t, has_bias>(
-            /*   A */ mat1 + mb_start * K,
-            /*   B */ mat2 + nb_start * packed_row_size /* nb * BLOCK_N * (K + 4) */,
-            /*   C */ out + mb_start * N + nb_start,
-            /* Ctmp*/ Ctmp,
-            /*  As */ scales1 + mb_start,
-            /*  Bs */ scales2 + nb_start,
-            /* bias*/ bias + nb_start,
-            /*   M */ mb_size,
-            /*   N */ nb_size,
-            /*   K */ K,
-            /* lda */ K,
-            /* ldb */ nb_size,
-            /* ldc */ N,
-            /* brg */ use_brgemm);
-
-        // move to the next index
-        data_index_step(mb, MB, nb, NB);
-      }
-
-      if (use_brgemm) {
-        at::native::cpublas::brgemm_release();
-      }
-    });
-  });
-}
-
-} // anonymous namespace
-
-// tinygemm interface
-template <typename scalar_t>
-void tinygemm_kernel(const uint8_t* __restrict__ A, const int8_t* __restrict__ B, scalar_t* __restrict__ C,
-    int32_t* __restrict__ Ctmp,  const float* __restrict__ As, const float* __restrict__ Bs,
-    int64_t M, int64_t N, int64_t K, int64_t lda, int64_t ldb, int64_t ldc, bool brg) {
-  tinygemm_kernel<scalar_t, false>(A, B, C, Ctmp, As, Bs, nullptr, M, N, K, lda, ldb, ldc, brg);
-}
-
-#define INSTANTIATE_TINYGEMM_TEMPLATE(TYPE)                                                     \
-    template void tinygemm_kernel<TYPE>(                                                        \
-        const uint8_t* __restrict__ A, const int8_t* __restrict__ B, TYPE* __restrict__ C,      \
-        int32_t* __restrict__ Ctmp, const float* __restrict__ As, const float* __restrict__ Bs, \
-        int64_t M, int64_t N, int64_t K, int64_t lda, int64_t ldb, int64_t ldc, bool brg)
-
-INSTANTIATE_TINYGEMM_TEMPLATE(at::BFloat16);
-INSTANTIATE_TINYGEMM_TEMPLATE(at::Half);
-
-std::tuple<at::Tensor, at::Tensor> per_token_quant_int8_cpu(at::Tensor& A) {
-  RECORD_FUNCTION("sgl-kernel::per_token_quant_int8_cpu", std::vector<c10::IValue>({A}));
-
-  CHECK_LAST_DIM_CONTIGUOUS_INPUT(A);
-  CHECK_DIM(2, A);
-
-  int64_t M = A.size(0);
-  int64_t K = A.size(1);
-  int64_t lda = A.stride(0);
-
-  const auto st = A.scalar_type();
-  TORCH_CHECK(st == at::kBFloat16 || st == at::kHalf,
-      "per_token_quant_int8: expect A to be bfloat16 or half.");
-
-  auto Aq = at::empty({M, K}, A.options().dtype(at::kByte));
-  auto As = at::empty({M}, A.options().dtype(at::kFloat));
-
-  AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "per_token_quant_int8", [&] {
-    uint8_t* __restrict__ Aq_data = Aq.data_ptr<uint8_t>();
-    float* __restrict__ As_data = As.data_ptr<float>();
-    const scalar_t* __restrict__ A_data = A.data_ptr<scalar_t>();
-
-    at::parallel_for(0, M, 0, [&] (int64_t begin, int64_t end) {
-      for (int64_t m = begin; m < end; ++m) {
-        quantize_row_int8<scalar_t>(
-            Aq_data + m * K,
-            As_data[m],
-            A_data + m * lda,
-            K);
-      }
-    });
-  });
-  return std::make_tuple(Aq, As);
-}
-
-// weight     :  static, per-channel, symmetric
-// activation : dynamic,   per-token, symmetric
-//
-// mat1    : [M, K]
-// mat2    : [N, K]
-// scales1 : [M]
-// scales2 : [N]
-// bias    : [N]
-// out     : [M, N]
-//
-at::Tensor int8_scaled_mm_cpu(at::Tensor& mat1, at::Tensor& mat2,
-    at::Tensor& scales1, at::Tensor& scales2,
-    std::optional<at::Tensor>& bias, at::ScalarType out_dtype, bool is_vnni) {
-  RECORD_FUNCTION("sgl-kernel::int8_scaled_mm_cpu", std::vector<c10::IValue>({mat1, mat2, scales1, scales2, bias}));
-
-  auto packed_w = is_vnni ? mat2 : convert_weight_packed(mat2);
-
-  CHECK_INPUT(mat1);
-  CHECK_INPUT(mat2);
-  CHECK_INPUT(scales1);
-  CHECK_INPUT(scales2);
-  CHECK_DIM(2, mat1);
-  CHECK_DIM(2, mat2);
-
-  int64_t M = mat1.size(0);
-  int64_t N = mat2.size(0);
-  int64_t K = mat1.size(1);
-
-  // see [NOTE]: s8s8 igemm compensation in avx512-vnni
-  CHECK_EQ(mat2.size(1), (int64_t)(is_vnni ? K + sizeof(int32_t) : K));
-  CHECK_EQ(scales1.numel(), M);
-  CHECK_EQ(scales2.numel(), N);
-
-  TORCH_CHECK(mat1.scalar_type() == at::kByte, "int8_scaled_mm: expect mat1 to be uint8.");
-  TORCH_CHECK(mat2.scalar_type() == at::kChar, "int8_scaled_mm: expect mat2 to be int8.");
-  TORCH_CHECK(scales1.scalar_type() == at::kFloat && scales2.scalar_type() == at::kFloat,
-      "int8_scaled_mm: expect scales to be float32.");
-
-  auto out = at::empty({M, N}, mat1.options().dtype(out_dtype));
-
-  const bool has_bias = bias.has_value();
-  const float* bias_data = nullptr;
-  if (has_bias) {
-    CHECK_EQ(bias.value().size(0), N);
-    bias_data = bias.value().data_ptr<float>();
-  }
-
-  AT_DISPATCH_REDUCED_FLOATING_TYPES(out_dtype, "int8_scaled_mm_kernel_impl", [&] {
-    int8_scaled_mm_kernel_impl<scalar_t>(
-        out.data_ptr<scalar_t>(),
-        mat1.data_ptr<uint8_t>(),
-        packed_w.data_ptr<int8_t>(),
-        scales1.data_ptr<float>(),
-        scales2.data_ptr<float>(),
-        bias_data,
-        M,
-        N,
-        K);
-  });
-  return out;
-}
-
-// fused `per_token_quant_int8_cpu` and `int8_scaled_mm_cpu`
-at::Tensor int8_scaled_mm_with_quant(at::Tensor& mat1, at::Tensor& mat2, at::Tensor& scales2,
-    const std::optional<at::Tensor>& bias, at::ScalarType out_dtype, bool is_vnni) {
-  RECORD_FUNCTION("sgl-kernel::int8_scaled_mm_cpu", std::vector<c10::IValue>({mat1, mat2, scales2, bias}));
-
-  auto packed_w = is_vnni ? mat2 : convert_weight_packed(mat2);
-
-  CHECK_LAST_DIM_CONTIGUOUS_INPUT(mat1);
-  CHECK_INPUT(mat2);
-  CHECK_INPUT(scales2);
-  CHECK_DIM(2, mat1);
-  CHECK_DIM(2, mat2);
-
-  int64_t M = mat1.size(0);
-  int64_t N = mat2.size(0);
-  int64_t K = mat1.size(1);
-  int64_t lda = mat1.stride(0);
-
-  // see [NOTE]: s8s8 igemm compensation in avx512-vnni
-  CHECK_EQ(mat2.size(1), (int64_t)(is_vnni ? K + sizeof(int32_t) : K));
-  CHECK_EQ(scales2.numel(), N);
-
-  const auto st = mat1.scalar_type();
-  TORCH_CHECK(st == at::kBFloat16 || st == at::kHalf,
-      "int8_scaled_mm_with_quant: expect A to be bfloat16 or half.");
-  TORCH_CHECK(st == out_dtype,
-      "int8_scaled_mm_with_quant: expect A has same dtype with out_dtype.");
-  TORCH_CHECK(mat2.scalar_type() == at::kChar,
-      "int8_scaled_mm_with_quant: expect mat2 to be int8.");
-  TORCH_CHECK(scales2.scalar_type() == at::kFloat,
-      "int8_scaled_mm_with_quant: expect scales to be float32.");
-
-  const int64_t buffer_size = M * K + M * sizeof(float);
-  auto buffer = at::empty({buffer_size}, mat1.options().dtype(at::kByte));
-  auto out = at::empty({M, N}, mat1.options().dtype(out_dtype));
-
-  const bool has_bias = bias.has_value();
-  const float* bias_data = nullptr;
-  if (has_bias) {
-    CHECK_EQ(bias.value().size(0), N);
-    bias_data = bias.value().data_ptr<float>();
-  }
-
-  AT_DISPATCH_REDUCED_FLOATING_TYPES(out_dtype, "int8_scaled_mm_with_quant_kernel_impl", [&] {
-    uint8_t* __restrict__ Aq_data = buffer.data_ptr<uint8_t>();
-    float* __restrict__ As_data = (float*)((void*)(Aq_data + M * K));
-    const scalar_t* __restrict__ A_data = mat1.data_ptr<scalar_t>();
-
-    at::parallel_for(0, M, 0, [&] (int64_t begin, int64_t end) {
-      for (int64_t m = begin; m < end; ++m) {
-        quantize_row_int8<scalar_t>(
-            Aq_data + m * K,
-            As_data[m],
-            A_data + m * lda,
-            K);
-      }
-    });
-
-    int8_scaled_mm_kernel_impl<scalar_t>(
-        out.data_ptr<scalar_t>(),
-        Aq_data,
-        packed_w.data_ptr<int8_t>(),
-        As_data,
-        scales2.data_ptr<float>(),
-        bias_data,
-        M,
-        N,
-        K);
-  });
-  return out;
-}

csrc/cpu/sgl-kernels/moe_fp8.cpp

@@ -1,502 +0,0 @@
-// Adapted from
-// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu
-
-#include "common.h"
-#include "gemm.h"
-#include "vec.h"
-
-// clang-format off
-
-namespace {
-
-template <typename scalar_t>
-inline void copy_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, int64_t size) {
-  using Vec = at::vec::Vectorized<scalar_t>;
-  // no remainder
-  #pragma GCC unroll 4
-  for (int64_t d = 0; d < size; d += Vec::size()) {
-    Vec data = Vec::loadu(input + d);
-    data.store(out + d);
-  }
-}
-
-template <typename scalar_t>
-inline void copy_mul_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, float weight, int64_t size) {
-  using bVec = at::vec::Vectorized<scalar_t>;
-  using fVec = at::vec::Vectorized<float>;
-  constexpr int kVecSize = bVec::size();
-  const fVec weight_vec = fVec(weight);
-  int64_t d;
-  #pragma GCC unroll 4
-  for (d = 0; d <= size - kVecSize; d += kVecSize) {
-    bVec x = bVec::loadu(input + d);
-    fVec x0, x1;
-    std::tie(x0, x1) = at::vec::convert_to_float(x);
-    x0 = x0 * weight_vec;
-    x1 = x1 * weight_vec;
-    bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1);
-    out_vec.store(out + d);
-  }
-  for (; d < size; ++d) {
-    out[d] = static_cast<scalar_t>(input[d] * weight);
-  }
-}
-
-// acc from [topk, K] to [K]
-template <typename scalar_t>
-inline void sum_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, int64_t topk, int64_t K) {
-  using bVec = at::vec::Vectorized<scalar_t>;
-  using fVec = at::vec::Vectorized<float>;
-  constexpr int kVecSize = bVec::size();
-  if (topk == 1) {
-    // do copy for topk = 1
-    copy_stub(out, input, K);
-  } else {
-    // do sum for topk != 1
-    int64_t d;
-    #pragma GCC unroll 4
-    for (d = 0; d <= K - kVecSize; d += kVecSize) {
-      fVec sum_fvec0 = fVec(0.f);
-      fVec sum_fvec1 = fVec(0.f);
-      for (int t = 0; t < topk; ++t) {
-        bVec x_bvec = bVec::loadu(input + t * K + d);
-        fVec x_fvec0, x_fvec1;
-        std::tie(x_fvec0, x_fvec1) = at::vec::convert_to_float(x_bvec);
-
-        sum_fvec0 += x_fvec0;
-        sum_fvec1 += x_fvec1;
-      }
-      bVec out_bvec = convert_from_float_ext<scalar_t>(sum_fvec0, sum_fvec1);
-      out_bvec.store(out + d);
-    }
-    for (; d < K; ++d) {
-      float sum_val = 0.f;
-      for (int t = 0; t < topk; ++t) {
-        sum_val += static_cast<float>(input[t * K + d]);
-      }
-      out[d] = static_cast<scalar_t>(sum_val);
-    }
-  }
-}
-
-// out = input + input2 * scale
-template <typename scalar_t>
-inline void add_mul_stub(
-    scalar_t* __restrict__ out,
-    const scalar_t* __restrict__ input,
-    const scalar_t* __restrict__ input2,
-    float scale,
-    int64_t size) {
-  using bVec = at::vec::Vectorized<scalar_t>;
-  using fVec = at::vec::Vectorized<float>;
-  constexpr int kVecSize = bVec::size();
-  const fVec s_vec = fVec(scale);
-
-  int64_t d;
-#pragma GCC unroll 4
-  for (d = 0; d <= size - kVecSize; d += kVecSize) {
-    bVec x_bvec = bVec::loadu(input + d);
-    fVec x0, x1;
-    std::tie(x0, x1) = at::vec::convert_to_float(x_bvec);
-
-    bVec y_bvec = bVec::loadu(input2 + d);
-    fVec y0, y1;
-    std::tie(y0, y1) = at::vec::convert_to_float(y_bvec);
-
-    x0 = x0 + y0 * s_vec;
-    x1 = x1 + y1 * s_vec;
-    bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1);
-    out_vec.store(out + d);
-  }
-  for (; d < size; ++d) {
-    out[d] = static_cast<scalar_t>(input[d] + float(input2[d]) * scale);
-  }
-}
-
-template <typename scalar_t>
-inline void silu_and_mul_stub(
-    scalar_t* __restrict__ out,
-    const scalar_t* __restrict__ input,
-    const scalar_t* __restrict__ input2,
-    int64_t size) {
-  using bVec = at::vec::Vectorized<scalar_t>;
-  using fVec = at::vec::Vectorized<float>;
-  const fVec one = fVec(1.f);
-
-  // no remainder
-#pragma GCC unroll 4
-  for (int64_t d = 0; d < size; d += bVec::size()) {
-    bVec x = bVec::loadu(input + d);
-    fVec x0, x1;
-    std::tie(x0, x1) = at::vec::convert_to_float(x);
-    bVec y = bVec::loadu(input2 + d);
-    fVec y0, y1;
-    std::tie(y0, y1) = at::vec::convert_to_float(y);
-    x0 = x0 / (one + x0.neg().exp_u20());
-    x1 = x1 / (one + x1.neg().exp_u20());
-    x0 = x0 * y0;
-    x1 = x1 * y1;
-    bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1);
-    out_vec.store(out + d);
-  }
-}
-
-} // anonymous namespace
-
-template <typename scalar_t>
-void fused_experts_fp8_kernel_impl(
-    scalar_t* __restrict__ output,
-    scalar_t* __restrict__ ic0,
-    scalar_t* __restrict__ ic1,
-    scalar_t* __restrict__ ic2,
-    scalar_t* __restrict__ A_tmp,
-    scalar_t* __restrict__ B_tmp,
-    float* __restrict__ C_tmp,
-    const scalar_t* __restrict__ input,
-    const at::Float8_e4m3fn* __restrict__ packed_w1,
-    const at::Float8_e4m3fn* __restrict__ packed_w2,
-    const float* __restrict__ w1s,
-    const float* __restrict__ w2s,
-    int64_t block_size_N,
-    int64_t block_size_K,
-    const float* __restrict__ topk_weights,
-    const int32_t* __restrict__ sorted_ids,
-    const int32_t* __restrict__ expert_ids,
-    const int32_t* __restrict__ offsets,
-    int64_t M,
-    int64_t N,
-    int64_t K,
-    int64_t E,
-    int64_t topk,
-    int64_t num_tokens_post_pad) {
-
-  constexpr int64_t BLOCK_M = block_size_m();
-  constexpr int64_t BLOCK_N = block_size_n();
-
-  // stage 1: intermediate_cache0 = hidden_states @ w1
-  const int64_t MB = div_up(num_tokens_post_pad, BLOCK_M);
-  const int64_t NB = div_up(2 * N, BLOCK_N);
-  int64_t scale_size_N = div_up(2 * N, block_size_N);
-  int64_t scale_size_K = div_up(K, block_size_K);
-  int64_t blocks_n_per_group = block_size_N / BLOCK_N;
-
-  const int64_t stride_e = 2 * N * K;
-  const int64_t stride_n = K;
-
-  // here we only parallel on half of 2N to fuse silu_and_mul with gemm
-  at::parallel_for(0, MB * NB, 0, [&](int64_t begin, int64_t end) {
-    // get local pointers
-    int tid = at::get_thread_num();
-    scalar_t* __restrict__ A = A_tmp + tid * BLOCK_M * K;
-
-    bool is_brgemm_used = false;
-
-    for (int64_t i = begin; i < end; ++i) {
-      int64_t mb = i / NB;
-      int64_t nb = i % NB;
-
-      int64_t n_size = std::min(2 * N - nb * BLOCK_N, BLOCK_N);
-
-      // B shape [K, n_size] in vnni format
-      int32_t expert_id = expert_ids[mb];
-      const at::Float8_e4m3fn* __restrict__ B = packed_w1 + expert_id * stride_e + nb * BLOCK_N * stride_n;
-      const float* __restrict__ Bs = w1s + expert_id * scale_size_N * scale_size_K + (nb / blocks_n_per_group) * scale_size_K;
-
-      // 1.a load A
-      const int32_t* A_ids = sorted_ids + mb * BLOCK_M;
-      int64_t m_size = offsets[mb + 1] - offsets[mb];
-
-      const bool use_brgemm = can_use_brgemm<at::Float8_e4m3fn>(m_size);
-      is_brgemm_used = is_brgemm_used || use_brgemm;
-
-      for (int64_t m = 0; m < m_size; ++m) {
-        int32_t index = A_ids[m] / topk;
-        copy_stub(A + m * K, input + index * K, K);
-      }
-
-      const int64_t offset = offsets[mb];
-      tinygemm_kernel<scalar_t>(
-          /*   A            */ A,
-          /*   B            */ B,
-          /*   C            */ ic0 + offset * 2 * N + nb * BLOCK_N,
-          /*   Btmp         */ B_tmp + tid * BLOCK_N * std::max(K, N),
-          /*   Ctmp         */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N,
-          /*   scale        */ Bs,
-          /*   M            */ m_size,
-          /*   N            */ n_size,
-          /*   K            */ K,
-          /*   lda          */ K,
-          /*   ldb          */ n_size,
-          /*   ldc          */ 2 * N,
-          /*   brg          */ use_brgemm,
-          /*   block_size_K */ block_size_K);
-    }
-
-    if (is_brgemm_used) {
-      at::native::cpublas::brgemm_release();
-    }
-  });
-
-  // stage 1.5: intermediate_cache1 = silu(intermediate_cache0)
-  at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
-    for (int64_t m = begin; m < end; ++m) {
-      silu_and_mul_stub(
-          ic1 + m * N,
-          ic0 + m * 2 * N,
-          ic0 + m * 2 * N + N,
-          N);
-    }
-  });
-
-  // stage 2: intermediate_cache2 = intermediate_cache1 @ w2
-  //   w2 : [E, K, N] as [E, OC, IC]
-  const int64_t OC = K;  // rename K as OC
-  const int64_t IC = N;  // rename N as IC
-  const int64_t MB2 = MB;
-  const int64_t NB2 = div_up(OC, BLOCK_N);
-  scale_size_N = div_up(K, block_size_N);
-  scale_size_K = div_up(N, block_size_K);
-  const int64_t stride_e2 = OC * IC;
-  const int64_t stride_oc = IC;
-
-  // parallel on [MB2, NB2]
-  at::parallel_for(0, MB2 * NB2, 0, [&](int64_t begin, int64_t end) {
-    int tid = at::get_thread_num();
-    alignas(64) scalar_t C[BLOCK_M * BLOCK_K];
-
-    bool is_brgemm_used = false;
-
-    for (int64_t i = begin; i < end; ++i) {
-      int64_t mb = i / NB2;
-      int64_t nb = i % NB2;
-
-      int64_t m_size = offsets[mb + 1] - offsets[mb];
-      int64_t n_size = std::min(OC - nb * BLOCK_N, BLOCK_N);
-
-      const bool use_brgemm = can_use_brgemm<at::Float8_e4m3fn>(m_size);
-      is_brgemm_used = is_brgemm_used || use_brgemm;
-
-      // A ptr from ic1 of [M * topk, N] in sorted order
-      // so as to avoid copy A to tmp buffer again
-      const scalar_t* __restrict__ A = ic1 + offsets[mb] * N;
-      const int32_t* A_ids = sorted_ids + mb * BLOCK_M;
-
-      // B shape [IC, n_size] in vnni format
-      int32_t expert_id = expert_ids[mb];
-      const at::Float8_e4m3fn* __restrict__ B = packed_w2 + expert_id * stride_e2 + nb * BLOCK_N * stride_oc;
-      const float* __restrict__ Bs = w2s + expert_id * scale_size_N * scale_size_K + (nb / blocks_n_per_group) * scale_size_K;
-
-      tinygemm_kernel<scalar_t>(
-          /*   A            */ A,
-          /*   B            */ B,
-          /*   C            */ C,
-          /*   Btmp         */ B_tmp + tid * BLOCK_N * std::max(K, N),
-          /*   Ctmp         */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N,
-          /*   scale        */ Bs,
-          /*   M            */ m_size,
-          /*   N            */ n_size,
-          /*   K            */ IC,
-          /*   lda          */ IC,
-          /*   ldb          */ n_size,
-          /*   ldc          */ BLOCK_N,
-          /*   brg          */ use_brgemm,
-          /*   block_size_K */ block_size_K);
-
-      // 2.b copy from C to ic2 in original order
-      //   and also mul topk_weights in float32
-      for (int64_t m = 0; m < m_size; ++m) {
-        int32_t index = A_ids[m];
-        float weight = topk_weights[index];
-        copy_mul_stub(ic2 + index * K + nb * BLOCK_N, C + m * BLOCK_N, weight, n_size);
-      }
-    }
-
-    if (is_brgemm_used) {
-      at::native::cpublas::brgemm_release();
-    }
-  });
-
-  // stage 3: out = intermediate_cache2.sum(dim=1)
-  //   from [M, topk, K] to [M, K]
-  at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
-    for (int64_t m = begin; m < end; ++m) {
-      sum_stub(output + m * K, ic2 + m * topk * K, topk, K);
-    }
-  });
-}
-
-#define INSTANTIATE_MOE_FP8_TEMPLATE(TYPE)             \
-  template void fused_experts_fp8_kernel_impl<TYPE>(   \
-      TYPE* __restrict__ output,                       \
-      TYPE* __restrict__ ic0,                          \
-      TYPE* __restrict__ ic1,                          \
-      TYPE* __restrict__ ic2,                          \
-      TYPE* __restrict__ A_tmp,                        \
-      TYPE* __restrict__ B_tmp,                        \
-      float* __restrict__ C_tmp,                       \
-      const TYPE* __restrict__ input,                  \
-      const at::Float8_e4m3fn* __restrict__ packed_w1, \
-      const at::Float8_e4m3fn* __restrict__ packed_w2, \
-      const float* __restrict__ w1s,                   \
-      const float* __restrict__ w2s,                   \
-      int64_t block_size_N,                            \
-      int64_t block_size_K,                            \
-      const float* __restrict__ topk_weights,          \
-      const int32_t* __restrict__ sorted_ids,          \
-      const int32_t* __restrict__ expert_ids,          \
-      const int32_t* __restrict__ offsets,             \
-      int64_t M,                                       \
-      int64_t N,                                       \
-      int64_t K,                                       \
-      int64_t E,                                       \
-      int64_t topk,                                    \
-      int64_t num_tokens_post_pad)
-
-INSTANTIATE_MOE_FP8_TEMPLATE(at::BFloat16);
-INSTANTIATE_MOE_FP8_TEMPLATE(at::Half);
-
-template <typename scalar_t>
-void shared_expert_fp8_kernel_impl(
-    scalar_t* __restrict__ output,
-    scalar_t* __restrict__ ic0,
-    scalar_t* __restrict__ ic1,
-    scalar_t* __restrict__ B_tmp,
-    float* __restrict__ C_tmp,
-    const scalar_t* __restrict__ input,
-    const at::Float8_e4m3fn* __restrict__ packed_w1,
-    const at::Float8_e4m3fn* __restrict__ packed_w2,
-    const float* __restrict__ w1s,
-    const float* __restrict__ w2s,
-    int64_t block_size_N,
-    int64_t block_size_K,
-    const scalar_t* __restrict__ fused_experts_out,
-    float routed_scaling_factor,
-    int64_t M,
-    int64_t N,
-    int64_t K) {
-
-  constexpr int64_t BLOCK_M = block_size_m();
-  constexpr int64_t BLOCK_N = block_size_n();
-
-  // stage 1: intermediate_cache0 = hidden_states @ w1
-  const int64_t MB = div_up(M, BLOCK_M);
-  const int64_t NB = div_up(2 * N, BLOCK_N);
-  int64_t scale_size_K = div_up(K, block_size_K);
-  int64_t blocks_n_per_group = block_size_N / BLOCK_N;
-
-  const bool use_brgemm = can_use_brgemm<at::Float8_e4m3fn>(M);
-
-  at::parallel_for(0, MB * NB, 0, [&](int64_t begin, int64_t end) {
-    int tid = at::get_thread_num();
-
-    for (int64_t i = begin; i < end; ++i) {
-      int64_t mb = i / NB;
-      int64_t nb = i % NB;
-      int64_t m_size = std::min(M - mb * BLOCK_M, BLOCK_M);
-      int64_t n_size = std::min(2 * N - nb * BLOCK_N, BLOCK_N);
-
-      tinygemm_kernel<scalar_t>(
-          /*   A            */ input + mb * BLOCK_M * K,
-          /*   B            */ packed_w1 + nb * BLOCK_N * K,
-          /*   C            */ ic0 + mb * BLOCK_M * 2 * N + nb * BLOCK_N,
-          /*   Btmp         */ B_tmp + tid * BLOCK_N * std::max(K, N),
-          /*   Ctmp         */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N,
-          /*   scale        */ w1s + (nb / blocks_n_per_group) * scale_size_K,
-          /*   M            */ m_size,
-          /*   N            */ n_size,
-          /*   K            */ K,
-          /*   lda          */ K,
-          /*   ldb          */ n_size,
-          /*   ldc          */ 2 * N,
-          /*   brg          */ use_brgemm,
-          /*   block_size_K */ block_size_K);
-    }
-
-    if (use_brgemm) {
-      at::native::cpublas::brgemm_release();
-    }
-  });
-
-  // stage 1.5: intermediate_cache1 = silu(intermediate_cache0)
-  at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
-    for (int64_t m = begin; m < end; ++m) {
-      silu_and_mul_stub(
-          ic1 + m * N,
-          ic0 + m * 2 * N,
-          ic0 + m * 2 * N + N,
-          N);
-    }
-  });
-
-  // stage 2: intermediate_cache2 = intermediate_cache1 @ w2
-  //   w2 : [K, N] as [OC, IC]
-  const int64_t OC = K;  // rename K as OC
-  const int64_t IC = N;  // rename N as IC
-  const int64_t MB2 = MB;
-  const int64_t NB2 = div_up(K, BLOCK_N);
-  scale_size_K = div_up(N, block_size_K);
-
-  // parallel on [MB2, NB2]
-  at::parallel_for(0, MB2 * NB2, 0, [&](int64_t begin, int64_t end) {
-    int tid = at::get_thread_num();
-    alignas(64) scalar_t C[BLOCK_M * BLOCK_K];
-
-    for (int64_t i = begin; i < end; ++i) {
-      int64_t mb = i / NB2;
-      int64_t nb = i % NB2;
-      int64_t m_size = std::min(M - mb * BLOCK_M, BLOCK_M);
-      int64_t n_size = std::min(OC - nb * BLOCK_N, BLOCK_N);
-
-      // 2.a gemm: C = A @ B
-      tinygemm_kernel<scalar_t>(
-          /*   A            */ ic1 + mb * BLOCK_M * N,
-          /*   B            */ packed_w2 + nb * BLOCK_N * N,
-          /*   C            */ C,
-          /*   Btmp         */ B_tmp + tid * BLOCK_N * std::max(K, N),
-          /*   Ctmp         */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N,
-          /*   scale        */ w2s + (nb / blocks_n_per_group) * scale_size_K,
-          /*   M            */ m_size,
-          /*   N            */ n_size,
-          /*   K            */ IC,
-          /*   lda          */ IC,
-          /*   ldb          */ n_size,
-          /*   ldc          */ BLOCK_N,
-          /*   brg          */ use_brgemm,
-          /*   block_size_K */ block_size_K);
-
-      // 2.b copy from C to output and add fused_experts_out
-      scalar_t* __restrict__ out = output + mb * BLOCK_M * K + nb * BLOCK_N;
-      const scalar_t* __restrict__ fused_out = fused_experts_out + mb * BLOCK_M * K + nb * BLOCK_N;
-      for (int64_t m = 0; m < m_size; ++m) {
-        add_mul_stub(out + m * K, C + m * BLOCK_N, fused_out + m * K, routed_scaling_factor, n_size);
-      }
-    }
-  });
-
-  if (use_brgemm) {
-    at::native::cpublas::brgemm_release();
-  }
-}
-
-#define INSTANTIATE_SHARED_EXPERT_FP8_TEMPLATE(TYPE)   \
-  template void shared_expert_fp8_kernel_impl<TYPE>(   \
-      TYPE* __restrict__ output,                       \
-      TYPE* __restrict__ ic0,                          \
-      TYPE* __restrict__ ic1,                          \
-      TYPE* __restrict__ B_tmp,                        \
-      float* __restrict__ C_tmp,                       \
-      const TYPE* __restrict__ input,                  \
-      const at::Float8_e4m3fn* __restrict__ packed_w1, \
-      const at::Float8_e4m3fn* __restrict__ packed_w2, \
-      const float* __restrict__ w1s,                   \
-      const float* __restrict__ w2s,                   \
-      int64_t block_size_N,                            \
-      int64_t block_size_K,                            \
-      const TYPE* __restrict__ fused_experts_out,      \
-      float routed_scaling_factor,                     \
-      int64_t M,                                       \
-      int64_t N,                                       \
-      int64_t K)
-
-INSTANTIATE_SHARED_EXPERT_FP8_TEMPLATE(at::BFloat16);
-INSTANTIATE_SHARED_EXPERT_FP8_TEMPLATE(at::Half);

csrc/cpu/sgl-kernels/moe_int8.cpp

@@ -1,769 +0,0 @@
-// Adapted from
-// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu
-
-#include "common.h"
-#include "vec.h"
-#include "gemm.h"
-
-// clang-format off
-
-namespace {
-
-template <typename scalar_t>
-inline void copy_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, int64_t size) {
-  using Vec = at::vec::Vectorized<scalar_t>;
-  // no remainder
-  #pragma GCC unroll 4
-  for (int64_t d = 0; d < size; d += Vec::size()) {
-    Vec data = Vec::loadu(input + d);
-    data.store(out + d);
-  }
-}
-
-template <>
-inline void copy_stub<uint8_t>(uint8_t* __restrict__ out, const uint8_t* __restrict__ input, int64_t size) {
-  // size might be 64x + 32
-  std::memcpy(out, input, size * sizeof(uint8_t));
-}
-
-template <typename scalar_t>
-inline void copy_mul_stub(scalar_t* __restrict__ out, const float* __restrict__ input, float weight, int64_t size) {
-  using bVec = at::vec::Vectorized<scalar_t>;
-  using fVec = at::vec::Vectorized<float>;
-  constexpr int kVecSize = bVec::size();
-  const fVec weight_vec = fVec(weight);
-  int64_t d;
-  #pragma GCC unroll 4
-  for (d = 0; d <= size - kVecSize; d += kVecSize) {
-    fVec data0 = fVec::loadu(input + d) * weight_vec;
-    fVec data1 = fVec::loadu(input + d + fVec::size()) * weight_vec;
-    bVec out_vec = convert_from_float_ext<scalar_t>(data0, data1);
-    out_vec.store(out + d);
-  }
-  for (; d < size; ++d) {
-    out[d] = static_cast<scalar_t>(input[d] * weight);
-  }
-}
-
-// acc from [topk, K] to [K]
-template <typename scalar_t>
-inline void sum_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, int64_t topk, int64_t K) {
-  using bVec = at::vec::Vectorized<scalar_t>;
-  using fVec = at::vec::Vectorized<float>;
-  constexpr int kVecSize = bVec::size();
-  if (topk == 1) {
-    // do copy for topk = 1
-    copy_stub(out, input, K);
-  } else {
-    // do sum for topk != 1
-    int64_t d;
-    #pragma GCC unroll 4
-    for (d = 0; d <= K - kVecSize; d += kVecSize) {
-      fVec sum_fvec0 = fVec(0.f);
-      fVec sum_fvec1 = fVec(0.f);
-      for (int t = 0; t < topk; ++t) {
-        bVec x_bvec = bVec::loadu(input + t * K + d);
-        fVec x_fvec0, x_fvec1;
-        std::tie(x_fvec0, x_fvec1) = at::vec::convert_to_float(x_bvec);
-
-        sum_fvec0 += x_fvec0;
-        sum_fvec1 += x_fvec1;
-      }
-      bVec out_bvec = convert_from_float_ext<scalar_t>(sum_fvec0, sum_fvec1);
-      out_bvec.store(out + d);
-    }
-    for (; d < K; ++d) {
-      float sum_val = 0.f;
-      for (int t = 0; t < topk; ++t) {
-        sum_val += static_cast<float>(input[t * K + d]);
-      }
-      out[d] = static_cast<scalar_t>(sum_val);
-    }
-  }
-}
-
-// out = input + input2 * scale
-template <typename scalar_t>
-inline void add_mul_stub(scalar_t* __restrict__ out, const float* __restrict__ input,
-    const scalar_t* __restrict__ input2, float scale, int64_t size) {
-
-  using bVec = at::vec::Vectorized<scalar_t>;
-  using fVec = at::vec::Vectorized<float>;
-  constexpr int kVecSize = bVec::size();
-  const fVec s_vec = fVec(scale);
-  int64_t d;
-  #pragma GCC unroll 4
-  for (d = 0; d <= size - kVecSize; d += kVecSize) {
-    fVec x0 = fVec::loadu(input + d);
-    fVec x1 = fVec::loadu(input + d + fVec::size());
-
-    bVec y_bvec = bVec::loadu(input2 + d);
-    fVec y0, y1;
-    std::tie(y0, y1) = at::vec::convert_to_float(y_bvec);
-
-    x0 = x0 + y0 * s_vec;
-    x1 = x1 + y1 * s_vec;
-    bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1);
-    out_vec.store(out + d);
-  }
-  for (; d < size; ++d) {
-    out[d] = static_cast<scalar_t>(input[d] + float(input2[d]) * scale);
-  }
-}
-
-/// gemm for w13
-template <typename scalar_t, int BLOCK_M, int BLOCK_N>
-struct tinygemm_kernel_vnni {
-  static inline void apply(
-      const uint8_t* __restrict__ A, const int8_t* __restrict__ B0, const int8_t* __restrict__ B1, scalar_t* __restrict__ C,
-      const float* __restrict__ As, const float* __restrict__ Bs0, const float* __restrict__ Bs1,
-      const int32_t* __restrict__ Bcomp0, const int32_t* __restrict__ Bcomp1,
-      int64_t K, int64_t lda, int64_t ldb, int64_t ldc) {
-    TORCH_CHECK(false, "tinygemm_kernel_nn: scalar path not implemented!");
-  }
-};
-
-#if defined(CPU_CAPABILITY_AVX512)
-template <int BLOCK_M, int BLOCK_N>
-struct tinygemm_kernel_vnni<at::BFloat16, BLOCK_M, BLOCK_N> {
-  static inline void apply(
-      const uint8_t* __restrict__ A, const int8_t* __restrict__ B0, const int8_t* __restrict__ B1, at::BFloat16* __restrict__ C,
-      const float* __restrict__ As, const float* __restrict__ Bs0, const float* __restrict__ Bs1,
-      const int32_t* __restrict__ Bcomp0, const int32_t* __restrict__ Bcomp1,
-      int64_t K, int64_t lda, int64_t ldb, int64_t ldc) {
-
-    constexpr int ROWS = BLOCK_M;
-    constexpr int COLS = BLOCK_N / 16;
-    static_assert(COLS % 2 == 0);
-
-    __m512i va;
-    __m512i vb0[COLS];
-    __m512i vb1[COLS];
-    __m512i vc0[ROWS * COLS];
-    __m512i vc1[ROWS * COLS];
-    __m512i vcomp0[COLS];
-    __m512i vcomp1[COLS];
-    __m512  was;
-    __m512  vbs0[COLS];
-    __m512  vbs1[COLS];
-
-    auto loadc = [&](auto i) {
-      vc0[i] = _mm512_set1_epi32(0);
-      vc1[i] = _mm512_set1_epi32(0);
-    };
-    Unroll<ROWS * COLS>{}(loadc);
-
-    const int64_t K4 = K >> 2;
-    const int64_t lda4 = lda >> 2;
-    const int64_t ldb4 = ldb; // ldb * 4 >> 2;
-    const int32_t* a_ptr = reinterpret_cast<const int32_t*>(A);
-    const int32_t* b0_ptr = reinterpret_cast<const int32_t*>(B0);
-    const int32_t* b1_ptr = reinterpret_cast<const int32_t*>(B1);
-
-    auto compute = [&](auto i, int64_t k) {
-      constexpr int row = i / COLS;
-      constexpr int col = i % COLS;
-
-      if constexpr (col == 0) {
-        va = _mm512_set1_epi32(a_ptr[row * lda4 + k]);
-      }
-      if constexpr (row == 0) {
-        vb0[col] = _mm512_loadu_si512(b0_ptr + k * ldb4 + col * 16);
-        vb1[col] = _mm512_loadu_si512(b1_ptr + k * ldb4 + col * 16);
-      }
-      vc0[i] = _mm512_dpbusd_epi32(vc0[i], va, vb0[col]);
-      vc1[i] = _mm512_dpbusd_epi32(vc1[i], va, vb1[col]);
-    };
-    for (int64_t k = 0; k < K4; ++k) {
-      Unroll<ROWS * COLS>{}(compute, k);
-    }
-
-    auto scalec = [&](auto i) {
-      constexpr int row = i / COLS;
-      constexpr int col = i % COLS;
-
-      // load a scale
-      if constexpr(col == 0) {
-        was = _mm512_set1_ps(As[row]);
-      }
-      // load b scale and vcomp
-      if constexpr (row == 0) {
-        vbs0[col] = _mm512_loadu_ps(Bs0 + col * 16);
-        vbs1[col] = _mm512_loadu_ps(Bs1 + col * 16);
-        vcomp0[col] = _mm512_loadu_si512(Bcomp0 + col * 16);
-        vcomp1[col] = _mm512_loadu_si512(Bcomp1 + col * 16);
-      }
-      __m512 c0 = _mm512_cvtepi32_ps(_mm512_sub_epi32(vc0[i], vcomp0[col]));
-      __m512 c1 = _mm512_cvtepi32_ps(_mm512_sub_epi32(vc1[i], vcomp1[col]));
-      vc0[i] = _mm512_castps_si512(_mm512_mul_ps(_mm512_mul_ps(c0, was), vbs0[col]));
-      vc1[i] = _mm512_castps_si512(_mm512_mul_ps(_mm512_mul_ps(c1, was), vbs1[col]));
-    };
-    Unroll<ROWS * COLS>{}(scalec);
-
-    using Vec = at::vec::Vectorized<float>;
-    const Vec one = Vec(1.f);
-    auto storec = [&](auto i) {
-      constexpr int row = i / COLS;
-      constexpr int col = i % COLS;
-      // for COLS = 2, 4 use 512bit store
-      if constexpr (col % 2 == 0) {
-        Vec x0 = _mm512_castsi512_ps(vc0[row * COLS + col + 0]);
-        Vec x1 = _mm512_castsi512_ps(vc0[row * COLS + col + 1]);
-        Vec y0 = _mm512_castsi512_ps(vc1[row * COLS + col + 0]);
-        Vec y1 = _mm512_castsi512_ps(vc1[row * COLS + col + 1]);
-        // silu
-        x0 = x0 / (one + x0.neg().exp_u20());
-        x1 = x1 / (one + x1.neg().exp_u20());
-        // mul
-        x0 = x0 * y0;
-        x1 = x1 * y1;
-
-        _mm512_storeu_si512(
-            reinterpret_cast<__m512i*>((C + row * ldc + col * 16)),
-            (__m512i)(_mm512_cvtne2ps_pbh(__m512(x1), __m512(x0))));
-        }
-    };
-    Unroll<ROWS * COLS>{}(storec);
-  }
-};
-#endif
-
-#define LAUNCH_TINYGEMM_KERNEL_VNNI(MB_SIZE, NB_SIZE)                        \
-    tinygemm_kernel_vnni<scalar_t, MB_SIZE, NB_SIZE>::apply(                 \
-        A + mb_start * lda, B0 + nb_start * 4, B1 + nb_start * 4,            \
-        C + mb_start * ldc + nb_start, As + mb_start,                        \
-        Bs0 + nb_start, Bs1 + nb_start, Bcomp0 + nb_start, Bcomp1 + nb_start,\
-        K, lda, ldb, ldc);
-
-template <typename scalar_t>
-void tinygemm_kernel(
-    const uint8_t* __restrict__ A,
-    const int8_t* __restrict__ B0,
-    const int8_t* __restrict__ B1,
-    scalar_t* __restrict__ C,
-    const float* __restrict__ As,
-    const float* __restrict__ Bs0,
-    const float* __restrict__ Bs1,
-    int64_t M,
-    int64_t N,
-    int64_t K,
-    int64_t lda,
-    int64_t ldb,
-    int64_t ldc) {
-
-  const int32_t* Bcomp0 = reinterpret_cast<const int32_t*>(B0 + block_size_n() * K);
-  const int32_t* Bcomp1 = reinterpret_cast<const int32_t*>(B1 + block_size_n() * K);
-
-  // pattern: 1-(2+2)-(8+8)
-  constexpr int64_t BLOCK_M = 4;
-  constexpr int64_t BLOCK_N = 32;
-  const int64_t MB = div_up(M, BLOCK_M);
-  const int64_t NB = div_up(N, BLOCK_N);
-  for (int mb = 0; mb < MB; ++mb) {
-    int64_t mb_start = mb * BLOCK_M;
-    int64_t mb_size = std::min(BLOCK_M, M - mb_start);
-    for (int64_t nb = 0; nb < NB; ++nb) {
-      int64_t nb_start = nb * BLOCK_N;
-      int64_t nb_size = std::min(BLOCK_N, N - nb_start);
-
-      switch(mb_size << 4 | nb_size >> 4) {
-        case 0x12: LAUNCH_TINYGEMM_KERNEL_VNNI(1, 32); break;
-        case 0x22: LAUNCH_TINYGEMM_KERNEL_VNNI(2, 32); break;
-        case 0x32: LAUNCH_TINYGEMM_KERNEL_VNNI(3, 32); break;
-        case 0x42: LAUNCH_TINYGEMM_KERNEL_VNNI(4, 32); break;
-        default: TORCH_CHECK(false, "Unexpected block size, ", mb_size, "x", "nb_size");
-      }
-    }
-  }
-}
-
-/// gemm for w2
-template <typename scalar_t, int BLOCK_M, int BLOCK_N>
-struct tinygemm_kernel_vnni2 {
-  static inline void apply(
-      const uint8_t* __restrict__ A, const int8_t* __restrict__ B, float* __restrict__ C,
-      const float* __restrict__ As, const float* __restrict__ Bs, const int32_t* __restrict__ Bcomp,
-      int64_t K, int64_t lda, int64_t ldb, int64_t ldc) {
-    TORCH_CHECK(false, "tinygemm_kernel_nn: scalar path not implemented!");
-  }
-};
-
-#if defined(CPU_CAPABILITY_AVX512)
-template <int BLOCK_M, int BLOCK_N>
-struct tinygemm_kernel_vnni2<at::BFloat16, BLOCK_M, BLOCK_N> {
-  static inline void apply(
-      const uint8_t* __restrict__ A, const int8_t* __restrict__ B, float* __restrict__ C,
-      const float* __restrict__ As, const float* __restrict__ Bs, const int32_t* __restrict__ Bcomp,
-      int64_t K, int64_t lda, int64_t ldb, int64_t ldc) {
-
-    constexpr int ROWS = BLOCK_M;
-    constexpr int COLS = BLOCK_N / 16;
-    static_assert(COLS % 2 == 0);
-
-    __m512i va;
-    __m512i vb[COLS];
-    __m512i vc[ROWS * COLS];
-    __m512i vcomp[COLS];
-    __m512  was;
-    __m512  vbs[COLS];
-
-    auto loadc = [&](auto i) {
-      vc[i] = _mm512_set1_epi32(0);
-    };
-    Unroll<ROWS * COLS>{}(loadc);
-
-    const int64_t K4 = K >> 2;
-    const int64_t lda4 = lda >> 2;
-    const int64_t ldb4 = ldb; // ldb * 4 >> 2;
-    const int32_t* a_ptr = reinterpret_cast<const int32_t*>(A);
-    const int32_t* b_ptr = reinterpret_cast<const int32_t*>(B);
-
-    auto compute = [&](auto i, int64_t k) {
-      constexpr int row = i / COLS;
-      constexpr int col = i % COLS;
-
-      if constexpr (col == 0) {
-        va = _mm512_set1_epi32(a_ptr[row * lda4 + k]);
-      }
-      if constexpr (row == 0) {
-        vb[col] = _mm512_loadu_si512(b_ptr + k * ldb4 + col * 16);
-      }
-      vc[i] = _mm512_dpbusd_epi32(vc[i], va, vb[col]);
-    };
-    for (int64_t k = 0; k < K4; ++k) {
-      Unroll<ROWS * COLS>{}(compute, k);
-    }
-
-    auto storec = [&](auto i) {
-      constexpr int row = i / COLS;
-      constexpr int col = i % COLS;
-
-      // load a scale
-      if constexpr(col == 0) {
-        was = _mm512_set1_ps(As[row]);
-      }
-      // load b scale and vcomp per 2 vectors
-      // also load bias if any
-      if constexpr (row == 0) {
-        if constexpr (col % 2 == 0) {
-          vbs[col + 0] = _mm512_loadu_ps(Bs + col * 16);
-          vbs[col + 1] = _mm512_loadu_ps(Bs + col * 16 + 16);
-          vcomp[col + 0] = _mm512_loadu_si512(Bcomp + col * 16);
-          vcomp[col + 1] = _mm512_loadu_si512(Bcomp + col * 16 + 16);
-        }
-      }
-      __m512 x = _mm512_cvtepi32_ps(_mm512_sub_epi32(vc[i], vcomp[col]));
-      x = _mm512_mul_ps(_mm512_mul_ps(x, was), vbs[col]);
-      _mm512_storeu_ps(reinterpret_cast<__m512*>(C + row * ldc + col * 16), x);
-    };
-    Unroll<ROWS * COLS>{}(storec);
-  }
-};
-#endif
-
-#define LAUNCH_TINYGEMM_KERNEL_VNNI2(MB_SIZE, NB_SIZE)                       \
-    tinygemm_kernel_vnni2<scalar_t, MB_SIZE, NB_SIZE>::apply(                \
-        A + mb_start * lda, B + nb_start * 4, C + mb_start * ldc + nb_start, \
-        As + mb_start, Bs + nb_start, Bcomp + nb_start,                      \
-        K, lda, ldb, ldc);
-
-template <typename scalar_t>
-void tinygemm_kernel(
-    const uint8_t* __restrict__ A,
-    const int8_t* __restrict__ B,
-    float* __restrict__ C,
-    const float* __restrict__ As,
-    const float* __restrict__ Bs,
-    int64_t M,
-    int64_t N,
-    int64_t K,
-    int64_t lda,
-    int64_t ldb,
-    int64_t ldc) {
-
-  // B compensation
-  const int32_t* Bcomp = reinterpret_cast<const int32_t*>(B + block_size_n() * K);
-
-  // pattern: 1-4-16
-  constexpr int64_t BLOCK_M = 4;
-  constexpr int64_t BLOCK_N = 64;
-  const int64_t MB = div_up(M, BLOCK_M);
-  const int64_t NB = div_up(N, BLOCK_N);
-  for (int64_t mb = 0; mb < MB; ++mb) {
-    int64_t mb_start = mb * BLOCK_M;
-    int64_t mb_size = std::min(BLOCK_M, M - mb_start);
-    for (int64_t nb = 0; nb < NB; ++nb) {
-      int64_t nb_start = nb * BLOCK_N;
-      int64_t nb_size = std::min(BLOCK_N, N - nb_start);
-
-      switch(mb_size << 4 | nb_size >> 4) {
-        case 0x12: LAUNCH_TINYGEMM_KERNEL_VNNI2(1, 32); break;
-        case 0x22: LAUNCH_TINYGEMM_KERNEL_VNNI2(2, 32); break;
-        case 0x32: LAUNCH_TINYGEMM_KERNEL_VNNI2(3, 32); break;
-        case 0x42: LAUNCH_TINYGEMM_KERNEL_VNNI2(4, 32); break;
-        default: TORCH_CHECK(false, "Unexpected block size, ", mb_size, "x", "nb_size");
-      }
-    }
-  }
-}
-
-} // anonymous namespace
-
-template <typename scalar_t>
-void fused_experts_int8_kernel_impl(
-    scalar_t* __restrict__ output,
-    scalar_t* __restrict__ ic1,
-    scalar_t* __restrict__ ic2,
-    uint8_t* __restrict__ A_tmp,
-    float* __restrict__ C_tmp,
-    uint8_t* __restrict__ Aq_tmp,
-    float* __restrict__ As_tmp,
-    const scalar_t* __restrict__ input,
-    const int8_t* __restrict__ packed_w1,
-    const int8_t* __restrict__ packed_w2,
-    const float* __restrict__ w1s,
-    const float* __restrict__ w2s,
-    const float* __restrict__ topk_weights,
-    const int32_t* __restrict__ sorted_ids,
-    const int32_t* __restrict__ expert_ids,
-    const int32_t* __restrict__ offsets,
-    int64_t M,
-    int64_t N,
-    int64_t K,
-    int64_t E,
-    int64_t topk,
-    int64_t num_tokens_post_pad) {
-
-  // handle 2 tiles per block
-  constexpr int64_t BLOCK_M = block_size_m();
-  constexpr int64_t BLOCK_N = block_size_n();
-
-  // stage 0: quantize input to uint8, [M, K]
-  at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
-    for (int64_t m = begin; m < end; ++m) {
-      quantize_row_int8<scalar_t>(
-          Aq_tmp + m * K,
-          As_tmp[m],
-          input + m * K,
-          K);
-    }
-  });
-
-  // stage 1: intermediate_cache1 = silu(hidden_states @ w1)
-  const int64_t MB = div_up(num_tokens_post_pad, BLOCK_M);
-  const int64_t NB = div_up(N, BLOCK_N);
-
-  // strides for w1: [E, 2N, K]
-  TORCH_CHECK(N % BLOCK_N == 0, "Fixme when N is not multiples of ", BLOCK_N);
-
-  // K and N are packed for int8
-  const int64_t packed_K = get_row_size<int8_t>(K);
-  const int64_t packed_N = get_row_size<int8_t>(N);
-
-  const int64_t stride_e = 2 * N * packed_K;
-  const int64_t stride_n = packed_K;
-  // here we only parallel on half of 2N to fuse silu_and_mul with gemm
-  at::parallel_for(0, MB * NB, 0, [&](int64_t begin, int64_t end) {
-    // get local pointers
-    int tid = at::get_thread_num();
-    uint8_t* __restrict__ A = A_tmp + tid * BLOCK_M * K;
-
-    alignas(64) float As[BLOCK_M];
-
-    for (int64_t i = begin; i < end; ++i) {
-      int64_t mb = i / NB;
-      int64_t nb = i % NB;
-
-      // nb0 from top half and nb1 from bottom half
-      int64_t nb0 = nb, nb1 = nb + NB;
-      int64_t n_size = std::min(N - nb0 * BLOCK_N, BLOCK_N);
-
-      // B shape [K, n_size] in vnni format
-      int32_t expert_id = expert_ids[mb];
-      const int8_t* __restrict__ B0 = packed_w1 + expert_id * stride_e + nb0 * BLOCK_N * stride_n;
-      const int8_t* __restrict__ B1 = packed_w1 + expert_id * stride_e + nb1 * BLOCK_N * stride_n;
-      const float* __restrict__ Bs0 = w1s + expert_id * 2 * N + nb0 * BLOCK_N;
-      const float* __restrict__ Bs1 = w1s + expert_id * 2 * N + nb1 * BLOCK_N;
-
-      // 1.a load A
-      const int32_t* A_ids = sorted_ids + mb * BLOCK_M;
-      int64_t m_size = offsets[mb + 1] - offsets[mb];
-
-      for (int64_t m = 0; m < m_size; ++m) {
-        int32_t index = A_ids[m] / topk;
-        copy_stub(A + m * K, Aq_tmp + index * K, K);
-        As[m] = As_tmp[index];
-      }
-
-      // fused 1.b: silu_and_mul(A @ B0, A @ B1)
-      const int64_t offset = offsets[mb];
-      tinygemm_kernel(
-          /* A     */ A,
-          /* B0    */ B0,
-          /* B1    */ B1,
-          /* C     */ ic1 + offset * N + nb * BLOCK_N,
-          /* As    */ As,
-          /* Bs0   */ Bs0,
-          /* Bs1   */ Bs1,
-          /* M     */ m_size,
-          /* N     */ n_size,
-          /* K     */ K,
-          /* lda   */ K,
-          /* ldb   */ n_size,
-          /* ldc   */ N);
-    }
-  });
-
-  // stage 1.5: quantize ic1 to uint8, [M * topk, N]
-  at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
-    for (int64_t m = begin; m < end; ++m) {
-      quantize_row_int8<scalar_t>(
-          Aq_tmp + m * N,
-          As_tmp[m],
-          ic1 + m * N,
-          N);
-    }
-  });
-
-  // stage 2: intermediate_cache2 = intermediate_cache1 @ w2
-  //   w2 : [E, K, N] as [E, OC, IC]
-  const int64_t OC = K;  // rename K as OC
-  const int64_t IC = N;  // rename N as IC
-  const int64_t MB2 = MB;
-  const int64_t NB2 = div_up(OC, BLOCK_N);
-  const int64_t stride_e2 = OC * packed_N;
-  const int64_t stride_oc = packed_N;
-
-  // parallel on [MB2, NB2]
-  at::parallel_for(0, MB2 * NB2, 0, [&](int64_t begin, int64_t end) {
-    // get local pointers
-    int tid = at::get_thread_num();
-    // we won't be using C1 for gemm2
-    float* __restrict__ C = C_tmp + tid * 2 * BLOCK_M * BLOCK_N;
-
-    for (int64_t i = begin; i < end; ++i) {
-      int64_t mb = i / NB2;
-      int64_t nb = i % NB2;
-
-      int64_t m_size = offsets[mb + 1] - offsets[mb];
-      int64_t n_size = std::min(OC - nb * BLOCK_N, BLOCK_N);
-
-      // A ptr from ic1 of [M * topk, N] in sorted order
-      // so as to avoid copy A to tmp buffer again
-      const uint8_t* __restrict__ A = Aq_tmp + offsets[mb] * N;
-      const float* __restrict__ As = As_tmp + offsets[mb];
-      const int32_t* A_ids = sorted_ids + mb * BLOCK_M;
-
-      // B shape [IC, n_size] in vnni format
-      int32_t expert_id = expert_ids[mb];
-      const int8_t* __restrict__ B = packed_w2 + expert_id * stride_e2 + nb * BLOCK_N * stride_oc;
-      const float* __restrict__ Bs = w2s + expert_id * K + nb * BLOCK_N;
-
-      // 2.a gemm: C = A @ B
-      tinygemm_kernel<scalar_t>(
-          /* A     */ A,
-          /* B     */ B,
-          /* C     */ C,
-          /* As    */ As,
-          /* Bs    */ Bs,
-          /* M     */ m_size,
-          /* N     */ n_size,
-          /* K     */ IC,
-          /* lda   */ IC,
-          /* ldb   */ n_size,
-          /* ldc   */ BLOCK_N);
-
-      // 2.b copy from C to ic2 in original order
-      //   and also mul topk_weights in float32
-      for (int64_t m = 0; m < m_size; ++m) {
-        int32_t index = A_ids[m];
-        float weight = topk_weights[index];
-        copy_mul_stub(ic2 + index * K + nb * BLOCK_N, C + m * BLOCK_N, weight, n_size);
-      }
-    }
-  });
-
-  // stage 3: out = intermediate_cache2.sum(dim=1)
-  //   from [M, topk, K] to [M, K]
-  at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
-    for (int64_t m = begin; m < end; ++m) {
-      sum_stub(output + m * K, ic2 + m * topk * K, topk, K);
-    }
-  });
-}
-
-#define INSTANTIATE_MOE_INT8_TEMPLATE(TYPE)                                                  \
-  template void fused_experts_int8_kernel_impl<TYPE> (                                       \
-      TYPE* __restrict__ output, TYPE* __restrict__ ic1,                                     \
-      TYPE* __restrict__ ic2, uint8_t* __restrict__ A_tmp,                                   \
-      float* __restrict__ C_tmp, uint8_t* __restrict__ Aq_tmp,                               \
-      float* __restrict__ As_tmp, const TYPE* __restrict__ input,                            \
-      const int8_t* __restrict__ packed_w1, const int8_t* __restrict__ packed_w2,            \
-      const float* __restrict__ w1s, const float* __restrict__ w2s,                          \
-      const float* __restrict__ topk_weights, const int32_t* __restrict__ sorted_ids,        \
-      const int32_t* __restrict__ expert_ids, const int32_t* __restrict__ offsets,           \
-      int64_t M, int64_t N, int64_t K, int64_t E, int64_t topk, int64_t num_tokens_post_pad)
-
-INSTANTIATE_MOE_INT8_TEMPLATE(at::BFloat16);
-INSTANTIATE_MOE_INT8_TEMPLATE(at::Half);
-
-template <typename scalar_t>
-void shared_expert_int8_kernel_impl(
-    scalar_t* __restrict__ output,
-    scalar_t* __restrict__ ic1,
-    float* __restrict__ C_tmp,
-    uint8_t* __restrict__ Aq_tmp,
-    float* __restrict__ As_tmp,
-    const scalar_t* __restrict__ input,
-    const int8_t* __restrict__ packed_w1,
-    const int8_t* __restrict__ packed_w2,
-    const float* __restrict__ w1s,
-    const float* __restrict__ w2s,
-    const scalar_t* __restrict__ fused_experts_out,
-    float routed_scaling_factor,
-    int64_t M,
-    int64_t N,
-    int64_t K) {
-
-  // handle 2 tiles per block
-  constexpr int64_t BLOCK_M = block_size_m();
-  constexpr int64_t BLOCK_N = block_size_n();
-
-  // stage 0: quantize input to uint8, [M, K]
-  at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
-    for (int64_t m = begin; m < end; ++m) {
-      quantize_row_int8<scalar_t>(
-          Aq_tmp + m * K,
-          As_tmp[m],
-          input + m * K,
-          K);
-    }
-  });
-
-   // stage 1: intermediate_cache1 = silu(hidden_states @ w1)
-  const int64_t MB = div_up(M, BLOCK_M);
-  const int64_t NB = div_up(N, BLOCK_N);
-
-  TORCH_CHECK(N % BLOCK_N == 0, "Fixme when N is not multiples of ", BLOCK_N);
-
-  // K and N are packed for int8
-  const int64_t packed_K = get_row_size<int8_t>(K);
-  const int64_t packed_N = get_row_size<int8_t>(N);
-  const int64_t stride_n = packed_K;
-
-  // here we only parallel on half of 2N to fuse silu_and_mul with gemm
-  at::parallel_for(0, MB * NB, 0, [&](int64_t begin, int64_t end) {
-    for (int64_t i = begin; i < end; ++i) {
-      int64_t mb = i / NB;
-      int64_t nb = i % NB;
-
-      // nb0 from top half and nb1 from bottom half
-      int64_t nb0 = nb, nb1 = nb + NB;
-      int64_t n_size = std::min(N - nb0 * BLOCK_N, BLOCK_N);
-      int64_t m_size = std::min(M - mb * BLOCK_M, BLOCK_M);
-
-      // A shape [m_size, K]
-      const uint8_t* A = Aq_tmp + mb * BLOCK_M * K;
-      const float* As = As_tmp + mb * BLOCK_M;
-
-      // B shape [K, n_size] in vnni format
-      const int8_t* __restrict__ B0 = packed_w1 + nb0 * BLOCK_N * stride_n;
-      const int8_t* __restrict__ B1 = packed_w1 + nb1 * BLOCK_N * stride_n;
-      const float* __restrict__ Bs0 = w1s + nb0 * BLOCK_N;
-      const float* __restrict__ Bs1 = w1s + nb1 * BLOCK_N;
-
-      // fused 1.b: silu_and_mul(A @ B0, A @ B1)
-      tinygemm_kernel(
-          /* A     */ A,
-          /* B0    */ B0,
-          /* B1    */ B1,
-          /* C     */ ic1 + mb * BLOCK_M * N + nb * BLOCK_N,
-          /* As    */ As,
-          /* Bs0   */ Bs0,
-          /* Bs1   */ Bs1,
-          /* M     */ m_size,
-          /* N     */ n_size,
-          /* K     */ K,
-          /* lda   */ K,
-          /* ldb   */ n_size,
-          /* ldc   */ N);
-    }
-  });
-
-  // stage 1.5: quantize ic1 to uint8, [M * topk, N]
-  at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
-    for (int64_t m = begin; m < end; ++m) {
-      quantize_row_int8<scalar_t>(
-          Aq_tmp + m * N,
-          As_tmp[m],
-          ic1 + m * N,
-          N);
-    }
-  });
-
-  // stage 2: intermediate_cache2 = intermediate_cache1 @ w2
-  //   w2 : [K, N] as [OC, IC]
-  const int64_t OC = K;  // rename K as OC
-  const int64_t IC = N;  // rename N as IC
-  const int64_t MB2 = MB;
-  const int64_t NB2 = div_up(OC, BLOCK_N);
-  const int64_t stride_oc = packed_N;
-
-  // parallel on [MB2, NB2]
-  at::parallel_for(0, MB2 * NB2, 0, [&](int64_t begin, int64_t end) {
-    // get local pointers
-    int tid = at::get_thread_num();
-    // we won't be using C1 for gemm2
-    float* __restrict__ C = C_tmp + tid * 2 * BLOCK_M * BLOCK_N;
-
-    for (int64_t i = begin; i < end; ++i) {
-      int64_t mb = i / NB2;
-      int64_t nb = i % NB2;
-
-      int64_t m_size = std::min(M - mb * BLOCK_M, BLOCK_M);
-      int64_t n_size = std::min(OC - nb * BLOCK_N, BLOCK_N);
-
-      // A shape [m_size, IC]
-      const uint8_t* __restrict__ A = Aq_tmp + mb * BLOCK_M * N;
-      const float* __restrict__ As = As_tmp + mb * BLOCK_M;
-
-      // B shape [IC, n_size] in vnni format
-      const int8_t* __restrict__ B = packed_w2 + nb * BLOCK_N * stride_oc;
-      const float* __restrict__ Bs = w2s + nb * BLOCK_N;
-
-      // 2.a gemm: C = A @ B
-      tinygemm_kernel<scalar_t>(
-          /* A     */ A,
-          /* B     */ B,
-          /* C     */ C,
-          /* As    */ As,
-          /* Bs    */ Bs,
-          /* M     */ m_size,
-          /* N     */ n_size,
-          /* K     */ IC,
-          /* lda   */ IC,
-          /* ldb   */ n_size,
-          /* ldc   */ BLOCK_N);
-
-      // 2.b copy from C to output and add fused_experts_out
-      scalar_t* __restrict__ out = output + mb * BLOCK_M * K + nb * BLOCK_N;
-      const scalar_t* __restrict__ fused_out = fused_experts_out + mb * BLOCK_M * K + nb * BLOCK_N;
-      for (int64_t m = 0; m < m_size; ++m) {
-        add_mul_stub(out + m * K, C + m * BLOCK_N, fused_out + m * K, routed_scaling_factor, n_size);
-      }
-    }
-  });
-}
-
-#define INSTANTIATE_SHARED_EXPERT_INT8_TEMPLATE(TYPE)                                        \
-  template void shared_expert_int8_kernel_impl<TYPE> (                                       \
-      TYPE* __restrict__ output, TYPE* __restrict__ ic1,                                     \
-      float* __restrict__ C_tmp, uint8_t* __restrict__ Aq_tmp,                               \
-      float* __restrict__ As_tmp, const TYPE* __restrict__ input,                            \
-      const int8_t* __restrict__ packed_w1, const int8_t* __restrict__ packed_w2,            \
-      const float* __restrict__ w1s, const float* __restrict__ w2s,                          \
-      const TYPE* __restrict__ fused_experts_out, float routed_scaling_factor,               \
-      int64_t M, int64_t N, int64_t K)
-
-INSTANTIATE_SHARED_EXPERT_INT8_TEMPLATE(at::BFloat16);
-INSTANTIATE_SHARED_EXPERT_INT8_TEMPLATE(at::Half);

csrc/cpu/sgl-kernels/vec.h

@@ -1,308 +0,0 @@
-// Adapted from
-// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu
-
-#pragma once
-
-// clang-format off
-
-#if defined(__AVX512F__) && defined(__AVX512BF16__) && defined(__AMX_BF16__)
-#define CPU_CAPABILITY_AVX512
-#endif
-
-#include <ATen/cpu/vec/functional.h>
-#include <ATen/cpu/vec/vec.h>
-
-namespace {
-
-using namespace at::vec;
-
-template <typename scalar_t,
-          typename std::enable_if_t<is_reduced_floating_point_v<scalar_t>, int> = 0>
-inline Vectorized<scalar_t> convert_from_float_ext(const Vectorized<float>& a, const Vectorized<float>& b) {
-  return at::vec::convert_from_float<scalar_t>(a, b);
-}
-
-#if defined(CPU_CAPABILITY_AVX512)
-
-// `at::vec::convert_from_float<>` from PyTorch doesn't have avx512-bf16 intrinsics
-// use native instruction for bfloat16->float32 conversion
-template <>
-inline Vectorized<at::BFloat16> convert_from_float_ext<at::BFloat16>(const Vectorized<float>& a, const Vectorized<float>& b) {
-  return (__m512i)(_mm512_cvtne2ps_pbh(__m512(b), __m512(a)));
-}
-
-#define CVT_BF16_TO_FP32(a) \
-    _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(a), 16))
-
-#define CVT_FP16_TO_FP32(a) \
-    _mm512_cvtps_ph(a, (_MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC))
-
-// this doesn't hanel NaN.
-inline __m512bh cvt_e4m3_bf16_intrinsic_no_nan(__m256i fp8_vec) {
-  const __m512i x = _mm512_cvtepu8_epi16(fp8_vec);
-
-  const __m512i mant = _mm512_slli_epi16(_mm512_and_si512(x, _mm512_set1_epi16(0x07)), 4);
-  const __m512i raw_exp = _mm512_srli_epi16(_mm512_and_si512(x, _mm512_set1_epi16(0x78)), 3);
-  const __m512i exp = _mm512_slli_epi16(_mm512_add_epi16(raw_exp, _mm512_set1_epi16(120)), 7);
-  const __m512i nonsign = _mm512_or_si512(exp, mant);
-
-  const __m512i sign = _mm512_slli_epi16(_mm512_and_si512(x, _mm512_set1_epi16(0x80)), 8);
-  const __m512i combined = _mm512_or_si512(nonsign, sign);
-
-  const __mmask32 is_nonzero = _mm512_cmpneq_epi16_mask(x, _mm512_setzero_si512());
-  return (__m512bh)_mm512_maskz_mov_epi16(is_nonzero, combined);
-}
-
-inline __m512bh cvt_e4m3_bf16_intrinsic_without_denorm(__m256i fp8_vec) {
-  // The following conversion is without denorm behavior, that is to say,
-  //   Max subnorm   : S.0000.111 = 0.875 ∗ 2**(−6)
-  //   Min subnorm   : S.0000.001 = 2**(−9)
-  // 0.0019 ~ 0.0137 cannot be converted correctly.
-  __m512i x = _mm512_cvtepu8_epi16(fp8_vec);
-  auto mask = _mm512_cmpneq_epi16_mask(
-      _mm512_and_si512(x, _mm512_set1_epi16(127)),
-      _mm512_setzero_si512());  // mask = x & 0x7f
-  auto mask_nan = _mm512_cmpneq_epi16_mask(
-      _mm512_and_si512(x, _mm512_set1_epi16(127)),
-      _mm512_set1_epi16(127));                                                      // mask_nan = x & 0x7f
-  auto mantissa = _mm512_slli_epi16(_mm512_and_si512(x, _mm512_set1_epi16(7)), 4);  // mantissa = (x & 7) << 4
-  auto exponent = _mm512_add_epi16(
-      _mm512_srli_epi16(_mm512_and_si512(x, _mm512_set1_epi16(120)), 3),
-      _mm512_set1_epi16(120));  // exponent = (((x >> 3) & 15) + 120)
-  auto nonsign = _mm512_maskz_mov_epi16(mask, _mm512_or_si512(mantissa, _mm512_slli_epi16(exponent, 7)));
-  nonsign = _mm512_mask_mov_epi16(_mm512_set1_epi16(0x7fff), mask_nan, nonsign);  // deal with Nan
-  return (__m512bh)(_mm512_or_si512(
-      nonsign,
-      _mm512_slli_epi16(
-          _mm512_and_si512(x, _mm512_set1_epi16(128)),
-          8)));  // add sign (x & 128) << 8
-}
-
-inline __m512bh cvt_e4m3_bf16_intrinsic_with_denorm(__m256i fp8_vec) {
-  __m512i x = _mm512_cvtepu8_epi16(fp8_vec);
-  __m512i lg2mant = _mm512_mask_mov_epi16(
-      _mm512_mask_mov_epi16(
-          _mm512_setzero_si512(), _mm512_test_epi16_mask(x, _mm512_set1_epi16(2)), _mm512_set1_epi16(1)),
-      _mm512_test_epi16_mask(x, _mm512_set1_epi16(4)),
-      _mm512_set1_epi16(2));
-  return (__m512bh)(_mm512_or_si512(
-      _mm512_maskz_mov_epi16(
-          _mm512_cmpneq_epi16_mask(_mm512_and_si512(x, _mm512_set1_epi16(127)), _mm512_setzero_si512()),
-          _mm512_mask_blend_epi16(
-              _mm512_test_epi16_mask(x, _mm512_set1_epi16(120)),
-              _mm512_or_si512(
-                  _mm512_and_si512(
-                      _mm512_sllv_epi16(
-                          _mm512_and_si512(x, _mm512_set1_epi16(3)), _mm512_sub_epi16(_mm512_set1_epi16(7), lg2mant)),
-                      _mm512_set1_epi16(0x007f)),
-                  _mm512_slli_epi16(_mm512_add_epi16(lg2mant, _mm512_set1_epi16(118)), 7)),
-              _mm512_or_si512(
-                  _mm512_slli_epi16(_mm512_and_si512(x, _mm512_set1_epi16(7)), 4),
-                  _mm512_slli_epi16(
-                      _mm512_add_epi16(
-                          _mm512_srli_epi16(_mm512_and_si512(x, _mm512_set1_epi16(120)), 3), _mm512_set1_epi16(120)),
-                      7)))),
-      _mm512_slli_epi16(_mm512_and_si512(x, _mm512_set1_epi16(128)), 8)));
-}
-
-inline __m512bh CVT_FP8_TO_BF16(__m256i a) {
-#ifdef SGLANG_CPU_FP8_CVT_FTZ
-  return cvt_e4m3_bf16_intrinsic_no_nan(a);
-#else
-  return cvt_e4m3_bf16_intrinsic_with_denorm(a);
-#endif
-}
-
-#endif
-
-// vector to scalar reduction
-#if defined(CPU_CAPABILITY_AVX512) && 0
-inline float vec_reduce_sum(const Vectorized<float>& a) {
-  return _mm512_reduce_add_ps(__m512(a));
-}
-
-inline float vec_reduce_max(const Vectorized<float>& a) {
-  return _mm512_reduce_max_ps(__m512(a));
-}
-#else
-inline float vec_reduce_sum(const Vectorized<float>& a) {
-  return vec_reduce_all([](Vectorized<float>& x, Vectorized<float>& y) { return x + y; }, a);
-}
-
-inline float vec_reduce_max(const Vectorized<float>& a) {
-  return vec_reduce_all([](Vectorized<float>& x, Vectorized<float>& y) { return maximum(x, y); }, a);
-}
-#endif
-
-// https://github.com/InternLM/lmdeploy/blob/086481ed84b59bee3b8e4274e5fc69620040c048/lmdeploy/pytorch/kernels/cuda/w8a8_triton_kernels.py#L282
-template <typename scalar_t>
-inline void quantize_row_int8(uint8_t* __restrict__ Aq, float& As,
-    const scalar_t* __restrict__ A, int64_t K, float eps = 1e-7) {
-
-  float amax = 0.f; // absolute max
-  for (int64_t k = 0; k < K; ++k) {
-    const float val = static_cast<float>(A[k]);
-    amax = std::max(amax, std::abs(val));
-  }
-
-  amax = std::max(amax, eps);
-  const float scale = amax / 127;
-  const float inv_scale = 127 / amax;
-
-  for (int64_t k = 0; k < K; ++k) {
-    const float val = static_cast<float>(A[k]) * inv_scale;
-    Aq[k] = (uint8_t)(std::round(val)) + 128;
-  }
-  As = scale;
-}
-
-#if defined(CPU_CAPABILITY_AVX512)
-template <>
-inline void quantize_row_int8<at::BFloat16>(uint8_t* __restrict__ Aq, float& As,
-    const at::BFloat16* __restrict__ A, int64_t K, float eps) {
-
-  const __m512 signBit = _mm512_set1_ps(-0.0f);
-  const __m512i off = _mm512_set1_epi32(128);
-
-  // K is 32x, no remainder
-  float amax = 0.f;
-  __m512 vamax0 = _mm512_set1_ps(0.f);
-  __m512 vamax1 = _mm512_set1_ps(0.f);
-  for (int64_t k = 0; k < K; k += 32) {
-    __m512i va = _mm512_loadu_si512((void*)(A + k));
-    __m512 va0 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(va, 0));
-    __m512 va1 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(va, 1));
-    vamax0 = _mm512_max_ps(vamax0, _mm512_andnot_ps(signBit, va0));
-    vamax1 = _mm512_max_ps(vamax1, _mm512_andnot_ps(signBit, va1));
-  }
-  amax = _mm512_reduce_max_ps(_mm512_max_ps(vamax0, vamax1));
-  amax = std::max(amax, eps);
-  const float scale = amax / 127;
-  const float inv_scale = 127 / amax;
-  const __m512 vd = _mm512_set1_ps(inv_scale);
-
-  for (int64_t k = 0; k < K; k += 32) {
-    __m512i va = _mm512_loadu_si512((void*)(A + k));
-    __m512 va0 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(va, 0));
-    __m512 va1 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(va, 1));
-    va0 = _mm512_mul_ps(va0, vd);
-    va1 = _mm512_mul_ps(va1, vd);
-    va0 = _mm512_roundscale_ps(va0, (_MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC));
-    va1 = _mm512_roundscale_ps(va1, (_MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC));
-    __m128i i0 = _mm512_cvtepi32_epi8(_mm512_add_epi32(_mm512_cvtps_epi32(va0), off));
-    __m128i i1 = _mm512_cvtepi32_epi8(_mm512_add_epi32(_mm512_cvtps_epi32(va1), off));
-    _mm256_storeu_si256(reinterpret_cast<__m256i*>(Aq + k), _mm256_set_m128i(i1, i0));
-  }
-  As = scale;
-}
-#endif
-
-// transpose utils
-// taken from my PR in ggml: https://github.com/ggml-org/llama.cpp/pull/8998
-#if defined(CPU_CAPABILITY_AVX512)
-inline void transpose_16x16_32bit(__m512i * v) {
-  __m512i v1[16];
-  v1[0] = _mm512_unpacklo_epi32(v[0], v[1]);
-  v1[1] = _mm512_unpackhi_epi32(v[0], v[1]);
-  v1[2] = _mm512_unpacklo_epi32(v[2], v[3]);
-  v1[3] = _mm512_unpackhi_epi32(v[2], v[3]);
-  v1[4] = _mm512_unpacklo_epi32(v[4], v[5]);
-  v1[5] = _mm512_unpackhi_epi32(v[4], v[5]);
-  v1[6] = _mm512_unpacklo_epi32(v[6], v[7]);
-  v1[7] = _mm512_unpackhi_epi32(v[6], v[7]);
-  v1[8] = _mm512_unpacklo_epi32(v[8], v[9]);
-  v1[9] = _mm512_unpackhi_epi32(v[8], v[9]);
-  v1[10] = _mm512_unpacklo_epi32(v[10], v[11]);
-  v1[11] = _mm512_unpackhi_epi32(v[10], v[11]);
-  v1[12] = _mm512_unpacklo_epi32(v[12], v[13]);
-  v1[13] = _mm512_unpackhi_epi32(v[12], v[13]);
-  v1[14] = _mm512_unpacklo_epi32(v[14], v[15]);
-  v1[15] = _mm512_unpackhi_epi32(v[14], v[15]);
-
-  v[0] = _mm512_unpacklo_epi64(v1[0], v1[2]);
-  v[1] = _mm512_unpackhi_epi64(v1[0], v1[2]);
-  v[2] = _mm512_unpacklo_epi64(v1[1], v1[3]);
-  v[3] = _mm512_unpackhi_epi64(v1[1], v1[3]);
-  v[4] = _mm512_unpacklo_epi64(v1[4], v1[6]);
-  v[5] = _mm512_unpackhi_epi64(v1[4], v1[6]);
-  v[6] = _mm512_unpacklo_epi64(v1[5], v1[7]);
-  v[7] = _mm512_unpackhi_epi64(v1[5], v1[7]);
-  v[8] = _mm512_unpacklo_epi64(v1[8], v1[10]);
-  v[9] = _mm512_unpackhi_epi64(v1[8], v1[10]);
-  v[10] = _mm512_unpacklo_epi64(v1[9], v1[11]);
-  v[11] = _mm512_unpackhi_epi64(v1[9], v1[11]);
-  v[12] = _mm512_unpacklo_epi64(v1[12], v1[14]);
-  v[13] = _mm512_unpackhi_epi64(v1[12], v1[14]);
-  v[14] = _mm512_unpacklo_epi64(v1[13], v1[15]);
-  v[15] = _mm512_unpackhi_epi64(v1[13], v1[15]);
-
-  v1[0] = _mm512_shuffle_i32x4(v[0], v[4], 0x88);
-  v1[1] = _mm512_shuffle_i32x4(v[1], v[5], 0x88);
-  v1[2] = _mm512_shuffle_i32x4(v[2], v[6], 0x88);
-  v1[3] = _mm512_shuffle_i32x4(v[3], v[7], 0x88);
-  v1[4] = _mm512_shuffle_i32x4(v[0], v[4], 0xdd);
-  v1[5] = _mm512_shuffle_i32x4(v[1], v[5], 0xdd);
-  v1[6] = _mm512_shuffle_i32x4(v[2], v[6], 0xdd);
-  v1[7] = _mm512_shuffle_i32x4(v[3], v[7], 0xdd);
-  v1[8] = _mm512_shuffle_i32x4(v[8], v[12], 0x88);
-  v1[9] = _mm512_shuffle_i32x4(v[9], v[13], 0x88);
-  v1[10] = _mm512_shuffle_i32x4(v[10], v[14], 0x88);
-  v1[11] = _mm512_shuffle_i32x4(v[11], v[15], 0x88);
-  v1[12] = _mm512_shuffle_i32x4(v[8], v[12], 0xdd);
-  v1[13] = _mm512_shuffle_i32x4(v[9], v[13], 0xdd);
-  v1[14] = _mm512_shuffle_i32x4(v[10], v[14], 0xdd);
-  v1[15] = _mm512_shuffle_i32x4(v[11], v[15], 0xdd);
-
-  v[0] = _mm512_shuffle_i32x4(v1[0], v1[8], 0x88);
-  v[1] = _mm512_shuffle_i32x4(v1[1], v1[9], 0x88);
-  v[2] = _mm512_shuffle_i32x4(v1[2], v1[10], 0x88);
-  v[3] = _mm512_shuffle_i32x4(v1[3], v1[11], 0x88);
-  v[4] = _mm512_shuffle_i32x4(v1[4], v1[12], 0x88);
-  v[5] = _mm512_shuffle_i32x4(v1[5], v1[13], 0x88);
-  v[6] = _mm512_shuffle_i32x4(v1[6], v1[14], 0x88);
-  v[7] = _mm512_shuffle_i32x4(v1[7], v1[15], 0x88);
-  v[8] = _mm512_shuffle_i32x4(v1[0], v1[8], 0xdd);
-  v[9] = _mm512_shuffle_i32x4(v1[1], v1[9], 0xdd);
-  v[10] = _mm512_shuffle_i32x4(v1[2], v1[10], 0xdd);
-  v[11] = _mm512_shuffle_i32x4(v1[3], v1[11], 0xdd);
-  v[12] = _mm512_shuffle_i32x4(v1[4], v1[12], 0xdd);
-  v[13] = _mm512_shuffle_i32x4(v1[5], v1[13], 0xdd);
-  v[14] = _mm512_shuffle_i32x4(v1[6], v1[14], 0xdd);
-  v[15] = _mm512_shuffle_i32x4(v1[7], v1[15], 0xdd);
-}
-
-// remove warning : ignoring attributes on template argument ‘__m512i’ [-Wignored-attributes]
-#pragma GCC diagnostic push
-#pragma GCC diagnostic ignored "-Wignored-attributes"
-
-// transpose from [2, 32] to [32, 2]
-inline std::tuple<__m512i, __m512i> transpose_2x32_16bit(__m512i r0, __m512i r1) {
-  // r0: {a0, a1, ..., a31}
-  // r1: {b0, b1, ..., b31}
-  //
-  // d0: {a0,   b0, ..., a15, b15}
-  // d1: {a16, b16, ..., a31, b31}
-  //
-  __m512i d0 = _mm512_unpacklo_epi16(r0, r1);
-  __m512i d1 = _mm512_unpackhi_epi16(r0, r1);
-  r0 = _mm512_shuffle_i32x4(d0, d1, 0x88);
-  r1 = _mm512_shuffle_i32x4(d0, d1, 0xdd);
-  d0 = _mm512_shuffle_i32x4(r0, r1, 0x88);
-  d1 = _mm512_shuffle_i32x4(r0, r1, 0xdd);
-  return std::make_tuple(d0, d1);
-}
-#pragma GCC diagnostic pop
-
-#endif
-
-// TODO: debug print, remove me later
-template<typename scalar_t>
-void print_array(scalar_t* ptr, int size) {
-  for (int d = 0; d < size; ++d) {
-    if (d % 16 == 0) { std::cout << std::endl; }
-    std::cout << ptr[d] << " ";
-  }
-  std::cout << std::endl;
-}
-
-} // anonymous namespace

csrc/cpu/shm.cpp

@@ -7,10 +7,9 @@
 
 namespace {
 #define MAX_SHM_RANK_NUM 8
-#define PER_THREAD_SHM_BUFFER_BYTES (2 * 1024 * 1024)
-static_assert(PER_THREAD_SHM_BUFFER_BYTES % 2 == 0);
-#define PER_THREAD_SHM_BUFFER_OFFSET (PER_THREAD_SHM_BUFFER_BYTES >> 1)
-#define MIN_THREAD_PROCESS_SIZE (256)
+#define MAX_THREAD_NUM 12
+#define PER_THREAD_SHM_BUFFER_BYTES (4 * 1024 * 1024)
+#define MIN_THREAD_PROCESS_SIZE (8 * 1024)
 #define MAX_P2P_SEND_TENSOR_NUM 8
 
 template <typename scalar_t>
@@ -33,10 +32,10 @@ struct KernelVecType<c10::Half> {
   using scalar_vec_t = vec_op::FP16Vec16;
 };
 
+enum class ThreadSHMStat : char { THREAD_READY = 0, SHM_DATA_READY, DONE };
+
 struct ThreadSHMContext {
-  volatile char _curr_thread_stamp;
-  volatile char _ready_thread_stamp;
-  char _padding1[6];
+  volatile ThreadSHMStat thread_stats[MAX_SHM_RANK_NUM];
   int thread_id;
   int thread_num;
   int rank;
@@ -45,19 +44,14 @@ struct ThreadSHMContext {
   int swizzled_ranks[MAX_SHM_RANK_NUM];
   void* thread_shm_ptrs[MAX_SHM_RANK_NUM];
   ThreadSHMContext* shm_contexts[MAX_SHM_RANK_NUM];
-  size_t _thread_buffer_mask;
-  char _padding2[56];
 
   ThreadSHMContext(const int thread_id, const int thread_num, const int rank,
                    const int group_size, void* thread_shm_ptr)
-      : _curr_thread_stamp(1),
-        _ready_thread_stamp(0),
-        thread_id(thread_id),
+      : thread_id(thread_id),
         thread_num(thread_num),
         rank(rank),
         group_size(group_size),
-        _spinning_count(0),
-        _thread_buffer_mask(0) {
+        _spinning_count(0) {
     static_assert(sizeof(ThreadSHMContext) % 64 == 0);
     TORCH_CHECK(group_size <= MAX_SHM_RANK_NUM);
     TORCH_CHECK((size_t)this % 64 == 0);
@@ -66,6 +60,7 @@ struct ThreadSHMContext {
       shm_contexts[i] = nullptr;
       thread_shm_ptrs[i] = nullptr;
       swizzled_ranks[i] = (i + rank) % group_size;
+      thread_stats[i] = ThreadSHMStat::DONE;
     }
     set_context(rank, this, thread_shm_ptr);
   }
@@ -82,66 +77,59 @@ struct ThreadSHMContext {
 
   template <typename T>
   T* get_thread_shm_ptr(int rank) {
-    return reinterpret_cast<T*>(
-        reinterpret_cast<int8_t*>(thread_shm_ptrs[rank]) +
-        (PER_THREAD_SHM_BUFFER_OFFSET & _thread_buffer_mask));
-  }
-
-  void next_buffer() { _thread_buffer_mask ^= 0xFFFFFFFFFFFFFFFF; }
-
-  char get_curr_stamp() const { return _curr_thread_stamp; }
-
-  char get_ready_stamp() const { return _ready_thread_stamp; }
-
-  void next_stamp() {
-    _mm_mfence();
-    _curr_thread_stamp += 1;
-  }
-
-  void commit_ready_stamp() {
-    _mm_mfence();
-    _ready_thread_stamp = _curr_thread_stamp;
+    return reinterpret_cast<T*>(thread_shm_ptrs[rank]);
   }
 
   int get_swizzled_rank(int idx) { return swizzled_ranks[idx]; }
 
-  template <typename Cond>
-  void wait_for_all(Cond&& cond) {
-    for (int idx = 1; idx < group_size; ++idx) {
+  void wait_for_all(ThreadSHMStat prev_stat) {
+    for (int idx = 0; idx < group_size; ++idx) {
       int rank = get_swizzled_rank(idx);
-      wait_for_one(rank, std::forward<Cond>(cond));
+      while (thread_stats[rank] == prev_stat) {
+        ++_spinning_count;
+        _mm_pause();
+      }
     }
+    vec_op::mem_barrier();
   }
 
-  template <typename Cond>
-  void wait_for_one(int rank, Cond&& cond) {
-    ThreadSHMContext* rank_ctx = shm_contexts[rank];
-    for (;;) {
-      char local_curr_stamp = get_curr_stamp();
-      char local_ready_stamp = get_ready_stamp();
-      char rank_curr_stamp = rank_ctx->get_curr_stamp();
-      char rank_ready_stamp = rank_ctx->get_ready_stamp();
-      if (cond(local_curr_stamp, local_ready_stamp, rank_curr_stamp,
-               rank_ready_stamp)) {
-        break;
-      }
+  void wait_for_one(int rank, ThreadSHMStat prev_stat) {
+    while (thread_stats[rank] == prev_stat) {
       ++_spinning_count;
       _mm_pause();
     }
+    vec_op::mem_barrier();
   }
 
-  static bool check_no_buffer_conflict(char local_curr_stamp,
-                                       char local_ready_stamp,
-                                       char rank_curr_stamp,
-                                       char rank_ready_stamp) {
-    char temp = rank_curr_stamp + 2;
-    return local_curr_stamp != temp;
+  void set_thread_stat(ThreadSHMStat stat) {
+    for (int idx = 0; idx < group_size; ++idx) {
+      int rank = get_swizzled_rank(idx);
+      shm_contexts[rank]->thread_stats[this->rank] = stat;
+    }
   }
 
-  static bool check_stamp_ready(char local_curr_stamp, char local_ready_stamp,
-                                char rank_curr_stamp, char rank_ready_stamp) {
-    char temp = local_curr_stamp + 1;
-    return (local_curr_stamp == rank_ready_stamp) || (temp == rank_ready_stamp);
+  void set_thread_stat(int target_rank, ThreadSHMStat stat) {
+    for (int idx = 0; idx < group_size; ++idx) {
+      int rank = get_swizzled_rank(idx);
+      shm_contexts[rank]->thread_stats[target_rank] = stat;
+    }
+  }
+
+  // barrier for all ranks in the group, used for all2all ops
+  // DONE -> THREAD_READY -> SHM_DATA_READY -> DONE -> ...
+  void barrier(ThreadSHMStat next_stat) {
+    if (next_stat == ThreadSHMStat::THREAD_READY) {
+      set_thread_stat(ThreadSHMStat::THREAD_READY);
+      wait_for_all(ThreadSHMStat::DONE);
+    } else if (next_stat == ThreadSHMStat::SHM_DATA_READY) {
+      set_thread_stat(ThreadSHMStat::SHM_DATA_READY);
+      wait_for_all(ThreadSHMStat::THREAD_READY);
+    } else if (next_stat == ThreadSHMStat::DONE) {
+      set_thread_stat(ThreadSHMStat::DONE);
+      wait_for_all(ThreadSHMStat::SHM_DATA_READY);
+    } else {
+      TORCH_CHECK(false, "Invalid next_stat to barrier.");
+    }
   }
 
   std::string to_string() const {
@@ -176,7 +164,7 @@ class SHMManager {
                       const int group_size)
       : _rank(rank),
         _group_size(group_size),
-        _thread_num(torch::get_num_threads()),
+        _thread_num(std::min(torch::get_num_threads(), MAX_THREAD_NUM)),
         _shm_names({""}),
         _shared_mem_ptrs({nullptr}),
         _shm_ctx(nullptr) {
@@ -338,8 +326,7 @@ void shm_cc_loop(ThreadSHMContext* ctx, int64_t elem_num, F&& inner_func) {
       (total_units_num + thread_num - 1) / thread_num;
   int64_t per_unit_elem_num = MIN_THREAD_PROCESS_SIZE / sizeof(scalar_t);
   int64_t max_per_thread_iteration_elem_num =
-      (PER_THREAD_SHM_BUFFER_BYTES >> 1) /
-      sizeof(scalar_t);  // Note: double buffer
+      PER_THREAD_SHM_BUFFER_BYTES / sizeof(scalar_t);
   int64_t per_thread_elem_num = per_unit_elem_num * per_thread_units_num;
 
 #pragma omp parallel for schedule(static, 1)
@@ -349,13 +336,10 @@ void shm_cc_loop(ThreadSHMContext* ctx, int64_t elem_num, F&& inner_func) {
     int64_t curr_elem_num =
         std::min(max_per_thread_iteration_elem_num, end - offset);
     ThreadSHMContext* thread_ctx = ctx + i;
-    bool fast_mode = ((end - offset) <= max_per_thread_iteration_elem_num);
 
     while (curr_elem_num > 0) {
-      inner_func(thread_ctx, offset, curr_elem_num, fast_mode);
+      inner_func(thread_ctx, offset, curr_elem_num);
 
-      thread_ctx->next_stamp();
-      thread_ctx->next_buffer();
       offset += max_per_thread_iteration_elem_num;
       curr_elem_num = std::min(max_per_thread_iteration_elem_num, end - offset);
     }
@@ -413,7 +397,7 @@ void all_reduce_sum_impl(ThreadSHMContext* ctx, scalar_t* data,
   shm_cc_ops::shm_cc_loop<scalar_t>(
       ctx, elem_num,
       [&](ThreadSHMContext* thread_ctx, int64_t data_offset,
-          int64_t data_elem_num, bool fast_mode) {
+          int64_t data_elem_num) {
         int rank = thread_ctx->rank;
         scalar_t* thread_shm_ptr =
             thread_ctx->get_thread_shm_ptr<scalar_t>(rank);
@@ -426,17 +410,16 @@ void all_reduce_sum_impl(ThreadSHMContext* ctx, scalar_t* data,
               thread_ctx->get_swizzled_rank(idx + 1));
         });
 
-        if (!fast_mode) {
-          thread_ctx->wait_for_all(ThreadSHMContext::check_no_buffer_conflict);
-        }
+        thread_ctx->barrier(ThreadSHMStat::THREAD_READY);
 
         shm_cc_ops::memcpy_to_shm(thread_shm_ptr, thread_data_ptr,
                                   thread_data_elem_num);
-        thread_ctx->commit_ready_stamp();
+
+        thread_ctx->barrier(ThreadSHMStat::SHM_DATA_READY);
+
         int64_t aligned_data_elem_num =
             (data_elem_num / vec_elem_num) * vec_elem_num;
         int64_t i = 0;
-        thread_ctx->wait_for_all(ThreadSHMContext::check_stamp_ready);
 #pragma GCC unroll 4
         for (; i < aligned_data_elem_num; i += vec_elem_num) {
           vec_t local_data(thread_data_ptr + i);  // load from cache
@@ -464,6 +447,8 @@ void all_reduce_sum_impl(ThreadSHMContext* ctx, scalar_t* data,
           reduced_data.save(thread_data_ptr + i,
                             data_elem_num - aligned_data_elem_num);
         }
+
+        thread_ctx->barrier(ThreadSHMStat::DONE);
       });
 
   return;
@@ -503,18 +488,18 @@ void shm_gather_impl(ThreadSHMContext* ctx, scalar_t* data, size_t elem_num,
   shm_cc_ops::shm_cc_loop<scalar_t>(
       ctx, elem_num,
       [&](ThreadSHMContext* thread_ctx, int64_t data_offset,
-          int64_t data_elem_num, bool fast_mode) {
+          int64_t data_elem_num) {
         int rank = thread_ctx->rank;
         scalar_t* thread_shm_ptr =
             thread_ctx->get_thread_shm_ptr<scalar_t>(rank);
 
-        if (!fast_mode) {
-          thread_ctx->wait_for_all(ThreadSHMContext::check_no_buffer_conflict);
-        }
+        thread_ctx->barrier(ThreadSHMStat::THREAD_READY);
+
+        shm_cc_ops::memcpy_to_shm(thread_shm_ptr, data + data_offset,
+                                  data_elem_num * sizeof(scalar_t));
+
+        thread_ctx->barrier(ThreadSHMStat::SHM_DATA_READY);
 
-        shm_cc_ops::memcpy(thread_shm_ptr, data + data_offset,
-                           data_elem_num * sizeof(scalar_t));
-        thread_ctx->commit_ready_stamp();
         if (rank == dst) {
           shm_cc_ops::memcpy(outputs[rank] + data_offset, data + data_offset,
                              data_elem_num * sizeof(scalar_t));
@@ -523,12 +508,12 @@ void shm_gather_impl(ThreadSHMContext* ctx, scalar_t* data, size_t elem_num,
             scalar_t* src_ptr =
                 thread_ctx->get_thread_shm_ptr<scalar_t>(src_rank);  // shm
             scalar_t* dst_ptr = outputs[src_rank] + data_offset;
-            thread_ctx->wait_for_one(src_rank,
-                                     ThreadSHMContext::check_stamp_ready);
-            shm_cc_ops::memcpy(dst_ptr, src_ptr,
-                               data_elem_num * sizeof(scalar_t));
+            shm_cc_ops::memcpy_from_shm(dst_ptr, src_ptr,
+                                        data_elem_num * sizeof(scalar_t));
           }
         }
+
+        thread_ctx->barrier(ThreadSHMStat::DONE);
       });
 
   return;
@@ -614,7 +599,7 @@ struct TensorListMeta {
   int8_t _padding[40];
 };
 
-void shm_send_tensor_list_impl(ThreadSHMContext* ctx, int64_t dst,
+void shm_send_tensor_list_impl(ThreadSHMContext* ctx,
                                const std::vector<torch::Tensor>& tensor_list) {
   CPU_KERNEL_GUARD_IN(shm_send_tensor_list_impl)
   std::vector<torch::Tensor> tensor_list_with_metadata;
@@ -635,11 +620,12 @@ void shm_send_tensor_list_impl(ThreadSHMContext* ctx, int64_t dst,
   shm_cc_ops::shm_cc_loop<int8_t>(
       ctx, metadata->total_bytes,
       [&](ThreadSHMContext* thread_ctx, int64_t data_offset,
-          int64_t data_elem_num, bool fast_mode) {
+          int64_t data_elem_num) {
         int rank = thread_ctx->rank;
+        // Wait until the receiver set the stat to DONE
+        thread_ctx->wait_for_one(rank, ThreadSHMStat::SHM_DATA_READY);
+
         int64_t curr_shm_offset = 0;
-        thread_ctx->wait_for_one(dst,
-                                 ThreadSHMContext::check_no_buffer_conflict);
         while (curr_shm_offset < data_elem_num) {
           MemPiece frag = metadata->get_data(data_offset + curr_shm_offset);
           frag.size = std::min(frag.size, data_elem_num - curr_shm_offset);
@@ -648,7 +634,8 @@ void shm_send_tensor_list_impl(ThreadSHMContext* ctx, int64_t dst,
               frag.ptr, frag.size);
           curr_shm_offset += frag.size;
         }
-        thread_ctx->commit_ready_stamp();
+
+        thread_ctx->set_thread_stat(rank, ThreadSHMStat::SHM_DATA_READY);
       });
 }
 
@@ -659,7 +646,8 @@ std::vector<torch::Tensor> shm_recv_tensor_list_impl(ThreadSHMContext* ctx,
   torch::Tensor metadata_tensor =
       torch::empty({sizeof(TensorListMeta)}, options);
 
-  ctx->wait_for_one(src, ThreadSHMContext::check_stamp_ready);
+  // Wait until the sender set the stat of the thread 0 to SHM_DATA_READY
+  ctx->wait_for_one(src, ThreadSHMStat::DONE);
   shm_cc_ops::memcpy(metadata_tensor.data_ptr(),
                      ctx->get_thread_shm_ptr<void>(src),
                      sizeof(TensorListMeta));
@@ -676,8 +664,9 @@ std::vector<torch::Tensor> shm_recv_tensor_list_impl(ThreadSHMContext* ctx,
   shm_cc_ops::shm_cc_loop<int8_t>(
       ctx, metadata.total_bytes,
       [&](ThreadSHMContext* thread_ctx, int64_t data_offset,
-          int64_t data_elem_num, bool fast_mode) {
-        ctx->wait_for_one(src, ThreadSHMContext::check_stamp_ready);
+          int64_t data_elem_num) {
+        // Wait until the sender set the stat to SHM_DATA_READY
+        thread_ctx->wait_for_one(src, ThreadSHMStat::DONE);
         int64_t curr_shm_offset = 0;
         while (curr_shm_offset < data_elem_num) {
           MemPiece frag = metadata.get_data(data_offset + curr_shm_offset);
@@ -688,6 +677,8 @@ std::vector<torch::Tensor> shm_recv_tensor_list_impl(ThreadSHMContext* ctx,
               frag.size);
           curr_shm_offset += frag.size;
         }
+
+        thread_ctx->set_thread_stat(src, ThreadSHMStat::DONE);
       });
 
   std::vector<torch::Tensor> tensor_list;
@@ -765,8 +756,7 @@ void shm_send_tensor_list(int64_t handle,
                           int64_t dst) {
   CPU_KERNEL_GUARD_IN(shm_send_tensor_list)
   shm_send_tensor_list_impl(
-      SHMManager::get_singleton_instance(handle)->get_shm_ctx(), dst,
-      tensor_list);
+      SHMManager::get_singleton_instance(handle)->get_shm_ctx(), tensor_list);
   CPU_KERNEL_GUARD_OUT(shm_send_tensor_list)
 }
 
@@ -788,4 +778,4 @@ std::string join_shm_manager(int64_t handle, const std::string& name) {
   TORCH_CHECK(shm_manager);
   shm_manager->join(name);
   return shm_manager->get_shm_ctx()->to_string();
-}
+}

csrc/cpu/torch_bindings.cpp

@@ -50,27 +50,6 @@ void shm_send_tensor_list(int64_t handle,
 
 std::vector<torch::Tensor> shm_recv_tensor_list(int64_t handle, int64_t src);
 
-at::Tensor weight_packed_linear(at::Tensor& mat1, at::Tensor& mat2,
-                                const std::optional<at::Tensor>& bias,
-                                bool is_vnni);
-
-at::Tensor convert_weight_packed(at::Tensor& weight);
-
-at::Tensor fused_experts_cpu(
-    at::Tensor& hidden_states, at::Tensor& w1, at::Tensor& w2,
-    at::Tensor& topk_weights, at::Tensor& topk_ids, bool inplace,
-    bool use_int8_w8a8, bool use_fp8_w8a16,
-    const std::optional<at::Tensor>& w1_scale,
-    const std::optional<at::Tensor>& w2_scale,
-    const std::optional<std::vector<int64_t>> block_size,
-    const std::optional<at::Tensor>& a1_scale,
-    const std::optional<at::Tensor>& a2_scale, bool is_vnni);
-
-at::Tensor int8_scaled_mm_with_quant(at::Tensor& mat1, at::Tensor& mat2,
-                                     at::Tensor& scales2,
-                                     const std::optional<at::Tensor>& bias,
-                                     at::ScalarType out_dtype, bool is_vnni);
-
 TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   // vLLM custom ops
 
@@ -152,19 +131,16 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
 
   // Quantization
 #ifdef __AVX512F__
-  at::Tag stride_tag = at::Tag::needs_fixed_stride_order;
   // Compute int8 quantized tensor for given scaling factor.
   ops.def(
       "static_scaled_int8_quant(Tensor! out, Tensor input, Tensor scale,"
-      "Tensor? azp) -> ()",
-      {stride_tag});
+      "Tensor? azp) -> ()");
   ops.impl("static_scaled_int8_quant", torch::kCPU, &static_scaled_int8_quant);
 
   // Compute int8 quantized tensor and scaling factor
   ops.def(
       "dynamic_scaled_int8_quant(Tensor! out, Tensor input, Tensor! scale, "
-      "Tensor!? azp) -> ()",
-      {stride_tag});
+      "Tensor!? azp) -> ()");
   ops.impl("dynamic_scaled_int8_quant", torch::kCPU,
            &dynamic_scaled_int8_quant);
   // W8A8 GEMM, supporting symmetric per-tensor or per-row/column
@@ -172,8 +148,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.def(
       "cutlass_scaled_mm(Tensor! out, Tensor a,"
       "                  Tensor b, Tensor a_scales,"
-      "                  Tensor b_scales, Tensor? bias) -> ()",
-      {stride_tag});
+      "                  Tensor b_scales, Tensor? bias) -> ()");
   ops.impl("cutlass_scaled_mm", torch::kCPU, &int8_scaled_mm);
   // w8a8 GEMM, supporting asymmetric per-tensor or per-row/column
   // quantization.
@@ -181,8 +156,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       "cutlass_scaled_mm_azp(Tensor! out, Tensor a,"
       "                  Tensor b, Tensor a_scales,"
       "                  Tensor b_scales, Tensor azp_adj,"
-      "                  Tensor? azp, Tensor? bias) -> ()",
-      {stride_tag});
+      "                  Tensor? azp, Tensor? bias) -> ()");
   ops.impl("cutlass_scaled_mm_azp", torch::kCPU, &int8_scaled_mm_azp);
 #elif defined(__powerpc64__)
   // Compute int8 quantized tensor for given scaling factor.
@@ -235,28 +209,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.def("shm_recv_tensor_list(int handle, int src) -> Tensor[](a)",
           &shm_recv_tensor_list);
 #endif
-
-  // sgl-kernels
-#if defined(__AVX512BF16__) && defined(__AVX512F__) && defined(__AVX512VNNI__)
-  ops.def(
-      "weight_packed_linear(Tensor(a0!) mat1, Tensor(a1!) mat2, Tensor(a2!)? "
-      "bias, bool is_vnni) -> Tensor");
-  ops.impl("weight_packed_linear", torch::kCPU, &weight_packed_linear);
-  ops.def("convert_weight_packed(Tensor! weight) -> Tensor");
-  ops.impl("convert_weight_packed", torch::kCPU, &convert_weight_packed);
-  ops.def(
-      "fused_experts_cpu(Tensor! hidden_states, Tensor w1, Tensor w2, Tensor "
-      "topk_weights, Tensor topk_ids, bool inplace, bool use_int8_w8a8, bool "
-      "use_fp8_w8a16, Tensor? w1_scale, Tensor? w2_scale, SymInt[]? "
-      "block_size, Tensor? a1_scale, Tensor? a2_scale, bool is_vnni) -> "
-      "Tensor");
-  ops.impl("fused_experts_cpu", torch::kCPU, &fused_experts_cpu);
-  ops.def(
-      "int8_scaled_mm_with_quant(Tensor mat1, Tensor mat2, Tensor scales2, "
-      "Tensor? bias, ScalarType out_dtype, bool is_vnni) -> Tensor");
-  ops.impl("int8_scaled_mm_with_quant", torch::kCPU,
-           &int8_scaled_mm_with_quant);
-#endif
 }
 
 TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cache_ops), cache_ops) {

csrc/custom_quickreduce.cu

@@ -1,114 +0,0 @@
-#include <ATen/cuda/Exceptions.h>
-#include <c10/cuda/CUDAGuard.h>
-#include <c10/cuda/CUDAStream.h>
-#include <torch/all.h>
-
-#ifdef USE_ROCM
-
-  #include "quickreduce/quick_reduce.h"
-
-quickreduce::fptr_t init_custom_qr(int64_t rank, int64_t world_size,
-                                   std::optional<int64_t> qr_max_size) {
-  if (world_size > 8)
-    throw std::invalid_argument("world size > 8 is not supported");
-  if (world_size == 6)
-    throw std::invalid_argument("world size == 6 is not supported");
-  if (world_size % 2 != 0)
-    throw std::invalid_argument("Odd num gpus is not supported for now");
-  if (rank < 0 || rank >= world_size)
-    throw std::invalid_argument("invalid rank passed in");
-  quickreduce::DeviceComms* fptr = new quickreduce::DeviceComms();
-  fptr->init(world_size, rank, qr_max_size);
-  return (quickreduce::fptr_t)fptr;
-}
-
-void qr_destroy(quickreduce::fptr_t _fa) {
-  if (_fa) {
-    auto fa = reinterpret_cast<quickreduce::DeviceComms*>(_fa);
-    fa->destroy();
-    delete fa;
-  }
-}
-
-torch::Tensor qr_get_handle(quickreduce::fptr_t _fa) {
-  auto fa = reinterpret_cast<quickreduce::DeviceComms*>(_fa);
-  hipIpcMemHandle_t handle = fa->get_handle();
-  auto options =
-      torch::TensorOptions().dtype(torch::kUInt8).device(torch::kCPU);
-  auto data_handle =
-      torch::empty({static_cast<int64_t>(sizeof(hipIpcMemHandle_t))}, options);
-  std::memcpy(data_handle.data_ptr(), &handle, sizeof(hipIpcMemHandle_t));
-  return data_handle;
-}
-
-void qr_open_handles(quickreduce::fptr_t _fa,
-                     const std::vector<torch::Tensor>& handles) {
-  auto fa = reinterpret_cast<quickreduce::DeviceComms*>(_fa);
-  std::vector<hipIpcMemHandle_t> ipc_handles;
-  ipc_handles.reserve(handles.size());
-  for (auto& handle : handles) {
-    // Ensure the tensor is on the same device as the current device.
-    hipIpcMemHandle_t ipc_handle;
-    std::memcpy(&ipc_handle, handle.data_ptr(), sizeof(hipIpcMemHandle_t));
-    ipc_handles.push_back(ipc_handle);
-  }
-  fa->open_ipc_handles(ipc_handles);
-}
-
-void qr_all_reduce(quickreduce::fptr_t _fa, torch::Tensor& inp,
-                   torch::Tensor& out, int64_t quant_level, bool cast_bf2half) {
-  auto fa = reinterpret_cast<quickreduce::DeviceComms*>(_fa);
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(inp));
-  auto stream = at::cuda::getCurrentHIPStreamMasqueradingAsCUDA();
-
-  TORCH_CHECK_EQ(inp.scalar_type(), out.scalar_type());
-  TORCH_CHECK_EQ(inp.numel(), out.numel());
-  TORCH_CHECK_LE(out.numel(), fa->kMaxProblemSize);
-  if (out.scalar_type() == at::ScalarType::Half) {
-    fa->allreduce<half, false>(reinterpret_cast<half*>(inp.data_ptr()),
-                               reinterpret_cast<half*>(out.data_ptr()),
-                               out.numel(), quant_level, stream);
-  } else if (out.scalar_type() == at::ScalarType::BFloat16) {
-    if (cast_bf2half) {
-      fa->allreduce<half, true>(reinterpret_cast<half*>(inp.data_ptr()),
-                                reinterpret_cast<half*>(out.data_ptr()),
-                                out.numel(), quant_level, stream);
-    } else {
-      fa->allreduce<quickreduce::nv_bfloat16, false>(
-          reinterpret_cast<quickreduce::nv_bfloat16*>(inp.data_ptr()),
-          reinterpret_cast<quickreduce::nv_bfloat16*>(out.data_ptr()),
-          out.numel(), quant_level, stream);
-    }
-  } else {
-    throw std::runtime_error(
-        "quick allreduce only supports float16 and bfloat16");
-  }
-}
-
-int64_t qr_max_size() {
-  // The default is 2GB (2,147,483,648 bytes)
-  return static_cast<int64_t>(std::numeric_limits<int32_t>::max()) + 1;
-}
-
-  #define INSTANTIATE_FOR_WORLDSIZE(T, Codec, cast_bf2half)       \
-    template struct quickreduce::AllReduceTwoshot<T, Codec<T, 2>, \
-                                                  cast_bf2half>;  \
-    template struct quickreduce::AllReduceTwoshot<T, Codec<T, 4>, \
-                                                  cast_bf2half>;  \
-    template struct quickreduce::AllReduceTwoshot<T, Codec<T, 8>, cast_bf2half>;
-
-INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecFP, false)
-INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ4, false)
-INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ6, false)
-INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ8, false)
-INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecFP, true)
-INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ4, true)
-INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ6, true)
-INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ8, true)
-
-INSTANTIATE_FOR_WORLDSIZE(half, quickreduce::CodecFP, false)
-INSTANTIATE_FOR_WORLDSIZE(half, quickreduce::CodecQ4, false)
-INSTANTIATE_FOR_WORLDSIZE(half, quickreduce::CodecQ6, false)
-INSTANTIATE_FOR_WORLDSIZE(half, quickreduce::CodecQ8, false)
-
-#endif  // USE_ROCM

csrc/mamba/causal_conv1d/causal_conv1d.cu

@@ -185,7 +185,9 @@ void causal_conv1d_fwd(const at::Tensor &x, const at::Tensor &weight,
         params.conv_states_ptr = nullptr;
     }
 
-    const at::cuda::OptionalCUDAGuard device_guard(device_of(x));
+    // Otherwise the kernel will be launched from cuda:0 device
+    // Cast to char to avoid compiler warning about narrowing
+    at::cuda::CUDAGuard device_guard{(char)x.get_device()};
     auto stream = at::cuda::getCurrentCUDAStream().stream();
     DISPATCH_WTYPE_ITYPE_FLOAT_AND_HALF_AND_BF16(x.scalar_type(), "causal_conv1d_fwd", [&] {
             causal_conv1d_fwd_cuda<input_t, weight_t>(params, stream);
@@ -276,7 +278,9 @@ void causal_conv1d_update(const at::Tensor &x,
         params.conv_state_indices_ptr = nullptr;
     }
 
-    const at::cuda::OptionalCUDAGuard device_guard(device_of(x));
+    // Otherwise the kernel will be launched from cuda:0 device
+    // Cast to char to avoid compiler warning about narrowing
+    at::cuda::CUDAGuard device_guard{(char)x.get_device()};
     auto stream = at::cuda::getCurrentCUDAStream().stream();
     DISPATCH_WTYPE_ITYPE_FLOAT_AND_HALF_AND_BF16(x.scalar_type(), "causal_conv1d_update", [&] {
             causal_conv1d_update_cuda<input_t, weight_t>(params, stream);

csrc/mamba/mamba_ssm/selective_scan_fwd.cu

@@ -647,7 +647,9 @@ void selective_scan_fwd(const torch::Tensor &u, const torch::Tensor &delta,
                        );
 
     
-    const at::cuda::OptionalCUDAGuard device_guard(device_of(u));
+    // Otherwise the kernel will be launched from cuda:0 device
+    // Cast to char to avoid compiler warning about narrowing
+    at::cuda::CUDAGuard device_guard{(char)u.get_device()};
     auto stream = at::cuda::getCurrentCUDAStream().stream();
     DISPATCH_WTYPE_ITYPE_FLOAT_AND_HALF_AND_BF16(u.scalar_type(), "selective_scan_fwd", [&] {
         selective_scan_fwd_cuda<input_t, weight_t>(params, stream);

csrc/moe/marlin_moe_wna16/marlin_template.h

@@ -1255,6 +1255,8 @@ __global__ void Marlin(
     if constexpr (has_zp && !is_zp_float) {
       if (is_new_zp) {
         if constexpr (group_blocks == -1) is_first_matmul_in_slice = false;
+        FragB frag_zp_0;
+        FragB frag_zp_1;
         int zp_quant_0, zp_quant_1;
 
         if constexpr (w_type.size_bits() == 4) {

csrc/moe/moe_align_sum_kernels.cu

@@ -239,7 +239,7 @@ void moe_sum(torch::Tensor& input,   // [num_tokens, topk, hidden_size]
              torch::Tensor& output)  // [num_tokens, hidden_size]
 {
   const int hidden_size = input.size(-1);
-  const auto num_tokens = output.numel() / hidden_size;
+  const int num_tokens = output.numel() / hidden_size;
   const int topk = input.size(1);
 
   dim3 grid(num_tokens);

csrc/moe/topk_softmax_kernels.cu

@@ -492,7 +492,7 @@ void topk_softmax(
     torch::Tensor& gating_output)               // [num_tokens, num_experts]
 {
     const int num_experts = gating_output.size(-1);
-    const auto num_tokens = gating_output.numel() / num_experts;
+    const int num_tokens = gating_output.numel() / num_experts;
     const int topk = topk_weights.size(-1);
 
     const bool is_pow_2 = (num_experts != 0) && ((num_experts & (num_experts - 1)) == 0);

csrc/ops.h

@@ -360,14 +360,3 @@ std::tuple<int64_t, torch::Tensor> allocate_shared_buffer_and_handle(
     int64_t size);
 int64_t open_mem_handle(torch::Tensor& mem_handle);
 void free_shared_buffer(int64_t buffer);
-
-#ifdef USE_ROCM
-fptr_t init_custom_qr(int64_t rank, int64_t world_size,
-                      std::optional<int64_t> qr_max_size = std::nullopt);
-void qr_destroy(fptr_t _fa);
-torch::Tensor qr_get_handle(fptr_t _fa);
-void qr_open_handles(fptr_t _fa, const std::vector<torch::Tensor>& handles);
-void qr_all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out,
-                   int64_t quant_level, bool cast_bf2half = false);
-int64_t qr_max_size();
-#endif

csrc/quantization/cutlass_w8a8/c3x/scaled_mm.cuh

@@ -144,65 +144,4 @@ struct cutlass_3x_gemm_sm100 {
       Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue, void>;
 };
 
-template <typename ElementAB_, typename ElementD_,
-          template <typename, typename, typename> typename Epilogue_,
-          typename TileShape, typename ClusterShape, typename KernelSchedule,
-          typename EpilogueSchedule>
-struct cutlass_3x_gemm_sm120 {
-  using ElementAB = ElementAB_;
-  using LayoutA = cutlass::layout::RowMajor;
-  static constexpr int AlignmentA =
-      128 / cutlass::sizeof_bits<ElementAB>::value;
-
-  using LayoutB = cutlass::layout::ColumnMajor;
-  static constexpr int AlignmentB =
-      128 / cutlass::sizeof_bits<ElementAB>::value;
-
-  using ElementC = void;
-  using LayoutC = cutlass::layout::RowMajor;
-  static constexpr int AlignmentC =
-      128 / cutlass::sizeof_bits<ElementD_>::value;
-
-  using ElementD = ElementD_;
-  using LayoutD = cutlass::layout::RowMajor;
-  static constexpr int AlignmentD = AlignmentC;
-
-  using ElementAcc =
-      typename std::conditional<std::is_same_v<ElementAB, int8_t>, int32_t,
-                                float>::type;
-  using Epilogue = Epilogue_<ElementAcc, ElementD, TileShape>;
-
-  // MMA type
-  using ElementAccumulator = float;
-
-  // Epilogue types
-  using ElementBias = cutlass::half_t;
-  using ElementCompute = float;
-  using ElementAux = ElementD;
-  using LayoutAux = LayoutD;
-  using ElementAmax = float;
-
-  using EVTCompute = typename Epilogue::EVTCompute;
-
-  using CollectiveEpilogue =
-      typename cutlass::epilogue::collective::CollectiveBuilder<
-          cutlass::arch::Sm120, cutlass::arch::OpClassTensorOp, TileShape,
-          ClusterShape, cutlass::epilogue::collective::EpilogueTileAuto,
-          ElementAccumulator, ElementCompute, ElementC, LayoutC, AlignmentC,
-          ElementD, LayoutD, AlignmentD, EpilogueSchedule,
-          EVTCompute>::CollectiveOp;
-
-  using CollectiveMainloop =
-      typename cutlass::gemm::collective::CollectiveBuilder<
-          cutlass::arch::Sm120, cutlass::arch::OpClassTensorOp, ElementAB,
-          LayoutA, AlignmentA, ElementAB, LayoutB, AlignmentB,
-          ElementAccumulator, TileShape, ClusterShape,
-          cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
-              sizeof(typename CollectiveEpilogue::SharedStorage))>,
-          KernelSchedule>::CollectiveOp;
-
-  using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
-      Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue, void>;
-};
-
 }  // namespace vllm

csrc/quantization/cutlass_w8a8/c3x/scaled_mm_kernels.hpp

@@ -36,12 +36,6 @@ void cutlass_scaled_mm_sm100_fp8(torch::Tensor& out, torch::Tensor const& a,
                                  torch::Tensor const& b_scales,
                                  std::optional<torch::Tensor> const& bias);
 
-void cutlass_scaled_mm_sm120_fp8(torch::Tensor& out, torch::Tensor const& a,
-                                 torch::Tensor const& b,
-                                 torch::Tensor const& a_scales,
-                                 torch::Tensor const& b_scales,
-                                 std::optional<torch::Tensor> const& bias);
-
 void cutlass_scaled_mm_blockwise_sm100_fp8(torch::Tensor& out,
                                            torch::Tensor const& a,
                                            torch::Tensor const& b,

csrc/quantization/cutlass_w8a8/c3x/scaled_mm_sm100_fp8_dispatch.cuh

@@ -29,12 +29,26 @@ struct sm100_fp8_config_default {
 template <typename InType, typename OutType,
           template <typename, typename, typename> typename Epilogue>
 struct sm100_fp8_config_M256 {
-  // M in (64, 256]
+  // M in (128, 256]
   static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
   using KernelSchedule = cutlass::gemm::collective::KernelScheduleAuto;
   using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto;
   using TileShape = Shape<_128, _128, _128>;
-  using ClusterShape = Shape<_2, _1, _1>;
+  using ClusterShape = Shape<_2, _2, _1>;
+  using Cutlass3xGemm =
+      cutlass_3x_gemm_sm100<InType, OutType, Epilogue, TileShape, ClusterShape,
+                            KernelSchedule, EpilogueSchedule>;
+};
+
+template <typename InType, typename OutType,
+          template <typename, typename, typename> typename Epilogue>
+struct sm100_fp8_config_M128 {
+  // M in (64, 128]
+  static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
+  using KernelSchedule = cutlass::gemm::collective::KernelScheduleAuto;
+  using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto;
+  using TileShape = Shape<_128, _128, _256>;
+  using ClusterShape = Shape<_2, _4, _1>;
   using Cutlass3xGemm =
       cutlass_3x_gemm_sm100<InType, OutType, Epilogue, TileShape, ClusterShape,
                             KernelSchedule, EpilogueSchedule>;
@@ -43,26 +57,12 @@ struct sm100_fp8_config_M256 {
 template <typename InType, typename OutType,
           template <typename, typename, typename> typename Epilogue>
 struct sm100_fp8_config_M64 {
-  // M in (16, 64]
+  // M in [1, 64]
   static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
   using KernelSchedule = cutlass::gemm::collective::KernelScheduleAuto;
   using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto;
-  using TileShape = Shape<_64, _64, _128>;
-  using ClusterShape = Shape<_1, _1, _1>;
-  using Cutlass3xGemm =
-      cutlass_3x_gemm_sm100<InType, OutType, Epilogue, TileShape, ClusterShape,
-                            KernelSchedule, EpilogueSchedule>;
-};
-
-template <typename InType, typename OutType,
-          template <typename, typename, typename> typename Epilogue>
-struct sm100_fp8_config_M16 {
-  // M in [1, 16]
-  static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
-  using KernelSchedule = cutlass::gemm::collective::KernelScheduleAuto;
-  using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto;
-  using TileShape = Shape<_64, _64, _128>;
-  using ClusterShape = Shape<_1, _4, _1>;
+  using TileShape = Shape<_64, _64, _256>;
+  using ClusterShape = Shape<_1, _8, _1>;
   using Cutlass3xGemm =
       cutlass_3x_gemm_sm100<InType, OutType, Epilogue, TileShape, ClusterShape,
                             KernelSchedule, EpilogueSchedule>;
@@ -82,27 +82,27 @@ inline void cutlass_gemm_sm100_fp8_dispatch(torch::Tensor& out,
   using Cutlass3xGemmDefault =
       typename sm100_fp8_config_default<InType, OutType,
                                         Epilogue>::Cutlass3xGemm;
-  using Cutlass3xGemmM16 =
-      typename sm100_fp8_config_M16<InType, OutType, Epilogue>::Cutlass3xGemm;
   using Cutlass3xGemmM64 =
       typename sm100_fp8_config_M64<InType, OutType, Epilogue>::Cutlass3xGemm;
+  using Cutlass3xGemmM128 =
+      typename sm100_fp8_config_M128<InType, OutType, Epilogue>::Cutlass3xGemm;
   using Cutlass3xGemmM256 =
       typename sm100_fp8_config_M256<InType, OutType, Epilogue>::Cutlass3xGemm;
 
   uint32_t const m = a.size(0);
   uint32_t const mp2 =
-      std::max(static_cast<uint32_t>(16), next_pow_2(m));  // next power of 2
+      std::max(static_cast<uint32_t>(64), next_pow_2(m));  // next power of 2
 
-  if (mp2 <= 16) {
-    // m in [1, 16]
-    return cutlass_gemm_caller<Cutlass3xGemmM16>(
-        out, a, b, std::forward<EpilogueArgs>(args)...);
-  } else if (mp2 <= 64) {
-    // m in (16, 64]
+  if (mp2 <= 64) {
+    // m in [1, 64]
     return cutlass_gemm_caller<Cutlass3xGemmM64>(
         out, a, b, std::forward<EpilogueArgs>(args)...);
+  } else if (mp2 <= 128) {
+    // m in (64, 128]
+    return cutlass_gemm_caller<Cutlass3xGemmM128>(
+        out, a, b, std::forward<EpilogueArgs>(args)...);
   } else if (mp2 <= 256) {
-    // m in (64, 256]
+    // m in (128, 256]
     return cutlass_gemm_caller<Cutlass3xGemmM256>(
         out, a, b, std::forward<EpilogueArgs>(args)...);
   } else {

csrc/quantization/cutlass_w8a8/c3x/scaled_mm_sm120_fp8.cu

@@ -1,24 +0,0 @@
-#include "scaled_mm_kernels.hpp"
-#include "scaled_mm_sm120_fp8_dispatch.cuh"
-#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
-
-namespace vllm {
-
-void cutlass_scaled_mm_sm120_fp8(torch::Tensor& out, torch::Tensor const& a,
-                                 torch::Tensor const& b,
-                                 torch::Tensor const& a_scales,
-                                 torch::Tensor const& b_scales,
-                                 std::optional<torch::Tensor> const& bias) {
-  TORCH_CHECK(a_scales.is_contiguous() && b_scales.is_contiguous());
-  if (bias) {
-    TORCH_CHECK(bias->dtype() == out.dtype(),
-                "currently bias dtype must match output dtype ", out.dtype());
-    return cutlass_scaled_mm_sm120_fp8_epilogue<c3x::ScaledEpilogueBias>(
-        out, a, b, a_scales, b_scales, *bias);
-  } else {
-    return cutlass_scaled_mm_sm120_fp8_epilogue<c3x::ScaledEpilogue>(
-        out, a, b, a_scales, b_scales);
-  }
-}
-
-}  // namespace vllm

csrc/quantization/cutlass_w8a8/c3x/scaled_mm_sm120_fp8_dispatch.cuh

@@ -1,67 +0,0 @@
-#pragma once
-
-#include "scaled_mm.cuh"
-#include "cutlass_gemm_caller.cuh"
-
-/**
- * This file defines Gemm kernel configurations for SM120 (fp8) based on the
- * Gemm shape.
- */
-
-namespace vllm {
-
-using c3x::cutlass_gemm_caller;
-
-template <typename InType, typename OutType,
-          template <typename, typename, typename> typename Epilogue>
-struct sm120_fp8_config_default {
-  static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
-  using KernelSchedule = cutlass::gemm::collective::KernelScheduleAuto;
-  using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto;
-  using TileShape = Shape<_128, _128, _128>;
-  using ClusterShape = Shape<_1, _1, _1>;  // Only work with Shape<_1, _1, _1>
-  using Cutlass3xGemm =
-      cutlass_3x_gemm_sm120<InType, OutType, Epilogue, TileShape, ClusterShape,
-                            KernelSchedule, EpilogueSchedule>;
-};
-
-template <typename InType, typename OutType,
-          template <typename, typename, typename> typename Epilogue,
-          typename... EpilogueArgs>
-inline void cutlass_gemm_sm120_fp8_dispatch(torch::Tensor& out,
-                                            torch::Tensor const& a,
-                                            torch::Tensor const& b,
-                                            EpilogueArgs&&... args) {
-  static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
-  TORCH_CHECK(a.dtype() == torch::kFloat8_e4m3fn);
-  TORCH_CHECK(b.dtype() == torch::kFloat8_e4m3fn);
-
-  using Cutlass3xGemmDefault =
-      typename sm120_fp8_config_default<InType, OutType,
-                                        Epilogue>::Cutlass3xGemm;
-  return cutlass_gemm_caller<Cutlass3xGemmDefault>(
-      out, a, b, std::forward<EpilogueArgs>(args)...);
-}
-
-template <template <typename, typename, typename> typename Epilogue,
-          typename... EpilogueArgs>
-void cutlass_scaled_mm_sm120_fp8_epilogue(torch::Tensor& out,
-                                          torch::Tensor const& a,
-                                          torch::Tensor const& b,
-                                          EpilogueArgs&&... epilogue_args) {
-  TORCH_CHECK(a.dtype() == torch::kFloat8_e4m3fn);
-  TORCH_CHECK(b.dtype() == torch::kFloat8_e4m3fn);
-
-  if (out.dtype() == torch::kBFloat16) {
-    return cutlass_gemm_sm120_fp8_dispatch<cutlass::float_e4m3_t,
-                                           cutlass::bfloat16_t, Epilogue>(
-        out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
-  } else {
-    TORCH_CHECK(out.dtype() == torch::kFloat16);
-    return cutlass_gemm_sm120_fp8_dispatch<cutlass::float_e4m3_t,
-                                           cutlass::half_t, Epilogue>(
-        out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
-  }
-}
-
-}  // namespace vllm

csrc/quantization/cutlass_w8a8/scaled_mm_c3x_sm120.cu

@@ -1,34 +0,0 @@
-#include <cudaTypedefs.h>
-#include "c3x/scaled_mm_kernels.hpp"
-
-#include "cuda_utils.h"
-
-/*
-   This file defines quantized GEMM operations using the CUTLASS 3.x API, for
-   NVIDIA GPUs with sm120 (Blackwell Geforce).
-*/
-
-#if defined ENABLE_SCALED_MM_SM120 && ENABLE_SCALED_MM_SM120
-
-void cutlass_scaled_mm_sm120(torch::Tensor& c, torch::Tensor const& a,
-                             torch::Tensor const& b,
-                             torch::Tensor const& a_scales,
-                             torch::Tensor const& b_scales,
-                             std::optional<torch::Tensor> const& bias) {
-  TORCH_CHECK(a_scales.dtype() == torch::kFloat32);
-  TORCH_CHECK(b_scales.dtype() == torch::kFloat32);
-
-  int M = a.size(0), N = b.size(1), K = a.size(1);
-  TORCH_CHECK(
-      (a_scales.numel() == 1 || a_scales.numel() == a.size(0)) &&
-          (b_scales.numel() == 1 || b_scales.numel() == b.size(1)),
-      "Currently, block scaled fp8 gemm is not implemented for Blackwell");
-
-  // Standard per-tensor/per-token/per-channel scaling
-  TORCH_CHECK(a_scales.is_contiguous() && b_scales.is_contiguous());
-  TORCH_CHECK(a.dtype() == torch::kFloat8_e4m3fn,
-              "Currently, only fp8 gemm is implemented for Blackwell");
-  vllm::cutlass_scaled_mm_sm120_fp8(c, a, b, a_scales, b_scales, bias);
-}
-
-#endif

csrc/quantization/cutlass_w8a8/scaled_mm_entry.cu

@@ -41,14 +41,6 @@ void cutlass_moe_mm_sm90(
 
 #endif
 
-#if defined ENABLE_SCALED_MM_SM120 && ENABLE_SCALED_MM_SM120
-void cutlass_scaled_mm_sm120(torch::Tensor& c, torch::Tensor const& a,
-                             torch::Tensor const& b,
-                             torch::Tensor const& a_scales,
-                             torch::Tensor const& b_scales,
-                             std::optional<torch::Tensor> const& bias);
-#endif
-
 #if defined ENABLE_SCALED_MM_SM100 && ENABLE_SCALED_MM_SM100
 void cutlass_scaled_mm_sm100(torch::Tensor& c, torch::Tensor const& a,
                              torch::Tensor const& b,
@@ -176,15 +168,8 @@ void cutlass_scaled_mm(torch::Tensor& c, torch::Tensor const& a,
   at::cuda::OptionalCUDAGuard const device_guard(device_of(a));
   int32_t version_num = get_sm_version_num();
 
-#if defined ENABLE_SCALED_MM_SM120 && ENABLE_SCALED_MM_SM120
-  if (version_num >= 120) {
-    cutlass_scaled_mm_sm120(c, a, b, a_scales, b_scales, bias);
-    return;
-  }
-#endif
-
 #if defined ENABLE_SCALED_MM_SM100 && ENABLE_SCALED_MM_SM100
-  if (version_num >= 100 && version_num < 120) {
+  if (version_num >= 100) {
     cutlass_scaled_mm_sm100(c, a, b, a_scales, b_scales, bias);
     return;
   }
@@ -256,7 +241,7 @@ void get_cutlass_moe_mm_data(
   // mm to run it for.
   int32_t version_num = get_sm_version_num();
 #if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
-    (defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100)
+    (defined ENABLE_SCALED_MM_SM100 && ENABLE_SCALED_MM_SM90)
   get_cutlass_moe_mm_data_caller(topk_ids, expert_offsets, problem_sizes1,
                                  problem_sizes2, input_permutation,
                                  output_permutation, num_experts, n, k,
@@ -267,7 +252,7 @@ void get_cutlass_moe_mm_data(
       false,
       "No compiled get_cutlass_moe_mm_data: no cutlass_scaled_mm kernel for "
       "CUDA device capability: ",
-      version_num, ". Required capability: 90 or 100");
+      version_num, ". Required capability: 90");
 }
 
 void get_cutlass_pplx_moe_mm_data(torch::Tensor& expert_offsets,
@@ -280,8 +265,7 @@ void get_cutlass_pplx_moe_mm_data(torch::Tensor& expert_offsets,
   // This function currently gets compiled only if we have a valid cutlass moe
   // mm to run it for.
   int32_t version_num = get_sm_version_num();
-#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
-    (defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100)
+#if defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90
   get_cutlass_pplx_moe_mm_data_caller(expert_offsets, problem_sizes1,
                                       problem_sizes2, expert_num_tokens,
                                       num_local_experts, padded_m, n, k);
@@ -291,7 +275,7 @@ void get_cutlass_pplx_moe_mm_data(torch::Tensor& expert_offsets,
       false,
       "No compiled get_cutlass_pplx_moe_mm_data: no cutlass_scaled_mm kernel "
       "for CUDA device capability: ",
-      version_num, ". Required capability: 90 or 100");
+      version_num, ". Required capability: 90");
 }
 
 void cutlass_scaled_mm_azp(torch::Tensor& c, torch::Tensor const& a,

csrc/quantization/fp4/nvfp4_blockwise_moe_kernel.cu

@@ -335,10 +335,8 @@ void run_fp4_blockwise_scaled_group_mm(
   TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to run GEMM");
 }
 
-#if defined ENABLE_NVFP4 && ENABLE_NVFP4
 constexpr auto FLOAT4_E2M1X2 = at::ScalarType::Byte;
 constexpr auto SF_DTYPE = at::ScalarType::Float8_e4m3fn;
-#endif
 
 #define CHECK_TYPE(x, st, m) \
   TORCH_CHECK(x.scalar_type() == st, ": Inconsistency of Tensor type:", m)

csrc/quantization/fp4/nvfp4_experts_quant.cu

@@ -561,7 +561,7 @@ void scaled_fp4_experts_quant_sm100a(
   TORCH_CHECK(output_scale.size(1) * 4 == padded_k);
 
   auto in_dtype = input.dtype();
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
+  at::cuda::CUDAGuard device_guard{(char)input.get_device()};
   const cudaStream_t stream =
       at::cuda::getCurrentCUDAStream(input.get_device());
   if (in_dtype == at::ScalarType::Half) {
@@ -579,4 +579,4 @@ void scaled_fp4_experts_quant_sm100a(
   } else {
     TORCH_CHECK(false, "Expected input data type to be half or bfloat16");
   }
-}
+}

csrc/quantization/fp4/nvfp4_quant_kernels.cu

@@ -347,7 +347,7 @@ void scaled_fp4_quant_sm100a(torch::Tensor const& output,
   auto input_sf_ptr = static_cast<float const*>(input_sf.data_ptr());
   auto sf_out = static_cast<int32_t*>(output_sf.data_ptr());
   auto output_ptr = static_cast<int64_t*>(output.data_ptr());
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
+  at::cuda::CUDAGuard device_guard{(char)input.get_device()};
   auto stream = at::cuda::getCurrentCUDAStream(input.get_device());
 
   // We don't support e8m0 scales at this moment.

csrc/quantization/fp4/nvfp4_scaled_mm_kernels.cu

@@ -267,7 +267,7 @@ void cutlass_scaled_fp4_mm_sm100a(torch::Tensor& D, torch::Tensor const& A,
               B_sf.sizes()[1], ")");
 
   auto out_dtype = D.dtype();
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(A));
+  at::cuda::CUDAGuard device_guard{(char)A.get_device()};
   const cudaStream_t stream = at::cuda::getCurrentCUDAStream(A.get_device());
 
   if (out_dtype == at::ScalarType::Half) {

csrc/quantization/gptq_marlin/marlin_template.h

@@ -1113,6 +1113,8 @@ __global__ void Marlin(
     if constexpr (has_zp && !is_zp_float) {
       if (is_new_zp) {
         if constexpr (group_blocks == -1) is_first_matmul_in_slice = false;
+        FragB frag_zp_0;
+        FragB frag_zp_1;
         int zp_quant_0, zp_quant_1;
 
         if constexpr (w_type.size_bits() == 4) {

csrc/quickreduce/base.h

@@ -1,338 +0,0 @@
-#pragma once
-
-#include <cstdint>
-#include <hip/hip_runtime.h>
-#include <hip/hip_fp16.h>
-#include <hip/hip_bf16.h>
-
-#define __quickreduce_device_inline__ __device__ __forceinline__
-#define __quickreduce_launch_bounds_two_shot__ __launch_bounds__(256, 4)
-#define __quickreduce_launch_bounds_one_shot__ __launch_bounds__(512, 4)
-
-namespace quickreduce {
-
-typedef __hip_bfloat16 nv_bfloat16;
-typedef __hip_bfloat162 nv_bfloat162;
-
-using int32x2_t = __attribute__((__vector_size__(2 * sizeof(int)))) int;
-using int32x4_t = __attribute__((__vector_size__(4 * sizeof(int)))) int;
-
-// Setup acquire-release semantics for vector memory reads (mubuf instruction)
-// as per architecture.
-#if defined(__gfx942__)
-// CDNA3: Scope bits sc0, sc1
-  #define MUBUF_ACQUIRE 16
-  #define MUBUF_RELEASE 16
-#elif (defined(__gfx908__) || defined(__gfx90a__))
-// CDNA1 and CDNA2 - glc bit
-  #define MUBUF_ACQUIRE 1
-  #define MUBUF_RELEASE 0
-#endif
-
-static constexpr int kNegOne = 0xBC00BC00;  // {-1, -1}, fp16x2_t
-
-// Number of atoms (4xf16x2_t) processed by a single thread
-static constexpr int kAtoms = 8;
-
-// We use a workgroup of 256 threads
-static constexpr int kBlockSize = 256;
-static constexpr int kAtomStride = kBlockSize;
-
-// Size and atom stride of source/destination data that the block will
-// process.
-// Workgroup scope = Tile = (256 threads x 8 atoms x 16B)
-static constexpr int kTileSize = kBlockSize * kAtoms * sizeof(int32x4_t);
-
-// Max number of blocks. 304 CUs on MI300
-static constexpr int kMaxNumBlocks = 304 * 4;
-
-// Standard CDNA wavefront size.
-static constexpr int kWavefront = 64;
-
-// 256 thread, 4 wavefronts.
-static dim3 constexpr kBlockTwoShot = {kWavefront, kBlockSize / kWavefront, 1};
-
-// Number of threads in a group for quantization
-// It corresponds to 32 F16 elements in quantization block
-static constexpr int kThreadGroupSize = 8;
-
-// Methods
-__quickreduce_device_inline__ __host__ unsigned long divceil(unsigned long x,
-                                                             unsigned long y) {
-  return ((x + y - 1) / y);
-}
-
-union BufferResource {
-  __quickreduce_device_inline__ constexpr BufferResource()
-      : config(0x00020000U) {}
-
-  __quickreduce_device_inline__ constexpr BufferResource(void* buffer_address,
-                                                         uint32_t buffer_size)
-      : address(buffer_address), range(buffer_size), config(0x00020000U) {}
-
-  int32x4_t descriptor;
-  struct {
-    void* address;  // 8B, out of which first 48b is address, and 16b is stride
-    // (unused)
-    uint32_t range;   // Byte range for the buffer resource
-    uint32_t config;  // Constant, DFMT=32b
-  };
-};
-
-__quickreduce_device_inline__ static int32x4_t buffer_load_dwordx4(
-    int32x4_t srsrc, int32_t voffset, int32_t soffset,
-    int32_t aux) __asm("llvm.amdgcn.raw.buffer.load.v4i32");
-
-__quickreduce_device_inline__ static void buffer_store_dwordx4(
-    int32x4_t data, int32x4_t srsrc, int32_t voffset, int32_t soffset,
-    int32_t aux) __asm("llvm.amdgcn.raw.buffer.store.v4i32");
-
-__quickreduce_device_inline__ static void set_fp16_ovfl(bool const value) {
-#if defined(__gfx942__)
-  if (value) {
-    asm volatile("s_setreg_imm32_b32 0xdc1, 1;" ::);
-  } else {
-    asm volatile("s_setreg_imm32_b32 0xdc1, 0;" ::);
-  }
-#endif
-}
-union bf162_int_union {
-  int i;
-  nv_bfloat162 bf2;
-};
-
-template <typename T>
-__quickreduce_device_inline__ void packed_assign_add(int32x4_t* A,
-                                                     int32x4_t* B);
-
-template <>
-__quickreduce_device_inline__ void packed_assign_add<half>(int32x4_t* A,
-                                                           int32x4_t* B) {
-  int32x4_t& tR_fragment = A[0];
-  int32x4_t& tA_fragment = B[0];
-
-  asm volatile("v_pk_add_f16 %0, %1, %2"
-               : "=v"(tR_fragment[0])
-               : "v"(tR_fragment[0]), "v"(tA_fragment[0]));
-  asm volatile("v_pk_add_f16 %0, %1, %2"
-               : "=v"(tR_fragment[1])
-               : "v"(tR_fragment[1]), "v"(tA_fragment[1]));
-  asm volatile("v_pk_add_f16 %0, %1, %2"
-               : "=v"(tR_fragment[2])
-               : "v"(tR_fragment[2]), "v"(tA_fragment[2]));
-  asm volatile("v_pk_add_f16 %0, %1, %2"
-               : "=v"(tR_fragment[3])
-               : "v"(tR_fragment[3]), "v"(tA_fragment[3]));
-}
-
-template <>
-__quickreduce_device_inline__ void packed_assign_add<nv_bfloat16>(
-    int32x4_t* A, int32x4_t* B) {
-  nv_bfloat162* tA = reinterpret_cast<nv_bfloat162*>(A);
-  nv_bfloat162* tB = reinterpret_cast<nv_bfloat162*>(B);
-#pragma unroll
-  for (int i = 0; i < 4; i++) {
-    tA[i] = __hadd2(tA[i], tB[i]);
-  }
-}
-
-template <typename T>
-__quickreduce_device_inline__ int packed_max(int a, int b);
-
-template <>
-__quickreduce_device_inline__ int packed_max<half>(int a, int b) {
-  int result;
-  asm volatile("v_pk_max_f16 %0, %1, %2" : "=v"(result) : "v"(a), "v"(b));
-  return result;
-}
-
-template <>
-__quickreduce_device_inline__ int packed_max<nv_bfloat16>(int a, int b) {
-  bf162_int_union A, B, R;
-  A.i = a;
-  B.i = b;
-  R.bf2 = __hmax2(A.bf2, B.bf2);
-  return R.i;
-}
-
-template <typename T>
-__quickreduce_device_inline__ int packed_min(int a, int b);
-
-template <>
-__quickreduce_device_inline__ int packed_min<half>(int a, int b) {
-  int result;
-  asm volatile("v_pk_min_f16 %0, %1, %2" : "=v"(result) : "v"(a), "v"(b));
-  return result;
-}
-
-template <>
-__quickreduce_device_inline__ int packed_min<nv_bfloat16>(int a, int b) {
-  bf162_int_union A, B, R;
-  A.i = a;
-  B.i = b;
-  R.bf2 = __hmin2(A.bf2, B.bf2);
-  return R.i;
-}
-
-template <typename T>
-__quickreduce_device_inline__ int packed_abs_max(int a, int b);
-
-template <>
-__quickreduce_device_inline__ int packed_abs_max<half>(int a, int b) {
-  half2 wmaxh2 = __builtin_bit_cast(half2, a);
-  half2 wminh2 = __builtin_bit_cast(half2, b);
-  half2 wblockmaxh2;
-
-  wblockmaxh2.x =
-      __hgt(__habs(wmaxh2.x), __habs(wminh2.x)) ? wmaxh2.x : wminh2.x;
-  wblockmaxh2.y =
-      __hgt(__habs(wmaxh2.y), __habs(wminh2.y)) ? wmaxh2.y : wminh2.y;
-  return __builtin_bit_cast(int, wblockmaxh2);
-}
-
-template <>
-__quickreduce_device_inline__ int packed_abs_max<nv_bfloat16>(int a, int b) {
-  bf162_int_union A, B, R;
-  A.i = a;
-  B.i = b;
-  R.bf2.x = __hgt(__habs(A.bf2.x), __habs(B.bf2.x)) ? A.bf2.x : B.bf2.x;
-  R.bf2.y = __hgt(__habs(A.bf2.y), __habs(B.bf2.y)) ? A.bf2.y : B.bf2.y;
-  return R.i;
-}
-
-template <typename T>
-__quickreduce_device_inline__ int packed_add(int a, int b);
-
-template <>
-__quickreduce_device_inline__ int packed_add<half>(int a, int b) {
-  int result;
-  asm volatile("v_pk_add_f16 %0, %1, %2" : "=v"(result) : "v"(a), "v"(b));
-  return result;
-}
-
-template <>
-__quickreduce_device_inline__ int packed_add<nv_bfloat16>(int a, int b) {
-  bf162_int_union A, B, R;
-  A.i = a;
-  B.i = b;
-  R.bf2 = __hadd2(A.bf2, B.bf2);
-  return R.i;
-}
-
-template <>
-__quickreduce_device_inline__ int packed_add<int16_t>(int a, int b) {
-  int result;
-  asm volatile("v_pk_add_i16 %0, %1, %2" : "=v"(result) : "v"(a), "v"(b));
-  return result;
-}
-
-template <typename T>
-__quickreduce_device_inline__ int packed_sub(int a, int b);
-
-template <>
-__quickreduce_device_inline__ int packed_sub<half>(int a, int b) {
-  int result;
-
-  // MI300 lacks packed fp16 sub instruction. So we do -1 * min + max
-  asm volatile("v_pk_fma_f16 %0, %1, %2 %3"
-               : "=v"(result)
-               : "v"(kNegOne), "v"(b), "v"(a));
-  return result;
-}
-
-template <>
-__quickreduce_device_inline__ int packed_sub<nv_bfloat16>(int a, int b) {
-  bf162_int_union A, B, R;
-  A.i = a;
-  B.i = b;
-  R.bf2 = __hsub2(A.bf2, B.bf2);
-  return R.i;
-}
-
-template <typename T>
-__quickreduce_device_inline__ int packed_mul(int a, int b);
-
-template <>
-__quickreduce_device_inline__ int packed_mul<half>(int a, int b) {
-  int result;
-  asm volatile("v_pk_mul_f16 %0, %1, %2" : "=v"(result) : "v"(a), "v"(b));
-  return result;
-}
-
-template <>
-__quickreduce_device_inline__ int packed_mul<nv_bfloat16>(int a, int b) {
-  nv_bfloat162* tA = reinterpret_cast<nv_bfloat162*>(&a);
-  nv_bfloat162* tB = reinterpret_cast<nv_bfloat162*>(&b);
-  nv_bfloat162 tR = __hmul2(*tA, *tB);
-  return *(reinterpret_cast<int*>(&tR));
-}
-
-template <typename T>
-__quickreduce_device_inline__ int packed_rcp(int a);
-
-template <>
-__quickreduce_device_inline__ int packed_rcp<half>(int a) {
-  return __builtin_bit_cast(int, h2rcp(__builtin_bit_cast(half2, a)));
-}
-
-template <>
-__quickreduce_device_inline__ int packed_rcp<nv_bfloat16>(int a) {
-  bf162_int_union A, R;
-  A.i = a;
-  R.bf2 = h2rcp(A.bf2);
-  return R.i;
-}
-
-// changes dtype
-__quickreduce_device_inline__ float T2float_cast(half a) {
-  return __half2float(a);
-}
-
-__quickreduce_device_inline__ float T2float_cast(nv_bfloat16 a) {
-  return __bfloat162float(a);
-}
-
-template <typename T>
-__quickreduce_device_inline__ int group_abs_max(int32x4_t atom) {
-  const int group_leader = (threadIdx.x / kThreadGroupSize) * kThreadGroupSize;
-
-  int wmax, wmin, wblockmax;
-  int a, b;
-  a = packed_max<T>(atom[0], atom[1]);
-  b = packed_max<T>(atom[2], atom[3]);
-
-  wmax = packed_max<T>(a, b);
-
-  a = packed_min<T>(atom[0], atom[1]);
-  b = packed_min<T>(atom[2], atom[3]);
-
-  wmin = packed_min<T>(a, b);
-
-  // Reduce the max among a group of threads
-  // Note: This is basically 2 blocks of values setup as the
-  // upper/lower halves of the f16x2_t
-  for (int i = 1; i < kThreadGroupSize; i <<= 1) {
-    int x = __shfl_down(wmax, i);
-    wmax = packed_max<T>(wmax, x);
-
-    int y = __shfl_down(wmin, i);
-    wmin = packed_min<T>(wmin, y);
-  }
-  wblockmax = packed_abs_max<T>(wmax, wmin);
-  // Share with the cohort
-  wblockmax = __shfl(wblockmax, group_leader);
-  return wblockmax;
-}
-
-__quickreduce_device_inline__ void set_sync_flag(uint32_t* flag_ptr,
-                                                 uint32_t flag) {
-  __atomic_store_n(flag_ptr, flag, __ATOMIC_RELEASE);
-}
-
-__quickreduce_device_inline__ void wait_sync_flag(uint32_t* flag_ptr,
-                                                  uint32_t flag) {
-  while (__atomic_load_n(flag_ptr, __ATOMIC_RELAXED) != flag) {
-  }
-}
-
-}  // namespace quickreduce

csrc/quickreduce/quick_reduce.h

@@ -1,196 +0,0 @@
-#pragma once
-
-#include <vector>
-#include <hip/hip_runtime.h>
-#include "quick_reduce_impl.cuh"
-
-#define HIP_CHECK(err)                                                     \
-  do {                                                                     \
-    hipError_t err_ = (err);                                               \
-    if (err_ != hipSuccess) {                                              \
-      std::printf("HIP error %d at %s:%d. %s\n", err_, __FILE__, __LINE__, \
-                  hipGetErrorString(err_));                                \
-      throw std::runtime_error("HIP error");                               \
-    }                                                                      \
-  } while (0)
-
-namespace quickreduce {
-using fptr_t = int64_t;
-static_assert(sizeof(void*) == sizeof(fptr_t));
-
-template <typename AllReduceKernel, typename T>
-__global__ __quickreduce_launch_bounds_two_shot__ static void
-allreduce_prototype_twoshot(T const* A, T* B, uint32_t N, uint32_t num_blocks,
-                            int rank, uint8_t** dbuffer_list,
-                            uint32_t data_offset, uint32_t flag_color) {
-  int block = blockIdx.x;
-  int grid = gridDim.x;
-
-  while (block < num_blocks) {
-    AllReduceKernel::run(A, B, N, block, rank, dbuffer_list, data_offset,
-                         flag_color);
-    block += grid;
-    flag_color++;
-  }
-}
-
-#define TWOSHOT_DISPATCH(__codec)                                           \
-  if (world_size == 2) {                                                    \
-    using LineCodec = __codec<T, 2>;                                        \
-    using AllReduceKernel = AllReduceTwoshot<T, LineCodec, cast_bf2half>;   \
-    hipLaunchKernelGGL((allreduce_prototype_twoshot<AllReduceKernel, T>),   \
-                       dim3(grid), dim3(kBlockTwoShot), 0, stream, A, B, N, \
-                       num_blocks, rank, dbuffer_list, data_offset,         \
-                       flag_color);                                         \
-  } else if (world_size == 4) {                                             \
-    using LineCodec = __codec<T, 4>;                                        \
-    using AllReduceKernel = AllReduceTwoshot<T, LineCodec, cast_bf2half>;   \
-    hipLaunchKernelGGL((allreduce_prototype_twoshot<AllReduceKernel, T>),   \
-                       dim3(grid), dim3(kBlockTwoShot), 0, stream, A, B, N, \
-                       num_blocks, rank, dbuffer_list, data_offset,         \
-                       flag_color);                                         \
-  } else if (world_size == 8) {                                             \
-    using LineCodec = __codec<T, 8>;                                        \
-    using AllReduceKernel = AllReduceTwoshot<T, LineCodec, cast_bf2half>;   \
-    hipLaunchKernelGGL((allreduce_prototype_twoshot<AllReduceKernel, T>),   \
-                       dim3(grid), dim3(kBlockTwoShot), 0, stream, A, B, N, \
-                       num_blocks, rank, dbuffer_list, data_offset,         \
-                       flag_color);                                         \
-  }
-
-enum QuickReduceQuantLevel {
-  F16 = 0,
-  INT8 = 1,
-  INT6 = 2,
-  INT4 = 3,
-};
-
-struct DeviceComms {
-  // Max problem size is 2GB (in bytes) or half of uint32_t max value.
-  int64_t kMaxProblemSize =
-      static_cast<int64_t>(std::numeric_limits<int32_t>::max()) + 1;
-
-  // Max TP-8
-  static int constexpr kMaxWorldSize = 8;
-
-  bool initialized = false;
-  uint32_t flag_color = 1;
-  int world_size;
-  int rank;
-
-  uint8_t* dbuffer;
-  uint8_t** dbuffer_list;
-  hipIpcMemHandle_t buffer_ipc_handle;
-  std::vector<hipIpcMemHandle_t> all_buffer_ipc_handles;
-  std::vector<uint8_t*> buffer_list;
-  uint32_t data_offset;
-
-  DeviceComms() : initialized(false), world_size(1), rank(0) {}
-  ~DeviceComms() { destroy(); }
-
-  void init(int world_size, int rank,
-            std::optional<int64_t> max_problem_size = std::nullopt) {
-    destroy();
-    this->world_size = world_size;
-    this->rank = rank;
-    if (max_problem_size.has_value() && max_problem_size.value() > 0) {
-      this->kMaxProblemSize = max_problem_size.value();
-    }
-    // Allocate buffer size for worst case: F16 2-stage buffer.
-    uint32_t flags_buffer_size =
-        2 * world_size * kMaxNumBlocks * sizeof(uint32_t);
-    static int64_t data_buffer_size = 2 * this->kMaxProblemSize;
-    int64_t total_buffer_size = flags_buffer_size + data_buffer_size;
-    data_offset = flags_buffer_size;
-    HIP_CHECK(hipExtMallocWithFlags((void**)&dbuffer, total_buffer_size,
-                                    hipDeviceMallocUncached));
-
-    // Clear the flags buffer.
-    HIP_CHECK(hipMemset(dbuffer, 0, flags_buffer_size));
-
-    // Device-side list of IPC buffers.
-    buffer_list.resize(world_size);
-    HIP_CHECK(hipMalloc(&dbuffer_list, world_size * sizeof(uint8_t*)));
-
-    // Create IPC handles for rank's communication buffer.
-    all_buffer_ipc_handles.resize(world_size);
-    HIP_CHECK(hipIpcGetMemHandle(&buffer_ipc_handle, dbuffer));
-
-    initialized = true;
-  }
-  int get_world_size() { return world_size; }
-  int get_rank() { return rank; }
-  bool status() { return initialized; }
-  hipIpcMemHandle_t const get_handle() { return buffer_ipc_handle; }
-
-  void destroy() {
-    if (initialized) {
-      for (int i = 0; i < world_size; i++) {
-        if (i != rank) {
-          HIP_CHECK(hipIpcCloseMemHandle(dbuffer_list[i]));
-        }
-      }
-
-      HIP_CHECK(hipFree(dbuffer));
-      HIP_CHECK(hipFree(dbuffer_list));
-
-      initialized = false;
-    }
-  }
-
-  void open_ipc_handles(std::vector<hipIpcMemHandle_t> const& ipc_handles) {
-    assert(ipc_handles.size() == all_buffer_ipc_handles.size());
-    for (int i = 0; i < world_size; i++) {
-      all_buffer_ipc_handles[i] = ipc_handles[i];
-    }
-
-    // Open device memory access to the IPC communication buffers.
-    // Note: For our own rank, we do not need to open a handle.
-    for (int i = 0; i < world_size; i++) {
-      if (i != rank) {
-        HIP_CHECK(hipIpcOpenMemHandle((void**)&buffer_list[i],
-                                      all_buffer_ipc_handles[i],
-                                      hipIpcMemLazyEnablePeerAccess));
-      } else {
-        buffer_list[i] = dbuffer;
-      }
-    }
-
-    HIP_CHECK(hipMemcpy(dbuffer_list, buffer_list.data(),
-                        world_size * sizeof(uint8_t*), hipMemcpyHostToDevice));
-  }
-
-  template <typename T, bool cast_bf2half>
-  void allreduce(T const* A, T* B, uint32_t N, int quant_level,
-                 hipStream_t stream) {
-    if (world_size != 2 && world_size != 4 && world_size != 8) {
-      throw std::runtime_error("All Reduce not supported for world_size = " +
-                               std::to_string(world_size));
-    }
-
-    // Configuration.
-    uint32_t msg_size = N * sizeof(T);
-    uint32_t num_blocks = divceil(msg_size, kTileSize);
-    uint32_t grid = min(kMaxNumBlocks, num_blocks);
-    auto quant_level_ = static_cast<QuickReduceQuantLevel>(quant_level);
-    switch (quant_level_) {
-      case QuickReduceQuantLevel::INT8:
-        TWOSHOT_DISPATCH(CodecQ8)
-        break;
-      case QuickReduceQuantLevel::INT6:
-        TWOSHOT_DISPATCH(CodecQ6)
-        break;
-      case QuickReduceQuantLevel::INT4:
-        TWOSHOT_DISPATCH(CodecQ4)
-        break;
-      default:
-        TWOSHOT_DISPATCH(CodecFP)
-        break;
-    }
-    HIP_CHECK(cudaGetLastError());
-    // Rotate the flag color.
-    flag_color += divceil(N, grid);
-  }
-};
-
-}  // namespace quickreduce

csrc/quickreduce/quick_reduce_impl.cuh

@@ -1,698 +0,0 @@
-#pragma once
-
-#include <hip/hip_runtime.h>
-#include "base.h"
-
-namespace quickreduce {
-
-struct CodecBase {
-  const int thread;
-  const int rank;
-  const int group_leader;
-  __quickreduce_device_inline__ CodecBase(int thread, int rank)
-      : thread(thread),
-        rank(rank),
-        group_leader((threadIdx.x / kThreadGroupSize) * kThreadGroupSize) {
-    set_fp16_ovfl(true);
-  }
-};
-
-// Default full precision codec.
-template <typename T, int world_size>
-struct CodecFP : public CodecBase {
-  static constexpr int kWorldSize = world_size;
-  static constexpr int kRankAtoms = kAtoms / kWorldSize;
-
-  // Codec tile size process by this workgroup.
-  // Each thread processes atoms of f16x8_t (16B).
-  static constexpr int kRankTransmittedTileSize =
-      kBlockSize * kRankAtoms * sizeof(int32x4_t);
-  static_assert(kRankTransmittedTileSize % 16 == 0,
-                "kRankTransmittedTileSize must be 16B aligned.");
-
-  // Total tile size for the collective communication.
-  static constexpr int kTransmittedTileSize =
-      kRankTransmittedTileSize * kWorldSize;
-
-  __quickreduce_device_inline__ CodecFP(int thread, int rank)
-      : CodecBase(thread, rank) {}
-
-  __quickreduce_device_inline__ void send(int32x4_t* __restrict__ send_buffer,
-                                          const int32x4_t* __restrict__ data) {
-    for (int i = 0; i < kRankAtoms; i++) {
-      __builtin_nontemporal_store(data[i], send_buffer + thread);
-      send_buffer += kAtomStride;
-    }
-  }
-
-  __quickreduce_device_inline__ void recv(int32x4_t** __restrict__ recv_buffer,
-                                          int32x4_t* __restrict__ data) {
-    for (int i = 0; i < kRankAtoms; i++) {
-      data[i] = __builtin_nontemporal_load(*recv_buffer + thread);
-      *recv_buffer += kAtomStride;
-    }
-  }
-};
-
-// Int4 symmetric quantization codec.
-// We quantize the FP16 data to block-scaled Int4 in blocks of 4 *
-// kThreadGroupSize.
-template <typename T, int world_size>
-struct CodecQ4 : public CodecBase {
-  static constexpr int kWorldSize = world_size;
-
-  // Codec tile size process by this workgroup.
-  // Each threads processes a fragment of fp16x8_t (16B),
-  // into a int4x8_t (4B) and a fp16 scale shared among 32 values.
-  static constexpr int kRankAtoms = kAtoms / kWorldSize;
-  static constexpr int kRankTileStride = 1152;
-  static constexpr int kRankTileScaleOffset = 1024;
-  static constexpr int kRankTransmittedTileSize = kRankTileStride * kRankAtoms;
-  static_assert(kRankTransmittedTileSize % 16 == 0,
-                "kRankTransmittedTileSize must be 16B aligned.");
-
-  static constexpr int kRankBufferTileStride =
-      kRankTileStride / sizeof(int32x4_t);
-
-  // Total tile size for the collective communication.
-  static constexpr int kTransmittedTileSize =
-      kRankTransmittedTileSize * kWorldSize;
-
-  // Constants configuration
-
-  // {-1/8.0h, -1/8.0h}, f16x2_t
-  static constexpr int kScaleFactor =
-      std::is_same<T, half>::value ? 0xB000B000 : 0xBE00BE00;
-
-  // {1e-7, 1e-7}, f16x2_t
-  static constexpr int kScaleEpsilon =
-      std::is_same<T, half>::value ? 0x00010001 : 0x33D733D7;
-
-  // {-8, -8}, f16x2_t
-  static constexpr int kRangeMin =
-      std::is_same<T, half>::value ? 0xC800C800 : 0xC100C100;
-
-  // {+7, +7}, f16x2_t
-  static constexpr int kRangeMax =
-      std::is_same<T, half>::value ? 0x47004700 : 0x40E040E0;
-
-  // {+8, +8}, int16x2_t
-  static constexpr int kRangeBias = 0x00080008;
-
-  __quickreduce_device_inline__ CodecQ4(int thread, int rank)
-      : CodecBase(thread, rank) {}
-
-  __quickreduce_device_inline__ void send(int32x4_t* __restrict__ send_buffer,
-                                          const int32x4_t* __restrict__ data) {
-    for (int k = 0; k < kRankAtoms; k++) {
-      int32x4_t const atom = data[k];
-
-      // Compute the absolute maximum of the atom in the thread group
-      // In 2 blocks of values, upper/lower halves of the f16x2_t
-      int wblockmax = group_abs_max<T>(atom);
-
-      // Derive scales
-      int decoding_scale;
-      int encoding_scale;
-      decoding_scale = packed_mul<T>(wblockmax, kScaleFactor);
-      encoding_scale = packed_add<T>(decoding_scale, kScaleEpsilon);
-      encoding_scale = packed_rcp<T>(encoding_scale);
-
-      // Apply scales to get quantized values
-      int32x4_t w;
-      for (int i = 0; i < 4; i++) {
-        w[i] = packed_mul<T>(atom[i], encoding_scale);
-        w[i] = packed_max<T>(w[i], kRangeMin);
-        w[i] = packed_min<T>(w[i], kRangeMax);
-      }
-
-      // Convert from f16x2_t to uint16x2_t
-      int32x4_t q;
-      {
-        int16_t* qi = reinterpret_cast<int16_t*>(&q);
-        T* wh = reinterpret_cast<T*>(&w);
-        for (int i = 0; i < 8; i++) qi[i] = (int16_t)rintf(T2float_cast(wh[i]));
-
-        for (int i = 0; i < 4; i++) {
-          q[i] = packed_add<int16_t>(q[i], kRangeBias);
-        }
-      }
-
-      // Pack 8 x q4 into int32_t
-      int qw = q[0] | (q[1] << 4) | (q[2] << 8) | (q[3] << 12);
-
-      // Write quantized atom to send_buffer
-      // note: only the group leader stores the scale
-      uint8_t* atom_ptr =
-          reinterpret_cast<uint8_t*>(send_buffer + k * kRankBufferTileStride);
-      int32_t* qw_ptr = reinterpret_cast<int32_t*>(atom_ptr) + thread;
-      int* qs_ptr = reinterpret_cast<int*>(atom_ptr + kRankTileScaleOffset) +
-                    (thread / 8);
-
-      __builtin_nontemporal_store(qw, qw_ptr);
-      if (threadIdx.x == group_leader) {
-        __builtin_nontemporal_store(decoding_scale, qs_ptr);
-      }
-    }
-  }
-
-  __quickreduce_device_inline__ void recv(int32x4_t** __restrict__ recv_buffer,
-                                          int32x4_t* __restrict__ data) {
-    for (int k = 0; k < kRankAtoms; k++) {
-      // Directly read quantized atom from recv_buffer
-      uint8_t* atom_ptr = reinterpret_cast<uint8_t*>(*recv_buffer);
-      int32_t* qw_ptr = reinterpret_cast<int32_t*>(atom_ptr) + thread;
-      int* qs_ptr = reinterpret_cast<int*>(atom_ptr + kRankTileScaleOffset) +
-                    (thread / 8);
-
-      int32_t qw = __builtin_nontemporal_load(qw_ptr);
-      int qs = __builtin_nontemporal_load(qs_ptr);
-
-      *recv_buffer += kRankBufferTileStride;
-
-      // Unpack q4 into f16x8_t
-      int32x4_t w;
-      {
-        static constexpr uint kMask000F = 0x000F000F;
-        static constexpr uint kHalf2_1024 =
-            0x64006400;  // {1024.0, 1024.0}, fp16x2_t
-        static uint constexpr kHalf2_1032 =
-            0xE408E408;  // {-1032.0, -1032.0}, fp16x2_t
-
-        for (int i = 0; i < 4; i++) {
-          if constexpr (std::is_same<T, half>::value) {
-            int32_t q4 = ((qw >> (i * 4)) & kMask000F) | kHalf2_1024;
-            w[i] = packed_add<half>(q4, kHalf2_1032);
-          } else {
-            int32_t int16_2 = (qw >> (i * 4)) & kMask000F;
-            int16_t low = static_cast<int16_t>(int16_2 & 0xFFFF);
-            int16_t high = static_cast<int16_t>((int16_2 >> 16) & 0xFFFF);
-            nv_bfloat16 bf_low = __float2bfloat16(static_cast<float>(low));
-            nv_bfloat16 bf_high = __float2bfloat16(static_cast<float>(high));
-            nv_bfloat162 bf2 = __halves2bfloat162(bf_low, bf_high);
-            int32_t packed_bf16 = *reinterpret_cast<int32_t*>(&bf2);
-            w[i] = packed_add<nv_bfloat16>(packed_bf16, kRangeMin);
-          }
-        }
-      }
-
-      // Apply decoding scales
-      for (int i = 0; i < 4; i++) {
-        w[i] = packed_mul<T>(w[i], qs);
-      }
-
-      data[k] = w;
-    }
-  }
-};
-
-// Int6 symmetric quantization codec.
-// We quantize the FP16 data to block-scaled Int6 in blocks of 4 *
-// kThreadGroupSize.
-template <typename T, int world_size>
-struct CodecQ6 : public CodecBase {
-  static constexpr int kWorldSize = world_size;
-
-  // Codec tile size process by this workgroup.
-  // Each threads processes a fragment of fp16x8_t (16B),
-  // into a int6x8_t (4B + 2B) and a fp16 scale shared among 32 values.
-  static constexpr int kRankAtoms = kAtoms / kWorldSize;
-  static constexpr int kRankTileStride = 1664;
-  static constexpr int kRankTileQ2Offset = 1024;
-  static constexpr int kRankTileScaleOffset = 1536;
-  static constexpr int kRankTransmittedTileSize = kRankTileStride * kRankAtoms;
-  static_assert(kRankTransmittedTileSize % 16 == 0,
-                "kRankTransmittedTileSize must be 16B aligned.");
-
-  static constexpr int kRankBufferTileStride =
-      kRankTileStride / sizeof(int32x4_t);
-
-  // Total tile size for the collective communication.
-  static constexpr int kTransmittedTileSize =
-      kRankTransmittedTileSize * kWorldSize;
-
-  // Constants configuration
-
-  // {-1/32.0h, -1/32.0h}, fp16x2_t
-  static constexpr int kScaleFactor =
-      std::is_same<T, half>::value ? 0xA800A800 : 0xBD00BD00;
-
-  // {1e-7, 1e-7}, fp16x2_t
-  static constexpr int kScaleEpsilon =
-      std::is_same<T, half>::value ? 0x00010001 : 0x33D733D7;
-
-  // {-32, -32}, fp16x2_t
-  static constexpr int kRangeMin =
-      std::is_same<T, half>::value ? 0xD000D000 : 0xC200C200;
-
-  // {+31, +31}, fp16x2_t
-  static constexpr int kRangeMax =
-      std::is_same<T, half>::value ? 0x4FC04FC0 : 0x41F841F8;
-
-  // {+32, +32}, int16x2_t
-  static constexpr int kRangeBias = 0x00200020;
-
-  __quickreduce_device_inline__ CodecQ6(int thread, int rank)
-      : CodecBase(thread, rank) {}
-
-  __quickreduce_device_inline__ void send(int32x4_t* __restrict__ send_buffer,
-                                          const int32x4_t* __restrict__ data) {
-    for (int k = 0; k < kRankAtoms; k++) {
-      int32x4_t const atom = data[k];
-
-      // Compute the absolute maximum of the atom in the thread group
-      // In 2 blocks of values, upper/lower halves of the f16x2_t
-      int wblockmax = group_abs_max<T>(atom);
-
-      // Derive scales
-      int decoding_scale;
-      int encoding_scale;
-      decoding_scale = packed_mul<T>(wblockmax, kScaleFactor);
-      encoding_scale = packed_add<T>(decoding_scale, kScaleEpsilon);
-      encoding_scale = packed_rcp<T>(encoding_scale);
-
-      // Apply scales to get quantized values
-      int32x4_t w;
-      for (int i = 0; i < 4; i++) {
-        w[i] = packed_mul<T>(atom[i], encoding_scale);
-        w[i] = packed_max<T>(w[i], kRangeMin);
-        w[i] = packed_min<T>(w[i], kRangeMax);
-      }
-
-      // Convert from f16x2_t to uint16x2_t
-      int32x4_t q;
-      {
-        int16_t* qi = reinterpret_cast<int16_t*>(&q);
-        T* wh = reinterpret_cast<T*>(&w);
-        for (int i = 0; i < 8; i++) qi[i] = (int16_t)rintf(T2float_cast(wh[i]));
-
-        for (int i = 0; i < 4; i++) {
-          q[i] = packed_add<int16_t>(q[i], kRangeBias);
-        }
-      }
-
-      // Pack 8 x q6 into int32_t + int16_t
-      uint32_t q4w;
-      uint16_t q2w = 0;
-      q4w = (q[0] & 0x000F000F) | ((q[1] & 0x000F000F) << 4) |
-            ((q[2] & 0x000F000F) << 8) | ((q[3] & 0x000F000F) << 12);
-      {
-        int16_t* tw = reinterpret_cast<int16_t*>(&q);
-#pragma unroll
-        for (int i = 0; i < 8; i++) {
-          q2w |= (tw[i] >> 4) << (i * 2);
-        }
-      }
-      // Write quantized atom to send_buffer
-      // note: only the group leader stores the scale
-      uint8_t* atom_ptr =
-          reinterpret_cast<uint8_t*>(send_buffer + k * kRankBufferTileStride);
-      uint32_t* q4w_ptr = reinterpret_cast<uint32_t*>(atom_ptr) + thread;
-      uint16_t* q2w_ptr =
-          reinterpret_cast<uint16_t*>(atom_ptr + kRankTileQ2Offset) + thread;
-      int* qs_ptr = reinterpret_cast<int*>(atom_ptr + kRankTileScaleOffset) +
-                    (thread / 8);
-
-      __builtin_nontemporal_store(q4w, q4w_ptr);
-      __builtin_nontemporal_store(q2w, q2w_ptr);
-      if (threadIdx.x == group_leader) {
-        __builtin_nontemporal_store(decoding_scale, qs_ptr);
-      }
-    }
-  }
-
-  __quickreduce_device_inline__ void recv(int32x4_t** __restrict__ recv_buffer,
-                                          int32x4_t* __restrict__ data) {
-    for (int k = 0; k < kRankAtoms; k++) {
-      // Directly read quantized atom from recv_buffer
-      uint8_t* atom_ptr = reinterpret_cast<uint8_t*>(*recv_buffer);
-      uint32_t* q4w_ptr = reinterpret_cast<uint32_t*>(atom_ptr) + thread;
-      uint16_t* q2w_ptr =
-          reinterpret_cast<uint16_t*>(atom_ptr + kRankTileQ2Offset) + thread;
-      int* qs_ptr = reinterpret_cast<int*>(atom_ptr + kRankTileScaleOffset) +
-                    (thread / 8);
-
-      uint32_t q4w = __builtin_nontemporal_load(q4w_ptr);
-      uint16_t q2w = __builtin_nontemporal_load(q2w_ptr);
-      int qs = __builtin_nontemporal_load(qs_ptr);
-
-      *recv_buffer += kRankBufferTileStride;
-
-      // Unpack q6 into fp16x8_t
-      int32x4_t w;
-      {
-        static uint constexpr kMask000F = 0x000F000F;
-        static uint constexpr kHalf2_1024 =
-            0x64006400;  // {1024.0, 1024.0}, fp16x2_t
-        static uint constexpr kHalf2_1056 =
-            0xE420E420;  // {-1056.0, -1056.0}, fp16x2_t
-
-#pragma unroll
-        for (int i = 0; i < 4; i++) {
-          int32_t q4 = q4w & kMask000F;
-          int32_t q2 = (q2w & 0x3) | ((q2w & 0xC) << 14);
-          q4w >>= 4;
-          q2w >>= 4;
-          if constexpr (std::is_same<T, half>::value) {
-            int32_t q6 = q4 | (q2 << 4) | kHalf2_1024;
-            asm volatile("v_pk_add_f16 %0, %1, %2"
-                         : "=v"(w[i])
-                         : "v"(q6), "v"(kHalf2_1056));
-          } else {
-            int32_t int16_2 = q4 | (q2 << 4);
-            int16_t low = static_cast<int16_t>(int16_2 & 0xFFFF);
-            int16_t high = static_cast<int16_t>((int16_2 >> 16) & 0xFFFF);
-
-            nv_bfloat16 bf_low = __float2bfloat16(static_cast<float>(low));
-            nv_bfloat16 bf_high = __float2bfloat16(static_cast<float>(high));
-            nv_bfloat162 bf2 = __halves2bfloat162(bf_low, bf_high);
-            int32_t packed_bf16 = *reinterpret_cast<int32_t*>(&bf2);
-            w[i] = packed_add<nv_bfloat16>(packed_bf16, kRangeMin);
-          }
-        }
-      }
-
-      // Apply decoding scales
-      for (int i = 0; i < 4; i++) {
-        w[i] = packed_mul<T>(w[i], qs);
-      }
-
-      // That's pretty much it...
-      data[k] = w;
-    }
-  }
-};
-
-// Int8 symmetric quantization codec.
-// We quantize the FP16 data to block-scaled Int8 in blocks of 4 *
-// kThreadGroupSize.
-template <typename T, int world_size>
-struct CodecQ8 : public CodecBase {
-  static constexpr int kWorldSize = world_size;
-
-  // Codec tile size process by this workgroup.
-  // Each threads processes a fragment of f16x8_t (16B),
-  // into a int8x8_t (8B) and a f16 scale shared among 32 values.
-  static constexpr int kRankAtoms = kAtoms / kWorldSize;
-  static constexpr int kRankTileStride = 2176;
-  static constexpr int kRankTileScaleOffset = 2048;
-  static constexpr int kRankTransmittedTileSize = kRankTileStride * kRankAtoms;
-  static_assert(kRankTransmittedTileSize % 16 == 0,
-                "kRankTileSize must be 16B aligned.");
-
-  static constexpr int kRankBufferTileStride =
-      kRankTileStride / sizeof(int32x4_t);
-
-  // Total tile size for the collective communication.
-  static constexpr int kTransmittedTileSize =
-      kRankTransmittedTileSize * kWorldSize;
-
-  // Constants configuration
-
-  // {-1/128.0h, -1/128.0h}, f16x2_t
-  static constexpr int kScaleFactor =
-      std::is_same<T, half>::value ? 0xA000A000 : 0xBC00BC00;
-
-  // {1e-7, 1e-7}, f16x2_t
-  static constexpr int kScaleEpsilon =
-      std::is_same<T, half>::value ? 0x00010001 : 0x33D733D7;
-
-  // {-128, -128}, f16x2_t
-  static constexpr int kRangeMin =
-      std::is_same<T, half>::value ? 0xD800D800 : 0xC300C300;
-  // {+127, +127}, f16x2_t
-  static constexpr int kRangeMax =
-      std::is_same<T, half>::value ? 0x57F057F0 : 0x42FE42FE;
-
-  // {+128, +128}, int16x2_t
-  static constexpr int kRangeBias = 0x00800080;
-
-  __quickreduce_device_inline__ CodecQ8(int thread, int rank)
-      : CodecBase(thread, rank) {}
-
-  __quickreduce_device_inline__ void send(int32x4_t* __restrict__ send_buffer,
-                                          int32x4_t const* __restrict__ data) {
-    for (int k = 0; k < kRankAtoms; k++) {
-      int32x4_t const atom = data[k];
-      // Compute the absolute maximum of the atom in the thread group
-      // In 2 blocks of values, upper/lower halves of the f16x2_t
-      int wblockmax = group_abs_max<T>(atom);
-
-      // Derive scales
-      int decoding_scale;
-      int encoding_scale;
-      decoding_scale = packed_mul<T>(wblockmax, kScaleFactor);
-      encoding_scale = packed_add<T>(decoding_scale, kScaleEpsilon);
-      encoding_scale = packed_rcp<T>(encoding_scale);
-
-      // Apply scales to get quantized values
-      int32x4_t w;
-      for (int i = 0; i < 4; i++) {
-        w[i] = packed_mul<T>(atom[i], encoding_scale);
-        w[i] = packed_max<T>(w[i], kRangeMin);
-        w[i] = packed_min<T>(w[i], kRangeMax);
-      }
-
-      // Convert from f16x2_t to uint16x2_t
-      int32x4_t q;
-      {
-        int16_t* qi = reinterpret_cast<int16_t*>(&q);
-        T* wh = reinterpret_cast<T*>(&w);
-        for (int i = 0; i < 8; i++) qi[i] = (int16_t)rintf(T2float_cast(wh[i]));
-
-        for (int i = 0; i < 4; i++) {
-          q[i] = packed_add<int16_t>(q[i], kRangeBias);
-        }
-      }
-
-      // Pack 8 x q8 into int32x2_t
-      int32x2_t qw;
-      qw[0] = q[0] | (q[1] << 8);
-      qw[1] = q[2] | (q[3] << 8);
-
-      // Write quantized atom to send_buffer
-      // note: only the group leader stores the scale
-      uint8_t* atom_ptr =
-          reinterpret_cast<uint8_t*>(send_buffer + k * kRankBufferTileStride);
-      int32x2_t* qw_ptr = reinterpret_cast<int32x2_t*>(atom_ptr) + thread;
-      int* qs_ptr = reinterpret_cast<int*>(atom_ptr + kRankTileScaleOffset) +
-                    (thread / 8);
-
-      __builtin_nontemporal_store(qw, qw_ptr);
-      if (threadIdx.x == group_leader) {
-        __builtin_nontemporal_store(decoding_scale, qs_ptr);
-      }
-    }
-  }
-
-  __quickreduce_device_inline__ void recv(int32x4_t** __restrict__ recv_buffer,
-                                          int32x4_t* __restrict__ data) {
-    for (int k = 0; k < kRankAtoms; k++) {
-      // Directly read quantized atom from recv_buffer
-      uint8_t* atom_ptr = reinterpret_cast<uint8_t*>(*recv_buffer);
-      int32x2_t* qw_ptr = reinterpret_cast<int32x2_t*>(atom_ptr) + thread;
-      int* qs_ptr = reinterpret_cast<int*>(atom_ptr + kRankTileScaleOffset) +
-                    (thread / 8);
-
-      int32x2_t qw = __builtin_nontemporal_load(qw_ptr);
-      int qs = __builtin_nontemporal_load(qs_ptr);
-
-      *recv_buffer += kRankBufferTileStride;
-
-      // Unpack q8 into fp16x8_t
-      int32x4_t w;
-      {
-        static uint constexpr kMask00FF = 0x00FF00FF;
-
-        // {1024.0, 1024.0}, fp16x2_t
-        static uint constexpr kHalf2_1024 = 0x64006400;
-
-        // {-1152.0, -1152.0}, fp16x2_t
-        static uint constexpr kHalf2_1152 = 0xE480E480;
-
-#pragma unroll
-        for (int i = 0; i < 4; i++) {
-          if constexpr (std::is_same<T, half>::value) {
-            int32_t q8 =
-                ((qw[i / 2] >> ((i % 2) * 8)) & kMask00FF) | kHalf2_1024;
-            w[i] = packed_add<half>(q8, kHalf2_1152);
-          } else {
-            int32_t int16_2 = (qw[i / 2] >> ((i % 2) * 8)) & kMask00FF;
-            int16_t low = static_cast<int16_t>(int16_2 & 0xFFFF);
-            int16_t high = static_cast<int16_t>((int16_2 >> 16) & 0xFFFF);
-            nv_bfloat16 bf_low = __float2bfloat16(static_cast<float>(low));
-            nv_bfloat16 bf_high = __float2bfloat16(static_cast<float>(high));
-            nv_bfloat162 bf2 = __halves2bfloat162(bf_low, bf_high);
-            int32_t packed_bf16 = *reinterpret_cast<int32_t*>(&bf2);
-            w[i] = packed_add<nv_bfloat16>(packed_bf16, kRangeMin);
-          }
-        }
-      }
-
-      // Apply decoding scales
-      for (int i = 0; i < 4; i++) {
-        w[i] = packed_mul<T>(w[i], qs);
-      }
-
-      data[k] = w;
-    }
-  }
-};
-
-// Twoshot All Reduce
-template <typename T, class Codec, bool cast_bf2half>
-struct AllReduceTwoshot {
-  static_assert(sizeof(T) == 2);
-
-  static constexpr int kWorldSize = Codec::kWorldSize;
-
-  __device__ static void run(
-      T const* __restrict__ input, T* __restrict__ output,
-      uint32_t const N,                    // number of elements
-      int const block,                     // block index
-      int const rank,                      // rank index
-      uint8_t** __restrict__ buffer_list,  // communication buffers
-      uint32_t const data_offset,          // offset to start of the data buffer
-      uint32_t flag_color) {
-    // Topology
-    int thread = threadIdx.x + threadIdx.y * kWavefront;
-    uint8_t* rank_buffer = buffer_list[rank];
-    Codec codec(thread, rank);
-    int block_id = blockIdx.x;
-    int grid_size = gridDim.x;
-    // --------------------------------------------------------
-    // Read input into registers
-    int32x4_t tA[kAtoms];
-
-    BufferResource src_buffer(const_cast<T*>(input), N * sizeof(T));
-    uint32_t src_offset = block * kTileSize + thread * sizeof(int32x4_t);
-
-    for (int i = 0; i < kAtoms; i++) {
-      tA[i] = buffer_load_dwordx4(src_buffer.descriptor, src_offset, 0, 0);
-      src_offset += kAtomStride * sizeof(int32x4_t);
-      if constexpr (cast_bf2half) {
-        const nv_bfloat162* bf_buf =
-            reinterpret_cast<const nv_bfloat162*>(&tA[i]);
-        half2 half_buf[4];
-#pragma unroll
-        for (int j = 0; j < 4; ++j) {
-          float2 f = __bfloat1622float2(bf_buf[j]);
-          half_buf[j] = __float22half2_rn(f);
-        }
-        tA[i] = *reinterpret_cast<const int32x4_t*>(half_buf);
-      }
-    }
-
-    // --------------------------------------------------------
-    // Phase-1A: Write segment data into the communication buffer of the target
-    // rank responsible for this segment.
-    uint32_t comm_data0_offset =
-        data_offset + block_id * Codec::kTransmittedTileSize;
-    uint32_t comm_data1_offset =
-        grid_size * Codec::kTransmittedTileSize + comm_data0_offset;
-
-    uint32_t comm_flags0_offset = block_id * (kWorldSize * sizeof(uint32_t));
-    uint32_t comm_flags1_offset =
-        grid_size * (kWorldSize * sizeof(uint32_t)) + comm_flags0_offset;
-
-    for (int r = 0; r < kWorldSize; r++) {
-      int32x4_t* send_buffer =
-          reinterpret_cast<int32x4_t*>(buffer_list[r] + comm_data0_offset +
-                                       rank * Codec::kRankTransmittedTileSize);
-      codec.send(send_buffer, &tA[r * Codec::kRankAtoms]);
-    }
-
-    __syncthreads();
-    if (thread < kWorldSize) {
-      int r = thread;
-      uint32_t* flag_ptr = reinterpret_cast<uint32_t*>(
-          buffer_list[r] + comm_flags0_offset + rank * sizeof(uint32_t));
-      set_sync_flag(flag_ptr, flag_color);
-    }
-    // --------------------------------------------------------
-    // Phase-1B: Reduce the segment data from the communication buffers.
-    int32x4_t tR[Codec::kRankAtoms] = {};
-    {
-      // Read the data from the communication buffer.
-      int32x4_t* recv_buffer =
-          reinterpret_cast<int32x4_t*>(rank_buffer + comm_data0_offset);
-      uint32_t* flag_ptr =
-          reinterpret_cast<uint32_t*>(rank_buffer + comm_flags0_offset);
-
-      for (int r = 0; r < kWorldSize; r++) {
-        // Wait for the flags to be set.
-        if (thread == 0) {
-          wait_sync_flag(&flag_ptr[r], flag_color);
-        }
-        __syncthreads();
-
-        // note: we reuse tA as temp buffer here
-        codec.recv(&recv_buffer, tA);
-
-        for (int i = 0; i < Codec::kRankAtoms; i++) {
-          packed_assign_add<T>(&tR[i], &tA[i]);
-        }
-      }
-    }
-
-    // Phase-2: Write the reduced segment to every other rank
-    for (int r = 0; r < kWorldSize; r++) {
-      int32x4_t* send_buffer =
-          reinterpret_cast<int32x4_t*>(buffer_list[r] + comm_data1_offset +
-                                       rank * Codec::kRankTransmittedTileSize);
-      codec.send(send_buffer, tR);
-    }
-
-    __syncthreads();
-    if (thread < kWorldSize) {
-      int r = thread;
-      uint32_t* flag_ptr = reinterpret_cast<uint32_t*>(
-          buffer_list[r] + comm_flags1_offset + rank * sizeof(uint32_t));
-      set_sync_flag(flag_ptr, flag_color);
-    }
-
-    // Phase-2: Read the gather segments from the rank's communication buffer.
-    {
-      // Read the data from the communication buffer.
-      int32x4_t* recv_buffer =
-          reinterpret_cast<int32x4_t*>(rank_buffer + comm_data1_offset);
-      uint32_t* flag_ptr =
-          reinterpret_cast<uint32_t*>(rank_buffer + comm_flags1_offset);
-
-      for (int r = 0; r < kWorldSize; r++) {
-        // Wait for the flags to be set.
-        if (thread == 0) {
-          wait_sync_flag(&flag_ptr[r], flag_color);
-        }
-        __syncthreads();
-
-        // Gather all reduced and final rank segments into tA.
-        codec.recv(&recv_buffer, &tA[r * Codec::kRankAtoms]);
-      }
-    }
-
-    // --------------------------------------------------------
-    // Write the result to output.
-    BufferResource dst_buffer(output, N * sizeof(T));
-    uint32_t dst_offset = block * kTileSize + thread * sizeof(int32x4_t);
-
-    for (int i = 0; i < kAtoms; i++) {
-      if constexpr (cast_bf2half) {
-        const half2* half_buf = reinterpret_cast<const half2*>(&tA[i]);
-        nv_bfloat162 bf16_buf[4];
-#pragma unroll
-        for (int j = 0; j < 4; ++j) {
-          float2 f = __half22float2(half_buf[j]);
-          bf16_buf[j] = __float22bfloat162_rn(f);
-        }
-        buffer_store_dwordx4(*reinterpret_cast<const int32x4_t*>(bf16_buf),
-                             dst_buffer.descriptor, dst_offset, 0, 0);
-      } else {
-        buffer_store_dwordx4(tA[i], dst_buffer.descriptor, dst_offset, 0, 0);
-      }
-      dst_offset += kAtomStride * sizeof(int32x4_t);
-    }
-  }
-};
-
-}  // namespace quickreduce

csrc/rocm/attention.cu

@@ -1598,6 +1598,7 @@ __launch_bounds__(NUM_THREADS, 3) void paged_attention_ll4mi_QKV_mfma16_kernel(
   const int warpid = threadIdx.x / WARP_SIZE;
   const int laneid = threadIdx.x % WARP_SIZE;
   const int lane2id = laneid % 2;
+  const int lane4id = laneid % 4;
   const int lane16id = laneid % 16;
   const int rowid = laneid / 16;
 
@@ -1744,6 +1745,7 @@ __launch_bounds__(NUM_THREADS, 3) void paged_attention_ll4mi_QKV_mfma16_kernel(
     const cache_t* k_ptr2 = k_ptr + kblock_number * kv_block_stride;
     const int klocal_token_idx =
         TOKENS_PER_WARP * warpid + token_depth * 16 + lane16id;
+    const int kglobal_token_idx = partition_start_token_idx + klocal_token_idx;
     const int kphysical_block_offset = klocal_token_idx % BLOCK_SIZE;
     const cache_t* k_ptr3 = k_ptr2 + kphysical_block_offset * KX;
 
@@ -2366,6 +2368,7 @@ __launch_bounds__(NUM_THREADS, 3) void paged_attention_ll4mi_QKV_mfma16_kernel(
   const int warpid = threadIdx.x / WARP_SIZE;
   const int laneid = threadIdx.x % WARP_SIZE;
   const int lane2id = laneid % 2;
+  const int lane4id = laneid % 4;
   const int lane16id = laneid % 16;
   const int rowid = laneid / 16;
 
@@ -2511,6 +2514,7 @@ __launch_bounds__(NUM_THREADS, 3) void paged_attention_ll4mi_QKV_mfma16_kernel(
     const cache_t* k_ptr2 = k_ptr + kblock_number * kv_block_stride;
     const int klocal_token_idx =
         TOKENS_PER_WARP * warpid + token_depth * 16 + lane16id;
+    const int kglobal_token_idx = partition_start_token_idx + klocal_token_idx;
     const int kphysical_block_offset = klocal_token_idx % BLOCK_SIZE;
     const cache_t* k_ptr3 = k_ptr2 + kphysical_block_offset * KX;
 

csrc/torch_bindings.cpp

@@ -725,24 +725,6 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _custom_ar), custom_ar) {
   custom_ar.impl("open_mem_handle", torch::kCPU, &open_mem_handle);
 
   custom_ar.def("free_shared_buffer", &free_shared_buffer);
-#ifdef USE_ROCM
-  // Quick Reduce all-reduce kernels
-  custom_ar.def(
-      "qr_all_reduce(int fa, Tensor inp, Tensor out, int quant_level, bool "
-      "cast_bf2half) -> ()");
-  custom_ar.impl("qr_all_reduce", torch::kCUDA, &qr_all_reduce);
-
-  custom_ar.def("init_custom_qr", &init_custom_qr);
-  custom_ar.def("qr_destroy", &qr_destroy);
-
-  custom_ar.def("qr_get_handle", &qr_get_handle);
-
-  custom_ar.def("qr_open_handles(int _fa, Tensor[](b!) handles) -> ()");
-  custom_ar.impl("qr_open_handles", torch::kCPU, &qr_open_handles);
-
-  // Max input size in bytes
-  custom_ar.def("qr_max_size", &qr_max_size);
-#endif
 }
 
 REGISTER_EXTENSION(TORCH_EXTENSION_NAME)

docker/Dockerfile

@@ -6,106 +6,30 @@
 # docs/assets/contributing/dockerfile-stages-dependency.png
 
 ARG CUDA_VERSION=12.8.1
-ARG PYTHON_VERSION=3.12
-
-# By parameterizing the base images, we allow third-party to use their own
-# base images. One use case is hermetic builds with base images stored in
-# private registries that use a different repository naming conventions.
-#
-# Example:
-# docker build --build-arg BUILD_BASE_IMAGE=registry.acme.org/mirror/nvidia/cuda:${CUDA_VERSION}-devel-ubuntu20.04
-ARG BUILD_BASE_IMAGE=nvidia/cuda:${CUDA_VERSION}-devel-ubuntu20.04
-ARG FINAL_BASE_IMAGE=nvidia/cuda:${CUDA_VERSION}-devel-ubuntu22.04
-
-# By parameterizing the Deadsnakes repository URL, we allow third-party to use
-# their own mirror. When doing so, we don't benefit from the transparent
-# installation of the GPG key of the PPA, as done by add-apt-repository, so we
-# also need a URL for the GPG key.
-ARG DEADSNAKES_MIRROR_URL
-ARG DEADSNAKES_GPGKEY_URL
-
-# The PyPA get-pip.py script is a self contained script+zip file, that provides
-# both the installer script and the pip base85-encoded zip archive. This allows
-# bootstrapping pip in environment where a dsitribution package does not exist.
-#
-# By parameterizing the URL for get-pip.py installation script, we allow
-# third-party to use their own copy of the script stored in a private mirror.
-# We set the default value to the PyPA owned get-pip.py script.
-#
-# Reference: https://pip.pypa.io/en/stable/installation/#get-pip-py
-ARG GET_PIP_URL="https://bootstrap.pypa.io/get-pip.py"
-
-# PIP supports fetching the packages from custom indexes, allowing third-party
-# to host the packages in private mirrors. The PIP_INDEX_URL and
-# PIP_EXTRA_INDEX_URL are standard PIP environment variables to override the
-# default indexes. By letting them empty by default, PIP will use its default
-# indexes if the build process doesn't override the indexes.
-#
-# Uv uses different variables. We set them by default to the same values as
-# PIP, but they can be overridden.
-ARG PIP_INDEX_URL
-ARG PIP_EXTRA_INDEX_URL
-ARG UV_INDEX_URL=${PIP_INDEX_URL}
-ARG UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
-
-# PyTorch provides its own indexes for standard and nightly builds
-ARG PYTORCH_CUDA_INDEX_BASE_URL=https://download.pytorch.org/whl
-ARG PYTORCH_CUDA_NIGHTLY_INDEX_BASE_URL=https://download.pytorch.org/whl/nightly
-
-# PIP supports multiple authentication schemes, including keyring
-# By parameterizing the PIP_KEYRING_PROVIDER variable and setting it to
-# disabled by default, we allow third-party to use keyring authentication for
-# their private Python indexes, while not changing the default behavior which
-# is no authentication.
-#
-# Reference: https://pip.pypa.io/en/stable/topics/authentication/#keyring-support
-ARG PIP_KEYRING_PROVIDER=disabled
-ARG UV_KEYRING_PROVIDER=${PIP_KEYRING_PROVIDER}
-
 #################### BASE BUILD IMAGE ####################
 # prepare basic build environment
-FROM ${BUILD_BASE_IMAGE} AS base
-ARG CUDA_VERSION
-ARG PYTHON_VERSION
+FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu20.04 AS base
+ARG CUDA_VERSION=12.8.1
+ARG PYTHON_VERSION=3.12
 ARG TARGETPLATFORM
 ENV DEBIAN_FRONTEND=noninteractive
 
-ARG DEADSNAKES_MIRROR_URL
-ARG DEADSNAKES_GPGKEY_URL
-ARG GET_PIP_URL
-
 # Install Python and other dependencies
 RUN echo 'tzdata tzdata/Areas select America' | debconf-set-selections \
     && echo 'tzdata tzdata/Zones/America select Los_Angeles' | debconf-set-selections \
     && apt-get update -y \
     && apt-get install -y ccache software-properties-common git curl sudo \
-    && if [ ! -z ${DEADSNAKES_MIRROR_URL} ] ; then \
-        if [ ! -z "${DEADSNAKES_GPGKEY_URL}" ] ; then \
-            mkdir -p -m 0755 /etc/apt/keyrings ; \
-            curl -L ${DEADSNAKES_GPGKEY_URL} | gpg --dearmor > /etc/apt/keyrings/deadsnakes.gpg ; \
-            sudo chmod 644 /etc/apt/keyrings/deadsnakes.gpg ; \
-            echo "deb [signed-by=/etc/apt/keyrings/deadsnakes.gpg] ${DEADSNAKES_MIRROR_URL} $(lsb_release -cs) main" > /etc/apt/sources.list.d/deadsnakes.list ; \
-        fi ; \
-    else \
-        for i in 1 2 3; do \
-            add-apt-repository -y ppa:deadsnakes/ppa && break || \
-            { echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \
-        done ; \
-    fi \
+    && for i in 1 2 3; do \
+        add-apt-repository -y ppa:deadsnakes/ppa && break || \
+        { echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \
+    done \
     && apt-get update -y \
     && apt-get install -y python${PYTHON_VERSION} python${PYTHON_VERSION}-dev python${PYTHON_VERSION}-venv \
     && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 \
     && update-alternatives --set python3 /usr/bin/python${PYTHON_VERSION} \
     && ln -sf /usr/bin/python${PYTHON_VERSION}-config /usr/bin/python3-config \
-    && curl -sS ${GET_PIP_URL} | python${PYTHON_VERSION} \
+    && curl -sS https://bootstrap.pypa.io/get-pip.py | python${PYTHON_VERSION} \
     && python3 --version && python3 -m pip --version
-
-ARG PIP_INDEX_URL UV_INDEX_URL
-ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
-ARG PYTORCH_CUDA_INDEX_BASE_URL
-ARG PYTORCH_CUDA_NIGHTLY_INDEX_BASE_URL
-ARG PIP_KEYRING_PROVIDER UV_KEYRING_PROVIDER
-
 # Install uv for faster pip installs
 RUN --mount=type=cache,target=/root/.cache/uv \
     python3 -m pip install uv
@@ -139,25 +63,21 @@ WORKDIR /workspace
 # after this step
 RUN --mount=type=cache,target=/root/.cache/uv \
     if [ "$TARGETPLATFORM" = "linux/arm64" ]; then \
-        uv pip install --system \
-            --index-url ${PYTORCH_CUDA_NIGHTLY_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
-            "torch==2.8.0.dev20250318+cu128" "torchvision==0.22.0.dev20250319";  \
-        uv pip install --system \
-            --index-url ${PYTORCH_CUDA_NIGHTLY_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
-            --pre pytorch_triton==3.3.0+gitab727c40; \
+        uv pip install --system --index-url https://download.pytorch.org/whl/nightly/cu128 "torch==2.8.0.dev20250318+cu128" "torchvision==0.22.0.dev20250319";  \
+        uv pip install --system --index-url https://download.pytorch.org/whl/nightly/cu128 --pre pytorch_triton==3.3.0+gitab727c40; \
     fi
 
 COPY requirements/common.txt requirements/common.txt
 COPY requirements/cuda.txt requirements/cuda.txt
 RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip install --system -r requirements/cuda.txt \
-    --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
+    --extra-index-url https://download.pytorch.org/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
 
 # cuda arch list used by torch
 # can be useful for both `dev` and `test`
 # explicitly set the list to avoid issues with torch 2.2
 # see https://github.com/pytorch/pytorch/pull/123243
-ARG torch_cuda_arch_list='7.0 7.5 8.0 8.9 9.0 10.0 12.0'
+ARG torch_cuda_arch_list='7.0 7.5 8.0 8.9 9.0 10.0+PTX'
 ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list}
 # Override the arch list for flash-attn to reduce the binary size
 ARG vllm_fa_cmake_gpu_arches='80-real;90-real'
@@ -168,10 +88,6 @@ ENV VLLM_FA_CMAKE_GPU_ARCHES=${vllm_fa_cmake_gpu_arches}
 FROM base AS build
 ARG TARGETPLATFORM
 
-ARG PIP_INDEX_URL UV_INDEX_URL
-ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
-ARG PYTORCH_CUDA_INDEX_BASE_URL
-
 # install build dependencies
 COPY requirements/build.txt requirements/build.txt
 
@@ -182,7 +98,7 @@ ENV UV_INDEX_STRATEGY="unsafe-best-match"
 
 RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip install --system -r requirements/build.txt \
-    --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
+    --extra-index-url https://download.pytorch.org/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
 
 COPY . .
 ARG GIT_REPO_CHECK=0
@@ -197,8 +113,6 @@ ARG nvcc_threads=8
 ENV NVCC_THREADS=$nvcc_threads
 
 ARG USE_SCCACHE
-ARG SCCACHE_DOWNLOAD_URL=https://github.com/mozilla/sccache/releases/download/v0.8.1/sccache-v0.8.1-x86_64-unknown-linux-musl.tar.gz
-ARG SCCACHE_ENDPOINT
 ARG SCCACHE_BUCKET_NAME=vllm-build-sccache
 ARG SCCACHE_REGION_NAME=us-west-2
 ARG SCCACHE_S3_NO_CREDENTIALS=0
@@ -207,11 +121,10 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     --mount=type=bind,source=.git,target=.git \
     if [ "$USE_SCCACHE" = "1" ]; then \
         echo "Installing sccache..." \
-        && curl -L -o sccache.tar.gz ${SCCACHE_DOWNLOAD_URL} \
+        && curl -L -o sccache.tar.gz https://github.com/mozilla/sccache/releases/download/v0.8.1/sccache-v0.8.1-x86_64-unknown-linux-musl.tar.gz \
         && tar -xzf sccache.tar.gz \
         && sudo mv sccache-v0.8.1-x86_64-unknown-linux-musl/sccache /usr/bin/sccache \
         && rm -rf sccache.tar.gz sccache-v0.8.1-x86_64-unknown-linux-musl \
-        && if [ ! -z ${SCCACHE_ENDPOINT} ] ; then export SCCACHE_ENDPOINT=${SCCACHE_ENDPOINT} ; fi \
         && export SCCACHE_BUCKET=${SCCACHE_BUCKET_NAME} \
         && export SCCACHE_REGION=${SCCACHE_REGION_NAME} \
         && export SCCACHE_S3_NO_CREDENTIALS=${SCCACHE_S3_NO_CREDENTIALS} \
@@ -249,10 +162,6 @@ RUN if [ "$RUN_WHEEL_CHECK" = "true" ]; then \
 #################### DEV IMAGE ####################
 FROM base as dev
 
-ARG PIP_INDEX_URL UV_INDEX_URL
-ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
-ARG PYTORCH_CUDA_INDEX_BASE_URL
-
 # This timeout (in seconds) is necessary when installing some dependencies via uv since it's likely to time out
 # Reference: https://github.com/astral-sh/uv/pull/1694
 ENV UV_HTTP_TIMEOUT=500
@@ -267,25 +176,21 @@ COPY requirements/test.txt requirements/test.txt
 COPY requirements/dev.txt requirements/dev.txt
 RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip install --system -r requirements/dev.txt \
-    --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
+    --extra-index-url https://download.pytorch.org/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
 #################### DEV IMAGE ####################
 
 #################### vLLM installation IMAGE ####################
 # image with vLLM installed
 # TODO: Restore to base image after FlashInfer AOT wheel fixed
-FROM ${FINAL_BASE_IMAGE} AS vllm-base
-ARG CUDA_VERSION
-ARG PYTHON_VERSION
+FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu22.04 AS vllm-base
+ARG CUDA_VERSION=12.8.1
+ARG PYTHON_VERSION=3.12
 WORKDIR /vllm-workspace
 ENV DEBIAN_FRONTEND=noninteractive
 ARG TARGETPLATFORM
 
 SHELL ["/bin/bash", "-c"]
 
-ARG DEADSNAKES_MIRROR_URL
-ARG DEADSNAKES_GPGKEY_URL
-ARG GET_PIP_URL
-
 RUN PYTHON_VERSION_STR=$(echo ${PYTHON_VERSION} | sed 's/\.//g') && \
     echo "export PYTHON_VERSION_STR=${PYTHON_VERSION_STR}" >> /etc/environment
 
@@ -295,33 +200,17 @@ RUN echo 'tzdata tzdata/Areas select America' | debconf-set-selections \
     && apt-get update -y \
     && apt-get install -y ccache software-properties-common git curl wget sudo vim python3-pip \
     && apt-get install -y ffmpeg libsm6 libxext6 libgl1 \
-    && if [ ! -z ${DEADSNAKES_MIRROR_URL} ] ; then \
-        if [ ! -z "${DEADSNAKES_GPGKEY_URL}" ] ; then \
-            mkdir -p -m 0755 /etc/apt/keyrings ; \
-            curl -L ${DEADSNAKES_GPGKEY_URL} | gpg --dearmor > /etc/apt/keyrings/deadsnakes.gpg ; \
-            sudo chmod 644 /etc/apt/keyrings/deadsnakes.gpg ; \
-            echo "deb [signed-by=/etc/apt/keyrings/deadsnakes.gpg] ${DEADSNAKES_MIRROR_URL} $(lsb_release -cs) main" > /etc/apt/sources.list.d/deadsnakes.list ; \
-        fi ; \
-    else \
-        for i in 1 2 3; do \
-            add-apt-repository -y ppa:deadsnakes/ppa && break || \
-            { echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \
-        done ; \
-    fi \
+    && for i in 1 2 3; do \
+        add-apt-repository -y ppa:deadsnakes/ppa && break || \
+        { echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \
+    done \
     && apt-get update -y \
     && apt-get install -y python${PYTHON_VERSION} python${PYTHON_VERSION}-dev python${PYTHON_VERSION}-venv libibverbs-dev \
     && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 \
     && update-alternatives --set python3 /usr/bin/python${PYTHON_VERSION} \
     && ln -sf /usr/bin/python${PYTHON_VERSION}-config /usr/bin/python3-config \
-    && curl -sS ${GET_PIP_URL} | python${PYTHON_VERSION} \
+    && curl -sS https://bootstrap.pypa.io/get-pip.py | python${PYTHON_VERSION} \
     && python3 --version && python3 -m pip --version
-
-ARG PIP_INDEX_URL UV_INDEX_URL
-ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
-ARG PYTORCH_CUDA_INDEX_BASE_URL
-ARG PYTORCH_CUDA_NIGHTLY_INDEX_BASE_URL
-ARG PIP_KEYRING_PROVIDER UV_KEYRING_PROVIDER
-
 # Install uv for faster pip installs
 RUN --mount=type=cache,target=/root/.cache/uv \
     python3 -m pip install uv
@@ -343,23 +232,19 @@ RUN ldconfig /usr/local/cuda-$(echo $CUDA_VERSION | cut -d. -f1,2)/compat/
 # after this step
 RUN --mount=type=cache,target=/root/.cache/uv \
     if [ "$TARGETPLATFORM" = "linux/arm64" ]; then \
-        uv pip install --system \
-            --index-url ${PYTORCH_CUDA_NIGHTLY_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
-            "torch==2.8.0.dev20250318+cu128" "torchvision==0.22.0.dev20250319" ; \
-        uv pip install --system \
-            --index-url ${PYTORCH_CUDA_NIGHTLY_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
-            --pre pytorch_triton==3.3.0+gitab727c40 ; \
+        uv pip install --system --index-url https://download.pytorch.org/whl/nightly/cu128 "torch==2.8.0.dev20250318+cu128" "torchvision==0.22.0.dev20250319";  \
+        uv pip install --system --index-url https://download.pytorch.org/whl/nightly/cu128 --pre pytorch_triton==3.3.0+gitab727c40; \
     fi
 
 # Install vllm wheel first, so that torch etc will be installed.
 RUN --mount=type=bind,from=build,src=/workspace/dist,target=/vllm-workspace/dist \
     --mount=type=cache,target=/root/.cache/uv \
     uv pip install --system dist/*.whl --verbose \
-        --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
+    --extra-index-url https://download.pytorch.org/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
 
 # If we need to build FlashInfer wheel before its release:
 # $ # Note we remove 7.0 from the arch list compared to the list below, since FlashInfer only supports sm75+
-# $ export TORCH_CUDA_ARCH_LIST='7.5 8.0 8.9 9.0a 10.0a 12.0'
+# $ export TORCH_CUDA_ARCH_LIST='7.5 8.0 8.9 9.0a 10.0a'
 # $ git clone https://github.com/flashinfer-ai/flashinfer.git --recursive
 # $ cd flashinfer
 # $ git checkout v0.2.6.post1
@@ -369,20 +254,15 @@ RUN --mount=type=bind,from=build,src=/workspace/dist,target=/vllm-workspace/dist
 # -rw-rw-r-- 1 mgoin mgoin 205M Jun  9 18:03 flashinfer_python-0.2.6.post1-cp39-abi3-linux_x86_64.whl
 # $ # upload the wheel to a public location, e.g. https://wheels.vllm.ai/flashinfer/v0.2.6.post1/flashinfer_python-0.2.6.post1-cp39-abi3-linux_x86_64.whl
 
-# Allow specifying a version, Git revision or local .whl file
-ARG FLASHINFER_CUDA128_INDEX_URL="https://download.pytorch.org/whl/cu128/flashinfer"
-ARG FLASHINFER_CUDA128_WHEEL="flashinfer_python-0.2.6.post1%2Bcu128torch2.7-cp39-abi3-linux_x86_64.whl"
-ARG FLASHINFER_GIT_REPO="https://github.com/flashinfer-ai/flashinfer.git"
-ARG FLASHINFER_GIT_REF="v0.2.6.post1"
 RUN --mount=type=cache,target=/root/.cache/uv \
 . /etc/environment && \
 if [ "$TARGETPLATFORM" != "linux/arm64" ]; then \
     # FlashInfer already has a wheel for PyTorch 2.7.0 and CUDA 12.8. This is enough for CI use
     if [[ "$CUDA_VERSION" == 12.8* ]]; then \
-        uv pip install --system ${FLASHINFER_CUDA128_INDEX_URL}/${FLASHINFER_CUDA128_WHEEL} ; \
+        uv pip install --system https://download.pytorch.org/whl/cu128/flashinfer/flashinfer_python-0.2.6.post1%2Bcu128torch2.7-cp39-abi3-linux_x86_64.whl; \
     else \
-        export TORCH_CUDA_ARCH_LIST='7.5 8.0 8.9 9.0a 10.0a 12.0' && \
-        git clone ${FLASHINFER_GIT_REPO} --single-branch --branch ${FLASHINFER_GIT_REF} --recursive && \
+        export TORCH_CUDA_ARCH_LIST='7.5 8.0 8.9 9.0a 10.0a' && \
+        git clone https://github.com/flashinfer-ai/flashinfer.git --single-branch --branch v0.2.6.post1 --recursive && \
         # Needed to build AOT kernels
         (cd flashinfer && \
             python3 -m flashinfer.aot && \
@@ -406,7 +286,7 @@ uv pip list
 COPY requirements/build.txt requirements/build.txt
 RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip install --system -r requirements/build.txt \
-        --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
+    --extra-index-url https://download.pytorch.org/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
 
 #################### vLLM installation IMAGE ####################
 
@@ -417,11 +297,6 @@ FROM vllm-base AS test
 
 ADD . /vllm-workspace/
 
-ARG PYTHON_VERSION
-
-ARG PIP_INDEX_URL UV_INDEX_URL
-ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
-
 # This timeout (in seconds) is necessary when installing some dependencies via uv since it's likely to time out
 # Reference: https://github.com/astral-sh/uv/pull/1694
 ENV UV_HTTP_TIMEOUT=500
@@ -432,7 +307,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip install --system --no-build-isolation "git+https://github.com/state-spaces/mamba@v2.2.4"
 
 # install development dependencies (for testing)
-RUN --mount=type=cache,target=/root/.cache/uv \
+RUN --mount=type=cache,target=/root/.cache/uv \   
     CUDA_MAJOR="${CUDA_VERSION%%.*}"; \
     if [ "$CUDA_MAJOR" -ge 12 ]; then \
         uv pip install --system -r requirements/dev.txt; \
@@ -448,7 +323,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 ENV HF_HUB_ENABLE_HF_TRANSFER 1
 
 # Copy in the v1 package for testing (it isn't distributed yet)
-COPY vllm/v1 /usr/local/lib/python${PYTHON_VERSION}/dist-packages/vllm/v1
+COPY vllm/v1 /usr/local/lib/python3.12/dist-packages/vllm/v1
 
 # doc requires source code
 # we hide them inside `test_docs/` , so that this source code
@@ -465,9 +340,6 @@ RUN mv mkdocs.yaml test_docs/
 FROM vllm-base AS vllm-openai-base
 ARG TARGETPLATFORM
 
-ARG PIP_INDEX_URL UV_INDEX_URL
-ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
-
 # This timeout (in seconds) is necessary when installing some dependencies via uv since it's likely to time out
 # Reference: https://github.com/astral-sh/uv/pull/1694
 ENV UV_HTTP_TIMEOUT=500

docker/Dockerfile.cpu

@@ -66,7 +66,7 @@ ENV VLLM_CPU_DISABLE_AVX512=${VLLM_CPU_DISABLE_AVX512}
 WORKDIR /workspace/vllm
 
 RUN --mount=type=cache,target=/root/.cache/uv \
-    --mount=type=bind,src=requirements/cpu-build.txt,target=requirements/build.txt \
+    --mount=type=bind,src=requirements/build.txt,target=requirements/build.txt \
     uv pip install -r requirements/build.txt
 
 COPY . .
@@ -79,22 +79,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     --mount=type=bind,source=.git,target=.git \
     VLLM_TARGET_DEVICE=cpu python3 setup.py bdist_wheel 
 
-######################### TEST DEPS #########################
-FROM base AS vllm-test-deps
-
-WORKDIR /workspace/vllm
-
-RUN --mount=type=bind,src=requirements/test.in,target=requirements/test.in \
-    cp requirements/test.in requirements/cpu-test.in && \
-    sed -i '/mamba_ssm/d' requirements/cpu-test.in && \
-    sed -i 's/torch==.*/torch==2.6.0/g' requirements/cpu-test.in && \
-    sed -i 's/torchaudio.*/torchaudio/g' requirements/cpu-test.in && \
-    sed -i 's/torchvision.*/torchvision/g' requirements/cpu-test.in && \
-    uv pip compile requirements/cpu-test.in -o requirements/cpu-test.txt --index-strategy unsafe-best-match --torch-backend cpu
-
-RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install -r requirements/cpu-test.txt 
-
 ######################### DEV IMAGE #########################
 FROM vllm-build AS vllm-dev
 
@@ -113,19 +97,28 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     --mount=type=bind,source=.git,target=.git \
     VLLM_TARGET_DEVICE=cpu python3 setup.py develop 
 
-COPY --from=vllm-test-deps /workspace/vllm/requirements/cpu-test.txt requirements/test.txt
-
 RUN --mount=type=cache,target=/root/.cache/uv \
+    --mount=type=bind,src=requirements/test.in,target=requirements/test.in \
+    cp requirements/test.in requirements/test-cpu.in && \
+    sed -i '/mamba_ssm/d' requirements/test-cpu.in && \
+    uv pip compile requirements/test-cpu.in -o requirements/test.txt && \
     uv pip install -r requirements/dev.txt && \
     pre-commit install --hook-type pre-commit --hook-type commit-msg
 
 ENTRYPOINT ["bash"]
 
 ######################### TEST IMAGE #########################
-FROM vllm-test-deps AS vllm-test
+FROM base AS vllm-test
 
 WORKDIR /workspace/
 
+RUN --mount=type=cache,target=/root/.cache/uv \
+    --mount=type=bind,src=requirements/test.in,target=requirements/test.in \
+    cp requirements/test.in requirements/test-cpu.in && \
+    sed -i '/mamba_ssm/d' requirements/test-cpu.in && \
+    uv pip compile requirements/test-cpu.in -o requirements/cpu-test.txt && \
+    uv pip install -r requirements/cpu-test.txt
+
 RUN --mount=type=cache,target=/root/.cache/uv \
     --mount=type=bind,from=vllm-build,src=/workspace/vllm/dist,target=dist \
     uv pip install dist/*.whl

docker/Dockerfile.xpu

@@ -35,7 +35,6 @@ RUN --mount=type=bind,source=.git,target=.git \
     if [ "$GIT_REPO_CHECK" != 0 ]; then bash tools/check_repo.sh; fi
 
 ENV VLLM_TARGET_DEVICE=xpu
-ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
 
 RUN --mount=type=cache,target=/root/.cache/pip \
     --mount=type=bind,source=.git,target=.git \

docs/.nav.yml

@@ -39,7 +39,6 @@ nav:
       - models/generative_models.md
       - models/pooling_models.md
       - models/extensions
-      - Hardware Supported Models: models/hardware_supported_models
     - Features:
       - features/compatibility_matrix.md
       - features/*
@@ -49,12 +48,7 @@ nav:
     - General:
       - glob: contributing/*
         flatten_single_child_sections: true
-    - Model Implementation: 
-      - contributing/model/README.md
-      - contributing/model/basic.md
-      - contributing/model/registration.md
-      - contributing/model/tests.md
-      - contributing/model/multimodal.md
+    - Model Implementation: contributing/model
     - Design Documents:
       - V0: design
       - V1: design/v1

docs/README.md

@@ -1,8 +1,7 @@
 # Welcome to vLLM
 
 <figure markdown="span">
-  ![](./assets/logos/vllm-logo-text-light.png){ align="center" alt="vLLM Light" class="logo-light" width="60%" }
-  ![](./assets/logos/vllm-logo-text-dark.png){ align="center" alt="vLLM Dark" class="logo-dark" width="60%" }
+  ![](./assets/logos/vllm-logo-text-light.png){ align="center" alt="vLLM" class="no-scaled-link" width="60%" }
 </figure>
 
 <p style="text-align:center">
@@ -41,7 +40,7 @@ vLLM is flexible and easy to use with:
 - OpenAI-compatible API server
 - Support NVIDIA GPUs, AMD CPUs and GPUs, Intel CPUs, Gaudi® accelerators and GPUs, IBM Power CPUs, TPU, and AWS Trainium and Inferentia Accelerators.
 - Prefix caching support
-- Multi-LoRA support
+- Multi-lora support
 
 For more information, check out the following:
 

docs/configuration/engine_args.md

@@ -6,7 +6,7 @@ title: Engine Arguments
 Engine arguments control the behavior of the vLLM engine.
 
 - For [offline inference][offline-inference], they are part of the arguments to [LLM][vllm.LLM] class.
-- For [online serving][serving-openai-compatible-server], they are part of the arguments to `vllm serve`.
+- For [online serving][openai-compatible-server], they are part of the arguments to `vllm serve`.
 
 You can look at [EngineArgs][vllm.engine.arg_utils.EngineArgs] and [AsyncEngineArgs][vllm.engine.arg_utils.AsyncEngineArgs] to see the available engine arguments.
 

docs/contributing/README.md

@@ -29,8 +29,6 @@ See <gh-file:LICENSE>.
 Depending on the kind of development you'd like to do (e.g. Python, CUDA), you can choose to build vLLM with or without compilation.
 Check out the [building from source][build-from-source] documentation for details.
 
-For an optimized workflow when iterating on C++/CUDA kernels, see the [Incremental Compilation Workflow](./incremental_build.md) for recommendations.
-
 ### Building the docs with MkDocs
 
 #### Introduction to MkDocs
@@ -151,14 +149,6 @@ the terms of the DCO.
 
 Using `-s` with `git commit` will automatically add this header.
 
-!!! tip
-    You can enable automatic sign-off via your IDE:
-  
-    - **PyCharm**: Click on the `Show Commit Options` icon to the right of the `Commit and Push...` button in the `Commit` window.
-      It will bring up a `git` window where you can modify the `Author` and enable `Sign-off commit`.
-    - **VSCode**: Open the [Settings editor](https://code.visualstudio.com/docs/configure/settings)
-      and enable the `Git: Always Sign Off` (`git.alwaysSignOff`) field.
-
 ### PR Title and Classification
 
 Only specific types of PRs will be reviewed. The PR title is prefixed
@@ -198,7 +188,6 @@ The PR needs to meet the following code quality standards:
 
 ### Adding or Changing Kernels
 
-When actively developing or modifying kernels, using the [Incremental Compilation Workflow](./incremental_build.md) is highly recommended for faster build times.
 Each custom kernel needs a schema and one or more implementations to be registered with PyTorch.
 
 - Make sure custom ops are registered following PyTorch guidelines:

docs/contributing/deprecation_policy.md

@@ -37,14 +37,14 @@ multiple Y releases:
 - **Timeline**: A removal version is explicitly stated in the deprecation
 warning (e.g., "This will be removed in v0.10.0").
 - **Communication**: Deprecation is noted in the following, as applicable:
-    - Help strings
-    - Log output
-    - API responses
-    - `/metrics` output (for metrics features)
-    - User-facing documentation
-    - Release notes
-    - GitHub Issue (RFC) for feedback
-    - Documentation and use of the `@typing_extensions.deprecated` decorator for Python APIs
+  - Help strings
+  - Log output
+  - API responses
+  - `/metrics` output (for metrics features)
+  - User-facing documentation
+  - Release notes
+  - GitHub Issue (RFC) for feedback
+  - Documentation and use of the `@typing_extensions.deprecated` decorator for Python APIs
 
 **2.Deprecated (Off By Default)**
 

docs/contributing/incremental_build.md

@@ -1,138 +0,0 @@
-# Incremental Compilation Workflow
-
-When working on vLLM's C++/CUDA kernels located in the `csrc/` directory, recompiling the entire project with `uv pip install -e .` for every change can be time-consuming. An incremental compilation workflow using CMake allows for faster iteration by only recompiling the necessary components after an initial setup. This guide details how to set up and use such a workflow, which complements your editable Python installation.
-
-## Prerequisites
-
-Before setting up the incremental build:
-
-1. **vLLM Editable Install:** Ensure you have vLLM installed from source in an editable mode. Using pre-compiled wheels for the initial editable setup can be faster, as the CMake workflow will handle subsequent kernel recompilations.
-
-    ```console
-    uv venv --python 3.12 --seed
-    source .venv/bin/activate
-    VLLM_USE_PRECOMPILED=1 uv pip install -U -e . --torch-backend=auto
-    ```
-
-2. **CUDA Toolkit:** Verify that the NVIDIA CUDA Toolkit is correctly installed and `nvcc` is accessible in your `PATH`. CMake relies on `nvcc` to compile CUDA code. You can typically find `nvcc` in `$CUDA_HOME/bin/nvcc` or by running `which nvcc`. If you encounter issues, refer to the [official CUDA Toolkit installation guides](https://developer.nvidia.com/cuda-toolkit-archive) and vLLM's main [GPU installation documentation](../getting_started/installation/gpu/cuda.inc.md#troubleshooting) for troubleshooting. The `CMAKE_CUDA_COMPILER` variable in your `CMakeUserPresets.json` should also point to your `nvcc` binary.
-
-3. **Build Tools:** It is highly recommended to install `ccache` for fast rebuilds by caching compilation results (e.g., `sudo apt install ccache` or `conda install ccache`). Also, ensure the core build dependencies like `cmake` and `ninja` are installed. These are installable through `requirements/build.txt` or your system's package manager.
-
-    ```console
-    uv pip install -r requirements/build.txt --torch-backend=auto
-    ```
-
-## Setting up the CMake Build Environment
-
-The incremental build process is managed through CMake. You can configure your build settings using a `CMakeUserPresets.json` file at the root of the vLLM repository.
-
-### Generate `CMakeUserPresets.json` using the helper script
-
-To simplify the setup, vLLM provides a helper script that attempts to auto-detect your system's configuration (like CUDA path, Python environment, and CPU cores) and generates the `CMakeUserPresets.json` file for you.
-
-**Run the script:**
-
-Navigate to the root of your vLLM clone and execute the following command:
-
-```console
-python tools/generate_cmake_presets.py
-```
-
-The script will prompt you if it cannot automatically determine certain paths (e.g., `nvcc` or a specific Python executable for your vLLM development environment). Follow the on-screen prompts. If an existing `CMakeUserPresets.json` is found, the script will ask for confirmation before overwriting it.
-
-After running the script, a `CMakeUserPresets.json` file will be created in the root of your vLLM repository.
-
-### Example `CMakeUserPresets.json`
-
-Below is an example of what the generated `CMakeUserPresets.json` might look like. The script will tailor these values based on your system and any input you provide.
-
-```json
-{
-    "version": 6,
-    "cmakeMinimumRequired": {
-        "major": 3,
-        "minor": 26,
-        "patch": 1
-    },
-    "configurePresets": [
-        {
-            "name": "release",
-            "generator": "Ninja",
-            "binaryDir": "${sourceDir}/cmake-build-release",
-            "cacheVariables": {
-                "CMAKE_CUDA_COMPILER": "/usr/local/cuda/bin/nvcc",
-                "CMAKE_C_COMPILER_LAUNCHER": "ccache",
-                "CMAKE_CXX_COMPILER_LAUNCHER": "ccache",
-                "CMAKE_CUDA_COMPILER_LAUNCHER": "ccache",
-                "CMAKE_BUILD_TYPE": "Release",
-                "VLLM_PYTHON_EXECUTABLE": "/home/user/venvs/vllm/bin/python",
-                "CMAKE_INSTALL_PREFIX": "${sourceDir}",
-                "CMAKE_CUDA_FLAGS": "",
-                "NVCC_THREADS": "4",
-                "CMAKE_JOB_POOLS": "compile=32"
-            }
-        }
-    ],
-    "buildPresets": [
-        {
-            "name": "release",
-            "configurePreset": "release",
-            "jobs": 32
-        }
-    ]
-}
-```
-
-**What do the various configurations mean?**
-- `CMAKE_CUDA_COMPILER`: Path to your `nvcc` binary. The script attempts to find this automatically.
-- `CMAKE_C_COMPILER_LAUNCHER`, `CMAKE_CXX_COMPILER_LAUNCHER`, `CMAKE_CUDA_COMPILER_LAUNCHER`: Setting these to `ccache` (or `sccache`) significantly speeds up rebuilds by caching compilation results. Ensure `ccache` is installed (e.g., `sudo apt install ccache` or `conda install ccache`). The script sets these by default.
-- `VLLM_PYTHON_EXECUTABLE`: Path to the Python executable in your vLLM development environment. The script will prompt for this, defaulting to the current Python environment if suitable.
-- `CMAKE_INSTALL_PREFIX: "${sourceDir}"`: Specifies that the compiled components should be installed back into your vLLM source directory. This is crucial for the editable install, as it makes the newly built kernels immediately available to your Python environment.
-- `CMAKE_JOB_POOLS` and `jobs` in build presets: Control the parallelism of the build. The script sets these based on the number of CPU cores detected on your system.
-- `binaryDir`: Specifies where the build artifacts will be stored (e.g., `cmake-build-release`).
-
-## Building and Installing with CMake
-
-Once your `CMakeUserPresets.json` is configured:
-
-1. **Initialize the CMake build environment:**
-   This step configures the build system according to your chosen preset (e.g., `release`) and creates the build directory at `binaryDir`
-
-   ```console
-   cmake --preset release
-   ```
-
-2. **Build and install the vLLM components:**
-   This command compiles the code and installs the resulting binaries into your vLLM source directory, making them available to your editable Python installation.
-
-   ```console
-   cmake --build --preset release --target install
-   ```
-
-3. **Make changes and repeat!**
-    Now you start using your editable install of vLLM, testing and making changes as needed. If you need to build again to update based on changes, simply run the CMake command again to build only the affected files.
-
-    ```console
-    cmake --build --preset release --target install
-    ```
-
-## Verifying the Build
-
-After a successful build, you will find a populated build directory (e.g., `cmake-build-release/` if you used the `release` preset and the example configuration).
-
-```console
-> ls cmake-build-release/
-bin             cmake_install.cmake      _deps                                machete_generation.log
-build.ninja     CPackConfig.cmake        detect_cuda_compute_capabilities.cu  marlin_generation.log
-_C.abi3.so      CPackSourceConfig.cmake  detect_cuda_version.cc               _moe_C.abi3.so
-CMakeCache.txt  ctest                    _flashmla_C.abi3.so                  moe_marlin_generation.log
-CMakeFiles      cumem_allocator.abi3.so  install_local_manifest.txt           vllm-flash-attn
-```
-
-The `cmake --build ... --target install` command copies the compiled shared libraries (like `_C.abi3.so`, `_moe_C.abi3.so`, etc.) into the appropriate `vllm` package directory within your source tree. This updates your editable installation with the newly compiled kernels.
-
-## Additional Tips
-
-- **Adjust Parallelism:** Fine-tune the `CMAKE_JOB_POOLS` in `configurePresets` and `jobs` in `buildPresets` in your `CMakeUserPresets.json`. Too many jobs can overload systems with limited RAM or CPU cores, leading to slower builds or system instability. Too few won't fully utilize available resources.
-- **Clean Builds When Necessary:** If you encounter persistent or strange build errors, especially after significant changes or switching branches, consider removing the CMake build directory (e.g., `rm -rf cmake-build-release`) and re-running the `cmake --preset` and `cmake --build` commands.
-- **Specific Target Builds:** For even faster iterations when working on a specific module, you can sometimes build a specific target instead of the full `install` target, though `install` ensures all necessary components are updated in your Python environment. Refer to CMake documentation for more advanced target management.

docs/contributing/model/README.md

@@ -1,23 +1,21 @@
 ---
-title: Summary
+title: Adding a New Model
 ---
 [](){ #new-model }
 
-!!! important
-    Many decoder language models can now be automatically loaded using the [Transformers backend][transformers-backend] without having to implement them in vLLM. See if `vllm serve <model>` works first!
+This section provides more information on how to integrate a [PyTorch](https://pytorch.org/) model into vLLM.
 
-vLLM models are specialized [PyTorch](https://pytorch.org/) models that take advantage of various [features][compatibility-matrix] to optimize their performance.
+Contents:
 
-The complexity of integrating a model into vLLM depends heavily on the model's architecture.
-The process is considerably straightforward if the model shares a similar architecture with an existing model in vLLM.
-However, this can be more complex for models that include new operators (e.g., a new attention mechanism).
+- [Basic](basic.md)
+- [Registration](registration.md)
+- [Tests](tests.md)
+- [Multimodal](multimodal.md)
 
-Read through these pages for a step-by-step guide:
-
-- [Basic Model](basic.md)
-- [Registering a Model](registration.md)
-- [Unit Testing](tests.md)
-- [Multi-Modal Support](multimodal.md)
+!!! note
+    The complexity of adding a new model depends heavily on the model's architecture.
+    The process is considerably straightforward if the model shares a similar architecture with an existing model in vLLM.
+    However, for models that include new operators (e.g., a new attention mechanism), the process can be a bit more complex.
 
 !!! tip
     If you are encountering issues while integrating your model into vLLM, feel free to open a [GitHub issue](https://github.com/vllm-project/vllm/issues)

docs/contributing/model/basic.md

@@ -1,5 +1,5 @@
 ---
-title: Basic Model
+title: Implementing a Basic Model
 ---
 [](){ #new-model-basic }
 

docs/contributing/model/multimodal.md

@@ -538,13 +538,11 @@ return a schema of the tensors outputted by the HF processor that are related to
             prompt: str,
             mm_data: Mapping[str, object],
             mm_kwargs: Mapping[str, object],
-            tok_kwargs: Mapping[str, object],
         ) -> BatchFeature:
             processed_outputs = super()._call_hf_processor(
                 prompt=prompt,
                 mm_data=mm_data,
                 mm_kwargs=mm_kwargs,
-                tok_kwargs=tok_kwargs,
             )
 
             image_patches = processed_outputs.get("image_patches")
@@ -568,11 +566,6 @@ return a schema of the tensors outputted by the HF processor that are related to
         Our [actual code](gh-file:vllm/model_executor/models/fuyu.py) has special handling
         for text-only inputs to prevent unnecessary warnings from HF processor.
 
-    !!! note
-        The `_call_hf_processor` method specifies both `mm_kwargs` and `tok_kwargs` for
-        processing. `mm_kwargs` is used to both initialize and call the huggingface
-        processor, whereas `tok_kwargs` is only used to call the huggingface processor.
-
     This lets us override [_get_mm_fields_config][vllm.multimodal.processing.BaseMultiModalProcessor._get_mm_fields_config] as follows:
 
     ```python

docs/contributing/model/registration.md

@@ -1,5 +1,5 @@
 ---
-title: Registering a Model
+title: Registering a Model to vLLM
 ---
 [](){ #new-model-registration }
 

docs/contributing/model/tests.md

@@ -1,5 +1,5 @@
 ---
-title: Unit Testing
+title: Writing Unit Tests
 ---
 [](){ #new-model-tests }
 

docs/deployment/frameworks/helm.md

@@ -5,9 +5,9 @@ title: Helm
 
 A Helm chart to deploy vLLM for Kubernetes
 
-Helm is a package manager for Kubernetes. It helps automate the deployment of vLLM applications on Kubernetes. With Helm, you can deploy the same framework architecture with different configurations to multiple namespaces by overriding variable values.
+Helm is a package manager for Kubernetes. It will help you to deploy vLLM on k8s and automate the deployment of vLLM Kubernetes applications. With Helm, you can deploy the same framework architecture with different configurations to multiple namespaces by overriding variable values.
 
-This guide will walk you through the process of deploying vLLM with Helm, including the necessary prerequisites, steps for Helm installation and documentation on architecture and values file.
+This guide will walk you through the process of deploying vLLM with Helm, including the necessary prerequisites, steps for helm installation and documentation on architecture and values file.
 
 ## Prerequisites
 
@@ -16,23 +16,17 @@ Before you begin, ensure that you have the following:
 - A running Kubernetes cluster
 - NVIDIA Kubernetes Device Plugin (`k8s-device-plugin`): This can be found at [https://github.com/NVIDIA/k8s-device-plugin](https://github.com/NVIDIA/k8s-device-plugin)
 - Available GPU resources in your cluster
-- An S3 with the model which will be deployed
+- S3 with the model which will be deployed
 
 ## Installing the chart
 
 To install the chart with the release name `test-vllm`:
 
 ```bash
-helm upgrade --install --create-namespace \
-  --namespace=ns-vllm test-vllm . \
-  -f values.yaml \
-  --set secrets.s3endpoint=$ACCESS_POINT \
-  --set secrets.s3bucketname=$BUCKET \
-  --set secrets.s3accesskeyid=$ACCESS_KEY \
-  --set secrets.s3accesskey=$SECRET_KEY
+helm upgrade --install --create-namespace --namespace=ns-vllm test-vllm . -f values.yaml --set secrets.s3endpoint=$ACCESS_POINT --set secrets.s3bucketname=$BUCKET --set secrets.s3accesskeyid=$ACCESS_KEY --set secrets.s3accesskey=$SECRET_KEY
 ```
 
-## Uninstalling the chart
+## Uninstalling the Chart
 
 To uninstall the `test-vllm` deployment:
 
@@ -45,59 +39,57 @@ chart **including persistent volumes** and deletes the release.
 
 ## Architecture
 
-![helm deployment architecture](../../assets/deployment/architecture_helm_deployment.png)
+![](../../assets/deployment/architecture_helm_deployment.png)
 
 ## Values
 
-The following table describes configurable parameters of the chart in `values.yaml`:
-
-| Key | Type | Default | Description |
-|-----|------|---------|-------------|
-| autoscaling | object | {"enabled":false,"maxReplicas":100,"minReplicas":1,"targetCPUUtilizationPercentage":80} | Autoscaling configuration |
-| autoscaling.enabled | bool | false | Enable autoscaling |
-| autoscaling.maxReplicas | int | 100 | Maximum replicas |
-| autoscaling.minReplicas | int | 1 | Minimum replicas |
-| autoscaling.targetCPUUtilizationPercentage | int | 80 | Target CPU utilization for autoscaling |
-| configs | object | {} | Configmap |
-| containerPort | int | 8000 | Container port |
-| customObjects | list | [] | Custom Objects configuration |
-| deploymentStrategy | object | {} | Deployment strategy configuration |
-| externalConfigs | list | [] | External configuration |
-| extraContainers | list | [] | Additional containers configuration |
-| extraInit | object | {"pvcStorage":"1Gi","s3modelpath":"relative_s3_model_path/opt-125m", "awsEc2MetadataDisabled": true} | Additional configuration for the init container |
-| extraInit.pvcStorage | string | "1Gi" | Storage size of the s3 |
-| extraInit.s3modelpath | string | "relative_s3_model_path/opt-125m" | Path of the model on the s3 which hosts model weights and config files |
-| extraInit.awsEc2MetadataDisabled | boolean | true | Disables the use of the Amazon EC2 instance metadata service |
-| extraPorts | list | [] | Additional ports configuration |
-| gpuModels | list | ["TYPE_GPU_USED"] | Type of gpu used |
-| image | object | {"command":["vllm","serve","/data/","--served-model-name","opt-125m","--host","0.0.0.0","--port","8000"],"repository":"vllm/vllm-openai","tag":"latest"} | Image configuration |
-| image.command | list | ["vllm","serve","/data/","--served-model-name","opt-125m","--host","0.0.0.0","--port","8000"] | Container launch command |
-| image.repository | string | "vllm/vllm-openai" | Image repository |
-| image.tag | string | "latest" | Image tag |
-| livenessProbe | object | {"failureThreshold":3,"httpGet":{"path":"/health","port":8000},"initialDelaySeconds":15,"periodSeconds":10} | Liveness probe configuration |
-| livenessProbe.failureThreshold | int | 3 | Number of times after which if a probe fails in a row, Kubernetes considers that the overall check has failed: the container is not alive |
-| livenessProbe.httpGet | object | {"path":"/health","port":8000} | Configuration of the kubelet http request on the server |
-| livenessProbe.httpGet.path | string | "/health" | Path to access on the HTTP server |
-| livenessProbe.httpGet.port | int | 8000 | Name or number of the port to access on the container, on which the server is listening |
-| livenessProbe.initialDelaySeconds | int | 15 | Number of seconds after the container has started before liveness probe is initiated |
-| livenessProbe.periodSeconds | int | 10 | How often (in seconds) to perform the liveness probe |
-| maxUnavailablePodDisruptionBudget | string | "" | Disruption Budget Configuration |
-| readinessProbe | object | {"failureThreshold":3,"httpGet":{"path":"/health","port":8000},"initialDelaySeconds":5,"periodSeconds":5} | Readiness probe configuration |
-| readinessProbe.failureThreshold | int | 3 | Number of times after which if a probe fails in a row, Kubernetes considers that the overall check has failed: the container is not ready |
-| readinessProbe.httpGet | object | {"path":"/health","port":8000} | Configuration of the kubelet http request on the server |
-| readinessProbe.httpGet.path | string | "/health" | Path to access on the HTTP server |
-| readinessProbe.httpGet.port | int | 8000 | Name or number of the port to access on the container, on which the server is listening |
-| readinessProbe.initialDelaySeconds | int | 5 | Number of seconds after the container has started before readiness probe is initiated |
-| readinessProbe.periodSeconds | int | 5 | How often (in seconds) to perform the readiness probe |
-| replicaCount | int | 1 | Number of replicas |
-| resources | object | {"limits":{"cpu":4,"memory":"16Gi","nvidia.com/gpu":1},"requests":{"cpu":4,"memory":"16Gi","nvidia.com/gpu":1}} | Resource configuration |
-| resources.limits."nvidia.com/gpu" | int | 1 | Number of GPUs used |
-| resources.limits.cpu | int | 4 | Number of CPUs |
-| resources.limits.memory | string | "16Gi" | CPU memory configuration |
-| resources.requests."nvidia.com/gpu" | int | 1 | Number of GPUs used |
-| resources.requests.cpu | int | 4 | Number of CPUs |
-| resources.requests.memory | string | "16Gi" | CPU memory configuration |
-| secrets | object | {} | Secrets configuration |
-| serviceName | string | "" | Service name |
-| servicePort | int | 80 | Service port |
-| labels.environment | string | test | Environment name |
+| Key                                        | Type    | Default                                                                                                                                                  | Description                                                                                                                               |
+|--------------------------------------------|---------|----------------------------------------------------------------------------------------------------------------------------------------------------------|-------------------------------------------------------------------------------------------------------------------------------------------|
+| autoscaling                                | object  | {"enabled":false,"maxReplicas":100,"minReplicas":1,"targetCPUUtilizationPercentage":80}                                                                  | Autoscaling configuration                                                                                                                 |
+| autoscaling.enabled                        | bool    | false                                                                                                                                                    | Enable autoscaling                                                                                                                        |
+| autoscaling.maxReplicas                    | int     | 100                                                                                                                                                      | Maximum replicas                                                                                                                          |
+| autoscaling.minReplicas                    | int     | 1                                                                                                                                                        | Minimum replicas                                                                                                                          |
+| autoscaling.targetCPUUtilizationPercentage | int     | 80                                                                                                                                                       | Target CPU utilization for autoscaling                                                                                                    |
+| configs                                    | object  | {}                                                                                                                                                       | Configmap                                                                                                                                 |
+| containerPort                              | int     | 8000                                                                                                                                                     | Container port                                                                                                                            |
+| customObjects                              | list    | []                                                                                                                                                       | Custom Objects configuration                                                                                                              |
+| deploymentStrategy                         | object  | {}                                                                                                                                                       | Deployment strategy configuration                                                                                                         |
+| externalConfigs                            | list    | []                                                                                                                                                       | External configuration                                                                                                                    |
+| extraContainers                            | list    | []                                                                                                                                                       | Additional containers configuration                                                                                                       |
+| extraInit                                  | object  | {"pvcStorage":"1Gi","s3modelpath":"relative_s3_model_path/opt-125m", "awsEc2MetadataDisabled": true}                                                     | Additional configuration for the init container                                                                                           |
+| extraInit.pvcStorage                       | string  | "50Gi"                                                                                                                                                   | Storage size of the s3                                                                                                                    |
+| extraInit.s3modelpath                      | string  | "relative_s3_model_path/opt-125m"                                                                                                                        | Path of the model on the s3 which hosts model weights and config files                                                                    |
+| extraInit.awsEc2MetadataDisabled           | boolean | true                                                                                                                                                     | Disables the use of the Amazon EC2 instance metadata service                                                                              |
+| extraPorts                                 | list    | []                                                                                                                                                       | Additional ports configuration                                                                                                            |
+| gpuModels                                  | list    | ["TYPE_GPU_USED"]                                                                                                                                        | Type of gpu used                                                                                                                          |
+| image                                      | object  | {"command":["vllm","serve","/data/","--served-model-name","opt-125m","--host","0.0.0.0","--port","8000"],"repository":"vllm/vllm-openai","tag":"latest"} | Image configuration                                                                                                                       |
+| image.command                              | list    | ["vllm","serve","/data/","--served-model-name","opt-125m","--host","0.0.0.0","--port","8000"]                                                            | Container launch command                                                                                                                  |
+| image.repository                           | string  | "vllm/vllm-openai"                                                                                                                                       | Image repository                                                                                                                          |
+| image.tag                                  | string  | "latest"                                                                                                                                                 | Image tag                                                                                                                                 |
+| livenessProbe                              | object  | {"failureThreshold":3,"httpGet":{"path":"/health","port":8000},"initialDelaySeconds":15,"periodSeconds":10}                                              | Liveness probe configuration                                                                                                              |
+| livenessProbe.failureThreshold             | int     | 3                                                                                                                                                        | Number of times after which if a probe fails in a row, Kubernetes considers that the overall check has failed: the container is not alive |
+| livenessProbe.httpGet                      | object  | {"path":"/health","port":8000}                                                                                                                           | Configuration of the Kubelet http request on the server                                                                                   |
+| livenessProbe.httpGet.path                 | string  | "/health"                                                                                                                                                | Path to access on the HTTP server                                                                                                         |
+| livenessProbe.httpGet.port                 | int     | 8000                                                                                                                                                     | Name or number of the port to access on the container, on which the server is listening                                                   |
+| livenessProbe.initialDelaySeconds          | int     | 15                                                                                                                                                       | Number of seconds after the container has started before liveness probe is initiated                                                      |
+| livenessProbe.periodSeconds                | int     | 10                                                                                                                                                       | How often (in seconds) to perform the liveness probe                                                                                      |
+| maxUnavailablePodDisruptionBudget          | string  | ""                                                                                                                                                       | Disruption Budget Configuration                                                                                                           |
+| readinessProbe                             | object  | {"failureThreshold":3,"httpGet":{"path":"/health","port":8000},"initialDelaySeconds":5,"periodSeconds":5}                                                | Readiness probe configuration                                                                                                             |
+| readinessProbe.failureThreshold            | int     | 3                                                                                                                                                        | Number of times after which if a probe fails in a row, Kubernetes considers that the overall check has failed: the container is not ready |
+| readinessProbe.httpGet                     | object  | {"path":"/health","port":8000}                                                                                                                           | Configuration of the Kubelet http request on the server                                                                                   |
+| readinessProbe.httpGet.path                | string  | "/health"                                                                                                                                                | Path to access on the HTTP server                                                                                                         |
+| readinessProbe.httpGet.port                | int     | 8000                                                                                                                                                     | Name or number of the port to access on the container, on which the server is listening                                                   |
+| readinessProbe.initialDelaySeconds         | int     | 5                                                                                                                                                        | Number of seconds after the container has started before readiness probe is initiated                                                     |
+| readinessProbe.periodSeconds               | int     | 5                                                                                                                                                        | How often (in seconds) to perform the readiness probe                                                                                     |
+| replicaCount                               | int     | 1                                                                                                                                                        | Number of replicas                                                                                                                        |
+| resources                                  | object  | {"limits":{"cpu":4,"memory":"16Gi","nvidia.com/gpu":1},"requests":{"cpu":4,"memory":"16Gi","nvidia.com/gpu":1}}                                          | Resource configuration                                                                                                                    |
+| resources.limits."nvidia.com/gpu"          | int     | 1                                                                                                                                                        | Number of gpus used                                                                                                                       |
+| resources.limits.cpu                       | int     | 4                                                                                                                                                        | Number of CPUs                                                                                                                            |
+| resources.limits.memory                    | string  | "16Gi"                                                                                                                                                   | CPU memory configuration                                                                                                                  |
+| resources.requests."nvidia.com/gpu"        | int     | 1                                                                                                                                                        | Number of gpus used                                                                                                                       |
+| resources.requests.cpu                     | int     | 4                                                                                                                                                        | Number of CPUs                                                                                                                            |
+| resources.requests.memory                  | string  | "16Gi"                                                                                                                                                   | CPU memory configuration                                                                                                                  |
+| secrets                                    | object  | {}                                                                                                                                                       | Secrets configuration                                                                                                                     |
+| serviceName                                | string  | Service name                                                                                                                                             |                                                                                                                                           |
+| servicePort                                | int     | 80                                                                                                                                                       | Service port                                                                                                                              |
+| labels.environment                         | string  | test                                                                                                                                                     | Environment name                                                                                                                          |

docs/design/arch_overview.md

@@ -74,7 +74,7 @@ python -m vllm.entrypoints.openai.api_server --model <model>
 
 That code can be found in <gh-file:vllm/entrypoints/openai/api_server.py>.
 
-More details on the API server can be found in the [OpenAI-Compatible Server][serving-openai-compatible-server] document.
+More details on the API server can be found in the [OpenAI-Compatible Server][openai-compatible-server] document.
 
 ## LLM Engine
 

docs/design/v1/prefix_caching.md

@@ -117,8 +117,8 @@ There are two design points to highlight:
 
 1. We allocate all KVCacheBlock when initializing the KV cache manager to be a block pool. This avoids Python object creation overheads and can easily track all blocks all the time.  
 2. We introduce doubly linked list pointers directly in the KVCacheBlock, so that we could construct a free queue directly. This gives us two benefits:  
-    1. We could have O(1) complexity moving elements in the middle to the tail.  
-    2. We could avoid introducing another Python queue (e.g., `deque`) which has a wrapper to the elements.
+   1. We could have O(1) complexity moving elements in the middle to the tail.  
+   2. We could avoid introducing another Python queue (e.g., `deque`) which has a wrapper to the elements.
 
 As a result, we will have the following components when the KV cache manager is initialized:
 
@@ -135,19 +135,19 @@ As a result, we will have the following components when the KV cache manager is
 
 **New request:** Workflow for the scheduler to schedule a new request with KV cache block allocation:
 
-1. The scheduler calls `kv_cache_manager.get_computed_blocks()` to get a sequence of blocks that have already been computed. This is done by hashing the prompt tokens in the request and looking up cache blocks.  
+1. The scheduler calls `kv_cache_manager.get_computed_blocks()` to get a sequence of blocks that have already been computed. This is done by hashing the prompt tokens in the request and looking up Cache Blocks.  
 2. The scheduler calls `kv_cache_manager.allocate_slots()`. It does the following steps:  
-    1. Compute the number of new required blocks, and return if there are no sufficient blocks to allocate.  
-    2. “Touch” the computed blocks. It increases the reference count of the computed block by one, and removes the block from the free queue if the block wasn’t used by other requests. This is to avoid these computed blocks being evicted. See the example in the next section for illustration.  
-    3. Allocate new blocks by popping the heads of the free queue. If the head block is a cached block, this also “evicts” the block so that no other requests can reuse it anymore from now on.  
-    4. If an allocated block is already full of tokens, we immediately add it to the cache block, so that the block can be reused by other requests in the same batch.
+   1. Compute the number of new required blocks, and return if there are no sufficient blocks to allocate.  
+   2. “Touch” the computed blocks. It increases the reference count of the computed block by one, and removes the block from the free queue if the block wasn’t used by other requests. This is to avoid these computed blocks being evicted. See the example in the next section for illustration.  
+   3. Allocate new blocks by popping the heads of the free queue. If the head block is a cached block, this also “evicts” the block so that no other requests can reuse it anymore from now on.  
+   4. If an allocated block is already full of tokens, we immediately add it to the Cache Block, so that the block can be reused by other requests in the same batch.
 
 **Running request:** Workflow for the scheduler to schedule a running request with KV cache block allocation:
 
 1. The scheduler calls `kv_cache_manager.allocate_slots()`. It does the following steps:  
-    1. Compute the number of new required blocks, and return if there are no sufficient blocks to allocate.  
-    2. Allocate new blocks by popping the heads of the free queue. If the head block is a cached block, this also “evicts” the block so that no other requests can reuse it anymore from now on.  
-    3. Append token IDs to the slots in existing blocks as well as the new blocks. If a block is full, we add it to the cache block to cache it.
+   1. Compute the number of new required blocks, and return if there are no sufficient blocks to allocate.  
+   2. Allocate new blocks by popping the heads of the free queue. If the head block is a cached block, this also “evicts” the block so that no other requests can reuse it anymore from now on.  
+   3. Append token IDs to the slots in existing blocks as well as the new blocks. If a block is full, we add it to the Cache Block to cache it.
 
 **Duplicated blocks**  
 Assuming block size is 4 and you send a request (Request 1\) with prompt ABCDEF and decoding length 3:
@@ -199,7 +199,7 @@ When a request is finished, we free all its blocks if no other requests are usin
 When the head block (least recently used block) of the free queue is cached, we have to evict the block to prevent it from being used by other requests. Specifically, eviction involves the following steps:
 
 1. Pop the block from the head of the free queue. This is the LRU block to be evicted.  
-2. Remove the block ID from the cache block.  
+2. Remove the block ID from the Cache Block.  
 3. Remove the block hash.
 
 ## Example

docs/features/compatibility_matrix.md

@@ -59,23 +59,23 @@ th:not(:first-child) {
 
 ## Feature x Hardware
 
-| Feature                                                   | Volta               | Turing    | Ampere    | Ada    | Hopper     | CPU                | AMD    | TPU |
-|-----------------------------------------------------------|---------------------|-----------|-----------|--------|------------|--------------------|--------|-----|
-| [CP][chunked-prefill]                                     | [❌](gh-issue:2729) | ✅        | ✅        | ✅     | ✅        | ✅                  | ✅     | ✅ |
-| [APC][automatic-prefix-caching]                           | [❌](gh-issue:3687) | ✅        | ✅        | ✅     | ✅        | ✅                  | ✅     | ✅ |
-| [LoRA][lora-adapter]                                      | ✅                  | ✅        | ✅        | ✅     | ✅        | ✅                  | ✅     | ✅ |
-| <abbr title="Prompt Adapter">prmpt adptr</abbr>           | ✅                  | ✅        | ✅        | ✅     | ✅        | [❌](gh-issue:8475) | ✅     | ❌ |
-| [SD][spec-decode]                                         | ✅                  | ✅        | ✅        | ✅     | ✅        | ✅                  | ✅     | ❌ |
-| CUDA graph                                                | ✅                  | ✅        | ✅        | ✅     | ✅        | ❌                  | ✅     | ❌ |
-| <abbr title="Pooling Models">pooling</abbr>               | ✅                  | ✅        | ✅        | ✅     | ✅        | ✅                  | ❔     | ❌ |
-| <abbr title="Encoder-Decoder Models">enc-dec</abbr>       | ✅                  | ✅        | ✅        | ✅     | ✅        | ✅                  | ❌     | ❌ |
-| <abbr title="Multimodal Inputs">mm</abbr>                 | ✅                  | ✅        | ✅        | ✅     | ✅        | ✅                  | ✅     | ❌ |
-| <abbr title="Logprobs">logP</abbr>                        | ✅                  | ✅        | ✅        | ✅     | ✅        | ✅                  | ✅     | ❌ |
-| <abbr title="Prompt Logprobs">prmpt logP</abbr>           | ✅                  | ✅        | ✅        | ✅     | ✅        | ✅                  | ✅     | ❌ |
-| <abbr title="Async Output Processing">async output</abbr> | ✅                  | ✅        | ✅        | ✅     | ✅        | ❌                  | ❌     | ❌ |
-| multi-step                                                | ✅                  | ✅        | ✅        | ✅     | ✅        | [❌](gh-issue:8477) | ✅     | ❌ |
-| best-of                                                   | ✅                  | ✅        | ✅        | ✅     | ✅        | ✅                  | ✅     | ❌ |
-| beam-search                                               | ✅                  | ✅        | ✅        | ✅     | ✅        | ✅                  | ✅     | ❌ |
+| Feature                                                   | Volta              | Turing   | Ampere   | Ada   | Hopper   | CPU                | AMD   |
+|-----------------------------------------------------------|--------------------|----------|----------|-------|----------|--------------------|-------|
+| [CP][chunked-prefill]                                     | [❌](gh-issue:2729) | ✅        | ✅        | ✅     | ✅        | ✅                  | ✅     |
+| [APC][automatic-prefix-caching]                           | [❌](gh-issue:3687) | ✅        | ✅        | ✅     | ✅        | ✅                  | ✅     |
+| [LoRA][lora-adapter]                                      | ✅                  | ✅        | ✅        | ✅     | ✅        | ✅                  | ✅     |
+| <abbr title="Prompt Adapter">prmpt adptr</abbr>           | ✅                  | ✅        | ✅        | ✅     | ✅        | [❌](gh-issue:8475) | ✅     |
+| [SD][spec-decode]                                         | ✅                  | ✅        | ✅        | ✅     | ✅        | ✅                  | ✅     |
+| CUDA graph                                                | ✅                  | ✅        | ✅        | ✅     | ✅        | ❌                  | ✅     |
+| <abbr title="Pooling Models">pooling</abbr>               | ✅                  | ✅        | ✅        | ✅     | ✅        | ✅                  | ❔     |
+| <abbr title="Encoder-Decoder Models">enc-dec</abbr>       | ✅                  | ✅        | ✅        | ✅     | ✅        | ✅                  | ❌     |
+| <abbr title="Multimodal Inputs">mm</abbr>                 | ✅                  | ✅        | ✅        | ✅     | ✅        | ✅                  | ✅     |
+| <abbr title="Logprobs">logP</abbr>                        | ✅                  | ✅        | ✅        | ✅     | ✅        | ✅                  | ✅     |
+| <abbr title="Prompt Logprobs">prmpt logP</abbr>           | ✅                  | ✅        | ✅        | ✅     | ✅        | ✅                  | ✅     |
+| <abbr title="Async Output Processing">async output</abbr> | ✅                  | ✅        | ✅        | ✅     | ✅        | ❌                  | ❌     |
+| multi-step                                                | ✅                  | ✅        | ✅        | ✅     | ✅        | [❌](gh-issue:8477) | ✅     |
+| best-of                                                   | ✅                  | ✅        | ✅        | ✅     | ✅        | ✅                  | ✅     |
+| beam-search                                               | ✅                  | ✅        | ✅        | ✅     | ✅        | ✅                  | ✅     |
 
 !!! note
     Please refer to [Feature support through NxD Inference backend][feature-support-through-nxd-inference-backend] for features supported on AWS Neuron hardware

docs/features/structured_outputs.md

@@ -21,7 +21,7 @@ The following parameters are supported, which must be added as extra parameters:
 - `guided_grammar`: the output will follow the context free grammar.
 - `structural_tag`: Follow a JSON schema within a set of specified tags within the generated text.
 
-You can see the complete list of supported parameters on the [OpenAI-Compatible Server][serving-openai-compatible-server] page.
+You can see the complete list of supported parameters on the [OpenAI-Compatible Server][openai-compatible-server] page.
 
 Structured outputs are supported by default in the OpenAI-Compatible Server. You
 may choose to specify the backend to use by setting the

docs/features/tool_calling.md

@@ -53,7 +53,7 @@ Next, make a request to the model that should result in it using the available t
     tool_call = response.choices[0].message.tool_calls[0].function
     print(f"Function called: {tool_call.name}")
     print(f"Arguments: {tool_call.arguments}")
-    print(f"Result: {tool_functions[tool_call.name](**json.loads(tool_call.arguments))}")
+    print(f"Result: {get_weather(**json.loads(tool_call.arguments))}")
     ```
 
 Example output:
@@ -99,14 +99,6 @@ vLLM supports the `tool_choice='required'` option in the chat completion API. Si
 
 When tool_choice='required' is set, the model is guaranteed to generate one or more tool calls based on the specified tool list in the `tools` parameter. The number of tool calls depends on the user's query. The output format strictly follows the schema defined in the `tools` parameter.
 
-## None Function Calling
-
-vLLM supports the `tool_choice='none'` option in the chat completion API. When this option is set, the model will not generate any tool calls and will respond with regular text content only, even if tools are defined in the request.
-
-By default, when `tool_choice='none'` is specified, vLLM excludes tool definitions from the prompt to optimize context usage. To include tool definitions even with `tool_choice='none'`, use the `--expand-tools-even-if-tool-choice-none` option.
-
-Note: This behavior will change in v0.10.0, where tool definitions will be included by default even with `tool_choice='none'`.
-
 ## Automatic Function Calling
 
 To enable this feature, you should set the following flags:

docs/getting_started/installation/cpu.md

@@ -118,7 +118,6 @@ vLLM CPU backend supports the following vLLM features:
 - `VLLM_CPU_OMP_THREADS_BIND`: specify the CPU cores dedicated to the OpenMP threads. For example, `VLLM_CPU_OMP_THREADS_BIND=0-31` means there will be 32 OpenMP threads bound on 0-31 CPU cores. `VLLM_CPU_OMP_THREADS_BIND=0-31|32-63` means there will be 2 tensor parallel processes, 32 OpenMP threads of rank0 are bound on 0-31 CPU cores, and the OpenMP threads of rank1 are bound on 32-63 CPU cores. By setting to `auto`, the OpenMP threads of each rank are bound to the CPU cores in each NUMA node. By setting to `all`, the OpenMP threads of each rank uses all CPU cores available on the system. Default value is `auto`.
 - `VLLM_CPU_NUM_OF_RESERVED_CPU`: specify the number of CPU cores which are not dedicated to the OpenMP threads for each rank. The variable only takes effect when VLLM_CPU_OMP_THREADS_BIND is set to `auto`. Default value is `0`.
 - `VLLM_CPU_MOE_PREPACK`: whether to use prepack for MoE layer. This will be passed to `ipex.llm.modules.GatedMLPMOE`. Default is `1` (True). On unsupported CPUs, you might need to set this to `0` (False).
-- `VLLM_CPU_SGL_KERNEL` (Experimental): whether to use small-batch optimized kernels for linear layer and MoE layer, especially for low-latency requirements like online serving. The kernels require AMX instruction set, BFloat16 weight type and weight shapes divisible by 32. Default is `0` (False).
 
 ## Performance tips
 

docs/getting_started/installation/google_tpu.md

@@ -58,9 +58,9 @@ assigned to your Google Cloud project for your immediate exclusive use.
 
 For more information about using TPUs with GKE, see:
 
-- [About TPUs in GKE](https://cloud.google.com/kubernetes-engine/docs/concepts/tpus)
-- [Deploy TPU workloads in GKE Standard](https://cloud.google.com/kubernetes-engine/docs/how-to/tpus)
-- [Plan for TPUs in GKE](https://cloud.google.com/kubernetes-engine/docs/concepts/plan-tpus)
+- <https://cloud.google.com/kubernetes-engine/docs/how-to/tpus>
+- <https://cloud.google.com/kubernetes-engine/docs/concepts/tpus>
+- <https://cloud.google.com/kubernetes-engine/docs/concepts/plan-tpus>
 
 ## Configure a new environment
 

docs/getting_started/installation/gpu/cuda.inc.md

@@ -151,9 +151,6 @@ pip install -e .
     [sccache](https://github.com/mozilla/sccache) works similarly to `ccache`, but has the capability to utilize caching in remote storage environments.
     The following environment variables can be set to configure the vLLM `sccache` remote: `SCCACHE_BUCKET=vllm-build-sccache SCCACHE_REGION=us-west-2 SCCACHE_S3_NO_CREDENTIALS=1`. We also recommend setting `SCCACHE_IDLE_TIMEOUT=0`.
 
-!!! note "Faster Kernel Development"
-    For frequent C++/CUDA kernel changes, after the initial `pip install -e .` setup, consider using the [Incremental Compilation Workflow](../../contributing/incremental_build.md) for significantly faster rebuilds of only the modified kernel code.
-
 ##### Use an existing PyTorch installation
 
 There are scenarios where the PyTorch dependency cannot be easily installed via pip, e.g.:

docs/getting_started/installation/gpu/xpu.inc.md

@@ -22,7 +22,7 @@ Currently, there are no pre-built XPU wheels.
 # --8<-- [end:pre-built-wheels]
 # --8<-- [start:build-wheel-from-source]
 
-- First, install required [driver](https://dgpu-docs.intel.com/driver/installation.html#installing-gpu-drivers) and [Intel OneAPI](https://www.intel.com/content/www/us/en/developer/tools/oneapi/base-toolkit.html) 2025.0 or later.
+- First, install required driver and Intel OneAPI 2025.0 or later.
 - Second, install Python packages for vLLM XPU backend building:
 
 ```bash

docs/getting_started/installation/intel_gaudi.md

@@ -110,7 +110,7 @@ docker run \
 ### Supported features
 
 - [Offline inference][offline-inference]
-- Online serving via [OpenAI-Compatible Server][serving-openai-compatible-server]
+- Online serving via [OpenAI-Compatible Server][openai-compatible-server]
 - HPU autodetection - no need to manually select device within vLLM
 - Paged KV cache with algorithms enabled for Intel Gaudi accelerators
 - Custom Intel Gaudi implementations of Paged Attention, KV cache ops,

docs/mkdocs/javascript/slack_and_forum.js

@@ -1,56 +0,0 @@
-/**
- * slack_and_forum.js
- *
- * Adds a custom Slack and Forum button to the MkDocs Material header.
- *
- */
-
-window.addEventListener('DOMContentLoaded', () => {
-  const headerInner = document.querySelector('.md-header__inner');
-
-  if (headerInner) {
-    const slackButton = document.createElement('button');
-    slackButton.className = 'slack-button';
-    slackButton.title = 'Join us on Slack';
-    slackButton.style.border = 'none';
-    slackButton.style.background = 'transparent';
-    slackButton.style.cursor = 'pointer';
-
-    slackButton.innerHTML = `
-      <img src="https://a.slack-edge.com/80588/marketing/img/icons/icon_slack_hash_colored.png" 
-           style="height: 1.1rem;" 
-           alt="Slack">
-    `;
-
-    slackButton.addEventListener('click', () => {
-      window.open('https://slack.vllm.ai', '_blank', 'noopener');
-    });
-
-    const forumButton = document.createElement('button');
-    forumButton.className = 'forum-button';
-    forumButton.title = 'Join the Forum';
-    forumButton.style.border = 'none';
-    forumButton.style.background = 'transparent';
-    forumButton.style.cursor = 'pointer';
-
-    forumButton.innerHTML = `
-      <svg
-        xmlns="http://www.w3.org/2000/svg"
-        viewBox="0 -960 960 960"
-        fill="currentColor"
-      >
-        <path d="M817.85-198.15 698.46-317.54H320q-24.48 0-41.47-16.99T261.54-376v-11.69h424.61q25.39 0 43.47-18.08 18.07-18.08 18.07-43.46v-268.92h11.69q24.48 0 41.47 16.99 17 16.99 17 41.47v461.54ZM179.08-434.69l66.84-66.85h363.31q10.77 0 17.69-6.92 6.93-6.92 6.93-17.69v-246.77q0-10.77-6.93-17.7-6.92-6.92-17.69-6.92H203.69q-10.77 0-17.69 6.92-6.92 6.93-6.92 17.7v338.23Zm-36.93 89.46v-427.69q0-25.39 18.08-43.46 18.08-18.08 43.46-18.08h405.54q25.39 0 43.46 18.08 18.08 18.07 18.08 43.46v246.77q0 25.38-18.08 43.46-18.07 18.07-43.46 18.07H261.54L142.15-345.23Zm36.93-180.92V-797.54v271.39Z"/>
-      </svg>
-    `;
-
-    forumButton.addEventListener('click', () => {
-      window.open('https://discuss.vllm.ai/', '_blank', 'noopener');
-    });
-
-    const githubSource = document.querySelector('.md-header__source');
-    if (githubSource) {
-      githubSource.parentNode.insertBefore(slackButton, githubSource.nextSibling);
-      githubSource.parentNode.insertBefore(forumButton, slackButton.nextSibling);
-    }
-  }
-});

docs/mkdocs/stylesheets/extra.css

@@ -108,38 +108,3 @@ body[data-md-color-scheme="slate"] .md-nav__item--section > label.md-nav__link .
 .md-content__button-wrapper a:hover {
   color: var(--md-accent-fg-color);
 }
-
-/* Slack and Forum css */
-.slack-button, 
-.forum-button {
-  display: inline-flex;
-  align-items: center;
-  justify-content: center;
-  margin-left: 0.4rem;
-  height: 24px;
-}
-
-.slack-button img {
-  height: 18px;
-  filter: none !important;
-}
-
-.slack-button:hover,
-.forum-button:hover {
-  opacity: 0.7;
-}
-
-.forum-button svg {
-  height: 28px;
-  opacity: 0.9;
-  transform: translateY(2px);
-}
-
-/* For logo css */
-[data-md-color-scheme="default"] .logo-dark {
-  display: none;
-}
-
-[data-md-color-scheme="slate"] .logo-light {
-  display: none;
-}

docs/models/generative_models.md

@@ -134,7 +134,7 @@ outputs = llm.chat(conversation, chat_template=custom_template)
 
 ## Online Serving
 
-Our [OpenAI-Compatible Server][serving-openai-compatible-server] provides endpoints that correspond to the offline APIs:
+Our [OpenAI-Compatible Server][openai-compatible-server] provides endpoints that correspond to the offline APIs:
 
 - [Completions API][completions-api] is similar to `LLM.generate` but only accepts text.
 - [Chat API][chat-api]  is similar to `LLM.chat`, accepting both text and [multi-modal inputs][multimodal-inputs] for models with a chat template.

docs/models/hardware_supported_models/tpu.md

@@ -1,36 +0,0 @@
----
-title: TPU
----
-[](){ #tpu-supported-models }
-
-# TPU Supported Models
-## Text-only Language Models
-
-| Model                                               | Architecture                   | Supported |
-|-----------------------------------------------------|--------------------------------|-----------|
-| mistralai/Mixtral-8x7B-Instruct-v0.1                | MixtralForCausalLM             | 🟨 |
-| mistralai/Mistral-Small-24B-Instruct-2501           | MistralForCausalLM             | ✅ |
-| mistralai/Codestral-22B-v0.1                        | MistralForCausalLM             | ✅ |
-| mistralai/Mixtral-8x22B-Instruct-v0.1               | MixtralForCausalLM             | ❌ |
-| meta-llama/Llama-3.3-70B-Instruct                   | LlamaForCausalLM               | ✅ |
-| meta-llama/Llama-3.1-8B-Instruct                    | LlamaForCausalLM               | ✅ |
-| meta-llama/Llama-3.1-70B-Instruct                   | LlamaForCausalLM               | ✅ |
-| meta-llama/Llama-4-*                                | Llama4ForConditionalGeneration | ❌ |
-| microsoft/Phi-3-mini-128k-instruct                  | Phi3ForCausalLM                | 🟨 |
-| microsoft/phi-4                                     | Phi3ForCausalLM                | ❌ |
-| google/gemma-3-27b-it                               | Gemma3ForConditionalGeneration | 🟨 |
-| google/gemma-3-4b-it                                | Gemma3ForConditionalGeneration | ❌ |
-| deepseek-ai/DeepSeek-R1                             | DeepseekV3ForCausalLM          | ❌ |
-| deepseek-ai/DeepSeek-V3                             | DeepseekV3ForCausalLM          | ❌ |
-| RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w8a8  | LlamaForCausalLM               | ✅ |
-| RedHatAI/Meta-Llama-3.1-70B-Instruct-quantized.w8a8 | LlamaForCausalLM               | ✅ |
-| Qwen/Qwen3-8B                                       | Qwen3ForCausalLM               | ✅ |
-| Qwen/Qwen3-32B                                      | Qwen3ForCausalLM               | ✅ |
-| Qwen/Qwen2.5-7B-Instruct                            | Qwen2ForCausalLM               | ✅ |
-| Qwen/Qwen2.5-32B                                    | Qwen2ForCausalLM               | ✅ |
-| Qwen/Qwen2.5-14B-Instruct                           | Qwen2ForCausalLM               | ✅ |
-| Qwen/Qwen2.5-1.5B-Instruct                          | Qwen2ForCausalLM               | 🟨 |
-
-✅ Runs and optimized.  
-🟨 Runs and correct but not optimized to green yet.  
-❌ Does not pass accuracy test or does not run.  

docs/models/pooling_models.md

@@ -113,7 +113,7 @@ A code example can be found here: <gh-file:examples/offline_inference/basic/scor
 
 ## Online Serving
 
-Our [OpenAI-Compatible Server][serving-openai-compatible-server] provides endpoints that correspond to the offline APIs:
+Our [OpenAI-Compatible Server][openai-compatible-server] provides endpoints that correspond to the offline APIs:
 
 - [Pooling API][pooling-api] is similar to `LLM.encode`, being applicable to all types of pooling models.
 - [Embeddings API][embeddings-api] is similar to `LLM.embed`, accepting both text and [multi-modal inputs][multimodal-inputs] for embedding models.

docs/models/supported_models.md

@@ -34,7 +34,7 @@ llm.apply_model(lambda model: print(type(model)))
 If it is `TransformersForCausalLM` then it means it's based on Transformers!
 
 !!! tip
-    You can force the use of `TransformersForCausalLM` by setting `model_impl="transformers"` for [offline-inference][offline-inference] or `--model-impl transformers` for the [openai-compatible-server][serving-openai-compatible-server].
+    You can force the use of `TransformersForCausalLM` by setting `model_impl="transformers"` for [offline-inference][offline-inference] or `--model-impl transformers` for the [openai-compatible-server][openai-compatible-server].
 
 !!! note
     vLLM may not fully optimise the Transformers implementation so you may see degraded performance if comparing a native model to a Transformers model in vLLM.
@@ -53,8 +53,8 @@ For a model to be compatible with the Transformers backend for vLLM it must:
 
 If the compatible model is:
 
-- on the Hugging Face Model Hub, simply set `trust_remote_code=True` for [offline-inference][offline-inference] or `--trust-remote-code` for the [openai-compatible-server][serving-openai-compatible-server].
-- in a local directory, simply pass directory path to `model=<MODEL_DIR>` for [offline-inference][offline-inference] or `vllm serve <MODEL_DIR>` for the [openai-compatible-server][serving-openai-compatible-server].
+- on the Hugging Face Model Hub, simply set `trust_remote_code=True` for [offline-inference][offline-inference] or `--trust-remote-code` for the [openai-compatible-server][openai-compatible-server].
+- in a local directory, simply pass directory path to `model=<MODEL_DIR>` for [offline-inference][offline-inference] or `vllm serve <MODEL_DIR>` for the [openai-compatible-server][openai-compatible-server].
 
 This means that, with the Transformers backend for vLLM, new models can be used before they are officially supported in Transformers or vLLM!
 
@@ -329,9 +329,6 @@ Specified using `--task generate`.
 | `DeepseekForCausalLM`                             | DeepSeek                                            | `deepseek-ai/deepseek-llm-67b-base`, `deepseek-ai/deepseek-llm-7b-chat` etc.                                                                                                 |                        | ✅︎                          | ✅︎                     |
 | `DeepseekV2ForCausalLM`                           | DeepSeek-V2                                         | `deepseek-ai/DeepSeek-V2`, `deepseek-ai/DeepSeek-V2-Chat` etc.                                                                                                               |                        | ✅︎                          | ✅︎                     |
 | `DeepseekV3ForCausalLM`                           | DeepSeek-V3                                         | `deepseek-ai/DeepSeek-V3-Base`, `deepseek-ai/DeepSeek-V3` etc.                                                                                                               |                        | ✅︎                          | ✅︎                     |
-| `Dots1ForCausalLM`                                | dots.llm1                                           | `rednote-hilab/dots.llm1.base`, `rednote-hilab/dots.llm1.inst` etc.                                                                                                               |                        | ✅︎                          | ✅︎                     |
-| `Ernie4_5_ForCausalLM`                            | Ernie4.5                                            | `baidu/ERNIE-4.5-0.3B-PT`,etc.                                                                                                               |                        | ✅︎                          | ✅︎                     |
-| `Ernie4_5_MoeForCausalLM`                         | Ernie4.5MoE                                         | `baidu/ERNIE-4.5-21B-A3B-PT`,  `baidu/ERNIE-4.5-300B-A47B-PT`, etc.                                                                                                               |                        | ✅︎                          | ✅︎                     |
 | `ExaoneForCausalLM`                               | EXAONE-3                                            | `LGAI-EXAONE/EXAONE-3.0-7.8B-Instruct`, etc.                                                                                                                                 | ✅︎                     | ✅︎                          | ✅︎                     |
 | `FalconForCausalLM`                               | Falcon                                              | `tiiuae/falcon-7b`, `tiiuae/falcon-40b`, `tiiuae/falcon-rw-7b`, etc.                                                                                                         |                        | ✅︎                          | ✅︎                     |
 | `FalconMambaForCausalLM`                          | FalconMamba                                         | `tiiuae/falcon-mamba-7b`, `tiiuae/falcon-mamba-7b-instruct`, etc.                                                                                                            |                        | ✅︎                          | ✅︎                     |
@@ -339,7 +336,6 @@ Specified using `--task generate`.
 | `GemmaForCausalLM`                                | Gemma                                               | `google/gemma-2b`, `google/gemma-1.1-2b-it`, etc.                                                                                                                            | ✅︎                     | ✅︎                          | ✅︎                     |
 | `Gemma2ForCausalLM`                               | Gemma 2                                             | `google/gemma-2-9b`, `google/gemma-2-27b`, etc.                                                                                                                              | ✅︎                     | ✅︎                          | ✅︎                     |
 | `Gemma3ForCausalLM`                               | Gemma 3                                             | `google/gemma-3-1b-it`, etc.                                                                                                                                                 | ✅︎                     | ✅︎                          | ✅︎                     |
-| `Gemma3nForConditionalGeneration`                  | Gemma 3n                                             | `google/gemma-3n-E2B-it`, `google/gemma-3n-E4B-it`, etc.                                                                                                                                                 |                      |                           | ✅︎                     |
 | `GlmForCausalLM`                                  | GLM-4                                               | `THUDM/glm-4-9b-chat-hf`, etc.                                                                                                                                               | ✅︎                     | ✅︎                          | ✅︎                     |
 | `Glm4ForCausalLM`                                 | GLM-4-0414                                          | `THUDM/GLM-4-32B-0414`, etc.                                                                                                                                                 | ✅︎                     | ✅︎                          | ✅︎                     |
 | `GPT2LMHeadModel`                                 | GPT-2                                               | `gpt2`, `gpt2-xl`, etc.                                                                                                                                                      |                        | ✅︎                          | ✅︎                     |
@@ -352,7 +348,6 @@ Specified using `--task generate`.
 | `GraniteMoeSharedForCausalLM`                     | Granite MoE Shared                                  | `ibm-research/moe-7b-1b-active-shared-experts` (test model)                                                                                                                  | ✅︎                     | ✅︎                          | ✅︎                     |
 | `GritLM`                                          | GritLM                                              | `parasail-ai/GritLM-7B-vllm`.                                                                                                                                                | ✅︎                     | ✅︎                          |                       |
 | `Grok1ModelForCausalLM`                           | Grok1                                               | `hpcai-tech/grok-1`.                                                                                                                                                         | ✅︎                     | ✅︎                          | ✅︎                     |
-| `HunYuanMoEV1ForCausalLM`                         | Hunyuan-80B-A13B                                    | `tencent/Hunyuan-A13B-Instruct`,  `tencent/Hunyuan-A13B-Pretrain`, `tencent/Hunyuan-A13B-Instruct-FP8`etc.                                                                  |                        |                             | ✅︎                     |
 | `InternLMForCausalLM`                             | InternLM                                            | `internlm/internlm-7b`, `internlm/internlm-chat-7b`, etc.                                                                                                                    | ✅︎                     | ✅︎                          | ✅︎                     |
 | `InternLM2ForCausalLM`                            | InternLM2                                           | `internlm/internlm2-7b`, `internlm/internlm2-chat-7b`, etc.                                                                                                                  | ✅︎                     | ✅︎                          | ✅︎                     |
 | `InternLM3ForCausalLM`                            | InternLM3                                           | `internlm/internlm3-8b-instruct`, etc.                                                                                                                                       | ✅︎                     | ✅︎                          | ✅︎                     |
@@ -390,16 +385,13 @@ Specified using `--task generate`.
 | `TeleChat2ForCausalLM`                            | TeleChat2                                           | `Tele-AI/TeleChat2-3B`, `Tele-AI/TeleChat2-7B`, `Tele-AI/TeleChat2-35B`, etc.                                                                                                | ✅︎                     | ✅︎                          | ✅︎                     |
 | `TeleFLMForCausalLM`                              | TeleFLM                                             | `CofeAI/FLM-2-52B-Instruct-2407`, `CofeAI/Tele-FLM`, etc.                                                                                                                    | ✅︎                     | ✅︎                          | ✅︎                     |
 | `XverseForCausalLM`                               | XVERSE                                              | `xverse/XVERSE-7B-Chat`, `xverse/XVERSE-13B-Chat`, `xverse/XVERSE-65B-Chat`, etc.                                                                                            | ✅︎                     | ✅︎                          | ✅︎                     |
-| `MiniMaxM1ForCausalLM`                        | MiniMax-Text                                        | `MiniMaxAI/MiniMax-M1-40k`, `MiniMaxAI/MiniMax-M1-80k`etc.                                                                                                                       |                        |                             |                       |
+| `MiniMaxM1ForCausalLM`                        | MiniMax-Text                                        | `MiniMaxAI/MiniMax-M1-40k`, `MiniMaxAI/MiniMax-M1-80k`etc.                                                                                                                                  |                        |                             |                       |
 | `MiniMaxText01ForCausalLM`                        | MiniMax-Text                                        | `MiniMaxAI/MiniMax-Text-01`, etc.                                                                                                                                            |                        |                             |                       |
 | `Zamba2ForCausalLM`                               | Zamba2                                              | `Zyphra/Zamba2-7B-instruct`, `Zyphra/Zamba2-2.7B-instruct`, `Zyphra/Zamba2-1.2B-instruct`, etc.                                                                              |                        |                             |                       |
 
 !!! note
     Currently, the ROCm version of vLLM supports Mistral and Mixtral only for context lengths up to 4096.
 
-!!! note
-    Only text inputs are currently supported for `Gemma3nForConditionalGeneration`. To use this model, please upgrade Hugging Face Transformers to version 4.53.0.
-
 ### Pooling Models
 
 See [this page](./pooling_models.md) for more information on how to use pooling models.
@@ -435,7 +427,7 @@ Specified using `--task embed`.
     See [relevant issue on HF Transformers](https://github.com/huggingface/transformers/issues/34882).
 
 !!! note
-    `jinaai/jina-embeddings-v3` supports multiple tasks through LoRA, while vllm temporarily only supports text-matching tasks by merging LoRA weights.
+    `jinaai/jina-embeddings-v3` supports multiple tasks through lora, while vllm temporarily only supports text-matching tasks by merging lora weights.
 
 !!! note
     The second-generation GTE model (mGTE-TRM) is named `NewModel`. The name `NewModel` is too generic, you should set `--hf-overrides '{"architectures": ["GteNewModel"]}'` to specify the use of the `GteNewModel` architecture.
@@ -556,12 +548,10 @@ Specified using `--task generate`.
 | `FuyuForCausalLM`                            | Fuyu                                                                     | T + I                                                                 | `adept/fuyu-8b` etc.                                                                                                                                    |                        | ✅︎                          | ✅︎                    |
 | `Gemma3ForConditionalGeneration`             | Gemma 3                                                                  | T + I<sup>+</sup>                                                     | `google/gemma-3-4b-it`, `google/gemma-3-27b-it`, etc.                                                                                                   | ✅︎                     | ✅︎                          | ⚠️                    |
 | `GLM4VForCausalLM`<sup>^</sup>               | GLM-4V                                                                   | T + I                                                                 | `THUDM/glm-4v-9b`, `THUDM/cogagent-9b-20241220` etc.                                                                                                    | ✅︎                     | ✅︎                          | ✅︎                    |
-| `Glm4vForConditionalGeneration`              | GLM-4.1V-Thinking                                                        | T + I<sup>E+</sup> + V<sup>E+</sup>                                   | `THUDM/GLM-4.1V-9B-Thinkg`,  etc.                                                                                                                       | ✅︎                     | ✅︎                          | ✅︎                    |
 | `GraniteSpeechForConditionalGeneration`      | Granite Speech                                                           | T + A                                                                 | `ibm-granite/granite-speech-3.3-8b`                                                                                                                     | ✅︎                     | ✅︎                          | ✅︎                    |
 | `H2OVLChatModel`                             | H2OVL                                                                    | T + I<sup>E+</sup>                                                    | `h2oai/h2ovl-mississippi-800m`, `h2oai/h2ovl-mississippi-2b`, etc.                                                                                      |                        | ✅︎                          | ✅︎\*                  |
 | `Idefics3ForConditionalGeneration`           | Idefics3                                                                 | T + I                                                                 | `HuggingFaceM4/Idefics3-8B-Llama3` etc.                                                                                                                 | ✅︎                     |                             |  ✅︎                   |
 | `InternVLChatModel`                          | InternVL 3.0, InternVideo 2.5, InternVL 2.5, Mono-InternVL, InternVL 2.0 | T + I<sup>E+</sup> + (V<sup>E+</sup>)                                 | `OpenGVLab/InternVL3-9B`, `OpenGVLab/InternVideo2_5_Chat_8B`, `OpenGVLab/InternVL2_5-4B`, `OpenGVLab/Mono-InternVL-2B`, `OpenGVLab/InternVL2-4B`, etc.  | ✅︎                     | ✅︎                          | ✅︎                    |
-| `KeyeForConditionalGeneration`               | Keye-VL-8B-Preview                                                       | T + I<sup>E+</sup> + V<sup>E+</sup>                                   | `Kwai-Keye/Keye-VL-8B-Preview`                                                                                                                          |                        |                             | ✅︎                    |
 | `KimiVLForConditionalGeneration`             | Kimi-VL-A3B-Instruct, Kimi-VL-A3B-Thinking                               | T + I<sup>+</sup>                                                     | `moonshotai/Kimi-VL-A3B-Instruct`, `moonshotai/Kimi-VL-A3B-Thinking`                                                                                    |                        |                             | ✅︎                    |
 | `Llama4ForConditionalGeneration`             | Llama 4                                                                  | T + I<sup>+</sup>                                                     | `meta-llama/Llama-4-Scout-17B-16E-Instruct`, `meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8`, `meta-llama/Llama-4-Maverick-17B-128E-Instruct`, etc. |                        | ✅︎                          | ✅︎                    |
 | `LlavaForConditionalGeneration`              | LLaVA-1.5                                                                | T + I<sup>E+</sup>                                                    | `llava-hf/llava-1.5-7b-hf`, `TIGER-Lab/Mantis-8B-siglip-llama3` (see note), etc.                                                                        |                        | ✅︎                          | ✅︎                    |
@@ -588,7 +578,7 @@ Specified using `--task generate`.
 | `SkyworkR1VChatModel`                        | Skywork-R1V-38B                                                          | T + I                                                                 | `Skywork/Skywork-R1V-38B`                                                                                                                               |                        | ✅︎                          | ✅︎                    |
 | `SmolVLMForConditionalGeneration`            | SmolVLM2                                                                 | T + I                                                                 | `SmolVLM2-2.2B-Instruct`                                                                                                                                | ✅︎                     |                             | ✅︎                    |
 | `TarsierForConditionalGeneration`            | Tarsier                                                                  | T + I<sup>E+</sup>                                                    | `omni-search/Tarsier-7b`,`omni-search/Tarsier-34b`                                                                                                      |                        | ✅︎                          | ✅︎                    |
-| `Tarsier2ForConditionalGeneration`<sup>^</sup> | Tarsier2                                                                 | T + I<sup>E+</sup> + V<sup>E+</sup>                                                | `omni-research/Tarsier2-Recap-7b`,`omni-research/Tarsier2-7b-0115`                                                                                      |                        | ✅︎                          | ✅︎                    |
+| `Tarsier2ForConditionalGeneration`<sup>^</sup>            | Tarsier2                                                                  | T + I<sup>E+</sup> + V<sup>E+</sup>                                                    | `omni-research/Tarsier2-Recap-7b`,`omni-research/Tarsier2-7b-0115`                                                                                                      |                        | ✅︎                          | ✅︎                    |
 
 <sup>^</sup> You need to set the architecture name via `--hf-overrides` to match the one in vLLM.  
 &nbsp;&nbsp;&nbsp;&nbsp;• For example, to use DeepSeek-VL2 series models:  

docs/serving/distributed_serving.md

@@ -100,50 +100,7 @@ vllm serve /path/to/the/model/in/the/container \
      --tensor-parallel-size 16
 ```
 
-To make tensor parallel performant, you should make sure the communication between nodes is efficient, e.g. using high-speed network cards like InfiniBand. To correctly set up the cluster to use InfiniBand, append additional arguments like `--privileged -e NCCL_IB_HCA=mlx5` to the `run_cluster.sh` script. Please contact your system administrator for more information on how to set up the flags. One way to confirm if the InfiniBand is working is to run vLLM with `NCCL_DEBUG=TRACE` environment variable set, e.g. `NCCL_DEBUG=TRACE vllm serve ...` and check the logs for the NCCL version and the network used. If you find `[send] via NET/Socket` in the logs, it means NCCL uses raw TCP Socket, which is not efficient for cross-node tensor parallel. If you find `[send] via NET/IB/GDRDMA` in the logs, it means NCCL uses InfiniBand with GPUDirect RDMA, which is efficient.
-
-### GPUDirect RDMA
-
-To enable GPUDirect RDMA with vLLM, specific configuration tweaks are needed. This setup ensures:
-
-- `IPC_LOCK` Security Context: Add the `IPC_LOCK` capability to the container’s security context to lock memory pages and prevent swapping to disk.
-- Shared Memory with `/dev/shm`: Mount `/dev/shm` in the pod spec to provide shared memory for IPC.
-
-When using Docker, you can set up the container as follows:
-
-```bash
-docker run --gpus all \
-    --ipc=host \
-    --shm-size=16G \
-    -v /dev/shm:/dev/shm \
-    vllm/vllm-openai
-```
-
-When using Kubernetes, you can set up the pod spec as follows:
-
-```yaml
-...
-spec:
-  containers:
-    - name: vllm
-      image: vllm/vllm-openai
-      securityContext:
-        capabilities:
-          add: ["IPC_LOCK"]
-      volumeMounts:
-        - mountPath: /dev/shm
-          name: dshm
-      resources:
-        limits:
-          nvidia.com/gpu: 8
-        requests:
-          nvidia.com/gpu: 8
-  volumes:
-    - name: dshm
-      emptyDir:
-        medium: Memory
-...
-```
+To make tensor parallel performant, you should make sure the communication between nodes is efficient, e.g. using high-speed network cards like Infiniband. To correctly set up the cluster to use Infiniband, append additional arguments like `--privileged -e NCCL_IB_HCA=mlx5` to the `run_cluster.sh` script. Please contact your system administrator for more information on how to set up the flags. One way to confirm if the Infiniband is working is to run vLLM with `NCCL_DEBUG=TRACE` environment variable set, e.g. `NCCL_DEBUG=TRACE vllm serve ...` and check the logs for the NCCL version and the network used. If you find `[send] via NET/Socket` in the logs, it means NCCL uses raw TCP Socket, which is not efficient for cross-node tensor parallel. If you find `[send] via NET/IB/GDRDMA` in the logs, it means NCCL uses Infiniband with GPU-Direct RDMA, which is efficient.
 
 !!! warning
     After you start the Ray cluster, you'd better also check the GPU-GPU communication between nodes. It can be non-trivial to set up. Please refer to the [sanity check script][troubleshooting-incorrect-hardware-driver] for more information. If you need to set some environment variables for the communication configuration, you can append them to the `run_cluster.sh` script, e.g. `-e NCCL_SOCKET_IFNAME=eth0`. Note that setting environment variables in the shell (e.g. `NCCL_SOCKET_IFNAME=eth0 vllm serve ...`) only works for the processes in the same node, not for the processes in the other nodes. Setting environment variables when you create the cluster is the recommended way. See <gh-issue:6803> for more information.

docs/serving/openai_compatible_server.md

@@ -1,7 +1,7 @@
 ---
 title: OpenAI-Compatible Server
 ---
-[](){ #serving-openai-compatible-server }
+[](){ #openai-compatible-server }
 
 vLLM provides an HTTP server that implements OpenAI's [Completions API](https://platform.openai.com/docs/api-reference/completions), [Chat API](https://platform.openai.com/docs/api-reference/chat), and more! This functionality lets you serve models and interact with them using an HTTP client.
 
@@ -57,8 +57,6 @@ We currently support the following OpenAI APIs:
     - Only applicable to [embedding models](../models/pooling_models.md) (`--task embed`).
 - [Transcriptions API][transcriptions-api] (`/v1/audio/transcriptions`)
     - Only applicable to Automatic Speech Recognition (ASR) models (OpenAI Whisper) (`--task generate`).
-- [Translation API][translations-api] (`/v1/audio/translations`)
-    - Only applicable to Automatic Speech Recognition (ASR) models (OpenAI Whisper) (`--task generate`).
 
 In addition, we have the following custom APIs:
 
@@ -146,6 +144,11 @@ completion = client.chat.completions.create(
 Only `X-Request-Id` HTTP request header is supported for now. It can be enabled
 with `--enable-request-id-headers`.
 
+> Note that enablement of the headers can impact performance significantly at high QPS
+> rates. We recommend implementing HTTP headers at the router level (e.g. via Istio),
+> rather than within the vLLM layer for this reason.
+> See [this PR](https://github.com/vllm-project/vllm/pull/11529) for more details.
+
 ??? Code
 
     ```python
@@ -371,34 +374,6 @@ The following extra parameters are supported:
     ```python
     --8<-- "vllm/entrypoints/openai/protocol.py:transcription-extra-params"
     ```
-  
-[](){ #translations-api }
-
-### Translations API
-
-Our Translation API is compatible with [OpenAI's Translations API](https://platform.openai.com/docs/api-reference/audio/createTranslation);
-you can use the [official OpenAI Python client](https://github.com/openai/openai-python) to interact with it.
-Whisper models can translate audio from one of the 55 non-English supported languages into English.
-Please mind that the popular `openai/whisper-large-v3-turbo` model does not support translating.
-
-!!! note
-    To use the Translation API, please install with extra audio dependencies using `pip install vllm[audio]`.
-
-Code example: <gh-file:examples/online_serving/openai_translation_client.py>
-
-#### Extra Parameters
-
-The following [sampling parameters][sampling-params] are supported.
-
-```python
---8<-- "vllm/entrypoints/openai/protocol.py:translation-sampling-params"
-```
-
-The following extra parameters are supported:
-
-```python
---8<-- "vllm/entrypoints/openai/protocol.py:translation-extra-params"
-```
 
 [](){ #tokenizer-api }
 

docs/usage/troubleshooting.md

@@ -273,27 +273,6 @@ But you are sure that the model is in the [list of supported models][supported-m
 
 If you see an error like `RuntimeError: Failed to infer device type`, it means that vLLM failed to infer the device type of the runtime environment. You can check [the code](gh-file:vllm/platforms/__init__.py) to see how vLLM infers the device type and why it is not working as expected. After [this PR](gh-pr:14195), you can also set the environment variable `VLLM_LOGGING_LEVEL=DEBUG` to see more detailed logs to help debug the issue.
 
-## NCCL error: unhandled system error during `ncclCommInitRank`
-
-If your serving workload uses GPUDirect RDMA for distributed serving across multiple nodes and encounters an error during `ncclCommInitRank`, with no clear error message even with `NCCL_DEBUG=INFO` set, it might look like this:
-
-```text
-Error executing method 'init_device'. This might cause deadlock in distributed execution.
-Traceback (most recent call last):
-...
-   File "/usr/local/lib/python3.12/dist-packages/vllm/distributed/device_communicators/pynccl.py", line 99, in __init__
-     self.comm: ncclComm_t = self.nccl.ncclCommInitRank(
-                             ^^^^^^^^^^^^^^^^^^^^^^^^^^^
-   File "/usr/local/lib/python3.12/dist-packages/vllm/distributed/device_communicators/pynccl_wrapper.py", line 277, in ncclCommInitRank
-     self.NCCL_CHECK(self._funcs["ncclCommInitRank"](ctypes.byref(comm),
-   File "/usr/local/lib/python3.12/dist-packages/vllm/distributed/device_communicators/pynccl_wrapper.py", line 256, in NCCL_CHECK
-     raise RuntimeError(f"NCCL error: {error_str}")
- RuntimeError: NCCL error: unhandled system error (run with NCCL_DEBUG=INFO for details)
-...
-```
-
-This indicates vLLM failed to initialize the NCCL communicator, possibly due to a missing `IPC_LOCK` linux capability  or an unmounted `/dev/shm`. Refer to [Distributed Inference and Serving](../serving/distributed_serving.md#running-vllm-on-multiple-nodes) for guidance on properly configuring the environment for distributed serving.
-
 ## Known Issues
 
 - In `v0.5.2`, `v0.5.3`, and `v0.5.3.post1`, there is a bug caused by [zmq](https://github.com/zeromq/pyzmq/issues/2000) , which can occasionally cause vLLM to hang depending on the machine configuration. The solution is to upgrade to the latest version of `vllm` to include the [fix](gh-pr:6759).

examples/offline_inference/data_parallel.py

@@ -64,18 +64,6 @@ def parse_args():
     parser.add_argument(
         "--trust-remote-code", action="store_true", help="Trust remote code."
     )
-    parser.add_argument(
-        "--max-num-seqs",
-        type=int,
-        default=64,
-        help=("Maximum number of sequences to be processed in a single iteration."),
-    )
-    parser.add_argument(
-        "--gpu-memory-utilization",
-        type=float,
-        default=0.8,
-        help=("Fraction of GPU memory vLLM is allowed to allocate (0.0, 1.0]."),
-    )
     return parser.parse_args()
 
 
@@ -89,8 +77,6 @@ def main(
     GPUs_per_dp_rank,
     enforce_eager,
     trust_remote_code,
-    max_num_seqs,
-    gpu_memory_utilization,
 ):
     os.environ["VLLM_DP_RANK"] = str(global_dp_rank)
     os.environ["VLLM_DP_RANK_LOCAL"] = str(local_dp_rank)
@@ -141,8 +127,6 @@ def main(
         enforce_eager=enforce_eager,
         enable_expert_parallel=True,
         trust_remote_code=trust_remote_code,
-        max_num_seqs=max_num_seqs,
-        gpu_memory_utilization=gpu_memory_utilization,
     )
     outputs = llm.generate(prompts, sampling_params)
     # Print the outputs.
@@ -197,8 +181,6 @@ if __name__ == "__main__":
                 tp_size,
                 args.enforce_eager,
                 args.trust_remote_code,
-                args.max_num_seqs,
-                args.gpu_memory_utilization,
             ),
         )
         proc.start()

examples/offline_inference/eagle.py

@@ -0,0 +1,144 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import argparse
+import json
+import os
+
+from transformers import AutoTokenizer
+
+from vllm import LLM, SamplingParams
+from vllm.v1.metrics.reader import Counter, Vector
+
+
+def load_prompts(dataset_path, num_prompts):
+    if os.path.exists(dataset_path):
+        prompts = []
+        try:
+            with open(dataset_path) as f:
+                for line in f:
+                    data = json.loads(line)
+                    prompts.append(data["turns"][0])
+        except Exception as e:
+            print(f"Error reading dataset: {e}")
+            return []
+    else:
+        prompts = ["The future of AI is", "The president of the United States is"]
+
+    return prompts[:num_prompts]
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--dataset",
+        type=str,
+        default="./examples/data/gsm8k.jsonl",
+        help="downloaded from the eagle repo "
+        "https://github.com/SafeAILab/EAGLE/blob/main/eagle/data/",
+    )
+    parser.add_argument(
+        "--method", type=str, default="eagle", choices=["eagle", "eagle3"]
+    )
+    parser.add_argument("--max_num_seqs", type=int, default=8)
+    parser.add_argument("--num_prompts", type=int, default=80)
+    parser.add_argument("--num_spec_tokens", type=int, default=2)
+    parser.add_argument("--tp", type=int, default=1)
+    parser.add_argument("--draft_tp", type=int, default=1)
+    parser.add_argument("--enforce_eager", action="store_true")
+    parser.add_argument("--enable_chunked_prefill", action="store_true")
+    parser.add_argument("--max_num_batched_tokens", type=int, default=2048)
+    parser.add_argument("--temp", type=float, default=0)
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+
+    model_dir = "meta-llama/Llama-3.1-8B-Instruct"
+
+    if args.method == "eagle":
+        eagle_dir = "yuhuili/EAGLE-LLaMA3.1-Instruct-8B"
+    elif args.method == "eagle3":
+        eagle_dir = "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B"
+    else:
+        raise ValueError(f"unknown method: {args.method}")
+
+    max_model_len = 2048
+
+    tokenizer = AutoTokenizer.from_pretrained(model_dir)
+
+    prompts = load_prompts(args.dataset, args.num_prompts)
+
+    prompt_ids = [
+        tokenizer.apply_chat_template(
+            [{"role": "user", "content": prompt}], add_generation_prompt=True
+        )
+        for prompt in prompts
+    ]
+
+    llm = LLM(
+        model=model_dir,
+        trust_remote_code=True,
+        tensor_parallel_size=args.tp,
+        enable_chunked_prefill=args.enable_chunked_prefill,
+        max_num_batched_tokens=args.max_num_batched_tokens,
+        enforce_eager=args.enforce_eager,
+        max_model_len=max_model_len,
+        max_num_seqs=args.max_num_seqs,
+        gpu_memory_utilization=0.8,
+        speculative_config={
+            "method": args.method,
+            "model": eagle_dir,
+            "num_speculative_tokens": args.num_spec_tokens,
+            "draft_tensor_parallel_size": args.draft_tp,
+            "max_model_len": max_model_len,
+        },
+        disable_log_stats=False,
+    )
+
+    sampling_params = SamplingParams(temperature=args.temp, max_tokens=256)
+
+    outputs = llm.generate(prompt_token_ids=prompt_ids, sampling_params=sampling_params)
+
+    # print the generated text
+    for output in outputs:
+        print("-" * 50)
+        print(f"prompt: {output.prompt}")
+        print(f"generated text: {output.outputs[0].text}")
+        print("-" * 50)
+
+    try:
+        metrics = llm.get_metrics()
+    except AssertionError:
+        print("Metrics are not supported in the V0 engine.")
+        return
+
+    num_drafts = num_accepted = 0
+    acceptance_counts = [0] * args.num_spec_tokens
+    for metric in metrics:
+        if metric.name == "vllm:spec_decode_num_drafts":
+            assert isinstance(metric, Counter)
+            num_drafts += metric.value
+        elif metric.name == "vllm:spec_decode_num_accepted_tokens":
+            assert isinstance(metric, Counter)
+            num_accepted += metric.value
+        elif metric.name == "vllm:spec_decode_num_accepted_tokens_per_pos":
+            assert isinstance(metric, Vector)
+            for pos in range(len(metric.values)):
+                acceptance_counts[pos] += metric.values[pos]
+
+    print("-" * 50)
+    print(f"mean acceptance length: {1 + (num_accepted / num_drafts):.2f}")
+    print("-" * 50)
+
+    # print acceptance at each token position
+    for i in range(len(acceptance_counts)):
+        print(f"acceptance at token {i}:{acceptance_counts[i] / num_drafts:.2f}")
+
+
+if __name__ == "__main__":
+    print(
+        "[WARNING] Use examples/offline_inference/spec_decode.py"
+        " instead of this script."
+    )
+    main()

examples/offline_inference/spec_decode.py

@@ -16,24 +16,29 @@ def parse_args():
     parser = FlexibleArgumentParser()
     add_dataset_parser(parser)
     parser.add_argument(
-        "--method",
+        "--dataset",
         type=str,
-        default="eagle",
-        choices=["ngram", "eagle", "eagle3", "mtp"],
+        default="./examples/data/gsm8k.jsonl",
+        help="downloaded from the eagle repo "
+        "https://github.com/SafeAILab/EAGLE/blob/main/eagle/data/",
     )
+    parser.add_argument(
+        "--method", type=str, default="eagle", choices=["ngram", "eagle", "eagle3"]
+    )
+    parser.add_argument("--max-num-seqs", type=int, default=8)
     parser.add_argument("--num-spec-tokens", type=int, default=2)
     parser.add_argument("--prompt-lookup-max", type=int, default=5)
     parser.add_argument("--prompt-lookup-min", type=int, default=2)
     parser.add_argument("--tp", type=int, default=1)
+    parser.add_argument("--draft-tp", type=int, default=1)
     parser.add_argument("--enforce-eager", action="store_true")
     parser.add_argument("--enable-chunked-prefill", action="store_true")
+    parser.add_argument("--max-num-batched-tokens", type=int, default=2048)
     parser.add_argument("--temp", type=float, default=0)
     parser.add_argument("--top-p", type=float, default=1.0)
     parser.add_argument("--top-k", type=int, default=-1)
     parser.add_argument("--print-output", action="store_true")
     parser.add_argument("--output-len", type=int, default=256)
-    parser.add_argument("--model-dir", type=str, default=None)
-    parser.add_argument("--eagle-dir", type=str, default=None)
     return parser.parse_args()
 
 
@@ -41,10 +46,9 @@ def main():
     args = parse_args()
     args.endpoint_type = "openai-chat"
 
-    model_dir = args.model_dir
-    if args.model_dir is None:
-        model_dir = "meta-llama/Llama-3.1-8B-Instruct"
+    model_dir = "meta-llama/Llama-3.1-8B-Instruct"
     tokenizer = AutoTokenizer.from_pretrained(model_dir)
+    max_model_len = 2048
 
     prompts = get_samples(args, tokenizer)
     # add_special_tokens is False to avoid adding bos twice when using chat templates
@@ -53,16 +57,16 @@ def main():
     ]
 
     if args.method == "eagle" or args.method == "eagle3":
-        eagle_dir = args.eagle_dir
-        if args.method == "eagle" and eagle_dir is None:
+        if args.method == "eagle":
             eagle_dir = "yuhuili/EAGLE-LLaMA3.1-Instruct-8B"
-
-        elif args.method == "eagle3" and eagle_dir is None:
+        elif args.method == "eagle3":
             eagle_dir = "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B"
         speculative_config = {
             "method": args.method,
             "model": eagle_dir,
             "num_speculative_tokens": args.num_spec_tokens,
+            "draft_tensor_parallel_size": args.draft_tp,
+            "max_model_len": max_model_len,
         }
     elif args.method == "ngram":
         speculative_config = {
@@ -70,6 +74,7 @@ def main():
             "num_speculative_tokens": args.num_spec_tokens,
             "prompt_lookup_max": args.prompt_lookup_max,
             "prompt_lookup_min": args.prompt_lookup_min,
+            "max_model_len": max_model_len,
         }
     else:
         raise ValueError(f"unknown method: {args.method}")
@@ -79,7 +84,10 @@ def main():
         trust_remote_code=True,
         tensor_parallel_size=args.tp,
         enable_chunked_prefill=args.enable_chunked_prefill,
+        max_num_batched_tokens=args.max_num_batched_tokens,
         enforce_eager=args.enforce_eager,
+        max_model_len=max_model_len,
+        max_num_seqs=args.max_num_seqs,
         gpu_memory_utilization=0.8,
         speculative_config=speculative_config,
         disable_log_stats=False,
@@ -102,41 +110,27 @@ def main():
         print("Metrics are not supported in the V0 engine.")
         return
 
-    total_num_output_tokens = sum(
-        len(output.outputs[0].token_ids) for output in outputs
-    )
-    num_drafts = 0
-    num_draft_tokens = 0
-    num_accepted_tokens = 0
+    num_drafts = num_accepted = 0
     acceptance_counts = [0] * args.num_spec_tokens
     for metric in metrics:
         if metric.name == "vllm:spec_decode_num_drafts":
             assert isinstance(metric, Counter)
             num_drafts += metric.value
-        elif metric.name == "vllm:spec_decode_num_draft_tokens":
-            assert isinstance(metric, Counter)
-            num_draft_tokens += metric.value
         elif metric.name == "vllm:spec_decode_num_accepted_tokens":
             assert isinstance(metric, Counter)
-            num_accepted_tokens += metric.value
+            num_accepted += metric.value
         elif metric.name == "vllm:spec_decode_num_accepted_tokens_per_pos":
             assert isinstance(metric, Vector)
             for pos in range(len(metric.values)):
                 acceptance_counts[pos] += metric.values[pos]
 
     print("-" * 50)
-    print(f"total_num_output_tokens: {total_num_output_tokens}")
-    print(f"num_drafts: {num_drafts}")
-    print(f"num_draft_tokens: {num_draft_tokens}")
-    print(f"num_accepted_tokens: {num_accepted_tokens}")
-    acceptance_length = 1 + (num_accepted_tokens / num_drafts) if num_drafts > 0 else 1
-    print(f"mean acceptance length: {acceptance_length:.2f}")
+    print(f"mean acceptance length: {1 + (num_accepted / num_drafts):.2f}")
     print("-" * 50)
 
     # print acceptance at each token position
     for i in range(len(acceptance_counts)):
-        acceptance_rate = acceptance_counts[i] / num_drafts if num_drafts > 0 else 0
-        print(f"acceptance at token {i}: {acceptance_rate:.2f}")
+        print(f"acceptance at token {i}:{acceptance_counts[i] / num_drafts:.2f}")
 
 
 if __name__ == "__main__":

examples/offline_inference/vision_language.py

@@ -248,42 +248,6 @@ def run_glm4v(questions: list[str], modality: str) -> ModelRequestData:
     )
 
 
-# GLM-4.1V
-def run_glm4_1v(questions: list[str], modality: str) -> ModelRequestData:
-    model_name = "THUDM/GLM-4.1V-9B-Thinking"
-
-    engine_args = EngineArgs(
-        model=model_name,
-        max_model_len=4096,
-        max_num_seqs=2,
-        mm_processor_kwargs={
-            "size": {"shortest_edge": 12544, "longest_edge": 47040000},
-            "fps": 1,
-        },
-        limit_mm_per_prompt={modality: 1},
-        enforce_eager=True,
-    )
-
-    if modality == "image":
-        placeholder = "<|begin_of_image|><|image|><|end_of_image|>"
-    elif modality == "video":
-        placeholder = "<|begin_of_video|><|video|><|end_of_video|>"
-
-    prompts = [
-        (
-            "[gMASK]<sop><|system|>\nYou are a helpful assistant.<|user|>\n"
-            f"{placeholder}"
-            f"{question}<|assistant|>assistant\n"
-        )
-        for question in questions
-    ]
-
-    return ModelRequestData(
-        engine_args=engine_args,
-        prompts=prompts,
-    )
-
-
 # H2OVL-Mississippi
 def run_h2ovl(questions: list[str], modality: str) -> ModelRequestData:
     assert modality == "image"
@@ -429,37 +393,6 @@ def run_internvl(questions: list[str], modality: str) -> ModelRequestData:
     )
 
 
-# Keye-VL
-def run_keye_vl(questions: list[str], modality: str) -> ModelRequestData:
-    model_name = "Kwai-Keye/Keye-VL-8B-Preview"
-
-    engine_args = EngineArgs(
-        model=model_name,
-        max_model_len=8192,
-        trust_remote_code=True,
-        limit_mm_per_prompt={modality: 1},
-    )
-
-    if modality == "image":
-        placeholder = "<|image_pad|>"
-    elif modality == "video":
-        placeholder = "<|video_pad|>"
-
-    prompts = [
-        (
-            f"<|im_start|>user\n<|vision_start|>{placeholder}<|vision_end|>"
-            f"{question}<|im_end|>\n"
-            "<|im_start|>assistant\n"
-        )
-        for question in questions
-    ]
-
-    return ModelRequestData(
-        engine_args=engine_args,
-        prompts=prompts,
-    )
-
-
 # Kimi-VL
 def run_kimi_vl(questions: list[str], modality: str) -> ModelRequestData:
     assert modality == "image"
@@ -1181,11 +1114,9 @@ model_example_map = {
     "fuyu": run_fuyu,
     "gemma3": run_gemma3,
     "glm4v": run_glm4v,
-    "glm4_1v": run_glm4_1v,
     "h2ovl_chat": run_h2ovl,
     "idefics3": run_idefics3,
     "internvl_chat": run_internvl,
-    "keye_vl": run_keye_vl,
     "kimi_vl": run_kimi_vl,
     "llava": run_llava,
     "llava-next": run_llava_next,
@@ -1241,11 +1172,10 @@ def get_multi_modal_input(args):
     if args.modality == "video":
         # Input video and question
         video = VideoAsset(name="baby_reading", num_frames=args.num_frames).np_ndarrays
-        metadata = VideoAsset(name="baby_reading", num_frames=args.num_frames).metadata
         vid_questions = ["Why is this video funny?"]
 
         return {
-            "data": [(video, metadata)] if args.model_type == "glm4_1v" else video,
+            "data": video,
             "questions": vid_questions,
         }
 

examples/offline_inference/vision_language_multi_image.py

@@ -423,43 +423,6 @@ def load_llama4(question: str, image_urls: list[str]) -> ModelRequestData:
     )
 
 
-def load_keye_vl(question: str, image_urls: list[str]) -> ModelRequestData:
-    model_name = "Kwai-Keye/Keye-VL-8B-Preview"
-
-    engine_args = EngineArgs(
-        model=model_name,
-        trust_remote_code=True,
-        max_model_len=8192,
-        max_num_seqs=5,
-        limit_mm_per_prompt={"image": len(image_urls)},
-    )
-
-    placeholders = [{"type": "image", "image": url} for url in image_urls]
-    messages = [
-        {
-            "role": "user",
-            "content": [
-                *placeholders,
-                {"type": "text", "text": question},
-            ],
-        },
-    ]
-
-    processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True)
-
-    prompt = processor.apply_chat_template(
-        messages, tokenize=False, add_generation_prompt=True
-    )
-
-    image_data = [fetch_image(url) for url in image_urls]
-
-    return ModelRequestData(
-        engine_args=engine_args,
-        prompt=prompt,
-        image_data=image_data,
-    )
-
-
 def load_kimi_vl(question: str, image_urls: list[str]) -> ModelRequestData:
     model_name = "moonshotai/Kimi-VL-A3B-Instruct"
 
@@ -899,7 +862,6 @@ model_example_map = {
     "h2ovl_chat": load_h2ovl,
     "idefics3": load_idefics3,
     "internvl_chat": load_internvl,
-    "keye_vl": load_keye_vl,
     "kimi_vl": load_kimi_vl,
     "llava": load_llava,
     "llava-next": load_llava_next,

examples/online_serving/disaggregated_serving_p2p_nccl_xpyd/disagg_example_p2p_nccl_xpyd.sh

@@ -1,245 +0,0 @@
-#!/bin/bash
-
-# =============================================================================
-# vLLM Disaggregated Serving Script - P2P NCCL XpYd Architecture
-# =============================================================================
-# This script demonstrates disaggregated prefill and decode serving using
-# P2P NCCL communication. The architecture supports various XpYd configurations:
-#
-# - 1P3D: 1 Prefill server + 3 Decode servers (current default)
-# - 3P1D: 3 Prefill servers + 1 Decode server
-# - etc.
-#
-# Configuration can be customized via environment variables:
-#   MODEL: Model to serve
-#   PREFILL_GPUS: Comma-separated GPU IDs for prefill servers
-#   DECODE_GPUS: Comma-separated GPU IDs for decode servers
-#   PREFILL_PORTS: Comma-separated ports for prefill servers
-#   DECODE_PORTS: Comma-separated ports for decode servers
-#   PROXY_PORT: Proxy server port used to setup XpYd connection.
-#   TIMEOUT_SECONDS: Server startup timeout
-# =============================================================================
-
-# Configuration - can be overridden via environment variables
-MODEL=${MODEL:-meta-llama/Llama-3.1-8B-Instruct}
-TIMEOUT_SECONDS=${TIMEOUT_SECONDS:-1200}
-PROXY_PORT=${PROXY_PORT:-30001}
-
-# Default 1P3D configuration (1 Prefill + 3 Decode)
-PREFILL_GPUS=${PREFILL_GPUS:-0}
-DECODE_GPUS=${DECODE_GPUS:-1,2,3}
-PREFILL_PORTS=${PREFILL_PORTS:-20003}
-DECODE_PORTS=${DECODE_PORTS:-20005,20007,20009} 
-
-echo "Warning: P2P NCCL disaggregated prefill XpYd support for vLLM v1 is experimental and subject to change."
-echo ""
-echo "Architecture Configuration:"
-echo "  Model: $MODEL"
-echo "  Prefill GPUs: $PREFILL_GPUS, Ports: $PREFILL_PORTS"
-echo "  Decode GPUs: $DECODE_GPUS, Ports: $DECODE_PORTS"
-echo "  Proxy Port: $PROXY_PORT"
-echo "  Timeout: ${TIMEOUT_SECONDS}s"
-echo ""
-
-PIDS=()
-
-# Switch to the directory of the current script
-cd "$(dirname "${BASH_SOURCE[0]}")"
-
-check_required_files() {
-    local files=("disagg_proxy_p2p_nccl_xpyd.py")
-    for file in "${files[@]}"; do
-        if [[ ! -f "$file" ]]; then
-            echo "Required file $file not found in $(pwd)"
-            exit 1
-        fi
-    done
-}
-
-check_hf_token() {
-    if [ -z "$HF_TOKEN" ]; then
-        echo "HF_TOKEN is not set. Please set it to your Hugging Face token."
-        echo "Example: export HF_TOKEN=your_token_here"
-        exit 1
-    fi
-    if [[ "$HF_TOKEN" != hf_* ]]; then
-        echo "HF_TOKEN is not a valid Hugging Face token. Please set it to your Hugging Face token."
-        exit 1
-    fi
-    echo "HF_TOKEN is set and valid."
-}
-
-check_num_gpus() {
-    # Check if the number of GPUs are >=2 via nvidia-smi
-    num_gpus=$(nvidia-smi --query-gpu=name --format=csv,noheader | wc -l)
-    if [ "$num_gpus" -lt 2 ]; then
-        echo "You need at least 2 GPUs to run disaggregated prefill."
-        exit 1
-    else
-        echo "Found $num_gpus GPUs."
-    fi
-}
-
-ensure_python_library_installed() {
-    echo "Checking if $1 is installed..."
-    if ! python3 -c "import $1" > /dev/null 2>&1; then
-        echo "$1 is not installed. Please install it via pip install $1."
-        exit 1
-    else
-        echo "$1 is installed."
-    fi
-}
-
-cleanup() {
-    echo "Stopping everything…"
-    trap - INT TERM        # prevent re-entrancy
-    kill -- -$$            # negative PID  ==  "this whole process-group"
-    wait                   # reap children so we don't leave zombies
-    exit 0
-}
-
-wait_for_server() {
-  local port=$1
-  local timeout_seconds=$TIMEOUT_SECONDS
-  local start_time=$(date +%s)
-
-  echo "Waiting for server on port $port..."
-
-  while true; do
-    if curl -s "localhost:${port}/v1/completions" > /dev/null; then
-      echo "Server on port $port is ready."
-      return 0
-    fi
-
-    local now=$(date +%s)
-    if (( now - start_time >= timeout_seconds )); then
-      echo "Timeout waiting for server on port $port"
-      return 1
-    fi
-
-    sleep 1
-  done
-}
-
-main() {
-    check_required_files
-    check_hf_token
-    check_num_gpus
-    ensure_python_library_installed pandas
-    ensure_python_library_installed datasets
-    ensure_python_library_installed vllm
-    ensure_python_library_installed quart
-
-    trap cleanup INT
-    trap cleanup USR1
-    trap cleanup TERM
-
-    echo "Launching disaggregated serving components..."
-    echo "Please check the log files for detailed output:"
-    echo "  - prefill*.log: Prefill server logs"
-    echo "  - decode*.log: Decode server logs"
-    echo "  - proxy.log: Proxy server log"
-
-    # =============================================================================
-    # Launch Proxy Server
-    # =============================================================================
-    echo ""
-    echo "Starting proxy server on port $PROXY_PORT..."
-    python3 disagg_proxy_p2p_nccl_xpyd.py &
-    PIDS+=($!)
-
-    # Parse GPU and port arrays
-    IFS=',' read -ra PREFILL_GPU_ARRAY <<< "$PREFILL_GPUS"
-    IFS=',' read -ra DECODE_GPU_ARRAY <<< "$DECODE_GPUS"
-    IFS=',' read -ra PREFILL_PORT_ARRAY <<< "$PREFILL_PORTS"
-    IFS=',' read -ra DECODE_PORT_ARRAY <<< "$DECODE_PORTS"
-
-    # =============================================================================
-    # Launch Prefill Servers (X Producers)
-    # =============================================================================
-    echo ""
-    echo "Starting ${#PREFILL_GPU_ARRAY[@]} prefill server(s)..."
-    for i in "${!PREFILL_GPU_ARRAY[@]}"; do
-        local gpu_id=${PREFILL_GPU_ARRAY[$i]}
-        local port=${PREFILL_PORT_ARRAY[$i]}
-        local kv_port=$((21001 + i))
-        
-        echo "  Prefill server $((i+1)): GPU $gpu_id, Port $port, KV Port $kv_port"
-        CUDA_VISIBLE_DEVICES=$gpu_id VLLM_USE_V1=1 vllm serve $MODEL \
-        --enforce-eager \
-        --host 0.0.0.0 \
-        --port $port \
-        --tensor-parallel-size 1 \
-        --seed 1024 \
-        --dtype float16 \
-        --max-model-len 10000 \
-        --max-num-batched-tokens 10000 \
-        --max-num-seqs 256 \
-        --trust-remote-code \
-        --gpu-memory-utilization 0.9 \
-        --disable-log-request \
-        --kv-transfer-config \
-        "{\"kv_connector\":\"P2pNcclConnector\",\"kv_role\":\"kv_producer\",\"kv_buffer_size\":\"1e1\",\"kv_port\":\"$kv_port\",\"kv_connector_extra_config\":{\"proxy_ip\":\"0.0.0.0\",\"proxy_port\":\"$PROXY_PORT\",\"http_port\":\"$port\",\"send_type\":\"PUT_ASYNC\",\"nccl_num_channels\":\"16\"}}" > prefill$((i+1)).log 2>&1 &
-        PIDS+=($!)
-    done
-
-    # =============================================================================
-    # Launch Decode Servers (Y Decoders)
-    # =============================================================================
-    echo ""
-    echo "Starting ${#DECODE_GPU_ARRAY[@]} decode server(s)..."
-    for i in "${!DECODE_GPU_ARRAY[@]}"; do
-        local gpu_id=${DECODE_GPU_ARRAY[$i]}
-        local port=${DECODE_PORT_ARRAY[$i]}
-        local kv_port=$((22001 + i))
-        
-        echo "  Decode server $((i+1)): GPU $gpu_id, Port $port, KV Port $kv_port"
-        VLLM_USE_V1=1 CUDA_VISIBLE_DEVICES=$gpu_id vllm serve $MODEL \
-        --enforce-eager \
-        --host 0.0.0.0 \
-        --port $port \
-        --tensor-parallel-size 1 \
-        --seed 1024 \
-        --dtype float16 \
-        --max-model-len 10000 \
-        --max-num-batched-tokens 10000 \
-        --max-num-seqs 256 \
-        --trust-remote-code \
-        --gpu-memory-utilization 0.7 \
-        --disable-log-request \
-        --kv-transfer-config \
-        "{\"kv_connector\":\"P2pNcclConnector\",\"kv_role\":\"kv_consumer\",\"kv_buffer_size\":\"8e9\",\"kv_port\":\"$kv_port\",\"kv_connector_extra_config\":{\"proxy_ip\":\"0.0.0.0\",\"proxy_port\":\"$PROXY_PORT\",\"http_port\":\"$port\",\"send_type\":\"PUT_ASYNC\",\"nccl_num_channels\":\"16\"}}" > decode$((i+1)).log 2>&1 &
-        PIDS+=($!)
-    done
-
-    # =============================================================================
-    # Wait for All Servers to Start
-    # =============================================================================
-    echo ""
-    echo "Waiting for all servers to start..."
-    for port in "${PREFILL_PORT_ARRAY[@]}" "${DECODE_PORT_ARRAY[@]}"; do
-        if ! wait_for_server $port; then
-            echo "Failed to start server on port $port"
-            cleanup
-            exit 1
-        fi
-    done
-
-    echo ""
-    echo "All servers are up. Starting benchmark..."
-
-    # =============================================================================
-    # Run Benchmark
-    # =============================================================================
-    cd ../../../benchmarks/
-    python3 benchmark_serving.py --port 10001 --seed $(date +%s) \
-        --model $MODEL \
-        --dataset-name random --random-input-len 7500 --random-output-len 200 \
-        --num-prompts 200 --burstiness 100 --request-rate 2 | tee benchmark.log
-
-    echo "Benchmarking done. Cleaning up..."
-
-    cleanup
-}
-
-main

examples/online_serving/openai_transcription_client.py

@@ -19,17 +19,30 @@ The script performs:
 """
 
 import asyncio
+import json
 
-from openai import AsyncOpenAI, OpenAI
+import httpx
+from openai import OpenAI
 
 from vllm.assets.audio import AudioAsset
 
+mary_had_lamb = AudioAsset("mary_had_lamb").get_local_path()
+winning_call = AudioAsset("winning_call").get_local_path()
 
-def sync_openai(audio_path: str, client: OpenAI):
+# Modify OpenAI's API key and API base to use vLLM's API server.
+openai_api_key = "EMPTY"
+openai_api_base = "http://localhost:8000/v1"
+client = OpenAI(
+    api_key=openai_api_key,
+    base_url=openai_api_base,
+)
+
+
+def sync_openai():
     """
     Perform synchronous transcription using OpenAI-compatible API.
     """
-    with open(audio_path, "rb") as f:
+    with open(str(mary_had_lamb), "rb") as f:
         transcription = client.audio.transcriptions.create(
             file=f,
             model="openai/whisper-large-v3",
@@ -45,52 +58,45 @@ def sync_openai(audio_path: str, client: OpenAI):
         print("transcription result:", transcription.text)
 
 
-async def stream_openai_response(audio_path: str, client: AsyncOpenAI):
+# OpenAI Transcription API client does not support streaming.
+async def stream_openai_response():
     """
-    Perform asynchronous transcription using OpenAI-compatible API.
+    Perform streaming transcription using vLLM's raw HTTP streaming API.
     """
-    print("\ntranscription result:", end=" ")
-    with open(audio_path, "rb") as f:
-        transcription = await client.audio.transcriptions.create(
-            file=f,
-            model="openai/whisper-large-v3",
-            language="en",
-            response_format="json",
-            temperature=0.0,
-            # Additional sampling params not provided by OpenAI API.
-            extra_body=dict(
-                seed=420,
-                top_p=0.6,
-            ),
-            stream=True,
-        )
-        async for chunk in transcription:
-            if chunk.choices:
-                content = chunk.choices[0].get("delta", {}).get("content")
-                print(content, end="", flush=True)
-
+    data = {
+        "language": "en",
+        "stream": True,
+        "model": "openai/whisper-large-v3",
+    }
+    url = openai_api_base + "/audio/transcriptions"
+    headers = {"Authorization": f"Bearer {openai_api_key}"}
+    print("transcription result:", end=" ")
+    async with httpx.AsyncClient() as client:
+        with open(str(winning_call), "rb") as f:
+            async with client.stream(
+                "POST", url, files={"file": f}, data=data, headers=headers
+            ) as response:
+                async for line in response.aiter_lines():
+                    # Each line is a JSON object prefixed with 'data: '
+                    if line:
+                        if line.startswith("data: "):
+                            line = line[len("data: ") :]
+                        # Last chunk, stream ends
+                        if line.strip() == "[DONE]":
+                            break
+                        # Parse the JSON response
+                        chunk = json.loads(line)
+                        # Extract and print the content
+                        content = chunk["choices"][0].get("delta", {}).get("content")
+                        print(content, end="")
     print()  # Final newline after stream ends
 
 
 def main():
-    mary_had_lamb = str(AudioAsset("mary_had_lamb").get_local_path())
-    winning_call = str(AudioAsset("winning_call").get_local_path())
+    sync_openai()
 
-    # Modify OpenAI's API key and API base to use vLLM's API server.
-    openai_api_key = "EMPTY"
-    openai_api_base = "http://localhost:8000/v1"
-    client = OpenAI(
-        api_key=openai_api_key,
-        base_url=openai_api_base,
-    )
-
-    sync_openai(mary_had_lamb, client)
     # Run the asynchronous function
-    client = AsyncOpenAI(
-        api_key=openai_api_key,
-        base_url=openai_api_base,
-    )
-    asyncio.run(stream_openai_response(winning_call, client))
+    asyncio.run(stream_openai_response())
 
 
 if __name__ == "__main__":

examples/online_serving/openai_translation_client.py

@@ -1,75 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import asyncio
-import json
-
-import httpx
-from openai import OpenAI
-
-from vllm.assets.audio import AudioAsset
-
-
-def sync_openai(audio_path: str, client: OpenAI):
-    with open(audio_path, "rb") as f:
-        translation = client.audio.translations.create(
-            file=f,
-            model="openai/whisper-large-v3",
-            response_format="json",
-            temperature=0.0,
-            # Additional params not provided by OpenAI API.
-            extra_body=dict(
-                language="it",
-                seed=4419,
-                repetition_penalty=1.3,
-            ),
-        )
-        print("translation result:", translation.text)
-
-
-async def stream_openai_response(audio_path: str, base_url: str, api_key: str):
-    data = {
-        "language": "it",
-        "stream": True,
-        "model": "openai/whisper-large-v3",
-    }
-    url = base_url + "/audio/translations"
-    headers = {"Authorization": f"Bearer {api_key}"}
-    print("translation result:", end=" ")
-    # OpenAI translation API client does not support streaming.
-    async with httpx.AsyncClient() as client:
-        with open(audio_path, "rb") as f:
-            async with client.stream(
-                "POST", url, files={"file": f}, data=data, headers=headers
-            ) as response:
-                async for line in response.aiter_lines():
-                    # Each line is a JSON object prefixed with 'data: '
-                    if line:
-                        if line.startswith("data: "):
-                            line = line[len("data: ") :]
-                        # Last chunk, stream ends
-                        if line.strip() == "[DONE]":
-                            break
-                        # Parse the JSON response
-                        chunk = json.loads(line)
-                        # Extract and print the content
-                        content = chunk["choices"][0].get("delta", {}).get("content")
-                        print(content, end="")
-
-
-def main():
-    foscolo = str(AudioAsset("azacinto_foscolo").get_local_path())
-
-    # Modify OpenAI's API key and API base to use vLLM's API server.
-    openai_api_key = "EMPTY"
-    openai_api_base = "http://localhost:8000/v1"
-    client = OpenAI(
-        api_key=openai_api_key,
-        base_url=openai_api_base,
-    )
-    sync_openai(foscolo, client)
-    # Run the asynchronous function
-    asyncio.run(stream_openai_response(foscolo, openai_api_base, openai_api_key))
-
-
-if __name__ == "__main__":
-    main()

mkdocs.yaml

@@ -127,7 +127,6 @@ extra_javascript:
   - mkdocs/javascript/run_llm_widget.js
   - https://cdn.mathjax.org/mathjax/latest/MathJax.js?config=TeX-AMS_HTML
   - mkdocs/javascript/edit_and_feedback.js
-  - mkdocs/javascript/slack_and_forum.js
 
 # Makes the url format end in .html rather than act as a dir
 # So index.md generates as index.html and is available under URL /index.html