Sparse MLA on Hopper: Use SGLang's kernel for the sparse mla low latency runs

Signed-off-by: Alexander Matveev <amatveev@redhat.com>

.buildkite/hardware_tests/cpu.yaml

@@ -46,7 +46,7 @@ steps:
   - tests/models/language/pooling/
   commands:
     - |
-      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
+      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 40m "
       pytest -x -v -s tests/models/language/generation -m cpu_model
       pytest -x -v -s tests/models/language/pooling -m cpu_model"
 
@@ -99,7 +99,7 @@ steps:
     - |
       bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 45m "
       pytest -x -v -s tests/models/multimodal/generation --ignore=tests/models/multimodal/generation/test_pixtral.py -m cpu_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB"
-  parallelism: 2
+  parallelism: 3
 
 - label: "Arm CPU Test"
   depends_on: []

.buildkite/image_build/image_build_torch_nightly.sh

@@ -0,0 +1,68 @@
+#!/bin/bash
+set -euo pipefail
+
+# Build a vLLM test image with PyTorch nightly installed.
+# Called by the pipeline generator's "vLLM Against PyTorch Nightly" group.
+
+if [[ $# -lt 5 ]]; then
+  echo "Usage: $0 <registry> <repo> <commit> <branch> <image_tag>"
+  exit 1
+fi
+
+REGISTRY=$1
+REPO=$2
+BUILDKITE_COMMIT=$3
+BRANCH=$4
+IMAGE_TAG=$5
+
+# --- Arguments ---
+echo "--- :mag: Arguments"
+echo "REGISTRY: ${REGISTRY}"
+echo "REPO: ${REPO}"
+echo "BUILDKITE_COMMIT: ${BUILDKITE_COMMIT}"
+echo "BRANCH: ${BRANCH}"
+echo "IMAGE_TAG: ${IMAGE_TAG}"
+
+# --- ECR login ---
+echo "--- :key: ECR login"
+aws ecr-public get-login-password --region us-east-1 \
+  | docker login --username AWS --password-stdin "$REGISTRY"
+aws ecr get-login-password --region us-east-1 \
+  | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com
+
+# --- Set up buildx ---
+echo "--- :docker: Setting up buildx"
+docker buildx create --name vllm-builder --driver docker-container --use || true
+docker buildx inspect --bootstrap
+docker buildx ls
+
+# --- Skip if image already exists ---
+echo "--- :mag: Checking if image already exists"
+if docker manifest inspect "$IMAGE_TAG" >/dev/null 2>&1; then
+  echo "Image found: $IMAGE_TAG — skipping build"
+  exit 0
+fi
+echo "Image not found, proceeding with build..."
+
+# --- CUDA 13.0 for nightly builds ---
+# Nightly CI uses CUDA 13.0 while regular CI stays on CUDA 12.9
+NIGHTLY_CUDA_VERSION="13.0.0"
+NIGHTLY_BUILD_BASE_IMAGE="nvidia/cuda:${NIGHTLY_CUDA_VERSION}-devel-ubuntu22.04"
+NIGHTLY_FINAL_BASE_IMAGE="nvidia/cuda:${NIGHTLY_CUDA_VERSION}-base-ubuntu22.04"
+
+echo "--- :docker: Building torch nightly image (CUDA ${NIGHTLY_CUDA_VERSION})"
+docker buildx build --file docker/Dockerfile \
+  --build-arg max_jobs=16 \
+  --build-arg buildkite_commit="$BUILDKITE_COMMIT" \
+  --build-arg USE_SCCACHE=1 \
+  --build-arg PYTORCH_NIGHTLY=1 \
+  --build-arg CUDA_VERSION="${NIGHTLY_CUDA_VERSION}" \
+  --build-arg BUILD_BASE_IMAGE="${NIGHTLY_BUILD_BASE_IMAGE}" \
+  --build-arg FINAL_BASE_IMAGE="${NIGHTLY_FINAL_BASE_IMAGE}" \
+  --build-arg torch_cuda_arch_list="8.0 8.9 9.0 10.0 12.0" \
+  --tag "$IMAGE_TAG" \
+  --push \
+  --target test \
+  --progress plain .
+
+echo "--- :white_check_mark: Torch nightly image build complete: $IMAGE_TAG"

.buildkite/release-pipeline.yaml

@@ -98,8 +98,15 @@ steps:
     commands:
       - "bash .buildkite/scripts/generate-and-upload-nightly-index.sh"
 
+  - block: "Unblock to build release Docker images"
+    depends_on: ~
+    key: block-build-release-images
+    if: build.env("NIGHTLY") != "1"
+
   - group: "Build release Docker images"
     key: "build-release-images"
+    depends_on: block-build-release-images
+    allow_dependency_failure: true
     steps:
       - label: "Build release image - x86_64 - CUDA 12.9"
         depends_on: ~
@@ -617,6 +624,8 @@ steps:
   - label: ":docker: Build release image - x86_64 - ROCm"
     id: build-rocm-release-image
     depends_on:
+      - step: block-build-release-images
+        allow_failure: true
       - step: build-rocm-base-wheels
         allow_failure: false
     agents:

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

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

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

@@ -51,6 +51,7 @@ function cpu_tests() {
     set -e
     pytest -x -v -s tests/kernels/test_onednn.py
     pytest -x -v -s tests/kernels/attention/test_cpu_attn.py
+    pytest -x -v -s tests/kernels/core/test_cpu_activation.py
     pytest -x -v -s tests/kernels/moe/test_moe.py -k test_cpu_fused_moe_basic"
 
   # basic online serving

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

@@ -16,5 +16,5 @@ echo "--- :docker: Building Docker image"
 docker build --progress plain --tag "$IMAGE_NAME" --target vllm-test -f docker/Dockerfile.cpu .
 
 # Run the image, setting --shm-size=4g for tensor parallel.
-docker run --rm --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN -e VLLM_CPU_KVCACHE_SPACE=16 -e VLLM_CPU_CI_ENV=1 -e VLLM_CPU_SIM_MULTI_NUMA=1 --shm-size=4g "$IMAGE_NAME" \
+docker run --rm --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN -e VLLM_CPU_KVCACHE_SPACE=16 -e VLLM_CPU_CI_ENV=1 -e VLLM_CPU_SIM_MULTI_NUMA=1 -e VLLM_CPU_ATTN_SPLIT_KV=0 --shm-size=4g "$IMAGE_NAME" \
         timeout "$TIMEOUT_VAL" bash -c "set -euox pipefail; echo \"--- Print packages\"; pip list; echo \"--- Running tests\"; ${TEST_COMMAND}"

.buildkite/test-amd.yaml

@@ -769,7 +769,7 @@ steps:
   - tests/kernels/helion/
   - vllm/platforms/rocm.py
   commands:
-  - pip install helion==0.3.3
+  - pip install helion==1.0.0
   - pytest -v -s kernels/helion/
 
 

.buildkite/test_areas/distributed.yaml

@@ -196,7 +196,8 @@ steps:
     - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_async_new_apis.py
     - VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
     - pytest -v -s tests/v1/distributed/test_dbo.py
-    - TP_SIZE=1 DP_SIZE=2 pytest -v -s tests/v1/distributed/test_eagle_dp.py
+    - VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s tests/distributed/test_weight_transfer.py
+    - pytest -v -s tests/distributed/test_packed_tensor.py
 
 - label: Distributed Tests (2 GPUs)(B200)
   device: b200

.buildkite/test_areas/kernels.yaml

@@ -155,7 +155,7 @@ steps:
   - vllm/utils/import_utils.py
   - tests/kernels/helion/
   commands:
-    - pip install helion==0.3.3
+    - pip install helion==1.0.0
     - pytest -v -s kernels/helion/
 
  
@@ -200,7 +200,14 @@ steps:
   timeout_in_minutes: 90
   device: h100
   num_devices: 2
-  optional: true
+  source_file_dependencies:
+  - csrc/quantization/cutlass_w8a8/moe/
+  - csrc/moe/
+  - tests/kernels/moe
+  - vllm/model_executor/layers/fused_moe/
+  - vllm/model_executor/layers/quantization/
+  - vllm/distributed/device_communicators/
+  - vllm/config
   commands:
     - pytest -v -s kernels/moe/test_moe_layer.py
 
@@ -209,6 +216,13 @@ steps:
   timeout_in_minutes: 90
   device: b200
   num_devices: 2
-  optional: true
+  source_file_dependencies:
+  - csrc/quantization/cutlass_w8a8/moe/
+  - csrc/moe/
+  - tests/kernels/moe
+  - vllm/model_executor/layers/fused_moe/
+  - vllm/model_executor/layers/quantization/
+  - vllm/distributed/device_communicators/
+  - vllm/config
   commands:
     - pytest -v -s kernels/moe/test_moe_layer.py

.buildkite/test_areas/lm_eval.yaml

@@ -91,6 +91,16 @@ steps:
     - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt
 
 
+- label: LM Eval TurboQuant KV Cache
+  timeout_in_minutes: 75
+  source_file_dependencies:
+  - vllm/model_executor/layers/quantization/turboquant/
+  - vllm/v1/attention/backends/turboquant_attn.py
+  - vllm/v1/attention/ops/triton_turboquant_decode.py
+  - vllm/v1/attention/ops/triton_turboquant_store.py
+  commands:
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/models-turboquant.txt
+
 - label: GPQA Eval (GPT-OSS) (H100)
   timeout_in_minutes: 120
   device: h100

.buildkite/test_areas/misc.yaml

@@ -224,6 +224,7 @@ steps:
     - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
     - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA]
     - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[FLASH_ATTN]
+    - pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py
   
 - label: Acceptance Length Test (Large Models) # optional
   timeout_in_minutes: 25

.buildkite/test_areas/models_basic.yaml

@@ -1,10 +1,9 @@
 group: Models - Basic
-depends_on: 
+depends_on:
   - image-build
 steps:
 - label: Basic Models Tests (Initialization)
   timeout_in_minutes: 45
-  device: h200_18gb
   torch_nightly: true
   source_file_dependencies:
   - vllm/
@@ -13,10 +12,11 @@ steps:
   commands:
     # Run a subset of model initialization tests
     - pytest -v -s models/test_initialization.py::test_can_initialize_small_subset
+  mirror:
+    torch_nightly: {}
 
 - label: Basic Models Tests (Extra Initialization) %N
   timeout_in_minutes: 45
-  torch_nightly: true
   source_file_dependencies:
   - vllm/model_executor/models/
   - tests/models/test_initialization.py
@@ -27,6 +27,8 @@ steps:
     # test.) Also run if model initialization test file is modified
     - pytest -v -s models/test_initialization.py -k 'not test_can_initialize_small_subset' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
   parallelism: 2
+  mirror:
+    torch_nightly: {}
 
 - label: Basic Models Tests (Other)
   timeout_in_minutes: 45
@@ -42,10 +44,10 @@ steps:
       device: mi325_1
       depends_on:
       - image-build-amd
-    
+
 
 - label: Basic Models Test (Other CPU) # 5min
-  depends_on: 
+  depends_on:
   - image-build-cpu
   timeout_in_minutes: 10
   source_file_dependencies:
@@ -70,3 +72,18 @@ steps:
     - python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl
     # Whisper needs spawn method to avoid deadlock
     - VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper
+
+- label: Transformers Backward Compatibility Models Test
+  working_dir: "/vllm-workspace/"
+  optional: true
+  soft_fail: true
+  commands:
+    - pip install transformers==4.57.5
+    - pytest -v -s tests/models/test_initialization.py
+    - pytest -v -s tests/models/test_transformers.py
+    - pytest -v -s tests/models/multimodal/processing/
+    - pytest -v -s tests/models/multimodal/test_mapping.py
+    - python3 examples/offline_inference/basic/chat.py
+    - python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl
+    # Whisper needs spawn method to avoid deadlock
+    - VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper

.buildkite/test_areas/models_language.yaml

@@ -1,10 +1,9 @@
 group: Models - Language
-depends_on: 
+depends_on:
   - image-build
 steps:
 - label: Language Models Tests (Standard)
   timeout_in_minutes: 25
-  torch_nightly: true
   source_file_dependencies:
   - vllm/
   - tests/models/language
@@ -12,10 +11,11 @@ steps:
     # Test standard language models, excluding a subset of slow tests
     - pip freeze | grep -E 'torch'
     - pytest -v -s models/language -m 'core_model and (not slow_test)'
+  mirror:
+    torch_nightly: {}
 
 - label: Language Models Tests (Extra Standard) %N
   timeout_in_minutes: 45
-  torch_nightly: true
   source_file_dependencies:
   - vllm/model_executor/models/
   - tests/models/language/pooling/test_embedding.py
@@ -27,10 +27,11 @@ steps:
     - pip freeze | grep -E 'torch'
     - pytest -v -s models/language -m 'core_model and slow_test' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
   parallelism: 2
+  mirror:
+    torch_nightly: {}
 
 - label: Language Models Tests (Hybrid) %N
   timeout_in_minutes: 75
-  torch_nightly: true
   source_file_dependencies:
   - vllm/
   - tests/models/language/generation
@@ -42,6 +43,8 @@ steps:
     # Shard hybrid language model tests
     - pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
   parallelism: 2
+  mirror:
+    torch_nightly: {}
 
 - label: Language Models Test (Extended Generation) # 80min
   timeout_in_minutes: 110
@@ -62,7 +65,7 @@ steps:
       - image-build-amd
       commands:
       - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr'
-      - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.5.2'
+      - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
       - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'
 
 - label: Language Models Test (PPL)

.buildkite/test_areas/spec_decode.yaml

@@ -42,3 +42,16 @@ steps:
     - tests/v1/e2e/spec_decode/
   commands:
     - pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
+
+- label: DFlash Speculators Correctness
+  timeout_in_minutes: 30
+  device: h100
+  optional: true
+  num_devices: 1
+  source_file_dependencies:
+    - vllm/v1/spec_decode/
+    - vllm/model_executor/models/qwen3_dflash.py
+    - tests/v1/spec_decode/test_speculators_dflash.py
+  commands:
+    - export VLLM_ALLOW_INSECURE_SERIALIZATION=1
+    - pytest -v -s v1/spec_decode/test_speculators_dflash.py -m slow_test

.github/CODEOWNERS

@@ -44,8 +44,9 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson
 /vllm/pooling_params.py @noooop @DarkLight1337
 /vllm/tokenizers @DarkLight1337 @njhill
 /vllm/renderers @DarkLight1337 @njhill
-/vllm/reasoning @aarnphm @chaunceyjiang
-/vllm/tool_parsers @aarnphm @chaunceyjiang
+/vllm/reasoning @aarnphm @chaunceyjiang @sfeng33 @bbrowning
+/vllm/tool_parsers @aarnphm @chaunceyjiang @sfeng33 @bbrowning
+/vllm/parser @aarnphm @chaunceyjiang @sfeng33 @bbrowning
 
 # vLLM V1
 /vllm/v1/attention @LucasWilkinson @MatthewBonanni
@@ -91,7 +92,10 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson
 /tests/v1/kv_connector/nixl_integration @NickLucche
 /tests/v1/kv_connector @ApostaC @orozery
 /tests/v1/kv_offload @ApostaC @orozery
-/tests/v1/determinism @yewentao256 
+/tests/v1/determinism @yewentao256
+/tests/reasoning @aarnphm @chaunceyjiang @sfeng33 @bbrowning
+/tests/tool_parsers @aarnphm @chaunceyjiang @sfeng33 @bbrowning
+/tests/tool_use @aarnphm @chaunceyjiang @sfeng33 @bbrowning
 
 # Transformers modeling backend
 /vllm/model_executor/models/transformers @hmellor

.github/PULL_REQUEST_TEMPLATE.md

@@ -15,7 +15,6 @@ PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTT
 - [ ] The test plan, such as providing test command.
 - [ ] The test results, such as pasting the results comparison before and after, or e2e results
 - [ ] (Optional) The necessary documentation update, such as updating `supported_models.md` and `examples` for a new model.
-- [ ] (Optional) Release notes update. If your change is user facing, please update the release notes draft in the [Google Doc](https://docs.google.com/document/d/1YyVqrgX4gHTtrstbq8oWUImOyPCKSGnJ7xtTpmXzlRs/edit?tab=t.0).
 </details>
 
 **BEFORE SUBMITTING, PLEASE READ <https://docs.vllm.ai/en/latest/contributing>** (anything written below this line will be removed by GitHub Actions)

.github/mergify.yml

@@ -264,6 +264,7 @@ pull_request_rules:
       - files=\.buildkite/ci_config_intel.yaml
       - files=vllm/model_executor/layers/fused_moe/xpu_fused_moe.py
       - files=vllm/model_executor/kernels/linear/mixed_precision/xpu.py
+      - files=vllm/model_executor/kernels/linear/mxfp8/xpu.py
       - files=vllm/model_executor/kernels/linear/scaled_mm/xpu.py
       - files=vllm/distributed/device_communicators/xpu_communicator.py
       - files=vllm/v1/attention/backends/mla/xpu_mla_sparse.py
@@ -271,6 +272,7 @@ pull_request_rules:
       - files=vllm/v1/worker/xpu_worker.py
       - files=vllm/v1/worker/xpu_model_runner.py
       - files=vllm/_xpu_ops.py
+      - files=vllm/kernels/xpu_ops.py
       - files~=^vllm/lora/ops/xpu_ops
       - files=vllm/lora/punica_wrapper/punica_xpu.py
       - files=vllm/platforms/xpu.py
@@ -278,7 +280,6 @@ pull_request_rules:
       - title~=(?i)XPU
       - title~=(?i)Intel
       - title~=(?i)BMG
-      - title~=(?i)Arc
   actions:
     label:
       add:

.github/workflows/macos-smoke-test.yml

@@ -45,6 +45,7 @@ jobs:
       - name: Smoke test vllm serve
         run: |
           # Start server in background
+          VLLM_CPU_KVCACHE_SPACE=1 \
           vllm serve Qwen/Qwen3-0.6B \
             --max-model-len=2K \
             --load-format=dummy \

.github/workflows/new_pr_bot.yml

@@ -62,14 +62,14 @@ jobs:
             const prAuthor = context.payload.pull_request.user.login;
 
             const { data: searchResults } = await github.rest.search.issuesAndPullRequests({
-              q: `repo:${owner}/${repo} type:pr author:${prAuthor}`,
+              q: `repo:${owner}/${repo} type:pr is:merged author:${prAuthor}`,
               per_page: 1,
             });
 
-            const authorPRCount = searchResults.total_count;
-            console.log(`Found ${authorPRCount} PRs by ${prAuthor}`);
+            const mergedPRCount = searchResults.total_count;
+            console.log(`Found ${mergedPRCount} merged PRs by ${prAuthor}`);
 
-            if (authorPRCount === 1) {
+            if (mergedPRCount === 0) {
               console.log(`Posting welcome comment for first-time contributor: ${prAuthor}`);
               await github.rest.issues.createComment({
                 owner,
@@ -98,5 +98,5 @@ jobs:
                 ].join('\n'),
               });
             } else {
-              console.log(`Skipping comment for ${prAuthor} - not their first PR (${authorPRCount} PRs found)`);
+              console.log(`Skipping comment for ${prAuthor} - not a first-time contributor (${mergedPRCount} merged PRs)`);
             }

CMakeLists.txt

@@ -923,6 +923,14 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
       SRCS "${SRCS}"
       CUDA_ARCHS "${FP4_ARCHS}")
     list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
+    # nvfp4_kv_cache_kernels uses non-stable torch API and is called directly
+    # from cache_kernels.cu, so it belongs in _C rather than _C_stable.
+    set(NVFP4_KV_SRC "csrc/nvfp4_kv_cache_kernels.cu")
+    set_gencode_flags_for_srcs(
+      SRCS "${NVFP4_KV_SRC}"
+      CUDA_ARCHS "${FP4_ARCHS}")
+    target_sources(_C PRIVATE ${NVFP4_KV_SRC})
+    target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM120=1)
     list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM120=1")
     list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM120=1")
     message(STATUS "Building NVFP4 for archs: ${FP4_ARCHS}")
@@ -949,6 +957,12 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
       SRCS "${SRCS}"
       CUDA_ARCHS "${FP4_ARCHS}")
     list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
+    set(NVFP4_KV_SRC "csrc/nvfp4_kv_cache_kernels.cu")
+    set_gencode_flags_for_srcs(
+      SRCS "${NVFP4_KV_SRC}"
+      CUDA_ARCHS "${FP4_ARCHS}")
+    target_sources(_C PRIVATE ${NVFP4_KV_SRC})
+    target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM100=1)
     list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM100=1")
     list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
     message(STATUS "Building NVFP4 for archs: ${FP4_ARCHS}")
@@ -1226,6 +1240,7 @@ endif()
 if (VLLM_GPU_LANG STREQUAL "CUDA")
     include(cmake/external_projects/deepgemm.cmake)
     include(cmake/external_projects/flashmla.cmake)
+    include(cmake/external_projects/cutlass_fa3.cmake)
     include(cmake/external_projects/qutlass.cmake)
 
     # vllm-flash-attn should be last as it overwrites some CMake functions

cmake/cpu_extension.cmake

@@ -30,6 +30,21 @@ else()
     list(APPEND CXX_COMPILE_FLAGS
         "-fopenmp"
         "-DVLLM_CPU_EXTENSION")
+
+    # locate PyTorch's libgomp (e.g. site-packages/torch.libs/libgomp-947d5fa1.so.1.0.0)
+    # and create a local shim dir with it
+    vllm_prepare_torch_gomp_shim(VLLM_TORCH_GOMP_SHIM_DIR)
+
+    find_library(OPEN_MP
+        NAMES gomp
+        PATHS ${VLLM_TORCH_GOMP_SHIM_DIR}
+        NO_DEFAULT_PATH
+        REQUIRED
+    )
+    # Set LD_LIBRARY_PATH to include the shim dir at build time to use the same libgomp as PyTorch
+    if (OPEN_MP)
+        set(ENV{LD_LIBRARY_PATH} "${VLLM_TORCH_GOMP_SHIM_DIR}:$ENV{LD_LIBRARY_PATH}")
+    endif()
 endif()
 
 if (NOT MACOSX_FOUND)
@@ -175,20 +190,6 @@ if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND
         if(NOT NPROC)
             set(NPROC 4)
         endif()
-        # locate PyTorch's libgomp (e.g. site-packages/torch.libs/libgomp-947d5fa1.so.1.0.0)
-        # and create a local shim dir with it
-        vllm_prepare_torch_gomp_shim(VLLM_TORCH_GOMP_SHIM_DIR)
-
-        find_library(OPEN_MP
-            NAMES gomp
-            PATHS ${VLLM_TORCH_GOMP_SHIM_DIR}
-            NO_DEFAULT_PATH
-            REQUIRED
-        )
-        # Set LD_LIBRARY_PATH to include the shim dir at build time to use the same libgomp as PyTorch
-        if (OPEN_MP)
-            set(ENV{LD_LIBRARY_PATH} "${VLLM_TORCH_GOMP_SHIM_DIR}:$ENV{LD_LIBRARY_PATH}")
-        endif()
 
         # Fetch and populate ACL
         if(DEFINED ENV{ACL_ROOT_DIR} AND IS_DIRECTORY "$ENV{ACL_ROOT_DIR}")
@@ -360,6 +361,7 @@ set(VLLM_EXT_SRC
 if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND)
     set(VLLM_EXT_SRC
         "csrc/cpu/shm.cpp"
+        "csrc/cpu/activation_lut_bf16.cpp"
         ${VLLM_EXT_SRC})
 endif()
 

cmake/external_projects/cutlass_fa3.cmake

@@ -0,0 +1,163 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# CUTLASS FA3 MLA Sparse Attention — requires CUDA >= 12.4, SM90a
+#
+# Vendors the sgl-attn CUTLASS FlashAttention3 kernel from SGLang into vLLM
+# as a self-contained extension (_cutlass_fa3_C). This provides a high-
+# performance sparse MLA attention kernel for SM90 (Hopper) GPUs.
+#
+# Source: https://github.com/sgl-project/sgl-attn (commit bcf72ccc)
+# CUTLASS: https://github.com/NVIDIA/cutlass (commit 57e3cfb4)
+
+# Guard: CUDA >= 12.4 required for SM90a features used by FA3
+if(NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL "12.4")
+  message(STATUS "Skipping CUTLASS FA3: requires CUDA >= 12.4")
+  # Create empty target so setup.py doesn't fail on unsupported systems
+  add_custom_target(_cutlass_fa3_C)
+  return()
+endif()
+
+# Guard: SM90 architecture required
+set(CUTLASS_FA3_SUPPORT_ARCHS)
+list(APPEND CUTLASS_FA3_SUPPORT_ARCHS "9.0a")
+cuda_archs_loose_intersection(
+  CUTLASS_FA3_ARCHS "${CUTLASS_FA3_SUPPORT_ARCHS}" "${CUDA_ARCHS}")
+if(NOT CUTLASS_FA3_ARCHS)
+  message(STATUS "Skipping CUTLASS FA3: requires SM90 (CUDA_ARCHS=${CUDA_ARCHS})")
+  add_custom_target(_cutlass_fa3_C)
+  return()
+endif()
+
+include(FetchContent)
+
+# Fetch sgl-attn (Flash Attention 3 kernels from SGLang)
+# We only need the source files, not the build system, so we use
+# FetchContent_Populate to download without building.
+if (DEFINED ENV{SGL_ATTN_SRC_DIR})
+  set(SGL_ATTN_SRC_DIR $ENV{SGL_ATTN_SRC_DIR})
+endif()
+if(SGL_ATTN_SRC_DIR)
+  FetchContent_Declare(cutlass_fa3
+    SOURCE_DIR ${SGL_ATTN_SRC_DIR})
+else()
+  FetchContent_Declare(cutlass_fa3
+    GIT_REPOSITORY https://github.com/sgl-project/sgl-attn.git
+    GIT_TAG bcf72ccc6816b36a5fae2c5a3c027604629785e0
+    GIT_PROGRESS TRUE
+    GIT_SHALLOW FALSE)
+endif()
+FetchContent_GetProperties(cutlass_fa3)
+if(NOT cutlass_fa3_POPULATED)
+  FetchContent_Populate(cutlass_fa3)
+endif()
+message(STATUS "CUTLASS FA3 sgl-attn source: ${cutlass_fa3_SOURCE_DIR}")
+
+# Fetch CUTLASS for FA3 (headers only, separate from vLLM's main CUTLASS
+# to avoid version conflicts). Use FetchContent_Populate to avoid running
+# CUTLASS's own CMakeLists.txt which would create conflicting targets.
+if (DEFINED ENV{CUTLASS_FA3_CUTLASS_SRC_DIR})
+  set(CUTLASS_FA3_CUTLASS_SRC_DIR $ENV{CUTLASS_FA3_CUTLASS_SRC_DIR})
+endif()
+if(CUTLASS_FA3_CUTLASS_SRC_DIR)
+  FetchContent_Declare(cutlass_for_fa3
+    SOURCE_DIR ${CUTLASS_FA3_CUTLASS_SRC_DIR})
+else()
+  FetchContent_Declare(cutlass_for_fa3
+    GIT_REPOSITORY https://github.com/NVIDIA/cutlass.git
+    GIT_TAG 57e3cfb47a2d9e0d46eb6335c3dc411498efa198
+    GIT_PROGRESS TRUE
+    GIT_SHALLOW FALSE)
+endif()
+FetchContent_GetProperties(cutlass_for_fa3)
+if(NOT cutlass_for_fa3_POPULATED)
+  FetchContent_Populate(cutlass_for_fa3)
+endif()
+message(STATUS "CUTLASS FA3 cutlass source: ${cutlass_for_fa3_SOURCE_DIR}")
+
+set(FA3_SRC "${cutlass_fa3_SOURCE_DIR}/hopper")
+
+# flash_api.cpp dispatches to all head dimensions + dtypes (BF16, FP16, FP8)
+# at compile time. With FLASHATTENTION_DISABLE_SM8x, only SM90 instantiations
+# are needed. We exclude hdimall_* (fails on CUDA 13+) and backward files.
+file(GLOB FA3_INSTANTIATION_SOURCES
+  # BF16 instantiations
+  "${FA3_SRC}/instantiations/flash_fwd_hdim64_bf16*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdim96_bf16*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdim128_bf16*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdim192_bf16*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdim256_bf16*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdimdiff_bf16*_sm90.cu"
+  # FP16 instantiations
+  "${FA3_SRC}/instantiations/flash_fwd_hdim64_fp16*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdim96_fp16*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdim128_fp16*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdim192_fp16*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdim256_fp16*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdimdiff_fp16*_sm90.cu"
+  # FP8 (e4m3) instantiations
+  "${FA3_SRC}/instantiations/flash_fwd_hdim64_e4m3*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdim96_e4m3*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdim128_e4m3*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdim192_e4m3*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdim256_e4m3*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdimdiff_e4m3*_sm90.cu")
+
+set(FA3_CORE_SOURCES
+  "${FA3_SRC}/flash_api.cpp"
+  "${FA3_SRC}/flash_prepare_scheduler.cu"
+  "${FA3_SRC}/flash_fwd_combine.cu")
+
+set(FA3_ALL_SOURCES
+  "${CMAKE_CURRENT_SOURCE_DIR}/csrc/cutlass_fa3_extension.cc"
+  ${FA3_CORE_SOURCES}
+  ${FA3_INSTANTIATION_SOURCES})
+
+set(FA3_INCLUDE_DIRS
+  ${FA3_SRC}
+  ${cutlass_fa3_SOURCE_DIR}/include
+  ${cutlass_for_fa3_SOURCE_DIR}/include
+  ${cutlass_for_fa3_SOURCE_DIR}/tools/util/include
+  ${CMAKE_CURRENT_SOURCE_DIR}/csrc)
+
+# Set SM90a gencode flags for all FA3 CUDA sources
+set_gencode_flags_for_srcs(
+  SRCS "${FA3_ALL_SOURCES}"
+  CUDA_ARCHS "${CUTLASS_FA3_ARCHS}")
+
+define_extension_target(_cutlass_fa3_C
+  DESTINATION vllm
+  LANGUAGE ${VLLM_GPU_LANG}
+  SOURCES ${FA3_ALL_SOURCES}
+  COMPILE_FLAGS ${VLLM_GPU_FLAGS}
+  ARCHITECTURES ${VLLM_GPU_ARCHES}
+  INCLUDE_DIRECTORIES ${FA3_INCLUDE_DIRS}
+  USE_SABI 3
+  WITH_SOABI)
+
+# FA3-specific compile options for CUDA and C++ source files:
+# - C++17 required by CUTLASS
+# - Fast math for performance
+# - Relaxed constexpr for CUTLASS template metaprogramming
+# - Disable backward pass, dropout, uneven K (not needed for inference)
+# - Enable varlen-only mode (all our use cases are variable-length)
+target_compile_options(_cutlass_fa3_C PRIVATE
+  $<$<COMPILE_LANGUAGE:CUDA>:-UPy_LIMITED_API>
+  $<$<COMPILE_LANGUAGE:CXX>:-UPy_LIMITED_API>
+  $<$<COMPILE_LANGUAGE:CUDA>:-std=c++17>
+  $<$<COMPILE_LANGUAGE:CXX>:-std=c++17>
+  $<$<COMPILE_LANGUAGE:CUDA>:--use_fast_math>
+  $<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr>)
+
+target_compile_definitions(_cutlass_fa3_C PRIVATE
+  CUTE_USE_PACKED_TUPLE=1
+  CUTLASS_ENABLE_GDC_FOR_SM90
+  CUTE_SM90_EXTENDED_MMA_SHAPES_ENABLED
+  CUTLASS_ENABLE_TENSOR_CORE_MMA=1
+  FLASHATTENTION_DISABLE_BACKWARD
+  FLASHATTENTION_DISABLE_DROPOUT
+  FLASHATTENTION_DISABLE_UNEVEN_K
+  FLASHATTENTION_DISABLE_SM8x
+  FLASHATTENTION_VARLEN_ONLY)
+
+message(STATUS "CUTLASS FA3 MLA Sparse: enabled for SM90 (${CUTLASS_FA3_ARCHS})")

csrc/cache_kernels.cu

@@ -724,6 +724,28 @@ void reshape_and_cache_flash(
   int num_tokens = slot_mapping.size(0);
   int num_heads = key.size(1);
   int head_size = key.size(2);
+
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(key));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+  if (kv_cache_dtype == "nvfp4") {
+#if defined(ENABLE_NVFP4_SM100) || defined(ENABLE_NVFP4_SM120)
+    // NVFP4 dispatch is compiled separately for SM100+.
+    extern void reshape_and_cache_nvfp4_dispatch(
+        torch::Tensor & key, torch::Tensor & value, torch::Tensor & key_cache,
+        torch::Tensor & value_cache, torch::Tensor & slot_mapping,
+        torch::Tensor & k_scale, torch::Tensor & v_scale);
+    reshape_and_cache_nvfp4_dispatch(key, value, key_cache, value_cache,
+                                     slot_mapping, k_scale, v_scale);
+    return;
+#else
+    TORCH_CHECK(false,
+                "NVFP4 KV cache requires SM100+ (Blackwell). "
+                "Please rebuild vllm with a Blackwell-compatible CUDA target.");
+#endif
+  }
+
+  // Original FP8/auto path.
   int block_size = key_cache.size(1);
 
   int64_t key_stride = key.stride(0);
@@ -741,8 +763,6 @@ void reshape_and_cache_flash(
 
   dim3 grid(num_tokens);
   dim3 block(std::min(num_heads * head_size, 512));
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(key));
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
 
   DISPATCH_BY_KV_CACHE_DTYPE(key.dtype(), kv_cache_dtype,
                              CALL_RESHAPE_AND_CACHE_FLASH);

csrc/cpu/activation_lut_bf16.cpp

@@ -0,0 +1,71 @@
+#include "cpu_types.hpp"
+
+#include <array>
+#include <cstdint>
+#include <mutex>
+#include <string>
+
+#include <ATen/ops/empty.h>
+#include <ATen/ops/gelu.h>
+#include <c10/util/BFloat16.h>
+
+constexpr uint32_t ActivationLutSize = 1u << 16;
+
+at::Tensor gelu_reference(const at::Tensor& x) { return at::gelu(x, "none"); }
+
+void maybe_init_activation_lut_bf16(
+    uint16_t* lut, std::once_flag& once,
+    at::Tensor (*activation)(const at::Tensor&)) {
+  std::call_once(once, [&]() {
+    auto lut_input =
+        at::empty({static_cast<int64_t>(ActivationLutSize)},
+                  at::TensorOptions().device(at::kCPU).dtype(at::kFloat));
+    auto* lut_input_ptr = lut_input.data_ptr<float>();
+#pragma omp parallel for
+    for (uint32_t i = 0; i < ActivationLutSize; ++i) {
+      lut_input_ptr[i] = c10::detail::f32_from_bits(static_cast<uint16_t>(i));
+    }
+
+    auto lut_output = activation(lut_input);
+    const auto* lut_output_ptr = lut_output.data_ptr<float>();
+#pragma omp parallel for
+    for (uint32_t i = 0; i < ActivationLutSize; ++i) {
+      lut[i] = c10::detail::round_to_nearest_even(lut_output_ptr[i]);
+    }
+  });
+}
+
+void activation_lut_bf16(torch::Tensor& out, torch::Tensor& input,
+                         const uint16_t* lut, const char* op_name) {
+  TORCH_CHECK(input.scalar_type() == at::kBFloat16, op_name,
+              ": input must be bfloat16");
+  TORCH_CHECK(out.scalar_type() == at::kBFloat16, op_name,
+              ": out must be bfloat16");
+  TORCH_CHECK(input.is_contiguous(), op_name, ": input must be contiguous");
+  TORCH_CHECK(out.is_contiguous(), op_name, ": out must be contiguous");
+
+  const auto* src =
+      reinterpret_cast<const uint16_t*>(input.data_ptr<at::BFloat16>());
+  auto* dst = reinterpret_cast<uint16_t*>(out.data_ptr<at::BFloat16>());
+  const int64_t n = input.numel();
+
+  CPU_KERNEL_GUARD_IN(activation_lut_bf16_impl)
+#pragma omp parallel for
+  for (int64_t i = 0; i < n; ++i) {
+    dst[i] = lut[src[i]];
+  }
+  CPU_KERNEL_GUARD_OUT(activation_lut_bf16_impl)
+}
+
+void activation_lut_bf16(torch::Tensor& out, torch::Tensor& input,
+                         const std::string& activation) {
+  if (activation == "gelu") {
+    static std::array<uint16_t, ActivationLutSize> lut{};
+    static std::once_flag once;
+    maybe_init_activation_lut_bf16(lut.data(), once, gelu_reference);
+    activation_lut_bf16(out, input, lut.data(), "gelu_lut");
+    return;
+  }
+
+  TORCH_CHECK(false, "Unsupported activation: ", activation);
+}

csrc/cpu/cpu_attn_impl.hpp

@@ -147,6 +147,9 @@ struct AttentionMetadata {
       case ISA::NEON:
         ss << "NEON, ";
         break;
+      case ISA::VXE:
+        ss << "VXE, ";
+        break;
     }
     ss << "workitem_group_num: " << workitem_group_num
        << ", reduction_item_num: " << reduction_item_num

csrc/cpu/torch_bindings.cpp

@@ -85,6 +85,9 @@ at::Tensor int4_scaled_mm_cpu(at::Tensor& x, at::Tensor& w, at::Tensor& w_zeros,
                               at::Tensor& w_scales,
                               std::optional<at::Tensor> bias);
 
+void activation_lut_bf16(torch::Tensor& out, torch::Tensor& input,
+                         const std::string& activation);
+
 torch::Tensor get_scheduler_metadata(
     const int64_t num_req, const int64_t num_heads_q,
     const int64_t num_heads_kv, const int64_t head_dim,
@@ -138,6 +141,8 @@ void compute_slot_mapping_kernel_impl(const torch::Tensor query_start_loc,
                                       torch::Tensor slot_mapping,
                                       const int64_t block_size);
 
+void init_cpu_memory_env(std::vector<int64_t> node_ids);
+
 namespace cpu_utils {
 void eagle_prepare_inputs_padded_kernel_impl(
     const torch::Tensor& cu_num_draft_tokens,
@@ -231,6 +236,15 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.def("gelu_quick(Tensor! out, Tensor input) -> ()");
   ops.impl("gelu_quick", torch::kCPU, &gelu_quick);
 
+#if (defined(__aarch64__) && !defined(__APPLE__))
+
+  ops.def(
+      "activation_lut_bf16(Tensor! out, Tensor input, str activation)"
+      " -> ()");
+  ops.impl("activation_lut_bf16", torch::kCPU, &activation_lut_bf16);
+
+#endif  // (defined(__aarch64__) && !defined(__APPLE__))
+
   // Layernorm
   // Apply Root Mean Square (RMS) Normalization to the input tensor.
   ops.def(
@@ -419,6 +433,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       "block_size) -> ()",
       &compute_slot_mapping_kernel_impl);
 
+  ops.def("init_cpu_memory_env(SymInt[] node_ids) -> ()", &init_cpu_memory_env);
+
   // Speculative decoding kernels
   ops.def(
       "eagle_prepare_inputs_padded_kernel_impl(Tensor cu_num_draft_tokens, "

csrc/cpu/utils.cpp

@@ -13,13 +13,80 @@
 #include "cpu/utils.hpp"
 
 #ifdef VLLM_NUMA_DISABLED
-std::string init_cpu_threads_env(const std::string& cpu_ids) {
-  return std::string(
-      "Warning: NUMA is not enabled in this build. `init_cpu_threads_env` has "
-      "no effect to setup thread affinity.");
-}
+void init_cpu_memory_env(std::vector<int64_t> node_ids) {}
+#else
+void init_cpu_memory_env(std::vector<int64_t> node_ids) {
+  // Memory node binding
+  if (numa_available() != -1) {
+    // Concatenate all node_ids into a single comma-separated string
+    if (!node_ids.empty()) {
+      std::string node_ids_str;
+      for (const int node_id : node_ids) {
+        if (!node_ids_str.empty()) {
+          node_ids_str += ",";
+        }
+        node_ids_str += std::to_string(node_id);
+      }
 
-#endif
+      bitmask* mask = numa_parse_nodestring(node_ids_str.c_str());
+      bitmask* src_mask = numa_get_mems_allowed();
+
+      int pid = getpid();
+
+      if (mask && src_mask) {
+        // move all existing pages to the specified numa node.
+        *(src_mask->maskp) = *(src_mask->maskp) ^ *(mask->maskp);
+        int page_num = numa_migrate_pages(pid, src_mask, mask);
+        if (page_num == -1) {
+          TORCH_WARN("numa_migrate_pages failed. errno: " +
+                     std::to_string(errno));
+        }
+
+        // Restrict memory allocation to the selected NUMA node(s).
+        // Enhances memory locality for the threads bound to those NUMA CPUs.
+        if (node_ids.size() > 1) {
+          errno = 0;
+          numa_set_interleave_mask(mask);
+          if (errno != 0) {
+            TORCH_WARN("numa_set_interleave_mask failed. errno: " +
+                       std::to_string(errno));
+          } else {
+            TORCH_WARN(
+                "NUMA binding: Using INTERLEAVE policy for memory "
+                "allocation across multiple NUMA nodes (nodes: " +
+                node_ids_str +
+                "). Memory allocations will be "
+                "interleaved across the specified NUMA nodes.");
+          }
+        } else {
+          errno = 0;
+          numa_set_membind(mask);
+          if (errno != 0) {
+            TORCH_WARN("numa_set_membind failed. errno: " +
+                       std::to_string(errno));
+          } else {
+            TORCH_WARN(
+                "NUMA binding: Using MEMBIND policy for memory "
+                "allocation on the NUMA nodes (" +
+                node_ids_str +
+                "). Memory allocations will be "
+                "strictly bound to these NUMA nodes.");
+          }
+        }
+
+        numa_set_strict(1);
+
+        numa_free_nodemask(mask);
+        numa_free_nodemask(src_mask);
+      } else {
+        TORCH_WARN(
+            "numa_parse_nodestring or numa_get_run_node_mask failed. errno: " +
+            std::to_string(errno));
+      }
+    }
+  }
+}
+#endif  // VLLM_NUMA_DISABLED
 
 namespace cpu_utils {
 ScratchPadManager::ScratchPadManager() : size_(0), ptr_(nullptr) {

csrc/cpu/utils.hpp

@@ -54,12 +54,34 @@ struct Counter {
 };
 
 inline int64_t get_available_l2_size() {
+#if defined(__s390x__)
+  static int64_t size = []() {
+    uint32_t l2_cache_size = 0;
+    auto caps = at::cpu::get_cpu_capabilities();
+    auto it = caps.find("l2_cache_size");
+    if (it != caps.end()) {
+      l2_cache_size = static_cast<uint32_t>(it->second.toInt());
+    }
+    if (l2_cache_size == 0) {
+      long sys_l2 = sysconf(_SC_LEVEL2_CACHE_SIZE);
+      if (sys_l2 > 0) {
+        l2_cache_size = static_cast<uint32_t>(sys_l2);
+      }
+    }
+    if (l2_cache_size == 0) {
+      l2_cache_size = 256 * 1024;
+    }
+    return static_cast<int64_t>(l2_cache_size) >> 1;  // use 50% of L2 cache
+  }();
+  return size;
+#else
   static int64_t size = []() {
     auto caps = at::cpu::get_cpu_capabilities();
     const uint32_t l2_cache_size = caps.at("l2_cache_size").toInt();
     return l2_cache_size >> 1;  // use 50% of L2 cache
   }();
   return size;
+#endif
 }
 
 template <int32_t alignment_v, typename T>

csrc/cutlass_fa3_extension.cc

@@ -0,0 +1,72 @@
+/* SPDX-License-Identifier: Apache-2.0
+ * SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+ *
+ * Vendored CUTLASS FA3 MLA attention kernel binding for vLLM.
+ * Based on sgl-kernel/csrc/flash_extension.cc from SGLang.
+ *
+ * This registers the FA3 forward pass as a PyTorch C++ extension under
+ * the _cutlass_fa3_C namespace, enabling torch.ops._cutlass_fa3_C.fwd().
+ *
+ * Original source:
+ *   https://github.com/sgl-project/sgl-attn (commit bcf72ccc)
+ *   sgl-kernel/csrc/flash_extension.cc
+ */
+#include <Python.h>
+#include <ATen/core/dispatch/Dispatcher.h>
+#include <torch/all.h>
+#include <torch/library.h>
+
+#include "sgl_flash_kernel_ops.h"
+
+TORCH_LIBRARY_FRAGMENT(_cutlass_fa3_C, m) {
+  /*
+   * CUTLASS FA3 MLA forward pass.
+   * Signature matches sgl-attn's mha_fwd() exactly.
+   */
+  m.def(
+      "fwd(Tensor   q,"
+      "    Tensor   k,"
+      "    Tensor   v,"
+      "    Tensor?  k_new,"
+      "    Tensor?  v_new,"
+      "    Tensor?  q_v,"
+      "    Tensor?  out,"
+      "    Tensor?  cu_seqlens_q,"
+      "    Tensor?  cu_seqlens_k,"
+      "    Tensor?  cu_seqlens_k_new,"
+      "    Tensor?  seqused_q,"
+      "    Tensor?  seqused_k,"
+      "    int?     max_seqlen_q,"
+      "    int?     max_seqlen_k,"
+      "    Tensor?  page_table,"
+      "    Tensor?  kv_batch_idx,"
+      "    Tensor?  leftpad_k,"
+      "    Tensor?  rotary_cos,"
+      "    Tensor?  rotary_sin,"
+      "    Tensor?  seqlens_rotary,"
+      "    Tensor?  q_descale,"
+      "    Tensor?  k_descale,"
+      "    Tensor?  v_descale,"
+      "    float?   softmax_scale,"
+      "    bool     is_causal,"
+      "    int      window_size_left,"
+      "    int      window_size_right,"
+      "    int      attention_chunk,"
+      "    float    softcap,"
+      "    bool     is_rotary_interleaved,"
+      "    Tensor?  scheduler_metadata,"
+      "    int      num_splits,"
+      "    bool?    pack_gqa,"
+      "    int      sm_margin,"
+      "    Tensor?  sinks"
+      ") -> (Tensor, Tensor, Tensor, Tensor)");
+
+  m.impl("fwd", torch::kCUDA, make_pytorch_shim(&mha_fwd));
+}
+
+// Python module initialization for _cutlass_fa3_C
+PyMODINIT_FUNC PyInit__cutlass_fa3_C() {
+  static struct PyModuleDef module = {PyModuleDef_HEAD_INIT, "_cutlass_fa3_C",
+                                      nullptr, 0, nullptr};
+  return PyModule_Create(&module);
+}

csrc/nvfp4_kv_cache_kernels.cu

@@ -0,0 +1,275 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+// NVFP4 KV cache store kernel.
+// Quantizes bf16 key/value to packed FP4 + FP8 block scales and writes them
+// into the paged KV cache.
+//
+// Per page layout: [K_data | K_scale | V_data | V_scale]
+// Both data and scale regions are contiguous per head, enabling direct
+// TMA descriptor use.
+//
+// Reuses device functions from nvfp4_utils.cuh:
+//   - cvt_warp_fp16_to_fp4()  for bf16 → fp4 quantization + block scale
+//   - pack_fp4()              for packing float pairs to fp4
+//   - reciprocal_approximate_ftz() for fast reciprocal
+
+#define NVFP4_ENABLE_ELTS16 1
+#include "libtorch_stable/quantization/fp4/nvfp4_utils.cuh"
+
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <torch/all.h>
+
+#include "dispatch_utils.h"
+
+namespace vllm {
+
+// Compute swizzled scale offset for SM100 trtllm-gen MHA kernel.
+// The swizzle pattern for HND layout is:
+//   [T//4, 4, 4, S//4] → permute(0, 2, 3, 1) → reshape to [T, S]
+// where T = block_size (page_size), S = scale_dim = head_size // 16.
+//
+// For a linear (t, s) position, the swizzled position is:
+//   swizzled_t = (t / 4) * 4 + (s / (S / 4))
+//   swizzled_s = (s % (S / 4)) * 4 + (t % 4)
+__device__ __forceinline__ int swizzle_scale_offset(int t, int s,
+                                                    int scale_dim) {
+  int s_group = scale_dim / 4;
+  int swizzled_t = (t / 4) * 4 + (s / s_group);
+  int swizzled_s = (s % s_group) * 4 + (t % 4);
+  return swizzled_t * scale_dim + swizzled_s;
+}
+
+// Kernel: quantize bf16 key/value to NVFP4 and store in paged KV cache.
+//
+// Takes separate data and scale cache pointers for K and V.
+// Within each KV side, data and scale are separate contiguous regions.
+//
+// Threading: one CUDA block per token, threads process heads and
+// groups of 16 elements within each head.
+template <typename scalar_t>
+__global__ void reshape_and_cache_nvfp4_kernel(
+    const scalar_t* __restrict__ key,      // [num_tokens, num_heads, head_size]
+    const scalar_t* __restrict__ value,    // [num_tokens, num_heads, head_size]
+    uint8_t* __restrict__ key_data_cache,  // data region for K
+    uint8_t* __restrict__ value_data_cache,    // data region for V
+    uint8_t* __restrict__ key_scale_cache,     // scale region for K
+    uint8_t* __restrict__ value_scale_cache,   // scale region for V
+    const int64_t* __restrict__ slot_mapping,  // [num_actual_tokens]
+    const float* __restrict__ k_scale_ptr,     // pointer to checkpoint k_scale
+    const float* __restrict__ v_scale_ptr,     // pointer to checkpoint v_scale
+    const int64_t key_stride,                  // key.stride(0) in elements
+    const int64_t value_stride,                // value.stride(0) in elements
+    const int num_heads, const int head_size, const int block_size,
+    const int64_t data_block_stride,         // data cache stride for dim 0
+    const int64_t data_head_stride,          // data cache stride for heads
+    const int64_t data_block_offset_stride,  // data cache stride for tokens
+    const int64_t scale_block_stride,        // scale cache stride for dim 0
+    const int64_t scale_head_stride,         // scale cache stride for heads
+    const int64_t scale_block_offset_stride  // scale cache stride for tokens
+) {
+  using CudaType = typename CUDATypeConverter<scalar_t>::Type;
+  using PVec = PackedVec<CudaType, CVT_FP4_PACK16>;
+
+  static constexpr int ELTS = CVT_FP4_ELTS_PER_THREAD;  // 16 or 8
+  static constexpr int THREADS_PER_SF = CVT_FP4_SF_VEC_SIZE / ELTS;
+
+  const int64_t token_idx = blockIdx.x;
+  const int64_t slot_idx = slot_mapping[token_idx];
+  if (slot_idx < 0) return;
+
+  const int64_t block_idx = slot_idx / block_size;
+  const int block_offset = static_cast<int>(slot_idx % block_size);
+
+  const int scale_dim = head_size / 16;
+  const int groups_per_head = head_size / CVT_FP4_SF_VEC_SIZE;
+
+  const int total_groups = num_heads * groups_per_head;
+  const int tid = threadIdx.x;
+  const int num_thread_groups = blockDim.x / THREADS_PER_SF;
+  const int tg_id = tid / THREADS_PER_SF;
+  const int tg_lane = tid % THREADS_PER_SF;
+
+  // Process both K (kv=0) and V (kv=1)
+#pragma unroll
+  for (int kv = 0; kv < 2; kv++) {
+    const scalar_t* __restrict__ src = (kv == 0) ? key : value;
+    const float global_scale = 1.0f / ((kv == 0) ? *k_scale_ptr : *v_scale_ptr);
+    const int64_t src_stride = (kv == 0) ? key_stride : value_stride;
+    uint8_t* __restrict__ data_cache =
+        (kv == 0) ? key_data_cache : value_data_cache;
+    uint8_t* __restrict__ sc_cache =
+        (kv == 0) ? key_scale_cache : value_scale_cache;
+
+    // Source pointer for this token (use actual stride, not assumed contiguous)
+    const CudaType* __restrict__ token_src =
+        reinterpret_cast<const CudaType*>(src) + token_idx * src_stride;
+
+    // Destination bases in data and scale caches for this token's block
+    uint8_t* __restrict__ data_block =
+        data_cache + block_idx * data_block_stride;
+    uint8_t* __restrict__ scale_block =
+        sc_cache + block_idx * scale_block_stride;
+
+    for (int g = tg_id; g < total_groups; g += num_thread_groups) {
+      const int head = g / groups_per_head;
+      const int group_in_head = g % groups_per_head;
+
+      // Load 16 (or 8) bf16 elements from source
+      PVec in_vec;
+      const CudaType* __restrict__ src_ptr =
+          token_src + head * head_size + group_in_head * CVT_FP4_SF_VEC_SIZE +
+          tg_lane * ELTS;
+
+#pragma unroll
+      for (int i = 0; i < ELTS / 2; i++) {
+        in_vec.elts[i] = reinterpret_cast<
+            const typename PackedTypeConverter<CudaType>::Type*>(src_ptr)[i];
+      }
+
+      // Quantize: produces packed fp4 and writes scale factor.
+      uint8_t sf_val;
+      uint8_t* sf_out_ptr = (tg_lane == 0) ? &sf_val : nullptr;
+
+      fp4_packed_t packed = cvt_warp_fp16_to_fp4<CudaType, THREADS_PER_SF>(
+          in_vec, global_scale, sf_out_ptr);
+
+      // Write packed FP4 data to data cache
+      uint8_t* __restrict__ data_dst = data_block + head * data_head_stride +
+                                       block_offset * data_block_offset_stride;
+
+#if CVT_FP4_PACK16
+      {
+        // 16 elements → 8 bytes (u32x2)
+        int data_byte_offset = group_in_head * 8;
+        reinterpret_cast<uint64_t*>(data_dst + data_byte_offset)[0] =
+            (uint64_t(packed.hi) << 32) | uint64_t(packed.lo);
+      }
+#else
+      {
+        // 8 elements → 4 bytes (uint32_t)
+        int data_byte_offset =
+            group_in_head * CVT_FP4_SF_VEC_SIZE / 2 + tg_lane * ELTS / 2;
+        reinterpret_cast<uint32_t*>(data_dst + data_byte_offset)[0] = packed;
+      }
+#endif
+
+      // Write block scale to scale cache.
+      // K (kv==0): linear layout (no swizzle).
+      // V (kv==1): swizzled layout for SM100 trtllm-gen MHA kernel.
+      if (sf_out_ptr != nullptr) {
+        int scale_idx = group_in_head;
+        uint8_t* __restrict__ scale_dst;
+        if (kv == 0) {
+          scale_dst = scale_block + head * scale_head_stride +
+                      block_offset * scale_block_offset_stride + scale_idx;
+        } else {
+          int swizzled_offset =
+              swizzle_scale_offset(block_offset, scale_idx, scale_dim);
+          int swizzled_t = swizzled_offset / scale_dim;
+          int swizzled_s = swizzled_offset % scale_dim;
+          scale_dst = scale_block + head * scale_head_stride +
+                      swizzled_t * scale_block_offset_stride + swizzled_s;
+        }
+        *scale_dst = sf_val;
+      }
+    }
+  }
+}
+
+}  // namespace vllm
+
+// Non-template entry point callable from cache_kernels.cu.
+// Receives key_cache/value_cache as kv_cache[:, 0] and kv_cache[:, 1].
+// Each KV side contains both data and scale:
+//   page = [K_data | K_scale | V_data | V_scale]
+void reshape_and_cache_nvfp4_dispatch(torch::Tensor& key, torch::Tensor& value,
+                                      torch::Tensor& key_cache,
+                                      torch::Tensor& value_cache,
+                                      torch::Tensor& slot_mapping,
+                                      torch::Tensor& k_scale,
+                                      torch::Tensor& v_scale) {
+  int num_tokens = slot_mapping.size(0);
+  int num_heads = key.size(1);
+  int head_size = key.size(2);
+  int data_dim = head_size / 2;
+  int scale_dim = head_size / 16;
+  int full_dim = data_dim + scale_dim;
+
+  // key_cache is kv_cache[:, 0] with shape
+  // [num_blocks, block_size, num_heads, full_dim] in logical order.
+  // Strides encode the physical layout (HND or NHD).
+  TORCH_CHECK(key_cache.dim() == 4, "key_cache must be 4D");
+  TORCH_CHECK(key_cache.size(3) == full_dim,
+              "key_cache last dim must be data_dim + scale_dim, got ",
+              key_cache.size(3), " expected ", full_dim);
+
+  int block_size = key_cache.size(1);
+
+  TORCH_CHECK(head_size % 16 == 0,
+              "head_size must be divisible by 16 for NVFP4 KV cache");
+  TORCH_CHECK(block_size % 4 == 0,
+              "block_size must be divisible by 4 for NVFP4 KV cache swizzle");
+
+  // Detect physical layout from strides (based on full_dim).
+  // HND: head stride > block_offset stride.
+  bool is_hnd = key_cache.stride(2) > key_cache.stride(1);
+
+  int64_t data_block_stride = key_cache.stride(0);  // page_bytes
+  int64_t data_head_stride, data_block_offset_stride;
+  if (is_hnd) {
+    data_head_stride = (int64_t)block_size * data_dim;
+    data_block_offset_stride = data_dim;
+  } else {
+    data_head_stride = data_dim;
+    data_block_offset_stride = (int64_t)num_heads * data_dim;
+  }
+
+  // Page layout: [K_data | K_scale | V_data | V_scale]
+  // Scale follows data within each KV side.
+  int64_t data_per_kv = (int64_t)num_heads * block_size * data_dim;
+
+  uint8_t* key_scale_ptr = key_cache.data_ptr<uint8_t>() + data_per_kv;
+  uint8_t* value_scale_ptr = value_cache.data_ptr<uint8_t>() + data_per_kv;
+
+  // Scale strides: same page stride, inner strides from layout.
+  int64_t scale_block_stride = data_block_stride;
+  int64_t scale_head_stride, scale_block_offset_stride;
+  if (is_hnd) {
+    scale_head_stride = (int64_t)block_size * scale_dim;
+    scale_block_offset_stride = scale_dim;
+  } else {
+    scale_head_stride = scale_dim;
+    scale_block_offset_stride = (int64_t)num_heads * scale_dim;
+  }
+
+  const float* k_scale_ptr = k_scale.data_ptr<float>();
+  const float* v_scale_ptr = v_scale.data_ptr<float>();
+
+  int groups_per_head = head_size / CVT_FP4_SF_VEC_SIZE;
+  int total_groups = num_heads * groups_per_head;
+  constexpr int THREADS_PER_SF = CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD;
+  int num_threads = std::min(total_groups * THREADS_PER_SF, 512);
+  num_threads = ((num_threads + 31) / 32) * 32;
+
+  dim3 grid(num_tokens);
+  dim3 block(num_threads);
+
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(key));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+  AT_DISPATCH_REDUCED_FLOATING_TYPES(
+      key.scalar_type(), "reshape_and_cache_nvfp4", [&] {
+        vllm::reshape_and_cache_nvfp4_kernel<scalar_t>
+            <<<grid, block, 0, stream>>>(
+                key.data_ptr<scalar_t>(), value.data_ptr<scalar_t>(),
+                key_cache.data_ptr<uint8_t>(), value_cache.data_ptr<uint8_t>(),
+                key_scale_ptr, value_scale_ptr,
+                slot_mapping.data_ptr<int64_t>(), k_scale_ptr, v_scale_ptr,
+                key.stride(0), value.stride(0), num_heads, head_size,
+                block_size, data_block_stride, data_head_stride,
+                data_block_offset_stride, scale_block_stride, scale_head_stride,
+                scale_block_offset_stride);
+      });
+}

csrc/ops.h

@@ -1,6 +1,7 @@
 #pragma once
 
 #include <optional>
+#include <string>
 #include <torch/library.h>
 #include <tuple>
 

csrc/sgl_flash_kernel_ops.h

@@ -0,0 +1,45 @@
+/* SPDX-License-Identifier: Apache-2.0
+ * SPDX-FileCopyrightText: Copyright 2025 SGLang Team. All Rights Reserved.
+ *
+ * Vendored from sgl-kernel/include/sgl_flash_kernel_ops.h (commit bcf72ccc).
+ * Declares the mha_fwd() C++ function signature for CUTLASS FA3 kernels.
+ * NO MODIFICATIONS from the original (except removing unused macros).
+ */
+
+#pragma once
+
+#include <ATen/ATen.h>
+#include <ATen/Tensor.h>
+#include <torch/library.h>
+#include <torch/torch.h>
+
+#include <vector>
+
+#include "sgl_kernel_torch_shim.h"
+
+/*
+ * From flash-attention (sgl-attn fork)
+ */
+std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor> mha_fwd(
+    at::Tensor q,  // (b, s_q, h, d) or (total_q, h, d) if there is cu_seqlens_q
+    at::Tensor k,  // (b_k, s_k, h_k, d) or (total_k, h_k, d) or paged
+    at::Tensor v,  // (b_k, s_k, h_k, dv) or (total_k, h_k, dv) or paged
+    std::optional<at::Tensor> k_new_, std::optional<at::Tensor> v_new_,
+    std::optional<at::Tensor> q_v_,  // MLA value projection query
+    std::optional<at::Tensor> out_, std::optional<at::Tensor> cu_seqlens_q_,
+    std::optional<at::Tensor> cu_seqlens_k_,
+    std::optional<at::Tensor> cu_seqlens_k_new_,
+    std::optional<at::Tensor> seqused_q_, std::optional<at::Tensor> seqused_k_,
+    std::optional<int64_t> max_seqlen_q_, std::optional<int64_t> max_seqlen_k_,
+    std::optional<at::Tensor> page_table_,
+    std::optional<at::Tensor> kv_batch_idx_,
+    std::optional<at::Tensor> leftpad_k_, std::optional<at::Tensor> rotary_cos_,
+    std::optional<at::Tensor> rotary_sin_,
+    std::optional<at::Tensor> seqlens_rotary_,
+    std::optional<at::Tensor> q_descale_, std::optional<at::Tensor> k_descale_,
+    std::optional<at::Tensor> v_descale_, std::optional<double> softmax_scale_,
+    bool is_causal, int64_t window_size_left, int64_t window_size_right,
+    int64_t attention_chunk, double softcap, bool is_rotary_interleaved,
+    std::optional<at::Tensor> scheduler_metadata_, int64_t num_splits,
+    std::optional<bool> pack_gqa_, int64_t sm_margin,
+    std::optional<const at::Tensor>& sinks_);

csrc/sgl_kernel_torch_shim.h

@@ -0,0 +1,121 @@
+/* Adapted from:
+ * https://github.com/neuralmagic/vllm-flash-attention/blob/90eacc1af2a7c3de62ea249e929ed5faccf38954/csrc/common/pytorch_shim.h
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ * SPDX-FileCopyrightText: Copyright 2025 SGLang Team. All Rights Reserved.
+ *
+ * Vendored from sgl-kernel/include/sgl_kernel_torch_shim.h (commit bcf72ccc).
+ * Provides make_pytorch_shim() template for PyTorch op registration type
+ * conversion. NO MODIFICATIONS from the original.
+ */
+
+#pragma once
+
+#include <torch/library.h>
+
+/**
+ * Unfortunately, the type signatures of the flash_attn ops are not compatible
+ * with the PyTorch library bindings. To get around that we use
+ * `make_pytorch_shim` which creates a lambda that exposes the API using
+ * PyTorch compatible types to the types, then converts them to the types
+ * expected by the flash_attn ops. This shims allows us to make minimal changes
+ * to `flash_api.cpp` making it easier to synchronize with upstream changes.
+ *
+ * The `pytorch_library_compatible_type` struct is used to map from the
+ * flash_attn ops types to a PyTorch library compatible one. The main issues is
+ * that the following types are not support by PyTorch library bindings:
+ *  - `int`
+ *  - `float`
+ *  - `std::optional<T> &`
+ *  - `std::optional<const at::Tensor> &`
+ * So we convert them to (respectively):
+ *  - `int64_t`
+ *  - `double`
+ *  - `const std::optional<T>&`
+ *  - `const std::optional<at::Tensor>&`
+ */
+
+template <typename T>
+struct pytorch_library_compatible_type {
+  using type = T;
+  static T convert_from_type(T arg) { return arg; }
+};
+
+template <typename T>
+using pytorch_library_compatible_type_t =
+    typename pytorch_library_compatible_type<T>::type;
+
+template <typename T>
+T convert_from_pytorch_compatible_type(
+    pytorch_library_compatible_type_t<T> arg) {
+  return pytorch_library_compatible_type<T>::convert_from_type(arg);
+}
+
+// Map `c10::optional<T> &` -> `const c10::optional<T>&`
+//  (NOTE: this is bit unsafe but non of the ops in flash_attn mutate
+//   the optional container)
+template <typename T>
+struct pytorch_library_compatible_type<c10::optional<T>&> {
+  using type = const c10::optional<T>&;
+  static c10::optional<T>& convert_from_type(const c10::optional<T>& arg) {
+    return const_cast<c10::optional<T>&>(arg);
+  }
+};
+
+// Map `c10::optional<T>` ->
+//          `c10::optional<pytorch_library_compatible_type_t<T>>`
+//  (NOTE: tested for `c10::optional<int>` -> `c10::optional<int64_t>`)
+template <typename T>
+struct pytorch_library_compatible_type<c10::optional<T>> {
+  using type = c10::optional<pytorch_library_compatible_type_t<T>>;
+  static c10::optional<pytorch_library_compatible_type_t<T>> convert_from_type(
+      c10::optional<T> arg) {
+    return arg;
+  }
+};
+
+// Map `c10::optional<const at::Tensor>&` -> `const c10::optional<at::Tensor>&`
+template <>
+struct pytorch_library_compatible_type<c10::optional<const at::Tensor>&> {
+  using type = const c10::optional<at::Tensor>&;
+  static c10::optional<const at::Tensor>& convert_from_type(
+      const c10::optional<at::Tensor>& arg) {
+    return const_cast<c10::optional<const at::Tensor>&>(
+        reinterpret_cast<const c10::optional<const at::Tensor>&>(arg));
+  }
+};
+
+// Map `int` -> `int64_t`
+template <>
+struct pytorch_library_compatible_type<int> {
+  using type = int64_t;
+  static int convert_from_type(int64_t arg) {
+    TORCH_CHECK(arg <= std::numeric_limits<int>::max(),
+                "int64_t value is too large to be converted to int");
+    TORCH_CHECK(arg >= std::numeric_limits<int>::min(),
+                "int64_t value is too small to be converted to int");
+    return arg;
+  }
+};
+
+// Map `float` -> `double`
+template <>
+struct pytorch_library_compatible_type<float> {
+  using type = double;
+  static float convert_from_type(double arg) {
+    TORCH_CHECK(std::abs(arg) <= std::numeric_limits<float>::max(),
+                "double value is too large to be converted to float");
+    return arg;
+  }
+};
+
+//
+//  Shim Utils
+//
+
+template <typename Ret, typename... Args>
+auto make_pytorch_shim(Ret (*fun)(Args... args)) {
+  return [fun](pytorch_library_compatible_type_t<Args>... args) {
+    return fun(convert_from_pytorch_compatible_type<Args>(args)...);
+  };
+}

csrc/standalone/custom_ar/bench_moe_allgather.py

@@ -1,200 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Benchmark: Lamport all-gather vs NCCL."""
-
-import ctypes
-import os
-import sys
-
-import torch
-import torch.distributed as dist
-
-_cudart = ctypes.CDLL("libcudart.so")
-IPC = 64
-
-
-def _cc(r):
-    if r:
-        raise RuntimeError(f"err {r}")
-
-
-def ipc_buf(sz, rank, ws):
-    p = ctypes.c_void_p()
-    _cc(_cudart.cudaMalloc(ctypes.byref(p), sz))
-    _cc(_cudart.cudaMemset(p, 0, sz))
-    _cc(_cudart.cudaDeviceSynchronize())
-    h = (ctypes.c_byte * IPC)()
-    _cc(_cudart.cudaIpcGetMemHandle(ctypes.byref(h), p))
-    ah = [None] * ws
-    dist.all_gather_object(ah, bytes(h))
-    ptrs = []
-    for i in range(ws):
-        if i == rank:
-            ptrs.append(p.value)
-        else:
-            hh = (ctypes.c_byte * IPC)(*ah[i])
-            pp = ctypes.c_void_p()
-            _cc(_cudart.cudaIpcOpenMemHandle(ctypes.byref(pp), hh, ctypes.c_uint(1)))
-            ptrs.append(pp.value)
-    return ptrs
-
-
-def gpu_timer(fn, warmup=20, repeats=200):
-    for _ in range(warmup):
-        fn()
-    torch.cuda.synchronize()
-    s = torch.cuda.Event(enable_timing=True)
-    e = torch.cuda.Event(enable_timing=True)
-    s.record()
-    for _ in range(repeats):
-        fn()
-    e.record()
-    torch.cuda.synchronize()
-    return s.elapsed_time(e) / repeats * 1000
-
-
-def gpu_timer_graph(fn, warmup=20, repeats=200):
-    """Time with CUDA graph to exclude CPU overhead."""
-    for _ in range(warmup):
-        fn()
-    torch.cuda.synchronize()
-    g = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(g):
-        fn()
-    for _ in range(5):
-        g.replay()
-    torch.cuda.synchronize()
-    s = torch.cuda.Event(enable_timing=True)
-    e = torch.cuda.Event(enable_timing=True)
-    s.record()
-    for _ in range(repeats):
-        g.replay()
-    e.record()
-    torch.cuda.synchronize()
-    return s.elapsed_time(e) / repeats * 1000
-
-
-def main():
-    dist.init_process_group("nccl")
-    rank = dist.get_rank()
-    ws = dist.get_world_size()
-    torch.cuda.set_device(rank)
-    dev = f"cuda:{rank}"
-
-    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-    if rank == 0:
-        from moe_allgather import _load_lib
-
-        lib = _load_lib()
-    dist.barrier()
-    if rank != 0:
-        from moe_allgather import _load_lib
-
-        lib = _load_lib()
-    dist.barrier()
-    from moe_allgather import MoeAllGather
-
-    max_size = 8 * 1024 * 1024
-    bp = ipc_buf(max_size, rank, ws)
-    dist.barrier()
-
-    class FakeCA:
-        pass
-
-    ca = FakeCA()
-    ca.rank = rank
-    ca.world_size = ws
-    ca.device = torch.device(dev)
-    ca.buffer_ptrs = bp
-    ca.max_size = max_size
-    ag = MoeAllGather(ca)
-    dist.barrier()
-
-    configs = [
-        ("1tok", 1),
-        ("2tok", 2),
-        ("4tok", 4),
-        ("8tok", 8),
-        ("16tok", 16),
-        ("32tok", 32),
-        ("64tok", 64),
-        ("128tok", 128),
-        ("256tok", 256),
-    ]
-    topk = 8
-    hd = 3584
-    sd = 448
-
-    if rank == 0:
-        print(f"world_size={ws}, max_per_rank={ag.max_per_rank} bytes")
-        print(f"{'config':<12} {'lamport_graph':>10} {'nccl_graph':>10} {'speedup':>8}")
-        print("-" * 65)
-
-    for name, N in configs:
-        # Check if data fits in buffer.
-        cursor = 0
-        per_tok = topk * 4 + topk * 4 + hd + sd
-        cursor = N * per_tok
-        cursor = (cursor + 15) & ~15
-        if cursor > ag.max_per_rank:
-            if rank == 0:
-                print(f"{name:<12} {'skip (too large)':>40}")
-            continue
-
-        ids = torch.randint(0, 256, (N, topk), dtype=torch.int32, device=dev)
-        wt = torch.randn(N, topk, dtype=torch.float32, device=dev).abs()
-        hs = torch.randint(0, 255, (N, hd), dtype=torch.uint8, device=dev)
-        sc = torch.randint(0, 255, (N, sd), dtype=torch.uint8, device=dev)
-        inputs = [ids, wt, hs, sc]
-
-        # Custom Lamport kernel.
-        c_outs = [
-            torch.empty(N * ws, *t.shape[1:], dtype=t.dtype, device=dev) for t in inputs
-        ]
-
-        def run_lamport():
-            lib.moe_all_gather(
-                ag._buf_ptrs_ptr,
-                ag._counters_ptr,
-                rank,
-                ws,
-                ag.seg_capacity,
-                ag.rank_stride,
-                inputs,
-                c_outs,
-            )
-
-        # Lamport with CUDA graph.
-        try:
-            lam_g_us = gpu_timer_graph(run_lamport)
-        except Exception as ex:
-            lam_g_us = float("nan")
-            if rank == 0:
-                print(f"  [graph capture failed: {ex}]")
-
-        # NCCL 1×AG (concat into one tensor).
-        cat_inp = torch.cat(
-            [t.reshape(N, -1).contiguous().view(torch.uint8) for t in inputs],
-            dim=1,
-        ).contiguous()
-        cat_out = torch.empty(N * ws, cat_inp.shape[1], dtype=torch.uint8, device=dev)
-
-        def run_nccl():
-            dist.all_gather_into_tensor(cat_out, cat_inp)
-
-        # NCCL with CUDA graph.
-        try:
-            nccl_g_us = gpu_timer_graph(run_nccl)
-        except Exception:
-            nccl_g_us = float("nan")
-
-        if rank == 0:
-            speedup = nccl_g_us / lam_g_us if lam_g_us > 0 else float("nan")
-            print(f"{name:<12} {lam_g_us:>9.1f}µ {nccl_g_us:>9.1f}µ {speedup:>7.2f}x")
-
-    dist.barrier()
-    dist.destroy_process_group()
-
-
-if __name__ == "__main__":
-    main()

csrc/standalone/custom_ar/moe_allgather.cu

@@ -1,301 +0,0 @@
-// Lamport-based MoE all-gather kernel for EP dispatch.
-//
-// Replaces the flag-barrier approach with a Lamport sentinel protocol
-// (inspired by FlashInfer's trtllm_allreduce_fusion).
-//
-// Key advantages over the flag-barrier approach:
-//   - No explicit barriers (sentinels provide per-element synchronization).
-//   - Push model: NVLink writes (fire-and-forget) instead of NVLink reads.
-//   - Triple buffering: no end barrier needed.
-//
-// Gathers the MoE dispatch tensors from all EP ranks:
-//   - topk_ids      [N, topk]     int32
-//   - topk_weights  [N, topk]     float32 / bfloat16
-//   - hidden_states [N, D_h]      uint8 (NVFP4) / bfloat16
-//   - quant_scales  [N, D_s]      (optional)
-//
-// Double-buffer layout in each rank's IPC buffer:
-//   [Segment 0][Segment 1]
-//   Each segment: [Rank 0 slot][Rank 1 slot]...[Rank N-1 slot]
-//   Each rank slot: packed tensors at 16-byte aligned offsets.
-//
-// Sentinel: 0x80000000 (negative-zero in float32).  The writer replaces
-// any data word matching the sentinel with 0 before pushing.  The reader
-// spin-loads (volatile) until no sentinel words remain in the vector.
-
-#include <cuda.h>
-#include <cuda_runtime.h>
-#include <torch/extension.h>
-
-#include <c10/cuda/CUDAGuard.h>
-#include <c10/cuda/CUDAStream.h>
-
-#define DINLINE __device__ __forceinline__
-
-constexpr uint32_t SENTINEL = 0x80000000u;
-constexpr int kMaxBlocks = 36;
-
-// ---------------------------------------------------------------------------
-// Volatile 128-bit load/store and sentinel helpers
-// ---------------------------------------------------------------------------
-
-static DINLINE int4 ld128v(const void* addr) {
-  int4 v;
-  asm volatile("ld.volatile.global.v4.b32 {%0,%1,%2,%3}, [%4];"
-               : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w)
-               : "l"(addr));
-  return v;
-}
-
-static DINLINE bool has_sentinel(int4 v) {
-  return reinterpret_cast<uint32_t&>(v.x) == SENTINEL |
-         reinterpret_cast<uint32_t&>(v.y) == SENTINEL |
-         reinterpret_cast<uint32_t&>(v.z) == SENTINEL |
-         reinterpret_cast<uint32_t&>(v.w) == SENTINEL;
-}
-
-static DINLINE int4 remove_sentinel(int4 v) {
-  if (reinterpret_cast<uint32_t&>(v.x) == SENTINEL) v.x = 0;
-  if (reinterpret_cast<uint32_t&>(v.y) == SENTINEL) v.y = 0;
-  if (reinterpret_cast<uint32_t&>(v.z) == SENTINEL) v.z = 0;
-  if (reinterpret_cast<uint32_t&>(v.w) == SENTINEL) v.w = 0;
-  return v;
-}
-
-// ---------------------------------------------------------------------------
-// Lamport all-gather kernel
-// ---------------------------------------------------------------------------
-//
-// Phase 1 — PUSH: each rank writes its packed data to ALL peers' current
-//           segment via regular stores (NVLink push, fire-and-forget).
-// Phase 2 — CLEAR: each rank writes sentinels to the OLDEST segment of
-//           its own buffer, preparing it for reuse.
-// Phase 3 — POLL + SCATTER: each rank volatile-loads from its own current
-//           segment, spinning until sentinels disappear, then scatters
-//           directly to per-tensor output arrays.
-// Phase 4 — ADVANCE: one thread advances the triple-buffer ring counter.
-
-template <int ngpus, int nbufs>
-__global__ void __launch_bounds__(512, 1) moe_allgather_lamport_kernel(
-    int64_t* buf_ptrs,  // [ngpus] IPC buffer base addresses (device)
-    int* counters,      // [0] = unused, [1] = ring (0/1/2), [2] = prev total_sz
-    int rank,
-    int seg_capacity,  // bytes per segment
-    int rank_stride,   // bytes per rank-slot within a segment
-    int total_sz,      // int4 units of actual packed data per rank
-    // inputs (up to 4)
-    const void* inp0, const void* inp1, const void* inp2, const void* inp3,
-    int off0, int sz0, int off1, int sz1, int off2, int sz2, int off3, int sz3,
-    // outputs (up to 4)
-    void* out0, void* out1, void* out2, void* out3) {
-  using V = int4;
-  const int tid = blockIdx.x * blockDim.x + threadIdx.x;
-  const int stride = gridDim.x * blockDim.x;
-
-  // Read segment index and previous clear size.
-  const int seg = counters[1];            // 0 or 1
-  const int prev_total_sz = counters[2];  // set by previous invocation
-  const int cur_seg = seg;
-  const int old_seg = 1 - seg;
-
-  char* bufs[ngpus];
-#pragma unroll
-  for (int r = 0; r < ngpus; r++)
-    bufs[r] = reinterpret_cast<char*>(buf_ptrs[r]) + cur_seg * seg_capacity;
-
-  // Sentinel vector for clearing.
-  V sent;
-  sent.x = sent.y = sent.z = sent.w = static_cast<int>(SENTINEL);
-
-  // ---- Phase 1: PUSH local data to ALL peers ----
-  // Write to peer_r's buffer at [rank * rank_stride + off_i].
-
-#define PUSH(idx, inp_ptr, off_val, sz_val)                                 \
-  if constexpr (nbufs > (idx)) {                                            \
-    const V* src = reinterpret_cast<const V*>(inp_ptr);                     \
-    for (int i = tid; i < (sz_val); i += stride) {                          \
-      V val = remove_sentinel(src[i]);                                      \
-      _Pragma("unroll") for (int r = 0; r < ngpus; r++) {                   \
-        reinterpret_cast<V*>(bufs[r] + rank * rank_stride + (off_val))[i] = \
-            val;                                                            \
-      }                                                                     \
-    }                                                                       \
-  }
-
-  PUSH(0, inp0, off0, sz0)
-  PUSH(1, inp1, off1, sz1)
-  PUSH(2, inp2, off2, sz2)
-  PUSH(3, inp3, off3, sz3)
-#undef PUSH
-
-  // ---- Phase 2: CLEAR only the previously-written data in oldest segment ----
-  // Only clear what the previous invocation actually wrote (per rank-slot).
-  if (prev_total_sz > 0) {
-    char* clr_base =
-        reinterpret_cast<char*>(buf_ptrs[rank]) + old_seg * seg_capacity;
-#pragma unroll
-    for (int r = 0; r < ngpus; r++) {
-      V* clr = reinterpret_cast<V*>(clr_base + r * rank_stride);
-      for (int i = tid; i < prev_total_sz; i += stride) clr[i] = sent;
-    }
-  }
-
-  // ---- Phase 3: POLL + SCATTER ----
-  // Volatile-load from own buffer; spin until sentinel gone; scatter to output.
-  char* my = bufs[rank];
-
-#define POLL(idx, out_ptr, off_val, sz_val)                                \
-  if constexpr (nbufs > (idx)) {                                           \
-    for (int i = tid; i < (sz_val); i += stride) {                         \
-      _Pragma("unroll") for (int s = 0; s < ngpus; s++) {                  \
-        V val;                                                             \
-        do {                                                               \
-          val = ld128v(                                                    \
-              reinterpret_cast<V*>(my + s * rank_stride + (off_val)) + i); \
-        } while (has_sentinel(val));                                       \
-        reinterpret_cast<V*>(out_ptr)[s * (sz_val) + i] = val;             \
-      }                                                                    \
-    }                                                                      \
-  }
-
-  POLL(0, out0, off0, sz0)
-  POLL(1, out1, off1, sz1)
-  POLL(2, out2, off2, sz2)
-  POLL(3, out3, off3, sz3)
-#undef POLL
-
-  // ---- Phase 4: ADVANCE ring counter + store clear size for next call ----
-  // Stream serialization ensures the next kernel sees these updates.
-  if (blockIdx.x == 0 && threadIdx.x == 0) {
-    counters[1] = 1 - seg;
-    counters[2] = total_sz;
-  }
-}
-
-// ---------------------------------------------------------------------------
-// Sentinel initialization kernel
-// ---------------------------------------------------------------------------
-
-__global__ void lamport_init_kernel(uint32_t* buf, int n) {
-  int tid = blockIdx.x * blockDim.x + threadIdx.x;
-  int stride = gridDim.x * blockDim.x;
-  for (int i = tid; i < n; i += stride) buf[i] = SENTINEL;
-}
-
-// ---------------------------------------------------------------------------
-// Host launcher
-// ---------------------------------------------------------------------------
-
-struct TensorDesc {
-  void* inp;
-  int off;
-  int sz;
-  int64_t nbytes;
-};
-
-static TensorDesc make_desc(torch::Tensor& inp, int64_t& cursor) {
-  TORCH_CHECK(inp.is_contiguous(), "input must be contiguous");
-  int64_t nbytes = inp.numel() * inp.element_size();
-  TORCH_CHECK(nbytes % 16 == 0, "tensor byte size must be multiple of 16, got ",
-              nbytes);
-  cursor = (cursor + 15) & ~15;
-  TensorDesc d;
-  d.inp = inp.data_ptr();
-  d.off = static_cast<int>(cursor);
-  d.sz = static_cast<int>(nbytes / 16);
-  d.nbytes = nbytes;
-  cursor += nbytes;
-  return d;
-}
-
-void lamport_init(int64_t buf_ptr, int64_t nbytes) {
-  auto stream = c10::cuda::getCurrentCUDAStream().stream();
-  int n = static_cast<int>(nbytes / 4);
-  lamport_init_kernel<<<256, 256, 0, stream>>>(
-      reinterpret_cast<uint32_t*>(buf_ptr), n);
-}
-
-void moe_all_gather(int64_t buf_ptrs_ptr, int64_t counters_ptr, int64_t rank,
-                    int64_t world_size, int64_t seg_capacity,
-                    int64_t rank_stride, std::vector<torch::Tensor>& inputs,
-                    std::vector<torch::Tensor>& outputs) {
-  auto stream = c10::cuda::getCurrentCUDAStream().stream();
-  int n = static_cast<int>(inputs.size());
-  TORCH_CHECK(n >= 2 && n <= 4, "2-4 input tensors required");
-  TORCH_CHECK(inputs.size() == outputs.size());
-
-  int64_t cursor = 0;
-  TensorDesc descs[4] = {};
-  for (int i = 0; i < n; i++) descs[i] = make_desc(inputs[i], cursor);
-  TORCH_CHECK(cursor % 16 == 0);
-  int total_sz = static_cast<int>(cursor / 16);
-  TORCH_CHECK(cursor <= rank_stride, "packed data (", cursor,
-              " bytes) exceeds rank_stride (", rank_stride, " bytes)");
-
-  int ws = static_cast<int>(world_size);
-  for (int i = 0; i < n; i++) {
-    TORCH_CHECK(outputs[i].is_contiguous());
-    TORCH_CHECK(outputs[i].numel() == inputs[i].numel() * ws);
-  }
-
-  int r = static_cast<int>(rank);
-  int threads = 512;
-  int blocks =
-      std::max(1, std::min(kMaxBlocks, (total_sz + threads - 1) / threads));
-
-  void *inps[4] = {}, *outs[4] = {};
-  int offs[4] = {}, szs[4] = {};
-  for (int i = 0; i < n; i++) {
-    inps[i] = descs[i].inp;
-    offs[i] = descs[i].off;
-    szs[i] = descs[i].sz;
-    outs[i] = outputs[i].data_ptr();
-  }
-
-  auto* bp = reinterpret_cast<int64_t*>(buf_ptrs_ptr);
-  auto* ct = reinterpret_cast<int*>(counters_ptr);
-  int sc = static_cast<int>(seg_capacity);
-  int rs = static_cast<int>(rank_stride);
-
-#define KL(ng, nb)                                                        \
-  moe_allgather_lamport_kernel<ng, nb><<<blocks, threads, 0, stream>>>(   \
-      bp, ct, r, sc, rs, total_sz, inps[0], inps[1], inps[2], inps[3],    \
-      offs[0], szs[0], offs[1], szs[1], offs[2], szs[2], offs[3], szs[3], \
-      outs[0], outs[1], outs[2], outs[3]);
-
-#define GPU_CASE(ng) \
-  case ng:           \
-    switch (n) {     \
-      case 2:        \
-        KL(ng, 2);   \
-        break;       \
-      case 3:        \
-        KL(ng, 3);   \
-        break;       \
-      case 4:        \
-        KL(ng, 4);   \
-        break;       \
-    }                \
-    break;
-
-  switch (ws) {
-    GPU_CASE(2)
-    GPU_CASE(4)
-    GPU_CASE(6)
-    GPU_CASE(8)
-    default:
-      TORCH_CHECK(false, "world_size must be 2, 4, 6, or 8");
-  }
-#undef GPU_CASE
-#undef KL
-}
-
-// ---------------------------------------------------------------------------
-// Python binding
-// ---------------------------------------------------------------------------
-
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-  m.def("moe_all_gather", &moe_all_gather, "Lamport MoE all-gather");
-  m.def("lamport_init", &lamport_init,
-        "Initialize Lamport buffer with sentinels");
-}

csrc/standalone/custom_ar/moe_allgather.py

@@ -1,116 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Lamport-based fused MoE all-gather for EP dispatch.
-
-JIT-compiles the CUDA kernel on first use (cached afterwards).
-Uses a Lamport sentinel protocol (push writes + per-element sync)
-with triple-buffered IPC regions — no explicit barriers.
-
-Usage:
-    ag = MoeAllGather(custom_allreduce)
-    ids_g, wt_g, hs_g, sc_g = ag.gather(topk_ids, topk_weights, hidden, scales)
-"""
-
-from __future__ import annotations
-
-import os
-from pathlib import Path
-
-import torch
-
-_lib = None
-
-
-def _load_lib():
-    global _lib
-    if _lib is not None:
-        return _lib
-
-    from torch.utils.cpp_extension import load
-
-    src = str(Path(__file__).with_name("moe_allgather.cu"))
-    _lib = load(
-        name="moe_allgather_kernel",
-        sources=[src],
-        extra_cuda_cflags=["-O3", "--use_fast_math"],
-        verbose=os.environ.get("MOE_AG_VERBOSE", "") == "1",
-    )
-    return _lib
-
-
-class MoeAllGather:
-    """Lamport-based MoE dispatch all-gather with triple buffering."""
-
-    def __init__(self, ca_comm):
-        self.rank = ca_comm.rank
-        self.world_size = ca_comm.world_size
-        self.device = ca_comm.device
-        self.buffer_ptrs = ca_comm.buffer_ptrs
-        self.max_size = ca_comm.max_size
-
-        ws = self.world_size
-        # Double-buffer layout: 2 segments, each with ws rank-slots.
-        # Safe because kernels in the same stream are serialized, and the
-        # Use the FIRST half of the IPC buffer (second half reserved for
-        # MoeReduceScatter) to avoid overlapping writes.
-        half_size = (self.max_size // 2) & ~15
-
-        # Double-buffer layout within our half: 2 segments, each ws rank-slots.
-        # seg_capacity and rank_stride are 16-byte aligned.
-        self.seg_capacity = (half_size // 2) & ~15
-        self.rank_stride = (self.seg_capacity // ws) & ~15
-        self.max_per_rank = self.rank_stride  # max packed bytes per rank
-
-        # Buffer pointer array on device (no offset — first half).
-        self._buf_ptrs = torch.zeros(
-            8, dtype=torch.int64, device=f"cuda:{self.device.index}"
-        )
-        for i in range(ws):
-            self._buf_ptrs[i] = self.buffer_ptrs[i]
-        self._buf_ptrs_ptr = self._buf_ptrs.data_ptr()
-
-        # Counters on device: [0]=unused, [1]=seg (0/1), [2]=prev_total_sz.
-        self._counters = torch.zeros(
-            3, dtype=torch.int32, device=f"cuda:{self.device.index}"
-        )
-        self._counters_ptr = self._counters.data_ptr()
-
-        # Initialize our half with sentinel values.
-        lib = _load_lib()
-        lib.lamport_init(self.buffer_ptrs[self.rank], half_size)
-        torch.accelerator.synchronize(self.device)
-
-    def gather(
-        self,
-        topk_ids: torch.Tensor,
-        topk_weights: torch.Tensor,
-        hidden_states: torch.Tensor,
-        scales: torch.Tensor | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]:
-        lib = _load_lib()
-        ws = self.world_size
-
-        inputs = [topk_ids, topk_weights, hidden_states]
-        if scales is not None:
-            inputs.append(scales)
-
-        outputs = [
-            torch.empty((t.shape[0] * ws, *t.shape[1:]), dtype=t.dtype, device=t.device)
-            for t in inputs
-        ]
-
-        lib.moe_all_gather(
-            self._buf_ptrs_ptr,
-            self._counters_ptr,
-            self.rank,
-            self.world_size,
-            self.seg_capacity,
-            self.rank_stride,
-            inputs,
-            outputs,
-        )
-
-        if scales is not None:
-            return outputs[0], outputs[1], outputs[2], outputs[3]
-        return outputs[0], outputs[1], outputs[2], None

csrc/standalone/custom_ar/moe_reduce_scatter.cu

@@ -1,261 +0,0 @@
-// Lamport-based MoE reduce-scatter kernel for EP combine.
-//
-// JIT-compilable via torch.utils.cpp_extension — no vLLM build required.
-//
-// Reduce-scatters a bf16 tensor [N_total, D] across EP ranks.  Each rank
-// contributes its partial MoE output; the kernel sums all contributions
-// and each rank receives its own slice of the result.
-//
-// Protocol (same as the all-gather variant):
-//   1. PUSH: write own data to all peers' Lamport buffers (NVLink push).
-//   2. CLEAR: write sentinels to old segment of own buffer.
-//   3. POLL + REDUCE: volatile-load all peers' data for own slice,
-//      accumulate in fp32, convert back to bf16, store to output.
-//   4. ADVANCE: toggle double-buffer index.
-//
-// Sentinel: 0x80000000 (two bf16 negative-zeros packed in uint32).
-// For bf16 reduce, replacing -0 with +0 is lossless.
-
-#include <cuda.h>
-#include <cuda_bf16.h>
-#include <cuda_runtime.h>
-#include <torch/extension.h>
-
-#include <c10/cuda/CUDAGuard.h>
-#include <c10/cuda/CUDAStream.h>
-
-#define DINLINE __device__ __forceinline__
-
-constexpr uint32_t SENTINEL = 0x80000000u;
-constexpr int kMaxBlocks = 36;
-
-// ---------------------------------------------------------------------------
-// Volatile 128-bit load and sentinel helpers
-// ---------------------------------------------------------------------------
-
-static DINLINE int4 ld128v(const void* addr) {
-  int4 v;
-  asm volatile("ld.volatile.global.v4.b32 {%0,%1,%2,%3}, [%4];"
-               : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w)
-               : "l"(addr));
-  return v;
-}
-
-static DINLINE bool has_sentinel(int4 v) {
-  return reinterpret_cast<uint32_t&>(v.x) == SENTINEL |
-         reinterpret_cast<uint32_t&>(v.y) == SENTINEL |
-         reinterpret_cast<uint32_t&>(v.z) == SENTINEL |
-         reinterpret_cast<uint32_t&>(v.w) == SENTINEL;
-}
-
-static DINLINE int4 remove_sentinel(int4 v) {
-  if (reinterpret_cast<uint32_t&>(v.x) == SENTINEL) v.x = 0;
-  if (reinterpret_cast<uint32_t&>(v.y) == SENTINEL) v.y = 0;
-  if (reinterpret_cast<uint32_t&>(v.z) == SENTINEL) v.z = 0;
-  if (reinterpret_cast<uint32_t&>(v.w) == SENTINEL) v.w = 0;
-  return v;
-}
-
-// ---------------------------------------------------------------------------
-// bf16 ↔ fp32 helpers for int4 (8 bf16 values = 16 bytes)
-// ---------------------------------------------------------------------------
-
-// Accumulate 8 bf16 values from an int4 into 8 fp32 accumulators.
-static DINLINE void accumulate_bf16(float* acc, int4 v) {
-  const __nv_bfloat16* bp = reinterpret_cast<const __nv_bfloat16*>(&v);
-#pragma unroll
-  for (int k = 0; k < 8; k++) acc[k] += __bfloat162float(bp[k]);
-}
-
-// Convert 8 fp32 accumulators to bf16 and pack into int4.
-static DINLINE int4 fp32_to_bf16_int4(const float* acc) {
-  int4 out;
-  __nv_bfloat16* bp = reinterpret_cast<__nv_bfloat16*>(&out);
-#pragma unroll
-  for (int k = 0; k < 8; k++) bp[k] = __float2bfloat16(acc[k]);
-  return out;
-}
-
-// ---------------------------------------------------------------------------
-// Lamport reduce-scatter kernel
-// ---------------------------------------------------------------------------
-
-template <int ngpus>
-__global__ void __launch_bounds__(512, 1) moe_rs_lamport_kernel(
-    int64_t* buf_ptrs,  // [ngpus] IPC buffer base addresses (device)
-    int* counters,      // [0] = unused, [1] = seg (0/1), [2] = prev total_sz
-    int rank,
-    int seg_capacity,   // bytes per segment
-    int rank_stride,    // bytes per rank-slot within a segment
-    const void* input,  // [N_total, D] bf16 — full input
-    void* output,       // [N_per_rank, D] bf16 — this rank's reduced slice
-    int total_sz,       // int4 units of full input per rank
-    int slice_off,      // int4 offset to this rank's slice within packed data
-    int slice_sz) {     // int4 units of this rank's slice
-  using V = int4;
-  const int tid = blockIdx.x * blockDim.x + threadIdx.x;
-  const int stride = gridDim.x * blockDim.x;
-
-  // Read segment index and previous clear size.
-  const int seg = counters[1];
-  const int prev_total_sz = counters[2];
-  const int cur_seg = seg;
-  const int old_seg = 1 - seg;
-
-  char* bufs[ngpus];
-#pragma unroll
-  for (int r = 0; r < ngpus; r++)
-    bufs[r] = reinterpret_cast<char*>(buf_ptrs[r]) + cur_seg * seg_capacity;
-
-  V sent;
-  sent.x = sent.y = sent.z = sent.w = static_cast<int>(SENTINEL);
-
-  // ---- Phase 1: PUSH full input to ALL peers ----
-  {
-    const V* src = reinterpret_cast<const V*>(input);
-    for (int i = tid; i < total_sz; i += stride) {
-      V val = remove_sentinel(src[i]);
-#pragma unroll
-      for (int r = 0; r < ngpus; r++)
-        reinterpret_cast<V*>(bufs[r] + rank * rank_stride)[i] = val;
-    }
-  }
-
-  // ---- Phase 2: CLEAR old segment ----
-  if (prev_total_sz > 0) {
-    char* clr_base =
-        reinterpret_cast<char*>(buf_ptrs[rank]) + old_seg * seg_capacity;
-#pragma unroll
-    for (int r = 0; r < ngpus; r++) {
-      V* clr = reinterpret_cast<V*>(clr_base + r * rank_stride);
-      for (int i = tid; i < prev_total_sz; i += stride) clr[i] = sent;
-    }
-  }
-
-  // ---- Phase 3: POLL + REDUCE for own slice ----
-  // Read all ranks' data at [slice_off, slice_off + slice_sz) from own buffer,
-  // sum in fp32, store bf16 result.
-  {
-    char* my = bufs[rank];
-    V* dst = reinterpret_cast<V*>(output);
-
-    for (int i = tid; i < slice_sz; i += stride) {
-      float acc[8] = {0, 0, 0, 0, 0, 0, 0, 0};
-
-#pragma unroll
-      for (int s = 0; s < ngpus; s++) {
-        V val;
-        do {
-          val = ld128v(reinterpret_cast<V*>(my + s * rank_stride) + slice_off +
-                       i);
-        } while (has_sentinel(val));
-        accumulate_bf16(acc, val);
-      }
-
-      dst[i] = fp32_to_bf16_int4(acc);
-    }
-  }
-
-  // ---- Phase 4: ADVANCE ----
-  if (blockIdx.x == 0 && threadIdx.x == 0) {
-    counters[1] = 1 - seg;
-    counters[2] = total_sz;
-  }
-}
-
-// ---------------------------------------------------------------------------
-// Sentinel initialization
-// ---------------------------------------------------------------------------
-
-__global__ void lamport_init_kernel(uint32_t* buf, int n) {
-  int tid = blockIdx.x * blockDim.x + threadIdx.x;
-  int stride = gridDim.x * blockDim.x;
-  for (int i = tid; i < n; i += stride) buf[i] = SENTINEL;
-}
-
-// ---------------------------------------------------------------------------
-// Host launcher
-// ---------------------------------------------------------------------------
-
-void lamport_init(int64_t buf_ptr, int64_t nbytes) {
-  auto stream = c10::cuda::getCurrentCUDAStream().stream();
-  int n = static_cast<int>(nbytes / 4);
-  lamport_init_kernel<<<256, 256, 0, stream>>>(
-      reinterpret_cast<uint32_t*>(buf_ptr), n);
-}
-
-void moe_reduce_scatter(int64_t buf_ptrs_ptr, int64_t counters_ptr,
-                        int64_t rank, int64_t world_size, int64_t seg_capacity,
-                        int64_t rank_stride, torch::Tensor input,
-                        torch::Tensor output) {
-  auto stream = c10::cuda::getCurrentCUDAStream().stream();
-  TORCH_CHECK(input.is_contiguous(), "input must be contiguous");
-  TORCH_CHECK(output.is_contiguous(), "output must be contiguous");
-  TORCH_CHECK(input.scalar_type() == torch::kBFloat16,
-              "input must be bf16, got ", input.scalar_type());
-  TORCH_CHECK(output.scalar_type() == torch::kBFloat16, "output must be bf16");
-
-  int ws = static_cast<int>(world_size);
-  int r = static_cast<int>(rank);
-
-  // Input: [N_total, D], Output: [N_per_rank, D]
-  int64_t N_total = input.size(0);
-  int64_t D = input.size(1);
-  TORCH_CHECK(N_total % ws == 0, "N_total must be divisible by world_size");
-  int64_t N_per_rank = N_total / ws;
-  TORCH_CHECK(output.size(0) == N_per_rank);
-  TORCH_CHECK(output.size(1) == D);
-
-  int64_t input_bytes = input.numel() * input.element_size();
-  TORCH_CHECK(input_bytes % 16 == 0,
-              "input byte size must be multiple of 16, got ", input_bytes);
-  TORCH_CHECK(input_bytes <= rank_stride, "input (", input_bytes,
-              " bytes) exceeds rank_stride (", rank_stride, " bytes)");
-
-  int total_sz = static_cast<int>(input_bytes / 16);
-  int slice_sz = total_sz / ws;
-  int slice_off = r * slice_sz;
-
-  int threads = 512;
-  int blocks =
-      std::max(1, std::min(kMaxBlocks, (total_sz + threads - 1) / threads));
-
-  auto* bp = reinterpret_cast<int64_t*>(buf_ptrs_ptr);
-  auto* ct = reinterpret_cast<int*>(counters_ptr);
-  int sc = static_cast<int>(seg_capacity);
-  int rs = static_cast<int>(rank_stride);
-
-#define LAUNCH(ng)                                                      \
-  moe_rs_lamport_kernel<ng><<<blocks, threads, 0, stream>>>(            \
-      bp, ct, r, sc, rs, input.data_ptr(), output.data_ptr(), total_sz, \
-      slice_off, slice_sz);
-
-  switch (ws) {
-    case 2:
-      LAUNCH(2);
-      break;
-    case 4:
-      LAUNCH(4);
-      break;
-    case 6:
-      LAUNCH(6);
-      break;
-    case 8:
-      LAUNCH(8);
-      break;
-    default:
-      TORCH_CHECK(false, "world_size must be 2, 4, 6, or 8");
-  }
-#undef LAUNCH
-}
-
-// ---------------------------------------------------------------------------
-// Python binding
-// ---------------------------------------------------------------------------
-
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-  m.def("moe_reduce_scatter", &moe_reduce_scatter,
-        "Lamport MoE reduce-scatter");
-  m.def("lamport_init", &lamport_init,
-        "Initialize Lamport buffer with sentinels");
-}

csrc/standalone/custom_ar/moe_reduce_scatter.py

@@ -1,102 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Lamport-based MoE reduce-scatter for EP combine.
-
-JIT-compiles the CUDA kernel on first use (cached afterwards).
-Uses the same Lamport sentinel protocol as the all-gather kernel.
-"""
-
-from __future__ import annotations
-
-import os
-from pathlib import Path
-
-import torch
-
-_lib = None
-
-
-def _load_lib():
-    global _lib
-    if _lib is not None:
-        return _lib
-
-    from torch.utils.cpp_extension import load
-
-    src = str(Path(__file__).with_name("moe_reduce_scatter.cu"))
-    _lib = load(
-        name="moe_reduce_scatter_kernel",
-        sources=[src],
-        extra_cuda_cflags=["-O3", "--use_fast_math"],
-        verbose=os.environ.get("MOE_RS_VERBOSE", "") == "1",
-    )
-    return _lib
-
-
-class MoeReduceScatter:
-    """Lamport-based MoE combine reduce-scatter with double buffering."""
-
-    def __init__(self, ca_comm):
-        self.rank = ca_comm.rank
-        self.world_size = ca_comm.world_size
-        self.device = ca_comm.device
-        self.buffer_ptrs = ca_comm.buffer_ptrs
-        self.max_size = ca_comm.max_size
-
-        ws = self.world_size
-        # Use the SECOND half of the IPC buffer (first half reserved for
-        # MoeAllGather) to avoid overlapping writes.
-        half_size = (self.max_size // 2) & ~15
-        self.buffer_offset = half_size
-
-        # Double-buffer layout within our half: 2 segments, each ws rank-slots.
-        self.seg_capacity = (half_size // 2) & ~15
-        self.rank_stride = (self.seg_capacity // ws) & ~15
-        self.max_per_rank = self.rank_stride
-
-        # Buffer pointer array on device — offset to our half.
-        self._buf_ptrs = torch.zeros(
-            8, dtype=torch.int64, device=f"cuda:{self.device.index}"
-        )
-        for i in range(ws):
-            self._buf_ptrs[i] = self.buffer_ptrs[i] + self.buffer_offset
-        self._buf_ptrs_ptr = self._buf_ptrs.data_ptr()
-
-        # Counters: [0]=unused, [1]=seg (0/1), [2]=prev_total_sz.
-        self._counters = torch.zeros(
-            3, dtype=torch.int32, device=f"cuda:{self.device.index}"
-        )
-        self._counters_ptr = self._counters.data_ptr()
-
-        # Initialize our half with sentinels.
-        lib = _load_lib()
-        lib.lamport_init(self.buffer_ptrs[self.rank] + self.buffer_offset, half_size)
-        torch.accelerator.synchronize(self.device)
-
-    def reduce_scatter(
-        self,
-        input: torch.Tensor,
-    ) -> torch.Tensor:
-        """Reduce-scatter input [N_total, D] bf16 → output [N_per_rank, D] bf16."""
-        lib = _load_lib()
-        ws = self.world_size
-
-        assert input.dim() == 2
-        N_total, D = input.shape
-        assert N_total % ws == 0
-        N_per_rank = N_total // ws
-
-        output = torch.empty((N_per_rank, D), dtype=input.dtype, device=input.device)
-
-        lib.moe_reduce_scatter(
-            self._buf_ptrs_ptr,
-            self._counters_ptr,
-            self.rank,
-            self.world_size,
-            self.seg_capacity,
-            self.rank_stride,
-            input,
-            output,
-        )
-        return output

csrc/standalone/custom_ar/moe_rs_fused.cu

@@ -1,308 +0,0 @@
-// Lamport reduce-scatter fused with residual add + RMSNorm.
-//
-// Replaces three separate kernels (RS + residual_add + RMSNorm) with one:
-//   1. PUSH: write MoE output to all peers' Lamport buffers.
-//   2. CLEAR: write sentinels to old segment.
-//   3. POLL+REDUCE+FUSE (per-token):
-//      a. Volatile-load from all peers, sum in fp32.
-//      b. Add residual.
-//      c. Compute RMSNorm (block reduction for variance).
-//      d. Store normed output + updated residual.
-//
-// Saves: one kernel launch (~3-5µs) + one global memory round-trip per layer.
-
-#include <cuda.h>
-#include <cuda_bf16.h>
-#include <cuda_runtime.h>
-#include <torch/extension.h>
-
-#include <c10/cuda/CUDAGuard.h>
-#include <c10/cuda/CUDAStream.h>
-
-#define DINLINE __device__ __forceinline__
-
-constexpr uint32_t SENTINEL = 0x80000000u;
-constexpr int kMaxBlocks = 36;
-
-// Each token has D=7168 bf16 values = 896 int4 vectors.
-// With 512 threads: ceil(896/512) = 2 int4 per thread = 16 fp32 values.
-constexpr int kMaxValsPerThread = 16;
-
-static DINLINE int4 ld128v(const void* addr) {
-  int4 v;
-  asm volatile("ld.volatile.global.v4.b32 {%0,%1,%2,%3}, [%4];"
-               : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w)
-               : "l"(addr));
-  return v;
-}
-
-static DINLINE bool has_sentinel(int4 v) {
-  return reinterpret_cast<uint32_t&>(v.x) == SENTINEL |
-         reinterpret_cast<uint32_t&>(v.y) == SENTINEL |
-         reinterpret_cast<uint32_t&>(v.z) == SENTINEL |
-         reinterpret_cast<uint32_t&>(v.w) == SENTINEL;
-}
-
-static DINLINE int4 remove_sentinel(int4 v) {
-  if (reinterpret_cast<uint32_t&>(v.x) == SENTINEL) v.x = 0;
-  if (reinterpret_cast<uint32_t&>(v.y) == SENTINEL) v.y = 0;
-  if (reinterpret_cast<uint32_t&>(v.z) == SENTINEL) v.z = 0;
-  if (reinterpret_cast<uint32_t&>(v.w) == SENTINEL) v.w = 0;
-  return v;
-}
-
-// bf16 helpers
-static DINLINE void accumulate_bf16(float* acc, int4 v) {
-  const __nv_bfloat16* bp = reinterpret_cast<const __nv_bfloat16*>(&v);
-#pragma unroll
-  for (int k = 0; k < 8; k++) acc[k] += __bfloat162float(bp[k]);
-}
-
-static DINLINE void add_bf16_to_fp32(float* dst, int4 v) {
-  const __nv_bfloat16* bp = reinterpret_cast<const __nv_bfloat16*>(&v);
-#pragma unroll
-  for (int k = 0; k < 8; k++) dst[k] += __bfloat162float(bp[k]);
-}
-
-static DINLINE int4 fp32_to_bf16_int4(const float* vals) {
-  int4 out;
-  __nv_bfloat16* bp = reinterpret_cast<__nv_bfloat16*>(&out);
-#pragma unroll
-  for (int k = 0; k < 8; k++) bp[k] = __float2bfloat16(vals[k]);
-  return out;
-}
-
-// Block-level tree reduction in shared memory.
-static DINLINE float block_reduce_sum(float val, float* smem) {
-  smem[threadIdx.x] = val;
-  __syncthreads();
-  for (int s = blockDim.x / 2; s > 0; s >>= 1) {
-    if (threadIdx.x < s) smem[threadIdx.x] += smem[threadIdx.x + s];
-    __syncthreads();
-  }
-  return smem[0];
-}
-
-// ---------------------------------------------------------------------------
-// Fused reduce-scatter + residual + RMSNorm kernel
-// ---------------------------------------------------------------------------
-
-template <int ngpus>
-__global__ void __launch_bounds__(512, 1) moe_rs_fused_kernel(
-    int64_t* buf_ptrs, int* counters, int rank, int seg_capacity,
-    int rank_stride,
-    const void* input,        // [N_total, D] bf16 — MoE output
-    const void* residual_in,  // [N_per_rank, D] bf16 — skip connection
-    const void* gamma,        // [D] bf16 — RMSNorm weight
-    void* normed_out,         // [N_per_rank, D] bf16 — normed result
-    void* residual_out,       // [N_per_rank, D] bf16 — updated residual
-    int total_sz,             // int4 units of full input
-    int slice_off,            // int4 offset to this rank's slice
-    int slice_sz,             // int4 units of this rank's slice
-    int D_int4,               // int4 units per token (hidden_dim * 2 / 16)
-    int N_per_rank,           // tokens in this rank's slice
-    float eps) {              // RMSNorm epsilon
-  using V = int4;
-  const int tid = blockIdx.x * blockDim.x + threadIdx.x;
-  const int stride = gridDim.x * blockDim.x;
-
-  const int seg = counters[1];
-  const int prev_total_sz = counters[2];
-  const int cur_seg = seg;
-  const int old_seg = 1 - seg;
-
-  char* bufs[ngpus];
-#pragma unroll
-  for (int r = 0; r < ngpus; r++)
-    bufs[r] = reinterpret_cast<char*>(buf_ptrs[r]) + cur_seg * seg_capacity;
-
-  V sent;
-  sent.x = sent.y = sent.z = sent.w = static_cast<int>(SENTINEL);
-
-  // ---- Phase 1: PUSH full MoE output to ALL peers ----
-  {
-    const V* src = reinterpret_cast<const V*>(input);
-    for (int i = tid; i < total_sz; i += stride) {
-      V val = remove_sentinel(src[i]);
-#pragma unroll
-      for (int r = 0; r < ngpus; r++)
-        reinterpret_cast<V*>(bufs[r] + rank * rank_stride)[i] = val;
-    }
-  }
-
-  // ---- Phase 2: CLEAR old segment ----
-  if (prev_total_sz > 0) {
-    char* clr_base =
-        reinterpret_cast<char*>(buf_ptrs[rank]) + old_seg * seg_capacity;
-#pragma unroll
-    for (int r = 0; r < ngpus; r++) {
-      V* clr = reinterpret_cast<V*>(clr_base + r * rank_stride);
-      for (int i = tid; i < prev_total_sz; i += stride) clr[i] = sent;
-    }
-  }
-
-  // ---- Phase 3: Fused POLL + REDUCE + RESIDUAL + RMSNORM ----
-  // Each block handles one token. Only first N_per_rank blocks participate.
-  if (blockIdx.x < N_per_rank) {
-    const int token = blockIdx.x;
-    const int token_off = slice_off + token * D_int4;  // in the full buffer
-    char* my = bufs[rank];
-
-    extern __shared__ float smem[];
-
-    // Register storage for intermediate fp32 values.
-    float local_vals[kMaxValsPerThread];
-    int n_vals = 0;
-    float partial_sum_sq = 0.0f;
-
-    // Pass 1: poll + reduce + add residual + compute sum_sq
-    for (int pos = threadIdx.x; pos < D_int4; pos += blockDim.x) {
-      float acc[8] = {0, 0, 0, 0, 0, 0, 0, 0};
-
-      // Poll all ranks' data for this position.
-#pragma unroll
-      for (int s = 0; s < ngpus; s++) {
-        V val;
-        do {
-          val = ld128v(reinterpret_cast<V*>(my + s * rank_stride) + token_off +
-                       pos);
-        } while (has_sentinel(val));
-        accumulate_bf16(acc, val);
-      }
-
-      // Add residual.
-      V res = reinterpret_cast<const V*>(residual_in)[token * D_int4 + pos];
-      add_bf16_to_fp32(acc, res);
-
-      // Store in registers and accumulate sum_sq.
-#pragma unroll
-      for (int k = 0; k < 8; k++) {
-        local_vals[n_vals++] = acc[k];
-        partial_sum_sq += acc[k] * acc[k];
-      }
-    }
-
-    // Block-level reduction: total sum of squares.
-    float total_sum_sq = block_reduce_sum(partial_sum_sq, smem);
-    float rms_scale = rsqrtf(total_sum_sq / (D_int4 * 8) + eps);
-
-    // Pass 2: apply RMSNorm, store outputs.
-    n_vals = 0;
-    const V* gamma_v = reinterpret_cast<const V*>(gamma);
-    for (int pos = threadIdx.x; pos < D_int4; pos += blockDim.x) {
-      V gv = gamma_v[pos];
-      const __nv_bfloat16* gp = reinterpret_cast<const __nv_bfloat16*>(&gv);
-
-      // Build normed output and residual output.
-      float normed_fp32[8], res_fp32[8];
-#pragma unroll
-      for (int k = 0; k < 8; k++) {
-        float val = local_vals[n_vals++];
-        res_fp32[k] = val;  // residual_out
-        normed_fp32[k] = val * rms_scale * __bfloat162float(gp[k]);  // normed
-      }
-
-      reinterpret_cast<V*>(normed_out)[token * D_int4 + pos] =
-          fp32_to_bf16_int4(normed_fp32);
-      reinterpret_cast<V*>(residual_out)[token * D_int4 + pos] =
-          fp32_to_bf16_int4(res_fp32);
-    }
-  }
-
-  // ---- Phase 4: ADVANCE ----
-  if (blockIdx.x == 0 && threadIdx.x == 0) {
-    counters[1] = 1 - seg;
-    counters[2] = total_sz;
-  }
-}
-
-// ---------------------------------------------------------------------------
-// Sentinel init + host launcher
-// ---------------------------------------------------------------------------
-
-__global__ void lamport_init_kernel(uint32_t* buf, int n) {
-  int tid = blockIdx.x * blockDim.x + threadIdx.x;
-  int stride = gridDim.x * blockDim.x;
-  for (int i = tid; i < n; i += stride) buf[i] = SENTINEL;
-}
-
-void lamport_init(int64_t buf_ptr, int64_t nbytes) {
-  auto stream = c10::cuda::getCurrentCUDAStream().stream();
-  int n = static_cast<int>(nbytes / 4);
-  lamport_init_kernel<<<256, 256, 0, stream>>>(
-      reinterpret_cast<uint32_t*>(buf_ptr), n);
-}
-
-void moe_rs_fused(int64_t buf_ptrs_ptr, int64_t counters_ptr, int64_t rank,
-                  int64_t world_size, int64_t seg_capacity, int64_t rank_stride,
-                  torch::Tensor input,         // [N_total, D] bf16
-                  torch::Tensor residual_in,   // [N_per_rank, D] bf16
-                  torch::Tensor gamma,         // [D] bf16
-                  torch::Tensor normed_out,    // [N_per_rank, D] bf16
-                  torch::Tensor residual_out,  // [N_per_rank, D] bf16
-                  double eps) {
-  auto stream = c10::cuda::getCurrentCUDAStream().stream();
-  TORCH_CHECK(input.is_contiguous() && residual_in.is_contiguous());
-  TORCH_CHECK(gamma.is_contiguous() && normed_out.is_contiguous());
-  TORCH_CHECK(residual_out.is_contiguous());
-  TORCH_CHECK(input.scalar_type() == torch::kBFloat16);
-
-  int ws = static_cast<int>(world_size);
-  int r = static_cast<int>(rank);
-  int64_t N_total = input.size(0);
-  int64_t D = input.size(1);
-  TORCH_CHECK(N_total % ws == 0);
-  int N_per_rank = static_cast<int>(N_total / ws);
-
-  int64_t input_bytes = input.numel() * input.element_size();
-  TORCH_CHECK(input_bytes % 16 == 0);
-  TORCH_CHECK(input_bytes <= rank_stride);
-
-  int total_sz = static_cast<int>(input_bytes / 16);
-  int D_int4 = static_cast<int>(D * 2 / 16);  // bf16 elements → int4 units
-  int slice_sz = total_sz / ws;
-  int slice_off = r * slice_sz;
-
-  int threads = 512;
-  // Need at least N_per_rank blocks for Phase 3 (one per token).
-  int blocks = std::max(
-      N_per_rank, std::min(kMaxBlocks, (total_sz + threads - 1) / threads));
-
-  auto* bp = reinterpret_cast<int64_t*>(buf_ptrs_ptr);
-  auto* ct = reinterpret_cast<int*>(counters_ptr);
-  int sc = static_cast<int>(seg_capacity);
-  int rs = static_cast<int>(rank_stride);
-  int smem = threads * sizeof(float);
-
-#define LAUNCH(ng)                                                      \
-  moe_rs_fused_kernel<ng><<<blocks, threads, smem, stream>>>(           \
-      bp, ct, r, sc, rs, input.data_ptr(), residual_in.data_ptr(),      \
-      gamma.data_ptr(), normed_out.data_ptr(), residual_out.data_ptr(), \
-      total_sz, slice_off, slice_sz, D_int4, N_per_rank,                \
-      static_cast<float>(eps));
-
-  switch (ws) {
-    case 2:
-      LAUNCH(2);
-      break;
-    case 4:
-      LAUNCH(4);
-      break;
-    case 6:
-      LAUNCH(6);
-      break;
-    case 8:
-      LAUNCH(8);
-      break;
-    default:
-      TORCH_CHECK(false, "world_size must be 2, 4, 6, or 8");
-  }
-#undef LAUNCH
-}
-
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-  m.def("moe_rs_fused", &moe_rs_fused,
-        "Fused Lamport reduce-scatter + residual + RMSNorm");
-  m.def("lamport_init", &lamport_init,
-        "Initialize Lamport buffer with sentinels");
-}

csrc/standalone/custom_ar/moe_rs_fused.py

@@ -1,29 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-JIT wrapper for the fused reduce-scatter + residual + RMSNorm kernel.
-"""
-
-from __future__ import annotations
-
-import os
-from pathlib import Path
-
-_lib = None
-
-
-def _load_lib():
-    global _lib
-    if _lib is not None:
-        return _lib
-
-    from torch.utils.cpp_extension import load
-
-    src = str(Path(__file__).with_name("moe_rs_fused.cu"))
-    _lib = load(
-        name="moe_rs_fused_kernel",
-        sources=[src],
-        extra_cuda_cflags=["-O3", "--use_fast_math"],
-        verbose=os.environ.get("MOE_RS_FUSED_VERBOSE", "") == "1",
-    )
-    return _lib

csrc/standalone/custom_ar/test_moe_allgather.py

@@ -1,137 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Test for the Lamport-based MoE all-gather kernel."""
-
-import ctypes
-import os
-import sys
-
-import torch
-import torch.distributed as dist
-
-_cudart = ctypes.CDLL("libcudart.so")
-IPC = 64
-
-
-def _cc(r):
-    if r:
-        raise RuntimeError(f"CUDA err {r}")
-
-
-def ipc_buf(sz, rank, ws):
-    p = ctypes.c_void_p()
-    _cc(_cudart.cudaMalloc(ctypes.byref(p), sz))
-    _cc(_cudart.cudaMemset(p, 0, sz))
-    _cc(_cudart.cudaDeviceSynchronize())
-    h = (ctypes.c_byte * IPC)()
-    _cc(_cudart.cudaIpcGetMemHandle(ctypes.byref(h), p))
-    ah = [None] * ws
-    dist.all_gather_object(ah, bytes(h))
-    ptrs = []
-    for i in range(ws):
-        if i == rank:
-            ptrs.append(p.value)
-        else:
-            hh = (ctypes.c_byte * IPC)(*ah[i])
-            pp = ctypes.c_void_p()
-            _cc(_cudart.cudaIpcOpenMemHandle(ctypes.byref(pp), hh, ctypes.c_uint(1)))
-            ptrs.append(pp.value)
-    return ptrs
-
-
-def main():
-    dist.init_process_group("nccl")
-    rank = dist.get_rank()
-    ws = dist.get_world_size()
-    torch.cuda.set_device(rank)
-    dev = f"cuda:{rank}"
-
-    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-    if rank == 0:
-        from moe_allgather import MoeAllGather, _load_lib
-
-        _load_lib()
-    dist.barrier()
-    from moe_allgather import MoeAllGather, _load_lib
-
-    _load_lib()
-    dist.barrier()
-
-    max_size = 8 * 1024 * 1024
-    bp = ipc_buf(max_size, rank, ws)
-    dist.barrier()
-
-    # Build a fake ca_comm-like object.
-    class FakeCA:
-        pass
-
-    ca = FakeCA()
-    ca.rank = rank
-    ca.world_size = ws
-    ca.device = torch.device(dev)
-    ca.buffer_ptrs = bp
-    ca.max_size = max_size
-    # meta_ptrs not needed for Lamport approach
-    ag = MoeAllGather(ca)
-    dist.barrier()
-
-    errors = 0
-
-    # Test with various token counts.
-    for N in [1, 4, 16, 64]:
-        topk = 8
-        hd = 3584
-        sd = 448
-        ids = (
-            torch.arange(N * topk, dtype=torch.int32, device=dev) + rank * 1000
-        ).reshape(N, topk)
-        wt = torch.ones(N, topk, dtype=torch.float32, device=dev) * (rank + 1) * 0.1
-        hs = torch.full((N, hd), rank + 1, dtype=torch.uint8, device=dev)
-        sc = torch.full((N, sd), rank + 1, dtype=torch.uint8, device=dev)
-
-        ids_g, wt_g, hs_g, sc_g = ag.gather(ids, wt, hs, sc)
-
-        for src in range(ws):
-            s, e = src * N, (src + 1) * N
-            exp_ids = (
-                torch.arange(N * topk, dtype=torch.int32, device=dev) + src * 1000
-            ).reshape(N, topk)
-            if not torch.equal(ids_g[s:e], exp_ids):
-                print(f"[{rank}] FAIL ids src={src} N={N}")
-                errors += 1
-            exp_wt = torch.full(
-                (N, topk), (src + 1) * 0.1, dtype=torch.float32, device=dev
-            )
-            if not torch.allclose(wt_g[s:e], exp_wt):
-                print(f"[{rank}] FAIL wt src={src} N={N}")
-                errors += 1
-            exp_hs = torch.full((N, hd), src + 1, dtype=torch.uint8, device=dev)
-            if not torch.equal(hs_g[s:e], exp_hs):
-                print(f"[{rank}] FAIL hs src={src} N={N}")
-                errors += 1
-            exp_sc = torch.full((N, sd), src + 1, dtype=torch.uint8, device=dev)
-            if not torch.equal(sc_g[s:e], exp_sc):
-                print(f"[{rank}] FAIL sc src={src} N={N}")
-                errors += 1
-
-        # Without scales.
-        ids_g2, wt_g2, hs_g2, _ = ag.gather(ids, wt, hs)
-        for src in range(ws):
-            s, e = src * N, (src + 1) * N
-            exp_ids = (
-                torch.arange(N * topk, dtype=torch.int32, device=dev) + src * 1000
-            ).reshape(N, topk)
-            if not torch.equal(ids_g2[s:e], exp_ids):
-                print(f"[{rank}] FAIL no-sc ids src={src} N={N}")
-                errors += 1
-
-    dist.barrier()
-    print(
-        f"[rank {rank}] {'PASSED' if errors == 0 else f'FAILED ({errors})'} (ws={ws})"
-    )
-    dist.destroy_process_group()
-    return errors
-
-
-if __name__ == "__main__":
-    sys.exit(main())

csrc/standalone/custom_ar/test_moe_allgather_stress.py

@@ -1,173 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Stress test: random data, check bitwise correctness against NCCL."""
-
-import ctypes
-import os
-import sys
-
-import torch
-import torch.distributed as dist
-
-_cudart = ctypes.CDLL("libcudart.so")
-IPC = 64
-
-
-def _cc(r):
-    if r:
-        raise RuntimeError(f"CUDA err {r}")
-
-
-def ipc_buf(sz, rank, ws):
-    p = ctypes.c_void_p()
-    _cc(_cudart.cudaMalloc(ctypes.byref(p), sz))
-    _cc(_cudart.cudaMemset(p, 0, sz))
-    _cc(_cudart.cudaDeviceSynchronize())
-    h = (ctypes.c_byte * IPC)()
-    _cc(_cudart.cudaIpcGetMemHandle(ctypes.byref(h), p))
-    ah = [None] * ws
-    dist.all_gather_object(ah, bytes(h))
-    ptrs = []
-    for i in range(ws):
-        if i == rank:
-            ptrs.append(p.value)
-        else:
-            hh = (ctypes.c_byte * IPC)(*ah[i])
-            pp = ctypes.c_void_p()
-            _cc(_cudart.cudaIpcOpenMemHandle(ctypes.byref(pp), hh, ctypes.c_uint(1)))
-            ptrs.append(pp.value)
-    return ptrs
-
-
-def main():
-    dist.init_process_group("nccl")
-    rank = dist.get_rank()
-    ws = dist.get_world_size()
-    torch.cuda.set_device(rank)
-    dev = f"cuda:{rank}"
-
-    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-    if rank == 0:
-        from moe_allgather import _load_lib
-
-        _load_lib()
-    dist.barrier()
-    from moe_allgather import MoeAllGather, _load_lib
-
-    _load_lib()
-    dist.barrier()
-
-    max_size = 8 * 1024 * 1024
-    bp = ipc_buf(max_size, rank, ws)
-    dist.barrier()
-
-    class FakeCA:
-        pass
-
-    ca = FakeCA()
-    ca.rank = rank
-    ca.world_size = ws
-    ca.device = torch.device(dev)
-    ca.buffer_ptrs = bp
-    ca.max_size = max_size
-    ag = MoeAllGather(ca)
-    dist.barrier()
-
-    topk = 8
-    hd = 3584
-    sd = 448
-    errors = 0
-    total_checks = 0
-    sentinel_collisions = 0
-
-    for trial in range(200):
-        # All ranks must use the same N for NCCL reference.
-        N_tensor = torch.randint(1, 65, (1,), device=dev)
-        dist.broadcast(N_tensor, src=0)
-        N = N_tensor.item()
-        # Random data including possible sentinel values
-        ids = torch.randint(0, 256, (N, topk), dtype=torch.int32, device=dev)
-        wt = torch.randn(N, topk, dtype=torch.float32, device=dev)
-        hs = torch.randint(0, 256, (N, hd), dtype=torch.uint8, device=dev)
-        sc = torch.randint(0, 256, (N, sd), dtype=torch.uint8, device=dev)
-
-        # Count sentinel patterns in hidden_states (as uint32 view)
-        hs_u32 = hs.view(torch.int32)
-        sentinel_collisions += (hs_u32 == 0x80000000).sum().item()
-
-        # Custom kernel
-        ids_g, wt_g, hs_g, sc_g = ag.gather(ids, wt, hs, sc)
-
-        # NCCL reference
-        ids_ref = torch.empty(N * ws, topk, dtype=torch.int32, device=dev)
-        wt_ref = torch.empty(N * ws, topk, dtype=torch.float32, device=dev)
-        hs_ref = torch.empty(N * ws, hd, dtype=torch.uint8, device=dev)
-        sc_ref = torch.empty(N * ws, sd, dtype=torch.uint8, device=dev)
-        dist.all_gather_into_tensor(ids_ref, ids)
-        dist.all_gather_into_tensor(wt_ref, wt)
-        dist.all_gather_into_tensor(hs_ref, hs)
-        dist.all_gather_into_tensor(sc_ref, sc)
-
-        # Compare
-        if not torch.equal(ids_g, ids_ref):
-            mismatches = (ids_g != ids_ref).sum().item()
-            if trial < 5 or mismatches > 0:
-                print(f"[{rank}] trial={trial} ids MISMATCH: {mismatches} elements")
-            errors += 1
-        if not torch.equal(wt_g, wt_ref):
-            # Check for -0 vs +0 differences
-            bit_diff = wt_g.view(torch.int32) != wt_ref.view(torch.int32)
-            neg_zero_mask = wt_ref.view(torch.int32) == 0x80000000
-            real_errors = bit_diff & ~neg_zero_mask
-            if real_errors.any():
-                print(
-                    f"[{rank}] trial={trial} wt MISMATCH (non-negzero): {real_errors.sum().item()}"
-                )
-                errors += 1
-        if not torch.equal(hs_g, hs_ref):
-            mismatches = (hs_g != hs_ref).sum().item()
-            # Check if mismatches are due to sentinel collision
-            hs_g_u32 = hs_g.view(torch.int32)
-            hs_ref_u32 = hs_ref.view(torch.int32)
-            diff_mask = hs_g_u32 != hs_ref_u32
-            sentinel_mask = (hs_ref_u32 == 0x80000000) & diff_mask
-            non_sentinel = diff_mask & ~sentinel_mask
-            if non_sentinel.any():
-                print(
-                    f"[{rank}] trial={trial} hs NON-SENTINEL MISMATCH: {non_sentinel.sum().item()}"
-                )
-                errors += 1
-            elif sentinel_mask.any():
-                if trial < 3:
-                    print(
-                        f"[{rank}] trial={trial} hs sentinel collision: "
-                        f"{sentinel_mask.sum().item()} words (expected rare)"
-                    )
-        if not torch.equal(sc_g, sc_ref):
-            mismatches = (sc_g != sc_ref).sum().item()
-            sc_g_u32 = sc_g.view(torch.int32) if sc_g.numel() % 4 == 0 else None
-            if sc_g_u32 is not None:
-                sc_ref_u32 = sc_ref.view(torch.int32)
-                diff_mask = sc_g_u32 != sc_ref_u32
-                sentinel_mask = (sc_ref_u32 == 0x80000000) & diff_mask
-                non_sentinel = diff_mask & ~sentinel_mask
-                if non_sentinel.any():
-                    print(
-                        f"[{rank}] trial={trial} sc NON-SENTINEL MISMATCH: {non_sentinel.sum().item()}"
-                    )
-                    errors += 1
-
-        total_checks += 1
-
-    dist.barrier()
-    print(
-        f"[rank {rank}] {total_checks} trials, {errors} real errors, "
-        f"{sentinel_collisions} sentinel patterns in hs data. "
-        f"{'PASSED' if errors == 0 else 'FAILED'}"
-    )
-    dist.destroy_process_group()
-    return errors
-
-
-if __name__ == "__main__":
-    sys.exit(main())

csrc/standalone/custom_ar/test_moe_reduce_scatter.py

@@ -1,204 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Test + benchmark for Lamport MoE reduce-scatter."""
-
-import ctypes
-import os
-import sys
-
-import torch
-import torch.distributed as dist
-
-_cudart = ctypes.CDLL("libcudart.so")
-IPC = 64
-
-
-def _cc(r):
-    if r:
-        raise RuntimeError(f"CUDA err {r}")
-
-
-def ipc_buf(sz, rank, ws):
-    p = ctypes.c_void_p()
-    _cc(_cudart.cudaMalloc(ctypes.byref(p), sz))
-    _cc(_cudart.cudaMemset(p, 0, sz))
-    _cc(_cudart.cudaDeviceSynchronize())
-    h = (ctypes.c_byte * IPC)()
-    _cc(_cudart.cudaIpcGetMemHandle(ctypes.byref(h), p))
-    ah = [None] * ws
-    dist.all_gather_object(ah, bytes(h))
-    ptrs = []
-    for i in range(ws):
-        if i == rank:
-            ptrs.append(p.value)
-        else:
-            hh = (ctypes.c_byte * IPC)(*ah[i])
-            pp = ctypes.c_void_p()
-            _cc(_cudart.cudaIpcOpenMemHandle(ctypes.byref(pp), hh, ctypes.c_uint(1)))
-            ptrs.append(pp.value)
-    return ptrs
-
-
-def gpu_timer_graph(fn, warmup=20, repeats=200):
-    for _ in range(warmup):
-        fn()
-    torch.cuda.synchronize()
-    g = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(g):
-        fn()
-    for _ in range(5):
-        g.replay()
-    torch.cuda.synchronize()
-    s = torch.cuda.Event(enable_timing=True)
-    e = torch.cuda.Event(enable_timing=True)
-    s.record()
-    for _ in range(repeats):
-        g.replay()
-    e.record()
-    torch.cuda.synchronize()
-    return s.elapsed_time(e) / repeats * 1000
-
-
-def main():
-    dist.init_process_group("nccl")
-    rank = dist.get_rank()
-    ws = dist.get_world_size()
-    torch.cuda.set_device(rank)
-    dev = f"cuda:{rank}"
-
-    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-    if rank == 0:
-        from moe_reduce_scatter import _load_lib
-
-        _load_lib()
-    dist.barrier()
-    if rank != 0:
-        from moe_reduce_scatter import _load_lib
-
-        _load_lib()
-    dist.barrier()
-    from moe_reduce_scatter import MoeReduceScatter
-
-    max_size = 8 * 1024 * 1024
-    bp = ipc_buf(max_size, rank, ws)
-    dist.barrier()
-
-    class FakeCA:
-        pass
-
-    ca = FakeCA()
-    ca.rank = rank
-    ca.world_size = ws
-    ca.device = torch.device(dev)
-    ca.buffer_ptrs = bp
-    ca.max_size = max_size
-    rs = MoeReduceScatter(ca)
-    dist.barrier()
-
-    D = 7168  # DeepSeek V3 hidden_dim
-    errors = 0
-
-    # ---- Correctness tests ----
-    for N_per_rank in [1, 4, 16]:
-        N_total = N_per_rank * ws
-
-        # Each rank gets a deterministic input.
-        torch.manual_seed(42)
-        # All ranks create the SAME "ground truth" inputs for each rank.
-        all_inputs = [
-            torch.randn(N_total, D, dtype=torch.bfloat16, device=dev) for _ in range(ws)
-        ]
-        # This rank's input is all_inputs[rank].
-        my_input = all_inputs[rank]
-
-        # Custom reduce-scatter.
-        custom_out = rs.reduce_scatter(my_input)
-
-        # NCCL reference: reduce_scatter_tensor.
-        nccl_out = torch.empty(N_per_rank, D, dtype=torch.bfloat16, device=dev)
-        dist.reduce_scatter_tensor(nccl_out, my_input)
-
-        # bf16 summation order differs between our kernel and NCCL,
-        # giving ~1-2 ULP differences. Use generous tolerance.
-        max_diff = (custom_out.float() - nccl_out.float()).abs().max().item()
-        if not torch.allclose(custom_out, nccl_out, atol=0.125, rtol=0.01):
-            mismatches = (
-                ((custom_out.float() - nccl_out.float()).abs() > 0.125).sum().item()
-            )
-            print(
-                f"[{rank}] N_per_rank={N_per_rank} MISMATCH: "
-                f"max_diff={max_diff:.6f}, mismatches={mismatches}"
-            )
-            errors += 1
-        else:
-            if rank == 0:
-                print(f"  N_per_rank={N_per_rank}: PASS (max_diff={max_diff:.6f})")
-
-    # ---- Benchmark ----
-    if rank == 0:
-        print(f"\nworld_size={ws}, max_per_rank={rs.max_per_rank} bytes")
-        print(
-            f"{'config':<12} {'lamport':>10} {'lamp_graph':>10} "
-            f"{'nccl':>10} {'nccl_graph':>10} {'speedup':>8}"
-        )
-        print("-" * 65)
-
-    configs = [
-        ("1tok", 1),
-        ("2tok", 2),
-        ("4tok", 4),
-        ("8tok", 8),
-        ("16tok", 16),
-        ("32tok", 32),
-        ("64tok", 64),
-        ("128tok", 128),
-        ("256tok", 256),
-    ]
-
-    for name, N_per_rank in configs:
-        N_total = N_per_rank * ws
-        input_bytes = N_total * D * 2  # bf16
-        if input_bytes > rs.max_per_rank:
-            if rank == 0:
-                print(f"{name:<12} {'skip (too large)':>40}")
-            continue
-
-        inp = torch.randn(N_total, D, dtype=torch.bfloat16, device=dev)
-        c_out = torch.empty(N_per_rank, D, dtype=torch.bfloat16, device=dev)
-        n_out = torch.empty(N_per_rank, D, dtype=torch.bfloat16, device=dev)
-
-        def run_lamport():
-            rs.reduce_scatter(inp)
-
-        def run_nccl():
-            dist.reduce_scatter_tensor(n_out, inp)
-
-        from bench_moe_allgather import gpu_timer
-
-        lam_us = gpu_timer(run_lamport)
-        try:
-            lam_g_us = gpu_timer_graph(run_lamport)
-        except Exception:
-            lam_g_us = float("nan")
-
-        nccl_us = gpu_timer(run_nccl)
-        try:
-            nccl_g_us = gpu_timer_graph(run_nccl)
-        except Exception:
-            nccl_g_us = float("nan")
-
-        if rank == 0:
-            speedup = nccl_g_us / lam_g_us if lam_g_us > 0 else float("nan")
-            print(
-                f"{name:<12} {lam_us:>9.1f}µ {lam_g_us:>9.1f}µ "
-                f"{nccl_us:>9.1f}µ {nccl_g_us:>9.1f}µ {speedup:>7.2f}x"
-            )
-
-    dist.barrier()
-    print(f"[rank {rank}] {'PASSED' if errors == 0 else f'FAILED ({errors})'}")
-    dist.destroy_process_group()
-    return errors
-
-
-if __name__ == "__main__":
-    sys.exit(main())

csrc/standalone/custom_ar/test_rs_fused.py

@@ -1,270 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Test + benchmark: fused RS + residual + RMSNorm vs separate kernels."""
-
-import ctypes
-import os
-import sys
-
-import torch
-import torch.distributed as dist
-
-_cudart = ctypes.CDLL("libcudart.so")
-IPC = 64
-
-
-def _cc(r):
-    if r:
-        raise RuntimeError(f"CUDA err {r}")
-
-
-def ipc_buf(sz, rank, ws):
-    p = ctypes.c_void_p()
-    _cc(_cudart.cudaMalloc(ctypes.byref(p), sz))
-    _cc(_cudart.cudaMemset(p, 0, sz))
-    _cc(_cudart.cudaDeviceSynchronize())
-    h = (ctypes.c_byte * IPC)()
-    _cc(_cudart.cudaIpcGetMemHandle(ctypes.byref(h), p))
-    ah = [None] * ws
-    dist.all_gather_object(ah, bytes(h))
-    ptrs = []
-    for i in range(ws):
-        if i == rank:
-            ptrs.append(p.value)
-        else:
-            hh = (ctypes.c_byte * IPC)(*ah[i])
-            pp = ctypes.c_void_p()
-            _cc(_cudart.cudaIpcOpenMemHandle(ctypes.byref(pp), hh, ctypes.c_uint(1)))
-            ptrs.append(pp.value)
-    return ptrs
-
-
-def rms_norm_ref(x, gamma, eps):
-    """Reference RMSNorm in fp32."""
-    xf = x.float()
-    rms = torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + eps)
-    return (xf * rms * gamma.float()).to(x.dtype)
-
-
-def gpu_timer(fn, warmup=20, repeats=200):
-    for _ in range(warmup):
-        fn()
-    torch.cuda.synchronize()
-    s = torch.cuda.Event(enable_timing=True)
-    e = torch.cuda.Event(enable_timing=True)
-    s.record()
-    for _ in range(repeats):
-        fn()
-    e.record()
-    torch.cuda.synchronize()
-    return s.elapsed_time(e) / repeats * 1000
-
-
-def gpu_timer_graph(fn, warmup=20, repeats=200):
-    for _ in range(warmup):
-        fn()
-    torch.cuda.synchronize()
-    g = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(g):
-        fn()
-    for _ in range(5):
-        g.replay()
-    torch.cuda.synchronize()
-    s = torch.cuda.Event(enable_timing=True)
-    e = torch.cuda.Event(enable_timing=True)
-    s.record()
-    for _ in range(repeats):
-        g.replay()
-    e.record()
-    torch.cuda.synchronize()
-    return s.elapsed_time(e) / repeats * 1000
-
-
-def main():
-    dist.init_process_group("nccl")
-    rank = dist.get_rank()
-    ws = dist.get_world_size()
-    torch.cuda.set_device(rank)
-    dev = f"cuda:{rank}"
-
-    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-
-    # Compile fused kernel
-    if rank == 0:
-        from torch.utils.cpp_extension import load
-
-        load(
-            name="moe_rs_fused_kernel",
-            sources=[os.path.join(os.path.dirname(__file__), "moe_rs_fused.cu")],
-            extra_cuda_cflags=["-O3", "--use_fast_math"],
-            verbose=False,
-        )
-    dist.barrier()
-    from torch.utils.cpp_extension import load
-
-    fused_lib = load(
-        name="moe_rs_fused_kernel",
-        sources=[os.path.join(os.path.dirname(__file__), "moe_rs_fused.cu")],
-        extra_cuda_cflags=["-O3", "--use_fast_math"],
-        verbose=False,
-    )
-
-    # Also compile separate RS kernel for comparison
-    from moe_reduce_scatter import MoeReduceScatter
-
-    max_size = 8 * 1024 * 1024
-    bp = ipc_buf(max_size, rank, ws)
-    dist.barrier()
-
-    class FakeCA:
-        pass
-
-    ca = FakeCA()
-    ca.rank = rank
-    ca.world_size = ws
-    ca.device = torch.device(dev)
-    ca.buffer_ptrs = bp
-    ca.max_size = max_size
-
-    rs_separate = MoeReduceScatter(ca)
-
-    # Fused kernel setup (uses same buffer layout as MoeReduceScatter)
-    half_size = (max_size // 2) & ~15
-    buf_offset = half_size
-    fused_seg_cap = (half_size // 2) & ~15
-    fused_rank_stride = (fused_seg_cap // ws) & ~15
-
-    fused_buf_ptrs = torch.zeros(8, dtype=torch.int64, device=dev)
-    for i in range(ws):
-        fused_buf_ptrs[i] = bp[i] + buf_offset
-    fused_counters = torch.zeros(3, dtype=torch.int32, device=dev)
-
-    # Init sentinels for fused kernel's buffer region
-    fused_lib.lamport_init(bp[rank] + buf_offset, half_size)
-    torch.cuda.synchronize()
-    dist.barrier()
-
-    D = 7168
-    eps = 1e-6
-    gamma = torch.randn(D, dtype=torch.bfloat16, device=dev).abs() + 0.5
-    errors = 0
-
-    # ---- Correctness ----
-    for N_per_rank in [1, 4]:
-        N_total = N_per_rank * ws
-        torch.manual_seed(42 + rank)
-        moe_out = torch.randn(N_total, D, dtype=torch.bfloat16, device=dev)
-        residual = torch.randn(N_per_rank, D, dtype=torch.bfloat16, device=dev)
-
-        # Reference: NCCL RS + add + norm
-        rs_ref = torch.empty(N_per_rank, D, dtype=torch.bfloat16, device=dev)
-        dist.reduce_scatter_tensor(rs_ref, moe_out)
-        ref_residual = residual + rs_ref
-        ref_normed = rms_norm_ref(ref_residual, gamma, eps)
-
-        # Fused kernel
-        normed_out = torch.empty(N_per_rank, D, dtype=torch.bfloat16, device=dev)
-        residual_out = torch.empty(N_per_rank, D, dtype=torch.bfloat16, device=dev)
-        fused_lib.moe_rs_fused(
-            fused_buf_ptrs.data_ptr(),
-            fused_counters.data_ptr(),
-            rank,
-            ws,
-            fused_seg_cap,
-            fused_rank_stride,
-            moe_out,
-            residual,
-            gamma,
-            normed_out,
-            residual_out,
-            eps,
-        )
-        torch.cuda.synchronize()
-
-        # Compare
-        max_diff_res = (residual_out.float() - ref_residual.float()).abs().max().item()
-        max_diff_norm = (normed_out.float() - ref_normed.float()).abs().max().item()
-
-        ok = max_diff_res < 0.125 and max_diff_norm < 0.125
-        if rank == 0:
-            print(
-                f"  N_per_rank={N_per_rank}: {'PASS' if ok else 'FAIL'} "
-                f"(res_diff={max_diff_res:.4f}, norm_diff={max_diff_norm:.4f})"
-            )
-        if not ok:
-            errors += 1
-
-    # ---- Benchmark ----
-    if rank == 0:
-        print(f"\nBenchmark: D={D}, world_size={ws}")
-        print(
-            f"{'config':<10} {'fused':>10} {'fused_g':>10} "
-            f"{'RS+norm':>10} {'RS+norm_g':>10} {'speedup':>8}"
-        )
-        print("-" * 58)
-
-    for N_per_rank in [1, 2, 4, 8]:
-        N_total = N_per_rank * ws
-        input_bytes = N_total * D * 2
-        if input_bytes > fused_rank_stride:
-            if rank == 0:
-                print(f"{N_per_rank}tok      skip (too large)")
-            continue
-
-        moe_out = torch.randn(N_total, D, dtype=torch.bfloat16, device=dev)
-        residual = torch.randn(N_per_rank, D, dtype=torch.bfloat16, device=dev)
-        normed_out = torch.empty_like(residual)
-        residual_out = torch.empty_like(residual)
-        rs_out = torch.empty_like(residual)
-
-        # Fused
-        def run_fused():
-            fused_lib.moe_rs_fused(
-                fused_buf_ptrs.data_ptr(),
-                fused_counters.data_ptr(),
-                rank,
-                ws,
-                fused_seg_cap,
-                fused_rank_stride,
-                moe_out,
-                residual,
-                gamma,
-                normed_out,
-                residual_out,
-                eps,
-            )
-
-        # Separate: RS + add + norm (our Lamport RS + triton-like ops)
-        def run_separate():
-            rs_separate.reduce_scatter(moe_out)
-            # Simulate add + RMSNorm (in practice this is a fused triton kernel)
-            tmp = residual + rs_out
-            torch.rsqrt(tmp.float().pow(2).mean(-1, keepdim=True) + eps)
-
-        fused_us = gpu_timer(run_fused)
-        try:
-            fused_g = gpu_timer_graph(run_fused)
-        except Exception:
-            fused_g = float("nan")
-
-        sep_us = gpu_timer(run_separate)
-        try:
-            sep_g = gpu_timer_graph(run_separate)
-        except Exception:
-            sep_g = float("nan")
-
-        if rank == 0:
-            speedup = sep_g / fused_g if fused_g > 0 else float("nan")
-            print(
-                f"{N_per_rank}tok     {fused_us:>9.1f}µ {fused_g:>9.1f}µ "
-                f"{sep_us:>9.1f}µ {sep_g:>9.1f}µ {speedup:>7.2f}x"
-            )
-
-    dist.barrier()
-    print(f"[rank {rank}] {'PASSED' if errors == 0 else f'FAILED ({errors})'}")
-    dist.destroy_process_group()
-    return errors
-
-
-if __name__ == "__main__":
-    sys.exit(main())

docker/Dockerfile

@@ -642,7 +642,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     else \
         BITSANDBYTES_VERSION="${BITSANDBYTES_VERSION_X86}"; \
     fi; \
-    uv pip install --system accelerate hf_transfer modelscope \
+    uv pip install --system accelerate modelscope \
         "bitsandbytes>=${BITSANDBYTES_VERSION}" "timm${TIMM_VERSION}" "runai-model-streamer[s3,gcs,azure]${RUNAI_MODEL_STREAMER_VERSION}"
 
 # ============================================================
@@ -756,9 +756,10 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip install --system -e tests/vllm_test_utils
 
 # enable fast downloads from hf (for testing)
-RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --system hf_transfer
-ENV HF_HUB_ENABLE_HF_TRANSFER 1
+ENV HF_XET_HIGH_PERFORMANCE 1
+
+# increase timeout for hf downloads (for testing)
+ENV HF_HUB_DOWNLOAD_TIMEOUT 60
 
 # 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

docker/Dockerfile.cpu

@@ -173,7 +173,8 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 COPY --from=vllm-test-deps /vllm-workspace/requirements/test/cpu.txt requirements/test/cpu.txt
 
 RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install -r requirements/dev.txt && \
+    uv pip install -r requirements/lint.txt && \
+    uv pip install -r requirements/test/cpu.txt && \
     pre-commit install --hook-type pre-commit --hook-type commit-msg
 
 ENTRYPOINT ["bash"]
@@ -197,6 +198,12 @@ ADD ./.buildkite/ ./.buildkite/
 RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip install -e tests/vllm_test_utils
 
+# enable fast downloads from hf (for testing)
+ENV HF_XET_HIGH_PERFORMANCE 1
+
+# increase timeout for hf downloads (for testing)
+ENV HF_HUB_DOWNLOAD_TIMEOUT 60
+
 ######################### RELEASE IMAGE #########################
 FROM base AS vllm-openai
 

docker/Dockerfile.nightly_torch

@@ -272,9 +272,10 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip install --system -e tests/vllm_test_utils
 
 # enable fast downloads from hf (for testing)
-RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --system hf_transfer
-ENV HF_HUB_ENABLE_HF_TRANSFER 1
+ENV HF_XET_HIGH_PERFORMANCE 1
+
+# increase timeout for hf downloads (for testing)
+ENV HF_HUB_DOWNLOAD_TIMEOUT 60
 
 RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip install --system -r requirements/test/nightly-torch.txt

docker/Dockerfile.rocm

@@ -365,9 +365,10 @@ RUN cd /vllm-workspace \
     && python3 -m pip install pytest-shard
 
 # enable fast downloads from hf (for testing)
-RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --system hf_transfer
-ENV HF_HUB_ENABLE_HF_TRANSFER=1
+ENV HF_XET_HIGH_PERFORMANCE=1
+
+# increase timeout for hf downloads (for testing)
+ENV HF_HUB_DOWNLOAD_TIMEOUT 60
 
 # install audio decode package `torchcodec` from source (required due to 
 # ROCm and torch version mismatch) for tests with datasets package

docker/Dockerfile.s390x

@@ -42,7 +42,7 @@ FROM python-install AS pyarrow
 # Build Apache Arrow
 WORKDIR /tmp
 RUN --mount=type=cache,target=/root/.cache/uv \
-    git clone https://github.com/apache/arrow.git && \
+    git clone https://github.com/apache/arrow.git  -b maint-19.0.1 && \
     cd arrow/cpp && \
     mkdir release && cd release && \
     cmake -DCMAKE_BUILD_TYPE=Release \
@@ -68,19 +68,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip install -r requirements-build.txt && \
     python setup.py build_ext --build-type=$ARROW_BUILD_TYPE --bundle-arrow-cpp bdist_wheel
 
-FROM python-install AS numa-build
-# Install numactl (needed for numa.h dependency)
-WORKDIR /tmp
-RUN curl -LO https://github.com/numactl/numactl/archive/refs/tags/v2.0.16.tar.gz && \
-    tar -xvzf v2.0.16.tar.gz && \
-    cd numactl-2.0.16 && \
-    ./autogen.sh && \
-    ./configure && \
-    make
-
-# Set include path
-ENV C_INCLUDE_PATH="/usr/local/include:$C_INCLUDE_PATH"
-
 FROM python-install AS rust
 ENV CARGO_HOME=/root/.cargo
 ENV RUSTUP_HOME=/root/.rustup
@@ -91,6 +78,18 @@ RUN curl https://sh.rustup.rs -sSf | sh -s -- -y && \
     rustup default stable && \
     rustup show
 
+FROM python-install AS numa-build
+WORKDIR /tmp
+RUN curl -LO https://github.com/numactl/numactl/archive/refs/tags/v2.0.19.tar.gz && \
+    tar -xvzf v2.0.19.tar.gz && \
+    cd numactl-2.0.19 && \
+    ./autogen.sh && \
+    ./configure && \
+    make
+
+# Set include path
+ENV C_INCLUDE_PATH="/usr/local/include:$C_INCLUDE_PATH"
+
 FROM python-install AS torch-vision
 # Install torchvision
 ARG TORCH_VISION_VERSION=v0.26.0
@@ -133,7 +132,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     git clone --recursive https://github.com/numba/llvmlite.git -b v0.44.0 && \
     git clone --recursive https://github.com/numba/numba.git -b ${NUMBA_VERSION} && \
     cd llvm-project && mkdir build && cd  build && \
-    uv pip install 'cmake<4' setuptools numpy && \
+    uv pip install 'cmake<4' 'setuptools<70' numpy && \
     export PREFIX=/usr/local && CMAKE_ARGS="${CMAKE_ARGS} -DLLVM_ENABLE_PROJECTS=lld;libunwind;compiler-rt" \
     CFLAGS="$(echo $CFLAGS | sed 's/-fno-plt //g')" \
     CXXFLAGS="$(echo $CXXFLAGS | sed 's/-fno-plt //g')" \
@@ -193,27 +192,22 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     cd opencv-python && \
     python -m build --wheel --installer=uv --outdir /tmp/opencv-python/dist
 
-# Build Outlines Core
-FROM python-install AS outlines-core-builder
+## Todo(r3hankhan123): Remove guidance-builder stage once vLLM upgrades to new version of llguidance that fixes s390x issues. See https://github.com/guidance-ai/llguidance/issues/330
+FROM python-install AS guidance-builder
 WORKDIR /tmp
 ENV CARGO_HOME=/root/.cargo
 ENV RUSTUP_HOME=/root/.rustup
 ENV PATH="$CARGO_HOME/bin:$RUSTUP_HOME/bin:$PATH"
-COPY requirements/common.txt /tmp/requirements/common.txt
-ARG OUTLINES_CORE_VERSION
 RUN --mount=type=cache,target=/root/.cache/uv \
     --mount=type=bind,from=rust,source=/root/.cargo,target=/root/.cargo,rw \
     --mount=type=bind,from=rust,source=/root/.rustup,target=/root/.rustup,rw \
-    OUTLINES_CORE_VERSION=${OUTLINES_CORE_VERSION:-$(grep -E '^outlines_core\s*==\s*[0-9.]+' /tmp/requirements/common.txt | grep -Eo '[0-9.]+')} && \
-    if [ -z "${OUTLINES_CORE_VERSION}" ]; then echo "ERROR: Could not determine outlines_core version"; exit 1; fi && \
-    git clone https://github.com/dottxt-ai/outlines-core.git && \
-    cd outlines-core && \
-    git checkout tags/${OUTLINES_CORE_VERSION} && \
-    sed -i "s/version = \"0.0.0\"/version = \"${OUTLINES_CORE_VERSION}\"/" Cargo.toml && \
+    git clone https://github.com/guidance-ai/llguidance.git && \
+    cd llguidance && \
+    git checkout s390x-fix-v2  && \
     uv pip install maturin && \
-    python -m maturin build --release --out dist
+    python -m maturin build --release --out dist --compatibility linux
 
-# Final build stage
+# # Final build stage
 FROM python-install AS vllm-cpu
 ARG PYTHON_VERSION
 ARG PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/cpu"
@@ -229,10 +223,12 @@ ENV PKG_CONFIG_PATH="/opt/rh/gcc-toolset-14/root/usr/lib64/pkgconfig:/usr/local/
 ENV PATH="${VIRTUAL_ENV:+${VIRTUAL_ENV}/bin}:/opt/rh/gcc-toolset-14/root/usr/bin:/usr/local/bin:$CARGO_HOME/bin:$RUSTUP_HOME/bin:$PATH"
 ENV PIP_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
 ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
+# Force pure Python protobuf to avoid s390x C++ extension crashes
+ENV PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python
 COPY . /workspace/vllm
 WORKDIR /workspace/vllm
 
-RUN --mount=type=bind,from=numa-build,src=/tmp/numactl-2.0.16,target=/numactl \
+RUN --mount=type=bind,from=numa-build,src=/tmp/numactl-2.0.19,target=/numactl \
     make -C /numactl install
 
 # Install dependencies, including PyTorch and Apache Arrow
@@ -245,22 +241,22 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     --mount=type=bind,from=numba-builder,source=/tmp/llvmlite/dist,target=/tmp/llvmlite-wheels/ \
     --mount=type=bind,from=numba-builder,source=/tmp/numba/dist,target=/tmp/numba-wheels/ \
     --mount=type=bind,from=opencv-builder,source=/tmp/opencv-python/dist,target=/tmp/opencv-wheels/ \
-    --mount=type=bind,from=outlines-core-builder,source=/tmp/outlines-core/dist,target=/tmp/outlines-core/dist/ \
-     ARROW_WHL_FILE=$(ls /tmp/arrow-wheels/pyarrow-*.whl) && \
+    --mount=type=bind,from=guidance-builder,source=/tmp/llguidance/dist,target=/tmp/guidance-wheels/ \
+     ARROW_WHL_FILE=$(ls /tmp/arrow-wheels/*.whl) && \
      VISION_WHL_FILE=$(ls /tmp/vision-wheels/*.whl) && \
      HF_XET_WHL_FILE=$(ls /tmp/hf-xet-wheels/*.whl) && \
      LLVM_WHL_FILE=$(ls /tmp/llvmlite-wheels/*.whl) && \
      NUMBA_WHL_FILE=$(ls /tmp/numba-wheels/*.whl) && \
      OPENCV_WHL_FILE=$(ls /tmp/opencv-wheels/*.whl) && \
-     OUTLINES_CORE_WHL_FILE=$(ls /tmp/outlines-core/dist/*.whl) && \
-     uv pip install -v \
-        $ARROW_WHL_FILE  \
+     GUIDANCE_WHL_FILE=$(ls /tmp/guidance-wheels/*.whl) && \
+     uv pip install -v \    
+        $ARROW_WHL_FILE \
         $VISION_WHL_FILE \
         $HF_XET_WHL_FILE \
         $LLVM_WHL_FILE \
         $NUMBA_WHL_FILE \
         $OPENCV_WHL_FILE \
-        $OUTLINES_CORE_WHL_FILE \
+        $GUIDANCE_WHL_FILE \
         --index-strategy unsafe-best-match \
         -r requirements/build/cpu.txt \
         -r requirements/cpu.txt
@@ -271,6 +267,10 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     VLLM_TARGET_DEVICE=cpu VLLM_CPU_MOE_PREPACK=0 python setup.py bdist_wheel && \
     uv pip install "$(echo dist/*.whl)[tensorizer]"
 
+# Remove protobuf C++ extension that crashes on s390x
+RUN rm -rf /opt/vllm/lib64/python${PYTHON_VERSION}/site-packages/google/_upb/*.so \
+           /opt/vllm/lib64/python${PYTHON_VERSION}/site-packages/google/protobuf/pyext/*.so 2>/dev/null || true
+
 # setup non-root user for vllm
 RUN umask 002 && \
     /usr/sbin/useradd --uid 2000 --gid 0 vllm && \

docs/benchmarking/cli.md

@@ -37,6 +37,7 @@ th {
 | HuggingFace-Blazedit | ✅ | ✅ | `vdaita/edit_5k_char`, `vdaita/edit_10k_char` |
 | HuggingFace-ASR | ✅ | ✅ | `openslr/librispeech_asr`, `facebook/voxpopuli`,  `LIUM/tedlium`, `edinburghcstr/ami`,        `speechcolab/gigaspeech`,        `kensho/spgispeech` |
 | Spec Bench | ✅ | ✅ | `wget https://raw.githubusercontent.com/hemingkx/Spec-Bench/refs/heads/main/data/spec_bench/question.jsonl` |
+| SPEED-Bench | ✅ | ✅ | `curl -LsSf https://raw.githubusercontent.com/NVIDIA-NeMo/Skills/refs/heads/main/nemo_skills/dataset/speed-bench/prepare.py \| python3 -` |
 | Custom | ✅ | ✅ | Local file: `data.jsonl` |
 | Custom MM | ✅ | ✅ | Local file: `mm_data.jsonl` |
 
@@ -239,6 +240,69 @@ vllm bench serve \
     --spec-bench-category "summarization"
 ```
 
+#### SPEED-Bench Benchmark with Speculative Decoding
+
+[SPEED-Bench](https://huggingface.co/datasets/nvidia/SPEED-Bench) is a unified and diverse dataset for speculative decoding, supporting acceptance rate and length measurements using the Qualitative split and throughput measurements using the Throughput splits in 5 configuration of input sequence length (1k, 2k, 8k, 16k, 32k).
+
+!!! note
+    This dataset is governed by the [NVIDIA Evaluation Dataset License Agreement](https://huggingface.co/datasets/nvidia/SPEED-Bench/blob/main/License.pdf). For each dataset a user elects to use, the user is responsible for checking if the dataset license is fit for the intended purpose. The `prepare.py` script automatically fetches data from all the source datasets.
+
+First, download the dataset to a folder, using this one liner:
+
+```bash
+curl -LsSf https://raw.githubusercontent.com/NVIDIA-NeMo/Skills/refs/heads/main/nemo_skills/dataset/speed-bench/prepare.py \| python3 -
+```
+
+The command supports also the following arguments:
+
+- `--config`: download only a subset of the dataset: `qualitative`, `throughput_1k`, `throughput_2k`, `throughput_8k`, `throughput_16k` and `throughput_32k`. By default, it will download all subsets.
+- `--output_dir`: download to a specified folder. By default, it will download to the current directory.
+
+Start a server with speculative decoding:
+
+```bash
+vllm serve meta-llama/Llama-3.3-70B-Instruct \
+    --speculative-config $'{"method": "eagle3",
+    "num_speculative_tokens": 3,
+    "model": "nvidia/Llama-3.3-70B-Instruct-Eagle3"}'
+```
+
+Run all categories in the Qualitative split:
+
+```bash
+vllm bench serve \
+    --model meta-llama/Llama-3.3-70B-Instruct \
+    --dataset-name speed_bench \
+    --dataset-path "<YOUR_DOWNLOADED_PATH>/data/speed_bench" \
+    --num-prompts -1
+```
+
+Available categories include `[writing, roleplay, reasoning, math, coding, stem, humanities, multilingual, summarization, qa, rag]`.
+
+Run only a specific category like "multilingual":
+
+```bash
+vllm bench serve \
+    --model meta-llama/Llama-3.3-70B-Instruct \
+    --dataset-name speed_bench \
+    --dataset-path "<YOUR_DOWNLOADED_PATH>/data/speed_bench" \
+    --num-prompts -1
+    --speed-bench-category "multilingual"
+```
+
+Run all categories in the Throughput split (2k ISL):
+
+```bash
+vllm bench serve \
+    --model meta-llama/Llama-3.3-70B-Instruct \
+    --dataset-name speed_bench \
+    --speed-bench-dataset-subset throughput_2k
+    --dataset-path "<YOUR_DOWNLOADED_PATH>/data/speed_bench/" \
+    --num-prompts -1
+```
+
+Available categories include `[high_entropy, mixed, low_entropy]`, where high entropy data contains unstructued data such as creative writing while low entropy data contains more structured data such as coding, more details are in the dataset card.
+
 #### Other HuggingFaceDataset Examples
 
 ```bash

docs/design/attention_backends.md

@@ -106,6 +106,7 @@ Priority is **1 = highest** (tried first).
 | 2 | `FLASH_ATTN` |
 | 3 | `TRITON_ATTN` |
 | 4 | `FLEX_ATTENTION` |
+| 5 | `TURBOQUANT` |
 
 **Ampere/Hopper (SM 8.x-9.x):**
 
@@ -115,6 +116,7 @@ Priority is **1 = highest** (tried first).
 | 2 | `FLASHINFER` |
 | 3 | `TRITON_ATTN` |
 | 4 | `FLEX_ATTENTION` |
+| 5 | `TURBOQUANT` |
 
 ### MLA Attention (DeepSeek-style)
 
@@ -130,16 +132,6 @@ Priority is **1 = highest** (tried first).
 | 6 | `FLASHINFER_MLA_SPARSE`**\*** |
 | 7 | `FLASHMLA_SPARSE` |
 
-**Ampere/Hopper (SM 8.x-9.x):**
-
-| Priority | Backend |
-| -------- | ------- |
-| 1 | `FLASH_ATTN_MLA` |
-| 2 | `FLASHMLA` |
-| 3 | `FLASHINFER_MLA` |
-| 4 | `TRITON_MLA` |
-| 5 | `FLASHMLA_SPARSE` |
-
 > **\*** For sparse MLA, FP8 KV cache always prefers `FLASHINFER_MLA_SPARSE`. With BF16 KV cache, `FLASHINFER_MLA_SPARSE` is preferred for low query-head counts (<= 16), while `FLASHMLA_SPARSE` is preferred otherwise.
 >
 > **Note:** ROCm and CPU platforms have their own selection logic. See the platform-specific documentation for details.
@@ -178,6 +170,7 @@ Priority is **1 = highest** (tried first).
 | `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ❌ | ✅ | ❌ | Decoder, Encoder, Encoder Only | N/A |
 | `TREE_ATTN` | | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | Decoder | Any |
 | `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `int8_per_token_head`, `fp8_per_token_head` | %16 | Any | ✅ | ✅ | ❌ | All | Any |
+| `TURBOQUANT` | | fp16, bf16 | `turboquant_k8v4`, `turboquant_4bit_nc`, `turboquant_k3v4_nc`, `turboquant_3bit_nc` | 16, 32, 64, 128 | Any | ❌ | ❌ | ❌ | Decoder | Any |
 
 > **†** FlashInfer uses TRTLLM attention on Blackwell (SM100), which supports sinks. Disable via `--attention-config.use_trtllm_attention=0`.
 >
@@ -206,6 +199,7 @@ configuration.
 
 | Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. |
 | ------- | ------ | --------- | ----------- | ---------- | ---- | ------ | --------- | --- | --------------- | ------------ |
+| `CUTLASS_FA3_MLA_SPARSE` | bf16 | `auto` | 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x |
 | `CUTLASS_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
 | `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
 | `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x |

docs/design/cuda_graphs_multimodal.md

@@ -28,6 +28,7 @@ Multiple CUDA Graphs are pre-captured at different **token budget** levels (e.g.
 class BudgetGraphMetadata:
     token_budget: int
     max_batch_size: int
+    max_frames_per_batch: int
     graph: torch.cuda.CUDAGraph
     input_buffer: torch.Tensor       # e.g. pixel_values
     metadata_buffers: dict[str, torch.Tensor]  # e.g. embeddings, seq metadata
@@ -51,6 +52,15 @@ For each graph replay:
 
 When `mm_encoder_tp_mode="data"`, the manager distributes images across TP ranks using load-balanced assignment via `get_load_balance_assignment`, executes locally on each rank, then gathers results back in the original order via `tensor_model_parallel_all_gather`.
 
+### Video inference support (experimental)
+
+Following <https://github.com/vllm-project/vllm/pull/35963> (ViT full CUDA graph support for image inference), <https://github.com/vllm-project/vllm/pull/38061> extends the encoder CUDA graph framework to support video inference for Qwen3-VL. Previously, the CUDA graph capture/replay path only handled image inputs (`pixel_values` + `image_grid_thw`). Video inputs use different keys (`pixel_values_videos` + `video_grid_thw`) and require larger `cu_seqlens` buffers because each video item contributes multiple frames (`T` attention sequences). This PR generalizes the protocol and manager to handle both modalities through a single shared graph manager.
+
+!!! note
+    Video CUDA graphs are automatically disabled when EVS (Efficient Video Sampling) pruning is enabled, since EVS makes the token count data-dependent and incompatible with CUDA graph capture.
+
+    Currently, we only support image-only or video-only inputs when enabling CUDA graph, mixed inputs (image + video) are not supported yet (we will work on it in the near future). Thus, it's recommended to turn off the image modality by `--limit-mm-per-prompt '{"image": 0}'` for video-only inputs.
+
 ## Model integration via `SupportsEncoderCudaGraph`
 
 Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGraph][vllm.model_executor.models.interfaces.SupportsEncoderCudaGraph] protocol. This protocol encapsulates all model-specific logic so that the manager remains model-agnostic. The protocol defines the following methods:
@@ -65,12 +75,17 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra
 * `prepare_encoder_cudagraph_replay_buffers(...)` — computes new buffer values from actual batch inputs before replay.
 * `encoder_cudagraph_forward(...)` — forward pass using precomputed buffers (called during capture and replay).
 * `encoder_eager_forward(...)` — fallback eager forward when no graph fits.
-
-Currently supported: **Qwen3-VL** (see `vllm/model_executor/models/qwen3_vl.py`).
+* `get_input_modality(...)` - return the modality of the inputs.
 
 !!! note
     The `SupportsEncoderCudaGraph` protocol is designed to be model-agnostic. New vision encoder models can opt-in by implementing the protocol methods without modifying the manager.
 
+**Supported models:**
+
+| Architecture | Models | CG for Image | CG for Video |
+| ------------ | ------ | ------------ | ------------ |
+| `Qwen3VLForConditionalGeneration` | `Qwen3-VL` | ✅︎ | ✅︎ |
+
 !!! note
     Encoder CUDA Graphs have currently been tested with `--mm-encoder-attn-backend=FLASH_ATTN` and `--mm-encoder-attn-backend=FLASHINFER` on Blackwell GPUs.
 
@@ -80,10 +95,13 @@ Three fields in `CompilationConfig` control encoder CUDA Graphs:
 
 * `cudagraph_mm_encoder` (`bool`, default `False`) — enable CUDA Graph capture for multimodal encoder. When enabled, captures the full encoder forward as a CUDA Graph for each token budget level.
 * `encoder_cudagraph_token_budgets` (`list[int]`, default `[]`) — token budget levels for capture. If empty (default), auto-inferred from model architecture as power-of-2 levels. User-provided values override auto-inference.
-* `encoder_cudagraph_max_images_per_batch` (`int`, default `0`) — maximum number of images per batch during capture. If 0 (default), auto-inferred as `max_budget // min_budget`.
+* `encoder_cudagraph_max_vision_items_per_batch` (`int`, default `0`) — maximum number of images/videos per batch during capture. If 0 (default), auto-inferred as `max_budget // min_budget`.
+* `encoder_cudagraph_max_frames_per_batch` (`int`, default `0`) — maximum number of video frames per batch during capture. If 0 (default), auto-inferred as `encoder_cudagraph_max_vision_items_per_batch * 2` (to be optimized).
 
 ## Usage guide
 
+### Image inference
+
 Enable encoder CUDA Graphs via `compilation_config`:
 
 ```bash
@@ -95,7 +113,7 @@ With explicit budgets:
 
 ```bash
 vllm serve Qwen/Qwen3-VL-32B \
-  --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [2048, 4096, 8192, 13824], "encoder_cudagraph_max_images_per_batch": 8}'
+  --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [2048, 4096, 8192, 13824], "encoder_cudagraph_max_vision_items_per_batch": 8}'
 ```
 
 Python example:
@@ -107,7 +125,7 @@ compilation_config = {
     "cudagraph_mm_encoder": True,
     # Optional: override auto-inferred budgets
     # "encoder_cudagraph_token_budgets": [2048, 4096, 8192, 13824],
-    # "encoder_cudagraph_max_images_per_batch": 8,
+    # "encoder_cudagraph_max_vision_items_per_batch": 8,
 }
 
 model = vllm.LLM(
@@ -118,6 +136,44 @@ model = vllm.LLM(
 
 The manager tracks hit/miss statistics and logs them periodically. A "hit" means an image was processed via CUDA Graph replay; a "miss" means eager fallback (image exceeded all budgets).
 
+### Video inference
+
+Enable encoder CUDA Graphs via `compilation_config`:
+
+```bash
+vllm serve Qwen/Qwen3-VL-32B \
+  --limit-mm-per-prompt '{"image": 0}' \
+  --compilation-config '{"cudagraph_mm_encoder": true}'
+```
+
+With explicit budgets:
+
+```bash
+vllm serve Qwen/Qwen3-VL-32B \
+  --limit-mm-per-prompt '{"image": 0}' \
+  --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [2048, 4096, 8192, 13824], "encoder_cudagraph_max_vision_items_per_batch": 8, "encoder_cudagraph_max_frames_per_batch": 64}'
+```
+
+Python example:
+
+```python
+import vllm
+
+compilation_config = {
+    "cudagraph_mm_encoder": True,
+    # Optional: override auto-inferred budgets
+    # "encoder_cudagraph_token_budgets": [2048, 4096, 8192, 13824],
+    # "encoder_cudagraph_max_vision_items_per_batch": 8,
+    # "encoder_cudagraph_max_frames_per_batch": 64,
+}
+
+model = vllm.LLM(
+    model="Qwen/Qwen3-VL-32B",
+    limit_mm_per_prompt='{"image": 0}',
+    compilation_config=compilation_config,
+)
+```
+
 ## About the Performance
 
 The following benchmarks were run on Blackwell GPUs (GB200) using `vllm bench mm-processor`. See [#35963](https://github.com/vllm-project/vllm/pull/35963) for full details.
@@ -140,7 +196,7 @@ vllm bench mm-processor \
   --num-prompts 3000 --num-warmups 300 \
   --max-model-len 32768 --seed 42 \
   --mm-encoder-attn-backend FLASH_ATTN \
-  --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [512, 1024, 1536, 2048, 2560, 3072, 3584, 4096, 4864], "encoder_cudagraph_max_images_per_batch": 8}'
+  --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [512, 1024, 1536, 2048, 2560, 3072, 3584, 4096, 4864], "encoder_cudagraph_max_vision_items_per_batch": 8}'
 ```
 
 ### Multi-GPU (4x GB200, TP=4, DP=4)
@@ -165,5 +221,8 @@ vllm bench mm-processor \
   --max-model-len 8192 --seed 42 \
   --mm-encoder-attn-backend FLASHINFER \
   --tensor-parallel-size 4 --mm-encoder-tp-mode data \
-  --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [512, 1024, 1536, 2048, 2560, 3072, 3584, 4096, 4864], "encoder_cudagraph_max_images_per_batch": 8}'
+  --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [512, 1024, 1536, 2048, 2560, 3072, 3584, 4096, 4864], "encoder_cudagraph_max_vision_items_per_batch": 8}'
 ```
+
+!!! note
+    Find more details about benchmarks on GPUs (A100) for video inference at [#38061](https://github.com/vllm-project/vllm/pull/38061).

docs/design/metrics.md

@@ -42,7 +42,7 @@ These are documented under [Inferencing and Serving -> Production Metrics](../us
 
 ### Grafana Dashboard
 
-vLLM also provides [a reference example](../../examples/online_serving/prometheus_grafana/README.md) for how to collect and store these metrics using Prometheus and visualize them using a Grafana dashboard.
+vLLM also provides [a reference example](../../examples/observability/prometheus_grafana/README.md) for how to collect and store these metrics using Prometheus and visualize them using a Grafana dashboard.
 
 The subset of metrics exposed in the Grafana dashboard gives us an indication of which metrics are especially important:
 
@@ -657,7 +657,7 @@ vLLM has support for OpenTelemetry tracing:
 - Added by <https://github.com/vllm-project/vllm/pull/4687> and reinstated by <https://github.com/vllm-project/vllm/pull/20372>
 - Configured with `--oltp-traces-endpoint` and `--collect-detailed-traces`
 - [OpenTelemetry blog post](https://opentelemetry.io/blog/2024/llm-observability/)
-- [User-facing docs](../../examples/online_serving/opentelemetry/README.md)
+- [User-facing docs](../../examples/observability/opentelemetry/README.md)
 - [Blog post](https://medium.com/@ronen.schaffer/follow-the-trail-supercharging-vllm-with-opentelemetry-distributed-tracing-aa655229b46f)
 - [IBM product docs](https://www.ibm.com/docs/en/instana-observability/current?topic=mgaa-monitoring-large-language-models-llms-vllm-public-preview)
 

docs/design/moe_kernel_features.md

@@ -86,7 +86,7 @@ To be used with a particular `FusedMoEPrepareAndFinalizeModular` subclass, MoE k
 | cutlass_fp4 | standard,</br>batched | nvfp4 | A,T | silu | Y | Y | [`CutlassExpertsFp4`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassExpertsFp4] |
 | cutlass_fp8 | standard,</br>batched | fp8 | A,T | silu, gelu | Y | Y | [`CutlassExpertsFp8`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassExpertsFp8],</br>[`CutlasBatchedExpertsFp8`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassBatchedExpertsFp8] |
 | flashinfer | standard | nvfp4,</br>fp8 | T | <sup>5</sup> | N | Y | [`FlashInferExperts`][vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe.FlashInferExperts] |
-| gpt oss triton | standard | N/A | N/A | <sup>5</sup> | Y | Y | [`triton_kernel_fused_experts`][vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe.triton_kernel_fused_experts],</br>[`OAITritonExperts`][vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe.OAITritonExperts] |
+| gpt oss triton | standard | N/A | N/A | <sup>5</sup> | Y | Y | [`triton_kernel_fused_experts`][vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe.triton_kernel_fused_experts],</br>[`OAITritonExperts`][vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe.OAITritonExperts] |
 | marlin | standard,</br>batched | <sup>3</sup> / N/A | <sup>3</sup> / N/A | silu,</br>swigluoai | Y | Y | [`fused_marlin_moe`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.fused_marlin_moe],</br>[`MarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.MarlinExperts],</br>[`BatchedMarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.BatchedMarlinExperts] |
 | trtllm | standard | mxfp4,</br>nvfp4 | G(16),G(32) | <sup>5</sup> | N | Y | [`TrtLlmMxfp4ExpertsMonolithic`][vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe.TrtLlmMxfp4ExpertsMonolithic],</br>[`TrtLlmMxfp4ExpertsModular`][vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe.TrtLlmMxfp4ExpertsModular],</br>[`TrtLlmNvFp4ExpertsMonolithic`][vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe.TrtLlmNvFp4ExpertsMonolithic],</br>[`TrtLlmNvfp4ExpertsModular`][vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe.TrtLlmNvFp4ExpertsModular] |
 | rocm aiter moe | standard | mxfp4,</br>fp8 | G(32),G(128),A,T | silu, gelu,</br>swigluoai | Y | N | `rocm_aiter_fused_experts`,</br>`AiterExperts` |

docs/features/quantization/README.md

@@ -16,6 +16,7 @@ The following are the supported quantization formats for vLLM:
 - [INT8 W8A8](int8.md)
 - [FP8 W8A8](fp8.md)
 - [NVIDIA Model Optimizer](modelopt.md)
+- [Online Quantization](online.md)
 - [AMD Quark](quark.md)
 - [Quantized KV Cache](quantized_kvcache.md)
 - [TorchAO](torchao.md)

docs/features/quantization/online.md

@@ -0,0 +1,94 @@
+# Online Quantization
+
+Online quantization lets you take a BF16/FP16 model and quantize its Linear
+and MoE weights to lower precision (such as FP8) at load time, without needing
+a pre-quantized checkpoint or calibration data. Weights are converted during
+model loading and activations are dynamically scaled during each forward pass.
+
+## Quick Start
+
+Pass a scheme name to the `quantization` parameter:
+
+```python
+from vllm import LLM
+
+# Per-tensor FP8 quantization (one scale per weight tensor)
+llm = LLM("meta-llama/Llama-3.1-8B", quantization="fp8_per_tensor")
+
+# Per-block FP8 quantization (128x128 block scaling for weights and 1x128 block scaling for activations)
+llm = LLM("meta-llama/Llama-3.1-8B", quantization="fp8_per_block")
+```
+
+Or with the CLI:
+
+```bash
+vllm serve meta-llama/Llama-3.1-8B --quantization fp8_per_tensor
+vllm serve meta-llama/Llama-3.1-8B --quantization fp8_per_block
+```
+
+## Supported Schemes
+
+| Scheme | Weight recipe | Activation recipe | Notes |
+| ------ | ------------- | ------------------ | ----- |
+| `fp8_per_tensor` | fp8_e4m3 data, fp32 per-tensor scale | fp8_e4m3 data, fp32 per-tensor scale | On some GPUs (Ada, Hopper) linear activations use per-token scaling for better performance |
+| `fp8_per_block` | fp8_e4m3 data, fp32 per-128x128-block scale | fp8_e4m3 data, fp32 per-1x128-block scale | |
+
+Support for additional schemes will be added in future versions of vllm.
+
+## Advanced Configuration
+
+For fine-grained control, use a `quantization_config` dictionary.
+
+### Separate Schemes for Dense and MoE Layers
+
+You can apply different quantization schemes to dense linear layers and MoE expert layers:
+
+```python
+from vllm import LLM
+
+llm = LLM(
+    "ibm-granite/granite-3.0-1b-a400m-base",
+    quantization="fp8_per_tensor",
+    quantization_config={
+        "linear_scheme_override": "fp8_per_block",
+    },
+)
+```
+
+Or,
+
+```python
+from vllm import LLM
+
+llm = LLM(
+    "ibm-granite/granite-3.0-1b-a400m-base",
+    quantization="fp8_per_tensor",
+    quantization_config={
+        "moe_scheme_override": "fp8_per_block",
+    },
+)
+```
+
+### Excluding Layers from Quantization
+
+Use the `ignore` parameter to skip specific layers. It accepts exact layer names and regex patterns (prefixed with `re:`):
+
+```python
+from vllm import LLM
+
+llm = LLM(
+    "ibm-granite/granite-3.0-1b-a400m-base",
+    quantization="fp8_per_tensor",
+    quantization_config={
+        "ignore": [
+            # exact layer name
+            "model.layers.1.self_attn.o_proj",
+            # regex: skip all QKV projections
+            "re:.*[qkv]_proj",
+        ],
+    },
+)
+```
+
+!!! note
+    For fused layers (e.g., `qkv_proj` which fuses `q_proj`, `k_proj`, `v_proj`), the ignore pattern must match the **unfused** shard names (`q_proj`, `k_proj`, `v_proj`), not the fused name.

docs/getting_started/installation/cpu.s390x.inc.md

@@ -3,15 +3,15 @@
 
 vLLM has experimental support for s390x architecture on IBM Z platform. For now, users must build from source to natively run on IBM Z platform.
 
-Currently, the CPU implementation for s390x architecture supports FP32 datatype only.
+Currently, the CPU implementation for s390x architecture supports FP32, BF16 and FP16.
 
 --8<-- [end:installation]
 --8<-- [start:requirements]
 
 - OS: `Linux`
-- SDK: `gcc/g++ >= 12.3.0` or later with Command Line Tools
+- SDK: `gcc/g++ >= 14.0.0` or later with Command Line Tools
 - Instruction Set Architecture (ISA): VXE support is required. Works with Z14 and above.
-- Build install python packages: `pyarrow`, `torch` and `torchvision`
+- Build install python packages: `torchvision`, `llvmlite`, `numba`, `pyarrow (for testing)`, `opencv-headless`
 
 --8<-- [end:requirements]
 --8<-- [start:set-up-using-python]
@@ -24,13 +24,14 @@ Currently, there are no pre-built IBM Z CPU wheels.
 --8<-- [end:pre-built-wheels]
 --8<-- [start:build-wheel-from-source]
 
-Install the following packages from the package manager before building the vLLM. For example on RHEL 9.4:
+Install the following packages from the package manager before building the vLLM. For example on RHEL 9.6:
 
 ```bash
 dnf install -y \
-    which procps findutils tar vim git gcc g++ make patch make cython zlib-devel \
+    which procps findutils tar vim git gcc-toolset-14 gcc-toolset-14-binutils gcc-toolset-14-libatomic-devel zlib-devel \
     libjpeg-turbo-devel libtiff-devel libpng-devel libwebp-devel freetype-devel harfbuzz-devel \
-    openssl-devel openblas openblas-devel wget autoconf automake libtool cmake numactl-devel
+    openssl-devel openblas openblas-devel autoconf automake libtool cmake numpy libsndfile \
+    clang llvm-devel llvm-static clang-devel
 ```
 
 Install rust>=1.80 which is needed for `outlines-core` and `uvloop` python packages installation.
@@ -43,13 +44,13 @@ curl https://sh.rustup.rs -sSf | sh -s -- -y && \
 Execute the following commands to build and install vLLM from source.
 
 !!! tip
-    Please build the following dependencies, `torchvision`, `pyarrow` from source before building vLLM.
+    Please build the following dependencies, `torchvision`, `llvmlite`, `numba`, `llguidance`, `pyarrow`, `opencv-headless` from source before building vLLM.
 
 ```bash
-    sed -i '/^torch/d' requirements/build/cuda.txt    # remove torch from requirements/build/cuda.txt since we use nightly builds
     uv pip install -v \
+        --extra-index-url https://download.pytorch.org/whl/cpu \
         --torch-backend auto \
-        -r requirements/build/cuda.txt \
+        -r requirements/build/cpu.txt \
         -r requirements/cpu.txt \
     VLLM_TARGET_DEVICE=cpu python setup.py bdist_wheel && \
         uv pip install dist/*.whl
@@ -57,10 +58,9 @@ Execute the following commands to build and install vLLM from source.
 
 ??? console "pip"
     ```bash
-        sed -i '/^torch/d' requirements/build/cuda.txt    # remove torch from requirements/build/cuda.txt since we use nightly builds
         pip install -v \
-            --extra-index-url https://download.pytorch.org/whl/nightly/cpu \
-            -r requirements/build/cuda.txt \
+            --extra-index-url https://download.pytorch.org/whl/cpu \
+            -r requirements/build/cpu.txt \
             -r requirements/cpu.txt \
         VLLM_TARGET_DEVICE=cpu python setup.py bdist_wheel && \
             pip install dist/*.whl

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

@@ -240,7 +240,7 @@ uv pip install vllm==${VLLM_VERSION} \
         # Install dependencies
         pip install --upgrade numba \
             scipy \
-            huggingface-hub[cli,hf_transfer] \
+            huggingface-hub[cli] \
             setuptools_scm
         pip install -r requirements/rocm.txt
 

docs/governance/committers.md

@@ -14,6 +14,7 @@ Sorted alphabetically by GitHub handle:
 - [@aarnphm](https://github.com/aarnphm): Structured output
 - [@alexm-redhat](https://github.com/alexm-redhat): Performance
 - [@ApostaC](https://github.com/ApostaC): Connectors, offloading
+- [@bbrowning](https://github.com/bbrowning): Tool use and reasoning parser
 - [@benchislett](https://github.com/benchislett): Engine core and spec decode
 - [@bigPYJ1151](https://github.com/bigPYJ1151): Intel CPU/XPU integration
 - [@chaunceyjiang](https://github.com/chaunceyjiang): Tool use and reasoning parser
@@ -31,6 +32,7 @@ Sorted alphabetically by GitHub handle:
 - [@LucasWilkinson](https://github.com/LucasWilkinson): Kernels and performance
 - [@luccafong](https://github.com/luccafong): Llama models, speculative decoding, distributed
 - [@markmc](https://github.com/markmc): Observability
+- [@MatthewBonanni](https://github.com/MatthewBonanni): Kernels and performance
 - [@mgoin](https://github.com/mgoin): Quantization and performance
 - [@NickLucche](https://github.com/NickLucche): KV connector
 - [@njhill](https://github.com/njhill): Distributed, API server, engine core
@@ -41,6 +43,7 @@ Sorted alphabetically by GitHub handle:
 - [@robertgshaw2-redhat](https://github.com/robertgshaw2-redhat): Core, distributed, disagg
 - [@ruisearch42](https://github.com/ruisearch42): Pipeline parallelism, Ray Support
 - [@russellb](https://github.com/russellb): Structured output, engine core, security
+- [@sfeng33](https://github.com/sfeng33): Tool use and reasoning parser
 - [@sighingnow](https://github.com/sighingnow): Qwen models, new model support
 - [@simon-mo](https://github.com/simon-mo): Project lead, API entrypoints, community
 - [@tdoublep](https://github.com/tdoublep): State space models
@@ -86,7 +89,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r
 - AsyncLLM: the zmq based protocol hosting engine core and making it accessible for entrypoints
     - @robertgshaw2-redhat, @njhill, @russellb
 - ModelRunner, Executor, Worker: the abstractions for engine wrapping model implementation
-    - @WoosukKwon, @tlrmchlsmth, @heheda12345, @LucasWilkinson, @ProExpertProg
+    - @WoosukKwon, @tlrmchlsmth, @heheda12345, @LucasWilkinson, @ProExpertProg, @MatthewBonanni
 - KV Connector: Connector interface and implementation for KV cache offload and transfer
     - @robertgshaw2-redhat, @njhill, @KuntaiDu, @NickLucche, @ApostaC
 - Distributed, Parallelism, Process Management: Process launchers managing each worker, and assign them to the right DP/TP/PP/EP ranks
@@ -105,7 +108,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r
 - Custom Layers: Utility layers in vLLM such as rotary embedding and rms norms
     - @ProExpertProg
 - Attention: Attention interface for paged attention
-    - @WoosukKwon, @LucasWilkinson, @heheda12345
+    - @WoosukKwon, @LucasWilkinson, @heheda12345, @MatthewBonanni
 - FusedMoE: FusedMoE kernel, Modular kernel framework, EPLB
     - @tlrmchlsmth
 - Quantization: Various quantization config, weight loading, and kernel.
@@ -119,7 +122,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r
 - State space models: The state space models implementation in vLLM
     - @tdoublep, @tlrmchlsmth
 - Reasoning and tool calling parsers
-    - @chaunceyjiang, @aarnphm
+    - @chaunceyjiang, @aarnphm, @sfeng33, @bbrowning
 
 ### Entrypoints
 
@@ -133,7 +136,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r
 ### Features
 
 - Spec Decode: Covers model definition, attention, sampler, and scheduler related to n-grams, EAGLE, and MTP.
-    - @WoosukKwon, @benchislett, @luccafong
+    - @WoosukKwon, @benchislett, @luccafong, @MatthewBonanni
 - Structured Output: The structured output implementation
     - @russellb, @aarnphm
 - RL: The RL related features such as collective rpc, sleep mode, etc.
@@ -153,8 +156,8 @@ If you have PRs touching the area, please feel free to ping the area owner for r
 
 ### External Kernels Integration
 
-- FlashAttention: @LucasWilkinson
-- FlashInfer: @LucasWilkinson, @mgoin, @WoosukKwon
+- FlashAttention: @LucasWilkinson, @MatthewBonanni
+- FlashInfer: @LucasWilkinson, @mgoin, @WoosukKwon, @MatthewBonanni
 - Blackwell Kernels: @mgoin, @yewentao256
 - DeepEP/DeepGEMM: @mgoin, @yewentao256
 

docs/models/pooling_models/README.md

@@ -59,6 +59,16 @@ please refer to [IO Processor Plugins](../../design/io_processor_plugins.md).
     Within classification tasks, there is a specialized subcategory: Cross-encoder (aka reranker) models. These models
 are a subset of classification models that accept two prompts as input and output num_labels equal to 1.
 
+### Pooling Types
+
+| Pooling Tasks  | Granularity   | Description                                                                                                                                                                                       |
+|----------------|---------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
+| `CLS` pooling  | Sequence-wise | For BERT‑like (bidirectional self‑attention) models, CLS pooling is used by default. This means the last_hidden_states corresponding to the first token (the [CLS] token) is taken as the output. |
+| `LAST` pooling | Sequence-wise | For GPT‑like (causal self‑attention) models, LAST pooling is used by default. This means the last_hidden_states corresponding to the last token is taken as the output.                           |
+| `MEAN` pooling | Sequence-wise | Many studies have shown that averaging the last_hidden_states over all input tokens performs better on certain downstream tasks. Therefore, more and more models are using MEAN pooling.          |
+| `ALL` pooling  | Token-wise    | Outputs the last_hidden_states for all input tokens.                                                                                                                                              |
+| `STEP` pooling | Token-wise    | Filters and outputs the last_hidden_states corresponding to the token IDs returned by returned_token_ids.                                                                                         |
+
 ### Score Types
 
 The scoring models is designed to compute similarity scores between two input prompts. It supports three model types

docs/models/pooling_models/embed.md

@@ -45,6 +45,7 @@ You can compute pairwise similarity scores to build a similarity matrix using th
 | `GritLM` | GritLM | `parasail-ai/GritLM-7B-vllm`. | ✅︎ | ✅︎ |
 | `GteModel` | Arctic-Embed-2.0-M | `Snowflake/snowflake-arctic-embed-m-v2.0`. | | |
 | `GteNewModel` | mGTE-TRM (see note) | `Alibaba-NLP/gte-multilingual-base`, etc. | | |
+| `JinaEmbeddingsV5Model`<sup>C</sup> | Qwen3-based with task-specific LoRA adapters | `jinaai/jina-embeddings-v5-text-small` (see note) | ✅︎ | ✅︎ |
 | `LlamaBidirectionalModel`<sup>C</sup> | Llama-based with bidirectional attention | `nvidia/llama-nemotron-embed-1b-v2`, etc. | ✅︎ | ✅︎ |
 | `LlamaModel`<sup>C</sup>, `LlamaForCausalLM`<sup>C</sup>, `MistralModel`<sup>C</sup>, etc. | Llama-based | `intfloat/e5-mistral-7b-instruct`, etc. | ✅︎ | ✅︎ |
 | `ModernBertModel` | ModernBERT-based | `Alibaba-NLP/gte-modernbert-base`, etc. | | |
@@ -73,6 +74,12 @@ You can compute pairwise similarity scores to build a similarity matrix using th
 !!! note
     `jinaai/jina-embeddings-v3` supports multiple tasks through LoRA, while vllm temporarily only supports text-matching tasks by merging LoRA weights.
 
+!!! note
+    `jinaai/jina-embeddings-v5-text-small` ships with four task-specific LoRA adapters
+    (`retrieval`, `text-matching`, `classification`, `clustering`). vLLM merges the
+    selected adapter into the base weights at load time. Choose the task with
+    `--hf-overrides '{"jina_task": "<task>"}'`; the default is `retrieval`.
+
 ### Multimodal Models
 
 !!! note

docs/models/pooling_models/scoring.md

@@ -160,6 +160,8 @@ The following Score API parameters are supported:
 --8<-- "vllm/entrypoints/pooling/base/protocol.py:pooling-common-params"
 --8<-- "vllm/entrypoints/pooling/base/protocol.py:pooling-common-extra-params"
 --8<-- "vllm/entrypoints/pooling/base/protocol.py:classify-extra-params"
+--8<-- "vllm/entrypoints/pooling/scoring/protocol.py:scoring-common-params"
+--8<-- "vllm/entrypoints/pooling/scoring/protocol.py:score-request-params"
 ```
 
 #### Examples
@@ -370,6 +372,8 @@ The following rerank api parameters are supported:
 --8<-- "vllm/entrypoints/pooling/base/protocol.py:pooling-common-params"
 --8<-- "vllm/entrypoints/pooling/base/protocol.py:pooling-common-extra-params"
 --8<-- "vllm/entrypoints/pooling/base/protocol.py:classify-extra-params"
+--8<-- "vllm/entrypoints/pooling/scoring/protocol.py:scoring-common-params"
+--8<-- "vllm/entrypoints/pooling/scoring/protocol.py:rerank-request-params"
 ```
 
 #### Examples

docs/models/pooling_models/token_classify.md

@@ -68,7 +68,7 @@ If your model is not in the above list, we will try to automatically convert the
     Forced alignment usage requires `--hf-overrides '{"architectures": ["Qwen3ASRForcedAlignerForTokenClassification"]}'`.
     Please refer to [examples/pooling/token_classify/forced_alignment_offline.py](../../../examples/pooling/token_classify/forced_alignment_offline.py).
 
-### As Reward Models
+### Reward Models
 
 Using token classification models as reward models. For details on reward models, see [Reward Models](reward.md).
 

docs/models/supported_models.md

@@ -400,6 +400,7 @@ th {
 | `Gemma2ForCausalLM` | Gemma 2 | `google/gemma-2-9b`, `google/gemma-2-27b`, etc. | ✅︎ | ✅︎ |
 | `Gemma3ForCausalLM` | Gemma 3 | `google/gemma-3-1b-it`, etc. | ✅︎ | ✅︎ |
 | `Gemma3nForCausalLM` | Gemma 3n | `google/gemma-3n-E2B-it`, `google/gemma-3n-E4B-it`, etc. | | |
+| `Gemma4ForCausalLM` | Gemma 4 | `google/gemma-4-E2B-it`, etc. | ✅︎ | ✅︎ |
 | `GlmForCausalLM` | GLM-4 | `zai-org/glm-4-9b-chat-hf`, etc. | ✅︎ | ✅︎ |
 | `Glm4ForCausalLM` | GLM-4-0414 | `zai-org/GLM-4-32B-0414`, etc. | ✅︎ | ✅︎ |
 | `Glm4MoeForCausalLM` | GLM-4.5, GLM-4.6, GLM-4.7 | `zai-org/GLM-4.5`, etc. | ✅︎ | ✅︎ |
@@ -554,6 +555,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
 | `FuyuForCausalLM` | Fuyu | T + I | `adept/fuyu-8b`, etc. | | ✅︎ |
 | `Gemma3ForConditionalGeneration` | Gemma 3 | T + I<sup>E+</sup> | `google/gemma-3-4b-it`, `google/gemma-3-27b-it`, etc. | ✅︎ | ✅︎ |
 | `Gemma3nForConditionalGeneration` | Gemma 3n | T + I + A | `google/gemma-3n-E2B-it`, `google/gemma-3n-E4B-it`, etc. | | |
+| `Gemma4ForConditionalGeneration` | Gemma 4 | T + I<sup>+</sup> + V + A<sup>*</sup> | `google/gemma-4-E2B-it`, etc. | | ✅︎ |
 | `GLM4VForCausalLM`<sup>^</sup> | GLM-4V | T + I | `zai-org/glm-4v-9b`, `zai-org/cogagent-9b-20241220`, etc. | ✅︎ | ✅︎ |
 | `Glm4vForConditionalGeneration` | GLM-4.1V-Thinking | T + I<sup>E+</sup> + V<sup>E+</sup> | `zai-org/GLM-4.1V-9B-Thinking`, etc. | ✅︎ | ✅︎ |
 | `Glm4vMoeForConditionalGeneration` | GLM-4.5V | T + I<sup>E+</sup> + V<sup>E+</sup> | `zai-org/GLM-4.5V`, etc. | ✅︎ | ✅︎ |
@@ -633,6 +635,7 @@ Some models are supported only via the [Transformers modeling backend](#transfor
 <sup>^</sup> You need to set the architecture name via `--hf-overrides` to match the one in vLLM.</br>
 <sup>E</sup> Pre-computed embeddings can be inputted for this modality.</br>
 <sup>+</sup> Multiple items can be inputted per text prompt for this modality.
+<sup>*</sup> Only specific variants of the model support this modality (see notes below).</br>
 
 !!! note
     `Gemma3nForConditionalGeneration` is only supported on V1 due to shared KV caching and it depends on `timm>=1.0.17` to make use of its
@@ -643,6 +646,11 @@ Some models are supported only via the [Transformers modeling backend](#transfor
     - Both audio and vision MM encoders use `transformers.AutoModel` implementation.  
     - There's no PLE caching or out-of-memory swapping support, as described in [Google's blog](https://developers.googleblog.com/en/introducing-gemma-3n/). These features might be too model-specific for vLLM, and swapping in particular may be better suited for constrained setups.
 
+!!! note
+    For `Gemma4ForConditionalGeneration`:
+    - audio input is only supported by the `gemma-4-E2B` and `gemma-4-E4B` variants.
+    - The model does not ingest videos directly. However, vLLM’s Gemma 4 implementation supports video inputs by handling video processing internally. Users can send videos directly in the message structure to vLLM, where they are converted into text and image frames before being passed to the model.
+
 !!! note
     For `InternVLChatModel`, only InternVL2.5 with Qwen2.5 text backbone (`OpenGVLab/InternVL2.5-1B` etc.), InternVL3 and InternVL3.5 have video inputs support currently.
 

docs/serving/openai_compatible_server.md

@@ -467,28 +467,11 @@ It consists of two endpoints:
 - `/tokenize` corresponds to calling `tokenizer.encode()`.
 - `/detokenize` corresponds to calling `tokenizer.decode()`.
 
-### Score API
-
-#### Score Template
-
-Some scoring models require a specific prompt format to work correctly. You can specify a custom score template using the `--chat-template` parameter (see [Chat Template](#chat-template)).
-
-Score templates are supported for **cross-encoder** models only. If you are using an **embedding** model for scoring, vLLM does not apply a score template.
-
-Like chat templates, the score template receives a `messages` list. For scoring, each message has a `role` attribute—either `"query"` or `"document"`. For the usual kind of point-wise cross-encoder, you can expect exactly two messages: one query and one document. To access the query and document content, use Jinja's `selectattr` filter:
-
-- **Query**: `{{ (messages | selectattr("role", "eq", "query") | first).content }}`
-- **Document**: `{{ (messages | selectattr("role", "eq", "document") | first).content }}`
-
-This approach is more robust than index-based access (`messages[0]`, `messages[1]`) because it selects messages by their semantic role. It also avoids assumptions about message ordering if additional message types are added to `messages` in the future.
-
-Example template file: [examples/pooling/score/template/nemotron-rerank.jinja](../../examples/pooling/score/template/nemotron-rerank.jinja)
-
 ### Generative Scoring API
 
 The `/generative_scoring` endpoint uses a CausalLM model (e.g., Llama, Qwen, Mistral) to compute the probability of specified token IDs appearing as the next token. Each item (document) is concatenated with the query to form a prompt, and the model predicts how likely each label token is as the next token after that prompt. This lets you score items against a query — for example, asking "Is this the capital of France?" and scoring each city by how likely the model is to answer "Yes".
 
-This endpoint is automatically available when the server is started with a generative model (task `"generate"`). It is separate from the pooling-based [Score API](#score-api), which uses cross-encoder, bi-encoder, or late-interaction models.
+This endpoint is automatically available when the server is started with a generative model (task `"generate"`). It is separate from the pooling-based [Score API](../models/pooling_models/scoring.md#score-api), which uses cross-encoder, bi-encoder, or late-interaction models.
 
 **Requirements:**
 

examples/observability/dashboards/README.md

@@ -74,8 +74,8 @@ percli apply -f perses/performance_statistics.yaml
 
 For detailed deployment instructions and platform-specific options, see:
 
-- **[Grafana Documentation](./grafana)** - JSON dashboards, operator usage, manual import
-- **[Perses Documentation](./perses)** - YAML specs, CLI usage, operator wrapping
+- **[Grafana Documentation](grafana)** - JSON dashboards, operator usage, manual import
+- **[Perses Documentation](perses)** - YAML specs, CLI usage, operator wrapping
 
 ## Contributing
 

pyproject.toml

@@ -170,6 +170,9 @@ eles = "eles"
 datas = "datas"
 ser = "ser"
 ure = "ure"
+# Walsh-Hadamard Transform
+wht = "wht"
+WHT = "WHT"
 
 [tool.uv]
 no-build-isolation-package = ["torch"]

requirements/common.txt

@@ -7,7 +7,7 @@ requests >= 2.26.0
 tqdm
 blake3
 py-cpuinfo
-transformers >= 4.56.0, < 5
+transformers >= 4.56.0, != 5.0.*, != 5.1.*, != 5.2.*, != 5.3.*, != 5.4.*, != 5.5.0
 tokenizers >= 0.21.1  # Required for fast incremental detokenization.
 protobuf >= 5.29.6, !=6.30.*, !=6.31.*, !=6.32.*, !=6.33.0.*, !=6.33.1.*, !=6.33.2.*, !=6.33.3.*, !=6.33.4.* # Required by LlamaTokenizer, gRPC. CVE-2026-0994
 fastapi[standard] >= 0.115.0 # Required by FastAPI's form models in the OpenAI API server's audio transcriptions endpoint.
@@ -19,7 +19,7 @@ pillow  # Required for image processing
 prometheus-fastapi-instrumentator >= 7.0.0
 tiktoken >= 0.6.0  # Required for DBRX tokenizer
 lm-format-enforcer == 0.11.3
-llguidance >= 1.3.0, < 1.4.0; platform_machine == "x86_64" or platform_machine == "arm64" or platform_machine == "aarch64" or platform_machine == "s390x" or platform_machine == "ppc64le"
+llguidance >= 1.3.0, < 1.4.0; platform_machine == "x86_64" or platform_machine == "arm64" or platform_machine == "aarch64" or platform_machine == "ppc64le"
 outlines_core == 0.2.11
 # required for outlines backend disk cache
 diskcache == 5.6.3
@@ -32,12 +32,14 @@ pyzmq >= 25.0.0
 msgspec
 gguf >= 0.17.0
 mistral_common[image] >= 1.11.0
+av  # required for audio in video IO
 opencv-python-headless >= 4.13.0    # required for video IO
+soundfile  # required for audio IO
 pyyaml
 six>=1.16.0; python_version > '3.11' # transitive dependency of pandas that needs to be the latest version for python 3.12
 setuptools>=77.0.3,<81.0.0; python_version > '3.11' # Setuptools is used by triton, we need to ensure a modern version is installed for 3.12+ so that it does not try to import distutils, which was removed in 3.12
 einops # Required for Qwen2-VL.
-compressed-tensors == 0.14.0.1 # required for compressed-tensors
+compressed-tensors == 0.15.0.1 # required for compressed-tensors
 depyf==0.20.0 # required for profiling and debugging with compilation config
 cloudpickle # allows pickling lambda functions in model_executor/models/registry.py
 watchfiles # required for http server to monitor the updates of TLS files

requirements/kv_connectors.txt

@@ -1,3 +1,5 @@
 lmcache >= 0.3.9
-nixl[cu13] >= 0.7.1, < 0.10.0 # Required for disaggregated prefill
+nixl[cu13] >= 0.7.1, <= 0.10.1 # Required for disaggregated prefill
+nixl-cu12 >= 0.7.1, <= 0.10.1
+nixl-cu13 >= 0.7.1, <= 0.10.1
 mooncake-transfer-engine >= 0.3.8

requirements/test/cuda.in

@@ -18,10 +18,9 @@ httpx
 librosa # required for audio tests
 vector_quantize_pytorch # required for minicpmo_26 test
 vocos # required for minicpmo_26 test
-peft>=0.15.0 # required for phi-4-mm test
+peft>=0.18.1 # required for phi-4-mm test
 pqdm
 ray[cgraph,default]>=2.48.0 # Ray Compiled Graph, required by pipeline parallelism tests
-resampy # required for audio tests
 sentence-transformers>=5.2.0 # required for embedding tests
 soundfile # required for audio tests
 jiwer # required for audio tests
@@ -39,8 +38,8 @@ opencv-python-headless >= 4.13.0 # required for video test
 datamodel_code_generator # required for minicpm3 test
 lm-eval[api]>=0.4.11 # required for model evaluation test
 mteb[bm25s]>=2, <3 # required for mteb test
-transformers==4.57.5
-tokenizers==0.22.0
+transformers==5.5.3
+tokenizers==0.22.2
 schemathesis>=3.39.15 # Required for openai schema test.
 # quantization
 bitsandbytes==0.49.2

requirements/test/cuda.txt

@@ -4,7 +4,7 @@ absl-py==2.1.0
     # via
     #   rouge-score
     #   tensorboard
-accelerate==1.0.1
+accelerate==1.13.0
     # via peft
 aenum==3.1.16
     # via lightly
@@ -248,7 +248,6 @@ filelock==3.16.1
     #   huggingface-hub
     #   ray
     #   torch
-    #   transformers
     #   virtualenv
 fiona==1.10.1
     # via torchgeo
@@ -331,7 +330,7 @@ h5py==3.13.0
     # via terratorch
 harfile==0.3.0
     # via schemathesis
-hf-xet==1.1.7
+hf-xet==1.4.3
     # via huggingface-hub
 hiredis==3.0.0
     # via tensorizer
@@ -345,9 +344,10 @@ httpx==0.27.2
     # via
     #   -r requirements/test/cuda.in
     #   diffusers
+    #   huggingface-hub
     #   perceptron
     #   schemathesis
-huggingface-hub==0.36.2
+huggingface-hub==1.10.2
     # via
     #   accelerate
     #   datasets
@@ -555,7 +555,6 @@ numba==0.61.2
     #   -c requirements/cuda.txt
     #   -r requirements/test/cuda.in
     #   librosa
-    #   resampy
 numpy==2.2.6
     # via
     #   -r requirements/test/cuda.in
@@ -596,7 +595,6 @@ numpy==2.2.6
     #   pyogrio
     #   pywavelets
     #   rasterio
-    #   resampy
     #   rioxarray
     #   rouge-score
     #   runai-model-streamer
@@ -756,7 +754,7 @@ pathvalidate==3.2.1
     # via pytablewriter
 patsy==1.0.1
     # via statsmodels
-peft==0.16.0
+peft==0.18.1
     # via -r requirements/test/cuda.in
 perceptron==0.1.4
     # via -r requirements/test/cuda.in
@@ -982,7 +980,7 @@ referencing==0.35.1
     # via
     #   jsonschema
     #   jsonschema-specifications
-regex==2024.9.11
+regex==2026.2.28
     # via
     #   diffusers
     #   nltk
@@ -1002,7 +1000,6 @@ requests==2.32.3
     #   google-api-core
     #   google-cloud-storage
     #   gpt-oss
-    #   huggingface-hub
     #   lightly
     #   lm-eval
     #   mistral-common
@@ -1015,10 +1012,7 @@ requests==2.32.3
     #   starlette-testclient
     #   tacoreader
     #   tiktoken
-    #   transformers
     #   wandb
-resampy==0.4.3
-    # via -r requirements/test/cuda.in
 responses==0.25.3
     # via genai-perf
 rfc3339-validator==0.1.4
@@ -1216,7 +1210,7 @@ timm==1.0.17
     #   segmentation-models-pytorch
     #   terratorch
     #   torchgeo
-tokenizers==0.22.0
+tokenizers==0.22.2
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/cuda.in
@@ -1295,7 +1289,7 @@ tqdm==4.67.3
     #   tacoreader
     #   terratorch
     #   transformers
-transformers==4.57.5
+transformers==5.5.3
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/cuda.in
@@ -1317,7 +1311,9 @@ typepy==1.3.2
 typer==0.15.2
     # via
     #   fastsafetensors
+    #   huggingface-hub
     #   perceptron
+    #   transformers
 types-python-dateutil==2.9.0.20241206
     # via arrow
 typeshed-client==2.8.2

requirements/test/nightly-torch.txt

@@ -29,8 +29,8 @@ opencv-python-headless >= 4.13.0 # required for video test
 datamodel_code_generator # required for minicpm3 test
 lm-eval[api]>=0.4.11 # required for model evaluation test
 mteb[bm25s]>=2, <3 # required for mteb test
-transformers==4.57.5
-tokenizers==0.22.0
+transformers==5.5.3
+tokenizers==0.22.2
 schemathesis>=3.39.15 # Required for openai schema test.
 # quantization
 bitsandbytes>=0.49.2

requirements/test/rocm.in

@@ -23,7 +23,6 @@ vocos # required for minicpmo_26 test
 peft>=0.15.0 # required for phi-4-mm test
 pqdm
 ray[cgraph,default]>=2.48.0 # Ray Compiled Graph, required by pipeline parallelism tests
-resampy # required for audio tests
 sentence-transformers>=5.2.0 # required for embedding tests
 soundfile # required for audio tests
 jiwer # required for audio tests
@@ -38,8 +37,8 @@ opencv-python-headless>=4.13.0 # required for video test
 datamodel_code_generator # required for minicpm3 test
 lm-eval[api]>=0.4.11 # required for model evaluation test
 mteb[bm25s]>=2, <3 # required for mteb test
-transformers==4.57.5
-tokenizers==0.22.0
+transformers==5.5.3
+tokenizers==0.22.2
 schemathesis>=3.39.15 # Required for openai schema test
 # quantization
 bitsandbytes==0.49.2
@@ -82,4 +81,3 @@ plotly # required for perf comparison html report
 rapidfuzz
 torchgeo==0.7.0
 multiprocess==0.70.16
-huggingface-hub==0.36.2

requirements/test/rocm.txt

@@ -39,7 +39,7 @@ annotated-doc==0.0.4
     #   typer
 annotated-types==0.7.0
     # via pydantic
-anthropic==0.89.0
+anthropic==0.93.0
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
@@ -76,7 +76,9 @@ attrs==26.1.0
 audioread==3.0.1
     # via librosa
 av==16.1.0
-    # via -r requirements/test/rocm.in
+    # via
+    #   -r requirements/test/../common.txt
+    #   -r requirements/test/rocm.in
 azure-core==1.39.0
     # via
     #   azure-identity
@@ -172,7 +174,7 @@ colorful==0.5.8
     # via ray
 colorlog==6.10.1
     # via optuna
-compressed-tensors==0.14.0.1
+compressed-tensors==0.15.0.1
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
@@ -269,9 +271,9 @@ fastapi==0.135.2
     #   model-hosting-container-standards
 fastapi-cli==0.0.24
     # via fastapi
-fastapi-cloud-cli==0.15.1
+fastapi-cloud-cli==0.16.1
     # via fastapi-cli
-fastar==0.9.0
+fastar==0.10.0
     # via fastapi-cloud-cli
 fastparquet==2026.3.0
     # via genai-perf
@@ -290,7 +292,6 @@ filelock==3.25.2
     #   python-discovery
     #   ray
     #   torch
-    #   transformers
     #   virtualenv
 fiona==1.10.1
     # via torchgeo
@@ -384,7 +385,7 @@ h5py==3.16.0
     # via terratorch
 harfile==0.4.0
     # via schemathesis
-hf-xet==1.4.2
+hf-xet==1.4.3
     # via huggingface-hub
 hiredis==3.3.1
     # via tensorizer
@@ -403,6 +404,7 @@ httpx==0.27.2
     #   diffusers
     #   fastapi
     #   fastapi-cloud-cli
+    #   huggingface-hub
     #   mcp
     #   model-hosting-container-standards
     #   openai
@@ -410,9 +412,8 @@ httpx==0.27.2
     #   schemathesis
 httpx-sse==0.4.3
     # via mcp
-huggingface-hub==0.36.2
+huggingface-hub==1.10.2
     # via
-    #   -r requirements/test/rocm.in
     #   accelerate
     #   datasets
     #   diffusers
@@ -484,7 +485,7 @@ jinja2==3.1.6
     #   genai-perf
     #   lm-eval
     #   torch
-jiter==0.13.0
+jiter==0.14.0
     # via
     #   anthropic
     #   openai
@@ -631,7 +632,7 @@ msgpack==1.1.2
     # via
     #   librosa
     #   ray
-msgspec==0.20.0
+msgspec==0.21.0
     # via -r requirements/test/../common.txt
 mteb==2.11.5
     # via -r requirements/test/rocm.in
@@ -663,7 +664,6 @@ numba==0.61.2
     #   -c requirements/rocm.txt
     #   -r requirements/test/rocm.in
     #   librosa
-    #   resampy
 numkong==7.1.1
     # via albucore
 numpy==2.2.6
@@ -709,7 +709,6 @@ numpy==2.2.6
     #   pytrec-eval-terrier
     #   pywavelets
     #   rasterio
-    #   resampy
     #   rioxarray
     #   rouge-score
     #   runai-model-streamer
@@ -742,7 +741,7 @@ omegaconf==2.3.0
     #   lightning
 open-clip-torch==2.32.0
     # via -r requirements/test/rocm.in
-openai==2.30.0
+openai==2.31.0
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
@@ -1093,7 +1092,7 @@ python-dotenv==1.2.2
     #   uvicorn
 python-json-logger==4.1.0
     # via -r requirements/test/../common.txt
-python-multipart==0.0.22
+python-multipart==0.0.26
     # via
     #   fastapi
     #   mcp
@@ -1180,7 +1179,6 @@ requests==2.32.5
     #   google-api-core
     #   google-cloud-storage
     #   gpt-oss
-    #   huggingface-hub
     #   lightly
     #   lm-eval
     #   mistral-common
@@ -1194,10 +1192,7 @@ requests==2.32.5
     #   starlette-testclient
     #   tacoreader
     #   tiktoken
-    #   transformers
     #   wandb
-resampy==0.4.3
-    # via -r requirements/test/rocm.in
 responses==0.26.0
     # via genai-perf
 rfc3339-validator==0.1.4
@@ -1338,6 +1333,7 @@ sortedcontainers==2.4.0
     # via hypothesis
 soundfile==0.13.1
     # via
+    #   -r requirements/test/../common.txt
     #   -r requirements/test/rocm.in
     #   genai-perf
     #   librosa
@@ -1428,7 +1424,7 @@ timm==1.0.17
     #   segmentation-models-pytorch
     #   terratorch
     #   torchgeo
-tokenizers==0.22.0
+tokenizers==0.22.2
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
@@ -1471,7 +1467,7 @@ tqdm==4.67.3
     #   tacoreader
     #   terratorch
     #   transformers
-transformers==4.57.5
+transformers==5.5.3
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
@@ -1498,7 +1494,9 @@ typer==0.24.1
     #   fastapi-cli
     #   fastapi-cloud-cli
     #   fastsafetensors
+    #   huggingface-hub
     #   perceptron
+    #   transformers
 typeshed-client==2.9.0
     # via jsonargparse
 typing-extensions==4.15.0

requirements/test/xpu.in

@@ -13,7 +13,6 @@ pytest-shard
 absl-py
 accelerate
 arctic-inference
-hf_transfer
 lm_eval[api]
 modelscope
 

requirements/test/xpu.txt

@@ -19,7 +19,9 @@ aiosignal==1.4.0
 albumentations==1.4.6
     # via -r requirements/test/xpu.in
 annotated-doc==0.0.4
-    # via fastapi
+    # via
+    #   fastapi
+    #   typer
 annotated-types==0.7.0
     # via pydantic
 anyio==4.13.0
@@ -64,6 +66,7 @@ click==8.3.1
     #   jiwer
     #   nltk
     #   schemathesis
+    #   typer
     #   uvicorn
 colorama==0.4.6
     # via sacrebleu
@@ -112,7 +115,6 @@ filelock==3.25.2
     #   huggingface-hub
     #   modelscope
     #   torch
-    #   transformers
 frozenlist==1.8.0
     # via
     #   aiohttp
@@ -133,9 +135,7 @@ h11==0.16.0
     #   uvicorn
 harfile==0.4.0
     # via schemathesis
-hf-transfer==0.1.9
-    # via -r requirements/test/xpu.in
-hf-xet==1.4.2
+hf-xet==1.4.3
     # via huggingface-hub
 html2text==2025.4.15
     # via gpt-oss
@@ -144,8 +144,9 @@ httpcore==1.0.9
 httpx==0.28.1
     # via
     #   datasets
+    #   huggingface-hub
     #   schemathesis
-huggingface-hub==0.36.2
+huggingface-hub==1.10.2
     # via
     #   accelerate
     #   datasets
@@ -515,7 +516,6 @@ requests==2.33.1
     #   docker
     #   evaluate
     #   gpt-oss
-    #   huggingface-hub
     #   lm-eval
     #   mistral-common
     #   modelscope
@@ -524,11 +524,11 @@ requests==2.33.1
     #   schemathesis
     #   starlette-testclient
     #   tiktoken
-    #   transformers
 rich==14.3.3
     # via
     #   mteb
     #   schemathesis
+    #   typer
 rouge-score==0.1.2
     # via lm-eval
 rpds-py==0.30.0
@@ -572,6 +572,8 @@ setuptools==80.10.2
     #   modelscope
     #   pytablewriter
     #   torch
+shellingham==1.5.4
+    # via typer
 six==1.17.0
     # via
     #   -c requirements/common.txt
@@ -665,7 +667,7 @@ tqdm==4.67.3
     #   pqdm
     #   sentence-transformers
     #   transformers
-transformers==4.57.6
+transformers==5.5.3
     # via
     #   -c requirements/common.txt
     #   sentence-transformers
@@ -676,6 +678,10 @@ typepy==1.3.4
     #   dataproperty
     #   pytablewriter
     #   tabledata
+typer==0.24.1
+    # via
+    #   huggingface-hub
+    #   transformers
 typing-extensions==4.15.0
     # via
     #   -c requirements/common.txt

setup.py

@@ -693,6 +693,12 @@ class precompiled_wheel_utils:
                 flash_attn_regex = re.compile(
                     r"vllm/vllm_flash_attn/(?:[^/.][^/]*/)*(?!\.)[^/]*\.py"
                 )
+                # __init__.py and flash_attn_interface.py are source-controlled
+                # in vllm and should not be overwritten (matches cmake exclusions)
+                flash_attn_files_to_skip = {
+                    "vllm/vllm_flash_attn/__init__.py",
+                    "vllm/vllm_flash_attn/flash_attn_interface.py",
+                }
                 triton_kernels_regex = re.compile(
                     r"vllm/third_party/triton_kernels/(?:[^/.][^/]*/)*(?!\.)[^/]*\.py"
                 )
@@ -705,7 +711,11 @@ class precompiled_wheel_utils:
                     filter(lambda x: x.filename in files_to_copy, wheel.filelist)
                 )
                 file_members += list(
-                    filter(lambda x: flash_attn_regex.match(x.filename), wheel.filelist)
+                    filter(
+                        lambda x: flash_attn_regex.match(x.filename)
+                        and x.filename not in flash_attn_files_to_skip,
+                        wheel.filelist,
+                    )
                 )
                 file_members += list(
                     filter(
@@ -1082,10 +1092,7 @@ setup(
         "instanttensor": ["instanttensor >= 0.1.5"],
         "runai": ["runai-model-streamer[s3,gcs,azure] >= 0.15.7"],
         "audio": [
-            "av",
-            "resampy",
             "scipy",
-            "soundfile",
             "mistral_common[audio]",
         ],  # Required for audio processing
         "video": [],  # Kept for backwards compatibility
@@ -1094,7 +1101,7 @@ setup(
         # NOTE: When updating helion version, also update CI files:
         #   - .buildkite/test_areas/kernels.yaml
         #   - .buildkite/test-amd.yaml
-        "helion": ["helion==0.3.3"],
+        "helion": ["helion==1.0.0"],
         # Optional deps for gRPC server (vllm serve --grpc)
         "grpc": ["smg-grpc-servicer[vllm] >= 0.5.0"],
         # Optional deps for OpenTelemetry tracing

tests/basic_correctness/test_cumem.py

@@ -13,6 +13,8 @@ from vllm.utils.mem_constants import GiB_bytes
 
 from ..utils import create_new_process_for_each_test, requires_fp8
 
+DEVICE_TYPE = current_platform.device_type
+
 
 @create_new_process_for_each_test("fork" if not current_platform.is_rocm() else "spawn")
 def test_python_error():
@@ -26,13 +28,13 @@ def test_python_error():
     tensors = []
     with allocator.use_memory_pool():
         # allocate 70% of the total memory
-        x = torch.empty(alloc_bytes, dtype=torch.uint8, device="cuda")
+        x = torch.empty(alloc_bytes, dtype=torch.uint8, device=DEVICE_TYPE)
         tensors.append(x)
     # release the memory
     allocator.sleep()
 
     # allocate more memory than the total memory
-    y = torch.empty(alloc_bytes, dtype=torch.uint8, device="cuda")
+    y = torch.empty(alloc_bytes, dtype=torch.uint8, device=DEVICE_TYPE)
     tensors.append(y)
     with pytest.raises(RuntimeError):
         # when the allocator is woken up, it should raise an error
@@ -44,17 +46,17 @@ def test_python_error():
 def test_basic_cumem():
     # some tensors from default memory pool
     shape = (1024, 1024)
-    x = torch.empty(shape, device="cuda")
+    x = torch.empty(shape, device=DEVICE_TYPE)
     x.zero_()
 
     # some tensors from custom memory pool
     allocator = CuMemAllocator.get_instance()
     with allocator.use_memory_pool():
         # custom memory pool
-        y = torch.empty(shape, device="cuda")
+        y = torch.empty(shape, device=DEVICE_TYPE)
         y.zero_()
         y += 1
-        z = torch.empty(shape, device="cuda")
+        z = torch.empty(shape, device=DEVICE_TYPE)
         z.zero_()
         z += 2
 
@@ -77,16 +79,16 @@ def test_basic_cumem():
 def test_cumem_with_cudagraph():
     allocator = CuMemAllocator.get_instance()
     with allocator.use_memory_pool():
-        weight = torch.eye(1024, device="cuda")
+        weight = torch.eye(1024, device=DEVICE_TYPE)
     with allocator.use_memory_pool(tag="discard"):
-        cache = torch.empty(1024, 1024, device="cuda")
+        cache = torch.empty(1024, 1024, device=DEVICE_TYPE)
 
     def model(x):
         out = x @ weight
         cache[: out.size(0)].copy_(out)
         return out + 1
 
-    x = torch.empty(128, 1024, device="cuda")
+    x = torch.empty(128, 1024, device=DEVICE_TYPE)
 
     # warmup
     model(x)

tests/benchmarks/test_sampling_params.py

@@ -0,0 +1,258 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import numpy as np
+import pytest
+
+from vllm.benchmarks.datasets.utils import get_sampling_params
+from vllm.tokenizers import TokenizerLike
+
+
+class _FakeTokenizer(TokenizerLike):
+    """Minimal tokenizer implementing the TokenizerLike protocol
+    for testing get_sampling_params."""
+
+    def __init__(self, vocab_size: int = 1000, num_special_tokens: int = 0) -> None:
+        self._vocab_size = vocab_size
+        self._num_special_tokens = num_special_tokens
+
+    # -- Properties required by TokenizerLike --
+
+    @classmethod
+    def from_pretrained(cls, path_or_repo_id, *a, **kw):  # type: ignore[override]
+        return cls()
+
+    @property
+    def vocab_size(self) -> int:
+        return self._vocab_size
+
+    @property
+    def all_special_tokens(self) -> list[str]:
+        return []
+
+    @property
+    def all_special_ids(self) -> list[int]:
+        return []
+
+    @property
+    def bos_token_id(self) -> int:
+        return 0
+
+    @property
+    def eos_token_id(self) -> int:
+        return 1
+
+    @property
+    def pad_token_id(self) -> int:
+        return 2
+
+    @property
+    def is_fast(self) -> bool:
+        return False
+
+    @property
+    def max_token_id(self) -> int:
+        return self._vocab_size - 1
+
+    @property
+    def max_chars_per_token(self) -> int:
+        return 4
+
+    @property
+    def truncation_side(self) -> str:
+        return "right"
+
+    def num_special_tokens_to_add(self) -> int:
+        return self._num_special_tokens
+
+    def __call__(self, text, text_pair=None, **kw):  # type: ignore[override]
+        raise NotImplementedError
+
+    def get_vocab(self) -> dict[str, int]:
+        return {}
+
+    def get_added_vocab(self) -> dict[str, int]:
+        return {}
+
+    def encode(self, text, **kw) -> list[int]:  # type: ignore[override]
+        raise NotImplementedError
+
+    def apply_chat_template(self, messages, **kw):  # type: ignore[override]
+        raise NotImplementedError
+
+    def convert_tokens_to_ids(self, tokens):  # type: ignore[override]
+        raise NotImplementedError
+
+    def convert_tokens_to_string(self, tokens: list[str]) -> str:
+        raise NotImplementedError
+
+    def decode(self, ids, skip_special_tokens: bool = False) -> str:  # type: ignore[override]
+        raise NotImplementedError
+
+    def convert_ids_to_tokens(  # type: ignore[override]
+        self, ids, skip_special_tokens: bool = False
+    ) -> list[str]:
+        raise NotImplementedError
+
+
+class TestGetSamplingParams:
+    """Tests for ``get_sampling_params`` in ``vllm.benchmarks.datasets.shared``."""
+
+    # -- helpers --
+
+    @staticmethod
+    def _tok(vocab_size: int = 1000, num_special: int = 0) -> _FakeTokenizer:
+        return _FakeTokenizer(vocab_size=vocab_size, num_special_tokens=num_special)
+
+    # -- return shape / dtype --
+
+    def test_returns_three_arrays(self):
+        rng = np.random.default_rng(0)
+        result = get_sampling_params(rng, 5, 0.0, 100, 50, self._tok())
+        assert len(result) == 3
+        for arr in result:
+            assert isinstance(arr, np.ndarray)
+
+    @pytest.mark.parametrize("n", [1, 10, 100])
+    def test_output_length_matches_num_requests(self, n: int):
+        rng = np.random.default_rng(42)
+        input_lens, output_lens, offsets = get_sampling_params(
+            rng, n, 0.0, 64, 32, self._tok()
+        )
+        assert input_lens.shape == (n,)
+        assert output_lens.shape == (n,)
+        assert offsets.shape == (n,)
+
+    # -- fixed lengths (range_ratio = 0) --
+
+    def test_zero_range_ratio_gives_constant_lengths(self):
+        rng = np.random.default_rng(7)
+        input_lens, output_lens, _ = get_sampling_params(
+            rng, 20, 0.0, 128, 64, self._tok()
+        )
+        assert np.all(input_lens == 128)
+        assert np.all(output_lens == 64)
+
+    def test_special_tokens_subtracted_from_input_only(self):
+        rng = np.random.default_rng(7)
+        input_lens, output_lens, _ = get_sampling_params(
+            rng, 10, 0.0, 100, 50, self._tok(num_special=4)
+        )
+        # real_input_len = 100 - 4 = 96, range_ratio 0 → all 96
+        assert np.all(input_lens == 96)
+        # special tokens are not subtracted from output length
+        assert np.all(output_lens == 50)
+
+    # -- range ratios --
+
+    def test_input_range_bounds(self):
+        rng = np.random.default_rng(0)
+        ratio = 0.5
+        base = 200
+        input_lens, _, _ = get_sampling_params(
+            rng, 500, {"input": ratio, "output": 0.0}, base, 50, self._tok()
+        )
+        lo = int(np.floor(base * (1 - ratio)))
+        hi = int(np.ceil(base * (1 + ratio)))
+        assert np.all(input_lens >= lo)
+        assert np.all(input_lens <= hi)
+
+    def test_output_range_bounds(self):
+        rng = np.random.default_rng(0)
+        ratio = 0.3
+        base = 100
+        _, output_lens, _ = get_sampling_params(
+            rng, 500, {"input": 0.0, "output": ratio}, 50, base, self._tok()
+        )
+        lo = max(1, int(np.floor(base * (1 - ratio))))
+        hi = int(np.ceil(base * (1 + ratio)))
+        assert np.all(output_lens >= lo)
+        assert np.all(output_lens <= hi)
+
+    def test_output_low_clamped_to_one(self):
+        """Even with a high ratio that would push output_low to 0,
+        the function clamps it to 1."""
+        rng = np.random.default_rng(0)
+        # output_len=1, ratio=0.99 → floor(1*0.01)=0, should clamp to 1
+        _, output_lens, _ = get_sampling_params(
+            rng, 50, {"input": 0.0, "output": 0.99}, 100, 1, self._tok()
+        )
+        assert np.all(output_lens >= 1)
+
+    # -- offsets bounded by vocab_size --
+
+    @pytest.mark.parametrize("vocab", [100, 32000, 128256])
+    def test_offsets_within_vocab(self, vocab: int):
+        rng = np.random.default_rng(0)
+        _, _, offsets = get_sampling_params(
+            rng, 200, 0.0, 64, 32, self._tok(vocab_size=vocab)
+        )
+        assert np.all(offsets >= 0)
+        assert np.all(offsets < vocab)
+
+    # -- reproducibility --
+
+    def test_same_seed_same_results(self):
+        tok = self._tok()
+        rr = {"input": 0.3, "output": 0.2}
+        a = get_sampling_params(np.random.default_rng(42), 50, rr, 256, 64, tok)
+        b = get_sampling_params(np.random.default_rng(42), 50, rr, 256, 64, tok)
+        for arr_a, arr_b in zip(a, b):
+            np.testing.assert_array_equal(arr_a, arr_b)
+
+    def test_different_seed_different_results(self):
+        tok = self._tok()
+        rr = {"input": 0.3, "output": 0.2}
+        a = get_sampling_params(np.random.default_rng(0), 50, rr, 256, 64, tok)
+        b = get_sampling_params(np.random.default_rng(1), 50, rr, 256, 64, tok)
+        # Extremely unlikely all three arrays match with different seeds
+        assert not all(np.array_equal(arr_a, arr_b) for arr_a, arr_b in zip(a, b))
+
+    # -- validation / error paths --
+
+    @pytest.mark.parametrize("bad_ratio", [-0.1, 1.0, 1.5])
+    def test_invalid_input_range_ratio(self, bad_ratio: float):
+        rng = np.random.default_rng(0)
+        with pytest.raises(ValueError, match="input_range_ratio"):
+            get_sampling_params(
+                rng, 10, {"input": bad_ratio, "output": 0.0}, 100, 50, self._tok()
+            )
+
+    @pytest.mark.parametrize("bad_ratio", [-0.1, 1.0, 1.5])
+    def test_invalid_output_range_ratio(self, bad_ratio: float):
+        rng = np.random.default_rng(0)
+        with pytest.raises(ValueError, match="output_range_ratio"):
+            get_sampling_params(
+                rng, 10, {"input": 0.0, "output": bad_ratio}, 100, 50, self._tok()
+            )
+
+    def test_invalid_dict_missing_keys(self):
+        rng = np.random.default_rng(0)
+        with pytest.raises(ValueError, match="input.*output"):
+            get_sampling_params(rng, 10, {"input": 0.1}, 100, 50, self._tok())
+
+    def test_input_len_zero_with_special_tokens(self):
+        """input_len < num_special_tokens → real_input_len = 0, which is fine
+        (range [0, 0])."""
+        rng = np.random.default_rng(0)
+        input_lens, _, _ = get_sampling_params(
+            rng, 5, 0.0, 5, 50, self._tok(num_special=10)
+        )
+        # real_input_len = max(0, 5 - 10) = 0
+        assert np.all(input_lens == 0)
+
+    # -- edge cases --
+
+    def test_single_request(self):
+        rng = np.random.default_rng(0)
+        i, o, off = get_sampling_params(rng, 1, 0.0, 100, 50, self._tok())
+        assert i.shape == (1,)
+        assert o.shape == (1,)
+        assert off.shape == (1,)
+
+    def test_large_num_requests(self):
+        rng = np.random.default_rng(0)
+        i, o, off = get_sampling_params(rng, 10_000, 0.5, 512, 128, self._tok())
+        assert i.shape == (10_000,)
+        assert o.shape == (10_000,)
+        assert off.shape == (10_000,)

tests/benchmarks/test_txt_slices_dataset.py

@@ -0,0 +1,68 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import json
+from pathlib import Path
+
+import pytest
+from transformers import AutoTokenizer, PreTrainedTokenizerBase
+
+from vllm.benchmarks.datasets import CustomDataset
+from vllm.benchmarks.datasets.create_txt_slices_dataset import create_txt_slices_jsonl
+
+
+@pytest.fixture(scope="session")
+def hf_tokenizer() -> PreTrainedTokenizerBase:
+    # Use a small, commonly available tokenizer
+    return AutoTokenizer.from_pretrained("gpt2")
+
+
+text_content = """
+Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor
+incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud
+exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat.
+Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat
+nulla pariatur. Excepteur sint occaecat cupidatat non proident,
+sunt in culpa qui officia deserunt mollit anim id est laborum.
+"""
+
+
+@pytest.mark.benchmark
+def test_create_txt_slices_jsonl(
+    hf_tokenizer: PreTrainedTokenizerBase, tmp_path: Path
+) -> None:
+    """Test that create_txt_slices_jsonl produces valid JSONL for CustomDataset."""
+    txt_path = tmp_path / "input.txt"
+    jsonl_path = tmp_path / "input.txt.jsonl"
+
+    txt_path.write_text(text_content)
+
+    create_txt_slices_jsonl(
+        input_path=str(txt_path),
+        output_path=str(jsonl_path),
+        tokenizer_name="gpt2",
+        num_prompts=10,
+        input_len=10,
+        output_len=10,
+    )
+
+    # Verify the JSONL file is valid and has the expected structure
+    records = [json.loads(line) for line in jsonl_path.read_text().splitlines()]
+
+    assert len(records) == 10
+    for record in records:
+        assert "prompt" in record
+        assert "output_tokens" in record
+        assert isinstance(record["prompt"], str)
+        assert record["output_tokens"] == 10
+
+    # Verify the JSONL file can be loaded by CustomDataset
+    dataset = CustomDataset(dataset_path=str(jsonl_path))
+    samples = dataset.sample(
+        tokenizer=hf_tokenizer,
+        num_requests=10,
+        output_len=10,
+        skip_chat_template=True,
+    )
+
+    assert len(samples) == 10
+    assert all(sample.expected_output_len == 10 for sample in samples)

tests/compile/passes/distributed/test_async_tp.py

@@ -31,6 +31,7 @@ from vllm.platforms import current_platform
 from vllm.utils.system_utils import update_environment_variables
 from vllm.utils.torch_utils import set_random_seed
 
+DEVICE_TYPE = current_platform.device_type
 FP8_DTYPE = current_platform.fp8_dtype()
 
 prompts = [
@@ -299,7 +300,7 @@ def async_tp_pass_on_test_model(
 ):
     set_random_seed(0)
 
-    device = torch.device(f"cuda:{local_rank}")
+    device = torch.device(f"{DEVICE_TYPE}:{local_rank}")
     torch.accelerator.set_device_index(device)
     torch.set_default_device(device)
     torch.set_default_dtype(dtype)
@@ -324,7 +325,7 @@ def async_tp_pass_on_test_model(
             fuse_gemm_comms=True,
         ),
     )
-    vllm_config.device_config = DeviceConfig(device=torch.device("cuda"))
+    vllm_config.device_config = DeviceConfig(device=torch.device(DEVICE_TYPE))
 
     # this is a fake model name to construct the model config
     # in the vllm_config, it's not really used.

tests/compile/passes/distributed/test_fusion_all_reduce.py

@@ -37,6 +37,8 @@ from vllm.platforms import current_platform
 from vllm.utils.system_utils import update_environment_variables
 from vllm.utils.torch_utils import set_random_seed
 
+DEVICE_TYPE = current_platform.device_type
+
 
 class TestAllReduceRMSNormModel(torch.nn.Module):
     def __init__(
@@ -268,7 +270,7 @@ def all_reduce_fusion_pass_on_test_model(
 ):
     set_random_seed(0)
 
-    device = torch.device(f"cuda:{local_rank}")
+    device = torch.device(f"{DEVICE_TYPE}:{local_rank}")
     torch.accelerator.set_device_index(device)
     torch.set_default_device(device)
     torch.set_default_dtype(dtype)
@@ -300,7 +302,7 @@ def all_reduce_fusion_pass_on_test_model(
     vllm_config.compilation_config.pass_config = PassConfig(
         fuse_allreduce_rms=True, eliminate_noops=True
     )
-    vllm_config.device_config = DeviceConfig(device=torch.device("cuda"))
+    vllm_config.device_config = DeviceConfig(device=torch.device(DEVICE_TYPE))
     vllm_config.parallel_config.rank = local_rank  # Setup rank for debug path
 
     # this is a fake model name to construct the model config

tests/compile/passes/distributed/test_sequence_parallelism.py

@@ -35,6 +35,8 @@ from vllm.platforms import current_platform
 from vllm.utils.system_utils import update_environment_variables
 from vllm.utils.torch_utils import set_random_seed
 
+DEVICE_TYPE = current_platform.device_type
+
 pytestmark = pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test CUDA")
 
 FP8_DTYPE = current_platform.fp8_dtype()
@@ -228,7 +230,7 @@ def sequence_parallelism_pass_on_test_model(
 ):
     set_random_seed(0)
 
-    device = torch.device(f"cuda:{local_rank}")
+    device = torch.device(f"{DEVICE_TYPE}:{local_rank}")
     torch.accelerator.set_device_index(device)
     torch.set_default_device(device)
     torch.set_default_dtype(dtype)
@@ -258,7 +260,7 @@ def sequence_parallelism_pass_on_test_model(
             eliminate_noops=True,
         ),
     )  # NoOp needed for fusion
-    device_config = DeviceConfig(device=torch.device("cuda"))
+    device_config = DeviceConfig(device=torch.device(DEVICE_TYPE))
 
     # this is a fake model name to construct the model config
     # in the vllm_config, it's not really used.

tests/compile/passes/test_fusion_attn.py

@@ -41,6 +41,7 @@ from vllm.v1.attention.backend import AttentionMetadata
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
 from vllm.v1.kv_cache_interface import AttentionSpec, get_kv_quant_mode
 
+DEVICE_TYPE = current_platform.device_type
 FP8_DTYPE = current_platform.fp8_dtype()
 FP4_DTYPE = torch.uint8
 
@@ -300,7 +301,7 @@ def test_attention_quant_pattern(
 
     custom_ops_list = custom_ops.split(",") if custom_ops else []
 
-    device = torch.device("cuda:0")
+    device = torch.device(f"{DEVICE_TYPE}:0")
     torch.set_default_dtype(dtype)
     torch.manual_seed(42)
 

tests/compile/passes/test_mla_attn_quant_fusion.py

@@ -45,6 +45,7 @@ from vllm.v1.kv_cache_interface import MLAAttentionSpec
 
 FP8_DTYPE = current_platform.fp8_dtype()
 FP4_DTYPE = torch.uint8
+DEVICE_TYPE = current_platform.device_type
 
 
 class MLAAttentionQuantPatternModel(torch.nn.Module):
@@ -356,7 +357,7 @@ def test_mla_attention_quant_pattern(
 
     custom_ops_list = custom_ops.split(",") if custom_ops else []
 
-    device = torch.device("cuda:0")
+    device = torch.device(f"{DEVICE_TYPE}:0")
     torch.set_default_dtype(dtype)
     torch.manual_seed(42)
 

tests/compile/passes/test_noop_elimination.py

@@ -8,6 +8,9 @@ import vllm
 from tests.compile.backend import TestBackend
 from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass
 from vllm.config import CompilationConfig, CompilationMode, PassConfig, VllmConfig
+from vllm.platforms import current_platform
+
+DEVICE_TYPE = current_platform.device_type
 
 
 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
@@ -17,7 +20,7 @@ from vllm.config import CompilationConfig, CompilationMode, PassConfig, VllmConf
 )
 @pytest.mark.parametrize("hidden_size", [64, 4096])
 def test_noop_elimination(dtype, num_tokens, hidden_size, buffer_size):
-    torch.set_default_device("cuda")
+    torch.set_default_device(DEVICE_TYPE)
     torch.set_default_dtype(dtype)
     torch.manual_seed(1)
 
@@ -88,7 +91,7 @@ def test_non_noop_slice_preserved():
     Regression test for a bug where end=-1 was treated like an inferred
     dimension (reshape semantics) leading to incorrect elimination.
     """
-    torch.set_default_device("cuda")
+    torch.set_default_device(DEVICE_TYPE)
     x = torch.randn(16, 16)
 
     class SliceModel(torch.nn.Module):

tests/compile/passes/test_scatter_split_replace.py

@@ -13,6 +13,9 @@ from vllm.compilation.passes.utility.scatter_split_replace import (
 from vllm.compilation.passes.utility.split_coalescing import SplitCoalescingPass
 from vllm.config import CompilationConfig, CompilationMode, VllmConfig
 from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
+from vllm.platforms import current_platform
+
+DEVICE_TYPE = current_platform.device_type
 
 
 class ScatterSplitReplacementModel(nn.Module):
@@ -61,7 +64,7 @@ class ScatterSplitReplacementModel(nn.Module):
 
 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
 def test_scatter_split_replace(dtype):
-    torch.set_default_device("cuda")
+    torch.set_default_device(DEVICE_TYPE)
     torch.set_default_dtype(dtype)
     torch.manual_seed(0)