[BugFix] Fix false assertion with spec-decode=[2,4,..] and TP>2 (#29036 )

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> (cherry picked from commit 8f4f77a727)
[BugFix] Ray with multiple nodes (#28873 )
2025-11-19 14:11:21 -08:00 · 2025-11-19 14:11:08 -08:00 · 2025-11-19 14:10:50 -08:00 · 2025-11-19 14:10:37 -08:00
1013 changed files with 17632 additions and 44575 deletions
@@ -23,8 +23,8 @@ To download the wheel (by version):
 aws s3 cp s3://vllm-wheels/${RELEASE_VERSION}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux1_x86_64.whl .
 aws s3 cp s3://vllm-wheels/${RELEASE_VERSION}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux2014_aarch64.whl .

+aws s3 cp s3://vllm-wheels/${RELEASE_VERSION}+cu126/vllm-${RELEASE_VERSION}+cu126-cp38-abi3-manylinux1_x86_64.whl .
 aws s3 cp s3://vllm-wheels/${RELEASE_VERSION}+cu129/vllm-${RELEASE_VERSION}+cu129-cp38-abi3-manylinux1_x86_64.whl .
-aws s3 cp s3://vllm-wheels/${RELEASE_VERSION}+cu130/vllm-${RELEASE_VERSION}+cu130-cp38-abi3-manylinux1_x86_64.whl .
 \`\`\`

 To download and upload the image:
@@ -45,10 +45,9 @@ docker tag vllm/vllm-openai:aarch64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
 docker push vllm/vllm-openai:latest-aarch64
 docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64

-docker manifest rm vllm/vllm-openai:latest
-docker manifest create vllm/vllm-openai:latest vllm/vllm-openai:latest-x86_64 vllm/vllm-openai:latest-aarch64
-docker manifest create vllm/vllm-openai:v${RELEASE_VERSION} vllm/vllm-openai:v${RELEASE_VERSION}-x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
+docker manifest create vllm/vllm-openai:latest vllm/vllm-openai:latest-x86_64 vllm/vllm-openai:latest-aarch64 --amend
+docker manifest create vllm/vllm-openai:v${RELEASE_VERSION} vllm/vllm-openai:v${RELEASE_VERSION}-x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64 --amend
 docker manifest push vllm/vllm-openai:latest
 docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}
 \`\`\`
-EOF 
+EOF 
@@ -1,62 +0,0 @@
-#!/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
-
-# allow to bind to different cores
-CORE_RANGE=${CORE_RANGE:-0-16}
-OMP_CORE_RANGE=${OMP_CORE_RANGE:-0-16}
-
-export CMAKE_BUILD_PARALLEL_LEVEL=16
-
-# Setup cleanup
-remove_docker_container() {
-    set -e;
-    docker rm -f cpu-test || true;
-}
-trap remove_docker_container EXIT
-remove_docker_container
-
-# Try building the docker image
-docker build --tag cpu-test --target vllm-test -f docker/Dockerfile.cpu .
-
-# Run the image
-docker run -itd --cpuset-cpus="$CORE_RANGE" --entrypoint /bin/bash -v ~/.cache/huggingface:/root/.cache/huggingface -e HF_TOKEN --env VLLM_CPU_KVCACHE_SPACE=16 --env VLLM_CPU_CI_ENV=1 -e E2E_OMP_THREADS="$OMP_CORE_RANGE" --shm-size=4g --name cpu-test cpu-test
-
-function cpu_tests() {
-  set -e
-
-  docker exec cpu-test bash -c "
-    set -e
-    pip list"
-
-  # offline inference
-  docker exec cpu-test bash -c "
-    set -e
-    python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m"
-
-  # Run kernel tests
-  docker exec cpu-test bash -c "
-    set -e
-    pytest -x -v -s tests/kernels/test_onednn.py
-    pytest -x -v -s tests/kernels/attention/test_cpu_attn.py"
-
-  # basic online serving
-  docker exec cpu-test bash -c '
-    set -e
-    VLLM_CPU_OMP_THREADS_BIND=$E2E_OMP_THREADS vllm serve Qwen/Qwen3-0.6B --max-model-len 2048 &
-    server_pid=$!
-    timeout 600 bash -c "until curl localhost:8000/v1/models; do sleep 1; done" || exit 1
-    vllm bench serve \
-      --backend vllm \
-      --dataset-name random \
-      --model Qwen/Qwen3-0.6B \
-      --num-prompts 20 \
-      --endpoint /v1/completions
-    kill -s SIGTERM $server_pid &'
-}
-
-# All of CPU tests are expected to be finished less than 40 mins.
-export -f cpu_tests
-timeout 2h bash -c cpu_tests
@@ -25,22 +25,20 @@ function cpu_tests() {

  # offline inference
  podman exec -it "$container_id" bash -c "
-    export TORCH_COMPILE_DISABLE=1
    set -xve
    python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m" >> $HOME/test_basic.log

  # Run basic model test
  podman exec -it "$container_id" bash -c "
-    export TORCH_COMPILE_DISABLE=1
    set -evx
    pip install pytest pytest-asyncio einops peft Pillow soundfile transformers_stream_generator matplotlib
-    pip install sentence-transformers datamodel_code_generator tblib 
+    pip install sentence-transformers datamodel_code_generator

    # Note: disable Bart until supports V1
    # pytest -v -s tests/models/language/generation/test_bart.py -m cpu_model
-    pytest -v -s tests/models/language/generation/test_common.py::test_models[False-False-5-32-openai-community/gpt2]
-    pytest -v -s tests/models/language/generation/test_common.py::test_models[False-False-5-32-facebook/opt-125m]
-    pytest -v -s tests/models/language/generation/test_common.py::test_models[False-False-5-32-google/gemma-1.1-2b-it]
+    pytest -v -s tests/models/language/generation/test_common.py::test_models[False-5-32-openai-community/gpt2]
+    pytest -v -s tests/models/language/generation/test_common.py::test_models[False-5-32-facebook/opt-125m]
+    pytest -v -s tests/models/language/generation/test_common.py::test_models[False-5-32-google/gemma-1.1-2b-it]
    pytest -v -s tests/models/language/pooling/test_classification.py::test_models[float-jason9693/Qwen2.5-1.5B-apeach]
    # TODO: Below test case tests/models/language/pooling/test_embedding.py::test_models[True-ssmits/Qwen2-7B-Instruct-embed-base] fails on ppc64le. Disabling it for time being.
    # pytest -v -s tests/models/language/pooling/test_embedding.py -m cpu_model" >> $HOME/test_rest.log
@@ -21,8 +21,8 @@ trap remove_docker_container EXIT
 remove_docker_container

 # Try building the docker image
-numactl -C "$CORE_RANGE" -N "$NUMA_NODE" docker build --progress plain --tag cpu-test-"$NUMA_NODE" --target vllm-test -f docker/Dockerfile.cpu .
-numactl -C "$CORE_RANGE" -N "$NUMA_NODE" docker build --progress plain --build-arg VLLM_CPU_DISABLE_AVX512="true" --tag cpu-test-"$NUMA_NODE"-avx2 --target vllm-test -f docker/Dockerfile.cpu .
+numactl -C "$CORE_RANGE" -N "$NUMA_NODE" docker build --tag cpu-test-"$NUMA_NODE" --target vllm-test -f docker/Dockerfile.cpu .
+numactl -C "$CORE_RANGE" -N "$NUMA_NODE" docker build --build-arg VLLM_CPU_DISABLE_AVX512="true" --tag cpu-test-"$NUMA_NODE"-avx2 --target vllm-test -f docker/Dockerfile.cpu .

 # Run the image, setting --shm-size=4g for tensor parallel.
 docker run -itd --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" --entrypoint /bin/bash -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN --env VLLM_CPU_KVCACHE_SPACE=16 --env VLLM_CPU_CI_ENV=1 -e E2E_OMP_THREADS="$OMP_CORE_RANGE" --shm-size=4g --name cpu-test-"$NUMA_NODE" cpu-test-"$NUMA_NODE"
@@ -73,11 +73,12 @@ function cpu_tests() {
    pytest -x -s -v \
    tests/quantization/test_compressed_tensors.py::test_compressed_tensors_w8a8_logprobs"

-  # Run AWQ/GPTQ test
-  docker exec cpu-test-"$NUMA_NODE" bash -c "
-    set -e
-    pytest -x -s -v \
-    tests/quantization/test_cpu_wna16.py"
+  # Note: disable it until supports V1
+  # Run AWQ test
+  # docker exec cpu-test-"$NUMA_NODE" bash -c "
+  #   set -e
+  #   pytest -x -s -v \
+  #   tests/quantization/test_ipex_quant.py"

  # Run multi-lora tests
  docker exec cpu-test-"$NUMA_NODE" bash -c "
@@ -35,7 +35,7 @@ docker run \
    echo $ZE_AFFINITY_MASK
    pip install tblib==3.1.0
    python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager
-    python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 -O3 -cc.cudagraph_mode=NONE
+    python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 -O3 -O.cudagraph_mode=NONE
    python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager -tp 2 --distributed-executor-backend ray
    python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager -tp 2 --distributed-executor-backend mp
    VLLM_ATTENTION_BACKEND=TRITON_ATTN python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager
@@ -17,17 +17,7 @@ wait_for_server() {
 }

 MODEL="deepseek-ai/DeepSeek-V2-lite"
-
-# Set BACKENDS based on platform
-if command -v rocm-smi &> /dev/null || [[ -d /opt/rocm ]] || [[ -n "${ROCM_PATH:-}" ]]; then
-  # ROCm platform
-  BACKENDS=("allgather_reducescatter")
-  # Disable MOE padding for ROCm since it is causing eplb to fail
-  export VLLM_ROCM_MOE_PADDING=0
-else
-  # Non-ROCm platform (CUDA/other)
-  BACKENDS=("deepep_high_throughput" "deepep_low_latency")
-fi
+BACKENDS=("deepep_high_throughput" "deepep_low_latency")

 cleanup() {
  if [[ -n "${SERVER_PID:-}" ]] && kill -0 "${SERVER_PID}" 2>/dev/null; then
@@ -1,12 +1,10 @@
 #!/usr/bin/env bash
 set -euxo pipefail

-# args: [THRESHOLD] [NUM_QUESTIONS] [START_PORT] [DATA_PARALLEL_SIZE] [TENSOR_PARALLEL_SIZE]
+# args: [THRESHOLD] [NUM_QUESTIONS] [START_PORT]
 THRESHOLD=${1:-0.8}
 NUM_Q=${2:-1319}
 PORT=${3:-8020}
-DATA_PARALLEL_SIZE=${4:-2}
-TENSOR_PARALLEL_SIZE=${5:-2}
 OUT_DIR=${OUT_DIR:-/tmp/vllm-scheduled}
 mkdir -p "${OUT_DIR}"

@@ -19,16 +17,7 @@ wait_for_server() {
 }

 MODEL="QWen/Qwen3-30B-A3B-FP8"
-# Set BACKENDS based on platform
-if command -v rocm-smi &> /dev/null || [[ -d /opt/rocm ]] || [[ -n "${ROCM_PATH:-}" ]]; then
-  # ROCm platform
-  BACKENDS=("allgather_reducescatter")
-  # Disable MOE padding for ROCm since it is causing eplb to fail
-  export VLLM_ROCM_MOE_PADDING=0
-else
-  # Non-ROCm platform (CUDA/other)
-  BACKENDS=("deepep_high_throughput" "deepep_low_latency")
-fi
+BACKENDS=("deepep_high_throughput" "deepep_low_latency")

 cleanup() {
  if [[ -n "${SERVER_PID:-}" ]] && kill -0 "${SERVER_PID}" 2>/dev/null; then
@@ -47,10 +36,8 @@ for BACK in "${BACKENDS[@]}"; do
  VLLM_ALL2ALL_BACKEND=$BACK \
  vllm serve "$MODEL" \
    --enforce-eager \
-    --enable-eplb \
-    --eplb-config '{"window_size":10, "step_interval":100, "num_redundant_experts":0, "log_balancedness":true}' \
-    --tensor-parallel-size ${TENSOR_PARALLEL_SIZE} \
-    --data-parallel-size ${DATA_PARALLEL_SIZE} \
+    --tensor-parallel-size 2 \
+    --data-parallel-size 2 \
    --enable-expert-parallel \
    --trust-remote-code \
    --max-model-len 2048 \
@@ -61,8 +61,8 @@ steps:
  - pytest -v -s -m 'not cpu_test' multimodal
  - pytest -v -s utils_

- label: Async Engine, Inputs, Utils, Worker, Config Test (CPU) # 15min
-  timeout_in_minutes: 20
+- label: Async Engine, Inputs, Utils, Worker Test (CPU) # 4 mins
+  timeout_in_minutes: 10
  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi325_1
  # grade: Blocking
@@ -72,18 +72,14 @@ steps:
  - tests/test_outputs.py
  - tests/multimodal
  - tests/standalone_tests/lazy_imports.py
-  - tests/tokenizers_
  - tests/transformers_utils
-  - tests/config
  no_gpu: true
  commands:
  - python3 standalone_tests/lazy_imports.py
  - pytest -v -s test_inputs.py
  - pytest -v -s test_outputs.py
  - pytest -v -s -m 'cpu_test' multimodal
-  - pytest -v -s tokenizers_
  - pytest -v -s transformers_utils
-  - pytest -v -s config

 - label: Python-only Installation Test # 10min
  timeout_in_minutes: 20
@@ -191,7 +187,7 @@ steps:
  - tests/distributed/test_utils
  - tests/distributed/test_pynccl
  - tests/distributed/test_events
-  - tests/compile/fullgraph/test_basic_correctness.py
+  - tests/compile/test_basic_correctness
  - examples/offline_inference/rlhf.py
  - examples/offline_inference/rlhf_colocate.py
  - tests/examples/offline_inference/data_parallel.py
@@ -219,7 +215,7 @@ steps:
  - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py
  - pytest -v -s v1/engine/test_engine_core_client.py::test_kv_cache_events_dp
  - pytest -v -s distributed/test_utils.py
-  - pytest -v -s compile/fullgraph/test_basic_correctness.py
+  - pytest -v -s compile/test_basic_correctness.py
  - pytest -v -s distributed/test_pynccl.py
  - pytest -v -s distributed/test_events.py
  - pytest -v -s distributed/test_symm_mem_allreduce.py
@@ -310,20 +306,23 @@ steps:
  - pytest -v -s test_regression.py
  working_dir: "/vllm-workspace/tests" # optional

- label: Engine Test # 9min
-  timeout_in_minutes: 15
+- label: Engine Test # 25min
+  timeout_in_minutes: 40
  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi325_1
  # grade: Blocking
  source_file_dependencies:
  - vllm/
  - tests/engine
+  - tests/tokenization
  - tests/test_sequence
  - tests/test_config
  - tests/test_logger
  - tests/test_vllm_port
  commands:
  - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py
+  # OOM in the CI unless we run this separately
+  - pytest -v -s tokenization

 - label: V1 Test e2e + engine # 30min
  timeout_in_minutes: 45
@@ -391,15 +390,6 @@ steps:
  commands:
    - pytest -v -s v1/attention

- label: V1 Test attention (B200) # 10min
-  timeout_in_minutes: 30
-  gpu: b200
-  source_file_dependencies:
-    - vllm/v1/attention
-    - tests/v1/attention
-  commands:
-    - VLLM_DISABLE_FLASHINFER_PREFILL=1 pytest -v -s v1/attention # TODO: FI prefill is bugged and causes incorrectness, fix this
-
 - label: V1 Test others (CPU) # 5 mins
  mirror_hardwares: [amdexperimental, amdproduction]
  agent_pool: mi325_1
@@ -503,12 +493,17 @@ steps:
    - vllm/
    - tests/compile
  commands:
-  # Run unit tests defined directly under compile/,
-  # not including subdirectories, which are usually heavier
-  # tests covered elsewhere.
-  # Use `find` to launch multiple instances of pytest so that
-  # they do not suffer from https://github.com/vllm-project/vllm/issues/28965
-  - "find compile/ -maxdepth 1 -name 'test_*.py' -exec pytest -s -v {} \\\\;"
+    - pytest -v -s compile/test_pass_manager.py
+    - pytest -v -s compile/test_fusion.py
+    - pytest -v -s compile/test_fusion_attn.py
+    - pytest -v -s compile/test_functionalization.py
+    - pytest -v -s compile/test_silu_mul_quant_fusion.py
+  #  - pytest -v -s compile/test_sequence_parallelism.py
+  #  - pytest -v -s compile/test_async_tp.py
+    - pytest -v -s compile/test_fusion_all_reduce.py
+    - pytest -v -s compile/test_decorator.py
+    - pytest -v -s compile/test_noop_elimination.py
+    - pytest -v -s compile/test_aot_compile.py

 - label: PyTorch Fullgraph Smoke Test # 15min
  timeout_in_minutes: 30
@@ -520,11 +515,9 @@ steps:
  - vllm/
  - tests/compile
  commands:
-  # Run smoke tests under fullgraph directory, except test_full_graph.py
-  # as it is a heavy test that is covered in other steps.
-  # Use `find` to launch multiple instances of pytest so that
-  # they do not suffer from https://github.com/vllm-project/vllm/issues/28965
-  - "find compile/fullgraph/ -name 'test_*.py' -not -name 'test_full_graph.py' -exec pytest -s -v {} \\\\;"
+  - pytest -v -s compile/test_basic_correctness.py
+  - pytest -v -s compile/test_multimodal_compile.py
+  - pytest -v -s compile/piecewise/

 - label: PyTorch Fullgraph Test # 27min
  timeout_in_minutes: 40
@@ -536,10 +529,10 @@ steps:
  - vllm/
  - tests/compile
  commands:
-  - pytest -v -s compile/fullgraph/test_full_graph.py -k 'not test_fp8_kv_scale_compile'
+  - pytest -v -s compile/test_full_graph.py -k 'not test_fp8_kv_scale_compile'
    # Limit to no custom ops to reduce running time
    # Wrap with quotes to escape yaml and avoid starting -k string with a -
-  - "pytest -v -s compile/distributed/test_fusions_e2e.py -k 'TRITON and not +quant_fp8 and not Llama-4'"
+  - "pytest -v -s compile/test_fusions_e2e.py -k 'TRITON and -quant_fp8'"

 - label: Cudagraph test
  timeout_in_minutes: 20
@@ -704,7 +697,7 @@ steps:
  - vllm/model_executor/models/whisper.py
  commands: # LMEval
  # Transcription WER check is skipped because encoder-decoder models are not supported on ROCm, see https://github.com/vllm-project/vllm/issues/27442
-  - pytest -s entrypoints/openai/correctness/
+  - pytest -s entrypoints/openai/correctness/  --ignore entrypoints/openai/correctness/test_transcription_api_correctness.py

 - label: OpenAI-Compatible Tool Use # 23 min
  timeout_in_minutes: 35
@@ -753,7 +746,6 @@ steps:
  torch_nightly: true
  source_file_dependencies:
  - vllm/model_executor/models/
-  - vllm/transformers_utils/
  - tests/models/test_initialization.py
  commands:
    # Only when vLLM model source is modified - test initialization of a large
@@ -1006,12 +998,12 @@ steps:
  optional: true
  commands:
    - pip install --upgrade git+https://github.com/huggingface/transformers
-    - pytest -v -s tests/models/test_initialization.py -k 'not (Gemma3 or ModernBert or Qwen2_5_VL or Qwen2_5vl or Qwen2VL or TransformersMultiModalEmbeddingModel or TransformersMultiModalForSequenceClassification or Ultravox or Phi4Multimodal or LlavaNextVideo or MiniCPMO or Lfm2Moe or PaliGemma or RobertaForSequenceClassification or Ovis2_5 or Fuyu or DeepseekOCR or KimiVL)'
+    - pytest -v -s tests/models/test_initialization.py
    - pytest -v -s tests/models/test_transformers.py
-    # - pytest -v -s tests/models/multimodal/processing/
-    - pytest -v -s tests/models/multimodal/test_mapping.py -k 'not (Gemma3 or Qwen2VL or Qwen2_5_VL)'
+    - pytest -v -s tests/models/multimodal/processing/
+    - pytest -v -s tests/models/multimodal/test_mapping.py
    - python3 examples/offline_inference/basic/chat.py
-    # - python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl
+    - python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl
    # Whisper needs spawn method to avoid deadlock
    - VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper

@@ -1056,7 +1048,7 @@ steps:
    - pytest -v -s tests/kernels/moe/test_ocp_mx_moe.py
    - pytest -v -s tests/kernels/moe/test_flashinfer.py

- label: Blackwell Fusion and Compile Tests # 30 min
+- label: Blackwell Fusion Tests # 30 min
  timeout_in_minutes: 40
  working_dir: "/vllm-workspace/"
  gpu: b200
@@ -1074,12 +1066,10 @@ steps:
    - pytest -v -s tests/compile/test_fusion_attn.py
    - pytest -v -s tests/compile/test_silu_mul_quant_fusion.py
    # this runner has 2 GPUs available even though num_gpus=2 is not set
-    - pytest -v -s tests/compile/distributed/test_fusion_all_reduce.py
+    - pytest -v -s tests/compile/test_fusion_all_reduce.py
    # Limit to Inductor partition, no custom ops, and allreduce & attn fusion to reduce running time
    # Wrap with quotes to escape yaml
-    - "pytest -v -s tests/compile/distributed/test_fusions_e2e.py::test_tp2_attn_quant_allreduce_rmsnorm -k 'True and not +quant_fp8 and not +rms_norm'"
-    # test_fp8_kv_scale_compile requires FlashAttention (not supported on default L4/L40)
-    - pytest -v -s tests/compile/distributed/test_full_graph.py::test_fp8_kv_scale_compile
+    - "pytest -v -s tests/compile/test_fusions_e2e.py::test_tp2_attn_quant_allreduce_rmsnorm -k 'True and Llama-3.1 and -quant_fp8 and -rms_norm'"

 - label: Blackwell Fusion E2E Tests # 30 min
  timeout_in_minutes: 40
@@ -1096,18 +1086,20 @@ steps:
  - vllm/model_executor/layers/layernorm.py
  - vllm/model_executor/layers/activation.py
  - vllm/model_executor/layers/quantization/input_quant_fp8.py
-  - tests/compile/distributed/test_fusions_e2e.py
-  - tests/compile/fullgraph/test_full_graph.py
+  - tests/compile/test_fusions_e2e.py
+  - tests/compile/test_full_graph.py
  commands:
    - nvidia-smi
    # Run all e2e fusion tests
    - pytest -v -s tests/compile/test_fusions_e2e.py
+    # test_fp8_kv_scale_compile requires FlashAttention (not supported on default L4/L40)
+    - pytest -v -s tests/compile/test_full_graph.py::test_fp8_kv_scale_compile

 - label: ROCm GPT-OSS Eval
  timeout_in_minutes: 60
  working_dir: "/vllm-workspace/"
  agent_pool: mi325_1
-  mirror_hardwares: [amdexperimental, amdproduction]
+  mirror_hardwares: [amdproduction]
  optional: true # run on nightlies
  source_file_dependencies:
  - tests/evals/gpt_oss
@@ -1206,7 +1198,7 @@ steps:
  - vllm/worker/worker_base.py
  - vllm/v1/engine/
  - vllm/v1/worker/
-  - tests/compile/fullgraph/test_basic_correctness.py
+  - tests/compile/test_basic_correctness.py
  - tests/compile/test_wrapper.py
  - tests/distributed/
  - tests/entrypoints/llm/test_collective_rpc.py
@@ -1219,7 +1211,7 @@ steps:
  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py
  - DP_SIZE=2 pytest -v -s v1/entrypoints/openai/test_multi_api_servers.py
  - pytest -v -s entrypoints/llm/test_collective_rpc.py
-  - pytest -v -s ./compile/fullgraph/test_basic_correctness.py
+  - pytest -v -s ./compile/test_basic_correctness.py
  - pytest -v -s ./compile/test_wrapper.py
  - VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
  - VLLM_TEST_SAME_HOST=1 VLLM_TEST_WITH_DEFAULT_DEVICE_SET=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
@@ -1319,10 +1311,7 @@ steps:
    - pytest -v -s -x lora/test_llama_tp.py
    - pytest -v -s -x lora/test_llm_with_multi_loras.py
    - pytest -v -s -x lora/test_olmoe_tp.py
-
-    # Disabled for now because MXFP4 backend on non-cuda platform 
-    # doesn't support LoRA yet
-    #- pytest -v -s -x lora/test_gptoss_tp.py
+    - pytest -v -s -x lora/test_gptoss_tp.py


 - label: Weight Loading Multiple GPU Test  # 33min
@@ -1337,7 +1326,7 @@ steps:
  - vllm/
  - tests/weight_loading
  commands:
-    - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-amd.txt
+    - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models.txt

 - label: Weight Loading Multiple GPU Test - Large Models # optional
  mirror_hardwares: [amdexperimental]
@@ -1345,12 +1334,13 @@ steps:
  # grade: Blocking
  working_dir: "/vllm-workspace/tests"
  num_gpus: 2
+  gpu: a100
  optional: true
  source_file_dependencies:
  - vllm/
  - tests/weight_loading
  commands:
-    - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-large-amd.txt
+    - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-large.txt

 - label: NixlConnector PD accuracy tests (Distributed) # 30min
  mirror_hardwares: [amdexperimental]
@@ -1427,12 +1417,10 @@ steps:
  working_dir: "/vllm-workspace/"
  num_gpus: 2
  commands:
-    - pytest -v -s tests/compile/distributed/test_async_tp.py
-    - pytest -v -s tests/compile/distributed/test_sequence_parallelism.py
-    - pytest -v -s tests/compile/distributed/test_fusion_all_reduce.py
-    #- pytest -v -s tests/compile/distributed/test_fusions_e2e.py::test_tp2_attn_quant_allreduce_rmsnorm
-    - "pytest -v -s tests/compile/distributed/test_fusions_e2e.py -k 'not Llama-4'"
-    - pytest -v -s tests/compile/distributed/test_sequence_parallel.py
+    - pytest -v -s tests/compile/test_async_tp.py
+    - pytest -v -s tests/compile/test_sequence_parallelism.py
+    - pytest -v -s tests/compile/test_fusion_all_reduce.py
+    - pytest -v -s tests/compile/test_fusions_e2e.py::test_tp2_attn_quant_allreduce_rmsnorm
    - pytest -v -s tests/distributed/test_context_parallel.py
    - CUDA_VISIBLE_DEVICES=1,2 VLLM_ALL2ALL_BACKEND=deepep_high_throughput VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model Qwen/Qwen1.5-MoE-A2.7B --tp-size=1  --dp-size=2 --max-model-len 2048
    - pytest -v -s tests/v1/distributed/test_dbo.py
@@ -1485,4 +1473,4 @@ steps:
  num_gpus: 4
  working_dir: "/vllm-workspace"
  commands:
-  - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020
+  - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep.sh 0.8 200 8020
@@ -57,15 +57,14 @@ steps:
  - pytest -v -s -m 'not cpu_test' multimodal
  - pytest -v -s utils_

- label: Async Engine, Inputs, Utils, Worker, Config Test (CPU) # 15min
-  timeout_in_minutes: 20
+- label: Async Engine, Inputs, Utils, Worker, Config Test (CPU) # 4 mins
+  timeout_in_minutes: 10
  source_file_dependencies:
  - vllm/
  - tests/test_inputs.py
  - tests/test_outputs.py
  - tests/multimodal
  - tests/standalone_tests/lazy_imports.py
-  - tests/tokenizers_
  - tests/transformers_utils
  - tests/config
  no_gpu: true
@@ -74,7 +73,6 @@ steps:
  - pytest -v -s test_inputs.py
  - pytest -v -s test_outputs.py
  - pytest -v -s -m 'cpu_test' multimodal
-  - pytest -v -s tokenizers_
  - pytest -v -s transformers_utils
  - pytest -v -s config

@@ -169,7 +167,7 @@ steps:
  - tests/distributed/test_utils
  - tests/distributed/test_pynccl
  - tests/distributed/test_events
-  - tests/compile/fullgraph/test_basic_correctness.py
+  - tests/compile/test_basic_correctness
  - examples/offline_inference/rlhf.py
  - examples/offline_inference/rlhf_colocate.py
  - tests/examples/offline_inference/data_parallel.py
@@ -194,13 +192,12 @@ steps:
  # test with internal dp
  - python3 ../examples/offline_inference/data_parallel.py --enforce-eager
  - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py
-  - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
  - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py
  - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_internal_lb_dp.py
  - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py
  - pytest -v -s v1/engine/test_engine_core_client.py::test_kv_cache_events_dp
  - pytest -v -s distributed/test_utils.py
-  - pytest -v -s compile/fullgraph/test_basic_correctness.py
+  - pytest -v -s compile/test_basic_correctness.py
  - pytest -v -s distributed/test_pynccl.py
  - pytest -v -s distributed/test_events.py
  - pytest -v -s distributed/test_symm_mem_allreduce.py
@@ -215,7 +212,6 @@ steps:
  timeout_in_minutes: 10
  gpu: h100
  num_gpus: 8
-  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - examples/offline_inference/torchrun_dp_example.py
@@ -279,18 +275,21 @@ steps:
  - pytest -v -s test_regression.py
  working_dir: "/vllm-workspace/tests" # optional

- label: Engine Test # 9min
-  timeout_in_minutes: 15
+- label: Engine Test # 25min
+  timeout_in_minutes: 40
  mirror_hardwares: [amdexperimental]
  source_file_dependencies:
  - vllm/
  - tests/engine
+  - tests/tokenization
  - tests/test_sequence
  - tests/test_config
  - tests/test_logger
  - tests/test_vllm_port
  commands:
  - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py
+  # OOM in the CI unless we run this separately
+  - pytest -v -s tokenization

 - label: V1 Test e2e + engine # 30min
  timeout_in_minutes: 45
@@ -347,18 +346,6 @@ steps:
  commands:
    - pytest -v -s v1/attention

- label: Batch Invariance Tests (H100) # 10min
-  timeout_in_minutes: 25
-  gpu: h100
-  source_file_dependencies:
-    - vllm/
-    - 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
-    - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
-
 - label: V1 Test attention (B200) # 10min
  timeout_in_minutes: 30
  gpu: b200
@@ -458,12 +445,18 @@ steps:
    - vllm/
    - tests/compile
  commands:
-  # Run unit tests defined directly under compile/,
-  # not including subdirectories, which are usually heavier
-  # tests covered elsewhere.
-  # Use `find` to launch multiple instances of pytest so that
-  # they do not suffer from https://github.com/vllm-project/vllm/issues/28965
-  - "find compile/ -maxdepth 1 -name 'test_*.py' -exec pytest -s -v {} \\\\;"
+    - pytest -v -s compile/test_graph_partition.py
+    - pytest -v -s compile/test_config.py
+    - pytest -v -s compile/test_pass_manager.py
+    - pytest -v -s compile/test_fusion.py
+    - pytest -v -s compile/test_fusion_attn.py
+    - pytest -v -s compile/test_functionalization.py
+    - pytest -v -s compile/test_silu_mul_quant_fusion.py
+    - pytest -v -s compile/test_fusion_all_reduce.py
+    - pytest -v -s compile/test_decorator.py
+    - pytest -v -s compile/test_noop_elimination.py
+    - pytest -v -s compile/test_aot_compile.py
+    - pytest -v -s compile/test_qk_norm_rope_fusion.py

 - label: PyTorch Fullgraph Smoke Test # 15min
  timeout_in_minutes: 30
@@ -473,11 +466,9 @@ steps:
  - vllm/
  - tests/compile
  commands:
-  # Run smoke tests under fullgraph directory, except test_full_graph.py
-  # as it is a heavy test that is covered in other steps.
-  # Use `find` to launch multiple instances of pytest so that
-  # they do not suffer from https://github.com/vllm-project/vllm/issues/28965
-  - "find compile/fullgraph/ -name 'test_*.py' -not -name 'test_full_graph.py' -exec pytest -s -v {} \\\\;"
+  - pytest -v -s compile/test_basic_correctness.py
+  - pytest -v -s compile/test_multimodal_compile.py
+  - pytest -v -s compile/piecewise/

 - label: PyTorch Fullgraph Test # 27min
  timeout_in_minutes: 40
@@ -488,10 +479,10 @@ steps:
  - tests/compile
  commands:
    # fp8 kv scales not supported on sm89, tested on Blackwell instead
-  - pytest -v -s compile/fullgraph/test_full_graph.py -k 'not test_fp8_kv_scale_compile'
+  - pytest -v -s compile/test_full_graph.py -k 'not test_fp8_kv_scale_compile'
    # Limit to no custom ops to reduce running time
    # Wrap with quotes to escape yaml and avoid starting -k string with a -
-  - "pytest -v -s compile/distributed/test_fusions_e2e.py -k 'TRITON and not +quant_fp8 and not Llama-4'"
+  - "pytest -v -s compile/test_fusions_e2e.py -k 'TRITON and not +quant_fp8 and not Llama-4'"

 - label: Cudagraph test
  timeout_in_minutes: 20
@@ -563,25 +554,6 @@ steps:
  commands:
    - pytest -v -s kernels/mamba

- label: Kernels DeepGEMM Test (H100)
-  timeout_in_minutes: 45
-  gpu: h100
-  num_gpus: 1
-  source_file_dependencies:
-  - tools/install_deepgemm.sh
-  - vllm/utils/deep_gemm.py
-  - vllm/model_executor/layers/fused_moe
-  - vllm/model_executor/layers/quantization
-  - tests/kernels/quantization/test_block_fp8.py
-  - tests/kernels/moe/test_deepgemm.py
-  - tests/kernels/moe/test_batched_deepgemm.py
-  - tests/kernels/attention/test_deepgemm_attention.py
-  commands:
-    - pytest -v -s kernels/quantization/test_block_fp8.py -k deep_gemm
-    - pytest -v -s kernels/moe/test_deepgemm.py
-    - pytest -v -s kernels/moe/test_batched_deepgemm.py
-    - pytest -v -s kernels/attention/test_deepgemm_attention.py
-
 - label: Model Executor Test # 23min
  timeout_in_minutes: 35
  torch_nightly: true
@@ -632,7 +604,6 @@ steps:
  # we can only upgrade after this is resolved
  # TODO(jerryzh168): resolve the above comment
  - uv pip install --system torchao==0.13.0 --index-url https://download.pytorch.org/whl/cu129
-  - uv pip install --system conch-triton-kernels
  - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py

 - label: LM Eval Small Models # 53min
@@ -693,7 +664,6 @@ steps:
  torch_nightly: true
  source_file_dependencies:
  - vllm/model_executor/models/
-  - vllm/transformers_utils/
  - tests/models/test_initialization.py
  commands:
    # Only when vLLM model source is modified - test initialization of a large
@@ -820,24 +790,14 @@ steps:
  commands:
    - pytest -v -s models/language/pooling_mteb_test

- label: Multi-Modal Processor Test (CPU)
-  timeout_in_minutes: 60
-  source_file_dependencies:
-  - vllm/
-  - tests/models/multimodal
-  no_gpu: true
-  commands:
-    - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
-    - pytest -v -s models/multimodal/processing --ignore models/multimodal/processing/test_tensor_schema.py
-
- label: Multi-Modal Processor Test
+- label: Multi-Modal Processor Test # 44min
  timeout_in_minutes: 60
  source_file_dependencies:
  - vllm/
  - tests/models/multimodal
  commands:
    - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
-    - pytest -v -s models/multimodal/processing/test_tensor_schema.py
+    - pytest -v -s models/multimodal/processing

 - label: Multi-Modal Models Test (Standard) # 60min
  timeout_in_minutes: 80
@@ -914,15 +874,14 @@ steps:
 - label: Transformers Nightly Models Test
  working_dir: "/vllm-workspace/"
  optional: true
-  soft_fail: true
  commands:
    - pip install --upgrade git+https://github.com/huggingface/transformers
-    - pytest -v -s tests/models/test_initialization.py
+    - pytest -v -s tests/models/test_initialization.py -k 'not (Gemma3 or ModernBert or Qwen2_5_VL or Qwen2_5vl or Qwen2VL or TransformersMultiModalEmbeddingModel or TransformersMultiModalForSequenceClassification or Ultravox or Phi4Multimodal or LlavaNextVideo or MiniCPMO or Lfm2Moe or PaliGemma or RobertaForSequenceClassification or Ovis2_5 or Fuyu or DeepseekOCR or KimiVL)'
    - pytest -v -s tests/models/test_transformers.py
-    - pytest -v -s tests/models/multimodal/processing/
-    - pytest -v -s tests/models/multimodal/test_mapping.py
+    # - pytest -v -s tests/models/multimodal/processing/
+    - pytest -v -s tests/models/multimodal/test_mapping.py -k 'not (Gemma3 or Qwen2VL or Qwen2_5_VL)'
    - python3 examples/offline_inference/basic/chat.py
-    - python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl
+    # - python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl
    # Whisper needs spawn method to avoid deadlock
    - VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper

@@ -966,7 +925,6 @@ steps:
    - pytest -v -s tests/kernels/moe/test_nvfp4_moe.py
    - pytest -v -s tests/kernels/moe/test_ocp_mx_moe.py
    - pytest -v -s tests/kernels/moe/test_flashinfer.py
-    - pytest -v -s tests/kernels/moe/test_cutedsl_moe.py

 - label: Blackwell Fusion and Compile Tests # 30 min
  timeout_in_minutes: 40
@@ -976,29 +934,22 @@ steps:
  - csrc/quantization/fp4/
  - vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
  - vllm/v1/attention/backends/flashinfer.py
-  - vllm/v1/worker/
-  - vllm/v1/cudagraph_dispatcher.py
  - vllm/compilation/
  # can affect pattern matching
  - vllm/model_executor/layers/layernorm.py
  - vllm/model_executor/layers/activation.py
  - vllm/model_executor/layers/quantization/input_quant_fp8.py
-  - tests/compile/test_fusion_attn.py
-  - tests/compile/test_silu_mul_quant_fusion.py
-  - tests/compile/distributed/test_fusion_all_reduce.py
-  - tests/compile/distributed/test_fusions_e2e.py
-  - tests/compile/fullgraph/test_full_graph.py
  commands:
    - nvidia-smi
    - pytest -v -s tests/compile/test_fusion_attn.py
    - pytest -v -s tests/compile/test_silu_mul_quant_fusion.py
    # this runner has 2 GPUs available even though num_gpus=2 is not set
-    - pytest -v -s tests/compile/distributed/test_fusion_all_reduce.py
+    - pytest -v -s tests/compile/test_fusion_all_reduce.py
    # Limit to Inductor partition, no custom ops, and allreduce & attn fusion to reduce running time
    # Wrap with quotes to escape yaml
-    - "pytest -v -s tests/compile/distributed/test_fusions_e2e.py::test_tp2_attn_quant_allreduce_rmsnorm -k 'True and not +quant_fp8 and not +rms_norm'"
+    - "pytest -v -s tests/compile/test_fusions_e2e.py::test_tp2_attn_quant_allreduce_rmsnorm -k 'True and not +quant_fp8 and not +rms_norm'"
    # test_fp8_kv_scale_compile requires FlashAttention (not supported on default L4/L40)
-    - pytest -v -s tests/compile/fullgraph/test_full_graph.py::test_fp8_kv_scale_compile
+    - pytest -v -s tests/compile/test_full_graph.py::test_fp8_kv_scale_compile

 - label: Blackwell Fusion E2E Tests # 30 min
  timeout_in_minutes: 40
@@ -1015,11 +966,12 @@ steps:
  - vllm/model_executor/layers/layernorm.py
  - vllm/model_executor/layers/activation.py
  - vllm/model_executor/layers/quantization/input_quant_fp8.py
-  - tests/compile/distributed/test_fusions_e2e.py
+  - tests/compile/test_fusions_e2e.py
+  - tests/compile/test_full_graph.py
  commands:
    - nvidia-smi
    # Run all e2e fusion tests
-    - pytest -v -s tests/compile/distributed/test_fusions_e2e.py
+    - pytest -v -s tests/compile/test_fusions_e2e.py

 - label: Blackwell GPT-OSS Eval
  timeout_in_minutes: 60
@@ -1117,7 +1069,7 @@ steps:
  - vllm/worker/worker_base.py
  - vllm/v1/engine/
  - vllm/v1/worker/
-  - tests/compile/fullgraph/test_basic_correctness.py
+  - tests/compile/test_basic_correctness.py
  - tests/compile/test_wrapper.py
  - tests/distributed/
  - tests/entrypoints/llm/test_collective_rpc.py
@@ -1129,11 +1081,10 @@ steps:
  # https://github.com/NVIDIA/nccl/issues/1838
  - export NCCL_CUMEM_HOST_ENABLE=0
  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py
-  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py
  - DP_SIZE=2 pytest -v -s v1/entrypoints/openai/test_multi_api_servers.py
  - pytest -v -s entrypoints/llm/test_collective_rpc.py
-  - pytest -v -s ./compile/fullgraph/test_basic_correctness.py
+  - pytest -v -s ./compile/test_basic_correctness.py
  - pytest -v -s ./compile/test_wrapper.py
  - VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
  - VLLM_TEST_SAME_HOST=1 VLLM_TEST_WITH_DEFAULT_DEVICE_SET=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
@@ -1313,11 +1264,11 @@ steps:
  working_dir: "/vllm-workspace/"
  num_gpus: 2
  commands:
-    - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/distributed/test_async_tp.py
-    - pytest -v -s tests/compile/distributed/test_sequence_parallelism.py
-    - pytest -v -s tests/compile/distributed/test_fusion_all_reduce.py
-    - "VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/distributed/test_fusions_e2e.py -k 'not Llama-4'"
-    - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/distributed/test_sequence_parallel.py
+    - pytest -v -s tests/compile/test_async_tp.py
+    - pytest -v -s tests/compile/test_sequence_parallelism.py
+    - pytest -v -s tests/compile/test_fusion_all_reduce.py
+    - "pytest -v -s tests/compile/test_fusions_e2e.py -k 'not Llama-4'"
+    - pytest -v -s tests/distributed/test_sequence_parallel.py
    - pytest -v -s tests/distributed/test_context_parallel.py
    - CUDA_VISIBLE_DEVICES=1,2 VLLM_ALL2ALL_BACKEND=deepep_high_throughput VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model Qwen/Qwen1.5-MoE-A2.7B --tp-size=1  --dp-size=2 --max-model-len 2048
    - pytest -v -s tests/v1/distributed/test_dbo.py
@@ -1354,20 +1305,11 @@ steps:
  commands:
  - bash .buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_ep_eplb.sh 0.25 200 8010

- label: Qwen3-30B-A3B-FP8-block Accuracy (H100)
+- label: Qwen3-30B-A3B-FP8-block Accuracy
  timeout_in_minutes: 60
  gpu: h100
  optional: true
  num_gpus: 4
  working_dir: "/vllm-workspace"
  commands:
-  - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020
-
- label: Qwen3-30B-A3B-FP8-block Accuracy (B200)
-  timeout_in_minutes: 60
-  gpu: b200
-  optional: true
-  num_gpus: 2
-  working_dir: "/vllm-workspace"
-  commands:
-  - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020 2 1
+  - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep.sh 0.8 200 8020
@@ -9,7 +9,6 @@
 /vllm/model_executor/layers/quantization @mgoin @robertgshaw2-redhat @tlrmchlsmth @yewentao256 @pavanimajety
 /vllm/model_executor/layers/mamba @tdoublep
 /vllm/model_executor/model_loader @22quinn
-/vllm/model_executor/layers/batch_invariant.py @yewentao256 
 /vllm/multimodal @DarkLight1337 @ywang96 @NickLucche @tjtanaa
 /vllm/vllm_flash_attn @LucasWilkinson
 /vllm/lora @jeejeelee
@@ -36,9 +35,6 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson
 /vllm/v1/kv_cache_interface.py @heheda12345
 /vllm/v1/offloading @ApostaC

-# Model runner V2
-/vllm/v1/worker/gpu @WoosukKwon
-
 # Test ownership
 /.buildkite/lm-eval-harness @mgoin 
 /tests/distributed/test_multi_node_assignment.py @youkaichao
@@ -60,7 +56,6 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson
 /tests/v1/kv_connector/nixl_integration @NickLucche
 /tests/v1/kv_connector @ApostaC
 /tests/v1/offloading @ApostaC
-/tests/v1/determinism @yewentao256 

 # Transformers modeling backend
 /vllm/model_executor/models/transformers @hmellor
@@ -149,7 +144,6 @@ mkdocs.yaml @hmellor
 /examples/*/pooling/ @noooop
 /tests/models/*/pooling* @noooop
 /tests/entrypoints/pooling @noooop
-/vllm/entrypoints/pooling @aarnphm @chaunceyjiang @noooop
 /vllm/config/pooler.py @noooop
 /vllm/pooling_params.py @noooop
 /vllm/model_executor/layers/pooler.py @noooop
@@ -13,7 +13,7 @@ jobs:

    steps:
      - name: Checkout repository
-        uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # v6.0.0
+        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2

      - name: Set up Python
        uses: actions/setup-python@e797f83bcb11b83ae66e0230d6156d7c80228e7c # v6.0.0
@@ -105,31 +105,6 @@ jobs:
                  }
                ],
              },
-              cpu: {
-                // Keyword search - matches whole words only (with word boundaries)
-                keywords: [
-                  {
-                    term: "CPU Backend",
-                    searchIn: "title"
-                  },
-                  {
-                    term: "x86",
-                    searchIn: "title"
-                  },
-                  {
-                    term: "ARM",
-                    searchIn: "title"
-                  },
-                  {
-                    term: "Apple Silicon",
-                    searchIn: "title"
-                  },
-                  {
-                    term: "IBM Z",
-                    searchIn: "title"
-                  },
-                ],
-              },
              // Add more label configurations here as needed
              // example: {
              //   keywords: [...],
@@ -9,10 +9,10 @@ on:
 jobs:
  macos-m1-smoke-test:
    runs-on: macos-latest
-    timeout-minutes: 30
+    timeout-minutes: 20

    steps:
-      - uses: actions/checkout@v6
+      - uses: actions/checkout@v4

      - uses: astral-sh/setup-uv@v7
        with:
@@ -37,14 +37,15 @@ jobs:
      - name: Verify installation
        run: |
          python -c "import vllm; print(f'vLLM version: {vllm.__version__}')"
+          python -c "import torch; print(f'PyTorch: {torch.__version__}')"

      - name: Smoke test vllm serve
+        timeout-minutes: 10
        run: |
          # Start server in background
          vllm serve Qwen/Qwen3-0.6B \
-            --max-model-len=2K \
+            --max-model-len=2048 \
            --load-format=dummy \
-            --hf-overrides '{"num_hidden_layers": 2}' \
            --enforce-eager \
            --port 8000 &

@@ -16,7 +16,7 @@ jobs:
  pre-commit:
    runs-on: ubuntu-latest
    steps:
-    - uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # v6.0.0
+    - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
    - uses: actions/setup-python@e797f83bcb11b83ae66e0230d6156d7c80228e7c # v6.0.0
      with:
        python-version: "3.12"
@@ -4,9 +4,6 @@
 # vllm-flash-attn built from source
 vllm/vllm_flash_attn/*

-# OpenAI triton kernels copied from source
-vllm/third_party/triton_kernels/*
-
 # triton jit
 .triton

@@ -136,7 +136,7 @@ elseif(HIP_FOUND)

  # ROCm 5.X and 6.X
  if (ROCM_VERSION_DEV_MAJOR GREATER_EQUAL 5 AND
-      Torch_VERSION VERSION_LESS ${TORCH_SUPPORTED_VERSION_ROCM})
+      NOT Torch_VERSION VERSION_EQUAL ${TORCH_SUPPORTED_VERSION_ROCM})
    message(WARNING "Pytorch version >= ${TORCH_SUPPORTED_VERSION_ROCM} "
      "expected for ROCm build, saw ${Torch_VERSION} instead.")
  endif()
@@ -307,7 +307,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
  SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")

  # Set CUTLASS_REVISION. Used for FetchContent. Also fixes some bogus messages when building.
-  set(CUTLASS_REVISION "v4.2.1")
+  set(CUTLASS_REVISION "v4.2.1" CACHE STRING "CUTLASS revision to use")

  # Use the specified CUTLASS source directory for compilation if VLLM_CUTLASS_SRC_DIR is provided
  if (DEFINED ENV{VLLM_CUTLASS_SRC_DIR})
@@ -354,17 +354,8 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
  # Only build Marlin kernels if we are building for at least some compatible archs.
  # Keep building Marlin for 9.0 as there are some group sizes and shapes that
  # are not supported by Machete yet.
-
-  # marlin arches for fp16 output
-  cuda_archs_loose_intersection(MARLIN_ARCHS "8.0+PTX" "${CUDA_ARCHS}")
-  # marlin arches for bf16 output (we need 9.0 for bf16 atomicAdd PTX)
-  cuda_archs_loose_intersection(MARLIN_BF16_ARCHS "8.0+PTX;9.0+PTX" "${CUDA_ARCHS}")
-  # marlin arches for fp8 input
-  # - sm80 doesn't support fp8 computation
-  # - sm90 and sm100 don't support QMMA.16832.F32.E4M3.E4M3 SAAS instruction
-  # so we only enable fp8 computation for SM89 (e.g. RTX 40x0)  and 12.0 (e.g. RTX 50x0)
-  cuda_archs_loose_intersection(MARLIN_FP8_ARCHS "8.9;12.0" "${CUDA_ARCHS}")
-
+  # 9.0 for latest bf16 atomicAdd PTX
+  cuda_archs_loose_intersection(MARLIN_ARCHS "8.0+PTX;9.0+PTX" "${CUDA_ARCHS}")
  if (MARLIN_ARCHS)

    #
@@ -374,18 +365,16 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
    set(MARLIN_GEN_SCRIPT
      ${CMAKE_CURRENT_SOURCE_DIR}/csrc/quantization/gptq_marlin/generate_kernels.py)
    file(MD5 ${MARLIN_GEN_SCRIPT} MARLIN_GEN_SCRIPT_HASH)
-    list(JOIN CUDA_ARCHS "," CUDA_ARCHS_STR)
-    set(MARLIN_GEN_SCRIPT_HASH_AND_ARCH "${MARLIN_GEN_SCRIPT_HASH}(ARCH:${CUDA_ARCHS_STR})")

-    message(STATUS "Marlin generation script hash: ${MARLIN_GEN_SCRIPT_HASH_AND_ARCH}")
-    message(STATUS "Last run Marlin generate script hash: $CACHE{MARLIN_GEN_SCRIPT_HASH_AND_ARCH}")
+    message(STATUS "Marlin generation script hash: ${MARLIN_GEN_SCRIPT_HASH}")
+    message(STATUS "Last run Marlin generate script hash: $CACHE{MARLIN_GEN_SCRIPT_HASH}")

-    if (NOT DEFINED CACHE{MARLIN_GEN_SCRIPT_HASH_AND_ARCH}
-        OR NOT $CACHE{MARLIN_GEN_SCRIPT_HASH_AND_ARCH} STREQUAL ${MARLIN_GEN_SCRIPT_HASH_AND_ARCH})
+    if (NOT DEFINED CACHE{MARLIN_GEN_SCRIPT_HASH}
+        OR NOT $CACHE{MARLIN_GEN_SCRIPT_HASH} STREQUAL ${MARLIN_GEN_SCRIPT_HASH})
      execute_process(
        COMMAND ${CMAKE_COMMAND} -E env
        PYTHONPATH=$PYTHONPATH
-          ${Python_EXECUTABLE} ${MARLIN_GEN_SCRIPT} ${CUDA_ARCHS_STR}
+          ${Python_EXECUTABLE} ${MARLIN_GEN_SCRIPT}
        RESULT_VARIABLE marlin_generation_result
        OUTPUT_VARIABLE marlin_generation_result
        OUTPUT_FILE ${CMAKE_CURRENT_BINARY_DIR}/marlin_generation.log
@@ -398,15 +387,15 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
                            "\nCheck the log for details: "
                            "${CMAKE_CURRENT_BINARY_DIR}/marlin_generation.log")
      else()
-        set(MARLIN_GEN_SCRIPT_HASH_AND_ARCH ${MARLIN_GEN_SCRIPT_HASH_AND_ARCH}
-            CACHE STRING "Last run Marlin generate script hash and arch" FORCE)
+        set(MARLIN_GEN_SCRIPT_HASH ${MARLIN_GEN_SCRIPT_HASH}
+            CACHE STRING "Last run Marlin generate script hash" FORCE)
        message(STATUS "Marlin generation completed successfully.")
      endif()
    else()
      message(STATUS "Marlin generation script has not changed, skipping generation.")
    endif()

-    file(GLOB MARLIN_TEMPLATE_KERNEL_SRC "csrc/quantization/gptq_marlin/sm80_kernel_*_float16.cu")
+    file(GLOB MARLIN_TEMPLATE_KERNEL_SRC "csrc/quantization/gptq_marlin/kernel_*.cu")
    set_gencode_flags_for_srcs(
      SRCS "${MARLIN_TEMPLATE_KERNEL_SRC}"
      CUDA_ARCHS "${MARLIN_ARCHS}")
@@ -414,34 +403,12 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
      set_source_files_properties(${MARLIN_TEMPLATE_KERNEL_SRC}
        PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false")
    endif()
+
    list(APPEND VLLM_EXT_SRC ${MARLIN_TEMPLATE_KERNEL_SRC})

-    file(GLOB MARLIN_TEMPLATE_BF16_KERNEL_SRC "csrc/quantization/gptq_marlin/sm80_kernel_*_bfloat16.cu")
-    set_gencode_flags_for_srcs(
-      SRCS "${MARLIN_TEMPLATE_BF16_KERNEL_SRC}"
-      CUDA_ARCHS "${MARLIN_BF16_ARCHS}")
-    if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
-      set_source_files_properties(${MARLIN_TEMPLATE_BF16_KERNEL_SRC}
-        PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false")
-    endif()
-    list(APPEND VLLM_EXT_SRC ${MARLIN_TEMPLATE_BF16_KERNEL_SRC})
-
-    if (MARLIN_FP8_ARCHS) 
-      file(GLOB MARLIN_TEMPLATE_FP8_KERNEL_SRC "csrc/quantization/gptq_marlin/sm89_kernel_*.cu")
-      set_gencode_flags_for_srcs(
-        SRCS "${MARLIN_TEMPLATE_FP8_KERNEL_SRC}"
-        CUDA_ARCHS "${MARLIN_FP8_ARCHS}")
-      if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
-        set_source_files_properties(${MARLIN_TEMPLATE_FP8_KERNEL_SRC}
-          PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false")
-      endif()
-      list(APPEND VLLM_EXT_SRC ${MARLIN_TEMPLATE_FP8_KERNEL_SRC})
-    endif()
-
    set(MARLIN_SRCS
       "csrc/quantization/marlin/sparse/marlin_24_cuda_kernel.cu"
       "csrc/quantization/gptq_marlin/gptq_marlin.cu"
-       "csrc/quantization/gptq_marlin/marlin_int4_fp8_preprocess.cu"
       "csrc/quantization/gptq_marlin/gptq_marlin_repack.cu"
       "csrc/quantization/gptq_marlin/awq_marlin_repack.cu")
    set_gencode_flags_for_srcs(
@@ -637,15 +604,12 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
    set(SRCS
      "csrc/quantization/fp4/nvfp4_quant_kernels.cu"
      "csrc/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu"
-      "csrc/quantization/fp4/nvfp4_experts_quant.cu"
-      "csrc/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu"
-      "csrc/quantization/fp4/nvfp4_blockwise_moe_kernel.cu")
+      "csrc/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu")
    set_gencode_flags_for_srcs(
      SRCS "${SRCS}"
      CUDA_ARCHS "${FP4_ARCHS}")
    list(APPEND VLLM_EXT_SRC "${SRCS}")
    list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM120=1")
-    list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM120=1")
    message(STATUS "Building NVFP4 for archs: ${FP4_ARCHS}")
  else()
    message(STATUS "Not building NVFP4 as no compatible archs were found.")
@@ -974,15 +938,8 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
    CUDA_ARCHS "${CUDA_ARCHS}")

  list(APPEND VLLM_MOE_EXT_SRC "${VLLM_MOE_WNA16_SRC}")
-  # moe marlin arches
-  # note that we always set `use_atomic_add=False` for moe marlin now,
-  # so we don't need 9.0 for bf16 atomicAdd PTX
-  cuda_archs_loose_intersection(MARLIN_MOE_ARCHS "8.0+PTX" "${CUDA_ARCHS}")
-  # moe marlin arches for fp8 input
-  # - sm80 doesn't support fp8 computation
-  # - sm90 and sm100 don't support QMMA.16832.F32.E4M3.E4M3 SAAS instruction
-  # so we only enable fp8 computation for SM89 (e.g. RTX 40x0)  and 12.0 (e.g. RTX 50x0)
-  cuda_archs_loose_intersection(MARLIN_MOE_FP8_ARCHS "8.9;12.0" "${CUDA_ARCHS}")
+  # 9.0 for latest bf16 atomicAdd PTX
+  cuda_archs_loose_intersection(MARLIN_MOE_ARCHS "8.0+PTX;9.0+PTX" "${CUDA_ARCHS}")
  if (MARLIN_MOE_ARCHS)

    #
@@ -992,18 +949,16 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
    set(MOE_MARLIN_GEN_SCRIPT
      ${CMAKE_CURRENT_SOURCE_DIR}/csrc/moe/marlin_moe_wna16/generate_kernels.py)
    file(MD5 ${MOE_MARLIN_GEN_SCRIPT} MOE_MARLIN_GEN_SCRIPT_HASH)
-    list(JOIN CUDA_ARCHS "," CUDA_ARCHS_STR)
-    set(MOE_MARLIN_GEN_SCRIPT_HASH_AND_ARCH "${MOE_MARLIN_GEN_SCRIPT_HASH}(ARCH:${CUDA_ARCHS_STR})")

-    message(STATUS "Marlin MOE generation script hash with arch: ${MOE_MARLIN_GEN_SCRIPT_HASH_AND_ARCH}")
-    message(STATUS "Last run Marlin MOE generate script hash with arch: $CACHE{MOE_MARLIN_GEN_SCRIPT_HASH_AND_ARCH}")
+    message(STATUS "Marlin MOE generation script hash: ${MOE_MARLIN_GEN_SCRIPT_HASH}")
+    message(STATUS "Last run Marlin MOE generate script hash: $CACHE{MOE_MARLIN_GEN_SCRIPT_HASH}")

-    if (NOT DEFINED CACHE{MOE_MARLIN_GEN_SCRIPT_HASH_AND_ARCH}
-        OR NOT $CACHE{MOE_MARLIN_GEN_SCRIPT_HASH_AND_ARCH} STREQUAL ${MOE_MARLIN_GEN_SCRIPT_HASH_AND_ARCH})
+    if (NOT DEFINED CACHE{MOE_MARLIN_GEN_SCRIPT_HASH}
+        OR NOT $CACHE{MOE_MARLIN_GEN_SCRIPT_HASH} STREQUAL ${MOE_MARLIN_GEN_SCRIPT_HASH})
      execute_process(
        COMMAND ${CMAKE_COMMAND} -E env
        PYTHONPATH=$PYTHONPATH
-          ${Python_EXECUTABLE} ${MOE_MARLIN_GEN_SCRIPT} ${CUDA_ARCHS_STR}
+          ${Python_EXECUTABLE} ${MOE_MARLIN_GEN_SCRIPT}
        RESULT_VARIABLE moe_marlin_generation_result
        OUTPUT_VARIABLE moe_marlin_generation_output
        OUTPUT_FILE ${CMAKE_CURRENT_BINARY_DIR}/moe_marlin_generation.log
@@ -1016,7 +971,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
                            "\nCheck the log for details: "
                            "${CMAKE_CURRENT_BINARY_DIR}/moe_marlin_generation.log")
      else()
-        set(MOE_MARLIN_GEN_SCRIPT_HASH_AND_ARCH ${MOE_MARLIN_GEN_SCRIPT_HASH_AND_ARCH}
+        set(MOE_MARLIN_GEN_SCRIPT_HASH ${MOE_MARLIN_GEN_SCRIPT_HASH}
            CACHE STRING "Last run Marlin MOE generate script hash" FORCE)
        message(STATUS "Marlin MOE generation completed successfully.")
      endif()
@@ -1024,28 +979,16 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
      message(STATUS "Marlin MOE generation script has not changed, skipping generation.")
    endif()

-    file(GLOB MARLIN_MOE_SRC "csrc/moe/marlin_moe_wna16/sm80_kernel_*.cu")
-    list(APPEND MARLIN_MOE_SRC "csrc/moe/marlin_moe_wna16/ops.cu")
+    file(GLOB MOE_WNAA16_MARLIN_SRC "csrc/moe/marlin_moe_wna16/*.cu")
    set_gencode_flags_for_srcs(
-      SRCS "${MARLIN_MOE_SRC}"
+      SRCS "${MOE_WNAA16_MARLIN_SRC}"
      CUDA_ARCHS "${MARLIN_MOE_ARCHS}")
    if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
-      set_source_files_properties(${MARLIN_MOE_SRC}
+      set_source_files_properties(${MOE_WNAA16_MARLIN_SRC}
        PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false")
    endif()
-    list(APPEND VLLM_MOE_EXT_SRC ${MARLIN_MOE_SRC})

-    if (MARLIN_MOE_FP8_ARCHS)
-      file(GLOB MARLIN_MOE_FP8_SRC "csrc/moe/marlin_moe_wna16/sm89_kernel_*.cu")
-      set_gencode_flags_for_srcs(
-        SRCS "${MARLIN_MOE_FP8_SRC}"
-        CUDA_ARCHS "${MARLIN_MOE_FP8_ARCHS}")
-      if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
-        set_source_files_properties(${MARLIN_MOE_FP8_SRC}
-          PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false")
-      endif()
-      list(APPEND VLLM_MOE_EXT_SRC ${MARLIN_MOE_FP8_SRC})
-    endif()
+    list(APPEND VLLM_MOE_EXT_SRC ${MOE_WNAA16_MARLIN_SRC})

    message(STATUS "Building Marlin MOE kernels for archs: ${MARLIN_MOE_ARCHS}")
  else()
@@ -1087,11 +1030,6 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
    WITH_SOABI)
 endif()

-# For CUDA and HIP builds also build the triton_kernels external package.
-if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
-    include(cmake/external_projects/triton_kernels.cmake)
-endif()
-
 # For CUDA we also build and ship some external projects.
 if (VLLM_GPU_LANG STREQUAL "CUDA")
    include(cmake/external_projects/flashmla.cmake)
@@ -21,7 +21,6 @@ Join us at the [PyTorch Conference, October 22-23](https://events.linuxfoundatio

 *Latest News* 🔥

- [2025/11] We hosted [vLLM Bangkok Meetup](https://luma.com/v0f647nv). We explored vLLM and LMCache inference and low-resource language adaptation with speakers from Embedded LLM, AMD, and Red Hat. Please find the meetup slides [here](https://drive.google.com/drive/folders/1H0DS57F8HQ5q3kSOSoRmucPJWL3E0A_X?usp=sharing).
 - [2025/11] We hosted [the first vLLM Europe Meetup in Zurich](https://luma.com/0gls27kb) focused on quantization, distributed inference, and reinforcement learning at scale with speakers from Mistral, IBM, and Red Hat. Please find the meetup slides [here](https://docs.google.com/presentation/d/1UC9PTLCHYXQpOmJDSFg6Sljra3iVXzc09DeEI7dnxMc/edit?usp=sharing) and recording [here](https://www.youtube.com/watch?v=6m6ZE6yVEDI)
 - [2025/11] We hosted [vLLM Beijing Meetup](https://mp.weixin.qq.com/s/xSrYXjNgr1HbCP4ExYNG1w) focusing on distributed inference and diverse accelerator support with vLLM! Please find the meetup slides [here](https://drive.google.com/drive/folders/1nQJ8ZkLSjKxvu36sSHaceVXtttbLvvu-?usp=drive_link).
 - [2025/10] We hosted [vLLM Shanghai Meetup](https://mp.weixin.qq.com/s/__xb4OyOsImz-9eAVrdlcg) focused on hands-on vLLM inference optimization! Please find the meetup slides [here](https://drive.google.com/drive/folders/1KqwjsFJLfEsC8wlDugnrR61zsWHt94Q6).
@@ -83,7 +83,7 @@ MIN_CACHE_HIT_PCT=0
 MAX_LATENCY_ALLOWED_MS=100000000000 # A very large number
 ```

-### 2. Maximize Throughput with a Latency Requirement
+#### 2. Maximize Throughput with a Latency Requirement

 - **Goal**: Find the best server parameters when P99 end-to-end latency must be below 500ms.
 - **Configuration**:
@@ -96,7 +96,7 @@ MIN_CACHE_HIT_PCT=0
 MAX_LATENCY_ALLOWED_MS=500
 ```

-### 3. Maximize Throughput with Prefix Caching and Latency Requirements
+#### 3. Maximize Throughput with Prefix Caching and Latency Requirements

 - **Goal**: Find the best server parameters assuming a 60% prefix cache hit rate and a latency requirement of 500ms.
 - **Configuration**:
@@ -620,7 +620,7 @@ def get_tokenizer(
        kwargs["use_fast"] = False
    if tokenizer_mode == "mistral":
        try:
-            from vllm.tokenizers import MistralTokenizer
+            from vllm.transformers_utils.tokenizer import MistralTokenizer
        except ImportError as e:
            raise ImportError(
                "MistralTokenizer requires vllm package.\n"
@@ -5,12 +5,11 @@ import argparse
 import asyncio
 import logging
 import os
-import time
-import uuid
-from urllib.parse import urlparse

 import aiohttp
 from quart import Quart, Response, make_response, request
+from rate_limiter import RateLimiter
+from request_queue import RequestQueue

 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -25,8 +24,26 @@ def parse_args():
    parser.add_argument(
        "--timeout",
        type=float,
-        default=6 * 60 * 60,
-        help="Timeout for backend service requests in seconds (default: 21600)",
+        default=300,
+        help="Timeout for backend service requests in seconds (default: 300)",
+    )
+    parser.add_argument(
+        "--max-concurrent",
+        type=int,
+        default=100,
+        help="Maximum concurrent requests to backend services (default: 100)",
+    )
+    parser.add_argument(
+        "--queue-size",
+        type=int,
+        default=500,
+        help="Maximum number of requests in the queue (default: 500)",
+    )
+    parser.add_argument(
+        "--rate-limit",
+        type=int,
+        default=40,
+        help="Maximum requests per second (default: 40)",
    )
    parser.add_argument(
        "--port",
@@ -37,32 +54,14 @@ def parse_args():
    parser.add_argument(
        "--prefill-url",
        type=str,
-        default="http://localhost:8100",
-        help="Prefill service base URL (protocol + host[:port])",
+        default="http://localhost:8100/v1/completions",
+        help="Prefill service endpoint URL",
    )
    parser.add_argument(
        "--decode-url",
        type=str,
-        default="http://localhost:8200",
-        help="Decode service base URL (protocol + host[:port])",
-    )
-    parser.add_argument(
-        "--kv-host",
-        type=str,
-        default="localhost",
-        help="Hostname or IP used by KV transfer (default: localhost)",
-    )
-    parser.add_argument(
-        "--prefill-kv-port",
-        type=int,
-        default=14579,
-        help="Prefill KV port (default: 14579)",
-    )
-    parser.add_argument(
-        "--decode-kv-port",
-        type=int,
-        default=14580,
-        help="Decode KV port (default: 14580)",
+        default="http://localhost:8200/v1/completions",
+        help="Decode service endpoint URL",
    )

    return parser.parse_args()
@@ -74,129 +73,70 @@ def main():

    # Initialize configuration using command line parameters
    AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=args.timeout)
+    MAX_CONCURRENT_REQUESTS = args.max_concurrent
+    REQUEST_QUEUE_SIZE = args.queue_size
+    RATE_LIMIT = args.rate_limit
    PREFILL_SERVICE_URL = args.prefill_url
    DECODE_SERVICE_URL = args.decode_url
    PORT = args.port

-    PREFILL_KV_ADDR = f"{args.kv_host}:{args.prefill_kv_port}"
-    DECODE_KV_ADDR = f"{args.kv_host}:{args.decode_kv_port}"
-
-    logger.info(
-        "Proxy resolved KV addresses -> prefill: %s, decode: %s",
-        PREFILL_KV_ADDR,
-        DECODE_KV_ADDR,
-    )
-
    app = Quart(__name__)

-    # Attach the configuration object to the application instance so helper
-    # coroutines can read the resolved backend URLs and timeouts without using
-    # globals.
+    # Initialize the rate limiter and request queue
+    rate_limiter = RateLimiter(RATE_LIMIT)
+    request_queue = RequestQueue(MAX_CONCURRENT_REQUESTS, REQUEST_QUEUE_SIZE)
+
+    # Attach the configuration object to the application instance
    app.config.update(
        {
            "AIOHTTP_TIMEOUT": AIOHTTP_TIMEOUT,
+            "rate_limiter": rate_limiter,
+            "request_queue": request_queue,
            "PREFILL_SERVICE_URL": PREFILL_SERVICE_URL,
            "DECODE_SERVICE_URL": DECODE_SERVICE_URL,
-            "PREFILL_KV_ADDR": PREFILL_KV_ADDR,
-            "DECODE_KV_ADDR": DECODE_KV_ADDR,
        }
    )

-    def _normalize_base_url(url: str) -> str:
-        """Remove any trailing slash so path joins behave predictably."""
-        return url.rstrip("/")
+    # Start queue processing on app startup
+    @app.before_serving
+    async def startup():
+        """Start request processing task when app starts serving"""
+        asyncio.create_task(request_queue.process())

-    def _get_host_port(url: str) -> str:
-        """Return the hostname:port portion for logging and KV headers."""
-        parsed = urlparse(url)
-        host = parsed.hostname or "localhost"
-        port = parsed.port
-        if port is None:
-            port = 80 if parsed.scheme == "http" else 443
-        return f"{host}:{port}"
+    async def forward_request(url, data):
+        """Forward request to backend service with rate limiting and error handling"""
+        headers = {"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}"}

-    PREFILL_BASE = _normalize_base_url(PREFILL_SERVICE_URL)
-    DECODE_BASE = _normalize_base_url(DECODE_SERVICE_URL)
-    KV_TARGET = _get_host_port(DECODE_SERVICE_URL)
-
-    def _build_headers(request_id: str) -> dict[str, str]:
-        """Construct the headers expected by vLLM's P2P disagg connector."""
-        headers: dict[str, str] = {"X-Request-Id": request_id, "X-KV-Target": KV_TARGET}
-        api_key = os.environ.get("OPENAI_API_KEY")
-        if api_key:
-            headers["Authorization"] = f"Bearer {api_key}"
-        return headers
-
-    async def _run_prefill(
-        request_path: str,
-        payload: dict,
-        headers: dict[str, str],
-        request_id: str,
-    ):
-        url = f"{PREFILL_BASE}{request_path}"
-        start_ts = time.perf_counter()
-        logger.info("[prefill] start request_id=%s url=%s", request_id, url)
-        try:
-            async with (
-                aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session,
-                session.post(url=url, json=payload, headers=headers) as resp,
-            ):
-                if resp.status != 200:
-                    error_text = await resp.text()
-                    raise RuntimeError(
-                        f"Prefill backend error {resp.status}: {error_text}"
-                    )
-                await resp.read()
-                logger.info(
-                    "[prefill] done request_id=%s status=%s elapsed=%.2fs",
-                    request_id,
-                    resp.status,
-                    time.perf_counter() - start_ts,
-                )
-        except asyncio.TimeoutError as exc:
-            raise RuntimeError(f"Prefill service timeout at {url}") from exc
-        except aiohttp.ClientError as exc:
-            raise RuntimeError(f"Prefill service unavailable at {url}") from exc
-
-    async def _stream_decode(
-        request_path: str,
-        payload: dict,
-        headers: dict[str, str],
-        request_id: str,
-    ):
-        url = f"{DECODE_BASE}{request_path}"
-        # Stream tokens from the decode service once the prefill stage has
-        # materialized KV caches on the target workers.
-        logger.info("[decode] start request_id=%s url=%s", request_id, url)
-        try:
-            async with (
-                aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session,
-                session.post(url=url, json=payload, headers=headers) as resp,
-            ):
-                if resp.status != 200:
-                    error_text = await resp.text()
-                    logger.error(
-                        "Decode backend error %s - %s", resp.status, error_text
-                    )
-                    err_msg = (
-                        '{"error": "Decode backend error ' + str(resp.status) + '"}'
-                    )
-                    yield err_msg.encode()
-                    return
-                logger.info(
-                    "[decode] streaming response request_id=%s status=%s",
-                    request_id,
-                    resp.status,
-                )
-                async for chunk_bytes in resp.content.iter_chunked(1024):
-                    yield chunk_bytes
-                logger.info("[decode] finished streaming request_id=%s", request_id)
-        except asyncio.TimeoutError:
-            logger.error("Decode service timeout at %s", url)
-            yield b'{"error": "Decode service timeout"}'
-        except aiohttp.ClientError as exc:
-            logger.error("Decode service error at %s: %s", url, exc)
-            yield b'{"error": "Decode service unavailable"}'
+        # Use rate limiter as context manager
+        async with (
+            rate_limiter,
+            aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session,
+        ):
+            try:
+                async with session.post(
+                    url=url, json=data, headers=headers
+                ) as response:
+                    if response.status == 200:
+                        # Stream response chunks
+                        async for chunk_bytes in response.content.iter_chunked(1024):
+                            yield chunk_bytes
+                    else:
+                        # Handle backend service errors
+                        error_text = await response.text()
+                        logger.error(
+                            "Backend service error: %s - %s",
+                            response.status,
+                            error_text,
+                        )
+                        yield b'{"error": "Backend service error"}'
+            except aiohttp.ClientError as e:
+                # Handle connection errors
+                logger.error("Connection error to %s: %s", url, str(e))
+                yield b'{"error": "Service unavailable"}'
+            except asyncio.TimeoutError:
+                # Handle timeout errors
+                logger.error("Timeout connecting to %s", url)
+                yield b'{"error": "Service timeout"}'

    async def process_request():
        """Process a single request through prefill and decode stages"""
@@ -206,27 +146,13 @@ def main():
            # Create prefill request (max_tokens=1)
            prefill_request = original_request_data.copy()
            prefill_request["max_tokens"] = 1
-            if "max_completion_tokens" in prefill_request:
-                prefill_request["max_completion_tokens"] = 1

            # Execute prefill stage
-            # The request id encodes both KV socket addresses so the backend can
-            # shuttle tensors directly via NCCL once the prefill response
-            # completes.
-            request_id = (
-                f"___prefill_addr_{PREFILL_KV_ADDR}___decode_addr_"
-                f"{DECODE_KV_ADDR}_{uuid.uuid4().hex}"
-            )
-
-            headers = _build_headers(request_id)
-            await _run_prefill(request.path, prefill_request, headers, request_id)
+            async for _ in forward_request(PREFILL_SERVICE_URL, prefill_request):
+                continue

            # Execute decode stage and stream response
-            # Pass the unmodified user request so the decode phase can continue
-            # sampling with the already-populated KV cache.
-            generator = _stream_decode(
-                request.path, original_request_data, headers, request_id
-            )
+            generator = forward_request(DECODE_SERVICE_URL, original_request_data)
            response = await make_response(generator)
            response.timeout = None  # Disable timeout for streaming response
            return response
@@ -242,10 +168,23 @@ def main():
    @app.route("/v1/completions", methods=["POST"])
    async def handle_request():
        """Handle incoming API requests with concurrency and rate limiting"""
+        # Create task for request processing
+        task = asyncio.create_task(process_request())
+
+        # Enqueue request or reject if queue is full
+        if not await request_queue.enqueue(task):
+            return Response(
+                response=b'{"error": "Server busy, try again later"}',
+                status=503,
+                content_type="application/json",
+            )
+
        try:
-            return await process_request()
+            # Return the response from the processing task
+            return await task
        except asyncio.CancelledError:
-            logger.warning("Request cancelled")
+            # Handle task cancellation (timeout or queue full)
+            logger.warning("Request cancelled due to timeout or queue full")
            return Response(
                response=b'{"error": "Request cancelled"}',
                status=503,
@@ -255,8 +255,8 @@ def bench_run(
        torch.cuda.synchronize()

        # Timing
-        start_event = torch.Event(enable_timing=True)
-        end_event = torch.Event(enable_timing=True)
+        start_event = torch.cuda.Event(enable_timing=True)
+        end_event = torch.cuda.Event(enable_timing=True)

        latencies = []
        for _ in range(num_iters):
@@ -237,7 +237,6 @@ def marlin_create_bench_fn(bt: BenchmarkTensors) -> Callable:
            b_q_weight=w_q,
            b_bias=None,
            b_scales=w_s,
-            a_scales=None,
            global_scale=None,
            b_zeros=w_zp,
            g_idx=g_idx,
@@ -263,7 +263,7 @@ def bench_run(

    results.append(
        benchmark.Timer(
-            stmt="output = gptq_marlin_gemm(a, None, marlin_q_w, marlin_s, None, marlin_s2, marlin_zp, marlin_g_idx, marlin_sort_indices, marlin_workspace.scratch, quant_type, size_m, size_n, size_k, is_k_full, False, False, False)",  # noqa: E501
+            stmt="output = gptq_marlin_gemm(a, None, marlin_q_w, marlin_s, marlin_s2, marlin_zp, marlin_g_idx, marlin_sort_indices, marlin_workspace.scratch, quant_type, size_m, size_n, size_k, is_k_full, False, False, False)",  # noqa: E501
            globals=globals,
            label=label,
            sub_label=sub_label,
@@ -273,7 +273,7 @@ def bench_run(

    results.append(
        benchmark.Timer(
-            stmt="output = gptq_marlin_gemm(a, None, marlin_q_w, marlin_s, None, marlin_s2, marlin_zp, marlin_g_idx, marlin_sort_indices, marlin_workspace.scratch, quant_type, size_m, size_n, size_k, is_k_full, False, True, False)",  # noqa: E501
+            stmt="output = gptq_marlin_gemm(a, None, marlin_q_w, marlin_s, marlin_s2, marlin_zp, marlin_g_idx, marlin_sort_indices, marlin_workspace.scratch, quant_type, size_m, size_n, size_k, is_k_full, False, True, False)",  # noqa: E501
            globals=globals,
            label=label,
            sub_label=sub_label,
@@ -185,8 +185,8 @@ def benchmark_config(
        graph.replay()
    torch.cuda.synchronize()

-    start_event = torch.Event(enable_timing=True)
-    end_event = torch.Event(enable_timing=True)
+    start_event = torch.cuda.Event(enable_timing=True)
+    end_event = torch.cuda.Event(enable_timing=True)

    latencies: list[float] = []
    for i in range(num_iters):
@@ -105,8 +105,8 @@ def benchmark_permute(
        graph.replay()
    torch.cuda.synchronize()

-    start_event = torch.Event(enable_timing=True)
-    end_event = torch.Event(enable_timing=True)
+    start_event = torch.cuda.Event(enable_timing=True)
+    end_event = torch.cuda.Event(enable_timing=True)

    latencies: list[float] = []
    for i in range(num_iters):
@@ -241,8 +241,8 @@ def benchmark_unpermute(
        graph.replay()
    torch.cuda.synchronize()

-    start_event = torch.Event(enable_timing=True)
-    end_event = torch.Event(enable_timing=True)
+    start_event = torch.cuda.Event(enable_timing=True)
+    end_event = torch.cuda.Event(enable_timing=True)

    latencies: list[float] = []
    for i in range(num_iters):
@@ -6,7 +6,7 @@
 #
 # The CSV file (named with current date/time) contains these columns:
 # model_name, tp_size, num_tokens, num_heads, num_kv_heads, head_dim, max_position,
-# is_neox_style, rope_parameters, dtype, torch_mean, torch_median, torch_p99,
+# rope_theta, is_neox_style, rope_scaling, dtype, torch_mean, torch_median, torch_p99,
 # torch_min, torch_max, triton_mean, triton_median, triton_p99, triton_min, triton_max,
 # speedup
 #
@@ -86,8 +86,9 @@ def benchmark_mrope(
    num_heads: int,
    num_kv_heads: int,
    max_position: int = 8192,
+    rope_theta: float = 10000,
    is_neox_style: bool = True,
-    rope_parameters: dict[str, Any] | None = None,
+    rope_scaling: dict[str, Any] = None,
    dtype: torch.dtype = torch.bfloat16,
    seed: int = 0,
    warmup_iter: int = 10,
@@ -101,8 +102,9 @@ def benchmark_mrope(
        head_size=head_dim,
        rotary_dim=head_dim,
        max_position=max_position,
+        base=rope_theta,
        is_neox_style=is_neox_style,
-        rope_parameters=rope_parameters,
+        rope_scaling=rope_scaling,
        dtype=dtype,
    ).to(device=device)

@@ -201,8 +203,9 @@ def benchmark_mrope(
            num_kv_heads,
            head_dim,
            max_position,
+            rope_theta,
            is_neox_style,
-            str(rope_parameters),
+            str(rope_scaling),
            str(dtype).split(".")[-1],
            torch_stats["mean"],
            torch_stats["median"],
@@ -252,8 +255,9 @@ if __name__ == "__main__":
            "num_kv_heads",
            "head_dim",
            "max_position",
+            "rope_theta",
            "is_neox_style",
-            "rope_parameters",
+            "rope_scaling",
            "dtype",
            "torch_mean",
            "torch_median",
@@ -299,7 +303,7 @@ if __name__ == "__main__":
                q_size = num_heads * head_dim
                kv_size = num_kv_heads * head_dim
                is_neox_style = True
-                rope_parameters = config.rope_parameters
+                rope_theta = config.rope_theta
                max_position = config.max_position_embeddings

                for num_tokens in num_tokens_list:
@@ -311,8 +315,9 @@ if __name__ == "__main__":
                        num_heads=num_heads,
                        num_kv_heads=num_kv_heads,
                        max_position=max_position,
+                        rope_theta=rope_theta,
                        is_neox_style=is_neox_style,
-                        rope_parameters=rope_parameters,
+                        rope_scaling=config.rope_scaling,
                        dtype=getattr(torch, args.dtype),
                        seed=args.seed,
                        warmup_iter=args.warmup_iter,
@@ -30,8 +30,8 @@ def _time_cuda(
        fn()
    torch.cuda.synchronize()

-    start = torch.Event(enable_timing=True)
-    end = torch.Event(enable_timing=True)
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)

    start.record()
    for _ in range(bench_iters):
@@ -253,8 +253,8 @@ def benchmark(
        )
    torch.cuda.synchronize()

-    start_event = torch.Event(enable_timing=True)
-    end_event = torch.Event(enable_timing=True)
+    start_event = torch.cuda.Event(enable_timing=True)
+    end_event = torch.cuda.Event(enable_timing=True)

    # Benchmark
    latencies: list[float] = []
@@ -127,8 +127,8 @@ def benchmark_decode(

    def time_fn(fn, warmup=10, trials=20):
        torch.cuda.synchronize()
-        start = torch.Event(enable_timing=True)
-        end = torch.Event(enable_timing=True)
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
        times = []
        for i in range(warmup):
            fn()
@@ -139,8 +139,8 @@ def benchmark_prefill(

    def time_fn(fn, warmup=10, trials=20):
        torch.cuda.synchronize()
-        start = torch.Event(enable_timing=True)
-        end = torch.Event(enable_timing=True)
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
        times = []
        for i in range(warmup):
            fn()
@@ -183,8 +183,8 @@ def benchmark_config(
        run()
    torch.cuda.synchronize()

-    start_event = torch.Event(enable_timing=True)
-    end_event = torch.Event(enable_timing=True)
+    start_event = torch.cuda.Event(enable_timing=True)
+    end_event = torch.cuda.Event(enable_timing=True)

    latencies: list[float] = []
    for i in range(num_iters):
@@ -2,7 +2,7 @@

 This directory includes benchmarks between DeepSeek's DeepGEMM block fp8 kernels against vLLM's existing triton and CUTLASS-based kernels.

-Currently, this just includes dense GEMMs and only works on Hopper GPUs.
+Currently this just includes dense GEMMs and only works on Hopper GPUs.

 ## Setup

@@ -55,10 +55,6 @@ output_num_chunks  166.0    99.01   11.80    79.00    90.00    98.00   108.75
 ----------------------------------------------------------------------------------------------------
 ```

-If you run with `--warmup-step`, the summary will also include `warmup_runtime_sec`
-and `total_runtime_incl_warmup_sec` (while `runtime_sec` continues to reflect the
-benchmark-only runtime so the reported throughput stays comparable).
-
 ### JSON configuration file for synthetic conversations generation

 The input flag `--input-file` is used to determine the input conversations for the benchmark.<br/>
@@ -1076,7 +1076,6 @@ def process_statistics(
    verbose: bool,
    gen_conv_args: GenConvArgs | None = None,
    excel_output: bool = False,
-    warmup_runtime_sec: float | None = None,
 ) -> None:
    if len(client_metrics) == 0:
        logger.info("No samples to process")
@@ -1170,13 +1169,8 @@ def process_statistics(
        # Convert milliseconds to seconds
        runtime_sec = runtime_sec / 1000.0
        requests_per_sec = float(len(df)) / runtime_sec
-        params = {
-            "runtime_sec": runtime_sec,
-            "requests_per_sec": requests_per_sec,
-        }
-        if warmup_runtime_sec is not None:
-            params["warmup_runtime_sec"] = warmup_runtime_sec
-            params["total_runtime_incl_warmup_sec"] = runtime_sec + warmup_runtime_sec
+
+        params = {"runtime_sec": runtime_sec, "requests_per_sec": requests_per_sec}

        # Generate a summary of relevant metrics (and drop irrelevant data)
        df = df.drop(columns=exclude).describe(percentiles=percentiles).transpose()
@@ -1558,8 +1552,6 @@ async def main() -> None:
        url=args.url, num_clients=args.num_clients, early_stop=not args.no_early_stop
    )

-    warmup_runtime_sec: float | None = None
-
    # Warm-up step
    if args.warmup_step:
        # Only send a single user prompt from every conversation.
@@ -1574,56 +1566,26 @@ async def main() -> None:
        # all clients should finish their work before exiting
        warmup_bench_args = bench_args._replace(early_stop=False)

-        logger.info("%sWarmup start%s", Color.PURPLE, Color.RESET)
-        warmup_start_ns = time.perf_counter_ns()
+        logger.info(f"{Color.PURPLE}Warmup start{Color.RESET}")
        conversations, _ = await main_mp(
            warmup_client_args, req_args, warmup_bench_args, tokenizer, conversations
        )
-        warmup_runtime_sec = nanosec_to_sec(time.perf_counter_ns() - warmup_start_ns)
-        logger.info(
-            "%sWarmup runtime: %.3f sec (%.3f ms)%s",
-            Color.PURPLE,
-            warmup_runtime_sec,
-            warmup_runtime_sec * 1000,
-            Color.RESET,
-        )
-        logger.info("%sWarmup done%s", Color.PURPLE, Color.RESET)
+        logger.info(f"{Color.PURPLE}Warmup done{Color.RESET}")

    # Run the benchmark
-    benchmark_start_ns = time.perf_counter_ns()
+    start_time = time.perf_counter_ns()
    client_convs, client_metrics = await main_mp(
        client_args, req_args, bench_args, tokenizer, conversations
    )
-    benchmark_runtime_sec = nanosec_to_sec(time.perf_counter_ns() - benchmark_start_ns)
+    total_runtime_ms = nanosec_to_millisec(time.perf_counter_ns() - start_time)

    # Calculate requests per second
-    requests_per_sec = len(client_metrics) / benchmark_runtime_sec
-    benchmark_runtime_ms = benchmark_runtime_sec * 1000.0
+    total_runtime_sec = total_runtime_ms / 1000.0
+    rps = len(client_metrics) / total_runtime_sec
    logger.info(
-        "%sAll clients finished, benchmark runtime: %.3f sec (%.3f ms), "
-        "requests per second: %.3f%s",
-        Color.GREEN,
-        benchmark_runtime_sec,
-        benchmark_runtime_ms,
-        requests_per_sec,
-        Color.RESET,
+        f"{Color.GREEN}All clients finished, total runtime: {total_runtime_sec:.3f} sec"
+        f" ({total_runtime_ms:.3f} ms), requests per second: {rps:.3f}{Color.RESET}"
    )
-    if warmup_runtime_sec is not None:
-        total_runtime_sec = benchmark_runtime_sec + warmup_runtime_sec
-        logger.info(
-            "%sWarmup runtime: %.3f sec (%.3f ms)%s",
-            Color.GREEN,
-            warmup_runtime_sec,
-            warmup_runtime_sec * 1000,
-            Color.RESET,
-        )
-        logger.info(
-            "%sTotal runtime (including warmup): %.3f sec (%.3f ms)%s",
-            Color.GREEN,
-            total_runtime_sec,
-            total_runtime_sec * 1000,
-            Color.RESET,
-        )

    # Benchmark parameters
    params = {
@@ -1648,7 +1610,6 @@ async def main() -> None:
        verbose=args.verbose,
        gen_conv_args=gen_conv_args,
        excel_output=args.excel_output,
-        warmup_runtime_sec=warmup_runtime_sec,
    )

    if args.output_file is not None:
@@ -375,7 +375,6 @@ set(VLLM_EXT_SRC
 if (AVX512_FOUND AND NOT AVX512_DISABLED)
    set(VLLM_EXT_SRC
        "csrc/cpu/shm.cpp"
-        "csrc/cpu/cpu_wna16.cpp"
        ${VLLM_EXT_SRC})
    if (ENABLE_AVX512BF16 AND ENABLE_AVX512VNNI)
        set(VLLM_EXT_SRC
@@ -1,53 +0,0 @@
-# Install OpenAI triton_kernels from https://github.com/triton-lang/triton/tree/main/python/triton_kernels
-
-set(DEFAULT_TRITON_KERNELS_TAG "v3.5.0")
-
-# Set TRITON_KERNELS_SRC_DIR for use with local development with vLLM. We expect TRITON_KERNELS_SRC_DIR to
-# be directly set to the triton_kernels python directory. 
-if (DEFINED ENV{TRITON_KERNELS_SRC_DIR})
-  message(STATUS "[triton_kernels] Fetch from $ENV{TRITON_KERNELS_SRC_DIR}")
-  FetchContent_Declare(
-          triton_kernels
-          SOURCE_DIR $ENV{TRITON_KERNELS_SRC_DIR}
-  )
-
-else()
-  set(TRITON_GIT "https://github.com/triton-lang/triton.git")
-  message (STATUS "[triton_kernels] Fetch from ${TRITON_GIT}:${DEFAULT_TRITON_KERNELS_TAG}")
-  FetchContent_Declare(
-          triton_kernels
-          # TODO (varun) : Fetch just the triton_kernels directory from Triton
-          GIT_REPOSITORY https://github.com/triton-lang/triton.git
-          GIT_TAG ${DEFAULT_TRITON_KERNELS_TAG}
-          GIT_PROGRESS TRUE
-          SOURCE_SUBDIR python/triton_kernels/triton_kernels
-  )
-endif()
-
-# Fetch content 
-FetchContent_MakeAvailable(triton_kernels)
-
-if (NOT triton_kernels_SOURCE_DIR)
-  message (FATAL_ERROR "[triton_kernels] Cannot resolve triton_kernels_SOURCE_DIR")
-endif()
-
-if (DEFINED ENV{TRITON_KERNELS_SRC_DIR})
-  set(TRITON_KERNELS_PYTHON_DIR "${triton_kernels_SOURCE_DIR}/")
-else()
-  set(TRITON_KERNELS_PYTHON_DIR "${triton_kernels_SOURCE_DIR}/python/triton_kernels/triton_kernels/")
-endif()
-
-message (STATUS "[triton_kernels] triton_kernels is available at ${TRITON_KERNELS_PYTHON_DIR}")
-
-add_custom_target(triton_kernels)
-
-# Ensure the vllm/third_party directory exists before installation
-install(CODE "file(MAKE_DIRECTORY \"\${CMAKE_INSTALL_PREFIX}/vllm/third_party/triton_kernels\")")
-
-## Copy .py files to install directory.
-install(DIRECTORY
-        ${TRITON_KERNELS_PYTHON_DIR}
-        DESTINATION 
-        vllm/third_party/triton_kernels/
-        COMPONENT triton_kernels
-        FILES_MATCHING PATTERN "*.py")
@@ -38,7 +38,7 @@ else()
  FetchContent_Declare(
          vllm-flash-attn
          GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
-          GIT_TAG 86f8f157cf82aa2342743752b97788922dd7de43
+          GIT_TAG 58e0626a692f09241182582659e3bf8f16472659
          GIT_PROGRESS TRUE
          # Don't share the vllm-flash-attn build between build types
          BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
@@ -495,13 +495,7 @@ function (define_extension_target MOD_NAME)
    set(SOABI_KEYWORD "")
  endif()

-  run_python(IS_FREETHREADED_PYTHON
-    "import sysconfig; print(1 if sysconfig.get_config_var(\"Py_GIL_DISABLED\") else 0)"
-    "Failed to determine whether interpreter is free-threaded")
-
-  # Free-threaded Python doesn't yet support the stable ABI (see PEP 803/809),
-  # so avoid using the stable ABI under free-threading only.
-  if (ARG_USE_SABI AND NOT IS_FREETHREADED_PYTHON)
+  if (ARG_USE_SABI)
    Python_add_library(${MOD_NAME} MODULE USE_SABI ${ARG_USE_SABI} ${SOABI_KEYWORD} "${ARG_SOURCES}")
  else()
    Python_add_library(${MOD_NAME} MODULE ${SOABI_KEYWORD} "${ARG_SOURCES}")
@@ -16,8 +16,7 @@ __global__ void merge_attn_states_kernel(
    scalar_t* output, float* output_lse, const scalar_t* prefix_output,
    const float* prefix_lse, const scalar_t* suffix_output,
    const float* suffix_lse, const uint num_tokens, const uint num_heads,
-    const uint head_size, const uint prefix_head_stride,
-    const uint output_head_stride) {
+    const uint head_size) {
  using pack_128b_t = uint4;
  const uint pack_size = 16 / sizeof(scalar_t);
  const uint threads_per_head = head_size / pack_size;
@@ -35,13 +34,11 @@ __global__ void merge_attn_states_kernel(
  const uint head_idx = token_head_idx % num_heads;

  const uint pack_offset = pack_idx * pack_size;  // (0~15)*8, etc.
-  const uint src_head_offset = token_idx * num_heads * prefix_head_stride +
-                               head_idx * prefix_head_stride;
-  const uint dst_head_offset = token_idx * num_heads * output_head_stride +
-                               head_idx * output_head_stride;
-  const scalar_t* prefix_head_ptr = prefix_output + src_head_offset;
-  const scalar_t* suffix_head_ptr = suffix_output + src_head_offset;
-  scalar_t* output_head_ptr = output + dst_head_offset;
+  const uint head_offset =
+      token_idx * num_heads * head_size + head_idx * head_size;
+  const scalar_t* prefix_head_ptr = prefix_output + head_offset;
+  const scalar_t* suffix_head_ptr = suffix_output + head_offset;
+  scalar_t* output_head_ptr = output + head_offset;

  float p_lse = prefix_lse[head_idx * num_tokens + token_idx];
  float s_lse = suffix_lse[head_idx * num_tokens + token_idx];
@@ -143,7 +140,7 @@ __global__ void merge_attn_states_kernel(
            reinterpret_cast<float*>(prefix_lse.data_ptr()),                \
            reinterpret_cast<scalar_t*>(suffix_output.data_ptr()),          \
            reinterpret_cast<float*>(suffix_lse.data_ptr()), num_tokens,    \
-            num_heads, head_size, prefix_head_stride, output_head_stride);  \
+            num_heads, head_size);                                          \
  }

 /*@brief Merges the attention states from prefix and suffix
@@ -169,11 +166,17 @@ void merge_attn_states_launcher(torch::Tensor& output,
  const uint num_tokens = output.size(0);
  const uint num_heads = output.size(1);
  const uint head_size = output.size(2);
-  const uint prefix_head_stride = prefix_output.stride(1);
-  const uint output_head_stride = output.stride(1);
  const uint pack_size = 16 / sizeof(scalar_t);
  TORCH_CHECK(head_size % pack_size == 0,
              "headsize must be multiple of pack_size:", pack_size);
+  TORCH_CHECK(output.stride(-2) == head_size && output.stride(-1) == 1,
+              "output heads must be contiguous in memory");
+  TORCH_CHECK(
+      prefix_output.stride(-2) == head_size && prefix_output.stride(-1) == 1,
+      "prefix_output heads must be contiguous in memory");
+  TORCH_CHECK(
+      suffix_output.stride(-2) == head_size && suffix_output.stride(-1) == 1,
+      "suffix_output heads must be contiguous in memory");
  float* output_lse_ptr = nullptr;
  if (output_lse.has_value()) {
    output_lse_ptr = output_lse.value().data_ptr<float>();
@@ -41,12 +41,11 @@ void convert_fp8(torch::Tensor& dst_cache, torch::Tensor& src_cache,
                 const double scale, const std::string& kv_cache_dtype);

 void gather_and_maybe_dequant_cache(
-    torch::Tensor const& src_cache,     // [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...]
-    torch::Tensor const& dst,           // [TOT_TOKENS, ENTRIES...]
-    torch::Tensor const& block_table,   // [BATCH, BLOCK_INDICES]
-    torch::Tensor const& cu_seq_lens,   // [BATCH+1]
-    torch::Tensor const& token_to_seq,  // [MAX_TOKEN_ACROSS_CHUNKS]
-    int64_t num_tokens, const std::string& kv_cache_dtype,
+    torch::Tensor const& src_cache,    // [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...]
+    torch::Tensor const& dst,          // [TOT_TOKENS, ENTRIES...]
+    torch::Tensor const& block_table,  // [BATCH, BLOCK_INDICES]
+    torch::Tensor const& cu_seq_lens,  // [BATCH+1]
+    int64_t batch_size, const std::string& kv_cache_dtype,
    torch::Tensor const& scale,
    std::optional<torch::Tensor> seq_starts = std::nullopt);

@@ -552,11 +552,7 @@ __global__ void indexer_k_quant_and_cache_kernel(
 #ifndef USE_ROCM
  __syncwarp();
 #endif
-#if defined(__gfx942__)
-  float scale = fmaxf(amax, 1e-4) / 224.0f;
-#else
  float scale = fmaxf(amax, 1e-4) / 448.0f;
-#endif
  if (use_ue8m0) {
    scale = exp2f(ceilf(log2f(scale)));
  }
@@ -905,80 +901,87 @@ void convert_fp8(torch::Tensor& dst_cache, torch::Tensor& src_cache,
 namespace vllm {

 // grid is launched with dimensions (batch, num_splits)
-template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt,
-          int ENTRY_SIZE, int CTA_SIZE>
+template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
 __global__ void gather_and_maybe_dequant_cache(
-    const cache_t* __restrict__ src_cache,     // [NUM_BLOCKS, BLOCK_SIZE,
-                                               // ENTRIES...]
-    scalar_t* __restrict__ dst,                // [TOT_TOKENS, ENTRIES...]
-    const int32_t* __restrict__ block_table,   // [BATCH, BLOCK_INDICES]
-    const int32_t* __restrict__ cu_seq_lens,   // [BATCH+1]
-    const int32_t* __restrict__ token_to_seq,  // [MAX_TOKEN_ACROSS_CHUNK]
-    const int32_t num_tokens, const int32_t block_size,
+    const cache_t* __restrict__ src_cache,    // [NUM_BLOCKS, BLOCK_SIZE,
+                                              // ENTRIES...]
+    scalar_t* __restrict__ dst,               // [TOT_TOKENS, ENTRIES...]
+    const int32_t* __restrict__ block_table,  // [BATCH, BLOCK_INDICES]
+    const int32_t* __restrict__ cu_seq_lens,  // [BATCH+1]
+    const int32_t block_size, const int32_t entry_size,
    const int64_t block_table_stride, const int64_t cache_block_stride,
    const int64_t cache_entry_stride, const int64_t dst_entry_stride,
    const float* __restrict__ scale,
    const int32_t* __restrict__ seq_starts) {  // Optional: starting offsets per
                                               // batch
-  constexpr int vec_size = sizeof(float4) / sizeof(scalar_t);
-  using ltype = vllm::vec_n_t<cache_t, vec_size>;
-  using stype = vllm::vec_n_t<scalar_t, vec_size>;
-  // We are adding this for code readability which will be optimized out when
-  // build in release.
-  assert(CTA_SIZE == blockDim.x);

-#pragma unroll
-  for (int token_id = blockIdx.x; token_id < num_tokens;
-       token_id += gridDim.x) {
-    int64_t batch_id = token_to_seq[token_id];
-    int64_t batch_start = cu_seq_lens[batch_id];
-    int64_t batch_end = cu_seq_lens[batch_id + 1];
-    int32_t batch_offset = token_id - batch_start;
+  const int64_t bid = blockIdx.x;  // Batch ID
+  const int32_t num_splits = gridDim.y;
+  const int32_t split = blockIdx.y;
+  const int32_t seq_start = cu_seq_lens[bid];
+  const int32_t seq_end = cu_seq_lens[bid + 1];
+  const int32_t seq_len = seq_end - seq_start;
+  const int32_t tot_blocks = cuda_utils::ceil_div(seq_len, block_size);
+  const int32_t split_blocks = cuda_utils::ceil_div(tot_blocks, num_splits);

-    if (token_id >= batch_end) return;
-    int32_t offset = 0;
-    if (seq_starts != nullptr) {
-      offset = seq_starts[batch_id];
-    }
-    batch_offset += offset;
-    int32_t block_table_id = batch_offset / block_size;
-    int32_t slot_id = batch_offset % block_size;
-    int32_t block_table_offset = batch_id * block_table_stride + block_table_id;
-    int32_t block_id = block_table[block_table_offset];
-    int64_t cache_offset =
-        block_id * cache_block_stride + slot_id * cache_entry_stride;
-    constexpr int32_t vec_iter_cnt = ENTRY_SIZE / vec_size;
-    scalar_t* dst_ = dst + token_id * dst_entry_stride;
-    cache_t* src_ = const_cast<cache_t*>(src_cache) + cache_offset;
+  const int32_t split_start = split * split_blocks;
+  const int32_t split_end = min((split + 1) * split_blocks, tot_blocks);

-#pragma unroll
-    for (int idx = threadIdx.x; idx < vec_iter_cnt; idx += CTA_SIZE) {
+  const bool is_active_split = (split_start < tot_blocks);
+  const bool is_last_split = (split_end == tot_blocks);
+
+  if (!is_active_split) return;
+
+  int32_t full_blocks_end = split_end;
+  int32_t partial_block_size = 0;
+
+  // Adjust the pointer for the block_table for this batch.
+  // If seq_starts is provided, compute an offset based on (seq_starts[bid] /
+  // page_size)
+  const int32_t batch_offset = bid * block_table_stride;
+  int32_t offset = 0;
+  if (seq_starts != nullptr) {
+    offset = seq_starts[bid] / block_size;
+  }
+  const int32_t* batch_block_table = block_table + batch_offset + offset;
+
+  // Adjust dst pointer based on the cumulative sequence lengths.
+  dst += seq_start * dst_entry_stride;
+
+  if (is_last_split) {
+    partial_block_size = seq_len % block_size;
+    if (partial_block_size) full_blocks_end -= 1;
+  }
+
+  auto copy_entry = [&](const cache_t* __restrict__ _src,
+                        scalar_t* __restrict__ _dst) {
+    for (int i = threadIdx.x; i < entry_size; i += blockDim.x) {
      if constexpr (kv_dt == Fp8KVCacheDataType::kAuto) {
-        reinterpret_cast<stype*>(dst_)[idx] =
-            static_cast<stype>(reinterpret_cast<ltype*>(src_)[idx]);
+        _dst[i] = static_cast<scalar_t>(_src[i]);
      } else {
-        ltype loaded_val = reinterpret_cast<ltype*>(src_)[idx];
-        stype store_val;
-#pragma unroll
-        for (int j = 0; j < vec_size; ++j) {
-          store_val.val[j] = fp8::scaled_convert<scalar_t, cache_t, kv_dt>(
-              loaded_val.val[j], *scale);
-        }
-        reinterpret_cast<stype*>(dst_)[idx] = store_val;
+        _dst[i] =
+            fp8::scaled_convert<scalar_t, cache_t, kv_dt>(_src[i], *scale);
      }
    }
-    // process tail
-    constexpr int32_t tail_cnt = ENTRY_SIZE % vec_size;
-    dst_ = dst_ + ENTRY_SIZE - tail_cnt;
-    src_ = src_ + ENTRY_SIZE - tail_cnt;
-#pragma unroll
-    for (int idx = threadIdx.x; idx < tail_cnt; idx += CTA_SIZE) {
-      if constexpr (kv_dt == Fp8KVCacheDataType::kAuto) {
-        dst_[idx] = static_cast<scalar_t>(src_[idx]);
-      } else {
-        dst_[idx] =
-            fp8::scaled_convert<scalar_t, cache_t, kv_dt>(src_[idx], *scale);
-      }
+  };
+
+  for (int pid = split_start; pid < full_blocks_end; ++pid) {
+    auto block_id = batch_block_table[pid];
+    auto block_start_ptr = src_cache + block_id * cache_block_stride;
+    auto block_dst_ptr = dst + pid * block_size * dst_entry_stride;
+    for (int eid = 0; eid < block_size; ++eid) {
+      copy_entry(block_start_ptr + eid * cache_entry_stride,
+                 block_dst_ptr + eid * dst_entry_stride);
+    }
+  }
+
+  if (partial_block_size) {
+    auto block_id = batch_block_table[full_blocks_end];
+    auto block_start_ptr = src_cache + block_id * cache_block_stride;
+    auto block_dst_ptr = dst + full_blocks_end * block_size * dst_entry_stride;
+    for (int eid = 0; eid < partial_block_size; ++eid) {
+      copy_entry(block_start_ptr + eid * cache_entry_stride,
+                 block_dst_ptr + eid * dst_entry_stride);
    }
  }
 }
@@ -989,38 +992,34 @@ __global__ void gather_and_maybe_dequant_cache(
 // SCALAR_T is the data type of the destination tensor.
 // CACHE_T is the stored data type of kv-cache.
 // KV_DTYPE is the real data type of kv-cache.
-#define CALL_GATHER_CACHE(SCALAR_T, CACHE_T, KV_DTYPE)                        \
-  vllm::gather_and_maybe_dequant_cache<SCALAR_T, CACHE_T, KV_DTYPE, 576,      \
-                                       thread_block_size>                     \
-      <<<grid, block, 0, stream>>>(                                           \
-          reinterpret_cast<CACHE_T*>(src_cache.data_ptr()),                   \
-          reinterpret_cast<SCALAR_T*>(dst.data_ptr()),                        \
-          block_table.data_ptr<int32_t>(), cu_seq_lens.data_ptr<int32_t>(),   \
-          token_to_seq.data_ptr<int32_t>(), num_tokens, block_size,           \
-          block_table_stride, cache_block_stride, cache_entry_stride,         \
-          dst_entry_stride, reinterpret_cast<const float*>(scale.data_ptr()), \
-          seq_starts_ptr);
+#define CALL_GATHER_CACHE(SCALAR_T, CACHE_T, KV_DTYPE)                      \
+  vllm::gather_and_maybe_dequant_cache<SCALAR_T, CACHE_T, KV_DTYPE>         \
+      <<<grid, block, 0, stream>>>(                                         \
+          reinterpret_cast<CACHE_T*>(src_cache.data_ptr()),                 \
+          reinterpret_cast<SCALAR_T*>(dst.data_ptr()),                      \
+          block_table.data_ptr<int32_t>(), cu_seq_lens.data_ptr<int32_t>(), \
+          block_size, entry_size, block_table_stride, cache_block_stride,   \
+          cache_entry_stride, dst_entry_stride,                             \
+          reinterpret_cast<const float*>(scale.data_ptr()), seq_starts_ptr);

 // Gather sequences from the cache into the destination tensor.
 //  - cu_seq_lens contains the cumulative sequence lengths for each batch
 //  - block_table contains the cache block indices for each sequence
-//  - token_to_seq contains the back mapping from token_id to batch_id
 //  - Optionally, seq_starts (if provided) offsets the starting block index by
 //  (seq_starts[bid] / page_size)
 void gather_and_maybe_dequant_cache(
-    torch::Tensor const& src_cache,     // [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...]
-    torch::Tensor const& dst,           // [TOT_TOKENS, ENTRIES...]
-    torch::Tensor const& block_table,   // [BATCH, BLOCK_INDICES]
-    torch::Tensor const& cu_seq_lens,   // [BATCH+1]
-    torch::Tensor const& token_to_seq,  // [MAX_TOKEN_ACROSS_CHUNKS]
-    int64_t num_tokens, const std::string& kv_cache_dtype,
+    torch::Tensor const& src_cache,    // [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...]
+    torch::Tensor const& dst,          // [TOT_TOKENS, ENTRIES...]
+    torch::Tensor const& block_table,  // [BATCH, BLOCK_INDICES]
+    torch::Tensor const& cu_seq_lens,  // [BATCH+1]
+    int64_t batch_size, const std::string& kv_cache_dtype,
    torch::Tensor const& scale,
    std::optional<torch::Tensor> seq_starts = std::nullopt) {
  at::cuda::OptionalCUDAGuard device_guard(src_cache.device());
  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();

  int32_t block_size = src_cache.size(1);
-  int32_t head_dim = dst.size(-1);
+  int32_t entry_size = src_cache.flatten(2, -1).size(2);

  TORCH_CHECK(block_table.dtype() == torch::kInt32,
              "block_table must be int32");
@@ -1030,9 +1029,6 @@ void gather_and_maybe_dequant_cache(
    TORCH_CHECK(seq_starts.value().dtype() == torch::kInt32,
                "seq_starts must be int32");
  }
-  TORCH_CHECK(head_dim == 576,
-              "gather_and_maybe_dequant_cache only support the head_dim to 576 "
-              "for better performance")

  TORCH_CHECK(src_cache.device() == dst.device(),
              "src_cache and dst must be on the same device");
@@ -1050,9 +1046,10 @@ void gather_and_maybe_dequant_cache(
  int64_t cache_entry_stride = src_cache.stride(1);
  int64_t dst_entry_stride = dst.stride(0);

-  constexpr int32_t thread_block_size = 64;
-  dim3 grid(num_tokens);
-  dim3 block(thread_block_size);
+  // Decide on the number of splits based on the batch size.
+  int num_splits = batch_size > 128 ? 2 : batch_size > 64 ? 4 : 16;
+  dim3 grid(batch_size, num_splits);
+  dim3 block(1024);

  const int32_t* seq_starts_ptr =
      seq_starts.has_value() ? seq_starts.value().data_ptr<int32_t>() : nullptr;
@@ -13,18 +13,6 @@
  #define AMX_DISPATCH(...) case cpu_attention::ISA::AMX:
 #endif

-#ifdef __aarch64__
-  #include "cpu_attn_neon.hpp"
-  #define NEON_DISPATCH(...)                                                   \
-    case cpu_attention::ISA::NEON: {                                           \
-      using attn_impl = cpu_attention::AttentionImpl<cpu_attention::ISA::NEON, \
-                                                     scalar_t, head_dim>;      \
-      return __VA_ARGS__();                                                    \
-    }
-#else
-  #define NEON_DISPATCH(...) case cpu_attention::ISA::NEON:
-#endif  // #ifdef __aarch64__
-
 #define CPU_ATTN_DISPATCH_CASE(HEAD_DIM, ...) \
  case HEAD_DIM: {                            \
    constexpr size_t head_dim = HEAD_DIM;     \
@@ -53,7 +41,6 @@
  [&] {                                                                       \
    switch (ISA_TYPE) {                                                       \
      AMX_DISPATCH(__VA_ARGS__)                                               \
-      NEON_DISPATCH(__VA_ARGS__)                                              \
      case cpu_attention::ISA::VEC: {                                         \
        using attn_impl =                                                     \
            cpu_attention::AttentionImpl<cpu_attention::ISA::VEC, scalar_t,   \
@@ -86,8 +73,6 @@ torch::Tensor get_scheduler_metadata(
    isa = cpu_attention::ISA::VEC;
  } else if (isa_hint == "vec16") {
    isa = cpu_attention::ISA::VEC16;
-  } else if (isa_hint == "neon") {
-    isa = cpu_attention::ISA::NEON;
  } else {
    TORCH_CHECK(false, "Unsupported CPU attention ISA hint: " + isa_hint);
  }
@@ -173,8 +158,6 @@ void cpu_attn_reshape_and_cache(
      return cpu_attention::ISA::VEC;
    } else if (isa == "vec16") {
      return cpu_attention::ISA::VEC16;
-    } else if (isa == "neon") {
-      return cpu_attention::ISA::NEON;
    } else {
      TORCH_CHECK(false, "Invalid ISA type: " + isa);
    }
@@ -1,6 +1,7 @@
 #ifndef CPU_ATTN_HPP
 #define CPU_ATTN_HPP

+#include <unistd.h>
 #include <type_traits>
 #include <cstddef>

@@ -11,10 +12,9 @@
 #include "cpu_types.hpp"
 #include "scratchpad_manager.h"
 #include "cpu_attn_macros.h"
-#include "utils.hpp"

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

 template <ISA isa, typename scalar_t, int64_t head_dim>
 class AttentionImpl {};
@@ -143,12 +143,6 @@ struct AttentionMetadata {
      case ISA::VEC:
        ss << "VEC, ";
        break;
-      case ISA::VEC16:
-        ss << "VEC16, ";
-        break;
-      case ISA::NEON:
-        ss << "NEON, ";
-        break;
    }
    ss << "workitem_group_num: " << workitem_group_num
       << ", reduction_item_num: " << reduction_item_num
@@ -847,7 +841,7 @@ struct VecTypeTrait<c10::BFloat16> {
 };
 #endif

-#if !defined(__powerpc__) && !defined(__s390x__)
+#if !defined(__powerpc__)
 template <>
 struct VecTypeTrait<c10::Half> {
  using vec_t = vec_op::FP16Vec16;
@@ -1,386 +0,0 @@
-#ifndef CPU_ATTN_NEON_HPP
-#define CPU_ATTN_NEON_HPP
-
-#include "cpu_attn_impl.hpp"
-#include <arm_neon.h>
-#include <type_traits>
-namespace cpu_attention {
-
-namespace {
-
-#define BLOCK_SIZE_ALIGNMENT 32
-#define HEAD_SIZE_ALIGNMENT 32
-#define MAX_Q_HEAD_NUM_PER_ITER 16
-
-// These do not use vectorized class for loading / converting
-// because csrc/cpu/cpu_types_arm.hpp does not have fallback options
-// for vec_op::BF16Vec* / vec_op::BF16Vec* on Arm HW that
-// doesn't support BF16.
-// We don't use vec_op::FP32Vec* or vec_op::FP16Vec* for consistency.
-template <typename kv_cache_t>
-FORCE_INLINE void load_row8_B_as_f32(const kv_cache_t* p, float32x4_t& b0,
-                                     float32x4_t& b1);
-
-template <>
-FORCE_INLINE void load_row8_B_as_f32<float>(const float* p, float32x4_t& b0,
-                                            float32x4_t& b1) {
-  b0 = vld1q_f32(p + 0);
-  b1 = vld1q_f32(p + 4);
-}
-
-template <>
-FORCE_INLINE void load_row8_B_as_f32<c10::Half>(const c10::Half* p,
-                                                float32x4_t& b0,
-                                                float32x4_t& b1) {
-  const float16_t* h = reinterpret_cast<const float16_t*>(p);
-  float16x8_t v = vld1q_f16(h);
-  b0 = vcvt_f32_f16(vget_low_f16(v));
-  b1 = vcvt_f32_f16(vget_high_f16(v));
-}
-
-template <>
-FORCE_INLINE void load_row8_B_as_f32<c10::BFloat16>(const c10::BFloat16* p,
-                                                    float32x4_t& b0,
-                                                    float32x4_t& b1) {
-  const uint16_t* u = reinterpret_cast<const uint16_t*>(p);
-#ifdef ARM_BF16_SUPPORT
-  uint16x8_t u0 = vld1q_u16(u);
-  bfloat16x8_t bf0 = vreinterpretq_bf16_u16(u0);
-  b0 = vcvtq_low_f32_bf16(bf0);
-  b1 = vcvtq_high_f32_bf16(bf0);
-#else
-  uint16x8_t x0 = vld1q_u16(u);
-  uint32x4_t lo = vshlq_n_u32(vmovl_u16(vget_low_u16(x0)), 16);
-  uint32x4_t hi = vshlq_n_u32(vmovl_u16(vget_high_u16(x0)), 16);
-  b0 = vreinterpretq_f32_u32(lo);
-  b1 = vreinterpretq_f32_u32(hi);
-#endif
-}
-
-// Mx8, with 1 <= M <= 8 , K streamed, unroll-by-4 with NEON FMLAs
-// #Loads = (K // 4) * (M + 4 * sizeof(kv_cache_t) / 2)
-// #FMLAs = (K // 4) * (4 * 2 * M)
-// We have (4 * 2 * M) FMLAs for (M + 4 * sizeof(kv_cache_t) / 2) loads
-template <int32_t M, typename kv_cache_t>
-FORCE_INLINE void gemm_micro_neon_fmla_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) {
-  // kernel supports max M of 8, as it'd spill for larger M
-  static_assert(1 <= M && M <= 8, "M must be in [1,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 2 accumulators per row of M
-#define DECL_ACC(i) float32x4_t acc##i##_0, acc##i##_1;
-  ROWS_APPLY(DECL_ACC)
-#undef DECL_ACC
-
-  // initialize accumulators
-#define INIT_ACC(i)                              \
-  IF_M(i) {                                      \
-    if (accumulate) {                            \
-      acc##i##_0 = vld1q_f32(C + (i) * ldc + 0); \
-      acc##i##_1 = vld1q_f32(C + (i) * ldc + 4); \
-    } else {                                     \
-      acc##i##_0 = vdupq_n_f32(0.f);             \
-      acc##i##_1 = vdupq_n_f32(0.f);             \
-    }                                            \
-  }
-  ROWS_APPLY(INIT_ACC)
-#undef INIT_ACC
-
-  int32_t k = 0;
-
-  // K unrolled by 4
-  for (; k + 3 < K; k += 4) {
-    // load A[k..k+3] for each active row (M)
-#define LOAD_A4(i)     \
-  float32x4_t a##i##v; \
-  IF_M(i) a##i##v = vld1q_f32(a##i + k);
-    ROWS_APPLY(LOAD_A4)
-#undef LOAD_A4
-
-    // helper: FMA lane L from aiv
-#define FMAS_LANE(i, aiv, L)                              \
-  IF_M(i) {                                               \
-    acc##i##_0 = vfmaq_laneq_f32(acc##i##_0, b0, aiv, L); \
-    acc##i##_1 = vfmaq_laneq_f32(acc##i##_1, b1, aiv, L); \
-  }
-
-    // k + 0
-    {
-      float32x4_t b0, b1;
-      load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 0) * ldb, b0, b1);
-#define STEP_K0(i) FMAS_LANE(i, a##i##v, 0)
-      ROWS_APPLY(STEP_K0)
-#undef STEP_K0
-    }
-    // k + 1
-    {
-      float32x4_t b0, b1;
-      load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 1) * ldb, b0, b1);
-#define STEP_K1(i) FMAS_LANE(i, a##i##v, 1)
-      ROWS_APPLY(STEP_K1)
-#undef STEP_K1
-    }
-    // k + 2
-    {
-      float32x4_t b0, b1;
-      load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 2) * ldb, b0, b1);
-#define STEP_K2(i) FMAS_LANE(i, a##i##v, 2)
-      ROWS_APPLY(STEP_K2)
-#undef STEP_K2
-    }
-    // k + 3
-    {
-      float32x4_t b0, b1;
-      load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 3) * ldb, b0, b1);
-#define STEP_K3(i) FMAS_LANE(i, a##i##v, 3)
-      ROWS_APPLY(STEP_K3)
-#undef STEP_K3
-    }
-#undef FMAS_LANE
-  }
-
-  // K tail
-  for (; k < K; ++k) {
-    float32x4_t b0, b1;
-    load_row8_B_as_f32<kv_cache_t>(B + (int64_t)k * ldb, b0, b1);
-#define TAIL_ROW(i)                             \
-  IF_M(i) {                                     \
-    float32x4_t ai = vdupq_n_f32(*(a##i + k));  \
-    acc##i##_0 = vfmaq_f32(acc##i##_0, b0, ai); \
-    acc##i##_1 = vfmaq_f32(acc##i##_1, b1, ai); \
-  }
-    ROWS_APPLY(TAIL_ROW)
-#undef TAIL_ROW
-  }
-
-  // store accumulators to C
-#define STORE_ROW(i)                          \
-  IF_M(i) {                                   \
-    vst1q_f32(C + (i) * ldc + 0, acc##i##_0); \
-    vst1q_f32(C + (i) * ldc + 4, acc##i##_1); \
-  }
-  ROWS_APPLY(STORE_ROW)
-#undef STORE_ROW
-
-#undef ROWS_APPLY
-#undef IF_M
-}
-
-template <int32_t N, typename kv_cache_t>
-FORCE_INLINE void gemm_macro_neon_fmla_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) {
-  // micro kernel is Mx8
-  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_neon_fmla_Mx8_Ku4<8, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
-                                                      K, accumulate);
-          break;
-        case 4:
-          gemm_micro_neon_fmla_Mx8_Ku4<4, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
-                                                      K, accumulate);
-          break;
-        case 2:
-          gemm_micro_neon_fmla_Mx8_Ku4<2, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
-                                                      K, accumulate);
-          break;
-        default:
-          gemm_micro_neon_fmla_Mx8_Ku4<1, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
-                                                      K, accumulate);
-          break;
-      }
-    }
-    // no tail loop for N as it's guaranteed to be a multiple of 8
-    m += mb;
-  }
-}
-
-template <typename kv_cache_t>
-class TileGemmNeonFMLA {
- 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_neon_fmla_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_neon_fmla_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
-
-// this is similar to "ISA::VEC" at the moment
-template <typename scalar_t, int64_t head_dim>
-class AttentionImpl<ISA::NEON, scalar_t, head_dim> {
- 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;  // KV token num unit of QK and PV phases
-  constexpr static int64_t HeadDimAlignment =
-      HEAD_SIZE_ALIGNMENT;  // headdim num unit of PV phase
-  constexpr static int64_t MaxQHeadNumPerIteration = MAX_Q_HEAD_NUM_PER_ITER;
-  constexpr static int64_t HeadDim = head_dim;
-  constexpr static ISA ISAType = ISA::NEON;
-  constexpr static bool scale_on_logits = false;  // apply scale on q_buffer
-
-  static_assert(HeadDim % HeadDimAlignment == 0);
-  // the gemm micro kernel is Mx8
-  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<TileGemmNeonFMLA<kv_cache_t>> attention_iteration;
-    attention_iteration(CPU_ATTENTION_PARAMS);
-  }
-
-  // k_cache_token_group_stride: stride of K cache when move to next
-  // BlockSizeAlignment tokens in a block
-  constexpr static int64_t k_cache_token_group_stride(
-      const int32_t block_size) {
-    return BlockSizeAlignment;  // layout of k_cache block is [head_dim,
-                                // block_size], row-major
-  }
-
-  // v_cache_token_group_stride: stride of V cache when move to next
-  // BlockSizeAlignment tokens in a block
-  constexpr static int64_t v_cache_token_group_stride(
-      const int32_t block_size) {
-    return head_dim * BlockSizeAlignment;  // layout of v_cache is [block_size,
-                                           // head_dim], row-major
-  }
-
-  // v_cache_head_group_stride: stride of V cache when move to next
-  // HeadDimAlignment head dims in a block
-  constexpr static int64_t v_cache_head_group_stride(const int32_t block_size) {
-    return HeadDimAlignment;  // layout of v_cache is [block_size, head_dim],
-                              // row-major
-  }
-
-  // Copy q to q_buffer and cast it to fp32
-  static void copy_q_heads_tile(
-      scalar_t* __restrict__ src,  // [q_num, q_heads_per_kv, head_size]
-      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;
-        });
-      }
-    }
-  }
-
-  // reshape K as column-major and V as row-major
-  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) {
-#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) {
-          // skip
-          continue;
-        }
-
-        const int64_t block_idx = pos / block_size;
-        const int64_t block_offset = pos % block_size;
-        {
-          // Write Key
-          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;
-
-#pragma GCC unroll 8
-          for (int64_t i = 0, j = 0; i < head_dim; ++i, j += block_size) {
-            key_cache_start_ptr[j] = key_start_ptr[i];
-          }
-        }
-        {
-          // Write Value
-          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
-
-#endif  // #ifndef CPU_ATTN_NEON_HPP
@@ -26,6 +26,10 @@ namespace vec_op {

 #define FORCE_INLINE __attribute__((always_inline)) inline

+#define __max(a, b) ((a) > (b) ? (a) : (b))
+#define __min(a, b) ((a) < (b) ? (a) : (b))
+#define __abs(a) ((a) < (0) ? (0 - a) : (a))
+
 typedef struct f16x8_t {
  uint16_t val[8];
 } f16x8_t;
@@ -95,7 +99,7 @@ struct FP16Vec16 : public Vec<FP16Vec16> {
  void save(void* ptr) const { *reinterpret_cast<f16x16_t*>(ptr) = reg; }

  void save(void* ptr, const int elem_num) const {
-    int num = std::min(elem_num, VEC_ELEM_NUM);
+    int num = __min(elem_num, VEC_ELEM_NUM);
    std::memcpy(ptr, &(reg.val[0]), num * sizeof(uint16_t));
  }
 };
@@ -124,7 +128,7 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
  void save(void* ptr) const { *reinterpret_cast<f16x16_t*>(ptr) = reg; }

  void save(void* ptr, const int elem_num) const {
-    int num = std::min(elem_num, VEC_ELEM_NUM);
+    int num = __min(elem_num, VEC_ELEM_NUM);
    std::memcpy(ptr, &(reg.val[0]), num * sizeof(uint16_t));
  }
 };
@@ -139,9 +143,9 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
  explicit BF16Vec32(f16x32_t data) : reg(data) {};

  explicit BF16Vec32(BF16Vec8& vec8_data) {
-    unroll_loop<int, VEC_ELEM_NUM>([&vec8_data, this](int i) {
+    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
      reg.val[i] = vec8_data.reg.val[i % BF16Vec8::VEC_ELEM_NUM];
-    });
+    }
  }

  void save(void* ptr) const { *reinterpret_cast<f16x32_t*>(ptr) = reg; }
@@ -153,11 +157,15 @@ struct FP32Vec4 : public Vec<FP32Vec4> {
  f32x4_t reg;

  explicit FP32Vec4(float v) {
-    unroll_loop<int, VEC_ELEM_NUM>([&v, this](int i) { reg.val[i] = v; });
+    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
+      reg.val[i] = v;
+    }
  }

  explicit FP32Vec4() {
-    unroll_loop<int, VEC_ELEM_NUM>([this](int i) { reg.val[i] = 0.0f; });
+    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
+      reg.val[i] = 0.0f;
+    }
  }

  explicit FP32Vec4(const float* ptr)
@@ -174,11 +182,15 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
  f32x8_t reg;

  explicit FP32Vec8(float v) {
-    unroll_loop<int, VEC_ELEM_NUM>([&v, this](int i) { reg.val[i] = v; });
+    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
+      reg.val[i] = v;
+    }
  }

  explicit FP32Vec8() {
-    unroll_loop<int, VEC_ELEM_NUM>([this](int i) { reg.val[i] = 0.0f; });
+    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
+      reg.val[i] = 0.0f;
+    }
  }

  explicit FP32Vec8(const float* ptr)
@@ -189,68 +201,78 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
  explicit FP32Vec8(const FP32Vec8& data) : reg(data.reg) {};

  explicit FP32Vec8(const FP16Vec8& v) {
-    unroll_loop<int, VEC_ELEM_NUM>(
-        [&v, this](int i) { reg.val[i] = fp16_to_float(v.reg.val[i]); });
+    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
+      reg.val[i] = fp16_to_float(v.reg.val[i]);
+    }
  }

  FP32Vec8(const BF16Vec8& v) {
-    unroll_loop<int, VEC_ELEM_NUM>(
-        [&v, this](int i) { reg.val[i] = bf16_to_float(v.reg.val[i]); });
+    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
+      reg.val[i] = bf16_to_float(v.reg.val[i]);
+    }
  }

  float reduce_sum() const {
    float result = 0;
-    unroll_loop<int, VEC_ELEM_NUM>(
-        [&result, this](int i) { result += reg.val[i]; });
+    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
+      result += reg.val[i];
+    }
    return result;
  }

  FP32Vec8 exp() const {
    f32x8_t ret;
-    unroll_loop<int, VEC_ELEM_NUM>(
-        [&ret, this](int i) { ret.val[i] = expf(reg.val[i]); });
+    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
+      ret.val[i] = expf(reg.val[i]);
+    }
    return FP32Vec8(ret);
  }

  FP32Vec8 tanh() const {
    f32x8_t ret;
-    unroll_loop<int, VEC_ELEM_NUM>(
-        [&ret, this](int i) { ret.val[i] = tanhf(reg.val[i]); });
+    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
+      ret.val[i] = tanhf(reg.val[i]);
+    }
    return FP32Vec8(ret);
  }

  FP32Vec8 er() const {
    f32x8_t ret;
-    unroll_loop<int, VEC_ELEM_NUM>(
-        [&ret, this](int i) { ret.val[i] = erf(reg.val[i]); });
+    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
+      ret.val[i] = erf(reg.val[i]);
+    }
    return FP32Vec8(ret);
  }

  FP32Vec8 operator*(const FP32Vec8& b) const {
    f32x8_t ret;
-    unroll_loop<int, VEC_ELEM_NUM>(
-        [&ret, &b, this](int i) { ret.val[i] = reg.val[i] * b.reg.val[i]; });
+    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
+      ret.val[i] = reg.val[i] * b.reg.val[i];
+    }
    return FP32Vec8(ret);
  }

  FP32Vec8 operator+(const FP32Vec8& b) const {
    f32x8_t ret;
-    unroll_loop<int, VEC_ELEM_NUM>(
-        [&ret, &b, this](int i) { ret.val[i] = reg.val[i] + b.reg.val[i]; });
+    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
+      ret.val[i] = reg.val[i] + b.reg.val[i];
+    }
    return FP32Vec8(ret);
  }

  FP32Vec8 operator-(const FP32Vec8& b) const {
    f32x8_t ret;
-    unroll_loop<int, VEC_ELEM_NUM>(
-        [&ret, &b, this](int i) { ret.val[i] = reg.val[i] - b.reg.val[i]; });
+    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
+      ret.val[i] = reg.val[i] - b.reg.val[i];
+    }
    return FP32Vec8(ret);
  }

  FP32Vec8 operator/(const FP32Vec8& b) const {
    f32x8_t ret;
-    unroll_loop<int, VEC_ELEM_NUM>(
-        [&ret, &b, this](int i) { ret.val[i] = reg.val[i] / b.reg.val[i]; });
+    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
+      ret.val[i] = reg.val[i] / b.reg.val[i];
+    }
    return FP32Vec8(ret);
  }

@@ -262,11 +284,15 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
  f32x16_t reg;

  explicit FP32Vec16(float v) {
-    unroll_loop<int, VEC_ELEM_NUM>([&v, this](int i) { reg.val[i] = v; });
+    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
+      reg.val[i] = v;
+    }
  }

  explicit FP32Vec16() {
-    unroll_loop<int, VEC_ELEM_NUM>([this](int i) { reg.val[i] = 0.0f; });
+    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
+      reg.val[i] = 0.0f;
+    }
  }

  explicit FP32Vec16(const float* ptr)
@@ -275,27 +301,29 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
  explicit FP32Vec16(f32x16_t data) : reg(data) {};

  FP32Vec16(const FP32Vec4& data) {
-    unroll_loop<int, VEC_ELEM_NUM>([&data, this](int i) {
+    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
      reg.val[i] = data.reg.val[i % FP32Vec4::VEC_ELEM_NUM];
-    });
+    }
  }

  FP32Vec16(const FP32Vec8& data) {
-    unroll_loop<int, VEC_ELEM_NUM>([&data, this](int i) {
+    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
      reg.val[i] = data.reg.val[i % FP32Vec8::VEC_ELEM_NUM];
-    });
+    }
  }

  FP32Vec16(const FP32Vec16& data) : reg(data.reg) {};

  explicit FP32Vec16(const FP16Vec16& v) {
-    unroll_loop<int, VEC_ELEM_NUM>(
-        [&v, this](int i) { reg.val[i] = fp16_to_float(v.reg.val[i]); });
+    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
+      reg.val[i] = fp16_to_float(v.reg.val[i]);
+    }
  }

  explicit FP32Vec16(const BF16Vec16& v) {
-    unroll_loop<int, VEC_ELEM_NUM>(
-        [&v, this](int i) { reg.val[i] = bf16_to_float(v.reg.val[i]); });
+    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
+      reg.val[i] = bf16_to_float(v.reg.val[i]);
+    }
  }

  explicit FP32Vec16(const FP16Vec8& v) : FP32Vec16(FP32Vec8(v)) {};
@@ -303,74 +331,82 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
  FP32Vec16(const BF16Vec8& v) : FP32Vec16(FP32Vec8(v)) {};

  FP32Vec16 operator*(const FP32Vec16& b) const {
-    f32x16_t ret;
-    unroll_loop<int, VEC_ELEM_NUM>(
-        [&ret, &b, this](int i) { ret.val[i] = reg.val[i] * b.reg.val[i]; });
-    return FP32Vec16(ret);
+    FP32Vec16 result(0.0f);
+    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
+      result.reg.val[i] = reg.val[i] * b.reg.val[i];
+    }
+    return result;
  }

  FP32Vec16 operator+(const FP32Vec16& b) const {
-    f32x16_t ret;
-    unroll_loop<int, VEC_ELEM_NUM>(
-        [&ret, &b, this](int i) { ret.val[i] = reg.val[i] + b.reg.val[i]; });
-    return FP32Vec16(ret);
+    FP32Vec16 result(0.0f);
+    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
+      result.reg.val[i] = reg.val[i] + b.reg.val[i];
+    }
+    return result;
  }

  FP32Vec16 operator-(const FP32Vec16& b) const {
-    f32x16_t ret;
-    unroll_loop<int, VEC_ELEM_NUM>(
-        [&ret, &b, this](int i) { ret.val[i] = reg.val[i] - b.reg.val[i]; });
-    return FP32Vec16(ret);
+    FP32Vec16 result(0.0f);
+    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
+      result.reg.val[i] = reg.val[i] - b.reg.val[i];
+    }
+    return result;
  }

  FP32Vec16 operator/(const FP32Vec16& b) const {
-    f32x16_t ret;
-    unroll_loop<int, VEC_ELEM_NUM>(
-        [&ret, &b, this](int i) { ret.val[i] = reg.val[i] / b.reg.val[i]; });
-    return FP32Vec16(ret);
+    FP32Vec16 result(0.0f);
+    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
+      result.reg.val[i] = reg.val[i] / b.reg.val[i];
+    }
+    return result;
  }

  FP32Vec16 max(const FP32Vec16& b) const {
-    f32x16_t ret;
-    unroll_loop<int, VEC_ELEM_NUM>([&ret, &b, this](int i) {
-      ret.val[i] = std::max(reg.val[i], b.reg.val[i]);
-    });
-    return FP32Vec16(ret);
+    FP32Vec16 result(0.0f);
+    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
+      result.reg.val[i] = __max(reg.val[i], b.reg.val[i]);
+    }
+    return result;
  }

  FP32Vec16 min(const FP32Vec16& b) const {
-    f32x16_t ret;
-    unroll_loop<int, VEC_ELEM_NUM>([&ret, &b, this](int i) {
-      ret.val[i] = std::min(reg.val[i], b.reg.val[i]);
-    });
-    return FP32Vec16(ret);
+    FP32Vec16 result(0.0f);
+    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
+      result.reg.val[i] = __min(reg.val[i], b.reg.val[i]);
+    }
+    return result;
  }

  FP32Vec16 abs() const {
-    f32x16_t ret;
-    unroll_loop<int, VEC_ELEM_NUM>(
-        [&ret, this](int i) { ret.val[i] = std::abs(reg.val[i]); });
-    return FP32Vec16(ret);
+    FP32Vec16 result(0.0f);
+    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
+      result.reg.val[i] = __abs(reg.val[i]);
+    }
+    return result;
  }

  float reduce_sum() const {
    float result = 0.0f;
-    unroll_loop<int, VEC_ELEM_NUM>(
-        [&result, this](int i) { result += reg.val[i]; });
+    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
+      result += reg.val[i];
+    }
    return result;
  }

  float reduce_max() const {
-    float result = std::numeric_limits<float>::lowest();
-    unroll_loop<int, VEC_ELEM_NUM>(
-        [&result, this](int i) { result = std::max(reg.val[i], result); });
+    float result = reg.val[0];
+    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
+      result = __max(reg.val[i], result);
+    }
    return result;
  }

  float reduce_min() const {
-    float result = std::numeric_limits<float>::max();
-    unroll_loop<int, VEC_ELEM_NUM>(
-        [&result, this](int i) { result = std::min(reg.val[i], result); });
+    float result = reg.val[0];
+    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
+      result = __min(reg.val[i], result);
+    }
    return result;
  }

@@ -378,9 +414,13 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
  float reduce_sub_sum(int idx) {
    static_assert(VEC_ELEM_NUM % group_size == 0);
    float sum = 0.0;
-    const int start = idx * group_size;
-    unroll_loop<int, group_size>(
-        [&sum, &start, this](int i) { sum += reg.val[start + i]; });
+    int start = idx * group_size;
+    int end = (idx + 1) * group_size;
+
+    for (; (start < VEC_ELEM_NUM) && (start < end); ++start) {
+      sum += reg.val[start];
+    }
+
    return sum;
  }

@@ -437,13 +477,17 @@ inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
 }

 inline FP16Vec16::FP16Vec16(const FP32Vec16& v) {
-  unroll_loop<int, FP16Vec16::VEC_ELEM_NUM>(
-      [&v, this](int i) { reg.val[i] = float_to_fp16(v.reg.val[i]); });
+  int i = 0;
+  for (i = 0; i < FP16Vec16::VEC_ELEM_NUM; ++i) {
+    reg.val[i] = float_to_fp16(v.reg.val[i]);
+  }
 }

 inline FP16Vec8 ::FP16Vec8(const FP32Vec8& v) {
-  unroll_loop<int, FP16Vec8::VEC_ELEM_NUM>(
-      [&v, this](int i) { reg.val[i] = float_to_fp16(v.reg.val[i]); });
+  int i = 0;
+  for (i = 0; i < FP16Vec8::VEC_ELEM_NUM; ++i) {
+    reg.val[i] = float_to_fp16(v.reg.val[i]);
+  }
 }

 inline void fma(FP32Vec16& acc, FP32Vec16& a, FP32Vec16& b) {
@@ -451,13 +495,17 @@ inline void fma(FP32Vec16& acc, FP32Vec16& a, FP32Vec16& b) {
 }

 inline BF16Vec8::BF16Vec8(const FP32Vec8& v) {
-  unroll_loop<int, BF16Vec8::VEC_ELEM_NUM>(
-      [&v, this](int i) { reg.val[i] = float_to_bf16(v.reg.val[i]); });
+  int i = 0;
+  for (i = 0; i < BF16Vec8::VEC_ELEM_NUM; ++i) {
+    reg.val[i] = float_to_bf16(v.reg.val[i]);
+  }
 }

 inline BF16Vec16::BF16Vec16(const FP32Vec16& v) {
-  unroll_loop<int, BF16Vec16::VEC_ELEM_NUM>(
-      [&v, this](int i) { reg.val[i] = float_to_bf16(v.reg.val[i]); });
+  int i = 0;
+  for (i = 0; i < BF16Vec16::VEC_ELEM_NUM; ++i) {
+    reg.val[i] = float_to_bf16(v.reg.val[i]);
+  }
 }

 inline void prefetch(const void* addr) { __builtin_prefetch(addr, 0, 3); }
@@ -4,7 +4,6 @@

 #include <vecintrin.h>
 #include <cmath>
-#include <limits>
 #include <torch/all.h>
 namespace vec_op {

@@ -175,9 +174,8 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
  }

  explicit FP32Vec8(const BF16Vec8& v) {
-    // On big-endian s390x, place BF16 first to get correct byte order
-    reg.val[0] = (__vector float)vec_mergeh(v.reg, zero);
-    reg.val[1] = (__vector float)vec_mergel(v.reg, zero);
+    reg.val[0] = (__vector float)vec_mergeh(zero, v.reg);
+    reg.val[1] = (__vector float)vec_mergel(zero, v.reg);
  }

  float reduce_sum() const {
@@ -191,257 +189,51 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
  }

  FP32Vec8 exp() const {
-    f32x4x2_t out;
-
-    const __vector float log2e = vec_splats(1.44269504088896341f);
-    const __vector float one = vec_splats(1.0f);
-    const __vector float min_x = vec_splats(-87.3f);
-    const __vector float max_x = vec_splats(88.7f);
-
-    // 5th-degree minimax polynomial for 2^r (r in [0,1))
-    const __vector float c1 = vec_splats(0.6931471805599453f);
-    const __vector float c2 = vec_splats(0.240226506959101f);
-    const __vector float c3 = vec_splats(0.05550410866482158f);
-    const __vector float c4 = vec_splats(0.009618129107628477f);
-    const __vector float c5 = vec_splats(0.0013333558146428443f);
-
-    for (int i = 0; i < 2; i++) {
-      __vector float x = reg.val[i];
-
-      x = vec_max(x, min_x);
-      x = vec_min(x, max_x);
-
-      __vector float y = vec_mul(x, log2e);
-
-      __vector float kf = vec_floor(y);
-      __vector float r = vec_sub(y, kf);
-
-      __vector signed int k = vec_signed(kf);
-      const __vector signed int min_k = vec_splats((signed int)-126);
-      const __vector signed int max_k = vec_splats((signed int)127);
-      k = vec_min(vec_max(k, min_k), max_k);
-
-      // Build 2^k from exponent bits
-      __vector signed int exp_int = vec_add(k, vec_splats((signed int)127));
-      __vector unsigned int bits = (__vector unsigned int)exp_int;
-      bits = vec_sl(bits, vec_splats((unsigned int)23));
-      __vector float pow2k = (__vector float)bits;
-
-      // Improved minimax polynomial
-      __vector float poly = vec_madd(c5, r, c4);
-      poly = vec_madd(poly, r, c3);
-      poly = vec_madd(poly, r, c2);
-      poly = vec_madd(poly, r, c1);
-      poly = vec_madd(poly, r, one);
-
-      out.val[i] = vec_mul(pow2k, poly);
-    }
-
-    return FP32Vec8(out);
+    // TODO: Vectorize this
+    AliasReg ar;
+    ar.reg = reg;
+    f32x4x4_t ret;
+    ret.val[0][0] = std::exp(ar.values[0]);
+    ret.val[0][1] = std::exp(ar.values[1]);
+    ret.val[0][2] = std::exp(ar.values[2]);
+    ret.val[0][3] = std::exp(ar.values[3]);
+    ret.val[1][0] = std::exp(ar.values[4]);
+    ret.val[1][1] = std::exp(ar.values[5]);
+    ret.val[1][2] = std::exp(ar.values[6]);
+    ret.val[1][3] = std::exp(ar.values[7]);
+    return FP32Vec8(f32x4x2_t({ret.val[0], ret.val[1]}));
  }

  FP32Vec8 tanh() const {
-    // tanh(x) = (exp(2x) - 1) / (exp(2x) + 1)
-    const __vector float one = vec_splats(1.0f);
-    const __vector float two = vec_splats(2.0f);
-    const __vector float zero = vec_splats(0.0f);
-    const __vector float sat =
-        vec_splats(9.0f);  // beyond this, tanh(x) ~ sign(x)
-
-    f32x4x2_t out;
-
-    for (int i = 0; i < 2; i++) {
-      __vector float x = reg.val[i];
-      __vector float ax = vec_abs(x);
-
-      // sign(x): +1 or -1
-      __vector float sign = vec_sel(vec_splats(-1.0f), one, vec_cmpgt(x, zero));
-
-      // saturation mask: |x| > sat
-      __vector __bool int saturated = vec_cmpgt(ax, sat);
-
-      // 2x
-      __vector float two_x = vec_mul(x, two);
-
-      // Build a temporary FP32Vec8 with both lanes = 2x, reuse exp()
-      f32x4x2_t tmp;
-      tmp.val[0] = two_x;
-      tmp.val[1] = two_x;
-      FP32Vec8 exp_2x_vec(tmp);
-
-      FP32Vec8 e2x = exp_2x_vec.exp();
-      __vector float e = e2x.reg.val[i];
-
-      // tanh(x) = (e - 1) / (e + 1)
-      __vector float num = vec_sub(e, one);
-      __vector float den = vec_add(e, one);
-
-      __vector float t = vec_div(num, den);
-
-      // For large |x|, clamp to sign(x)
-      out.val[i] = vec_sel(t, sign, saturated);
-    }
-
-    return FP32Vec8(out);
+    // TODO: Vectorize this
+    AliasReg ar;
+    ar.reg = reg;
+    f32x4x4_t ret;
+    ret.val[0][0] = std::tanh(ar.values[0]);
+    ret.val[0][1] = std::tanh(ar.values[1]);
+    ret.val[0][2] = std::tanh(ar.values[2]);
+    ret.val[0][3] = std::tanh(ar.values[3]);
+    ret.val[1][0] = std::tanh(ar.values[4]);
+    ret.val[1][1] = std::tanh(ar.values[5]);
+    ret.val[1][2] = std::tanh(ar.values[6]);
+    ret.val[1][3] = std::tanh(ar.values[7]);
+    return FP32Vec8(f32x4x2_t({ret.val[0], ret.val[1]}));
  }

  FP32Vec8 er() const {
-    // A&S 7.1.26 approximation:
-    // erf(x) = sign(x) * (1 - ((((a5*t + a4)*t + a3)*t + a2)*t + a1) * t *
-    // exp(-x^2)) t = 1 / (1 + p*|x|),  p = 0.3275911
-
-    const __vector float one = vec_splats(1.0f);
-    const __vector float zero = vec_splats(0.0f);
-    const __vector float p = vec_splats(0.3275911f);
-
-    // Polynomial coeffs
-    const __vector float a1 = vec_splats(0.254829592f);
-    const __vector float a2 = vec_splats(-0.284496736f);
-    const __vector float a3 = vec_splats(1.421413741f);
-    const __vector float a4 = vec_splats(-1.453152027f);
-    const __vector float a5 = vec_splats(1.061405429f);
-
-    // Threshold where erf(x) ~ sign(x)
-    const __vector float sat = vec_splats(6.0f);
-
-    f32x4x2_t out;
-
-    for (int lane = 0; lane < 2; lane++) {
-      __vector float x = reg.val[lane];
-      __vector float ax = vec_abs(x);
-
-      // sign(x)
-      __vector float sign = vec_sel(vec_splats(-1.0f), one, vec_cmpgt(x, zero));
-
-      // |x| > 6 → erf(x) = ±1
-      __vector __bool int saturated = vec_cmpgt(ax, sat);
-
-      // t = 1 / (1 + p * |x|)
-      __vector float t = vec_madd(p, ax, one);
-      t = vec_div(one, t);
-
-      // poly = a5
-      __vector float poly = a5;
-      poly = vec_madd(poly, t, a4);
-      poly = vec_madd(poly, t, a3);
-      poly = vec_madd(poly, t, a2);
-      poly = vec_madd(poly, t, a1);
-
-      // full polynomial: poly = poly * t
-      poly = vec_mul(poly, t);
-
-      // Compute exp(-x^2)
-      __vector float x2 = vec_mul(x, x);
-      __vector float neg_x2 = vec_neg(x2);
-
-      f32x4x2_t tmp;
-      tmp.val[0] = neg_x2;
-      tmp.val[1] = neg_x2;
-      FP32Vec8 exp_neg_x2(tmp);
-
-      FP32Vec8 e = exp_neg_x2.exp();
-      __vector float ex = e.reg.val[lane];
-
-      // erf(x) = sign * (1 - poly * exp(-x^2))
-      __vector float term = vec_mul(poly, ex);
-      __vector float y = vec_sub(one, term);
-      y = vec_mul(y, sign);
-
-      // saturated → ±1
-      __vector float sat_val = vec_mul(sign, one);
-      out.val[lane] = vec_sel(y, sat_val, saturated);
-    }
-
-    return FP32Vec8(out);
-  }
-  // Elementwise sigmoid(x) = 1 / (1 + exp(-x))
-  FP32Vec8 sigmoid() const {
-    const __vector float one = vec_splats(1.0f);
-
-    f32x4x2_t neg;
-    for (int i = 0; i < 2; ++i) {
-      neg.val[i] = vec_neg(reg.val[i]);
-    }
-
-    FP32Vec8 neg_x(neg);
-    FP32Vec8 e = neg_x.exp();  // exp(-x)
-
-    f32x4x2_t denom;
-    for (int i = 0; i < 2; ++i) {
-      denom.val[i] = vec_add(one, e.reg.val[i]);
-    }
-
-    FP32Vec8 denom_vec(denom);
-    FP32Vec8 one_vec(1.0f);
-
-    return one_vec / denom_vec;
-  }
-
-  // Tanh-based GELU:
-  // gelu(x) = 0.5 * x * (1 + tanh(√(2/π) * (x + 0.044715 * x^3)))
-  FP32Vec8 gelu_tanh() const {
-    const __vector float k_s2pi = vec_splats(0.7978845608028654f);  // √(2/π)
-    const __vector float k_0_0447 = vec_splats(0.044715f);
-
-    f32x4x2_t x2, x3, inner;
-    for (int i = 0; i < 2; ++i) {
-      __vector float x = reg.val[i];
-      x2.val[i] = vec_mul(x, x);                            // x^2
-      x3.val[i] = vec_mul(x2.val[i], x);                    // x^3
-      __vector float t = vec_madd(k_0_0447, x3.val[i], x);  // x + 0.044715*x^3
-      inner.val[i] = vec_mul(k_s2pi, t);                    // √(2/π)*(...)
-    }
-
-    FP32Vec8 inner_vec(inner);
-    FP32Vec8 t = inner_vec.tanh();  // tanh part
-
-    FP32Vec8 one_vec(1.0f);
-    FP32Vec8 half_vec(0.5f);
-
-    FP32Vec8 x_vec(*this);
-    return x_vec * half_vec * (one_vec + t);
-  }
-
-  // Erf-based GELU:
-  // gelu(x) = 0.5 * x * (1 + erf(x / √2))
-  FP32Vec8 gelu_erf() const {
-    const __vector float inv_sqrt2 = vec_splats(0.7071067811865476f);  // 1/√2
-    FP32Vec8 x_vec(*this);
-
-    f32x4x2_t scaled;
-    for (int i = 0; i < 2; ++i) {
-      scaled.val[i] = vec_mul(reg.val[i], inv_sqrt2);
-    }
-    FP32Vec8 x_scaled(scaled);
-
-    FP32Vec8 erf_x = x_scaled.er();
-
-    FP32Vec8 one_vec(1.0f);
-    FP32Vec8 half_vec(0.5f);
-
-    return x_vec * half_vec * (one_vec + erf_x);
-  }
-
-  // Elementwise reciprocal: 1/x (scalar per lane, for correctness)
-  FP32Vec8 rcp() const {
-    AliasReg in, out;
-    in.reg = reg;
-
-    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
-      out.values[i] = 1.0f / in.values[i];
-    }
-    return FP32Vec8(out.reg);
-  }
-
-  // Elementwise rsqrt(x) = 1 / sqrt(x) (scalar per lane, for correctness)
-  FP32Vec8 rsqrt() const {
-    AliasReg in, out;
-    in.reg = reg;
-
-    for (int i = 0; i < VEC_ELEM_NUM; ++i) {
-      out.values[i] = 1.0f / std::sqrt(in.values[i]);
-    }
-    return FP32Vec8(out.reg);
+    // TODO: Vectorize this
+    AliasReg ar;
+    ar.reg = reg;
+    f32x4x4_t ret;
+    ret.val[0][0] = std::erf(ar.values[0]);
+    ret.val[0][1] = std::erf(ar.values[1]);
+    ret.val[0][2] = std::erf(ar.values[2]);
+    ret.val[0][3] = std::erf(ar.values[3]);
+    ret.val[1][0] = std::erf(ar.values[4]);
+    ret.val[1][1] = std::erf(ar.values[5]);
+    ret.val[1][2] = std::erf(ar.values[6]);
+    ret.val[1][3] = std::erf(ar.values[7]);
+    return FP32Vec8(f32x4x2_t({ret.val[0], ret.val[1]}));
  }

  FP32Vec8 operator*(const FP32Vec8& b) const {
@@ -524,11 +316,10 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
  }

  explicit FP32Vec16(const BF16Vec16& v) {
-    // On big-endian s390x, place BF16 first to get correct byte order
-    reg.val[0] = (__vector float)vec_mergeh(v.reg.val[0], zero);
-    reg.val[1] = (__vector float)vec_mergel(v.reg.val[0], zero);
-    reg.val[2] = (__vector float)vec_mergeh(v.reg.val[1], zero);
-    reg.val[3] = (__vector float)vec_mergel(v.reg.val[1], zero);
+    reg.val[0] = (__vector float)vec_mergeh(zero, v.reg.val[0]);
+    reg.val[1] = (__vector float)vec_mergel(zero, v.reg.val[0]);
+    reg.val[2] = (__vector float)vec_mergeh(zero, v.reg.val[1]);
+    reg.val[3] = (__vector float)vec_mergel(zero, v.reg.val[1]);
  }

  explicit FP32Vec16(const BF16Vec8& v) : FP32Vec16(FP32Vec8(v)) {}
@@ -585,23 +376,6 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
    return result;
  }

-  FP32Vec16 max(const FP32Vec16& b) const {
-    return FP32Vec16(f32x4x4_t({vec_max(reg.val[0], b.reg.val[0]),
-                                vec_max(reg.val[1], b.reg.val[1]),
-                                vec_max(reg.val[2], b.reg.val[2]),
-                                vec_max(reg.val[3], b.reg.val[3])}));
-  }
-
-  float reduce_max() const {
-    AliasReg ar;
-    ar.reg = reg;
-    float result = ar.values[0];
-    unroll_loop<int, VEC_ELEM_NUM>([&result, &ar](int i) {
-      if (ar.values[i] > result) result = ar.values[i];
-    });
-    return result;
-  }
-
  void save(float* ptr) const {
    vec_xst(reg.val[0], 0, ptr);
    vec_xst(reg.val[1], 16, ptr);
@@ -628,14 +402,15 @@ struct VecType<c10::BFloat16> {
  using vec_type = BF16Vec8;
 };

-// On s390x, FP16 (Half) is not natively supported, use FP32 vectors instead
-using FP16Vec16 = FP32Vec16;
-
 template <typename T>
 void storeFP32(float v, T* ptr) {
  *ptr = v;
 }

+inline void fma(FP32Vec16& acc, FP32Vec16& a, FP32Vec16& b) {
+  acc = acc + a * b;
+}
+
 namespace c10 {
 struct BFloat16 {
  uint16_t value;  // Assume BFloat16 is defined as a struct containing a 16-bit
@@ -654,79 +429,6 @@ inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
  #define __VEC_CLASS_FP_NAN (1 << 6)
 #endif

-// Optimized FMA (Fused Multiply-Add) implementations using IBM Z vector
-// intrinsics
-
-// FP32Vec4 FMA: acc = acc + (a * b) or equivalently acc = fma(a, b, acc)
-FORCE_INLINE void fma(FP32Vec4& acc, const FP32Vec4& a, const FP32Vec4& b) {
-  acc.reg = vec_madd(a.reg, b.reg, acc.reg);
-}
-
-// FP32Vec8 FMA: acc = acc + (a * b)
-FORCE_INLINE void fma(FP32Vec8& acc, const FP32Vec8& a, const FP32Vec8& b) {
-  acc.reg.val[0] = vec_madd(a.reg.val[0], b.reg.val[0], acc.reg.val[0]);
-  acc.reg.val[1] = vec_madd(a.reg.val[1], b.reg.val[1], acc.reg.val[1]);
-}
-
-// FP32Vec16 FMA: acc = acc + (a * b)
-FORCE_INLINE void fma(FP32Vec16& acc, const FP32Vec16& a, const FP32Vec16& b) {
-  acc.reg.val[0] = vec_madd(a.reg.val[0], b.reg.val[0], acc.reg.val[0]);
-  acc.reg.val[1] = vec_madd(a.reg.val[1], b.reg.val[1], acc.reg.val[1]);
-  acc.reg.val[2] = vec_madd(a.reg.val[2], b.reg.val[2], acc.reg.val[2]);
-  acc.reg.val[3] = vec_madd(a.reg.val[3], b.reg.val[3], acc.reg.val[3]);
-}
-
-// Multiply-Subtract: acc = acc - (a * b)
-FORCE_INLINE void fms(FP32Vec4& acc, const FP32Vec4& a, const FP32Vec4& b) {
-  acc.reg = vec_msub(a.reg, b.reg, acc.reg);
-}
-
-FORCE_INLINE void fms(FP32Vec8& acc, const FP32Vec8& a, const FP32Vec8& b) {
-  acc.reg.val[0] = vec_msub(a.reg.val[0], b.reg.val[0], acc.reg.val[0]);
-  acc.reg.val[1] = vec_msub(a.reg.val[1], b.reg.val[1], acc.reg.val[1]);
-}
-
-FORCE_INLINE void fms(FP32Vec16& acc, const FP32Vec16& a, const FP32Vec16& b) {
-  acc.reg.val[0] = vec_msub(a.reg.val[0], b.reg.val[0], acc.reg.val[0]);
-  acc.reg.val[1] = vec_msub(a.reg.val[1], b.reg.val[1], acc.reg.val[1]);
-  acc.reg.val[2] = vec_msub(a.reg.val[2], b.reg.val[2], acc.reg.val[2]);
-  acc.reg.val[3] = vec_msub(a.reg.val[3], b.reg.val[3], acc.reg.val[3]);
-}
-
-// Negative Multiply-Add: acc = -(a * b) + acc
-FORCE_INLINE void nfma(FP32Vec4& acc, const FP32Vec4& a, const FP32Vec4& b) {
-  acc.reg = vec_nmadd(a.reg, b.reg, acc.reg);
-}
-
-FORCE_INLINE void nfma(FP32Vec8& acc, const FP32Vec8& a, const FP32Vec8& b) {
-  acc.reg.val[0] = vec_nmadd(a.reg.val[0], b.reg.val[0], acc.reg.val[0]);
-  acc.reg.val[1] = vec_nmadd(a.reg.val[1], b.reg.val[1], acc.reg.val[1]);
-}
-
-FORCE_INLINE void nfma(FP32Vec16& acc, const FP32Vec16& a, const FP32Vec16& b) {
-  acc.reg.val[0] = vec_nmadd(a.reg.val[0], b.reg.val[0], acc.reg.val[0]);
-  acc.reg.val[1] = vec_nmadd(a.reg.val[1], b.reg.val[1], acc.reg.val[1]);
-  acc.reg.val[2] = vec_nmadd(a.reg.val[2], b.reg.val[2], acc.reg.val[2]);
-  acc.reg.val[3] = vec_nmadd(a.reg.val[3], b.reg.val[3], acc.reg.val[3]);
-}
-
-// Negative Multiply-Subtract: acc = -(a * b) - acc
-FORCE_INLINE void nfms(FP32Vec4& acc, const FP32Vec4& a, const FP32Vec4& b) {
-  acc.reg = vec_nmsub(a.reg, b.reg, acc.reg);
-}
-
-FORCE_INLINE void nfms(FP32Vec8& acc, const FP32Vec8& a, const FP32Vec8& b) {
-  acc.reg.val[0] = vec_nmsub(a.reg.val[0], b.reg.val[0], acc.reg.val[0]);
-  acc.reg.val[1] = vec_nmsub(a.reg.val[1], b.reg.val[1], acc.reg.val[1]);
-}
-
-FORCE_INLINE void nfms(FP32Vec16& acc, const FP32Vec16& a, const FP32Vec16& b) {
-  acc.reg.val[0] = vec_nmsub(a.reg.val[0], b.reg.val[0], acc.reg.val[0]);
-  acc.reg.val[1] = vec_nmsub(a.reg.val[1], b.reg.val[1], acc.reg.val[1]);
-  acc.reg.val[2] = vec_nmsub(a.reg.val[2], b.reg.val[2], acc.reg.val[2]);
-  acc.reg.val[3] = vec_nmsub(a.reg.val[3], b.reg.val[3], acc.reg.val[3]);
-}
-
 const static __vector unsigned char omask = {2,  3,  6,  7,  10, 11, 14, 15,
                                             18, 19, 22, 23, 26, 27, 30, 31};
 const static __vector unsigned int bias = {0x00007fff, 0x00007fff, 0x00007fff,
@@ -739,24 +441,13 @@ const static __vector unsigned int one = {1, 1, 1, 1};
 inline BF16Vec8::BF16Vec8(const FP32Vec8& v) {
  __vector unsigned int inp0 = (__vector unsigned int)(v.reg.val[0]);
  __vector unsigned int inp1 = (__vector unsigned int)(v.reg.val[1]);
-  __vector unsigned int lsb0 = inp0 >> sh16;
-  __vector unsigned int lsb1 = inp1 >> sh16;
-  lsb0 = lsb0 & one;
-  lsb1 = lsb1 & one;
-  __vector unsigned int rnd0 = lsb0 + bias;
-  __vector unsigned int rnd1 = lsb1 + bias;
-  inp0 = inp0 + rnd0;
-  inp1 = inp1 + rnd1;
  int cc;
  __vector __bool int sel0 =
      vec_fp_test_data_class(v.reg.val[0], __VEC_CLASS_FP_NAN, &cc);
  __vector __bool int sel1 =
      vec_fp_test_data_class(v.reg.val[1], __VEC_CLASS_FP_NAN, &cc);
-  inp0 = vec_sel(inp0, nan, sel0);
-  inp1 = vec_sel(inp1, nan, sel1);
-  inp0 = inp0 >> sh16;
-  inp1 = inp1 >> sh16;
-
+  inp0 = vec_sel(inp0, nan, sel0) >> sh16;
+  inp1 = vec_sel(inp1, nan, sel1) >> sh16;
  reg = (__vector signed short)vec_perm(inp0, inp1, omask);
 }

@@ -765,22 +456,6 @@ inline BF16Vec16::BF16Vec16(const FP32Vec16& v) {
  __vector unsigned int inp1 = (__vector unsigned int)(v.reg.val[1]);
  __vector unsigned int inp2 = (__vector unsigned int)(v.reg.val[2]);
  __vector unsigned int inp3 = (__vector unsigned int)(v.reg.val[3]);
-  __vector unsigned int lsb0 = inp0 >> sh16;
-  __vector unsigned int lsb1 = inp1 >> sh16;
-  __vector unsigned int lsb2 = inp2 >> sh16;
-  __vector unsigned int lsb3 = inp3 >> sh16;
-  lsb0 = lsb0 & one;
-  lsb1 = lsb1 & one;
-  lsb2 = lsb2 & one;
-  lsb3 = lsb3 & one;
-  __vector unsigned int rnd0 = lsb0 + bias;
-  __vector unsigned int rnd1 = lsb1 + bias;
-  __vector unsigned int rnd2 = lsb2 + bias;
-  __vector unsigned int rnd3 = lsb3 + bias;
-  inp0 = inp0 + rnd0;
-  inp1 = inp1 + rnd1;
-  inp2 = inp2 + rnd2;
-  inp3 = inp3 + rnd3;
  int cc;
  __vector __bool int sel0 =
      vec_fp_test_data_class(v.reg.val[0], __VEC_CLASS_FP_NAN, &cc);
@@ -790,164 +465,15 @@ inline BF16Vec16::BF16Vec16(const FP32Vec16& v) {
      vec_fp_test_data_class(v.reg.val[2], __VEC_CLASS_FP_NAN, &cc);
  __vector __bool int sel3 =
      vec_fp_test_data_class(v.reg.val[3], __VEC_CLASS_FP_NAN, &cc);
-  inp0 = vec_sel(inp0, nan, sel0);
-  inp1 = vec_sel(inp1, nan, sel1);
-  inp2 = vec_sel(inp2, nan, sel2);
-  inp3 = vec_sel(inp3, nan, sel3);
-  inp0 = inp0 >> sh16;
-  inp1 = inp1 >> sh16;
-  inp2 = inp2 >> sh16;
-  inp3 = inp3 >> sh16;
-
+  inp0 = vec_sel(inp0, nan, sel0) >> sh16;
+  inp1 = vec_sel(inp1, nan, sel1) >> sh16;
+  inp2 = vec_sel(inp2, nan, sel2) >> sh16;
+  inp3 = vec_sel(inp3, nan, sel3) >> sh16;
  reg.val[0] = (__vector signed short)vec_perm(inp0, inp1, omask);
  reg.val[1] = (__vector signed short)vec_perm(inp2, inp3, omask);
 }

-// 1D softmax over `n` elements in `input`, writes result to `output`.
-// Uses FP32Vec8 for main body, scalar tail handling.
-// Requirement: n > 0
-FORCE_INLINE void softmax_fp32vec8(float* output, const float* input, int n) {
-  if (n <= 0) return;
-
-  // ---------- Pass 1: find max ----------
-  float max_val = -std::numeric_limits<float>::infinity();
-  int i = 0;
-
-  for (; i + FP32Vec8::VEC_ELEM_NUM <= n; i += FP32Vec8::VEC_ELEM_NUM) {
-    FP32Vec8 v(input + i);
-    FP32Vec8::AliasReg ar;
-    ar.reg = v.reg;
-    for (int j = 0; j < FP32Vec8::VEC_ELEM_NUM; ++j) {
-      if (ar.values[j] > max_val) max_val = ar.values[j];
-    }
-  }
-  for (; i < n; ++i) {
-    if (input[i] > max_val) max_val = input[i];
-  }
-
-  // ---------- Pass 2: compute exp(x - max) and sum ----------
-  float sum = 0.0f;
-  i = 0;
-
-  for (; i + FP32Vec8::VEC_ELEM_NUM <= n; i += FP32Vec8::VEC_ELEM_NUM) {
-    float tmp[FP32Vec8::VEC_ELEM_NUM];
-    for (int j = 0; j < FP32Vec8::VEC_ELEM_NUM; ++j) {
-      tmp[j] = input[i + j] - max_val;
-    }
-
-    FP32Vec8 v(tmp);
-    FP32Vec8 e = v.exp();
-
-    FP32Vec8::AliasReg ar;
-    ar.reg = e.reg;
-    for (int j = 0; j < FP32Vec8::VEC_ELEM_NUM; ++j) {
-      output[i + j] = ar.values[j];
-      sum += ar.values[j];
-    }
-  }
-
-  // Tail
-  for (; i < n; ++i) {
-    float x = input[i] - max_val;
-    float ex = std::exp(x);  // scalar tail
-    output[i] = ex;
-    sum += ex;
-  }
-
-  // ---------- Pass 3: normalize ----------
-  float inv_sum = 1.0f / sum;
-  i = 0;
-
-  for (; i + FP32Vec8::VEC_ELEM_NUM <= n; i += FP32Vec8::VEC_ELEM_NUM) {
-    float tmp[FP32Vec8::VEC_ELEM_NUM];
-    for (int j = 0; j < FP32Vec8::VEC_ELEM_NUM; ++j) {
-      tmp[j] = output[i + j] * inv_sum;
-    }
-    FP32Vec8 v(tmp);
-    v.save(output + i);
-  }
-
-  for (; i < n; ++i) {
-    output[i] *= inv_sum;
-  }
-}
-
-// 1D RMSNorm kernel:
-//   input:  x[0..n-1]
-//   weight: w[0..n-1] (gamma), may be nullptr
-//   output: y[i] = x[i] * inv_rms * (weight[i] if weight != nullptr else 1)
-//   eps: small epsilon for numerical stability
-FORCE_INLINE void rmsnorm_fp32vec8(float* output, const float* input,
-                                   const float* weight, int n, float eps) {
-  if (n <= 0) return;
-
-  // ---------- Pass 1: compute sum of squares ----------
-  float sum_sq = 0.0f;
-  int i = 0;
-
-  for (; i + FP32Vec8::VEC_ELEM_NUM <= n; i += FP32Vec8::VEC_ELEM_NUM) {
-    FP32Vec8 x_vec(input + i);
-
-    FP32Vec8 sq = x_vec * x_vec;
-
-    FP32Vec8::AliasReg ar;
-    ar.reg = sq.reg;
-    for (int j = 0; j < FP32Vec8::VEC_ELEM_NUM; ++j) {
-      sum_sq += ar.values[j];
-    }
-  }
-
-  // Tail
-  for (; i < n; ++i) {
-    float v = input[i];
-    sum_sq += v * v;
-  }
-
-  float mean_sq = sum_sq / static_cast<float>(n);
-  float inv_rms = 1.0f / std::sqrt(mean_sq + eps);
-
-  // ---------- Pass 2: scale (and apply weight if given) ----------
-  const float inv_rms_f = inv_rms;
-  i = 0;
-
-  if (weight) {
-    // with gamma
-    for (; i + FP32Vec8::VEC_ELEM_NUM <= n; i += FP32Vec8::VEC_ELEM_NUM) {
-      FP32Vec8 x_vec(input + i);
-
-      float wtmp[FP32Vec8::VEC_ELEM_NUM];
-      for (int j = 0; j < FP32Vec8::VEC_ELEM_NUM; ++j) {
-        wtmp[j] = weight[i + j];
-      }
-      FP32Vec8 w_vec(wtmp);
-
-      FP32Vec8 scale_vec(inv_rms_f);
-      FP32Vec8 y = x_vec * scale_vec * w_vec;
-      y.save(output + i);
-    }
-
-    for (; i < n; ++i) {
-      output[i] = input[i] * inv_rms_f * weight[i];
-    }
-  } else {
-    // without gamma
-    for (; i + FP32Vec8::VEC_ELEM_NUM <= n; i += FP32Vec8::VEC_ELEM_NUM) {
-      FP32Vec8 x_vec(input + i);
-      FP32Vec8 scale_vec(inv_rms_f);
-      FP32Vec8 y = x_vec * scale_vec;
-      y.save(output + i);
-    }
-
-    for (; i < n; ++i) {
-      output[i] = input[i] * inv_rms_f;
-    }
-  }
-}
-
-// Prefetch data to cache for better memory access performance
-FORCE_INLINE void prefetch(const void* addr) {
-  __builtin_prefetch(addr, 0, 3);  // 0=read, 3=high temporal locality
-}
+inline void prefetch(const void* addr) { void __dcbt(const void* addr); }

 };  // namespace vec_op

@@ -104,8 +104,6 @@ struct FP16Vec16 : public Vec<FP16Vec16> {
  explicit FP16Vec16(bool, void* ptr)
      : reg(_mm256_stream_load_si256((__m256i*)ptr)) {}

-  explicit FP16Vec16(const c10::Half v) : reg(_mm256_set1_epi16(v.x)) {}
-
  explicit FP16Vec16(const FP32Vec16&);

  void save(void* ptr) const { _mm256_storeu_si256((__m256i*)ptr, reg); }
@@ -143,8 +141,6 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
  explicit BF16Vec16(bool, void* ptr)
      : reg(_mm256_stream_load_si256((__m256i*)ptr)) {}

-  explicit BF16Vec16(const c10::BFloat16 v) : reg(_mm256_set1_epi16(v.x)) {}
-
  explicit BF16Vec16(const FP32Vec16&);

  void save(void* ptr) const { _mm256_storeu_si256((__m256i*)ptr, reg); }
@@ -354,22 +350,6 @@ struct FP32Vec16 : public Vec<FP32Vec16> {

  explicit FP32Vec16(__m512 data) : reg(data) {}

-  // de-pack 4 bit values
-  explicit FP32Vec16(int64_t value, const FP32Vec16& lut) {
-    int64_t mask_0 = 0x0F0F0F0F0F0F0F0F;
-    int64_t mask_1 = 0xF0F0F0F0F0F0F0F0;
-    int64_t value_0 = value & mask_0;
-    int64_t value_1 = value & mask_1;
-    __m128i vec_0 = _mm_movpi64_epi64((__m64)value_0);
-    __m128i vec_1 = _mm_movpi64_epi64((__m64)value_1);
-    vec_0 = _mm_cvtepu8_epi16(vec_0);
-    vec_1 = _mm_cvtepu8_epi16(vec_1);
-    vec_1 = _mm_slli_epi16(vec_1, 4);
-    __m128i vec = _mm_or_si128(vec_0, vec_1);
-    __m512i vec_i32 = _mm512_cvtepu8_epi32(vec);
-    reg = _mm512_permutexvar_ps(vec_i32, lut.reg);
-  }
-
  explicit FP32Vec16(const FP32Vec4& data)
      : reg((__m512)_mm512_inserti32x4(
            _mm512_inserti32x4(
@@ -446,6 +426,14 @@ struct FP32Vec16 : public Vec<FP32Vec16> {

  float get_last_elem() const { return _mm512_cvtss_f32(reg); }

+  template <int group_size>
+  float reduce_sub_sum(int idx) {
+    static_assert(VEC_ELEM_NUM % group_size == 0);
+    constexpr uint32_t base_mask = (0xFFFF >> (16 - group_size));
+    __mmask16 mask = _cvtu32_mask16(base_mask << (idx * group_size));
+    return _mm512_mask_reduce_add_ps(mask, reg);
+  }
+
  void save(float* ptr) const { _mm512_storeu_ps(ptr, reg); }

  void save(float* ptr, const int elem_num) const {
@@ -767,25 +755,6 @@ inline void non_temporal_save(BF16Vec16& vec, void* ptr) {
 inline void non_temporal_save(FP32Vec16& vec, void* ptr) {
  _mm512_stream_ps((float*)ptr, vec.reg);
 }
-
-static void interleave_save(const BF16Vec16& vec0, const BF16Vec16& vec1,
-                            void* ptr) {
-  __m512i vec_0 = _mm512_cvtepu16_epi32(vec0.reg);
-  __m512i vec_1 = _mm512_cvtepu16_epi32(vec1.reg);
-  vec_1 = _mm512_slli_epi32(vec_1, 16);
-  vec_0 = _mm512_or_si512(vec_0, vec_1);
-  _mm512_storeu_epi32(ptr, vec_0);
-}
-
-static void interleave_save(const FP16Vec16& vec0, const FP16Vec16& vec1,
-                            void* ptr) {
-  __m512i vec_0 = _mm512_cvtepu16_epi32(vec0.reg);
-  __m512i vec_1 = _mm512_cvtepu16_epi32(vec1.reg);
-  vec_1 = _mm512_slli_epi32(vec_1, 16);
-  vec_0 = _mm512_or_si512(vec_0, vec_1);
-  _mm512_storeu_epi32(ptr, vec_0);
-}
-
 #endif

 inline void mem_barrier() { _mm_mfence(); }
@@ -1,402 +0,0 @@
-#include "cpu_types.hpp"
-#include "scratchpad_manager.h"
-#include "utils.hpp"
-
-#ifdef CPU_CAPABILITY_AMXBF16
-  #include "cpu/micro_gemm/cpu_micro_gemm_amx.hpp"
-#endif
-#include "cpu/micro_gemm/cpu_micro_gemm_vec.hpp"
-
-#define VLLM_DISPATCH_CASE_16B_TYPES(...)                 \
-  AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) \
-  AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__)
-
-#define VLLM_DISPATCH_16B_TYPES(TYPE, NAME, ...) \
-  AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_16B_TYPES(__VA_ARGS__))
-
-template <typename T>
-void print_logits(const char* name, T* ptr, int32_t row, int32_t col,
-                  int32_t stride) {
-  std::stringstream ss;
-  ss << std::fixed << std::setprecision(5) << name << ": [\n";
-  auto* curr_logits_buffer = ptr;
-  for (int32_t m = 0; m < row; ++m) {
-    for (int32_t n = 0; n < col; ++n) {
-      ss << curr_logits_buffer[n] << ", ";
-    }
-    ss << "\n";
-    curr_logits_buffer += stride;
-  }
-  ss << "]\n";
-  std::printf("%s", ss.str().c_str());
-}
-
-namespace {
-using cpu_utils::ISA;
-using cpu_utils::VecTypeTrait;
-
-template <typename scalar_t, ISA isa, bool has_zp, bool use_desc_act>
-class Dequantizer4b {
- public:
-  constexpr static int32_t pack_num = 32 / 4;
-  using scalar_vec_t = typename VecTypeTrait<scalar_t>::vec_t;
-
- public:
-  static void dequant(int32_t* __restrict__ q_weight,
-                      scalar_t* __restrict__ weight,
-                      scalar_t* __restrict__ scales,
-                      int32_t* __restrict__ zeros, int32_t* __restrict__ g_idx,
-                      const int64_t scales_stride, const int64_t zeros_stride,
-                      const int32_t k_size, const int32_t group_size) {
-    vec_op::FP32Vec16 lut;
-    if constexpr (has_zp) {
-      // AWQ
-      alignas(64) static const float LUT[16] = {
-          0.0f, 1.0f, 2.0f,  3.0f,  4.0f,  5.0f,  6.0f,  7.0f,
-          8.0f, 9.0f, 10.0f, 11.0f, 12.0f, 13.0f, 14.0f, 15.0f};
-      lut = vec_op::FP32Vec16(LUT);
-    } else {
-      // GPTQ
-      alignas(64) static const float LUT[16] = {
-          -8.0f, -7.0f, -6.0f, -5.0f, -4.0f, -3.0f, -2.0f, -1.0f,
-          0.0f,  1.0f,  2.0f,  3.0f,  4.0f,  5.0f,  6.0f,  7.0f};
-      lut = vec_op::FP32Vec16(LUT);
-    }
-
-    // per 64-bits elem contains 16 output channels
-    int64_t* __restrict__ curr_q_weight = reinterpret_cast<int64_t*>(q_weight);
-    int64_t* __restrict__ curr_zeros = reinterpret_cast<int64_t*>(zeros);
-    scalar_t* __restrict__ curr_weight = weight;
-    scalar_t* __restrict__ curr_scale = scales;
-    vec_op::FP32Vec16 scale_0;
-    vec_op::FP32Vec16 scale_1;
-    vec_op::FP32Vec16 zero_0;
-    vec_op::FP32Vec16 zero_1;
-    int32_t group_counter = 0;
-    for (int32_t k_idx = 0; k_idx < k_size; k_idx += 2) {
-      int64_t qwb_0 = *curr_q_weight;
-      int64_t qwb_1 = *(curr_q_weight + 1);
-      vec_op::FP32Vec16 wb_0(qwb_0, lut);
-      vec_op::FP32Vec16 wb_1(qwb_1, lut);
-
-      if constexpr (!use_desc_act) {
-        if (group_counter == 0) {
-          scale_0 = vec_op::FP32Vec16(scalar_vec_t(curr_scale));
-          scale_1 = vec_op::FP32Vec16(scale_0);
-          curr_scale += scales_stride;
-
-          if constexpr (has_zp) {
-            zero_0 = vec_op::FP32Vec16(*curr_zeros, lut);
-            zero_1 = vec_op::FP32Vec16(zero_0);
-            curr_zeros += zeros_stride / 2;
-          }
-        }
-      } else {
-        int32_t g_idx_0 = g_idx[k_idx];
-        int32_t g_idx_1 = g_idx[k_idx + 1];
-        scale_0 = vec_op::FP32Vec16(
-            scalar_vec_t(curr_scale + g_idx_0 * scales_stride));
-        scale_1 = vec_op::FP32Vec16(
-            scalar_vec_t(curr_scale + g_idx_1 * scales_stride));
-        if constexpr (has_zp) {
-          zero_0 = vec_op::FP32Vec16(*(curr_zeros + g_idx_0 * zeros_stride / 2),
-                                     lut);
-          zero_1 = vec_op::FP32Vec16(*(curr_zeros + g_idx_1 * zeros_stride / 2),
-                                     lut);
-        }
-      }
-
-      if constexpr (has_zp) {
-        wb_0 = wb_0 - zero_0;
-        wb_1 = wb_1 - zero_1;
-      }
-
-      wb_0 = wb_0 * scale_0;
-      wb_1 = wb_1 * scale_1;
-
-      scalar_vec_t output_vec_0(wb_0);
-      scalar_vec_t output_vec_1(wb_1);
-
-      // AMX needs to interlave K elements to pack as 32 bits
-      if constexpr (isa == ISA::AMX) {
-        vec_op::interleave_save(output_vec_0, output_vec_1, curr_weight);
-      } else {
-        output_vec_0.save(curr_weight);
-        output_vec_1.save(curr_weight + 16);
-      }
-
-      // update
-      curr_q_weight += 2;
-      curr_weight += 32;
-      if constexpr (!use_desc_act) {
-        group_counter += 2;
-        if (group_counter == group_size) {
-          group_counter = 0;
-        }
-      }
-    }
-  }
-};
-};  // namespace
-
-template <typename scalar_t, typename dequantizer_t, typename gemm_t>
-void cpu_gemm_wna16_impl(
-    scalar_t* __restrict__ input, int32_t* __restrict__ q_weight,
-    scalar_t* __restrict__ output, scalar_t* __restrict__ scales,
-    int32_t* __restrict__ zeros, int32_t* __restrict__ g_idx,
-    scalar_t* __restrict__ bias, const int32_t m_size, const int32_t n_size,
-    const int32_t k_size, const int64_t input_stride,
-    const int64_t output_stride, const int64_t scales_group_stride,
-    const int64_t zeros_group_stride, const int32_t group_num,
-    const int32_t group_size, const int64_t pack_factor) {
-  constexpr int32_t gemm_n_tile_size = gemm_t::NSize;
-  constexpr int32_t gemm_m_tile_size = gemm_t::MaxMSize;
-  constexpr int32_t n_block_size = 16;
-  static_assert(gemm_n_tile_size % n_block_size == 0);
-  const int32_t thread_num = omp_get_max_threads();
-
-  // a simple schedule policy, just to hold more B tiles in L2 and make sure
-  // each thread has tasks
-  const int32_t n_partition_size = [&]() {
-    const int64_t cache_size = cpu_utils::get_l2_size();
-    int64_t ps_cache_limit = cache_size / (k_size * sizeof(scalar_t));
-    int64_t ps_thread_limit = n_size / thread_num;
-    ps_cache_limit =
-        std::max((ps_cache_limit / gemm_n_tile_size) * gemm_n_tile_size,
-                 (int64_t)gemm_n_tile_size);
-    ps_thread_limit =
-        std::max((ps_thread_limit / gemm_n_tile_size) * gemm_n_tile_size,
-                 (int64_t)gemm_n_tile_size);
-    return std::min(ps_cache_limit, ps_thread_limit);
-  }();
-  const int32_t task_num = (n_size + n_partition_size - 1) / n_partition_size;
-
-  // get buffer size
-  const int64_t b_buffer_size =
-      (((n_partition_size * k_size * sizeof(scalar_t) + 63) / 64) * 64);
-  const int64_t c_buffer_size =
-      (((gemm_m_tile_size * gemm_n_tile_size * sizeof(float) + 63) / 64) * 64);
-  const int64_t b_buffer_offset = 0;
-  const int64_t c_buffer_offset = b_buffer_size;
-  const int64_t buffer_size = b_buffer_size + c_buffer_size;
-  DNNLScratchPadManager::get_dnnl_scratchpad_manager()->realloc(buffer_size *
-                                                                thread_num);
-
-  alignas(64) cpu_utils::Counter counter;
-  cpu_utils::Counter* counter_ptr = &counter;
-
-#pragma omp parallel for schedule(static, 1)
-  for (int32_t thread_id = 0; thread_id < thread_num; ++thread_id) {
-    scalar_t* __restrict__ b_buffer = nullptr;
-    float* __restrict__ c_buffer = nullptr;
-    {
-      uint8_t* buffer_ptr = DNNLScratchPadManager::get_dnnl_scratchpad_manager()
-                                ->get_data<uint8_t>() +
-                            thread_id * buffer_size;
-      b_buffer = reinterpret_cast<scalar_t*>(buffer_ptr + b_buffer_offset);
-      c_buffer = reinterpret_cast<float*>(buffer_ptr + c_buffer_offset);
-    }
-
-    const int64_t q_weight_block_stride = n_block_size / pack_factor * k_size;
-    const int64_t b_buffer_block_stride = n_block_size * k_size;
-    const int32_t zeros_block_stride = n_block_size / pack_factor;
-
-    gemm_t gemm;
-
-    for (;;) {
-      int32_t task_id = counter_ptr->acquire_counter();
-
-      if (task_id >= task_num) {
-        break;
-      }
-
-      const int32_t n_start_idx = task_id * n_partition_size;
-      const int32_t n_block_start_idx = n_start_idx / n_block_size;
-      const int32_t n_num = std::min(n_partition_size, n_size - n_start_idx);
-      const int32_t n_block_num = n_num / n_block_size;
-      // std::printf("thread_id: %d, task_id: %d, n_start_idx: %d, n_num: %d\n",
-      // thread_id, task_id, n_start_idx, n_num);
-
-      // dequant weight
-      {
-        int32_t* __restrict__ curr_q_weight =
-            q_weight + n_block_start_idx * q_weight_block_stride;
-        scalar_t* __restrict__ curr_b_buffer = b_buffer;
-        scalar_t* __restrict__ curr_scales = scales + n_start_idx;
-        int32_t* __restrict__ curr_zeros = zeros + n_start_idx / pack_factor;
-        for (int32_t block_idx = 0; block_idx < n_block_num; ++block_idx) {
-          dequantizer_t::dequant(curr_q_weight, curr_b_buffer, curr_scales,
-                                 curr_zeros, g_idx, scales_group_stride,
-                                 zeros_group_stride, k_size, group_size);
-
-          // if (block_idx == 0 && n_start_idx == 0) {
-          //     print_logits("depacked weight", curr_b_buffer, k_size,
-          //     n_block_size, n_block_size);
-          // }
-
-          // update
-          curr_q_weight += q_weight_block_stride;
-          curr_b_buffer += b_buffer_block_stride;
-          curr_scales += n_block_size;
-          curr_zeros += zeros_block_stride;
-        }
-      }
-
-      // compute loop
-      {
-        const int32_t n_tile_num = n_num / gemm_n_tile_size;
-        scalar_t* __restrict__ curr_input = input;
-        scalar_t* __restrict__ init_bias = bias;
-        if (bias != nullptr) {
-          init_bias += n_start_idx;
-        }
-        scalar_t* __restrict__ init_output = output + n_start_idx;
-        for (int32_t m_idx = 0; m_idx < m_size; m_idx += gemm_m_tile_size) {
-          const int32_t curr_m_size =
-              std::min(gemm_m_tile_size, m_size - m_idx);
-          scalar_t* __restrict__ curr_b_buffer = b_buffer;
-          scalar_t* __restrict__ curr_bias = init_bias;
-          scalar_t* __restrict__ curr_output = init_output;
-          for (int32_t n_tile_idx = 0; n_tile_idx < n_tile_num; ++n_tile_idx) {
-            gemm.gemm(curr_input, curr_b_buffer, c_buffer, curr_m_size, k_size,
-                      input_stride, b_buffer_block_stride, gemm_n_tile_size,
-                      false);
-
-            if (bias != nullptr) {
-              cpu_micro_gemm::bias_epilogue<gemm_n_tile_size>(
-                  c_buffer, curr_output, curr_bias, curr_m_size,
-                  gemm_n_tile_size, output_stride);
-              curr_bias += gemm_n_tile_size;
-            } else {
-              cpu_micro_gemm::default_epilogue<gemm_n_tile_size>(
-                  c_buffer, curr_output, curr_m_size, gemm_n_tile_size,
-                  output_stride);
-            }
-
-            curr_b_buffer +=
-                b_buffer_block_stride * (gemm_n_tile_size / n_block_size);
-            curr_output += gemm_n_tile_size;
-          }
-          curr_input += gemm_m_tile_size * input_stride;
-          init_output += gemm_m_tile_size * output_stride;
-        }
-      }
-    }
-  }
-}
-
-void cpu_gemm_wna16(
-    const torch::Tensor& input,  // [M, K]
-    const torch::Tensor&
-        q_weight,           // [N / 16, K * 16 / pack_factor], packed as int32
-    torch::Tensor& output,  // [M, N]
-    const torch::Tensor& scales,  // [group_num, N]
-    const std::optional<torch::Tensor>&
-        zeros,  // [group_num, N / pack_factor], packed as int32
-    const std::optional<torch::Tensor>& g_idx,  // [K]
-    const std::optional<torch::Tensor>& bias,   // [N]
-    const int64_t pack_factor, const std::string& isa_hint) {
-  using cpu_utils::ISA;
-  TORCH_CHECK_EQ(pack_factor, 8);  // only supports 4bits
-  const int32_t a_m_size = input.size(0);
-  const int32_t a_k_size = input.size(1);
-  const int64_t a_m_stride = input.stride(0);
-  const int32_t b_n_size = q_weight.size(0) * 16;
-  TORCH_CHECK_EQ(a_k_size % 32, 0);
-  TORCH_CHECK_EQ(b_n_size % 32, 0);
-  const int32_t group_num = scales.size(0);
-  const int32_t group_size = a_k_size / group_num;
-  TORCH_CHECK_EQ(group_size % 2, 0);
-  const int64_t scales_group_stride = scales.stride(0);
-  const int64_t output_m_stride = output.stride(0);
-
-  bool has_zp = zeros.has_value();
-  bool use_desc_act = g_idx.has_value();
-  TORCH_CHECK(!(has_zp && use_desc_act));
-
-  ISA isa = [&]() {
-    if (isa_hint == "amx") {
-      return ISA::AMX;
-    } else if (isa_hint == "vec") {
-      return ISA::VEC;
-    } else {
-      TORCH_CHECK(false, "unsupported isa hint: " + isa_hint);
-    }
-  }();
-
-  int32_t* zeros_ptr = has_zp ? zeros->data_ptr<int32_t>() : nullptr;
-  const int64_t zeros_group_stride = has_zp ? zeros->stride(0) : 0;
-  int32_t* g_idx_ptr = use_desc_act ? g_idx->data_ptr<int32_t>() : nullptr;
-
-  VLLM_DISPATCH_16B_TYPES(input.scalar_type(), "cpu_gemm_wna16", [&]() {
-    if (isa == ISA::AMX) {
-      using gemm_t = cpu_micro_gemm::MicroGemm<ISA::AMX, scalar_t>;
-      if (has_zp) {
-        using dequantizer_t = Dequantizer4b<scalar_t, ISA::AMX, true, false>;
-        cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
-            input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
-            output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
-            g_idx_ptr, bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr,
-            a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
-            scales_group_stride, zeros_group_stride, group_num, group_size,
-            pack_factor);
-        return;
-      }
-      if (use_desc_act) {
-        using dequantizer_t = Dequantizer4b<scalar_t, ISA::AMX, false, true>;
-        cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
-            input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
-            output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
-            g_idx_ptr, bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr,
-            a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
-            scales_group_stride, zeros_group_stride, group_num, group_size,
-            pack_factor);
-        return;
-      } else {
-        using dequantizer_t = Dequantizer4b<scalar_t, ISA::AMX, false, false>;
-        cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
-            input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
-            output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
-            g_idx_ptr, bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr,
-            a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
-            scales_group_stride, zeros_group_stride, group_num, group_size,
-            pack_factor);
-        return;
-      }
-    } else if (isa == ISA::VEC) {
-      using gemm_t = cpu_micro_gemm::MicroGemm<ISA::VEC, scalar_t>;
-      if (has_zp) {
-        using dequantizer_t = Dequantizer4b<scalar_t, ISA::VEC, true, false>;
-        cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
-            input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
-            output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
-            g_idx_ptr, bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr,
-            a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
-            scales_group_stride, zeros_group_stride, group_num, group_size,
-            pack_factor);
-        return;
-      }
-      if (use_desc_act) {
-        using dequantizer_t = Dequantizer4b<scalar_t, ISA::VEC, false, true>;
-        cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
-            input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
-            output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
-            g_idx_ptr, bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr,
-            a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
-            scales_group_stride, zeros_group_stride, group_num, group_size,
-            pack_factor);
-        return;
-      } else {
-        using dequantizer_t = Dequantizer4b<scalar_t, ISA::VEC, false, false>;
-        cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
-            input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
-            output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
-            g_idx_ptr, bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr,
-            a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
-            scales_group_stride, zeros_group_stride, group_num, group_size,
-            pack_factor);
-        return;
-      }
-    }
-  });
-}
@@ -396,9 +396,9 @@ MatMulPrimitiveHandler::MatMulPrimitiveHandler(const Args& args)
    : DNNLMatMulPrimitiveHandler(
          static_cast<DNNLMatMulPrimitiveHandler::Args>(args), args.ab_type),
      m_size_cache_(nullptr) {
-  assert(b_type_ == dnnl::memory::data_type::f32 ||
-         b_type_ == dnnl::memory::data_type::bf16 ||
-         b_type_ == dnnl::memory::data_type::f16);
+  assert(ab_type_ == dnnl::memory::data_type::f32 ||
+         ab_type_ == dnnl::memory::data_type::bf16 ||
+         ab_type_ == dnnl::memory::data_type::f16);

  dnnl::memory::desc original_b_md({b_k_size_, b_n_size_}, b_type_,
                                   {b_k_stride_, b_n_stride_});
@@ -1,245 +0,0 @@
-#ifndef CPU_MICRO_GEMM_AMX_HPP
-#define CPU_MICRO_GEMM_AMX_HPP
-#include "cpu/micro_gemm/cpu_micro_gemm_impl.hpp"
-
-namespace cpu_micro_gemm {
-namespace {
-// AMX specific
-constexpr static int64_t AMX_TILE_ROW_BYTES = 64;
-constexpr static int64_t AMX_TILE_ROW_NUM = 16;
-constexpr static int64_t AMX_TILE_BYTES = AMX_TILE_ROW_BYTES * AMX_TILE_ROW_NUM;
-
-typedef struct __tile_config {
-  uint8_t palette_id = 1;
-  uint8_t start_row = 0;
-  uint8_t reserved_0[14] = {0};
-  uint16_t colsb[16] = {0};
-  uint8_t rows[16] = {0};
-} __tilecfg;
-
-// 2-2-4 pattern, for 16 < m <= 32
-// TILE 0, 1: load A matrix, row num should be 16, m - 16
-// TILE 2, 3: load B matrix, row num should be 16
-// TILE 4, 5, 6, 7: store results C matrix, row num should be 16, 16, m - 16, m
-// - 16
-template <typename scalar_t>
-class TileGemm224 {
- public:
-  FORCE_INLINE static void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
-    TORCH_CHECK(false, "Unsupported data type for TileGemm224");
-  }
-
-  FORCE_INLINE static void init_tile_config(int32_t m, __tilecfg& config) {
-    TORCH_CHECK(false, "Unsupported data type for TileGemm224");
-  }
-};
-
-template <>
-class TileGemm224<c10::BFloat16> {
- public:
-  using scalar_t = c10::BFloat16;
-  FORCE_INLINE static void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
-    const int32_t k_times = k / (AMX_TILE_ROW_NUM * 4 / sizeof(c10::BFloat16));
-    c10::BFloat16* __restrict__ a_tile_0 = a_ptr;
-    c10::BFloat16* __restrict__ a_tile_1 = a_ptr + lda * AMX_TILE_ROW_NUM;
-    const int64_t a_tile_stride = lda * sizeof(c10::BFloat16);
-
-    // B is always packed as 16 output channels block
-    c10::BFloat16* __restrict__ b_tile_2 = b_ptr;
-    c10::BFloat16* __restrict__ b_tile_3 = b_ptr + b_n_group_stride;
-    const int32_t b_tile_stride = AMX_TILE_ROW_BYTES;
-
-    float* __restrict__ c_tile_4 = c_ptr;
-    float* __restrict__ c_tile_5 =
-        c_tile_4 + AMX_TILE_ROW_BYTES / sizeof(float);
-    float* __restrict__ c_tile_6 = c_ptr + AMX_TILE_ROW_NUM * ldc;
-    float* __restrict__ c_tile_7 =
-        c_tile_6 + AMX_TILE_ROW_BYTES / sizeof(float);
-    const int32_t c_tile_stride = ldc * sizeof(float);
-
-    if (accum_c) {
-      _tile_loadd(4, c_tile_4, c_tile_stride);
-      _tile_loadd(5, c_tile_5, c_tile_stride);
-      _tile_loadd(6, c_tile_6, c_tile_stride);
-      _tile_loadd(7, c_tile_7, c_tile_stride);
-    } else {
-      _tile_zero(4);
-      _tile_zero(5);
-      _tile_zero(6);
-      _tile_zero(7);
-    }
-
-    for (int32_t k = 0; k < k_times; ++k) {
-      _tile_loadd(0, a_tile_0, a_tile_stride);
-      _tile_stream_loadd(2, b_tile_2, b_tile_stride);
-      _tile_dpbf16ps(4, 0, 2);
-      _tile_stream_loadd(3, b_tile_3, b_tile_stride);
-      _tile_dpbf16ps(5, 0, 3);
-      _tile_loadd(1, a_tile_1, a_tile_stride);
-      _tile_dpbf16ps(6, 1, 2);
-      _tile_dpbf16ps(7, 1, 3);
-
-      // update ptrs
-      a_tile_0 += AMX_TILE_ROW_BYTES / sizeof(c10::BFloat16);
-      a_tile_1 += AMX_TILE_ROW_BYTES / sizeof(c10::BFloat16);
-      b_tile_2 += AMX_TILE_BYTES / sizeof(c10::BFloat16);
-      b_tile_3 += AMX_TILE_BYTES / sizeof(c10::BFloat16);
-    }
-
-    _tile_stored(4, c_tile_4, c_tile_stride);
-    _tile_stored(5, c_tile_5, c_tile_stride);
-    _tile_stored(6, c_tile_6, c_tile_stride);
-    _tile_stored(7, c_tile_7, c_tile_stride);
-  }
-
-  FORCE_INLINE static void init_tile_config(int32_t m, __tilecfg& config) {
-    const int32_t m_0 = AMX_TILE_ROW_NUM;
-    const int32_t m_1 = m - AMX_TILE_ROW_NUM;
-    config.rows[0] = m_0;
-    config.rows[1] = m_1;
-    config.rows[2] = AMX_TILE_ROW_NUM;
-    config.rows[3] = AMX_TILE_ROW_NUM;
-    config.rows[4] = m_0;
-    config.rows[5] = m_0;
-    config.rows[6] = m_1;
-    config.rows[7] = m_1;
-    _tile_loadconfig(&config);
-  }
-};
-
-// 1-2-2 pattern, for 0 < m <= 16
-// TILE 0, (1): load A matrix, use extra 1 tile for prefetch, row num should be
-// m, m
-// TILE 2, 3, (4, 5): load B matrix, use extra 2 tiles for prefetch, row
-// num should be 16
-// TILE 6, 7, (6, 7): store results C matrix, row num should be
-// m
-template <typename scalar_t>
-class TileGemm122 {
- public:
-  FORCE_INLINE static void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
-    TORCH_CHECK(false, "Unsupported data type for TileGemm122");
-  }
-
-  FORCE_INLINE static void init_tile_config(int32_t m, __tilecfg& config) {
-    TORCH_CHECK(false, "Unsupported data type for TileGemm122");
-  }
-};
-
-template <>
-class TileGemm122<c10::BFloat16> {
- public:
-  using scalar_t = c10::BFloat16;
-  FORCE_INLINE static void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
-    c10::BFloat16* __restrict__ a_tile_0 = a_ptr;
-    c10::BFloat16* __restrict__ a_tile_1 =
-        a_ptr + AMX_TILE_ROW_BYTES / sizeof(c10::BFloat16);
-    const int64_t a_tile_stride = lda * sizeof(c10::BFloat16);
-
-    c10::BFloat16* __restrict__ b_tile_2 = b_ptr;
-    c10::BFloat16* __restrict__ b_tile_3 = b_ptr + b_n_group_stride;
-    c10::BFloat16* __restrict__ b_tile_4 =
-        b_tile_2 + AMX_TILE_BYTES / sizeof(c10::BFloat16);
-    c10::BFloat16* __restrict__ b_tile_5 =
-        b_tile_3 + AMX_TILE_BYTES / sizeof(c10::BFloat16);
-    int64_t b_stride = AMX_TILE_ROW_BYTES;
-
-    float* __restrict__ c_tile_6 = c_ptr;
-    float* __restrict__ c_tile_7 = c_ptr + AMX_TILE_ROW_BYTES / sizeof(float);
-    int64_t c_stride = ldc * sizeof(float);
-
-    const int32_t k_times = k / (AMX_TILE_ROW_NUM * 4 / sizeof(c10::BFloat16));
-    const int32_t k_group_times = k_times / 2;
-    const bool has_tail = (k_times % 2 == 1);
-
-    if (accum_c) {
-      _tile_loadd(6, c_tile_6, c_stride);
-      _tile_loadd(7, c_tile_7, c_stride);
-    } else {
-      _tile_zero(6);
-      _tile_zero(7);
-    }
-
-    for (int32_t k = 0; k < k_group_times; ++k) {
-      _tile_loadd(0, a_tile_0, a_tile_stride);
-      _tile_stream_loadd(2, b_tile_2, b_stride);
-      _tile_dpbf16ps(6, 0, 2);
-      _tile_stream_loadd(3, b_tile_3, b_stride);
-      _tile_dpbf16ps(7, 0, 3);
-      _tile_loadd(1, a_tile_1, a_tile_stride);
-      _tile_stream_loadd(4, b_tile_4, b_stride);
-      _tile_dpbf16ps(6, 1, 4);
-      _tile_stream_loadd(5, b_tile_5, b_stride);
-      _tile_dpbf16ps(7, 1, 5);
-
-      // update ptrs
-      a_tile_0 += 2 * AMX_TILE_ROW_BYTES / sizeof(c10::BFloat16);
-      a_tile_1 += 2 * AMX_TILE_ROW_BYTES / sizeof(c10::BFloat16);
-      b_tile_2 += 2 * AMX_TILE_BYTES / sizeof(c10::BFloat16);
-      b_tile_3 += 2 * AMX_TILE_BYTES / sizeof(c10::BFloat16);
-      b_tile_4 += 2 * AMX_TILE_BYTES / sizeof(c10::BFloat16);
-      b_tile_5 += 2 * AMX_TILE_BYTES / sizeof(c10::BFloat16);
-    }
-
-    if (has_tail) {
-      _tile_loadd(0, a_tile_0, a_tile_stride);
-      _tile_stream_loadd(2, b_tile_2, b_stride);
-      _tile_dpbf16ps(6, 0, 2);
-      _tile_stream_loadd(3, b_tile_3, b_stride);
-      _tile_dpbf16ps(7, 0, 3);
-    }
-
-    _tile_stored(6, c_tile_6, c_stride);
-    _tile_stored(7, c_tile_7, c_stride);
-  }
-
-  FORCE_INLINE static void init_tile_config(int32_t m, __tilecfg& config) {
-    config.rows[0] = m;
-    config.rows[1] = m;
-    config.rows[2] = AMX_TILE_ROW_NUM;
-    config.rows[3] = AMX_TILE_ROW_NUM;
-    config.rows[4] = AMX_TILE_ROW_NUM;
-    config.rows[5] = AMX_TILE_ROW_NUM;
-    config.rows[6] = m;
-    config.rows[7] = m;
-    _tile_loadconfig(&config);
-  }
-};
-}  // namespace
-
-// Gemm kernel uses AMX, requires B matrix to be packed
-template <typename scalar_t>
-class MicroGemm<cpu_utils::ISA::AMX, scalar_t> {
- public:
-  static constexpr int32_t MaxMSize = 32;
-  static constexpr int32_t NSize = 32;
-
- public:
-  MicroGemm() : curr_m_(-1) {
-    vec_op::unroll_loop<int, 8>([&](int i) { amx_tile_config_.colsb[i] = 64; });
-  }
-
-  void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
-    if (m > AMX_TILE_ROW_NUM) {
-      if (m != curr_m_) {
-        curr_m_ = m;
-        TileGemm224<scalar_t>::init_tile_config(m, amx_tile_config_);
-      }
-      TileGemm224<scalar_t>::gemm(CPU_MICRO_GEMM_PARAMS);
-    } else {
-      if (m != curr_m_) {
-        curr_m_ = m;
-        TileGemm122<scalar_t>::init_tile_config(m, amx_tile_config_);
-      }
-      TileGemm122<scalar_t>::gemm(CPU_MICRO_GEMM_PARAMS);
-    }
-  }
-
- private:
-  alignas(64) __tilecfg amx_tile_config_;
-  int32_t curr_m_;
-};
-
-}  // namespace cpu_micro_gemm
-
-#endif
@@ -1,91 +0,0 @@
-#ifndef CPU_MICRO_GEMM_IMPL_HPP
-#define CPU_MICRO_GEMM_IMPL_HPP
-#include "cpu/utils.hpp"
-#include "cpu/cpu_types.hpp"
-
-namespace cpu_micro_gemm {
-#define DEFINE_CPU_MICRO_GEMM_PARAMS                                        \
-  scalar_t *__restrict__ a_ptr, scalar_t *__restrict__ b_ptr,               \
-      float *__restrict__ c_ptr, const int32_t m, const int32_t k,          \
-      const int64_t lda, const int64_t b_n_group_stride, const int64_t ldc, \
-      const bool accum_c
-
-#define CPU_MICRO_GEMM_PARAMS \
-  a_ptr, b_ptr, c_ptr, m, k, lda, b_n_group_stride, ldc, accum_c
-
-template <cpu_utils::ISA isa, typename scalar_t>
-class MicroGemm {
- public:
-  static constexpr int32_t MaxMSize = 16;
-  static constexpr int32_t NSize = 16;
-
- public:
-  void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
-    TORCH_CHECK(false, "Unimplemented MicroGemm.");
-  }
-};
-
-template <int32_t n_size, typename scalar_t>
-FORCE_INLINE void default_epilogue(float* __restrict__ c_ptr,
-                                   scalar_t* __restrict__ d_ptr,
-                                   const int32_t m, const int64_t ldc,
-                                   const int64_t ldd) {
-  using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
-  static_assert(n_size % 16 == 0);
-
-  float* __restrict__ curr_c = c_ptr;
-  scalar_t* __restrict__ curr_d = d_ptr;
-  for (int32_t i = 0; i < m; ++i) {
-    float* __restrict__ curr_c_iter = curr_c;
-    scalar_t* __restrict__ curr_d_iter = curr_d;
-    vec_op::unroll_loop<int32_t, n_size / 16>([&](int32_t n_g_idx) {
-      vec_op::FP32Vec16 c_vec_fp32(curr_c_iter);
-      scalar_vec_t c_vec(c_vec_fp32);
-      c_vec.save(curr_d_iter);
-      curr_c_iter += 16;
-      curr_d_iter += 16;
-    });
-    curr_c += ldc;
-    curr_d += ldd;
-  }
-}
-
-template <int32_t n_size, typename scalar_t>
-FORCE_INLINE void bias_epilogue(float* __restrict__ c_ptr,
-                                scalar_t* __restrict__ d_ptr,
-                                scalar_t* __restrict__ bias_ptr,
-                                const int32_t m, const int64_t ldc,
-                                const int64_t ldd) {
-  using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
-  static_assert(n_size % 16 == 0);
-  constexpr int32_t n_group_num = n_size / 16;
-  static_assert(n_group_num <= 16);
-
-  vec_op::FP32Vec16 bias_vecs[n_group_num];
-  scalar_t* __restrict__ curr_bias = bias_ptr;
-  vec_op::unroll_loop<int32_t, n_group_num>([&](int32_t i) {
-    scalar_vec_t vec(curr_bias);
-    bias_vecs[i] = vec_op::FP32Vec16(vec);
-    curr_bias += 16;
-  });
-
-  float* __restrict__ curr_c = c_ptr;
-  scalar_t* __restrict__ curr_d = d_ptr;
-  for (int32_t i = 0; i < m; ++i) {
-    float* __restrict__ curr_c_iter = curr_c;
-    scalar_t* __restrict__ curr_d_iter = curr_d;
-    vec_op::unroll_loop<int32_t, n_group_num>([&](int32_t n_g_idx) {
-      vec_op::FP32Vec16 c_vec_fp32(curr_c_iter);
-      c_vec_fp32 = c_vec_fp32 + bias_vecs[n_g_idx];
-      scalar_vec_t c_vec(c_vec_fp32);
-      c_vec.save(curr_d_iter);
-      curr_c_iter += 16;
-      curr_d_iter += 16;
-    });
-    curr_c += ldc;
-    curr_d += ldd;
-  }
-}
-}  // namespace cpu_micro_gemm
-
-#endif
@@ -1,115 +0,0 @@
-#ifndef CPU_MICRO_GEMM_VEC_HPP
-#define CPU_MICRO_GEMM_VEC_HPP
-#include "cpu/micro_gemm/cpu_micro_gemm_impl.hpp"
-
-namespace cpu_micro_gemm {
-namespace {
-// 8-2-16 pattern, 8 regs for A, 2 regs for B, 16 regs for C, [8, K] @ [k, 32]
-template <typename scalar_t>
-class TileGemm82 {
- public:
-  FORCE_INLINE static void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
-    switch (m) {
-      case 1:
-        gemm_micro<1>(CPU_MICRO_GEMM_PARAMS);
-        break;
-      case 2:
-        gemm_micro<2>(CPU_MICRO_GEMM_PARAMS);
-        break;
-      case 3:
-        gemm_micro<3>(CPU_MICRO_GEMM_PARAMS);
-        break;
-      case 4:
-        gemm_micro<4>(CPU_MICRO_GEMM_PARAMS);
-        break;
-      case 5:
-        gemm_micro<5>(CPU_MICRO_GEMM_PARAMS);
-        break;
-      case 6:
-        gemm_micro<6>(CPU_MICRO_GEMM_PARAMS);
-        break;
-      case 7:
-        gemm_micro<7>(CPU_MICRO_GEMM_PARAMS);
-        break;
-      case 8:
-        gemm_micro<8>(CPU_MICRO_GEMM_PARAMS);
-        break;
-    }
-  }
-
-  template <int32_t M>
-  static void gemm_micro(DEFINE_CPU_MICRO_GEMM_PARAMS) {
-    static_assert(0 < M <= 8);
-    using load_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
-
-    scalar_t* __restrict__ curr_b_0 = b_ptr;
-    scalar_t* __restrict__ curr_b_1 = b_ptr + b_n_group_stride;
-    float* __restrict__ curr_c_0 = c_ptr;
-    float* __restrict__ curr_c_1 = c_ptr + 16;
-
-    vec_op::FP32Vec16 c_regs[M * 2];
-    if (accum_c) {
-      float* __restrict__ curr_m_c_0 = curr_c_0;
-      float* __restrict__ curr_m_c_1 = curr_c_1;
-      vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
-        c_regs[i * 2] = vec_op::FP32Vec16(curr_m_c_0);
-        c_regs[i * 2 + 1] = vec_op::FP32Vec16(curr_m_c_1);
-
-        // update
-        curr_m_c_0 += ldc;
-        curr_m_c_1 += ldc;
-      });
-    }
-
-    scalar_t* __restrict__ curr_a = a_ptr;
-    for (int32_t k_idx = 0; k_idx < k; ++k_idx) {
-      load_vec_t b_0_reg(curr_b_0);
-      vec_op::FP32Vec16 fp32_b_0_reg(b_0_reg);
-      load_vec_t b_1_reg(curr_b_1);
-      vec_op::FP32Vec16 fp32_b_1_reg(b_1_reg);
-
-      scalar_t* __restrict__ curr_m_a = curr_a;
-      vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
-        scalar_t v = *curr_m_a;
-        load_vec_t a_reg_original(v);
-        vec_op::FP32Vec16 a_reg(a_reg_original);
-        c_regs[i * 2] = c_regs[i * 2] + a_reg * fp32_b_0_reg;
-        c_regs[i * 2 + 1] = c_regs[i * 2 + 1] + a_reg * fp32_b_1_reg;
-
-        // update
-        curr_m_a += lda;
-      });
-
-      // update
-      curr_a += 1;
-      curr_b_0 += 16;
-      curr_b_1 += 16;
-    }
-
-    vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
-      c_regs[i * 2].save(curr_c_0);
-      c_regs[i * 2 + 1].save(curr_c_1);
-
-      // update
-      curr_c_0 += ldc;
-      curr_c_1 += ldc;
-    });
-  }
-};
-}  // namespace
-
-// Gemm kernel uses vector instructions, requires B matrix to be packed
-template <typename scalar_t>
-class MicroGemm<cpu_utils::ISA::VEC, scalar_t> {
- public:
-  static constexpr int32_t MaxMSize = 8;
-  static constexpr int32_t NSize = 32;
-
- public:
-  void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
-    TileGemm82<scalar_t>::gemm(CPU_MICRO_GEMM_PARAMS);
-  }
-};
-}  // namespace cpu_micro_gemm
-
-#endif
@@ -103,13 +103,6 @@ void cpu_attention_with_kv_cache(
 // Note: just for avoiding importing errors
 void placeholder_op() { TORCH_CHECK(false, "Unimplemented"); }

-void cpu_gemm_wna16(const torch::Tensor& input, const torch::Tensor& q_weight,
-                    torch::Tensor& output, const torch::Tensor& scales,
-                    const std::optional<torch::Tensor>& zeros,
-                    const std::optional<torch::Tensor>& g_idx,
-                    const std::optional<torch::Tensor>& bias,
-                    const int64_t pack_factor, const std::string& isa_hint);
-
 TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  // vLLM custom ops

@@ -172,6 +165,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  // Quantization
 #if defined(__AVX512F__) || (defined(__aarch64__) && !defined(__APPLE__)) || \
    defined(__powerpc64__)
+  at::Tag stride_tag = at::Tag::needs_fixed_stride_order;
  // Helper function to release oneDNN handlers
  ops.def("release_dnnl_matmul_handler(int handler) -> ()",
          &release_dnnl_matmul_handler);
@@ -207,13 +201,15 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  // Compute int8 quantized tensor for given scaling factor.
  ops.def(
      "static_scaled_int8_quant(Tensor! out, Tensor input, Tensor scale,"
-      "Tensor? azp) -> ()");
+      "Tensor? azp) -> ()",
+      {stride_tag});
  ops.impl("static_scaled_int8_quant", torch::kCPU, &static_scaled_int8_quant);

  // Compute int8 quantized tensor and scaling factor
  ops.def(
      "dynamic_scaled_int8_quant(Tensor! out, Tensor input, Tensor! scale, "
-      "Tensor!? azp) -> ()");
+      "Tensor!? azp) -> ()",
+      {stride_tag});
  ops.impl("dynamic_scaled_int8_quant", torch::kCPU,
           &dynamic_scaled_int8_quant);
 #endif
@@ -287,15 +283,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  ops.def("static_scaled_fp8_quant() -> ()", placeholder_op);
  ops.def("dynamic_scaled_fp8_quant() -> ()", placeholder_op);
  ops.def("dynamic_per_token_scaled_fp8_quant() -> ()", placeholder_op);
-
-  // WNA16
-#if defined(__AVX512F__)
-  ops.def(
-      "cpu_gemm_wna16(Tensor input, Tensor q_weight, Tensor(a2!) output, "
-      "Tensor scales, Tensor? zeros, Tensor? g_idx, Tensor? bias, SymInt "
-      "pack_factor, str isa_hint) -> ()");
-  ops.impl("cpu_gemm_wna16", torch::kCPU, &cpu_gemm_wna16);
-#endif
 }

 TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _utils), utils) {
@@ -45,55 +45,31 @@ std::string init_cpu_threads_env(const std::string& cpu_ids) {
  // Memory node binding
  if (numa_available() != -1) {
    int mem_node_id = numa_node_of_cpu(omp_cpu_ids.front());
-    std::set<int> node_ids;
-    for (const auto& cpu_id : omp_cpu_ids) {
-      int node_id = numa_node_of_cpu(cpu_id);
-      if (node_id != -1) {
-        node_ids.insert(node_id);
-      }
-      if (node_id != mem_node_id) {
-        TORCH_WARN("CPU ", cpu_id, " is on NUMA node ", node_id, ", but CPU ",
-                   omp_cpu_ids.front(), " is on NUMA node ", mem_node_id,
-                   ". All CPUs should be on the same NUMA node for optimal "
-                   "performance. Memory will be bound to NUMA node ",
-                   mem_node_id, ".");
-      }
+    // Verify all CPUs are on the same NUMA node
+    for (size_t i = 1; i < omp_cpu_ids.size(); ++i) {
+      int node_id = numa_node_of_cpu(omp_cpu_ids[i]);
+      TORCH_CHECK(node_id == mem_node_id, "CPU ", omp_cpu_ids[i],
+                  " is on NUMA node ", node_id, ", but CPU ",
+                  omp_cpu_ids.front(), " is on NUMA node ", mem_node_id,
+                  ". All CPUs should be on the same NUMA node for optimal "
+                  "performance. Memory will be bound to NUMA node ",
+                  mem_node_id, ".");
    }
-    // Concatenate all node_ids into a single comma-separated string
-    if (!node_ids.empty()) {
-      std::string node_ids_str;
-      for (const int node_id : node_ids) {
-        if (!node_ids_str.empty()) {
-          node_ids_str += ",";
-        }
-        node_ids_str += std::to_string(node_id);
-      }
+    bitmask* mask = numa_parse_nodestring(std::to_string(mem_node_id).c_str());
+    bitmask* src_mask = numa_get_membind();

-      bitmask* mask = numa_parse_nodestring(node_ids_str.c_str());
-      bitmask* src_mask = numa_get_membind();
+    int pid = getpid();

-      int pid = getpid();
-
-      if (mask && src_mask) {
-        // move all existing pages to the specified numa node.
-        *(src_mask->maskp) = *(src_mask->maskp) ^ *(mask->maskp);
-        int page_num = numa_migrate_pages(pid, src_mask, mask);
-        if (page_num == -1) {
-          TORCH_WARN("numa_migrate_pages failed. errno: " +
-                     std::to_string(errno));
-        }
-
-        // restrict memory allocation node.
-        numa_set_membind(mask);
-        numa_set_strict(1);
-
-        numa_free_nodemask(mask);
-        numa_free_nodemask(src_mask);
-      } else {
-        TORCH_WARN("numa_parse_nodestring or numa_get_membind failed. errno: " +
-                   std::to_string(errno));
-      }
+    // move all existing pages to the specified numa node.
+    *(src_mask->maskp) = *(src_mask->maskp) ^ *(mask->maskp);
+    int page_num = numa_migrate_pages(pid, src_mask, mask);
+    if (page_num == -1) {
+      TORCH_WARN("numa_migrate_pages failed. errno: " + std::to_string(errno));
    }
+
+    // restrict memory allocation node.
+    numa_set_membind(mask);
+    numa_set_strict(1);
  }

  // OMP threads binding
@@ -1,73 +0,0 @@
-#ifndef UTILS_HPP
-#define UTILS_HPP
-
-#include <atomic>
-#include <cassert>
-#include <cstdint>
-#include <unistd.h>
-
-#if defined(__APPLE__)
-  #include <sys/sysctl.h>
-#endif
-
-#include "cpu_types.hpp"
-
-namespace cpu_utils {
-enum class ISA { AMX, VEC };
-
-template <typename T>
-struct VecTypeTrait {
-  using vec_t = void;
-};
-
-template <>
-struct VecTypeTrait<float> {
-  using vec_t = vec_op::FP32Vec16;
-};
-
-#if !defined(__aarch64__) || defined(ARM_BF16_SUPPORT)
-template <>
-struct VecTypeTrait<c10::BFloat16> {
-  using vec_t = vec_op::BF16Vec16;
-};
-#endif
-
-template <>
-struct VecTypeTrait<c10::Half> {
-  using vec_t = vec_op::FP16Vec16;
-};
-
-struct Counter {
-  std::atomic<int64_t> counter;
-  char _padding[56];
-
-  Counter() : counter(0) {}
-
-  void reset_counter() { counter.store(0); }
-
-  int64_t acquire_counter() { return counter++; }
-};
-
-inline int64_t get_l2_size() {
-  static int64_t size = []() {
-#if defined(__APPLE__)
-    // macOS doesn't have _SC_LEVEL2_CACHE_SIZE. Use sysctlbyname.
-    int64_t l2_cache_size = 0;
-    size_t len = sizeof(l2_cache_size);
-    if (sysctlbyname("hw.l2cachesize", &l2_cache_size, &len, NULL, 0) == 0 &&
-        l2_cache_size > 0) {
-      return l2_cache_size >> 1;  // use 50% of L2 cache
-    }
-    // Fallback if sysctlbyname fails
-    return 128LL * 1024 >> 1;  // use 50% of 128KB
-#else
-    long l2_cache_size = sysconf(_SC_LEVEL2_CACHE_SIZE);
-    assert(l2_cache_size != -1);
-    return l2_cache_size >> 1;  // use 50% of L2 cache
-#endif
-  }();
-  return size;
-}
-}  // namespace cpu_utils
-
-#endif
@@ -22,10 +22,15 @@ torch::Tensor get_cuda_view_from_cpu_tensor(torch::Tensor& cpu_tensor) {
  auto strides = cpu_tensor.strides();
  auto options = cpu_tensor.options().device(torch::kCUDA);

-  // use default no-op deleter, since the memory is owned by the original CPU
-  // tensor
+  // from_blob signature: from_blob(void *data, IntArrayRef sizes, ..., Deleter,
+  // const TensorOptions &) Provide a no-op deleter. The CPU tensor holds the
+  // memory, so we don't free it here.
+  auto deleter = [](void*) {
+    // no-op, since the memory is owned by the original CPU tensor
+  };
+
  torch::Tensor cuda_tensor =
-      torch::from_blob(device_ptr, sizes, strides, options);
+      torch::from_blob(device_ptr, sizes, strides, deleter, options);

  TORCH_CHECK(cuda_tensor.device().is_cuda(),
              "Resulting tensor is not on CUDA device");
@@ -117,24 +117,3 @@
      break;                                  \
    }                                         \
  }
-
-#define VLLM_DISPATCH_RANK234(NUM_DIMS, ...)                                   \
-  switch (NUM_DIMS) {                                                          \
-    case 2: {                                                                  \
-      constexpr int tensor_rank = 2;                                           \
-      __VA_ARGS__();                                                           \
-      break;                                                                   \
-    }                                                                          \
-    case 3: {                                                                  \
-      constexpr int tensor_rank = 3;                                           \
-      __VA_ARGS__();                                                           \
-      break;                                                                   \
-    }                                                                          \
-    case 4: {                                                                  \
-      constexpr int tensor_rank = 4;                                           \
-      __VA_ARGS__();                                                           \
-      break;                                                                   \
-    }                                                                          \
-    default:                                                                   \
-      TORCH_CHECK(false, "Expects rank 2, 3 or 4 tensors but got ", NUM_DIMS); \
-  }
@@ -10,38 +10,16 @@
 namespace vllm {

 // TODO(woosuk): Further optimize this kernel.
-template <typename scalar_t, int VEC_SIZE, int NUM_DIMS>
+template <typename scalar_t, int VEC_SIZE>
 __global__ void rms_norm_kernel(
-    scalar_t* __restrict__ out,           // [..., hidden_size]
-    const scalar_t* __restrict__ input,   // [..., hidden_size]
-    const int64_t input_stride_d2,        // input.stride(-2)
-    const int64_t input_stride_d3,        // input.stride(-3)
-    const int64_t input_stride_d4,        // input.stride(-4)
-    const int64_t input_shape_d2,         // input.size(-2)
-    const int64_t input_shape_d3,         // input.size(-3)
+    scalar_t* __restrict__ out,          // [..., hidden_size]
+    const scalar_t* __restrict__ input,  // [..., hidden_size]
+    const int64_t input_stride,
    const scalar_t* __restrict__ weight,  // [hidden_size]
    const float epsilon, const int num_tokens, const int hidden_size) {
  __shared__ float s_variance;
  float variance = 0.0f;
-  const scalar_t* input_row;
-  if constexpr (NUM_DIMS == 2) {
-    // 2D for layernorm normal case [batch_size, hidden]
-    input_row = input + blockIdx.x * input_stride_d2;
-  } else if constexpr (NUM_DIMS == 3) {
-    // 3D for q/k norm [batch_size, num_heads, head_size]
-    int batch_idx = blockIdx.x / input_shape_d2;
-    int head_idx = blockIdx.x % input_shape_d2;
-    input_row =
-        input + batch_idx * input_stride_d3 + head_idx * input_stride_d2;
-  } else if constexpr (NUM_DIMS == 4) {
-    // 4D for transformers model_impl qk norm [batch, seq, head, head_dim]
-    int batch_idx = blockIdx.x / (input_shape_d3 * input_shape_d2);
-    int remaining = blockIdx.x % (input_shape_d3 * input_shape_d2);
-    int seq_idx = remaining / input_shape_d2;
-    int head_idx = remaining % input_shape_d2;
-    input_row = input + batch_idx * input_stride_d4 +
-                seq_idx * input_stride_d3 + head_idx * input_stride_d2;
-  }
+  const scalar_t* input_row = input + blockIdx.x * input_stride;

  auto vec_op = [&variance](const vec_n_t<scalar_t, VEC_SIZE>& vec) {
 #pragma unroll
@@ -186,44 +164,38 @@ void rms_norm(torch::Tensor& out,     // [..., hidden_size]
              torch::Tensor& weight,  // [hidden_size]
              double epsilon) {
  TORCH_CHECK(out.is_contiguous());
-  if (input.stride(-1) != 1) {
-    input = input.contiguous();
-  }
  TORCH_CHECK(input.stride(-1) == 1);
  TORCH_CHECK(weight.is_contiguous());

  int hidden_size = input.size(-1);

-  int num_tokens = input.numel() / hidden_size;
-  int num_dims = input.dim();
-  int64_t input_stride_d2 = input.stride(-2);
-  int64_t input_stride_d3 = (num_dims >= 3) ? input.stride(-3) : 0;
-  int64_t input_stride_d4 = (num_dims >= 4) ? input.stride(-4) : 0;
-  int64_t input_shape_d2 = (num_dims >= 3) ? input.size(-2) : 0;
-  int64_t input_shape_d3 = (num_dims >= 4) ? input.size(-3) : 0;
+  // We cannot just use `input.stride(-2)` if the tensor is not row-major.
+  // Instead, we use a 2d view to get the second-innermost stride.
+  // That way the dimensions (except the last one) can be arbitrarily permuted.
+  torch::Tensor input_view = input.view({-1, hidden_size});
+
+  int num_tokens = input_view.numel() / hidden_size;
+  int64_t input_stride = input_view.stride(-2);

  // For large num_tokens, use smaller blocks to increase SM concurrency.
  const int max_block_size = (num_tokens < 256) ? 1024 : 256;
  dim3 grid(num_tokens);
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(input_view));
  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-  VLLM_DISPATCH_RANK234(num_dims, [&] {
-    VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "rms_norm_kernel", [&] {
-      const int calculated_vec_size =
-          std::gcd(16 / sizeof(scalar_t), hidden_size);
-      const int block_size =
-          std::min(hidden_size / calculated_vec_size, max_block_size);
-      dim3 block(block_size);
-      VLLM_DISPATCH_VEC_SIZE(calculated_vec_size, [&] {
-        vllm::rms_norm_kernel<scalar_t, vec_size, tensor_rank>
-            <<<grid, block, 0, stream>>>(
-                out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(),
-                input_stride_d2, input_stride_d3, input_stride_d4,
-                input_shape_d2, input_shape_d3, weight.data_ptr<scalar_t>(),
-                epsilon, num_tokens, hidden_size);
+  VLLM_DISPATCH_FLOATING_TYPES(
+      input_view.scalar_type(), "rms_norm_kernel", [&] {
+        const int calculated_vec_size =
+            std::gcd(16 / sizeof(scalar_t), hidden_size);
+        const int block_size =
+            std::min(hidden_size / calculated_vec_size, max_block_size);
+        dim3 block(block_size);
+        VLLM_DISPATCH_VEC_SIZE(calculated_vec_size, [&] {
+          vllm::rms_norm_kernel<scalar_t, vec_size><<<grid, block, 0, stream>>>(
+              out.data_ptr<scalar_t>(), input_view.data_ptr<scalar_t>(),
+              input_stride, weight.data_ptr<scalar_t>(), epsilon, num_tokens,
+              hidden_size);
+        });
      });
-    });
-  });
 }

 #define LAUNCH_FUSED_ADD_RMS_NORM(width)                                    \
@@ -1,2 +1 @@
-sm*_kernel_*.cu
-kernel_selector.h
+kernel_*.cu
@@ -4,282 +4,134 @@ import glob
 import itertools
 import os
 import subprocess
-import sys

 import jinja2

-ARCHS = []
-SUPPORT_FP8 = False
-for arch in sys.argv[1].split(","):
-    arch = arch[: arch.index(".") + 2].replace(".", "")
-    arch = int(arch)
-    # 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 in [89, 120]:
-        SUPPORT_FP8 = True
-
-FILE_HEAD_COMMENT = """
-// auto generated by generate_kernels.py
+FILE_HEAD = """
+// auto generated by generate.py
 // clang-format off
-""".lstrip()

-FILE_HEAD = (
-    FILE_HEAD_COMMENT
-    + """
 #include "kernel.h"
 #include "marlin_template.h"

 namespace MARLIN_NAMESPACE_NAME {
-"""
-)
+""".strip()

 TEMPLATE = (
    "template __global__ void Marlin<"
-    "{{a_type_id}}, "
-    "{{b_type_id}}, "
-    "{{c_type_id}}, "
+    "{{scalar_t}}, "
+    "{{w_type_id}}, "
    "{{s_type_id}}, "
    "{{threads}}, "
    "{{thread_m_blocks}}, "
    "{{thread_n_blocks}}, "
    "{{thread_k_blocks}}, "
-    "{{m_block_size_8}}, "
+    "{{'true' if m_block_size_8 else 'false'}}, "
    "{{stages}}, "
    "{{group_blocks}}, "
-    "{{is_zp_float}}>"
+    "{{'true' if is_zp_float else 'false'}}>"
    "( MARLIN_KERNEL_PARAMS );"
 )

+# int8 with zero point case (vllm::kU8) is also supported,
+# we don't add it to reduce wheel size.
+SCALAR_TYPES = [
+    "vllm::kU4",
+    "vllm::kU4B8",
+    "vllm::kU8B128",
+    "vllm::kFE4M3fn",
+    "vllm::kFE2M1f",
+]
 THREAD_CONFIGS = [(128, 128, 256), (64, 256, 256), (64, 128, 128)]

 THREAD_M_BLOCKS = [0.5, 1, 2, 3, 4]
-
-QUANT_CONFIGS = [
-    # AWQ-INT4
-    {
-        "b_type": "kU4",
-        "thread_configs": THREAD_CONFIGS,
-        "thread_m_blocks": THREAD_M_BLOCKS,
-        "group_blocks": [-1, 2, 4, 8],
-    },
-    # GPTQ-INT4
-    {
-        "b_type": "kU4B8",
-        "thread_configs": THREAD_CONFIGS,
-        "thread_m_blocks": THREAD_M_BLOCKS,
-        "group_blocks": [-1, 0, 2, 4, 8],
-    },
-    # AWQ-INT8
-    {
-        "b_type": "kU8B128",
-        "thread_configs": THREAD_CONFIGS,
-        "thread_m_blocks": THREAD_M_BLOCKS,
-        "group_blocks": [-1, 0, 2, 4, 8],
-    },
-    # FP8
-    {
-        "b_type": "kFE4M3fn",
-        "thread_configs": THREAD_CONFIGS,
-        "thread_m_blocks": THREAD_M_BLOCKS,
-        "group_blocks": [-1, 8],
-    },
-    # NVFP4
-    {
-        "b_type": "kFE2M1f",
-        "s_type": "kFE4M3fn",
-        "thread_configs": THREAD_CONFIGS,
-        "thread_m_blocks": THREAD_M_BLOCKS,
-        "group_blocks": [1],
-    },
-    # MXFP4
-    {
-        "a_type": ["kBFloat16"],
-        "b_type": "kFE2M1f",
-        "s_type": "kFE8M0fnu",
-        "thread_configs": THREAD_CONFIGS,
-        "thread_m_blocks": THREAD_M_BLOCKS,
-        "group_blocks": [2],
-    },
-    # AWQ-INT4 with INT8 activation
-    {
-        "a_type": ["kS8"],
-        "b_type": "kU4",
-        "thread_configs": THREAD_CONFIGS,
-        "thread_m_blocks": [1, 2, 3, 4],
-        "group_blocks": [-1, 2, 4, 8],
-    },
-    # GPTQ-INT4 with INT8 activation
-    {
-        "a_type": ["kS8"],
-        "b_type": "kU4B8",
-        "thread_configs": THREAD_CONFIGS,
-        "thread_m_blocks": [1, 2, 3, 4],
-        "group_blocks": [-1, 2, 4, 8],
-    },
-    # GPTQ-INT4 with FP8 activation
-    {
-        "a_type": ["kFE4M3fn"],
-        "b_type": "kU4B8",
-        "thread_configs": THREAD_CONFIGS,
-        "thread_m_blocks": [1, 2, 3, 4],
-        "group_blocks": [-1, 2, 4, 8],
-    },
-    # AWQ-INT4 with FP8 activation
-    {
-        "a_type": ["kFE4M3fn"],
-        "b_type": "kU4",
-        "thread_configs": THREAD_CONFIGS,
-        "thread_m_blocks": [1, 2, 3, 4],
-        "group_blocks": [-1, 2, 4, 8],
-    },
-    # MXFP4 with FP8 activation
-    {
-        "a_type": ["kFE4M3fn"],
-        "b_type": "kFE2M1f",
-        "c_type": ["kBFloat16"],
-        "s_type": "kFE8M0fnu",
-        "thread_configs": THREAD_CONFIGS,
-        "thread_m_blocks": [1, 2, 3, 4],
-        "group_blocks": [2],
-    },
-]
+# group_blocks:
+#   = 0 : act order case
+#   = -1 : channelwise quantization
+#   > 0 : group_size=16*group_blocks
+GROUP_BLOCKS = [0, -1, 1, 2, 4, 8]
+DTYPES = ["fp16", "bf16"]


 def remove_old_kernels():
-    for filename in glob.glob(os.path.dirname(__file__) + "/*kernel_*.cu"):
+    for filename in glob.glob(os.path.dirname(__file__) + "/kernel_*.cu"):
        subprocess.call(["rm", "-f", filename])

-    filename = os.path.dirname(__file__) + "/kernel_selector.h"
-    subprocess.call(["rm", "-f", filename])
-

 def generate_new_kernels():
-    result_dict = {}
-
-    for quant_config in QUANT_CONFIGS:
-        c_types = quant_config.get("c_type", ["kFloat16", "kBFloat16"])
-        a_types = quant_config.get("a_type", ["kFloat16", "kBFloat16"])
-        b_type = quant_config["b_type"]
-        all_group_blocks = quant_config["group_blocks"]
-        all_m_blocks = quant_config["thread_m_blocks"]
-        all_thread_configs = quant_config["thread_configs"]
-
-        for a_type, c_type in itertools.product(a_types, c_types):
-            if not SUPPORT_FP8 and a_type == "kFE4M3fn":
-                continue
-            if "16" in a_type and "16" in c_type and a_type != c_type:
-                continue
-            s_type = quant_config.get("s_type", c_type)
-            if (a_type, b_type, c_type) not in result_dict:
-                result_dict[(a_type, b_type, c_type)] = []
-
-            for group_blocks, m_blocks, thread_configs in itertools.product(
-                all_group_blocks, all_m_blocks, all_thread_configs
-            ):
-                thread_k, thread_n, threads = thread_configs
-
-                if threads == 256:
-                    # for small batch (m_blocks == 1),
-                    #     we only need (128, 128, 256)
-                    # for large batch (m_blocks > 1),
-                    #     we only need (64, 256, 256)
-                    if m_blocks <= 1 and (thread_k, thread_n) != (128, 128):
-                        continue
-                    if m_blocks > 1 and (thread_k, thread_n) != (64, 256):
-                        continue
-
-                config = {
-                    "threads": threads,
-                    "s_type": s_type,
-                    "thread_m_blocks": max(m_blocks, 1),
-                    "thread_k_blocks": thread_k // 16,
-                    "thread_n_blocks": thread_n // 16,
-                    "m_block_size_8": "true" if m_blocks == 0.5 else "false",
-                    "stages": "pipe_stages",
-                    "group_blocks": group_blocks,
-                    "is_zp_float": "false",
-                }
-
-                result_dict[(a_type, b_type, c_type)].append(config)
-
-    kernel_selector_str = FILE_HEAD_COMMENT
-
-    for (a_type, b_type, c_type), config_list in result_dict.items():
+    for scalar_type, dtype in itertools.product(SCALAR_TYPES, DTYPES):
        all_template_str_list = []
-        for config in config_list:
-            s_type = config["s_type"]
+
+        for group_blocks, m_blocks, thread_configs in itertools.product(
+            GROUP_BLOCKS, THREAD_M_BLOCKS, THREAD_CONFIGS
+        ):
+            # act order case only support gptq-int4 and gptq-int8
+            if group_blocks == 0 and scalar_type not in [
+                "vllm::kU4B8",
+                "vllm::kU8B128",
+            ]:
+                continue
+            if thread_configs[2] == 256:
+                # for small batch (m_blocks == 1), we only need (128, 128, 256)
+                # for large batch (m_blocks > 1), we only need (64, 256, 256)
+                if m_blocks <= 1 and thread_configs[0] != 128:
+                    continue
+                if m_blocks > 1 and thread_configs[0] != 64:
+                    continue
+
+            # we only support channelwise quantization and group_size == 128
+            # for fp8
+            if scalar_type == "vllm::kFE4M3fn" and group_blocks not in [-1, 8]:
+                continue
+            # nvfp4 only supports group_size == 16
+            # mxfp4 only supports group_size == 32
+            if scalar_type == "vllm::kFE2M1f" and group_blocks not in [1, 2]:
+                continue
+            # other quantization methods don't support group_size = 16
+            if scalar_type != "vllm::kFE2M1f" and group_blocks == 1:
+                continue
+
+            k_blocks = thread_configs[0] // 16
+            n_blocks = thread_configs[1] // 16
+            threads = thread_configs[2]
+
+            c_dtype = "half" if dtype == "fp16" else "nv_bfloat16"
+
+            if scalar_type == "vllm::kFE2M1f" and group_blocks == 1:
+                s_type = "vllm::kFE4M3fn"
+            elif scalar_type == "vllm::kFE2M1f" and group_blocks == 2:
+                s_type = "vllm::kFE8M0fnu"
+                if dtype == "fp16":
+                    # we cannot safely dequantize e8m0 to fp16, so skip this
+                    continue
+            elif dtype == "fp16":
+                s_type = "vllm::kFloat16"
+            elif dtype == "bf16":
+                s_type = "vllm::kBFloat16"
+
            template_str = jinja2.Template(TEMPLATE).render(
-                a_type_id=f"vllm::{a_type}.id()",
-                b_type_id=f"vllm::{b_type}.id()",
-                c_type_id=f"vllm::{c_type}.id()",
-                s_type_id=f"vllm::{s_type}.id()",
-                **config,
+                scalar_t=c_dtype,
+                w_type_id=scalar_type + ".id()",
+                s_type_id=s_type + ".id()",
+                threads=threads,
+                thread_m_blocks=max(m_blocks, 1),
+                thread_n_blocks=n_blocks,
+                thread_k_blocks=k_blocks,
+                m_block_size_8=m_blocks == 0.5,
+                stages="pipe_stages",
+                group_blocks=group_blocks,
+                is_zp_float=False,
            )
+
            all_template_str_list.append(template_str)

-            conditions = [
-                f"a_type == vllm::{a_type}",
-                f"b_type == vllm::{b_type}",
-                f"c_type == vllm::{c_type}",
-                f"s_type == vllm::{s_type}",
-                f"threads == {config['threads']}",
-                f"thread_m_blocks == {config['thread_m_blocks']}",
-                f"thread_n_blocks == {config['thread_n_blocks']}",
-                f"thread_k_blocks == {config['thread_k_blocks']}",
-                f"m_block_size_8 == {config['m_block_size_8']}",
-                f"group_blocks == {config['group_blocks']}",
-                f"is_zp_float == {config['is_zp_float']}",
-            ]
-            conditions = " && ".join(conditions)
-
-            if kernel_selector_str == FILE_HEAD_COMMENT:
-                kernel_selector_str += f"if ({conditions})\n  kernel = "
-            else:
-                kernel_selector_str += f"else if ({conditions})\n  kernel = "
-
-            kernel_template2 = (
-                "Marlin<{{a_type_id}}, {{b_type_id}}, {{c_type_id}}, "
-                "{{s_type_id}}, {{threads}}, {{thread_m_blocks}}, "
-                "{{thread_n_blocks}}, {{thread_k_blocks}}, "
-                "{{m_block_size_8}}, {{stages}}, {{group_blocks}}, "
-                "{{is_zp_float}}>;"
-            )
-
-            kernel_selector_str += (
-                jinja2.Template(kernel_template2).render(
-                    a_type_id=f"vllm::{a_type}.id()",
-                    b_type_id=f"vllm::{b_type}.id()",
-                    c_type_id=f"vllm::{c_type}.id()",
-                    s_type_id=f"vllm::{s_type}.id()",
-                    **config,
-                )
-                + "\n"
-            )
-
        file_content = FILE_HEAD + "\n\n"
        file_content += "\n\n".join(all_template_str_list) + "\n\n}\n"
-        if a_type == "kFE4M3fn":
-            filename = f"sm89_kernel_{a_type[1:]}_{b_type[1:]}_{c_type[1:]}.cu"
-        else:
-            filename = f"sm80_kernel_{a_type[1:]}_{b_type[1:]}_{c_type[1:]}.cu"
-
-        filename = filename.lower()
+        filename = f"kernel_{dtype}_{scalar_type[6:].lower()}.cu"

        with open(os.path.join(os.path.dirname(__file__), filename), "w") as f:
            f.write(file_content)

-    if not SUPPORT_FP8 and kernel_selector_str != FILE_HEAD_COMMENT:
-        kernel_selector_str += (
-            "else if (a_type == vllm::kFE4M3fn)\n"
-            "  TORCH_CHECK(false, "
-            '"marlin kernel with fp8 activation is not built.");'
-        )
-
-    with open(os.path.join(os.path.dirname(__file__), "kernel_selector.h"), "w") as f:
-        f.write(kernel_selector_str)
-

 if __name__ == "__main__":
    remove_old_kernels()
@@ -11,9 +11,8 @@
  const int4 *__restrict__ A, const int4 *__restrict__ B,                     \
      int4 *__restrict__ C, int4 *__restrict__ C_tmp,                         \
      const int4 *__restrict__ b_bias_ptr,                                    \
-      const float *__restrict__ a_scales_ptr,                                 \
      const int4 *__restrict__ scales_ptr,                                    \
-      const uint16_t *__restrict__ global_scale_ptr,                          \
+      const uint16_t *__restrict__ scale2_ptr,                                \
      const int4 *__restrict__ zp_ptr, const int *__restrict__ g_idx,         \
      const int32_t *__restrict__ sorted_token_ids_ptr,                       \
      const int32_t *__restrict__ expert_ids_ptr,                             \
@@ -21,13 +20,12 @@
      const float *__restrict__ topk_weights_ptr, int top_k,                  \
      bool mul_topk_weights, bool is_ep, int num_groups, int prob_m,          \
      int prob_n, int prob_k, int *locks, bool has_bias, bool use_atomic_add, \
-      bool use_fp32_reduce
+      bool use_fp32_reduce, int max_shared_mem

 namespace MARLIN_NAMESPACE_NAME {
-template <const vllm::ScalarTypeId a_type_id,  // A ScalarType id
-          const vllm::ScalarTypeId b_type_id,  // B ScalarType id
-          const vllm::ScalarTypeId c_type_id,  // C ScalarType id
-          const vllm::ScalarTypeId s_type_id,  // B_SCALE ScalarType id
+template <typename scalar_t,  // compute dtype, half or nv_float16
+          const vllm::ScalarTypeId w_type_id,  // weight ScalarType id
+          const vllm::ScalarTypeId s_type_id,  // weight scale ScalarType id
          const int threads,          // number of threads in a threadblock
          const int thread_m_blocks,  // number of 16x16 blocks in the m
                                      // dimension (batchsize) of the
@@ -37,6 +37,39 @@ __global__ void MarlinDefault(MARLIN_KERNEL_PARAMS){};

 using MarlinFuncPtr = void (*)(MARLIN_KERNEL_PARAMS);

+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
+
+template <int moe_block_size>
+__global__ void permute_cols_kernel(
+    int4 const* __restrict__ a_int4_ptr, int const* __restrict__ perm_int_ptr,
+    int4* __restrict__ out_int4_ptr,
+    const int32_t* __restrict__ sorted_token_ids_ptr,
+    const int32_t* __restrict__ expert_ids_ptr,
+    const int32_t* __restrict__ num_tokens_past_padded_ptr, int size_m,
+    int size_k, int top_k) {};
+
+}  // namespace marlin
+
+torch::Tensor moe_wna16_marlin_gemm(
+    torch::Tensor& a, std::optional<torch::Tensor> c_or_none,
+    torch::Tensor& b_q_weight,
+    std::optional<torch::Tensor> const& b_bias_or_none, torch::Tensor& b_scales,
+    std::optional<torch::Tensor> const& b_zeros_or_none,
+    std::optional<torch::Tensor> const& g_idx_or_none,
+    std::optional<torch::Tensor> const& perm_or_none, torch::Tensor& workspace,
+    torch::Tensor& sorted_token_ids, torch::Tensor& expert_ids,
+    torch::Tensor& num_tokens_past_padded, torch::Tensor& topk_weights,
+    int64_t moe_block_size, int64_t top_k, bool mul_topk_weights, bool is_ep,
+    vllm::ScalarTypeId const& b_q_type_id, int64_t size_m, int64_t size_n,
+    int64_t size_k, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce,
+    bool is_zp_float) {
+  TORCH_CHECK_NOT_IMPLEMENTED(false,
+                              "marlin_gemm(..) requires CUDA_ARCH >= 8.0");
+  return torch::empty({1, 1});
+}
+
+#else
+
 // For a given "a" of size [M,K] performs a permutation of the K columns based
 // on the given "perm" indices.
 template <int moe_block_size>
@@ -174,7 +207,7 @@ int get_kernel_cache_size(thread_config_t const& th_config, bool m_block_size_8,
                          int thread_m_blocks, int prob_m, int prob_n,
                          int prob_k, int num_bits, int group_size,
                          bool has_act_order, bool is_k_full, int has_zp,
-                          int is_zp_float, bool is_a_8bit) {
+                          int is_zp_float) {
  int pack_factor = 32 / num_bits;

  // Get B size
@@ -184,8 +217,8 @@ int get_kernel_cache_size(thread_config_t const& th_config, bool m_block_size_8,

  // shm size for block_sorted_ids/rd_block_sorted_ids/block_topk_weights
  // both of them requires tb_m * 4 bytes (tb_m * int32 or tb_m * float32)
-  int sh_block_meta_size = tb_m * 16;
-  int sh_a_size = pipe_stages * (tb_m * tb_k) * (is_a_8bit ? 1 : 2);
+  int sh_block_meta_size = tb_m * 4;
+  int sh_a_size = pipe_stages * (tb_m * tb_k) * 2;
  int sh_b_size = pipe_stages * (tb_k * tb_n / pack_factor) * 4;
  int sh_red_size = tb_m * (tb_n + 8) * 2;
  int sh_bias_size = tb_n * 2;
@@ -217,7 +250,7 @@ bool is_valid_config(thread_config_t const& th_config, bool m_block_size_8,
                     int thread_m_blocks, int prob_m, int prob_n, int prob_k,
                     int num_bits, int group_size, bool has_act_order,
                     bool is_k_full, int has_zp, int is_zp_float,
-                     int max_shared_mem, bool is_a_8bit) {
+                     int max_shared_mem) {
  // Sanity
  if (th_config.thread_k == -1 || th_config.thread_n == -1 ||
      th_config.num_threads == -1) {
@@ -240,34 +273,188 @@ bool is_valid_config(thread_config_t const& th_config, bool m_block_size_8,
  }

  // Check that pipeline fits into cache
-  int cache_size =
-      get_kernel_cache_size(th_config, m_block_size_8, thread_m_blocks, prob_m,
-                            prob_n, prob_k, num_bits, group_size, has_act_order,
-                            is_k_full, has_zp, is_zp_float, is_a_8bit);
-  return cache_size <= max_shared_mem;
+  int cache_size = get_kernel_cache_size(
+      th_config, m_block_size_8, thread_m_blocks, prob_m, prob_n, prob_k,
+      num_bits, group_size, has_act_order, is_k_full, has_zp, is_zp_float);
+  return cache_size + 512 <= max_shared_mem;
 }

-MarlinFuncPtr get_marlin_kernel(
-    const vllm::ScalarType a_type, const vllm::ScalarType b_type,
-    const vllm::ScalarType c_type, const vllm::ScalarType s_type,
-    int thread_m_blocks, int thread_n_blocks, int thread_k_blocks,
-    bool m_block_size_8, bool has_act_order, bool has_zp, int group_blocks,
-    int threads, bool is_zp_float) {
-  int num_bits = b_type.size_bits();
-  auto kernel = MarlinDefault;
+  #define _GET_IF(W_TYPE, THREAD_M_BLOCKS, THREAD_N_BLOCKS, THREAD_K_BLOCKS,   \
+                  M_BLOCK_SIZE_8, GROUP_BLOCKS, NUM_THREADS, IS_ZP_FLOAT)      \
+    else if (q_type == W_TYPE && thread_m_blocks == THREAD_M_BLOCKS &&         \
+             thread_n_blocks == THREAD_N_BLOCKS &&                             \
+             thread_k_blocks == THREAD_K_BLOCKS &&                             \
+             m_block_size_8 == M_BLOCK_SIZE_8 &&                               \
+             group_blocks == GROUP_BLOCKS && num_threads == NUM_THREADS &&     \
+             is_zp_float == IS_ZP_FLOAT) {                                     \
+      constexpr auto S_TYPE =                                                  \
+          W_TYPE == vllm::kFE2M1f                                              \
+              ? (GROUP_BLOCKS == 1 ? vllm::kFE4M3fn : vllm::kFE8M0fnu)         \
+              : (std::is_same<scalar_t, half>::value ? vllm::kFloat16          \
+                                                     : vllm::kBFloat16);       \
+      kernel = Marlin<scalar_t, W_TYPE.id(), S_TYPE.id(), NUM_THREADS,         \
+                      THREAD_M_BLOCKS, THREAD_N_BLOCKS, THREAD_K_BLOCKS,       \
+                      M_BLOCK_SIZE_8, pipe_stages, GROUP_BLOCKS, IS_ZP_FLOAT>; \
+    }

-#include "kernel_selector.h"
+  // COMMON: cases for (group_blocks in [-1, 2, 4, 8] and is_zp_float == false)
+  //         this is the most common cases
+  // BIGGROUP: cases for big group size (group_blocks in [-1, 8])
+  // FZP: cases for float-zero-point (is_zp_float = true)
+  // ACT: cases for act order case (group_blocks == 0)
+  // FP4: cases for nvfp4(e2m1) (group_blocks == 1)
+  #define COMMON_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)       \
+    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, -1, NUM_THREADS, false)  \
+    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 2, NUM_THREADS, false)   \
+    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 4, NUM_THREADS, false)   \
+    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 8, NUM_THREADS, false)   \
+    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
+    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false)  \
+    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false)  \
+    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)
+
+  #define COMMON_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)     \
+    _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
+    _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false)  \
+    _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false)  \
+    _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)  \
+                                                                          \
+    _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
+    _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false)  \
+    _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false)  \
+    _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)  \
+                                                                          \
+    _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
+    _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false)  \
+    _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false)  \
+    _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)
+
+  #define COMMON_GET_IF(W_TYPE)            \
+    COMMON_GET_IF_M1(W_TYPE, 8, 8, 256)    \
+    COMMON_GET_IF_M1(W_TYPE, 8, 4, 128)    \
+    COMMON_GET_IF_M234(W_TYPE, 16, 4, 256) \
+    COMMON_GET_IF_M234(W_TYPE, 8, 4, 128)
+
+  #define BIGGROUP_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)     \
+    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, -1, NUM_THREADS, false)  \
+    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 8, NUM_THREADS, false)   \
+    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
+    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)
+
+  #define BIGGROUP_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)   \
+    _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
+    _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)  \
+    _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
+    _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)  \
+    _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
+    _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)
+
+  #define BIGGROUP_GET_IF(W_TYPE)            \
+    BIGGROUP_GET_IF_M1(W_TYPE, 8, 8, 256)    \
+    BIGGROUP_GET_IF_M1(W_TYPE, 8, 4, 128)    \
+    BIGGROUP_GET_IF_M234(W_TYPE, 16, 4, 256) \
+    BIGGROUP_GET_IF_M234(W_TYPE, 8, 4, 128)
+
+  #define NVFP4_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)      \
+    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 1, NUM_THREADS, false) \
+    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false)
+
+  #define NVFP4_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)     \
+    _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false) \
+    _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false) \
+    _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false)
+
+  #define NVFP4_GET_IF(W_TYPE)            \
+    NVFP4_GET_IF_M1(W_TYPE, 8, 8, 256)    \
+    NVFP4_GET_IF_M1(W_TYPE, 8, 4, 128)    \
+    NVFP4_GET_IF_M234(W_TYPE, 16, 4, 256) \
+    NVFP4_GET_IF_M234(W_TYPE, 8, 4, 128)
+
+  #define MXFP4_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)      \
+    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 2, NUM_THREADS, false) \
+    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false)
+
+  #define MXFP4_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)     \
+    _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \
+    _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \
+    _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false)
+
+  #define MXFP4_GET_IF(W_TYPE)            \
+    MXFP4_GET_IF_M1(W_TYPE, 8, 8, 256)    \
+    MXFP4_GET_IF_M1(W_TYPE, 8, 4, 128)    \
+    MXFP4_GET_IF_M234(W_TYPE, 16, 4, 256) \
+    MXFP4_GET_IF_M234(W_TYPE, 8, 4, 128)
+
+  // We currently have 4-bit models only with group_blocks == 4
+  #define FZP_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)       \
+    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 4, NUM_THREADS, true) \
+    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true)
+
+  #define FZP_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)      \
+    _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true) \
+    _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true) \
+    _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true)
+
+  #define FZP_GET_IF(W_TYPE)            \
+    FZP_GET_IF_M1(W_TYPE, 8, 8, 256)    \
+    FZP_GET_IF_M1(W_TYPE, 8, 4, 128)    \
+    FZP_GET_IF_M234(W_TYPE, 16, 4, 256) \
+    FZP_GET_IF_M234(W_TYPE, 8, 4, 128)
+
+  // We currently have 4-bit models only with group_blocks == 4
+  #define ACT_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)        \
+    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 0, NUM_THREADS, false) \
+    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false)
+
+  #define ACT_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)       \
+    _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false) \
+    _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false) \
+    _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false)
+
+  #define ACT_GET_IF(W_TYPE)            \
+    ACT_GET_IF_M1(W_TYPE, 8, 8, 256)    \
+    ACT_GET_IF_M1(W_TYPE, 8, 4, 128)    \
+    ACT_GET_IF_M234(W_TYPE, 16, 4, 256) \
+    ACT_GET_IF_M234(W_TYPE, 8, 4, 128)
+
+template <typename scalar_t>
+MarlinFuncPtr get_marlin_kernel(const vllm::ScalarType q_type,
+                                int thread_m_blocks, int thread_n_blocks,
+                                int thread_k_blocks, bool m_block_size_8,
+                                bool has_act_order, bool has_zp,
+                                int group_blocks, int num_threads,
+                                bool is_zp_float) {
+  int num_bits = q_type.size_bits();
+  auto kernel = MarlinDefault;
+  if (false) {
+  }
+
+  COMMON_GET_IF(vllm::kU4)
+  COMMON_GET_IF(vllm::kU4B8)
+  COMMON_GET_IF(vllm::kU8B128)
+
+  NVFP4_GET_IF(vllm::kFE2M1f)
+
+  BIGGROUP_GET_IF(vllm::kFE4M3fn)
+
+  ACT_GET_IF(vllm::kU4B8)
+  ACT_GET_IF(vllm::kU8B128)
+  if (std::is_same<scalar_t, nv_bfloat16>::value) {
+    if (false) {
+    }
+    MXFP4_GET_IF(vllm::kFE2M1f)
+  }

  return kernel;
 }

-exec_config_t determine_exec_config(
-    const vllm::ScalarType& a_type, const vllm::ScalarType& b_type,
-    const vllm::ScalarType& c_type, const vllm::ScalarType& s_type, int prob_m,
-    int prob_n, int prob_k, int num_experts, int top_k, int thread_m_blocks,
-    bool m_block_size_8, int num_bits, int group_size, bool has_act_order,
-    bool is_k_full, bool has_zp, bool is_zp_float, int max_shared_mem, int sms,
-    bool is_a_8bit) {
+template <typename scalar_t>
+exec_config_t determine_exec_config(const vllm::ScalarType& q_type, int prob_m,
+                                    int prob_n, int prob_k, int thread_m_blocks,
+                                    bool m_block_size_8, int num_bits,
+                                    int group_size, bool has_act_order,
+                                    bool is_k_full, bool has_zp,
+                                    bool is_zp_float, int max_shared_mem) {
  exec_config_t exec_cfg = exec_config_t{1, thread_config_t{-1, -1, -1}};
  thread_config_t* thread_configs = thread_m_blocks > 1
                                        ? large_batch_thread_configs
@@ -284,69 +471,73 @@ exec_config_t determine_exec_config(

    if (!is_valid_config(th_config, m_block_size_8, thread_m_blocks, prob_m,
                         prob_n, prob_k, num_bits, group_size, has_act_order,
-                         is_k_full, has_zp, is_zp_float, max_shared_mem - 512,
-                         is_a_8bit)) {
+                         is_k_full, has_zp, is_zp_float, max_shared_mem)) {
      continue;
    }

    int cache_size = get_kernel_cache_size(
        th_config, m_block_size_8, thread_m_blocks, prob_m, prob_n, prob_k,
-        num_bits, group_size, has_act_order, is_k_full, has_zp, is_zp_float,
-        is_a_8bit);
+        num_bits, group_size, has_act_order, is_k_full, has_zp, is_zp_float);

    int group_blocks = 0;
    if (!has_act_order) {
      group_blocks = group_size == -1 ? -1 : (group_size / 16);
    }

-    auto kernel =
-        get_marlin_kernel(a_type, b_type, c_type, s_type, thread_m_blocks,
-                          th_config.thread_n / 16, th_config.thread_k / 16,
-                          m_block_size_8, has_act_order, has_zp, group_blocks,
-                          th_config.num_threads, is_zp_float);
+    auto kernel = get_marlin_kernel<scalar_t>(
+        q_type, thread_m_blocks, th_config.thread_n / 16,
+        th_config.thread_k / 16, m_block_size_8, has_act_order, has_zp,
+        group_blocks, th_config.num_threads, is_zp_float);

    if (kernel == MarlinDefault) continue;

-    cudaFuncAttributes attr;
-    cudaFuncGetAttributes(&attr, kernel);
-    int reg_size = max(attr.numRegs, 1) * th_config.num_threads * 4;
-    int allow_count = min(device_max_reg_size / reg_size,
-                          max_shared_mem / (cache_size + 1536));
-    if (thread_m_blocks == 1)
+    if (thread_m_blocks > 1) {
+      exec_cfg = {1, th_config};
+      break;
+    } else {
+      cudaFuncAttributes attr;
+      cudaFuncGetAttributes(&attr, kernel);
+      int reg_size = max(attr.numRegs, 1) * th_config.num_threads * 4;
+      int allow_count = min(device_max_reg_size / reg_size,
+                            max_shared_mem / (cache_size + 1024));
      allow_count = max(min(allow_count, 4), 1);
-    else
-      allow_count = max(min(allow_count, 2), 1);
-
-    if (prob_n / th_config.thread_n * prob_m * top_k * 4 < sms * allow_count) {
-      allow_count =
-          max(prob_n / th_config.thread_n * prob_m * top_k * 4 / sms, 1);
+      if (allow_count > count) {
+        count = allow_count;
+        exec_cfg = {count, th_config};
+      };
    }
-
-    if (allow_count > count) {
-      count = allow_count;
-      exec_cfg = {count, th_config};
-    };
  }

  return exec_cfg;
 }

+template <typename scalar_t>
 void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
-               void* a_s, void* b_s, void* g_s, void* zp, void* g_idx,
-               void* perm, void* a_tmp, void* sorted_token_ids,
-               void* expert_ids, void* num_tokens_past_padded,
-               void* topk_weights, int moe_block_size, int num_experts,
-               int top_k, bool mul_topk_weights, bool is_ep, int prob_m,
-               int prob_n, int prob_k, void* workspace,
-               vllm::ScalarType const& a_type, vllm::ScalarType const& b_type,
-               vllm::ScalarType const& c_type, vllm::ScalarType const& s_type,
-               bool has_bias, bool has_act_order, bool is_k_full, bool has_zp,
-               int num_groups, int group_size, int dev, cudaStream_t stream,
-               int thread_k, int thread_n, int sms, int blocks_per_sm,
-               bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) {
+               void* s, void* s2, void* zp, void* g_idx, void* perm,
+               void* a_tmp, void* sorted_token_ids, void* expert_ids,
+               void* num_tokens_past_padded, void* topk_weights,
+               int moe_block_size, int top_k, bool mul_topk_weights, bool is_ep,
+               int prob_m, int prob_n, int prob_k, void* workspace,
+               vllm::ScalarType const& q_type, bool has_bias,
+               bool has_act_order, bool is_k_full, bool has_zp, int num_groups,
+               int group_size, int dev, cudaStream_t stream, int thread_k,
+               int thread_n, int sms, bool use_atomic_add, bool use_fp32_reduce,
+               bool is_zp_float) {
  int thread_m_blocks = div_ceil(moe_block_size, 16);
  bool m_block_size_8 = moe_block_size == 8;
-  bool is_a_8bit = a_type.size_bits() == 8;
+
+  if (has_zp) {
+    TORCH_CHECK(
+        q_type == vllm::kU4 || q_type == vllm::kU8,
+        "q_type must be u4 or u8 when has_zp = True. Got = ", q_type.str());
+  } else {
+    TORCH_CHECK(
+        q_type == vllm::kU4B8 || q_type == vllm::kU8B128 ||
+            q_type == vllm::kFE4M3fn || q_type == vllm::kFE2M1f,
+        "q_type must be uint4b8, uint8b128, float8_e4m3fn or float4_e2m1f when "
+        "has_zp = False. Got = ",
+        q_type.str());
+  }

  TORCH_CHECK(prob_m > 0 && prob_n > 0 && prob_k > 0, "Invalid MNK = [", prob_m,
              ", ", prob_n, ", ", prob_k, "]");
@@ -372,15 +563,14 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
    }
  }

-  int num_bits = b_type.size_bits();
+  int num_bits = q_type.size_bits();
  const int4* A_ptr = (const int4*)A;
  const int4* B_ptr = (const int4*)B;
  int4* C_ptr = (int4*)C;
  int4* C_tmp_ptr = (int4*)C_tmp;
  const int4* bias_ptr = (const int4*)b_bias;
-  const float* a_s_ptr = (const float*)a_s;
-  const int4* b_s_ptr = (const int4*)b_s;
-  const uint16_t* g_s_ptr = (const uint16_t*)g_s;
+  const int4* s_ptr = (const int4*)s;
+  const uint16_t* s2_ptr = (const uint16_t*)s2;
  const int4* zp_ptr = (const int4*)zp;
  const int* g_idx_ptr = (const int*)g_idx;
  const int* perm_ptr = (const int*)perm;
@@ -428,41 +618,22 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
                         cudaDevAttrMaxSharedMemoryPerBlockOptin, dev);
  TORCH_CHECK(max_shared_mem > 0);

-  int major_capability, minor_capability;
-  cudaDeviceGetAttribute(&major_capability, cudaDevAttrComputeCapabilityMajor,
-                         dev);
-  cudaDeviceGetAttribute(&minor_capability, cudaDevAttrComputeCapabilityMinor,
-                         dev);
-  TORCH_CHECK(major_capability * 10 + minor_capability >= 80,
-              "marlin kernel only support Ampere or newer GPUs.");
-  if (a_type == vllm::kFE4M3fn) {
-    TORCH_CHECK(major_capability * 10 + minor_capability >= 89,
-                "FP8 only support Ada Lovelace or newer GPUs.");
-    TORCH_CHECK(
-        major_capability * 10 + minor_capability == 89 ||
-            major_capability * 10 + minor_capability == 120,
-        "Marlin W4A8-FP8 only support SM89 or SM120 device (It is slower than "
-        "Marlin W4A16 on other devices).");
-  }
-
  // Set thread config
  exec_config_t exec_cfg;
  thread_config_t thread_tfg;
  if (thread_k != -1 && thread_n != -1) {
-    thread_tfg = thread_config_t{thread_k, thread_n, thread_k * thread_n / 64};
-    if (blocks_per_sm == -1) blocks_per_sm = 1;
-    exec_cfg = exec_config_t{blocks_per_sm, thread_tfg};
+    thread_tfg = thread_config_t{thread_k, thread_n, default_threads};
+    exec_cfg = exec_config_t{1, thread_tfg};
    TORCH_CHECK(prob_n % thread_n == 0, "prob_n = ", prob_n,
                " is not divisible by thread_n = ", thread_n);
    TORCH_CHECK(prob_k % thread_k == 0, "prob_k = ", prob_k,
                " is not divisible by thread_k = ", thread_k);
  } else {
    // Auto config
-    exec_cfg = determine_exec_config(
-        a_type, b_type, c_type, s_type, prob_m, prob_n, prob_k, num_experts,
-        top_k, thread_m_blocks, m_block_size_8, num_bits, group_size,
-        has_act_order, is_k_full, has_zp, is_zp_float, max_shared_mem, sms,
-        is_a_8bit);
+    exec_cfg = determine_exec_config<scalar_t>(
+        q_type, prob_m, prob_n, prob_k, thread_m_blocks, m_block_size_8,
+        num_bits, group_size, has_act_order, is_k_full, has_zp, is_zp_float,
+        max_shared_mem);
    thread_tfg = exec_cfg.tb_cfg;
  }

@@ -476,29 +647,22 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
  int thread_k_blocks = thread_k / 16;
  int thread_n_blocks = thread_n / 16;

-  TORCH_CHECK(is_valid_config(thread_tfg, m_block_size_8, thread_m_blocks,
-                              prob_m, prob_n, prob_k, num_bits, group_size,
-                              has_act_order, is_k_full, has_zp, is_zp_float,
-                              max_shared_mem, is_a_8bit),
-              "Invalid thread config: thread_m_blocks = ", thread_m_blocks,
-              ", thread_k = ", thread_tfg.thread_k,
-              ", thread_n = ", thread_tfg.thread_n,
-              ", num_threads = ", thread_tfg.num_threads, " for MKN = [",
-              prob_m, ", ", prob_k, ", ", prob_n, "] and num_bits = ", num_bits,
-              ", group_size = ", group_size,
-              ", has_act_order = ", has_act_order, ", is_k_full = ", is_k_full,
-              ", has_zp = ", has_zp, ", is_zp_float = ", is_zp_float,
-              ", max_shared_mem = ", max_shared_mem);
+  TORCH_CHECK(
+      is_valid_config(thread_tfg, m_block_size_8, thread_m_blocks, prob_m,
+                      prob_n, prob_k, num_bits, group_size, has_act_order,
+                      is_k_full, has_zp, is_zp_float, max_shared_mem),
+      "Invalid thread config: thread_m_blocks = ", thread_m_blocks,
+      ", thread_k = ", thread_tfg.thread_k,
+      ", thread_n = ", thread_tfg.thread_n,
+      ", num_threads = ", thread_tfg.num_threads, " for MKN = [", prob_m, ", ",
+      prob_k, ", ", prob_n, "] and num_bits = ", num_bits,
+      ", group_size = ", group_size, ", has_act_order = ", has_act_order,
+      ", is_k_full = ", is_k_full, ", has_zp = ", has_zp,
+      ", is_zp_float = ", is_zp_float, ", max_shared_mem = ", max_shared_mem);

-  int sh_cache_size =
-      get_kernel_cache_size(thread_tfg, m_block_size_8, thread_m_blocks, prob_m,
-                            prob_n, prob_k, num_bits, group_size, has_act_order,
-                            is_k_full, has_zp, is_zp_float, is_a_8bit);
-
-  auto kernel = get_marlin_kernel(
-      a_type, b_type, c_type, s_type, thread_m_blocks, thread_n_blocks,
-      thread_k_blocks, m_block_size_8, has_act_order, has_zp, group_blocks,
-      num_threads, is_zp_float);
+  auto kernel = get_marlin_kernel<scalar_t>(
+      q_type, thread_m_blocks, thread_n_blocks, thread_k_blocks, m_block_size_8,
+      has_act_order, has_zp, group_blocks, num_threads, is_zp_float);

  if (kernel == MarlinDefault) {
    TORCH_CHECK(false, "Unsupported shapes: MNK = [", prob_m, ", ", prob_n,
@@ -515,20 +679,19 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
  // avoid ">>>" being formatted to "> > >"
  // clang-format off
  kernel<<<blocks, num_threads, max_shared_mem, stream>>>(
-      A_ptr, B_ptr, C_ptr, C_tmp_ptr, bias_ptr, a_s_ptr, b_s_ptr, g_s_ptr, zp_ptr, g_idx_ptr,
+      A_ptr, B_ptr, C_ptr, C_tmp_ptr, bias_ptr, s_ptr, s2_ptr, zp_ptr, g_idx_ptr,
      sorted_token_ids_ptr, expert_ids_ptr, num_tokens_past_padded_ptr,
      topk_weights_ptr, top_k, mul_topk_weights, is_ep, num_groups, prob_m,
-      prob_n, prob_k, locks, has_bias, use_atomic_add, use_fp32_reduce);
+      prob_n, prob_k, locks, has_bias, use_atomic_add, use_fp32_reduce, max_shared_mem);
  // clang-format on
 }

 }  // namespace MARLIN_NAMESPACE_NAME

 torch::Tensor moe_wna16_marlin_gemm(
-    torch::Tensor& a, std::optional<torch::Tensor> c_or_none,
+    torch::Tensor& a, std::optional<torch::Tensor> const& c_or_none,
    torch::Tensor& b_q_weight,
    std::optional<torch::Tensor> const& b_bias_or_none, torch::Tensor& b_scales,
-    std::optional<torch::Tensor> const& a_scales_or_none,
    std::optional<torch::Tensor> const& global_scale_or_none,
    std::optional<torch::Tensor> const& b_zeros_or_none,
    std::optional<torch::Tensor> const& g_idx_or_none,
@@ -536,70 +699,11 @@ torch::Tensor moe_wna16_marlin_gemm(
    torch::Tensor& sorted_token_ids, torch::Tensor& expert_ids,
    torch::Tensor& num_tokens_past_padded, torch::Tensor& topk_weights,
    int64_t moe_block_size, int64_t top_k, bool mul_topk_weights, bool is_ep,
-    vllm::ScalarTypeId const& b_type_id, int64_t size_m, int64_t size_n,
+    vllm::ScalarTypeId const& b_q_type_id, int64_t size_m, int64_t size_n,
    int64_t size_k, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce,
-    bool is_zp_float, int64_t thread_k, int64_t thread_n,
-    int64_t blocks_per_sm) {
-  vllm::ScalarTypeId a_type_id, c_type_id, s_type_id;
-
-  auto c_dtype = a.dtype();
-  if (a.scalar_type() == at::ScalarType::Half) {
-    a_type_id = vllm::kFloat16.id();
-    c_type_id = vllm::kFloat16.id();
-  } else if (a.scalar_type() == at::ScalarType::BFloat16) {
-    a_type_id = vllm::kBFloat16.id();
-    c_type_id = vllm::kBFloat16.id();
-  } else {
-    c_dtype = b_scales.dtype();
-    if (b_scales.scalar_type() == at::ScalarType::Half) {
-      c_type_id = vllm::kFloat16.id();
-    } else if (b_scales.scalar_type() == at::ScalarType::BFloat16) {
-      c_type_id = vllm::kBFloat16.id();
-    } else {
-      c_type_id = vllm::kBFloat16.id();
-
-      TORCH_CHECK(c_or_none.has_value(), "c must be passed for W4A8-FP4");
-      torch::Tensor c = c_or_none.value();
-      c_dtype = c.dtype();
-
-      if (c.scalar_type() == at::ScalarType::Half) {
-        c_type_id = vllm::kFloat16.id();
-      } else if (c.scalar_type() == at::ScalarType::BFloat16) {
-        c_type_id = vllm::kBFloat16.id();
-      } else {
-        TORCH_CHECK(false, "unsupported c dtype");
-      }
-    }
-
-    if (a.scalar_type() == at::ScalarType::Float8_e4m3fn) {
-      a_type_id = vllm::kFE4M3fn.id();
-    } else if (a.scalar_type() == at::ScalarType::Char) {
-      a_type_id = vllm::kS8.id();
-    } else {
-      TORCH_CHECK(false, "unsupported `a` scalar_type");
-    }
-  }
-
-  s_type_id = c_type_id;
-  if (b_type_id == vllm::kFE2M1f.id()) {
-    if (b_scales.scalar_type() == at::ScalarType::Float8_e4m3fn) {
-      s_type_id = vllm::kFE4M3fn.id();
-    } else if (b_scales.scalar_type() == at::ScalarType::Float8_e8m0fnu) {
-      s_type_id = vllm::kFE8M0fnu.id();
-    } else {
-      TORCH_CHECK(false,
-                  "When b_type = float4_e2m1f, b_scale scalar type must be",
-                  "float8_e4m3fn (for NVFP4) or float8_e8m0fnu (for MXFP4).");
-    }
-  }
-
-  vllm::ScalarType a_type = vllm::ScalarType::from_id(a_type_id);
-  vllm::ScalarType b_type = vllm::ScalarType::from_id(b_type_id);
-  vllm::ScalarType c_type = vllm::ScalarType::from_id(c_type_id);
-  vllm::ScalarType s_type = vllm::ScalarType::from_id(s_type_id);
-
-  int pack_factor = 32 / b_type.size_bits();
-  int num_experts = b_q_weight.size(0);
+    bool is_zp_float) {
+  vllm::ScalarType const b_q_type = vllm::ScalarType::from_id(b_q_type_id);
+  int pack_factor = 32 / b_q_type.size_bits();

  if (moe_block_size != 8) {
    TORCH_CHECK(moe_block_size % 16 == 0,
@@ -641,27 +745,19 @@ torch::Tensor moe_wna16_marlin_gemm(
  TORCH_CHECK(b_scales.device().is_cuda(), "b_scales is not on GPU");
  TORCH_CHECK(b_scales.is_contiguous(), "b_scales is not contiguous");

-  torch::Tensor a_scales;
-  auto options = torch::TensorOptions().dtype(c_dtype).device(a.device());
-  auto options_fp32 =
-      torch::TensorOptions().dtype(at::kFloat).device(a.device());
-
-  if (a_scales_or_none.has_value()) {
-    a_scales = a_scales_or_none.value();
-    TORCH_CHECK(a_type.size_bits() == 8,
-                "a_scales can only be used for 8bit activation.");
-  } else {
-    a_scales = torch::empty({0}, options_fp32);
-    TORCH_CHECK(a_type.size_bits() != 8,
-                "the a_scales parameter must be passed for 8bit activation.");
-  }
-
+  // thread_k: `k` size of a thread_tile in `weights` (can usually be left as
+  // auto -1)
+  int thread_k = -1;
+  // thread_n: `n` size of a thread_tile in `weights` (can usually be left as
+  // auto -1)
+  int thread_n = -1;
  // sms: number of SMs to use for the kernel
  int sms = -1;
  cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, a.get_device());

  // Alloc buffers
  const at::cuda::OptionalCUDAGuard device_guard(device_of(a));
+  auto options = torch::TensorOptions().dtype(a.dtype()).device(a.device());
  torch::Tensor c;
  if (c_or_none.has_value()) {
    c = c_or_none.value();
@@ -678,6 +774,8 @@ torch::Tensor moe_wna16_marlin_gemm(

  // Alloc C tmp buffer that is going to be used for the global reduce
  torch::Tensor c_tmp;
+  auto options_fp32 =
+      torch::TensorOptions().dtype(at::kFloat).device(a.device());
  if (use_fp32_reduce && !use_atomic_add) {
    // max num of threadblocks is sms * 4
    long max_c_tmp_size = min(
@@ -748,11 +846,11 @@ torch::Tensor moe_wna16_marlin_gemm(
  torch::Tensor global_scale;
  if (global_scale_or_none.has_value()) {
    global_scale = global_scale_or_none.value();
-    TORCH_CHECK(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn,
+    TORCH_CHECK(b_q_type == vllm::kFE2M1f && group_size == 16,
                "global_scale can only be used for nvfp4 format.");
  } else {
    global_scale = torch::empty({0}, options);
-    TORCH_CHECK(!(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn),
+    TORCH_CHECK(!(b_q_type == vllm::kFE2M1f && group_size == 16),
                "the global_scale parameter must be passed for nvfp4 format.");
  }

@@ -779,15 +877,15 @@ torch::Tensor moe_wna16_marlin_gemm(
  bool has_zp = b_zeros.size(-1) > 0;
  if (has_zp) {
    TORCH_CHECK(
-        b_type == vllm::kU4 || b_type == vllm::kU8,
-        "b_type must be u4 or u8 when has_zp = True. Got = ", b_type.str());
+        b_q_type == vllm::kU4 || b_q_type == vllm::kU8,
+        "b_q_type must be u4 or u8 when has_zp = True. Got = ", b_q_type.str());
  } else {
-    TORCH_CHECK(b_type == vllm::kU4B8 || b_type == vllm::kU8B128 ||
-                    b_type == vllm::kS4 || b_type == vllm::kS8 ||
-                    b_type == vllm::kFE4M3fn || b_type == vllm::kFE2M1f,
-                "b_type must be uint4b8, uint8b128, int4, int8, "
-                "float8_e4m3fn or float4_e2m1f when has_zp = False. Got = ",
-                b_type.str());
+    TORCH_CHECK(b_q_type == vllm::kU4B8 || b_q_type == vllm::kU8B128 ||
+                    b_q_type == vllm::kFE4M3fn || b_q_type == vllm::kFE2M1f,
+                "b_q_type must be uint4b8, uint8b128, float8_e4m3fn or "
+                "float4_e2m1f when "
+                "has_zp = False. Got = ",
+                b_q_type.str());
  }

  if (has_zp && is_zp_float) {
@@ -831,33 +929,71 @@ torch::Tensor moe_wna16_marlin_gemm(
              " is below min_workspace_size = ", min_workspace_size);

  int dev = a.get_device();
+  if (a.scalar_type() == at::ScalarType::Half) {
+    void* scales_ptr;
+    if (b_q_type == vllm::kFE2M1f) {
+      if (group_size == 16)
+        scales_ptr = b_scales.data_ptr<at::Float8_e4m3fn>();
+      else if (group_size == 32)
+        scales_ptr = b_scales.data_ptr<at::Float8_e8m0fnu>();
+      else
+        TORCH_CHECK(false,
+                    "float4_e2m1f only supports group_size == 16 (NVFP4) ",
+                    "and group_size == 32 (MXFP4)");
+    } else {
+      scales_ptr = b_scales.data_ptr<at::Half>();
+    }

-  TORCH_CHECK(a_scales.scalar_type() == at::ScalarType::Float,
-              "scalar type of a_scales must be float");
-  TORCH_CHECK(global_scale.scalar_type() == c.scalar_type(),
-              "scalar type of global_scale must be the same with c");
-  if (a_type.size_bits() == 16) {
-    TORCH_CHECK(
-        a.scalar_type() == c.scalar_type(),
-        "scalar type of a must be the same with c for 16 bit activation");
+    MARLIN_NAMESPACE_NAME::marlin_mm<half>(
+        a.data_ptr<at::Half>(), b_q_weight.data_ptr(), c.data_ptr<at::Half>(),
+        c_tmp.data_ptr<float>(), b_bias.data_ptr<at::Half>(), scales_ptr,
+        global_scale.data_ptr<at::Half>(), b_zeros.data_ptr(), g_idx.data_ptr(),
+        perm.data_ptr(), a_tmp.data_ptr<at::Half>(),
+        sorted_token_ids.data_ptr(), expert_ids.data_ptr(),
+        num_tokens_past_padded.data_ptr(), topk_weights.data_ptr(),
+        moe_block_size, top_k, mul_topk_weights, is_ep, size_m, size_n, size_k,
+        workspace.data_ptr(), b_q_type, has_bias, has_act_order, is_k_full,
+        has_zp, num_groups, group_size, dev,
+        at::cuda::getCurrentCUDAStream(dev), thread_k, thread_n, sms,
+        use_atomic_add, use_fp32_reduce, is_zp_float);
+  } else if (a.scalar_type() == at::ScalarType::BFloat16) {
+    void* scales_ptr;
+    if (b_q_type == vllm::kFE2M1f) {
+      if (group_size == 16)
+        scales_ptr = b_scales.data_ptr<at::Float8_e4m3fn>();
+      else if (group_size == 32)
+        scales_ptr = b_scales.data_ptr<at::Float8_e8m0fnu>();
+      else
+        TORCH_CHECK(false,
+                    "float4_e2m1f only supports group_size == 16 (NVFP4) ",
+                    "and group_size == 32 (MXFP4)");
+    } else {
+      scales_ptr = b_scales.data_ptr<at::BFloat16>();
+    }
+
+    MARLIN_NAMESPACE_NAME::marlin_mm<nv_bfloat16>(
+        a.data_ptr<at::BFloat16>(), b_q_weight.data_ptr(),
+        c.data_ptr<at::BFloat16>(), c_tmp.data_ptr<float>(),
+        b_bias.data_ptr<at::BFloat16>(), scales_ptr,
+        global_scale.data_ptr<at::BFloat16>(), b_zeros.data_ptr(),
+        g_idx.data_ptr(), perm.data_ptr(), a_tmp.data_ptr<at::BFloat16>(),
+        sorted_token_ids.data_ptr(), expert_ids.data_ptr(),
+        num_tokens_past_padded.data_ptr(), topk_weights.data_ptr(),
+        moe_block_size, top_k, mul_topk_weights, is_ep, size_m, size_n, size_k,
+        workspace.data_ptr(), b_q_type, has_bias, has_act_order, is_k_full,
+        has_zp, num_groups, group_size, dev,
+        at::cuda::getCurrentCUDAStream(dev), thread_k, thread_n, sms,
+        use_atomic_add, use_fp32_reduce, is_zp_float);
+  } else {
+    TORCH_CHECK(false,
+                "moe_wna16_marlin_gemm only supports bfloat16 and float16");
  }

-  MARLIN_NAMESPACE_NAME::marlin_mm(
-      a.data_ptr(), b_q_weight.data_ptr(), c.data_ptr(), c_tmp.data_ptr(),
-      b_bias.data_ptr(), a_scales.data_ptr(), b_scales.data_ptr(),
-      global_scale.data_ptr(), b_zeros.data_ptr(), g_idx.data_ptr(),
-      perm.data_ptr(), a_tmp.data_ptr(), sorted_token_ids.data_ptr(),
-      expert_ids.data_ptr(), num_tokens_past_padded.data_ptr(),
-      topk_weights.data_ptr(), moe_block_size, num_experts, top_k,
-      mul_topk_weights, is_ep, size_m, size_n, size_k, workspace.data_ptr(),
-      a_type, b_type, c_type, s_type, has_bias, has_act_order, is_k_full,
-      has_zp, num_groups, group_size, dev, at::cuda::getCurrentCUDAStream(dev),
-      thread_k, thread_n, sms, blocks_per_sm, use_atomic_add, use_fp32_reduce,
-      is_zp_float);
-
  return c;
 }

+#endif
+
 TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
  m.impl("moe_wna16_marlin_gemm", &moe_wna16_marlin_gemm);
 }
@@ -63,18 +63,16 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
  m.def(
      "moe_wna16_marlin_gemm(Tensor! a, Tensor? c_or_none,"
      "Tensor! b_q_weight, Tensor? b_bias_or_none,"
-      "Tensor! b_scales, Tensor? a_scales, Tensor? global_scale, Tensor? "
+      "Tensor! b_scales, Tensor? global_scale, Tensor? "
      "b_zeros_or_none,"
      "Tensor? g_idx_or_none, Tensor? perm_or_none, Tensor! workspace,"
      "Tensor sorted_token_ids,"
      "Tensor! expert_ids, Tensor! num_tokens_past_padded,"
      "Tensor! topk_weights, int moe_block_size, int top_k, "
-      "bool mul_topk_weights, bool is_ep, int b_type_id,"
+      "bool mul_topk_weights, bool is_ep, int b_q_type_id,"
      "int size_m, int size_n, int size_k,"
      "bool is_full_k, bool use_atomic_add,"
-      "bool use_fp32_reduce, bool is_zp_float,"
-      "int thread_k, int thread_n, int blocks_per_sm) -> Tensor");
-
+      "bool use_fp32_reduce, bool is_zp_float) -> 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! "
@@ -52,13 +52,14 @@ void paged_attention_v2(
    const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
    const int64_t blocksparse_head_sliding_step);

+#ifndef USE_ROCM
 void merge_attn_states(torch::Tensor& output,
                       std::optional<torch::Tensor> output_lse,
                       const torch::Tensor& prefix_output,
                       const torch::Tensor& prefix_lse,
                       const torch::Tensor& suffix_output,
                       const torch::Tensor& suffix_lse);
-#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]
@@ -22,7 +22,6 @@
 #include <ATen/cuda/CUDAContext.h>
 #include <c10/cuda/CUDAGuard.h>
 #include <c10/cuda/CUDAStream.h>
-#include "cutlass_extensions/common.hpp"

 #include "cute/tensor.hpp"
 #include "cutlass/tensor_ref.h"
@@ -174,7 +173,7 @@ void run_get_group_gemm_starts(
 }

 template <typename OutType>
-void run_fp4_blockwise_scaled_group_mm_sm100(
+void run_fp4_blockwise_scaled_group_mm(
    torch::Tensor& output, const torch::Tensor& a, const torch::Tensor& b,
    const torch::Tensor& a_blockscale, const torch::Tensor& b_blockscales,
    const torch::Tensor& alphas, const torch::Tensor& problem_sizes,
@@ -344,225 +343,17 @@ void run_fp4_blockwise_scaled_group_mm_sm100(

  auto can_implement_status = gemm_op.can_implement(args);
  TORCH_CHECK(can_implement_status == cutlass::Status::kSuccess,
-              "Failed to implement GEMM: status=", (int)can_implement_status);
+              "Failed to implement GEMM");

  // Run the GEMM
  auto status = gemm_op.initialize(args, workspace.data_ptr());
-  TORCH_CHECK(status == cutlass::Status::kSuccess,
-              "Failed to initialize GEMM: status=", (int)status,
-              " workspace_size=", workspace_size, " num_experts=", num_experts,
-              " M=", M, " N=", N, " K=", K);
+  TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to initialize GEMM");

  status = gemm_op.run(args, workspace.data_ptr(), stream);
  TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to run GEMM");
 }

-void run_fp4_blockwise_scaled_group_mm_sm120(
-    torch::Tensor& output, const torch::Tensor& a, const torch::Tensor& b,
-    const torch::Tensor& a_blockscale, const torch::Tensor& b_blockscales,
-    const torch::Tensor& alphas, const torch::Tensor& problem_sizes,
-    const torch::Tensor& expert_offsets, const torch::Tensor& sf_offsets, int M,
-    int N, int K) {
-  using ProblemShape =
-      cutlass::gemm::GroupProblemShape<Shape<int32_t, int32_t, int32_t>>;
-  using ElementType = cutlass::float_e2m1_t;
-  using ElementSFType = cutlass::float_ue4m3_t;
-  using ElementA = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
-  using ElementB = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
-
-  // NOTE: For SM120 it seems templating the output type is not supported and
-  // we need to hardcode the output type to bfloat16
-  using ElementC = cutlass::bfloat16_t;
-  using ElementD = ElementC;
-  using ElementAccumulator = float;
-  // Layout definitions
-  using LayoutA = cutlass::layout::RowMajor;
-  using LayoutB = cutlass::layout::ColumnMajor;
-  using LayoutC = cutlass::layout::RowMajor;
-  using LayoutD = LayoutC;
-
-  // Alignment constraints
-  static constexpr int AlignmentA = 32;
-  static constexpr int AlignmentB = 32;
-  static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
-  static constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
-
-  // Architecture definitions
-  using ArchTag = cutlass::arch::Sm120;
-  using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp;
-
-  using ClusterShape = Shape<_1, _1, _1>;
-  using MmaTileShape = Shape<_128, _128, _128>;
-
-  using FusionOperation = cutlass::epilogue::fusion::LinearCombination<
-      ElementD, ElementAccumulator, ElementC, ElementAccumulator>;
-
-  using CollectiveEpilogue =
-      typename cutlass::epilogue::collective::CollectiveBuilder<
-          ArchTag, OperatorClass, MmaTileShape, ClusterShape,
-          cutlass::epilogue::collective::EpilogueTileAuto, ElementAccumulator,
-          ElementAccumulator, ElementC, LayoutC*, AlignmentC, ElementD,
-          LayoutD*, AlignmentD,
-          cutlass::epilogue::collective::EpilogueScheduleAuto,
-          FusionOperation>::CollectiveOp;
-
-  using CollectiveMainloop =
-      typename cutlass::gemm::collective::CollectiveBuilder<
-          ArchTag, OperatorClass, ElementA, LayoutA*, AlignmentA, ElementB,
-          LayoutB*, AlignmentB, ElementAccumulator, MmaTileShape, ClusterShape,
-          cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
-              sizeof(typename CollectiveEpilogue::SharedStorage))>,
-          cutlass::gemm::collective::KernelScheduleAuto>::CollectiveOp;
-
-  using GemmKernel =
-      cutlass::gemm::kernel::GemmUniversal<ProblemShape, CollectiveMainloop,
-                                           CollectiveEpilogue>;
-
-  using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
-  using StrideA = typename Gemm::GemmKernel::InternalStrideA;
-  using StrideB = typename Gemm::GemmKernel::InternalStrideB;
-  using StrideC = typename Gemm::GemmKernel::InternalStrideC;
-  using StrideD = typename Gemm::GemmKernel::InternalStrideD;
-
-  using LayoutSFA =
-      typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFA;
-  using LayoutSFB =
-      typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFB;
-  using ScaleConfig =
-      typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig;
-
-  using UnderlyingProblemShape = ProblemShape::UnderlyingProblemShape;
-  int num_experts = static_cast<int>(expert_offsets.size(0));
-  auto options_int =
-      torch::TensorOptions().dtype(torch::kInt64).device(a.device());
-
-  torch::Tensor a_ptrs = torch::empty(num_experts, options_int);
-  torch::Tensor b_ptrs = torch::empty(num_experts, options_int);
-  torch::Tensor out_ptrs = torch::empty(num_experts, options_int);
-  torch::Tensor a_scales_ptrs = torch::empty(num_experts, options_int);
-  torch::Tensor b_scales_ptrs = torch::empty(num_experts, options_int);
-  torch::Tensor alpha_ptrs = torch::empty(num_experts, options_int);
-  torch::Tensor layout_sfa = torch::empty({num_experts, 5}, options_int);
-  torch::Tensor layout_sfb = torch::empty({num_experts, 5}, options_int);
-  torch::Tensor c_strides1 =
-      torch::full({num_experts}, output.stride(0), options_int);
-  torch::Tensor a_strides1 =
-      torch::full({num_experts}, a.stride(0) * 2, options_int);
-  torch::Tensor b_strides1 =
-      torch::full({num_experts}, b.stride(1) * 2, options_int);
-
-  run_get_group_gemm_starts<LayoutSFA, LayoutSFB, ScaleConfig>(
-      a_ptrs, b_ptrs, out_ptrs, a_scales_ptrs, b_scales_ptrs, alpha_ptrs,
-      layout_sfa, layout_sfb, a, b, output, a_blockscale, b_blockscales, alphas,
-      expert_offsets, sf_offsets, problem_sizes, M, N, K);
-
-  // Create an instance of the GEMM
-  Gemm gemm_op;
-
-  // Initialize problem_sizes_as_shapes correctly
-  UnderlyingProblemShape* problem_sizes_as_shapes =
-      static_cast<UnderlyingProblemShape*>(problem_sizes.data_ptr());
-
-  // Set the Scheduler info
-  cutlass::KernelHardwareInfo hw_info;
-  using RasterOrderOptions = cutlass::gemm::kernel::detail::RasterOrderOptions;
-  typename Gemm::GemmKernel::TileSchedulerArguments scheduler;
-  scheduler.raster_order = RasterOrderOptions::AlongM;
-  hw_info.device_id = a.get_device();
-  static std::unordered_map<int, int> cached_sm_counts;
-  if (cached_sm_counts.find(hw_info.device_id) == cached_sm_counts.end()) {
-    cached_sm_counts[hw_info.device_id] =
-        cutlass::KernelHardwareInfo::query_device_multiprocessor_count(
-            hw_info.device_id);
-  }
-  hw_info.sm_count = min(cached_sm_counts[hw_info.device_id], INT_MAX);
-
-  // Mainloop Arguments
-  typename GemmKernel::MainloopArguments mainloop_args{
-      static_cast<const ElementType**>(a_ptrs.data_ptr()),
-      static_cast<StrideA*>(a_strides1.data_ptr()),
-      static_cast<const ElementType**>(b_ptrs.data_ptr()),
-      static_cast<StrideB*>(b_strides1.data_ptr()),
-      static_cast<const ElementSFType**>(a_scales_ptrs.data_ptr()),
-      reinterpret_cast<LayoutSFA*>(layout_sfa.data_ptr()),
-      static_cast<const ElementSFType**>(b_scales_ptrs.data_ptr()),
-      reinterpret_cast<LayoutSFB*>(layout_sfb.data_ptr())};
-
-  // Epilogue Arguments
-  typename GemmKernel::EpilogueArguments epilogue_args{
-      {},  // epilogue.thread
-      nullptr,
-      static_cast<StrideC*>(c_strides1.data_ptr()),
-      static_cast<ElementD**>(out_ptrs.data_ptr()),
-      static_cast<StrideC*>(c_strides1.data_ptr())};
-  auto& fusion_args = epilogue_args.thread;
-  fusion_args.alpha_ptr_array =
-      reinterpret_cast<float**>(alpha_ptrs.data_ptr());
-  fusion_args.dAlpha = {_0{}, _0{}, 1};
-  fusion_args.beta = 0.0f;
-
-  // Gemm Arguments
-  typename GemmKernel::Arguments args{
-      cutlass::gemm::GemmUniversalMode::kGrouped,
-      {num_experts, problem_sizes_as_shapes, nullptr},
-      mainloop_args,
-      epilogue_args,
-      hw_info,
-      scheduler};
-
-  size_t workspace_size = Gemm::get_workspace_size(args);
-  auto const workspace_options =
-      torch::TensorOptions().dtype(torch::kUInt8).device(a.device());
-  auto workspace = torch::empty(workspace_size, workspace_options);
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream(a.get_device());
-
-  auto can_implement_status = gemm_op.can_implement(args);
-  TORCH_CHECK(can_implement_status == cutlass::Status::kSuccess,
-              "Failed to implement GEMM: status=", (int)can_implement_status);
-
-  // Run the GEMM
-  auto status = gemm_op.initialize(args, workspace.data_ptr());
-  TORCH_CHECK(status == cutlass::Status::kSuccess,
-              "Failed to initialize GEMM: status=", (int)status,
-              " workspace_size=", workspace_size, " num_experts=", num_experts,
-              " M=", M, " N=", N, " K=", K);
-
-  status = gemm_op.run(args, workspace.data_ptr(), stream);
-  TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to run GEMM");
-}
-
-template <typename OutType>
-void run_fp4_blockwise_scaled_group_mm(
-    torch::Tensor& output, const torch::Tensor& a, const torch::Tensor& b,
-    const torch::Tensor& a_blockscale, const torch::Tensor& b_blockscales,
-    const torch::Tensor& alphas, const torch::Tensor& problem_sizes,
-    const torch::Tensor& expert_offsets, const torch::Tensor& sf_offsets, int M,
-    int N, int K) {
-  int32_t version_num = get_sm_version_num();
-#if defined ENABLE_NVFP4_SM120 && ENABLE_NVFP4_SM120
-  if (version_num >= 120 && version_num < 130) {
-    run_fp4_blockwise_scaled_group_mm_sm120(
-        output, a, b, a_blockscale, b_blockscales, alphas, problem_sizes,
-        expert_offsets, sf_offsets, M, N, K);
-    return;
-  }
-#endif
 #if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
-  if (version_num >= 100 && version_num < 120) {
-    run_fp4_blockwise_scaled_group_mm_sm100<OutType>(
-        output, a, b, a_blockscale, b_blockscales, alphas, problem_sizes,
-        expert_offsets, sf_offsets, M, N, K);
-    return;
-  }
-#endif
-  TORCH_CHECK_NOT_IMPLEMENTED(
-      false,
-      "No compiled cutlass_fp4_group_mm kernel for CUDA device capability: ",
-      version_num, ". Required capability: 100 or 120");
-}
-
-#if (defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100) || \
-    (defined ENABLE_NVFP4_SM120 && ENABLE_NVFP4_SM120)
 constexpr auto FLOAT4_E2M1X2 = at::ScalarType::Byte;
 constexpr auto SF_DTYPE = at::ScalarType::Float8_e4m3fn;
 #endif
@@ -583,8 +374,7 @@ void cutlass_fp4_group_mm(
    const torch::Tensor& a_blockscale, const torch::Tensor& b_blockscales,
    const torch::Tensor& alphas, const torch::Tensor& problem_sizes,
    const torch::Tensor& expert_offsets, const torch::Tensor& sf_offsets) {
-#if (defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100) || \
-    (defined ENABLE_NVFP4_SM120 && ENABLE_NVFP4_SM120)
+#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
  // Input validation
  CHECK_INPUT(a, FLOAT4_E2M1X2, "a");
  CHECK_INPUT(b, FLOAT4_E2M1X2, "b");
@@ -618,14 +408,6 @@ void cutlass_fp4_group_mm(
        output, a, b, a_blockscale, b_blockscales, alphas, problem_sizes,
        expert_offsets, sf_offsets, M, N, K);
  } else {
-  #if defined ENABLE_NVFP4_SM120 && ENABLE_NVFP4_SM120
-    int32_t version_num = get_sm_version_num();
-    if (version_num >= 120 && version_num < 130) {
-      TORCH_CHECK_NOT_IMPLEMENTED(
-          false, "SM120 NVFP4 MOE only supports bfloat16 output, got: ",
-          output.scalar_type());
-    }
-  #endif
    run_fp4_blockwise_scaled_group_mm<cutlass::half_t>(
        output, a, b, a_blockscale, b_blockscales, alphas, problem_sizes,
        expert_offsets, sf_offsets, M, N, K);
@@ -634,8 +416,8 @@ void cutlass_fp4_group_mm(
  TORCH_CHECK_NOT_IMPLEMENTED(
      false,
      "No compiled cutlass_fp4_group_mm kernel, vLLM must "
-      "be compiled with ENABLE_NVFP4_SM100 or ENABLE_NVFP4_SM120 for SM100/120 "
-      "and CUDA 12.8 or above.");
+      "be compiled with ENABLE_NVFP4_SM100 for SM100+ and CUDA "
+      "12.8 or above.");
 #endif
 }

@@ -307,7 +307,7 @@ constexpr auto FLOAT = at::ScalarType::Float;
 constexpr auto INT = at::ScalarType::Int;
 constexpr auto UINT8 = at::ScalarType::Byte;

-void scaled_fp4_experts_quant_sm1xxa(
+void scaled_fp4_experts_quant_sm100a(
    torch::Tensor& output, torch::Tensor& output_scale,
    torch::Tensor const& input, torch::Tensor const& input_global_scale,
    torch::Tensor const& input_offset_by_experts,
@@ -24,9 +24,8 @@ void scaled_fp4_quant_sm1xxa(torch::Tensor const& output,
                             torch::Tensor const& input_sf);
 #endif

-#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
-    (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
-void scaled_fp4_experts_quant_sm1xxa(
+#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
+void scaled_fp4_experts_quant_sm100a(
    torch::Tensor& output, torch::Tensor& output_scale,
    torch::Tensor const& input, torch::Tensor const& input_global_scale,
    torch::Tensor const& input_offset_by_experts,
@@ -55,9 +54,8 @@ void scaled_fp4_experts_quant(
    torch::Tensor const& input, torch::Tensor const& input_global_scale,
    torch::Tensor const& input_offset_by_experts,
    torch::Tensor const& output_scale_offset_by_experts) {
-#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
-    (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
-  return scaled_fp4_experts_quant_sm1xxa(
+#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
+  return scaled_fp4_experts_quant_sm100a(
      output, output_scale, input, input_global_scale, input_offset_by_experts,
      output_scale_offset_by_experts);
 #endif
@@ -437,10 +437,10 @@ struct ComputeTile_W8A16_PerC_MtilexNtilex32_multistage_SM8x_SplitK {
      for (int n_idx = 0; n_idx < WARP_NITER; ++n_idx) {
  #pragma unroll
        for (int k_idx = 0; k_idx < 2; ++k_idx) {
-          FType low16 = MarlinScalarType2<FType>::float2num(
-              C_frag[m_idx][n_idx][k_idx * 2]);
-          FType high16 = MarlinScalarType2<FType>::float2num(
-              C_frag[m_idx][n_idx][k_idx * 2 + 1]);
+          FType low16 =
+              ScalarType<FType>::float2num(C_frag[m_idx][n_idx][k_idx * 2]);
+          FType high16 =
+              ScalarType<FType>::float2num(C_frag[m_idx][n_idx][k_idx * 2 + 1]);
          uint32_t tmp = (reinterpret_cast<uint32_t&>(low16) & 0xffff) |
                         (reinterpret_cast<uint32_t&>(high16) << 16);
          int sts_offset =
@@ -8,7 +8,7 @@
 #include <cuda_bf16.h>
 #include <iostream>
 #include "../gptq_marlin/marlin_dtypes.cuh"
-using marlin::MarlinScalarType2;
+using marlin::ScalarType;

 namespace allspark {

@@ -72,10 +72,10 @@ __global__ void f16_gemm_splitk_reduce_kernel(const FType* C_split, FType* C,

  int n_mat = N_MATRIX > 0 ? N_MATRIX : (int)n_matrix;
  for (int i = 0; i < n_mat; ++i) {
-    sum += MarlinScalarType2<FType>::num2float(C_split[idx + i * matrix_size]);
+    sum += ScalarType<FType>::num2float(C_split[idx + i * matrix_size]);
  }

-  C[idx] = MarlinScalarType2<FType>::float2num(sum);
+  C[idx] = ScalarType<FType>::float2num(sum);
 }

 template <typename FType>
@@ -1,2 +1 @@
-sm*_kernel_*.cu
-kernel_selector.h
+kernel_*.cu
@@ -4,16 +4,14 @@

 namespace marlin {

-template <int const num_threads, int const num_bits, bool is_a_8bit>
+template <int const num_threads, int const num_bits>
 __global__ void awq_marlin_repack_kernel(
    uint32_t const* __restrict__ b_q_weight_ptr, uint32_t* __restrict__ out_ptr,
    int size_k, int size_n) {
  constexpr int pack_factor = 32 / num_bits;

-  constexpr int target_tile_n_size = tile_n_size / (is_a_8bit ? 2 : 1);
-  constexpr int target_tile_k_size = tile_k_size * (is_a_8bit ? 2 : 1);
-  int k_tiles = size_k / target_tile_k_size;
-  int n_tiles = size_n / target_tile_n_size;
+  int k_tiles = size_k / tile_k_size;
+  int n_tiles = size_n / tile_n_size;
  int block_k_tiles = div_ceil(k_tiles, gridDim.x);

  auto start_k_tile = blockIdx.x * block_k_tiles;
@@ -35,10 +33,10 @@ __global__ void awq_marlin_repack_kernel(

  extern __shared__ int4 sh[];

-  constexpr int tile_n_ints = target_tile_n_size / pack_factor;
+  constexpr int tile_n_ints = tile_n_size / pack_factor;

  constexpr int stage_n_threads = tile_n_ints / 4;
-  constexpr int stage_k_threads = target_tile_k_size;
+  constexpr int stage_k_threads = tile_k_size;
  constexpr int stage_size = stage_k_threads * stage_n_threads;

  auto fetch_to_shared = [&](int pipe, int k_tile_id, int n_tile_id) {
@@ -47,7 +45,7 @@ __global__ void awq_marlin_repack_kernel(
      return;
    }

-    int first_n = n_tile_id * target_tile_n_size;
+    int first_n = n_tile_id * tile_n_size;
    int first_n_packed = first_n / pack_factor;

    int4* sh_ptr = sh + stage_size * pipe;
@@ -56,7 +54,7 @@ __global__ void awq_marlin_repack_kernel(
      auto k_id = threadIdx.x / stage_n_threads;
      auto n_id = threadIdx.x % stage_n_threads;

-      int first_k = k_tile_id * target_tile_k_size;
+      int first_k = k_tile_id * tile_k_size;

      cp_async4(&sh_ptr[k_id * stage_n_threads + n_id],
                reinterpret_cast<int4 const*>(
@@ -80,11 +78,11 @@ __global__ void awq_marlin_repack_kernel(
    }

    int tc_col = th_id / 4;
-    int tc_row = (th_id % 4) * (is_a_8bit ? 4 : 2);
+    int tc_row = (th_id % 4) * 2;

    constexpr int tc_offsets[4] = {0, 1, 8, 9};

-    int cur_n = (warp_id / (is_a_8bit ? 2 : 1)) * 16 + tc_col;
+    int cur_n = warp_id * 16 + tc_col;
    int cur_n_packed = cur_n / pack_factor;
    int cur_n_pos = cur_n % pack_factor;

@@ -107,50 +105,23 @@ __global__ void awq_marlin_repack_kernel(
    uint32_t vals[8];
 #pragma unroll
    for (int i = 0; i < 4; i++) {
-      if constexpr (is_a_8bit) {
-        int cur_elem = tc_row + i;
+      int cur_elem = tc_row + tc_offsets[i];

-        int packed_src_0 =
-            sh_stage_int_ptr[cur_n_packed + (8 / pack_factor) * (warp_id % 2) +
-                             sh_stride * cur_elem];
-        int packed_src_1 =
-            sh_stage_int_ptr[cur_n_packed + (8 / pack_factor) * (warp_id % 2) +
-                             sh_stride * (cur_elem + 16)];
+      int packed_src_0 = sh_stage_int_ptr[cur_n_packed + sh_stride * cur_elem];
+      int packed_src_1 = sh_stage_int_ptr[cur_n_packed + (8 / pack_factor) +
+                                          sh_stride * cur_elem];

-        vals[i] = (packed_src_0 >> (cur_n_pos_unpacked * num_bits)) & mask;
-        vals[4 + i] = (packed_src_1 >> (cur_n_pos_unpacked * num_bits)) & mask;
-      } else {
-        int cur_elem = tc_row + tc_offsets[i];
-
-        int packed_src_0 =
-            sh_stage_int_ptr[cur_n_packed + sh_stride * cur_elem];
-        int packed_src_1 = sh_stage_int_ptr[cur_n_packed + (8 / pack_factor) +
-                                            sh_stride * cur_elem];
-
-        vals[i] = (packed_src_0 >> (cur_n_pos_unpacked * num_bits)) & mask;
-        vals[4 + i] = (packed_src_1 >> (cur_n_pos_unpacked * num_bits)) & mask;
-      }
+      vals[i] = (packed_src_0 >> (cur_n_pos_unpacked * num_bits)) & mask;
+      vals[4 + i] = (packed_src_1 >> (cur_n_pos_unpacked * num_bits)) & mask;
    }

-    constexpr int tile_size =
-        target_tile_k_size * target_tile_n_size / pack_factor;
+    constexpr int tile_size = tile_k_size * tile_n_size / pack_factor;
    int out_offset = (k_tile_id * n_tiles + n_tile_id) * tile_size;

    // Result of:
    // https://github.com/NVIDIA/FasterTransformer/blob/main/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h
-    if constexpr (!is_a_8bit && num_bits == 4) {
-      int pack_idx[8] = {0, 2, 4, 6, 1, 3, 5, 7};
-
-      uint32_t res = 0;
-#pragma unroll
-      for (int i = 0; i < 8; i++) {
-        res |= vals[pack_idx[i]] << (i * 4);
-      }
-
-      out_ptr[out_offset + th_id * 4 + warp_id] = res;
-
-    } else if constexpr (is_a_8bit && num_bits == 4) {
-      int pack_idx[8] = {0, 4, 1, 5, 2, 6, 3, 7};
+    if constexpr (num_bits == 4) {
+      constexpr int pack_idx[8] = {0, 2, 4, 6, 1, 3, 5, 7};

      uint32_t res = 0;
 #pragma unroll
@@ -167,9 +138,8 @@ __global__ void awq_marlin_repack_kernel(
      uint32_t res2 = 0;
 #pragma unroll
      for (int i = 0; i < 4; i++) {
-        const int ii = is_a_8bit ? i : pack_idx[i];
-        res1 |= vals[ii] << (i * 8);
-        res2 |= vals[4 + ii] << (i * 8);
+        res1 |= vals[pack_idx[i]] << (i * 8);
+        res2 |= vals[4 + pack_idx[i]] << (i * 8);
      }

      out_ptr[out_offset + th_id * 8 + (warp_id * 2) + 0] = res1;
@@ -206,21 +176,18 @@ __global__ void awq_marlin_repack_kernel(

 }  // namespace marlin

-#define CALL_IF(NUM_BITS, IS_A_8BIT)                                       \
-  else if (num_bits == NUM_BITS && is_a_8bit == IS_A_8BIT) {               \
-    cudaFuncSetAttribute(                                                  \
-        marlin::awq_marlin_repack_kernel<marlin::repack_threads, NUM_BITS, \
-                                         IS_A_8BIT>,                       \
-        cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem);      \
-    marlin::awq_marlin_repack_kernel<marlin::repack_threads, NUM_BITS,     \
-                                     IS_A_8BIT>                            \
-        <<<blocks, marlin::repack_threads, max_shared_mem, stream>>>(      \
-            b_q_weight_ptr, out_ptr, size_k, size_n);                      \
+#define CALL_IF(NUM_BITS)                                                   \
+  else if (num_bits == NUM_BITS) {                                          \
+    cudaFuncSetAttribute(                                                   \
+        marlin::awq_marlin_repack_kernel<marlin::repack_threads, NUM_BITS>, \
+        cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem);       \
+    marlin::awq_marlin_repack_kernel<marlin::repack_threads, NUM_BITS>      \
+        <<<blocks, marlin::repack_threads, max_shared_mem, stream>>>(       \
+            b_q_weight_ptr, out_ptr, size_k, size_n);                       \
  }

 torch::Tensor awq_marlin_repack(torch::Tensor& b_q_weight, int64_t size_k,
-                                int64_t size_n, int64_t num_bits,
-                                bool is_a_8bit) {
+                                int64_t size_n, int64_t num_bits) {
  // Verify compatibility with marlin tile of 16x64
  TORCH_CHECK(size_k % marlin::tile_k_size == 0, "size_k = ", size_k,
              " is not divisible by tile_k_size = ", marlin::tile_k_size);
@@ -271,13 +238,10 @@ torch::Tensor awq_marlin_repack(torch::Tensor& b_q_weight, int64_t size_k,

  if (false) {
  }
-  CALL_IF(4, false)
-  CALL_IF(8, false)
-  CALL_IF(4, true)
-  CALL_IF(8, true)
+  CALL_IF(4)
+  CALL_IF(8)
  else {
-    TORCH_CHECK(false, "Unsupported repack config: num_bits = ", num_bits,
-                ", is_a_8bit = ", is_a_8bit);
+    TORCH_CHECK(false, "Unsupported repack config: num_bits = ", num_bits);
  }

  return out;
@@ -470,50 +470,6 @@ __device__ inline void dequant<nv_bfloat162, vllm::kFE2M1f.id(), false>(
  frag_b[0] = __hmul2(frag_b[0], bias_reg);
 }

-template <>
-__device__ inline void dequant<__nv_fp8x4_e4m3, vllm::kFE2M1f.id(), true>(
-    int q, __nv_fp8x4_e4m3* frag_b) {
-  // Constants for FP4 (E2M1) and FP16 formats
-  constexpr int FP4_EXPONENT = 2, FP8_EXPONENT = 4;
-  constexpr int RIGHT_SHIFT = FP8_EXPONENT - FP4_EXPONENT;
-  constexpr int MASK = 0x70707070;
-
-  // Extract and shift FP4 values to FP16 format
-  int Out1 = (q & 0x80808080) | ((q & MASK) >> RIGHT_SHIFT);
-  q <<= 4;
-  int Out2 = (q & 0x80808080) | ((q & MASK) >> RIGHT_SHIFT);
-
-  // Note1: reverse indexing is intentional because weights are permuted
-  // Note2: when dequant to 8bit type, we write to `frag_b[2]` instead of
-  //        `frag_b[1]` to fit the layout of tensorcore
-  frag_b[1] = *reinterpret_cast<const __nv_fp8x4_e4m3*>(&Out1);
-  frag_b[0] = *reinterpret_cast<const __nv_fp8x4_e4m3*>(&Out2);
-}
-
-template <>
-__device__ inline void dequant<int32_t, vllm::kU4B8.id(), true>(
-    int q, int32_t* frag_b) {
-  constexpr int repeated_zp = 0x08080808;
-  constexpr int MASK = 0x80808080;
-
-  frag_b[0] = ((q & 0x0F0F0F0F | MASK) - repeated_zp) ^ MASK;
-  q >>= 4;
-  frag_b[1] = ((q & 0x0F0F0F0F | MASK) - repeated_zp) ^ MASK;
-}
-
-template <>
-__device__ inline void dequant<__nv_fp8x4_e4m3, vllm::kU4B8.id(), true>(
-    int q, __nv_fp8x4_e4m3* frag_b) {
-  int s = q & 0x08080808;
-  int Out1 = ((q & 0x07070707) | (s << 4)) + (s >> 3);
-  q >>= 4;
-  s = q & 0x08080808;
-  int Out2 = ((q & 0x07070707) | (s << 4)) + (s >> 3);
-
-  frag_b[0] = *reinterpret_cast<const __nv_fp8x4_e4m3*>(&Out1);
-  frag_b[1] = *reinterpret_cast<const __nv_fp8x4_e4m3*>(&Out2);
-}
-
 template <typename scalar_t2, vllm::ScalarTypeId s_type_id>
 __device__ inline void dequant_fp8_scales(int q, scalar_t2* frag_b);

@@ -559,49 +515,6 @@ __device__ inline void dequant_fp8_scales<nv_bfloat162, vllm::kFE8M0fnu.id()>(
  // Note: reverse indexing is intentional because weights are permuted
  frag_b[1] = *reinterpret_cast<const nv_bfloat162*>(&Out1);
  frag_b[0] = *reinterpret_cast<const nv_bfloat162*>(&Out2);
-};
-
-// subtract zero point in quanted format and then dequant
-template <typename scalar_t2, vllm::ScalarTypeId w_type_id,
-          bool skip_flop = false>
-__device__ inline void sub_zp_and_dequant(int q, scalar_t2* frag_b, int zp);
-
-template <>
-__device__ inline void sub_zp_and_dequant<int32_t, vllm::kU4.id(), true>(
-    int q, int32_t* frag_b, int zp) {
-  // INT4 with zp -> INT8
-  // see https://github.com/vllm-project/vllm/pull/24722
-  int repeated_zp = 0x01010101 * zp;
-  int MASK = 0x80808080;
-
-  frag_b[0] = ((q & 0x0F0F0F0F | MASK) - repeated_zp) ^ MASK;
-  q >>= 4;
-  frag_b[1] = ((q & 0x0F0F0F0F | MASK) - repeated_zp) ^ MASK;
-}
-
-template <>
-__device__ inline void sub_zp_and_dequant<__nv_fp8x4_e4m3, vllm::kU4.id(),
-                                          true>(int q, __nv_fp8x4_e4m3* frag_b,
-                                                int zp) {
-  // INT4 with zp -> FP8
-  // see https://github.com/vllm-project/vllm/pull/24722
-  uint32_t u_q = *reinterpret_cast<uint32_t*>(&q);
-  uint32_t u_zp = *reinterpret_cast<uint32_t*>(&zp);
-  uint32_t u_zp1 = u_zp + 1;
-  uint32_t repeated_zp = 0x01010101 * u_zp;
-
-  uint32_t q0, s;
-  q0 = (u_q & 0x0F0F0F0F) | 0x70707070;
-  s = (q0 + repeated_zp) & 0x80808080;
-  uint32_t Out1 = (q0 + (s >> 7) * u_zp1) & 0x0F0F0F0F | s;
-
-  u_q >>= 4;
-  q0 = (u_q & 0x0F0F0F0F) | 0x70707070;
-  s = (q0 + repeated_zp) & 0x80808080;
-  uint32_t Out2 = (q0 + (s >> 7) * u_zp1) & 0x0F0F0F0F | s;
-
-  frag_b[0] = *reinterpret_cast<const __nv_fp8x4_e4m3*>(&Out1);
-  frag_b[1] = *reinterpret_cast<const __nv_fp8x4_e4m3*>(&Out2);
 }

 #endif
@@ -4,292 +4,141 @@ import glob
 import itertools
 import os
 import subprocess
-import sys

 import jinja2

-ARCHS = []
-SUPPORT_FP8 = False
-for arch in sys.argv[1].split(","):
-    arch = arch[: arch.index(".") + 2].replace(".", "")
-    arch = int(arch)
-    # 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 in [89, 120]:
-        SUPPORT_FP8 = True
-
-FILE_HEAD_COMMENT = """
-// auto generated by generate_kernels.py
+FILE_HEAD = """
+// auto generated by generate.py
 // clang-format off
-""".lstrip()

-FILE_HEAD = (
-    FILE_HEAD_COMMENT
-    + """
 #include "kernel.h"
 #include "marlin_template.h"

 namespace MARLIN_NAMESPACE_NAME {
-"""
-)
+""".strip()

 TEMPLATE = (
    "template __global__ void Marlin<"
-    "{{a_type_id}}, "
-    "{{b_type_id}}, "
-    "{{c_type_id}}, "
+    "{{scalar_t}}, "
+    "{{w_type_id}}, "
    "{{s_type_id}}, "
    "{{threads}}, "
    "{{thread_m_blocks}}, "
    "{{thread_n_blocks}}, "
    "{{thread_k_blocks}}, "
-    "{{m_block_size_8}}, "
+    "{{'true' if m_block_size_8 else 'false'}}, "
    "{{stages}}, "
    "{{group_blocks}}, "
-    "{{is_zp_float}}>"
+    "{{'true' if is_zp_float else 'false'}}>"
    "( MARLIN_KERNEL_PARAMS );"
 )

+# int8 with zero point case (vllm::kU8) is also supported,
+# we don't add it to reduce wheel size.
+SCALAR_TYPES = [
+    "vllm::kU4",
+    "vllm::kU4B8",
+    "vllm::kU8B128",
+    "vllm::kFE4M3fn",
+    "vllm::kFE2M1f",
+]
 THREAD_CONFIGS = [(128, 128, 256), (64, 256, 256), (64, 128, 128), (128, 64, 128)]

 THREAD_M_BLOCKS = [0.5, 1, 2, 3, 4]
-
-QUANT_CONFIGS = [
-    # AWQ-INT4
-    {
-        "b_type": "kU4",
-        "thread_configs": THREAD_CONFIGS,
-        "thread_m_blocks": THREAD_M_BLOCKS,
-        "group_blocks": [-1, 2, 4, 8],
-    },
-    # HQQ
-    {
-        "a_type": ["kFloat16"],
-        "b_type": "kU4",
-        "thread_configs": THREAD_CONFIGS,
-        "thread_m_blocks": THREAD_M_BLOCKS,
-        "group_blocks": [4],
-        "is_zp_float": True,
-    },
-    # GPTQ-INT4
-    {
-        "b_type": "kU4B8",
-        "thread_configs": THREAD_CONFIGS,
-        "thread_m_blocks": THREAD_M_BLOCKS,
-        "group_blocks": [-1, 0, 2, 4, 8],
-    },
-    # GPTQ-INT8
-    {
-        "b_type": "kU8B128",
-        "thread_configs": THREAD_CONFIGS,
-        "thread_m_blocks": THREAD_M_BLOCKS,
-        "group_blocks": [-1, 0, 2, 4, 8],
-    },
-    # FP8
-    {
-        "b_type": "kFE4M3fn",
-        "thread_configs": THREAD_CONFIGS,
-        "thread_m_blocks": THREAD_M_BLOCKS,
-        "group_blocks": [-1, 8],
-    },
-    # NVFP4
-    {
-        "b_type": "kFE2M1f",
-        "s_type": "kFE4M3fn",
-        "thread_configs": THREAD_CONFIGS,
-        "thread_m_blocks": THREAD_M_BLOCKS,
-        "group_blocks": [1],
-    },
-    # MXFP4
-    {
-        "a_type": ["kBFloat16"],
-        "b_type": "kFE2M1f",
-        "s_type": "kFE8M0fnu",
-        "thread_configs": THREAD_CONFIGS,
-        "thread_m_blocks": THREAD_M_BLOCKS,
-        "group_blocks": [2],
-    },
-    # AWQ-INT4 with INT8 activation
-    {
-        "a_type": ["kS8"],
-        "b_type": "kU4",
-        "thread_configs": THREAD_CONFIGS,
-        "thread_m_blocks": [1, 2, 3, 4],
-        "group_blocks": [-1, 2, 4, 8],
-    },
-    # GPTQ-INT4 with INT8 activation
-    {
-        "a_type": ["kS8"],
-        "b_type": "kU4B8",
-        "thread_configs": THREAD_CONFIGS,
-        "thread_m_blocks": [1, 2, 3, 4],
-        "group_blocks": [-1, 2, 4, 8],
-    },
-    # GPTQ-INT4 with FP8 activation
-    {
-        "a_type": ["kFE4M3fn"],
-        "b_type": "kU4B8",
-        "thread_configs": THREAD_CONFIGS,
-        "thread_m_blocks": [1, 2, 3, 4],
-        "group_blocks": [-1, 2, 4, 8],
-    },
-    # AWQ-INT4 with FP8 activation
-    {
-        "a_type": ["kFE4M3fn"],
-        "b_type": "kU4",
-        "thread_configs": THREAD_CONFIGS,
-        "thread_m_blocks": [1, 2, 3, 4],
-        "group_blocks": [-1, 2, 4, 8],
-    },
-    # MXFP4 with FP8 activation
-    {
-        "a_type": ["kFE4M3fn"],
-        "b_type": "kFE2M1f",
-        "c_type": ["kBFloat16"],
-        "s_type": "kFE8M0fnu",
-        "thread_configs": THREAD_CONFIGS,
-        "thread_m_blocks": [1, 2, 3, 4],
-        "group_blocks": [2],
-    },
-]
+# group_blocks:
+#   = 0 : act order case
+#   = -1 : channelwise quantization
+#   > 0 : group_size=16*group_blocks
+GROUP_BLOCKS = [0, 1, -1, 2, 4, 8]
+DTYPES = ["fp16", "bf16"]


 def remove_old_kernels():
-    for filename in glob.glob(os.path.dirname(__file__) + "/*kernel_*.cu"):
+    for filename in glob.glob(os.path.dirname(__file__) + "/kernel_*.cu"):
        subprocess.call(["rm", "-f", filename])

-    filename = os.path.dirname(__file__) + "/kernel_selector.h"
-    subprocess.call(["rm", "-f", filename])
-

 def generate_new_kernels():
-    result_dict = {}
-
-    for quant_config in QUANT_CONFIGS:
-        c_types = quant_config.get("c_type", ["kFloat16", "kBFloat16"])
-        a_types = quant_config.get("a_type", ["kFloat16", "kBFloat16"])
-        b_type = quant_config["b_type"]
-        is_zp_float = quant_config.get("is_zp_float", False)
-        all_group_blocks = quant_config["group_blocks"]
-        all_m_blocks = quant_config["thread_m_blocks"]
-        all_thread_configs = quant_config["thread_configs"]
-
-        for a_type, c_type in itertools.product(a_types, c_types):
-            if not SUPPORT_FP8 and a_type == "kFE4M3fn":
-                continue
-            if "16" in a_type and "16" in c_type and a_type != c_type:
-                continue
-            s_type = quant_config.get("s_type", c_type)
-            if (a_type, b_type, c_type) not in result_dict:
-                result_dict[(a_type, b_type, c_type)] = []
-
-            for group_blocks, m_blocks, thread_configs in itertools.product(
-                all_group_blocks, all_m_blocks, all_thread_configs
-            ):
-                thread_k, thread_n, threads = thread_configs
-
-                if threads == 256:
-                    # for small batch (m_blocks == 1),
-                    #     we only need (128, 128, 256)
-                    # for large batch (m_blocks > 1),
-                    #     we only need (64, 256, 256)
-                    if m_blocks <= 1 and (thread_k, thread_n) != (128, 128):
-                        continue
-                    if m_blocks > 1 and (thread_k, thread_n) != (64, 256):
-                        continue
-
-                config = {
-                    "threads": threads,
-                    "s_type": s_type,
-                    "thread_m_blocks": max(m_blocks, 1),
-                    "thread_k_blocks": thread_k // 16,
-                    "thread_n_blocks": thread_n // 16,
-                    "m_block_size_8": "true" if m_blocks == 0.5 else "false",
-                    "stages": "pipe_stages",
-                    "group_blocks": group_blocks,
-                    "is_zp_float": "true" if is_zp_float else "false",
-                }
-
-                result_dict[(a_type, b_type, c_type)].append(config)
-
-    kernel_selector_str = FILE_HEAD_COMMENT
-
-    for (a_type, b_type, c_type), config_list in result_dict.items():
+    for scalar_type, dtype in itertools.product(SCALAR_TYPES, DTYPES):
        all_template_str_list = []
-        for config in config_list:
-            s_type = config["s_type"]
-            template_str = jinja2.Template(TEMPLATE).render(
-                a_type_id=f"vllm::{a_type}.id()",
-                b_type_id=f"vllm::{b_type}.id()",
-                c_type_id=f"vllm::{c_type}.id()",
-                s_type_id=f"vllm::{s_type}.id()",
-                **config,
-            )
-            all_template_str_list.append(template_str)

-            conditions = [
-                f"a_type == vllm::{a_type}",
-                f"b_type == vllm::{b_type}",
-                f"c_type == vllm::{c_type}",
-                f"s_type == vllm::{s_type}",
-                f"threads == {config['threads']}",
-                f"thread_m_blocks == {config['thread_m_blocks']}",
-                f"thread_n_blocks == {config['thread_n_blocks']}",
-                f"thread_k_blocks == {config['thread_k_blocks']}",
-                f"m_block_size_8 == {config['m_block_size_8']}",
-                f"group_blocks == {config['group_blocks']}",
-                f"is_zp_float == {config['is_zp_float']}",
-            ]
-            conditions = " && ".join(conditions)
+        for group_blocks, m_blocks, thread_configs in itertools.product(
+            GROUP_BLOCKS, THREAD_M_BLOCKS, THREAD_CONFIGS
+        ):
+            # act order case only support gptq-int4 and gptq-int8
+            if group_blocks == 0 and scalar_type not in [
+                "vllm::kU4B8",
+                "vllm::kU8B128",
+            ]:
+                continue
+            if thread_configs[2] == 256:
+                # for small batch (m_blocks == 1), we only need (128, 128, 256)
+                # for large batch (m_blocks > 1), we only need (64, 256, 256)
+                if m_blocks <= 1 and thread_configs[0] != 128:
+                    continue
+                if m_blocks > 1 and thread_configs[0] != 64:
+                    continue

-            if kernel_selector_str == FILE_HEAD_COMMENT:
-                kernel_selector_str += f"if ({conditions})\n  kernel = "
-            else:
-                kernel_selector_str += f"else if ({conditions})\n  kernel = "
+            # we only support channelwise quantization and group_size == 128
+            # for fp8
+            if scalar_type == "vllm::kFE4M3fn" and group_blocks not in [-1, 8]:
+                continue
+            # nvfp4 only supports group_size == 16
+            # mxfp4 only supports group_size == 32
+            if scalar_type == "vllm::kFE2M1f" and group_blocks not in [1, 2]:
+                continue
+            # other quantization methods don't support group_size = 16
+            if scalar_type != "vllm::kFE2M1f" and group_blocks == 1:
+                continue

-            kernel_template2 = (
-                "Marlin<{{a_type_id}}, {{b_type_id}}, {{c_type_id}}, "
-                "{{s_type_id}}, {{threads}}, {{thread_m_blocks}}, "
-                "{{thread_n_blocks}}, {{thread_k_blocks}}, "
-                "{{m_block_size_8}}, {{stages}}, {{group_blocks}}, "
-                "{{is_zp_float}}>;"
-            )
+            k_blocks = thread_configs[0] // 16
+            n_blocks = thread_configs[1] // 16
+            threads = thread_configs[2]

-            kernel_selector_str += (
-                jinja2.Template(kernel_template2).render(
-                    a_type_id=f"vllm::{a_type}.id()",
-                    b_type_id=f"vllm::{b_type}.id()",
-                    c_type_id=f"vllm::{c_type}.id()",
-                    s_type_id=f"vllm::{s_type}.id()",
-                    **config,
+            c_dtype = "half" if dtype == "fp16" else "nv_bfloat16"
+
+            is_zp_float_list = [False]
+            if dtype == "fp16" and scalar_type == "vllm::kU4" and group_blocks == 4:
+                # HQQ (is_zp_float = true) only supports
+                # 4bit quantization and fp16
+                is_zp_float_list.append(True)
+
+            if scalar_type == "vllm::kFE2M1f" and group_blocks == 1:
+                s_type = "vllm::kFE4M3fn"
+            elif scalar_type == "vllm::kFE2M1f" and group_blocks == 2:
+                s_type = "vllm::kFE8M0fnu"
+                if dtype == "fp16":
+                    # we cannot safely dequantize e8m0 to fp16, so skip this
+                    continue
+            elif dtype == "fp16":
+                s_type = "vllm::kFloat16"
+            elif dtype == "bf16":
+                s_type = "vllm::kBFloat16"
+
+            for is_zp_float in is_zp_float_list:
+                template_str = jinja2.Template(TEMPLATE).render(
+                    scalar_t=c_dtype,
+                    w_type_id=scalar_type + ".id()",
+                    s_type_id=s_type + ".id()",
+                    threads=threads,
+                    thread_m_blocks=max(m_blocks, 1),
+                    thread_n_blocks=n_blocks,
+                    thread_k_blocks=k_blocks,
+                    m_block_size_8=m_blocks == 0.5,
+                    stages="pipe_stages",
+                    group_blocks=group_blocks,
+                    is_zp_float=is_zp_float,
                )
-                + "\n"
-            )
+
+                all_template_str_list.append(template_str)

        file_content = FILE_HEAD + "\n\n"
        file_content += "\n\n".join(all_template_str_list) + "\n\n}\n"
-        if a_type == "kFE4M3fn":
-            filename = f"sm89_kernel_{a_type[1:]}_{b_type[1:]}_{c_type[1:]}.cu"
-        else:
-            filename = f"sm80_kernel_{a_type[1:]}_{b_type[1:]}_{c_type[1:]}.cu"
-
-        filename = filename.lower()
+        filename = f"kernel_{dtype}_{scalar_type[6:].lower()}.cu"

        with open(os.path.join(os.path.dirname(__file__), filename), "w") as f:
            f.write(file_content)

-    if not SUPPORT_FP8 and kernel_selector_str != FILE_HEAD_COMMENT:
-        kernel_selector_str += (
-            "else if (a_type == vllm::kFE4M3fn)\n"
-            "  TORCH_CHECK(false, "
-            '"marlin kernel with fp8 activation is not built.");'
-        )
-
-    with open(os.path.join(os.path.dirname(__file__), "kernel_selector.h"), "w") as f:
-        f.write(kernel_selector_str)
-

 if __name__ == "__main__":
    remove_old_kernels()
@@ -53,7 +53,7 @@ torch::Tensor gptq_marlin_gemm(
    std::optional<torch::Tensor> const& b_zeros_or_none,
    std::optional<torch::Tensor> const& g_idx_or_none,
    std::optional<torch::Tensor> const& perm_or_none, torch::Tensor& workspace,
-    vllm::ScalarTypeId const& b_type_id, int64_t size_m, int64_t size_n,
+    vllm::ScalarTypeId const& b_q_type_id, int64_t size_m, int64_t size_n,
    int64_t size_k, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce,
    bool is_zp_float) {
  TORCH_CHECK_NOT_IMPLEMENTED(false,
@@ -243,29 +243,204 @@ bool is_valid_config(thread_config_t const& th_config, int thread_m_blocks,
  int cache_size = get_kernel_cache_size(
      th_config, thread_m_blocks, prob_m, prob_n, prob_k, num_bits, group_size,
      has_act_order, is_k_full, has_zp, is_zp_float);
-  return cache_size <= max_shared_mem;
+  return cache_size + 512 <= max_shared_mem;
 }

-MarlinFuncPtr get_marlin_kernel(
-    const vllm::ScalarType a_type, const vllm::ScalarType b_type,
-    const vllm::ScalarType c_type, const vllm::ScalarType s_type,
-    int thread_m_blocks, int thread_n_blocks, int thread_k_blocks,
-    bool m_block_size_8, bool has_act_order, bool has_zp, int group_blocks,
-    int threads, bool is_zp_float) {
-  int num_bits = b_type.size_bits();
-  auto kernel = MarlinDefault;
+  #define _GET_IF(W_TYPE, THREAD_M_BLOCKS, THREAD_N_BLOCKS, THREAD_K_BLOCKS,   \
+                  M_BLOCK_SIZE_8, GROUP_BLOCKS, NUM_THREADS, IS_ZP_FLOAT)      \
+    else if (q_type == W_TYPE && thread_m_blocks == THREAD_M_BLOCKS &&         \
+             thread_n_blocks == THREAD_N_BLOCKS &&                             \
+             thread_k_blocks == THREAD_K_BLOCKS &&                             \
+             m_block_size_8 == M_BLOCK_SIZE_8 &&                               \
+             group_blocks == GROUP_BLOCKS && num_threads == NUM_THREADS &&     \
+             is_zp_float == IS_ZP_FLOAT) {                                     \
+      constexpr auto S_TYPE =                                                  \
+          W_TYPE == vllm::kFE2M1f                                              \
+              ? (GROUP_BLOCKS == 1 ? vllm::kFE4M3fn : vllm::kFE8M0fnu)         \
+              : (std::is_same<scalar_t, half>::value ? vllm::kFloat16          \
+                                                     : vllm::kBFloat16);       \
+      kernel = Marlin<scalar_t, W_TYPE.id(), S_TYPE.id(), NUM_THREADS,         \
+                      THREAD_M_BLOCKS, THREAD_N_BLOCKS, THREAD_K_BLOCKS,       \
+                      M_BLOCK_SIZE_8, pipe_stages, GROUP_BLOCKS, IS_ZP_FLOAT>; \
+    }

-  #include "kernel_selector.h"
+  // COMMON: cases for (group_blocks in [-1, 2, 4, 8] and is_zp_float == false)
+  //         this is the most common cases
+  // BIGGROUP: cases for big group size (group_blocks in [-1, 8])
+  // FZP: cases for float-zero-point (is_zp_float = true)
+  // ACT: cases for act order case (group_blocks == 0)
+  // FP4: cases for nvfp4(e2m1) (group_blocks == 1)
+  #define COMMON_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)       \
+    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, -1, NUM_THREADS, false)  \
+    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 2, NUM_THREADS, false)   \
+    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 4, NUM_THREADS, false)   \
+    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 8, NUM_THREADS, false)   \
+    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
+    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false)  \
+    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false)  \
+    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)
+
+  #define COMMON_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)     \
+    _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
+    _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false)  \
+    _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false)  \
+    _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)  \
+                                                                          \
+    _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
+    _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false)  \
+    _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false)  \
+    _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)  \
+                                                                          \
+    _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
+    _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false)  \
+    _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false)  \
+    _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)
+
+  #define COMMON_GET_IF(W_TYPE)            \
+    COMMON_GET_IF_M1(W_TYPE, 8, 8, 256)    \
+    COMMON_GET_IF_M1(W_TYPE, 8, 4, 128)    \
+    COMMON_GET_IF_M1(W_TYPE, 4, 8, 128)    \
+    COMMON_GET_IF_M234(W_TYPE, 16, 4, 256) \
+    COMMON_GET_IF_M234(W_TYPE, 8, 4, 128)  \
+    COMMON_GET_IF_M234(W_TYPE, 4, 8, 128)
+
+  #define BIGGROUP_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)     \
+    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, -1, NUM_THREADS, false)  \
+    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 8, NUM_THREADS, false)   \
+    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
+    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)
+
+  #define BIGGROUP_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)   \
+    _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
+    _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)  \
+    _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
+    _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)  \
+    _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
+    _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)
+
+  #define BIGGROUP_GET_IF(W_TYPE)            \
+    BIGGROUP_GET_IF_M1(W_TYPE, 8, 8, 256)    \
+    BIGGROUP_GET_IF_M1(W_TYPE, 8, 4, 128)    \
+    BIGGROUP_GET_IF_M1(W_TYPE, 4, 8, 128)    \
+    BIGGROUP_GET_IF_M234(W_TYPE, 16, 4, 256) \
+    BIGGROUP_GET_IF_M234(W_TYPE, 8, 4, 128)  \
+    BIGGROUP_GET_IF_M234(W_TYPE, 4, 8, 128)
+
+  #define NVFP4_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)      \
+    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 1, NUM_THREADS, false) \
+    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false)
+
+  #define NVFP4_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)     \
+    _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false) \
+    _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false) \
+    _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false)
+
+  #define NVFP4_GET_IF(W_TYPE)            \
+    NVFP4_GET_IF_M1(W_TYPE, 8, 8, 256)    \
+    NVFP4_GET_IF_M1(W_TYPE, 8, 4, 128)    \
+    NVFP4_GET_IF_M1(W_TYPE, 4, 8, 128)    \
+    NVFP4_GET_IF_M234(W_TYPE, 16, 4, 256) \
+    NVFP4_GET_IF_M234(W_TYPE, 8, 4, 128)  \
+    NVFP4_GET_IF_M234(W_TYPE, 4, 8, 128)
+
+  #define MXFP4_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)      \
+    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 2, NUM_THREADS, false) \
+    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false)
+
+  #define MXFP4_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)     \
+    _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \
+    _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \
+    _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false)
+
+  #define MXFP4_GET_IF(W_TYPE)            \
+    MXFP4_GET_IF_M1(W_TYPE, 8, 8, 256)    \
+    MXFP4_GET_IF_M1(W_TYPE, 8, 4, 128)    \
+    MXFP4_GET_IF_M1(W_TYPE, 4, 8, 128)    \
+    MXFP4_GET_IF_M234(W_TYPE, 16, 4, 256) \
+    MXFP4_GET_IF_M234(W_TYPE, 8, 4, 128)  \
+    MXFP4_GET_IF_M234(W_TYPE, 4, 8, 128)
+
+  // We currently have 4-bit models only with group_blocks == 4
+  #define FZP_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)       \
+    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 4, NUM_THREADS, true) \
+    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true)
+
+  #define FZP_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)      \
+    _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true) \
+    _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true) \
+    _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true)
+
+  #define FZP_GET_IF(W_TYPE)            \
+    FZP_GET_IF_M1(W_TYPE, 8, 8, 256)    \
+    FZP_GET_IF_M1(W_TYPE, 8, 4, 128)    \
+    FZP_GET_IF_M1(W_TYPE, 4, 8, 128)    \
+    FZP_GET_IF_M234(W_TYPE, 16, 4, 256) \
+    FZP_GET_IF_M234(W_TYPE, 8, 4, 128)  \
+    FZP_GET_IF_M234(W_TYPE, 4, 8, 128)
+
+  // We currently have 4-bit models only with group_blocks == 4
+  #define ACT_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)        \
+    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 0, NUM_THREADS, false) \
+    _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false)
+
+  #define ACT_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS)       \
+    _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false) \
+    _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false) \
+    _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false)
+
+  #define ACT_GET_IF(W_TYPE)            \
+    ACT_GET_IF_M1(W_TYPE, 8, 8, 256)    \
+    ACT_GET_IF_M1(W_TYPE, 8, 4, 128)    \
+    ACT_GET_IF_M1(W_TYPE, 4, 8, 128)    \
+    ACT_GET_IF_M234(W_TYPE, 16, 4, 256) \
+    ACT_GET_IF_M234(W_TYPE, 8, 4, 128)  \
+    ACT_GET_IF_M234(W_TYPE, 4, 8, 128)
+
+template <typename scalar_t>
+MarlinFuncPtr get_marlin_kernel(const vllm::ScalarType q_type,
+                                int thread_m_blocks, int thread_n_blocks,
+                                int thread_k_blocks, bool m_block_size_8,
+                                bool has_act_order, bool has_zp,
+                                int group_blocks, int num_threads,
+                                bool is_zp_float) {
+  int num_bits = q_type.size_bits();
+  auto kernel = MarlinDefault;
+  if (false) {
+  }
+
+  COMMON_GET_IF(vllm::kU4)
+  COMMON_GET_IF(vllm::kU4B8)
+  COMMON_GET_IF(vllm::kU8B128)
+
+  NVFP4_GET_IF(vllm::kFE2M1f)
+
+  BIGGROUP_GET_IF(vllm::kFE4M3fn)
+
+  ACT_GET_IF(vllm::kU4B8)
+  ACT_GET_IF(vllm::kU8B128)
+
+  if (std::is_same<scalar_t, half>::value) {
+    if (false) {
+    }
+    FZP_GET_IF(vllm::kU4)
+  }
+  if (std::is_same<scalar_t, nv_bfloat16>::value) {
+    if (false) {
+    }
+    MXFP4_GET_IF(vllm::kFE2M1f)
+  }

  return kernel;
 }

-exec_config_t determine_exec_config(
-    const vllm::ScalarType& a_type, const vllm::ScalarType& b_type,
-    const vllm::ScalarType& c_type, const vllm::ScalarType& s_type, int prob_m,
-    int prob_n, int prob_k, int thread_m_blocks, bool m_block_size_8,
-    int num_bits, int group_size, bool has_act_order, bool is_k_full,
-    bool has_zp, bool is_zp_float, int max_shared_mem, int sms) {
+template <typename scalar_t>
+exec_config_t determine_exec_config(const vllm::ScalarType& q_type, int prob_m,
+                                    int prob_n, int prob_k, int thread_m_blocks,
+                                    bool m_block_size_8, int num_bits,
+                                    int group_size, bool has_act_order,
+                                    bool is_k_full, bool has_zp,
+                                    bool is_zp_float, int max_shared_mem,
+                                    int sms) {
  exec_config_t exec_cfg = exec_config_t{1, thread_config_t{-1, -1, -1}};
  thread_config_t* thread_configs = thread_m_blocks > 1
                                        ? large_batch_thread_configs
@@ -280,7 +455,7 @@ exec_config_t determine_exec_config(

    if (!is_valid_config(th_config, thread_m_blocks, prob_m, prob_n, prob_k,
                         num_bits, group_size, has_act_order, is_k_full, has_zp,
-                         is_zp_float, max_shared_mem - 512)) {
+                         is_zp_float, max_shared_mem)) {
      continue;
    }

@@ -293,11 +468,10 @@ exec_config_t determine_exec_config(
      group_blocks = group_size == -1 ? -1 : group_size / 16;
    }

-    auto kernel =
-        get_marlin_kernel(a_type, b_type, c_type, s_type, thread_m_blocks,
-                          th_config.thread_n / 16, th_config.thread_k / 16,
-                          m_block_size_8, has_act_order, has_zp, group_blocks,
-                          th_config.num_threads, is_zp_float);
+    auto kernel = get_marlin_kernel<scalar_t>(
+        q_type, thread_m_blocks, th_config.thread_n / 16,
+        th_config.thread_k / 16, m_block_size_8, has_act_order, has_zp,
+        group_blocks, th_config.num_threads, is_zp_float);

    if (kernel == MarlinDefault) continue;

@@ -311,16 +485,28 @@ exec_config_t determine_exec_config(
  return exec_cfg;
 }

+template <typename scalar_t>
 void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
-               void* a_s, void* b_s, void* g_s, void* zp, void* g_idx,
-               void* perm, void* a_tmp, int prob_m, int prob_n, int prob_k,
-               int lda, void* workspace, vllm::ScalarType const& a_type,
-               vllm::ScalarType const& b_type, vllm::ScalarType const& c_type,
-               vllm::ScalarType const& s_type, bool has_bias,
+               void* s, void* s2, void* zp, void* g_idx, void* perm,
+               void* a_tmp, int prob_m, int prob_n, int prob_k, int lda,
+               void* workspace, vllm::ScalarType const& q_type, bool has_bias,
               bool has_act_order, bool is_k_full, bool has_zp, int num_groups,
               int group_size, int dev, cudaStream_t stream, int thread_k_init,
               int thread_n_init, int sms, bool use_atomic_add,
               bool use_fp32_reduce, bool is_zp_float) {
+  if (has_zp) {
+    TORCH_CHECK(
+        q_type == vllm::kU4 || q_type == vllm::kU8,
+        "q_type must be u4 or u8 when has_zp = True. Got = ", q_type.str());
+  } else {
+    TORCH_CHECK(
+        q_type == vllm::kU4B8 || q_type == vllm::kU8B128 ||
+            q_type == vllm::kFE4M3fn || q_type == vllm::kFE2M1f,
+        "q_type must be uint4b8, uint8b128, float8_e4m3fn or float4_e2m1f when "
+        "has_zp = False. Got = ",
+        q_type.str());
+  }
+
  TORCH_CHECK(prob_m > 0 && prob_n > 0 && prob_k > 0, "Invalid MNK = [", prob_m,
              ", ", prob_n, ", ", prob_k, "]");

@@ -345,21 +531,19 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
    }
  }

-  int num_bits = b_type.size_bits();
+  int num_bits = q_type.size_bits();
  const int4* A_ptr = (const int4*)A;
  const int4* B_ptr = (const int4*)B;
  int4* C_ptr = (int4*)C;
  int4* C_tmp_ptr = (int4*)C_tmp;
-
  const int4* bias_ptr = (const int4*)b_bias;
-  const float* a_s_ptr = (const float*)a_s;
-  const int4* b_s_ptr = (const int4*)b_s;
-  const uint16_t* g_s_ptr = (const uint16_t*)g_s;
-
+  const int4* s_ptr = (const int4*)s;
+  const uint16_t* s2_ptr = (const uint16_t*)s2;
  const int4* zp_ptr = (const int4*)zp;
  const int* g_idx_ptr = (const int*)g_idx;
  const int* perm_ptr = (const int*)perm;
  int4* a_tmp_ptr = (int4*)a_tmp;
+
  int* locks = (int*)workspace;

  if (has_act_order) {
@@ -384,21 +568,6 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
                         cudaDevAttrMaxSharedMemoryPerBlockOptin, dev);
  TORCH_CHECK(max_shared_mem > 0);

-  int major_capability, minor_capability;
-  cudaDeviceGetAttribute(&major_capability, cudaDevAttrComputeCapabilityMajor,
-                         dev);
-  cudaDeviceGetAttribute(&minor_capability, cudaDevAttrComputeCapabilityMinor,
-                         dev);
-  TORCH_CHECK(major_capability * 10 + minor_capability >= 80,
-              "marlin kernel only support Ampere or newer GPUs.");
-  if (a_type == vllm::kFE4M3fn) {
-    TORCH_CHECK(
-        major_capability * 10 + minor_capability == 89 ||
-            major_capability * 10 + minor_capability == 120,
-        "Marlin W4A8-FP8 only support SM89 or SM120 device (It is slower than "
-        "Marlin W4A16 on other devices).");
-  }
-
  int max_par = 16;
  if (prob_n <= 4096) max_par = 16 * 8;
  int max_shared_mem_new = max_shared_mem;
@@ -414,7 +583,7 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
    int thread_n = thread_n_init;

    int thread_m_blocks = min(div_ceil(prob_m_split, 16), max_thread_m_blocks);
-    int m_block_size_8 = prob_m_split <= 8 && a_type.size_bits() == 16;
+    int m_block_size_8 = prob_m_split <= 8;

    // Set thread config
    exec_config_t exec_cfg;
@@ -428,25 +597,11 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
                  " is not divisible by thread_k = ", thread_k);
    } else {
      // Auto config
-      exec_cfg = determine_exec_config(
-          a_type, b_type, c_type, s_type, prob_m_split, prob_n, prob_k,
-          thread_m_blocks, m_block_size_8, num_bits, group_size, has_act_order,
-          is_k_full, has_zp, is_zp_float, max_shared_mem, sms);
+      exec_cfg = determine_exec_config<scalar_t>(
+          q_type, prob_m_split, prob_n, prob_k, thread_m_blocks, m_block_size_8,
+          num_bits, group_size, has_act_order, is_k_full, has_zp, is_zp_float,
+          max_shared_mem, sms);
      thread_tfg = exec_cfg.tb_cfg;
-      if (thread_tfg.thread_n != -1) {
-        if (prob_n / thread_tfg.thread_n *
-                div_ceil(prob_m_split, thread_m_blocks * 16) * 4 <=
-            sms) {
-          if (is_valid_config({128, 64, 128}, thread_m_blocks, prob_m_split,
-                              prob_n, prob_k, num_bits, group_size,
-                              has_act_order, is_k_full, has_zp, is_zp_float,
-                              max_shared_mem_new)) {
-            thread_tfg = {128, 64, 128};
-            exec_cfg = {1, thread_tfg};
-          }
-        }
-      }
-
      if (thread_tfg.thread_k == -1 && max_thread_m_blocks > 1) {
        max_thread_m_blocks--;
        continue;
@@ -477,10 +632,10 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
        ", has_zp = ", has_zp, ", is_zp_float = ", is_zp_float,
        ", max_shared_mem_new = ", max_shared_mem_new);

-    auto kernel = get_marlin_kernel(
-        a_type, b_type, c_type, s_type, thread_m_blocks, thread_n_blocks,
-        thread_k_blocks, m_block_size_8, has_act_order, has_zp, group_blocks,
-        num_threads, is_zp_float);
+    auto kernel = get_marlin_kernel<scalar_t>(
+        q_type, thread_m_blocks, thread_n_blocks, thread_k_blocks,
+        m_block_size_8, has_act_order, has_zp, group_blocks, num_threads,
+        is_zp_float);

    if (kernel == MarlinDefault) {
      TORCH_CHECK(false, "Unsupported shapes: MNK = [", prob_m, ", ", prob_n,
@@ -502,15 +657,13 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
    // avoid ">>>" being formatted to "> > >"
    // clang-format off
    kernel<<<blocks, num_threads, max_shared_mem_new, stream>>>(
-        A_ptr, B_ptr, C_ptr, C_tmp_ptr, bias_ptr, a_s_ptr, b_s_ptr, g_s_ptr, zp_ptr,
+        A_ptr, B_ptr, C_ptr, C_tmp_ptr, bias_ptr, s_ptr, s2_ptr, zp_ptr,
        g_idx_ptr, num_groups,
        prob_m_split, prob_n, prob_k, lda, locks, has_bias, part_use_atomic_add,
        use_fp32_reduce, max_shared_mem_new);
    // clang-format on

-    bool is_a_8bit = a_type.size_bits() == 8;
-    A_ptr += prob_m_split * (lda / (is_a_8bit ? 16 : 8));
-    a_s_ptr += prob_m_split;
+    A_ptr += prob_m_split * (lda / 8);
    C_ptr += prob_m_split * (prob_n / 8);
    rest_m -= prob_m_split;
  }
@@ -522,73 +675,15 @@ torch::Tensor gptq_marlin_gemm(
    torch::Tensor& a, std::optional<torch::Tensor> c_or_none,
    torch::Tensor& b_q_weight,
    std::optional<torch::Tensor> const& b_bias_or_none, torch::Tensor& b_scales,
-    std::optional<torch::Tensor> const& a_scales_or_none,
    std::optional<torch::Tensor> const& global_scale_or_none,
    std::optional<torch::Tensor> const& b_zeros_or_none,
    std::optional<torch::Tensor> const& g_idx_or_none,
    std::optional<torch::Tensor> const& perm_or_none, torch::Tensor& workspace,
-    vllm::ScalarTypeId const& b_type_id, int64_t size_m, int64_t size_n,
+    vllm::ScalarTypeId const& b_q_type_id, int64_t size_m, int64_t size_n,
    int64_t size_k, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce,
    bool is_zp_float) {
-  vllm::ScalarTypeId a_type_id, c_type_id, s_type_id;
-
-  auto c_dtype = a.dtype();
-  if (a.scalar_type() == at::ScalarType::Half) {
-    a_type_id = vllm::kFloat16.id();
-    c_type_id = vllm::kFloat16.id();
-  } else if (a.scalar_type() == at::ScalarType::BFloat16) {
-    a_type_id = vllm::kBFloat16.id();
-    c_type_id = vllm::kBFloat16.id();
-  } else {
-    c_dtype = b_scales.dtype();
-    if (b_scales.scalar_type() == at::ScalarType::Half) {
-      c_type_id = vllm::kFloat16.id();
-    } else if (b_scales.scalar_type() == at::ScalarType::BFloat16) {
-      c_type_id = vllm::kBFloat16.id();
-    } else {
-      c_type_id = vllm::kBFloat16.id();
-
-      TORCH_CHECK(c_or_none.has_value(), "c must be passed for W4A8-FP4");
-      torch::Tensor c = c_or_none.value();
-      c_dtype = c.dtype();
-
-      if (c.scalar_type() == at::ScalarType::Half) {
-        c_type_id = vllm::kFloat16.id();
-      } else if (c.scalar_type() == at::ScalarType::BFloat16) {
-        c_type_id = vllm::kBFloat16.id();
-      } else {
-        TORCH_CHECK(false, "unsupported c dtype");
-      }
-    }
-
-    if (a.scalar_type() == at::ScalarType::Float8_e4m3fn) {
-      a_type_id = vllm::kFE4M3fn.id();
-    } else if (a.scalar_type() == at::ScalarType::Char) {
-      a_type_id = vllm::kS8.id();
-    } else {
-      TORCH_CHECK(false, "unsupported `a` scalar_type");
-    }
-  }
-
-  s_type_id = c_type_id;
-  if (b_type_id == vllm::kFE2M1f.id()) {
-    if (b_scales.scalar_type() == at::ScalarType::Float8_e4m3fn) {
-      s_type_id = vllm::kFE4M3fn.id();
-    } else if (b_scales.scalar_type() == at::ScalarType::Float8_e8m0fnu) {
-      s_type_id = vllm::kFE8M0fnu.id();
-    } else {
-      TORCH_CHECK(false,
-                  "When b_type = float4_e2m1f, b_scale scalar type must be",
-                  "float8_e4m3fn (for NVFP4) or float8_e8m0fnu (for MXFP4).");
-    }
-  }
-
-  vllm::ScalarType a_type = vllm::ScalarType::from_id(a_type_id);
-  vllm::ScalarType b_type = vllm::ScalarType::from_id(b_type_id);
-  vllm::ScalarType c_type = vllm::ScalarType::from_id(c_type_id);
-  vllm::ScalarType s_type = vllm::ScalarType::from_id(s_type_id);
-
-  int pack_factor = 32 / b_type.size_bits();
+  vllm::ScalarType const b_q_type = vllm::ScalarType::from_id(b_q_type_id);
+  int pack_factor = 32 / b_q_type.size_bits();

  // Verify A
  TORCH_CHECK(a.size(0) == size_m, "Shape mismatch: a.size(0) = ", a.size(0),
@@ -626,21 +721,6 @@ torch::Tensor gptq_marlin_gemm(
  TORCH_CHECK(b_scales.device().is_cuda(), "b_scales is not on GPU");
  TORCH_CHECK(b_scales.is_contiguous(), "b_scales is not contiguous");

-  torch::Tensor a_scales;
-  auto options = torch::TensorOptions().dtype(c_dtype).device(a.device());
-  auto options_fp32 =
-      torch::TensorOptions().dtype(at::kFloat).device(a.device());
-
-  if (a_scales_or_none.has_value()) {
-    a_scales = a_scales_or_none.value();
-    TORCH_CHECK(a_type.size_bits() == 8,
-                "a_scales can only be used for 8bit activation.");
-  } else {
-    a_scales = torch::empty({0}, options_fp32);
-    TORCH_CHECK(a_type.size_bits() != 8,
-                "the a_scales parameter must be passed for 8bit activation.");
-  }
-
  // thread_k: `k` size of a thread_tile in `weights` (can usually be left as
  // auto -1)
  int thread_k = -1;
@@ -653,6 +733,7 @@ torch::Tensor gptq_marlin_gemm(

  // Alloc buffers
  const at::cuda::OptionalCUDAGuard device_guard(device_of(a));
+  auto options = torch::TensorOptions().dtype(a.dtype()).device(a.device());
  torch::Tensor c;
  if (c_or_none.has_value()) {
    c = c_or_none.value();
@@ -669,6 +750,8 @@ torch::Tensor gptq_marlin_gemm(

  // Alloc C tmp buffer that is going to be used for the global reduce
  torch::Tensor c_tmp;
+  auto options_fp32 =
+      torch::TensorOptions().dtype(at::kFloat).device(a.device());
  if (use_fp32_reduce) {
    int max_m_block_size = (size_m + 16 - 1) / 16 * 16;
    max_m_block_size = min(max_m_block_size, 64);
@@ -738,11 +821,11 @@ torch::Tensor gptq_marlin_gemm(
  torch::Tensor global_scale;
  if (global_scale_or_none.has_value()) {
    global_scale = global_scale_or_none.value();
-    TORCH_CHECK(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn,
+    TORCH_CHECK(b_q_type == vllm::kFE2M1f && group_size == 16,
                "global_scale can only be used for nvfp4 format.");
  } else {
    global_scale = torch::empty({0}, options);
-    TORCH_CHECK(!(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn),
+    TORCH_CHECK(!(b_q_type == vllm::kFE2M1f && group_size == 16),
                "the global_scale parameter must be passed for nvfp4 format.");
  }

@@ -769,15 +852,15 @@ torch::Tensor gptq_marlin_gemm(
  bool has_zp = b_zeros.size(-1) > 0;
  if (has_zp) {
    TORCH_CHECK(
-        b_type == vllm::kU4 || b_type == vllm::kU8,
-        "b_type must be u4 or u8 when has_zp = True. Got = ", b_type.str());
+        b_q_type == vllm::kU4 || b_q_type == vllm::kU8,
+        "b_q_type must be u4 or u8 when has_zp = True. Got = ", b_q_type.str());
  } else {
-    TORCH_CHECK(b_type == vllm::kU4B8 || b_type == vllm::kU8B128 ||
-                    b_type == vllm::kS4 || b_type == vllm::kS8 ||
-                    b_type == vllm::kFE4M3fn || b_type == vllm::kFE2M1f,
-                "b_type must be uint4b8, uint8b128, int4, int8, "
-                "float8_e4m3fn or float4_e2m1f when has_zp = False. Got = ",
-                b_type.str());
+    TORCH_CHECK(b_q_type == vllm::kU4B8 || b_q_type == vllm::kU8B128 ||
+                    b_q_type == vllm::kFE4M3fn || b_q_type == vllm::kFE2M1f,
+                "b_q_type must be uint4b8, uint8b128, float8_e4m3fn or "
+                "float4_e2m1f when "
+                "has_zp = False. Got = ",
+                b_q_type.str());
  }

  if (has_zp && is_zp_float) {
@@ -819,27 +902,59 @@ torch::Tensor gptq_marlin_gemm(
              " is below min_workspace_size = ", min_workspace_size);

  int dev = a.get_device();
+  if (a.scalar_type() == at::ScalarType::Half) {
+    void* scales_ptr;
+    if (b_q_type == vllm::kFE2M1f) {
+      if (group_size == 16)
+        scales_ptr = b_scales.data_ptr<at::Float8_e4m3fn>();
+      else if (group_size == 32)
+        scales_ptr = b_scales.data_ptr<at::Float8_e8m0fnu>();
+      else
+        TORCH_CHECK(false,
+                    "float4_e2m1f only supports group_size == 16 (NVFP4) ",
+                    "and group_size == 32 (MXFP4)");
+    } else {
+      scales_ptr = b_scales.data_ptr<at::Half>();
+    }

-  TORCH_CHECK(a_scales.scalar_type() == at::ScalarType::Float,
-              "scalar type of a_scales must be float");
-  TORCH_CHECK(global_scale.scalar_type() == c.scalar_type(),
-              "scalar type of global_scale must be the same with c");
-  if (a_type.size_bits() == 16) {
-    TORCH_CHECK(
-        a.scalar_type() == c.scalar_type(),
-        "scalar type of a must be the same with c for 16 bit activation");
+    marlin::marlin_mm<half>(
+        a.data_ptr<at::Half>(), b_q_weight.data_ptr(), c.data_ptr<at::Half>(),
+        c_tmp.data_ptr<float>(), b_bias.data_ptr<at::Half>(), scales_ptr,
+        global_scale.data_ptr<at::Half>(), b_zeros.data_ptr(), g_idx.data_ptr(),
+        perm.data_ptr(), a_tmp.data_ptr<at::Half>(), size_m, size_n, size_k,
+        a.stride(0), workspace.data_ptr(), b_q_type, has_bias, has_act_order,
+        is_k_full, has_zp, num_groups, group_size, dev,
+        at::cuda::getCurrentCUDAStream(dev), thread_k, thread_n, sms,
+        use_atomic_add, use_fp32_reduce, is_zp_float);
+  } else if (a.scalar_type() == at::ScalarType::BFloat16) {
+    void* scales_ptr;
+    if (b_q_type == vllm::kFE2M1f) {
+      if (group_size == 16)
+        scales_ptr = b_scales.data_ptr<at::Float8_e4m3fn>();
+      else if (group_size == 32)
+        scales_ptr = b_scales.data_ptr<at::Float8_e8m0fnu>();
+      else
+        TORCH_CHECK(false,
+                    "float4_e2m1f only supports group_size == 16 (NVFP4) ",
+                    "and group_size == 32 (MXFP4)");
+    } else {
+      scales_ptr = b_scales.data_ptr<at::BFloat16>();
+    }
+
+    marlin::marlin_mm<nv_bfloat16>(
+        a.data_ptr<at::BFloat16>(), b_q_weight.data_ptr(),
+        c.data_ptr<at::BFloat16>(), c_tmp.data_ptr<float>(),
+        b_bias.data_ptr<at::BFloat16>(), scales_ptr,
+        global_scale.data_ptr<at::BFloat16>(), b_zeros.data_ptr(),
+        g_idx.data_ptr(), perm.data_ptr(), a_tmp.data_ptr<at::BFloat16>(),
+        size_m, size_n, size_k, a.stride(0), workspace.data_ptr(), b_q_type,
+        has_bias, has_act_order, is_k_full, has_zp, num_groups, group_size, dev,
+        at::cuda::getCurrentCUDAStream(dev), thread_k, thread_n, sms,
+        use_atomic_add, use_fp32_reduce, is_zp_float);
+  } else {
+    TORCH_CHECK(false, "gpt_marlin_gemm only supports bfloat16 and float16");
  }

-  marlin::marlin_mm(
-      a.data_ptr(), b_q_weight.data_ptr(), c.data_ptr(), c_tmp.data_ptr(),
-      b_bias.data_ptr(), a_scales.data_ptr(), b_scales.data_ptr(),
-      global_scale.data_ptr(), b_zeros.data_ptr(), g_idx.data_ptr(),
-      perm.data_ptr(), a_tmp.data_ptr(), size_m, size_n, size_k, a.stride(0),
-      workspace.data_ptr(), a_type, b_type, c_type, s_type, has_bias,
-      has_act_order, is_k_full, has_zp, num_groups, group_size, dev,
-      at::cuda::getCurrentCUDAStream(dev), thread_k, thread_n, sms,
-      use_atomic_add, use_fp32_reduce, is_zp_float);
-
  return c;
 }

@@ -4,18 +4,15 @@

 namespace marlin {

-template <int const num_threads, int const num_bits, bool const has_perm,
-          bool is_a_8bit>
+template <int const num_threads, int const num_bits, bool const has_perm>
 __global__ void gptq_marlin_repack_kernel(
    uint32_t const* __restrict__ b_q_weight_ptr,
    uint32_t const* __restrict__ perm_ptr, uint32_t* __restrict__ out_ptr,
    int size_k, int size_n) {
  constexpr int pack_factor = 32 / num_bits;

-  constexpr int target_tile_n_size = tile_n_size / (is_a_8bit ? 2 : 1);
-  constexpr int target_tile_k_size = tile_k_size * (is_a_8bit ? 2 : 1);
-  int k_tiles = size_k / target_tile_k_size;
-  int n_tiles = size_n / target_tile_n_size;
+  int k_tiles = size_k / tile_k_size;
+  int n_tiles = size_n / tile_n_size;
  int block_k_tiles = div_ceil(k_tiles, gridDim.x);

  auto start_k_tile = blockIdx.x * block_k_tiles;
@@ -37,7 +34,7 @@ __global__ void gptq_marlin_repack_kernel(

  extern __shared__ int4 sh[];

-  constexpr int perm_size = target_tile_k_size / 4;
+  constexpr int perm_size = tile_k_size / 4;

  int4* sh_perm_ptr = sh;
  int4* sh_pipe_ptr = sh_perm_ptr;
@@ -45,14 +42,14 @@ __global__ void gptq_marlin_repack_kernel(
    sh_pipe_ptr += perm_size;
  }

-  constexpr int tile_ints = target_tile_k_size / pack_factor;
+  constexpr int tile_ints = tile_k_size / pack_factor;

-  constexpr int stage_n_threads = target_tile_n_size / 4;
-  constexpr int stage_k_threads = has_perm ? target_tile_k_size : tile_ints;
+  constexpr int stage_n_threads = tile_n_size / 4;
+  constexpr int stage_k_threads = has_perm ? tile_k_size : tile_ints;
  constexpr int stage_size = stage_k_threads * stage_n_threads;

  auto load_perm_to_shared = [&](int k_tile_id) {
-    int first_k_int4 = (k_tile_id * target_tile_k_size) / 4;
+    int first_k_int4 = (k_tile_id * tile_k_size) / 4;

    int4 const* perm_int4_ptr = reinterpret_cast<int4 const*>(perm_ptr);

@@ -68,7 +65,7 @@ __global__ void gptq_marlin_repack_kernel(
      return;
    }

-    int first_n = n_tile_id * target_tile_n_size;
+    int first_n = n_tile_id * tile_n_size;

    int4* sh_ptr = sh_pipe_ptr + stage_size * pipe;

@@ -94,7 +91,7 @@ __global__ void gptq_marlin_repack_kernel(
        auto k_id = threadIdx.x / stage_n_threads;
        auto n_id = threadIdx.x % stage_n_threads;

-        int first_k = k_tile_id * target_tile_k_size;
+        int first_k = k_tile_id * tile_k_size;
        int first_k_packed = first_k / pack_factor;

        cp_async4(&sh_ptr[k_id * stage_n_threads + n_id],
@@ -120,13 +117,13 @@ __global__ void gptq_marlin_repack_kernel(
    }

    int tc_col = th_id / 4;
-    int tc_row = (th_id % 4) * (is_a_8bit ? 4 : 2);
+    int tc_row = (th_id % 4) * 2;

    constexpr int tc_offsets[4] = {0, 1, 8, 9};

-    int cur_n = (warp_id / (is_a_8bit ? 2 : 1)) * 16 + tc_col;
+    int cur_n = warp_id * 16 + tc_col;

-    constexpr int sh_stride = target_tile_n_size;
+    constexpr int sh_stride = 64;
    constexpr uint32_t mask = (1 << num_bits) - 1;

    int4* sh_stage_ptr = sh_pipe_ptr + stage_size * pipe;
@@ -137,7 +134,6 @@ __global__ void gptq_marlin_repack_kernel(
    uint32_t vals[8];

    if constexpr (has_perm) {
-      static_assert(!is_a_8bit);
      for (int i = 0; i < 4; i++) {
        int k_idx = tc_row + tc_offsets[i];

@@ -160,49 +156,28 @@ __global__ void gptq_marlin_repack_kernel(

 #pragma unroll
      for (int i = 0; i < tile_ints; i++) {
-        if constexpr (is_a_8bit) {
-          b1_vals[i] =
-              sh_stage_int_ptr[cur_n + sh_stride * i + (warp_id % 2) * 8];
-        } else {
-          b1_vals[i] = sh_stage_int_ptr[cur_n + sh_stride * i];
-          b2_vals[i] = sh_stage_int_ptr[cur_n + 8 + sh_stride * i];
-        }
+        b1_vals[i] = sh_stage_int_ptr[cur_n + sh_stride * i];
+        b2_vals[i] = sh_stage_int_ptr[cur_n + 8 + sh_stride * i];
      }

 #pragma unroll
      for (int i = 0; i < 4; i++) {
-        int cur_elem = tc_row + (is_a_8bit ? i : tc_offsets[i]);
+        int cur_elem = tc_row + tc_offsets[i];
        int cur_int = cur_elem / pack_factor;
        int cur_pos = cur_elem % pack_factor;

        vals[i] = (b1_vals[cur_int] >> (cur_pos * num_bits)) & mask;
-        if constexpr (is_a_8bit)
-          vals[4 + i] =
-              (b1_vals[cur_int + tile_ints / 2] >> (cur_pos * num_bits)) & mask;
-        else
-          vals[4 + i] = (b2_vals[cur_int] >> (cur_pos * num_bits)) & mask;
+        vals[4 + i] = (b2_vals[cur_int] >> (cur_pos * num_bits)) & mask;
      }
    }

-    constexpr int tile_size =
-        target_tile_k_size * target_tile_n_size / pack_factor;
+    constexpr int tile_size = tile_k_size * tile_n_size / pack_factor;
    int out_offset = (k_tile_id * n_tiles + n_tile_id) * tile_size;

    // Result of:
    // https://github.com/NVIDIA/FasterTransformer/blob/main/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h
-    if constexpr (!is_a_8bit && num_bits == 4) {
-      int pack_idx[8] = {0, 2, 4, 6, 1, 3, 5, 7};
-
-      uint32_t res = 0;
-#pragma unroll
-      for (int i = 0; i < 8; i++) {
-        res |= vals[pack_idx[i]] << (i * 4);
-      }
-
-      out_ptr[out_offset + th_id * 4 + warp_id] = res;
-
-    } else if constexpr (is_a_8bit && num_bits == 4) {
-      int pack_idx[8] = {0, 4, 1, 5, 2, 6, 3, 7};
+    if constexpr (num_bits == 4) {
+      constexpr int pack_idx[8] = {0, 2, 4, 6, 1, 3, 5, 7};

      uint32_t res = 0;
 #pragma unroll
@@ -219,9 +194,8 @@ __global__ void gptq_marlin_repack_kernel(
      uint32_t res2 = 0;
 #pragma unroll
      for (int i = 0; i < 4; i++) {
-        const int ii = is_a_8bit ? i : pack_idx[i];
-        res1 |= vals[ii] << (i * 8);
-        res2 |= vals[4 + ii] << (i * 8);
+        res1 |= vals[pack_idx[i]] << (i * 8);
+        res2 |= vals[4 + pack_idx[i]] << (i * 8);
      }

      out_ptr[out_offset + th_id * 8 + (warp_id * 2) + 0] = res1;
@@ -262,22 +236,21 @@ __global__ void gptq_marlin_repack_kernel(

 }  // namespace marlin

-#define CALL_IF(NUM_BITS, HAS_PERM, IS_A_8BIT)                              \
-  else if (num_bits == NUM_BITS && has_perm == HAS_PERM &&                  \
-           is_a_8bit == IS_A_8BIT) {                                        \
+#define CALL_IF(NUM_BITS, HAS_PERM)                                         \
+  else if (num_bits == NUM_BITS && has_perm == HAS_PERM) {                  \
    cudaFuncSetAttribute(                                                   \
        marlin::gptq_marlin_repack_kernel<marlin::repack_threads, NUM_BITS, \
-                                          HAS_PERM, IS_A_8BIT>,             \
+                                          HAS_PERM>,                        \
        cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem);       \
    marlin::gptq_marlin_repack_kernel<marlin::repack_threads, NUM_BITS,     \
-                                      HAS_PERM, IS_A_8BIT>                  \
+                                      HAS_PERM>                             \
        <<<blocks, marlin::repack_threads, max_shared_mem, stream>>>(       \
            b_q_weight_ptr, perm_ptr, out_ptr, size_k, size_n);             \
  }

 torch::Tensor gptq_marlin_repack(torch::Tensor& b_q_weight, torch::Tensor& perm,
                                 int64_t size_k, int64_t size_n,
-                                 int64_t num_bits, bool is_a_8bit) {
+                                 int64_t num_bits) {
  // Verify compatibility with marlin tile of 16x64
  TORCH_CHECK(size_k % marlin::tile_k_size == 0, "size_k = ", size_k,
              " is not divisible by tile_k_size = ", marlin::tile_k_size);
@@ -336,17 +309,13 @@ torch::Tensor gptq_marlin_repack(torch::Tensor& b_q_weight, torch::Tensor& perm,

  if (false) {
  }
-  CALL_IF(4, false, false)
-  CALL_IF(4, true, false)
-  CALL_IF(8, false, false)
-  CALL_IF(8, true, false)
-
-  CALL_IF(4, false, true)
-  CALL_IF(8, false, true)
-
+  CALL_IF(4, false)
+  CALL_IF(4, true)
+  CALL_IF(8, false)
+  CALL_IF(8, true)
  else {
    TORCH_CHECK(false, "Unsupported repack config: num_bits = ", num_bits,
-                ", has_perm = ", has_perm, ", is_a_8bit = ", is_a_8bit);
+                ", has_perm = ", has_perm);
  }

  return out;
@@ -11,19 +11,17 @@
  const int4 *__restrict__ A, const int4 *__restrict__ B,                      \
      int4 *__restrict__ C, int4 *__restrict__ C_tmp,                          \
      const int4 *__restrict__ b_bias_ptr,                                     \
-      const float *__restrict__ a_scales_ptr,                                  \
      const int4 *__restrict__ scales_ptr,                                     \
-      const uint16_t *__restrict__ global_scale_ptr,                           \
+      const uint16_t *__restrict__ scale2_ptr,                                 \
      const int4 *__restrict__ zp_ptr, const int *__restrict__ g_idx,          \
      int num_groups, int prob_m, int prob_n, int prob_k, int lda, int *locks, \
      bool has_bias, bool use_atomic_add, bool use_fp32_reduce,                \
      int max_shared_mem

 namespace MARLIN_NAMESPACE_NAME {
-template <const vllm::ScalarTypeId a_type_id,  // A ScalarType id
-          const vllm::ScalarTypeId b_type_id,  // B ScalarType id
-          const vllm::ScalarTypeId c_type_id,  // C ScalarType id
-          const vllm::ScalarTypeId s_type_id,  // B_SCALE ScalarType id
+template <typename scalar_t,  // compute dtype, half or nv_float16
+          const vllm::ScalarTypeId w_type_id,  // weight ScalarType id
+          const vllm::ScalarTypeId s_type_id,  // weight ScalarType id
          const int threads,          // number of threads in a threadblock
          const int thread_m_blocks,  // number of 16x16 blocks in the m
                                      // dimension (batchsize) of the
@@ -55,45 +55,6 @@ constexpr int div_ceil(int a, int b) { return (a + b - 1) / b; }
 // No support for async
 #else

-__device__ inline void cp_async1_ca_pred(void* smem_ptr, const void* glob_ptr,
-                                         bool pred = true) {
-  const int BYTES = 4;
-  uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
-  asm volatile(
-      "{\n"
-      "   .reg .pred p;\n"
-      "   setp.ne.b32 p, %0, 0;\n"
-      "   @p cp.async.ca.shared.global [%1], [%2], %3;\n"
-      "}\n" ::"r"((int)pred),
-      "r"(smem), "l"(glob_ptr), "n"(BYTES));
-}
-
-__device__ inline void cp_async2_ca_pred(void* smem_ptr, const void* glob_ptr,
-                                         bool pred = true) {
-  const int BYTES = 8;
-  uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
-  asm volatile(
-      "{\n"
-      "   .reg .pred p;\n"
-      "   setp.ne.b32 p, %0, 0;\n"
-      "   @p cp.async.ca.shared.global [%1], [%2], %3;\n"
-      "}\n" ::"r"((int)pred),
-      "r"(smem), "l"(glob_ptr), "n"(BYTES));
-}
-
-__device__ inline void cp_async4_ca_pred(void* smem_ptr, const void* glob_ptr,
-                                         bool pred = true) {
-  const int BYTES = 16;
-  uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
-  asm volatile(
-      "{\n"
-      "   .reg .pred p;\n"
-      "   setp.ne.b32 p, %0, 0;\n"
-      "   @p cp.async.ca.shared.global [%1], [%2], %3;\n"
-      "}\n" ::"r"((int)pred),
-      "r"(smem), "l"(glob_ptr), "n"(BYTES));
-}
-
 __device__ inline void cp_async4_pred(void* smem_ptr, const void* glob_ptr,
                                      bool pred = true) {
  const int BYTES = 16;
@@ -2,10 +2,8 @@
 #ifndef _data_types_cuh
 #define _data_types_cuh
 #include "marlin.cuh"
-#include "core/scalar_type.hpp"
 #include <cuda_fp16.h>
 #include <cuda_bf16.h>
-#include <cuda_fp8.h>

 #ifndef MARLIN_NAMESPACE_NAME
  #define MARLIN_NAMESPACE_NAME marlin
@@ -13,16 +11,14 @@

 namespace MARLIN_NAMESPACE_NAME {

-template <long scalar_type_id>
-class MarlinScalarType {};
+template <typename scalar_t>
+class ScalarType {};

 template <>
-class MarlinScalarType<vllm::kFloat16.id()> {
+class ScalarType<half> {
 public:
  using scalar_t = half;
  using scalar_t2 = half2;
-  using scalar_t4 = half2;
-  using scalar_32bit_t = half2;

  // Matrix fragments for tensor core instructions; their precise layout is
  // documented here:
@@ -31,7 +27,6 @@ class MarlinScalarType<vllm::kFloat16.id()> {
  using FragB = Vec<half2, 2>;
  using FragC = Vec<float, 4>;
  using FragS = Vec<half2, 1>;
-  using FragS0 = Vec<__nv_fp8x2_e4m3, 1>;
  using FragZP = Vec<half2, 4>;

  static __device__ float inline num2float(const half x) {
@@ -49,25 +44,18 @@ class MarlinScalarType<vllm::kFloat16.id()> {
  static __host__ __device__ half inline float2num(const float x) {
    return __float2half(x);
  }
-
-  static __host__ __device__ float2 inline num22float2(const half2 x) {
-    return __half22float2(x);
-  }
 };

 template <>
-class MarlinScalarType<vllm::kBFloat16.id()> {
+class ScalarType<nv_bfloat16> {
 public:
  using scalar_t = nv_bfloat16;
  using scalar_t2 = nv_bfloat162;
-  using scalar_t4 = nv_bfloat162;
-  using scalar_32bit_t = nv_bfloat162;

  using FragA = Vec<nv_bfloat162, 4>;
  using FragB = Vec<nv_bfloat162, 2>;
  using FragC = Vec<float, 4>;
  using FragS = Vec<nv_bfloat162, 1>;
-  using FragS0 = Vec<__nv_fp8x2_e4m3, 1>;
  using FragZP = Vec<nv_bfloat162, 4>;

 #if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800
@@ -87,63 +75,9 @@ class MarlinScalarType<vllm::kBFloat16.id()> {
  static __host__ __device__ nv_bfloat16 inline float2num(const float x) {
    return __float2bfloat16(x);
  }
-
-  static __host__ __device__ float2 inline num22float2(const nv_bfloat162 x) {
-    return __bfloat1622float2(x);
-  }
 #endif
 };

-template <>
-class MarlinScalarType<vllm::kFE4M3fn.id()> {
- public:
-  using scalar_t = __nv_fp8_e4m3;
-  using scalar_t2 = __nv_fp8x2_e4m3;
-  using scalar_t4 = __nv_fp8x4_e4m3;
-  using scalar_32bit_t = __nv_fp8x4_e4m3;
-
-  using FragA = Vec<__nv_fp8x4_e4m3, 4>;
-  using FragB = Vec<__nv_fp8x4_e4m3, 2>;
-  using FragC = Vec<float, 4>;
-  using FragZP = Vec<__nv_fp8x2_e4m3, 4>;
-
-  static __host__ __device__
-      float2 inline num22float2(const __nv_fp8x2_e4m3 x) {
-    return (float2)x;
-  }
-};
-
-template <>
-class MarlinScalarType<vllm::kS8.id()> {
- public:
-  using scalar_t = int8_t;
-  using scalar_t2 = int16_t;
-  using scalar_t4 = int32_t;
-  using scalar_32bit_t = int32_t;
-
-  using FragA = Vec<int32_t, 4>;
-  using FragB = Vec<int32_t, 2>;
-  using FragC = Vec<float, 4>;
-  using FragZP = Vec<int16_t, 4>;
-};
-
-template <typename scalar_t>
-class MarlinScalarType2 {};
-
-template <>
-class MarlinScalarType2<half> : public MarlinScalarType<vllm::kFloat16.id()> {};
-
-template <>
-class MarlinScalarType2<nv_bfloat16>
-    : public MarlinScalarType<vllm::kBFloat16.id()> {};
-
-template <>
-class MarlinScalarType2<__nv_fp8_e4m3>
-    : public MarlinScalarType<vllm::kFE4M3fn.id()> {};
-
-template <>
-class MarlinScalarType2<int8_t> : public MarlinScalarType<vllm::kS8.id()> {};
-
 }  // namespace MARLIN_NAMESPACE_NAME

 #endif
@@ -1,106 +0,0 @@
-
-
-#include "marlin.cuh"
-
-#include "core/registration.h"
-
-// for only non-zp format (like gptq)
-__global__ void marlin_int4_fp8_preprocess_kernel_without_zp(
-    // qweight: (size_k * size_n // 8,)
-    const int32_t* __restrict__ qweight,
-    // output: same shape with qweight
-    int32_t* __restrict__ output) {
-  int32_t val = qweight[blockIdx.x * 32 + threadIdx.x];
-  int32_t new_val = 0;
-
-#pragma unroll
-  for (int32_t i = 0; i < 8; i++) {
-    int32_t single_val = val & 0xF;
-    single_val = single_val >= 8 ? single_val - 8 : 15 - single_val;
-    new_val |= single_val << (i * 4);
-    val >>= 4;
-  }
-
-  output[blockIdx.x * 32 + threadIdx.x] = new_val;
-}
-
-// for awq format only (with zp and with awq weight layout)
-__global__ void marlin_int4_fp8_preprocess_kernel_awq(
-    // AWQ qweight: (size_k, size_n // 8)
-    const int32_t* __restrict__ qweight,
-    // output: same shape with qweight
-    int32_t* __restrict__ output,
-    // AWQ zeros: (size_k // group_size, size_n // 8)
-    const int32_t* __restrict__ qzeros, int32_t size_n, int32_t size_k,
-    int32_t group_size) {
-  int32_t val =
-      qweight[(blockIdx.x * 32 + threadIdx.x) * size_n / 8 + blockIdx.y];
-  int32_t zero =
-      qzeros[(blockIdx.x * 32 + threadIdx.x) / group_size * size_n / 8 +
-             blockIdx.y];
-  int32_t new_val = 0;
-
-#pragma unroll
-  for (int32_t i = 0; i < 8; i++) {
-    int32_t single_val = val & 0xF;
-    int32_t single_zero = zero & 0xF;
-
-    single_val =
-        single_val >= single_zero ? single_val - single_zero : 15 - single_val;
-    new_val |= single_val << (i * 4);
-    val >>= 4;
-    zero >>= 4;
-  }
-
-  output[(blockIdx.x * 32 + threadIdx.x) * size_n / 8 + blockIdx.y] = new_val;
-}
-
-torch::Tensor marlin_int4_fp8_preprocess(
-    torch::Tensor& qweight, std::optional<torch::Tensor> qzeros_or_none,
-    bool inplace) {
-  TORCH_CHECK(qweight.device().is_cuda(), "qweight is not on GPU");
-  TORCH_CHECK(qweight.scalar_type() == at::ScalarType::Int,
-              "qweight.dtype != torch.int32");
-
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(qweight));
-
-  torch::Tensor output = inplace ? qweight : torch::empty_like(qweight);
-
-  if (!qzeros_or_none.has_value()) {
-    TORCH_CHECK(qweight.numel() * 8 % 256 == 0,
-                "qweight.numel() * 8 % 256 != 0");
-
-    int blocks = qweight.numel() * 8 / 256;
-    marlin_int4_fp8_preprocess_kernel_without_zp<<<blocks, 32>>>(
-        (const int32_t*)qweight.data_ptr(), (int32_t*)output.data_ptr());
-  } else {
-    int32_t size_k = qweight.size(0);
-    int32_t size_n = qweight.size(1) * 8;
-    torch::Tensor qzeros = qzeros_or_none.value();
-
-    TORCH_CHECK(size_k % 32 == 0, "size_k % 32 != 0");
-    TORCH_CHECK(qzeros.device().is_cuda(), "qzeros is not on GPU");
-    TORCH_CHECK(qzeros.scalar_type() == at::ScalarType::Int,
-                "qweight.dtype != torch.int32");
-    TORCH_CHECK(device_of(qweight) == device_of(qzeros),
-                "qzeros is not on the same device with qweight");
-
-    int32_t group_size = qweight.size(0) / qzeros.size(0);
-    TORCH_CHECK(qweight.size(1) == qzeros.size(1),
-                "qweight.size(1) != qzeros.size(1)");
-    TORCH_CHECK(qweight.size(0) % qzeros.size(0) == 0,
-                "qweight.size(0) % qzeros.size(0) != 0");
-    TORCH_CHECK(group_size % 8 == 0, "group_size % 8 != 0");
-
-    dim3 blocks(size_k / 32, size_n / 8);
-    marlin_int4_fp8_preprocess_kernel_awq<<<blocks, 32>>>(
-        (const int32_t*)qweight.data_ptr(), (int32_t*)output.data_ptr(),
-        (const int32_t*)qzeros.data_ptr(), size_n, size_k, group_size);
-  }
-
-  return output;
-}
-
-TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
-  m.impl("marlin_int4_fp8_preprocess", &marlin_int4_fp8_preprocess);
-}
@@ -67,9 +67,9 @@ void cutlass_scaled_mm_sm100(torch::Tensor& c, torch::Tensor const& a,
                             std::optional<torch::Tensor> const& bias);
 #endif

-#if (defined(ENABLE_CUTLASS_MOE_SM90) && ENABLE_CUTLASS_MOE_SM90) ||   \
-    (defined(ENABLE_CUTLASS_MOE_SM100) && ENABLE_CUTLASS_MOE_SM100) || \
-    (defined(ENABLE_CUTLASS_MOE_SM120) && ENABLE_CUTLASS_MOE_SM120)
+#if defined(ENABLE_SCALED_MM_SM90) && ENABLE_SCALED_MM_SM90 ||   \
+    defined(ENABLE_SCALED_MM_SM100) && ENABLE_SCALED_MM_SM100 || \
+    defined(ENABLE_SCALED_MM_SM120) && ENABLE_SCALED_MM_SM120
 void get_cutlass_moe_mm_data_caller(
    const torch::Tensor& topk_ids, torch::Tensor& expert_offsets,
    torch::Tensor& problem_sizes1, torch::Tensor& problem_sizes2,
@@ -284,9 +284,8 @@ void get_cutlass_moe_mm_data(
  // This function currently gets compiled only if we have a valid cutlass moe
  // mm to run it for.
  int32_t version_num = get_sm_version_num();
-#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) ||   \
-    (defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100) || \
-    (defined ENABLE_CUTLASS_MOE_SM120 && ENABLE_CUTLASS_MOE_SM120)
+#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
+    (defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100)
  get_cutlass_moe_mm_data_caller(topk_ids, expert_offsets, problem_sizes1,
                                 problem_sizes2, input_permutation,
                                 output_permutation, num_experts, n, k,
@@ -297,7 +296,7 @@ void get_cutlass_moe_mm_data(
      false,
      "No compiled get_cutlass_moe_mm_data: no cutlass_scaled_mm kernel for "
      "CUDA device capability: ",
-      version_num, ". Required capability: 90, 100, or 120");
+      version_num, ". Required capability: 90 or 100");
 }

 void get_cutlass_moe_mm_problem_sizes(
@@ -305,9 +304,8 @@ void get_cutlass_moe_mm_problem_sizes(
    torch::Tensor& problem_sizes2, const int64_t num_experts, const int64_t n,
    const int64_t k, const std::optional<torch::Tensor>& blockscale_offsets) {
  int32_t version_num = get_sm_version_num();
-#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) ||   \
-    (defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100) || \
-    (defined ENABLE_CUTLASS_MOE_SM120 && ENABLE_CUTLASS_MOE_SM120)
+#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
+    (defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100)
  get_cutlass_moe_mm_problem_sizes_caller(topk_ids, problem_sizes1,
                                          problem_sizes2, num_experts, n, k,
                                          blockscale_offsets);
@@ -317,7 +315,7 @@ void get_cutlass_moe_mm_problem_sizes(
      false,
      "No compiled get_cutlass_moe_mm_problem_sizes: no cutlass_scaled_mm "
      "kernel for CUDA device capability: ",
-      version_num, ". Required capability: 90, 100, or 120");
+      version_num, ". Required capability: 90 or 100");
 }

 void get_cutlass_pplx_moe_mm_data(torch::Tensor& expert_offsets,
@@ -330,9 +328,8 @@ void get_cutlass_pplx_moe_mm_data(torch::Tensor& expert_offsets,
  // This function currently gets compiled only if we have a valid cutlass moe
  // mm to run it for.
  int32_t version_num = get_sm_version_num();
-#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) ||   \
-    (defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100) || \
-    (defined ENABLE_CUTLASS_MOE_SM120 && ENABLE_CUTLASS_MOE_SM120)
+#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
+    (defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100)
  get_cutlass_pplx_moe_mm_data_caller(expert_offsets, problem_sizes1,
                                      problem_sizes2, expert_num_tokens,
                                      num_local_experts, padded_m, n, k);
@@ -342,7 +339,7 @@ void get_cutlass_pplx_moe_mm_data(torch::Tensor& expert_offsets,
      false,
      "No compiled get_cutlass_pplx_moe_mm_data: no cutlass_scaled_mm kernel "
      "for CUDA device capability: ",
-      version_num, ". Required capability: 90, 100, or 120");
+      version_num, ". Required capability: 90 or 100");
 }

 void cutlass_scaled_mm_azp(torch::Tensor& c, torch::Tensor const& a,
@@ -20,6 +20,18 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  // vLLM custom ops
  //

+  // The default behavior in PyTorch 2.6 was changed to "requires_contiguous",
+  // so we need
+  // to override this for many GEMMs with the following tag. Otherwise,
+  // torch.compile will force all input tensors to be contiguous(), which
+  // will break many custom ops that require column-major weight matrices.
+  // This was a bug and PyTorch 2.7 has since fixed this.
+#if TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 6
+  #define stride_tag at::Tag::needs_fixed_stride_order
+#else
+  #define stride_tag
+#endif
+
  ops.def(
      "persistent_masked_m_silu_mul_quant(Tensor input, Tensor counts, Tensor! "
      "y_q, Tensor! y_s,"
@@ -63,6 +75,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
      "    int blocksparse_head_sliding_step) -> ()");
  ops.impl("paged_attention_v2", torch::kCUDA, &paged_attention_v2);

+#ifndef USE_ROCM
  // Merge attn states
  // Implements section 2.2 of https://www.arxiv.org/pdf/2501.01005
  // can be used to combine partial attention results (in the split-KV case)
@@ -75,7 +88,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
      "    Tensor suffix_output,"
      "    Tensor suffix_lse) -> ()");
  ops.impl("merge_attn_states", torch::kCUDA, &merge_attn_states);
-#ifndef USE_ROCM
+
  ops.def(
      "convert_vertical_slash_indexes("
      "   Tensor! block_count, Tensor! block_offset, "
@@ -228,13 +241,15 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  // Quantized GEMM for AWQ.
  ops.def(
      "awq_gemm(Tensor _in_feats, Tensor _kernel, Tensor _scaling_factors, "
-      "Tensor _zeros, SymInt split_k_iters) -> Tensor");
+      "Tensor _zeros, SymInt split_k_iters) -> Tensor",
+      {stride_tag});
  ops.impl("awq_gemm", torch::kCUDA, &awq_gemm);

  // Dequantization for AWQ.
  ops.def(
      "awq_dequantize(Tensor _kernel, Tensor _scaling_factors, "
-      "Tensor _zeros, SymInt split_k_iters, int thx, int thy) -> Tensor");
+      "Tensor _zeros, SymInt split_k_iters, int thx, int thy) -> Tensor",
+      {stride_tag});
  ops.impl("awq_dequantize", torch::kCUDA, &awq_dequantize);

  // Note about marlin kernel 'workspace' arguments:
@@ -256,7 +271,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
      "gptq_marlin_24_gemm(Tensor a, Tensor b_q_weight, Tensor b_meta, "
      "Tensor b_scales, Tensor workspace, "
      "int b_q_type, "
-      "SymInt size_m, SymInt size_n, SymInt size_k) -> Tensor");
+      "SymInt size_m, SymInt size_n, SymInt size_k) -> Tensor",
+      {stride_tag});
  //  conditionally compiled so impl in source file

  // Machete (Dense) Optimized Mixed Precision GEMM for Hopper.
@@ -282,7 +298,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
      "   Tensor? channel_scales,"
      "   Tensor? token_scales,"
      "   str?    schedule"
-      ") -> Tensor");
+      ") -> Tensor",
+      {stride_tag});
  ops.def(
      "machete_prepack_B("
      "   Tensor B,"
@@ -298,30 +315,24 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  // gptq_marlin Optimized Quantized GEMM for GPTQ.
  ops.def(
      "gptq_marlin_gemm(Tensor a, Tensor? c_or_none, Tensor b_q_weight, "
-      "Tensor? b_bias_or_none,Tensor b_scales, "
-      "Tensor? a_scales, Tensor? global_scale, Tensor? b_zeros_or_none, "
-      "Tensor? "
-      "g_idx_or_none, Tensor? perm_or_none, Tensor workspace, int b_type_id, "
+      "Tensor? b_bias_or_none,"
+      "Tensor b_scales, Tensor? global_scale, Tensor? b_zeros_or_none, Tensor? "
+      "g_idx_or_none, Tensor? perm_or_none, Tensor workspace, int b_q_type, "
      "SymInt size_m, SymInt size_n, SymInt size_k, bool is_k_full, "
-      "bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) -> Tensor");
+      "bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) -> Tensor",
+      {stride_tag});
  // conditionally compiled so impl registration is in source file

  // gptq_marlin repack from GPTQ.
  ops.def(
      "gptq_marlin_repack(Tensor b_q_weight, Tensor perm, "
-      "SymInt size_k, SymInt size_n, int num_bits, bool is_a_8bit) -> Tensor");
+      "SymInt size_k, SymInt size_n, int num_bits) -> Tensor");
  // conditionally compiled so impl registrations are in source file

  // awq_marlin repack from AWQ.
  ops.def(
      "awq_marlin_repack(Tensor b_q_weight, SymInt size_k, "
-      "SymInt size_n, int num_bits, bool is_a_8bit) -> Tensor");
-  // conditionally compiled so impl registrations are in source file
-
-  // preprocess W-int4A-fp8 weight for marlin kernel
-  ops.def(
-      "marlin_int4_fp8_preprocess(Tensor qweight, "
-      "Tensor? qzeros_or_none, bool inplace) -> Tensor");
+      "SymInt size_n, int num_bits) -> Tensor");
  // conditionally compiled so impl registrations are in source file

  // CUTLASS w4a8 GEMM
@@ -335,7 +346,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
      "   Tensor token_scales,"
      "   ScalarType? out_type,"
      "   str?   maybe_schedule"
-      ") -> Tensor");
+      ") -> Tensor",
+      {stride_tag});
  // pack scales
  ops.def("cutlass_pack_scale_fp8(Tensor scales) -> Tensor");
  // encode and reorder weight matrix
@@ -382,21 +394,24 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  ops.def(
      "cutlass_scaled_fp4_mm(Tensor! out, Tensor a, Tensor b,"
      "                      Tensor block_scale_a, Tensor block_scale_b,"
-      "                      Tensor alpha) -> ()");
+      "                      Tensor alpha) -> ()",
+      {stride_tag});
  ops.impl("cutlass_scaled_fp4_mm", torch::kCUDA, &cutlass_scaled_fp4_mm);

  // cutlass blockwise scaledgroup GEMM
  ops.def(
      "cutlass_blockwise_scaled_grouped_mm(Tensor! output, Tensor a, Tensor b, "
      "Tensor scales_a, Tensor scales_b, "
-      "Tensor problem_sizes, Tensor expert_offsets) -> ()");
+      "Tensor problem_sizes, Tensor expert_offsets) -> ()",
+      {stride_tag});
  // conditionally compiled so impl registration is in source file

  // cutlass nvfp4 block scaled group GEMM
  ops.def(
      "cutlass_fp4_group_mm(Tensor! out, Tensor a, Tensor b,"
      " Tensor a_blockscale, Tensor b_blockscales, Tensor alphas,"
-      " Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()");
+      " Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()",
+      {stride_tag});
  // conditionally compiled so impl registration is in source file

  // CUTLASS w8a8 GEMM, supporting symmetric per-tensor or per-row/column
@@ -404,7 +419,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  ops.def(
      "cutlass_scaled_mm(Tensor! out, Tensor a,"
      "                  Tensor b, Tensor a_scales,"
-      "                  Tensor b_scales, Tensor? bias) -> ()");
+      "                  Tensor b_scales, Tensor? bias) -> ()",
+      {stride_tag});
  ops.impl("cutlass_scaled_mm", torch::kCUDA, &cutlass_scaled_mm);

  // CUTLASS w8a8 GEMM, supporting asymmetric per-tensor or per-row/column
@@ -413,7 +429,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
      "cutlass_scaled_mm_azp(Tensor! out, Tensor a,"
      "                  Tensor b, Tensor a_scales,"
      "                  Tensor b_scales, Tensor azp_adj,"
-      "                  Tensor? azp, Tensor? bias) -> ()");
+      "                  Tensor? azp, Tensor? bias) -> ()",
+      {stride_tag});
  ops.impl("cutlass_scaled_mm_azp", torch::kCUDA, &cutlass_scaled_mm_azp);

  // Check if cutlass scaled_mm is supported for CUDA devices of the given
@@ -432,7 +449,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
      "               Tensor a_scales, Tensor b_scales, Tensor expert_offsets, "
      "               Tensor problem_sizes, Tensor a_strides, "
      "               Tensor b_strides, Tensor c_strides, bool per_act_token, "
-      "               bool per_out_ch) -> ()");
+      "               bool per_out_ch) -> ()",
+      {stride_tag});
  ops.impl("cutlass_moe_mm", torch::kCUDA, &cutlass_moe_mm);

  // A function that computes data required to run fused MoE with w8a8 grouped
@@ -446,8 +464,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
      "                        Tensor! problem_sizes1, Tensor! problem_sizes2, "
      "                        Tensor! input_permutation, "
      "                        Tensor! output_permutation, int num_experts, "
-      "                        int n, int k, Tensor? blockscale_offsets) -> "
-      "()");
+      "                        int n, int k, Tensor? blockscale_offsets) -> ()",
+      {stride_tag});
  ops.impl("get_cutlass_moe_mm_data", torch::kCUDA, &get_cutlass_moe_mm_data);

  // A function that computes problem sizes for each expert's multiplication
@@ -458,7 +476,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
      "                                 Tensor! problem_sizes1, "
      "                                 Tensor! problem_sizes2, "
      "                                 int num_experts, int n, int k, "
-      "                                 Tensor? blockscale_offsets) -> ()");
+      "                                 Tensor? blockscale_offsets) -> ()",
+      {stride_tag});
  ops.impl("get_cutlass_moe_mm_problem_sizes", torch::kCUDA,
           &get_cutlass_moe_mm_problem_sizes);

@@ -473,7 +492,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
      "                             Tensor! problem_sizes2, "
      "                             Tensor expert_num_tokens, "
      "                             int num_local_experts, int padded_m, "
-      "                             int n, int k) -> ()");
+      "                             int n, int k) -> ()",
+      {stride_tag});
  ops.impl("get_cutlass_pplx_moe_mm_data", torch::kCUDA,
           &get_cutlass_pplx_moe_mm_data);

@@ -497,7 +517,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
      "cutlass_scaled_sparse_mm(Tensor! out, Tensor a,"
      "                         Tensor bt_nzs,"
      "                         Tensor bt_meta, Tensor a_scales,"
-      "                         Tensor b_scales, Tensor? bias) -> ()");
+      "                         Tensor b_scales, Tensor? bias) -> ()",
+      {stride_tag});
  ops.impl("cutlass_scaled_sparse_mm", torch::kCUDA, &cutlass_scaled_sparse_mm);

  // CUTLASS sparse matrix compressor
@@ -546,7 +567,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
      "gptq_gemm(Tensor a, Tensor b_q_weight, Tensor b_gptq_qzeros, "
      "Tensor b_gptq_scales, Tensor b_g_idx, bool use_exllama, bool "
      "use_v2_format, int bit) "
-      "-> Tensor");
+      "-> Tensor",
+      {stride_tag});
  ops.impl("gptq_gemm", torch::kCUDA, &gptq_gemm);

  // Post processing for GPTQ.
@@ -701,8 +723,7 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cache_ops), cache_ops) {
  cache_ops.def(
      "gather_and_maybe_dequant_cache(Tensor src_cache, Tensor! dst, "
      "                               Tensor block_table, Tensor cu_seq_lens, "
-      "                               Tensor token_to_seq, "
-      "                               int num_tokens, "
+      "                               int batch_size, "
      "                               str kv_cache_dtype, "
      "                               Tensor scale, Tensor? seq_starts) -> ()");
  cache_ops.impl("gather_and_maybe_dequant_cache", torch::kCUDA,
@@ -20,8 +20,8 @@ ARG PYTHON_VERSION=3.12
 # glibc version is baked into the distro, and binaries built with one glibc
 # version are not backwards compatible with OSes that use an earlier version.
 ARG BUILD_BASE_IMAGE=nvidia/cuda:${CUDA_VERSION}-devel-ubuntu20.04
-# Using cuda base image with minimal dependencies necessary for JIT compilation (FlashInfer, DeepGEMM, EP kernels)
-ARG FINAL_BASE_IMAGE=nvidia/cuda:${CUDA_VERSION}-base-ubuntu22.04
+# TODO: Restore to base image after FlashInfer AOT wheel fixed
+ARG FINAL_BASE_IMAGE=nvidia/cuda:${CUDA_VERSION}-devel-ubuntu22.04

 # By parameterizing the Deadsnakes repository URL, we allow third-party to use
 # their own mirror. When doing so, we don't benefit from the transparent
@@ -56,6 +56,7 @@ ARG UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}

 # PyTorch provides its own indexes for standard and nightly builds
 ARG PYTORCH_CUDA_INDEX_BASE_URL=https://download.pytorch.org/whl
+ARG PYTORCH_CUDA_NIGHTLY_INDEX_BASE_URL=https://download.pytorch.org/whl/nightly

 # PIP supports multiple authentication schemes, including keyring
 # By parameterizing the PIP_KEYRING_PROVIDER variable and setting it to
@@ -85,20 +86,7 @@ ARG GET_PIP_URL
 RUN echo 'tzdata tzdata/Areas select America' | debconf-set-selections \
    && echo 'tzdata tzdata/Zones/America select Los_Angeles' | debconf-set-selections \
    && apt-get update -y \
-    && apt-get install -y --no-install-recommends \
-        ccache \
-        software-properties-common \
-        git \
-        curl \
-        sudo \
-        python3-pip \
-        libibverbs-dev \
-        # Upgrade to GCC 10 to avoid https://gcc.gnu.org/bugzilla/show_bug.cgi?id=92519
-        # as it was causing spam when compiling the CUTLASS kernels
-        gcc-10 \
-        g++-10 \
-    && update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-10 110 --slave /usr/bin/g++ g++ /usr/bin/g++-10 \
-    && rm -rf /var/lib/apt/lists/* \
+    && apt-get install -y ccache software-properties-common git curl sudo python3-pip \
    && curl -LsSf https://astral.sh/uv/install.sh | sh \
    && $HOME/.local/bin/uv venv /opt/venv --python ${PYTHON_VERSION} \
    && rm -f /usr/bin/python3 /usr/bin/python3-config /usr/bin/pip \
@@ -110,6 +98,7 @@ RUN echo 'tzdata tzdata/Areas select America' | debconf-set-selections \
 ARG PIP_INDEX_URL UV_INDEX_URL
 ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
 ARG PYTORCH_CUDA_INDEX_BASE_URL
+ARG PYTORCH_CUDA_NIGHTLY_INDEX_BASE_URL
 ARG PIP_KEYRING_PROVIDER UV_KEYRING_PROVIDER

 # Activate virtual environment and add uv to PATH
@@ -123,6 +112,10 @@ ENV UV_INDEX_STRATEGY="unsafe-best-match"
 # Use copy mode to avoid hardlink failures with Docker cache mounts
 ENV UV_LINK_MODE=copy

+# Upgrade to GCC 10 to avoid https://gcc.gnu.org/bugzilla/show_bug.cgi?id=92519
+# as it was causing spam when compiling the CUTLASS kernels
+RUN apt-get install -y gcc-10 g++-10
+RUN update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-10 110 --slave /usr/bin/g++ g++ /usr/bin/g++-10
 RUN <<EOF
 gcc --version
 EOF
@@ -233,28 +226,6 @@ RUN --mount=type=cache,target=/root/.cache/ccache \
        python3 setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38; \
    fi

-# Install DeepGEMM from source
-ARG DEEPGEMM_GIT_REF
-COPY tools/install_deepgemm.sh /tmp/install_deepgemm.sh
-RUN --mount=type=cache,target=/root/.cache/uv \
-    VLLM_DOCKER_BUILD_CONTEXT=1 TORCH_CUDA_ARCH_LIST="9.0a 10.0a" /tmp/install_deepgemm.sh --cuda-version "${CUDA_VERSION}" ${DEEPGEMM_GIT_REF:+--ref "$DEEPGEMM_GIT_REF"} --wheel-dir /tmp/deepgemm/dist
-
-# Ensure the wheel dir exists so later-stage COPY won't fail when DeepGEMM is skipped
-RUN mkdir -p /tmp/deepgemm/dist && touch /tmp/deepgemm/dist/.deepgemm_skipped
-
-COPY tools/ep_kernels/install_python_libraries.sh /tmp/install_python_libraries.sh
-# Install EP kernels(pplx-kernels and DeepEP)
-ARG PPLX_COMMIT_HASH
-ARG DEEPEP_COMMIT_HASH
-RUN --mount=type=cache,target=/root/.cache/uv \
-    export TORCH_CUDA_ARCH_LIST='9.0a 10.0a' && \
-    /tmp/install_python_libraries.sh \
-        --workspace /tmp/ep_kernels_workspace \
-        --mode wheel \
-        ${PPLX_COMMIT_HASH:+--pplx-ref "$PPLX_COMMIT_HASH"} \
-        ${DEEPEP_COMMIT_HASH:+--deepep-ref "$DEEPEP_COMMIT_HASH"} && \
-    find /tmp/ep_kernels_workspace/nvshmem -name '*.a' -delete
-
 # Check the size of the wheel if RUN_WHEEL_CHECK is true
 COPY .buildkite/check-wheel-size.py check-wheel-size.py
 # sync the default value with .buildkite/check-wheel-size.py
@@ -283,7 +254,7 @@ ENV UV_INDEX_STRATEGY="unsafe-best-match"
 ENV UV_LINK_MODE=copy

 # Install libnuma-dev, required by fastsafetensors (fixes #20384)
-RUN apt-get update && apt-get install -y --no-install-recommends libnuma-dev && rm -rf /var/lib/apt/lists/*
+RUN apt-get update && apt-get install -y libnuma-dev && rm -rf /var/lib/apt/lists/*
 COPY requirements/lint.txt requirements/lint.txt
 COPY requirements/test.txt requirements/test.txt
 COPY requirements/dev.txt requirements/dev.txt
@@ -320,15 +291,8 @@ RUN PYTHON_VERSION_STR=$(echo ${PYTHON_VERSION} | sed 's/\.//g') && \
 RUN echo 'tzdata tzdata/Areas select America' | debconf-set-selections \
    && echo 'tzdata tzdata/Zones/America select Los_Angeles' | debconf-set-selections \
    && apt-get update -y \
-    && apt-get install -y --no-install-recommends \
-        software-properties-common \
-        curl \
-        sudo \
-        python3-pip \
-        ffmpeg \
-        libsm6 \
-        libxext6 \
-        libgl1 \
+    && apt-get install -y ccache software-properties-common git curl wget sudo vim python3-pip \
+    && apt-get install -y ffmpeg libsm6 libxext6 libgl1 \
    && if [ ! -z ${DEADSNAKES_MIRROR_URL} ] ; then \
        if [ ! -z "${DEADSNAKES_GPGKEY_URL}" ] ; then \
            mkdir -p -m 0755 /etc/apt/keyrings ; \
@@ -343,33 +307,17 @@ RUN echo 'tzdata tzdata/Areas select America' | debconf-set-selections \
        done ; \
    fi \
    && apt-get update -y \
-    && apt-get install -y --no-install-recommends \
-        python${PYTHON_VERSION} \
-        python${PYTHON_VERSION}-dev \
-        python${PYTHON_VERSION}-venv \
-        libibverbs-dev \
-    && rm -rf /var/lib/apt/lists/* \
+    && apt-get install -y python${PYTHON_VERSION} python${PYTHON_VERSION}-dev python${PYTHON_VERSION}-venv libibverbs-dev \
    && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 \
    && update-alternatives --set python3 /usr/bin/python${PYTHON_VERSION} \
    && ln -sf /usr/bin/python${PYTHON_VERSION}-config /usr/bin/python3-config \
    && curl -sS ${GET_PIP_URL} | python${PYTHON_VERSION} \
    && python3 --version && python3 -m pip --version

-# Install CUDA development tools and build essentials for runtime JIT compilation
-# (FlashInfer, DeepGEMM, EP kernels all require compilation at runtime)
-RUN CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-') && \
-    apt-get update -y && \
-    apt-get install -y --no-install-recommends \
-    cuda-nvcc-${CUDA_VERSION_DASH} \
-    cuda-cudart-${CUDA_VERSION_DASH} \
-    cuda-nvrtc-${CUDA_VERSION_DASH} \
-    cuda-cuobjdump-${CUDA_VERSION_DASH} \
-    libcublas-${CUDA_VERSION_DASH} && \
-    rm -rf /var/lib/apt/lists/*
-
 ARG PIP_INDEX_URL UV_INDEX_URL
 ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
 ARG PYTORCH_CUDA_INDEX_BASE_URL
+ARG PYTORCH_CUDA_NIGHTLY_INDEX_BASE_URL
 ARG PIP_KEYRING_PROVIDER UV_KEYRING_PROVIDER

 # Install uv for faster pip installs
@@ -389,6 +337,20 @@ ENV UV_LINK_MODE=copy
 # or future versions of triton.
 RUN ldconfig /usr/local/cuda-$(echo $CUDA_VERSION | cut -d. -f1,2)/compat/

+# arm64 (GH200) build follows the practice of "use existing pytorch" build,
+# we need to install torch and torchvision from the nightly builds first,
+# pytorch will not appear as a vLLM dependency in all of the following steps
+# after this step
+RUN --mount=type=cache,target=/root/.cache/uv \
+    if [ "$TARGETPLATFORM" = "linux/arm64" ]; then \
+        uv pip install --system \
+            --index-url ${PYTORCH_CUDA_NIGHTLY_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
+            "torch==2.8.0.dev20250318+cu128" "torchvision==0.22.0.dev20250319" ; \
+        uv pip install --system \
+            --index-url ${PYTORCH_CUDA_NIGHTLY_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
+            --pre pytorch_triton==3.3.0+gitab727c40 ; \
+    fi
+
 # Install vllm wheel first, so that torch etc will be installed.
 RUN --mount=type=bind,from=build,src=/workspace/dist,target=/vllm-workspace/dist \
    --mount=type=cache,target=/root/.cache/uv \
@@ -398,8 +360,8 @@ RUN --mount=type=bind,from=build,src=/workspace/dist,target=/vllm-workspace/dist
 # Install FlashInfer pre-compiled kernel cache and binaries
 # https://docs.flashinfer.ai/installation.html
 RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --system flashinfer-cubin==0.5.3 \
-    && uv pip install --system flashinfer-jit-cache==0.5.3 \
+    uv pip install --system flashinfer-cubin==0.5.2 \
+    && uv pip install --system flashinfer-jit-cache==0.5.2 \
        --extra-index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
    && flashinfer show-config

@@ -411,32 +373,36 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 . /etc/environment && \
 uv pip list

-# Install deepgemm wheel that has been built in the `build` stage
+# Even when we build Flashinfer with AOT mode, there's still
+# some issues w.r.t. JIT compilation. Therefore we need to
+# install build dependencies for JIT compilation.
+# TODO: Remove this once FlashInfer AOT wheel is fixed
+COPY requirements/build.txt requirements/build.txt
 RUN --mount=type=cache,target=/root/.cache/uv \
-    --mount=type=bind,from=build,source=/tmp/deepgemm/dist,target=/tmp/deepgemm/dist,ro \
-    sh -c 'if ls /tmp/deepgemm/dist/*.whl >/dev/null 2>&1; then \
-              uv pip install --system /tmp/deepgemm/dist/*.whl; \
-           else \
-              echo "No DeepGEMM wheels to install; skipping."; \
-           fi'
-
-# Pytorch now installs NVSHMEM, setting LD_LIBRARY_PATH (https://github.com/pytorch/pytorch/blob/d38164a545b4a4e4e0cf73ce67173f70574890b6/.ci/manywheel/build_cuda.sh#L141C14-L141C36)
-ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH
-
-# Install EP kernels wheels (pplx-kernels and DeepEP) that have been built in the `build` stage
-RUN --mount=type=bind,from=build,src=/tmp/ep_kernels_workspace/dist,target=/vllm-workspace/ep_kernels/dist \
-    --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --system ep_kernels/dist/*.whl --verbose \
+    uv pip install --system -r requirements/build.txt \
        --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')

-RUN --mount=type=bind,source=tools/install_gdrcopy.sh,target=/tmp/install_gdrcopy.sh,ro \
-    set -eux; \
+# Install DeepGEMM from source
+ARG DEEPGEMM_GIT_REF
+COPY tools/install_deepgemm.sh /tmp/install_deepgemm.sh
+RUN --mount=type=cache,target=/root/.cache/uv \
+    VLLM_DOCKER_BUILD_CONTEXT=1 TORCH_CUDA_ARCH_LIST="9.0a 10.0a" /tmp/install_deepgemm.sh --cuda-version "${CUDA_VERSION}" ${DEEPGEMM_GIT_REF:+--ref "$DEEPGEMM_GIT_REF"}
+
+COPY tools/install_gdrcopy.sh install_gdrcopy.sh
+RUN set -eux; \
    case "${TARGETPLATFORM}" in \
      linux/arm64) UUARCH="aarch64" ;; \
      linux/amd64) UUARCH="x64" ;; \
      *) echo "Unsupported TARGETPLATFORM: ${TARGETPLATFORM}" >&2; exit 1 ;; \
    esac; \
-    /tmp/install_gdrcopy.sh "${GDRCOPY_OS_VERSION}" "${GDRCOPY_CUDA_VERSION}" "${UUARCH}"
+    ./install_gdrcopy.sh "${GDRCOPY_OS_VERSION}" "${GDRCOPY_CUDA_VERSION}" "${UUARCH}"; \
+    rm ./install_gdrcopy.sh
+
+# Install EP kernels(pplx-kernels and DeepEP)
+COPY tools/ep_kernels/install_python_libraries.sh install_python_libraries.sh
+ENV CUDA_HOME=/usr/local/cuda
+RUN export TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST:-9.0a 10.0a+PTX}" \
+    && bash install_python_libraries.sh

 # CUDA image changed from /usr/local/nvidia to /usr/local/cuda in 12.8 but will
 # return to /usr/local/nvidia in 13.0 to allow container providers to mount drivers
@@ -466,11 +432,6 @@ ENV UV_INDEX_STRATEGY="unsafe-best-match"
 # Use copy mode to avoid hardlink failures with Docker cache mounts
 ENV UV_LINK_MODE=copy

-RUN echo 'tzdata tzdata/Areas select America' | debconf-set-selections \
-    && echo 'tzdata tzdata/Zones/America select Los_Angeles' | debconf-set-selections \
-    && apt-get update -y \
-    && apt-get install -y git
-
 # install development dependencies (for testing)
 RUN --mount=type=cache,target=/root/.cache/uv \
    CUDA_MAJOR="${CUDA_VERSION%%.*}"; \
@@ -511,11 +472,12 @@ ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
 # Reference: https://github.com/astral-sh/uv/pull/1694
 ENV UV_HTTP_TIMEOUT=500

+COPY requirements/kv_connectors.txt requirements/kv_connectors.txt
+
 # install additional dependencies for openai api server
 RUN --mount=type=cache,target=/root/.cache/uv \
-    --mount=type=bind,source=requirements/kv_connectors.txt,target=/tmp/kv_connectors.txt,ro \
    if [ "$INSTALL_KV_CONNECTORS" = "true" ]; then \
-        uv pip install --system -r /tmp/kv_connectors.txt; \
+        uv pip install --system -r requirements/kv_connectors.txt; \
    fi; \
    if [ "$TARGETPLATFORM" = "linux/arm64" ]; then \
        BITSANDBYTES_VERSION="0.42.0"; \
@@ -37,7 +37,6 @@ RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
    && update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-12 10 --slave /usr/bin/g++ g++ /usr/bin/g++-12 \
    && curl -LsSf https://astral.sh/uv/install.sh | sh

-ENV CC=/usr/bin/gcc-12 CXX=/usr/bin/g++-12
 ENV CCACHE_DIR=/root/.cache/ccache
 ENV CMAKE_CXX_COMPILER_LAUNCHER=ccache

@@ -123,16 +122,7 @@ WORKDIR /workspace/vllm
 RUN --mount=type=bind,src=requirements/test.in,target=requirements/test.in \
    cp requirements/test.in requirements/cpu-test.in && \
    sed -i '/mamba_ssm/d' requirements/cpu-test.in && \
-    remove_packages_not_supported_on_aarch64() { \
-      case "$(uname -m)" in \
-        aarch64|arm64) \
-          sed -i '/decord/d' requirements/cpu-test.in; \
-          sed -i '/terratorch/d' requirements/cpu-test.in; \
-          ;; \
-      esac; \
-    }; \
-    remove_packages_not_supported_on_aarch64 && \
-    sed -i 's/^torch==.*/torch==2.9.1/g' requirements/cpu-test.in && \
+    sed -i 's/^torch==.*/torch==2.8.0/g' requirements/cpu-test.in && \
    sed -i 's/torchaudio.*/torchaudio/g' requirements/cpu-test.in && \
    sed -i 's/torchvision.*/torchvision/g' requirements/cpu-test.in && \
    uv pip compile requirements/cpu-test.in -o requirements/cpu-test.txt --index-strategy unsafe-best-match --torch-backend cpu
@@ -76,6 +76,34 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 RUN --mount=type=cache,target=/root/.cache/uv \
 uv pip install --system -r requirements/common.txt

+# must put before installing xformers, so it can install the correct version of xfomrers.
+ARG torch_cuda_arch_list='8.0;8.6;8.9;9.0'
+ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list}
+
+# Build xformers with cuda and torch nightly
+# following official xformers guidance: https://github.com/facebookresearch/xformers#build
+# todo(elainewy): cache xformers build result for faster build
+ARG max_jobs=16
+ENV MAX_JOBS=${max_jobs}
+ARG XFORMERS_COMMIT=f2de641ef670510cadab099ce6954031f52f191c
+
+ENV CCACHE_DIR=/root/.cache/ccache
+RUN --mount=type=cache,target=/root/.cache/ccache \
+     --mount=type=cache,target=/root/.cache/uv \
+    echo 'git clone xformers...' \
+    && git clone https://github.com/facebookresearch/xformers.git --recursive \
+    && cd xformers \
+    && git checkout ${XFORMERS_COMMIT} \
+    && git submodule update --init --recursive \
+    && echo 'finish git clone xformers...' \
+    && rm -rf build \
+    && python3 setup.py bdist_wheel --dist-dir=../xformers-dist --verbose \
+    && cd .. \
+    && rm -rf xformers
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install --system xformers-dist/*.whl --verbose
+
 # build can take a long time, and the torch nightly version fetched from url can be different in next docker stage.
 # track the nightly torch version used in the build, when we set up runtime environment we can make sure the version is the same
 RUN uv pip freeze | grep -i '^torch\|^torchvision\|^torchaudio' > torch_build_versions.txt
@@ -205,6 +233,11 @@ RUN --mount=type=bind,from=build,src=/workspace/dist,target=/vllm-workspace/vllm
    --mount=type=cache,target=/root/.cache/uv \
    uv pip install --system vllm-dist/*.whl --verbose

+# install xformers again for the new environment
+RUN --mount=type=bind,from=base,src=/workspace/xformers-dist,target=/vllm-workspace/xformers-dist \
+    --mount=type=cache,target=/root/.cache/uv \
+    uv pip install --system /vllm-workspace/xformers-dist/*.whl --verbose
+
 ARG torch_cuda_arch_list='8.0;8.6;8.9;9.0'

 # install package for build flashinfer
@@ -274,7 +307,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
    uv pip install --system -r requirements/nightly_torch_test.txt

 # Logging to confirm the torch versions
-RUN pip freeze | grep -E 'torch|vllm|flashinfer'
+RUN pip freeze | grep -E 'torch|xformers|vllm|flashinfer'

 # Logging to confirm all the packages are installed
 RUN pip freeze
@@ -8,8 +8,8 @@ FROM registry.access.redhat.com/ubi9/ubi-minimal:${BASE_UBI_IMAGE_TAG} AS openbl

 ARG MAX_JOBS
 ARG OPENBLAS_VERSION=0.3.30
-RUN microdnf install -y dnf && dnf install -y gcc-toolset-14 make wget unzip \
-    && source /opt/rh/gcc-toolset-14/enable \
+RUN microdnf install -y dnf && dnf install -y gcc-toolset-13 make wget unzip \
+    && source /opt/rh/gcc-toolset-13/enable \
    && wget https://github.com/OpenMathLib/OpenBLAS/releases/download/v$OPENBLAS_VERSION/OpenBLAS-$OPENBLAS_VERSION.zip \
    && unzip OpenBLAS-$OPENBLAS_VERSION.zip \
    && cd OpenBLAS-$OPENBLAS_VERSION \
@@ -57,7 +57,7 @@ COPY --from=openblas-builder /tmp/control /dev/null
 RUN --mount=type=bind,from=openblas-builder,source=/OpenBLAS-$OPENBLAS_VERSION/,target=/openblas/,rw \
    dnf install -y openssl-devel \
    && dnf install -y \
-       git tar gcc-toolset-14 automake libtool \
+       git tar gcc-toolset-13 automake libtool \
       pkgconfig xsimd zeromq-devel kmod findutils protobuf* \
       libtiff-devel libjpeg-devel zlib-devel freetype-devel libwebp-devel \
       harfbuzz-devel libraqm-devel libimagequant-devel libxcb-devel \
@@ -84,7 +84,7 @@ ARG _GLIBCXX_USE_CXX11_ABI=1
 ARG OPENBLAS_VERSION=0.3.30

 RUN --mount=type=cache,target=/root/.cache/uv \
-    source /opt/rh/gcc-toolset-14/enable &&  \
+    source /opt/rh/gcc-toolset-13/enable &&  \
    git clone --recursive https://github.com/pytorch/pytorch.git -b v${TORCH_VERSION} && \
    cd pytorch && \
    uv pip install -r requirements.txt && \
@@ -97,7 +97,7 @@ ARG TORCHVISION_VERSION=0.22.0
 ARG TORCHVISION_USE_NVJPEG=0
 ARG TORCHVISION_USE_FFMPEG=0
 RUN --mount=type=cache,target=/root/.cache/uv \
-    source /opt/rh/gcc-toolset-14/enable && \
+    source /opt/rh/gcc-toolset-13/enable && \
    git clone --recursive https://github.com/pytorch/vision.git -b v${TORCHVISION_VERSION} && \
    cd vision && \
    MAX_JOBS=${MAX_JOBS:-$(nproc)} \
@@ -113,7 +113,7 @@ ARG USE_ROCM=0
 ARG USE_CUDA=0
 ARG TORCHAUDIO_TEST_ALLOW_SKIP_IF_NO_FFMPEG=1
 RUN --mount=type=cache,target=/root/.cache/uv \
-    source /opt/rh/gcc-toolset-14/enable && \
+    source /opt/rh/gcc-toolset-13/enable && \
    git clone --recursive https://github.com/pytorch/audio.git -b v${TORCHAUDIO_VERSION} && \
    cd audio && \
    MAX_JOBS=${MAX_JOBS:-$(nproc)} \
@@ -130,7 +130,7 @@ ARG MAX_JOBS
 ARG PYARROW_PARALLEL
 ARG PYARROW_VERSION=21.0.0
 RUN --mount=type=cache,target=/root/.cache/uv \
-    source /opt/rh/gcc-toolset-14/enable && \
+    source /opt/rh/gcc-toolset-13/enable && \
    git clone --recursive https://github.com/apache/arrow.git -b apache-arrow-${PYARROW_VERSION} && \
    cd arrow/cpp && \
    mkdir build && cd build && \
@@ -162,7 +162,7 @@ ARG OPENCV_VERSION=86
 ARG OPENCV_PATCH=97f3f39
 ARG ENABLE_HEADLESS=1
 RUN --mount=type=cache,target=/root/.cache/uv \
-    source /opt/rh/gcc-toolset-14/enable && \
+    source /opt/rh/gcc-toolset-13/enable && \
    git clone --recursive https://github.com/opencv/opencv-python.git -b ${OPENCV_VERSION} && \
    cd opencv-python && \
    sed -i -E -e 's/"setuptools.+",/"setuptools",/g' pyproject.toml && \
@@ -196,7 +196,7 @@ ARG MAX_JOBS
 ARG NUMBA_VERSION=0.61.2

 # Clone all required dependencies
-RUN dnf install ninja-build llvm15 llvm15-devel -y && source /opt/rh/gcc-toolset-14/enable && export PATH=$PATH:/usr/lib64/llvm15/bin && \
+RUN dnf install ninja-build llvm15 llvm15-devel -y && source /opt/rh/gcc-toolset-13/enable && export PATH=$PATH:/usr/lib64/llvm15/bin && \
    git clone --recursive https://github.com/numba/numba.git -b ${NUMBA_VERSION} && \
    cd ./numba && \
    if ! grep '#include "dynamic_annotations.h"' numba/_dispatcher.cpp; then \
@@ -211,9 +211,6 @@ RUN dnf install ninja-build llvm15 llvm15-devel -y && source /opt/rh/gcc-toolset

 FROM base-builder AS vllmcache-builder

-ENV LLVM_CONFIG=/usr/lib64/llvm15/bin/llvm-config
-ENV PATH=/usr/lib64/llvm15/bin:$PATH
-
 COPY --from=torch-builder /tmp/control /dev/null
 COPY --from=arrow-builder /tmp/control /dev/null
 COPY --from=cv-builder /tmp/control /dev/null
@@ -228,13 +225,10 @@ ARG GRPC_PYTHON_BUILD_SYSTEM_OPENSSL=1
 RUN --mount=type=cache,target=/root/.cache/uv \
    dnf install llvm15 llvm15-devel -y && \
    rpm -ivh --nodeps https://mirror.stream.centos.org/9-stream/CRB/ppc64le/os/Packages/protobuf-lite-devel-3.14.0-16.el9.ppc64le.rpm && \
-    source /opt/rh/gcc-toolset-14/enable && \
+    source /opt/rh/gcc-toolset-13/enable && \
    git clone https://github.com/huggingface/xet-core.git && cd xet-core/hf_xet/ && \
    uv pip install maturin && \
    uv build --wheel --out-dir /hf_wheels/
-
-ENV CXXFLAGS="-fno-lto -Wno-error=free-nonheap-object" \
-    CFLAGS="-fno-lto"
 RUN --mount=type=cache,target=/root/.cache/uv \
    --mount=type=bind,from=torch-builder,source=/torchwheels/,target=/torchwheels/,ro \
    --mount=type=bind,from=arrow-builder,source=/arrowwheels/,target=/arrowwheels/,ro \
@@ -242,7 +236,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
    --mount=type=bind,from=numa-builder,source=/numactl/,target=/numactl/,rw \
    --mount=type=bind,from=numba-builder,source=/numbawheels/,target=/numbawheels/,ro \
    --mount=type=bind,src=.,dst=/src/,rw \
-    source /opt/rh/gcc-toolset-14/enable && \
+    source /opt/rh/gcc-toolset-13/enable && \
    export PATH=$PATH:/usr/lib64/llvm15/bin && \
    uv pip install /opencvwheels/*.whl /arrowwheels/*.whl /torchwheels/*.whl /numbawheels/*.whl && \
    sed -i -e 's/.*torch.*//g' /src/pyproject.toml /src/requirements/*.txt && \
@@ -266,7 +260,7 @@ FROM base-builder AS lapack-builder
 ARG MAX_JOBS
 ARG LAPACK_VERSION=3.12.1
 RUN git clone --recursive https://github.com/Reference-LAPACK/lapack.git -b v${LAPACK_VERSION} \
-    && cd lapack && source /opt/rh/gcc-toolset-14/enable \
+    && cd lapack && source /opt/rh/gcc-toolset-13/enable \
    && cmake -B build -S . \
    && cmake --build build -j ${MAX_JOBS:-$(nproc)}

@@ -305,7 +299,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
    --mount=type=bind,from=openblas-builder,source=/OpenBLAS-$OPENBLAS_VERSION/,target=/openblas/,rw \
    rpm -ivh https://dl.fedoraproject.org/pub/epel/epel-release-latest-9.noarch.rpm && \
    microdnf install --nodocs -y \
-    libomp libicu tar findutils openssl llvm15 llvm15-devel \
+    libomp tar findutils openssl llvm15 llvm15-devel \
    pkgconfig xsimd g++ gcc-fortran libsndfile \
    libtiff libjpeg openjpeg2 zlib zeromq \
    freetype lcms2 libwebp tcl tk utf8proc \
@@ -7,8 +7,6 @@ FROM ${BASE_IMAGE} AS base

 ARG ARG_PYTORCH_ROCM_ARCH
 ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}}
-ENV RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES=1
-ENV RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES=1

 # Install some basic utilities
 RUN apt-get update -q -y && apt-get install -q -y \
@@ -65,8 +63,6 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tests /tests
 COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/examples /examples
 COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/
 COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite
-# Centralized v1 package - copied to both test and final stages
-COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1

 # -----------------------
 # Test vLLM image
@@ -90,22 +86,10 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm /vllm-workspace

 # install development dependencies (for testing)
 RUN cd /vllm-workspace \
+    && rm -rf vllm \
    && python3 -m pip install -e tests/vllm_test_utils \
    && python3 -m pip install pytest-shard

-# enable fast downloads from hf (for testing)
-RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --system hf_transfer
-ENV HF_HUB_ENABLE_HF_TRANSFER=1
-
-# Copy in the v1 package
-COPY --from=export_vllm /vllm_v1 /usr/local/lib/python${PYTHON_VERSION}/dist-packages/vllm/v1
-
-# Source code is used in the `python_only_compile.sh` test
-# We hide it inside `src/` so that this source code
-# will not be imported by other tests
-RUN mkdir src && mv vllm src/vllm
-
 # -----------------------
 # Final vLLM image
 FROM base AS final
@@ -130,9 +114,6 @@ RUN --mount=type=bind,from=export_vllm,src=/,target=/install \
    && pip uninstall -y vllm \
    && uv pip install --system *.whl

-# Copy in the v1 package
-COPY --from=export_vllm /vllm_v1 /usr/local/lib/python${PYTHON_VERSION}/dist-packages/vllm/v1
-
 ARG COMMON_WORKDIR

 # Copy over the benchmark scripts as well
@@ -140,6 +121,8 @@ COPY --from=export_vllm /benchmarks ${COMMON_WORKDIR}/vllm/benchmarks
 COPY --from=export_vllm /examples ${COMMON_WORKDIR}/vllm/examples
 COPY --from=export_vllm /docker ${COMMON_WORKDIR}/vllm/docker

+ENV RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES=1
+ENV RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES=1
 ENV TOKENIZERS_PARALLELISM=false

 # ENV that can improve safe tensor loading, and end-to-end time
@@ -5,8 +5,6 @@ ARG PYTORCH_BRANCH="1c57644d"
 ARG PYTORCH_VISION_BRANCH="v0.23.0"
 ARG PYTORCH_REPO="https://github.com/ROCm/pytorch.git"
 ARG PYTORCH_VISION_REPO="https://github.com/pytorch/vision.git"
-ARG PYTORCH_AUDIO_BRANCH="v2.9.0"
-ARG PYTORCH_AUDIO_REPO="https://github.com/pytorch/audio.git"
 ARG FA_BRANCH="0e60e394"
 ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git"
 ARG AITER_BRANCH="59bd8ff2"
@@ -25,7 +23,6 @@ ENV AITER_ROCM_ARCH=gfx942;gfx950
 ENV HSA_NO_SCRATCH_RECLAIM=1

 ARG PYTHON_VERSION=3.12
-ENV PYTHON_VERSION=${PYTHON_VERSION}

 RUN mkdir -p /app
 WORKDIR /app
@@ -48,7 +45,6 @@ RUN apt-get update -y \
    && python3 --version && python3 -m pip --version

 RUN pip install -U packaging 'cmake<4' ninja wheel 'setuptools<80' pybind11 Cython
-RUN apt-get update && apt-get install -y libjpeg-dev libsox-dev libsox-fmt-all sox && rm -rf /var/lib/apt/lists/*

 FROM base AS build_triton
 ARG TRITON_BRANCH
@@ -70,14 +66,11 @@ RUN mkdir -p /app/install && cp /opt/rocm/share/amd_smi/dist/*.whl /app/install
 FROM base AS build_pytorch
 ARG PYTORCH_BRANCH
 ARG PYTORCH_VISION_BRANCH
-ARG PYTORCH_AUDIO_BRANCH
 ARG PYTORCH_REPO
 ARG PYTORCH_VISION_REPO
-ARG PYTORCH_AUDIO_REPO
-
 RUN git clone ${PYTORCH_REPO} pytorch
-RUN cd pytorch && git checkout ${PYTORCH_BRANCH} \
-    && pip install -r requirements.txt && git submodule update --init --recursive \
+RUN cd pytorch && git checkout ${PYTORCH_BRANCH} && \
+    pip install -r requirements.txt && git submodule update --init --recursive \
    && python3 tools/amd_build/build_amd.py \
    && CMAKE_PREFIX_PATH=$(python3 -c 'import sys; print(sys.prefix)') python3 setup.py bdist_wheel --dist-dir=dist \
    && pip install dist/*.whl
@@ -85,15 +78,8 @@ RUN git clone ${PYTORCH_VISION_REPO} vision
 RUN cd vision && git checkout ${PYTORCH_VISION_BRANCH} \
    && python3 setup.py bdist_wheel --dist-dir=dist \
    && pip install dist/*.whl
-RUN git clone ${PYTORCH_AUDIO_REPO} audio
-RUN cd audio && git checkout ${PYTORCH_AUDIO_BRANCH} \
-    && git submodule update --init --recursive \
-    && pip install -r requirements.txt \
-    && python3 setup.py bdist_wheel --dist-dir=dist \
-    && pip install dist/*.whl
 RUN mkdir -p /app/install && cp /app/pytorch/dist/*.whl /app/install \
-    && cp /app/vision/dist/*.whl /app/install \
-    && cp /app/audio/dist/*.whl /app/install
+    && cp /app/vision/dist/*.whl /app/install

 FROM base AS build_fa
 ARG FA_BRANCH
@@ -144,8 +130,6 @@ ARG PYTORCH_BRANCH
 ARG PYTORCH_VISION_BRANCH
 ARG PYTORCH_REPO
 ARG PYTORCH_VISION_REPO
-ARG PYTORCH_AUDIO_BRANCH
-ARG PYTORCH_AUDIO_REPO
 ARG FA_BRANCH
 ARG FA_REPO
 ARG AITER_BRANCH
@@ -157,9 +141,7 @@ RUN echo "BASE_IMAGE: ${BASE_IMAGE}" > /app/versions.txt \
    && echo "PYTORCH_VISION_BRANCH: ${PYTORCH_VISION_BRANCH}" >> /app/versions.txt \
    && echo "PYTORCH_REPO: ${PYTORCH_REPO}" >> /app/versions.txt \
    && echo "PYTORCH_VISION_REPO: ${PYTORCH_VISION_REPO}" >> /app/versions.txt \
-    && echo "PYTORCH_AUDIO_BRANCH: ${PYTORCH_AUDIO_BRANCH}" >> /app/versions.txt \
-    && echo "PYTORCH_AUDIO_REPO: ${PYTORCH_AUDIO_REPO}" >> /app/versions.txt \
    && echo "FA_BRANCH: ${FA_BRANCH}" >> /app/versions.txt \
    && echo "FA_REPO: ${FA_REPO}" >> /app/versions.txt \
    && echo "AITER_BRANCH: ${AITER_BRANCH}" >> /app/versions.txt \
-    && echo "AITER_REPO: ${AITER_REPO}" >> /app/versions.txt
+    && echo "AITER_REPO: ${AITER_REPO}" >> /app/versions.txt
@@ -1,4 +1,4 @@
-FROM intel/deep-learning-essentials:2025.2.2-0-devel-ubuntu24.04 AS vllm-base
+FROM intel/deep-learning-essentials:2025.1.3-0-devel-ubuntu24.04 AS vllm-base

 RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB | gpg --dearmor | tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null && \
    echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | tee /etc/apt/sources.list.d/oneAPI.list && \
@@ -25,14 +25,10 @@ RUN apt clean && apt-get update -y && \
 RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.12 1
 RUN update-alternatives --install /usr/bin/python python /usr/bin/python3.12 1

-RUN apt install -y libze1 libze-dev libze-intel-gpu1 intel-opencl-icd libze-intel-gpu-raytracing intel-ocloc
-
-# This oneccl contains the BMG support which is not the case for default version of oneapi 2025.2.
-RUN wget https://github.com/uxlfoundation/oneCCL/releases/download/2021.15.6/intel-oneccl-2021.15.6.9_offline.sh
-RUN bash intel-oneccl-2021.15.6.9_offline.sh -a --silent --eula accept && \
-    echo "source /opt/intel/oneapi/setvars.sh --force" >> /root/.bashrc && \
-    echo "source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force" >> /root/.bashrc
+RUN apt install -y libze1 libze-dev libze-intel-gpu1 intel-opencl-icd libze-intel-gpu-raytracing

+RUN wget https://github.com/uxlfoundation/oneCCL/releases/download/2021.15.4/intel-oneccl-2021.15.4.11_offline.sh
+RUN bash intel-oneccl-2021.15.4.11_offline.sh -a --silent --eula accept && echo "source /opt/intel/oneapi/setvars.sh --force" >> /root/.bashrc
 SHELL ["bash", "-c"]
 CMD ["bash", "-c", "source /root/.bashrc && exec bash"]

@@ -76,7 +72,6 @@ RUN python3 -m pip install -e tests/vllm_test_utils
 ENV NIXL_VERSION=0.7.0
 RUN python3 /workspace/vllm/tools/install_nixl_from_source_ubuntu.py

-# remove torch bundled oneccl to avoid conflicts
 RUN --mount=type=cache,target=/root/.cache/pip \
    pip uninstall oneccl oneccl-devel -y

@@ -24,16 +24,14 @@ nav:
      - deployment/integrations
    - Training: training
    - Configuration:
+      - configuration/README.md
      - configuration/*
-      - TPU: https://docs.vllm.ai/projects/tpu/en/latest/
    - Models:
      - models/supported_models.md
      - models/generative_models.md
      - models/pooling_models.md
      - models/extensions
-      - Hardware Supported Models:
-        - models/hardware_supported_models/*
-        - TPU: https://docs.vllm.ai/projects/tpu/en/latest/recommended_models_features/
+      - Hardware Supported Models: models/hardware_supported_models
    - Features: features
  - Developer Guide:
    - contributing/README.md
@@ -52,11 +50,6 @@ nav:
      - Plugins:
        - design/*plugin*.md
      - design/*
-  - Benchmarking:
-      - benchmarking/README.md
-      - benchmarking/cli.md
-      - benchmarking/sweeps.md
-      - benchmarking/dashboard.md
  - API Reference:
    - api/README.md
    - api/vllm
@@ -1,7 +0,0 @@
-# Benchmark Suites
-
-vLLM provides comprehensive benchmarking tools for performance testing and evaluation:
-
- **[Benchmark CLI](./cli.md)**: `vllm bench` CLI tools and specialized benchmark scripts for interactive performance testing.
- **[Parameter Sweeps](./sweeps.md)**: Automate `vllm bench` runs for multiple configurations, useful for [optimization and tuning](../configuration/optimization.md).
- **[Performance Dashboard](./dashboard.md)**: Automated CI that publishes benchmarks on each commit.
@@ -1,58 +0,0 @@
-# Performance Dashboard
-
-The performance dashboard is used to confirm whether new changes improve/degrade performance under various workloads.
-It is updated by triggering benchmark runs on every commit with both the `perf-benchmarks` and `ready` labels, and when a PR is merged into vLLM.
-
-The results are automatically published to the public [vLLM Performance Dashboard](https://hud.pytorch.org/benchmark/llms?repoName=vllm-project%2Fvllm).
-
-## Manually Trigger the benchmark
-
-Use [vllm-ci-test-repo images](https://gallery.ecr.aws/q9t5s3a7/vllm-ci-test-repo) with vLLM benchmark suite.
-For CPU environment, please use the image with "-cpu" postfix.
-
-Here is an example for docker run command for CPU.
-
-```bash
-docker run -it --entrypoint /bin/bash -v /data/huggingface:/root/.cache/huggingface  -e HF_TOKEN=''  --shm-size=16g --name vllm-cpu-ci  public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:1da94e673c257373280026f75ceb4effac80e892-cpu
-```
-
-Then, run below command inside the docker instance.
-
-```bash
-bash .buildkite/performance-benchmarks/scripts/run-performance-benchmarks.sh
-```
-
-When run, benchmark script generates results under **benchmark/results** folder, along with the benchmark_results.md and benchmark_results.json.
-
-### Runtime environment variables
-
- `ON_CPU`: set the value to '1' on Intel® Xeon® Processors. Default value is 0.
- `SERVING_JSON`: JSON file to use for the serving tests. Default value is empty string (use default file).
- `LATENCY_JSON`: JSON file to use for the latency tests. Default value is empty string (use default file).
- `THROUGHPUT_JSON`: JSON file to use for the throughout tests. Default value is empty string (use default file).
- `REMOTE_HOST`: IP for the remote vLLM service to benchmark. Default value is empty string.
- `REMOTE_PORT`: Port for the remote vLLM service to benchmark. Default value is empty string.
-
-For more results visualization, check the [visualizing the results](https://github.com/intel-ai-tce/vllm/blob/more_cpu_models/.buildkite/nightly-benchmarks/README.md#visualizing-the-results).
-
-More information on the performance benchmarks and their parameters can be found in [Benchmark README](https://github.com/intel-ai-tce/vllm/blob/more_cpu_models/.buildkite/nightly-benchmarks/README.md) and [performance benchmark description](../../.buildkite/performance-benchmarks/performance-benchmarks-descriptions.md).
-
-## Continuous Benchmarking
-
-The continuous benchmarking provides automated performance monitoring for vLLM across different models and GPU devices. This helps track vLLM's performance characteristics over time and identify any performance regressions or improvements.
-
-### How It Works
-
-The continuous benchmarking is triggered via a [GitHub workflow CI](https://github.com/pytorch/pytorch-integration-testing/actions/workflows/vllm-benchmark.yml) in the PyTorch infrastructure repository, which runs automatically every 4 hours. The workflow executes three types of performance tests:
-
- **Serving tests**: Measure request handling and API performance
- **Throughput tests**: Evaluate token generation rates
- **Latency tests**: Assess response time characteristics
-
-### Benchmark Configuration
-
-The benchmarking currently runs on a predefined set of models configured in the [vllm-benchmarks directory](https://github.com/pytorch/pytorch-integration-testing/tree/main/vllm-benchmarks/benchmarks). To add new models for benchmarking:
-
-1. Navigate to the appropriate GPU directory in the benchmarks configuration
-2. Add your model specifications to the corresponding configuration files
-3. The new models will be included in the next scheduled benchmark run
@@ -1,178 +0,0 @@
-# Parameter Sweeps
-
-## Online Benchmark
-
-### Basic
-
-`vllm bench sweep serve` automatically starts `vllm serve` and runs `vllm bench serve` to evaluate vLLM over multiple configurations.
-
-Follow these steps to run the script:
-
-1. Construct the base command to `vllm serve`, and pass it to the `--serve-cmd` option.
-2. Construct the base command to `vllm bench serve`, and pass it to the `--bench-cmd` option.
-3. (Optional) If you would like to vary the settings of `vllm serve`, create a new JSON file and populate it with the parameter combinations you want to test. Pass the file path to `--serve-params`.
-
-    - Example: Tuning `--max-num-seqs` and `--max-num-batched-tokens`:
-
-    ```json
-    [
-        {
-            "max_num_seqs": 32,
-            "max_num_batched_tokens": 1024
-        },
-        {
-            "max_num_seqs": 64,
-            "max_num_batched_tokens": 1024
-        },
-        {
-            "max_num_seqs": 64,
-            "max_num_batched_tokens": 2048
-        },
-        {
-            "max_num_seqs": 128,
-            "max_num_batched_tokens": 2048
-        },
-        {
-            "max_num_seqs": 128,
-            "max_num_batched_tokens": 4096
-        },
-        {
-            "max_num_seqs": 256,
-            "max_num_batched_tokens": 4096
-        }
-    ]
-    ```
-
-4. (Optional) If you would like to vary the settings of `vllm bench serve`, create a new JSON file and populate it with the parameter combinations you want to test. Pass the file path to `--bench-params`.
-
-    - Example: Using different input/output lengths for random dataset:
-
-    ```json
-    [
-        {
-            "random_input_len": 128,
-            "random_output_len": 32
-        },
-        {
-            "random_input_len": 256,
-            "random_output_len": 64
-        },
-        {
-            "random_input_len": 512,
-            "random_output_len": 128
-        }
-    ]
-    ```
-
-5. Determine where you want to save the results, and pass that to `--output-dir`.
-
-Example command:
-
-```bash
-vllm bench sweep serve \
-    --serve-cmd 'vllm serve meta-llama/Llama-2-7b-chat-hf' \
-    --bench-cmd 'vllm bench serve --model meta-llama/Llama-2-7b-chat-hf --backend vllm --endpoint /v1/completions --dataset-name sharegpt --dataset-path benchmarks/ShareGPT_V3_unfiltered_cleaned_split.json' \
-    --serve-params benchmarks/serve_hparams.json \
-    --bench-params benchmarks/bench_hparams.json \
-    -o benchmarks/results
-```
-
-!!! important
-    If both `--serve-params` and `--bench-params` are passed, the script will iterate over the Cartesian product between them.
-    You can use `--dry-run` to preview the commands to be run.
-
-    We only start the server once for each `--serve-params`, and keep it running for multiple `--bench-params`.
-    Between each benchmark run, we call the `/reset_prefix_cache` and `/reset_mm_cache` endpoints to get a clean slate for the next run.
-    In case you are using a custom `--serve-cmd`, you can override the commands used for resetting the state by setting `--after-bench-cmd`.
-
-!!! note
-    By default, each parameter combination is run 3 times to make the results more reliable. You can adjust the number of runs by setting `--num-runs`.
-
-!!! tip
-    You can use the `--resume` option to continue the parameter sweep if one of the runs failed.
-  
-### SLA auto-tuner
-
-`vllm bench sweep serve_sla` is a wrapper over `vllm bench sweep serve` that tunes either the request rate or concurrency (choose using `--sla-variable`) in order to satisfy the SLA constraints given by `--sla-params`.
-
-For example, to ensure E2E latency within different target values for 99% of requests:
-
-```json
-[
-    {
-        "p99_e2el_ms": "<=200"
-    },
-    {
-        "p99_e2el_ms": "<=500"
-    },
-    {
-        "p99_e2el_ms": "<=1000"
-    },
-    {
-        "p99_e2el_ms": "<=2000"
-    }
-]
-```
-
-Example command:
-
-```bash
-vllm bench sweep serve_sla \
-    --serve-cmd 'vllm serve meta-llama/Llama-2-7b-chat-hf' \
-    --bench-cmd 'vllm bench serve --model meta-llama/Llama-2-7b-chat-hf --backend vllm --endpoint /v1/completions --dataset-name sharegpt --dataset-path benchmarks/ShareGPT_V3_unfiltered_cleaned_split.json' \
-    --serve-params benchmarks/serve_hparams.json \
-    --bench-params benchmarks/bench_hparams.json \
-    --sla-params benchmarks/sla_hparams.json \
-    --sla-variable max_concurrency \
-    -o benchmarks/results
-```
-
-The algorithm for adjusting the SLA variable is as follows:
-
-1. Run the benchmark with infinite QPS, and use the corresponding metrics to determine the initial value of the variable.
-    - For example, the initial request rate is set to the concurrency under infinite QPS.
-2. If the SLA is still satisfied, keep doubling the value until the SLA is no longer satisfied. This gives a relatively narrow window that contains the point where the SLA is barely satisfied.
-3. Apply binary search over the window to find the maximum value that still satisfies the SLA.
-
-!!! important
-    SLA tuning is applied over each combination of `--serve-params`, `--bench-params`, and `--sla-params`.
-
-    For a given combination of `--serve-params` and `--bench-params`, we share the benchmark results across `--sla-params` to avoid rerunning benchmarks with the same SLA variable value.
-
-## Visualization
-
-### Basic
-
-`vllm bench sweep plot` can be used to plot performance curves from parameter sweep results.
-
-Example command:
-
-```bash
-vllm bench sweep plot benchmarks/results/<timestamp> \
-    --var-x max_concurrency \
-    --row-by random_input_len \
-    --col-by random_output_len \
-    --curve-by api_server_count,max_num_batched_tokens \
-    --filter-by 'max_concurrency<=1024'
-```
-
-!!! tip
-    You can use `--dry-run` to preview the figures to be plotted.
-
-### Pareto chart
-
-`vllm bench sweep plot_pareto` helps pick configurations that balance per-user and per-GPU throughput.
-
-Higher concurrency or batch size can raise GPU efficiency (per-GPU), but can add per user latency; lower concurrency improves per-user rate but underutilizes GPUs; The Pareto frontier shows the best achievable pairs across your runs.
-
- x-axis: tokens/s/user = `output_throughput` ÷ concurrency (`--user-count-var`, default `max_concurrency`, fallback `max_concurrent_requests`).
- y-axis: tokens/s/GPU = `output_throughput` ÷ GPU count (`--gpu-count-var` if set; else gpu_count is TP×PP*DP).
- Output: a single figure at `OUTPUT_DIR/pareto/PARETO.png`.
- Show the configuration used in each data point `--label-by` (default: `max_concurrency,gpu_count`).
-
-Example:
-
-```bash
-vllm bench sweep plot_pareto benchmarks/results/<timestamp> \
-  --label-by max_concurrency,tensor_parallel_size,pipeline_parallel_size
-```
@@ -1,9 +0,0 @@
-# vllm bench sweep plot_pareto
-
-## JSON CLI Arguments
-
--8<-- "docs/cli/json_tip.inc.md"
-
-## Arguments
-
--8<-- "docs/argparse/bench_sweep_plot_pareto.inc.md"
@@ -10,7 +10,6 @@ Stay tuned for upcoming meetups! Follow us on [Twitter/X](https://x.com/vllm_pro

 Below you'll find slides and recordings from our previous meetups:

- [vLLM Bangkok Meetup](https://luma.com/v0f647nv), November 21st 2025. [[Slides]](https://drive.google.com/drive/folders/1H0DS57F8HQ5q3kSOSoRmucPJWL3E0A_X?usp=sharing)
 - [vLLM Zurich Meetup](https://luma.com/0gls27kb), November 6th 2025. [[Slides]](https://docs.google.com/presentation/d/1UC9PTLCHYXQpOmJDSFg6Sljra3iVXzc09DeEI7dnxMc/edit?usp=sharing) [[Recording]](https://www.youtube.com/watch?v=6m6ZE6yVEDI)
 - [vLLM Beijing Meetup](https://mp.weixin.qq.com/s/xSrYXjNgr1HbCP4ExYNG1w), November 1st 2025. [[Slides]](https://drive.google.com/drive/folders/1nQJ8ZkLSjKxvu36sSHaceVXtttbLvvu-?usp=drive_link)
 - [vLLM Shanghai Meetup](https://mp.weixin.qq.com/s/__xb4OyOsImz-9eAVrdlcg), October 25th 2025. [[Slides]](https://drive.google.com/drive/folders/1KqwjsFJLfEsC8wlDugnrR61zsWHt94Q6)
@@ -49,6 +49,9 @@ llm = LLM(model="adept/fuyu-8b", max_model_len=2048, max_num_seqs=2)

 By default, we optimize model inference using CUDA graphs which take up extra memory in the GPU.

+!!! warning
+    CUDA graph capture takes up more memory in V1 than in V0.
+
 You can adjust `compilation_config` to achieve a better balance between inference speed and memory usage:

 ??? code
@@ -7,6 +7,8 @@ vLLM uses the following environment variables to configure the system:

    All environment variables used by vLLM are prefixed with `VLLM_`. **Special care should be taken for Kubernetes users**: please do not name the service as `vllm`, otherwise environment variables set by Kubernetes might conflict with vLLM's environment variables, because [Kubernetes sets environment variables for each service with the capitalized service name as the prefix](https://kubernetes.io/docs/concepts/services-networking/service/#environment-variables).

-```python
--8<-- "vllm/envs.py:env-vars-definition"
-```
+??? code
+
+    ```python
+    --8<-- "vllm/envs.py:env-vars-definition"
+    ```
@@ -31,7 +31,9 @@ In vLLM V1, the default preemption mode is `RECOMPUTE` rather than `SWAP`, as re

 Chunked prefill allows vLLM to process large prefills in smaller chunks and batch them together with decode requests. This feature helps improve both throughput and latency by better balancing compute-bound (prefill) and memory-bound (decode) operations.

-In V1, **chunked prefill is enabled by default whenever possible**. With chunked prefill enabled, the scheduling policy prioritizes decode requests. It batches all pending decode requests before scheduling any prefill operations. When there are available tokens in the `max_num_batched_tokens` budget, it schedules pending prefills. If a pending prefill request cannot fit into `max_num_batched_tokens`, it automatically chunks it.
+In vLLM V1, **chunked prefill is always enabled by default**. This is different from vLLM V0, where it was conditionally enabled based on model characteristics.
+
+With chunked prefill enabled, the scheduling policy prioritizes decode requests. It batches all pending decode requests before scheduling any prefill operations. When there are available tokens in the `max_num_batched_tokens` budget, it schedules pending prefills. If a pending prefill request cannot fit into `max_num_batched_tokens`, it automatically chunks it.

 This policy has two benefits:

--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
Lucas WilkinsonandKevin H. Luu	275de34170	[BugFix] Fix false assertion with spec-decode=[2,4,..] and TP>2 (#29036 ) Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> (cherry picked from commit `8f4f77a727`)	2025-11-19 14:11:21 -08:00
Julien DenizeandKevin H. Luu	fa3ffb4365	[BugFix] Ray with multiple nodes (#28873 ) Signed-off-by: Julien Denize <julien.denize@mistral.ai> (cherry picked from commit `cdeec2e606`)	2025-11-19 14:11:08 -08:00
Lucas WilkinsonandKevin H. Luu	6d5974369c	[BugFix] Fix async-scheduling + FlashAttn MLA (#28990 ) Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> (cherry picked from commit `48fc8b1e59`)	2025-11-19 14:10:50 -08:00
JohnnyandKevin H. Luu	0ce9990d2c	[NVIDIA] Guard SM100 CUTLASS MoE macro to SM100 builds v2 (#28938 ) Signed-off-by: johnnynunez <johnnynuca14@gmail.com> Signed-off-by: Johnny <johnnynuca14@gmail.com> (cherry picked from commit `49ef847aa8`)	2025-11-19 14:10:37 -08:00