[ROCm][CI] Stage B gating (#42025 )

Signed-off-by: Andreas Karatzas <akaratza@amd.com>
DeepSeekV4-Pro enable cuda graph full and piecewise mode (#42604 )
2026-05-15 01:49:27 -07:00 · 2026-05-15 01:45:30 -07:00 · 2026-05-15 08:05:57 +00:00 · 2026-05-15 15:59:18 +08:00 · 2026-05-15 15:20:42 +08:00 · 2026-05-15 00:13:32 -07:00
474 changed files with 23811 additions and 7702 deletions
@@ -8,6 +8,7 @@ run_all_patterns:
  - "CMakeLists.txt"
  - "requirements/common.txt"
  - "requirements/cuda.txt"
+  - "requirements/kv_connectors.txt"
  - "requirements/build/cuda.txt"
  - "requirements/test/cuda.txt"
  - "setup.py"
@@ -61,7 +61,7 @@ steps:
  source_file_dependencies:
  - csrc/cpu/
  - vllm/model_executor/layers/quantization/cpu_wna16.py
-  - vllm/model_executor/layers/quantization/gptq_marlin.py
+  - vllm/model_executor/layers/quantization/auto_gptq.py
  - vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_int8.py
  - vllm/model_executor/layers/quantization/kernels/scaled_mm/cpu.py
  - vllm/model_executor/layers/quantization/kernels/mixed_precision/cpu.py
@@ -36,6 +36,7 @@ steps:
        python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --attention-backend=TRITON_ATTN &&
        python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --quantization fp8 &&
        python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --kv-cache-dtype fp8 &&
+        python3 examples/basic/offline_inference/generate.py --model nvidia/Llama-3.1-8B-Instruct-FP8 --block-size 64 --enforce-eager --quantization modelopt --kv-cache-dtype fp8 --attention-backend TRITON_ATTN --max-model-len 4096 &&
        python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4 --block-size 64 --enforce-eager --max-model-len 8192 &&
        python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2 &&
        python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2 --enable-expert-parallel'
@@ -62,4 +63,4 @@ steps:
        pytest -v -s v1/structured_output &&
        pytest -v -s v1/test_serial_utils.py &&
        pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py &&
-        pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --ignore=v1/kv_connector/unit/test_lmcache_integration.py --ignore=v1/kv_connector/unit/test_hf3fs_client.py --ignore=v1/kv_connector/unit/test_hf3fs_connector.py --ignore=v1/kv_connector/unit/test_hf3fs_metadata_server.py'
+        pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --ignore=v1/kv_connector/unit/test_lmcache_integration.py --ignore=v1/kv_connector/unit/test_hf3fs_client.py --ignore=v1/kv_connector/unit/test_hf3fs_connector.py --ignore=v1/kv_connector/unit/test_hf3fs_metadata_server.py --ignore=v1/kv_connector/unit/test_offloading_connector.py'
@@ -2,7 +2,7 @@
 # We can use this script to compute baseline accuracy on chartqa for vllm.
 #
 # Make sure you have lm-eval-harness installed:
-#   pip install "lm-eval[api]>=0.4.11"
+#   pip install "lm-eval[api]>=0.4.12"

 usage() {
    echo``
@@ -2,7 +2,7 @@
 # We can use this script to compute baseline accuracy on GSM for transformers.
 #
 # Make sure you have lm-eval-harness installed:
-#   pip install "lm-eval[api]>=0.4.11"
+#   pip install "lm-eval[api]>=0.4.12"

 usage() {
    echo``
@@ -3,7 +3,7 @@
 # We use this for fp8, which HF does not support.
 #
 # Make sure you have lm-eval-harness installed:
-#   pip install "lm-eval[api]>=0.4.11"
+#   pip install "lm-eval[api]>=0.4.12"

 usage() {
    echo``
@@ -3,7 +3,7 @@
 # We use this for fp8, which HF does not support.
 #
 # Make sure you have lm-eval-harness installed:
-#   pip install "lm-eval[api]>=0.4.11"
+#   pip install "lm-eval[api]>=0.4.12"

 usage() {
    echo``
@@ -31,30 +31,9 @@
    }
  },
  "tests": [
-    {
-      "test_name": "serving_llama8B_tp1_sharegpt",
-      "server_parameters": {
-        "tensor_parallel_size": 1
-      },
-      "client_parameters": {
-        "dataset_name": "sharegpt",
-        "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json"
-      }
-    },
-    {
-      "test_name": "serving_llama8B_tp2_sharegpt",
-      "server_parameters": {
-        "tensor_parallel_size": 2
-      },
-      "client_parameters": {
-        "dataset_name": "sharegpt",
-        "dataset_path": "./ShareGPT_V3_unfiltered_cleaned_split.json"
-      }
-    },
    {
      "test_name": "serving_llama8B_tp1_random_128_128",
      "server_parameters": {
-        "tensor_parallel_size": 1
      },
      "client_parameters": {
        "dataset_name": "random",
@@ -62,158 +41,10 @@
        "random-output-len": 128
      }
    },
-    {
-      "test_name": "serving_llama8B_tp2_random_128_128",
-      "server_parameters": {
-        "tensor_parallel_size": 2
-      },
-      "client_parameters": {
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "random-output-len": 128
-      }
-    },
-    {
-      "test_name": "serving_llama8B_tp4_random_128_128",
-      "server_parameters": {
-        "tensor_parallel_size": 4
-      },
-      "client_parameters": {
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "random-output-len": 128
-      }
-    },
-    {
-      "test_name": "serving_llama8B_tp1_random_128_2048",
-      "server_parameters": {
-        "tensor_parallel_size": 1
-      },
-      "client_parameters": {
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "random-output-len": 2048
-      }
-    },
-    {
-      "test_name": "serving_llama8B_tp2_random_128_2048",
-      "server_parameters": {
-        "tensor_parallel_size": 2
-      },
-      "client_parameters": {
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "random-output-len": 2048
-      }
-    },
-    {
-      "test_name": "serving_llama8B_tp4_random_128_2048",
-      "server_parameters": {
-        "tensor_parallel_size": 4
-      },
-      "client_parameters": {
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "random-output-len": 2048
-      }
-    },
-    {
-      "test_name": "serving_llama8B_tp1_random_2048_128",
-      "server_parameters": {
-        "tensor_parallel_size": 1
-      },
-      "client_parameters": {
-        "dataset_name": "random",
-        "random-input-len": 2048,
-        "random-output-len": 128
-      }
-    },
-    {
-      "test_name": "serving_llama8B_tp2_random_2048_128",
-      "server_parameters": {
-        "tensor_parallel_size": 2
-      },
-      "client_parameters": {
-        "dataset_name": "random",
-        "random-input-len": 2048,
-        "random-output-len": 128
-      }
-    },
-    {
-      "test_name": "serving_llama8B_tp4_random_2048_128",
-      "server_parameters": {
-        "tensor_parallel_size": 4
-      },
-      "client_parameters": {
-        "dataset_name": "random",
-        "random-input-len": 2048,
-        "random-output-len": 128
-      }
-    },
-    {
-      "test_name": "serving_llama8B_tp1_random_2048_2048",
-      "server_parameters": {
-        "tensor_parallel_size": 1
-      },
-      "client_parameters": {
-        "dataset_name": "random",
-        "random-input-len": 2048,
-        "random-output-len": 2048
-      }
-    },
-    {
-      "test_name": "serving_llama8B_tp2_random_2048_2048",
-      "server_parameters": {
-        "tensor_parallel_size": 2
-      },
-      "client_parameters": {
-        "dataset_name": "random",
-        "random-input-len": 2048,
-        "random-output-len": 2048
-      }
-    },
-    {
-      "test_name": "serving_llama8B_tp4_random_2048_2048",
-      "server_parameters": {
-        "tensor_parallel_size": 4
-      },
-      "client_parameters": {
-        "dataset_name": "random",
-        "random-input-len": 2048,
-        "random-output-len": 2048
-      }
-    },
    {
      "test_name": "serving_llama8B_int4_tp1_random_128_128",
      "server_parameters": {
-        "model": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
-        "tensor_parallel_size": 1
-      },
-      "client_parameters": {
-        "model": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "random-output-len": 128
-      }
-    },
-    {
-      "test_name": "serving_llama8B_int4_tp2_random_128_128",
-      "server_parameters": {
-        "model": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
-        "tensor_parallel_size": 2
-      },
-      "client_parameters": {
-        "model": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "random-output-len": 128
-      }
-    },
-    {
-      "test_name": "serving_llama8B_int4_tp4_random_128_128",
-      "server_parameters": {
-        "model": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
-        "tensor_parallel_size": 4
+        "model": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4"
      },
      "client_parameters": {
        "model": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
@@ -225,8 +56,7 @@
    {
      "test_name": "serving_llama8B_int8_tp1_random_128_128",
      "server_parameters": {
-        "model": "RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w8a8",
-        "tensor_parallel_size": 1
+        "model": "RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w8a8"
      },
      "client_parameters": {
        "model": "RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w8a8",
@@ -236,26 +66,12 @@
      }
    },
    {
-      "test_name": "serving_llama8B_int8_tp2_random_128_128",
+      "test_name": "serving_llama1B_tp1_random_128_128",
      "server_parameters": {
-        "model": "RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w8a8",
-        "tensor_parallel_size": 2
+        "model": "meta-llama/Llama-3.2-1B"
      },
      "client_parameters": {
-        "model": "RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w8a8",
-        "dataset_name": "random",
-        "random-input-len": 128,
-        "random-output-len": 128
-      }
-    },
-    {
-      "test_name": "serving_llama8B_int8_tp4_random_128_128",
-      "server_parameters": {
-        "model": "RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w8a8",
-        "tensor_parallel_size": 4
-      },
-      "client_parameters": {
-        "model": "RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w8a8",
+        "model": "meta-llama/Llama-3.2-1B",
        "dataset_name": "random",
        "random-input-len": 128,
        "random-output-len": 128
@@ -264,8 +80,7 @@
    {
      "test_name": "serving_llama3B_tp1_random_128_128",
      "server_parameters": {
-        "model": "meta-llama/Llama-3.2-3B-Instruct",
-        "tensor_parallel_size": 1
+        "model": "meta-llama/Llama-3.2-3B-Instruct"
      },
      "client_parameters": {
        "model": "meta-llama/Llama-3.2-3B-Instruct",
@@ -274,11 +89,22 @@
        "random-output-len": 128
      }
    },
+    {
+      "test_name": "serving_llama70B_tp1_random_128_128",
+      "server_parameters": {
+        "model": "meta-llama/Llama-3.3-70B-Instruct"
+      },
+      "client_parameters": {
+        "model": "meta-llama/Llama-3.3-70B-Instruct",
+        "dataset_name": "random",
+        "random-input-len": 128,
+        "random-output-len": 128
+      }
+    },
    {
      "test_name": "serving_granite2B_tp1_random_128_128",
      "server_parameters": {
-        "model": "ibm-granite/granite-3.2-2b-instruct",
-        "tensor_parallel_size": 1
+        "model": "ibm-granite/granite-3.2-2b-instruct"
      },
      "client_parameters": {
        "model": "ibm-granite/granite-3.2-2b-instruct",
@@ -290,8 +116,7 @@
    {
      "test_name": "serving_qwen1.7B_tp1_random_128_128",
      "server_parameters": {
-        "model": "Qwen/Qwen3-1.7B",
-        "tensor_parallel_size": 1
+        "model": "Qwen/Qwen3-1.7B"
      },
      "client_parameters": {
        "model": "Qwen/Qwen3-1.7B",
@@ -303,8 +128,7 @@
    {
      "test_name": "serving_qwen4B_tp1_random_128_128",
      "server_parameters": {
-        "model": "Qwen/Qwen3-4B",
-        "tensor_parallel_size": 1
+        "model": "Qwen/Qwen3-4B"
      },
      "client_parameters": {
        "model": "Qwen/Qwen3-4B",
@@ -316,8 +140,7 @@
    {
      "test_name": "serving_qwen8B_tp1_random_128_128",
      "server_parameters": {
-        "model": "Qwen/Qwen3-8B",
-        "tensor_parallel_size": 1
+        "model": "Qwen/Qwen3-8B"
      },
      "client_parameters": {
        "model": "Qwen/Qwen3-8B",
@@ -326,11 +149,34 @@
        "random-output-len": 128
      }
    },
+    {
+      "test_name": "serving_qwen14B_tp1_random_128_128",
+      "server_parameters": {
+        "model": "Qwen/Qwen3-14B"
+      },
+      "client_parameters": {
+        "model": "Qwen/Qwen3-14B",
+        "dataset_name": "random",
+        "random-input-len": 128,
+        "random-output-len": 128
+      }
+    },
+    {
+      "test_name": "serving_qwen30B_tp1_random_128_128",
+      "server_parameters": {
+        "model": "Qwen/Qwen3-30B-A3B"
+      },
+      "client_parameters": {
+        "model": "Qwen/Qwen3-30B-A3B",
+        "dataset_name": "random",
+        "random-input-len": 128,
+        "random-output-len": 128
+      }
+    },
    {
      "test_name": "serving_glm9B_tp1_random_128_128",
      "server_parameters": {
-        "model": "zai-org/glm-4-9b-hf",
-        "tensor_parallel_size": 1
+        "model": "zai-org/glm-4-9b-hf"
      },
      "client_parameters": {
        "model": "zai-org/glm-4-9b-hf",
@@ -342,8 +188,7 @@
    {
      "test_name": "serving_gemma7B_tp1_random_128_128",
      "server_parameters": {
-        "model": "google/gemma-7b",
-        "tensor_parallel_size": 1
+        "model": "google/gemma-7b"
      },
      "client_parameters": {
        "model": "google/gemma-7b",
@@ -351,6 +196,94 @@
        "random-input-len": 128,
        "random-output-len": 128
      }
+    },
+    {
+      "test_name": "serving_gemma3-4b_tp1_random_128_128",
+      "server_environment_variables": {
+	"VLLM_CPU_SGL_KERNEL": 0
+      },
+      "server_parameters": {
+        "model": "google/gemma-3-4b-it"
+      },
+      "client_parameters": {
+        "model": "google/gemma-3-4b-it",
+        "dataset_name": "random",
+        "random-input-len": 128,
+        "random-output-len": 128
+      }
+    },
+    {
+      "test_name": "serving_gemma3-12b_tp1_random_128_128",
+      "server_environment_variables": {
+	"VLLM_CPU_SGL_KERNEL": 0
+      },
+      "server_parameters": {
+        "model": "google/gemma-3-12b-it"
+      },
+      "client_parameters": {
+        "model": "google/gemma-3-12b-it",
+        "dataset_name": "random",
+        "random-input-len": 128,
+        "random-output-len": 128
+      }
+    },
+    {
+      "test_name": "serving_gemma4-4b_tp1_random_128_128",
+      "server_environment_variables": {
+	"VLLM_CPU_SGL_KERNEL": 0
+      },
+      "server_parameters": {
+        "model": "google/gemma-4-E4B-it"
+      },
+      "client_parameters": {
+        "model": "google/gemma-4-E4B-it",
+        "dataset_name": "random",
+        "random-input-len": 128,
+        "random-output-len": 128
+      }
+    },
+    {
+      "test_name": "serving_gemma4-2b_tp1_random_128_128",
+      "server_environment_variables": {
+	"VLLM_CPU_SGL_KERNEL": 0
+      },
+      "server_parameters": {
+        "model": "google/gemma-4-E2B-it"
+      },
+      "client_parameters": {
+        "model": "google/gemma-4-E2B-it",
+        "dataset_name": "random",
+        "random-input-len": 128,
+        "random-output-len": 128
+      }
+    },
+    {
+      "test_name": "serving_gemma4-26b_tp1_random_128_128",
+      "server_environment_variables": {
+	"VLLM_CPU_SGL_KERNEL": 0,
+        "VLLM_CPU_ATTN_SPLIT_KV": 0
+      },
+      "server_parameters": {
+        "model": "google/gemma-4-26B-A4B-it"
+      },
+      "client_parameters": {
+        "model": "google/gemma-4-26B-A4B-it",
+        "dataset_name": "random",
+        "random-input-len": 128,
+        "random-output-len": 128
+      }
+    },
+    {
+      "test_name": "serving_phi4_tp1_random_128_128",
+      "server_parameters": {
+        "model": "microsoft/Phi-4-reasoning"
+      },
+      "client_parameters": {
+        "model": "microsoft/Phi-4-reasoning",
+        "dataset_name": "random",
+        "random-input-len": 128,
+        "random-output-len": 128
+      }
    }
  ]
 }
@@ -28,6 +28,7 @@ steps:
          - "mkdir artifacts"
          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
          - "bash .buildkite/scripts/upload-nightly-wheels.sh"
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"'
        env:
          DOCKER_BUILDKIT: "1"

@@ -41,6 +42,7 @@ steps:
          - "mkdir artifacts"
          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
          - "bash .buildkite/scripts/upload-nightly-wheels.sh"
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"'
        env:
          DOCKER_BUILDKIT: "1"

@@ -54,6 +56,7 @@ steps:
          - "mkdir artifacts"
          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
          - "bash .buildkite/scripts/upload-nightly-wheels.sh"
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"'
        env:
          DOCKER_BUILDKIT: "1"

@@ -67,6 +70,7 @@ steps:
          - "mkdir artifacts"
          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
          - "bash .buildkite/scripts/upload-nightly-wheels.sh"
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"'
        env:
          DOCKER_BUILDKIT: "1"

@@ -80,6 +84,7 @@ steps:
          - "mkdir artifacts"
          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
          - "bash .buildkite/scripts/upload-nightly-wheels.sh"
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"'
        env:
          DOCKER_BUILDKIT: "1"

@@ -93,6 +98,7 @@ steps:
          - "mkdir artifacts"
          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
          - "bash .buildkite/scripts/upload-nightly-wheels.sh"
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"'
        env:
          DOCKER_BUILDKIT: "1"

@@ -138,6 +144,7 @@ steps:
          # re-tag to default image tag and push, just in case arm64 build fails
          - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"
          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"'

      - label: "Build release image - aarch64 - CUDA 13.0"
        depends_on: ~
@@ -160,6 +167,7 @@ steps:
              --progress plain \
              -f docker/Dockerfile .
          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"'

      - label: "Build release image - x86_64 - CUDA 12.9"
        depends_on: ~
@@ -184,6 +192,7 @@ steps:
          # re-tag to default image tag and push, just in case arm64 build fails
          - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129"
          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129"
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129"'

      - label: "Build release image - aarch64 - CUDA 12.9"
        depends_on: ~
@@ -205,6 +214,7 @@ steps:
              --progress plain \
              -f docker/Dockerfile .
          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129"
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129"'

      - label: "Build release image - x86_64 - CUDA 13.0 - Ubuntu 24.04"
        depends_on: ~
@@ -231,6 +241,7 @@ steps:
          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"
          - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"
          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"'

      - label: "Build release image - aarch64 - CUDA 13.0 - Ubuntu 24.04"
        depends_on: ~
@@ -255,6 +266,7 @@ steps:
              --progress plain \
              -f docker/Dockerfile .
          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"'

      - label: "Build release image - x86_64 - CUDA 12.9 - Ubuntu 24.04"
        depends_on: ~
@@ -280,6 +292,7 @@ steps:
          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404"
          - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404"
          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404"
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404"'

      - label: "Build release image - aarch64 - CUDA 12.9 - Ubuntu 24.04"
        depends_on: ~
@@ -303,6 +316,7 @@ steps:
              --progress plain \
              -f docker/Dockerfile .
          - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404"
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404"'

      - block: "Build release image for x86_64 CPU"
        key: block-cpu-release-image-build
@@ -320,6 +334,7 @@ steps:
          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_CPU_X86=true --tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:$(buildkite-agent meta-data get release-version) --tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:latest --progress plain --target vllm-openai -f docker/Dockerfile.cpu ."
          - "docker push public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:latest"
          - "docker push public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:$(buildkite-agent meta-data get release-version)"
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:$(buildkite-agent meta-data get release-version)"'
        env:
          DOCKER_BUILDKIT: "1"

@@ -339,6 +354,7 @@ steps:
          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --tag public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:$(buildkite-agent meta-data get release-version) --tag public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:latest --progress plain --target vllm-openai -f docker/Dockerfile.cpu ."
          - "docker push public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:latest"
          - "docker push public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:$(buildkite-agent meta-data get release-version)"
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:$(buildkite-agent meta-data get release-version)"'
        env:
          DOCKER_BUILDKIT: "1"

@@ -356,15 +372,7 @@ steps:
          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
          - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64 --amend"
          - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"
-
-      - label: "Annotate release workflow - CUDA 13.0"
-        depends_on:
-          - create-multi-arch-manifest
-        id: annotate-release-workflow
-        agents:
-          queue: small_cpu_queue_release
-        commands:
-          - "bash .buildkite/scripts/annotate-release.sh"
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: CUDA 13.0" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"'

      - label: "Create multi-arch manifest - CUDA 12.9"
        depends_on:
@@ -377,6 +385,7 @@ steps:
          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
          - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu129 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu129 --amend"
          - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129"
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: CUDA 12.9" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129"'

      - label: "Create multi-arch manifest - CUDA 13.0 - Ubuntu 24.04"
        depends_on:
@@ -389,6 +398,7 @@ steps:
          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
          - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-ubuntu2404 --amend"
          - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: CUDA 13.0 Ubuntu 24.04" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"'

      - label: "Create multi-arch manifest - CUDA 12.9 - Ubuntu 24.04"
        depends_on:
@@ -401,6 +411,7 @@ steps:
          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
          - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu129-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu129-ubuntu2404 --amend"
          - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404"
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: CUDA 12.9 Ubuntu 24.04" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404"'

      - label: "Publish nightly multi-arch image to DockerHub"
        depends_on:
@@ -438,59 +449,6 @@ steps:
          DOCKER_BUILDKIT: "1"
          DOCKERHUB_USERNAME: "vllmbot"

-      - block: "Publish release images to DockerHub"
-        key: block-publish-release-images
-        depends_on:
-          - create-multi-arch-manifest
-          - create-multi-arch-manifest-cuda-12-9
-          - create-multi-arch-manifest-ubuntu2404
-          - create-multi-arch-manifest-cuda-12-9-ubuntu2404
-          - build-rocm-release-image
-          - input-release-version
-          # Wait for CPU builds if their block steps were unblocked, so publish
-          # doesn't race the in-progress CPU build. allow_failure lets publish
-          # proceed when the operator legitimately leaves the CPU block steps
-          # unblocked or the CPU build fails.
-          - step: build-cpu-release-image-x86
-            allow_failure: true
-          - step: build-cpu-release-image-arm64
-            allow_failure: true
-        if: build.env("NIGHTLY") != "1"
-
-      - label: "Publish release images to DockerHub"
-        depends_on:
-          - block-publish-release-images
-        key: publish-release-images-dockerhub
-        agents:
-          queue: small_cpu_queue_release
-        commands:
-          - "bash .buildkite/scripts/publish-release-images.sh"
-        plugins:
-          - docker-login#v3.0.0:
-              username: vllmbot
-              password-env: DOCKERHUB_TOKEN
-        env:
-          DOCKER_BUILDKIT: "1"
-          DOCKERHUB_USERNAME: "vllmbot"
-
-  - group: "Publish wheels"
-    key: "publish-wheels"
-    steps:
-      - block: "Confirm update release wheels to PyPI (experimental, use with caution)?"
-        key: block-upload-release-wheels
-        depends_on:
-          - input-release-version
-          - build-wheels
-
-      - label: "Upload release wheels to PyPI"
-        depends_on:
-          - block-upload-release-wheels
-        id: upload-release-wheels
-        agents:
-          queue: small_cpu_queue_release
-        commands:
-          - "bash .buildkite/scripts/upload-release-wheels-pypi.sh"
-
  # =============================================================================
  # ROCm Release Pipeline (x86_64 only)
  # =============================================================================
@@ -604,7 +562,7 @@ steps:
          echo ""
          echo " Build complete - Image and wheels cached"
        fi
-        
+
    artifact_paths:
      - "artifacts/rocm-base-wheels/*.whl"
    env:
@@ -820,7 +778,7 @@ steps:
        
        # Push to ECR
        docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm
-        
+
        echo ""
        echo " Successfully built and pushed ROCm release image"
        echo "   Image: public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm"
@@ -847,3 +805,60 @@ steps:
    env:
      DOCKER_BUILDKIT: "1"
      DOCKERHUB_USERNAME: "vllmbot"
+
+  # =============================================================================
+  # Publish to DockerHub and PyPI (at the end so all builds complete first)
+  # =============================================================================
+
+  - block: "Publish release images to DockerHub"
+    key: block-publish-release-images
+    depends_on:
+      - create-multi-arch-manifest
+      - create-multi-arch-manifest-cuda-12-9
+      - create-multi-arch-manifest-ubuntu2404
+      - create-multi-arch-manifest-cuda-12-9-ubuntu2404
+      - build-rocm-release-image
+      - input-release-version
+      # Wait for CPU builds if their block steps were unblocked, so publish
+      # doesn't race the in-progress CPU build. allow_failure lets publish
+      # proceed when the operator legitimately leaves the CPU block steps
+      # unblocked or the CPU build fails.
+      - step: build-cpu-release-image-x86
+        allow_failure: true
+      - step: build-cpu-release-image-arm64
+        allow_failure: true
+    if: build.env("NIGHTLY") != "1"
+
+  - label: "Publish release images to DockerHub"
+    depends_on:
+      - block-publish-release-images
+    key: publish-release-images-dockerhub
+    agents:
+      queue: small_cpu_queue_release
+    commands:
+      - "bash .buildkite/scripts/publish-release-images.sh"
+    plugins:
+      - docker-login#v3.0.0:
+          username: vllmbot
+          password-env: DOCKERHUB_TOKEN
+    env:
+      DOCKER_BUILDKIT: "1"
+      DOCKERHUB_USERNAME: "vllmbot"
+
+  - group: "Publish wheels"
+    key: "publish-wheels"
+    steps:
+      - block: "Confirm update release wheels to PyPI (experimental, use with caution)?"
+        key: block-upload-release-wheels
+        depends_on:
+          - input-release-version
+          - build-wheels
+
+      - label: "Upload release wheels to PyPI"
+        depends_on:
+          - block-upload-release-wheels
+        id: upload-release-wheels
+        agents:
+          queue: small_cpu_queue_release
+        commands:
+          - "bash .buildkite/scripts/upload-release-wheels-pypi.sh"
@@ -0,0 +1,9 @@
+#!/bin/bash
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# Append a build artifact line to the Buildkite annotation.
+# Usage: annotate-build-artifact.sh <label> <value>
+set -e
+echo "- **${1}**: \`${2}\`" | \
+  buildkite-agent annotate --append --style 'info' --context 'release-artifacts'
@@ -1,27 +0,0 @@
-#!/bin/bash
-
-set -ex
-
-# Get release version, default to 1.0.0.dev for nightly/per-commit builds
-RELEASE_VERSION=$(buildkite-agent meta-data get release-version 2>/dev/null | sed 's/^v//')
-if [ -z "${RELEASE_VERSION}" ]; then
-  RELEASE_VERSION="1.0.0.dev"
-fi
-
-buildkite-agent annotate --style 'info' --context 'release-workflow' << EOF
-To download the wheel (by commit):
-\`\`\`
-aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux_2_35_x86_64.whl .
-aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux_2_35_aarch64.whl .
-
-(Optional) For CUDA 12.9:
-aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu129-cp38-abi3-manylinux_2_31_x86_64.whl .
-aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu129-cp38-abi3-manylinux_2_31_aarch64.whl .
-
-(Optional) For CPU:
-aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cpu-cp38-abi3-manylinux_2_35_x86_64.whl .
-aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cpu-cp38-abi3-manylinux_2_35_aarch64.whl .
-\`\`\`
-
-Docker images are published automatically by the "Publish release images to DockerHub" pipeline step.
-EOF
@@ -61,7 +61,7 @@ echo "Results will be stored in: $RESULTS_DIR"
 echo "--- Installing Python dependencies ---"
 python3 -m pip install --progress-bar off git+https://github.com/thuml/depyf.git \
    && python3 -m pip install --progress-bar off pytest pytest-asyncio tpu-info \
-    && python3 -m pip install --progress-bar off "lm-eval[api]>=0.4.11" \
+    && python3 -m pip install --progress-bar off "lm-eval[api]>=0.4.12" \
    && python3 -m pip install --progress-bar off hf-transfer tblib==3.1.0
 echo "--- Python dependencies installed ---"

@@ -61,7 +61,7 @@ echo "Results will be stored in: $RESULTS_DIR"
 echo "--- Installing Python dependencies ---"
 python3 -m pip install --progress-bar off git+https://github.com/thuml/depyf.git \
    && python3 -m pip install --progress-bar off pytest pytest-asyncio tpu-info \
-    && python3 -m pip install --progress-bar off "lm-eval[api]>=0.4.11" \
+    && python3 -m pip install --progress-bar off "lm-eval[api]>=0.4.12" \
    && python3 -m pip install --progress-bar off hf-transfer tblib==3.1.0
 echo "--- Python dependencies installed ---"

@@ -39,11 +39,17 @@ fi

 set -x # avoid printing secrets above

-# install twine and sdist build prerequisites from pypi
-python3 -m venv /tmp/vllm-release-env
+# install uv if not already available
+if ! command -v uv &> /dev/null; then
+  curl -LsSf https://astral.sh/uv/install.sh | UV_VERSION=0.11.14 sh
+  export PATH="$HOME/.local/bin:$PATH"
+fi
+
+# install twine and sdist build prerequisites using uv with Python 3.12
+uv venv --python 3.12 /tmp/vllm-release-env
 source /tmp/vllm-release-env/bin/activate
-pip install twine
-pip install -r requirements/build/cuda.txt
+uv pip install twine
+uv pip install -r requirements/build/cuda.txt
 python3 -m twine --version

 # copy release wheels to local directory
@@ -798,32 +798,6 @@ steps:
  commands:
  - pytest -v -s v1/attention

- label: V1 Core + KV + Metrics # TBD
-  timeout_in_minutes: 60
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
-  agent_pool: mi250_1
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/
-  - tests/v1/core
-  - tests/v1/executor
-  - tests/v1/kv_offload
-  - tests/v1/worker
-  - tests/v1/kv_connector/unit
-  - tests/v1/metrics
-  - tests/entrypoints/openai/correctness/test_lmeval.py
-  commands:
-  - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
-  - pytest -v -s -m 'not cpu_test' v1/core
-  - pytest -v -s v1/executor
-  - pytest -v -s v1/kv_offload
-  - pytest -v -s v1/worker
-  - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit
-  - pytest -v -s -m 'not cpu_test' v1/metrics
-  - pip install -U git+https://github.com/robertgshaw2-redhat/lm-evaluation-harness.git@streaming-api
-  - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
-
 - label: V1 Sample + Logits # TBD
  timeout_in_minutes: 60
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
@@ -1279,6 +1253,7 @@ steps:
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
+  optional: true
  fast_check: true
  torch_nightly: true
  working_dir: "/vllm-workspace/tests"
@@ -1299,6 +1274,7 @@ steps:
  agent_pool: mi300_1
  fast_check: true
  torch_nightly: true
+  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -1312,6 +1288,7 @@ steps:
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
+  optional: true
  fast_check: true
  torch_nightly: true
  working_dir: "/vllm-workspace/tests"
@@ -1328,6 +1305,7 @@ steps:
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
+  optional: true
  fast_check: true
  torch_nightly: true
  working_dir: "/vllm-workspace/tests"
@@ -1728,6 +1706,7 @@ steps:
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
+  optional: true
  parallelism: 2
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
@@ -1786,6 +1765,7 @@ steps:
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
+  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -51,7 +51,7 @@ steps:

 - label: Hybrid SSM NixlConnector PD accuracy tests (4 GPUs)
  key: hybrid-ssm-nixlconnector-pd-accuracy-tests-4-gpus
-  timeout_in_minutes: 20
+  timeout_in_minutes: 25
  working_dir: "/vllm-workspace/tests"
  num_devices: 4
  source_file_dependencies:
@@ -36,6 +36,12 @@ steps:
  commands:
    - pytest -v -s v1/engine/test_preprocess_error_handling.py
    - pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py
+  mirror:
+    amd:
+      device: mi300_1
+      timeout_in_minutes: 40
+      depends_on:
+      - image-build-amd

 - label: e2e Scheduling (1 GPU)
  key: e2e-scheduling-1-gpu
@@ -46,6 +52,12 @@ steps:
    - tests/v1/e2e/general/
  commands:
    - pytest -v -s v1/e2e/general/test_async_scheduling.py
+  mirror:
+    amd:
+      device: mi250_1
+      timeout_in_minutes: 60
+      depends_on:
+      - image-build-amd

 - label: e2e Core (1 GPU)
  key: e2e-core-1-gpu
@@ -43,6 +43,12 @@ steps:
  commands:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
  - pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py
+  mirror:
+    amd:
+      device: mi300_1
+      timeout_in_minutes: 80
+      depends_on:
+      - image-build-amd

 - label: Entrypoints Integration (API Server openai - Part 2)
  key: entrypoints-integration-api-server-openai-part-2
@@ -55,6 +61,12 @@ steps:
  commands:
  - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
  - pytest -v -s entrypoints/test_chat_utils.py
+  mirror:
+    amd:
+      device: mi300_1
+      timeout_in_minutes: 60
+      depends_on:
+      - image-build-amd

 - label: Entrypoints Integration (API Server openai - Part 3)
  key: entrypoints-integration-api-server-openai-part-3
@@ -67,7 +79,13 @@ steps:
  - tests/entrypoints/test_chat_utils
  commands:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/completion  --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/test_multi_api_servers.py
+  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/test_multi_api_servers.py
+  mirror:
+    amd:
+      device: mi300_1
+      timeout_in_minutes: 60
+      depends_on:
+      - image-build-amd

 - label: Entrypoints Integration (API Server 2)
  key: entrypoints-integration-api-server-2
@@ -83,6 +101,11 @@ steps:
  - pytest -v -s entrypoints/serve/instrumentator
  - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc
  - pytest -v -s tool_use
+  mirror:
+    amd:
+      device: mi300_1
+      depends_on:
+      - image-build-amd

 - label: Entrypoints Integration (Speech to Text)
  key: entrypoints-integration-speech_to_text
@@ -62,6 +62,19 @@ steps:
  commands:
    - pytest -v -s kernels/quantization --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
  parallelism: 2
+  mirror:
+    amd:
+      device: mi300_1
+      source_file_dependencies:
+      - csrc/quantization/
+      - vllm/model_executor/layers/quantization
+      - tests/kernels/quantization
+      - tests/kernels/quantization/test_rocm_skinny_gemms.py
+      - vllm/_aiter_ops.py
+      - vllm/platforms/rocm.py
+      - vllm/model_executor/kernels/
+      depends_on:
+      - image-build-amd

 - label: Kernels MoE Test %N
  key: kernels-moe-test
@@ -91,6 +104,7 @@ steps:

 - label: Kernels KDA Test
  timeout_in_minutes: 20
+  device: h200_18gb
  source_file_dependencies:
  - vllm/model_executor/layers/fla/ops/kda.py
  - vllm/model_executor/layers/fla/ops/chunk_delta_h.py
@@ -105,7 +119,11 @@ steps:
  device: h100
  num_devices: 1
  source_file_dependencies:
+  - cmake/external_projects/deepgemm.cmake
  - tools/install_deepgemm.sh
+  - tools/build_deepgemm_C.py
+  - tools/setup_deepgemm_pythons.sh
+  - tools/check_wheel_deepgemm.py
  - vllm/utils/deep_gemm.py
  - vllm/model_executor/layers/fused_moe
  - vllm/model_executor/layers/quantization
@@ -115,6 +133,7 @@ steps:
  - tests/kernels/attention/test_deepgemm_attention.py
  - tests/quantization/test_cutlass_w4a16.py
  commands:
+    - python3 ../tools/check_wheel_deepgemm.py
    - pytest -v -s kernels/quantization/test_block_fp8.py
    - pytest -v -s kernels/moe/test_deepgemm.py
    - pytest -v -s kernels/moe/test_batched_deepgemm.py
@@ -114,6 +114,7 @@ steps:
 - label: LM Eval TurboQuant KV Cache
  key: lm-eval-turboquant-kv-cache
  timeout_in_minutes: 75
+  device: h200_18gb
  source_file_dependencies:
  - vllm/model_executor/layers/quantization/turboquant/
  - vllm/v1/attention/backends/turboquant_attn.py
@@ -79,6 +79,7 @@ steps:
    - tests/v1/core
    - tests/v1/executor
    - tests/v1/kv_offload
+    - tests/v1/simple_kv_offload
    - tests/v1/worker
    - tests/v1/kv_connector/unit
    - tests/v1/metrics
@@ -90,6 +91,7 @@ steps:
    - pytest -v -s -m 'not cpu_test' v1/core
    - pytest -v -s v1/executor
    - pytest -v -s v1/kv_offload
+    - pytest -v -s v1/simple_kv_offload
    - pytest -v -s v1/worker
    - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit
    - pytest -v -s -m 'not cpu_test' v1/metrics
@@ -127,6 +129,19 @@ steps:
    - pytest -v -s -m 'cpu_test' v1/kv_connector/unit
    - pytest -v -s -m 'cpu_test' v1/metrics

+- label: Extract Hidden States Integration
+  key: extract-hidden-states-integration
+  timeout_in_minutes: 20
+  device: h200_18gb
+  source_file_dependencies:
+    - vllm/v1/spec_decode/extract_hidden_states.py
+    - vllm/model_executor/models/extract_hidden_states.py
+    - vllm/transformers_utils/configs/extract_hidden_states.py
+    - tests/v1/kv_connector/extract_hidden_states_integration
+  commands:
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - pytest -v -s v1/kv_connector/extract_hidden_states_integration
+
 - label: Regression
  key: regression
  timeout_in_minutes: 20
@@ -212,6 +227,7 @@ steps:
  depends_on: ~
  optional: true
  timeout_in_minutes: 20
+  device: h200_18gb
  source_file_dependencies:
  - tests/standalone_tests/python_only_compile.sh
  - setup.py
@@ -302,6 +318,20 @@ steps:
  - pytest -v -s transformers_utils
  - pytest -v -s config

+- label: Batch Invariance (A100)
+  key: batch-invariance-a100
+  timeout_in_minutes: 30
+  device: a100
+  source_file_dependencies:
+    - vllm/v1/attention
+    - vllm/model_executor/layers
+    - tests/v1/determinism/
+  commands:
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - pip install pytest-timeout pytest-forked
+    - pytest -v -s v1/determinism/test_batch_invariance.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]
+
 - label: Batch Invariance (H100)
  key: batch-invariance-h100
  timeout_in_minutes: 30
@@ -106,13 +106,12 @@ steps:
  - vllm/v1/worker/gpu/
  - vllm/v1/worker/gpu_worker.py
  - tests/v1/spec_decode/test_max_len.py
-  - tests/v1/spec_decode/test_probabilistic_rejection_sampler_utils.py
-  - tests/v1/spec_decode/test_synthetic_rejection_sampler_utils.py
+  - tests/v1/spec_decode/test_rejection_sampler_utils.py
  - tests/v1/e2e/spec_decode/test_spec_decode.py
  commands:
  - set -x
  - export VLLM_USE_V2_MODEL_RUNNER=1
  - pytest -v -s v1/spec_decode/test_max_len.py -k "eagle or mtp"
-  - pytest -v -s v1/spec_decode/test_probabilistic_rejection_sampler_utils.py
+  - pytest -v -s v1/spec_decode/test_rejection_sampler_utils.py
  - pytest -v -s v1/spec_decode/test_synthetic_rejection_sampler_utils.py
  - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "eagle or mtp"
@@ -5,6 +5,7 @@ steps:
 - label: Basic Models Tests (Initialization)
  key: basic-models-tests-initialization
  timeout_in_minutes: 45
+  device: h200_18gb
  torch_nightly: true
  source_file_dependencies:
  - vllm/
@@ -5,6 +5,7 @@ steps:
 - label: Language Models Tests (Standard)
  key: language-models-tests-standard
  timeout_in_minutes: 25
+  device: h200_18gb
  source_file_dependencies:
  - vllm/
  - tests/models/language
@@ -91,6 +92,12 @@ steps:
  - tests/models/language/pooling
  commands:
    - pytest -v -s models/language/pooling -m 'not core_model'
+  mirror:
+    amd:
+      device: mi300_1
+      timeout_in_minutes: 100
+      depends_on:
+      - image-build-amd

 - label: Language Models Test (MTEB)
  key: language-models-test-mteb
@@ -84,19 +84,19 @@ steps:
  commands:
    - pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"

- label: DFlash Speculators Correctness
-  key: dflash-speculators-correctness
-  timeout_in_minutes: 30
+- label: Speculators Correctness
+  key: speculators-correctness
+  timeout_in_minutes: 60
  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
+    - tests/v1/spec_decode/test_speculators_correctness.py
  commands:
    - export VLLM_ALLOW_INSECURE_SERIALIZATION=1
-    - pytest -v -s v1/spec_decode/test_speculators_dflash.py -m slow_test
+    - pytest -v -s v1/spec_decode/test_speculators_correctness.py -m slow_test

 - label: Spec Decode MTP hybrid (B200)
  timeout_in_minutes: 30
@@ -9,6 +9,9 @@ PATH=${cuda_home}/bin:$PATH
 LD_LIBRARY_PATH=${cuda_home}/lib64:$LD_LIBRARY_PATH

 # Install requirements
+if [ "$(echo $2 | cut -d. -f1)" = "12" ]; then
+    sed -i 's/^nvidia-cutlass-dsl\[cu13\]>=/nvidia-cutlass-dsl>=/' requirements/cuda.txt
+fi
 $python_executable -m pip install -r requirements/build/cuda.txt -r requirements/cuda.txt

 # Limit the number of parallel jobs to avoid OOM
@@ -369,7 +369,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
  FetchContent_MakeAvailable(cutlass)

  list(APPEND VLLM_EXT_SRC
-    "csrc/quantization/awq/gemm_kernels.cu"
    "csrc/cutlass_extensions/common.cpp")

  set_gencode_flags_for_srcs(
@@ -501,46 +500,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
                   " in CUDA target architectures")
  endif()

-  # Only build AllSpark kernels if we are building for at least some compatible archs.
-  cuda_archs_loose_intersection(ALLSPARK_ARCHS "8.0;8.6;8.7;8.9" "${CUDA_ARCHS}")
-  if (ALLSPARK_ARCHS)
-    set(ALLSPARK_SRCS
-       "csrc/quantization/gptq_allspark/allspark_repack.cu"
-       "csrc/quantization/gptq_allspark/allspark_qgemm_w8a16.cu")
-    set_gencode_flags_for_srcs(
-      SRCS "${ALLSPARK_SRCS}"
-      CUDA_ARCHS "${ALLSPARK_ARCHS}")
-    list(APPEND VLLM_EXT_SRC "${ALLSPARK_SRCS}")
-    message(STATUS "Building AllSpark kernels for archs: ${ALLSPARK_ARCHS}")
-  else()
-    message(STATUS "Not building AllSpark kernels as no compatible archs found"
-                   " in CUDA target architectures")
-  endif()
-
-  # CUTLASS MLA Archs and flags
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(MLA_ARCHS "10.0f;11.0f;12.0f" "${CUDA_ARCHS}")
-  else()
-    cuda_archs_loose_intersection(MLA_ARCHS "10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}")
-  endif()
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND MLA_ARCHS)
-    set(SRCS
-      "csrc/attention/mla/sm100_cutlass_mla_kernel.cu")
-    set_gencode_flags_for_srcs(
-      SRCS "${SRCS}"
-      CUDA_ARCHS "${MLA_ARCHS}")
-    list(APPEND VLLM_EXT_SRC "${SRCS}")
-    list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MLA=1")
-    # Add MLA-specific include directories only to MLA source files
-    set_source_files_properties(${SRCS}
-      PROPERTIES INCLUDE_DIRECTORIES "${CUTLASS_DIR}/examples/77_blackwell_fmha;${CUTLASS_DIR}/examples/common")
-    message(STATUS "Building CUTLASS MLA for archs: ${MLA_ARCHS}")
-  else()
-    message(STATUS "Not building CUTLASS MLA as no compatible archs were found.")
-    # clear MLA_ARCHS
-    set(MLA_ARCHS)
-  endif()
-
  # Expert-specialization MXFP8 blockscaled grouped kernels (SM100+).
  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
    cuda_archs_loose_intersection(ES_MXFP8_GROUPED_MM_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
@@ -568,24 +527,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
    endif()
  endif()

-  # DeepSeek V3 fused A GEMM kernel (requires SM 9.0+, Hopper and later)
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0f;11.0f" "${CUDA_ARCHS}")
-  else()
-    cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
-  endif()
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND DSV3_FUSED_A_GEMM_ARCHS)
-    set(DSV3_FUSED_A_GEMM_SRC "csrc/dsv3_fused_a_gemm.cu")
-    set_gencode_flags_for_srcs(
-      SRCS "${DSV3_FUSED_A_GEMM_SRC}"
-      CUDA_ARCHS "${DSV3_FUSED_A_GEMM_ARCHS}")
-    list(APPEND VLLM_EXT_SRC ${DSV3_FUSED_A_GEMM_SRC})
-    message(STATUS "Building dsv3_fused_a_gemm for archs: ${DSV3_FUSED_A_GEMM_ARCHS}")
-  else()
-    message(STATUS "Not building dsv3_fused_a_gemm as no compatible archs found "
-                   "in CUDA target architectures.")
-  endif()
-
  #
  # Machete kernels

@@ -657,16 +598,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
  endif()


-  # Hadacore kernels
-  cuda_archs_loose_intersection(HADACORE_ARCHS "8.0+PTX;9.0+PTX" "${CUDA_ARCHS}")
-  if(HADACORE_ARCHS)
-    set(SRCS "csrc/quantization/hadamard/hadacore/hadamard_transform_cuda.cu")
-    set_gencode_flags_for_srcs(
-      SRCS "${SRCS}"
-      CUDA_ARCHS "${HADACORE_ARCHS}")
-    list(APPEND VLLM_EXT_SRC "${SRCS}")
-    message(STATUS "Building hadacore")
-  endif()

 # if CUDA endif
 endif()
@@ -716,7 +647,8 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
    list(APPEND VLLM_STABLE_EXT_SRC
      "csrc/libtorch_stable/permute_cols.cu"
      "csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu"
-      "csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu")
+      "csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu"
+      "csrc/libtorch_stable/quantization/awq/gemm_kernels.cu")
  endif()

  if(VLLM_GPU_LANG STREQUAL "CUDA")
@@ -725,6 +657,40 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
      CUDA_ARCHS "${CUDA_ARCHS}")
  endif()

+  # DeepSeek V3 fused A GEMM kernel (requires SM 9.0+, Hopper and later)
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
+    cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0f;11.0f" "${CUDA_ARCHS}")
+  else()
+    cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
+  endif()
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND DSV3_FUSED_A_GEMM_ARCHS)
+    set(SRCS "csrc/libtorch_stable/dsv3_fused_a_gemm.cu")
+    set_gencode_flags_for_srcs(
+      SRCS "${SRCS}"
+      CUDA_ARCHS "${DSV3_FUSED_A_GEMM_ARCHS}")
+    list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
+    message(STATUS "Building dsv3_fused_a_gemm for archs: ${DSV3_FUSED_A_GEMM_ARCHS}")
+  else()
+    message(STATUS "Not building dsv3_fused_a_gemm as no compatible archs found "
+                   "in CUDA target architectures.")
+  endif()
+
+  # Only build AllSpark kernels if we are building for at least some compatible archs.
+  cuda_archs_loose_intersection(ALLSPARK_ARCHS "8.0;8.6;8.7;8.9" "${CUDA_ARCHS}")
+  if (ALLSPARK_ARCHS)
+    set(SRCS
+       "csrc/libtorch_stable/quantization/gptq_allspark/allspark_repack.cu"
+       "csrc/libtorch_stable/quantization/gptq_allspark/allspark_qgemm_w8a16.cu")
+    set_gencode_flags_for_srcs(
+      SRCS "${SRCS}"
+      CUDA_ARCHS "${ALLSPARK_ARCHS}")
+    list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
+    message(STATUS "Building AllSpark kernels for archs: ${ALLSPARK_ARCHS}")
+  else()
+    message(STATUS "Not building AllSpark kernels as no compatible archs found"
+                   " in CUDA target architectures")
+  endif()
+
  #
  # CUTLASS scaled_mm kernels (moved from _C to _C_stable_libtorch)
  #
@@ -1034,6 +1000,41 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
    endif()
  endif()

+  # CUTLASS MLA Archs and flags
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
+    cuda_archs_loose_intersection(MLA_ARCHS "10.0f;11.0f;12.0f" "${CUDA_ARCHS}")
+  else()
+    cuda_archs_loose_intersection(MLA_ARCHS "10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}")
+  endif()
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND MLA_ARCHS)
+    set(SRCS
+      "csrc/libtorch_stable/attention/mla/sm100_cutlass_mla_kernel.cu")
+    set_gencode_flags_for_srcs(
+      SRCS "${SRCS}"
+      CUDA_ARCHS "${MLA_ARCHS}")
+    list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
+    list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MLA=1")
+    # Add MLA-specific include directories only to MLA source files
+    set_source_files_properties(${SRCS}
+      PROPERTIES INCLUDE_DIRECTORIES "${CUTLASS_DIR}/examples/77_blackwell_fmha;${CUTLASS_DIR}/examples/common")
+    message(STATUS "Building CUTLASS MLA for archs: ${MLA_ARCHS}")
+  else()
+    message(STATUS "Not building CUTLASS MLA as no compatible archs were found.")
+    # clear MLA_ARCHS
+    set(MLA_ARCHS)
+  endif()
+
+  # Hadacore kernels
+  cuda_archs_loose_intersection(HADACORE_ARCHS "8.0+PTX;9.0+PTX" "${CUDA_ARCHS}")
+  if(HADACORE_ARCHS)
+    set(SRCS "csrc/libtorch_stable/quantization/hadamard/hadacore/hadamard_transform_cuda.cu")
+    set_gencode_flags_for_srcs(
+      SRCS "${SRCS}"
+      CUDA_ARCHS "${HADACORE_ARCHS}")
+    list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
+    message(STATUS "Building hadacore")
+  endif()
+
  message(STATUS "Enabling C_stable extension.")
  define_extension_target(
    _C_stable_libtorch
@@ -1221,6 +1222,16 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
      CUDA_ARCHS "${DSV3_ROUTER_GEMM_ARCHS}")
    list(APPEND VLLM_MOE_EXT_SRC "${DSV3_ROUTER_GEMM_SRC}")
    message(STATUS "Building DSV3 router GEMM kernel for archs: ${DSV3_ROUTER_GEMM_ARCHS}")
+
+    # DeepSeek V4 fused RMSNorm + router GEMV - same arch gating as DSV3.
+    set(DSV4_NORM_ROUTER_GEMM_SRC
+      "csrc/moe/dsv4_norm_router_gemm_entry.cu"
+      "csrc/moe/dsv4_norm_router_gemm_kernel.cu")
+    set_gencode_flags_for_srcs(
+      SRCS "${DSV4_NORM_ROUTER_GEMM_SRC}"
+      CUDA_ARCHS "${DSV3_ROUTER_GEMM_ARCHS}")
+    list(APPEND VLLM_MOE_EXT_SRC "${DSV4_NORM_ROUTER_GEMM_SRC}")
+    message(STATUS "Building DSV4 norm+router GEMV kernel for archs: ${DSV3_ROUTER_GEMM_ARCHS}")
  else()
    message(STATUS "Not building DSV3 router GEMM kernel as no compatible archs found"
                   " (requires SM90+ and CUDA >= 12.0)")
@@ -53,6 +53,7 @@ backends:
  - FLASHINFER_MLA
  - FLASH_ATTN_MLA  # Hopper only
  - FLASHMLA        # Hopper only
+  - TOKENSPEED_MLA  # Blackwell + R1 dims + FP8 KV (use --kv-cache-dtype fp8)

 device: "cuda:0"
 repeats: 100
@@ -3,6 +3,7 @@
 # Compares all available MLA prefill backends:
 #   FA backends:  fa2, fa3, fa4 (FlashAttention versions)
 #   Non-FA:       flashinfer, cudnn, trtllm (Blackwell-only, require flashinfer)
+#   CuTe DSL:     tokenspeed (Blackwell + R1 dims, requires tokenspeed_mla)
 #
 # Uses cutlass_mla as the decode backend for impl construction
 # (only the prefill path is exercised).
@@ -120,6 +121,7 @@ prefill_backends:
  - flashinfer
  - cudnn
  - trtllm
+  - tokenspeed

 device: "cuda:0"
 repeats: 20
@@ -29,6 +29,7 @@ from vllm.config import (
    VllmConfig,
    set_current_vllm_config,
 )
+from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum

 # ============================================================================
 # VllmConfig Creation
@@ -79,8 +80,8 @@ def create_minimal_vllm_config(
        index_topk: Optional topk value for sparse MLA backends. If provided,
                    the config will include index_topk for sparse attention.
        prefill_backend: Prefill backend name (e.g., "fa3", "fa4", "flashinfer",
-                        "cudnn", "trtllm"). Configures the attention config to
-                        force the specified prefill backend.
+                        "trtllm"). Configures the attention config to force
+                        the specified prefill backend.

    Returns:
        VllmConfig for benchmarking
@@ -179,19 +180,13 @@ def create_minimal_vllm_config(

    if prefill_backend is not None:
        prefill_cfg = get_prefill_backend_config(prefill_backend)
+        vllm_config.attention_config.mla_prefill_backend = prefill_cfg[
+            "mla_prefill_backend"
+        ]
        if prefill_cfg["flash_attn_version"] is not None:
            vllm_config.attention_config.flash_attn_version = prefill_cfg[
                "flash_attn_version"
            ]
-        vllm_config.attention_config.disable_flashinfer_prefill = prefill_cfg[
-            "disable_flashinfer_prefill"
-        ]
-        vllm_config.attention_config.use_cudnn_prefill = prefill_cfg[
-            "use_cudnn_prefill"
-        ]
-        vllm_config.attention_config.use_trtllm_ragged_deepseek_prefill = prefill_cfg[
-            "use_trtllm_ragged_deepseek_prefill"
-        ]

    return vllm_config

@@ -206,39 +201,27 @@ def create_minimal_vllm_config(
 _PREFILL_BACKEND_CONFIG: dict[str, dict] = {
    "fa2": {
        "flash_attn_version": 2,
-        "disable_flashinfer_prefill": True,
-        "use_cudnn_prefill": False,
-        "use_trtllm_ragged_deepseek_prefill": False,
+        "mla_prefill_backend": MLAPrefillBackendEnum.FLASH_ATTN,
    },
    "fa3": {
        "flash_attn_version": 3,
-        "disable_flashinfer_prefill": True,
-        "use_cudnn_prefill": False,
-        "use_trtllm_ragged_deepseek_prefill": False,
+        "mla_prefill_backend": MLAPrefillBackendEnum.FLASH_ATTN,
    },
    "fa4": {
        "flash_attn_version": 4,
-        "disable_flashinfer_prefill": True,
-        "use_cudnn_prefill": False,
-        "use_trtllm_ragged_deepseek_prefill": False,
+        "mla_prefill_backend": MLAPrefillBackendEnum.FLASH_ATTN,
    },
    "flashinfer": {
        "flash_attn_version": None,
-        "disable_flashinfer_prefill": False,
-        "use_cudnn_prefill": False,
-        "use_trtllm_ragged_deepseek_prefill": False,
-    },
-    "cudnn": {
-        "flash_attn_version": None,
-        "disable_flashinfer_prefill": True,
-        "use_cudnn_prefill": True,
-        "use_trtllm_ragged_deepseek_prefill": False,
+        "mla_prefill_backend": MLAPrefillBackendEnum.FLASHINFER,
    },
    "trtllm": {
        "flash_attn_version": None,
-        "disable_flashinfer_prefill": True,
-        "use_cudnn_prefill": False,
-        "use_trtllm_ragged_deepseek_prefill": True,
+        "mla_prefill_backend": MLAPrefillBackendEnum.TRTLLM_RAGGED,
+    },
+    "tokenspeed": {
+        "flash_attn_version": None,
+        "mla_prefill_backend": MLAPrefillBackendEnum.TOKENSPEED_MLA,
    },
 }

@@ -625,6 +608,21 @@ def _create_backend_impl(
    # Create mock layer
    layer = MockLayer(device, impl=impl, kv_cache_spec=kv_cache_spec)

+    # Attach a prefill backend (MLAAttention does this in __init__; the metadata
+    # builder reads layer.prefill_backend from static_forward_context).
+    from vllm.v1.attention.backends.mla.prefill import get_mla_prefill_backend
+
+    prefill_backend_cls = get_mla_prefill_backend(vllm_config)
+    layer.prefill_backend = prefill_backend_cls(
+        num_heads=mla_dims["num_q_heads"],
+        scale=(mla_dims["qk_nope_head_dim"] + mla_dims["qk_rope_head_dim"]) ** -0.5,
+        kv_lora_rank=mla_dims["kv_lora_rank"],
+        qk_nope_head_dim=mla_dims["qk_nope_head_dim"],
+        qk_rope_head_dim=mla_dims["qk_rope_head_dim"],
+        v_head_dim=mla_dims["v_head_dim"],
+        vllm_config=vllm_config,
+    )
+
    # Create builder instance if needed
    builder_instance = None
    if builder_class:
@@ -961,19 +959,6 @@ def _run_mla_benchmark_batched(
    results = []

    with set_current_vllm_config(vllm_config):
-        # Clear cached prefill backend detection functions so they re-evaluate
-        # with the current VllmConfig. These are @functools.cache decorated and
-        # would otherwise return stale results from a previous backend's config.
-        from vllm.model_executor.layers.attention.mla_attention import (
-            use_cudnn_prefill,
-            use_flashinfer_prefill,
-            use_trtllm_ragged_deepseek_prefill,
-        )
-
-        use_flashinfer_prefill.cache_clear()
-        use_cudnn_prefill.cache_clear()
-        use_trtllm_ragged_deepseek_prefill.cache_clear()
-
        # Create backend impl, layer, builder, and indexer (reused across benchmarks)
        impl, layer, builder_instance, indexer = _create_backend_impl(
            backend_cfg,
@@ -985,38 +970,36 @@ def _run_mla_benchmark_batched(
            kv_cache_dtype=kv_cache_dtype,
        )

-        # Verify the actual prefill backend matches what was requested
+        # Verify the actual prefill backend matches what was requested. The
+        # selector + impl construction already raise on misuse; here we just
+        # check the resolved class against the requested name as a sanity guard.
        if prefill_backend is not None:
-            prefill_cfg = get_prefill_backend_config(prefill_backend)
-            fa_version = prefill_cfg["flash_attn_version"]
-
-            if fa_version is not None:
-                # FA backend: verify the impl's FA version
-                actual_fa_version = getattr(impl, "vllm_flash_attn_version", None)
+            expected_class = {
+                "fa2": "FlashAttnPrefillBackend",
+                "fa3": "FlashAttnPrefillBackend",
+                "fa4": "FlashAttnPrefillBackend",
+                "flashinfer": "FlashInferPrefillBackend",
+                "trtllm": "TrtllmRaggedPrefillBackend",
+                "tokenspeed": "TokenspeedMLAPrefillBackend",
+            }.get(prefill_backend)
+            actual_class = type(getattr(layer, "prefill_backend", None)).__name__
+            if expected_class and actual_class != expected_class:
+                raise RuntimeError(
+                    f"Prefill backend '{prefill_backend}' requested "
+                    f"{expected_class}, got {actual_class}. Check "
+                    f"attention_config plumbing or installed deps."
+                )
+            if prefill_backend in {"fa2", "fa3", "fa4"}:
+                fa_version = int(prefill_backend[2:])
+                actual_fa_version = getattr(
+                    layer.prefill_backend, "vllm_flash_attn_version", None
+                )
                if actual_fa_version != fa_version:
                    raise RuntimeError(
                        f"Prefill backend '{prefill_backend}' requested FA "
-                        f"version {fa_version}, but the impl is using FA "
-                        f"version {actual_fa_version}. Check "
-                        f"vllm/v1/attention/backends/fa_utils.py."
+                        f"version {fa_version}, got "
+                        f"{actual_fa_version} on {actual_class}."
                    )
-            else:
-                # Non-FA backend: verify the builder picked the right path
-                expected_flags = {
-                    "flashinfer": "_use_fi_prefill",
-                    "cudnn": "_use_cudnn_prefill",
-                    "trtllm": "_use_trtllm_ragged_prefill",
-                }
-                flag_name = expected_flags.get(prefill_backend)
-                if flag_name and not getattr(builder_instance, flag_name, False):
-                    raise RuntimeError(
-                        f"Prefill backend '{prefill_backend}' was requested "
-                        f"but the metadata builder did not enable it. This "
-                        f"usually means a dependency is missing (e.g., "
-                        f"flashinfer not installed) or the platform doesn't "
-                        f"support it."
-                    )
-
        # Run each benchmark with the shared impl
        for config, threshold, num_splits in configs_with_params:
            # Set threshold for this benchmark (FlashAttn/FlashMLA only)
@@ -27,10 +27,10 @@ from vllm.model_executor.layers.fused_moe.config import (
    RoutingMethodType,
    _get_config_dtype_str,
 )
-from vllm.model_executor.layers.fused_moe.fused_moe import *
-from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import (
+from vllm.model_executor.layers.fused_moe.experts.triton_deep_gemm_moe import (
    TritonOrDeepGemmExperts,
 )
+from vllm.model_executor.layers.fused_moe.fused_moe import *
 from vllm.transformers_utils.config import get_config
 from vllm.triton_utils import triton
 from vllm.utils.argparse_utils import FlexibleArgumentParser
@@ -0,0 +1,183 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Benchmark and correctness check for ``ops.dsv4_norm_router_gemm``.
+
+Two implementations are compared:
+
+  1. ``unfused``   — ``vllm_ops.rms_norm`` then ``ops.dsv3_router_gemm``,
+                     i.e. the current vLLM hot path (two kernel launches).
+  2. ``fused``     — ``ops.dsv4_norm_router_gemm``, the new single-kernel
+                     fused path.
+
+Both produce ``(normed_x: bf16, router_logits: fp32)``.  The correctness
+check verifies that ``fused`` and ``unfused`` agree to within ~1 bf16
+ULP — that is the precision floor for this op.
+"""
+
+import argparse
+
+import torch
+
+from vllm import _custom_ops as vllm_ops
+from vllm.triton_utils import triton
+
+# The fused dsv4_norm_router_gemm kernel is templated only for DSV4-Pro
+# (hidden_size=7168, num_experts=384).  Other shapes fall back to the
+# unfused path on the Python side (NormGatedLinear), so benchmark only
+# the configuration that the fused kernel actually targets.
+HIDDEN_SIZE = 7168
+NUM_EXPERTS_CHOICES = (384,)
+RMS_EPS = 1e-6
+
+
+def unfused_norm_router_gemm(
+    x: torch.Tensor,
+    norm_weight: torch.Tensor,
+    gate_weight: torch.Tensor,
+    eps: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    # Call ``_C::rms_norm`` directly (mirroring ``_dsv4_pro_norm_gate``'s
+    # fallback path) so the benchmarked baseline doesn't inherit any
+    # Python wrapper overhead or risk falling through to the native
+    # eager-primitive ``RMSNorm.forward_native`` path.
+    normed = torch.empty_like(x)
+    torch.ops._C.rms_norm(normed, x, norm_weight, eps)
+    logits = vllm_ops.dsv3_router_gemm(normed, gate_weight, torch.float32)
+    return normed, logits
+
+
+def fused_norm_router_gemm(
+    x: torch.Tensor,
+    norm_weight: torch.Tensor,
+    gate_weight: torch.Tensor,
+    eps: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    return vllm_ops.dsv4_norm_router_gemm(x, norm_weight, gate_weight, eps)
+
+
+def _make_inputs(num_tokens: int, num_experts: int, hidden_size: int, seed: int = 0):
+    torch.manual_seed(seed)
+    device = "cuda"
+    x = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device=device)
+    norm_w = torch.randn(hidden_size, dtype=torch.bfloat16, device=device)
+    gate_w = torch.randn(num_experts, hidden_size, dtype=torch.bfloat16, device=device)
+    # Down-scale gate_w so the GEMV output stays in a representable range.
+    gate_w = gate_w / float(hidden_size) ** 0.5
+    norm_w = (norm_w * 0.1) + 1.0
+    return x, norm_w, gate_w
+
+
+def calculate_diff(
+    num_tokens: int,
+    num_experts: int,
+    hidden_size: int = HIDDEN_SIZE,
+    normed_atol: float = 2e-3,
+    logits_atol: float = 1e-2,
+    rtol: float = 1e-2,
+) -> None:
+    x, norm_w, gate_w = _make_inputs(num_tokens, num_experts, hidden_size)
+
+    normed_unfused, logits_unfused = unfused_norm_router_gemm(
+        x.clone(), norm_w, gate_w, RMS_EPS
+    )
+    normed_fused, logits_fused = fused_norm_router_gemm(
+        x.clone(), norm_w, gate_w, RMS_EPS
+    )
+
+    def _max_abs(a, b):
+        return (a.float() - b.float()).abs().max().item()
+
+    print(f"\n=== M={num_tokens} E={num_experts} H={hidden_size} ===")
+    print(f"normed_x  |fused - unfused| = {_max_abs(normed_fused, normed_unfused):.3e}")
+    print(f"logits    |fused - unfused| = {_max_abs(logits_fused, logits_unfused):.3e}")
+
+    ok_normed = torch.allclose(
+        normed_fused.float(),
+        normed_unfused.float(),
+        atol=normed_atol,
+        rtol=rtol,
+    )
+    ok_logits = torch.allclose(
+        logits_fused.float(),
+        logits_unfused.float(),
+        atol=logits_atol,
+        rtol=rtol,
+    )
+    if ok_normed and ok_logits:
+        print(
+            f"OK   fused vs unfused within "
+            f"normed_atol={normed_atol:.0e} logits_atol={logits_atol:.0e} "
+            f"rtol={rtol:.0e}"
+        )
+    else:
+        print(
+            f"FAIL normed_ok={ok_normed} logits_ok={ok_logits}; "
+            f"see max-abs values above"
+        )
+
+
+def get_benchmark():
+    # Only num_tokens varies (DSV4-Pro hard-codes E=384); single-axis
+    # sweep yields a clean line plot with M on the x-axis.
+    num_experts = NUM_EXPERTS_CHOICES[0]
+
+    @triton.testing.perf_report(
+        triton.testing.Benchmark(
+            x_names=["num_tokens"],
+            x_vals=list(range(1, 17)),
+            line_arg="provider",
+            line_vals=["unfused", "fused"],
+            line_names=["unfused (rms+dsv3)", "fused (dsv4)"],
+            styles=[("green", "-"), ("red", "-")],
+            ylabel="us",
+            plot_name=f"norm-router-gemm-E{num_experts}-H{HIDDEN_SIZE}",
+            args={},
+        )
+    )
+    def benchmark(num_tokens, provider):
+        x, norm_w, gate_w = _make_inputs(num_tokens, num_experts, HIDDEN_SIZE)
+
+        quantiles = [0.5, 0.2, 0.8]
+        if provider == "unfused":
+            fn = lambda: unfused_norm_router_gemm(  # noqa: E731
+                x, norm_w, gate_w, RMS_EPS
+            )
+        else:
+            fn = lambda: fused_norm_router_gemm(  # noqa: E731
+                x, norm_w, gate_w, RMS_EPS
+            )
+
+        ms, min_ms, max_ms = triton.testing.do_bench(fn, quantiles=quantiles)
+        return 1000 * ms, 1000 * max_ms, 1000 * min_ms
+
+    return benchmark
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--save-path",
+        type=str,
+        default="./configs/norm_router_gemm/",
+    )
+    parser.add_argument(
+        "--skip-bench",
+        action="store_true",
+        help="Run only the correctness check, not the perf sweep.",
+    )
+    args = parser.parse_args()
+
+    # Correctness sweep over the full fast-path range M=1..16.
+    for m in range(1, 17):
+        for e in NUM_EXPERTS_CHOICES:
+            calculate_diff(num_tokens=m, num_experts=e, hidden_size=HIDDEN_SIZE)
+
+    if args.skip_bench:
+        return
+
+    benchmark = get_benchmark()
+    benchmark.run(print_data=True, save_path=args.save_path)
+
+
+if __name__ == "__main__":
+    main()
@@ -12,7 +12,6 @@ from vllm._custom_ops import (
    cpu_attn_get_scheduler_metadata,
    cpu_attn_reshape_and_cache,
 )
-from vllm.platforms import CpuArchEnum, current_platform
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, set_random_seed
 from vllm.v1.attention.backends.cpu_attn import CPUAttentionBackend, _get_attn_isa
@@ -22,15 +21,14 @@ def get_attn_isa(
    block_size: int | None = None,
    dtype: torch.dtype | None = None,
 ):
-    if block_size and dtype:
-        return _get_attn_isa(dtype, block_size)
-    else:
-        if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
-            return "neon"
-        elif torch.cpu._is_amx_tile_supported():
-            return "amx"
-        else:
-            return "vec"
+    # Delegate to _get_attn_isa so the fallback path applies the same arch
+    # gating (e.g. RISC-V RVV is only chosen when the build's hardcoded
+    # VLEN=128 kernel is actually present; on VLEN=256 / scalar hosts it
+    # correctly falls through to vec/vec16).
+    return _get_attn_isa(
+        dtype if dtype is not None else torch.bfloat16,
+        block_size if block_size else 32,
+    )


 # rand number generation takes too much time, cache rand tensors
@@ -235,7 +233,7 @@ if __name__ == "__main__":
    )
    parser.add_argument("--use-sink", action="store_true")
    parser.add_argument(
-        "--isa", type=str, choices=["vec", "neon", "amx", "vec16"], default=None
+        "--isa", type=str, choices=["vec", "neon", "amx", "vec16", "rvv"], default=None
    )
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--iters", type=int, default=20)
@@ -195,10 +195,12 @@ elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64")
    endif()
    if(VLLM_RVV_VLEN AND VLLM_RVV_VLEN GREATER 0)
        message(STATUS "RISC-V RVV VLEN=${VLLM_RVV_VLEN}")
+        # Sources gate FP16/BF16 paths on the compiler-provided
+        # __riscv_zvfh / __riscv_zvfbfmin macros, which GCC and clang
+        # define automatically when those extensions appear in -march.
        if(RVV_BF16_FOUND)
            message(STATUS "BF16 extension detected")
            set(MARCH_FLAGS -march=rv64gcv_zvfh_zfbfmin_zvfbfmin_zvl${VLLM_RVV_VLEN}b -mrvv-vector-bits=zvl -mabi=lp64d)
-            add_compile_definitions(RISCV_BF16_SUPPORT)
        elseif(RVV_FP16_FOUND)
            message(WARNING "BF16 functionality is not available")
            set(MARCH_FLAGS -march=rv64gcv_zvfh_zvl${VLLM_RVV_VLEN}b -mrvv-vector-bits=zvl -mabi=lp64d)
@@ -53,48 +53,80 @@ cuda_archs_loose_intersection(DEEPGEMM_ARCHS
 if(DEEPGEMM_ARCHS)
  message(STATUS "DeepGEMM CUDA architectures: ${DEEPGEMM_ARCHS}")

-  find_package(CUDAToolkit REQUIRED)
-
  #
-  # Build the _C pybind11 extension from DeepGEMM's C++ source.
-  # This is a CXX-only module — CUDA kernels are JIT-compiled at runtime.
+  # DeepGEMM integration notes
+  # --------------------------
+  # We vendor DeepGEMM into vllm/third_party/deep_gemm/ and bundle a
+  # `_C.cpython-X.Y-*.so` for every CPython in `requires-python`. The
+  # per-Python build is delegated to tools/build_deepgemm_C.py.
+  #
+  # Why per-Python: DeepGEMM's binding uses PYBIND11_MODULE, which links
+  # private CPython symbols — a single `_C.abi3.so` is not viable today
+  # (see #41476 / #41512 for the failed attempt).
+  #
+  # TODOs (tracked in vllm-project/vllm#42431):
+  #   - Replace DeepGEMM's pybind11 binding with a TORCH_LIBRARY + shim
+  #     binding (cf. vllm-flash-attention/csrc/common/pytorch_shim.h) to
+  #     collapse to one `_C.abi3.so`. Needs either an upstream change or
+  #     a maintained binding fork in vLLM.
+  #   - AOT-compile DeepGEMM's CUDA kernels instead of runtime JIT to drop
+  #     the vendored CUTLASS/CCCL headers and the CUDA-toolkit-at-runtime
+  #     requirement.
  #
-  Python_add_library(_deep_gemm_C MODULE WITH_SOABI
-    "${deepgemm_SOURCE_DIR}/csrc/python_api.cpp")

-  # The pybind11 module name must be _C to match DeepGEMM's Python imports.
-  set_target_properties(_deep_gemm_C PROPERTIES OUTPUT_NAME "_C")
+  # DEEPGEMM_PYTHON_INTERPRETERS: ":"-separated target Python paths.
+  # Empty/unset → fall back to the build interpreter (editable installs).
+  # (Empty-but-set env vars test as DEFINED in cmake — treat as unset.)
+  if(NOT "$ENV{DEEPGEMM_PYTHON_INTERPRETERS}" STREQUAL "")
+    string(REPLACE ":" ";" _dg_pythons "$ENV{DEEPGEMM_PYTHON_INTERPRETERS}")
+  else()
+    set(_dg_pythons "${Python_EXECUTABLE}")
+  endif()
+  message(STATUS "DeepGEMM _C will be built for: ${_dg_pythons}")

-  target_compile_definitions(_deep_gemm_C PRIVATE
-    "-DTORCH_EXTENSION_NAME=_C")
+  # add_custom_command does no implicit header scanning; glob explicitly so
+  # header-only edits in DeepGEMM/cutlass/fmt re-trigger the rebuild.
+  file(GLOB_RECURSE _dg_headers
+    "${deepgemm_SOURCE_DIR}/csrc/*.h"
+    "${deepgemm_SOURCE_DIR}/csrc/*.hpp"
+    "${deepgemm_SOURCE_DIR}/deep_gemm/include/*.h"
+    "${deepgemm_SOURCE_DIR}/deep_gemm/include/*.hpp"
+    "${deepgemm_SOURCE_DIR}/deep_gemm/include/*.cuh")

-  target_include_directories(_deep_gemm_C PRIVATE
-    "${deepgemm_SOURCE_DIR}/csrc"
-    "${deepgemm_SOURCE_DIR}/deep_gemm/include"
-    "${deepgemm_SOURCE_DIR}/third-party/cutlass/include"
-    "${deepgemm_SOURCE_DIR}/third-party/cutlass/tools/util/include"
-    "${deepgemm_SOURCE_DIR}/third-party/fmt/include")
-
-  target_compile_options(_deep_gemm_C PRIVATE
-    $<$<COMPILE_LANGUAGE:CXX>:-O3>
-    $<$<COMPILE_LANGUAGE:CXX>:-Wno-psabi>
-    $<$<COMPILE_LANGUAGE:CXX>:-Wno-deprecated-declarations>)
-
-  # torch_python is required because DeepGEMM uses pybind11 type casters
-  # for at::Tensor (via PYBIND11_MODULE), unlike vLLM's own extensions which
-  # use torch::Library custom ops.
-  find_library(TORCH_PYTHON_LIBRARY torch_python
-    PATHS "${TORCH_INSTALL_PREFIX}/lib"
-    REQUIRED)
-
-  target_link_libraries(_deep_gemm_C PRIVATE
-    torch ${TORCH_LIBRARIES} "${TORCH_PYTHON_LIBRARY}"
-    CUDA::cudart CUDA::nvrtc)
-
-  # Install the shared library into the vendored package directory
-  install(TARGETS _deep_gemm_C
-    LIBRARY DESTINATION vllm/third_party/deep_gemm
-    COMPONENT _deep_gemm_C)
+  set(_dg_markers)
+  set(_dg_seen_soabis)
+  foreach(_pybin IN LISTS _dg_pythons)
+    execute_process(
+      COMMAND "${_pybin}" -c
+        "import sysconfig; print(sysconfig.get_config_var('SOABI'))"
+      OUTPUT_VARIABLE _dg_soabi
+      OUTPUT_STRIP_TRAILING_WHITESPACE
+      COMMAND_ERROR_IS_FATAL ANY)
+    # Dedup interpreters that resolve to the same CPython.
+    if(_dg_soabi IN_LIST _dg_seen_soabis)
+      continue()
+    endif()
+    list(APPEND _dg_seen_soabis "${_dg_soabi}")
+    set(_dg_dir "${CMAKE_CURRENT_BINARY_DIR}/deepgemm_C_${_dg_soabi}")
+    set(_dg_marker "${_dg_dir}/.built")
+    add_custom_command(
+      OUTPUT "${_dg_marker}"
+      COMMAND "${Python_EXECUTABLE}"
+              "${CMAKE_SOURCE_DIR}/tools/build_deepgemm_C.py"
+              "${deepgemm_SOURCE_DIR}" "${_dg_dir}" "${_pybin}"
+      COMMAND "${CMAKE_COMMAND}" -E touch "${_dg_marker}"
+      DEPENDS "${CMAKE_SOURCE_DIR}/tools/build_deepgemm_C.py"
+              "${deepgemm_SOURCE_DIR}/csrc/python_api.cpp"
+              ${_dg_headers}
+      COMMENT "Building DeepGEMM _C for ${_pybin}"
+      VERBATIM)
+    list(APPEND _dg_markers "${_dg_marker}")
+    install(DIRECTORY "${_dg_dir}/"
+      DESTINATION vllm/third_party/deep_gemm
+      COMPONENT _deep_gemm_C
+      FILES_MATCHING PATTERN "_C.cpython-*.so")
+  endforeach()
+  add_custom_target(_deep_gemm_C ALL DEPENDS ${_dg_markers})

  #
  # Vendor DeepGEMM Python package files
@@ -39,7 +39,7 @@ else()
  FetchContent_Declare(
          vllm-flash-attn
          GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
-          GIT_TAG f5bc33cfc02c744d24a2e9d50e6db656de40611c
+          GIT_TAG bce29425653ec0fbc579d329883030e832d15ada
          GIT_PROGRESS TRUE
          # Don't share the vllm-flash-attn build between build types
          BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
@@ -1,7 +1,13 @@
 #pragma once

-// For TORCH_CHECK
-#include <torch/library.h>
+#include <cstdint>
+#include <string>
+#include <tuple>
+#include <utility>
+#include <variant>
+
+// For STD_TORCH_CHECK
+#include <torch/headeronly/util/Exception.h>

 namespace vllm {

@@ -45,7 +51,7 @@ class ScalarType {
  // IEEE 754 compliant floating point type
  static constexpr ScalarType float_IEEE754(uint8_t exponent,
                                            uint8_t mantissa) {
-    TORCH_CHECK(mantissa > 0 && exponent > 0);
+    STD_TORCH_CHECK(mantissa > 0 && exponent > 0);
    return ScalarType(exponent, mantissa, true, 0, false, NAN_IEEE_754);
  }

@@ -53,11 +59,12 @@ class ScalarType {
  static constexpr ScalarType float_(uint8_t exponent, uint8_t mantissa,
                                     bool finite_values_only,
                                     NanRepr nan_repr) {
-    TORCH_CHECK(nan_repr < NAN_REPR_ID_MAX, "Invalid NanRepr");
-    TORCH_CHECK(mantissa > 0 && exponent > 0);
-    TORCH_CHECK(nan_repr != NAN_IEEE_754,
-                "use `float_IEEE754` constructor for floating point types that "
-                "follow IEEE 754 conventions");
+    STD_TORCH_CHECK(nan_repr < NAN_REPR_ID_MAX, "Invalid NanRepr");
+    STD_TORCH_CHECK(mantissa > 0 && exponent > 0);
+    STD_TORCH_CHECK(
+        nan_repr != NAN_IEEE_754,
+        "use `float_IEEE754` constructor for floating point types that "
+        "follow IEEE 754 conventions");
    return ScalarType(exponent, mantissa, true, 0, finite_values_only,
                      nan_repr);
  }
@@ -176,8 +183,8 @@ class ScalarType {

 private:
  double _floating_point_max() const {
-    TORCH_CHECK(mantissa <= 52 && exponent <= 11,
-                "Cannot represent max/min as a double for type ", str());
+    STD_TORCH_CHECK(mantissa <= 52 && exponent <= 11,
+                    "Cannot represent max/min as a double for type ", str());

    uint64_t max_mantissa = (uint64_t(1) << mantissa) - 1;
    if (nan_repr == NAN_EXTD_RANGE_MAX_MIN) {
@@ -186,8 +193,8 @@ class ScalarType {

    uint64_t max_exponent = (uint64_t(1) << exponent) - 2;
    if (nan_repr == NAN_EXTD_RANGE_MAX_MIN || nan_repr == NAN_NONE) {
-      TORCH_CHECK(exponent < 11,
-                  "Cannot represent max/min as a double for type ", str());
+      STD_TORCH_CHECK(exponent < 11,
+                      "Cannot represent max/min as a double for type ", str());
      max_exponent += 1;
    }

@@ -216,16 +223,17 @@ class ScalarType {
    if (is_floating_point()) {
      return {_floating_point_max()};
    } else {
-      TORCH_CHECK(size_bits() < 64 || size_bits() == 64 && is_signed(),
-                  "Cannot represent max as a int64_t");
+      STD_TORCH_CHECK(size_bits() < 64 || size_bits() == 64 && is_signed(),
+                      "Cannot represent max as a int64_t");
      return {(int64_t(1) << mantissa) - 1};
    }
  }

  constexpr std::variant<int64_t, double> _raw_min() const {
    if (is_floating_point()) {
-      TORCH_CHECK(is_signed(),
-                  "We currently assume all floating point types are signed");
+      STD_TORCH_CHECK(
+          is_signed(),
+          "We currently assume all floating point types are signed");
      constexpr uint64_t sign_bit_double = (uint64_t(1) << 63);

      double max = _floating_point_max();
@@ -233,8 +241,8 @@ class ScalarType {
      uint64_t min_raw = max_raw | sign_bit_double;
      return {*reinterpret_cast<double*>(&min_raw)};
    } else {
-      TORCH_CHECK(!is_signed() || size_bits() <= 64,
-                  "Cannot represent min as a int64_t");
+      STD_TORCH_CHECK(!is_signed() || size_bits() <= 64,
+                      "Cannot represent min as a int64_t");
      if (is_signed()) {
        // set the top bit to 1 (i.e. INT64_MIN) and the rest to 0
        // then perform an arithmetic shift right to set all the bits above
@@ -156,4 +156,20 @@

 #endif  // __aarch64__

+// RISC-V RVV
+#ifdef __riscv_v
+  #include <riscv_vector.h>
+
+  #ifdef __riscv_zihintpause
+    #define FAST_SPINNING __riscv_pause();
+  #endif
+
+  // FP32Vec16::exp() in cpu_types_riscv.hpp already implements the full
+  // polynomial approximation for RVV, so we simply delegate to it.
+  #define DEFINE_FAST_EXP                             \
+    auto fast_exp = [&](const vec_op::FP32Vec16& vec) \
+                        __attribute__((always_inline)) { return vec.exp(); };
+
+#endif  // __riscv_v
+
 #endif
@@ -29,6 +29,8 @@ torch::Tensor get_scheduler_metadata(
    isa = cpu_attention::ISA::NEON;
  } else if (isa_hint == "vxe") {
    isa = cpu_attention::ISA::VXE;
+  } else if (isa_hint == "rvv") {
+    isa = cpu_attention::ISA::RVV;
  } else if (isa_hint == "vsx") {
    isa = cpu_attention::ISA::VSX;
  } else {
@@ -131,6 +133,8 @@ void cpu_attn_reshape_and_cache(
      return cpu_attention::ISA::NEON;
    } else if (isa == "vxe") {
      return cpu_attention::ISA::VXE;
+    } else if (isa == "rvv") {
+      return cpu_attention::ISA::RVV;
    } else if (isa == "vsx") {
      return cpu_attention::ISA::VSX;
    } else {
@@ -12,7 +12,7 @@
 #include "cpu/utils.hpp"

 namespace cpu_attention {
-enum class ISA { AMX, VEC, VEC16, NEON, VXE, VSX };
+enum class ISA { AMX, VEC, VEC16, NEON, VXE, RVV, VSX };

 // Mirrors csrc/attention/dtype_fp8.cuh Fp8KVCacheDataType exactly.
 enum class Fp8KVCacheDataType {
@@ -164,6 +164,9 @@ struct AttentionMetadata {
      case ISA::VXE:
        ss << "VXE, ";
        break;
+      case ISA::RVV:
+        ss << "RVV, ";
+        break;
      case ISA::VSX:
        ss << "VSX, ";
        break;
@@ -0,0 +1,445 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+#ifndef CPU_ATTN_RVV_HPP
+#define CPU_ATTN_RVV_HPP
+
+// This kernel is currently hardcoded to VLEN=128 (m1/m2 intrinsics, vl=8).
+// The fixed-width typedefs below use `riscv_rvv_vector_bits(128)`, which
+// only matches `vfloat16m1_t`/`vuint16m1_t` register layout when VLEN==128;
+// at VLEN>=256 those typedefs fail to compile.  Scalar RISC-V builds
+// (-march=rv64gc) additionally don't have <riscv_vector.h>.  For both
+// cases we omit the file entirely and let the dispatcher fall back to the
+// scalar VEC / VEC16 implementations.  TODO: migrate to RVVI() macros +
+// semantic names in cpu_types_riscv_defs.hpp to support VLEN>=256 natively.
+#if defined(__riscv_v_min_vlen) && __riscv_v_min_vlen == 128
+
+  #include "cpu_attn_impl.hpp"
+  #include <riscv_vector.h>
+  #include <type_traits>
+
+namespace cpu_attention {
+
+namespace {
+
+// File-local concrete-LMUL typedefs.  The shared _defs.hpp exposes
+// VLEN-independent semantic names (fixed_fp32x8_t, fixed_fp16x8_t, ...),
+// but this kernel is currently hardcoded to VLEN=128 (m1/m2 intrinsics),
+// so keep the legacy concrete aliases scoped to this file.
+typedef vfloat16m1_t fixed_vfloat16m1_t
+    __attribute__((riscv_rvv_vector_bits(128)));
+typedef vfloat32m2_t fixed_vfloat32m2_t
+    __attribute__((riscv_rvv_vector_bits(256)));
+typedef vuint16m1_t fixed_vuint16m1_t
+    __attribute__((riscv_rvv_vector_bits(128)));
+typedef vuint32m2_t fixed_vuint32m2_t
+    __attribute__((riscv_rvv_vector_bits(256)));
+  #ifdef __riscv_zvfbfmin
+typedef vbfloat16m1_t fixed_vbfloat16m1_t
+    __attribute__((riscv_rvv_vector_bits(128)));
+  #endif
+
+  #define BLOCK_SIZE_ALIGNMENT 32
+  #define HEAD_SIZE_ALIGNMENT 32
+  #define MAX_Q_HEAD_NUM_PER_ITER 16
+
+// ============================================================================
+// B-matrix row loading: load 8 elements as FP32 (using m2 LMUL at VLEN=128)
+// ============================================================================
+
+template <typename kv_cache_t>
+FORCE_INLINE fixed_vfloat32m2_t load_row8_B_as_f32(const kv_cache_t* p);
+
+template <>
+FORCE_INLINE fixed_vfloat32m2_t load_row8_B_as_f32<float>(const float* p) {
+  return __riscv_vle32_v_f32m2(p, 8);
+}
+
+template <>
+FORCE_INLINE fixed_vfloat32m2_t
+load_row8_B_as_f32<c10::Half>(const c10::Half* p) {
+  #ifdef __riscv_zvfh
+  fixed_vfloat16m1_t h =
+      __riscv_vle16_v_f16m1(reinterpret_cast<const _Float16*>(p), 8);
+  return __riscv_vfwcvt_f_f_v_f32m2(h, 8);
+  #else
+  // Fallback for hardware without Zvfh: scalar half->float conversion.
+  // c10::Half provides operator float() so this is correct on any RVV CPU
+  // that has only the base V extension.  Slower than the Zvfh path, but
+  // keeps the kernel buildable on Zvfhmin-only / no-fp16 hardware.
+  alignas(16) float tmp[8];
+  for (int i = 0; i < 8; ++i) {
+    tmp[i] = static_cast<float>(p[i]);
+  }
+  return __riscv_vle32_v_f32m2(tmp, 8);
+  #endif
+}
+
+template <>
+FORCE_INLINE fixed_vfloat32m2_t
+load_row8_B_as_f32<c10::BFloat16>(const c10::BFloat16* p) {
+  #ifdef __riscv_zvfbfmin
+  fixed_vbfloat16m1_t bf =
+      __riscv_vle16_v_bf16m1(reinterpret_cast<const __bf16*>(p), 8);
+  return __riscv_vfwcvtbf16_f_f_v_f32m2(bf, 8);
+  #else
+  // Fallback: load as uint16, zero-extend to uint32, shift left by 16
+  fixed_vuint16m1_t raw =
+      __riscv_vle16_v_u16m1(reinterpret_cast<const uint16_t*>(p), 8);
+  fixed_vuint32m2_t wide = __riscv_vzext_vf2_u32m2(raw, 8);
+  fixed_vuint32m2_t shifted = __riscv_vsll_vx_u32m2(wide, 16, 8);
+  return __riscv_vreinterpret_v_u32m2_f32m2(shifted);
+  #endif
+}
+
+// ============================================================================
+// Micro kernel: Mx8 tile, K unrolled by 4, RVV scalar-broadcast FMA
+// ============================================================================
+//
+// NEON uses vfmaq_laneq_f32 (lane-indexed FMA from a preloaded A vector).
+// RVV has no lane-indexed FMA; instead we load A elements as scalars and
+// use __riscv_vfmacc_vf (scalar * vector + accumulator), which is equally
+// efficient and avoids the need for vrgather/vslidedown.
+//
+// At VLEN=128, m2 holds 8 x FP32, matching the 8-column tile width.
+// Register budget: M accumulators (m2 each) + 1 B temp = 2M+2 regs.
+// M=8 => 18 regs out of 32 available — no spills.
+
+template <int32_t M, typename kv_cache_t>
+FORCE_INLINE void gemm_micro_rvv_fma_Mx8_Ku4(
+    const float* __restrict A,       // [M x K]
+    const kv_cache_t* __restrict B,  // [K x 8]
+    float* __restrict C,             // [M x 8]
+    int64_t lda, int64_t ldb, int64_t ldc, int32_t K, bool accumulate) {
+  static_assert(1 <= M && M <= 8, "M must be in [1,8]");
+
+  constexpr size_t vl = 8;
+
+  // helpers for per-M codegen
+  #define ROWS_APPLY(OP) OP(0) OP(1) OP(2) OP(3) OP(4) OP(5) OP(6) OP(7)
+  #define IF_M(i) if constexpr (M > (i))
+
+    // A row base pointers
+  #define DECL_A(i) const float* a##i = A + (i) * lda;
+  ROWS_APPLY(DECL_A)
+  #undef DECL_A
+
+    // declare one m2 accumulator per row
+  #define DECL_ACC(i) fixed_vfloat32m2_t acc##i;
+  ROWS_APPLY(DECL_ACC)
+  #undef DECL_ACC
+
+    // initialize accumulators
+  #define INIT_ACC(i)                                      \
+    IF_M(i) {                                              \
+      if (accumulate) {                                    \
+        acc##i = __riscv_vle32_v_f32m2(C + (i) * ldc, vl); \
+      } else {                                             \
+        acc##i = __riscv_vfmv_v_f_f32m2(0.f, vl);          \
+      }                                                    \
+    }
+  ROWS_APPLY(INIT_ACC)
+  #undef INIT_ACC
+
+  int32_t k = 0;
+
+  // K unrolled by 4
+  for (; k + 3 < K; k += 4) {
+    // k + 0
+    {
+      fixed_vfloat32m2_t b =
+          load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 0) * ldb);
+  #define STEP_K0(i)                                                    \
+    IF_M(i) {                                                           \
+      acc##i = __riscv_vfmacc_vf_f32m2(acc##i, *(a##i + k + 0), b, vl); \
+    }
+      ROWS_APPLY(STEP_K0)
+  #undef STEP_K0
+    }
+    // k + 1
+    {
+      fixed_vfloat32m2_t b =
+          load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 1) * ldb);
+  #define STEP_K1(i)                                                    \
+    IF_M(i) {                                                           \
+      acc##i = __riscv_vfmacc_vf_f32m2(acc##i, *(a##i + k + 1), b, vl); \
+    }
+      ROWS_APPLY(STEP_K1)
+  #undef STEP_K1
+    }
+    // k + 2
+    {
+      fixed_vfloat32m2_t b =
+          load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 2) * ldb);
+  #define STEP_K2(i)                                                    \
+    IF_M(i) {                                                           \
+      acc##i = __riscv_vfmacc_vf_f32m2(acc##i, *(a##i + k + 2), b, vl); \
+    }
+      ROWS_APPLY(STEP_K2)
+  #undef STEP_K2
+    }
+    // k + 3
+    {
+      fixed_vfloat32m2_t b =
+          load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 3) * ldb);
+  #define STEP_K3(i)                                                    \
+    IF_M(i) {                                                           \
+      acc##i = __riscv_vfmacc_vf_f32m2(acc##i, *(a##i + k + 3), b, vl); \
+    }
+      ROWS_APPLY(STEP_K3)
+  #undef STEP_K3
+    }
+  }
+
+  // K tail
+  for (; k < K; ++k) {
+    fixed_vfloat32m2_t b = load_row8_B_as_f32<kv_cache_t>(B + (int64_t)k * ldb);
+  #define TAIL_ROW(i) \
+    IF_M(i) { acc##i = __riscv_vfmacc_vf_f32m2(acc##i, *(a##i + k), b, vl); }
+    ROWS_APPLY(TAIL_ROW)
+  #undef TAIL_ROW
+  }
+
+    // store accumulators to C
+  #define STORE_ROW(i) \
+    IF_M(i) { __riscv_vse32_v_f32m2(C + (i) * ldc, acc##i, vl); }
+  ROWS_APPLY(STORE_ROW)
+  #undef STORE_ROW
+
+  #undef ROWS_APPLY
+  #undef IF_M
+}
+
+// ============================================================================
+// Macro kernel: dispatch M tiles of {8,4,2,1}, step N by 8
+// ============================================================================
+
+template <int32_t N, typename kv_cache_t>
+FORCE_INLINE void gemm_macro_rvv_fma_Mx8_Ku4(const float* __restrict A,
+                                             const kv_cache_t* __restrict B,
+                                             float* __restrict C, int32_t M,
+                                             int32_t K, int64_t lda,
+                                             int64_t ldb, int64_t ldc,
+                                             bool accumulate) {
+  static_assert(N % 8 == 0, "N must be a multiple of 8");
+  for (int32_t m = 0; m < M;) {
+    int32_t mb = (M - m >= 8) ? 8 : (M - m >= 4) ? 4 : (M - m >= 2) ? 2 : 1;
+    const float* Ab = A + m * lda;
+    float* Cb = C + m * ldc;
+
+    for (int32_t n = 0; n < N; n += 8) {
+      const kv_cache_t* Bn = B + n;
+      float* Cn = Cb + n;
+      switch (mb) {
+        case 8:
+          gemm_micro_rvv_fma_Mx8_Ku4<8, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
+                                                    K, accumulate);
+          break;
+        case 4:
+          gemm_micro_rvv_fma_Mx8_Ku4<4, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
+                                                    K, accumulate);
+          break;
+        case 2:
+          gemm_micro_rvv_fma_Mx8_Ku4<2, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
+                                                    K, accumulate);
+          break;
+        default:
+          gemm_micro_rvv_fma_Mx8_Ku4<1, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
+                                                    K, accumulate);
+          break;
+      }
+    }
+    m += mb;
+  }
+}
+
+// ============================================================================
+// TileGemm wrapper — plugs into AttentionMainLoop
+// ============================================================================
+
+template <typename kv_cache_t>
+class TileGemmRVV {
+ public:
+  template <AttentionGemmPhase phase, int32_t k_size>
+  FORCE_INLINE static void gemm(const int32_t m_size,
+                                float* __restrict__ a_tile,
+                                kv_cache_t* __restrict__ b_tile,
+                                float* __restrict__ c_tile, const int64_t lda,
+                                const int64_t ldb, const int64_t ldc,
+                                const int32_t block_size,
+                                const int32_t dynamic_k_size,
+                                const bool accum_c) {
+    if constexpr (phase == AttentionGemmPhase::QK) {
+      gemm_macro_rvv_fma_Mx8_Ku4<BLOCK_SIZE_ALIGNMENT, kv_cache_t>(
+          a_tile, b_tile, c_tile, m_size, k_size, lda, ldb, ldc, accum_c);
+    } else {
+      gemm_macro_rvv_fma_Mx8_Ku4<HEAD_SIZE_ALIGNMENT, kv_cache_t>(
+          a_tile, b_tile, c_tile, m_size, dynamic_k_size, lda, ldb, ldc,
+          accum_c);
+    }
+  }
+};
+
+}  // namespace
+
+// ============================================================================
+// AttentionImpl<ISA::RVV> — mirrors ISA::NEON specialization
+// ============================================================================
+
+template <typename scalar_t, int64_t head_dim, typename kv_cache_scalar_t>
+class AttentionImpl<ISA::RVV, scalar_t, head_dim, kv_cache_scalar_t> {
+ public:
+  using query_t = scalar_t;
+  using q_buffer_t = float;
+  using kv_cache_t = scalar_t;
+  using logits_buffer_t = float;
+  using partial_output_buffer_t = float;
+  using prob_buffer_t = float;
+
+  constexpr static int64_t BlockSizeAlignment = BLOCK_SIZE_ALIGNMENT;
+  constexpr static int64_t HeadDimAlignment = HEAD_SIZE_ALIGNMENT;
+  constexpr static int64_t MaxQHeadNumPerIteration = MAX_Q_HEAD_NUM_PER_ITER;
+  constexpr static int64_t HeadDim = head_dim;
+  constexpr static ISA ISAType = ISA::RVV;
+  constexpr static bool scale_on_logits = false;
+
+  static_assert(HeadDim % HeadDimAlignment == 0);
+  static_assert(HeadDimAlignment % 8 == 0);
+  static_assert(BlockSizeAlignment % 8 == 0);
+
+ public:
+  template <template <typename tile_gemm_t> typename attention>
+  FORCE_INLINE void execute_attention(DEFINE_CPU_ATTENTION_PARAMS) {
+    attention<TileGemmRVV<kv_cache_t>> attention_iteration;
+    attention_iteration(CPU_ATTENTION_PARAMS);
+  }
+
+  constexpr static int64_t k_cache_token_group_stride(
+      const int32_t block_size) {
+    return BlockSizeAlignment;
+  }
+
+  constexpr static int64_t v_cache_token_group_stride(
+      const int32_t block_size) {
+    return head_dim * BlockSizeAlignment;
+  }
+
+  constexpr static int64_t v_cache_head_group_stride(const int32_t block_size) {
+    return HeadDimAlignment;
+  }
+
+  static void copy_q_heads_tile(scalar_t* __restrict__ src,
+                                float* __restrict__ q_buffer,
+                                const int32_t q_num,
+                                const int32_t q_heads_per_kv,
+                                const int64_t q_num_stride,
+                                const int64_t q_head_stride, float scale) {
+    static_assert(head_dim % 16 == 0);
+    constexpr int32_t unroll_size = head_dim / 16;
+    using load_vec_t = typename VecTypeTrait<scalar_t>::vec_t;
+
+    vec_op::FP32Vec16 scale_vec(scale);
+    for (int32_t q_num_idx = 0; q_num_idx < q_num; ++q_num_idx) {
+      for (int32_t q_head_idx = 0; q_head_idx < q_heads_per_kv; ++q_head_idx) {
+        scalar_t* __restrict__ curr_q =
+            src + q_num_idx * q_num_stride + q_head_idx * q_head_stride;
+        float* __restrict__ curr_q_buffer =
+            q_buffer + q_num_idx * q_heads_per_kv * head_dim +
+            q_head_idx * head_dim;
+
+        vec_op::unroll_loop<int32_t, unroll_size>([&](int32_t i) {
+          load_vec_t vec(curr_q);
+          vec_op::FP32Vec16 fp32_vec(vec);
+          fp32_vec = fp32_vec * scale_vec;
+          fp32_vec.save(curr_q_buffer);
+
+          curr_q += 16;
+          curr_q_buffer += 16;
+        });
+      }
+    }
+  }
+
+  static void reshape_and_cache(
+      const scalar_t* __restrict__ key, const scalar_t* __restrict__ value,
+      scalar_t* __restrict__ key_cache, scalar_t* __restrict__ value_cache,
+      const int64_t* __restrict__ slot_mapping, const int64_t token_num,
+      const int64_t key_token_num_stride, const int64_t value_token_num_stride,
+      const int64_t head_num, const int64_t key_head_num_stride,
+      const int64_t value_head_num_stride, const int64_t num_blocks,
+      const int64_t num_blocks_stride, const int64_t cache_head_num_stride,
+      const int64_t block_size, const int64_t block_size_stride,
+      const float /*k_inv*/ = 0.0f, const float /*v_inv*/ = 0.0f) {
+  #pragma omp parallel for collapse(2)
+    for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
+      for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
+        const int64_t pos = slot_mapping[token_idx];
+        if (pos < 0) {
+          continue;
+        }
+
+        const int64_t block_idx = pos / block_size;
+        const int64_t block_offset = pos % block_size;
+        {
+          // Write Key (transpose to column-major: [head_dim, block_size])
+          const scalar_t* key_start_ptr = key +
+                                          token_idx * key_token_num_stride +
+                                          head_idx * key_head_num_stride;
+          scalar_t* key_cache_start_ptr =
+              key_cache + block_idx * num_blocks_stride +
+              head_idx * cache_head_num_stride + block_offset;
+
+          // Strided vector store for efficient transpose.
+          // Load contiguous key elements, store with stride = block_size.
+          {
+            const ptrdiff_t byte_stride = block_size * sizeof(scalar_t);
+            int64_t i = 0;
+            for (; i < head_dim;) {
+              size_t vl;
+              if constexpr (std::is_same_v<scalar_t, float>) {
+                vl = __riscv_vsetvl_e32m2(head_dim - i);
+                vfloat32m2_t v = __riscv_vle32_v_f32m2(
+                    reinterpret_cast<const float*>(key_start_ptr + i), vl);
+                __riscv_vsse32_v_f32m2(
+                    reinterpret_cast<float*>(key_cache_start_ptr +
+                                             i * block_size),
+                    byte_stride, v, vl);
+              } else {
+                // Half and BFloat16 are both 16-bit types
+                vl = __riscv_vsetvl_e16m1(head_dim - i);
+                vuint16m1_t v = __riscv_vle16_v_u16m1(
+                    reinterpret_cast<const uint16_t*>(key_start_ptr + i), vl);
+                __riscv_vsse16_v_u16m1(
+                    reinterpret_cast<uint16_t*>(key_cache_start_ptr +
+                                                i * block_size),
+                    byte_stride, v, vl);
+              }
+              i += vl;
+            }
+          }
+        }
+        {
+          // Write Value (row-major: [block_size, head_dim])
+          const scalar_t* value_start_ptr = value +
+                                            token_idx * value_token_num_stride +
+                                            head_idx * value_head_num_stride;
+          scalar_t* value_cache_start_ptr =
+              value_cache + block_idx * num_blocks_stride +
+              head_idx * cache_head_num_stride + block_offset * head_dim;
+          std::memcpy(value_cache_start_ptr, value_start_ptr,
+                      sizeof(scalar_t) * head_dim);
+        }
+      }
+    }
+  }
+};
+
+}  // namespace cpu_attention
+
+  #undef BLOCK_SIZE_ALIGNMENT
+  #undef HEAD_SIZE_ALIGNMENT
+  #undef MAX_Q_HEAD_NUM_PER_ITER
+
+#endif  // __riscv_v_min_vlen == 128
+
+#endif  // CPU_ATTN_RVV_HPP
@@ -72,7 +72,7 @@ typedef RVVTYPE(vuint16, LMUL_512, _t) fixed_u16x32_t
    __attribute__((riscv_rvv_vector_bits(512)));

 // bfloat16
-#ifdef RISCV_BF16_SUPPORT
+#ifdef __riscv_zvfbfmin
 typedef RVVTYPE(vbfloat16, LMUL_128, _t) fixed_bf16x8_t
    __attribute__((riscv_rvv_vector_bits(128)));
 typedef RVVTYPE(vbfloat16, LMUL_256, _t) fixed_bf16x16_t
@@ -15,8 +15,17 @@
 #include <torch/all.h>
 namespace vec_op {

-// BFloat16 is always supported on RISC-V: natively when RISCV_BF16_SUPPORT
-// is defined, otherwise via the FP32-simulation fallback path.
+// FP8 KV cache is not supported on RISC-V. These tag types and the
+// corresponding BF16Vec32 stub constructors below exist solely so that
+// templates referencing vec_op::fp8_*_tag in their bodies (e.g. in
+// cpu_attn_vec.hpp) compile under GCC's -Wtemplate-body lookup. The
+// stubs are never instantiated by CPU_ATTN_DISPATCH on __riscv.
+struct fp8_e4m3_tag {};
+struct fp8_e5m2_tag {};
+
+// BFloat16 is always supported on RISC-V: natively when __riscv_zvfbfmin
+// is defined (compiler-provided when -march includes zvfbfmin), otherwise
+// via the FP32-simulation fallback path.
 #define VLLM_DISPATCH_CASE_FLOATING_TYPES(...)         \
  AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
  AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__)  \
@@ -106,7 +115,7 @@ struct FP16Vec16 : public Vec<FP16Vec16> {
 // BF16 Implementation
 // ============================================================================

-#ifdef RISCV_BF16_SUPPORT
+#ifdef __riscv_zvfbfmin

 FORCE_INLINE fixed_u16x8_t bf16_to_u16(fixed_bf16x8_t v) {
  return RVVI4(__riscv_vreinterpret_v_bf16, LMUL_128, _u16, LMUL_128)(v);
@@ -185,6 +194,13 @@ struct BF16Vec32 : public Vec<BF16Vec32> {

  explicit BF16Vec32(fixed_bf16x32_t data) : reg(data) {};

+  // FP8 KV cache stubs: never instantiated on RISC-V (CPU_ATTN_DISPATCH
+  // omits FP8 cases on __riscv); exist only so name lookup succeeds.
+  explicit BF16Vec32(const uint8_t* ptr, fp8_e4m3_tag)
+      : BF16Vec32(static_cast<const void*>(ptr)) {}
+  explicit BF16Vec32(const uint8_t* ptr, fp8_e5m2_tag)
+      : BF16Vec32(static_cast<const void*>(ptr)) {}
+
  explicit BF16Vec32(const BF16Vec8& v) {
    fixed_u16x8_t u16_val = bf16_to_u16(v.reg);
    fixed_u16x32_t u16_combined =
@@ -323,6 +339,13 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
    reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_1024)(tmp, 32);
  }

+  // FP8 KV cache stubs: never instantiated on RISC-V (CPU_ATTN_DISPATCH
+  // omits FP8 cases on __riscv); exist only so name lookup succeeds.
+  explicit BF16Vec32(const uint8_t* ptr, fp8_e4m3_tag)
+      : BF16Vec32(static_cast<const void*>(ptr)) {}
+  explicit BF16Vec32(const uint8_t* ptr, fp8_e5m2_tag)
+      : BF16Vec32(static_cast<const void*>(ptr)) {}
+
  explicit BF16Vec32(const BF16Vec8& v) {
    float tmp_small[8];
    RVVI(__riscv_vse32_v_f32, LMUL_256)(tmp_small, v.reg_fp32, 8);
@@ -410,7 +433,7 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
  explicit FP32Vec8(fixed_fp16x8_t v)
      : reg(RVVI(__riscv_vfwcvt_f_f_v_f32, LMUL_256)(v, VEC_ELEM_NUM)) {};

-#ifdef RISCV_BF16_SUPPORT
+#ifdef __riscv_zvfbfmin
  explicit FP32Vec8(fixed_bf16x8_t v)
      : reg(RVVI(__riscv_vfwcvtbf16_f_f_v_f32, LMUL_256)(v, VEC_ELEM_NUM)) {};
  explicit FP32Vec8(const BF16Vec8& v)
@@ -608,7 +631,7 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
  explicit FP32Vec16(const FP32Vec16& data) : reg(data.reg) {};
  explicit FP32Vec16(const FP16Vec16& v);

-#ifdef RISCV_BF16_SUPPORT
+#ifdef __riscv_zvfbfmin
  explicit FP32Vec16(fixed_bf16x16_t v)
      : reg(RVVI(__riscv_vfwcvtbf16_f_f_v_f32, LMUL_512)(v, VEC_ELEM_NUM)) {};
  explicit FP32Vec16(const BF16Vec16& v)
@@ -868,7 +891,7 @@ inline void fma(FP32Vec16& acc, const FP32Vec16& a, const FP32Vec16& b) {
  acc = acc.fma(a, b);
 }

-#ifdef RISCV_BF16_SUPPORT
+#ifdef __riscv_zvfbfmin
 template <>
 inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
  *ptr = static_cast<__bf16>(v);
@@ -20,7 +20,8 @@ ISA_TYPES = {
    "VEC16": 2,
    "NEON": 3,
    "VXE": 4,
-    "VSX": 5,
+    "RVV": 5,
+    "VSX": 6,
 }

 # KV cache index: 0 = auto (same as scalar_t), 1 = fp8_e4m3, 2 = fp8_e5m2
@@ -38,7 +39,7 @@ KV_CACHE_CPP_TYPES = {
 }

 # ISAs supported for head_dims divisible by 32
-ISA_FOR_32 = ["AMX", "NEON", "VEC", "VEC16", "VXE", "VSX"]
+ISA_FOR_32 = ["AMX", "NEON", "VEC", "VEC16", "VXE", "RVV", "VSX"]

 # ISAs supported for head_dims divisible by 16 only
 ISA_FOR_16 = ["VEC16"]
@@ -149,6 +150,15 @@ def generate_header_file() -> str:
  #include "cpu_attn_vxe.hpp"
 #endif

+// cpu_attn_rvv.hpp is hardcoded to VLEN==128 (m1/m2 intrinsics, vl=8) and
+// itself includes <riscv_vector.h>, which is unavailable on scalar
+// (-march=rv64gc) builds. Gate the include the same way as the dispatch
+// macro below, so non-128 / scalar RISC-V builds skip it entirely.
+#if defined(__riscv) && defined(__riscv_v_min_vlen) && \
+    __riscv_v_min_vlen == 128
+  #include "cpu_attn_rvv.hpp"
+#endif
+
 #ifdef __powerpc__
  #include "cpu_attn_vsx.hpp"
 #endif
@@ -212,6 +222,23 @@ def generate_header_file() -> str:
        ["VXE", "VEC", "VEC16"],
        fp8=False,
    )
+    # RISC-V with RVV.  cpu_attn_rvv.hpp is hardcoded to VLEN==128
+    # (riscv_rvv_vector_bits(128) typedefs + vl=8 m1/m2 intrinsics), so
+    # we split the dispatch into two top-level branches: VLEN==128 builds
+    # get the full RVV+VEC+VEC16 case set, other VLEN builds get a
+    # VEC/VEC16-only fallback.  Preprocessor directives cannot appear
+    # inside a #define body, so this duplication is necessary.
+    header += _macro_block(
+        "#elif defined(__riscv) && defined(__riscv_v_min_vlen) "
+        "&& __riscv_v_min_vlen == 128",
+        ["RVV", "VEC", "VEC16"],
+        fp8=False,
+    )
+    header += _macro_block(
+        "#elif defined(__riscv)",
+        ["VEC", "VEC16"],
+        fp8=False,
+    )
    header += _macro_block(
        "#elif defined(__powerpc__)",
        ["VSX", "VEC", "VEC16"],
@@ -233,8 +260,8 @@ def generate_header_file() -> str:
        fp8=False,
    )
    header += (
-        "#endif  /* CPU_CAPABILITY_AMXBF16 / __aarch64__ / "
-        "__s390x__ / __powerpc__ */\n\n"
+        "#endif  /* CPU_CAPABILITY_AMXBF16 / __aarch64__ / __s390x__ /"
+        " __riscv / __powerpc__ */\n\n"
        "#endif  // CPU_ATTN_DISPATCH_GENERATED_H\n"
    )

@@ -18,13 +18,12 @@ limitations under the License.
 * Taken from SGLANG PR https://github.com/sgl-project/sglang/pull/6929
 * by Alcanderian JieXin Liang
 */
-#include "core/registration.h"
+#include "libtorch_stable/torch_utils.h"
+
+#include <torch/csrc/stable/library.h>

-#include <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAGuard.h>
 #include <cutlass/cutlass.h>
 #include <cutlass/kernel_hardware_info.h>
-#include <torch/all.h>

 #include <cute/tensor.hpp>
 #include <iostream>
@@ -35,27 +34,27 @@ limitations under the License.
 // clang-format off
 #if !defined(CUDA_VERSION) || CUDA_VERSION < 12040
 void sm100_cutlass_mla_decode(
-    torch::Tensor const& out,
-    torch::Tensor const& lse,
-    torch::Tensor const& q_nope,
-    torch::Tensor const& q_pe,
-    torch::Tensor const& kv_c_and_k_pe_cache,
-    torch::Tensor const& seq_lens,
-    torch::Tensor const& page_table,
-    torch::Tensor const& workspace,
+    torch::stable::Tensor const& out,
+    torch::stable::Tensor const& lse,
+    torch::stable::Tensor const& q_nope,
+    torch::stable::Tensor const& q_pe,
+    torch::stable::Tensor const& kv_c_and_k_pe_cache,
+    torch::stable::Tensor const& seq_lens,
+    torch::stable::Tensor const& page_table,
+    torch::stable::Tensor const& workspace,
    double sm_scale,
    int64_t num_kv_splits) {
-  TORCH_CHECK(false, "CUDA version must be >= 12.4 for cutlass_mla_decode");
+  STD_TORCH_CHECK(false, "CUDA version must be >= 12.4 for cutlass_mla_decode");
 }
 int64_t sm100_cutlass_mla_get_workspace_size(int64_t max_seq_len, int64_t num_batches, int64_t sm_count, int64_t num_kv_splits) {
-  TORCH_CHECK(false, "CUDA version must be >= 12.4 for cutlass_mla_get_workspace_size");
+  STD_TORCH_CHECK(false, "CUDA version must be >= 12.4 for cutlass_mla_get_workspace_size");
 }
 #else

 #define CUTLASS_CHECK(status)                                                       \
  {                                                                                 \
    cutlass::Status error = status;                                                 \
-    TORCH_CHECK(error == cutlass::Status::kSuccess, cutlassGetStatusString(error)); \
+    STD_TORCH_CHECK(error == cutlass::Status::kSuccess, cutlassGetStatusString(error)); \
  }

 using namespace cute;
@@ -100,23 +99,23 @@ struct MlaSm100 {

 template <typename T>
 typename T::Fmha::Arguments args_from_options(
-    at::Tensor const& out,
-    at::Tensor const& lse,
-    at::Tensor const& q_nope,
-    at::Tensor const& q_pe,
-    at::Tensor const& kv_c_and_k_pe_cache,
-    at::Tensor const& seq_lens,
-    at::Tensor const& page_table,
+    torch::stable::Tensor const& out,
+    torch::stable::Tensor const& lse,
+    torch::stable::Tensor const& q_nope,
+    torch::stable::Tensor const& q_pe,
+    torch::stable::Tensor const& kv_c_and_k_pe_cache,
+    torch::stable::Tensor const& seq_lens,
+    torch::stable::Tensor const& page_table,
    double sm_scale,
    int64_t num_kv_splits) {
  cutlass::KernelHardwareInfo hw_info;
-  hw_info.device_id = q_nope.device().index();
+  hw_info.device_id = q_nope.get_device_index();
  hw_info.sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id);

-  int batches = q_nope.sizes()[0];
-  int page_count_per_seq = page_table.sizes()[1];
-  int page_count_total = kv_c_and_k_pe_cache.sizes()[0];
-  int page_size = kv_c_and_k_pe_cache.sizes()[1];
+  int batches = q_nope.size(0);
+  int page_count_per_seq = page_table.size(1);
+  int page_count_total = kv_c_and_k_pe_cache.size(0);
+  int page_size = kv_c_and_k_pe_cache.size(1);
  int max_seq_len = page_size * page_count_per_seq;
  using TileShapeH = typename T::TileShapeH;
  using TileShapeD = typename T::TileShapeD;
@@ -186,14 +185,14 @@ typename T::Fmha::Arguments args_from_options(

 template <typename Element, typename ElementOut, bool IsPaged128, typename PersistenceOption>
 void runMla(
-    at::Tensor const& out,
-    at::Tensor const& lse,
-    at::Tensor const& q_nope,
-    at::Tensor const& q_pe,
-    at::Tensor const& kv_c_and_k_pe_cache,
-    at::Tensor const& seq_lens,
-    at::Tensor const& page_table,
-    at::Tensor const& workspace,
+    torch::stable::Tensor const& out,
+    torch::stable::Tensor const& lse,
+    torch::stable::Tensor const& q_nope,
+    torch::stable::Tensor const& q_pe,
+    torch::stable::Tensor const& kv_c_and_k_pe_cache,
+    torch::stable::Tensor const& seq_lens,
+    torch::stable::Tensor const& page_table,
+    torch::stable::Tensor const& workspace,
    double sm_scale,
    int64_t num_kv_splits,
    cudaStream_t stream) {
@@ -220,37 +219,37 @@ void runMla(
  }()

 void sm100_cutlass_mla_decode(
-    torch::Tensor const& out,
-    torch::Tensor const& lse,
-    torch::Tensor const& q_nope,
-    torch::Tensor const& q_pe,
-    torch::Tensor const& kv_c_and_k_pe_cache,
-    torch::Tensor const& seq_lens,
-    torch::Tensor const& page_table,
-    torch::Tensor const& workspace,
+    torch::stable::Tensor const& out,
+    torch::stable::Tensor const& lse,
+    torch::stable::Tensor const& q_nope,
+    torch::stable::Tensor const& q_pe,
+    torch::stable::Tensor const& kv_c_and_k_pe_cache,
+    torch::stable::Tensor const& seq_lens,
+    torch::stable::Tensor const& page_table,
+    torch::stable::Tensor const& workspace,
    double sm_scale,
    int64_t num_kv_splits) {
-  auto in_dtype = q_nope.dtype();
-  at::cuda::CUDAGuard device_guard{(char)q_nope.get_device()};
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream(q_nope.get_device());
-  const int page_size = kv_c_and_k_pe_cache.sizes()[1];
-  
+  auto in_dtype = q_nope.scalar_type();
+  torch::stable::accelerator::DeviceGuard device_guard(q_nope.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream(q_nope.get_device_index());
+  const int page_size = kv_c_and_k_pe_cache.size(1);
+
  // NOTE(alcanderian): IsPersistent has bug with manual split_kv.
  // Kernel will hang if batch is too large with large num_kv_splits. (for example bs=8, num_kv_splits=8)
  // Maybe per batch split kv will fix this.
  DISPATCH_BOOL(page_size == 128, IsPaged128, [&] {
    DISPATCH_BOOL(num_kv_splits <= 1, NotManualSplitKV, [&] {
-      if (in_dtype == at::ScalarType::Half) {
+      if (in_dtype == torch::headeronly::ScalarType::Half) {
        runMla<cutlass::half_t, cutlass::half_t, IsPaged128, IsPersistent<NotManualSplitKV>>(
          out, lse, q_nope, q_pe, kv_c_and_k_pe_cache, seq_lens, page_table, workspace, sm_scale, num_kv_splits, stream);
-      } else if (in_dtype == at::ScalarType::BFloat16) {
+      } else if (in_dtype == torch::headeronly::ScalarType::BFloat16) {
        runMla<cutlass::bfloat16_t, cutlass::bfloat16_t, IsPaged128, IsPersistent<NotManualSplitKV>>(
          out, lse, q_nope, q_pe, kv_c_and_k_pe_cache, seq_lens, page_table, workspace, sm_scale, num_kv_splits, stream);
-      } else if (in_dtype == at::ScalarType::Float8_e4m3fn) {
+      } else if (in_dtype == torch::headeronly::ScalarType::Float8_e4m3fn) {
        runMla<cutlass::float_e4m3_t, cutlass::bfloat16_t, IsPaged128, IsPersistent<NotManualSplitKV>>(
          out, lse, q_nope, q_pe, kv_c_and_k_pe_cache, seq_lens, page_table, workspace, sm_scale, num_kv_splits, stream);
      } else {
-        TORCH_CHECK(false, "Unsupported input data type of MLA");
+        STD_TORCH_CHECK(false, "Unsupported input data type of MLA");
      }
      return true;
    });
@@ -280,12 +279,12 @@ int64_t sm100_cutlass_mla_get_workspace_size(int64_t max_seq_len, int64_t num_ba

 #endif

-TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
-  m.impl("sm100_cutlass_mla_decode", &sm100_cutlass_mla_decode);
+STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
+  m.impl("sm100_cutlass_mla_decode", TORCH_BOX(&sm100_cutlass_mla_decode));
 }

-TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CatchAll, m) {
-  m.impl("sm100_cutlass_mla_get_workspace_size", &sm100_cutlass_mla_get_workspace_size);
+STABLE_TORCH_LIBRARY_IMPL(_C, CompositeExplicitAutograd, m) {
+  m.impl("sm100_cutlass_mla_get_workspace_size", TORCH_BOX(&sm100_cutlass_mla_get_workspace_size));
 }

 // clang-format on
@@ -20,13 +20,15 @@
 * limitations under the License.
 */

-#include <ATen/ATen.h>
-#include <ATen/cuda/CUDAContext.h>
-#include <cuda_bf16.h>
-#include <cuda_runtime.h>
-#include <torch/all.h>
+#include <torch/csrc/stable/library.h>
+#include <torch/csrc/stable/tensor.h>
+#include <torch/headeronly/core/ScalarType.h>

 #include "core/registration.h"
+#include "libtorch_stable/torch_utils.h"
+
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>

 #include <cstdlib>
 #include <mutex>
@@ -34,7 +36,7 @@
 namespace {

 inline int getSMVersion() {
-  auto* props = at::cuda::getCurrentDeviceProperties();
+  auto* props = get_device_prop();
  return props->major * 10 + props->minor;
 }

@@ -700,37 +702,40 @@ template void invokeFusedAGemm<__nv_bfloat16, 7168, 2112, 16>(
    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens,
    cudaStream_t);

-void dsv3_fused_a_gemm(torch::Tensor& output, torch::Tensor const& mat_a,
-                       torch::Tensor const& mat_b) {
-  TORCH_CHECK(mat_a.dim() == 2 && mat_b.dim() == 2 && output.dim() == 2);
+void dsv3_fused_a_gemm(torch::stable::Tensor& output,
+                       torch::stable::Tensor const& mat_a,
+                       torch::stable::Tensor const& mat_b) {
+  STD_TORCH_CHECK(mat_a.dim() == 2 && mat_b.dim() == 2 && output.dim() == 2);
  int const num_tokens = mat_a.size(0);
  int const hd_in = mat_a.size(1);
  int const hd_out = mat_b.size(1);

  constexpr int kHdIn = 7168;
  constexpr int kHdOut = 2112;
-  TORCH_CHECK(num_tokens >= 1 && num_tokens <= 16,
-              "required 1 <= mat_a.shape[0] <= 16")
-  TORCH_CHECK(hd_in == kHdIn, "required mat_a.shape[1] == 7168")
-  TORCH_CHECK(hd_out == kHdOut, "required mat_b.shape[1] == 2112")
-  TORCH_CHECK(output.size(0) == num_tokens,
-              "required output.shape[0] == mat_a.shape[0]")
-  TORCH_CHECK(output.size(1) == hd_out,
-              "required output.shape[1] == mat_b.shape[1]")
+  STD_TORCH_CHECK(num_tokens >= 1 && num_tokens <= 16,
+                  "required 1 <= mat_a.shape[0] <= 16");
+  STD_TORCH_CHECK(hd_in == kHdIn, "required mat_a.shape[1] == 7168");
+  STD_TORCH_CHECK(hd_out == kHdOut, "required mat_b.shape[1] == 2112");
+  STD_TORCH_CHECK(output.size(0) == num_tokens,
+                  "required output.shape[0] == mat_a.shape[0]");
+  STD_TORCH_CHECK(output.size(1) == hd_out,
+                  "required output.shape[1] == mat_b.shape[1]");

-  TORCH_CHECK(mat_a.stride(1) == 1, "mat_a must be a row major tensor");
-  TORCH_CHECK(output.stride(1) == 1, "output must be a row major tensor");
-  TORCH_CHECK(mat_b.stride(0) == 1, "mat_b must be a column major tensor");
+  STD_TORCH_CHECK(mat_a.stride(1) == 1, "mat_a must be a row major tensor");
+  STD_TORCH_CHECK(output.stride(1) == 1, "output must be a row major tensor");
+  STD_TORCH_CHECK(mat_b.stride(0) == 1, "mat_b must be a column major tensor");

-  TORCH_CHECK(mat_a.scalar_type() == torch::kBFloat16 &&
-                  mat_b.scalar_type() == torch::kBFloat16,
-              "Only BFloat16 input dtype is supported")
-  TORCH_CHECK(output.scalar_type() == torch::kBFloat16,
-              "Only BFloat16 output dtype is supported")
+  STD_TORCH_CHECK(
+      mat_a.scalar_type() == torch::headeronly::ScalarType::BFloat16 &&
+          mat_b.scalar_type() == torch::headeronly::ScalarType::BFloat16,
+      "Only BFloat16 input dtype is supported");
+  STD_TORCH_CHECK(
+      output.scalar_type() == torch::headeronly::ScalarType::BFloat16,
+      "Only BFloat16 output dtype is supported");

-  TORCH_CHECK(getSMVersion() >= 90, "required CUDA ARCH >= SM_90");
+  STD_TORCH_CHECK(getSMVersion() >= 90, "required CUDA ARCH >= SM_90");

-  auto stream = at::cuda::getCurrentCUDAStream(mat_a.get_device());
+  auto stream = get_current_cuda_stream(mat_a.get_device_index());
  if (num_tokens <= 8) {
    invokeFusedAGemm<__nv_bfloat16, kHdIn, kHdOut, 8>(
        reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()),
@@ -746,6 +751,6 @@ void dsv3_fused_a_gemm(torch::Tensor& output, torch::Tensor const& mat_a,
  }
 }

-TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
-  m.impl("dsv3_fused_a_gemm", &dsv3_fused_a_gemm);
+STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
+  m.impl("dsv3_fused_a_gemm", TORCH_BOX(&dsv3_fused_a_gemm));
 }
@@ -143,4 +143,26 @@ void cutlass_mxfp4_group_mm(torch::stable::Tensor& output,
                            const torch::stable::Tensor& expert_offsets,
                            const torch::stable::Tensor& sf_offsets);

+// AWQ ops
+torch::stable::Tensor awq_gemm(torch::stable::Tensor _in_feats,
+                               torch::stable::Tensor _kernel,
+                               torch::stable::Tensor _scaling_factors,
+                               torch::stable::Tensor _zeros,
+                               int64_t split_k_iters);
+
+torch::stable::Tensor awq_dequantize(torch::stable::Tensor _kernel,
+                                     torch::stable::Tensor _scaling_factors,
+                                     torch::stable::Tensor _zeros,
+                                     int64_t split_k_iters, int64_t thx,
+                                     int64_t thy);
+
+// DSV3 fused A GEMM: conditionally compiled so declaration and impl
+// registration are in the source file (dsv3_fused_a_gemm.cu)
+
+// AllSpark ops: declarations are in the source files
+// (allspark_repack.cu and allspark_qgemm_w8a16.cu)
+
 #endif
+
+torch::stable::Tensor hadacore_transform(torch::stable::Tensor& x,
+                                         bool inplace);
@@ -7,10 +7,11 @@ Shang and Dang, Xingyu and Han, Song}, journal={arXiv}, year={2023}
 }
 */

-#include <torch/all.h>
-#include <c10/cuda/CUDAGuard.h>
+#include <torch/csrc/stable/tensor.h>
+#include <torch/csrc/stable/ops.h>
+#include "libtorch_stable/torch_utils.h"

-#include "dequantize.cuh"
+#include "libtorch_stable/quantization/awq/dequantize.cuh"

 #include <cuda_fp16.h>

@@ -410,10 +411,11 @@ __global__ void __launch_bounds__(64)
 }  // namespace awq
 }  // namespace vllm

-torch::Tensor awq_dequantize(torch::Tensor _kernel,
-                             torch::Tensor _scaling_factors,
-                             torch::Tensor _zeros, int64_t split_k_iters,
-                             int64_t thx, int64_t thy) {
+torch::stable::Tensor awq_dequantize(torch::stable::Tensor _kernel,
+                                     torch::stable::Tensor _scaling_factors,
+                                     torch::stable::Tensor _zeros,
+                                     int64_t split_k_iters, int64_t thx,
+                                     int64_t thy) {
  int in_c = _kernel.size(0);
  int qout_c = _kernel.size(1);
  int out_c = qout_c * 8;
@@ -437,23 +439,24 @@ torch::Tensor awq_dequantize(torch::Tensor _kernel,
    y_blocks = (int)(in_c / 8);
  }

-  const at::cuda::OptionalCUDAGuard device_guard(device_of(_scaling_factors));
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      _scaling_factors.get_device_index());

-  auto options = torch::TensorOptions()
-                     .dtype(_scaling_factors.dtype())
-                     .device(_scaling_factors.device());
-  at::Tensor _de_kernel = torch::empty({in_c, out_c}, options);
+  auto _de_kernel =
+      torch::stable::empty({in_c, out_c}, _scaling_factors.scalar_type(),
+                           std::nullopt, _scaling_factors.device());

-  auto kernel = reinterpret_cast<int*>(_kernel.data_ptr<int>());
-  auto de_kernel = reinterpret_cast<half*>(_de_kernel.data_ptr<at::Half>());
-  auto scaling_factors =
-      reinterpret_cast<half*>(_scaling_factors.data_ptr<at::Half>());
-  auto zeros = reinterpret_cast<int*>(_zeros.data_ptr<int>());
+  auto kernel = reinterpret_cast<int*>(_kernel.mutable_data_ptr<int>());
+  auto de_kernel = reinterpret_cast<half*>(
+      _de_kernel.mutable_data_ptr<torch::headeronly::Half>());
+  auto scaling_factors = reinterpret_cast<half*>(
+      _scaling_factors.mutable_data_ptr<torch::headeronly::Half>());
+  auto zeros = reinterpret_cast<int*>(_zeros.mutable_data_ptr<int>());

  dim3 num_blocks(x_blocks, y_blocks);
  dim3 threads_per_block(x_thread, y_thread);

-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  const cudaStream_t stream = get_current_cuda_stream();
  vllm::awq::dequantize_weights<<<num_blocks, threads_per_block, 0, stream>>>(
      kernel, scaling_factors, zeros, de_kernel, G);

@@ -466,27 +469,30 @@ torch::Tensor awq_dequantize(torch::Tensor _kernel,
 // zeros: IC // G, OC // 8 [int32] -> cast to IC // G, OC [uint4b]
 // assume that batch_size < 16 for now

-torch::Tensor awq_gemm(torch::Tensor _in_feats, torch::Tensor _kernel,
-                       torch::Tensor _scaling_factors, torch::Tensor _zeros,
-                       int64_t split_k_iters) {
+torch::stable::Tensor awq_gemm(torch::stable::Tensor _in_feats,
+                               torch::stable::Tensor _kernel,
+                               torch::stable::Tensor _scaling_factors,
+                               torch::stable::Tensor _zeros,
+                               int64_t split_k_iters) {
  int num_in_feats = _in_feats.size(0);
  int num_in_channels = _in_feats.size(1);
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(_in_feats));
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      _in_feats.get_device_index());

-  auto options = torch::TensorOptions()
-                     .dtype(_in_feats.dtype())
-                     .device(_in_feats.device());
-  at::Tensor _out_feats =
-      torch::empty({split_k_iters, num_in_feats, _kernel.size(1) * 8}, options);
+  auto _out_feats = torch::stable::empty(
+      {split_k_iters, num_in_feats, _kernel.size(1) * 8},
+      _in_feats.scalar_type(), std::nullopt, _in_feats.device());
  int num_out_feats = _out_feats.size(-2);
  int num_out_channels = _out_feats.size(-1);

-  auto in_feats = reinterpret_cast<half*>(_in_feats.data_ptr<at::Half>());
-  auto kernel = reinterpret_cast<int*>(_kernel.data_ptr<int>());
-  auto out_feats = reinterpret_cast<half*>(_out_feats.data_ptr<at::Half>());
-  auto scaling_factors =
-      reinterpret_cast<half*>(_scaling_factors.data_ptr<at::Half>());
-  auto zeros = reinterpret_cast<int*>(_zeros.data_ptr<int>());
+  auto in_feats = reinterpret_cast<half*>(
+      _in_feats.mutable_data_ptr<torch::headeronly::Half>());
+  auto kernel = reinterpret_cast<int*>(_kernel.mutable_data_ptr<int>());
+  auto out_feats = reinterpret_cast<half*>(
+      _out_feats.mutable_data_ptr<torch::headeronly::Half>());
+  auto scaling_factors = reinterpret_cast<half*>(
+      _scaling_factors.mutable_data_ptr<torch::headeronly::Half>());
+  auto zeros = reinterpret_cast<int*>(_zeros.mutable_data_ptr<int>());
  int group_size = num_in_channels / _scaling_factors.size(0);

  if (num_out_channels % 64 != 0)
@@ -498,7 +504,7 @@ torch::Tensor awq_gemm(torch::Tensor _in_feats, torch::Tensor _kernel,
  if (num_out_channels % group_size != 0)
    throw std::invalid_argument("OC is not multiple of Group size");

-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  const cudaStream_t stream = get_current_cuda_stream();
  if (num_out_channels % 128 == 0) {
    int j_factors1 = num_out_channels / 128 / 1;
    dim3 num_blocks((num_out_feats + 16 - 1) / 16 * j_factors1 * split_k_iters);
@@ -522,5 +528,5 @@ torch::Tensor awq_gemm(torch::Tensor _in_feats, torch::Tensor _kernel,
            group_size, split_k_iters, in_feats, kernel, scaling_factors, zeros,
            num_in_feats, num_in_channels, num_out_channels, out_feats);
  }
-  return _out_feats.sum(0);
+  return torch::stable::sum(_out_feats, 0);
 }
@@ -1,20 +1,28 @@
 #include "allspark_utils.cuh"
-#include <torch/all.h>
-#include "core/registration.h"
+
+#include <torch/csrc/stable/c/shim.h>
+#include <torch/csrc/stable/library.h>
+#include <torch/csrc/stable/tensor.h>
+#include <torch/headeronly/core/ScalarType.h>
+
 #include <cublas_v2.h>

-at::Tensor as_g_workspace;
+#include "core/registration.h"
+#include "libtorch_stable/torch_utils.h"
+
+torch::stable::Tensor as_g_workspace;

 #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800

-torch::Tensor allspark_w8a16_gemm(
-    torch::Tensor const& a, torch::Tensor const& b_qweight,
-    torch::Tensor const& b_scales, std::optional<torch::Tensor> const& b_qzeros,
-    int64_t n, int64_t group_size, int64_t sm_count, int64_t sm_version,
+torch::stable::Tensor allspark_w8a16_gemm(
+    torch::stable::Tensor const& a, torch::stable::Tensor const& b_qweight,
+    torch::stable::Tensor const& b_scales,
+    std::optional<torch::stable::Tensor> const& b_qzeros, int64_t n,
+    int64_t group_size, int64_t sm_count, int64_t sm_version,
    int64_t CUBLAS_M_THRESHOLD, bool has_zp, bool n32k16_reorder) {
-  TORCH_CHECK_NOT_IMPLEMENTED(
+  STD_TORCH_CHECK_NOT_IMPLEMENTED(
      false, "allspark_w8a16_gemm(..) requires CUDA_ARCH >= 8.0");
-  return torch::empty({1, 1});
+  return torch::stable::empty({1, 1});
 }

 #else
@@ -848,8 +856,8 @@ void restore_N32_K16_dequantize_rhs_w8a16(const QT* qdata, const FT* scales,
                                          const int N_32align, const int N,
                                          const int K, const int GroupSize,
                                          cudaStream_t stream) {
-  TORCH_CHECK(N % 8 == 0 && K % 16 == 0 && N_32align % 32 == 0,
-              "Unsupported shape");
+  STD_TORCH_CHECK(N % 8 == 0 && K % 16 == 0 && N_32align % 32 == 0,
+                  "Unsupported shape");
  if (GroupSize == -1) {
    const int BLOCK = 128;
    dim3 grid(N_32align / 32, ((K / 16) + 3) / 4);
@@ -859,7 +867,7 @@ void restore_N32_K16_dequantize_rhs_w8a16(const QT* qdata, const FT* scales,
  }
  // TODO: Support SubChannel
  else {
-    TORCH_CHECK(false, "Now only support PerChannel");
+    STD_TORCH_CHECK(false, "Now only support PerChannel");
  }
 }

@@ -916,24 +924,27 @@ void allspark_qgemm_w8a16_perc_ampere(

 }  // namespace allspark

-torch::Tensor allspark_w8a16_gemm(
-    torch::Tensor const& a, torch::Tensor const& b_qweight,
-    torch::Tensor const& b_scales, std::optional<torch::Tensor> const& b_qzeros,
-    int64_t n, int64_t group_size, int64_t sm_count, int64_t sm_version,
+torch::stable::Tensor allspark_w8a16_gemm(
+    torch::stable::Tensor const& a, torch::stable::Tensor const& b_qweight,
+    torch::stable::Tensor const& b_scales,
+    std::optional<torch::stable::Tensor> const& b_qzeros, int64_t n,
+    int64_t group_size, int64_t sm_count, int64_t sm_version,
    int64_t CUBLAS_M_THRESHOLD, bool has_zp, bool n32k16_reorder) {
  // Verify device and strides
-  TORCH_CHECK(a.device().is_cuda(), "A is not on GPU");
-  TORCH_CHECK(a.is_contiguous(), "A is not contiguous");
+  STD_TORCH_CHECK(a.device().is_cuda(), "A is not on GPU");
+  STD_TORCH_CHECK(a.is_contiguous(), "A is not contiguous");

-  TORCH_CHECK(b_qweight.device().is_cuda(), "b_qweight is not on GPU");
-  TORCH_CHECK(b_qweight.is_contiguous(), "b_qweight is not contiguous");
+  STD_TORCH_CHECK(b_qweight.device().is_cuda(), "b_qweight is not on GPU");
+  STD_TORCH_CHECK(b_qweight.is_contiguous(), "b_qweight is not contiguous");

-  TORCH_CHECK(b_scales.device().is_cuda(), "b_scales is not on GPU");
-  TORCH_CHECK(b_scales.is_contiguous(), "b_scales is not contiguous");
+  STD_TORCH_CHECK(b_scales.device().is_cuda(), "b_scales is not on GPU");
+  STD_TORCH_CHECK(b_scales.is_contiguous(), "b_scales is not contiguous");

  if (has_zp) {
-    TORCH_CHECK(b_qzeros.value().device().is_cuda(), "b_qzeros is not on GPU");
-    TORCH_CHECK(b_qzeros.value().is_contiguous(), "b_qzeros is not contiguous");
+    STD_TORCH_CHECK(b_qzeros.value().device().is_cuda(),
+                    "b_qzeros is not on GPU");
+    STD_TORCH_CHECK(b_qzeros.value().is_contiguous(),
+                    "b_qzeros is not contiguous");
  }

  int m = a.size(0);
@@ -941,16 +952,17 @@ torch::Tensor allspark_w8a16_gemm(
  int k = a.size(1);

  // Verify shape
-  TORCH_CHECK(b_qweight.size(0) == n_32align,
-              "Shape mismatch: b_qweight.size(0) = ", b_qweight.size(0),
-              ", n_32align = ", n_32align);
-  TORCH_CHECK(b_qweight.size(1) == k,
-              "Shape mismatch: b_qweight.size(1) = ", b_qweight.size(1),
-              ", k = ", k);
+  STD_TORCH_CHECK(b_qweight.size(0) == n_32align,
+                  "Shape mismatch: b_qweight.size(0) = ", b_qweight.size(0),
+                  ", n_32align = ", n_32align);
+  STD_TORCH_CHECK(b_qweight.size(1) == k,
+                  "Shape mismatch: b_qweight.size(1) = ", b_qweight.size(1),
+                  ", k = ", k);

-  TORCH_CHECK(group_size == -1, "Currently only supports group_size = -1");
+  STD_TORCH_CHECK(group_size == -1, "Currently only supports group_size = -1");

-  const at::cuda::OptionalCUDAGuard device_guard(device_of(a));
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      a.get_device_index());
  const void* a_ptr = reinterpret_cast<const void*>(a.data_ptr());
  const uint8_t* b_ptr = reinterpret_cast<const uint8_t*>(b_qweight.data_ptr());
  const void* b_scale_ptr = reinterpret_cast<const void*>(b_scales.data_ptr());
@@ -959,12 +971,12 @@ torch::Tensor allspark_w8a16_gemm(
    b_zero_ptr = reinterpret_cast<const void*>(b_qzeros.value().data_ptr());
  }

-  auto c_options = torch::TensorOptions().dtype(a.dtype()).device(a.device());
-  torch::Tensor c = torch::empty({m, n}, c_options);
-  void* c_ptr = reinterpret_cast<void*>(c.data_ptr());
+  auto c =
+      torch::stable::empty({m, n}, a.scalar_type(), std::nullopt, a.device());
+  void* c_ptr = reinterpret_cast<void*>(c.mutable_data_ptr());

-  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-  cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
+  cudaStream_t stream = get_current_cuda_stream();
+  cublasHandle_t handle = get_current_cuda_blas_handle();

  allspark::BlockTileSplitkParams fused_gemm_params;

@@ -976,14 +988,15 @@ torch::Tensor allspark_w8a16_gemm(
        m, n, k, sm_count, fused_gemm_params);
  }

-  auto ws_options = torch::TensorOptions().dtype(at::kChar).device(a.device());
  if (as_g_workspace.numel() <
      ws_size) {  // ws_options: kChar, so numel() is bytes
-    as_g_workspace = torch::empty({long(ws_size)}, ws_options);
+    as_g_workspace = torch::stable::empty({static_cast<int64_t>(ws_size)},
+                                          torch::headeronly::ScalarType::Char,
+                                          std::nullopt, a.device());
  }
  void* ws = reinterpret_cast<void*>(as_g_workspace.data_ptr());

-  if (a.dtype() == at::ScalarType::Half) {
+  if (a.scalar_type() == torch::headeronly::ScalarType::Half) {
    allspark::allspark_qgemm_w8a16_perc_ampere<__half, uint8_t>(
        reinterpret_cast<const __half*>(a_ptr), b_ptr,
        reinterpret_cast<const __half*>(b_scale_ptr),
@@ -991,7 +1004,7 @@ torch::Tensor allspark_w8a16_gemm(
        reinterpret_cast<__half*>(c_ptr), m, n_32align, n, k, ws,
        fused_gemm_params, group_size, CUBLAS_M_THRESHOLD, sm_version, stream,
        handle);
-  } else if (a.dtype() == at::ScalarType::BFloat16) {
+  } else if (a.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
    allspark::allspark_qgemm_w8a16_perc_ampere<__nv_bfloat16, uint8_t>(
        reinterpret_cast<const __nv_bfloat16*>(a_ptr), b_ptr,
        reinterpret_cast<const __nv_bfloat16*>(b_scale_ptr),
@@ -1006,6 +1019,6 @@ torch::Tensor allspark_w8a16_gemm(

 #endif

-TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
-  m.impl("allspark_w8a16_gemm", &allspark_w8a16_gemm);
+STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
+  m.impl("allspark_w8a16_gemm", TORCH_BOX(&allspark_w8a16_gemm));
 }
@@ -1,6 +1,11 @@
 #include "allspark_utils.cuh"
-#include <torch/all.h>
+
+#include <torch/csrc/stable/library.h>
+#include <torch/csrc/stable/tensor.h>
+#include <torch/headeronly/core/ScalarType.h>
+
 #include "core/registration.h"
+#include "libtorch_stable/torch_utils.h"

 namespace allspark {

@@ -99,36 +104,40 @@ void rearrange_kn_weight_as_n32k16_order_ldg16(
 }  // namespace allspark

 void rearrange_kn_weight_as_n32k16_order(
-    torch::Tensor const& b_qweight, torch::Tensor const& b_scales,
-    std::optional<torch::Tensor> const& b_zeros, bool has_zp,
-    torch::Tensor& b_qweight_reorder, torch::Tensor& b_scales_reorder,
-    std::optional<torch::Tensor> const& b_zeros_reorder, const int64_t K,
-    const int64_t N, const int64_t N_32align) {
+    torch::stable::Tensor const& b_qweight,
+    torch::stable::Tensor const& b_scales,
+    std::optional<torch::stable::Tensor> const& b_zeros, bool has_zp,
+    torch::stable::Tensor& b_qweight_reorder,
+    torch::stable::Tensor& b_scales_reorder,
+    std::optional<torch::stable::Tensor> const& b_zeros_reorder,
+    const int64_t K, const int64_t N, const int64_t N_32align) {
  // Verify device and strides
-  TORCH_CHECK(b_qweight.device().is_cuda(), "b_qweight is not on GPU");
-  TORCH_CHECK(b_qweight.is_contiguous(), "b_qweight is not contiguous");
+  STD_TORCH_CHECK(b_qweight.device().is_cuda(), "b_qweight is not on GPU");
+  STD_TORCH_CHECK(b_qweight.is_contiguous(), "b_qweight is not contiguous");

-  TORCH_CHECK(b_scales.device().is_cuda(), "b_scales is not on GPU");
-  TORCH_CHECK(b_scales.is_contiguous(), "b_scales is not contiguous");
+  STD_TORCH_CHECK(b_scales.device().is_cuda(), "b_scales is not on GPU");
+  STD_TORCH_CHECK(b_scales.is_contiguous(), "b_scales is not contiguous");

-  TORCH_CHECK(b_qweight_reorder.device().is_cuda(),
-              "b_qweight_reorder is not on GPU");
-  TORCH_CHECK(b_qweight_reorder.is_contiguous(),
-              "b_qweight_reorder is not contiguous");
+  STD_TORCH_CHECK(b_qweight_reorder.device().is_cuda(),
+                  "b_qweight_reorder is not on GPU");
+  STD_TORCH_CHECK(b_qweight_reorder.is_contiguous(),
+                  "b_qweight_reorder is not contiguous");

-  TORCH_CHECK(b_scales_reorder.device().is_cuda(),
-              "b_scales_reorder is not on GPU");
-  TORCH_CHECK(b_scales_reorder.is_contiguous(),
-              "b_scales_reorder is not contiguous");
+  STD_TORCH_CHECK(b_scales_reorder.device().is_cuda(),
+                  "b_scales_reorder is not on GPU");
+  STD_TORCH_CHECK(b_scales_reorder.is_contiguous(),
+                  "b_scales_reorder is not contiguous");

  if (has_zp) {
-    TORCH_CHECK(b_zeros.value().device().is_cuda(), "b_zeros is not on GPU");
-    TORCH_CHECK(b_zeros.value().is_contiguous(), "b_zeros is not contiguous");
+    STD_TORCH_CHECK(b_zeros.value().device().is_cuda(),
+                    "b_zeros is not on GPU");
+    STD_TORCH_CHECK(b_zeros.value().is_contiguous(),
+                    "b_zeros is not contiguous");

-    TORCH_CHECK(b_zeros_reorder.value().device().is_cuda(),
-                "b_zeros_reorder is not on GPU");
-    TORCH_CHECK(b_zeros_reorder.value().is_contiguous(),
-                "b_zeros_reorder is not contiguous");
+    STD_TORCH_CHECK(b_zeros_reorder.value().device().is_cuda(),
+                    "b_zeros_reorder is not on GPU");
+    STD_TORCH_CHECK(b_zeros_reorder.value().is_contiguous(),
+                    "b_zeros_reorder is not contiguous");
  }

  const uint8_t* matB = reinterpret_cast<const uint8_t*>(b_qweight.data_ptr());
@@ -136,18 +145,20 @@ void rearrange_kn_weight_as_n32k16_order(
  const void* b_zero = has_zp ? b_zeros.value().data_ptr() : nullptr;

  uint8_t* matB_reorder =
-      reinterpret_cast<uint8_t*>(b_qweight_reorder.data_ptr());
-  void* b_scale_reorder = b_scales_reorder.data_ptr();
-  void* b_zero_reorder = has_zp ? b_zeros_reorder.value().data_ptr() : nullptr;
+      reinterpret_cast<uint8_t*>(b_qweight_reorder.mutable_data_ptr());
+  void* b_scale_reorder = b_scales_reorder.mutable_data_ptr();
+  void* b_zero_reorder =
+      has_zp ? b_zeros_reorder.value().mutable_data_ptr() : nullptr;

-  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-  if (b_scales.dtype() == at::ScalarType::Half) {
+  cudaStream_t stream = get_current_cuda_stream();
+  if (b_scales.scalar_type() == torch::headeronly::ScalarType::Half) {
    allspark::rearrange_kn_weight_as_n32k16_order_ldg16<__half>(
        matB, reinterpret_cast<const __half*>(b_scale),
        reinterpret_cast<const __half*>(b_zero), matB_reorder,
        reinterpret_cast<__half*>(b_scale_reorder),
        reinterpret_cast<__half*>(b_zero_reorder), K, N, N_32align, stream);
-  } else if (b_scales.dtype() == at::ScalarType::BFloat16) {
+  } else if (b_scales.scalar_type() ==
+             torch::headeronly::ScalarType::BFloat16) {
    allspark::rearrange_kn_weight_as_n32k16_order_ldg16<__nv_bfloat16>(
        matB, reinterpret_cast<const __nv_bfloat16*>(b_scale),
        reinterpret_cast<const __nv_bfloat16*>(b_zero), matB_reorder,
@@ -157,7 +168,7 @@ void rearrange_kn_weight_as_n32k16_order(
  }
 }

-TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
+STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
  m.impl("rearrange_kn_weight_as_n32k16_order",
-         &rearrange_kn_weight_as_n32k16_order);
+         TORCH_BOX(&rearrange_kn_weight_as_n32k16_order));
 }
@@ -1,13 +1,12 @@
 #pragma once

-#include <torch/all.h>
-#include <c10/cuda/CUDAGuard.h>
-#include <ATen/cuda/CUDAContext.h>
-#include <cuda_runtime.h>
-#include <cuda_fp16.h>
 #include <cuda_bf16.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+
 #include <iostream>
-#include "../marlin/marlin_dtypes.cuh"
+
+#include "quantization/marlin/marlin_dtypes.cuh"
 using marlin::MarlinScalarType2;

 namespace allspark {
@@ -11,18 +11,16 @@ Redistribution and use in source and binary forms, with or without modification,
 THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS “AS IS” AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 ***********/

-#include <torch/all.h>
+#include "libtorch_stable/torch_utils.h"
+#include "libtorch_stable/dispatch_utils.h"
+
+#include <torch/csrc/stable/library.h>
+#include <torch/csrc/stable/macros.h>
+
 #include <stdint.h>
 #include <cuda_runtime.h>
 #include <mma.h>
 #include <cuda/annotated_ptr>
-#include <c10/cuda/CUDAException.h>
-
-#include <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAGuard.h>
-
-#include "core/registration.h"
-#include "dispatch_utils.h"

 namespace hadacore {

@@ -65,12 +63,12 @@ constexpr int launch_configs_big[7][3] = {
 };

 // a 4x2, b 2x2, c 2x2
-template <torch::ScalarType dtype>
+template <torch::headeronly::ScalarType dtype>
 __device__ __forceinline__ void mma_m16_n8_k16_b16_b16_b16_noacc(b32 a0, b32 a1, b32 a2, b32 a3, b32 b0, b32 b1, b32& c0, b32& c1){
-    static_assert(dtype == torch::ScalarType::Half || dtype == torch::ScalarType::BFloat16);
+    static_assert(dtype == torch::headeronly::ScalarType::Half || dtype == torch::headeronly::ScalarType::BFloat16);
    // d, a, b, c
    b32 zero = 0;
-    if constexpr(dtype == torch::ScalarType::Half) {
+    if constexpr(dtype == torch::headeronly::ScalarType::Half) {
        asm (
            "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 "
            "{%0, %1}, {%2, %3, %4, %5}, {%6, %7}, {%8, %9};\n\t"
@@ -89,7 +87,7 @@ __device__ __forceinline__ void mma_m16_n8_k16_b16_b16_b16_noacc(b32 a0, b32 a1,
 }

 // a 4x2, b 4x2, c 4x2
-template <torch::ScalarType dtype>
+template <torch::headeronly::ScalarType dtype>
 __device__ __forceinline__ void mma_m16_n16_k16_b16_b16_b16_noacc(b32 a0, b32 a1, b32 a2, b32 a3, b32 b0, b32 b1, b32 b2, b32 b3, b32& c0, b32& c1, b32& c2, b32& c3){
    mma_m16_n8_k16_b16_b16_b16_noacc<dtype>(a0, a1, a2, a3, b0, b1, c0, c1);
    mma_m16_n8_k16_b16_b16_b16_noacc<dtype>(a0, a1, a2, a3, b2, b3, c2, c3);
@@ -108,11 +106,11 @@ __device__ __forceinline__ void matrix_transpose_m8_n8_b16_inplace(b32& a0) {
 #define n_p(i) ((val_1n[i] & 0x0000FFFF) | val_1p[i] << 16)
 #define n_n(i) ((val_1n[i] & 0x0000FFFF) | val_1n[i] << 16)

-template<int64_t num_chunks, int64_t warps_per_block, int64_t log_had_size, int64_t blocks_per_sm, bool enable_mask, torch::ScalarType dtype>
+template<int64_t num_chunks, int64_t warps_per_block, int64_t log_had_size, int64_t blocks_per_sm, bool enable_mask, torch::headeronly::ScalarType dtype>
 __global__ void __launch_bounds__(32 * warps_per_block, blocks_per_sm)
 // a is column major, b is row major
 hadamard_transform_kernel(b16* a, b16* out, int total_num_chunks) {
-    static_assert(dtype == torch::ScalarType::Half || dtype == torch::ScalarType::BFloat16, "Only fp16 and bf16 supported currently");
+    static_assert(dtype == torch::headeronly::ScalarType::Half || dtype == torch::headeronly::ScalarType::BFloat16, "Only fp16 and bf16 supported currently");

    b32 b_frag_all[num_chunks][4]; // for all chunks, holds matrix fragment (which takes 4 regs of b16x2 * 32 threads)

@@ -162,8 +160,8 @@ hadamard_transform_kernel(b16* a, b16* out, int total_num_chunks) {
    constexpr b16 bf16_1p[4] = {0b0011111100110101, 0b0011111100000000, 0b0011111010110101, 0b0011111010000000};
    constexpr b16 bf16_1n[4] = {0b1011111100110101, 0b1011111100000000, 0b1011111010110101, 0b1011111010000000};

-    #define val_type_1p(i) (((dtype) == torch::ScalarType::Half) ? (fp16_1p[i]) : (bf16_1p[i]))
-    #define val_type_1n(i) (((dtype) == torch::ScalarType::Half) ? (fp16_1n[i]) : (bf16_1n[i]))
+    #define val_type_1p(i) (((dtype) == torch::headeronly::ScalarType::Half) ? (fp16_1p[i]) : (bf16_1p[i]))
+    #define val_type_1n(i) (((dtype) == torch::headeronly::ScalarType::Half) ? (fp16_1n[i]) : (bf16_1n[i]))
    constexpr b16 val_1p[4] = {val_type_1p(0), val_type_1p(1), val_type_1p(2), val_type_1p(3)};
    constexpr b16 val_1n[4] = {val_type_1n(0), val_type_1n(1), val_type_1n(2), val_type_1n(3)};

@@ -684,14 +682,14 @@ constexpr int64_t ceil_div(int64_t a, int64_t b) {
    return (a + b - 1) / b;
 }

-template <torch::ScalarType dtype, int64_t chunks_per_warp, int64_t warps_per_block, int64_t log_had_size, int64_t blocks_per_sm, bool check_masking = false>
+template <torch::headeronly::ScalarType dtype, int64_t chunks_per_warp, int64_t warps_per_block, int64_t log_had_size, int64_t blocks_per_sm, bool check_masking = false>
 void __forceinline__ run_kernel(b16* a_mat, b16* out, int64_t num_chunks, cudaStream_t stream) {
    int64_t shared_size = chunks_per_warp * warps_per_block * 128 * 4;
    dim3 block_size = 32 * warps_per_block;

    #define CHECK_SHARED_LIM() {                                                                              \
        if (shared_size > 48 * 1024) {                                                                        \
-            C10_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, 65536)); \
+            STD_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, 65536)); \
        }                                                                                                     \
    }                                                                                                         \

@@ -714,10 +712,10 @@ void __forceinline__ run_kernel(b16* a_mat, b16* out, int64_t num_chunks, cudaSt
        kernel<<<dim3(grid_size), dim3(block_size), shared_size, stream>>>(a_mat, out, num_chunks);
    }
    
-    C10_CUDA_KERNEL_LAUNCH_CHECK();
+    STD_CUDA_KERNEL_LAUNCH_CHECK();
 }

-template <torch::ScalarType dtype>
+template <torch::headeronly::ScalarType dtype>
 void run_fht(void* a_mat_ptr, void* out_ptr, int64_t numel, int64_t had_size, cudaStream_t stream) {
    int64_t num_chunks = numel / 256; // caller required to ensure divisible by 256
    // for size 256, use (2, 1)
@@ -764,54 +762,54 @@ void run_fht(void* a_mat_ptr, void* out_ptr, int64_t numel, int64_t had_size, cu
    }
 }

-template void run_fht<torch::ScalarType::Half>(void* a_mat_ptr, void* out_ptr, int64_t numel, int64_t had_size, cudaStream_t stream);
-template void run_fht<torch::ScalarType::BFloat16>(void* a_mat_ptr, void* out_ptr, int64_t numel, int64_t had_size, cudaStream_t stream);
+template void run_fht<torch::headeronly::ScalarType::Half>(void* a_mat_ptr, void* out_ptr, int64_t numel, int64_t had_size, cudaStream_t stream);
+template void run_fht<torch::headeronly::ScalarType::BFloat16>(void* a_mat_ptr, void* out_ptr, int64_t numel, int64_t had_size, cudaStream_t stream);

 }  // namespace hadacore

 constexpr bool is_power_of_two(int x) { return x && !(x & (x - 1)); }

-torch::Tensor hadacore_transform(torch::Tensor& x, bool inplace) {
+torch::stable::Tensor hadacore_transform(torch::stable::Tensor& x, bool inplace) {
    auto dtype = x.scalar_type();
-    TORCH_CHECK(dtype == torch::ScalarType::Half || dtype == torch::ScalarType::BFloat16, "Only fp16 and bf16 supported currently");
-    TORCH_CHECK(x.is_cuda());
-    
+    STD_TORCH_CHECK(dtype == torch::headeronly::ScalarType::Half || dtype == torch::headeronly::ScalarType::BFloat16, "Only fp16 and bf16 supported currently");
+    STD_TORCH_CHECK(x.is_cuda());
+
    const int had_size = x.size(-1);
-    TORCH_CHECK(is_power_of_two(had_size) && (had_size <= (1U << 15)),
+    STD_TORCH_CHECK(is_power_of_two(had_size) && (had_size <= (1U << 15)),
        "Only power of two Hadamard sizes up to 2^15 are supported, got ", had_size);
-    
+
    const auto res_shape = x.sizes();
-    x = x.reshape({-1, had_size});
-    
+    x = torch::stable::reshape(x, {-1, had_size});
+
    auto numel = x.numel();
    if (numel % 256 != 0) {
-        x = torch::nn::functional::pad(x, torch::nn::functional::PadFuncOptions({0, 0, 0, (256 - numel % 256) / had_size}));
+        x = torch::stable::pad(x, {0, 0, 0, (256 - numel % 256) / had_size});
    }
-    
+
    if (x.stride(-1) != 1) {
-        x = x.contiguous();
+        x = torch::stable::contiguous(x);
    }
-    torch::Tensor out = inplace ? x : torch::empty_like(x);
+    torch::stable::Tensor out = inplace ? x : torch::stable::empty_like(x);

-    at::cuda::CUDAGuard device_guard{(char)x.get_device()};
-    auto stream = at::cuda::getCurrentCUDAStream().stream();
+    torch::stable::accelerator::DeviceGuard device_guard(x.get_device_index());
+    auto stream = get_current_cuda_stream();

-    VLLM_DISPATCH_HALF_TYPES(x.scalar_type(), "hadacore_transform_runfht", [&] {
-      auto constexpr SCALAR_TYPE = c10::CppTypeToScalarType<scalar_t>::value;
+    VLLM_STABLE_DISPATCH_HALF_TYPES(x.scalar_type(), "hadacore_transform_runfht", [&] {
+      auto constexpr SCALAR_TYPE = torch::headeronly::CppTypeToScalarType<scalar_t>::value;
      hadacore::run_fht<SCALAR_TYPE>(x.data_ptr(), x.data_ptr(), x.numel(), had_size, stream);
    });

    if (numel % 256 != 0) {
-        out = out.narrow(0, 0, numel / had_size);
+        out = torch::stable::narrow(out, 0, 0, numel / had_size);
    }

    if (inplace && out.data_ptr() != x.data_ptr()) {
-        x.copy_(out.view(res_shape));
+        torch::stable::copy_(x, torch::stable::view(out, res_shape));
        return x;
    }
-    return out.reshape(res_shape);
+    return torch::stable::reshape(out, res_shape);
 }

-TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
-    m.impl("hadacore_transform", &hadacore_transform);
+STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
+    m.impl("hadacore_transform", TORCH_BOX(&hadacore_transform));
 }
@@ -156,6 +156,17 @@ inline int GetGroupsPerBlock(int64_t num_groups) {
  return 1;
 }

+// Largest divisor of padded_groups_per_row that is <= 16. ry = 16 / kx.
+inline int GetGroupsPerBlockX(int64_t padded_groups_per_row) {
+  if (padded_groups_per_row % 16 == 0) {
+    return 16;
+  }
+  if (padded_groups_per_row % 8 == 0) {
+    return 8;
+  }
+  return 4;
+}
+
 void per_token_group_quant_8bit(const torch::stable::Tensor& input,
                                torch::stable::Tensor& output_q,
                                torch::stable::Tensor& output_s,
@@ -247,11 +258,11 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input,
 //
 // Constraints: GROUP_SIZE % (THREADS_PER_GROUP * VEC_SIZE) == 0; for
 // THREADS_PER_GROUP=8 and bf16/fp16 (VEC_SIZE=16), this means GROUP_SIZE=128.
-template <typename T, typename DST_DTYPE, int GROUP_SIZE>
+template <typename T, typename DST_DTYPE, int GROUP_SIZE, int kGroupsPerBlockX,
+          int kRowsPerBlock>
 __global__ void per_token_group_quant_8bit_packed_register_kernel(
    const T* __restrict__ input, void* __restrict__ output_q,
-    unsigned int* __restrict__ output_s_packed, const int64_t num_groups_padded,
-    const int groups_per_block, const int padded_groups_per_row,
+    unsigned int* __restrict__ output_s_packed, const int padded_groups_per_row,
    const int groups_per_row, const int mn, const int output_q_mn_extent,
    const int tma_aligned_mn, const int64_t num_scale_elems, const float eps,
    const float min_8bit, const float max_8bit) {
@@ -260,27 +271,25 @@ __global__ void per_token_group_quant_8bit_packed_register_kernel(
  constexpr int VEC_SIZE = 32 / sizeof(T);  // 16 for bf16/fp16
  static_assert(GROUP_SIZE == THREADS_PER_GROUP * VEC_SIZE,
                "GROUP_SIZE must equal THREADS_PER_GROUP * VEC_SIZE");
-  // Each group's 8 threads must live in a single warp octet so the
-  // 0xffu << (threadIdx.x & 24u) shuffle mask selects exactly the lanes
-  // that share a group. Requires 32 % THREADS_PER_GROUP == 0 and the host
-  // to launch num_threads as a multiple of THREADS_PER_GROUP (which it does
-  // via num_threads = groups_per_block * THREADS_PER_GROUP).
  static_assert(32 % THREADS_PER_GROUP == 0,
                "THREADS_PER_GROUP must divide warp size for the shuffle "
                "mask to be valid");
+  static_assert(
+      kGroupsPerBlockX > 0 && (kGroupsPerBlockX & (kGroupsPerBlockX - 1)) == 0,
+      "kGroupsPerBlockX must be a positive power of 2");
+  static_assert(kRowsPerBlock > 0, "kRowsPerBlock must be positive");

  const int local_group_id = threadIdx.x / THREADS_PER_GROUP;
  const int lane_id = threadIdx.x % THREADS_PER_GROUP;

-  const int64_t block_group_id = blockIdx.x * groups_per_block;
-  const int64_t global_group_id = block_group_id + local_group_id;
-  if (global_group_id >= num_groups_padded) {
+  const int sf_k_local = local_group_id % kGroupsPerBlockX;
+  const int row_local = local_group_id / kGroupsPerBlockX;
+  const int sf_k_idx = blockIdx.x * kGroupsPerBlockX + sf_k_local;
+  const int mn_idx = blockIdx.y * kRowsPerBlock + row_local;
+
+  if (mn_idx >= tma_aligned_mn) {
    return;
  }
-
-  const int sf_k_idx =
-      static_cast<int>(global_group_id % padded_groups_per_row);
-  const int mn_idx = static_cast<int>(global_group_id / padded_groups_per_row);
  const bool is_valid_group = (mn_idx < mn) && (sf_k_idx < groups_per_row);

  // Load 16 input elements (32 B) into registers as two adjacent uint4
@@ -443,34 +452,53 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,

  constexpr int THREADS_PER_GROUP = 8;
  const int64_t padded_groups_per_row = k_num_packed_sfk * 4;
-  const int64_t num_groups_padded = tma_aligned_mn * padded_groups_per_row;
  const int64_t num_scale_elems = mn + (k_num_packed_sfk - 1) * tma_aligned_mn;
-  const int groups_per_block = GetGroupsPerBlock(num_groups_padded);
+
+  STD_TORCH_CHECK(padded_groups_per_row % 4 == 0,
+                  "padded_groups_per_row=", padded_groups_per_row,
+                  " is not a multiple of 4.");
+  const int kx = GetGroupsPerBlockX(padded_groups_per_row);
+  const int ry = 16 / kx;
+  const int64_t blocks_x = padded_groups_per_row / kx;
+  const int64_t blocks_y = (tma_aligned_mn + ry - 1) / ry;
+  const int num_threads = (kx * ry) * THREADS_PER_GROUP;
+  // CUDA caps grid.x and grid.y at 2^31 - 1; guard against pathological inputs.
+  STD_TORCH_CHECK(blocks_x <= static_cast<int64_t>(INT32_MAX) &&
+                      blocks_y <= static_cast<int64_t>(INT32_MAX),
+                  "per_token_group_quant_8bit_packed grid too large: (",
+                  blocks_x, ", ", blocks_y, ").");

  auto dst_type = output_q.scalar_type();
-  const int64_t num_blocks = num_groups_padded / groups_per_block;
-  const int num_threads = groups_per_block * THREADS_PER_GROUP;
-  // CUDA caps grid.x at 2^31 - 1; this fits any realistic shape but guard
-  // against pathological inputs.
-  STD_TORCH_CHECK(num_blocks <= static_cast<int64_t>(INT32_MAX),
-                  "per_token_group_quant_8bit_packed grid too large: ",
-                  num_blocks, " blocks (max ", INT32_MAX, ").");

-#define LAUNCH_REG_KERNEL(T, DST_DTYPE)                                   \
-  do {                                                                    \
-    dim3 grid(static_cast<unsigned int>(num_blocks));                     \
-    dim3 block(num_threads);                                              \
-    per_token_group_quant_8bit_packed_register_kernel<T, DST_DTYPE, 128>  \
-        <<<grid, block, 0, stream>>>(                                     \
-            static_cast<const T*>(input.data_ptr()), output_q.data_ptr(), \
-            reinterpret_cast<unsigned int*>(output_s_packed.data_ptr()),  \
-            num_groups_padded, groups_per_block,                          \
-            static_cast<int>(padded_groups_per_row),                      \
-            static_cast<int>(groups_per_row), static_cast<int>(mn),       \
-            static_cast<int>(output_q_mn_extent),                         \
-            static_cast<int>(tma_aligned_mn), num_scale_elems,            \
-            static_cast<float>(eps), static_cast<float>(min_8bit),        \
-            static_cast<float>(max_8bit));                                \
+#define LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, KX, RY)                         \
+  do {                                                                       \
+    dim3 grid(static_cast<unsigned int>(blocks_x),                           \
+              static_cast<unsigned int>(blocks_y));                          \
+    dim3 block(num_threads);                                                 \
+    per_token_group_quant_8bit_packed_register_kernel<T, DST_DTYPE, 128, KX, \
+                                                      RY>                    \
+        <<<grid, block, 0, stream>>>(                                        \
+            static_cast<const T*>(input.data_ptr()), output_q.data_ptr(),    \
+            reinterpret_cast<unsigned int*>(output_s_packed.data_ptr()),     \
+            static_cast<int>(padded_groups_per_row),                         \
+            static_cast<int>(groups_per_row), static_cast<int>(mn),          \
+            static_cast<int>(output_q_mn_extent),                            \
+            static_cast<int>(tma_aligned_mn), num_scale_elems,               \
+            static_cast<float>(eps), static_cast<float>(min_8bit),           \
+            static_cast<float>(max_8bit));                                   \
+  } while (0)
+
+#define LAUNCH_REG_KERNEL(T, DST_DTYPE)                    \
+  do {                                                     \
+    if (kx == 16) {                                        \
+      LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, 16, 1);         \
+    } else if (kx == 8) {                                  \
+      LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, 8, 2);          \
+    } else if (kx == 4) {                                  \
+      LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, 4, 4);          \
+    } else {                                               \
+      STD_TORCH_CHECK(false, "Unsupported kx value ", kx); \
+    }                                                      \
  } while (0)

  VLLM_STABLE_DISPATCH_HALF_TYPES(
@@ -488,6 +516,7 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
      }));

 #undef LAUNCH_REG_KERNEL
+#undef LAUNCH_REG_KERNEL_INST
 }

 void per_token_group_quant_fp8(const torch::stable::Tensor& input,
@@ -218,7 +218,54 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
  ops.def(
      "cutlass_encode_and_reorder_int4b_grouped(Tensor b_tensors) -> (Tensor, "
      "Tensor)");
+
+  // SM100 CUTLASS MLA decode
+  // conditionally compiled so impl registrations are in source file
+  ops.def(
+      "sm100_cutlass_mla_decode(Tensor! out, Tensor! lse, Tensor q_nope,"
+      "                         Tensor q_pe, Tensor kv_c_and_k_pe_cache,"
+      "                         Tensor seq_lens, Tensor page_table,"
+      "                         Tensor workspace, float scale,"
+      "                         int num_kv_splits) -> ()");
+
+  ops.def(
+      "sm100_cutlass_mla_get_workspace_size(int max_seq_len, int num_batches,"
+      "                                     int sm_count, int num_kv_splits) "
+      "-> int");
+  // Quantized GEMM for AWQ.
+  ops.def(
+      "awq_gemm(Tensor _in_feats, Tensor _kernel, Tensor _scaling_factors, "
+      "Tensor _zeros, SymInt split_k_iters) -> Tensor");
+
+  // Dequantization for AWQ.
+  ops.def(
+      "awq_dequantize(Tensor _kernel, Tensor _scaling_factors, "
+      "Tensor _zeros, SymInt split_k_iters, int thx, int thy) -> Tensor");
+
+  // DeepSeek V3 fused A GEMM (SM 9.0+, bf16 only, 1-16 tokens).
+  // conditionally compiled so impl registration is in source file
+  ops.def(
+      "dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
+
+  // reorder weight for AllSpark Ampere W8A16 Fused Gemm kernel
+  ops.def(
+      "rearrange_kn_weight_as_n32k16_order(Tensor b_qweight, Tensor b_scales, "
+      "Tensor? b_zeros, "
+      "bool has_zp, Tensor! b_qweight_reorder, Tensor! b_scales_reorder, "
+      "Tensor!? b_zeros_reorder, "
+      "int K, int N, int N_32align) -> ()");
+
+  // AllSpark quantization ops
+  ops.def(
+      "allspark_w8a16_gemm(Tensor a, Tensor b_qweight, Tensor b_scales, "
+      "Tensor? b_qzeros, "
+      "SymInt n, SymInt group_size, SymInt sm_count, SymInt sm_version, SymInt "
+      "CUBLAS_M_THRESHOLD, bool has_zp, bool n32k16_reorder) -> Tensor");
 #endif
+
+  // Hadamard transforms
+  // conditionally compiled so impl registration is in source file
+  ops.def("hadacore_transform(Tensor! x, bool inplace) -> Tensor");
 }

 STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
@@ -254,6 +301,16 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
  ops.impl("silu_and_mul_nvfp4_quant", TORCH_BOX(&silu_and_mul_nvfp4_quant));
  // mxfp4_experts_quant: registered in mxfp4_experts_quant.cu (SM100 only).
  // W4A8 ops: registered in w4a8_mm_entry.cu / w4a8_grouped_mm_entry.cu.
+
+  // AWQ ops
+  ops.impl("awq_gemm", TORCH_BOX(&awq_gemm));
+  ops.impl("awq_dequantize", TORCH_BOX(&awq_dequantize));
+
+  // DSV3 fused A GEMM: conditionally compiled so impl registration is in
+  // source file (dsv3_fused_a_gemm.cu)
+
+  // AllSpark ops: conditionally compiled so impl registrations are in source
+  // files (allspark_repack.cu and allspark_qgemm_w8a16.cu)
 #endif
 }

@@ -6,12 +6,71 @@
 #include <torch/csrc/stable/tensor.h>
 #include <torch/headeronly/util/shim_utils.h>

+#include <cublas_v2.h>
 #include <cuda_runtime.h>

+#include <deque>
+#include <mutex>
+#include <string>
+#include <vector>
+
 // Stable ABI equivalent of TORCH_CHECK_NOT_IMPLEMENTED.
 #define STD_TORCH_CHECK_NOT_IMPLEMENTED(cond, ...) \
  STD_TORCH_CHECK(cond, "NotImplementedError: ", __VA_ARGS__)

+// Device properties cache for stable ABI compatibility.
+// Uses raw CUDA/HIP APIs instead of ATen functions.
+// Using inline ensures a single instance across all translation units.
+inline std::deque<std::once_flag> device_flags;
+inline std::vector<cudaDeviceProp> device_properties;
+inline std::once_flag vectors_init_flag;
+
+inline void do_init_device_vectors() {
+  int device_count;
+  cudaError_t err = cudaGetDeviceCount(&device_count);
+  if (err != cudaSuccess) {
+    STD_TORCH_CHECK(false, "cudaGetDeviceCount failed: " +
+                               std::string(cudaGetErrorString(err)));
+  }
+  device_flags.resize(device_count);
+  device_properties.resize(device_count);
+}
+
+inline void initDeviceVectors() {
+  std::call_once(vectors_init_flag, do_init_device_vectors);
+}
+
+inline void initDeviceProperty(int device_index) {
+  cudaDeviceProp device_prop{};
+  cudaError_t err = cudaGetDeviceProperties(&device_prop, device_index);
+  if (err != cudaSuccess) {
+    STD_TORCH_CHECK(false, "cudaGetDeviceProperties failed: " +
+                               std::string(cudaGetErrorString(err)));
+  }
+  device_properties[device_index] = device_prop;
+}
+
+// Get device properties using raw CUDA/HIP APIs (stable ABI compatible).
+// Caches results per device so cudaGetDeviceProperties is called at most once
+// per device.
+inline cudaDeviceProp* get_device_prop() {
+  initDeviceVectors();
+  int device_index;
+  cudaError_t err = cudaGetDevice(&device_index);
+  if (err != cudaSuccess) {
+    STD_TORCH_CHECK(
+        false, "cudaGetDevice failed: " + std::string(cudaGetErrorString(err)));
+  }
+  STD_TORCH_CHECK(device_index >= 0 && static_cast<size_t>(device_index) <
+                                           device_properties.size(),
+                  "CUDA device index " + std::to_string(device_index) +
+                      " out of range [0, " +
+                      std::to_string(device_properties.size()) + ")");
+
+  std::call_once(device_flags[device_index], initDeviceProperty, device_index);
+  return &device_properties[device_index];
+}
+
 // Utility to get the current CUDA stream for a given device using stable APIs.
 // Returns a cudaStream_t for use in kernel launches.
 inline cudaStream_t get_current_cuda_stream(int32_t device_index = -1) {
@@ -20,3 +79,10 @@ inline cudaStream_t get_current_cuda_stream(int32_t device_index = -1) {
      aoti_torch_get_current_cuda_stream(device_index, &stream_ptr));
  return reinterpret_cast<cudaStream_t>(stream_ptr);
 }
+
+// Utility to get the current cuBLAS handle using stable APIs.
+inline cublasHandle_t get_current_cuda_blas_handle() {
+  void* blas_handle_ptr = nullptr;
+  TORCH_ERROR_CODE_CHECK(torch_get_current_cuda_blas_handle(&blas_handle_ptr));
+  return reinterpret_cast<cublasHandle_t>(blas_handle_ptr);
+}
@@ -0,0 +1,30 @@
+/*
+ * Fused RMSNorm + router GEMV for DeepSeek V4.
+ *
+ * Computes in a single kernel:
+ *   normed_x[m,k]   = x[m,k] * rsqrt(mean(x[m]^2) + eps) * norm_weight[k]
+ *   router_logits[m,n] = sum_k(normed_x[m,k] * gate_weight[n,k])
+ *
+ * The GEMV body mirrors the algorithm in csrc/moe/dsv3_router_gemm_*.cu
+ * (warp butterfly + smem cross-warp reduction, fp32 accumulation, PDL on
+ * SM90+).  Blocks 0..kNumTokens-1 each materialize one token's normed_x
+ * row to global memory using the algebraic identity
+ *      logits[m,n] = rsqrt[m] * sum_k(x[m,k] * nw[k] * gw[n,k])
+ * which lets every block produce its column of logits before normed_x
+ * exists in gmem.
+ *
+ * Logits output is fp32 only — DeepSeek V4 router gate is hard-coded to
+ * fp32 (vllm/model_executor/models/deepseek_v4.py:749).
+ */
+
+#pragma once
+
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+
+#include "dsv3_router_gemm_utils.h"
+
+template <typename T, int kNumTokens, int kNumExperts, int kHiddenDim>
+void invokeNormRouterGemm(float* logits, __nv_bfloat16* normed_x, T const* x,
+                          T const* norm_weight, T const* gate_weight, float eps,
+                          cudaStream_t stream);
@@ -0,0 +1,130 @@
+/*
+ * TORCH op entry for the fused RMSNorm + router GEMV kernel
+ * (DeepSeek V4 Pro).  This op is DSV4-Pro-specific: the kernel is
+ * instantiated only for ``num_experts == 384`` and ``hidden_dim ==
+ * 7168``.  Other configurations (e.g. DSV4-Flash with H=4096) must
+ * fall back to the unfused ``rms_norm`` + ``dsv3_router_gemm`` path.
+ */
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <torch/all.h>
+
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+
+#include "core/registration.h"
+#include "dsv4_norm_router_gemm.h"
+
+namespace {
+
+// DSV4-Pro hard-coded shape constants.  Renamed from the earlier
+// ``kKimiK2NumExperts`` to avoid the misleading impression that this
+// kernel targets Kimi K2 — 384 happens to match Kimi K2's gate but the
+// intent here is DSV4-Pro.
+constexpr int kDsv4NumExperts = 384;
+constexpr int kDsv4HiddenDim = 7168;
+
+template <int kBegin, int kEnd>
+struct LoopUnroller {
+  static void unroll(int num_tokens, float* logits, __nv_bfloat16* normed_x,
+                     __nv_bfloat16 const* x, __nv_bfloat16 const* norm_weight,
+                     __nv_bfloat16 const* gate_weight, float eps,
+                     cudaStream_t stream) {
+    if (num_tokens == kBegin) {
+      invokeNormRouterGemm<__nv_bfloat16, kBegin, kDsv4NumExperts,
+                           kDsv4HiddenDim>(logits, normed_x, x, norm_weight,
+                                           gate_weight, eps, stream);
+    } else {
+      LoopUnroller<kBegin + 1, kEnd>::unroll(num_tokens, logits, normed_x, x,
+                                             norm_weight, gate_weight, eps,
+                                             stream);
+    }
+  }
+};
+
+template <int kEnd>
+struct LoopUnroller<kEnd, kEnd> {
+  static void unroll(int num_tokens, float* logits, __nv_bfloat16* normed_x,
+                     __nv_bfloat16 const* x, __nv_bfloat16 const* norm_weight,
+                     __nv_bfloat16 const* gate_weight, float eps,
+                     cudaStream_t stream) {
+    if (num_tokens == kEnd) {
+      invokeNormRouterGemm<__nv_bfloat16, kEnd, kDsv4NumExperts,
+                           kDsv4HiddenDim>(logits, normed_x, x, norm_weight,
+                                           gate_weight, eps, stream);
+    } else {
+      throw std::invalid_argument(
+          "Invalid num_tokens, only supports 1 to 16 for "
+          "dsv4_norm_router_gemm");
+    }
+  }
+};
+
+}  // namespace
+
+void dsv4_norm_router_gemm(at::Tensor& logits,    // [num_tokens, E] fp32
+                           at::Tensor& normed_x,  // [num_tokens, H] bf16
+                           at::Tensor const& x,   // [num_tokens, H] bf16
+                           at::Tensor const& norm_weight,  // [H] bf16
+                           at::Tensor const& gate_weight,  // [E, H] bf16
+                           double eps) {
+  TORCH_CHECK(x.dim() == 2 && norm_weight.dim() == 1 && gate_weight.dim() == 2,
+              "x must be 2D, norm_weight 1D, gate_weight 2D");
+  TORCH_CHECK(logits.dim() == 2 && normed_x.dim() == 2,
+              "logits and normed_x must be 2D");
+
+  int const num_tokens = x.size(0);
+  int const hidden_dim = x.size(1);
+  int const num_experts = gate_weight.size(0);
+
+  TORCH_CHECK(hidden_dim == kDsv4HiddenDim,
+              "Expected hidden_dim=", kDsv4HiddenDim,
+              " (DSV4-Pro), but got hidden_dim=", hidden_dim);
+  TORCH_CHECK(gate_weight.size(1) == hidden_dim,
+              "gate_weight.shape[1] must equal x.shape[1]");
+  TORCH_CHECK(norm_weight.size(0) == hidden_dim,
+              "norm_weight.shape[0] must equal x.shape[1]");
+  TORCH_CHECK(num_experts == kDsv4NumExperts,
+              "Expected num_experts=", kDsv4NumExperts,
+              " (DSV4-Pro), but got num_experts=", num_experts);
+  TORCH_CHECK(num_tokens >= 1 && num_tokens <= 16,
+              "num_tokens must be in [1, 16] for dsv4_norm_router_gemm");
+
+  TORCH_CHECK(x.dtype() == at::kBFloat16, "x must be bf16");
+  TORCH_CHECK(norm_weight.dtype() == at::kBFloat16, "norm_weight must be bf16");
+  TORCH_CHECK(gate_weight.dtype() == at::kBFloat16, "gate_weight must be bf16");
+  TORCH_CHECK(normed_x.dtype() == at::kBFloat16, "normed_x must be bf16");
+  TORCH_CHECK(logits.dtype() == at::kFloat,
+              "logits must be float32 (DSV4 router output is hard-coded fp32)");
+
+  TORCH_CHECK(normed_x.size(0) == num_tokens && normed_x.size(1) == hidden_dim,
+              "normed_x must be [num_tokens, hidden_dim]");
+  TORCH_CHECK(logits.size(0) == num_tokens && logits.size(1) == num_experts,
+              "logits must be [num_tokens, num_experts]");
+
+  TORCH_CHECK(x.is_contiguous() && norm_weight.is_contiguous() &&
+                  gate_weight.is_contiguous() && normed_x.is_contiguous() &&
+                  logits.is_contiguous(),
+              "all tensors must be contiguous");
+
+  auto const sm = getSMVersion();
+  TORCH_CHECK(sm >= 90 && sm <= 103,
+              "dsv4_norm_router_gemm requires SM_90 <= CUDA ARCH <= SM_103");
+
+  cudaStream_t const stream = at::cuda::getCurrentCUDAStream();
+
+  auto* logits_ptr = reinterpret_cast<float*>(logits.mutable_data_ptr());
+  auto* nx_ptr = reinterpret_cast<__nv_bfloat16*>(normed_x.mutable_data_ptr());
+  auto* x_ptr = reinterpret_cast<__nv_bfloat16 const*>(x.data_ptr());
+  auto* nw_ptr = reinterpret_cast<__nv_bfloat16 const*>(norm_weight.data_ptr());
+  auto* gw_ptr = reinterpret_cast<__nv_bfloat16 const*>(gate_weight.data_ptr());
+  float const eps_f = static_cast<float>(eps);
+
+  LoopUnroller<1, 16>::unroll(num_tokens, logits_ptr, nx_ptr, x_ptr, nw_ptr,
+                              gw_ptr, eps_f, stream);
+}
+
+TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
+  m.impl("dsv4_norm_router_gemm", &dsv4_norm_router_gemm);
+}
@@ -0,0 +1,249 @@
+/*
+ * Fused RMSNorm + router GEMV for DeepSeek V4 (logits are fp32; bf16
+ * output is unsupported because DSV4 hard-codes fp32 logits).  See
+ * dsv4_norm_router_gemm.h for the math.
+ *
+ * The GEMV body mirrors csrc/moe/dsv3_router_gemm_float_out.cu (warp
+ * butterfly reduction + smem cross-warp reduction, fp32 accumulation,
+ * 128-thread block, PDL on SM90+).  RMSNorm is folded into the same
+ * pass via the identity
+ *   logits[m,n] = rsqrt[m] * sum_k(x[m,k] * nw[k] * gw[n,k])
+ * so x is read exactly once per block during the GEMV phase.  Blocks
+ * 0..kNumTokens-1 each materialize one row of normed_x for downstream
+ * experts / shared_experts to consume.
+ */
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+
+#include "dsv4_norm_router_gemm.h"
+
+namespace {
+
+// Convert 8 bf16 values packed in uint4 into 8 floats. Mirrors the helper
+// in dsv3_router_gemm_float_out.cu (kept local so the dsv3 file stays
+// untouched).
+template <int VPT>
+__device__ __forceinline__ void bf16_uint4_to_float8(uint4 const& vec,
+                                                     float* dst) {
+  __nv_bfloat16* bf16_ptr =
+      reinterpret_cast<__nv_bfloat16*>(const_cast<uint4*>(&vec));
+#pragma unroll
+  for (int i = 0; i < VPT; i++) {
+    dst[i] = __bfloat162float(bf16_ptr[i]);
+  }
+}
+
+template <typename T, int kBlockSize, int VPT, int kNumTokens, int kNumExperts,
+          int kHiddenDim>
+__global__ __launch_bounds__(128, 1) void norm_router_gemm_kernel(
+    float* __restrict__ logits, __nv_bfloat16* __restrict__ normed_x,
+    T const* __restrict__ x, T const* __restrict__ norm_weight,
+    T const* __restrict__ gate_weight, float eps) {
+  static_assert(kBlockSize == 128, "kernel assumes blockDim.x == 128");
+  static_assert(kHiddenDim % (VPT * kBlockSize) == 0,
+                "kHiddenDim must be a multiple of VPT * kBlockSize");
+
+  int const n_idx = blockIdx.x;
+  int const tid = threadIdx.x;
+  constexpr int kWarpSize = 32;
+  constexpr int kNumWarps = kBlockSize / kWarpSize;
+  constexpr int k_elems_per_iter = VPT * kBlockSize;
+  constexpr int k_iterations = kHiddenDim / k_elems_per_iter;
+
+  T const* gw_col = gate_weight + n_idx * kHiddenDim;
+
+  // Per-thread accumulators — fp32 throughout, matching dsv3 / layernorm.
+  float partial[kNumTokens] = {};
+  float ss[kNumTokens] = {};
+
+  // Cross-warp reduction scratch.
+  __shared__ float sm_partial[kNumTokens][kNumWarps];
+  __shared__ float sm_ss[kNumTokens][kNumWarps];
+  __shared__ float s_rsqrt[kNumTokens];
+
+  int k_bases[k_iterations];
+#pragma unroll
+  for (int ki = 0; ki < k_iterations; ki++) {
+    k_bases[ki] = ki * k_elems_per_iter + tid * VPT;
+  }
+
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
+  asm volatile("griddepcontrol.wait;");
+#endif
+
+  // ---- Phase 1: single pass over x, accumulate partial GEMV and ss. ----
+#pragma unroll
+  for (int ki = 0; ki < k_iterations; ki++) {
+    int const k_base = k_bases[ki];
+
+    uint4 nw_vec = *reinterpret_cast<uint4 const*>(norm_weight + k_base);
+    float nw_f[VPT];
+    bf16_uint4_to_float8<VPT>(nw_vec, nw_f);
+
+    uint4 b_vec = *reinterpret_cast<uint4 const*>(gw_col + k_base);
+    float b_f[VPT];
+    bf16_uint4_to_float8<VPT>(b_vec, b_f);
+
+#pragma unroll
+    for (int m = 0; m < kNumTokens; m++) {
+      uint4 a_vec =
+          *reinterpret_cast<uint4 const*>(x + m * kHiddenDim + k_base);
+      float a_f[VPT];
+      bf16_uint4_to_float8<VPT>(a_vec, a_f);
+
+#pragma unroll
+      for (int k = 0; k < VPT; k++) {
+        float a = a_f[k];
+        ss[m] += a * a;
+        partial[m] += a * nw_f[k] * b_f[k];
+      }
+    }
+  }
+
+  // ---- Phase 2: warp butterfly reduction for both ss[] and partial[]. ----
+  int const warpId = tid / kWarpSize;
+  int const laneId = tid % kWarpSize;
+
+#pragma unroll
+  for (int m = 0; m < kNumTokens; m++) {
+    float p = partial[m];
+    float s = ss[m];
+
+    p += __shfl_xor_sync(0xffffffff, p, 16);
+    s += __shfl_xor_sync(0xffffffff, s, 16);
+    p += __shfl_xor_sync(0xffffffff, p, 8);
+    s += __shfl_xor_sync(0xffffffff, s, 8);
+    p += __shfl_xor_sync(0xffffffff, p, 4);
+    s += __shfl_xor_sync(0xffffffff, s, 4);
+    p += __shfl_xor_sync(0xffffffff, p, 2);
+    s += __shfl_xor_sync(0xffffffff, s, 2);
+    p += __shfl_xor_sync(0xffffffff, p, 1);
+    s += __shfl_xor_sync(0xffffffff, s, 1);
+
+    if (laneId == 0) {
+      sm_partial[m][warpId] = p;
+      sm_ss[m][warpId] = s;
+    }
+  }
+
+  __syncthreads();
+
+  // ---- Phase 3: tid 0 finalises the reduction, writes logits, stashes
+  //               rsqrt[m] in smem for phase 4. ----
+  if (tid == 0) {
+#pragma unroll
+    for (int m = 0; m < kNumTokens; m++) {
+      float p_sum = 0.0f;
+      float s_sum = 0.0f;
+#pragma unroll
+      for (int w = 0; w < kNumWarps; w++) {
+        p_sum += sm_partial[m][w];
+        s_sum += sm_ss[m][w];
+      }
+      // Order matches layernorm_kernels.cu: rsqrtf(variance / H + eps).
+      // Use division (not multiply-by-reciprocal) to avoid an extra ULP
+      // mismatch with the reference RMSNorm.
+      float rs = rsqrtf(s_sum / static_cast<float>(kHiddenDim) + eps);
+      s_rsqrt[m] = rs;
+      logits[m * kNumExperts + n_idx] = p_sum * rs;
+    }
+  }
+
+  __syncthreads();
+
+  // ---- Phase 4: spread normed_x writes across blocks 0..kNumTokens-1.
+  //              Each writer block handles exactly one token row,
+  //              avoiding the long tail of block 0 doing all M rows.
+  //              Every block has every token's rsqrt[] in s_rsqrt
+  //              already (computed independently in phase 3), so no
+  //              cross-block synchronization is required. ----
+  if (n_idx < kNumTokens) {
+    int const m_writer = n_idx;
+    float const rs = s_rsqrt[m_writer];
+    __nv_bfloat16 const* x_row = x + m_writer * kHiddenDim;
+    __nv_bfloat16* normed_row = normed_x + m_writer * kHiddenDim;
+
+#pragma unroll
+    for (int ki = 0; ki < k_iterations; ki++) {
+      int const k_base = k_bases[ki];
+
+      uint4 nw_vec = *reinterpret_cast<uint4 const*>(norm_weight + k_base);
+      float nw_f[VPT];
+      bf16_uint4_to_float8<VPT>(nw_vec, nw_f);
+
+      uint4 a_vec = *reinterpret_cast<uint4 const*>(x_row + k_base);
+      float a_f[VPT];
+      bf16_uint4_to_float8<VPT>(a_vec, a_f);
+
+      uint4 normed_vec;
+      __nv_bfloat16* np = reinterpret_cast<__nv_bfloat16*>(&normed_vec);
+#pragma unroll
+      for (int k = 0; k < VPT; k++) {
+        np[k] = __float2bfloat16(a_f[k] * rs * nw_f[k]);
+      }
+      *reinterpret_cast<uint4*>(normed_row + k_base) = normed_vec;
+    }
+  }
+
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
+  asm volatile("griddepcontrol.launch_dependents;");
+#endif
+}
+
+}  // namespace
+
+template <typename T, int kNumTokens, int kNumExperts, int kHiddenDim>
+void invokeNormRouterGemm(float* logits, __nv_bfloat16* normed_x, T const* x,
+                          T const* norm_weight, T const* gate_weight, float eps,
+                          cudaStream_t stream) {
+  constexpr int VPT = 16 / sizeof(T);
+  constexpr int kBlockSize = 128;
+
+  cudaLaunchConfig_t config;
+  config.gridDim = kNumExperts;
+  config.blockDim = kBlockSize;
+  config.dynamicSmemBytes = 0;
+  config.stream = stream;
+
+  cudaLaunchAttribute attrs[1];
+  attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+  attrs[0].val.programmaticStreamSerializationAllowed = 1;
+  config.numAttrs = 1;
+  config.attrs = attrs;
+
+  cudaLaunchKernelEx(&config,
+                     norm_router_gemm_kernel<T, kBlockSize, VPT, kNumTokens,
+                                             kNumExperts, kHiddenDim>,
+                     logits, normed_x, x, norm_weight, gate_weight, eps);
+}
+
+// Template instantiations — DSV4-Pro is the only supported configuration:
+// num_experts=384, hidden_dim=7168.  Other shapes (e.g. DSV4-Flash with
+// hidden_dim=4096) fall back to the unfused path on the Python side.
+#define INSTANTIATE(M)                                                    \
+  template void invokeNormRouterGemm<__nv_bfloat16, M, 384, 7168>(        \
+      float*, __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, \
+      __nv_bfloat16 const*, float, cudaStream_t);
+
+INSTANTIATE(1)
+INSTANTIATE(2)
+INSTANTIATE(3)
+INSTANTIATE(4)
+INSTANTIATE(5)
+INSTANTIATE(6)
+INSTANTIATE(7)
+INSTANTIATE(8)
+INSTANTIATE(9)
+INSTANTIATE(10)
+INSTANTIATE(11)
+INSTANTIATE(12)
+INSTANTIATE(13)
+INSTANTIATE(14)
+INSTANTIATE(15)
+INSTANTIATE(16)
+
+#undef INSTANTIATE
@@ -75,4 +75,12 @@ void shuffle_rows(const torch::Tensor& input_tensor,
 // Supports num_tokens in [1, 16], num_experts in {256, 384}, hidden_dim = 7168
 void dsv3_router_gemm(torch::Tensor& output, const torch::Tensor& mat_a,
                      const torch::Tensor& mat_b);
+
+// Fused RMSNorm + router GEMV for DeepSeek V4. Produces both:
+//   normed_x[m,k]      = x[m,k] * rsqrt(mean(x[m]^2) + eps) * norm_weight[k]
+//   logits[m,n]        = sum_k(normed_x[m,k] * gate_weight[n,k])
+// in a single kernel launch. Same dim/dtype constraints as dsv3_router_gemm.
+void dsv4_norm_router_gemm(at::Tensor& logits, at::Tensor& normed_x,
+                           at::Tensor const& x, at::Tensor const& norm_weight,
+                           at::Tensor const& gate_weight, double eps);
 #endif
@@ -298,130 +298,131 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
    asm volatile("griddepcontrol.launch_dependents;");
 #endif
    return;
-  }
-
+  } else {
 #pragma unroll
-  for (int ii = 0; ii < VPT; ++ii) {
-    float val = row_chunk[ii];
-    float val_b = val * beta;
-    // Compute softplus: log(1 + exp(val)) with numerical stability
-    // When val > threshold, softplus(x) ≈ x to avoid exp overflow
-    val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
-    val = sqrtf(val);
-    if (correction_bias) {
-      const int group_id = ii / ELTS_PER_LDG;
-      const int local_id = ii % ELTS_PER_LDG;
-      const int expert_idx = first_elt_read_by_thread +
-                             group_id * THREADS_PER_ROW * ELTS_PER_LDG +
-                             local_id;
-      val = val + correction_bias[expert_idx];
+    for (int ii = 0; ii < VPT; ++ii) {
+      float val = row_chunk[ii];
+      float val_b = val * beta;
+      // Compute softplus: log(1 + exp(val)) with numerical stability
+      // When val > threshold, softplus(x) ≈ x to avoid exp overflow
+      val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
+      val = sqrtf(val);
+      if (correction_bias) {
+        const int group_id = ii / ELTS_PER_LDG;
+        const int local_id = ii % ELTS_PER_LDG;
+        const int expert_idx = first_elt_read_by_thread +
+                               group_id * THREADS_PER_ROW * ELTS_PER_LDG +
+                               local_id;
+        val = val + correction_bias[expert_idx];
+      }
+      row_chunk[ii] = val;
    }
-    row_chunk[ii] = val;
-  }

-  // Original TopK path: find top-k experts by score
-  // Now, sigmoid_res contains the sigmoid of the row chunk. Now, I want to find
-  // the topk elements in each row, along with the max index.
-  int start_col = first_elt_read_by_thread;
-  static constexpr int COLS_PER_GROUP_LDG = ELTS_PER_LDG * THREADS_PER_ROW;
+    // Original TopK path: find top-k experts by score
+    // Now, sigmoid_res contains the sigmoid of the row chunk. Now, I want to
+    // find the topk elements in each row, along with the max index.
+    int start_col = first_elt_read_by_thread;
+    static constexpr int COLS_PER_GROUP_LDG = ELTS_PER_LDG * THREADS_PER_ROW;

-  float selected_sum = 0.f;
-  for (int k_idx = 0; k_idx < k; ++k_idx) {
-    // First, each thread does the local argmax
-    float max_val = row_chunk[0];
-    int expert = start_col;
+    float selected_sum = 0.f;
+    for (int k_idx = 0; k_idx < k; ++k_idx) {
+      // First, each thread does the local argmax
+      float max_val = row_chunk[0];
+      int expert = start_col;
 #pragma unroll
-    for (int ldg = 0, col = start_col; ldg < LDG_PER_THREAD;
-         ++ldg, col += COLS_PER_GROUP_LDG) {
+      for (int ldg = 0, col = start_col; ldg < LDG_PER_THREAD;
+           ++ldg, col += COLS_PER_GROUP_LDG) {
 #pragma unroll
-      for (int ii = 0; ii < ELTS_PER_LDG; ++ii) {
-        float val = row_chunk[ldg * ELTS_PER_LDG + ii];
+        for (int ii = 0; ii < ELTS_PER_LDG; ++ii) {
+          float val = row_chunk[ldg * ELTS_PER_LDG + ii];

-        // No check on the experts here since columns with the smallest index
-        // are processed first and only updated if > (not >=)
-        if (val > max_val) {
-          max_val = val;
-          expert = col + ii;
+          // No check on the experts here since columns with the smallest index
+          // are processed first and only updated if > (not >=)
+          if (val > max_val) {
+            max_val = val;
+            expert = col + ii;
+          }
        }
      }
-    }

 // Now, we perform the argmax reduce. We use the butterfly pattern so threads
 // reach consensus about the max. This will be useful for K > 1 so that the
 // threads can agree on "who" had the max value. That thread can then blank out
 // their max with -inf and the warp can run more iterations...
 #pragma unroll
-    for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) {
-      float other_max =
-          VLLM_SHFL_XOR_SYNC_WIDTH(max_val, mask, THREADS_PER_ROW);
-      int other_expert =
-          VLLM_SHFL_XOR_SYNC_WIDTH(expert, mask, THREADS_PER_ROW);
+      for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) {
+        float other_max =
+            VLLM_SHFL_XOR_SYNC_WIDTH(max_val, mask, THREADS_PER_ROW);
+        int other_expert =
+            VLLM_SHFL_XOR_SYNC_WIDTH(expert, mask, THREADS_PER_ROW);

-      // We want lower indices to "win" in every thread so we break ties this
-      // way
-      if (other_max > max_val ||
-          (other_max == max_val && other_expert < expert)) {
-        max_val = other_max;
-        expert = other_expert;
+        // We want lower indices to "win" in every thread so we break ties this
+        // way
+        if (other_max > max_val ||
+            (other_max == max_val && other_expert < expert)) {
+          max_val = other_max;
+          expert = other_expert;
+        }
+      }
+
+      // Write the max for this k iteration to global memory.
+      if (thread_group_idx == 0) {
+        // Add a guard to ignore experts not included by this node
+        const bool node_uses_expert =
+            expert >= start_expert && expert < end_expert;
+        const bool should_process_row = row_is_active && node_uses_expert;
+
+        // The lead thread from each sub-group will write out the final results
+        // to global memory. (This will be a single) thread per row of the
+        // input/output matrices.
+        const int idx = k * thread_row + k_idx;
+        if (correction_bias != nullptr) {
+          max_val -= correction_bias[expert];
+        }
+        output[idx] = max_val;
+        indices[idx] =
+            should_process_row ? (expert - start_expert) : NUM_EXPERTS;
+        source_rows[idx] = k_idx * num_rows + thread_row;
+        if (renormalize) {
+          selected_sum += max_val;
+        }
+      }
+
+      // Finally, we clear the value in the thread with the current max if there
+      // is another iteration to run.
+      if (k_idx + 1 < k) {
+        const int ldg_group_for_expert = expert / COLS_PER_GROUP_LDG;
+        const int thread_to_clear_in_group =
+            (expert / ELTS_PER_LDG) % THREADS_PER_ROW;
+
+        // Only the thread in the group which produced the max will reset the
+        // "winning" value to -inf.
+        if (thread_group_idx == thread_to_clear_in_group) {
+          const int offset_for_expert = expert % ELTS_PER_LDG;
+          // Safe to set to any negative value since row_chunk values must be
+          // between 0 and 1.
+          row_chunk[ldg_group_for_expert * ELTS_PER_LDG + offset_for_expert] =
+              -10000.f;
+        }
      }
    }

-    // Write the max for this k iteration to global memory.
+    // Apply renormalization and routed scaling factor to final weights.
    if (thread_group_idx == 0) {
-      // Add a guard to ignore experts not included by this node
-      const bool node_uses_expert =
-          expert >= start_expert && expert < end_expert;
-      const bool should_process_row = row_is_active && node_uses_expert;
-
-      // The lead thread from each sub-group will write out the final results to
-      // global memory. (This will be a single) thread per row of the
-      // input/output matrices.
-      const int idx = k * thread_row + k_idx;
-      if (correction_bias != nullptr) {
-        max_val -= correction_bias[expert];
-      }
-      output[idx] = max_val;
-      indices[idx] = should_process_row ? (expert - start_expert) : NUM_EXPERTS;
-      source_rows[idx] = k_idx * num_rows + thread_row;
+      float scale = static_cast<float>(routed_scaling_factor);
      if (renormalize) {
-        selected_sum += max_val;
+        const float denom = selected_sum > 0.f ? selected_sum : 1.f;
+        scale /= denom;
+      }
+      for (int k_idx = 0; k_idx < k; ++k_idx) {
+        const int idx = k * thread_row + k_idx;
+        output[idx] = output[idx] * scale;
      }
    }
-
-    // Finally, we clear the value in the thread with the current max if there
-    // is another iteration to run.
-    if (k_idx + 1 < k) {
-      const int ldg_group_for_expert = expert / COLS_PER_GROUP_LDG;
-      const int thread_to_clear_in_group =
-          (expert / ELTS_PER_LDG) % THREADS_PER_ROW;
-
-      // Only the thread in the group which produced the max will reset the
-      // "winning" value to -inf.
-      if (thread_group_idx == thread_to_clear_in_group) {
-        const int offset_for_expert = expert % ELTS_PER_LDG;
-        // Safe to set to any negative value since row_chunk values must be
-        // between 0 and 1.
-        row_chunk[ldg_group_for_expert * ELTS_PER_LDG + offset_for_expert] =
-            -10000.f;
-      }
-    }
-  }
-
-  // Apply renormalization and routed scaling factor to final weights.
-  if (thread_group_idx == 0) {
-    float scale = static_cast<float>(routed_scaling_factor);
-    if (renormalize) {
-      const float denom = selected_sum > 0.f ? selected_sum : 1.f;
-      scale /= denom;
-    }
-    for (int k_idx = 0; k_idx < k; ++k_idx) {
-      const int idx = k * thread_row + k_idx;
-      output[idx] = output[idx] * scale;
-    }
-  }
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-  asm volatile("griddepcontrol.launch_dependents;");
+    asm volatile("griddepcontrol.launch_dependents;");
 #endif
+  }
 }

 namespace detail {
@@ -135,6 +135,12 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
  // DeepSeek V3 optimized router GEMM for SM90+
  m.def("dsv3_router_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
  // conditionally compiled so impl registration is in source file
+
+  // DeepSeek V4 fused RMSNorm + router GEMV for SM90+
+  m.def(
+      "dsv4_norm_router_gemm(Tensor! logits, Tensor! normed_x, Tensor x, "
+      "Tensor norm_weight, Tensor gate_weight, float eps) -> ()");
+  // conditionally compiled so impl registration is in source file
 #endif
 }

@@ -200,19 +200,6 @@ void cutlass_mla_decode(torch::Tensor const& out, torch::Tensor const& q_nope,

 torch::Tensor get_cuda_view_from_cpu_tensor(torch::Tensor& cpu_tensor);

-#ifndef USE_ROCM
-
-torch::Tensor awq_gemm(torch::Tensor _in_feats, torch::Tensor _kernel,
-                       torch::Tensor _scaling_factors, torch::Tensor _zeros,
-                       int64_t split_k_iters);
-
-torch::Tensor awq_dequantize(torch::Tensor _kernel,
-                             torch::Tensor _scaling_factors,
-                             torch::Tensor _zeros, int64_t split_k_iters,
-                             int64_t thx, int64_t thy);
-
-#endif
-
 torch::Tensor ggml_dequantize(torch::Tensor W, int64_t type, int64_t m,
                              int64_t n,
                              std::optional<at::ScalarType> const& dtype);
@@ -302,8 +289,6 @@ std::tuple<int64_t, torch::Tensor> allocate_shared_buffer_and_handle(
 int64_t open_mem_handle(torch::Tensor& mem_handle);
 void free_shared_buffer(int64_t buffer);

-torch::Tensor hadacore_transform(torch::Tensor& x, bool inplace);
-
 #ifdef USE_ROCM
 fptr_t init_custom_qr(int64_t rank, int64_t world_size,
                      std::optional<int64_t> qr_max_size = std::nullopt);
@@ -315,11 +300,6 @@ void qr_all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out,
 int64_t qr_max_size();
 #endif

-#ifndef USE_ROCM
-void dsv3_fused_a_gemm(torch::Tensor& output, torch::Tensor const& mat_a,
-                       torch::Tensor const& mat_b);
-#endif
-
 #ifndef USE_ROCM
 torch::Tensor minimax_allreduce_rms(torch::Tensor const& input,
                                    torch::Tensor const& norm_weight,
@@ -2,10 +2,14 @@

 #ifndef _marlin_cuh
  #define _marlin_cuh
-  #include <torch/all.h>
-
-  #include <ATen/cuda/CUDAContext.h>
-  #include <c10/cuda/CUDAGuard.h>
+  // These torch headers are only needed by non-stable callers (e.g. ops.cu).
+  // Guard them so that stable ABI targets can still include marlin.cuh
+  // for Vec, constants, and cp_async helpers without pulling in torch/all.h.
+  #ifndef TORCH_TARGET_VERSION
+    #include <torch/all.h>
+    #include <ATen/cuda/CUDAContext.h>
+    #include <c10/cuda/CUDAGuard.h>
+  #endif
  #include <cuda.h>
  #include <cuda_fp16.h>
  #include <cuda_runtime.h>
@@ -263,22 +263,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {

  // Quantization ops
 #ifndef USE_ROCM
-  // DeepSeek V3 fused A GEMM (SM 9.0+, bf16 only, 1-16 tokens).
-  ops.def(
-      "dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
-  // conditionally compiled so impl registration is in source file
-
-  // Quantized GEMM for AWQ.
-  ops.def(
-      "awq_gemm(Tensor _in_feats, Tensor _kernel, Tensor _scaling_factors, "
-      "Tensor _zeros, SymInt split_k_iters) -> Tensor");
-  ops.impl("awq_gemm", torch::kCUDA, &awq_gemm);
-
-  // Dequantization for AWQ.
-  ops.def(
-      "awq_dequantize(Tensor _kernel, Tensor _scaling_factors, "
-      "Tensor _zeros, SymInt split_k_iters, int thx, int thy) -> Tensor");
-  ops.impl("awq_dequantize", torch::kCUDA, &awq_dequantize);

  // Note about marlin kernel 'workspace' arguments:
  // Technically these should be mutable since they are modified by the kernel.
@@ -408,22 +392,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
      " -> ()");
  // conditionally compiled so impl registration is in source file

-  // SM100 CUTLASS MLA decode
-  ops.def(
-      "sm100_cutlass_mla_decode(Tensor! out, Tensor! lse, Tensor q_nope,"
-      "                         Tensor q_pe, Tensor kv_c_and_k_pe_cache,"
-      "                         Tensor seq_lens, Tensor page_table,"
-      "                         Tensor workspace, float scale,"
-      "                         int num_kv_splits) -> ()");
-  // conditionally compiled so impl in source file
-
-  // SM100 CUTLASS MLA workspace
-  ops.def(
-      "sm100_cutlass_mla_get_workspace_size(int max_seq_len, int num_batches,"
-      "                                     int sm_count, int num_kv_splits) "
-      "-> int");
-  // conditionally compiled so impl in source file
-
 #endif

  // Quantized GEMM for GPTQ.
@@ -496,26 +464,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
      "Tensor? last_chunk_indices) -> ()");
  ops.impl("selective_scan_fwd", torch::kCUDA, &selective_scan_fwd);

-  // Hadamard transforms
-  ops.def("hadacore_transform(Tensor! x, bool inplace) -> Tensor");
-
 #ifndef USE_ROCM
-  // reorder weight for AllSpark Ampere W8A16 Fused Gemm kernel
-  ops.def(
-      "rearrange_kn_weight_as_n32k16_order(Tensor b_qweight, Tensor b_scales, "
-      "Tensor? b_zeros, "
-      "bool has_zp, Tensor! b_qweight_reorder, Tensor! b_scales_reorder, "
-      "Tensor!? b_zeros_reorder, "
-      "int K, int N, int N_32align) -> ()");
-  //  conditionally compiled so impl in source file
-
-  // AllSpark quantization ops
-  ops.def(
-      "allspark_w8a16_gemm(Tensor a, Tensor b_qweight, Tensor b_scales, "
-      "Tensor? b_qzeros, "
-      "SymInt n, SymInt group_size, SymInt sm_count, SymInt sm_version, SymInt "
-      "CUBLAS_M_THRESHOLD, bool has_zp, bool n32k16_reorder) -> Tensor");
-
  ops.def(
      "minimax_allreduce_rms("
      "Tensor input,"
@@ -199,7 +199,10 @@ COPY requirements/cuda.txt requirements/cuda.txt
 COPY use_existing_torch.py use_existing_torch.py
 COPY pyproject.toml pyproject.toml
 RUN --mount=type=cache,target=/root/.cache/uv \
-    if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
+    if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "12" ]; then \
+        sed -i 's/^nvidia-cutlass-dsl\[cu13\]>=/nvidia-cutlass-dsl>=/' requirements/cuda.txt; \
+    fi \
+    && if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
        echo "Installing torch nightly..." \
        && uv pip install --python /opt/venv/bin/python3 torch torchaudio torchvision --pre \
        --index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/nightly/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
@@ -301,6 +304,14 @@ RUN --mount=type=cache,target=/root/.cache/uv \
        python3 use_existing_torch.py --prefix; \
    fi

+# Provision one bare Python per `requires-python` entry; cmake reads
+# DEEPGEMM_PYTHON_INTERPRETERS to build DeepGEMM `_C` for each. See
+# cmake/external_projects/deepgemm.cmake for the full picture.
+COPY tools/setup_deepgemm_pythons.sh tools/build_deepgemm_C.py tools/
+ENV DEEPGEMM_VENV_PREFIX=/opt/dgenv
+RUN --mount=type=cache,target=/root/.cache/uv \
+    tools/setup_deepgemm_pythons.sh > /tmp/dg_pythons.txt
+
 # Build the vLLM wheel
 # if USE_SCCACHE is set, use sccache to speed up compilation
 # AWS credentials mounted at ~/.aws/credentials for sccache S3 auth (optional)
@@ -328,6 +339,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
        && export VLLM_PRECOMPILED_WHEEL_COMMIT="${VLLM_MERGE_BASE_COMMIT}" \
        && export VLLM_MAIN_CUDA_VERSION="${VLLM_MAIN_CUDA_VERSION}" \
        && export VLLM_DOCKER_BUILD_CONTEXT=1 \
+        && export DEEPGEMM_PYTHON_INTERPRETERS=$(cat /tmp/dg_pythons.txt) \
        && sccache --show-stats \
        && python3 setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38 \
        && sccache --show-stats; \
@@ -345,6 +357,7 @@ RUN --mount=type=cache,target=/root/.cache/ccache \
        export VLLM_USE_PRECOMPILED="${VLLM_USE_PRECOMPILED}" && \
        export VLLM_PRECOMPILED_WHEEL_COMMIT="${VLLM_MERGE_BASE_COMMIT}" && \
        export VLLM_DOCKER_BUILD_CONTEXT=1 && \
+        export DEEPGEMM_PYTHON_INTERPRETERS=$(cat /tmp/dg_pythons.txt) && \
        python3 setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38; \
    fi

@@ -616,6 +629,9 @@ ARG PYTORCH_CUDA_INDEX_BASE_URL
 COPY requirements/common.txt /tmp/common.txt
 COPY requirements/cuda.txt /tmp/requirements-cuda.txt
 RUN --mount=type=cache,target=/root/.cache/uv \
+    if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "12" ]; then \
+        sed -i 's/^nvidia-cutlass-dsl\[cu13\]>=/nvidia-cutlass-dsl>=/' /tmp/requirements-cuda.txt; \
+    fi && \
    uv pip install --system -r /tmp/requirements-cuda.txt \
        --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') && \
    rm /tmp/requirements-cuda.txt /tmp/common.txt
@@ -829,9 +845,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
                libcublas-dev-${CUDA_VERSION_DASH} \
                libcusolver-dev-${CUDA_VERSION_DASH}"; \
    if [ "$INSTALL_KV_CONNECTORS" = "true" ]; then \
-        if [ "$CUDA_MAJOR" -ge 13 ]; then \
-            uv pip install --system nixl-cu13; \
-        fi; \
        uv pip install --system -r /tmp/kv_connectors.txt --no-build || ( \
            # if the above fails, install from source
            apt-get update -y && \
@@ -841,6 +854,8 @@ RUN --mount=type=cache,target=/root/.cache/uv \
            # clean up -dev packages, keep runtime libraries
            rm -rf /var/lib/apt/lists/* \
        ); \
+        # Force-reinstall the matching CUDA wheel so the correct nixl_ep_cpp.so is installed.
+        uv pip install --system --force-reinstall --no-deps nixl-cu${CUDA_MAJOR}; \
    fi

 ENV VLLM_USAGE_SOURCE production-docker-image
@@ -2,10 +2,18 @@
 ARG REMOTE_VLLM="0"
 ARG COMMON_WORKDIR=/app
 ARG BASE_IMAGE=rocm/vllm-dev:base
-# AMD NIC backend
-ARG NIC_BACKEND=none
+# NIC backend for MoRI RDMA support.
+# By default (all), drivers and userspace libraries for all supported NIC types
+# (ainic and bnxt) are installed; MoRI selects the appropriate one at runtime.
+# To install drivers for a single NIC type only, set NIC_BACKEND explicitly:
+#   --build-arg NIC_BACKEND=ainic   # AMD AINIC (Pensando) only
+#   --build-arg NIC_BACKEND=bnxt    # Broadcom Thor-2 only
+#   --build-arg NIC_BACKEND=none    # Install nothing.
+ARG NIC_BACKEND=all
 # AMD AINIC apt repo settings
-ARG AINIC_VERSION=1.117.5
+# Users can specify a custom version compatible with their host drivers.
+# The default version has been tested with ioinic-dkms=25.11.1.001
+ARG AINIC_VERSION=1.117.3-hydra
 ARG UBUNTU_CODENAME=jammy

 # Sccache configuration (only used in release pipeline)
@@ -231,39 +239,52 @@ ARG NIC_BACKEND
 ARG AINIC_VERSION
 ARG UBUNTU_CODENAME
 RUN /bin/bash -lc 'set -euo pipefail; \
+ \
+ install_ainic() { \
+   apt-get update && apt-get install -y --no-install-recommends ca-certificates curl gnupg apt-transport-https; \
+   rm -rf /var/lib/apt/lists/*; \
+   mkdir -p /etc/apt/keyrings; \
+   curl -fsSL https://repo.radeon.com/rocm/rocm.gpg.key | gpg --dearmor > /etc/apt/keyrings/amdainic.gpg; \
+   echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/amdainic.gpg] https://repo.radeon.com/amdainic/pensando/ubuntu/${AINIC_VERSION} ${UBUNTU_CODENAME} main" \
+     > /etc/apt/sources.list.d/amdainic.list; \
+   apt-get update && apt-get install -y --no-install-recommends \
+     libionic-dev \
+     ionic-common \
+   ; \
+   rm -rf /var/lib/apt/lists/*; \
+ }; \
+ \
+ # NOTE: requires FW 235.2.86.0 and kernel drivers on the host: \
+ #   bnxt-en-dkms=1.10.3.235.2.86.0 bnxt-re-dkms=235.2.86.0 (from packages.broadcom.com PPA) \
+ install_bnxt() { \
+   install -m 0755 -d /etc/apt/keyrings; \
+   curl -fsSL https://packages.broadcom.com/artifactory/api/security/keypair/PackagesKey/public \
+     -o /etc/apt/keyrings/broadcom-nic.asc; \
+   chmod a+r /etc/apt/keyrings/broadcom-nic.asc; \
+   echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/broadcom-nic.asc] https://packages.broadcom.com/artifactory/ethernet-nic-debian-public jammy main" \
+     > /etc/apt/sources.list.d/broadcom-nic.list; \
+   apt-get update && apt-get install -y --no-install-recommends \
+     bnxt-rocelib=235.2.86.0 \
+   ; \
+   cp -a /usr/local/lib/x86_64-linux-gnu/libbnxt_re* /usr/local/lib/; \
+   ldconfig; \
+   rm -rf /var/lib/apt/lists/*; \
+ }; \
+ \
 echo "[MORI] Install MoRI proxy deps"; \
 pip install --quiet --ignore-installed blinker && \
 pip install --quiet quart msgpack aiohttp pyzmq; \
 echo "[MORI] NIC_BACKEND=${NIC_BACKEND}"; \
-  \
-  # NIC backend deps — mori auto-detects NIC at runtime (MORI_DEVICE_NIC env var override).
-  # Only vendor packages are installed here for dlopen (e.g. libionic.so); no compile-time flags needed.
-  case "${NIC_BACKEND}" in \
-    # default: mlx5
-    none) \
-      ;; \
-    # AMD NIC
-    ainic) \
-      apt-get update && apt-get install -y --no-install-recommends ca-certificates curl gnupg apt-transport-https && \
-      rm -rf /var/lib/apt/lists/* && mkdir -p /etc/apt/keyrings; \
-      curl -fsSL https://repo.radeon.com/rocm/rocm.gpg.key | gpg --dearmor > /etc/apt/keyrings/amdainic.gpg; \
-      echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/amdainic.gpg] https://repo.radeon.com/amdainic/pensando/ubuntu/${AINIC_VERSION} ${UBUNTU_CODENAME} main" \
-        > /etc/apt/sources.list.d/amdainic.list; \
-      apt-get update && apt-get install -y --no-install-recommends \
-          libionic-dev \
-          ionic-common \
-      ; \
-      rm -rf /var/lib/apt/lists/*; \
-      ;; \
-    # TODO: Add Broadcom bnxt packages/repos here later.
-    # bnxt) \
-    #   echo "[MORI] Add Broadcom bnxt packages/repos here later."; \
-    #   ;; \
-    *) \
-      echo "ERROR: unknown NIC_BACKEND=${NIC_BACKEND}. Use one of: none, ainic"; \
-      exit 2; \
-      ;; \
-  esac;'
+ \
+ # NIC backend deps — mori auto-detects NIC at runtime (MORI_DEVICE_NIC env var override). \
+ # Only vendor packages are installed here for dlopen; no compile-time flags needed. \
+ case "${NIC_BACKEND}" in \
+   none)  ;; \
+   all)   install_ainic; install_bnxt ;; \
+   ainic) install_ainic ;; \
+   bnxt)  install_bnxt ;; \
+   *)     echo "ERROR: unknown NIC_BACKEND=${NIC_BACKEND}. Use one of: none, ainic, bnxt, all"; exit 2 ;; \
+ esac'

 # -----------------------
 # vLLM wheel release build stage (for building distributable wheels)
@@ -34,6 +34,7 @@ you may contact the following individuals:

 - Simon Mo - <simon.mo@hey.com>
 - Russell Bryant - <rbryant@redhat.com>
+- Juan Pérez de Algaba - <jperezde@redhat.com>
 - Huzaifa Sidhpurwala - <huzaifas@redhat.com>

 ## Slack Discussion
@@ -2,4 +2,4 @@

 vLLM can be deployed with [llm-d](https://github.com/llm-d/llm-d), a Kubernetes-native distributed inference serving stack providing well-lit paths for anyone to serve large generative AI models at scale. It helps achieve the fastest "time to state-of-the-art (SOTA) performance" for key OSS models across most hardware accelerators and infrastructure providers.

-You can use vLLM with llm-d directly by following [this guide](https://llm-d.ai/docs/guide) or via [KServe's LLMInferenceService](https://kserve.github.io/website/docs/model-serving/generative-inference/llmisvc/llmisvc-overview).
+You can use vLLM with llm-d directly by following [the official guides](https://llm-d.ai/docs/guides) or via [KServe's LLMInferenceService](https://kserve.github.io/website/docs/model-serving/generative-inference/llmisvc/llmisvc-overview).
@@ -125,12 +125,13 @@ Priority is **1 = highest** (tried first).
 | Priority | Backend |
 | -------- | ------- |
 | 1 | `FLASHINFER_MLA` |
-| 2 | `CUTLASS_MLA` |
-| 3 | `FLASH_ATTN_MLA` |
-| 4 | `FLASHMLA` |
-| 5 | `TRITON_MLA` |
-| 6 | `FLASHINFER_MLA_SPARSE`**\*** |
-| 7 | `FLASHMLA_SPARSE` |
+| 2 | `TOKENSPEED_MLA` |
+| 3 | `CUTLASS_MLA` |
+| 4 | `FLASH_ATTN_MLA` |
+| 5 | `FLASHMLA` |
+| 6 | `TRITON_MLA` |
+| 7 | `FLASHINFER_MLA_SPARSE`**\*** |
+| 8 | `FLASHMLA_SPARSE` |

 **Ampere/Hopper (SM 8.x-9.x):**

@@ -202,6 +203,7 @@ hardware and configuration.
 | `FLASH_ATTN`‡ | FlashAttention varlen (FA2/FA3/FA4) | fp16, bf16 | Any | FA4 on SM100+, FA3 on SM90, FA2 otherwise |
 | `TRTLLM_RAGGED` | TensorRT-LLM ragged attention | fp16, bf16 | 10.x | DeepSeek R1 dims only |
 | `FLASHINFER` | FlashInfer CUTLASS backend | fp16, bf16 | 10.x | DeepSeek R1 dims only |
+| `TOKENSPEED_MLA` | | fp16, bf16 | 10.x | DeepSeek R1 dims only |

 > **‡** TRT-LLM Ragged is the default on Blackwell (SM100).
 > On other GPUs, FlashAttention is used as the default.
@@ -222,5 +224,6 @@ MLA decode backends are selected using the standard
 | `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %1 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
 | `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 1, 64 | Any | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
 | `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
+| `TOKENSPEED_MLA` | fp16, bf16 | `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
 | `TRITON_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
 | `XPU_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | Any | 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | Any |
@@ -85,12 +85,14 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra

 | Architecture | Models | CG for Image | CG for Video |
 | ------------ | ------ | ------------ | ------------ |
-| `Qwen3VLForConditionalGeneration` | `Qwen3-VL` | ✅︎ | ✅︎ |
+| `Qwen2VLForConditionalGeneration` | `Qwen2-VL` | ✅︎ | ✅︎ |
 | `Qwen2_5_VLForConditionalGeneration` | `Qwen2.5-VL` | ✅︎ | ✅︎ |
+| `Qwen3VLForConditionalGeneration` | `Qwen3-VL` | ✅︎ | ✅︎ |
+| `Qwen3_5ForConditionalGeneration` | `Qwen3.5` | ✅︎ | ✅︎ |

 !!! note
    Encoder CUDA Graphs have currently been tested with `--mm-encoder-attn-backend=FLASH_ATTN` and `--mm-encoder-attn-backend=FLASHINFER` on Blackwell GPUs.
-    For Qwen2.5-VL only FA2 and FA3 has been tested.
+    For Qwen2-VL and Qwen2.5-VL only FA2 and FA3 has been tested.

 ## Configuration

@@ -81,7 +81,7 @@ To be used with a particular `FusedMoEPrepareAndFinalizeModular` subclass, MoE k
 | Kernel | Input act. format | Quant. types | Quant. format | Activation function | Apply Weight On Input | Modular | Source |
 | ------ | ----------------- | ------------ | ------------- | ------------------- | --------------------- | ------- | ------ |
 | triton | standard | all<sup>1</sup> | G,A,T | silu, gelu,</br>swigluoai,</br>silu_no_mul,</br>gelu_no_mul | Y | Y | [`fused_experts`][vllm.model_executor.layers.fused_moe.fused_moe.fused_experts],</br>[`TritonExperts`][vllm.model_executor.layers.fused_moe.experts.triton_moe.TritonExperts] |
-| triton (batched) | batched | all<sup>1</sup> | G,A,T | silu, gelu | <sup>6</sup> | Y | [`BatchedTritonExperts`][vllm.model_executor.layers.fused_moe.fused_batched_moe.BatchedTritonExperts] |
+| triton (batched) | batched | all<sup>1</sup> | G,A,T | silu, gelu | <sup>6</sup> | Y | [`BatchedTritonExperts`][vllm.model_executor.layers.fused_moe.experts.fused_batched_moe.BatchedTritonExperts] |
 | deep gemm | standard,</br>batched | fp8 | G(128),A,T | silu, gelu | <sup>6</sup> | Y | </br>[`DeepGemmExperts`][vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe.DeepGemmExperts],</br>[`BatchedDeepGemmExperts`][vllm.model_executor.layers.fused_moe.experts.batched_deep_gemm_moe.BatchedDeepGemmExperts] |
 | cutlass_fp4 | standard,</br>batched | nvfp4 | A,T | silu | Y | Y | [`CutlassExpertsFp4`][vllm.model_executor.layers.fused_moe.experts.cutlass_moe.CutlassExpertsFp4] |
 | cutlass_fp8 | standard,</br>batched | fp8 | A,T | silu, gelu | Y | Y | [`CutlassExpertsFp8`][vllm.model_executor.layers.fused_moe.experts.cutlass_moe.CutlassExpertsFp8],</br>[`CutlasBatchedExpertsFp8`][vllm.model_executor.layers.fused_moe.experts.cutlass_moe.CutlassBatchedExpertsFp8] |
@@ -91,7 +91,7 @@ To be used with a particular `FusedMoEPrepareAndFinalizeModular` subclass, MoE k
 | 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` |
 | cpu_fused_moe | standard | N/A | N/A | silu | N | N | [`CPUFusedMOE`][vllm.model_executor.layers.fused_moe.cpu_fused_moe.CPUFusedMOE] |
-| naive batched<sup>4</sup> | batched | int8,</br>fp8 | G,A,T | silu, gelu | <sup>6</sup> | Y | [`NaiveBatchedExperts`][vllm.model_executor.layers.fused_moe.fused_batched_moe.NaiveBatchedExperts] |
+| naive batched<sup>4</sup> | batched | int8,</br>fp8 | G,A,T | silu, gelu | <sup>6</sup> | Y | [`NaiveBatchedExperts`][vllm.model_executor.layers.fused_moe.experts.fused_batched_moe.NaiveBatchedExperts] |

 !!! info "Table key"
    1. All types: mxfp4, nvfp4, int4, int8, fp8
@@ -0,0 +1,161 @@
+# MooncakeStoreConnector Usage Guide
+
+MooncakeStoreConnector is a KV cache connector that uses [MooncakeDistributedStore](https://github.com/kvcache-ai/Mooncake) as a shared KV cache pool. Unlike `MooncakeConnector` which does direct point-to-point KV transfer between prefiller and decoder, MooncakeStoreConnector enables KV cache offloading to an external distributed store, supporting:
+
+- **CPU offloading**: Extend effective KV cache capacity by offloading to CPU memory via Mooncake's transfer engine.
+- **Prefix caching across instances**: Hash-based deduplication allows multiple vLLM instances to share cached KV blocks through the store.
+- **Single-node and multi-node deployment**: Works both as a standalone KV cache extension and in disaggregated prefill-decode setups.
+
+## Prerequisites
+
+### Install Mooncake
+
+Install mooncake through pip:
+
+```bash
+uv pip install mooncake-transfer-engine
+```
+
+Refer to the [Mooncake official repository](https://github.com/kvcache-ai/Mooncake) for more installation instructions and building from source.
+
+### Start the Mooncake Master Server
+
+The Mooncake master manages metadata and coordinates the distributed store. Start it before launching vLLM:
+
+```bash
+mooncake_master --port 50051
+```
+
+Default ports:
+
+- RPC: 50051
+
+Multiple vLLM instances can share the same master server.
+
+### Configure Mooncake
+
+Create a JSON configuration file (e.g., `mooncake_config.json`):
+
+```json
+{
+  "metadata_server": "P2PHANDSHAKE",
+  "master_server_address": "127.0.0.1:50051",
+  "global_segment_size": "80GB",
+  "local_buffer_size": "4GB",
+  "protocol": "rdma",
+  "device_name": ""
+}
+```
+
+- `protocol`: Use `"rdma"` for best performance. `"tcp"` works as a fallback.
+- `global_segment_size`: CPU memory contributed to the distributed pool (per GPU).
+- `local_buffer_size`: Private buffer for this node's own operations (per GPU).
+
+Set the config path via environment variable:
+
+```bash
+export MOONCAKE_CONFIG_PATH=/path/to/mooncake_config.json
+```
+
+## Usage
+
+### Single-Node KV Cache Offloading
+
+Use MooncakeStoreConnector to offload KV cache to CPU memory, extending the effective cache size:
+
+```bash
+MOONCAKE_CONFIG_PATH=mooncake_config.json \
+vllm serve meta-llama/Llama-3.1-8B-Instruct \
+    --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_both"}'
+```
+
+### Disaggregated Prefill-Decode (XpYd)
+
+In disaggregated prefill-decode mode, use `MultiConnector` to combine `MooncakeConnector` (point-to-point KV transfer) with `MooncakeStoreConnector` (shared KV cache pool). This enables both direct P2P transfer between prefiller and decoder, and cross-instance prefix cache sharing via the distributed store.
+**Prefiller Node:**
+
+```bash
+MOONCAKE_CONFIG_PATH=mooncake_config.json \
+VLLM_MOONCAKE_BOOTSTRAP_PORT=50052 \
+vllm serve meta-llama/Llama-3.1-8B-Instruct \
+    --port 8100 \
+    --kv-transfer-config '{
+        "kv_connector": "MultiConnector",
+        "kv_role": "kv_producer",
+        "kv_connector_extra_config": {
+            "connectors": [
+                {
+                    "kv_connector": "MooncakeConnector",
+                    "kv_role": "kv_producer"
+                },
+                {
+                    "kv_connector": "MooncakeStoreConnector",
+                    "kv_role": "kv_producer"
+                }
+            ]
+        }
+    }'
+```
+
+**Decoder Node:**
+
+```bash
+MOONCAKE_CONFIG_PATH=mooncake_config.json \
+VLLM_MOONCAKE_BOOTSTRAP_PORT=50053 \
+vllm serve meta-llama/Llama-3.1-8B-Instruct \
+    --port 8200 \
+    --kv-transfer-config '{
+        "kv_connector": "MultiConnector",
+        "kv_role": "kv_consumer",
+        "kv_connector_extra_config": {
+            "connectors": [
+                {
+                    "kv_connector": "MooncakeConnector",
+                    "kv_role": "kv_consumer"
+                },
+                {
+                    "kv_connector": "MooncakeStoreConnector",
+                    "kv_role": "kv_consumer"
+                }
+            ]
+        }
+    }'
+```
+
+**Proxy:**
+
+A disaggregation proxy is required to route requests between prefiller and decoder nodes. The proxy assigns `do_remote_prefill=True` / `do_remote_decode=True` to coordinate P2P transfer via `MooncakeConnector`. Refer to the [MooncakeConnector usage guide](mooncake_connector_usage.md) for proxy setup details.
+
+## Environment Variables
+
+| Variable | Description | Default |
+| --- | --- | --- |
+| `MOONCAKE_CONFIG_PATH` | Path to Mooncake JSON config file | (required) |
+| `VLLM_MOONCAKE_BOOTSTRAP_PORT` | Bootstrap port for MooncakeConnector P2P transfer (disagg mode only) | 8998 |
+
+## KV Transfer Config
+
+### KV Role Options
+
+- **kv_producer**: For prefiller instances that store KV caches to the pool.
+- **kv_consumer**: For decoder instances that load KV caches from the pool.
+- **kv_both**: The instance both stores and loads KV caches. Use this for single-node CPU offloading.
+
+### kv_connector_extra_config
+
+- `load_async` (bool): Enable asynchronous loading for better compute-I/O overlap. Default: `true`.
+- `enable_cross_layers_blocks` (bool): Enable cross-layer block packing for reduced store operations. Default: `false`.
+- `discard_partial_chunks` (bool): Discard partial block chunks during store. Default: `true`.
+- `lookup_rpc_port` (int): Custom port for the ZMQ lookup RPC socket. Default: `0`.
+
+## Notes
+
+### Cross-DP Prefix Cache Hits
+
+When running with data parallelism, set a fixed `PYTHONHASHSEED` so that block hashes are consistent across DP ranks:
+
+```bash
+PYTHONHASHSEED=0 vllm serve ...
+```
+
+Without this, identical prompts may produce different block hashes on different DP ranks, preventing cross-instance prefix cache hits.
@@ -84,7 +84,7 @@ Since simple RTN does not require data for weight quantization and the activatio
 Install `vllm` and `lm-evaluation-harness` for evaluation:

 ```bash
-pip install vllm "lm-eval[api]>=0.4.11"
+pip install vllm "lm-eval[api]>=0.4.12"
 ```

 Load and run the model in `vllm`:
@@ -18,7 +18,7 @@ pip install llmcompressor
 Additionally, install `vllm` and `lm-evaluation-harness` for evaluation:

 ```bash
-pip install vllm "lm-eval[api]>=0.4.11"
+pip install vllm "lm-eval[api]>=0.4.12"
 ```

 ## Quantization Process
@@ -23,7 +23,7 @@ pip install llmcompressor
 Additionally, install `vllm` and `lm-evaluation-harness` for evaluation:

 ```bash
-pip install vllm "lm-eval[api]>=0.4.11"
+pip install vllm "lm-eval[api]>=0.4.12"
 ```

 ## Quantization Process
@@ -42,18 +42,59 @@ vllm serve meta-llama/Llama-3.1-8B --quantization mxfp8

 For fine-grained control, use a `quantization_config` dictionary.

+### Schema
+
+```yaml
+quantization_config:
+  linear:
+    weight: <name>      # see QUANT_KEY_NAMES in vllm/config/quantization.py
+    activation: <name>
+  moe:
+    weight: <name>
+    activation: <name>
+  ignore: [<layer-name-or-regex>, ...]
+```
+
+`linear` and `moe` accept a full `{weight, activation}` dict, or a bare
+string. A string resolves first against the `--quantization` shorthands
+(taking the matching layer-kind slot), then against `QUANT_KEY_NAMES` as a
+weight name. Unset fields fall back to the `--quantization` shorthand's
+defaults, or for already-quantized checkpoints to whatever the checkpoint
+declares.
+
+The CLI accepts the same shape as JSON or as dotted keys:
+
+```bash
+vllm serve <model> --quantization-config '{"moe":{"activation":"mxfp8"}}'
+vllm serve <model> --quantization-config.moe.activation mxfp8
+```
+
+### Activation overrides on already-quantized checkpoints
+
+For checkpoint-quantized models, `quantization_config` lets you pick an
+activation format independently of the baked-in weights. The supported
+overrides are checkpoint-specific; today this is wired up for MXFP4 MoE
+checkpoints (gpt-oss) where you can opt into FP8 activations:
+
+```bash
+vllm serve openai/gpt-oss-20b --quantization-config.moe.activation mxfp8
+```
+
+Combine with `--moe-backend` to pin a specific kernel family.
+
 ### Separate Schemes for Dense and MoE Layers

-You can apply different quantization schemes to dense linear layers and MoE expert layers:
+You can apply different quantization schemes to dense linear layers and MoE expert layers via the `linear` and `moe` fields. Each accepts either a full spec dict, or a bare string naming an online shorthand (e.g. `"fp8_per_block"`) or weight format (e.g. `"fp8_per_block_static"`); fields not set fall back to the shorthand defaults.

 ```python
 from vllm import LLM

+# Linear: per-block FP8; MoE: per-tensor FP8 (inherited from the shorthand)
 llm = LLM(
    "ibm-granite/granite-3.0-1b-a400m-base",
    quantization="fp8_per_tensor",
    quantization_config={
-        "linear_scheme_override": "fp8_per_block",
+        "linear": "fp8_per_block",
    },
 )
 ```
@@ -63,11 +104,12 @@ Or,
 ```python
 from vllm import LLM

+# Linear: per-tensor FP8 (inherited); MoE: per-block FP8
 llm = LLM(
    "ibm-granite/granite-3.0-1b-a400m-base",
    quantization="fp8_per_tensor",
    quantization_config={
-        "moe_scheme_override": "fp8_per_block",
+        "moe": "fp8_per_block",
    },
 )
 ```
@@ -20,7 +20,7 @@ for more installation details.
 Additionally, install `vllm` and `lm-evaluation-harness` for evaluation:

 ```bash
-pip install vllm "lm-eval[api]>=0.4.11"
+pip install vllm "lm-eval[api]>=0.4.12"
 ```

 ## Quantization Process
@@ -15,6 +15,7 @@ vLLM supports a variety of methods of speculative decoding. Model-based methods
 - [Multi-Layer Perceptron](mlp.md)
 - [N-Gram](n_gram.md)
 - [Suffix Decoding](suffix.md)
+- [Custom Proposer Backend (Experimental)](#custom-proposer-backend-experimental)

 ## Method Selection at a Glance

@@ -30,11 +31,22 @@ depend on your model family, traffic pattern, hardware, and sampling settings.
 | MLP speculator | Medium to high gain | Medium gain | Good when compatible MLP speculators are available. |
 | N-gram | Low to medium gain | Medium gain | Lightweight and easy to enable. |
 | Suffix decoding | Low to medium gain | Medium gain | No extra draft model; dynamic speculation depth. |
+| Custom Proposer | Varies | Varies | Bring your own proposer class (experimental). |

 For reproducible measurements in your environment, use
 [`examples/features/speculative_decoding/spec_decode_offline.py`](../../../examples/features/speculative_decoding/spec_decode_offline.py)
 or the [benchmark CLI guide](../../benchmarking/cli.md).

+## Custom Proposer Backend (Experimental)
+
+You can plug in your own custom proposer class for speculative decoding by setting the method to `custom_class` and providing the full module path to your class.
+Your custom class must accept a `VllmConfig` upon instantiation and implement a `propose` method.
+
+**Example configuration:**
+
+- `speculative_config.method = "custom_class"`
+- `speculative_config.model = "your_module.YourCustomProposerClass"`
+
 ## `--speculative-config` schema

 Use `--speculative-config` to pass speculative decoding settings as a JSON
@@ -11,24 +11,50 @@

 ### Text-only Language Models

-| Model                                | Architecture                             | Supported |
+| Model | Architecture | Supported |
 | ------------------------------------ | ---------------------------------------- | --------- |
-| meta-llama/Llama-3.1-8B-Instruct     | LlamaForCausalLM                         | ✅        |
-| meta-llama/Llama-3.2-3B-Instruct     | LlamaForCausalLM                         | ✅        |
-| ibm-granite/granite-3.2-2b-instruct  | GraniteForCausalLM                       | ✅        |
-| Qwen/Qwen3-1.7B                      | Qwen3ForCausalLM                         | ✅        |
-| Qwen/Qwen3-4B                        | Qwen3ForCausalLM                         | ✅        |
-| Qwen/Qwen3-8B                        | Qwen3ForCausalLM                         | ✅        |
-| zai-org/glm-4-9b-hf                  | GLMForCausalLM                           | ✅        |
-| google/gemma-7b                      | GemmaForCausalLM                         | ✅        |
+| unsloth/gpt-oss-20b | GptOssForCausalLM | ✅ |
+| meta-llama/Llama-3.1-8B-Instruct | LlamaForCausalLM | ✅ |
+| meta-llama/Llama-3.2-1B | LlamaForCausalLM | ✅ |
+| meta-llama/Llama-3.2-3B-Instruct | LlamaForCausalLM | ✅ |
+| meta-llama/Llama-3.3-70B-Instruct | LlamaForCausalLM | ✅ |
+| RedHatAI/Meta-Llama-3.1-8B-quantized.w8a8 | LlamaForCausalLM | ✅ |
+| RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w8a8 | LlamaForCausalLM | ✅ |
+| RedHatAI/Llama-3.2-1B-Instruct-quantized.w8a8 | LlamaForCausalLM | ✅ |
+| RedHatAI/Llama-3.2-3B-Instruct-quantized.w8a8 | LlamaForCausalLM | ✅ |
+| RedHatAI/DeepSeek-R1-Distill-Llama-70B-quantized.w8a8 | LlamaForCausalLM | ✅ |
+| hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4 | LlamaForCausalLM | ✅ |
+| AMead10/Llama-3.2-1B-Instruct-AWQ | LlamaForCausalLM | ✅ |
+| AMead10/Llama-3.2-3B-Instruct-AWQ | LlamaForCausalLM | ✅ |
+| TheBloke/TinyLlama-1.1B-Chat-v1.0-AWQ | LlamaForCausalLM | ✅ |
+| TheBloke/TinyLlama-1.1B-Chat-v1.0-GPTQ | LlamaForCausalLM | ✅ |
+| ibm-granite/granite-3.2-2b-instruct | GraniteForCausalLM | ✅ |
+| Qwen/Qwen3-1.7B | Qwen3ForCausalLM | ✅ |
+| Qwen/Qwen3-4B | Qwen3ForCausalLM | ✅ |
+| Qwen/Qwen3-8B | Qwen3ForCausalLM | ✅ |
+| Qwen/Qwen3-14B | Qwen3ForCausalLM | ✅ |
+| Qwen/Qwen3-14B-AWQ | Qwen3ForCausalLM | ✅ |
+| Qwen/Qwen3-30B-A3B | Qwen3MoeForCausalLM | ✅ |
+| Qwen/QwQ-32B-AWQ | Qwen2ForCausalLM | ✅ |
+| Qwen/Qwen1.5-0.5B-Chat-GPTQ-Int4 | Qwen2ForCausalLM | ✅ |
+| RedHatAI/QwQ-32B-quantized.w8a8 | Qwen2ForCausalLM | ✅ |
+| zai-org/glm-4-9b-hf | GLMForCausalLM | ✅ |
+| google/gemma-7b | GemmaForCausalLM | ✅ |
+| microsoft/Phi-4-reasoning | Phi3ForCausalLM | ✅ |
+| TheBloke/Mistral-7B-Instruct-v0.2-AWQ | MistralForCausalLM | ✅ |

 ### Multimodal Language Models

-| Model                                | Architecture                             | Supported |
+| Model | Architecture | Supported |
 | ------------------------------------ | ---------------------------------------- | --------- |
-| Qwen/Qwen2.5-VL-7B-Instruct          | Qwen2VLForConditionalGeneration          | ✅        |
-| openai/whisper-large-v3              | WhisperForConditionalGeneration          | ✅        |
+| meta-llama/Llama-4-Scout-17B-16E-Instruct | Llama4ForConditionalGeneration | ✅ |
+| google/gemma-3-4b-it | Gemma3ForConditionalGeneration | ✅ |
+| google/gemma-3-12b-it | Gemma3ForConditionalGeneration | ✅ |
+| google/gemma-4-E4B-it | Gemma4ForConditionalGeneration | ✅ |
+| google/gemma-4-E2B-it | Gemma4ForConditionalGeneration | ✅ |
+| google/gemma-4-26B-A4B-it | Gemma4ForConditionalGeneration | ✅ |
+| microsoft/Phi-4-multimodal-instruct | Phi4MMForCausalLM | ✅ |
+| Qwen/Qwen2.5-VL-7B-Instruct | Qwen2VLForConditionalGeneration | ✅ |
+| openai/whisper-large-v3 | WhisperForConditionalGeneration | ✅ |

-✅ Runs and optimized.  
-🟨 Runs and correct but not optimized to green yet.  
-❌ Does not pass accuracy test or does not run.  
+✅ Runs and optimized.
@@ -131,24 +131,24 @@ enabling the corresponding APIs.

 ### `LLM.classify`

-The [classify][vllm.LLM.classify] method outputs a probability vector for each prompt.
+The [classify][vllm.entrypoints.pooling.offline.PoolingOfflineMixin.classify] method outputs a probability vector for each prompt.
 It is primarily designed for [classification models](classify.md).
 For more information about `LLM.embed`, see [this page](classify.md#offline-inference).

 ### `LLM.embed`

-The [embed][vllm.LLM.embed] method outputs an embedding vector for each prompt.
+The [embed][vllm.entrypoints.pooling.offline.PoolingOfflineMixin.embed] method outputs an embedding vector for each prompt.
 It is primarily designed for [embedding models](embed.md).
 For more information about `LLM.embed`, see [this page](embed.md#offline-inference).

 ### `LLM.score`

-The [score][vllm.LLM.score] method outputs similarity scores between sentence pairs.
+The [score][vllm.entrypoints.pooling.offline.PoolingOfflineMixin.score] method outputs similarity scores between sentence pairs.
 It is primarily designed for [score models](scoring.md).

 ### `LLM.encode`

-The [encode][vllm.LLM.encode] method is available to all pooling models in vLLM.
+The [encode][vllm.entrypoints.pooling.offline.PoolingOfflineMixin.encode] method is available to all pooling models in vLLM.

 Please use one of the more specific methods or set the task directly when using `LLM.encode`, refer to the [table above](#offline-apis-corresponding-to-pooling-usages).

@@ -77,7 +77,7 @@ The following [pooling parameters][vllm.PoolingParams] are supported.

 ### `LLM.classify`

-The [classify][vllm.LLM.classify] method outputs a probability vector for each prompt.
+The [classify][vllm.entrypoints.pooling.offline.PoolingOfflineMixin.classify] method outputs a probability vector for each prompt.

 ```python
 from vllm import LLM
@@ -93,7 +93,7 @@ A code example can be found here: [examples/basic/offline_inference/classify.py]

 ### `LLM.encode`

-The [encode][vllm.LLM.encode] method is available to all pooling models in vLLM.
+The [encode][vllm.entrypoints.pooling.offline.PoolingOfflineMixin.encode] method is available to all pooling models in vLLM.

 Set `pooling_task="classify"` when using `LLM.encode` for classification Models:

@@ -120,7 +120,7 @@ The following [pooling parameters][vllm.PoolingParams] are supported.

 ### `LLM.embed`

-The [embed][vllm.LLM.embed] method outputs an embedding vector for each prompt.
+The [embed][vllm.entrypoints.pooling.offline.PoolingOfflineMixin.embed] method outputs an embedding vector for each prompt.

 ```python
 from vllm import LLM
@@ -136,7 +136,7 @@ A code example can be found here: [examples/basic/offline_inference/embed.py](..

 ### `LLM.encode`

-The [encode][vllm.LLM.encode] method is available to all pooling models in vLLM.
+The [encode][vllm.entrypoints.pooling.offline.PoolingOfflineMixin.encode] method is available to all pooling models in vLLM.

 Set `pooling_task="embed"` when using `LLM.encode` for embedding Models:

@@ -152,7 +152,7 @@ print(f"Data: {data!r}")

 ### `LLM.score`

-The [score][vllm.LLM.score] method outputs similarity scores between sentence pairs.
+The [score][vllm.entrypoints.pooling.offline.PoolingOfflineMixin.score] method outputs similarity scores between sentence pairs.

 All models that support embedding task also support using the score API to compute similarity scores by calculating the cosine similarity of two input prompt's embeddings.

@@ -87,7 +87,7 @@ The following [pooling parameters][vllm.PoolingParams] are supported.

 ### `LLM.encode`

-The [encode][vllm.LLM.encode] method is available to all pooling models in vLLM.
+The [encode][vllm.entrypoints.pooling.offline.PoolingOfflineMixin.encode] method is available to all pooling models in vLLM.

 - Reward Models

@@ -130,7 +130,7 @@ The following [pooling parameters][vllm.PoolingParams] are only supported by cro

 ### `LLM.score`

-The [score][vllm.LLM.score] method outputs similarity scores between sentence pairs.
+The [score][vllm.entrypoints.pooling.offline.PoolingOfflineMixin.score] method outputs similarity scores between sentence pairs.

 ```python
 from vllm import LLM
@@ -87,7 +87,7 @@ The following [pooling parameters][vllm.PoolingParams] are supported.

 ### `LLM.encode`

-The [encode][vllm.LLM.encode] method is available to all pooling models in vLLM.
+The [encode][vllm.entrypoints.pooling.offline.PoolingOfflineMixin.encode] method is available to all pooling models in vLLM.

 Set `pooling_task="token_classify"` when using `LLM.encode` for token classification Models:

@@ -94,7 +94,7 @@ The following [pooling parameters][vllm.PoolingParams] are supported.

 ### `LLM.encode`

-The [encode][vllm.LLM.encode] method is available to all pooling models in vLLM.
+The [encode][vllm.entrypoints.pooling.offline.PoolingOfflineMixin.encode] method is available to all pooling models in vLLM.

 Set `pooling_task="token_embed"` when using `LLM.encode` for token embedding Models:

@@ -110,7 +110,7 @@ print(f"Data: {data!r}")

 ### `LLM.score`

-The [score][vllm.LLM.score] method outputs similarity scores between sentence pairs.
+The [score][vllm.entrypoints.pooling.offline.PoolingOfflineMixin.score] method outputs similarity scores between sentence pairs.

 All models that support token embedding task also support using the score API to compute similarity scores by calculating the late interaction of two input prompts.

@@ -575,6 +575,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
 | `IsaacForConditionalGeneration` | Isaac | T + I<sup>+</sup> | `PerceptronAI/Isaac-0.1` | ✅︎ | ✅︎ |
 | `InternS1ForConditionalGeneration` | Intern-S1 | T + I<sup>E+</sup> + V<sup>E+</sup> | `internlm/Intern-S1`, `internlm/Intern-S1-mini`, etc. | ✅︎ | ✅︎ |
 | `InternS1ProForConditionalGeneration` | Intern-S1-Pro | T + I<sup>E+</sup> + V<sup>E+</sup> | `internlm/Intern-S1-Pro`, etc. | ✅︎ | ✅︎ |
+| `InternS2PreviewForConditionalGeneration` | Intern-S2-Preview | T + I<sup>E+</sup> + V<sup>E+</sup> | `internlm/Intern-S2-Preview`, etc. | ✅︎ | ✅︎ |
 | `InternVLChatModel` | InternVL 3.5, InternVL 3.0, InternVideo 2.5, InternVL 2.5, Mono-InternVL, InternVL 2.0 | T + I<sup>E+</sup> + (V<sup>E+</sup>) | `OpenGVLab/InternVL3_5-14B`, `OpenGVLab/InternVL3-9B`, `OpenGVLab/InternVideo2_5_Chat_8B`, `OpenGVLab/InternVL2_5-4B`, `OpenGVLab/Mono-InternVL-2B`, `OpenGVLab/InternVL2-4B`, etc. | ✅︎ | ✅︎ |
 | `InternVLForConditionalGeneration` | InternVL 3.0 (HF format) | T + I<sup>E+</sup> + V<sup>E+</sup> | `OpenGVLab/InternVL3-1B-hf`, etc. | ✅︎ | ✅︎ |
 | `KananaVForConditionalGeneration` | Kanana-V | T + I<sup>+</sup> | `kakaocorp/kanana-1.5-v-3b-instruct`, etc. | | ✅︎ |
@@ -1,289 +0,0 @@
-# Routed Experts Replay
-
-## Overview
-
-Routed experts replay captures which MoE (Mixture of Experts) experts process each token during inference and returns this information alongside the generated text. This is essential for **reinforcement learning (RL) training pipelines** (such as GRPO and RLHF) where the training step needs to reconstruct expert routing decisions from the inference pass.
-
-When enabled, each API response includes:
-
- **`prompt_routed_experts`**: A `[prompt_len, num_moe_layers, top_k]` array of expert IDs for the prompt tokens (at the response level, shared across completions).
- **`routed_experts`**: A `[gen_len, num_moe_layers, top_k]` array of expert IDs for the generated tokens (per completion).
-
-For example, a model with 40 MoE layers and top-22 routing that processes a 100-token prompt and generates 50 tokens would return:
-
- `prompt_routed_experts`: shape `[100, 40, 22]`
- `routed_experts`: shape `[50, 40, 22]`
-
-Each value is an int16 expert ID in the range `[0, num_experts)`.
-
-## Quickstart
-
-### OpenAI API Server
-
-```bash
-vllm serve <MODEL> \
-    --enable-return-routed-experts \
-    --tensor-parallel-size 4 \
-    --enable-expert-parallel
-```
-
-Then query the `/v1/completions` endpoint as usual. The response includes routing data:
-
-```python
-import requests
-
-resp = requests.post("http://localhost:8000/v1/completions", json={
-    "model": "<MODEL>",
-    "prompt": "Explain quantum computing.",
-    "max_tokens": 64,
-    "temperature": 0.0,
-}).json()
-
-# Generation routing (per completion choice)
-gen_routing = resp["choices"][0]["routed_experts"]  # [gen_len, layers, top_k]
-
-# Prompt routing (shared across all choices)
-prompt_routing = resp["prompt_routed_experts"]      # [prompt_len, layers, top_k]
-
-print(f"Prompt routing shape: [{len(prompt_routing)}, "
-      f"{len(prompt_routing[0])}, {len(prompt_routing[0][0])}]")
-print(f"Gen routing shape: [{len(gen_routing)}, "
-      f"{len(gen_routing[0])}, {len(gen_routing[0][0])}]")
-```
-
-### Python SDK (Offline Inference)
-
-```python
-from vllm import LLM, SamplingParams
-
-llm = LLM(
-    model="<MODEL>",
-    enable_return_routed_experts=True,
-    tensor_parallel_size=4,
-    enable_expert_parallel=True,
-)
-
-outputs = llm.generate(
-    ["Explain quantum computing."],
-    SamplingParams(temperature=0, max_tokens=64),
-)
-
-result = outputs[0]
-
-# Prompt routing: numpy array, shape [prompt_len, num_moe_layers, top_k]
-prompt_routing = result.prompt_routed_experts
-print(f"Prompt routing: {prompt_routing.shape}, dtype={prompt_routing.dtype}")
-
-# Generation routing: numpy array, shape [gen_len, num_moe_layers, top_k]
-gen_routing = result.outputs[0].routed_experts
-print(f"Gen routing: {gen_routing.shape}, dtype={gen_routing.dtype}")
-```
-
-## Output Format
-
-### `CompletionOutput.routed_experts`
-
- **Type**: `numpy.ndarray` (Python SDK) or `list[list[list[int]]]` (JSON API)
- **Shape**: `[gen_len, num_moe_layers, top_k]`
- **Dtype**: `int16`
- **Content**: Expert IDs for **generated tokens only**. `gen_len` matches the number of generated tokens (i.e., `usage.completion_tokens` or fewer).
-
-### `RequestOutput.prompt_routed_experts`
-
- **Type**: `numpy.ndarray` (Python SDK) or `list[list[list[int]]]` (JSON API)
- **Shape**: `[prompt_len, num_moe_layers, top_k]`
- **Dtype**: `int16`
- **Content**: Expert IDs for **prompt tokens only**. `prompt_len` matches `usage.prompt_tokens`. This field lives on the request-level response (not per-choice), because prompt routing is shared across all completions when `n > 1`.
-
-### Why Separate Prompt and Generation Routing?
-
-When a request has multiple completions (`n > 1`), each completion shares the same prompt but produces different generated text. Storing prompt routing once on the `RequestOutput` (rather than duplicating it on every `CompletionOutput`) avoids redundant data. For RL training, the consumer typically needs:
-
-1. The prompt routing (once) to reconstruct the forward pass for the shared prefix.
-2. The per-completion generation routing to reconstruct each completion's forward pass.
-
-## Architecture
-
-### Data Flow
-
-```text
-Forward Pass                    Async D2H Pipeline              Output
-─────────────                   ──────────────────              ──────
-FusedMoE layer                  After forward pass:             On request finish:
-writes topk_ids    ──────►     D2H copy to pinned     ──────►  Extract from host cache
-to device buffer               staging buffer                   Split at prompt_len
-(L, N, K) int16                 (via CUDA stream)                Trim gen to output len
-                                Scatter to per-request           Serialize to API response
-                                host cache (numpy)
-```
-
-### Device Cache
-
-A pre-allocated GPU buffer with layout `(L, N, K)` where:
-
- `L` = number of MoE layers
- `N` = `max_num_batched_tokens`
- `K` = `num_experts_per_tok` (top-k)
-
-The `(L, N, K)` layout ensures that `buffer[layer_id]` gives a contiguous `(N, K)` view per layer. Each `FusedMoE` layer gets a persistent reference to its slice via `module._routing_replay_out = buffer[layer_id]`.
-
-**Dtype**: `int16` — sufficient for expert IDs (max ~512 experts in practice) and half the memory of `int32`.
-
-### Host Cache
-
-Per-request numpy arrays for accumulating routing data across decode steps. Each request gets a lazily allocated `(seq_len, L, K)` int16 buffer that grows as the sequence lengthens. Buffers are freed when a request completes.
-
-### Async D2H Pipeline
-
-After each forward pass, the model runner issues a non-blocking device-to-host copy on a dedicated CUDA stream:
-
-1. **Copy**: `pinned_staging[:, :total_tokens, :].copy_(device_buffer[:, :total_tokens, :])` on a separate stream, recorded with a CUDA event.
-2. **Scatter** (deferred to next step): On the *next* forward pass, synchronize the event (effectively free — an entire forward pass has elapsed) and scatter the staging data into per-request host cache buffers using the token positions.
-
-This design ensures the D2H copy overlaps with the next forward pass, minimizing GPU stall time.
-
-### CUDA Graph Compatibility
-
-CUDA graph compatibility requires two mechanisms:
-
-1. **Persistent tensor attribute**: Each `FusedMoE` layer stores a reference to its buffer slice as `module._routing_replay_out`. Because `torch.compile` captures module attributes by reference, graph replay always writes to the live buffer — not a stale snapshot.
-
-2. **Static marking**: Both the full `(L, N, K)` buffer and each per-layer `(N, K)` view are marked with `cudagraph_mark_tensor_static()`. This prevents CUDA graphs from snapshot/restore behavior that would zero the buffer on replay.
-
-### Multi-Node Support
-
-On multi-node tensor-parallel setups, all TP ranks allocate a device buffer (required for symmetric CUDA graph structure), but only TP rank 0 runs the D2H pipeline and host cache. Routing data flows from the model runner through `ModelRunnerOutput` via Ray DAG to the scheduler — no shared memory or file locks needed.
-
-### Routing Capture Path
-
-For the **non-monolithic (Triton) kernel path** (e.g., BF16 MoE), routing is captured after `select_experts()` in the MoE runner:
-
-```python
-routing_replay_out = getattr(layer, "_routing_replay_out", None)
-topk_weights, topk_ids = self.router.select_experts(...)
-
-if routing_replay_out is not None:
-    routing_replay_out[:topk_ids.shape[0]].copy_(topk_ids.to(torch.int16))
-```
-
-For the **monolithic kernel path** (e.g., FP8/MXFP8 via FlashInfer), `routing_replay_out` is threaded through the `apply_monolithic()` call chain and FlashInfer writes expert IDs directly during routing inside the fused kernel.
-
-### MTP (Multi-Token Prediction) Handling
-
-With MTP speculative decoding, the model captures routing for all tokens including speculative ones that may later be rejected. When a request finishes, the generation routing is trimmed to match the actual number of accepted output tokens:
-
-```python
-num_gen = self.detokenizer.num_output_tokens()
-if gen_routed_experts.shape[0] > num_gen and num_gen > 0:
-    gen_routed_experts = gen_routed_experts[:num_gen]
-```
-
-This ensures the routing array length always matches the token IDs in the response.
-
-## Design Decisions
-
-### Why Replace SharedMemory with Device Cache?
-
-The previous implementation used `multiprocessing.SharedMemory` with `fcntl` file locking to transfer routing data from GPU workers to the scheduler. This approach had fundamental problems:
-
- **Multi-node**: `SharedMemory` is node-local. On multi-node TP setups (required for 400B+ parameter models), the scheduler on node 0 cannot read shared memory from workers on other nodes.
- **Performance**: Synchronous `.cpu().numpy()` D2H transfers block the GPU. File-based locking adds further overhead.
- **CUDA graphs**: The callback-based capture mechanism bakes tensor references at trace time, causing stale data on graph replay.
-
-The device cache approach solves all three: data flows through Ray DAG (works multi-node), D2H is async (non-blocking), and persistent tensor attributes work with CUDA graphs.
-
-### Why `(L, N, K)` Layout Instead of `(N, L, K)`?
-
-FlashInfer's `routing_replay_out` parameter expects a contiguous `(N, K)` tensor per layer. With `(L, N, K)` layout, `buffer[layer_id]` gives a contiguous `(N, K)` view with zero-copy slicing. The previous `(N, L, K)` layout would require non-contiguous indexing or an explicit copy.
-
-### Why int16 Instead of int32?
-
-Expert IDs are small integers (typically 0-255 for models with up to 256 experts). `int16` supports up to 32,767 experts — far more than any current model — while halving GPU memory usage and D2H bandwidth compared to `int32`.
-
-### Why Split Prompt and Generation Routing?
-
-RL training pipelines process prompt and generation routing separately:
-
- Prompt routing reconstructs the shared forward pass for the input.
- Generation routing reconstructs each sampled trajectory.
-
-With `n > 1` completions, all completions share the same prompt routing. Duplicating it per completion would waste memory proportional to `n * prompt_len * L * K`. Instead, `prompt_routed_experts` is stored once on `RequestOutput` and shared.
-
-### Why Async D2H Instead of Synchronous Copy?
-
-A synchronous `.cpu()` call forces the GPU to drain its command queue before the copy can begin, stalling the pipeline. The async approach:
-
-1. Issues the copy on a separate CUDA stream (non-blocking to the main compute stream).
-2. Defers the host-side scatter to the *next* step, by which time the copy has finished.
-
-This means the D2H transfer overlaps entirely with the next forward pass, adding near-zero latency to the critical path.
-
-### Why All TP Ranks Get a Device Buffer?
-
-CUDA graph capture records the exact sequence of kernel calls and their arguments. If only rank 0 had a device buffer, the `FusedMoE` layer would take a different code path on rank 0 vs. other ranks (one writes to a buffer, others don't). This asymmetry causes different CUDA graph structures across ranks, which can lead to NCCL deadlocks during collective operations inside the graph. Giving all ranks a real buffer ensures symmetric graph structure. Only rank 0 does the D2H copy and host cache management.
-
-## Performance
-
-Routing replay adds a small overhead from the device buffer writes and async D2H copies. On tested configurations:
-
- **Throughput overhead** (random data, ISL=1024, OSL=1024): **~2%**
- **Memory overhead** (int16 buffer, 40 layers, 8192 tokens, top-22): **~14 MB per GPU**
- **Accuracy impact** (GSM8K): **Zero** (pass@1 identical with and without routing replay)
-
-The overhead is dominated by the per-layer `.copy_()` during the forward pass. The async D2H pipeline runs entirely in the background.
-
-## Supported Configurations
-
-| Configuration                            | Supported                                                 |
-|------------------------------------------|-----------------------------------------------------------|
-| BF16 Triton MoE (non-monolithic)         | Yes                                                       |
-| FP8/MXFP8 FlashInfer MoE (monolithic)    | Yes (requires FlashInfer with `routing_replay_out`)       |
-| CUDA graphs                              | Yes                                                       |
-| Multi-node tensor parallelism            | Yes                                                       |
-| Data parallelism (DP)                    | Yes                                                       |
-| Expert parallelism (EP)                  | Yes                                                       |
-| Prefix caching                           | Yes (cached positions marked with `-1` sentinel)          |
-| MTP speculative decoding                 | Yes (gen routing trimmed to accepted tokens)              |
-| `n > 1` (multiple completions)           | Yes (prompt routing shared, gen routing per-completion)   |
-
-## Limitations
-
- **Streaming**: Routing data is only available when the request finishes (not streamed incrementally).
- **V1 engine only**: Routing replay is implemented for the vLLM V1 engine.
- **Preempted requests**: When a request is preempted by the scheduler (and later resumed via re-prefill), any routing already accumulated in the worker's host cache for that request is dropped without being emitted. The consumer sees `routed_experts=None` for the resumed request with no other signal. Partial-rollout and async-RL pipelines that rely on routing for preempted requests should either disable preemption (`--no-enable-chunked-prefill` / sufficient KV headroom) or reconstruct routing on the resumed prefill.
- **Async scheduling**: Not supported; rejected at config time. The worker-side stop predicate reads `req_state.output_token_ids[-1]`, which under async scheduling is the placeholder `-1` until `AsyncGPUModelRunnerOutput` resolves the real sampled token, so EOS / stop-token finishes would silently drop routing. Use sync scheduling (the default when `--enable-return-routed-experts` is set, or set explicitly with the appropriate scheduler config).
- **Sequence parallelism / naive DP MoE dispatch**: Not supported on the FusedMoE layer; rejected at bind time. SP shards `topk_ids` along dim 0 across the TP group so each rank only captures `1/sp_size` of the rows; naive DP dispatch all-gathers tokens across DP ranks before routing, so `topk_ids.shape[0]` exceeds the per-rank buffer size. Both raise `NotImplementedError` from `bind_routing_capture_to_model`.
- **Pipeline / prefill-context / decode-context parallelism**: Not yet validated; rejected at config time.
-
-## CLI Reference
-
-| Flag                               | Description                                                            |
-|------------------------------------|------------------------------------------------------------------------|
-| `--enable-return-routed-experts`   | Enable routing replay capture and return expert IDs in API responses.  |
-
-## API Reference
-
-### Completions (`/v1/completions`)
-
-**Response-level field:**
-
-| Field                     | Type                                | Description                                                                 |
-|---------------------------|-------------------------------------|-----------------------------------------------------------------------------|
-| `prompt_routed_experts`   | `list[list[list[int]]]` or `null`   | Expert IDs for prompt tokens. Shape: `[prompt_len, num_moe_layers, top_k]`. |
-
-**Choice-level field:**
-
-| Field              | Type                                | Description                                                                   |
-|--------------------|-------------------------------------|-------------------------------------------------------------------------------|
-| `routed_experts`   | `list[list[list[int]]]` or `null`   | Expert IDs for generated tokens. Shape: `[gen_len, num_moe_layers, top_k]`.   |
-
-### Chat Completions (`/v1/chat/completions`)
-
-Same fields as above on `ChatCompletionResponse` and `ChatCompletionResponseChoice`.
-
-### Python SDK
-
-| Object               | Field                     | Type                     | Description                 |
-|----------------------|---------------------------|--------------------------|-----------------------------|
-| `RequestOutput`      | `prompt_routed_experts`   | `np.ndarray` or `None`   | `[prompt_len, L, K]` i16    |
-| `CompletionOutput`   | `routed_experts`          | `np.ndarray` or `None`   | `[gen_len, L, K]` int16     |
@@ -2179,6 +2179,92 @@ def run_qwen3_vl_moe(questions: list[str], modality: str) -> ModelRequestData:
    )


+# Qwen3.5-Dense
+def run_qwen3_5(questions: list[str], modality: str) -> ModelRequestData:
+    model_name = "Qwen/Qwen3.5-4B"
+
+    mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
+    engine_args = EngineArgs(
+        model=model_name,
+        max_model_len=4096,
+        max_num_seqs=5,
+        mm_processor_kwargs={
+            "min_pixels": 28 * 28,
+            "max_pixels": 1280 * 28 * 28,
+            "fps": 1,
+        },
+        limit_mm_per_prompt=mm_limit,
+    )
+
+    image_placeholder = "<|vision_start|><|image_pad|><|vision_end|>"
+    video_placeholder = "<|vision_start|><|video_pad|><|vision_end|>"
+
+    if modality == "image":
+        placeholder = image_placeholder
+    elif modality == "video":
+        placeholder = video_placeholder
+    elif modality == "image+video":
+        placeholder = image_placeholder + video_placeholder
+
+    prompts = [
+        (
+            "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
+            f"<|im_start|>user\n{placeholder}"
+            f"{question}<|im_end|>\n"
+            "<|im_start|>assistant\n"
+        )
+        for question in questions
+    ]
+
+    return ModelRequestData(
+        engine_args=engine_args,
+        prompts=prompts,
+    )
+
+
+# Qwen3.5-MoE
+def run_qwen3_5_moe(questions: list[str], modality: str) -> ModelRequestData:
+    model_name = "Qwen/Qwen3.5-35B-A3B"
+
+    mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
+    engine_args = EngineArgs(
+        model=model_name,
+        max_model_len=4096,
+        max_num_seqs=5,
+        mm_processor_kwargs={
+            "min_pixels": 28 * 28,
+            "max_pixels": 1280 * 28 * 28,
+            "fps": 1,
+        },
+        limit_mm_per_prompt=mm_limit,
+    )
+
+    image_placeholder = "<|vision_start|><|image_pad|><|vision_end|>"
+    video_placeholder = "<|vision_start|><|video_pad|><|vision_end|>"
+
+    if modality == "image":
+        placeholder = image_placeholder
+    elif modality == "video":
+        placeholder = video_placeholder
+    elif modality == "image+video":
+        placeholder = image_placeholder + video_placeholder
+
+    prompts = [
+        (
+            "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
+            f"<|im_start|>user\n{placeholder}"
+            f"{question}<|im_end|>\n"
+            "<|im_start|>assistant\n"
+        )
+        for question in questions
+    ]
+
+    return ModelRequestData(
+        engine_args=engine_args,
+        prompts=prompts,
+    )
+
+
 # R-4B
 def run_r_vl(questions: list[str], modality: str) -> ModelRequestData:
    assert modality == "image"
@@ -2442,6 +2528,8 @@ model_example_map = {
    "qwen2_5_omni": run_qwen2_5_omni,
    "qwen3_vl": run_qwen3_vl,
    "qwen3_vl_moe": run_qwen3_vl_moe,
+    "qwen3_5": run_qwen3_5,
+    "qwen3_5_moe": run_qwen3_5_moe,
    "rvl": run_r_vl,
    "skywork_chat": run_skyworkr1v,
    "smolvlm": run_smolvlm,
@@ -2460,13 +2548,18 @@ MODELS_NEED_VIDEO_METADATA = [
    "molmo2",
    "qwen3_vl",
    "qwen3_vl_moe",
+    "qwen3_5",
+    "qwen3_5_moe",
 ]


 MODELS_SUPPORT_VIT_CUDA_GRAPH = [
+    "qwen2_5_vl",
    "qwen3_vl",
    "qwen3_vl_moe",
-    "qwen2_5_vl",
+    "qwen2_vl",
+    "qwen3_5",
+    "qwen3_5_moe",
 ]


@@ -1,7 +1,7 @@
 <|im_start|>system
 Judge whether the Document meets the requirements based on the Query and the Instruct provided. Note that the answer can only be "yes" or "no".<|im_end|>
 <|im_start|>user
-<Instruct>: {{ messages | selectattr("role", "eq", "system") | map(attribute="content") | first | default("Given a web search query, retrieve relevant passages that answer the query") }}
+<Instruct>: {{ instruction | default(instruct | default(messages | selectattr("role", "eq", "system") | map(attribute="content") | first | default("Given a web search query, retrieve relevant passages that answer the query", true), true), true) }}
 <Query>: {{ messages | selectattr("role", "eq", "query") | map(attribute="content") | first }}
 <Document>: {{ messages | selectattr("role", "eq", "document") | map(attribute="content") | first }}<|im_end|>
 <|im_start|>assistant
@@ -1,13 +1,7 @@
 <|im_start|>system
 Judge whether the Document meets the requirements based on the Query and the Instruct provided. Note that the answer can only be "yes" or "no".<|im_end|>
 <|im_start|>user
-<Instruct>: {{
-    messages
-    | selectattr("role", "eq", "system")
-    | map(attribute="content")
-    | first
-    | default("Given a search query, retrieve relevant candidates that answer the query.")
-}}<Query>:{{
+<Instruct>: {{ instruction | default(instruct | default(messages | selectattr("role", "eq", "system") | map(attribute="content") | first | default("Given a search query, retrieve relevant candidates that answer the query.", true), true), true) }}<Query>:{{
    messages
    | selectattr("role", "eq", "query")
    | map(attribute="content")
@@ -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 == "ppc64le"
+llguidance >= 1.7.0, < 1.8.0; platform_machine == "x86_64" or platform_machine == "arm64" or platform_machine == "aarch64" or platform_machine == "ppc64le"
 outlines_core == 0.2.14
 # required for outlines backend disk cache
 diskcache == 5.6.3
@@ -21,5 +21,8 @@ nvidia-cudnn-frontend>=1.13.0,<1.19.0
 fastsafetensors >= 0.2.2

 # QuACK and Cutlass DSL for FA4 (cute-DSL implementation)
-nvidia-cutlass-dsl>=4.4.2
+nvidia-cutlass-dsl[cu13]==4.5.0
 quack-kernels>=0.3.3
+
+# Tokenspeed_MLA for faster mla with spec decode
+tokenspeed-mla==0.1.2
@@ -1,5 +1,3 @@
 lmcache >= 0.3.9
-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
+nixl >= 1.1.0 # Required for disaggregated prefill
 mooncake-transfer-engine >= 0.3.8
@@ -21,6 +21,3 @@ timm>=1.0.17
 # amd-quark: required for Quark quantization on ROCm 
 # To be consistent with test_quark.py
 amd-quark>=0.8.99
-# tilelang has to be installed for mhc module to be
-# imported correctly.
-tilelang==0.1.9
--- a/Show More
+++ b/Show More