Add tasks/version/logging_utils to models_basic GPU deps

Following the deeper audit pattern from #42219: - `vllm/tasks.py`: imported by `vllm/model_executor/models/registry.py` and other model-side code; a change to `POOLING_TASKS` / `SupportedTask` would otherwise silently skip these jobs. - `vllm/version.py`: imported by `vllm/v1/engine/core.py` and `vllm/engine/arg_utils.py`. Release-tooling bumps would otherwise skip the model init jobs. - `vllm/logging_utils/`: imported by `vllm/model_executor/models/registry.py` (`logtime`) and `vllm/v1/engine/core.py`. Applied to GPU jobs (Initialization, Other). The CPU job (test_utils, test_vision) does not transitively reach these modules — its tests only import `vllm.distributed`, `vllm.model_executor.models.{utils,vision}`, `vllm.platforms`, `vllm.utils`, none of which import the new edge modules — so the CPU job is unchanged. Signed-off-by: Kevin H. Luu <kevin@anyscale.com> Signed-off-by: khluu <khluu000@gmail.com>
Add sequence.py and config/ to models_basic deps
2026-05-10 03:48:11 -07:00 · 2026-05-10 03:35:03 -07:00 · 2026-05-10 03:27:42 -07:00
802 changed files with 12836 additions and 44453 deletions
@@ -8,7 +8,6 @@ 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"
@@ -12,7 +12,7 @@ steps:
  - vllm/_custom_ops.py
  - tests/kernels/attention/test_cpu_attn.py
  - tests/kernels/moe/test_cpu_fused_moe.py
-  - tests/kernels/moe/test_cpu_quant_fused_moe.py
+  - tests/kernels/moe/test_cpu_fp8_fused_moe.py
  - tests/kernels/test_onednn.py
  - tests/kernels/test_awq_int4_to_int8.py
  - tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
@@ -21,7 +21,7 @@ steps:
      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
      pytest -x -v -s tests/kernels/attention/test_cpu_attn.py
      pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py
-      pytest -x -v -s tests/kernels/moe/test_cpu_quant_fused_moe.py
+      pytest -x -v -s tests/kernels/moe/test_cpu_fp8_fused_moe.py
      pytest -x -v -s tests/kernels/test_onednn.py
      pytest -x -v -s tests/kernels/test_awq_int4_to_int8.py
      pytest -x -v -s tests/kernels/quantization/test_cpu_fp8_scaled_mm.py"
@@ -61,7 +61,7 @@ steps:
  source_file_dependencies:
  - csrc/cpu/
  - vllm/model_executor/layers/quantization/cpu_wna16.py
-  - vllm/model_executor/layers/quantization/auto_gptq.py
+  - vllm/model_executor/layers/quantization/gptq_marlin.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
@@ -70,7 +70,7 @@ steps:
  - tests/quantization/test_cpu_wna16.py
  commands:
    - |
-      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
+      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 20m "
      pytest -x -v -s tests/quantization/test_compressed_tensors.py::test_compressed_tensors_w8a8_logprobs
      pytest -x -v -s tests/quantization/test_cpu_wna16.py"
      
@@ -8,3 +8,10 @@ steps:
    commands: 
    - bash .buildkite/scripts/hardware_ci/run-hpu-test.sh

+  - label: "Intel GPU Test"
+    depends_on: []
+    soft_fail: true
+    device: intel_gpu
+    no_plugin: true
+    commands: 
+    - bash .buildkite/scripts/hardware_ci/run-xpu-test.sh
@@ -49,7 +49,7 @@ steps:
      'cd tests &&
      export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
      pytest -v -s lora/test_fused_moe_lora_kernel.py && 
-      pytest -v -s lora/test_moe_lora_align_sum.py --deselect="tests/lora/test_moe_lora_align_sum.py::test_moe_lora_align_block_size_mixed_base_and_lora[1]"'
+      pytest -v -s lora/test_moe_lora_align_sum.py'

 - label: LoRA Punica Kernels
  timeout_in_minutes: 45
@@ -36,7 +36,6 @@ 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'
@@ -63,4 +62,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 --ignore=v1/kv_connector/unit/test_offloading_connector.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'
@@ -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.12"
+#   pip install "lm-eval[api]>=0.4.11"

 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.12"
+#   pip install "lm-eval[api]>=0.4.11"

 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.12"
+#   pip install "lm-eval[api]>=0.4.11"

 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.12"
+#   pip install "lm-eval[api]>=0.4.11"

 usage() {
    echo``
@@ -31,9 +31,30 @@
    }
  },
  "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",
@@ -41,10 +62,158 @@
        "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"
+        "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
      },
      "client_parameters": {
        "model": "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
@@ -56,7 +225,8 @@
    {
      "test_name": "serving_llama8B_int8_tp1_random_128_128",
      "server_parameters": {
-        "model": "RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w8a8"
+        "model": "RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w8a8",
+        "tensor_parallel_size": 1
      },
      "client_parameters": {
        "model": "RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w8a8",
@@ -66,12 +236,26 @@
      }
    },
    {
-      "test_name": "serving_llama1B_tp1_random_128_128",
+      "test_name": "serving_llama8B_int8_tp2_random_128_128",
      "server_parameters": {
-        "model": "meta-llama/Llama-3.2-1B"
+        "model": "RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w8a8",
+        "tensor_parallel_size": 2
      },
      "client_parameters": {
-        "model": "meta-llama/Llama-3.2-1B",
+        "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",
        "dataset_name": "random",
        "random-input-len": 128,
        "random-output-len": 128
@@ -80,7 +264,8 @@
    {
      "test_name": "serving_llama3B_tp1_random_128_128",
      "server_parameters": {
-        "model": "meta-llama/Llama-3.2-3B-Instruct"
+        "model": "meta-llama/Llama-3.2-3B-Instruct",
+        "tensor_parallel_size": 1
      },
      "client_parameters": {
        "model": "meta-llama/Llama-3.2-3B-Instruct",
@@ -89,22 +274,11 @@
        "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"
+        "model": "ibm-granite/granite-3.2-2b-instruct",
+        "tensor_parallel_size": 1
      },
      "client_parameters": {
        "model": "ibm-granite/granite-3.2-2b-instruct",
@@ -116,7 +290,8 @@
    {
      "test_name": "serving_qwen1.7B_tp1_random_128_128",
      "server_parameters": {
-        "model": "Qwen/Qwen3-1.7B"
+        "model": "Qwen/Qwen3-1.7B",
+        "tensor_parallel_size": 1
      },
      "client_parameters": {
        "model": "Qwen/Qwen3-1.7B",
@@ -128,7 +303,8 @@
    {
      "test_name": "serving_qwen4B_tp1_random_128_128",
      "server_parameters": {
-        "model": "Qwen/Qwen3-4B"
+        "model": "Qwen/Qwen3-4B",
+        "tensor_parallel_size": 1
      },
      "client_parameters": {
        "model": "Qwen/Qwen3-4B",
@@ -140,7 +316,8 @@
    {
      "test_name": "serving_qwen8B_tp1_random_128_128",
      "server_parameters": {
-        "model": "Qwen/Qwen3-8B"
+        "model": "Qwen/Qwen3-8B",
+        "tensor_parallel_size": 1
      },
      "client_parameters": {
        "model": "Qwen/Qwen3-8B",
@@ -149,34 +326,11 @@
        "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"
+        "model": "zai-org/glm-4-9b-hf",
+        "tensor_parallel_size": 1
      },
      "client_parameters": {
        "model": "zai-org/glm-4-9b-hf",
@@ -188,7 +342,8 @@
    {
      "test_name": "serving_gemma7B_tp1_random_128_128",
      "server_parameters": {
-        "model": "google/gemma-7b"
+        "model": "google/gemma-7b",
+        "tensor_parallel_size": 1
      },
      "client_parameters": {
        "model": "google/gemma-7b",
@@ -196,94 +351,6 @@
        "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
-      }
    }
  ]
 }
@@ -7,9 +7,6 @@ env:
  CUDA_ARCH_AARCH64: "8.0 8.7 8.9 9.0 10.0 11.0 12.0+PTX"
  CUDA_ARCH_X86_CU129: "7.5 8.0 8.6 8.9 9.0 10.0 12.0"
  CUDA_ARCH_AARCH64_CU129: "8.0 8.7 8.9 9.0 10.0 12.0"
-  MOONCAKE_WHEEL_AARCH64_2_35: "https://vllm-wheels.s3.amazonaws.com/mooncake/mooncake_transfer_engine-0.3.10.post2-0da9dfea3-cp312-cp312-manylinux_2_35_aarch64.whl"
-  MOONCAKE_WHEEL_AARCH64_2_39: "https://vllm-wheels.s3.amazonaws.com/mooncake/mooncake_transfer_engine-0.3.10.post2-0da9dfea3-cp312-cp312-manylinux_2_39_aarch64.whl"
-  MOONCAKE_WHEEL_X86_64: "https://vllm-wheels.s3.amazonaws.com/mooncake/mooncake_transfer_engine-0.3.10.post2-0da9dfea3-cp312-cp312-manylinux_2_35_x86_64.whl"

 steps:
  - input: "Provide Release version here"
@@ -27,11 +24,10 @@ steps:
        agents:
          queue: arm64_cpu_queue_release
        commands:
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64_CU129}\" --build-arg BUILD_OS=manylinux --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda12.9 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64_CU129}\" --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
          - "mkdir artifacts"
          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
          - "bash .buildkite/scripts/upload-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"

@@ -45,7 +41,6 @@ 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"

@@ -59,7 +54,6 @@ 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"

@@ -69,11 +63,10 @@ steps:
        agents:
          queue: cpu_queue_release
        commands:
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86_CU129}\" --build-arg BUILD_OS=manylinux --build-arg BUILD_BASE_IMAGE=pytorch/manylinux2_28-builder:cuda12.9 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86_CU129}\" --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
          - "mkdir artifacts"
          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
          - "bash .buildkite/scripts/upload-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"

@@ -87,7 +80,6 @@ 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"

@@ -101,7 +93,6 @@ 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"

@@ -139,8 +130,6 @@ steps:
              --build-arg CUDA_VERSION=13.0.2 \
              --build-arg torch_cuda_arch_list="${CUDA_ARCH_X86}" \
              --build-arg INSTALL_KV_CONNECTORS=true \
-              --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_35}" \
-              --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \
              --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 \
              --target vllm-openai \
              --progress plain \
@@ -149,7 +138,6 @@ 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: ~
@@ -167,14 +155,11 @@ steps:
              --build-arg CUDA_VERSION=13.0.2 \
              --build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64}" \
              --build-arg INSTALL_KV_CONNECTORS=true \
-              --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_35}" \
-              --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \
              --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 \
              --target vllm-openai \
              --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: ~
@@ -192,8 +177,6 @@ steps:
              --build-arg CUDA_VERSION=12.9.1 \
              --build-arg torch_cuda_arch_list="${CUDA_ARCH_X86_CU129}" \
              --build-arg INSTALL_KV_CONNECTORS=true \
-              --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_35}" \
-              --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \
              --target vllm-openai \
              --progress plain \
              -f docker/Dockerfile .
@@ -201,7 +184,6 @@ 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: ~
@@ -219,13 +201,10 @@ steps:
              --build-arg CUDA_VERSION=12.9.1 \
              --build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64_CU129}" \
              --build-arg INSTALL_KV_CONNECTORS=true \
-              --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_35}" \
-              --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \
              --target vllm-openai \
              --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: ~
@@ -245,8 +224,6 @@ steps:
              --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \
              --build-arg torch_cuda_arch_list="${CUDA_ARCH_X86}" \
              --build-arg INSTALL_KV_CONNECTORS=true \
-              --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_39}" \
-              --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \
              --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu24.04 \
              --target vllm-openai \
              --progress plain \
@@ -254,7 +231,6 @@ 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: ~
@@ -274,14 +250,11 @@ steps:
              --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \
              --build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64}" \
              --build-arg INSTALL_KV_CONNECTORS=true \
-              --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_39}" \
-              --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \
              --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu24.04 \
              --target vllm-openai \
              --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: ~
@@ -301,15 +274,12 @@ steps:
              --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \
              --build-arg torch_cuda_arch_list="${CUDA_ARCH_X86_CU129}" \
              --build-arg INSTALL_KV_CONNECTORS=true \
-              --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_39}" \
-              --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \
              --target vllm-openai \
              --progress plain \
              -f docker/Dockerfile .
          - "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: ~
@@ -329,13 +299,10 @@ steps:
              --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \
              --build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64_CU129}" \
              --build-arg INSTALL_KV_CONNECTORS=true \
-              --build-arg MOONCAKE_WHEEL_AARCH64="${MOONCAKE_WHEEL_AARCH64_2_39}" \
-              --build-arg MOONCAKE_WHEEL_X86_64="${MOONCAKE_WHEEL_X86_64}" \
              --target vllm-openai \
              --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
@@ -353,7 +320,6 @@ 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"

@@ -373,7 +339,6 @@ 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"

@@ -391,7 +356,15 @@ 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"
-          - 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: CUDA 13.0" "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"

      - label: "Create multi-arch manifest - CUDA 12.9"
        depends_on:
@@ -404,7 +377,6 @@ 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:
@@ -417,7 +389,6 @@ 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:
@@ -430,7 +401,6 @@ 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:
@@ -468,6 +438,59 @@ 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)
  # =============================================================================
@@ -581,7 +604,7 @@ steps:
          echo ""
          echo " Build complete - Image and wheels cached"
        fi
-
+        
    artifact_paths:
      - "artifacts/rocm-base-wheels/*.whl"
    env:
@@ -797,7 +820,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"
@@ -824,60 +847,3 @@ 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"
@@ -1,9 +0,0 @@
-#!/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'
@@ -0,0 +1,27 @@
+#!/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
@@ -114,7 +114,8 @@ handle_pytest_exit() {
 # unquoted since they have no spaces and work fine.
 #
 # Already-quoted expressions (containing literal single quotes) are passed
-# through untouched to avoid double-quoting well-formed shell fragments.
+# through untouched to avoid double-quoting values injected by
+# apply_rocm_test_overrides.
 #
 # NOTE: This ONLY fixes -m/-k flags. It cannot recover arbitrary inner
 # double-quotes stripped by the calling shell (see header comment).
@@ -246,6 +247,102 @@ re_quote_pytest_markers() {
  echo "${output% }"
 }

+###############################################################################
+# ROCm-specific pytest command rewrites
+#
+# These apply ignore flags and environment overrides for tests that are not
+# yet supported or behave differently on ROCm hardware. Kept as a single
+# function so new exclusions are easy to add in one place.
+###############################################################################
+
+apply_rocm_test_overrides() {
+  local cmds="$1"
+
+  # --- Model registry filter ---
+  if [[ $cmds == *"pytest -v -s models/test_registry.py"* ]]; then
+    cmds=${cmds//"pytest -v -s models/test_registry.py"/"pytest -v -s models/test_registry.py -k 'not BambaForCausalLM and not GritLM and not Mamba2ForCausalLM and not Zamba2ForCausalLM'"}
+  fi
+
+  # --- LoRA: disable custom paged attention ---
+  if [[ $cmds == *"pytest -v -s lora"* ]]; then
+    cmds=${cmds//"pytest -v -s lora"/"pytest -v -s lora"}
+  fi
+
+  # --- Kernel ignores ---
+  if [[ $cmds == *" kernels/core"* ]]; then
+    cmds="${cmds} \
+    --ignore=kernels/core/test_fused_quant_layernorm.py \
+    --ignore=kernels/core/test_permute_cols.py"
+  fi
+
+  if [[ $cmds == *" kernels/attention"* ]]; then
+    cmds="${cmds} \
+    --ignore=kernels/attention/test_attention_selector.py \
+    --ignore=kernels/attention/test_encoder_decoder_attn.py \
+    --ignore=kernels/attention/test_flash_attn.py \
+    --ignore=kernels/attention/test_flashinfer.py \
+    --ignore=kernels/attention/test_prefix_prefill.py \
+    --ignore=kernels/attention/test_cascade_flash_attn.py \
+    --ignore=kernels/attention/test_mha_attn.py \
+    --ignore=kernels/attention/test_lightning_attn.py \
+    --ignore=kernels/attention/test_attention.py"
+  fi
+
+  if [[ $cmds == *" kernels/quantization"* ]]; then
+    cmds="${cmds} \
+    --ignore=kernels/quantization/test_int8_quant.py \
+    --ignore=kernels/quantization/test_machete_mm.py \
+    --ignore=kernels/quantization/test_block_fp8.py \
+    --ignore=kernels/quantization/test_block_int8.py \
+    --ignore=kernels/quantization/test_marlin_gemm.py \
+    --ignore=kernels/quantization/test_cutlass_scaled_mm.py \
+    --ignore=kernels/quantization/test_int8_kernel.py"
+  fi
+
+  if [[ $cmds == *" kernels/mamba"* ]]; then
+    cmds="${cmds} \
+    --ignore=kernels/mamba/test_mamba_mixer2.py \
+    --ignore=kernels/mamba/test_causal_conv1d.py \
+    --ignore=kernels/mamba/test_mamba_ssm_ssd.py"
+  fi
+
+  if [[ $cmds == *" kernels/moe"* ]]; then
+    cmds="${cmds} \
+    --ignore=kernels/moe/test_moe.py \
+    --ignore=kernels/moe/test_cutlass_moe.py"
+  fi
+
+  # --- Entrypoint ignores ---
+  if [[ $cmds == *" entrypoints/openai "* ]]; then
+    cmds=${cmds//" entrypoints/openai "/" entrypoints/openai \
+    --ignore=entrypoints/openai/chat_completion/test_audio.py \
+    --ignore=entrypoints/openai/completion/test_shutdown.py \
+    --ignore=entrypoints/openai/test_completion.py \
+    --ignore=entrypoints/openai/models/test_models.py \
+    --ignore=entrypoints/openai/test_return_tokens_as_ids.py \
+    --ignore=entrypoints/openai/chat_completion/test_root_path.py \
+    --ignore=entrypoints/openai/completion/test_prompt_validation.py "}
+  fi
+
+  if [[ $cmds == *" entrypoints/serve"* ]]; then
+    cmds="${cmds} \
+    --ignore=entrypoints/serve/lora/test_lora_adapters.py"
+  fi
+
+  if [[ $cmds == *" entrypoints/llm "* ]]; then
+    cmds=${cmds//" entrypoints/llm "/" entrypoints/llm \
+    --ignore=entrypoints/llm/test_chat.py \
+    --ignore=entrypoints/llm/test_accuracy.py \
+    --ignore=entrypoints/llm/test_init.py \
+    --ignore=entrypoints/llm/test_prompt_validation.py "}
+  fi
+
+  # Clean up escaped newlines from --ignore appends
+  cmds=$(echo "$cmds" | sed 's/ \\ / /g')
+
+  echo "$cmds"
+}
+
 ###############################################################################
 # Main
 ###############################################################################
@@ -312,6 +409,7 @@ else
  echo "Skipping re-quoting for VLLM_TEST_COMMANDS input"
 fi

+commands=$(apply_rocm_test_overrides "$commands")
 echo "Final commands: $commands"

 MYPYTHONPATH=".."
@@ -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.12" \
+    && python3 -m pip install --progress-bar off "lm-eval[api]>=0.4.11" \
    && 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.12" \
+    && python3 -m pip install --progress-bar off "lm-eval[api]>=0.4.11" \
    && python3 -m pip install --progress-bar off hf-transfer tblib==3.1.0
 echo "--- Python dependencies installed ---"

@@ -0,0 +1,53 @@
+#!/bin/bash
+
+# This script build the CPU docker image and run the offline inference inside the container.
+# It serves a sanity check for compilation and basic model usage.
+set -ex
+
+image_name="xpu/vllm-ci:${BUILDKITE_COMMIT}"
+container_name="xpu_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | head -c 10; echo)"
+
+# Try building the docker image
+docker build -t "${image_name}" -f docker/Dockerfile.xpu .
+
+# Setup cleanup
+remove_docker_container() {
+  docker rm -f "${container_name}" || true
+}
+trap remove_docker_container EXIT
+
+# Run the image and test offline inference/tensor parallel
+docker run \
+    --device /dev/dri:/dev/dri \
+    --net=host \
+    --ipc=host \
+    --privileged \
+    -v /dev/dri/by-path:/dev/dri/by-path \
+    --entrypoint="" \
+    -e "HF_TOKEN=${HF_TOKEN}" \
+    -e "ZE_AFFINITY_MASK=${ZE_AFFINITY_MASK}" \
+    --name "${container_name}" \
+    "${image_name}" \
+    bash -c '
+    set -e
+    echo $ZE_AFFINITY_MASK
+    python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager
+    python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 -O3 -cc.cudagraph_mode=NONE
+    python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager -tp 2 --distributed-executor-backend ray
+    python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager -tp 2 --distributed-executor-backend mp
+    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 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
+    python3 examples/basic/offline_inference/generate.py --model OPEA/Qwen2.5-0.5B-Instruct-int4-sym-inc --block-size 64 --enforce-eager --max-model-len 8192
+    cd tests
+    pytest -v -s v1/core --ignore=v1/core/test_reset_prefix_cache_e2e.py --ignore=v1/core/test_scheduler_e2e.py
+    pytest -v -s v1/engine
+    pytest -v -s v1/sample --ignore=v1/sample/test_logprobs.py --ignore=v1/sample/test_logprobs_e2e.py
+    pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py --ignore=v1/worker/test_worker_memory_snapshot.py
+    pytest -v -s v1/structured_output
+    pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_tree_attention.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/test_serial_utils.py
+'
@@ -39,17 +39,11 @@ fi

 set -x # avoid printing secrets above

-# 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
+# install twine and sdist build prerequisites from pypi
+python3 -m venv /tmp/vllm-release-env
 source /tmp/vllm-release-env/bin/activate
-uv pip install twine
-uv pip install -r requirements/build/cuda.txt
+pip install twine
+pip install -r requirements/build/cuda.txt
 python3 -m twine --version

 # copy release wheels to local directory
@@ -460,7 +460,7 @@ steps:
  - tests/lora
  - vllm/platforms/rocm.py
  commands:
-  - pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py --ignore=lora/test_qwen3_with_multi_loras.py --ignore=lora/test_olmoe_tp.py --ignore=lora/test_deepseekv2_tp.py --ignore=lora/test_gptoss_tp.py --ignore=lora/test_qwen3moe_tp.py --ignore=lora/test_qwen35_densemodel_lora.py
+  - pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py --ignore=lora/test_llm_with_multi_loras.py --ignore=lora/test_olmoe_tp.py --ignore=lora/test_deepseekv2_tp.py --ignore=lora/test_gptoss_tp.py --ignore=lora/test_qwen3moe_tp.py --ignore=lora/test_qwen35_densemodel_lora.py

 #------------------------------------------------------  mi250 · model_executor  -------------------------------------------------------#

@@ -798,6 +798,32 @@ 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]
@@ -903,7 +929,6 @@ steps:
  - tests/tokenizers_
  - tests/reasoning
  - tests/tool_parsers
-  - tests/parser
  - tests/transformers_utils
  - tests/config
  commands:
@@ -917,7 +942,6 @@ steps:
  - pytest -v -s tokenizers_
  - pytest -v -s reasoning --ignore=reasoning/test_seedoss_reasoning_parser.py --ignore=reasoning/test_glm4_moe_reasoning_parser.py
  - pytest -v -s tool_parsers
-  - pytest -v -s parser
  - pytest -v -s transformers_utils
  - pytest -v -s config

@@ -1253,7 +1277,6 @@ 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"
@@ -1274,7 +1297,6 @@ steps:
  agent_pool: mi300_1
  fast_check: true
  torch_nightly: true
-  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -1288,7 +1310,6 @@ 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,13 +1320,13 @@ steps:
  commands:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
  - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
+  - pytest -v -s entrypoints/openai/speech_to_text/
  - pytest -v -s entrypoints/test_chat_utils.py

 - label: Entrypoints Integration (API Server openai - Part 3) # TBD
  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"
@@ -1315,21 +1336,7 @@ 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
-
- label: Entrypoints Integration (Speech to Text) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  agent_pool: mi300_1
-  fast_check: true
-  torch_nightly: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/
-  - tests/entrypoints/speech_to_text
-  commands:
-  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/speech_to_text
+  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/speech_to_text/ --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/test_multi_api_servers.py

 - label: Entrypoints Integration (LLM) # TBD
  timeout_in_minutes: 180
@@ -1706,7 +1713,6 @@ 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:
@@ -1754,7 +1760,7 @@ steps:
  - export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
  - pytest -v -s -x lora/test_chatglm3_tp.py
  - pytest -v -s -x lora/test_llama_tp.py
-  - pytest -v -s -x lora/test_qwen3_with_multi_loras.py
+  - pytest -v -s -x lora/test_llm_with_multi_loras.py
  - pytest -v -s -x lora/test_olmoe_tp.py
  - pytest -v -s -x lora/test_gptoss_tp.py
  - pytest -v -s -x lora/test_qwen35_densemodel_lora.py
@@ -1765,7 +1771,6 @@ 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/
@@ -2761,6 +2766,7 @@ steps:
  commands:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
  - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
+  - pytest -v -s entrypoints/openai/speech_to_text/
  - pytest -v -s entrypoints/test_chat_utils.py

 - label: Entrypoints Integration (API Server openai - Part 3) # TBD
@@ -2776,21 +2782,7 @@ 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
-
- label: Entrypoints Integration (Speech to Text) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi355]
-  agent_pool: mi355_1
-  fast_check: true
-  torch_nightly: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/
-  - tests/entrypoints/speech_to_text
-  commands:
-  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/speech_to_text
+  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/speech_to_text/ --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/test_multi_api_servers.py

 - label: Entrypoints Integration (Pooling) # TBD
  timeout_in_minutes: 180
@@ -27,5 +27,4 @@ steps:
  - vllm/compilation
  commands:
    - pytest -v -s v1/cudagraph/test_cudagraph_dispatch.py
-    - pytest -v -s v1/cudagraph/test_cudagraph_mode.py
-    - pytest -v -s v1/cudagraph/test_breakable_cudagraph.py
+    - pytest -v -s v1/cudagraph/test_cudagraph_mode.py
@@ -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: 25
+  timeout_in_minutes: 20
  working_dir: "/vllm-workspace/tests"
  num_devices: 4
  source_file_dependencies:
@@ -36,12 +36,6 @@ 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
@@ -52,12 +46,6 @@ 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
@@ -11,7 +11,7 @@ steps:
  - tests/entrypoints/
  commands:
  - pytest -v -s entrypoints/openai/tool_parsers
-  - pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/rpc --ignore=entrypoints/sleep --ignore=entrypoints/serve/instrumentator --ignore=entrypoints/openai --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py  --ignore=entrypoints/pooling --ignore=entrypoints/speech_to_text
+  - pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/rpc --ignore=entrypoints/sleep --ignore=entrypoints/serve/instrumentator --ignore=entrypoints/openai --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py  --ignore=entrypoints/pooling

 - label: Entrypoints Integration (LLM)
  key: entrypoints-integration-llm
@@ -43,12 +43,7 @@ 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
@@ -60,13 +55,8 @@ steps:
  - tests/entrypoints/test_chat_utils
  commands:
  - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
+  - pytest -v -s entrypoints/openai/speech_to_text/
  - 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
@@ -79,13 +69,7 @@ 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
-  mirror:
-    amd:
-      device: mi300_1
-      timeout_in_minutes: 60
-      depends_on:
-      - image-build-amd
+  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/speech_to_text/ --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/test_multi_api_servers.py

 - label: Entrypoints Integration (API Server 2)
  key: entrypoints-integration-api-server-2
@@ -101,22 +85,6 @@ 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
-  timeout_in_minutes: 50
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/
-  - tests/entrypoints/speech_to_text
-  commands:
-  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/speech_to_text

 - label: Entrypoints Integration (Pooling)
  key: entrypoints-integration-pooling
@@ -147,5 +115,5 @@ steps:
  - csrc/
  - vllm/entrypoints/openai/
  - vllm/model_executor/models/whisper.py
-  commands: # LMEval
+  commands: # LMEval+Transcription WER check
  - pytest -s entrypoints/openai/correctness/
@@ -62,19 +62,6 @@ 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
@@ -104,7 +91,6 @@ 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
@@ -119,11 +105,7 @@ 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
@@ -133,7 +115,6 @@ 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
@@ -261,7 +242,7 @@ steps:
 - label: Kernels FusedMoE Layer Test (2 B200s)
  key: kernels-fusedmoe-layer-test-2-b200s
  timeout_in_minutes: 90
-  device: b200-k8s
+  device: b200
  num_devices: 2
  source_file_dependencies:
  - csrc/quantization/cutlass_w8a8/moe/
@@ -40,7 +40,7 @@ steps:
 - label: LM Eval Small Models (B200)
  key: lm-eval-small-models-b200
  timeout_in_minutes: 120
-  device: b200-k8s
+  device: b200
  optional: true
  source_file_dependencies:
  - csrc/
@@ -48,18 +48,6 @@ steps:
  commands:
  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-blackwell.txt

- label: LM Eval Large Models (B200, EP)
-  key: lm-eval-large-models-b200-ep
-  timeout_in_minutes: 120
-  device: b200-k8s
-  optional: true
-  num_devices: 2
-  source_file_dependencies:
-  - csrc/
-  - vllm/model_executor/layers/quantization
-  commands:
-  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-blackwell-ep.txt
-
 - label: LM Eval Qwen3.5 Models (B200)
  key: lm-eval-qwen3-5-models-b200
  timeout_in_minutes: 120
@@ -104,7 +92,7 @@ steps:

 - label: MoE Refactor Integration Test (B200 DP - TEMPORARY)
  key: moe-refactor-integration-test-b200-dp-temporary
-  device: b200-k8s
+  device: b200
  optional: true
  num_devices: 2
  commands:
@@ -114,7 +102,6 @@ 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
@@ -140,7 +127,7 @@ steps:
 - label: GPQA Eval (GPT-OSS) (B200)
  key: gpqa-eval-gpt-oss-b200
  timeout_in_minutes: 120
-  device: b200-k8s
+  device: b200
  optional: true
  num_devices: 2
  source_file_dependencies:
@@ -9,7 +9,7 @@ steps:
  - vllm/lora
  - tests/lora
  commands:
-    - pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py --ignore=lora/test_qwen3_with_multi_loras.py --ignore=lora/test_olmoe_tp.py --ignore=lora/test_deepseekv2_tp.py --ignore=lora/test_gptoss_tp.py --ignore=lora/test_qwen3moe_tp.py --ignore=lora/test_qwen35_densemodel_lora.py 
+    - pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py --ignore=lora/test_llm_with_multi_loras.py --ignore=lora/test_olmoe_tp.py --ignore=lora/test_deepseekv2_tp.py --ignore=lora/test_gptoss_tp.py --ignore=lora/test_qwen3moe_tp.py --ignore=lora/test_qwen35_densemodel_lora.py 
  parallelism: 4


@@ -31,7 +31,7 @@ steps:
    # requires multi-GPU testing for validation.
    - pytest -v -s -x lora/test_chatglm3_tp.py
    - pytest -v -s -x lora/test_llama_tp.py
-    - pytest -v -s -x lora/test_qwen3_with_multi_loras.py
+    - pytest -v -s -x lora/test_llm_with_multi_loras.py
    - pytest -v -s -x lora/test_olmoe_tp.py
    - pytest -v -s -x lora/test_gptoss_tp.py
    - pytest -v -s -x lora/test_qwen35_densemodel_lora.py
@@ -79,7 +79,6 @@ 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
@@ -91,7 +90,6 @@ 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
@@ -129,19 +127,6 @@ 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
@@ -225,9 +210,7 @@ steps:
 - label: Python-only Installation
  key: python-only-installation
  depends_on: ~
-  optional: true
  timeout_in_minutes: 20
-  device: h200_18gb
  source_file_dependencies:
  - tests/standalone_tests/python_only_compile.sh
  - setup.py
@@ -299,7 +282,6 @@ steps:
  - tests/tokenizers_
  - tests/reasoning
  - tests/tool_parsers
-  - tests/parser
  - tests/transformers_utils
  - tests/config
  device: cpu-small
@@ -314,24 +296,9 @@ steps:
  - pytest -v -s tokenizers_
  - pytest -v -s reasoning --ignore=reasoning/test_seedoss_reasoning_parser.py --ignore=reasoning/test_glm4_moe_reasoning_parser.py
  - pytest -v -s tool_parsers
-  - pytest -v -s parser
  - 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,12 +106,13 @@ steps:
  - vllm/v1/worker/gpu/
  - vllm/v1/worker/gpu_worker.py
  - tests/v1/spec_decode/test_max_len.py
-  - tests/v1/spec_decode/test_rejection_sampler_utils.py
+  - tests/v1/spec_decode/test_probabilistic_rejection_sampler_utils.py
+  - tests/v1/spec_decode/test_synthetic_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_rejection_sampler_utils.py
+  - pytest -v -s v1/spec_decode/test_probabilistic_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,10 +5,23 @@ 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/
+  - vllm/config/
+  - vllm/distributed/
+  - vllm/engine/
+  - vllm/inputs/
+  - vllm/logging_utils/
+  - vllm/model_executor/
+  - vllm/multimodal/
+  - vllm/platforms/
+  - vllm/sampling_params.py
+  - vllm/sequence.py
+  - vllm/tasks.py
+  - vllm/transformers_utils/
+  - vllm/utils/
+  - vllm/v1/
+  - vllm/version.py
  - tests/models/test_initialization.py
  - tests/models/registry.py
  commands:
@@ -37,7 +50,21 @@ steps:
  key: basic-models-tests-other
  timeout_in_minutes: 45
  source_file_dependencies:
-  - vllm/
+  - vllm/config/
+  - vllm/distributed/
+  - vllm/engine/
+  - vllm/inputs/
+  - vllm/logging_utils/
+  - vllm/model_executor/
+  - vllm/multimodal/
+  - vllm/platforms/
+  - vllm/sampling_params.py
+  - vllm/sequence.py
+  - vllm/tasks.py
+  - vllm/transformers_utils/
+  - vllm/utils/
+  - vllm/v1/
+  - vllm/version.py
  - tests/models/test_terratorch.py
  - tests/models/test_transformers.py
  - tests/models/test_registry.py
@@ -50,7 +77,11 @@ steps:
  - image-build-cpu
  timeout_in_minutes: 10
  source_file_dependencies:
-  - vllm/
+  - vllm/config/
+  - vllm/distributed/
+  - vllm/model_executor/
+  - vllm/platforms/
+  - vllm/utils/
  - tests/models/test_utils.py
  - tests/models/test_vision.py
  device: cpu-small
@@ -5,7 +5,6 @@ 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
@@ -92,12 +91,6 @@ 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
@@ -16,7 +16,7 @@ steps:
 - label: Spec Decode Eagle Nightly B200
  key: spec-decode-eagle-nightly-b200
  timeout_in_minutes: 30
-  device: b200-k8s
+  device: b200
  optional: true
  source_file_dependencies:
    - vllm/v1/spec_decode/
@@ -40,7 +40,7 @@ steps:
 - label: Spec Decode Speculators + MTP Nightly B200
  key: spec-decode-speculators-mtp-nightly-b200
  timeout_in_minutes: 30
-  device: b200-k8s
+  device: b200
  optional: true
  source_file_dependencies:
    - vllm/v1/spec_decode/
@@ -84,23 +84,23 @@ steps:
  commands:
    - pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"

- label: Speculators Correctness
-  key: speculators-correctness
-  timeout_in_minutes: 60
+- label: DFlash Speculators Correctness
+  key: dflash-speculators-correctness
+  timeout_in_minutes: 30
  device: h100
  optional: true
  num_devices: 1
  source_file_dependencies:
    - vllm/v1/spec_decode/
    - vllm/model_executor/models/qwen3_dflash.py
-    - tests/v1/spec_decode/test_speculators_correctness.py
+    - tests/v1/spec_decode/test_speculators_dflash.py
  commands:
    - export VLLM_ALLOW_INSECURE_SERIALIZATION=1
-    - pytest -v -s v1/spec_decode/test_speculators_correctness.py -m slow_test
+    - pytest -v -s v1/spec_decode/test_speculators_dflash.py -m slow_test

 - label: Spec Decode MTP hybrid (B200)
  timeout_in_minutes: 30
-  device: b200-k8s
+  device: b200
  optional: true
  source_file_dependencies:
    - vllm/v1/spec_decode/
@@ -31,8 +31,8 @@
 /vllm/entrypoints/cli @hmellor @mgoin @DarkLight1337 @russellb
 /vllm/entrypoints/mcp @heheda12345
 /vllm/entrypoints/openai @aarnphm @chaunceyjiang @DarkLight1337 @russellb
-/vllm/entrypoints/speech_to_text/realtime @njhill
-/vllm/entrypoints/speech_to_text @NickLucche
+/vllm/entrypoints/openai/realtime @njhill
+/vllm/entrypoints/openai/speech_to_text @NickLucche
 /vllm/entrypoints/pooling @noooop
 /vllm/entrypoints/sagemaker @DarkLight1337
 /vllm/entrypoints/serve @njhill
@@ -9,9 +9,6 @@ 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
@@ -9,7 +9,44 @@ build:
    python: "3.12"
  jobs:
    post_checkout:
-      - bash docs/pre_run_check.sh
+      - |
+        if [ "$READTHEDOCS_VERSION_TYPE" = "external" ]; then
+          MAX_WAIT=300
+          INTERVAL=60
+          ELAPSED=0
+          while :; do
+            RAW=$(curl -sS -w "\n%{http_code}" "https://api.github.com/repos/vllm-project/vllm/commits/${READTHEDOCS_GIT_COMMIT_HASH}/check-runs?check_name=pre-run-check&filter=latest")
+            HTTP_CODE=$(printf %s "$RAW" | tail -n1)
+            BODY=$(printf %s "$RAW" | head -n -1)
+            if [ "$HTTP_CODE" != "200" ]; then
+              echo "GitHub API returned HTTP $HTTP_CODE (likely rate-limited); skipping pre-run-check gate."
+              break
+            fi
+            STATUS=$(printf %s "$BODY" | python3 -c "import sys, json; r=json.load(sys.stdin).get(\"check_runs\",[]); print((r[0].get(\"status\") or \"\") if r else \"none\")")
+            CONCLUSION=$(printf %s "$BODY" | python3 -c "import sys, json; r=json.load(sys.stdin).get(\"check_runs\",[]); print((r[0].get(\"conclusion\") or \"\") if r else \"\")")
+            if [ "$STATUS" = "none" ]; then
+              echo "no pre-run-check found for this commit; skipping gate."
+              break
+            fi
+            if [ -n "$CONCLUSION" ]; then
+              echo "pre-run-check conclusion: $CONCLUSION"
+              if [ "$CONCLUSION" = "failure" ] || [ "$CONCLUSION" = "cancelled" ] || [ "$CONCLUSION" = "timed_out" ]; then
+                echo "pre-run-check did not pass; failing docs build."
+                exit 1
+              fi
+              break
+            fi
+            if [ "$ELAPSED" -ge "$MAX_WAIT" ]; then
+              echo "pre-run-check status=$STATUS after ${MAX_WAIT}s; skipping gate."
+              break
+            fi
+            echo "pre-run-check status=$STATUS; waiting ${INTERVAL}s..."
+            sleep "$INTERVAL"
+            ELAPSED=$((ELAPSED + INTERVAL))
+          done
+        else
+          echo "Not a PR build (version type=$READTHEDOCS_VERSION_TYPE); skipping pre-run-check gate."
+        fi
      - git fetch origin main --unshallow --no-tags --filter=blob:none || true
    pre_create_environment:
      - pip install uv
@@ -109,24 +109,6 @@ else()
  set(CUDA_SUPPORTED_ARCHS "7.0;7.5;8.0;8.6;8.7;8.9;9.0")
 endif()

-#
-# spinloop extension (pure CXX; must stay above the non-CUDA device branch so
-# CPU builds define the target before the early return)
-#
-set(VLLM_SPINLOOP_EXT_SRC "csrc/spinloop.cpp")
-set(SPINLOOP_COMPILE_FLAGS "")
-if(CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|amd64")
-  list(APPEND SPINLOOP_COMPILE_FLAGS "-mmwaitx")
-endif()
-define_extension_target(
-  spinloop
-  DESTINATION vllm
-  LANGUAGE CXX
-  SOURCES ${VLLM_SPINLOOP_EXT_SRC}
-  COMPILE_FLAGS ${SPINLOOP_COMPILE_FLAGS}
-  USE_SABI 3.11
-  WITH_SOABI)
-
 #
 # Forward the non-CUDA device extensions to external CMake scripts.
 #
@@ -311,6 +293,7 @@ set(VLLM_EXT_SRC
  "csrc/attention/paged_attention_v1.cu"
  "csrc/attention/paged_attention_v2.cu"
  "csrc/attention/merge_attn_states.cu"
+  "csrc/attention/vertical_slash_index.cu"
  "csrc/pos_encoding_kernels.cu"
  "csrc/activation_kernels.cu"
  "csrc/layernorm_kernels.cu"
@@ -368,6 +351,7 @@ 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(
@@ -499,6 +483,46 @@ 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}")
@@ -526,6 +550,24 @@ 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

@@ -597,6 +639,16 @@ 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()
@@ -646,8 +698,7 @@ 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/awq/gemm_kernels.cu")
+      "csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu")
  endif()

  if(VLLM_GPU_LANG STREQUAL "CUDA")
@@ -656,40 +707,6 @@ 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)
  #
@@ -999,41 +1016,6 @@ 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,16 +1203,6 @@ 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,7 +53,6 @@ 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,7 +3,6 @@
 # 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).
@@ -121,7 +120,6 @@ prefill_backends:
  - flashinfer
  - cudnn
  - trtllm
-  - tokenspeed

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

 # ============================================================================
 # VllmConfig Creation
@@ -80,8 +79,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",
-                        "trtllm"). Configures the attention config to force
-                        the specified prefill backend.
+                        "cudnn", "trtllm"). Configures the attention config to
+                        force the specified prefill backend.

    Returns:
        VllmConfig for benchmarking
@@ -180,13 +179,19 @@ 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

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

@@ -608,21 +625,6 @@ 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:
@@ -959,6 +961,19 @@ 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,
@@ -970,36 +985,38 @@ def _run_mla_benchmark_batched(
            kv_cache_dtype=kv_cache_dtype,
        )

-        # 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.
+        # Verify the actual prefill backend matches what was requested
        if prefill_backend is not 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
-                )
+            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)
                if actual_fa_version != fa_version:
                    raise RuntimeError(
                        f"Prefill backend '{prefill_backend}' requested FA "
-                        f"version {fa_version}, got "
-                        f"{actual_fa_version} on {actual_class}."
+                        f"version {fa_version}, but the impl is using FA "
+                        f"version {actual_fa_version}. Check "
+                        f"vllm/v1/attention/backends/fa_utils.py."
                    )
+            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)
@@ -115,39 +115,6 @@ class SampleRequest:
 def sample_requests(
    tokenizer: PreTrainedTokenizerBase, args: argparse.Namespace
 ) -> list[SampleRequest]:
-    def _apply_random_prefix(
-        tokenizer: PreTrainedTokenizerBase,
-        requests: list[SampleRequest],
-        prefix_len: int,
-        seed: int,
-    ) -> list[SampleRequest]:
-        if prefix_len <= 0:
-            return requests
-        rng = np.random.default_rng(seed)
-        vocab_size = tokenizer.vocab_size
-        prohibited = getattr(tokenizer, "all_special_ids", None) or []
-        allowed = np.array([i for i in range(vocab_size) if i not in prohibited])
-        if len(allowed) == 0:
-            return requests
-        prefix_ids = rng.integers(0, len(allowed), size=prefix_len)
-        prefix_token_ids = allowed[prefix_ids].tolist()
-        out = []
-        for req in requests:
-            prompt_ids = tokenizer(req.prompt, add_special_tokens=False).input_ids
-            full_ids = prefix_token_ids + prompt_ids
-            full_prompt = tokenizer.decode(full_ids, skip_special_tokens=False)
-            out.append(
-                SampleRequest(
-                    prompt=full_prompt,
-                    prompt_len=len(tokenizer(full_prompt).input_ids),
-                    expected_output_len=req.expected_output_len,
-                    schema=req.schema,
-                    structure_type=req.structure_type,
-                    completion=req.completion,
-                )
-            )
-        return out
-
    if args.dataset == "json" or args.dataset == "json-unique":
        if args.json_schema_path is None:
            dir_path = os.path.dirname(os.path.realpath(__file__))
@@ -294,9 +261,6 @@ def sample_requests(
                )
            )

-    requests = _apply_random_prefix(
-        tokenizer, requests, args.random_prefix_len, args.seed
-    )
    return requests


@@ -981,15 +945,6 @@ def create_argument_parser():
        "results in a more uniform arrival of requests.",
    )
    parser.add_argument("--seed", type=int, default=0)
-    parser.add_argument(
-        "--random-prefix-len",
-        type=int,
-        default=0,
-        help=(
-            "Number of prefix tokens to prepend to every prompt. "
-            "The same prefix is used for all prompts to enable prefix caching."
-        ),
-    )
    parser.add_argument(
        "--trust-remote-code",
        action="store_true",
@@ -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.experts.triton_deep_gemm_moe import (
+from vllm.model_executor.layers.fused_moe.fused_moe import *
+from vllm.model_executor.layers.fused_moe.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
@@ -1,183 +0,0 @@
-# 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,6 +12,7 @@ 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
@@ -21,14 +22,15 @@ def get_attn_isa(
    block_size: int | None = None,
    dtype: torch.dtype | None = None,
 ):
-    # 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,
-    )
+    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"


 # rand number generation takes too much time, cache rand tensors
@@ -233,7 +235,7 @@ if __name__ == "__main__":
    )
    parser.add_argument("--use-sink", action="store_true")
    parser.add_argument(
-        "--isa", type=str, choices=["vec", "neon", "amx", "vec16", "rvv"], default=None
+        "--isa", type=str, choices=["vec", "neon", "amx", "vec16"], default=None
    )
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--iters", type=int, default=20)
@@ -195,12 +195,10 @@ 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,80 +53,48 @@ cuda_archs_loose_intersection(DEEPGEMM_ARCHS
 if(DEEPGEMM_ARCHS)
  message(STATUS "DeepGEMM CUDA architectures: ${DEEPGEMM_ARCHS}")

-  #
-  # 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.
-  #
+  find_package(CUDAToolkit REQUIRED)

-  # 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}")
+  #
+  # Build the _C pybind11 extension from DeepGEMM's C++ source.
+  # This is a CXX-only module — CUDA kernels are JIT-compiled at runtime.
+  #
+  Python_add_library(_deep_gemm_C MODULE WITH_SOABI
+    "${deepgemm_SOURCE_DIR}/csrc/python_api.cpp")

-  # 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")
+  # The pybind11 module name must be _C to match DeepGEMM's Python imports.
+  set_target_properties(_deep_gemm_C PROPERTIES OUTPUT_NAME "_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})
+  target_compile_definitions(_deep_gemm_C PRIVATE
+    "-DTORCH_EXTENSION_NAME=_C")
+
+  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)

  #
  # 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 bce29425653ec0fbc579d329883030e832d15ada
+          GIT_TAG f5bc33cfc02c744d24a2e9d50e6db656de40611c
          GIT_PROGRESS TRUE
          # Don't share the vllm-flash-attn build between build types
          BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
@@ -18,12 +18,13 @@ limitations under the License.
 * Taken from SGLANG PR https://github.com/sgl-project/sglang/pull/6929
 * by Alcanderian JieXin Liang
 */
-#include "libtorch_stable/torch_utils.h"
-
-#include <torch/csrc/stable/library.h>
+#include "core/registration.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>
@@ -34,27 +35,27 @@ limitations under the License.
 // clang-format off
 #if !defined(CUDA_VERSION) || CUDA_VERSION < 12040
 void sm100_cutlass_mla_decode(
-    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,
+    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,
    double sm_scale,
    int64_t num_kv_splits) {
-  STD_TORCH_CHECK(false, "CUDA version must be >= 12.4 for cutlass_mla_decode");
+  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) {
-  STD_TORCH_CHECK(false, "CUDA version must be >= 12.4 for cutlass_mla_get_workspace_size");
+  TORCH_CHECK(false, "CUDA version must be >= 12.4 for cutlass_mla_get_workspace_size");
 }
 #else

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

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

 template <typename T>
 typename T::Fmha::Arguments args_from_options(
-    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,
+    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,
    double sm_scale,
    int64_t num_kv_splits) {
  cutlass::KernelHardwareInfo hw_info;
-  hw_info.device_id = q_nope.get_device_index();
+  hw_info.device_id = q_nope.device().index();
  hw_info.sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id);

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

 template <typename Element, typename ElementOut, bool IsPaged128, typename PersistenceOption>
 void runMla(
-    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,
+    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,
    double sm_scale,
    int64_t num_kv_splits,
    cudaStream_t stream) {
@@ -219,37 +220,37 @@ void runMla(
  }()

 void sm100_cutlass_mla_decode(
-    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,
+    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,
    double sm_scale,
    int64_t num_kv_splits) {
-  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);
-
+  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];
+  
  // 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 == torch::headeronly::ScalarType::Half) {
+      if (in_dtype == at::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 == torch::headeronly::ScalarType::BFloat16) {
+      } else if (in_dtype == at::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 == torch::headeronly::ScalarType::Float8_e4m3fn) {
+      } else if (in_dtype == at::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 {
-        STD_TORCH_CHECK(false, "Unsupported input data type of MLA");
+        TORCH_CHECK(false, "Unsupported input data type of MLA");
      }
      return true;
    });
@@ -279,12 +280,12 @@ int64_t sm100_cutlass_mla_get_workspace_size(int64_t max_seq_len, int64_t num_ba

 #endif

-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, CUDA, m) {
+  m.impl("sm100_cutlass_mla_decode", &sm100_cutlass_mla_decode);
 }

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

 // clang-format on
@@ -0,0 +1,401 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+#include <assert.h>
+
+#include <cuda.h>
+
+#include <torch/all.h>
+
+__device__ int64_t save_blocks(int* block_offset, int64_t range_start,
+                               int64_t range_end, int64_t block_size,
+                               int64_t input_block_count, int64_t kv_seqlen) {
+  if (range_start >= kv_seqlen) {
+    return input_block_count;
+  }
+  if (range_end > kv_seqlen) {
+    range_end = kv_seqlen;
+  }
+  int64_t current_block_count = input_block_count;
+  for (int idx = range_start; idx < range_end; idx += block_size) {
+    block_offset[current_block_count++] = idx;
+  }
+  return current_block_count;
+}
+
+__global__ void convert_vertical_slash_indexes_kernel(
+    const int* q_seqlens,         // [BATCH, ]
+    const int* kv_seqlens,        // [BATCH, ]
+    const int* vertical_indexes,  // [BATCH, N_HEADS, NNZ_V]
+    const int* slash_indexes,     // [BATCH, N_HEADS, NNZ_S]
+    int* block_count,             // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)]
+    int* block_offset,  // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_S]
+    int* column_count,  // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)]
+    int* column_index,  // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_V]
+    int64_t N_HEADS, int64_t N_ROWS, int64_t BLOCK_SIZE_M, int64_t BLOCK_SIZE_N,
+    int64_t NNZ_V, int64_t NNZ_S,
+    bool causal  // True for intra, False for succ
+) {
+  const int batch_idx = blockIdx.y;
+  const int head_idx = blockIdx.x;
+  const int group_idx = blockIdx.z;
+
+  int64_t q_seqlen = q_seqlens[batch_idx];
+  int64_t kv_seqlen = kv_seqlens[batch_idx];
+  int64_t block_idx_m = group_idx * blockDim.x + threadIdx.x;
+  int64_t start_m = block_idx_m * BLOCK_SIZE_M;
+  if (start_m >= q_seqlen) {
+    return;
+  }
+  int64_t end_m = start_m + BLOCK_SIZE_M;
+  vertical_indexes += (batch_idx * N_HEADS + head_idx) * NNZ_V;
+  slash_indexes += (batch_idx * N_HEADS + head_idx) * NNZ_S;
+  int64_t row_offset = (batch_idx * N_HEADS + head_idx) * N_ROWS + block_idx_m;
+  block_count += row_offset;
+  block_offset += row_offset * NNZ_S;
+  column_count += row_offset;
+  column_index += row_offset * NNZ_V;
+
+  bool has_slash = true;
+  int64_t tmp_col_cnt = 0, tmp_blk_cnt = 0;
+  int64_t s = 0, v = 0;
+  int64_t v_idx = vertical_indexes[v++];
+  int64_t s_idx = slash_indexes[s++];
+  if (causal) {
+    while (s_idx >= end_m + (kv_seqlen - q_seqlen) && s < NNZ_S) {
+      s_idx = slash_indexes[s++];
+    }
+    if (s_idx > end_m + (kv_seqlen - q_seqlen)) has_slash = false;
+    s_idx = max((kv_seqlen - q_seqlen) + end_m - s_idx, BLOCK_SIZE_M);
+  } else {
+    while (s_idx >= end_m + kv_seqlen && s < NNZ_S) {
+      s_idx = slash_indexes[s++];
+    }
+    if (s_idx > end_m + kv_seqlen) has_slash = false;
+    s_idx = max(kv_seqlen + end_m - s_idx, BLOCK_SIZE_M);
+  }
+
+  int64_t range_start = s_idx - BLOCK_SIZE_M, range_end = s_idx;
+  if (!has_slash) {
+    if (causal) {
+      range_start = (kv_seqlen - q_seqlen) + end_m;
+      range_end = (kv_seqlen - q_seqlen) + end_m + BLOCK_SIZE_N;
+    } else {
+      range_start = kv_seqlen;
+      range_end = kv_seqlen + BLOCK_SIZE_N;
+    }
+  }
+
+  bool slash_finished = false;
+  while (1) {
+    if (v_idx < range_end) {
+      if (v_idx < range_start) {
+        column_index[tmp_col_cnt++] = v_idx;
+      }
+      if (v < NNZ_V) {
+        v_idx = vertical_indexes[v++];
+      } else {
+        if (causal)
+          v_idx = end_m + BLOCK_SIZE_N + (kv_seqlen - q_seqlen);
+        else
+          v_idx = end_m + BLOCK_SIZE_N + kv_seqlen;
+      }
+    } else {
+      if ((s < NNZ_S && causal) ||
+          (s < NNZ_S && !causal && slash_indexes[s] >= start_m)) {
+        if (causal)
+          s_idx = max((kv_seqlen - q_seqlen) + end_m - slash_indexes[s++],
+                      BLOCK_SIZE_M);
+        else
+          s_idx = max(kv_seqlen + end_m - slash_indexes[s++], BLOCK_SIZE_M);
+      } else {
+        if (v == NNZ_V || (v_idx > range_start && causal)) {
+          // add the last vertical if no more slash
+          if (v == NNZ_V && !causal && v_idx < kv_seqlen) {
+            column_index[tmp_col_cnt++] = v_idx;
+          }
+          tmp_blk_cnt = save_blocks(block_offset, range_start, range_end,
+                                    BLOCK_SIZE_N, tmp_blk_cnt, kv_seqlen);
+          break;
+        } else {
+          if (causal) {
+            range_start = (kv_seqlen - q_seqlen) + end_m;
+            range_end = (kv_seqlen - q_seqlen) + end_m + BLOCK_SIZE_N;
+          } else {
+            // if slash_finished but there are vertical left, save current
+            // blocks
+            tmp_blk_cnt = save_blocks(block_offset, range_start, range_end,
+                                      BLOCK_SIZE_N, tmp_blk_cnt, kv_seqlen);
+            range_start = kv_seqlen;
+            range_end = kv_seqlen + BLOCK_SIZE_N;
+          }
+          slash_finished = true;
+        }
+      }
+      if (!slash_finished) {
+        if (s_idx > range_end + BLOCK_SIZE_M) {
+          tmp_blk_cnt = save_blocks(block_offset, range_start, range_end,
+                                    BLOCK_SIZE_N, tmp_blk_cnt, kv_seqlen);
+          range_start = s_idx - BLOCK_SIZE_M;
+          range_end = s_idx;
+        } else if (s_idx > range_end) {
+          range_end += BLOCK_SIZE_M;
+        }
+      }
+    }
+  }
+
+  block_count[0] = tmp_blk_cnt;
+  column_count[0] = tmp_col_cnt;
+}
+
+void convert_vertical_slash_indexes_64x64(
+    const int* q_seqlens,         // [BATCH, ]
+    const int* kv_seqlens,        // [BATCH, ]
+    const int* vertical_indexes,  // [BATCH, N_HEADS, NNZ_V]
+    const int* slash_indexes,     // [BATCH, N_HEADS, NNZ_S]
+    int* block_count,             // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)]
+    int* block_offset,  // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_S]
+    int* column_count,  // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)]
+    int* column_index,  // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_V]
+    int64_t BATCH_SIZE, int64_t N_HEADS, int64_t N_ROWS, int64_t BLOCK_SIZE_M,
+    int64_t BLOCK_SIZE_N, int64_t NNZ_V, int64_t NNZ_S, bool causal) {
+  const int N_THREADS = 64;
+  const dim3 dimBlock(N_THREADS);
+  const dim3 dimGrid(N_HEADS, BATCH_SIZE, (N_ROWS + N_THREADS - 1) / N_THREADS);
+  convert_vertical_slash_indexes_kernel<<<dimGrid, dimBlock>>>(
+      q_seqlens, kv_seqlens, vertical_indexes, slash_indexes, block_count,
+      block_offset, column_count, column_index, N_HEADS, N_ROWS, BLOCK_SIZE_M,
+      BLOCK_SIZE_N, NNZ_V, NNZ_S, causal);
+}
+
+/**
+ * Implements the Algorithm 4 in paper https://arxiv.org/abs/2407.02490.
+ *
+ * This function builds the index of each row of blocks from vertical indices
+ * and slash indices. The vertical indices are treated as points, while the
+ * slash indices are converted as ranges. The output consists of the merged
+ * ranges and separate column indices, where the ranges are represented by
+ * block indices.
+ *
+ * The implementation is referenced from the original MInference repo:
+ * https://github.com/microsoft/MInference/blob/main/csrc/vertical_slash_index.cu.
+ */
+void convert_vertical_slash_indexes(
+    torch::Tensor& block_count,      // [BATCH, N_HEADS, NUM_ROWS]
+    torch::Tensor& block_offset,     // [BATCH, N_HEADS, NUM_ROWS, NNZ_S]
+    torch::Tensor& column_count,     // [BATCH, N_HEADS, NUM_ROWS]
+    torch::Tensor& column_index,     // [BATCH, N_HEADS, NUM_ROWS, NNZ_V]
+    torch::Tensor q_seqlens,         // [BATCH, ]
+    torch::Tensor kv_seqlens,        // [BATCH, ]
+    torch::Tensor vertical_indexes,  // [BATCH, N_HEADS, NNZ_V]
+    torch::Tensor slash_indexes,     // [BATCH, N_HEADS, NNZ_S]
+    int64_t context_size, int64_t block_size_M, int64_t block_size_N,
+    bool causal) {
+  cudaSetDevice(q_seqlens.get_device());
+
+  int batch_size = slash_indexes.size(0);
+  int num_heads = slash_indexes.size(1);
+  int nnz_slash = slash_indexes.size(2);
+  int nnz_vertical = vertical_indexes.size(2);
+  int num_rows = (context_size + block_size_M - 1) / block_size_M;
+
+  convert_vertical_slash_indexes_64x64(
+      q_seqlens.data_ptr<int>(), kv_seqlens.data_ptr<int>(),
+      vertical_indexes.data_ptr<int>(), slash_indexes.data_ptr<int>(),
+      block_count.data_ptr<int>(), block_offset.data_ptr<int>(),
+      column_count.data_ptr<int>(), column_index.data_ptr<int>(), batch_size,
+      num_heads, num_rows, block_size_M, block_size_N, nnz_vertical, nnz_slash,
+      causal);
+}
+
+__global__ void convert_vertical_slash_indexes_kernel_mergehead(
+    const int* q_seqlens,         // [BATCH, ]
+    const int* kv_seqlens,        // [BATCH, ]
+    const int* vertical_indexes,  // [BATCH, N_HEADS, NNZ_V]
+    const int* slash_indexes,     // [BATCH, N_HEADS, NNZ_S]
+    const int* per_head_vertical_topkv, const int* per_head_slash_topkv,
+    int* block_count,   // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)]
+    int* block_offset,  // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_S]
+    int* column_count,  // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)]
+    int* column_index,  // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_V]
+    int64_t N_HEADS, int64_t N_ROWS, int64_t BLOCK_SIZE_M, int64_t BLOCK_SIZE_N,
+    int64_t NNZ_V, int64_t NNZ_S,
+    bool causal  // True for intra, False for succ
+) {
+  const int batch_idx = blockIdx.y;
+  const int head_idx = blockIdx.x;
+  const int group_idx = blockIdx.z;
+
+  int64_t q_seqlen = q_seqlens[batch_idx];
+  int64_t kv_seqlen = kv_seqlens[batch_idx];
+  int64_t block_idx_m = group_idx * blockDim.x + threadIdx.x;
+  int64_t start_m = block_idx_m * BLOCK_SIZE_M;
+  if (start_m >= q_seqlen) {
+    return;
+  }
+  int64_t end_m = start_m + BLOCK_SIZE_M;
+  vertical_indexes += (batch_idx * N_HEADS + head_idx) * NNZ_V;
+  slash_indexes += (batch_idx * N_HEADS + head_idx) * NNZ_S;
+  int64_t row_offset = (batch_idx * N_HEADS + head_idx) * N_ROWS + block_idx_m;
+  block_count += row_offset;
+  block_offset += row_offset * NNZ_S;
+  column_count += row_offset;
+  column_index += row_offset * NNZ_V;
+
+  // MergeHead: each head has it's unique max topk NNZ_V，NNZ_S. (NNZ_V，NNZ_S
+  // above is buffer size, use to compute offset)
+  NNZ_S = per_head_slash_topkv[head_idx];
+  NNZ_V = per_head_vertical_topkv[head_idx];
+
+  bool has_slash = true;
+  int64_t tmp_col_cnt = 0, tmp_blk_cnt = 0;
+  int64_t s = 0, v = 0;
+  int64_t v_idx = vertical_indexes[v++];
+  int64_t s_idx = slash_indexes[s++];
+  if (causal) {
+    while (s_idx >= end_m + (kv_seqlen - q_seqlen) && s < NNZ_S) {
+      s_idx = slash_indexes[s++];
+    }
+    if (s_idx > end_m + (kv_seqlen - q_seqlen)) has_slash = false;
+    s_idx = max((kv_seqlen - q_seqlen) + end_m - s_idx, BLOCK_SIZE_M);
+  } else {
+    while (s_idx >= end_m + kv_seqlen && s < NNZ_S) {
+      s_idx = slash_indexes[s++];
+    }
+    if (s_idx > end_m + kv_seqlen) has_slash = false;
+    s_idx = max(kv_seqlen + end_m - s_idx, BLOCK_SIZE_M);
+  }
+
+  int64_t range_start = s_idx - BLOCK_SIZE_M, range_end = s_idx;
+  if (!has_slash) {
+    if (causal) {
+      range_start = (kv_seqlen - q_seqlen) + end_m;
+      range_end = (kv_seqlen - q_seqlen) + end_m + BLOCK_SIZE_N;
+    } else {
+      range_start = kv_seqlen;
+      range_end = kv_seqlen + BLOCK_SIZE_N;
+    }
+  }
+
+  bool slash_finished = false;
+  while (1) {
+    if (v_idx < range_end) {
+      if (v_idx < range_start) {
+        column_index[tmp_col_cnt++] = v_idx;
+      }
+      if (v < NNZ_V) {
+        v_idx = vertical_indexes[v++];
+      } else {
+        if (causal)
+          v_idx = end_m + BLOCK_SIZE_N + (kv_seqlen - q_seqlen);
+        else
+          v_idx = end_m + BLOCK_SIZE_N + kv_seqlen;
+      }
+    } else {
+      if ((s < NNZ_S && causal) ||
+          (s < NNZ_S && !causal && slash_indexes[s] >= start_m)) {
+        if (causal)
+          s_idx = max((kv_seqlen - q_seqlen) + end_m - slash_indexes[s++],
+                      BLOCK_SIZE_M);
+        else
+          s_idx = max(kv_seqlen + end_m - slash_indexes[s++], BLOCK_SIZE_M);
+      } else {
+        if (v == NNZ_V || (v_idx > range_start && causal)) {
+          // add the last vertical if no more slash
+          if (v == NNZ_V && !causal && v_idx < kv_seqlen) {
+            column_index[tmp_col_cnt++] = v_idx;
+          }
+          tmp_blk_cnt = save_blocks(block_offset, range_start, range_end,
+                                    BLOCK_SIZE_N, tmp_blk_cnt, kv_seqlen);
+          break;
+        } else {
+          if (causal) {
+            range_start = (kv_seqlen - q_seqlen) + end_m;
+            range_end = (kv_seqlen - q_seqlen) + end_m + BLOCK_SIZE_N;
+          } else {
+            // if slash_finished but there are vertical left, save current
+            // blocks
+            tmp_blk_cnt = save_blocks(block_offset, range_start, range_end,
+                                      BLOCK_SIZE_N, tmp_blk_cnt, kv_seqlen);
+            range_start = kv_seqlen;
+            range_end = kv_seqlen + BLOCK_SIZE_N;
+          }
+          slash_finished = true;
+        }
+      }
+      if (!slash_finished) {
+        if (s_idx > range_end + BLOCK_SIZE_M) {
+          tmp_blk_cnt = save_blocks(block_offset, range_start, range_end,
+                                    BLOCK_SIZE_N, tmp_blk_cnt, kv_seqlen);
+          range_start = s_idx - BLOCK_SIZE_M;
+          range_end = s_idx;
+        } else if (s_idx > range_end) {
+          range_end += BLOCK_SIZE_M;
+        }
+      }
+    }
+  }
+
+  block_count[0] = tmp_blk_cnt;
+  column_count[0] = tmp_col_cnt;
+}
+
+void convert_vertical_slash_indexes_64x64_mergehead(
+    const int* q_seqlens,         // [BATCH, ]
+    const int* kv_seqlens,        // [BATCH, ]
+    const int* vertical_indexes,  // [BATCH, N_HEADS, NNZ_V]
+    const int* slash_indexes,     // [BATCH, N_HEADS, NNZ_S]
+    int* per_head_vertical_topkv, int* per_head_slash_topkv,
+    int* block_count,   // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)]
+    int* block_offset,  // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_S]
+    int* column_count,  // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)]
+    int* column_index,  // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_V]
+    int64_t BATCH_SIZE, int64_t N_HEADS, int64_t N_ROWS, int64_t BLOCK_SIZE_M,
+    int64_t BLOCK_SIZE_N, int64_t NNZ_V, int64_t NNZ_S, bool causal) {
+  const int N_THREADS = 64;
+  const dim3 dimBlock(N_THREADS);
+  const dim3 dimGrid(N_HEADS, BATCH_SIZE, (N_ROWS + N_THREADS - 1) / N_THREADS);
+  convert_vertical_slash_indexes_kernel_mergehead<<<dimGrid, dimBlock>>>(
+      q_seqlens, kv_seqlens, vertical_indexes, slash_indexes,
+      per_head_vertical_topkv, per_head_slash_topkv, block_count, block_offset,
+      column_count, column_index, N_HEADS, N_ROWS, BLOCK_SIZE_M, BLOCK_SIZE_N,
+      NNZ_V, NNZ_S, causal);
+}
+
+/**
+ * Implements the Algorithm 4 in paper https://arxiv.org/abs/2407.02490.
+ *
+ * Like the above convert_vertical_slash_indexes, but with
+ * pre-computed vertical and slash counts.
+ */
+void convert_vertical_slash_indexes_mergehead(
+    torch::Tensor& block_count,            // [BATCH, N_HEADS, NUM_ROWS]
+    torch::Tensor& block_offset,           // [BATCH, N_HEADS, NUM_ROWS, NNZ_S]
+    torch::Tensor& column_count,           // [BATCH, N_HEADS, NUM_ROWS]
+    torch::Tensor& column_index,           // [BATCH, N_HEADS, NUM_ROWS, NNZ_V]
+    torch::Tensor q_seqlens,               // [BATCH, ]
+    torch::Tensor kv_seqlens,              // [BATCH, ]
+    torch::Tensor vertical_indexes,        // [BATCH, N_HEADS, NNZ_V]
+    torch::Tensor slash_indexes,           // [BATCH, N_HEADS, NNZ_S]
+    torch::Tensor vertical_indices_count,  // [N_HEADS, ]
+    torch::Tensor slash_indices_count,     // [N_HEADS, ]
+    int64_t context_size, int64_t block_size_M, int64_t block_size_N,
+    bool causal) {
+  cudaSetDevice(q_seqlens.get_device());
+
+  int batch_size = slash_indexes.size(0);
+  int num_heads = slash_indexes.size(1);
+  int nnz_slash = slash_indexes.size(2);
+  int nnz_vertical = vertical_indexes.size(2);
+  int num_rows = (context_size + block_size_M - 1) / block_size_M;
+
+  convert_vertical_slash_indexes_64x64_mergehead(
+      q_seqlens.data_ptr<int>(), kv_seqlens.data_ptr<int>(),
+      vertical_indexes.data_ptr<int>(), slash_indexes.data_ptr<int>(),
+      vertical_indices_count.data_ptr<int>(),
+      slash_indices_count.data_ptr<int>(), block_count.data_ptr<int>(),
+      block_offset.data_ptr<int>(), column_count.data_ptr<int>(),
+      column_index.data_ptr<int>(), batch_size, num_heads, num_rows,
+      block_size_M, block_size_N, nnz_vertical, nnz_slash, causal);
+}
@@ -21,33 +21,28 @@ namespace vllm {

 // NOTE Be EXTRA careful with raw_kv_scalar_t, for __half and __nv_bfloat16 it's
 // using u16 as the backing type.
-template <typename qk_t, typename cos_sin_t, bool IS_NEOX,
-          typename raw_kv_scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
+template <typename qk_t, bool IS_NEOX, typename raw_kv_scalar_t,
+          typename cache_t, Fp8KVCacheDataType kv_dt>
 __global__ void concat_and_cache_mla_rope_fused_kernel(
    const int64_t* __restrict__ positions,  // [num_tokens]
    qk_t* __restrict__ q_pe,        // [num_tokens, num_q_heads, rot_dim]
    qk_t* __restrict__ k_pe,        // [num_tokens, rot_dim]
    const qk_t* __restrict__ kv_c,  // [num_tokens, kv_lora_rank]
-    const cos_sin_t* __restrict__ rope_cos_sin_cache,  // [max_position, 2,
-                                                       // rot_dim // 2]
+    const qk_t* __restrict__ rope_cos_sin_cache,  // [max_position, 2,
+                                                  // rot_dim // 2]
    const int rot_dim, const int64_t q_pe_stride_token,
    const int64_t q_pe_stride_head, const int64_t k_pe_stride,
    const int64_t kv_c_stride, const int num_q_heads,
    cache_t* __restrict__ kv_cache,  // [num_blocks, block_size, (kv_lora_rank +
                                     // rot_dim)]
-    const int64_t* __restrict__ slot_mapping,  // [num_tokens]
+    const int64_t* __restrict__ kv_cache_slot_mapping,  // [num_tokens]
    const int block_stride, const int entry_stride, const int kv_lora_rank,
    const int block_size, const float* kv_cache_quant_scale) {
  // Each thread block is responsible for one token.
  const int64_t token_idx = blockIdx.x;
-  const int64_t slot_idx = slot_mapping[token_idx];
-  // NOTE: slot_idx can be -1 if the token is padded
-  if (slot_idx < 0) {
-    return;
-  }
  const int64_t pos = positions[token_idx];

-  const cos_sin_t* cos_sin_ptr = rope_cos_sin_cache + pos * rot_dim;
+  const qk_t* cos_sin_ptr = rope_cos_sin_cache + pos * rot_dim;

  const int embed_dim = rot_dim / 2;

@@ -59,8 +54,8 @@ __global__ void concat_and_cache_mla_rope_fused_kernel(

    // NOTE: Would be nice to have interleaved sin/cos so we could just load
    // both at the same time.
-    qk_t cos = static_cast<qk_t>(VLLM_LDG(cos_sin_ptr + pair_idx));
-    qk_t sin = static_cast<qk_t>(VLLM_LDG(cos_sin_ptr + pair_idx + embed_dim));
+    qk_t cos = VLLM_LDG(cos_sin_ptr + pair_idx);
+    qk_t sin = VLLM_LDG(cos_sin_ptr + pair_idx + embed_dim);

    qk_t* q_pe_head_ptr =
        q_pe + token_idx * q_pe_stride_token + head_idx * q_pe_stride_head;
@@ -86,15 +81,21 @@ __global__ void concat_and_cache_mla_rope_fused_kernel(
    q_pe_head_ptr[pair_idx_y] = y_dst;
  }

+  const int64_t slot_idx = kv_cache_slot_mapping[token_idx];
  const int64_t block_idx = slot_idx / block_size;
  const int64_t entry_idx = slot_idx % block_size;

+  // NOTE: slot_idx can be -1 if the token is padded
+  if (slot_idx < 0) {
+    return;
+  }
+
  // K with 1 HEAD
  for (int i = threadIdx.x; i < embed_dim; i += blockDim.x) {
    int pair_idx = i;

-    qk_t cos = static_cast<qk_t>(VLLM_LDG(cos_sin_ptr + pair_idx));
-    qk_t sin = static_cast<qk_t>(VLLM_LDG(cos_sin_ptr + pair_idx + embed_dim));
+    qk_t cos = VLLM_LDG(cos_sin_ptr + pair_idx);
+    qk_t sin = VLLM_LDG(cos_sin_ptr + pair_idx + embed_dim);

    qk_t* k_pe_head_ptr = k_pe + token_idx * k_pe_stride;

@@ -164,43 +165,36 @@ __global__ void concat_and_cache_mla_rope_fused_kernel(

 }  // namespace vllm

-#define CALL_CONCAT_AND_CACHE_MLA_ROPE_FUSED(RAW_KV_T, CACHE_T, KV_DTYPE)     \
-  do {                                                                        \
-    VLLM_DISPATCH_FLOATING_TYPES(q_pe.scalar_type(), "qk_scalar_type", [&] {  \
-      using qk_t = scalar_t;                                                  \
-      VLLM_DISPATCH_FLOATING_TYPES(                                           \
-          rope_cos_sin_cache.scalar_type(), "rope_cos_sin_cache_scalar_type", \
-          [&] {                                                               \
-            using cos_sin_t = scalar_t;                                       \
-            if (rope_is_neox) {                                               \
-              vllm::concat_and_cache_mla_rope_fused_kernel<                   \
-                  qk_t, cos_sin_t, true, RAW_KV_T, CACHE_T, KV_DTYPE>         \
-                  <<<grid, block, 0, stream>>>(                               \
-                      positions.data_ptr<int64_t>(), q_pe.data_ptr<qk_t>(),   \
-                      k_pe.data_ptr<qk_t>(), kv_c.data_ptr<qk_t>(),           \
-                      rope_cos_sin_cache.data_ptr<cos_sin_t>(), rot_dim,      \
-                      q_pe_stride_token, q_pe_stride_head, k_pe_stride,       \
-                      kv_c_stride, num_q_heads,                               \
-                      reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()),        \
-                      slot_mapping.data_ptr<int64_t>(), block_stride,         \
-                      entry_stride, kv_lora_rank, block_size,                 \
-                      kv_cache_quant_scale.data_ptr<float>());                \
-            } else {                                                          \
-              vllm::concat_and_cache_mla_rope_fused_kernel<                   \
-                  qk_t, cos_sin_t, false, RAW_KV_T, CACHE_T, KV_DTYPE>        \
-                  <<<grid, block, 0, stream>>>(                               \
-                      positions.data_ptr<int64_t>(), q_pe.data_ptr<qk_t>(),   \
-                      k_pe.data_ptr<qk_t>(), kv_c.data_ptr<qk_t>(),           \
-                      rope_cos_sin_cache.data_ptr<cos_sin_t>(), rot_dim,      \
-                      q_pe_stride_token, q_pe_stride_head, k_pe_stride,       \
-                      kv_c_stride, num_q_heads,                               \
-                      reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()),        \
-                      slot_mapping.data_ptr<int64_t>(), block_stride,         \
-                      entry_stride, kv_lora_rank, block_size,                 \
-                      kv_cache_quant_scale.data_ptr<float>());                \
-            }                                                                 \
-          });                                                                 \
-    });                                                                       \
+#define CALL_CONCAT_AND_CACHE_MLA_ROPE_FUSED(RAW_KV_T, CACHE_T, KV_DTYPE)      \
+  do {                                                                         \
+    VLLM_DISPATCH_FLOATING_TYPES(q_pe.scalar_type(), "qk_scalar_type", [&] {   \
+      using qk_t = scalar_t;                                                   \
+      if (rope_is_neox) {                                                      \
+        vllm::concat_and_cache_mla_rope_fused_kernel<qk_t, true, RAW_KV_T,     \
+                                                     CACHE_T, KV_DTYPE>        \
+            <<<grid, block, 0, stream>>>(                                      \
+                positions.data_ptr<int64_t>(), q_pe.data_ptr<qk_t>(),          \
+                k_pe.data_ptr<qk_t>(), kv_c.data_ptr<qk_t>(),                  \
+                rope_cos_sin_cache.data_ptr<qk_t>(), rot_dim,                  \
+                q_pe_stride_token, q_pe_stride_head, k_pe_stride, kv_c_stride, \
+                num_q_heads, reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()),  \
+                kv_cache_slot_mapping.data_ptr<int64_t>(), block_stride,       \
+                entry_stride, kv_lora_rank, block_size,                        \
+                kv_cache_quant_scale.data_ptr<float>());                       \
+      } else {                                                                 \
+        vllm::concat_and_cache_mla_rope_fused_kernel<qk_t, false, RAW_KV_T,    \
+                                                     CACHE_T, KV_DTYPE>        \
+            <<<grid, block, 0, stream>>>(                                      \
+                positions.data_ptr<int64_t>(), q_pe.data_ptr<qk_t>(),          \
+                k_pe.data_ptr<qk_t>(), kv_c.data_ptr<qk_t>(),                  \
+                rope_cos_sin_cache.data_ptr<qk_t>(), rot_dim,                  \
+                q_pe_stride_token, q_pe_stride_head, k_pe_stride, kv_c_stride, \
+                num_q_heads, reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()),  \
+                kv_cache_slot_mapping.data_ptr<int64_t>(), block_stride,       \
+                entry_stride, kv_lora_rank, block_size,                        \
+                kv_cache_quant_scale.data_ptr<float>());                       \
+      }                                                                        \
+    });                                                                        \
  } while (false)

 // Executes RoPE on q_pe and k_pe, then writes k_pe and kv_c in the kv cache.
@@ -214,52 +208,43 @@ void concat_and_cache_mla_rope_fused(
    torch::Tensor& kv_c,                // [num_tokens, kv_lora_rank]
    torch::Tensor& rope_cos_sin_cache,  // [max_position, rot_dim]
    bool rope_is_neox,
-    torch::Tensor& slot_mapping,  // [num_tokens] or [num_actual_tokens]
+    torch::Tensor&
+        kv_cache_slot_mapping,  // [num_tokens] or [num_actual_tokens]
    torch::Tensor&
        kv_cache,  // [num_blocks, block_size, (kv_lora_rank + rot_dim)]
    const std::string& kv_cache_dtype, torch::Tensor& kv_cache_quant_scale) {
-  // NOTE(woosuk): In vLLM V1, query/key/position.size(0) can be different from
-  // slot_mapping.size(0) because of padding for CUDA graphs.
-  // In vLLM V0, key.size(0) is always equal to slot_mapping.size(0) because
-  // both include padding.
-  // In vLLM V1, however, key.size(0) can be larger than slot_mapping.size(0)
-  // since key includes padding for CUDA graphs, while slot_mapping does not.
-  // In this case, slot_mapping.size(0) represents the actual number of tokens
-  // before padding.
-  // For compatibility with both cases, we use slot_mapping.size(0) as the
-  // number of tokens.
-  int num_tokens = slot_mapping.size(0);
-  int num_padded_tokens = q_pe.size(0);
-  TORCH_CHECK_GE(num_padded_tokens, num_tokens);
+  const int64_t num_tokens = q_pe.size(0);

  const int num_q_heads = q_pe.size(1);
  const int rot_dim = q_pe.size(2);
  const int kv_lora_rank = kv_c.size(1);

-  TORCH_CHECK_EQ(positions.size(0), num_padded_tokens);
+  TORCH_CHECK(positions.size(0) >=
+              num_tokens);  // CUDA Graphs might pad this for us
  TORCH_CHECK_EQ(positions.dim(), 1);
  TORCH_CHECK_EQ(positions.scalar_type(), c10::ScalarType::Long);

-  TORCH_CHECK_EQ(q_pe.dim(), 3);
-  TORCH_CHECK_EQ(q_pe.size(0), num_padded_tokens);
+  TORCH_CHECK_EQ(q_pe.size(0), num_tokens);
  TORCH_CHECK_EQ(q_pe.size(1), num_q_heads);
  TORCH_CHECK_EQ(q_pe.size(2), rot_dim);
+  TORCH_CHECK_EQ(q_pe.dim(), 3);

-  TORCH_CHECK_EQ(k_pe.dim(), 2);
-  TORCH_CHECK_EQ(k_pe.size(0), num_padded_tokens);
+  TORCH_CHECK_EQ(k_pe.size(0), num_tokens);
  TORCH_CHECK_EQ(k_pe.size(1), rot_dim);
+  TORCH_CHECK_EQ(k_pe.dim(), 2);
  TORCH_CHECK_EQ(k_pe.scalar_type(), q_pe.scalar_type());

-  TORCH_CHECK_EQ(kv_c.dim(), 2);
-  TORCH_CHECK_EQ(kv_c.size(0), num_padded_tokens);
+  TORCH_CHECK_EQ(kv_c.size(0), num_tokens);
  TORCH_CHECK_EQ(kv_c.size(1), kv_lora_rank);
+  TORCH_CHECK_EQ(kv_c.dim(), 2);
  TORCH_CHECK_EQ(kv_c.scalar_type(), q_pe.scalar_type());
  TORCH_CHECK_EQ(kv_c.dtype(), q_pe.dtype());

  TORCH_CHECK_EQ(rope_cos_sin_cache.size(1), rot_dim);
+  TORCH_CHECK_EQ(rope_cos_sin_cache.scalar_type(), q_pe.scalar_type());

-  TORCH_CHECK_EQ(slot_mapping.size(0), num_tokens);
-  TORCH_CHECK_EQ(slot_mapping.scalar_type(), c10::ScalarType::Long);
+  TORCH_CHECK_EQ(kv_cache_slot_mapping.size(0), num_tokens);
+  TORCH_CHECK_EQ(kv_cache_slot_mapping.scalar_type(), c10::ScalarType::Long);

  TORCH_CHECK_EQ(kv_cache.size(2), kv_lora_rank + rot_dim);
  TORCH_CHECK_EQ(kv_cache.dim(), 3);
@@ -1,6 +1,7 @@
 #pragma once
 #include <cstdlib>
 #include <string>
+#include <cctype>

 namespace vllm {

@@ -1,13 +1,7 @@
 #pragma once

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

 namespace vllm {

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

@@ -59,12 +53,11 @@ class ScalarType {
  static constexpr ScalarType float_(uint8_t exponent, uint8_t mantissa,
                                     bool finite_values_only,
                                     NanRepr nan_repr) {
-    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");
+    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");
    return ScalarType(exponent, mantissa, true, 0, finite_values_only,
                      nan_repr);
  }
@@ -183,8 +176,8 @@ class ScalarType {

 private:
  double _floating_point_max() const {
-    STD_TORCH_CHECK(mantissa <= 52 && exponent <= 11,
-                    "Cannot represent max/min as a double for type ", str());
+    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) {
@@ -193,8 +186,8 @@ class ScalarType {

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

@@ -223,17 +216,16 @@ class ScalarType {
    if (is_floating_point()) {
      return {_floating_point_max()};
    } else {
-      STD_TORCH_CHECK(size_bits() < 64 || size_bits() == 64 && is_signed(),
-                      "Cannot represent max as a int64_t");
+      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()) {
-      STD_TORCH_CHECK(
-          is_signed(),
-          "We currently assume all floating point types are signed");
+      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();
@@ -241,8 +233,8 @@ class ScalarType {
      uint64_t min_raw = max_raw | sign_bit_double;
      return {*reinterpret_cast<double*>(&min_raw)};
    } else {
-      STD_TORCH_CHECK(!is_signed() || size_bits() <= 64,
-                      "Cannot represent min as a int64_t");
+      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,20 +156,4 @@

 #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,8 +29,6 @@ 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 {
@@ -133,8 +131,6 @@ 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, RVV, VSX };
+enum class ISA { AMX, VEC, VEC16, NEON, VXE, VSX };

 // Mirrors csrc/attention/dtype_fp8.cuh Fp8KVCacheDataType exactly.
 enum class Fp8KVCacheDataType {
@@ -164,9 +164,6 @@ struct AttentionMetadata {
      case ISA::VXE:
        ss << "VXE, ";
        break;
-      case ISA::RVV:
-        ss << "RVV, ";
-        break;
      case ISA::VSX:
        ss << "VSX, ";
        break;
@@ -459,7 +456,7 @@ class AttentionScheduler {
    const int64_t kv_len_per_thread =
        (((total_kv_len / thread_num) + kv_len_alignment - 1) /
         kv_len_alignment) *
-        kv_len_alignment;
+        kv_len_alignment * (use_gqa ? input.num_heads_kv : input.num_heads_q);
    std::vector<AttentionWorkItemGroup> workitems;
    std::vector<ReductionWorkItemGroup> reduce_workitems;
    workitems.reserve(1024);
@@ -1,445 +0,0 @@
-// 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_zvfbfmin
+#ifdef RISCV_BF16_SUPPORT
 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,17 +15,8 @@
 #include <torch/all.h>
 namespace vec_op {

-// 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.
+// BFloat16 is always supported on RISC-V: natively when RISCV_BF16_SUPPORT
+// is defined, 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__)  \
@@ -115,7 +106,7 @@ struct FP16Vec16 : public Vec<FP16Vec16> {
 // BF16 Implementation
 // ============================================================================

-#ifdef __riscv_zvfbfmin
+#ifdef RISCV_BF16_SUPPORT

 FORCE_INLINE fixed_u16x8_t bf16_to_u16(fixed_bf16x8_t v) {
  return RVVI4(__riscv_vreinterpret_v_bf16, LMUL_128, _u16, LMUL_128)(v);
@@ -194,13 +185,6 @@ 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 =
@@ -339,13 +323,6 @@ 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);
@@ -433,7 +410,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_zvfbfmin
+#ifdef RISCV_BF16_SUPPORT
  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)
@@ -631,7 +608,7 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
  explicit FP32Vec16(const FP32Vec16& data) : reg(data.reg) {};
  explicit FP32Vec16(const FP16Vec16& v);

-#ifdef __riscv_zvfbfmin
+#ifdef RISCV_BF16_SUPPORT
  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)
@@ -891,7 +868,7 @@ inline void fma(FP32Vec16& acc, const FP32Vec16& a, const FP32Vec16& b) {
  acc = acc.fma(a, b);
 }

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

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

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

 # ISAs supported for head_dims divisible by 16 only
 ISA_FOR_16 = ["VEC16"]
@@ -150,15 +149,6 @@ 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
@@ -222,23 +212,6 @@ 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"],
@@ -260,8 +233,8 @@ def generate_header_file() -> str:
        fp8=False,
    )
    header += (
-        "#endif  /* CPU_CAPABILITY_AMXBF16 / __aarch64__ / __s390x__ /"
-        " __riscv / __powerpc__ */\n\n"
+        "#endif  /* CPU_CAPABILITY_AMXBF16 / __aarch64__ / "
+        "__s390x__ / __powerpc__ */\n\n"
        "#endif  // CPU_ATTN_DISPATCH_GENERATED_H\n"
    )

@@ -348,13 +348,6 @@ inline int get_cache_blocks<at::Float8_e4m3fn>(int chunk_size) {
  return std::min(MAX_CACHE_BLOCK_SIZE, cache_block_size);
 }

-template <>
-inline int get_cache_blocks<uint8_t>(int chunk_size) {
-  // mxfp4 uses bf16 as accumulate type
-  int cache_block_size = get_cache_blocks<at::BFloat16>(chunk_size);
-  return std::min(MAX_CACHE_BLOCK_SIZE, cache_block_size);
-}
-
 // 2d sequential loop in range : [mb0, mb1), [nb0, nb1)
 template <typename T, typename func_t>
 inline void loop_2d(int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1, int64_t chunk_size, const func_t& f) {
@@ -213,7 +213,7 @@ struct tinygemm_kernel<at::BFloat16, K, BLOCK_N, has_bias, has_silu> {
      weight + nb_start * width,                                                             \
      out + bs * seqlen * dim + mb_start * dim + nb_start,                                   \
      has_bias ? bias + nb_start : nullptr,                                                  \
-      has_conv_states ? conv_states + conv_state_index * conv_state_slot_stride + nb_start : nullptr, \
+      has_conv_states ? conv_states + conv_state_index * (K - 1) * dim + nb_start : nullptr, \
      has_initial_states_value,                                                              \
      mb_size,                                                                               \
      dim,                                                                                   \
@@ -233,8 +233,7 @@ void causal_conv1d_fwd_kernel_impl(
    int64_t dim,
    int64_t seqlen,
    int64_t width,
-    int64_t num_seq_blocks,
-    int64_t conv_state_slot_stride) {
+    int64_t num_seq_blocks) {
  // handle 32 x 64 per block
  constexpr int64_t BLOCK_M = block_size_m();
  constexpr int64_t BLOCK_N = block_size_n() * 2;
@@ -283,7 +282,7 @@ void causal_conv1d_fwd_kernel_impl(
    at::parallel_for(0, batch, 0, [&](int64_t begin, int64_t end) {
      for (int64_t bs = begin; bs < end; ++bs) {
        update_conv_state(
-            conv_states + bs * conv_state_slot_stride, input + bs * seqlen * dim, width, dim, seqlen, has_initial_state[bs]);
+            conv_states + bs * (width - 1) * dim, input + bs * seqlen * dim, width, dim, seqlen, has_initial_state[bs]);
      }
    });
  }
@@ -317,8 +316,7 @@ void causal_conv1d_fwd_varlen_kernel_impl(
    int64_t batch,
    int64_t dim,
    int64_t width,
-    int64_t num_seq_blocks,
-    int64_t conv_state_slot_stride) {
+    int64_t num_seq_blocks) {
  // handle 32 x 64 per block
  constexpr int64_t BLOCK_M = block_size_m();
  constexpr int64_t BLOCK_N = block_size_n() * 2;
@@ -368,7 +366,7 @@ void causal_conv1d_fwd_varlen_kernel_impl(
        int32_t seqlen = query_start_loc[bs + 1] - query_start_loc[bs];
        int32_t batch_offset = query_start_loc[bs];
        update_conv_state(
-            conv_states + conv_state_index * conv_state_slot_stride,
+            conv_states + conv_state_index * (width - 1) * dim,
            input + batch_offset * dim,
            width,
            dim,
@@ -391,8 +389,7 @@ void causal_conv1d_update_kernel_impl(
    int64_t batch,
    int64_t dim,
    int64_t seqlen,
-    int64_t width,
-    int64_t conv_state_slot_stride) {
+    int64_t width) {
  // handle 32 x 64 per block
  constexpr int64_t BLOCK_M = block_size_m();
  constexpr int64_t BLOCK_N = block_size_n() * 2;
@@ -433,7 +430,7 @@ void causal_conv1d_update_kernel_impl(
    });
  });

-#define CONV_STATE_INDEXR(w) conv_states + conv_state_index*conv_state_slot_stride + (w) * dim
+#define CONV_STATE_INDEXR(w) conv_states + conv_state_index*(width - 1) * dim + (w) * dim

  // update conv_states
  at::parallel_for(0, batch, 0, [&](int64_t begin, int64_t end) {
@@ -595,9 +592,6 @@ at::Tensor causal_conv1d_fwd_cpu(
    }
  }

-// IMPORTANT: To make the kernal compatible with vLLM KV cache layout 
-  int64_t conv_state_slot_stride = conv_states->stride(0);
-
  // block size for sequence blocks, 32
  constexpr int64_t BLOCK_M = block_size_m();

@@ -624,8 +618,7 @@ at::Tensor causal_conv1d_fwd_cpu(
          batch,
          dim,
          width,
-          num_seq_blocks,
-          conv_state_slot_stride);
+          num_seq_blocks);
    } else {
      causal_conv1d_fwd_kernel_impl<scalar_t>(
          out.data_ptr<scalar_t>(),
@@ -640,8 +633,7 @@ at::Tensor causal_conv1d_fwd_cpu(
          dim,
          seqlen,
          width,
-          num_seq_blocks,
-          conv_state_slot_stride);
+          num_seq_blocks);
    }
  });
  return out;
@@ -698,8 +690,6 @@ at::Tensor causal_conv1d_update_cpu(
    conv_states.copy_(conv_states_copy);
  }

-  // IMPORTANT: To make the kernal compatible with vLLM KV cache layout 
-  int64_t conv_state_slot_stride = conv_states.stride(0);
  at::Tensor out = at::empty_like(x);
  AT_DISPATCH_REDUCED_FLOATING_TYPES(scalar_type, "causal_conv1d_update_kernel_impl", [&] {
    causal_conv1d_update_kernel_impl<scalar_t>(
@@ -713,8 +703,7 @@ at::Tensor causal_conv1d_update_cpu(
        batch,
        dim,
        seqlen,
-        width,
-        conv_state_slot_stride);
+        width);
  });
  return out;
 }
@@ -847,7 +847,6 @@ void fused_sigmoid_gating_delta_rule_update_kernel_impl(
    int64_t v_strideB,
    int64_t v_strideS,
    int64_t v_strideH,
-    int64_t state_slot_stride,
    bool use_qk_l2norm_in_kernel,
    double softplus_threshold) {
  using bVec = at::vec::Vectorized<scalar_t>;
@@ -908,7 +907,7 @@ void fused_sigmoid_gating_delta_rule_update_kernel_impl(
    data_index_init(begin, bi, batch_size, si, seq_len, ni, v_num_heads);
    for (int64_t i = begin; i < end; ++i) {
      int64_t cache_index = indices_ptr[bi];
-      int64_t state_offset = cache_index * state_slot_stride + ni * head_dim * v_head_dim;
+      int64_t state_offset = (cache_index * v_num_heads + ni) * head_dim * v_head_dim;
      float g_val = -std::exp(float(A_log_ptr[ni])) *
                    softplus(float(a_ptr[bi * v_num_heads + ni]) + float(dt_bias_ptr[ni]), softplus_threshold);
      float g_val_exp = std::exp(g_val);
@@ -1322,8 +1321,6 @@ at::Tensor fused_sigmoid_gating_delta_rule_update_cpu(
  int64_t v_strideB = v.stride(1);
  int64_t v_strideS = v.stride(0);
  int64_t v_strideH = v.stride(2);
-  // IMPORTANT: To make the kernal compatible with vLLM KV cache layout 
-  int64_t state_slot_stride = initial_state_source.stride(0);
  at::Tensor core_attn_out = at::empty({batch_size, seq_len, v_num_heads, v_head_dim}, q.options());
  at::Tensor qk_scale_buf = at::empty({2 * batch_size, seq_len, num_heads}, at::kFloat);

@@ -1356,7 +1353,6 @@ at::Tensor fused_sigmoid_gating_delta_rule_update_cpu(
            v_strideB,
            v_strideS,
            v_strideH,
-            state_slot_stride,
            use_qk_l2norm_in_kernel,
            softplus_threshold);
      });
@@ -72,17 +72,7 @@ inline int64_t get_row_size(int64_t K, bool use_int8_w8a8) {
  return use_int8_w8a8 ? K + sizeof(int32_t) : K;
 }

-enum class CPUAcTMethod : int { silu_and_mul = 0, swiglu = 1 };
-
-constexpr bool operator==(CPUAcTMethod a, int b) {
-  return static_cast<int>(a) == b;
-}
-
-constexpr bool operator==(int a, CPUAcTMethod b) {
-  return a == static_cast<int>(b);
-}
-
-enum class CPUQuantMethod : int64_t { BF16 = 0, INT8_W8A8 = 1, FP8_W8A16 = 2, INT4_W4A8 = 3, MXFP4 = 4 };
+enum class CPUQuantMethod : int64_t { BF16 = 0, INT8_W8A8 = 1, FP8_W8A16 = 2, INT4_W4A8 = 3 };

 constexpr bool operator==(CPUQuantMethod a, int64_t b) {
  return static_cast<int64_t>(a) == b;
@@ -109,9 +99,6 @@ inline int64_t get_4bit_block_k_size(int64_t group_size) {
 // pack weight to vnni format
 at::Tensor convert_weight_packed(at::Tensor& weight);

-// pack scale to blocked format for mxfp4
-at::Tensor convert_scale_packed(at::Tensor& scale);
-
 // pack weight to vnni format for int4
 std::tuple<at::Tensor, at::Tensor, at::Tensor>
 convert_weight_packed_scale_zp(at::Tensor qweight, at::Tensor qzeros, at::Tensor scales);
@@ -142,9 +129,9 @@ void fused_experts_int8_kernel_impl(
    int64_t topk,
    int64_t num_tokens_post_pad);

-// moe implementations for fp8 w8a16 and mxfp4
-template <typename scalar_t, typename packed_t, typename param_t, bool is_mxfp4>
-void fused_experts_fp_kernel_impl(
+// moe implementations for fp8 w8a16
+template <typename scalar_t>
+void fused_experts_fp8_kernel_impl(
    scalar_t* __restrict__ output,
    scalar_t* __restrict__ ic0,
    scalar_t* __restrict__ ic1,
@@ -153,12 +140,10 @@ void fused_experts_fp_kernel_impl(
    scalar_t* __restrict__ B_tmp,
    float* __restrict__ C_tmp,
    const scalar_t* __restrict__ input,
-    const packed_t* __restrict__ packed_w1,
-    const packed_t* __restrict__ packed_w2,
-    const float* __restrict__ w1_bias,
-    const float* __restrict__ w2_bias,
-    const param_t* __restrict__ w1s,
-    const param_t* __restrict__ w2s,
+    const at::Float8_e4m3fn* __restrict__ packed_w1,
+    const at::Float8_e4m3fn* __restrict__ packed_w2,
+    const float* __restrict__ w1s,
+    const float* __restrict__ w2s,
    int64_t block_size_N,
    int64_t block_size_K,
    const float* __restrict__ topk_weights,
@@ -170,11 +155,7 @@ void fused_experts_fp_kernel_impl(
    int64_t K,
    int64_t E,
    int64_t topk,
-    int64_t num_tokens_post_pad,
-    float alpha,
-    float limit,
-    CPUAcTMethod act_func,
-    bool with_bias);
+    int64_t num_tokens_post_pad);

 // shared expert implementation for int8 w8a8
 template <typename scalar_t>
@@ -285,7 +266,6 @@ void tinygemm_kernel(
    scalar_t* __restrict__ C,
    scalar_t* __restrict__ Btmp,
    float* __restrict__ Ctmp,
-    const float* __restrict__ Bbias,
    const float* __restrict__ scale,
    int64_t M,
    int64_t N,
@@ -314,26 +294,6 @@ void tinygemm_kernel(
    int64_t ldc,
    bool brg);

-// mxfp4
-template <typename scalar_t>
-void tinygemm_kernel(
-    const scalar_t* __restrict__ A,
-    const uint8_t* __restrict__ B,
-    scalar_t* __restrict__ C,
-    scalar_t* __restrict__ Btmp,
-    float* __restrict__ Ctmp,
-    const float* __restrict__ Bbias,
-    const uint8_t* __restrict__ scale,
-    int64_t M,
-    int64_t N,
-    int64_t K,
-    int64_t lda,
-    int64_t ldb,
-    int64_t ldc,
-    bool brg,
-    int64_t block_size_K,
-    bool do_unpack = true);
-
 template <typename scalar_t>
 void tinygemm_kernel(
    scalar_t* C,
@@ -67,23 +67,6 @@ inline void copy_mul_stub(scalar_t* __restrict__ out, const float* __restrict__
  }
 }

-template <>
-inline void
-copy_add_stub(float* __restrict__ out, const float* __restrict__ input, const float* __restrict__ bias, int64_t size) {
-  using fVec = at::vec::Vectorized<float>;
-  constexpr int kVecSize = fVec::size();
-
-  int64_t d;
-#pragma GCC unroll 4
-  for (d = 0; d <= size - kVecSize; d += kVecSize) {
-    fVec data = fVec::loadu(input + d) + fVec::loadu(bias + d);
-    data.store(out + d);
-  }
-  for (; d < size; ++d) {
-    out[d] = input[d] + bias[d];
-  }
-}
-
 inline void unpack_B(
    at::BFloat16* __restrict__ Btmp,
    const at::Float8_e4m3fn* __restrict__ packed_B,
@@ -360,6 +343,7 @@ struct tinygemm_kernel_nn<at::BFloat16, at::Float8_e4m3fn, float, has_bias, BLOC
    Unroll<ROWS * COLS>{}(storec);
  }
 };
+
 template <int BLOCK_M, int BLOCK_N>
 struct tinygemm_kernel_nn2<at::BFloat16, BLOCK_M, BLOCK_N> {
  static inline void apply(
@@ -934,7 +918,6 @@ void tinygemm_kernel(
    scalar_t* __restrict__ C,
    scalar_t* __restrict__ Btmp,
    float* __restrict__ Ctmp,
-    const float* __restrict__ Bbias,
    const float* __restrict__ scale,
    int64_t M,
    int64_t N,
@@ -945,11 +928,6 @@ void tinygemm_kernel(
    bool brg,
    int64_t block_size_K,
    bool do_unpack) {
-  if (Bbias != nullptr) {
-    tinygemm_kernel<scalar_t, at::Float8_e4m3fn, float, true>(
-        A, B, C, Btmp, Ctmp, scale, Bbias, M, N, K, lda, ldb, ldc, brg, block_size_K, do_unpack);
-    return;
-  }
  tinygemm_kernel<scalar_t, at::Float8_e4m3fn, float, false>(
      A, B, C, Btmp, Ctmp, scale, nullptr, M, N, K, lda, ldb, ldc, brg, block_size_K, do_unpack);
 }
@@ -971,6 +949,7 @@ void tinygemm_kernel(
    bool brg) {
  tinygemm_kernel2<scalar_t>(A, B, C, Btmp, Ctmp, scale, M, N, K, lda, ldb, ldc, brg);
 }
+
 template <typename scalar_t>
 void tinygemm_kernel(
    const scalar_t* __restrict__ A,
@@ -978,7 +957,6 @@ void tinygemm_kernel(
    scalar_t* __restrict__ C,
    scalar_t* __restrict__ Btmp,
    float* __restrict__ Ctmp,
-    const float* __restrict__ Bbias,
    const uint8_t* __restrict__ scale,
    int64_t M,
    int64_t N,
@@ -989,68 +967,10 @@ void tinygemm_kernel(
    bool brg,
    int64_t block_size_K,
    bool do_unpack) {
-  if (Bbias != nullptr) {
-    tinygemm_kernel<scalar_t, uint8_t, uint8_t, true>(
-        A, B, C, Btmp, Ctmp, scale, Bbias, M, N, K, lda, ldb, ldc, brg, block_size_K, do_unpack);
-    return;
-  }
  tinygemm_kernel<scalar_t, uint8_t, uint8_t, false>(
      A, B, C, Btmp, Ctmp, scale, nullptr, M, N, K, lda, ldb, ldc, brg, block_size_K, do_unpack);
 }

-// tinygemm interface
-template <typename scalar_t>
-void tinygemm_kernel(
-    const scalar_t* __restrict__ A,
-    const at::Float8_e4m3fn* __restrict__ B,
-    float* __restrict__ C,
-    scalar_t* __restrict__ Btmp,
-    const float* __restrict__ Bbias,
-    const float* __restrict__ scale,
-    int64_t M,
-    int64_t N,
-    int64_t K,
-    int64_t lda,
-    int64_t ldb,
-    int64_t ldc,
-    bool brg,
-    int64_t block_size_K,
-    bool do_unpack) {
-  if (Bbias != nullptr) {
-    tinygemm_kernel<scalar_t, at::Float8_e4m3fn, float, true>(
-        A, B, C, Btmp, scale, Bbias, M, N, K, lda, ldb, ldc, brg, block_size_K, do_unpack);
-    return;
-  }
-  tinygemm_kernel<scalar_t, at::Float8_e4m3fn, float, false>(
-      A, B, C, Btmp, scale, nullptr, M, N, K, lda, ldb, ldc, brg, block_size_K, do_unpack);
-}
-
-template <typename scalar_t>
-void tinygemm_kernel(
-    const scalar_t* __restrict__ A,
-    const uint8_t* __restrict__ B,
-    float* __restrict__ C,
-    scalar_t* __restrict__ Btmp,
-    const float* __restrict__ Bbias,
-    const uint8_t* __restrict__ scale,
-    int64_t M,
-    int64_t N,
-    int64_t K,
-    int64_t lda,
-    int64_t ldb,
-    int64_t ldc,
-    bool brg,
-    int64_t block_size_K,
-    bool do_unpack) {
-  if (Bbias != nullptr) {
-    tinygemm_kernel<scalar_t, uint8_t, uint8_t, true>(
-        A, B, C, Btmp, scale, Bbias, M, N, K, lda, ldb, ldc, brg, block_size_K, do_unpack);
-    return;
-  }
-  tinygemm_kernel<scalar_t, uint8_t, uint8_t, false>(
-      A, B, C, Btmp, scale, nullptr, M, N, K, lda, ldb, ldc, brg, block_size_K, do_unpack);
-}
-
 #define INSTANTIATE_TINYGEMM_TEMPLATE(TYPE_A, TYPE_B, TYPE_S) \
  template void tinygemm_kernel<TYPE_A>(                      \
      const TYPE_A* __restrict__ A,                           \
@@ -1058,7 +978,6 @@ void tinygemm_kernel(
      TYPE_A* __restrict__ C,                                 \
      TYPE_A* __restrict__ Btmp,                              \
      float* __restrict__ Ctmp,                               \
-      const float* __restrict__ Bbias,                        \
      const TYPE_S* __restrict__ scale,                       \
      int64_t M,                                              \
      int64_t N,                                              \
@@ -1101,6 +1020,7 @@ inline const float* get_bias_data(const std::optional<at::Tensor>& bias, int64_t
  }
  return nullptr;
 }
+
 // FP8 and MXFP4 WoQ uses the same pattern:
 //   Btmp : [T, BLOCK_N * K]
 //   Ctmp : [T, BLOCK_M * BLOCK_N]
@@ -812,7 +812,6 @@ void shared_expert_kernel_impl(
 static inline void check_moe_scales(
    bool use_int8_w8a8,
    bool use_fp8_w8a16,
-    bool use_mxfp4,
    const std::optional<at::Tensor>& w1_scale,
    const std::optional<at::Tensor>& w2_scale,
    const std::optional<std::vector<int64_t>> block_size) {
@@ -826,12 +825,6 @@ static inline void check_moe_scales(
    TORCH_CHECK(block_size.has_value(), "missing block_size for fp8 w8a16.");
    TORCH_CHECK(block_size.value().size() == 2, "expect block_size.size() to be 2.");
  }
-  if (use_mxfp4) {
-    TORCH_CHECK(w1_scale.has_value(), "missing w1_scale for mxfp4.");
-    TORCH_CHECK(w2_scale.has_value(), "missing w2_scale for mxfp4.");
-    TORCH_CHECK(w1_scale.value().scalar_type() == at::kByte, "expect w1_scale to be uint8.");
-    TORCH_CHECK(w2_scale.value().scalar_type() == at::kByte, "expect w2_scale to be uint8.");
-  }
 }

 #define CHECK_MOE_SCALES_FP8(DIM0, DIM1)                      \
@@ -846,8 +839,8 @@ static inline void check_moe_scales(
  TORCH_CHECK(w2s.size(DIM1) == div_up(N, block_size_K))

 // hidden_states: [M, K]
-// w1: [E, 2N, K] or [E, 2N, K / 2] for uint8
-// w2: [E, K, N] or [E, K, N / 2] for uint8
+// w1: [E, 2N, K]
+// w2: [E, K, N]
 // topk_weights: [M, topk]
 // topk_ids: [M, topk] (int32_t)
 //
@@ -865,10 +858,6 @@ at::Tensor fused_experts_cpu(
    const std::optional<at::Tensor>& w1_zero,
    const std::optional<at::Tensor>& w2_zero,
    const std::optional<std::vector<int64_t>> block_size,
-    const std::optional<at::Tensor>& w1_bias,
-    const std::optional<at::Tensor>& w2_bias,
-    const std::optional<double>& alpha,
-    const std::optional<double>& limit,
    bool is_vnni) {
  auto packed_w1 = is_vnni ? w1 : convert_weight_packed(w1);
  auto packed_w2 = is_vnni ? w2 : convert_weight_packed(w2);
@@ -908,12 +897,8 @@ at::Tensor fused_experts_cpu(
  int64_t topk = topk_weights_.size(1);

  // we use int32_t compensation for int8 w8a8
-  int64_t packed_K = moe_comp_method == CPUQuantMethod::MXFP4
-                         ? get_row_size<uint8_t>(K)
-                         : get_row_size(K, moe_comp_method == CPUQuantMethod::INT8_W8A8);
-  int64_t packed_N = moe_comp_method == CPUQuantMethod::MXFP4
-                         ? get_row_size<uint8_t>(N)
-                         : get_row_size(N, moe_comp_method == CPUQuantMethod::INT8_W8A8);
+  int64_t packed_K = get_row_size(K, moe_comp_method == CPUQuantMethod::INT8_W8A8);
+  int64_t packed_N = get_row_size(N, moe_comp_method == CPUQuantMethod::INT8_W8A8);

  // check weight shapes
  CHECK_EQ(w2.size(0), E);
@@ -926,7 +911,6 @@ at::Tensor fused_experts_cpu(
  check_moe_scales(
      moe_comp_method == CPUQuantMethod::INT8_W8A8,
      moe_comp_method == CPUQuantMethod::FP8_W8A16,
-      moe_comp_method == CPUQuantMethod::MXFP4,
      w1_scale,
      w2_scale,
      block_size);
@@ -981,7 +965,7 @@ at::Tensor fused_experts_cpu(
  //   5. Aq_tmp : [M, K] or [M * topk, N]
  //   6. As_tmp : [M * topk]
  //
-  // for fp8 w8a16 and mxfp4:
+  // for fp8 w8a16:
  //   7. intermediate_cache0 : [M * topk, 2N]
  //   8. B_tmp : [T, MAX_CACHE_BLOCK_SIZE, BLOCK_N, std::max(K, N)]
  //
@@ -994,7 +978,7 @@ at::Tensor fused_experts_cpu(
  if (moe_comp_method == CPUQuantMethod::INT8_W8A8) {
    buffer_size_nbytes += std::max(M * K, M * topk * N) + M * topk * sizeof(float);
  }
-  if (moe_comp_method == CPUQuantMethod::FP8_W8A16 || moe_comp_method == CPUQuantMethod::MXFP4) {
+  if (moe_comp_method == CPUQuantMethod::FP8_W8A16) {
    buffer_size_nbytes += M * topk * 2 * N * 2 + num_threads * MAX_CACHE_BLOCK_SIZE * BLOCK_N * std::max(K, N) * 2;
  }
  if (moe_comp_method == CPUQuantMethod::INT4_W4A8) {
@@ -1047,11 +1031,9 @@ at::Tensor fused_experts_cpu(
      float* __restrict__ C_tmp = (float*)((void*)(A_tmp + num_threads * BLOCK_M * K));
      scalar_t* __restrict__ intermediate_cache0 = (scalar_t*)((void*)(C_tmp + num_threads * 2 * BLOCK_M * BLOCK_N));
      scalar_t* __restrict__ B_tmp = (scalar_t*)((void*)(intermediate_cache0 + M * topk * 2 * N));
-      bool with_bias = w1_bias.has_value();
-      auto act_func = alpha.has_value() && limit.has_value() ? CPUAcTMethod::swiglu : CPUAcTMethod::silu_and_mul;

      CHECK_MOE_SCALES_FP8(1, 2);
-      fused_experts_fp_kernel_impl<scalar_t, at::Float8_e4m3fn, float, false>(
+      fused_experts_fp8_kernel_impl(
          out_hidden_states.data_ptr<scalar_t>(),
          intermediate_cache0,
          intermediate_cache1,
@@ -1062,8 +1044,6 @@ at::Tensor fused_experts_cpu(
          hidden_states.data_ptr<scalar_t>(),
          packed_w1.data_ptr<at::Float8_e4m3fn>(),
          packed_w2.data_ptr<at::Float8_e4m3fn>(),
-          with_bias ? w1_bias.value().data_ptr<float>() : nullptr,
-          with_bias ? w2_bias.value().data_ptr<float>() : nullptr,
          w1s.data_ptr<float>(),
          w2s.data_ptr<float>(),
          block_size_N,
@@ -1077,56 +1057,7 @@ at::Tensor fused_experts_cpu(
          K,
          E,
          topk,
-          num_tokens_post_pad,
-          alpha.has_value() ? float(alpha.value()) : 0,
-          limit.has_value() ? float(limit.value()) : 0,
-          act_func,
-          with_bias);
-    } else if (moe_comp_method == CPUQuantMethod::MXFP4) {
-      scalar_t* __restrict__ A_tmp = (scalar_t*)((void*)(intermediate_cache2 + M * topk * K));
-      float* __restrict__ C_tmp = (float*)((void*)(A_tmp + num_threads * BLOCK_M * K));
-      scalar_t* __restrict__ intermediate_cache0 = (scalar_t*)((void*)(C_tmp + num_threads * 2 * BLOCK_M * BLOCK_N));
-      scalar_t* __restrict__ B_tmp = (scalar_t*)((void*)(intermediate_cache0 + M * topk * 2 * N));
-      bool with_bias = w1_bias.has_value();
-      auto act_func = alpha.has_value() && limit.has_value() ? CPUAcTMethod::swiglu : CPUAcTMethod::silu_and_mul;
-
-      // mxfp4 supports only group size of 32 (2^5)
-      constexpr int64_t group_size = 32;
-      auto w1s = w1_scale.value();
-      auto w2s = w2_scale.value();
-      TORCH_CHECK(w1s.numel() == E * 2 * N * K / group_size, "w1_scale size mismatch");
-      TORCH_CHECK(w2s.numel() == E * K * N / group_size, "w2_scale size mismatch");
-      fused_experts_fp_kernel_impl<scalar_t, uint8_t, uint8_t, true>(
-          out_hidden_states.data_ptr<scalar_t>(),
-          intermediate_cache0,
-          intermediate_cache1,
-          intermediate_cache2,
-          A_tmp,
-          B_tmp,
-          C_tmp,
-          hidden_states.data_ptr<scalar_t>(),
-          packed_w1.data_ptr<uint8_t>(),
-          packed_w2.data_ptr<uint8_t>(),
-          with_bias ? w1_bias.value().data_ptr<float>() : nullptr,
-          with_bias ? w2_bias.value().data_ptr<float>() : nullptr,
-          w1s.data_ptr<uint8_t>(),
-          w2s.data_ptr<uint8_t>(),
-          /*block_size_N*/ 1,
-          /*block_size_K*/ group_size,
-          topk_weights_.data_ptr<float>(),
-          sorted_ids,
-          expert_ids,
-          offsets,
-          M,
-          N,
-          K,
-          E,
-          topk,
-          num_tokens_post_pad,
-          alpha.has_value() ? float(alpha.value()) : 0,
-          limit.has_value() ? float(limit.value()) : 0,
-          act_func,
-          with_bias);
+          num_tokens_post_pad);
    } else if (moe_comp_method == CPUQuantMethod::INT4_W4A8) {
      uint8_t* __restrict__ A_tmp = (uint8_t*)((void*)(intermediate_cache2 + M * topk * K));
      float* __restrict__ C_tmp = (float*)((void*)(A_tmp + num_threads * BLOCK_M * K));
@@ -1254,7 +1185,7 @@ at::Tensor shared_expert_cpu(
  CHECK_EQ(packed_w2.size(1), packed_N);

  // check scales
-  check_moe_scales(use_int8_w8a8, use_fp8_w8a16, false, w1_scale, w2_scale, block_size);
+  check_moe_scales(use_int8_w8a8, use_fp8_w8a16, w1_scale, w2_scale, block_size);

  at::Tensor out_hidden_states = inplace ? hidden_states : at::empty_like(hidden_states);

@@ -176,109 +176,3 @@ inline void silu_and_mul_stub(
    out_vec.store(out + d);
  }
 }
-
-template <typename scalar_t>
-inline void copy_mul_stub(scalar_t* __restrict__ out, const float* __restrict__ input, float weight, int64_t size) {
-  using bVec = at::vec::Vectorized<scalar_t>;
-  using fVec = at::vec::Vectorized<float>;
-  constexpr int kVecSize = bVec::size();
-  const fVec weight_vec = fVec(weight);
-  int64_t d;
-#pragma GCC unroll 4
-  for (d = 0; d <= size - kVecSize; d += kVecSize) {
-    fVec data0 = fVec::loadu(input + d) * weight_vec;
-    fVec data1 = fVec::loadu(input + d + fVec::size()) * weight_vec;
-    bVec out_vec = convert_from_float_ext<scalar_t>(data0, data1);
-    out_vec.store(out + d);
-  }
-  for (; d < size; ++d) {
-    out[d] = static_cast<scalar_t>(input[d] * weight);
-  }
-}
-
-// input = input + input2
-inline void add_bias_stub(float* __restrict__ input, const float* __restrict__ input2, int64_t size) {
-  using fVec = at::vec::Vectorized<float>;
-  constexpr int kVecSize = fVec::size();
-  int64_t d;
-#pragma GCC unroll 4
-  for (d = 0; d <= size - kVecSize; d += kVecSize) {
-    fVec x_fvec = fVec::loadu(input + d);
-    fVec y_fvec = fVec::loadu(input2 + d);
-    x_fvec = x_fvec + y_fvec;
-    x_fvec.store(input + d);
-  }
-  for (; d < size; ++d) {
-    input[d] = input[d] + input2[d];
-  }
-}
-
-template <typename scalar_t>
-inline void copy_mul_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, float weight, int64_t size) {
-  using bVec = at::vec::Vectorized<scalar_t>;
-  using fVec = at::vec::Vectorized<float>;
-  constexpr int kVecSize = bVec::size();
-  const fVec weight_vec = fVec(weight);
-  int64_t d;
-#pragma GCC unroll 4
-  for (d = 0; d <= size - kVecSize; d += kVecSize) {
-    bVec x = bVec::loadu(input + d);
-    fVec x0, x1;
-    std::tie(x0, x1) = at::vec::convert_to_float(x);
-    x0 = x0 * weight_vec;
-    x1 = x1 * weight_vec;
-    bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1);
-    out_vec.store(out + d);
-  }
-  for (; d < size; ++d) {
-    out[d] = static_cast<scalar_t>(input[d] * weight);
-  }
-}
-
-template <typename scalar_t>
-inline void clamp_sigmoid_and_mul_stub(
-    scalar_t* __restrict__ out,
-    const scalar_t* __restrict__ input,
-    int64_t size,
-    const float alpha,
-    const float limit) {
-  using bVec = at::vec::Vectorized<scalar_t>;
-  using fVec = at::vec::Vectorized<float>;
-  const fVec one = fVec(1.f);
-  const fVec zero = fVec(0.f);
-  const fVec limit_v = fVec(limit);
-  const fVec nlimit_v = fVec(-limit);
-  const fVec alpha_v = fVec(alpha);
-
-  // no remainder
-#pragma GCC unroll 4
-  for (int64_t d = 0; d < size; d += bVec::size()) {
-    bVec x = bVec::loadu(input + d);
-    fVec x0_, y0_;
-    std::tie(x0_, y0_) = at::vec::convert_to_float(x);
-    float tmp_buffer[fVec::size() * 2];  // 32
-    float tmp_glu[fVec::size()];         // 16
-    float tmp_linear[fVec::size()];      // 16
-    x0_.store(tmp_buffer);
-    y0_.store(tmp_buffer + fVec::size());
-    // interleaved: x[2i] = glu, x[2i+1] = linear
-    for (int j = 0; j < fVec::size(); ++j) {
-      // x0 [0,2,..30]
-      tmp_glu[j] = tmp_buffer[j * 2];
-      // y0 [1,3,...31]
-      tmp_linear[j] = tmp_buffer[j * 2 + 1];
-    }
-    fVec x0 = fVec::loadu(tmp_glu);
-    fVec y0 = fVec::loadu(tmp_linear);
-
-    // clamp
-    x0 = at::vec::minimum(x0, limit_v);
-    y0 = at::vec::minimum(limit_v, at::vec::maximum(nlimit_v, y0));
-    // x * sigmoid(x * alpha)
-    x0 = x0 / (one + (x0 * alpha_v).neg().exp_u20());
-    // (y + 1) * x
-    y0 = y0 + one;
-    x0 = x0 * y0;
-    convert_from_float_and_store<scalar_t>(out + d / 2, x0);
-  }
-}
@@ -7,8 +7,8 @@
 #include "gemm.h"
 #include "moe.h"

-template <typename scalar_t, typename packed_t, typename param_t, bool is_mxfp4>
-void fused_experts_fp_kernel_impl(
+template <typename scalar_t>
+void fused_experts_fp8_kernel_impl(
    scalar_t* __restrict__ output,
    scalar_t* __restrict__ ic0,
    scalar_t* __restrict__ ic1,
@@ -17,12 +17,10 @@ void fused_experts_fp_kernel_impl(
    scalar_t* __restrict__ B_tmp,
    float* __restrict__ C_tmp,
    const scalar_t* __restrict__ input,
-    const packed_t* __restrict__ packed_w1,
-    const packed_t* __restrict__ packed_w2,
-    const float* __restrict__ w1_bias,
-    const float* __restrict__ w2_bias,
-    const param_t* __restrict__ w1s,
-    const param_t* __restrict__ w2s,
+    const at::Float8_e4m3fn* __restrict__ packed_w1,
+    const at::Float8_e4m3fn* __restrict__ packed_w2,
+    const float* __restrict__ w1s,
+    const float* __restrict__ w2s,
    int64_t block_size_N,
    int64_t block_size_K,
    const float* __restrict__ topk_weights,
@@ -34,11 +32,7 @@ void fused_experts_fp_kernel_impl(
    int64_t K,
    int64_t E,
    int64_t topk,
-    int64_t num_tokens_post_pad,
-    float alpha,
-    float limit,
-    CPUAcTMethod act_func,
-    bool with_bias) {
+    int64_t num_tokens_post_pad) {
  constexpr int64_t BLOCK_M = block_size_m();
  constexpr int64_t BLOCK_N = block_size_n();

@@ -48,20 +42,12 @@ void fused_experts_fp_kernel_impl(
  int64_t scale_size_N = div_up(2 * N, block_size_N);
  int64_t scale_size_K = div_up(K, block_size_K);
  int64_t blocks_n_per_group = block_size_N / BLOCK_N;
-  std::function<int64_t(int64_t)> scale_offset_per_block;
-  if constexpr (is_mxfp4) {
-    scale_offset_per_block = [&](int64_t a) { return a * BLOCK_N; };
-  } else {
-    scale_offset_per_block = [&](int64_t a) { return a / blocks_n_per_group; };
-  }

-  const int64_t packed_K = get_row_size<packed_t>(K);
-
-  const int64_t stride_e = 2 * N * packed_K;
-  const int64_t stride_n = packed_K;
+  const int64_t stride_e = 2 * N * K;
+  const int64_t stride_n = K;

  int64_t avg_M = std::max(int64_t(1), M * topk / E);
-  const bool use_brgemm = can_use_brgemm<packed_t>(avg_M);
+  const bool use_brgemm = can_use_brgemm<at::Float8_e4m3fn>(avg_M);

  int64_t B_tmp_size_per_thread = MAX_CACHE_BLOCK_SIZE * BLOCK_N * std::max(K, N);

@@ -71,15 +57,14 @@ void fused_experts_fp_kernel_impl(
    int tid = get_thread_num();
    scalar_t* __restrict__ A = A_tmp + tid * BLOCK_M * K;

-    loop_2d<packed_t>(mb0, mb1, nb0, nb1, BLOCK_N * K, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
+    loop_2d<at::Float8_e4m3fn>(mb0, mb1, nb0, nb1, BLOCK_N * K, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
      int64_t n_size = std::min(2 * N - nb * BLOCK_N, BLOCK_N);

      // B shape [K, n_size] in vnni format
      int32_t expert_id = expert_ids[mb];
-      const packed_t* __restrict__ B = packed_w1 + expert_id * stride_e + nb * BLOCK_N * stride_n;
-      const param_t* __restrict__ Bs =
-          w1s + expert_id * scale_size_N * scale_size_K + scale_offset_per_block(nb) * scale_size_K;
-      const float* __restrict__ B_bias = with_bias ? w1_bias + expert_id * 2 * N + nb * BLOCK_N : nullptr;
+      const at::Float8_e4m3fn* __restrict__ B = packed_w1 + expert_id * stride_e + nb * BLOCK_N * stride_n;
+      const float* __restrict__ Bs =
+          w1s + expert_id * scale_size_N * scale_size_K + (nb / blocks_n_per_group) * scale_size_K;

      // do unpacking for the first row or a new expert
      int32_t pre_expert_id = mb == 0 ? -1 : expert_ids[mb - 1];
@@ -103,7 +88,6 @@ void fused_experts_fp_kernel_impl(
          /*   C            */ ic0 + offset * 2 * N + nb * BLOCK_N,
          /*   Btmp         */ B_tmp + tid * B_tmp_size_per_thread + nb_offset * BLOCK_N * K,
          /*   Ctmp         */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N,
-          /*   Bbias        */ B_bias,
          /*   scale        */ Bs,
          /*   M            */ m_size,
          /*   N            */ n_size,
@@ -122,20 +106,12 @@ void fused_experts_fp_kernel_impl(
  });

  // stage 1.5: intermediate_cache1 = silu(intermediate_cache0)
-  if (act_func == CPUAcTMethod::silu_and_mul) {
-    at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
-      for (int64_t m = begin; m < end; ++m) {
-        silu_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, ic0 + m * 2 * N + N, N);
-      }
-    });
-  } else if (act_func == CPUAcTMethod::swiglu) {
-    at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
-      for (int64_t m = begin; m < end; ++m) {
-        clamp_sigmoid_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, N, alpha, limit);
-        clamp_sigmoid_and_mul_stub(ic1 + m * N + N / 2, ic0 + m * 2 * N + N, N, alpha, limit);
-      }
-    });
-  }
+  at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
+    for (int64_t m = begin; m < end; ++m) {
+      silu_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, ic0 + m * 2 * N + N, N);
+    }
+  });
+
  // stage 2: intermediate_cache2 = intermediate_cache1 @ w2
  //   w2 : [E, K, N] as [E, OC, IC]
  const int64_t OC = K;  // rename K as OC
@@ -144,16 +120,15 @@ void fused_experts_fp_kernel_impl(
  const int64_t NB2 = div_up(OC, BLOCK_N);
  scale_size_N = div_up(K, block_size_N);
  scale_size_K = div_up(N, block_size_K);
-  const int64_t packed_IC = get_row_size<packed_t>(IC);
-  const int64_t stride_e2 = OC * packed_IC;
-  const int64_t stride_oc = packed_IC;
+  const int64_t stride_e2 = OC * IC;
+  const int64_t stride_oc = IC;

  // parallel on [MB2, NB2]
  parallel_2d(MB2, NB2, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
    int tid = get_thread_num();
    alignas(64) scalar_t C[BLOCK_M * BLOCK_K];

-    loop_2d<packed_t>(mb0, mb1, nb0, nb1, BLOCK_N * IC, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
+    loop_2d<at::Float8_e4m3fn>(mb0, mb1, nb0, nb1, BLOCK_N * IC, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
      int64_t m_size = offsets[mb + 1] - offsets[mb];
      int64_t n_size = std::min(OC - nb * BLOCK_N, BLOCK_N);

@@ -164,10 +139,9 @@ void fused_experts_fp_kernel_impl(

      // B shape [IC, n_size] in vnni format
      int32_t expert_id = expert_ids[mb];
-      const packed_t* __restrict__ B = packed_w2 + expert_id * stride_e2 + nb * BLOCK_N * stride_oc;
-      const param_t* __restrict__ Bs =
-          w2s + expert_id * scale_size_N * scale_size_K + scale_offset_per_block(nb) * scale_size_K;
-      const float* __restrict__ B_bias = with_bias ? w2_bias + expert_id * OC + nb * BLOCK_N : nullptr;
+      const at::Float8_e4m3fn* __restrict__ B = packed_w2 + expert_id * stride_e2 + nb * BLOCK_N * stride_oc;
+      const float* __restrict__ Bs =
+          w2s + expert_id * scale_size_N * scale_size_K + (nb / blocks_n_per_group) * scale_size_K;

      // do unpacking for the first row or a new expert
      int32_t pre_expert_id = mb == 0 ? -1 : expert_ids[mb - 1];
@@ -179,7 +153,6 @@ void fused_experts_fp_kernel_impl(
          /*   C            */ C,
          /*   Btmp         */ B_tmp + tid * B_tmp_size_per_thread + nb_offset * BLOCK_N * IC,
          /*   Ctmp         */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N,
-          /*   Bbias        */ B_bias,
          /*   scale        */ Bs,
          /*   M            */ m_size,
          /*   N            */ n_size,
@@ -204,6 +177,7 @@ void fused_experts_fp_kernel_impl(
      at::native::cpublas::brgemm_release();
    }
  });
+
  // stage 3: out = intermediate_cache2.sum(dim=1)
  //   from [M, topk, K] to [M, K]
  at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
@@ -213,43 +187,35 @@ void fused_experts_fp_kernel_impl(
  });
 }

-#define INSTANTIATE_MOE_FP_TEMPLATE(TYPE1, TYPE2, TYPE3, IS_MXFP4)           \
-  template void fused_experts_fp_kernel_impl<TYPE1, TYPE2, TYPE3, IS_MXFP4>( \
-      TYPE1* __restrict__ output,                                            \
-      TYPE1* __restrict__ ic0,                                               \
-      TYPE1* __restrict__ ic1,                                               \
-      TYPE1* __restrict__ ic2,                                               \
-      TYPE1* __restrict__ A_tmp,                                             \
-      TYPE1* __restrict__ B_tmp,                                             \
-      float* __restrict__ C_tmp,                                             \
-      const TYPE1* __restrict__ input,                                       \
-      const TYPE2* __restrict__ packed_w1,                                   \
-      const TYPE2* __restrict__ packed_w2,                                   \
-      const float* __restrict__ w1_bias,                                     \
-      const float* __restrict__ w2_bias,                                     \
-      const TYPE3* __restrict__ w1s,                                         \
-      const TYPE3* __restrict__ w2s,                                         \
-      int64_t block_size_N,                                                  \
-      int64_t block_size_K,                                                  \
-      const float* __restrict__ topk_weights,                                \
-      const int32_t* __restrict__ sorted_ids,                                \
-      const int32_t* __restrict__ expert_ids,                                \
-      const int32_t* __restrict__ offsets,                                   \
-      int64_t M,                                                             \
-      int64_t N,                                                             \
-      int64_t K,                                                             \
-      int64_t E,                                                             \
-      int64_t topk,                                                          \
-      int64_t num_tokens_post_pad,                                           \
-      float alpha,                                                           \
-      float limit,                                                           \
-      CPUAcTMethod act_func,                                                 \
-      bool with_bias)
+#define INSTANTIATE_MOE_FP8_TEMPLATE(TYPE)             \
+  template void fused_experts_fp8_kernel_impl<TYPE>(   \
+      TYPE* __restrict__ output,                       \
+      TYPE* __restrict__ ic0,                          \
+      TYPE* __restrict__ ic1,                          \
+      TYPE* __restrict__ ic2,                          \
+      TYPE* __restrict__ A_tmp,                        \
+      TYPE* __restrict__ B_tmp,                        \
+      float* __restrict__ C_tmp,                       \
+      const TYPE* __restrict__ input,                  \
+      const at::Float8_e4m3fn* __restrict__ packed_w1, \
+      const at::Float8_e4m3fn* __restrict__ packed_w2, \
+      const float* __restrict__ w1s,                   \
+      const float* __restrict__ w2s,                   \
+      int64_t block_size_N,                            \
+      int64_t block_size_K,                            \
+      const float* __restrict__ topk_weights,          \
+      const int32_t* __restrict__ sorted_ids,          \
+      const int32_t* __restrict__ expert_ids,          \
+      const int32_t* __restrict__ offsets,             \
+      int64_t M,                                       \
+      int64_t N,                                       \
+      int64_t K,                                       \
+      int64_t E,                                       \
+      int64_t topk,                                    \
+      int64_t num_tokens_post_pad)

-INSTANTIATE_MOE_FP_TEMPLATE(at::BFloat16, at::Float8_e4m3fn, float, false);
-INSTANTIATE_MOE_FP_TEMPLATE(at::Half, at::Float8_e4m3fn, float, false);
-INSTANTIATE_MOE_FP_TEMPLATE(at::BFloat16, uint8_t, uint8_t, true);
-INSTANTIATE_MOE_FP_TEMPLATE(at::Half, uint8_t, uint8_t, true);
+INSTANTIATE_MOE_FP8_TEMPLATE(at::BFloat16);
+INSTANTIATE_MOE_FP8_TEMPLATE(at::Half);

 template <typename scalar_t>
 void shared_expert_fp8_kernel_impl(
@@ -300,7 +266,6 @@ void shared_expert_fp8_kernel_impl(
          /*   C            */ ic0 + mb * BLOCK_M * 2 * N + nb * BLOCK_N,
          /*   Btmp         */ B_tmp + tid * B_tmp_size_per_thread + nb_offset * BLOCK_N * K,
          /*   Ctmp         */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N,
-          /*   Bbias        */ nullptr,
          /*   scale        */ w1s + (nb / blocks_n_per_group) * scale_size_K,
          /*   M            */ m_size,
          /*   N            */ n_size,
@@ -352,7 +317,6 @@ void shared_expert_fp8_kernel_impl(
          /*   C            */ C,
          /*   Btmp         */ B_tmp + tid * B_tmp_size_per_thread + nb_offset * BLOCK_N * IC,
          /*   Ctmp         */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N,
-          /*   Bbias        */ nullptr,
          /*   scale        */ w2s + (nb / blocks_n_per_group) * scale_size_K,
          /*   M            */ m_size,
          /*   N            */ n_size,
@@ -11,7 +11,7 @@

 #include <ATen/cpu/vec/functional.h>
 #include <ATen/cpu/vec/vec.h>
-#include <immintrin.h>
+
 namespace {

 using namespace at::vec;
@@ -21,15 +21,6 @@ inline Vectorized<scalar_t> convert_from_float_ext(const Vectorized<float>& a, c
  return at::vec::convert_from_float<scalar_t>(a, b);
 }

-template <typename scalar_t>
-inline void convert_from_float_and_store(scalar_t* out, const Vectorized<float>& a) {
-  float out_buffer[at::vec::Vectorized<float>::size()];
-  a.store(out_buffer);
-  for (int i = 0; i < 16; i++) {
-    out[i] = (scalar_t)out_buffer[i];
-  }
-}
-
 // allow f16, bf16
 template <typename scalar_t, typename std::enable_if_t<is_reduced_floating_point_v<scalar_t>, int> = 1>
 inline std::tuple<Vectorized<float>, Vectorized<float>> load_float_vec2(const scalar_t* __restrict__ data) {
@@ -59,11 +50,6 @@ convert_from_float_ext<at::BFloat16>(const Vectorized<float>& a, const Vectorize
  return (__m512i)(_mm512_cvtne2ps_pbh(__m512(b), __m512(a)));
 }

-template <>
-inline void convert_from_float_and_store<at::BFloat16>(at::BFloat16* out, const Vectorized<float>& a) {
-  _mm256_storeu_si256((__m256i*)out, (__m256i)(_mm512_cvtneps_pbh(__m512(a))));
-}
-
 #define CVT_BF16_TO_FP32(a) _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(a), 16))

 #define CVT_FP16_TO_FP32(a) _mm512_cvtph_ps(a)
@@ -139,6 +125,7 @@ inline __m512bh CVT_FP8_TO_BF16(__m256i a) {
  return cvt_e4m3_bf16_intrinsic_with_denorm(a);
 #endif
 }
+
 // faster version of float8_e4m3fn conversion to bfloat16
 //
 // we mapped cuda implementation from below link and vectorized with avx512:
@@ -64,8 +64,6 @@ at::Tensor weight_packed_linear(at::Tensor& mat1, at::Tensor& mat2,

 at::Tensor convert_weight_packed(at::Tensor& weight);

-at::Tensor convert_scale_packed(at::Tensor& scale);
-
 at::Tensor fused_experts_cpu(
    at::Tensor& hidden_states, at::Tensor& w1, at::Tensor& w2,
    at::Tensor& topk_weights, at::Tensor& topk_ids, bool inplace,
@@ -73,11 +71,7 @@ at::Tensor fused_experts_cpu(
    const std::optional<at::Tensor>& w2_scale,
    const std::optional<at::Tensor>& w1_zero,
    const std::optional<at::Tensor>& w2_zero,
-    const std::optional<std::vector<int64_t>> block_size,
-    const std::optional<at::Tensor>& w1_bias,
-    const std::optional<at::Tensor>& w2_bias,
-    const std::optional<double>& alpha, const std::optional<double>& limit,
-    bool is_vnni);
+    const std::optional<std::vector<int64_t>> block_size, bool is_vnni);

 at::Tensor int8_scaled_mm_with_quant(at::Tensor& mat1, at::Tensor& mat2,
                                     at::Tensor& scales2,
@@ -104,45 +98,6 @@ at::Tensor int4_scaled_mm_cpu(at::Tensor& x, at::Tensor& w, at::Tensor& w_zeros,
                              at::Tensor& w_scales,
                              std::optional<at::Tensor> bias);

-// Adapted from sglang: GDN
-std::tuple<at::Tensor, at::Tensor> chunk_gated_delta_rule_cpu(
-    const at::Tensor& query, const at::Tensor& key, const at::Tensor& value,
-    const at::Tensor& g, const at::Tensor& beta,
-    const at::Tensor& initial_state, bool output_final_state,
-    const at::Tensor& cu_seqlens, bool head_first, bool use_qk_l2norm_in_kernel,
-    double eps = 1e-5);
-
-at::Tensor fused_sigmoid_gating_delta_rule_update_cpu(
-    const at::Tensor& A_log, const at::Tensor& dt_bias, const at::Tensor& q,
-    const at::Tensor& k, const at::Tensor& v, const at::Tensor& a,
-    const at::Tensor& b, at::Tensor& initial_state_source,
-    const at::Tensor& initial_state_indices, const at::Tensor& cu_seqlens,
-    bool use_qk_l2norm_in_kernel, double softplus_beta = 1.0,
-    double softplus_threshold = 20.0);
-
-std::tuple<at::Tensor, at::Tensor> fused_gdn_gating_cpu(
-    const at::Tensor& A_log, const at::Tensor& a, const at::Tensor& b,
-    const at::Tensor& dt_bias);
-
-// Adapted from sglang: casual_conv1d kernels
-at::Tensor causal_conv1d_weight_pack(const at::Tensor& weight);
-
-at::Tensor causal_conv1d_fwd_cpu(
-    const at::Tensor& x, const at::Tensor& weight,
-    const std::optional<at::Tensor>& bias,
-    const std::optional<at::Tensor>& conv_states,
-    const std::optional<at::Tensor>& query_start_loc,
-    const std::optional<at::Tensor>& cache_indices,
-    const std::optional<at::Tensor>& has_initial_state, bool silu_activation,
-    int64_t pad_slot_id, bool is_vnni);
-
-at::Tensor causal_conv1d_update_cpu(
-    const at::Tensor& x, const at::Tensor& conv_states,
-    const at::Tensor& weight, const std::optional<at::Tensor>& bias,
-    bool silu_activation, const std::optional<at::Tensor>& cache_seqlens,
-    const std::optional<at::Tensor>& conv_state_indices, int64_t pad_slot_id,
-    bool is_vnni);
-
 void activation_lut_bf16(torch::Tensor& out, torch::Tensor& input,
                         const std::string& activation);

@@ -410,15 +365,11 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  ops.impl("weight_packed_linear", torch::kCPU, &weight_packed_linear);
  ops.def("convert_weight_packed(Tensor! weight) -> Tensor");
  ops.impl("convert_weight_packed", torch::kCPU, &convert_weight_packed);
-  ops.def("convert_scale_packed(Tensor! scale) -> Tensor");
-  ops.impl("convert_scale_packed", torch::kCPU, &convert_scale_packed);
  ops.def(
      "fused_experts_cpu(Tensor hidden_states, Tensor w1, Tensor w2, Tensor "
      "topk_weights, Tensor topk_ids, bool "
      "inplace, int moe_comp_method, Tensor? w1_scale, Tensor? w2_scale, "
-      "Tensor? w1_zero, Tensor? w2_zero, int[]? block_size, "
-      "Tensor? w1_bias, Tensor? w2_bias, float? alpha, float? limit, "
-      "bool is_vnni) -> "
+      "Tensor? w1_zero, Tensor? w2_zero, int[]? block_size, bool is_vnni) -> "
      "Tensor");
  ops.impl("fused_experts_cpu", torch::kCPU, &fused_experts_cpu);
  ops.def(
@@ -446,47 +397,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
      "scales2, SymInt[] block_size, Tensor? bias, ScalarType out_dtype, "
      "bool is_vnni) -> Tensor");
  ops.impl("fp8_scaled_mm_cpu", torch::kCPU, &fp8_scaled_mm_cpu);
-
-  // Adapted from sglang: GDN kernels
-  ops.def(
-      "chunk_gated_delta_rule_cpu(Tensor query, Tensor key, Tensor value, "
-      "Tensor g, Tensor beta, "
-      "Tensor initial_state, bool output_final_state, Tensor cu_seqlens, bool "
-      "head_first, "
-      "bool use_qk_l2norm_in_kernel, float eps=1e-5) -> (Tensor, Tensor)");
-  ops.impl("chunk_gated_delta_rule_cpu", torch::kCPU,
-           &chunk_gated_delta_rule_cpu);
-  ops.def(
-      "fused_sigmoid_gating_delta_rule_update_cpu(Tensor A_log, Tensor "
-      "dt_bias, Tensor q, Tensor k, Tensor v, Tensor "
-      "a, Tensor b, Tensor(a!) initial_state_source, Tensor "
-      "initial_state_indices, Tensor cu_seqlens, bool "
-      "use_qk_l2norm_in_kernel, float softplus_beta=1.0, float "
-      "softplus_threshold=20.0) -> Tensor");
-  ops.impl("fused_sigmoid_gating_delta_rule_update_cpu", torch::kCPU,
-           &fused_sigmoid_gating_delta_rule_update_cpu);
-  ops.def(
-      "fused_gdn_gating_cpu(Tensor A_log, Tensor a, Tensor b, Tensor dt_bias) "
-      "-> (Tensor, Tensor)");
-  ops.impl("fused_gdn_gating_cpu", torch::kCPU, &fused_gdn_gating_cpu);
-
-  // Adapted from sglang: casual_conv1d kernels
-  ops.def("causal_conv1d_weight_pack(Tensor weight) -> Tensor");
-  ops.impl("causal_conv1d_weight_pack", torch::kCPU,
-           &causal_conv1d_weight_pack);
-  ops.def(
-      "causal_conv1d_fwd_cpu(Tensor x, Tensor weight, Tensor? bias, Tensor? "
-      "conv_states, Tensor? query_start_loc,"
-      "Tensor? cache_indices, Tensor? has_initial_state, bool silu_activation, "
-      "int pad_slot_id, bool is_vnni) -> "
-      "Tensor");
-  ops.impl("causal_conv1d_fwd_cpu", torch::kCPU, &causal_conv1d_fwd_cpu);
-  ops.def(
-      "causal_conv1d_update_cpu(Tensor x, Tensor(a!) conv_states, Tensor "
-      "weight, Tensor? bias, bool silu_activation,"
-      "Tensor? cache_seqlens, Tensor? conv_state_indices, int pad_slot_id, "
-      "bool is_vnni) -> Tensor");
-  ops.impl("causal_conv1d_update_cpu", torch::kCPU, &causal_conv1d_update_cpu);
 #endif

  // CPU attention kernels
@@ -20,15 +20,13 @@
 * limitations under the License.
 */

-#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 <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
 #include <cuda_bf16.h>
 #include <cuda_runtime.h>
+#include <torch/all.h>
+
+#include "core/registration.h"

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

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

@@ -702,40 +700,37 @@ 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::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);
+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);
  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;
-  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(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(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.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.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(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(getSMVersion() >= 90, "required CUDA ARCH >= SM_90");
+  TORCH_CHECK(getSMVersion() >= 90, "required CUDA ARCH >= SM_90");

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

-STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
-  m.impl("dsv3_fused_a_gemm", TORCH_BOX(&dsv3_fused_a_gemm));
+TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
+  m.impl("dsv3_fused_a_gemm", &dsv3_fused_a_gemm);
 }
@@ -143,26 +143,4 @@ 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);
@@ -202,7 +202,7 @@ struct cutlass_3x_gemm_sm120 {
              sizeof(typename CollectiveEpilogue::SharedStorage))>,
          KernelSchedule>::CollectiveOp;

-  using GemmKernel = enable_sm120_family<cutlass::gemm::kernel::GemmUniversal<
+  using GemmKernel = enable_sm120_only<cutlass::gemm::kernel::GemmUniversal<
      Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue, void>>;
 };

@@ -1,6 +1,5 @@
 #include "scaled_mm_kernels.hpp"
 #include "scaled_mm_sm100_fp8_dispatch.cuh"
-#include "core/batch_invariant.hpp"

 namespace vllm {

@@ -14,17 +13,9 @@ void cutlass_scaled_mm_sm100_fp8(
    STD_TORCH_CHECK(bias->scalar_type() == out.scalar_type(),
                    "currently bias dtype must match output dtype ",
                    out.scalar_type());
-    if (vllm_is_batch_invariant()) {
-      return cutlass_scaled_mm_sm100_fp8_batch_invariant_epilogue<true>(
-          out, a, b, a_scales, b_scales, *bias);
-    }
    return cutlass_scaled_mm_sm100_fp8_epilogue<true>(out, a, b, a_scales,
                                                      b_scales, *bias);
  } else {
-    if (vllm_is_batch_invariant()) {
-      return cutlass_scaled_mm_sm100_fp8_batch_invariant_epilogue<false>(
-          out, a, b, a_scales, b_scales);
-    }
    return cutlass_scaled_mm_sm100_fp8_epilogue<false>(out, a, b, a_scales,
                                                       b_scales);
  }
@@ -294,34 +294,6 @@ inline void cutlass_gemm_sm100_fp8_dispatch(
  }
 }

-template <typename InType, typename OutType, bool EnableBias,
-          typename... EpilogueArgs>
-inline void cutlass_gemm_sm100_fp8_batch_invariant_dispatch(
-    torch::stable::Tensor& out, torch::stable::Tensor const& a,
-    torch::stable::Tensor const& b, torch::stable::Tensor const& a_scales,
-    torch::stable::Tensor const& b_scales, EpilogueArgs&&... args) {
-  static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
-  STD_TORCH_CHECK(a.scalar_type() ==
-                  torch::headeronly::ScalarType::Float8_e4m3fn);
-  STD_TORCH_CHECK(b.scalar_type() ==
-                  torch::headeronly::ScalarType::Float8_e4m3fn);
-
-  using Cutlass3xGemmM64SwapAB =
-      typename sm100_fp8_config_M64_swap_ab<InType, OutType,
-                                            EnableBias>::Cutlass3xGemm;
-  using Cutlass3xGemmM64 =
-      typename sm100_fp8_config_M64<InType, OutType, EnableBias>::Cutlass3xGemm;
-
-  // keep the CUTLASS config independent of M for batch invariance
-  uint32_t const k = a.size(1);
-  if (k < 4096) {
-    return cutlass_gemm_caller_sm100_fp8<Cutlass3xGemmM64>(
-        out, a, b, a_scales, b_scales, std::forward<EpilogueArgs>(args)...);
-  }
-  return cutlass_gemm_caller_sm100_fp8<Cutlass3xGemmM64SwapAB>(
-      out, a, b, b_scales, a_scales, std::forward<EpilogueArgs>(args)...);
-}
-
 template <bool EnableBias, typename... EpilogueArgs>
 void cutlass_scaled_mm_sm100_fp8_epilogue(torch::stable::Tensor& out,
                                          torch::stable::Tensor const& a,
@@ -348,28 +320,4 @@ void cutlass_scaled_mm_sm100_fp8_epilogue(torch::stable::Tensor& out,
  }
 }

-template <bool EnableBias, typename... EpilogueArgs>
-void cutlass_scaled_mm_sm100_fp8_batch_invariant_epilogue(
-    torch::stable::Tensor& out, torch::stable::Tensor const& a,
-    torch::stable::Tensor const& b, torch::stable::Tensor const& a_scales,
-    torch::stable::Tensor const& b_scales, EpilogueArgs&&... epilogue_args) {
-  STD_TORCH_CHECK(a.scalar_type() ==
-                  torch::headeronly::ScalarType::Float8_e4m3fn);
-  STD_TORCH_CHECK(b.scalar_type() ==
-                  torch::headeronly::ScalarType::Float8_e4m3fn);
-
-  if (out.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
-    return cutlass_gemm_sm100_fp8_batch_invariant_dispatch<
-        cutlass::float_e4m3_t, cutlass::bfloat16_t, EnableBias>(
-        out, a, b, a_scales, b_scales,
-        std::forward<EpilogueArgs>(epilogue_args)...);
-  } else {
-    STD_TORCH_CHECK(out.scalar_type() == torch::headeronly::ScalarType::Half);
-    return cutlass_gemm_sm100_fp8_batch_invariant_dispatch<
-        cutlass::float_e4m3_t, cutlass::half_t, EnableBias>(
-        out, a, b, a_scales, b_scales,
-        std::forward<EpilogueArgs>(epilogue_args)...);
-  }
-}
-
 }  // namespace vllm
@@ -1,6 +1,5 @@
 #include "scaled_mm_kernels.hpp"
 #include "scaled_mm_sm120_fp8_dispatch.cuh"
-#include "core/batch_invariant.hpp"
 #include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"

 namespace vllm {
@@ -15,17 +14,9 @@ void cutlass_scaled_mm_sm120_fp8(
    STD_TORCH_CHECK(bias->scalar_type() == out.scalar_type(),
                    "currently bias dtype must match output dtype ",
                    out.scalar_type());
-    if (vllm_is_batch_invariant()) {
-      return cutlass_scaled_mm_sm120_fp8_batch_invariant_epilogue<
-          c3x::ScaledEpilogueBias>(out, a, b, a_scales, b_scales, *bias);
-    }
    return cutlass_scaled_mm_sm120_fp8_epilogue<c3x::ScaledEpilogueBias>(
        out, a, b, a_scales, b_scales, *bias);
  } else {
-    if (vllm_is_batch_invariant()) {
-      return cutlass_scaled_mm_sm120_fp8_batch_invariant_epilogue<
-          c3x::ScaledEpilogue>(out, a, b, a_scales, b_scales);
-    }
    return cutlass_scaled_mm_sm120_fp8_epilogue<c3x::ScaledEpilogue>(
        out, a, b, a_scales, b_scales);
  }
@@ -72,7 +72,7 @@ struct cutlass_3x_gemm_sm120_custom {
              sizeof(typename CollectiveEpilogue::SharedStorage))>,
          KernelSchedule, void>::CollectiveOp;

-  using GemmKernel = enable_sm120_family<cutlass::gemm::kernel::GemmUniversal<
+  using GemmKernel = enable_sm120_only<cutlass::gemm::kernel::GemmUniversal<
      Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue, void>>;
 };

@@ -179,26 +179,6 @@ inline void cutlass_gemm_sm120_fp8_dispatch(torch::stable::Tensor& out,
      out, a, b, std::forward<EpilogueArgs>(args)...);
 }

-template <typename InType, typename OutType,
-          template <typename, typename, typename> typename Epilogue,
-          typename... EpilogueArgs>
-inline void cutlass_gemm_sm120_fp8_batch_invariant_dispatch(
-    torch::stable::Tensor& out, torch::stable::Tensor const& a,
-    torch::stable::Tensor const& b, EpilogueArgs&&... args) {
-  static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
-  STD_TORCH_CHECK(a.scalar_type() ==
-                  torch::headeronly::ScalarType::Float8_e4m3fn);
-  STD_TORCH_CHECK(b.scalar_type() ==
-                  torch::headeronly::ScalarType::Float8_e4m3fn);
-
-  using Cutlass3xGemmM64 =
-      typename sm120_fp8_config_M64<InType, OutType, Epilogue>::Cutlass3xGemm;
-
-  // keep the CUTLASS config independent of M for batch invariance
-  return cutlass_gemm_caller<Cutlass3xGemmM64>(
-      out, a, b, std::forward<EpilogueArgs>(args)...);
-}
-
 template <template <typename, typename, typename> typename Epilogue,
          typename... EpilogueArgs>
 void cutlass_scaled_mm_sm120_fp8_epilogue(torch::stable::Tensor& out,
@@ -222,26 +202,4 @@ void cutlass_scaled_mm_sm120_fp8_epilogue(torch::stable::Tensor& out,
  }
 }

-template <template <typename, typename, typename> typename Epilogue,
-          typename... EpilogueArgs>
-void cutlass_scaled_mm_sm120_fp8_batch_invariant_epilogue(
-    torch::stable::Tensor& out, torch::stable::Tensor const& a,
-    torch::stable::Tensor const& b, EpilogueArgs&&... epilogue_args) {
-  STD_TORCH_CHECK(a.scalar_type() ==
-                  torch::headeronly::ScalarType::Float8_e4m3fn);
-  STD_TORCH_CHECK(b.scalar_type() ==
-                  torch::headeronly::ScalarType::Float8_e4m3fn);
-
-  if (out.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
-    return cutlass_gemm_sm120_fp8_batch_invariant_dispatch<
-        cutlass::float_e4m3_t, cutlass::bfloat16_t, Epilogue>(
-        out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
-  } else {
-    STD_TORCH_CHECK(out.scalar_type() == torch::headeronly::ScalarType::Half);
-    return cutlass_gemm_sm120_fp8_batch_invariant_dispatch<
-        cutlass::float_e4m3_t, cutlass::half_t, Epilogue>(
-        out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
-  }
-}
-
 }  // namespace vllm
@@ -1,6 +1,5 @@
 #include "scaled_mm_kernels.hpp"
 #include "scaled_mm_sm90_fp8_dispatch.cuh"
-#include "core/batch_invariant.hpp"

 namespace vllm {

@@ -14,17 +13,9 @@ void cutlass_scaled_mm_sm90_fp8(
    STD_TORCH_CHECK(bias->scalar_type() == out.scalar_type(),
                    "currently bias dtype must match output dtype ",
                    out.scalar_type());
-    if (vllm_is_batch_invariant()) {
-      return cutlass_scaled_mm_sm90_fp8_batch_invariant_epilogue<true>(
-          out, a, b, a_scales, b_scales, *bias);
-    }
    return cutlass_scaled_mm_sm90_fp8_epilogue<true>(out, a, b, a_scales,
                                                     b_scales, *bias);
  } else {
-    if (vllm_is_batch_invariant()) {
-      return cutlass_scaled_mm_sm90_fp8_batch_invariant_epilogue<false>(
-          out, a, b, a_scales, b_scales);
-    }
    return cutlass_scaled_mm_sm90_fp8_epilogue<false>(out, a, b, a_scales,
                                                      b_scales);
  }
@@ -349,35 +349,6 @@ inline void cutlass_gemm_sm90_fp8_dispatch(
  }
 }

-template <typename InType, typename OutType, bool EnableBias,
-          typename... EpilogueArgs>
-inline void cutlass_gemm_sm90_fp8_batch_invariant_dispatch(
-    torch::stable::Tensor& out, torch::stable::Tensor const& a,
-    torch::stable::Tensor const& b, torch::stable::Tensor const& a_scales,
-    torch::stable::Tensor const& b_scales, EpilogueArgs&&... args) {
-  static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
-  STD_TORCH_CHECK(a.scalar_type() ==
-                  torch::headeronly::ScalarType::Float8_e4m3fn);
-  STD_TORCH_CHECK(b.scalar_type() ==
-                  torch::headeronly::ScalarType::Float8_e4m3fn);
-
-  using Cutlass3xGemmM64_N1280 =
-      typename sm90_fp8_config_M64_N1280<InType, OutType,
-                                         EnableBias>::Cutlass3xGemm;
-  using Cutlass3xGemmM64_N8192 =
-      typename sm90_fp8_config_M64_N8192<InType, OutType,
-                                         EnableBias>::Cutlass3xGemm;
-
-  // keep the CUTLASS config independent of M for batch invariance
-  uint32_t const n = b.size(1);
-  if (n <= 1280) {
-    return cutlass_gemm_caller_sm90_fp8<Cutlass3xGemmM64_N1280>(
-        out, a, b, b_scales, a_scales, std::forward<EpilogueArgs>(args)...);
-  }
-  return cutlass_gemm_caller_sm90_fp8<Cutlass3xGemmM64_N8192>(
-      out, a, b, b_scales, a_scales, std::forward<EpilogueArgs>(args)...);
-}
-
 template <bool EnableBias, typename... EpilogueArgs>
 void cutlass_scaled_mm_sm90_fp8_epilogue(torch::stable::Tensor& out,
                                         torch::stable::Tensor const& a,
@@ -404,28 +375,4 @@ void cutlass_scaled_mm_sm90_fp8_epilogue(torch::stable::Tensor& out,
  }
 }

-template <bool EnableBias, typename... EpilogueArgs>
-void cutlass_scaled_mm_sm90_fp8_batch_invariant_epilogue(
-    torch::stable::Tensor& out, torch::stable::Tensor const& a,
-    torch::stable::Tensor const& b, torch::stable::Tensor const& a_scales,
-    torch::stable::Tensor const& b_scales, EpilogueArgs&&... epilogue_args) {
-  STD_TORCH_CHECK(a.scalar_type() ==
-                  torch::headeronly::ScalarType::Float8_e4m3fn);
-  STD_TORCH_CHECK(b.scalar_type() ==
-                  torch::headeronly::ScalarType::Float8_e4m3fn);
-
-  if (out.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
-    return cutlass_gemm_sm90_fp8_batch_invariant_dispatch<
-        cutlass::float_e4m3_t, cutlass::bfloat16_t, EnableBias>(
-        out, a, b, a_scales, b_scales,
-        std::forward<EpilogueArgs>(epilogue_args)...);
-  } else {
-    STD_TORCH_CHECK(out.scalar_type() == torch::headeronly::ScalarType::Half);
-    return cutlass_gemm_sm90_fp8_batch_invariant_dispatch<
-        cutlass::float_e4m3_t, cutlass::half_t, EnableBias>(
-        out, a, b, a_scales, b_scales,
-        std::forward<EpilogueArgs>(epilogue_args)...);
-  }
-}
-
 }  // namespace vllm
@@ -9,7 +9,6 @@
 #include "scaled_mm_c2x_sm89_fp8_dispatch.cuh"
 #include "scaled_mm_c2x_sm89_int8_dispatch.cuh"

-#include "core/batch_invariant.hpp"
 #include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c2x.hpp"

 using namespace vllm;
@@ -192,17 +191,9 @@ void cutlass_scaled_mm_sm89(torch::stable::Tensor& out,
    STD_TORCH_CHECK(bias->scalar_type() == out.scalar_type(),
                    "currently bias dtype must match output dtype ",
                    out.scalar_type());
-    if (vllm_is_batch_invariant()) {
-      return cutlass_scaled_mm_sm89_fp8_batch_invariant_epilogue<
-          c2x::ScaledEpilogueBias>(out, a, b, a_scales, b_scales, *bias);
-    }
    return cutlass_scaled_mm_sm89_epilogue<c2x::ScaledEpilogueBias>(
        out, a, b, a_scales, b_scales, *bias);
  } else {
-    if (vllm_is_batch_invariant()) {
-      return cutlass_scaled_mm_sm89_fp8_batch_invariant_epilogue<
-          c2x::ScaledEpilogue>(out, a, b, a_scales, b_scales);
-    }
    return cutlass_scaled_mm_sm89_epilogue<c2x::ScaledEpilogue>(
        out, a, b, a_scales, b_scales);
  }
@@ -381,43 +381,4 @@ inline void cutlass_gemm_sm89_fp8_dispatch(torch::stable::Tensor& out,
  }
 }

-template <typename InType, typename OutType,
-          template <typename, typename> typename Epilogue,
-          typename... EpilogueArgs>
-inline void cutlass_gemm_sm89_fp8_batch_invariant_dispatch(
-    torch::stable::Tensor& out, torch::stable::Tensor const& a,
-    torch::stable::Tensor const& b, EpilogueArgs&&... args) {
-  static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
-  STD_TORCH_CHECK(a.scalar_type() ==
-                  torch::headeronly::ScalarType::Float8_e4m3fn);
-  STD_TORCH_CHECK(b.scalar_type() ==
-                  torch::headeronly::ScalarType::Float8_e4m3fn);
-
-  // keep the CUTLASS config independent of M for batch invariance
-  return sm89_fp8_config_M64::dispatch<InType, OutType, Epilogue>(
-      out, a, b, std::forward<EpilogueArgs>(args)...);
-}
-
-template <template <typename, typename> typename Epilogue,
-          typename... EpilogueArgs>
-void cutlass_scaled_mm_sm89_fp8_batch_invariant_epilogue(
-    torch::stable::Tensor& out, torch::stable::Tensor const& a,
-    torch::stable::Tensor const& b, EpilogueArgs&&... epilogue_args) {
-  STD_TORCH_CHECK(a.scalar_type() ==
-                  torch::headeronly::ScalarType::Float8_e4m3fn);
-  STD_TORCH_CHECK(b.scalar_type() ==
-                  torch::headeronly::ScalarType::Float8_e4m3fn);
-
-  if (out.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
-    return cutlass_gemm_sm89_fp8_batch_invariant_dispatch<
-        cutlass::float_e4m3_t, cutlass::bfloat16_t, Epilogue>(
-        out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
-  } else {
-    STD_TORCH_CHECK(out.scalar_type() == torch::headeronly::ScalarType::Half);
-    return cutlass_gemm_sm89_fp8_batch_invariant_dispatch<
-        cutlass::float_e4m3_t, cutlass::half_t, Epilogue>(
-        out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
-  }
-}
-
 }  // namespace vllm
@@ -156,17 +156,6 @@ 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,
@@ -258,11 +247,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, int kGroupsPerBlockX,
-          int kRowsPerBlock>
+template <typename T, typename DST_DTYPE, int GROUP_SIZE>
 __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 int padded_groups_per_row,
+    unsigned int* __restrict__ output_s_packed, const int64_t num_groups_padded,
+    const int groups_per_block, 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) {
@@ -271,25 +260,27 @@ __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 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) {
+  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) {
    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
@@ -452,53 +443,34 @@ 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;
-
-  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, ").");
+  const int groups_per_block = GetGroupsPerBlock(num_groups_padded);

  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_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); \
-    }                                                      \
+#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));                                \
  } while (0)

  VLLM_STABLE_DISPATCH_HALF_TYPES(
@@ -516,7 +488,6 @@ 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,54 +218,7 @@ 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) {
@@ -301,16 +254,6 @@ 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,71 +6,12 @@
 #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) {
@@ -79,10 +20,3 @@ 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);
-}
@@ -1,30 +0,0 @@
-/*
- * 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);
@@ -1,130 +0,0 @@
-/*
- * 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);
-}
@@ -1,249 +0,0 @@
-/*
- * 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
@@ -15,11 +15,11 @@ SUPPORT_SM80 = False
 for arch in sys.argv[1].split(","):
    arch = arch[: arch.index(".") + 2].replace(".", "")
    arch = int(arch)
-    # SM89 and the SM12x family (SM120 RTX 5090, SM121 DGX Spark GB10)
-    # fully support mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32.
+    # only SM89 and SM120 fully support
+    # mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32.
    # SM90 and SM100 can use this PTX, but it’s simulated
    # with FP16 MMA, so it cannot achieve any acceleration.
-    if arch == 89 or arch // 10 == 12:
+    if arch in [89, 120]:
        SUPPORT_FP8 = True
    if arch >= 80:
        SUPPORT_SM80 = True
@@ -448,8 +448,8 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
                "FP8 only support Ada Lovelace or newer GPUs.");
    TORCH_CHECK(
        major_capability * 10 + minor_capability == 89 ||
-            major_capability == 12,
-        "Marlin W4A8-FP8 only support SM89 or SM12x device (It is slower than "
+            major_capability * 10 + minor_capability == 120,
+        "Marlin W4A8-FP8 only support SM89 or SM120 device (It is slower than "
        "Marlin W4A16 on other devices).");
  }

@@ -390,13 +390,7 @@ __global__ void moe_lora_align_block_size_kernel(
    int32_t* __restrict__ token_mask, bool has_expert_map) {
  int lora_idx = blockIdx.x / 2;
  int lora_id = lora_ids[lora_idx];
-  // Output buffers are indexed by lora_id (in [0, max_loras)). The grid
-  // iterates one extra slot to accommodate the "-1" entry that
-  // active_lora_ids may hold in position 0 for mixed base + LoRA batches;
-  // guard against any other unexpected lora_id >= max_loras to avoid
-  // out-of-bounds writes. This mirrors the `lora_id >= max_loras` guard in
-  // the Triton _fused_moe_lora_kernel.
-  if (lora_id == -1 || lora_id >= max_loras || adapter_enabled[lora_id] == 0) {
+  if (lora_id == -1 || adapter_enabled[lora_id] == 0) {
    return;
  }

@@ -426,21 +420,10 @@ __global__ void lora_count_and_sort_expert_tokens_kernel(
    int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ cumsum_buffer,
    int32_t* __restrict__ expert_map, size_t numel, int32_t num_experts,
    int32_t max_num_tokens_padded, int32_t topk_num, int32_t* token_mask,
-    int32_t max_loras, int32_t* lora_ids, int32_t* adapter_enabled,
-    bool has_expert_map) {
+    int32_t* lora_ids, bool has_expert_map) {
  int lora_idx = blockIdx.x;
  int lora_id = lora_ids[lora_idx];
-  // Same guard rationale as moe_lora_align_block_size_kernel. Additionally
-  // skip disabled adapter slots: moe_lora_align_block_size_kernel early-returns
-  // for them and leaves token_mask[lora_id, :] uninitialized (token_mask is
-  // allocated with torch::empty), so running the sort loop here would traverse
-  // garbage mask bits and pollute this slot's rows of sorted_token_ids and
-  // cumsum_buffer. Downstream consumers already skip disabled slots, so the
-  // pollution is dormant today, but the check keeps behavior symmetric with
-  // the other two align kernels and avoids O(numel) wasted work per disabled
-  // slot. Short-circuit evaluation ensures adapter_enabled is only indexed
-  // after lora_id is confirmed to be in [0, max_loras).
-  if (lora_id == -1 || lora_id >= max_loras || adapter_enabled[lora_id] == 0) {
+  if (lora_id == -1) {
    return;
  }

@@ -463,8 +446,7 @@ __global__ void moe_lora_align_block_size_small_batch_expert_kernel(
    int32_t* token_mask, bool has_expert_map) {
  int lora_idx = blockIdx.x;
  int lora_id = lora_ids[lora_idx];
-  // Same guard rationale as moe_lora_align_block_size_kernel.
-  if (lora_id == -1 || lora_id >= max_loras || adapter_enabled[lora_id] == 0) {
+  if (lora_id == -1 || adapter_enabled[lora_id] == 0) {
    return;
  }

@@ -716,15 +698,7 @@ void moe_lora_align_block_size(
                  scalar_t, fill_threads>;
          AT_CUDA_CHECK(VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize(
              (void*)kernel, shared_mem));
-          // Grid size is (max_loras + 1) because active_lora_ids has length
-          // max_loras + 1: sorted-unique values of token_lora_mapping, which
-          // can include -1 (base-model tokens) in addition to up to max_loras
-          // real LoRA slots. Using max_loras would drop the real LoRA slot
-          // when -1 is present at position 0 and leave output buffers
-          // uninitialized, causing illegal memory accesses in downstream
-          // MoE-LoRA kernels. This mirrors the fix made for the Triton
-          // _fused_moe_lora_kernel grid in vllm-project/vllm#32277.
-          kernel<<<max_loras + 1, blockDim, shared_mem, stream>>>(
+          kernel<<<max_loras, blockDim, shared_mem, stream>>>(
              topk_ids.data_ptr<scalar_t>(),
              token_lora_mapping.data_ptr<int32_t>(), block_size,
              expert_map.data_ptr<int32_t>(), num_experts, max_loras,
@@ -748,17 +722,10 @@ void moe_lora_align_block_size(
          auto align_kernel =
              vllm::moe::moe_lora_align_block_size_kernel<scalar_t>;

-          // Launch two threadblocks per LoRA slot, across max_loras + 1 slots
-          // to cover the extra "-1" (base-model tokens) entry that
-          // active_lora_ids may contain in addition to up to max_loras real
-          // LoRA slots. Using max_loras would drop the real LoRA slot when -1
-          // occupies position 0 and leave the output buffers uninitialized,
-          // causing illegal memory accesses downstream. Mirrors the grid fix
-          // applied to _fused_moe_lora_kernel in vllm-project/vllm#32277.
+          // launch two threadblocks for each lora
          // blockIdx.x % 2 == 0: counting experts and aligning
          // blockIdx.x % 2 == 1: filling sorted_token_ids
-          align_kernel<<<(max_loras + 1) * 2, blockDim, shared_mem_size,
-                         stream>>>(
+          align_kernel<<<max_loras * 2, blockDim, shared_mem_size, stream>>>(
              topk_ids.data_ptr<scalar_t>(),
              token_lora_mapping.data_ptr<int32_t>(), block_size,
              expert_map.data_ptr<int32_t>(), num_experts, max_loras,
@@ -777,10 +744,7 @@ void moe_lora_align_block_size(
          const int max_blocks = 65535;
          const int actual_blocks = std::min(num_blocks, max_blocks);

-          // Same rationale as align_kernel above: iterate over max_loras + 1
-          // slots so the sort kernel processes the real LoRA slot even when
-          // active_lora_ids has -1 at position 0.
-          dim3 gridDims(max_loras + 1, actual_blocks);
+          dim3 gridDims(max_loras, actual_blocks);
          auto sort_kernel =
              vllm::moe::lora_count_and_sort_expert_tokens_kernel<scalar_t>;

@@ -789,8 +753,7 @@ void moe_lora_align_block_size(
              sorted_token_ids.data_ptr<int32_t>(), cumsum.data_ptr<int32_t>(),
              expert_map.data_ptr<int32_t>(), topk_ids.numel(), num_experts,
              max_num_tokens_padded, topk_num, token_mask.data_ptr<int32_t>(),
-              max_loras, lora_ids.data_ptr<int32_t>(),
-              adapter_enabled.data_ptr<int32_t>(), has_expert_map);
+              lora_ids.data_ptr<int32_t>(), has_expert_map);
        }
      });
 }
@@ -75,12 +75,4 @@ 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,131 +298,130 @@ __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];
-      }
-      row_chunk[ii] = val;
+  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;
+  }

-    // 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;
-        }
-      }
-
-      // 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;
-        }
+      // 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;
      }
    }

-    // Apply renormalization and routed scaling factor to final weights.
+    // Write the max for this k iteration to global memory.
    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;
+      // 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];
      }
-      for (int k_idx = 0; k_idx < k; ++k_idx) {
-        const int idx = k * thread_row + k_idx;
-        output[idx] = output[idx] * scale;
+      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;
      }
    }
-#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-    asm volatile("griddepcontrol.launch_dependents;");
-#endif
  }
+
+  // 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;");
+#endif
 }

 namespace detail {
@@ -92,6 +92,16 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
      "bool use_fp32_reduce, bool is_zp_float,"
      "int thread_k, int thread_n, int blocks_per_sm) -> Tensor");

+  m.def(
+      "marlin_gemm_moe(Tensor! a, Tensor! b_q_weights, Tensor! sorted_ids, "
+      "Tensor! topk_weights, Tensor! topk_ids, Tensor! b_scales, Tensor! "
+      "b_zeros, Tensor! g_idx, Tensor! perm, Tensor! workspace, "
+      "int b_q_type, SymInt size_m, "
+      "SymInt size_n, SymInt size_k, bool is_k_full, int num_experts, int "
+      "topk, "
+      "int moe_block_size, bool replicate_input, bool apply_weights)"
+      " -> Tensor");
+
  m.def(
      "moe_permute(Tensor input, Tensor topk_ids,"
      "Tensor token_expert_indices, Tensor? expert_map, int n_expert,"
@@ -125,12 +135,6 @@ 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
 }

@@ -60,6 +60,32 @@ void merge_attn_states(
    const torch::Tensor& suffix_output, const torch::Tensor& suffix_lse,
    const std::optional<int64_t> prefill_tokens_with_context,
    const std::optional<torch::Tensor>& output_scale = std::nullopt);
+#ifndef USE_ROCM
+void convert_vertical_slash_indexes(
+    torch::Tensor& block_count,      // [BATCH, N_HEADS, NUM_ROWS]
+    torch::Tensor& block_offset,     // [BATCH, N_HEADS, NUM_ROWS, NNZ_S]
+    torch::Tensor& column_count,     // [BATCH, N_HEADS, NUM_ROWS]
+    torch::Tensor& column_index,     // [BATCH, N_HEADS, NUM_ROWS, NNZ_V]
+    torch::Tensor q_seqlens,         // [BATCH, ]
+    torch::Tensor kv_seqlens,        // [BATCH, ]
+    torch::Tensor vertical_indexes,  // [BATCH, N_HEADS, NNZ_V]
+    torch::Tensor slash_indexes,     // [BATCH, N_HEADS, NNZ_S]
+    int64_t context_size, int64_t block_size_M, int64_t block_size_N,
+    bool causal);
+
+void convert_vertical_slash_indexes_mergehead(
+    torch::Tensor& block_count,            // [BATCH, N_HEADS, NUM_ROWS]
+    torch::Tensor& block_offset,           // [BATCH, N_HEADS, NUM_ROWS, NNZ_S]
+    torch::Tensor& column_count,           // [BATCH, N_HEADS, NUM_ROWS]
+    torch::Tensor& column_index,           // [BATCH, N_HEADS, NUM_ROWS, NNZ_V]
+    torch::Tensor q_seqlens,               // [BATCH, ]
+    torch::Tensor kv_seqlens,              // [BATCH, ]
+    torch::Tensor vertical_indexes,        // [BATCH, N_HEADS, NNZ_V]
+    torch::Tensor slash_indexes,           // [BATCH, N_HEADS, NNZ_S]
+    torch::Tensor vertical_indices_count,  // [N_HEADS, ]
+    torch::Tensor slash_indices_count, int64_t context_size,
+    int64_t block_size_M, int64_t block_size_N, bool causal);
+#endif

 void rms_norm(torch::Tensor& out, torch::Tensor& input, torch::Tensor& weight,
              double epsilon);
@@ -174,6 +200,19 @@ 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);
@@ -263,6 +302,8 @@ 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);
@@ -274,6 +315,11 @@ 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,
@@ -7,11 +7,10 @@ Shang and Dang, Xingyu and Han, Song}, journal={arXiv}, year={2023}
 }
 */

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

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

 #include <cuda_fp16.h>

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

-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) {
+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) {
  int in_c = _kernel.size(0);
  int qout_c = _kernel.size(1);
  int out_c = qout_c * 8;
@@ -439,24 +437,23 @@ torch::stable::Tensor awq_dequantize(torch::stable::Tensor _kernel,
    y_blocks = (int)(in_c / 8);
  }

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

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

-  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>());
+  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>());

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

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

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

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

-  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());
+  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);
  int num_out_feats = _out_feats.size(-2);
  int num_out_channels = _out_feats.size(-1);

-  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>());
+  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>());
  int group_size = num_in_channels / _scaling_factors.size(0);

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

-  const cudaStream_t stream = get_current_cuda_stream();
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
  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);
@@ -528,5 +522,5 @@ torch::stable::Tensor awq_gemm(torch::stable::Tensor _in_feats,
            group_size, split_k_iters, in_feats, kernel, scaling_factors, zeros,
            num_in_feats, num_in_channels, num_out_channels, out_feats);
  }
-  return torch::stable::sum(_out_feats, 0);
+  return _out_feats.sum(0);
 }
@@ -1,7 +1,7 @@
 #include <cstdint>

 /* Adapted from ./csrc/quantization/gguf/mmq.cuh
-   based on ./vllm/model_executor/layers/fused_moe/experts/triton_moe.py */
+   based on ./vllm/model_executor/layers/fused_moe/fused_moe.py */
 template <typename scalar_t, int qk, int qr, int qi, bool need_sum,
          typename block_q_t, int mmq_x, int mmq_y, int nwarps,
          allocate_tiles_cuda_t allocate_tiles, load_tiles_cuda_t load_tiles,
@@ -1,28 +1,20 @@
 #include "allspark_utils.cuh"
-
-#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 <torch/all.h>
+#include "core/registration.h"
 #include <cublas_v2.h>

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

 #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800

-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,
+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,
    int64_t CUBLAS_M_THRESHOLD, bool has_zp, bool n32k16_reorder) {
-  STD_TORCH_CHECK_NOT_IMPLEMENTED(
+  TORCH_CHECK_NOT_IMPLEMENTED(
      false, "allspark_w8a16_gemm(..) requires CUDA_ARCH >= 8.0");
-  return torch::stable::empty({1, 1});
+  return torch::empty({1, 1});
 }

 #else
@@ -856,8 +848,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) {
-  STD_TORCH_CHECK(N % 8 == 0 && K % 16 == 0 && N_32align % 32 == 0,
-                  "Unsupported shape");
+  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);
@@ -867,7 +859,7 @@ void restore_N32_K16_dequantize_rhs_w8a16(const QT* qdata, const FT* scales,
  }
  // TODO: Support SubChannel
  else {
-    STD_TORCH_CHECK(false, "Now only support PerChannel");
+    TORCH_CHECK(false, "Now only support PerChannel");
  }
 }

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

 }  // namespace allspark

-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,
+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,
    int64_t CUBLAS_M_THRESHOLD, bool has_zp, bool n32k16_reorder) {
  // Verify device and strides
-  STD_TORCH_CHECK(a.device().is_cuda(), "A is not on GPU");
-  STD_TORCH_CHECK(a.is_contiguous(), "A is not contiguous");
+  TORCH_CHECK(a.device().is_cuda(), "A is not on GPU");
+  TORCH_CHECK(a.is_contiguous(), "A 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_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_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_scales.device().is_cuda(), "b_scales is not on GPU");
+  TORCH_CHECK(b_scales.is_contiguous(), "b_scales is not contiguous");

  if (has_zp) {
-    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");
+    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");
  }

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

  // Verify shape
-  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(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(group_size == -1, "Currently only supports group_size = -1");
+  TORCH_CHECK(group_size == -1, "Currently only supports group_size = -1");

-  const torch::stable::accelerator::DeviceGuard device_guard(
-      a.get_device_index());
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(a));
  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());
@@ -971,12 +959,12 @@ torch::stable::Tensor allspark_w8a16_gemm(
    b_zero_ptr = reinterpret_cast<const void*>(b_qzeros.value().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());
+  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());

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

  allspark::BlockTileSplitkParams fused_gemm_params;

@@ -988,15 +976,14 @@ torch::stable::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::stable::empty({static_cast<int64_t>(ws_size)},
-                                          torch::headeronly::ScalarType::Char,
-                                          std::nullopt, a.device());
+    as_g_workspace = torch::empty({long(ws_size)}, ws_options);
  }
  void* ws = reinterpret_cast<void*>(as_g_workspace.data_ptr());

-  if (a.scalar_type() == torch::headeronly::ScalarType::Half) {
+  if (a.dtype() == at::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),
@@ -1004,7 +991,7 @@ torch::stable::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.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
+  } else if (a.dtype() == at::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),
@@ -1019,6 +1006,6 @@ torch::stable::Tensor allspark_w8a16_gemm(

 #endif

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

 namespace allspark {

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

 void rearrange_kn_weight_as_n32k16_order(
-    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) {
+    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) {
  // Verify device and strides
-  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_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_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_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_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_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_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");
+  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");

  if (has_zp) {
-    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.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_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");
+    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");
  }

  const uint8_t* matB = reinterpret_cast<const uint8_t*>(b_qweight.data_ptr());
@@ -145,20 +136,18 @@ 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.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;
+      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;

-  cudaStream_t stream = get_current_cuda_stream();
-  if (b_scales.scalar_type() == torch::headeronly::ScalarType::Half) {
+  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  if (b_scales.dtype() == at::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.scalar_type() ==
-             torch::headeronly::ScalarType::BFloat16) {
+  } else if (b_scales.dtype() == at::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,
@@ -168,7 +157,7 @@ void rearrange_kn_weight_as_n32k16_order(
  }
 }

-STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
+TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
  m.impl("rearrange_kn_weight_as_n32k16_order",
-         TORCH_BOX(&rearrange_kn_weight_as_n32k16_order));
+         &rearrange_kn_weight_as_n32k16_order);
 }
--- a/Show More
+++ b/Show More