Remove broken allocator pre-warming (initializes CUDA before fork)

The pre-warming called torch.cuda.mem_get_info() in the parent process before the EngineCore subprocess forks, which initialized CUDA. This caused "Cannot re-initialize CUDA in forked subprocess" in the child process, breaking tests that previously passed. The pre-warming also wouldn't have helped — PyTorch's allocator pool is per-process, so the child's pool starts empty regardless. Signed-off-by: khluu <khluu000@gmail.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
[CI] Migrate all remaining gpu_1_queue jobs to h200_18gb MIG
2026-05-13 02:26:43 -07:00 · 2026-05-12 18:50:08 -07:00 · 2026-05-12 15:57:19 -07:00 · 2026-05-12 02:00:37 -07:00 · 2026-05-12 08:45:51 +00:00 · 2026-05-12 07:47:47 +00:00
938 changed files with 86508 additions and 41882 deletions
@@ -12,15 +12,19 @@ steps:
  - vllm/_custom_ops.py
  - tests/kernels/attention/test_cpu_attn.py
  - tests/kernels/moe/test_cpu_fused_moe.py
+  - tests/kernels/moe/test_cpu_fp8_fused_moe.py
  - tests/kernels/test_onednn.py
  - tests/kernels/test_awq_int4_to_int8.py
+  - tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
  commands:
    - |
-      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 20m "
+      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
      pytest -x -v -s tests/kernels/attention/test_cpu_attn.py
      pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py
+      pytest -x -v -s tests/kernels/moe/test_cpu_fp8_fused_moe.py
      pytest -x -v -s tests/kernels/test_onednn.py
-      pytest -x -v -s tests/kernels/test_awq_int4_to_int8.py"
+      pytest -x -v -s tests/kernels/test_awq_int4_to_int8.py
+      pytest -x -v -s tests/kernels/quantization/test_cpu_fp8_scaled_mm.py"

 - label: CPU-Compatibility Tests
  depends_on: []
@@ -61,6 +65,7 @@ steps:
  - vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_int8.py
  - vllm/model_executor/layers/quantization/kernels/scaled_mm/cpu.py
  - vllm/model_executor/layers/quantization/kernels/mixed_precision/cpu.py
+  - vllm/model_executor/layers/fused_moe/experts/cpu_moe.py
  - tests/quantization/test_compressed_tensors.py
  - tests/quantization/test_cpu_wna16.py
  commands:
@@ -18,17 +18,18 @@ steps:
    - >-
      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
      'cd tests &&
+      export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
      pytest -v -s lora/test_layers.py &&
      pytest -v -s lora/test_lora_checkpoints.py &&
-      (pytest -v -s lora/test_lora_functions.py --deselect="tests/lora/test_lora_functions.py::test_lora_functions_sync" --deselect="tests/lora/test_lora_functions.py::test_lora_functions_async" || true) &&
+      pytest -v -s lora/test_lora_functions.py &&
      pytest -v -s lora/test_lora_huggingface.py &&
      pytest -v -s lora/test_lora_manager.py &&
      pytest -v -s lora/test_lora_utils.py &&
      pytest -v -s lora/test_peft_helper.py &&
      pytest -v -s lora/test_resolver.py &&
      pytest -v -s lora/test_utils.py &&
-      (pytest -v -s lora/test_add_lora.py --deselect="tests/lora/test_add_lora.py::test_add_lora" || true) &&
-      (pytest -v -s lora/test_worker.py --deselect="tests/lora/test_worker.py::test_worker_apply_lora" || true)'
+      pytest -v -s lora/test_add_lora.py  &&
+      pytest -v -s lora/test_worker.py'

 - label: LoRA Fused/MoE Kernels
  timeout_in_minutes: 45
@@ -46,6 +47,7 @@ steps:
    - >-
      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
      'cd tests &&
+      export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
      pytest -v -s lora/test_fused_moe_lora_kernel.py && 
      pytest -v -s lora/test_moe_lora_align_sum.py'

@@ -65,8 +67,9 @@ steps:
    - >-
      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
      'cd tests &&
+      export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
      set -o pipefail &&
-      pytest -v -s lora/test_punica_ops.py --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-2-2049-64-32-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[expand-0-xpu:0-dtype1-2-64000-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-128-1-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-256-1-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-256-8-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[expand-0-xpu:0-dtype0-3-2049-128-8-16]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-128-8-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels[expand-0-xpu:0-dtype1-1-2049-256-128-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype0-3-64256-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype1-2-29696-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype1-3-49408-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype0-2-16384-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[expand-0-xpu:0-dtype0-2-51328-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[expand-0-xpu:0-dtype1-1-102656-32-4-4]"'
+      pytest -v -s lora/test_punica_ops.py --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[expand-0-xpu:0-dtype0-3-43264-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype1-1-2049-64-128-16]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-128-1-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-256-1-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-256-8-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels[expand-0-xpu:0-dtype0-3-2049-128-8-16]" --deselect="tests/lora/test_punica_ops.py::test_kernels[shrink-0-xpu:0-dtype0-1-2049-128-8-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels[expand-0-xpu:0-dtype1-1-2049-256-128-32]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype0-3-64256-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype1-2-29696-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype1-3-49408-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[shrink-0-xpu:0-dtype0-2-16384-32-4-4]" --deselect="tests/lora/test_punica_ops.py::test_kernels_hidden_size[expand-0-xpu:0-dtype0-2-51328-32-4-4]"'

 - label: LoRA Punica FP8/XPU Ops
  timeout_in_minutes: 45
@@ -84,6 +87,7 @@ steps:
    - >-
      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
      'cd tests &&
+      export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
      pytest -v -s lora/test_punica_ops_fp8.py &&
      pytest -v -s lora/test_punica_xpu_ops.py'

@@ -103,10 +107,12 @@ steps:
    - >-
      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
      'cd tests &&
+      export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
      (pytest -v -s lora/test_mixtral.py --deselect="tests/lora/test_mixtral.py::test_mixtral_lora[4]" || true) &&
      pytest -v -s lora/test_quant_model.py --deselect="tests/lora/test_quant_model.py::test_quant_model_lora[model0]" --deselect="tests/lora/test_quant_model.py::test_quant_model_lora[model1]" --deselect="tests/lora/test_quant_model.py::test_quant_model_tp_equality[model0]" &&
-      pytest -v -s lora/test_qwen35_densemodel_lora.py &&
-      pytest -v -s lora/test_transformers_model.py'
+      pytest -v -s lora/test_transformers_model.py &&
+      pytest -v -s lora/test_chatglm3_tp.py &&
+      pytest -s -v lora/test_minicpmv_tp.py'

 - label: LoRA Multimodal
  timeout_in_minutes: 45
@@ -124,6 +130,6 @@ steps:
    - >-
      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
      'cd tests &&
+      export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
      pytest -v -s lora/test_default_mm_loras.py && 
-      (pytest -v -s lora/test_qwen3_unembed.py || true) &&
      pytest -v -s lora/test_whisper.py'
@@ -49,7 +49,7 @@ steps:
      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
      'export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
      cd tests &&
-      pytest -v -s v1/logits_processors &&
+      pytest -v -s v1/logits_processors --ignore=v1/logits_processors/test_custom_online.py --ignore=v1/logits_processors/test_custom_offline.py &&
      pytest -v -s v1/test_oracle.py &&
      pytest -v -s v1/test_request.py &&
      pytest -v -s v1/test_outputs.py'
@@ -61,5 +61,5 @@ steps:
        pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py --ignore=v1/worker/test_worker_memory_snapshot.py &&
        pytest -v -s v1/structured_output &&
        pytest -v -s v1/test_serial_utils.py &&
-        pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_tree_attention.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py &&
+        pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py &&
        pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --ignore=v1/kv_connector/unit/test_lmcache_integration.py --ignore=v1/kv_connector/unit/test_hf3fs_client.py --ignore=v1/kv_connector/unit/test_hf3fs_connector.py --ignore=v1/kv_connector/unit/test_hf3fs_metadata_server.py'
@@ -37,7 +37,7 @@ steps:
        agents:
          queue: arm64_cpu_queue_release
        commands:
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64}\" --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04  --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64}\" --build-arg BUILD_OS=manylinux --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda13.0 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
          - "mkdir artifacts"
          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
          - "bash .buildkite/scripts/upload-nightly-wheels.sh"
@@ -76,7 +76,7 @@ steps:
        agents:
          queue: cpu_queue_release
        commands:
-          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86}\" --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
+          - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86}\" --build-arg BUILD_OS=manylinux --build-arg BUILD_BASE_IMAGE=pytorch/manylinux2_28-builder:cuda13.0 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
          - "mkdir artifacts"
          - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
          - "bash .buildkite/scripts/upload-nightly-wheels.sh"
@@ -309,6 +309,7 @@ steps:
        depends_on: ~

      - label: "Build release image - x86_64 - CPU"
+        key: build-cpu-release-image-x86
        depends_on:
          - block-cpu-release-image-build
          - input-release-version
@@ -327,7 +328,8 @@ steps:
        depends_on: ~

      - label: "Build release image - arm64 - CPU"
-        depends_on: 
+        key: build-cpu-release-image-arm64
+        depends_on:
          - block-arm64-cpu-release-image-build
          - input-release-version
        agents:
@@ -436,6 +438,41 @@ steps:
          DOCKER_BUILDKIT: "1"
          DOCKERHUB_USERNAME: "vllmbot"

+      - block: "Publish release images to DockerHub"
+        key: block-publish-release-images
+        depends_on:
+          - create-multi-arch-manifest
+          - create-multi-arch-manifest-cuda-12-9
+          - create-multi-arch-manifest-ubuntu2404
+          - create-multi-arch-manifest-cuda-12-9-ubuntu2404
+          - build-rocm-release-image
+          - input-release-version
+          # Wait for CPU builds if their block steps were unblocked, so publish
+          # doesn't race the in-progress CPU build. allow_failure lets publish
+          # proceed when the operator legitimately leaves the CPU block steps
+          # unblocked or the CPU build fails.
+          - step: build-cpu-release-image-x86
+            allow_failure: true
+          - step: build-cpu-release-image-arm64
+            allow_failure: true
+        if: build.env("NIGHTLY") != "1"
+
+      - label: "Publish release images to DockerHub"
+        depends_on:
+          - block-publish-release-images
+        key: publish-release-images-dockerhub
+        agents:
+          queue: small_cpu_queue_release
+        commands:
+          - "bash .buildkite/scripts/publish-release-images.sh"
+        plugins:
+          - docker-login#v3.0.0:
+              username: vllmbot
+              password-env: DOCKERHUB_TOKEN
+        env:
+          DOCKER_BUILDKIT: "1"
+          DOCKERHUB_USERNAME: "vllmbot"
+
  - group: "Publish wheels"
    key: "publish-wheels"
    steps:
@@ -723,7 +760,7 @@ steps:
      - "bash tools/vllm-rocm/generate-rocm-wheels-root-index.sh"
    env:
      S3_BUCKET: "vllm-wheels"
-      VARIANT: "rocm721"
+      VARIANT: "rocm722"

  # ROCm Job 6: Build ROCm Release Docker Image
  - label: ":docker: Build release image - x86_64 - ROCm"
@@ -8,8 +8,6 @@ if [ -z "${RELEASE_VERSION}" ]; then
  RELEASE_VERSION="1.0.0.dev"
 fi

-ROCM_BASE_CACHE_KEY=$(.buildkite/scripts/cache-rocm-base-wheels.sh key)
-
 buildkite-agent annotate --style 'info' --context 'release-workflow' << EOF
 To download the wheel (by commit):
 \`\`\`
@@ -25,95 +23,5 @@ aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cpu-cp38-
 aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cpu-cp38-abi3-manylinux_2_35_aarch64.whl .
 \`\`\`

-
-To download and upload the image:
-
-\`\`\`
-# Download images:
-
-docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64
-docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64
-docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu129
-docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu129
-docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base
-docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm
-docker pull public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:v${RELEASE_VERSION}
-docker pull public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:v${RELEASE_VERSION}
-
-# Tag and push images:
-
-## CUDA
-
-docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64 vllm/vllm-openai:x86_64
-docker tag vllm/vllm-openai:x86_64 vllm/vllm-openai:latest-x86_64
-docker tag vllm/vllm-openai:x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64
-docker push vllm/vllm-openai:latest-x86_64
-docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64
-
-docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu129 vllm/vllm-openai:x86_64-cu129
-docker tag vllm/vllm-openai:x86_64-cu129 vllm/vllm-openai:latest-x86_64-cu129
-docker tag vllm/vllm-openai:x86_64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129
-docker push vllm/vllm-openai:latest-x86_64-cu129
-docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129
-
-docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64 vllm/vllm-openai:aarch64
-docker tag vllm/vllm-openai:aarch64 vllm/vllm-openai:latest-aarch64
-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 tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu129 vllm/vllm-openai:aarch64-cu129
-docker tag vllm/vllm-openai:aarch64-cu129 vllm/vllm-openai:latest-aarch64-cu129
-docker tag vllm/vllm-openai:aarch64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129
-docker push vllm/vllm-openai:latest-aarch64-cu129
-docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129
-
-## ROCm
-
-docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}
-docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT} vllm/vllm-openai-rocm:latest
-docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT} vllm/vllm-openai-rocm:v${RELEASE_VERSION}
-docker push vllm/vllm-openai-rocm:latest
-docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}
-
-docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base
-docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base vllm/vllm-openai-rocm:latest-base
-docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base
-docker push vllm/vllm-openai-rocm:latest-base
-docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base
-
-## CPU
-
-docker tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:v${RELEASE_VERSION} vllm/vllm-openai-cpu:x86_64
-docker tag vllm/vllm-openai-cpu:x86_64 vllm/vllm-openai-cpu:latest-x86_64
-docker tag vllm/vllm-openai-cpu:x86_64 vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64
-docker push vllm/vllm-openai-cpu:latest-x86_64
-docker push vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64
-
-docker tag public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:v${RELEASE_VERSION} vllm/vllm-openai-cpu:arm64
-docker tag vllm/vllm-openai-cpu:arm64 vllm/vllm-openai-cpu:latest-arm64
-docker tag vllm/vllm-openai-cpu:arm64 vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64
-docker push vllm/vllm-openai-cpu:latest-arm64
-docker push vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64
-
-# Create multi-arch manifest:
-
-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 push vllm/vllm-openai:latest
-docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}
-
-docker manifest rm vllm/vllm-openai:latest-cu129
-docker manifest create vllm/vllm-openai:latest-cu129 vllm/vllm-openai:latest-x86_64-cu129 vllm/vllm-openai:latest-aarch64-cu129
-docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129
-docker manifest push vllm/vllm-openai:latest-cu129
-docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu129
-
-docker manifest rm vllm/vllm-openai-cpu:latest || true
-docker manifest create vllm/vllm-openai-cpu:latest vllm/vllm-openai-cpu:latest-x86_64 vllm/vllm-openai-cpu:latest-arm64
-docker manifest create vllm/vllm-openai-cpu:v${RELEASE_VERSION} vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64 vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64
-docker manifest push vllm/vllm-openai-cpu:latest
-docker manifest push vllm/vllm-openai-cpu:v${RELEASE_VERSION}
-\`\`\`
+Docker images are published automatically by the "Publish release images to DockerHub" pipeline step.
 EOF
@@ -0,0 +1,55 @@
+#!/bin/bash
+# Usage: ./ci-fetch-log.sh <buildkite_job_url> [output_file]
+#        ./ci-fetch-log.sh <build_number> <job_uuid> [output_file]
+#
+# Downloads the raw log for a Buildkite job from the public, unauthenticated
+# /organizations/<org>/pipelines/<pipeline>/builds/<n>/jobs/<uuid>/download
+# endpoint, then strips ANSI/timestamps via ci-clean-log.sh.
+#
+# Find <build_number> and <job_uuid> via:
+#   gh pr checks <PR> --repo vllm-project/vllm
+# Each failing row's URL is .../builds/<build_number>#<job_uuid>.
+
+set -euo pipefail
+
+ORG="vllm"
+PIPELINE="ci"
+
+usage() {
+    echo "Usage: $0 <buildkite_job_url> [output_file]"
+    echo "       $0 <build_number> <job_uuid> [output_file]"
+    exit 1
+}
+
+if [ $# -lt 1 ]; then usage; fi
+
+if [[ "$1" == https://* ]]; then
+    BUILD=$(echo "$1" | sed -nE 's#.*/builds/([0-9]+).*#\1#p')
+    JOB=$(echo "$1" | grep -oE '[0-9a-f]{8}-[0-9a-f-]+' | head -n 1)
+    OUT="${2:-ci-${BUILD}-${JOB:0:8}.log}"
+else
+    if [ $# -lt 2 ]; then usage; fi
+    BUILD="$1"
+    JOB="$2"
+    OUT="${3:-ci-${BUILD}-${JOB:0:8}.log}"
+fi
+
+if [ -z "$BUILD" ] || [ -z "$JOB" ]; then
+    echo "Could not parse build number or job UUID from: $1" >&2
+    usage
+fi
+
+COOKIES=$(mktemp)
+trap 'rm -f "$COOKIES"' EXIT
+
+# Buildkite issues a session cookie on first hit; subsequent /download needs it.
+curl -fsSL -c "$COOKIES" -A "vllm-ci-fetch-log" \
+    "https://buildkite.com/${ORG}/${PIPELINE}/builds/${BUILD}" -o /dev/null
+
+curl -fsSL -b "$COOKIES" -A "vllm-ci-fetch-log" \
+    "https://buildkite.com/organizations/${ORG}/pipelines/${PIPELINE}/builds/${BUILD}/jobs/${JOB}/download" \
+    -o "$OUT"
+
+bash "$(dirname "$0")/ci-clean-log.sh" "$OUT"
+
+echo "$OUT"
@@ -378,9 +378,11 @@ HF_MOUNT="/root/.cache/huggingface"
 # double-quotes will have been stripped by the calling shell.
 if [[ -n "${VLLM_TEST_COMMANDS:-}" ]]; then
  commands="${VLLM_TEST_COMMANDS}"
+  commands_source="env"
  echo "Commands sourced from VLLM_TEST_COMMANDS (quoting preserved)"
 else
  commands="$*"
+  commands_source="argv"
  if [[ -z "$commands" ]]; then
    echo "Error: No test commands provided." >&2
    echo "Usage:" >&2
@@ -397,9 +399,15 @@ fi

 echo "Raw commands: $commands"

-# Fix quoting before ROCm overrides (so overrides see correct structure)
-commands=$(re_quote_pytest_markers "$commands")
-echo "After re-quoting: $commands"
+# Only try to repair stripped pytest -m/-k quoting in legacy argv mode.
+# VLLM_TEST_COMMANDS preserves inner quoting already, and re-quoting that path
+# can corrupt embedded echo strings or otherwise well-formed shell fragments.
+if [[ "$commands_source" == "argv" ]]; then
+  commands=$(re_quote_pytest_markers "$commands")
+  echo "After re-quoting: $commands"
+else
+  echo "Skipping re-quoting for VLLM_TEST_COMMANDS input"
+fi

 commands=$(apply_rocm_test_overrides "$commands")
 echo "Final commands: $commands"
@@ -31,6 +31,21 @@ function cpu_tests() {
    set -e
    pip list"

+  # 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
+    pytest -x -v -s tests/kernels/core/test_cpu_activation.py
+    pytest -x -v -s tests/kernels/moe/test_moe.py -k test_cpu_fused_moe_basic"
+
+  # skip tests requiring model downloads if HF_TOKEN is not set
+  # due to rate-limits
+  if [ -z "$HF_TOKEN" ]; then
+    echo "Warning: HF_TOKEN is not set. Skipping tests that require model downloads."
+    return
+  fi
+
  # offline inference
  docker exec cpu-test bash -c "
    set -e
@@ -46,13 +61,6 @@ function cpu_tests() {
    set -e
    pytest -x -v -s tests/quantization/test_compressed_tensors.py::test_compressed_tensors_w8a8_logprobs"

-  # 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
-    pytest -x -v -s tests/kernels/core/test_cpu_activation.py
-    pytest -x -v -s tests/kernels/moe/test_moe.py -k test_cpu_fused_moe_basic"

  # basic online serving
  docker exec cpu-test bash -c '
@@ -67,6 +75,21 @@ function cpu_tests() {
      --num-prompts 20 \
      --endpoint /v1/completions
    kill -s SIGTERM $server_pid &'
+
+  # smoke test for Gated DeltaNet
+  docker exec cpu-test bash -c '
+    set -e
+    VLLM_CPU_OMP_THREADS_BIND=$E2E_OMP_THREADS vllm serve Qwen/Qwen3.5-0.8B --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.5-0.8B \
+      --num-prompts 20 \
+      --endpoint /v1/completions
+    kill -s SIGTERM $server_pid &'
+
 }

 # All of CPU tests are expected to be finished less than 40 mins.
@@ -136,8 +136,6 @@ run_and_track_test 3 "test_accuracy.py::test_lm_eval_accuracy_v1_engine" \
    "python3 -m pytest -s -v /workspace/vllm/tests/entrypoints/llm/test_accuracy.py::test_lm_eval_accuracy_v1_engine"
 run_and_track_test 4 "test_quantization_accuracy.py" \
    "python3 -m pytest -s -v /workspace/vllm/tests/tpu/test_quantization_accuracy.py"
-run_and_track_test 5 "examples/offline_inference/tpu.py" \
-    "python3 /workspace/vllm/examples/offline_inference/tpu.py"
 run_and_track_test 6 "test_tpu_model_runner.py" \
    "python3 -m pytest -s -v /workspace/vllm/tests/v1/tpu/worker/test_tpu_model_runner.py"
 run_and_track_test 7 "test_sampler.py" \
@@ -0,0 +1,180 @@
+#!/bin/bash
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# Publish release Docker images from ECR to DockerHub.
+# Pulls per-arch images, tags with latest and versioned tags, pushes them,
+# then creates and pushes multi-arch manifests.
+
+set -euo pipefail
+
+RELEASE_VERSION=$(buildkite-agent meta-data get release-version --default "" | sed 's/^v//')
+if [ -z "${RELEASE_VERSION}" ]; then
+  echo "ERROR: release-version metadata not set"
+  exit 1
+fi
+
+COMMIT="$BUILDKITE_COMMIT"
+ROCM_BASE_CACHE_KEY=$(.buildkite/scripts/cache-rocm-base-wheels.sh key)
+
+echo "========================================"
+echo "Publishing release images v${RELEASE_VERSION}"
+echo "  Commit: ${COMMIT}"
+echo "  ROCm base cache key: ${ROCM_BASE_CACHE_KEY}"
+echo "========================================"
+
+# Login to ECR to pull staging images
+aws ecr-public get-login-password --region us-east-1 | \
+  docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
+
+# ---- CUDA (default: 13.0) ----
+
+docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64
+docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64
+
+docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64 vllm/vllm-openai:latest-x86_64
+docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64
+docker push vllm/vllm-openai:latest-x86_64
+docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64
+
+docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64 vllm/vllm-openai:latest-aarch64
+docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-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 || true
+docker manifest rm vllm/vllm-openai:v${RELEASE_VERSION} || true
+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 push vllm/vllm-openai:latest
+docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}
+
+# ---- CUDA 12.9 ----
+
+docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129
+docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129
+
+docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129 vllm/vllm-openai:latest-x86_64-cu129
+docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129
+docker push vllm/vllm-openai:latest-x86_64-cu129
+docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129
+
+docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129 vllm/vllm-openai:latest-aarch64-cu129
+docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129
+docker push vllm/vllm-openai:latest-aarch64-cu129
+docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129
+
+docker manifest rm vllm/vllm-openai:latest-cu129 || true
+docker manifest rm vllm/vllm-openai:v${RELEASE_VERSION}-cu129 || true
+docker manifest create vllm/vllm-openai:latest-cu129 vllm/vllm-openai:latest-x86_64-cu129 vllm/vllm-openai:latest-aarch64-cu129
+docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129
+docker manifest push vllm/vllm-openai:latest-cu129
+docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu129
+
+# ---- Ubuntu 24.04 (CUDA 13.0) ----
+
+docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-ubuntu2404
+docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-ubuntu2404
+
+docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-ubuntu2404 vllm/vllm-openai:latest-x86_64-ubuntu2404
+docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-ubuntu2404
+docker push vllm/vllm-openai:latest-x86_64-ubuntu2404
+docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-ubuntu2404
+
+docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-ubuntu2404 vllm/vllm-openai:latest-aarch64-ubuntu2404
+docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-ubuntu2404
+docker push vllm/vllm-openai:latest-aarch64-ubuntu2404
+docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-ubuntu2404
+
+docker manifest rm vllm/vllm-openai:latest-ubuntu2404 || true
+docker manifest rm vllm/vllm-openai:v${RELEASE_VERSION}-ubuntu2404 || true
+docker manifest create vllm/vllm-openai:latest-ubuntu2404 vllm/vllm-openai:latest-x86_64-ubuntu2404 vllm/vllm-openai:latest-aarch64-ubuntu2404
+docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-ubuntu2404
+docker manifest push vllm/vllm-openai:latest-ubuntu2404
+docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-ubuntu2404
+
+# ---- Ubuntu 24.04 (CUDA 12.9) ----
+
+docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129-ubuntu2404
+docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129-ubuntu2404
+
+docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129-ubuntu2404 vllm/vllm-openai:latest-x86_64-cu129-ubuntu2404
+docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129-ubuntu2404
+docker push vllm/vllm-openai:latest-x86_64-cu129-ubuntu2404
+docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129-ubuntu2404
+
+docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129-ubuntu2404 vllm/vllm-openai:latest-aarch64-cu129-ubuntu2404
+docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129-ubuntu2404
+docker push vllm/vllm-openai:latest-aarch64-cu129-ubuntu2404
+docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129-ubuntu2404
+
+docker manifest rm vllm/vllm-openai:latest-cu129-ubuntu2404 || true
+docker manifest rm vllm/vllm-openai:v${RELEASE_VERSION}-cu129-ubuntu2404 || true
+docker manifest create vllm/vllm-openai:latest-cu129-ubuntu2404 vllm/vllm-openai:latest-x86_64-cu129-ubuntu2404 vllm/vllm-openai:latest-aarch64-cu129-ubuntu2404
+docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu129-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129-ubuntu2404
+docker manifest push vllm/vllm-openai:latest-cu129-ubuntu2404
+docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu129-ubuntu2404
+
+# ---- ROCm ----
+
+docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-rocm
+docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base
+
+docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-rocm vllm/vllm-openai-rocm:latest
+docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-rocm vllm/vllm-openai-rocm:v${RELEASE_VERSION}
+docker push vllm/vllm-openai-rocm:latest
+docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}
+
+docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base vllm/vllm-openai-rocm:latest-base
+docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base
+docker push vllm/vllm-openai-rocm:latest-base
+docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base
+
+# ---- CPU ----
+# CPU images are behind separate block steps and may not have been built.
+# All-or-nothing: inspect both arches first, then either publish everything
+# (per-arch + multi-arch manifest) or skip everything. Publishing only one
+# arch would leave `:latest-x86_64` pointing at the new release while the
+# `:latest` multi-arch manifest still resolves to the previous release.
+
+CPU_X86_TAG=public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:v${RELEASE_VERSION}
+CPU_ARM_TAG=public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:v${RELEASE_VERSION}
+
+CPU_X86_AVAILABLE=false
+CPU_ARM_AVAILABLE=false
+docker manifest inspect "${CPU_X86_TAG}" >/dev/null 2>&1 && CPU_X86_AVAILABLE=true
+docker manifest inspect "${CPU_ARM_TAG}" >/dev/null 2>&1 && CPU_ARM_AVAILABLE=true
+
+if [ "$CPU_X86_AVAILABLE" = "true" ] && [ "$CPU_ARM_AVAILABLE" = "true" ]; then
+  docker pull "${CPU_X86_TAG}"
+  docker tag "${CPU_X86_TAG}" vllm/vllm-openai-cpu:latest-x86_64
+  docker tag "${CPU_X86_TAG}" vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64
+  docker push vllm/vllm-openai-cpu:latest-x86_64
+  docker push vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64
+
+  docker pull "${CPU_ARM_TAG}"
+  docker tag "${CPU_ARM_TAG}" vllm/vllm-openai-cpu:latest-arm64
+  docker tag "${CPU_ARM_TAG}" vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64
+  docker push vllm/vllm-openai-cpu:latest-arm64
+  docker push vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64
+
+  docker manifest rm vllm/vllm-openai-cpu:latest || true
+  docker manifest rm vllm/vllm-openai-cpu:v${RELEASE_VERSION} || true
+  docker manifest create vllm/vllm-openai-cpu:latest vllm/vllm-openai-cpu:latest-x86_64 vllm/vllm-openai-cpu:latest-arm64
+  docker manifest create vllm/vllm-openai-cpu:v${RELEASE_VERSION} vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64 vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64
+  docker manifest push vllm/vllm-openai-cpu:latest
+  docker manifest push vllm/vllm-openai-cpu:v${RELEASE_VERSION}
+elif [ "$CPU_X86_AVAILABLE" = "false" ] && [ "$CPU_ARM_AVAILABLE" = "false" ]; then
+  echo "WARNING: Neither CPU image found in ECR, skipping CPU publish (ensure block-cpu-release-image-build and block-arm64-cpu-release-image-build were unblocked and the builds finished pushing)"
+else
+  # Partial state: one arch built, the other did not. Fail loudly rather than
+  # ship a Docker Hub state where `:latest-${arch}` and `:latest` (multi-arch)
+  # disagree on which release they point at.
+  echo "ERROR: Partial CPU build detected (x86_64=${CPU_X86_AVAILABLE}, arm64=${CPU_ARM_AVAILABLE})."
+  echo "       Refusing to publish to avoid split-tag drift between per-arch and multi-arch tags."
+  echo "       Re-run the missing CPU build and retry, or manually publish if a single-arch release is intended."
+  exit 1
+fi
+
+echo ""
+echo "Successfully published release images for v${RELEASE_VERSION}"
@@ -51,6 +51,7 @@ vllm serve "$MODEL" \
  --offload-num-in-group 2 \
  --offload-prefetch-step 1 \
  --offload-params w13_weight w2_weight \
+  --generation-config vllm \
  --port "$PORT" \
  ${EXTRA_ARGS+"${EXTRA_ARGS[@]}"} &
 SERVER_PID=$!
@@ -39,10 +39,11 @@ fi

 set -x # avoid printing secrets above

-# install twine from pypi
+# install twine and sdist build prerequisites from pypi
 python3 -m venv /tmp/vllm-release-env
 source /tmp/vllm-release-env/bin/activate
 pip install twine
+pip install -r requirements/build/cuda.txt
 python3 -m twine --version

 # copy release wheels to local directory
@@ -230,7 +230,6 @@ steps:
  - tests/entrypoints/llm/test_collective_rpc.py
  - vllm/platforms/rocm.py
  commands:
-  - export TORCH_NCCL_BLOCKING_WAIT=1
  - pytest -v -s entrypoints/llm/test_collective_rpc.py
  - pytest -v -s ./compile/fullgraph/test_basic_correctness.py
  - pytest -v -s ./compile/test_wrapper.py
@@ -272,7 +271,6 @@ steps:
  - tests/v1/worker/test_worker_memory_snapshot.py
  - vllm/platforms/rocm.py
  commands:
-  - export TORCH_NCCL_BLOCKING_WAIT=1
  - 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'
  - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown
@@ -396,8 +394,8 @@ steps:
    - python3 pooling/embed/vision_embedding_offline.py --seed 0
    # Features demo
    - python3 features/automatic_prefix_caching/prefix_caching_offline.py
-    - python3 offline_inference/llm_engine_example.py
-    - python3 others/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 others/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
+    - python3 deployment/llm_engine_example.py
+    - python3 features/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 features/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
    - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
    - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536

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

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

@@ -590,7 +588,6 @@ steps:
  - vllm/platforms/rocm.py
  commands:
  - pip freeze | grep -E 'torch'
-  - export TORCH_NCCL_BLOCKING_WAIT=1
  - pytest -v -s models/language -m 'core_model and slow_test' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB

 - label: Multi-Modal Models (Extended Generation 2) # TBD
@@ -621,6 +618,7 @@ steps:
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
  agent_pool: mi250_1
  torch_nightly: true
+  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -864,7 +862,6 @@ steps:
  - tests/entrypoints/openai/test_multi_api_servers.py
  - vllm/platforms/rocm.py
  commands:
-  - export TORCH_NCCL_BLOCKING_WAIT=1
  - 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
@@ -883,7 +880,7 @@ steps:
  - vllm/platforms/rocm.py
  commands:
  - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
-  - ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh
+  - ATTENTION_BACKEND=ROCM_ATTN bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh

 - label: V1 e2e (2 GPUs) # TBD
  timeout_in_minutes: 180
@@ -930,7 +927,9 @@ steps:
  - tests/renderers
  - tests/standalone_tests/lazy_imports.py
  - tests/tokenizers_
+  - tests/reasoning
  - tests/tool_parsers
+  - tests/parser
  - tests/transformers_utils
  - tests/config
  commands:
@@ -942,8 +941,9 @@ steps:
  - pytest -v -s -m 'cpu_test' multimodal
  - pytest -v -s renderers
  - pytest -v -s tokenizers_
-  - pytest -v -s reasoning --ignore=reasoning/test_seedoss_reasoning_parser.py --ignore=reasoning/test_glm4_moe_reasoning_parser.py --ignore=reasoning/test_gemma4_reasoning_parser.py
+  - pytest -v -s reasoning --ignore=reasoning/test_seedoss_reasoning_parser.py --ignore=reasoning/test_glm4_moe_reasoning_parser.py
  - pytest -v -s tool_parsers
+  - pytest -v -s parser
  - pytest -v -s transformers_utils
  - pytest -v -s config

@@ -1102,12 +1102,13 @@ steps:
  - vllm/compilation/
  - vllm/model_executor/layers
  - tests/compile/passes/distributed/
+  - tests/compile/fusions_e2e/
  - vllm/_aiter_ops.py
  - vllm/platforms/rocm.py
  commands:
  - export VLLM_TEST_CLEAN_GPU_MEMORY=1
  - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/passes/distributed/test_async_tp.py
-  - pytest -v -s tests/compile/passes/distributed/test_sequence_parallelism.py
+  - pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py::test_tp2_ar_rms_fusions

 #-----------------------------------------------------------  mi300 · cuda  ------------------------------------------------------------#

@@ -1172,7 +1173,6 @@ steps:
  - vllm/_aiter_ops.py
  - vllm/platforms/rocm.py
  commands:
-  - export TORCH_NCCL_BLOCKING_WAIT=1
  - pytest -v -s tests/distributed/test_context_parallel.py
  - VLLM_LOGGING_LEVEL=DEBUG python3 examples/features/data_parallel/data_parallel_offline.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=allgather_reducescatter --disable-nccl-for-dp-synchronization

@@ -1186,7 +1186,6 @@ steps:
  source_file_dependencies:
  - vllm/
  commands:
-  - export TORCH_NCCL_BLOCKING_WAIT=1
  - pytest -v -s distributed/test_custom_all_reduce.py
  - torchrun --nproc_per_node=2 distributed/test_ca_buffer_sharing.py
  - TARGET_TEST_SUITE=A100 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
@@ -1206,7 +1205,6 @@ steps:
  - tests/examples/features/data_parallel/data_parallel_offline.py
  - vllm/platforms/rocm.py
  commands:
-  - export TORCH_NCCL_BLOCKING_WAIT=1
  - torchrun --nproc-per-node=4 distributed/test_torchrun_example.py
  - PP_SIZE=2 torchrun --nproc-per-node=4 distributed/test_torchrun_example.py
  - TP_SIZE=4 torchrun --nproc-per-node=4 distributed/test_torchrun_example_moe.py
@@ -1252,7 +1250,6 @@ steps:
  - vllm/platforms/rocm.py
  commands:
  - export VLLM_USE_RAY_V2_EXECUTOR_BACKEND=1
-  - export TORCH_NCCL_BLOCKING_WAIT=1
  - pytest -v -s distributed/test_ray_v2_executor.py
  - pytest -v -s distributed/test_ray_v2_executor_e2e.py
  - pytest -v -s distributed/test_pipeline_parallel.py -k "ray"
@@ -1274,7 +1271,6 @@ steps:
  - vllm/v1/worker/gpu_worker.py
  - vllm/platforms/rocm.py
  commands:
-  - export TORCH_NCCL_BLOCKING_WAIT=1
  - torchrun --nproc-per-node=8 ../examples/features/torchrun/torchrun_dp_example_offline.py --tp-size=2 --pp-size=1 --dp-size=4 --enable-ep

 #--------------------------------------------------------  mi300 · entrypoints  --------------------------------------------------------#
@@ -1326,7 +1322,6 @@ steps:
  commands:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
  - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
-  - pytest -v -s entrypoints/openai/speech_to_text/
  - pytest -v -s entrypoints/test_chat_utils.py

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

 - label: Entrypoints Integration (LLM) # TBD
  timeout_in_minutes: 180
@@ -1655,8 +1664,8 @@ steps:
    - python3 pooling/embed/vision_embedding_offline.py --seed 0
    # Features demo
    - python3 features/automatic_prefix_caching/prefix_caching_offline.py
-    - python3 offline_inference/llm_engine_example.py
-    - python3 others/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 others/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
+    - python3 deployment/llm_engine_example.py
+    - python3 features/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 features/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
    - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
    - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536

@@ -1766,7 +1775,7 @@ steps:
  - export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
  - pytest -v -s -x lora/test_chatglm3_tp.py
  - pytest -v -s -x lora/test_llama_tp.py
-  - pytest -v -s -x lora/test_llm_with_multi_loras.py
+  - pytest -v -s -x lora/test_qwen3_with_multi_loras.py
  - pytest -v -s -x lora/test_olmoe_tp.py
  - pytest -v -s -x lora/test_gptoss_tp.py
  - pytest -v -s -x lora/test_qwen35_densemodel_lora.py
@@ -1802,15 +1811,17 @@ steps:
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
+  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
  - tests/models/multimodal/generation
  - tests/models/multimodal/test_mapping.py
  commands:
-  - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
-  - pytest -v -s models/multimodal/generation -m 'not core_model' --ignore models/multimodal/generation/test_common.py
-  - pytest -v -s models/multimodal/test_mapping.py
+  - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@rocm-7.0-v2.3.0'
+  - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
+  - pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
+

 - label: Multi-Modal Models (Extended Generation 2) # TBD
  timeout_in_minutes: 180
@@ -1822,8 +1833,10 @@ steps:
  - vllm/
  - tests/models/multimodal/generation
  commands:
-  - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
-  - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=0) and not core_model'
+  - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@rocm-7.0-v2.3.0'
+  - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
+  - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'
+

 - label: Multi-Modal Models (Extended Generation 3) # TBD
  timeout_in_minutes: 180
@@ -1843,6 +1856,7 @@ steps:
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
  torch_nightly: true
+  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -2203,7 +2217,6 @@ steps:
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
-  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -2280,7 +2293,6 @@ steps:
  - tests/entrypoints/openai/test_multi_api_servers.py
  - vllm/platforms/rocm.py
  commands:
-  - export TORCH_NCCL_BLOCKING_WAIT=1
  - 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
@@ -2300,7 +2312,6 @@ steps:
  - vllm/_aiter_ops.py
  - vllm/platforms/rocm.py
  commands:
-  - export TORCH_NCCL_BLOCKING_WAIT=1
  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_async_new_apis.py
  - VLLM_LOGGING_LEVEL=DEBUG python3 examples/features/data_parallel/data_parallel_offline.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
  - pytest -v -s tests/v1/distributed/test_dbo.py
@@ -2363,7 +2374,6 @@ steps:
  - tests/distributed/test_utils
  - vllm/platforms/rocm.py
  commands:
-  - export TORCH_NCCL_BLOCKING_WAIT=1
  - 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
@@ -2493,7 +2503,6 @@ steps:
  - tests/entrypoints/llm/test_collective_rpc.py
  - vllm/platforms/rocm.py
  commands:
-  - export TORCH_NCCL_BLOCKING_WAIT=1
  - pytest -v -s entrypoints/llm/test_collective_rpc.py
  - pytest -v -s ./compile/fullgraph/test_basic_correctness.py
  - pytest -v -s ./compile/test_wrapper.py
@@ -2518,7 +2527,6 @@ steps:
  - tests/v1/worker/test_worker_memory_snapshot.py
  - vllm/platforms/rocm.py
  commands:
-  - export TORCH_NCCL_BLOCKING_WAIT=1
  - 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'
  - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown
@@ -2539,7 +2547,6 @@ steps:
  - tests/distributed/test_multiproc_executor.py
  - vllm/platforms/rocm.py
  commands:
-  - export TORCH_NCCL_BLOCKING_WAIT=1
  - pytest -v -s compile/fullgraph/test_basic_correctness.py
  - pytest -v -s distributed/test_pynccl.py
  - pytest -v -s distributed/test_events.py
@@ -2627,6 +2634,7 @@ steps:
  agent_pool: mi325_1
  torch_nightly: true
  parallelism: 2
+  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -2652,6 +2660,7 @@ steps:
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
  agent_pool: mi325_1
  torch_nightly: true
+  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -2717,7 +2726,6 @@ steps:
  - vllm/_aiter_ops.py
  - vllm/platforms/rocm.py
  commands:
-  - export TORCH_NCCL_BLOCKING_WAIT=1
  - pytest -v -s tests/distributed/test_context_parallel.py
  - pytest -v -s tests/v1/distributed/test_dbo.py

@@ -2748,6 +2756,7 @@ steps:
  agent_pool: mi355_1
  fast_check: true
  torch_nightly: true
+  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -2763,6 +2772,7 @@ steps:
  agent_pool: mi355_1
  fast_check: true
  torch_nightly: true
+  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -2771,7 +2781,6 @@ steps:
  commands:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
  - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
-  - pytest -v -s entrypoints/openai/speech_to_text/
  - pytest -v -s entrypoints/test_chat_utils.py

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

 - label: Entrypoints Integration (Pooling) # TBD
  timeout_in_minutes: 180
@@ -2938,8 +2961,8 @@ steps:
  - python3 pooling/embed/vision_embedding_offline.py --seed 0
  # Features demo
  - python3 features/automatic_prefix_caching/prefix_caching_offline.py
-  - python3 offline_inference/llm_engine_example.py
-  - python3 others/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 others/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
+  - python3 deployment/llm_engine_example.py
+  - python3 features/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 features/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
  - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
  - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536

@@ -3051,7 +3074,7 @@ steps:
  - vllm/
  - tests/models/language/generation
  commands:
-  - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr'
+  - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@rocm-7.0-v2.3.0'
  - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
  - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'

@@ -3059,6 +3082,7 @@ steps:
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
  agent_pool: mi355_1
+  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -3258,6 +3282,7 @@ steps:
  timeout_in_minutes: 60
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
  agent_pool: mi355_1
+  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -3283,6 +3308,7 @@ steps:
  timeout_in_minutes: 60
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
  agent_pool: mi355_1
+  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -3323,7 +3349,7 @@ steps:
  - vllm/platforms/rocm.py
  commands:
  - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
-  - ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh
+  - ATTENTION_BACKEND=ROCM_ATTN bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh

 - label: Distributed NixlConnector PD accuracy (4 GPUs) # TBD
  timeout_in_minutes: 180
@@ -17,7 +17,7 @@ steps:
 - label: V1 attention (B200)
  key: v1-attention-b200
  timeout_in_minutes: 30
-  device: b200
+  device: b200-k8s
  source_file_dependencies:
    - vllm/config/attention.py
    - vllm/model_executor/layers/attention
@@ -14,7 +14,7 @@ steps:

 - label: Attention Benchmarks Smoke Test (B200)
  key: attention-benchmarks-smoke-test-b200
-  device: b200
+  device: b200-k8s
  num_gpus: 2
  optional: true
  working_dir: "/vllm-workspace/"
@@ -43,7 +43,7 @@ steps:
  key: asynctp-correctness-tests-b200
  timeout_in_minutes: 50
  working_dir: "/vllm-workspace/"
-  device: b200
+  device: b200-k8s
  optional: true
  num_devices: 2
  commands:
@@ -68,7 +68,7 @@ steps:
  key: fusion-and-compile-unit-tests-2xb200
  timeout_in_minutes: 20
  working_dir: "/vllm-workspace/"
-  device: b200
+  device: b200-k8s
  source_file_dependencies:
  - csrc/quantization/fp4/
  - vllm/model_executor/layers/quantization/
@@ -137,7 +137,7 @@ steps:
  key: fusion-e2e-config-sweep-b200
  timeout_in_minutes: 30
  working_dir: "/vllm-workspace/"
-  device: b200
+  device: b200-k8s
  num_devices: 1
  optional: true
  commands:
@@ -209,7 +209,7 @@ steps:
  key: fusion-e2e-tp2-b200
  timeout_in_minutes: 20
  working_dir: "/vllm-workspace/"
-  device: b200
+  device: b200-k8s
  num_devices: 2
  source_file_dependencies:
    - csrc/quantization/
@@ -7,7 +7,11 @@ steps:
  timeout_in_minutes: 15
  device: h200_18gb
  source_file_dependencies:
-  - vllm/
+  - vllm/envs.py
+  - vllm/logger.py
+  - vllm/platforms/
+  - vllm/plugins/
+  - vllm/utils/
  - tests/cuda
  commands:
    - pytest -v -s cuda/test_cuda_context.py
@@ -16,6 +20,7 @@ steps:
 - label: Cudagraph
  key: cudagraph
  timeout_in_minutes: 20
+  device: h200_18gb
  source_file_dependencies:
  - tests/v1/cudagraph
  - vllm/v1/cudagraph_dispatcher.py
@@ -8,7 +8,7 @@ steps:
  working_dir: "/vllm-workspace/tests"
  num_devices: 4
  source_file_dependencies:
-    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+    - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
    - tests/v1/kv_connector/nixl_integration/
  commands:
    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
@@ -19,7 +19,7 @@ steps:
  working_dir: "/vllm-workspace/tests"
  num_devices: 4
  source_file_dependencies:
-    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+    - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
    - tests/v1/kv_connector/nixl_integration/
  commands:
    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
@@ -31,7 +31,7 @@ steps:
  working_dir: "/vllm-workspace/tests"
  num_devices: 4
  source_file_dependencies:
-    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+    - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
    - tests/v1/kv_connector/nixl_integration/
  commands:
    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
@@ -43,7 +43,7 @@ steps:
  working_dir: "/vllm-workspace/tests"
  num_devices: 4
  source_file_dependencies:
-    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+    - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
    - tests/v1/kv_connector/nixl_integration/
  commands:
    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
@@ -55,7 +55,7 @@ steps:
  working_dir: "/vllm-workspace/tests"
  num_devices: 4
  source_file_dependencies:
-    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+    - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
    - tests/v1/kv_connector/nixl_integration/
  commands:
    - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
@@ -67,7 +67,7 @@ steps:
  working_dir: "/vllm-workspace/tests"
  num_devices: 2
  source_file_dependencies:
-    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+    - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
    - vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
    - vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
    - vllm/distributed/kv_transfer/kv_connector/v1/offloading/
@@ -83,7 +83,7 @@ steps:
  working_dir: "/vllm-workspace/tests"
  num_devices: 2
  source_file_dependencies:
-    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+    - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
    - vllm/v1/worker/kv_connector_model_runner_mixin.py
    - tests/v1/kv_connector/nixl_integration/
  commands:
@@ -96,7 +96,7 @@ steps:
  working_dir: "/vllm-workspace/tests"
  num_devices: 2
  source_file_dependencies:
-    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+    - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
    - vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
    - vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
    - vllm/distributed/kv_transfer/kv_connector/v1/offloading/
@@ -212,7 +212,7 @@ steps:

 - label: Distributed Tests (2 GPUs)(B200)
  key: distributed-tests-2-gpus-b200
-  device: b200
+  device: b200-k8s
  optional: true
  working_dir: "/vllm-workspace/"
  num_devices: 2
@@ -25,7 +25,7 @@ steps:
 - label: Qwen3-30B-A3B-FP8-block Accuracy (B200)
  key: qwen3-30b-a3b-fp8-block-accuracy-b200
  timeout_in_minutes: 60
-  device: b200
+  device: b200-k8s
  optional: true
  num_devices: 2
  working_dir: "/vllm-workspace"
@@ -7,18 +7,30 @@ steps:
  timeout_in_minutes: 15
  device: h200_18gb
  source_file_dependencies:
-  - vllm/
+  - vllm/compilation/
+  - vllm/config/
+  - vllm/engine/
+  - vllm/entrypoints/logger.py
+  - vllm/envs.py
+  - vllm/logger.py
+  - vllm/logging_utils/
+  - vllm/platforms/
+  - vllm/sequence.py
+  - vllm/triton_utils/
+  - vllm/utils/
  - tests/engine
  - tests/test_sequence
  - tests/test_config
  - tests/test_logger
  - tests/test_vllm_port
+  - tests/test_jit_monitor.py
  commands:
-  - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py
+  - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py test_jit_monitor.py

 - label: Engine (1 GPU)
  key: engine-1-gpu
  timeout_in_minutes: 30
+  device: h200_18gb
  source_file_dependencies:
    - vllm/v1/engine/
    - tests/v1/engine/
@@ -39,6 +51,7 @@ steps:
 - label: e2e Core (1 GPU)
  key: e2e-core-1-gpu
  timeout_in_minutes: 30
+  device: h200_18gb
  source_file_dependencies:
    - vllm/v1/
    - tests/v1/e2e/general/
@@ -51,16 +64,27 @@ steps:
  optional: true
  num_devices: 2
  source_file_dependencies:
-    - vllm/
-    - tests/v1/e2e
+    - vllm/compilation/
+    - vllm/config/
+    - vllm/distributed/
+    - vllm/engine/
+    - vllm/envs.py
+    - vllm/forward_context.py
+    - vllm/inputs/
+    - vllm/logger.py
+    - vllm/logging_utils/
+    - vllm/model_executor/
+    - vllm/multimodal/
+    - vllm/platforms/
+    - vllm/sampling_params.py
+    - vllm/transformers_utils/
+    - vllm/triton_utils/
+    - vllm/utils/
+    - vllm/v1/
+    - tests/v1/e2e/spec_decode
  commands:
    # Only run tests that need exactly 2 GPUs
    - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "tensor_parallelism"
-  mirror:
-    amd:
-      device: mi325_2
-      depends_on:
-      - image-build-amd

 - label: V1 e2e (4 GPUs)
  key: v1-e2e-4-gpus
@@ -68,16 +92,27 @@ steps:
  optional: true
  num_devices: 4
  source_file_dependencies:
-    - vllm/
-    - tests/v1/e2e
+    - vllm/compilation/
+    - vllm/config/
+    - vllm/distributed/
+    - vllm/engine/
+    - vllm/envs.py
+    - vllm/forward_context.py
+    - vllm/inputs/
+    - vllm/logger.py
+    - vllm/logging_utils/
+    - vllm/model_executor/
+    - vllm/multimodal/
+    - vllm/platforms/
+    - vllm/sampling_params.py
+    - vllm/transformers_utils/
+    - vllm/triton_utils/
+    - vllm/utils/
+    - vllm/v1/
+    - tests/v1/e2e/spec_decode
  commands:
    # Only run tests that need 4 GPUs
    - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "eagle_correctness_heavy"
-  mirror:
-    amd:
-      device: mi325_4
-      depends_on:
-      - image-build-amd

 - label: V1 e2e (4xH100)
  key: v1-e2e-4xh100
@@ -6,17 +6,19 @@ steps:
  key: entrypoints-unit-tests
  timeout_in_minutes: 10
  working_dir: "/vllm-workspace/tests"
+  device: h200_18gb
  source_file_dependencies:
  - vllm/entrypoints
  - tests/entrypoints/
  commands:
  - pytest -v -s entrypoints/openai/tool_parsers
-  - pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/rpc --ignore=entrypoints/sleep --ignore=entrypoints/serve/instrumentator --ignore=entrypoints/openai --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py  --ignore=entrypoints/pooling
+  - pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/rpc --ignore=entrypoints/sleep --ignore=entrypoints/serve/instrumentator --ignore=entrypoints/openai --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py  --ignore=entrypoints/pooling --ignore=entrypoints/speech_to_text

 - label: Entrypoints Integration (LLM)
  key: entrypoints-integration-llm
  timeout_in_minutes: 40
  working_dir: "/vllm-workspace/tests"
+  device: h200_18gb
  source_file_dependencies:
  - vllm/
  - tests/entrypoints/llm
@@ -26,11 +28,17 @@ steps:
  - pytest -v -s entrypoints/llm --ignore=entrypoints/llm/test_generate.py --ignore=entrypoints/llm/test_collective_rpc.py
  - pytest -v -s entrypoints/llm/test_generate.py # it needs a clean process
  - pytest -v -s entrypoints/offline_mode # Needs to avoid interference with other tests
+  mirror:
+    amd:
+      device: mi300_1
+      depends_on:
+      - image-build-amd

 - label: Entrypoints Integration (API Server openai - Part 1)
  key: entrypoints-integration-api-server-openai-part-1
  timeout_in_minutes: 50
  working_dir: "/vllm-workspace/tests"
+  device: h200_18gb
  source_file_dependencies:
  - vllm/
  - tests/entrypoints/openai
@@ -38,30 +46,19 @@ steps:
  commands:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
  - pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py
-  mirror:
-    amd:
-      device: mi325_1
-      depends_on:
-      - image-build-amd
-

 - label: Entrypoints Integration (API Server openai - Part 2)
  key: entrypoints-integration-api-server-openai-part-2
  timeout_in_minutes: 50
  working_dir: "/vllm-workspace/tests"
+  device: h200_18gb
  source_file_dependencies:
  - vllm/
  - tests/entrypoints/openai
  - tests/entrypoints/test_chat_utils
  commands:
  - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
-  - pytest -v -s entrypoints/openai/speech_to_text/
  - pytest -v -s entrypoints/test_chat_utils.py
-  mirror:
-    amd:
-      device: mi325_1
-      depends_on:
-      - image-build-amd

 - label: Entrypoints Integration (API Server openai - Part 3)
  key: entrypoints-integration-api-server-openai-part-3
@@ -74,12 +71,13 @@ steps:
  - tests/entrypoints/test_chat_utils
  commands:
  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/speech_to_text/ --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/test_multi_api_servers.py
+  - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/completion  --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/test_multi_api_servers.py

 - label: Entrypoints Integration (API Server 2)
  key: entrypoints-integration-api-server-2
  timeout_in_minutes: 130
  working_dir: "/vllm-workspace/tests"
+  device: h200_18gb
  source_file_dependencies:
  - vllm/
  - tests/entrypoints/rpc
@@ -91,10 +89,23 @@ steps:
  - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc
  - pytest -v -s tool_use

+- label: Entrypoints Integration (Speech to Text)
+  key: entrypoints-integration-speech_to_text
+  timeout_in_minutes: 50
+  working_dir: "/vllm-workspace/tests"
+  device: h200_18gb
+  source_file_dependencies:
+  - vllm/
+  - tests/entrypoints/speech_to_text
+  commands:
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  - pytest -v -s entrypoints/speech_to_text
+
 - label: Entrypoints Integration (Pooling)
  key: entrypoints-integration-pooling
  timeout_in_minutes: 50
  working_dir: "/vllm-workspace/tests"
+  device: h200_18gb
  source_file_dependencies:
  - vllm/
  - tests/entrypoints/pooling
@@ -106,6 +117,7 @@ steps:
  key: entrypoints-integration-responses-api
  timeout_in_minutes: 50
  working_dir: "/vllm-workspace/tests"
+  device: h200_18gb
  source_file_dependencies:
  - vllm/
  - tests/entrypoints/openai/responses
@@ -120,5 +132,5 @@ steps:
  - csrc/
  - vllm/entrypoints/openai/
  - vllm/model_executor/models/whisper.py
-  commands: # LMEval+Transcription WER check
+  commands: # LMEval
  - pytest -s entrypoints/openai/correctness/
@@ -17,6 +17,7 @@ steps:
 - label: Kernels Core Operation Test
  key: kernels-core-operation-test
  timeout_in_minutes: 75
+  device: h200_18gb
  source_file_dependencies:
  - csrc/
  - tests/kernels/core
@@ -41,6 +42,7 @@ steps:
 - label: Kernels Attention Test %N
  key: kernels-attention-test
  timeout_in_minutes: 35
+  device: h200_18gb
  source_file_dependencies:
  - csrc/attention/
  - vllm/v1/attention
@@ -55,6 +57,7 @@ steps:
 - label: Kernels Quantization Test %N
  key: kernels-quantization-test
  timeout_in_minutes: 90
+  device: h200_18gb
  source_file_dependencies:
  - csrc/quantization/
  - vllm/model_executor/layers/quantization
@@ -66,6 +69,7 @@ steps:
 - label: Kernels MoE Test %N
  key: kernels-moe-test
  timeout_in_minutes: 25
+  device: h200_18gb
  source_file_dependencies:
  - csrc/quantization/cutlass_w8a8/moe/
  - csrc/moe/
@@ -82,6 +86,7 @@ steps:
 - label: Kernels Mamba Test
  key: kernels-mamba-test
  timeout_in_minutes: 45
+  device: h200_18gb
  source_file_dependencies:
  - csrc/mamba/
  - tests/kernels/mamba
@@ -91,6 +96,7 @@ steps:

 - label: Kernels KDA Test
  timeout_in_minutes: 20
+  device: h200_18gb
  source_file_dependencies:
  - vllm/model_executor/layers/fla/ops/kda.py
  - vllm/model_executor/layers/fla/ops/chunk_delta_h.py
@@ -125,7 +131,7 @@ steps:
  key: kernels-b200
  timeout_in_minutes: 30
  working_dir: "/vllm-workspace/"
-  device: b200
+  device: b200-k8s
  # optional: true
  source_file_dependencies:
  - csrc/quantization/fp4/
@@ -212,7 +218,7 @@ steps:
 - label: Kernels Fp4 MoE Test (B200)
  key: kernels-fp4-moe-test-b200
  timeout_in_minutes: 60
-  device: b200
+  device: b200-k8s
  num_devices: 1
  optional: true
  commands:
@@ -242,7 +248,7 @@ steps:
 - label: Kernels FusedMoE Layer Test (2 B200s)
  key: kernels-fusedmoe-layer-test-2-b200s
  timeout_in_minutes: 90
-  device: b200
+  device: b200-k8s
  num_devices: 2
  source_file_dependencies:
  - csrc/quantization/cutlass_w8a8/moe/
@@ -5,6 +5,7 @@ steps:
 - label: LM Eval Small Models
  key: lm-eval-small-models
  timeout_in_minutes: 75
+  device: h200_18gb
  source_file_dependencies:
  - csrc/
  - vllm/model_executor/layers/quantization
@@ -40,7 +41,7 @@ steps:
 - label: LM Eval Small Models (B200)
  key: lm-eval-small-models-b200
  timeout_in_minutes: 120
-  device: b200
+  device: b200-k8s
  optional: true
  source_file_dependencies:
  - csrc/
@@ -48,10 +49,22 @@ steps:
  commands:
  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-blackwell.txt

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

 - label: MoE Refactor Integration Test (B200 DP - TEMPORARY)
  key: moe-refactor-integration-test-b200-dp-temporary
-  device: b200
+  device: b200-k8s
  optional: true
  num_devices: 2
  commands:
@@ -102,6 +115,7 @@ steps:
 - label: LM Eval TurboQuant KV Cache
  key: lm-eval-turboquant-kv-cache
  timeout_in_minutes: 75
+  device: h200_18gb
  source_file_dependencies:
  - vllm/model_executor/layers/quantization/turboquant/
  - vllm/v1/attention/backends/turboquant_attn.py
@@ -127,7 +141,7 @@ steps:
 - label: GPQA Eval (GPT-OSS) (B200)
  key: gpqa-eval-gpt-oss-b200
  timeout_in_minutes: 120
-  device: b200
+  device: b200-k8s
  optional: true
  num_devices: 2
  source_file_dependencies:
@@ -137,3 +151,11 @@ steps:
  commands:
    - uv pip install --system 'gpt-oss[eval]==0.0.5'
    - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-b200.txt
+
+- label: MRCR Eval Small Models
+  timeout_in_minutes: 30
+  device: h200_18gb
+  source_file_dependencies:
+  - tests/evals/mrcr/
+  commands:
+    - pytest -s -v evals/mrcr/test_mrcr_correctness.py --config-list-file=evals/mrcr/configs/models-small.txt
@@ -5,11 +5,12 @@ steps:
 - label: LoRA %N
  key: lora
  timeout_in_minutes: 30
+  device: h200_18gb
  source_file_dependencies:
  - vllm/lora
  - tests/lora
  commands:
-    - pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py --ignore=lora/test_llm_with_multi_loras.py --ignore=lora/test_olmoe_tp.py --ignore=lora/test_deepseekv2_tp.py --ignore=lora/test_gptoss_tp.py --ignore=lora/test_qwen3moe_tp.py --ignore=lora/test_qwen35_densemodel_lora.py 
+    - pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py --ignore=lora/test_qwen3_with_multi_loras.py --ignore=lora/test_olmoe_tp.py --ignore=lora/test_deepseekv2_tp.py --ignore=lora/test_gptoss_tp.py --ignore=lora/test_qwen3moe_tp.py --ignore=lora/test_qwen35_densemodel_lora.py 
  parallelism: 4


@@ -19,6 +20,7 @@ steps:
  num_devices: 4
  source_file_dependencies:
  - vllm/lora
+  - vllm/model_executor/layers/fused_moe/
  - tests/lora
  commands:
    # FIXIT: find out which code initialize cuda before running the test
@@ -30,7 +32,7 @@ steps:
    # requires multi-GPU testing for validation.
    - pytest -v -s -x lora/test_chatglm3_tp.py
    - pytest -v -s -x lora/test_llama_tp.py
-    - pytest -v -s -x lora/test_llm_with_multi_loras.py
+    - pytest -v -s -x lora/test_qwen3_with_multi_loras.py
    - pytest -v -s -x lora/test_olmoe_tp.py
    - pytest -v -s -x lora/test_gptoss_tp.py
    - pytest -v -s -x lora/test_qwen35_densemodel_lora.py
@@ -5,25 +5,39 @@ steps:
 - label: V1 Spec Decode
  key: v1-spec-decode
  timeout_in_minutes: 30
+  device: h200_18gb
  source_file_dependencies:
-    - vllm/
+    - vllm/config/
+    - vllm/distributed/
+    - vllm/inputs/
+    - vllm/model_executor/
+    - vllm/platforms/
+    - vllm/sampling_params.py
+    - vllm/transformers_utils/
+    - vllm/utils/
+    - vllm/v1/
    - tests/v1/spec_decode
  commands:
    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
    # TODO: create another `optional` test group for slow tests
    - pytest -v -s -m 'not slow_test' v1/spec_decode
-  mirror:
-    amd:
-      device: mi325_1
-      depends_on:
-      - image-build-amd

 - label: V1 Sample + Logits
  key: v1-sample-logits
  timeout_in_minutes: 30
  device: h200_18gb
  source_file_dependencies:
-    - vllm/
+    - vllm/config/
+    - vllm/distributed/
+    - vllm/engine/
+    - vllm/inputs/
+    - vllm/logger.py
+    - vllm/model_executor/
+    - vllm/platforms/
+    - vllm/sampling_params.py
+    - vllm/transformers_utils/
+    - vllm/utils/
+    - vllm/v1/
    - tests/v1/sample
    - tests/v1/logits_processors
    - tests/v1/test_oracle.py
@@ -38,15 +52,32 @@ steps:
    - pytest -v -s v1/test_outputs.py
  mirror:
    amd:
-      device: mi325_1
+      device: mi300_1
      depends_on:
      - image-build-amd

 - label: V1 Core + KV + Metrics
  key: v1-core-kv-metrics
  timeout_in_minutes: 30
+  device: h200_18gb
  source_file_dependencies:
-    - vllm/
+    - vllm/config/
+    - vllm/distributed/
+    - vllm/engine/
+    - vllm/entrypoints/pooling/
+    - vllm/inputs/
+    - vllm/lora/
+    - vllm/model_executor/
+    - vllm/multimodal/
+    - vllm/outputs.py
+    - vllm/platforms/
+    - vllm/pooling_params.py
+    - vllm/profiler/
+    - vllm/sampling_params.py
+    - vllm/tokenizers/
+    - vllm/transformers_utils/
+    - vllm/utils/
+    - vllm/v1/
    - tests/v1/core
    - tests/v1/executor
    - tests/v1/kv_offload
@@ -67,18 +98,27 @@ steps:
    # Integration test for streaming correctness (requires special branch).
    - pip install -U git+https://github.com/robertgshaw2-redhat/lm-evaluation-harness.git@streaming-api
    - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
-  mirror:
-    amd:
-      device: mi325_1
-      depends_on:
-      - image-build-amd

 - label: V1 Others (CPU)
  key: v1-others-cpu
  depends_on:
    - image-build-cpu
  source_file_dependencies:
-    - vllm/
+    - vllm/config/
+    - vllm/distributed/
+    - vllm/engine/
+    - vllm/inputs/
+    - vllm/lora/
+    - vllm/multimodal/
+    - vllm/outputs.py
+    - vllm/platforms/
+    - vllm/pooling_params.py
+    - vllm/profiler/
+    - vllm/sampling_params.py
+    - vllm/tokenizers/
+    - vllm/transformers_utils/
+    - vllm/utils/
+    - vllm/v1/
    - tests/v1
  device: cpu-small
  commands:
@@ -94,7 +134,17 @@ steps:
  timeout_in_minutes: 20
  device: h200_18gb
  source_file_dependencies:
-  - vllm/
+  - vllm/config/
+  - vllm/distributed/
+  - vllm/engine/
+  - vllm/inputs/
+  - vllm/model_executor/
+  - vllm/multimodal/
+  - vllm/platforms/
+  - vllm/sampling_params.py
+  - vllm/transformers_utils/
+  - vllm/utils/
+  - vllm/v1/
  - tests/test_regression
  commands:
  - pip install modelscope
@@ -105,6 +155,7 @@ steps:
  key: examples
  timeout_in_minutes: 45
  working_dir: "/vllm-workspace/examples"
+  device: h200_18gb
  source_file_dependencies:
  - vllm/entrypoints
  - vllm/multimodal
@@ -127,8 +178,8 @@ steps:
    - python3 pooling/embed/vision_embedding_offline.py --seed 0
    # for features demo
    - python3 features/automatic_prefix_caching/prefix_caching_offline.py
-    - python3 offline_inference/llm_engine_example.py
-    - python3 others/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 others/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
+    - python3 deployment/llm_engine_example.py
+    - python3 features/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 features/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
    - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
    # https://github.com/vllm-project/vllm/pull/26682 uses slightly more memory in PyTorch 2.9+ causing this test to OOM in 1xL4 GPU
    - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536
@@ -138,7 +189,18 @@ steps:
  timeout_in_minutes: 20
  num_devices: 2
  source_file_dependencies:
-  - vllm/
+  - vllm/config/
+  - vllm/distributed/
+  - vllm/engine/
+  - vllm/inputs/
+  - vllm/model_executor/
+  - vllm/multimodal/
+  - vllm/platforms/
+  - vllm/sampling_params.py
+  - vllm/tracing/
+  - vllm/transformers_utils/
+  - vllm/utils/
+  - vllm/v1/
  - tests/v1/tracing
  commands:
  - "pip install \
@@ -151,7 +213,9 @@ steps:
 - label: Python-only Installation
  key: python-only-installation
  depends_on: ~
+  optional: true
  timeout_in_minutes: 20
+  device: h200_18gb
  source_file_dependencies:
  - tests/standalone_tests/python_only_compile.sh
  - setup.py
@@ -161,8 +225,21 @@ steps:
 - label: Async Engine, Inputs, Utils, Worker
  key: async-engine-inputs-utils-worker
  timeout_in_minutes: 50
+  device: h200_18gb
  source_file_dependencies:
-  - vllm/
+  - vllm/assets/
+  - vllm/config/
+  - vllm/distributed/
+  - vllm/engine/
+  - vllm/inputs/
+  - vllm/model_executor/
+  - vllm/multimodal/
+  - vllm/platforms/
+  - vllm/sampling_params.py
+  - vllm/tokenizers/
+  - vllm/transformers_utils/
+  - vllm/utils/
+  - vllm/v1/
  - tests/detokenizer
  - tests/multimodal
  - tests/utils_
@@ -177,7 +254,30 @@ steps:
  - image-build-cpu
  timeout_in_minutes: 30
  source_file_dependencies:
-  - vllm/
+  - vllm/assets/
+  - vllm/config/
+  - vllm/engine/arg_utils.py
+  - vllm/entrypoints/chat_utils.py
+  - vllm/entrypoints/mcp/
+  - vllm/entrypoints/openai/chat_completion/protocol.py
+  - vllm/entrypoints/openai/engine/protocol.py
+  - vllm/envs.py
+  - vllm/exceptions.py
+  - vllm/inputs/
+  - vllm/model_executor/layers/quantization/quark/
+  - vllm/multimodal/
+  - vllm/outputs.py
+  - vllm/platforms/
+  - vllm/pooling_params.py
+  - vllm/ray/
+  - vllm/reasoning/
+  - vllm/renderers/
+  - vllm/sampling_params.py
+  - vllm/tokenizers/
+  - vllm/tool_parsers/
+  - vllm/transformers_utils/
+  - vllm/utils/
+  - vllm/v1/
  - tests/test_inputs.py
  - tests/test_outputs.py
  - tests/test_pooling_params.py
@@ -188,6 +288,7 @@ steps:
  - tests/tokenizers_
  - tests/reasoning
  - tests/tool_parsers
+  - tests/parser
  - tests/transformers_utils
  - tests/config
  device: cpu-small
@@ -200,8 +301,9 @@ steps:
  - pytest -v -s -m 'cpu_test' multimodal
  - pytest -v -s renderers
  - pytest -v -s tokenizers_
-  - pytest -v -s reasoning --ignore=reasoning/test_seedoss_reasoning_parser.py --ignore=reasoning/test_glm4_moe_reasoning_parser.py --ignore=reasoning/test_gemma4_reasoning_parser.py
+  - pytest -v -s reasoning --ignore=reasoning/test_seedoss_reasoning_parser.py --ignore=reasoning/test_glm4_moe_reasoning_parser.py
  - pytest -v -s tool_parsers
+  - pytest -v -s parser
  - pytest -v -s transformers_utils
  - pytest -v -s config

@@ -224,7 +326,7 @@ steps:
 - label: Batch Invariance (B200)
  key: batch-invariance-b200
  timeout_in_minutes: 30
-  device: b200
+  device: b200-k8s
  source_file_dependencies:
    - vllm/v1/attention
    - vllm/model_executor/layers
@@ -5,6 +5,7 @@ steps:
 - label: Model Executor
  key: model-executor
  timeout_in_minutes: 35
+  device: h200_18gb
  source_file_dependencies:
  - vllm/engine/arg_utils.py
  - vllm/config/model.py
@@ -5,6 +5,7 @@ steps:
 - label: Model Runner V2 Core Tests
  key: model-runner-v2-core-tests
  timeout_in_minutes: 45
+  device: h200_18gb
  source_file_dependencies:
  - vllm/v1/worker/gpu/
  - vllm/v1/worker/gpu_worker.py
@@ -29,6 +30,7 @@ steps:
  key: model-runner-v2-examples
  timeout_in_minutes: 45
  working_dir: "/vllm-workspace/examples"
+  device: h200_18gb
  source_file_dependencies:
    - vllm/v1/worker/gpu/
    - vllm/v1/core/sched/
@@ -37,7 +39,7 @@ steps:
    - examples/generate/multimodal/
    - examples/features/
    - examples/pooling/embed/vision_embedding_offline.py
-    - examples/others/tensorize_vllm_model.py
+    - examples/features/tensorize_vllm_model.py
  commands:
    - set -x
    - export VLLM_USE_V2_MODEL_RUNNER=1
@@ -55,8 +57,8 @@ steps:
    - python3 pooling/embed/vision_embedding_offline.py --seed 0
    # for features demo
    - python3 features/automatic_prefix_caching/prefix_caching_offline.py
-    - python3 offline_inference/llm_engine_example.py
-    - python3 others/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 others/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
+    - python3 deployment/llm_engine_example.py
+    - python3 features/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 features/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
    - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
    # https://github.com/vllm-project/vllm/pull/26682 uses slightly more memory in PyTorch 2.9+ causing this test to OOM in 1xL4 GPU
    - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536
@@ -102,6 +104,7 @@ steps:
  key: model-runner-v2-spec-decode
  timeout_in_minutes: 30
  working_dir: "/vllm-workspace/tests"
+  device: h200_18gb
  source_file_dependencies:
  - vllm/v1/worker/gpu/
  - vllm/v1/worker/gpu_worker.py
@@ -6,6 +6,7 @@ steps:
  key: basic-models-tests-initialization
  timeout_in_minutes: 45
  torch_nightly: true
+  device: h200_18gb
  source_file_dependencies:
  - vllm/
  - tests/models/test_initialization.py
@@ -19,6 +20,7 @@ steps:
 - label: Basic Models Tests (Extra Initialization) %N
  key: basic-models-tests-extra-initialization
  timeout_in_minutes: 45
+  device: h200_18gb
  source_file_dependencies:
  - vllm/model_executor/models/
  - tests/models/test_initialization.py
@@ -35,6 +37,7 @@ steps:
 - label: Basic Models Tests (Other)
  key: basic-models-tests-other
  timeout_in_minutes: 45
+  device: h200_18gb
  source_file_dependencies:
  - vllm/
  - tests/models/test_terratorch.py
@@ -42,12 +45,6 @@ steps:
  - tests/models/test_registry.py
  commands:
    - pytest -v -s models/test_terratorch.py models/test_transformers.py models/test_registry.py
-  mirror:
-    amd:
-      device: mi325_1
-      depends_on:
-      - image-build-amd
-

 - label: Basic Models Test (Other CPU) # 5min
  key: basic-models-test-other-cpu
@@ -67,6 +64,7 @@ steps:
  working_dir: "/vllm-workspace/"
  optional: true
  soft_fail: true
+  device: h200_18gb
  commands:
    - pip install --upgrade git+https://github.com/huggingface/transformers
    - pytest -v -s tests/models/test_initialization.py
@@ -83,6 +81,7 @@ steps:
  working_dir: "/vllm-workspace/"
  optional: true
  soft_fail: true
+  device: h200_18gb
  commands:
    - pip install transformers==4.57.5
    - pytest -v -s tests/models/test_initialization.py
@@ -5,6 +5,7 @@ steps:
 - label: Language Models Tests (Standard)
  key: language-models-tests-standard
  timeout_in_minutes: 25
+  device: h200_18gb
  source_file_dependencies:
  - vllm/
  - tests/models/language
@@ -18,6 +19,7 @@ steps:
 - label: Language Models Tests (Extra Standard) %N
  key: language-models-tests-extra-standard
  timeout_in_minutes: 45
+  device: h200_18gb
  source_file_dependencies:
  - vllm/model_executor/models/
  - tests/models/language/pooling/test_embedding.py
@@ -35,6 +37,7 @@ steps:
 - label: Language Models Tests (Hybrid) %N
  key: language-models-tests-hybrid
  timeout_in_minutes: 75
+  device: h200_18gb
  source_file_dependencies:
  - vllm/
  - tests/models/language/generation
@@ -48,10 +51,19 @@ steps:
  parallelism: 2
  mirror:
    torch_nightly: {}
+    amd:
+      device: mi300_1
+      depends_on:
+      - image-build-amd
+      commands:
+      - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr'
+      - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
+      - pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB

 - label: Language Models Test (Extended Generation) # 80min
  key: language-models-test-extended-generation
  timeout_in_minutes: 110
+  device: h200_18gb
  optional: true
  source_file_dependencies:
  - vllm/
@@ -62,15 +74,6 @@ steps:
    - uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
    - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
    - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'
-  mirror:
-    amd:
-      device: mi325_1
-      depends_on:
-      - image-build-amd
-      commands:
-      - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr'
-      - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
-      - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'

 - label: Language Models Test (PPL)
  key: language-models-test-ppl
@@ -86,17 +89,13 @@ steps:
 - label: Language Models Test (Extended Pooling)  # 36min
  key: language-models-test-extended-pooling
  timeout_in_minutes: 50
+  device: h200_18gb
  optional: true
  source_file_dependencies:
  - vllm/
  - tests/models/language/pooling
  commands:
    - pytest -v -s models/language/pooling -m 'not core_model'
-  mirror:
-    amd:
-      device: mi325_1
-      depends_on:
-      - image-build-amd

 - label: Language Models Test (MTEB)
  key: language-models-test-mteb
@@ -15,7 +15,7 @@ steps:
    - pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model
  mirror:
    amd:
-      device: mi325_1
+      device: mi300_1
      depends_on:
      - image-build-amd

@@ -33,13 +33,14 @@ steps:
    - pytest -v -s models/multimodal/generation/test_vit_cudagraph.py -m core_model
  mirror:
    amd:
-      device: mi325_1
+      device: mi300_1
      depends_on:
      - image-build-amd

 - label: "Multi-Modal Models (Standard) 3: llava + qwen2_vl"
  key: multi-modal-models-standard-3-llava-qwen2-vl
  timeout_in_minutes: 45
+  device: h200_18gb
  source_file_dependencies:
  - vllm/
  - tests/models/multimodal
@@ -49,13 +50,14 @@ steps:
    - pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model
  mirror:
    amd:
-      device: mi325_1
+      device: mi300_1
      depends_on:
      - image-build-amd

 - label: "Multi-Modal Models (Standard) 4: other + whisper"
  key: multi-modal-models-standard-4-other-whisper
  timeout_in_minutes: 45
+  device: h200_18gb
  source_file_dependencies:
  - vllm/
  - tests/models/multimodal
@@ -64,11 +66,6 @@ steps:
    - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py  --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/processing
    - pytest models/multimodal/generation/test_memory_leak.py -m core_model
    - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model  # Otherwise, mp_method="spawn" doesn't work
-  mirror:
-    amd:
-      device: mi325_1
-      depends_on:
-      - image-build-amd

 - label: Multi-Modal Processor (CPU)
  key: multi-modal-processor-cpu
@@ -99,6 +96,7 @@ steps:
 - label: Multi-Modal Accuracy Eval (Small Models) # 50min
  key: multi-modal-accuracy-eval-small-models
  timeout_in_minutes: 70
+  device: h200_18gb
  working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
  source_file_dependencies:
  - vllm/multimodal/
@@ -110,6 +108,7 @@ steps:
 - label: Multi-Modal Models (Extended Generation 1)
  key: multi-modal-models-extended-generation-1
  optional: true
+  device: h200_18gb
  source_file_dependencies:
  - vllm/
  - tests/models/multimodal/generation
@@ -120,13 +119,14 @@ steps:
    - pytest -v -s models/multimodal/test_mapping.py
  mirror:
    amd:
-      device: mi325_1
+      device: mi300_1
      depends_on:
      - image-build-amd

 - label: Multi-Modal Models (Extended Generation 2)
  key: multi-modal-models-extended-generation-2
  optional: true
+  device: h200_18gb
  source_file_dependencies:
  - vllm/
  - tests/models/multimodal/generation
@@ -137,6 +137,7 @@ steps:
 - label: Multi-Modal Models (Extended Generation 3)
  key: multi-modal-models-extended-generation-3
  optional: true
+  device: h200_18gb
  source_file_dependencies:
  - vllm/
  - tests/models/multimodal/generation
@@ -5,8 +5,32 @@ steps:
 - label: PyTorch Compilation Unit Tests
  key: pytorch-compilation-unit-tests
  timeout_in_minutes: 10
+  device: h200_18gb
  source_file_dependencies:
-    - vllm/
+    - vllm/__init__.py
+    - vllm/_aiter_ops.py
+    - vllm/_custom_ops.py
+    - vllm/compilation/
+    - vllm/config/
+    - vllm/distributed/
+    - vllm/engine/
+    - vllm/env_override.py
+    - vllm/envs.py
+    - vllm/forward_context.py
+    - vllm/inputs/
+    - vllm/ir/
+    - vllm/kernels/
+    - vllm/logger.py
+    - vllm/model_executor/
+    - vllm/multimodal/
+    - vllm/platforms/
+    - vllm/plugins/
+    - vllm/sampling_params.py
+    - vllm/sequence.py
+    - vllm/transformers_utils/
+    - vllm/triton_utils/
+    - vllm/utils/
+    - vllm/v1/
    - tests/compile
  commands:
  # Run unit tests defined directly under compile/,
@@ -24,7 +48,30 @@ steps:
  device: h100
  num_devices: 1
  source_file_dependencies:
-    - vllm/
+    - vllm/__init__.py
+    - vllm/_aiter_ops.py
+    - vllm/_custom_ops.py
+    - vllm/compilation/
+    - vllm/config/
+    - vllm/distributed/
+    - vllm/engine/
+    - vllm/env_override.py
+    - vllm/envs.py
+    - vllm/forward_context.py
+    - vllm/inputs/
+    - vllm/ir/
+    - vllm/kernels/
+    - vllm/logger.py
+    - vllm/model_executor/
+    - vllm/multimodal/
+    - vllm/platforms/
+    - vllm/plugins/
+    - vllm/sampling_params.py
+    - vllm/sequence.py
+    - vllm/transformers_utils/
+    - vllm/triton_utils/
+    - vllm/utils/
+    - vllm/v1/
    - tests/compile/h100/
  commands:
  - "find compile/h100/ -name 'test_*.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"
@@ -32,8 +79,32 @@ steps:
 - label: PyTorch Compilation Passes Unit Tests
  key: pytorch-compilation-passes-unit-tests
  timeout_in_minutes: 20
+  device: h200_18gb
  source_file_dependencies:
-    - vllm/
+    - vllm/__init__.py
+    - vllm/_aiter_ops.py
+    - vllm/_custom_ops.py
+    - vllm/compilation/
+    - vllm/config/
+    - vllm/distributed/
+    - vllm/engine/
+    - vllm/env_override.py
+    - vllm/envs.py
+    - vllm/forward_context.py
+    - vllm/inputs/
+    - vllm/ir/
+    - vllm/kernels/
+    - vllm/logger.py
+    - vllm/model_executor/
+    - vllm/multimodal/
+    - vllm/platforms/
+    - vllm/plugins/
+    - vllm/sampling_params.py
+    - vllm/sequence.py
+    - vllm/transformers_utils/
+    - vllm/triton_utils/
+    - vllm/utils/
+    - vllm/v1/
    - tests/compile/passes
  commands:
  - pytest -s -v compile/passes --ignore compile/passes/distributed
@@ -41,8 +112,32 @@ steps:
 - label: PyTorch Fullgraph Smoke Test
  key: pytorch-fullgraph-smoke-test
  timeout_in_minutes: 35
+  device: h200_18gb
  source_file_dependencies:
-  - vllm/
+  - vllm/__init__.py
+  - vllm/_aiter_ops.py
+  - vllm/_custom_ops.py
+  - vllm/compilation/
+  - vllm/config/
+  - vllm/distributed/
+  - vllm/engine/
+  - vllm/env_override.py
+  - vllm/envs.py
+  - vllm/forward_context.py
+  - vllm/inputs/
+  - vllm/ir/
+  - vllm/kernels/
+  - vllm/logger.py
+  - vllm/model_executor/
+  - vllm/multimodal/
+  - vllm/platforms/
+  - vllm/plugins/
+  - vllm/sampling_params.py
+  - vllm/sequence.py
+  - vllm/transformers_utils/
+  - vllm/triton_utils/
+  - vllm/utils/
+  - vllm/v1/
  - tests/compile
  commands:
  # Run smoke tests under fullgraph directory, except test_full_graph.py
@@ -56,7 +151,30 @@ steps:
  timeout_in_minutes: 30
  device: h200_18gb
  source_file_dependencies:
-  - vllm/
+  - vllm/__init__.py
+  - vllm/_aiter_ops.py
+  - vllm/_custom_ops.py
+  - vllm/compilation/
+  - vllm/config/
+  - vllm/distributed/
+  - vllm/engine/
+  - vllm/env_override.py
+  - vllm/envs.py
+  - vllm/forward_context.py
+  - vllm/inputs/
+  - vllm/ir/
+  - vllm/kernels/
+  - vllm/logger.py
+  - vllm/model_executor/
+  - vllm/multimodal/
+  - vllm/platforms/
+  - vllm/plugins/
+  - vllm/sampling_params.py
+  - vllm/sequence.py
+  - vllm/transformers_utils/
+  - vllm/triton_utils/
+  - vllm/utils/
+  - vllm/v1/
  - tests/compile
  commands:
    # fp8 kv scales not supported on sm89, tested on Blackwell instead
@@ -5,6 +5,7 @@ steps:
 - label: Quantization
  key: quantization
  timeout_in_minutes: 90
+  device: h200_18gb
  source_file_dependencies:
  - csrc/
  - vllm/model_executor/layers/quantization
@@ -25,7 +26,7 @@ steps:
  key: quantized-moe-test-b200
  timeout_in_minutes: 60
  working_dir: "/vllm-workspace/"
-  device: b200
+  device: b200-k8s
  source_file_dependencies:
  - tests/quantization/test_blackwell_moe.py
  - vllm/model_executor/models/deepseek_v2.py
@@ -42,6 +43,7 @@ steps:
 - label: Quantized Models Test
  key: quantized-models-test
  timeout_in_minutes: 60
+  device: h200_18gb
  source_file_dependencies:
  - vllm/model_executor/layers/quantization
  - tests/models/quantization
@@ -5,6 +5,7 @@ steps:
 - label: Samplers Test
  key: samplers-test
  timeout_in_minutes: 75
+  device: h200_18gb
  source_file_dependencies:
  - vllm/model_executor/layers
  - vllm/sampling_metadata.py
@@ -17,7 +18,7 @@ steps:
    - VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers
  mirror:
    amd:
-      device: mi325_1
+      device: mi250_1
      depends_on:
      - image-build-amd
      commands:
@@ -16,7 +16,7 @@ steps:
 - label: Spec Decode Eagle Nightly B200
  key: spec-decode-eagle-nightly-b200
  timeout_in_minutes: 30
-  device: b200
+  device: b200-k8s
  optional: true
  source_file_dependencies:
    - vllm/v1/spec_decode/
@@ -40,7 +40,7 @@ steps:
 - label: Spec Decode Speculators + MTP Nightly B200
  key: spec-decode-speculators-mtp-nightly-b200
  timeout_in_minutes: 30
-  device: b200
+  device: b200-k8s
  optional: true
  source_file_dependencies:
    - vllm/v1/spec_decode/
@@ -75,7 +75,7 @@ steps:
 - label: Spec Decode Draft Model Nightly B200
  key: spec-decode-draft-model-nightly-b200
  timeout_in_minutes: 30
-  device: b200
+  device: b200-k8s
  optional: true
  source_file_dependencies:
    - vllm/v1/spec_decode/
@@ -100,7 +100,7 @@ steps:

 - label: Spec Decode MTP hybrid (B200)
  timeout_in_minutes: 30
-  device: b200
+  device: b200-k8s
  optional: true
  source_file_dependencies:
    - vllm/v1/spec_decode/
@@ -6,8 +6,8 @@
 /vllm/distributed/kv_transfer @NickLucche @ApostaC @orozery @xuechendi
 /vllm/lora @jeejeelee
 /vllm/model_executor/layers/attention @LucasWilkinson @MatthewBonanni
-/vllm/model_executor/layers/fused_moe @mgoin @pavanimajety
-/vllm/model_executor/layers/quantization @mgoin @robertgshaw2-redhat @tlrmchlsmth @yewentao256 @pavanimajety
+/vllm/model_executor/layers/fused_moe @mgoin @pavanimajety @zyongye
+/vllm/model_executor/layers/quantization @mgoin @robertgshaw2-redhat @tlrmchlsmth @yewentao256 @pavanimajety @zyongye
 /vllm/model_executor/layers/mamba @tdoublep @tomeras91
 /vllm/model_executor/layers/mamba/gdn_linear_attn.py @tdoublep @ZJY0516 @vadiklyutiy
 /vllm/model_executor/layers/rotary_embedding.py @vadiklyutiy
@@ -18,7 +18,8 @@
 /vllm/kernels/helion @ProExpertProg @zou3519
 /vllm/multimodal @DarkLight1337 @ywang96 @NickLucche @tjtanaa
 /vllm/vllm_flash_attn @LucasWilkinson @MatthewBonanni
-CMakeLists.txt @tlrmchlsmth @LucasWilkinson
+/CMakeLists.txt @tlrmchlsmth @LucasWilkinson @Harry-Chen
+/cmake @tlrmchlsmth @LucasWilkinson @Harry-Chen

 # Any change to the VllmConfig changes can have a large user-facing impact,
 # so spam a lot of people
@@ -30,8 +31,8 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson
 /vllm/entrypoints/cli @hmellor @mgoin @DarkLight1337 @russellb
 /vllm/entrypoints/mcp @heheda12345
 /vllm/entrypoints/openai @aarnphm @chaunceyjiang @DarkLight1337 @russellb
-/vllm/entrypoints/openai/realtime @njhill
-/vllm/entrypoints/openai/speech_to_text @NickLucche
+/vllm/entrypoints/speech_to_text/realtime @njhill
+/vllm/entrypoints/speech_to_text @NickLucche
 /vllm/entrypoints/pooling @noooop
 /vllm/entrypoints/sagemaker @DarkLight1337
 /vllm/entrypoints/serve @njhill
@@ -70,6 +71,10 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson
 /vllm/v1/worker/gpu @WoosukKwon @njhill
 /vllm/v1/worker/gpu/kv_connector.py @orozery

+# CI & building
+/.buildkite @Harry-Chen
+/docker/Dockerfile @Harry-Chen
+
 # Test ownership
 /.buildkite/lm-eval-harness @mgoin 
 /tests/distributed/test_multi_node_assignment.py @youkaichao
@@ -77,11 +82,11 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson
 /tests/distributed/test_same_node.py @youkaichao
 /tests/entrypoints @DarkLight1337 @robertgshaw2-redhat @aarnphm @NickLucche
 /tests/evals @mgoin @vadiklyutiy
-/tests/kernels @mgoin @tlrmchlsmth @WoosukKwon @yewentao256
+/tests/kernels @mgoin @tlrmchlsmth @WoosukKwon @yewentao256 @zyongye
 /tests/kernels/ir @ProExpertProg @tjtanaa
 /tests/models @DarkLight1337 @ywang96
 /tests/multimodal @DarkLight1337 @ywang96 @NickLucche
-/tests/quantization @mgoin @robertgshaw2-redhat @yewentao256 @pavanimajety
+/tests/quantization @mgoin @robertgshaw2-redhat @yewentao256 @pavanimajety @zyongye
 /tests/test_inputs.py @DarkLight1337 @ywang96
 /tests/entrypoints/llm/test_struct_output_generate.py @mgoin @russellb @aarnphm
 /tests/v1/structured_output @mgoin @russellb @aarnphm
@@ -147,6 +152,12 @@ mkdocs.yaml @hmellor
 # MTP-specific files
 /vllm/model_executor/models/deepseek_mtp.py @luccafong

+# DeepseekV4-specific files
+/vllm/v1/attention/ops/deepseek_v4_ops @zyongye
+/vllm/model_executor/layers/deepseek_compressor.py @zyongye
+/vllm/model_executor/layers/deepseek_v4_attention.py @zyongye
+/vllm/model_executor/layers/sparse_attn_indexer.py @zyongye
+
 # Mistral-specific files
 /vllm/model_executor/models/mistral*.py @patrickvonplaten
 /vllm/model_executor/models/mixtral*.py @patrickvonplaten
@@ -477,9 +477,7 @@ pull_request_rules:
  conditions:
    - label != stale
    - or:
-      - files~=^examples/online_serving/disaggregated[^/]*/.*
-      - files~=^examples/offline_inference/disaggregated[^/]*/.*
-      - files~=^examples/others/lmcache/
+      - files~=^examples/disaggregated/
      - files~=^tests/v1/kv_connector/
      - files~=^vllm/distributed/kv_transfer/
      - title~=(?i)\bP/?D\b
@@ -131,6 +131,16 @@ repos:
        --python-version, "3.12",
      ]
      files: ^requirements/(common|xpu|test/xpu)\.(in|txt)$
+    - id: pip-compile
+      alias: pip-compile-docs
+      name: pip-compile-docs
+      args: [
+        requirements/docs.in,
+        -o, requirements/docs.txt,
+        --python-platform, x86_64-manylinux_2_28,
+        --python-version, "3.12",
+      ]
+      files: ^requirements/docs\.(in|txt)$
 - repo: local
  hooks:
  - id: format-torch-nightly-test
@@ -9,7 +9,44 @@ build:
    python: "3.12"
  jobs:
    post_checkout:
-      # - bash docs/maybe_skip_pr_build.sh
+      - |
+        if [ "$READTHEDOCS_VERSION_TYPE" = "external" ]; then
+          MAX_WAIT=300
+          INTERVAL=60
+          ELAPSED=0
+          while :; do
+            RAW=$(curl -sS -w "\n%{http_code}" "https://api.github.com/repos/vllm-project/vllm/commits/${READTHEDOCS_GIT_COMMIT_HASH}/check-runs?check_name=pre-run-check&filter=latest")
+            HTTP_CODE=$(printf %s "$RAW" | tail -n1)
+            BODY=$(printf %s "$RAW" | head -n -1)
+            if [ "$HTTP_CODE" != "200" ]; then
+              echo "GitHub API returned HTTP $HTTP_CODE (likely rate-limited); skipping pre-run-check gate."
+              break
+            fi
+            STATUS=$(printf %s "$BODY" | python3 -c "import sys, json; r=json.load(sys.stdin).get(\"check_runs\",[]); print((r[0].get(\"status\") or \"\") if r else \"none\")")
+            CONCLUSION=$(printf %s "$BODY" | python3 -c "import sys, json; r=json.load(sys.stdin).get(\"check_runs\",[]); print((r[0].get(\"conclusion\") or \"\") if r else \"\")")
+            if [ "$STATUS" = "none" ]; then
+              echo "no pre-run-check found for this commit; skipping gate."
+              break
+            fi
+            if [ -n "$CONCLUSION" ]; then
+              echo "pre-run-check conclusion: $CONCLUSION"
+              if [ "$CONCLUSION" = "failure" ] || [ "$CONCLUSION" = "cancelled" ] || [ "$CONCLUSION" = "timed_out" ]; then
+                echo "pre-run-check did not pass; failing docs build."
+                exit 1
+              fi
+              break
+            fi
+            if [ "$ELAPSED" -ge "$MAX_WAIT" ]; then
+              echo "pre-run-check status=$STATUS after ${MAX_WAIT}s; skipping gate."
+              break
+            fi
+            echo "pre-run-check status=$STATUS; waiting ${INTERVAL}s..."
+            sleep "$INTERVAL"
+            ELAPSED=$((ELAPSED + INTERVAL))
+          done
+        else
+          echo "Not a PR build (version type=$READTHEDOCS_VERSION_TYPE); skipping pre-run-check gate."
+        fi
      - git fetch origin main --unshallow --no-tags --filter=blob:none || true
    pre_create_environment:
      - pip install uv
@@ -13,8 +13,12 @@ cmake_minimum_required(VERSION 3.26)
 # cmake --install . --component _C
 project(vllm_extensions LANGUAGES CXX)

-set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_CXX_STANDARD 20)
 set(CMAKE_CXX_STANDARD_REQUIRED ON)
+set(CMAKE_CUDA_STANDARD 20)
+set(CMAKE_CUDA_STANDARD_REQUIRED ON)
+set(CMAKE_HIP_STANDARD 20)
+set(CMAKE_HIP_STANDARD_REQUIRED ON)


 # CUDA by default, can be overridden by using -DVLLM_TARGET_DEVICE=... (used by setup.py)
@@ -105,6 +109,24 @@ else()
  set(CUDA_SUPPORTED_ARCHS "7.0;7.5;8.0;8.6;8.7;8.9;9.0")
 endif()

+#
+# spinloop extension (pure CXX; must stay above the non-CUDA device branch so
+# CPU builds define the target before the early return)
+#
+set(VLLM_SPINLOOP_EXT_SRC "csrc/spinloop.cpp")
+set(SPINLOOP_COMPILE_FLAGS "")
+if(CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|amd64")
+  list(APPEND SPINLOOP_COMPILE_FLAGS "-mmwaitx")
+endif()
+define_extension_target(
+  spinloop
+  DESTINATION vllm
+  LANGUAGE CXX
+  SOURCES ${VLLM_SPINLOOP_EXT_SRC}
+  COMPILE_FLAGS ${SPINLOOP_COMPILE_FLAGS}
+  USE_SABI 3.11
+  WITH_SOABI)
+
 #
 # Forward the non-CUDA device extensions to external CMake scripts.
 #
@@ -307,12 +329,12 @@ set(VLLM_EXT_SRC
  "csrc/quantization/activation_kernels.cu"
  "csrc/cuda_utils_kernels.cu"
  "csrc/custom_all_reduce.cu"
-  "csrc/torch_bindings.cpp")
+  "csrc/torch_bindings.cpp"
+  "csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")

 if(VLLM_GPU_LANG STREQUAL "CUDA")
  list(APPEND VLLM_EXT_SRC
-    "csrc/minimax_reduce_rms_kernel.cu"
-    "csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")
+    "csrc/minimax_reduce_rms_kernel.cu")

  SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")

@@ -1047,14 +1069,13 @@ endif()
 set(VLLM_MOE_EXT_SRC
  "csrc/moe/torch_bindings.cpp"
  "csrc/moe/moe_align_sum_kernels.cu"
-  "csrc/moe/topk_softmax_kernels.cu")
+  "csrc/moe/topk_softmax_kernels.cu"
+  "csrc/moe/topk_softplus_sqrt_kernels.cu")

 if(VLLM_GPU_LANG STREQUAL "CUDA")
  list(APPEND VLLM_MOE_EXT_SRC
    "csrc/moe/moe_wna16.cu"
-    "csrc/moe/grouped_topk_kernels.cu"
-    "csrc/moe/router_gemm.cu"
-    "csrc/moe/topk_softplus_sqrt_kernels.cu")
+    "csrc/moe/grouped_topk_kernels.cu")
 endif()

 if(VLLM_GPU_LANG STREQUAL "CUDA")
@@ -1473,6 +1473,12 @@ async def main() -> None:
        "(for example: --warmup-percentages=0%%,50%%)",
    )

+    parser.add_argument(
+        "--trust-remote-code",
+        action="store_true",
+        help="Trust remote code when loading the tokenizer.",
+    )
+
    args = parser.parse_args()

    logger.info(args)
@@ -1515,7 +1521,9 @@ async def main() -> None:
    np.random.seed(args.seed)

    logger.info("Loading tokenizer")
-    tokenizer = AutoTokenizer.from_pretrained(args.model)
+    tokenizer = AutoTokenizer.from_pretrained(
+        args.model, trust_remote_code=args.trust_remote_code
+    )

    await get_server_info(args.url)

@@ -1,7 +1,7 @@
 include(FetchContent)

 set(CMAKE_CXX_STANDARD_REQUIRED ON)
-set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_CXX_STANDARD 20)
 set(CMAKE_CXX_EXTENSIONS ON)
 set(CMAKE_EXPORT_COMPILE_COMMANDS ON)

@@ -32,18 +32,23 @@ else()
        "-DVLLM_CPU_EXTENSION")

    # locate PyTorch's libgomp (e.g. site-packages/torch.libs/libgomp-947d5fa1.so.1.0.0)
-    # and create a local shim dir with it
+    # and create a local shim dir with it. When PyTorch is built from source or packaged
+    # by a distro (common on RISC-V, s390x, Fedora/RHEL aarch64), no vendored libgomp
+    # exists and the shim dir is empty; fall back to the system libgomp in that case.
    vllm_prepare_torch_gomp_shim(VLLM_TORCH_GOMP_SHIM_DIR)

-    find_library(OPEN_MP
-        NAMES gomp
-        PATHS ${VLLM_TORCH_GOMP_SHIM_DIR}
-        NO_DEFAULT_PATH
-        REQUIRED
-    )
-    # Set LD_LIBRARY_PATH to include the shim dir at build time to use the same libgomp as PyTorch
-    if (OPEN_MP)
+    if(VLLM_TORCH_GOMP_SHIM_DIR)
+        find_library(OPEN_MP
+            NAMES gomp
+            PATHS "${VLLM_TORCH_GOMP_SHIM_DIR}"
+            NO_DEFAULT_PATH
+            REQUIRED
+        )
+        # Use the same libgomp as PyTorch at runtime
        set(ENV{LD_LIBRARY_PATH} "${VLLM_TORCH_GOMP_SHIM_DIR}:$ENV{LD_LIBRARY_PATH}")
+    else()
+        # Fall back to system / toolchain libgomp
+        find_library(OPEN_MP NAMES gomp REQUIRED)
    endif()
 endif()

@@ -321,14 +326,6 @@ if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND
    set(ONEDNN_VERBOSE "ON")
    set(CMAKE_POLICY_DEFAULT_CMP0077 NEW)

-    # TODO: Refactor this
-    if (ENABLE_X86_ISA)
-        # Note: only enable oneDNN for AVX512
-        list(APPEND DNNL_COMPILE_FLAGS ${CXX_COMPILE_FLAGS_AVX512})
-    else()
-        list(APPEND DNNL_COMPILE_FLAGS ${CXX_COMPILE_FLAGS})
-    endif()
-
    set(VLLM_BUILD_TYPE ${CMAKE_BUILD_TYPE})
    set(CMAKE_BUILD_TYPE "Release") # remove oneDNN debug symbols to reduce size
    FetchContent_MakeAvailable(oneDNN)
@@ -341,8 +338,14 @@ if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND
        PRIVATE ${oneDNN_SOURCE_DIR}/src
    )
    target_link_libraries(dnnl_ext dnnl torch)
-    target_compile_options(dnnl_ext PRIVATE ${DNNL_COMPILE_FLAGS} -fPIC)
+    if (ENABLE_X86_ISA)
+        target_compile_options(dnnl_ext PRIVATE ${CXX_COMPILE_FLAGS_AVX2} -fPIC)
+    else()
+        target_compile_options(dnnl_ext PRIVATE ${CXX_COMPILE_FLAGS} -fPIC)
+    endif()
    list(APPEND LIBS dnnl_ext)
+
+
    set(USE_ONEDNN ON)
 else()
    set(USE_ONEDNN OFF)
@@ -406,12 +409,15 @@ endif()

 if (ENABLE_X86_ISA)
    set(VLLM_EXT_SRC_SGL
+        "csrc/cpu/sgl-kernels/fla.cpp"
+        "csrc/cpu/sgl-kernels/conv.cpp"
        "csrc/cpu/sgl-kernels/gemm.cpp"
        "csrc/cpu/sgl-kernels/gemm_int8.cpp"
        "csrc/cpu/sgl-kernels/gemm_fp8.cpp"
        "csrc/cpu/sgl-kernels/gemm_int4.cpp"
        "csrc/cpu/sgl-kernels/moe.cpp"
        "csrc/cpu/sgl-kernels/moe_int8.cpp"
+        "csrc/cpu/sgl-kernels/moe_int4.cpp"
        "csrc/cpu/sgl-kernels/moe_fp8.cpp")

    set(VLLM_EXT_SRC_AVX512
@@ -430,10 +436,11 @@ if (ENABLE_X86_ISA)
        "csrc/cpu/pos_encoding.cpp"
        "csrc/moe/dynamic_4bit_int_moe_cpu.cpp") 

-    set(VLLM_EXT_SRC_AVX2 
+    set(VLLM_EXT_SRC_AVX2
        "csrc/cpu/utils.cpp"
        "csrc/cpu/spec_decode_utils.cpp"
        "csrc/cpu/cpu_attn.cpp"
+        "csrc/cpu/dnnl_kernels.cpp"
        "csrc/cpu/torch_bindings.cpp"
        # TODO: Remove these files
        "csrc/cpu/activation.cpp"
@@ -448,7 +455,7 @@ if (ENABLE_X86_ISA)

    set(_C_LIBS numa dnnl_ext)
    set(_C_AVX512_LIBS numa dnnl_ext)
-    set(_C_AVX2_LIBS numa)
+    set(_C_AVX2_LIBS numa dnnl_ext)

    # AMX + AVX512F + AVX512BF16 + AVX512VNNI
    define_extension_target(
@@ -76,7 +76,6 @@ if(DEEPGEMM_ARCHS)
    "${deepgemm_SOURCE_DIR}/third-party/fmt/include")

  target_compile_options(_deep_gemm_C PRIVATE
-    $<$<COMPILE_LANGUAGE:CXX>:-std=c++17>
    $<$<COMPILE_LANGUAGE:CXX>:-O3>
    $<$<COMPILE_LANGUAGE:CXX>:-Wno-psabi>
    $<$<COMPILE_LANGUAGE:CXX>:-Wno-deprecated-declarations>)
@@ -12,7 +12,8 @@ void swap_blocks(torch::Tensor& src, torch::Tensor& dst,

 void swap_blocks_batch(const torch::Tensor& src_ptrs,
                       const torch::Tensor& dst_ptrs,
-                       const torch::Tensor& sizes);
+                       const torch::Tensor& sizes,
+                       bool is_src_access_order_any);

 void reshape_and_cache(torch::Tensor& key, torch::Tensor& value,
                       torch::Tensor& key_cache, torch::Tensor& value_cache,
@@ -77,7 +77,8 @@ void swap_blocks(torch::Tensor& src, torch::Tensor& dst,

 void swap_blocks_batch(const torch::Tensor& src_ptrs,
                       const torch::Tensor& dst_ptrs,
-                       const torch::Tensor& sizes) {
+                       const torch::Tensor& sizes,
+                       bool is_src_access_order_any) {
  TORCH_CHECK(src_ptrs.device().is_cpu(), "src_ptrs must be on CPU");
  TORCH_CHECK(dst_ptrs.device().is_cpu(), "dst_ptrs must be on CPU");
  TORCH_CHECK(sizes.device().is_cpu(), "sizes must be on CPU");
@@ -97,13 +98,13 @@ void swap_blocks_batch(const torch::Tensor& src_ptrs,

  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();

-  // Use cuMemcpyBatchAsync (CUDA 12.8+) to submit all copies in a single
-  // driver call, amortizing per-copy submission overhead.
-  // int64_t and CUdeviceptr/size_t are both 8 bytes on 64-bit platforms,
-  // so we reinterpret_cast the tensor data directly to avoid copies.
-  static_assert(sizeof(CUdeviceptr) == sizeof(int64_t));
+  // Use cuMemcpyBatchAsync / hipMemcpyBatchAsync to submit all copies in a
+  // single driver call, amortizing per-copy submission overhead. int64_t
+  // and CUdeviceptr/void*/size_t are all 8 bytes on 64-bit platforms, so we
+  // reinterpret_cast the tensor data directly to avoid copies.
  static_assert(sizeof(size_t) == sizeof(int64_t));
 #if !defined(USE_ROCM) && defined(CUDA_VERSION) && CUDA_VERSION >= 12080
+  static_assert(sizeof(CUdeviceptr) == sizeof(int64_t));
  // Resolve cuMemcpyBatchAsync at runtime via cuGetProcAddress so that
  // binaries compiled with CUDA 12.8+ still work on older drivers, and
  // we avoid the CUDA 13.0 header remapping (#define to _v2 signature).
@@ -124,7 +125,12 @@ void swap_blocks_batch(const torch::Tensor& src_ptrs,

  if (batch_fn != nullptr) {
    CUmemcpyAttributes attr = {};
-    attr.srcAccessOrder = CU_MEMCPY_SRC_ACCESS_ORDER_STREAM;
+    // ANY lets the DMA engine prefetch source bytes out of stream order,
+    // which is only safe when no GPU stream is concurrently writing the
+    // source.
+    attr.srcAccessOrder = is_src_access_order_any
+                              ? CU_MEMCPY_SRC_ACCESS_ORDER_ANY
+                              : CU_MEMCPY_SRC_ACCESS_ORDER_STREAM;
    size_t attrs_idx = 0;
    size_t fail_idx = 0;
    CUresult result = batch_fn(reinterpret_cast<CUdeviceptr*>(dst_data),
@@ -134,12 +140,30 @@ void swap_blocks_batch(const torch::Tensor& src_ptrs,
                               &fail_idx, static_cast<CUstream>(stream));
    TORCH_CHECK(result == CUDA_SUCCESS, "cuMemcpyBatchAsync failed at index ",
                fail_idx, " with error ", result);
-  } else
+    return;
+  }
+#elif defined(USE_ROCM) && defined(HIP_VERSION) && HIP_VERSION >= 70100000
+  // ROCm 7.1+ exposes hipMemcpyBatchAsync. The 7.2.1 implementation early-
+  // returns hipErrorNotSupported whenever numAttrs > 0 (see ROCm/clr @
+  // rocm-7.2.1 hipamd/src/hip_memory.cpp:2819-2822), so call with
+  // numAttrs=0.
+  {
+    hipMemcpyAttributes attr = {};
+    size_t attrs_idx = 0;
+    size_t fail_idx = 0;
+    hipError_t result = hipMemcpyBatchAsync(
+        reinterpret_cast<void**>(dst_data), reinterpret_cast<void**>(src_data),
+        reinterpret_cast<size_t*>(size_data), static_cast<size_t>(n), &attr,
+        &attrs_idx, 0, &fail_idx, static_cast<hipStream_t>(stream));
+    TORCH_CHECK(result == hipSuccess, "hipMemcpyBatchAsync failed at index ",
+                fail_idx, " with error ", result);
+    return;
+  }
 #endif
  {
-    // Fallback for CUDA < 12.8, older drivers, and ROCm:
-    // individual async copies.
-    // cudaMemcpyDefault lets the driver infer direction from pointer types.
+    // Fallback for CUDA < 12.8, older CUDA drivers, and ROCm < 7.1:
+    // individual async copies. cudaMemcpyDefault lets the driver infer
+    // direction from pointer types.
    for (int64_t i = 0; i < n; i++) {
      cudaMemcpyAsync(reinterpret_cast<void*>(dst_data[i]),
                      reinterpret_cast<void*>(src_data[i]),
@@ -21,28 +21,33 @@ namespace vllm {

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

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

  const int embed_dim = rot_dim / 2;

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

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

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

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

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

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

    qk_t* k_pe_head_ptr = k_pe + token_idx * k_pe_stride;

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

 }  // namespace vllm

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

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

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

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

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

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

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

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

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

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

 namespace vllm {

@@ -29,6 +29,8 @@ torch::Tensor get_scheduler_metadata(
    isa = cpu_attention::ISA::NEON;
  } else if (isa_hint == "vxe") {
    isa = cpu_attention::ISA::VXE;
+  } else if (isa_hint == "vsx") {
+    isa = cpu_attention::ISA::VSX;
  } else {
    TORCH_CHECK(false, "Unsupported CPU attention ISA hint: " + isa_hint);
  }
@@ -129,6 +131,8 @@ void cpu_attn_reshape_and_cache(
      return cpu_attention::ISA::NEON;
    } else if (isa == "vxe") {
      return cpu_attention::ISA::VXE;
+    } else if (isa == "vsx") {
+      return cpu_attention::ISA::VSX;
    } else {
      TORCH_CHECK(false, "Invalid ISA type: " + isa);
    }
@@ -12,7 +12,7 @@
 #include "cpu/utils.hpp"

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

 // Mirrors csrc/attention/dtype_fp8.cuh Fp8KVCacheDataType exactly.
 enum class Fp8KVCacheDataType {
@@ -164,6 +164,9 @@ struct AttentionMetadata {
      case ISA::VXE:
        ss << "VXE, ";
        break;
+      case ISA::VSX:
+        ss << "VSX, ";
+        break;
    }
    ss << "workitem_group_num: " << workitem_group_num
       << ", reduction_item_num: " << reduction_item_num
@@ -27,8 +27,8 @@ FORCE_INLINE std::pair<vec_op::FP32Vec16, vec_op::FP32Vec16> load_b_pair_vec(
    return {vec_op::FP32Vec16(bf16_b_reg, 0), vec_op::FP32Vec16(bf16_b_reg, 1)};
  } else {
    using load_vec_t = typename VecTypeTrait<kv_cache_t>::vec_t;
-    return {vec_op::FP32Vec16(load_vec_t(ptr)),
-            vec_op::FP32Vec16(load_vec_t(ptr + 16))};
+    return std::make_pair(vec_op::FP32Vec16(load_vec_t(ptr)),
+                          vec_op::FP32Vec16(load_vec_t(ptr + 16)));
  }
 }

@@ -0,0 +1,359 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#ifndef CPU_ATTN_VSX_HPP
+#define CPU_ATTN_VSX_HPP
+
+#include "cpu_attn_impl.hpp"
+#include <altivec.h>
+#include <type_traits>
+
+namespace cpu_attention {
+
+namespace {
+
+// ppc64le Vector = 16 bytes (128 bits)
+#define BLOCK_SIZE_ALIGNMENT 32
+#define HEAD_SIZE_ALIGNMENT 32
+#define MAX_Q_HEAD_NUM_PER_ITER 16
+
+template <typename kv_cache_t>
+FORCE_INLINE void load_row8_B_as_f32(const kv_cache_t* p, __vector float& b0,
+                                     __vector float& b1);
+
+// [1] Float Specialization
+template <>
+FORCE_INLINE void load_row8_B_as_f32<float>(const float* p, __vector float& b0,
+                                            __vector float& b1) {
+  b0 = vec_xl(0, const_cast<float*>(p));
+  b1 = vec_xl(0, const_cast<float*>(p + 4));
+}
+
+// [2] BFloat16 Specialization (Little Endian ppc64le)
+// On ppc64le (LE): BF16 bits should land in the HIGH 16 bits of each float32.
+// Byte layout of float32 on LE: [byte0(LSB), byte1, byte2, byte3(MSB)]
+// We need BF16 in bytes2-3 (high half) with bytes0-1 zeroed.
+// vec_mergeh on LE interleaves elements 0..3: result_i = {a[i], b[i]}
+// So vec_mergeh(zeros_u16, raw_u16) gives for each uint16 pair:
+//   uint16[2i]   = zeros[i]  -> low 16 bits of uint32  -> zeroed mantissa LSBs
+//   uint16[2i+1] = raw[i]    -> high 16 bits of uint32 -> BF16 bits
+// Cast to float32 gives exactly (bf16_bits << 16) per element.
+template <>
+FORCE_INLINE void load_row8_B_as_f32<c10::BFloat16>(const c10::BFloat16* p,
+                                                    __vector float& b0,
+                                                    __vector float& b1) {
+  __vector unsigned short raw = vec_xl(
+      0, reinterpret_cast<unsigned short*>(const_cast<c10::BFloat16*>(p)));
+  __vector unsigned short zeros = vec_splat_u16(0);
+
+  // LE: zeros in low 16 bits, raw in high 16 bits → bf16 << 16 == float32
+  b0 = (__vector float)vec_mergeh(zeros, raw);
+  b1 = (__vector float)vec_mergel(zeros, raw);
+}
+
+// Note: c10::Half (FP16) is not supported on PowerPC architecture
+
+template <int32_t M, typename kv_cache_t>
+FORCE_INLINE void gemm_micro_ppc64le_Mx8_Ku4(
+    const float* __restrict A,       // [M x K]
+    const kv_cache_t* __restrict B,  // [K x 8]
+    float* __restrict C,             // [M x 8]
+    int64_t lda, int64_t ldb, int64_t ldc, int32_t K, bool accumulate) {
+  static_assert(1 <= M && M <= 8, "M must be in [1,8]");
+
+#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))
+
+  // 1. Define A pointers
+#define DECL_A(i) const float* a##i = A + (i) * lda;
+  ROWS_APPLY(DECL_A)
+#undef DECL_A
+
+  // 2. Define Accumulators (2 vectors covers 8 columns)
+#define DECL_ACC(i) __vector float acc##i##_0, acc##i##_1;
+  ROWS_APPLY(DECL_ACC)
+#undef DECL_ACC
+
+  // 3. Initialize Accumulators (Load C or Zero)
+#define INIT_ACC(i)                                                  \
+  IF_M(i) {                                                          \
+    if (accumulate) {                                                \
+      acc##i##_0 = vec_xl(0, const_cast<float*>(C + (i) * ldc + 0)); \
+      acc##i##_1 = vec_xl(0, const_cast<float*>(C + (i) * ldc + 4)); \
+    } else {                                                         \
+      acc##i##_0 = vec_splats(0.0f);                                 \
+      acc##i##_1 = vec_splats(0.0f);                                 \
+    }                                                                \
+  }
+  ROWS_APPLY(INIT_ACC)
+#undef INIT_ACC
+
+  int32_t k = 0;
+
+  for (; k + 3 < K; k += 4) {
+    // Load 4 values of A for each Row M: A[k...k+3]
+#define LOAD_A4(i)        \
+  __vector float a##i##v; \
+  IF_M(i) a##i##v = vec_xl(0, const_cast<float*>(a##i + k));
+    ROWS_APPLY(LOAD_A4)
+#undef LOAD_A4
+
+    // FMA for specific lane L of A
+    // ppc64le: vec_madd(b, vec_splat(a, lane), acc)
+#define FMAS_LANE(i, aiv, L)                        \
+  IF_M(i) {                                         \
+    __vector float a_broad = vec_splat(aiv, L);     \
+    acc##i##_0 = vec_madd(b0, a_broad, acc##i##_0); \
+    acc##i##_1 = vec_madd(b1, a_broad, acc##i##_1); \
+  }
+
+    // Unroll K=0..3
+    {
+      __vector float 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
+    }
+    {
+      __vector float 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
+    }
+    {
+      __vector float 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
+    }
+    {
+      __vector float 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
+  }
+
+  for (; k < K; ++k) {
+    __vector float b0, b1;
+    load_row8_B_as_f32<kv_cache_t>(B + (int64_t)k * ldb, b0, b1);
+#define TAIL_ROW(i)                              \
+  IF_M(i) {                                      \
+    __vector float ai = vec_splats(*(a##i + k)); \
+    acc##i##_0 = vec_madd(b0, ai, acc##i##_0);   \
+    acc##i##_1 = vec_madd(b1, ai, acc##i##_1);   \
+  }
+    ROWS_APPLY(TAIL_ROW)
+#undef TAIL_ROW
+  }
+
+#define STORE_ROW(i)                           \
+  IF_M(i) {                                    \
+    vec_xst(acc##i##_0, 0, C + (i) * ldc + 0); \
+    vec_xst(acc##i##_1, 0, C + (i) * ldc + 4); \
+  }
+  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_ppc64le_Mx8_Ku4(const float* __restrict A,
+                                             const kv_cache_t* __restrict B,
+                                             float* __restrict C, int32_t M,
+                                             int32_t K, int64_t lda,
+                                             int64_t ldb, int64_t ldc,
+                                             bool accumulate) {
+  static_assert(N % 8 == 0, "N must be a multiple of 8");
+  for (int32_t m = 0; m < M;) {
+    int32_t mb = (M - m >= 8) ? 8 : (M - m >= 4) ? 4 : (M - m >= 2) ? 2 : 1;
+    const float* Ab = A + m * lda;
+    float* Cb = C + m * ldc;
+
+    for (int32_t n = 0; n < N; n += 8) {
+      const kv_cache_t* Bn = B + n;
+      float* Cn = Cb + n;
+      switch (mb) {
+        case 8:
+          gemm_micro_ppc64le_Mx8_Ku4<8, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
+                                                    K, accumulate);
+          break;
+        case 4:
+          gemm_micro_ppc64le_Mx8_Ku4<4, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
+                                                    K, accumulate);
+          break;
+        case 2:
+          gemm_micro_ppc64le_Mx8_Ku4<2, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
+                                                    K, accumulate);
+          break;
+        default:
+          gemm_micro_ppc64le_Mx8_Ku4<1, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
+                                                    K, accumulate);
+          break;
+      }
+    }
+    m += mb;
+  }
+}
+
+template <typename kv_cache_t>
+class TileGemmPPC64 {
+ 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_ppc64le_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_ppc64le_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
+
+template <typename scalar_t, int64_t head_dim>
+class AttentionImpl<ISA::VSX, 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;
+  constexpr static int64_t HeadDimAlignment = HEAD_SIZE_ALIGNMENT;
+  constexpr static int64_t MaxQHeadNumPerIteration = MAX_Q_HEAD_NUM_PER_ITER;
+  constexpr static int64_t HeadDim = head_dim;
+  constexpr static ISA ISAType = ISA::VSX;
+  constexpr static bool scale_on_logits =
+      false;  // Scale is applied to Q during copy
+
+ public:
+  AttentionImpl() {}
+
+  template <template <typename tile_gemm_t> typename attention>
+  FORCE_INLINE void execute_attention(DEFINE_CPU_ATTENTION_PARAMS) {
+    attention<TileGemmPPC64<kv_cache_t>> attention_iteration;
+    attention_iteration(CPU_ATTENTION_PARAMS);
+  }
+
+  // Strides for Memory Layout
+  constexpr static int64_t k_cache_token_group_stride(
+      const int32_t block_size) {
+    return BlockSizeAlignment;  // [head_dim, block_size] layout
+  }
+
+  constexpr static int64_t v_cache_token_group_stride(
+      const int32_t block_size) {
+    return head_dim * BlockSizeAlignment;
+  }
+
+  constexpr static int64_t v_cache_head_group_stride(const int32_t block_size) {
+    return HeadDimAlignment;
+  }
+
+  static void copy_q_heads_tile(scalar_t* __restrict__ src,
+                                float* __restrict__ q_buffer,
+                                const int32_t q_num,
+                                const int32_t q_heads_per_kv,
+                                const int64_t q_num_stride,
+                                const int64_t q_head_stride, float scale) {
+    __vector float scale_vec = vec_splats(scale);
+    constexpr bool is_bf16 = std::is_same<scalar_t, c10::BFloat16>::value;
+
+    for (int32_t i = 0; i < q_num; ++i) {
+      for (int32_t h = 0; h < q_heads_per_kv; ++h) {
+        scalar_t* curr_src = src + i * q_num_stride + h * q_head_stride;
+        float* curr_dst =
+            q_buffer + i * q_heads_per_kv * head_dim + h * head_dim;
+
+        int32_t d = 0;
+        for (; d <= head_dim - 8; d += 8) {
+          __vector float v0, v1;
+          load_row8_B_as_f32<scalar_t>(curr_src + d, v0, v1);
+
+          v0 = vec_mul(v0, scale_vec);
+          v1 = vec_mul(v1, scale_vec);
+
+          vec_xst(v0, 0, curr_dst + d);
+          vec_xst(v1, 0, curr_dst + d + 4);
+        }
+
+        for (; d < head_dim; ++d) {
+          float val = static_cast<float>(curr_src[d]);
+          curr_dst[d] = val * scale;
+        }
+      }
+    }
+  }
+
+  static void reshape_and_cache(
+      const scalar_t* __restrict__ key, const scalar_t* __restrict__ value,
+      scalar_t* __restrict__ key_cache, scalar_t* __restrict__ value_cache,
+      const int64_t* __restrict__ slot_mapping, const int64_t token_num,
+      const int64_t key_token_num_stride, const int64_t value_token_num_stride,
+      const int64_t head_num, const int64_t key_head_num_stride,
+      const int64_t value_head_num_stride, const int64_t num_blocks,
+      const int64_t num_blocks_stride, const int64_t cache_head_num_stride,
+      const int64_t block_size, const int64_t block_size_stride,
+      const float k_inv = 0.0f, const float v_inv = 0.0f) {
+    // k_inv and v_inv are unused on VSX: FP8 KV cache is not supported on
+    // PowerPC. The parameters are present to match the common interface.
+#pragma omp parallel for collapse(2)
+    for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
+      for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
+        const int64_t pos = slot_mapping[token_idx];
+        if (pos < 0) continue;
+
+        const int64_t block_idx = pos / block_size;
+        const int64_t block_offset = pos % block_size;
+
+        {
+          const scalar_t* key_src = key + token_idx * key_token_num_stride +
+                                    head_idx * key_head_num_stride;
+          scalar_t* key_dst = key_cache + block_idx * num_blocks_stride +
+                              head_idx * cache_head_num_stride + block_offset;
+
+          for (int64_t i = 0, j = 0; i < head_dim; ++i, j += block_size) {
+            key_dst[j] = key_src[i];
+          }
+        }
+
+        {
+          const scalar_t* val_src = value + token_idx * value_token_num_stride +
+                                    head_idx * value_head_num_stride;
+          scalar_t* val_dst = value_cache + block_idx * num_blocks_stride +
+                              head_idx * cache_head_num_stride +
+                              block_offset * head_dim;
+
+          std::memcpy(val_dst, val_src, sizeof(scalar_t) * head_dim);
+        }
+      }
+    }
+  }
+};
+
+}  // namespace cpu_attention
+
+#undef BLOCK_SIZE_ALIGNMENT
+#undef HEAD_SIZE_ALIGNMENT
+#undef MAX_Q_HEAD_NUM_PER_ITER
+
+#endif  // CPU_ATTN_VSX_HPP
@@ -486,6 +486,10 @@ struct FP32Vec16 : public VectorizedRegWrapper<FP32Vec16, 4, float> {

  explicit FP32Vec16(const BF16Vec8& v) : FP32Vec16(FP32Vec8(v)) {};

+  // FP8 stub: dead code on ARM (fp8 KV cache is x86-only), needed for
+  // load_b_pair_vec template to compile on all platforms.
+  explicit FP32Vec16(const BF16Vec32&, int) : Base() {}
+
  explicit FP32Vec16(const FP16Vec16& v) {
    reg.val[0] = Vectorized<float>(vcvt_f32_f16(vget_low_f16(v.reg.val[0])));
    reg.val[1] = Vectorized<float>(vcvt_f32_f16(vget_high_f16(v.reg.val[0])));
@@ -6,6 +6,9 @@

 namespace vec_op {

+struct fp8_e4m3_tag {};
+struct fp8_e5m2_tag {};
+
 #define VLLM_DISPATCH_CASE_FLOATING_TYPES(...)            \
  AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__)    \
  AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) \
@@ -145,6 +148,9 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
  }

  void save(void* ptr) const { *reinterpret_cast<f16x32_t*>(ptr) = reg; }
+
+  explicit BF16Vec32(const uint8_t*, fp8_e4m3_tag) : reg{} {}
+  explicit BF16Vec32(const uint8_t*, fp8_e5m2_tag) : reg{} {}
 };

 struct FP32Vec4 : public Vec<FP32Vec4> {
@@ -302,6 +308,10 @@ struct FP32Vec16 : public Vec<FP32Vec16> {

  FP32Vec16(const BF16Vec8& v) : FP32Vec16(FP32Vec8(v)) {};

+  // FP8 stub: dead code on scalar path (fp8 KV cache is x86-only), needed for
+  // load_b_pair_vec template to compile on all platforms.
+  explicit FP32Vec16(const BF16Vec32&, int) : reg{} {}
+
  FP32Vec16 operator*(const FP32Vec16& b) const {
    f32x16_t ret;
    unroll_loop<int, VEC_ELEM_NUM>(
@@ -9,6 +9,10 @@

 namespace vec_op {

+// FP8 tag types for tag dispatch (see cpu_attn_vec.hpp)
+struct fp8_e4m3_tag {};
+struct fp8_e5m2_tag {};
+
 // FIXME: FP16 is not fully supported in Torch-CPU
 #define VLLM_DISPATCH_CASE_FLOATING_TYPES(...)         \
  AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
@@ -142,6 +146,9 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
      : reg({vec8_data.reg, vec8_data.reg, vec8_data.reg, vec8_data.reg}) {}

  void save(void* ptr) const { *reinterpret_cast<ss16x8x4_t*>(ptr) = reg; }
+
+  explicit BF16Vec32(const uint8_t*, fp8_e4m3_tag) : reg{} {}
+  explicit BF16Vec32(const uint8_t*, fp8_e5m2_tag) : reg{} {}
 };

 struct FP32Vec4 : public Vec<FP32Vec4> {
@@ -404,6 +411,10 @@ struct FP32Vec16 : public Vec<FP32Vec16> {

  explicit FP32Vec16(const BF16Vec8& v) : FP32Vec16(FP32Vec8(v)) {}

+  // FP8 stub: dead code on PowerPC (fp8 KV cache is x86-only), needed for
+  // load_b_pair_vec template to compile on all platforms.
+  explicit FP32Vec16(const BF16Vec32&, int) : reg{} {}
+
  explicit FP32Vec16(const INT32Vec16& v) {
    reg.val[0] = vec_ctf(v.reg.val[0], 0);
    reg.val[1] = vec_ctf(v.reg.val[1], 0);
@@ -688,6 +688,10 @@ struct FP32Vec16 : public Vec<FP32Vec16> {

  explicit FP32Vec16(const BF16Vec8& v) : FP32Vec16(FP32Vec8(v)) {}

+  // FP8 stub: dead code on s390x (fp8 KV cache is x86-only), needed for
+  // load_b_pair_vec template to compile on all platforms.
+  explicit FP32Vec16(const BF16Vec32&, int) : reg{} {}
+
  FP32Vec16 operator*(const FP32Vec16& b) const {
    return FP32Vec16(f32x4x4_t({vec_mul(reg.val[0], b.reg.val[0]),
                                vec_mul(reg.val[1], b.reg.val[1]),
@@ -122,9 +122,17 @@ struct FP16Vec16 : public Vec<FP16Vec16> {
  void save(void* ptr) const { _mm256_storeu_si256((__m256i*)ptr, reg); }

  void save(void* ptr, const int elem_num) const {
+#ifdef __AVX512BW__
    constexpr uint32_t M = 0xFFFFFFFF;
    __mmask16 mask = _cvtu32_mask16(M >> (32 - elem_num));
    _mm256_mask_storeu_epi16(ptr, mask, reg);
+#else
+    // Fallback for lack of 16-bit masked store
+    int16_t tmp[VEC_ELEM_NUM];
+    _mm256_storeu_si256((__m256i*)tmp, reg);
+    for (int i = 0; i < elem_num; ++i)
+      reinterpret_cast<int16_t*>(ptr)[i] = tmp[i];
+#endif
  }
 };

@@ -161,9 +169,17 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
  void save(void* ptr) const { _mm256_storeu_si256((__m256i*)ptr, reg); }

  void save(void* ptr, const int elem_num) const {
+#ifdef __AVX512BW__
    constexpr uint32_t M = 0xFFFFFFFF;
    __mmask16 mask = _cvtu32_mask16(M >> (32 - elem_num));
    _mm256_mask_storeu_epi16(ptr, mask, reg);
+#else
+    // Fallback for lack of 16-bit masked store
+    int16_t tmp[VEC_ELEM_NUM];
+    _mm256_storeu_si256((__m256i*)tmp, reg);
+    for (int i = 0; i < elem_num; ++i)
+      reinterpret_cast<int16_t*>(ptr)[i] = tmp[i];
+#endif
  }
 };

@@ -247,13 +263,12 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
  explicit BF16Vec32(__m256i low, __m256i high)
      : reg_low(low), reg_high(high) {}

+  explicit BF16Vec32()
+      : reg_low(_mm256_setzero_si256()), reg_high(_mm256_setzero_si256()) {}
+
  explicit BF16Vec32(BF16Vec8& vec8_data)
-      : reg_low((__m256i)_mm256_inserti32x4(
-            _mm256_castsi128_si256((__m128i)vec8_data.reg),
-            (__m128i)vec8_data.reg, 1)),
-        reg_high((__m256i)_mm256_inserti32x4(
-            _mm256_castsi128_si256((__m128i)vec8_data.reg),
-            (__m128i)vec8_data.reg, 1)) {}
+      : reg_low(_mm256_broadcastsi128_si256((__m128i)vec8_data.reg)),
+        reg_high(_mm256_broadcastsi128_si256((__m128i)vec8_data.reg)) {}

  // E4M3 decode (AVX2 path) — same bit-layout trick as the AVX512 variant
  // above.  Result = true_E4M3 * 2^-8; caller applies scale * 2^8.
@@ -674,6 +689,11 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
                     _mm256_sub_ps(reg_high, b.reg_high));
  }

+  FP32Vec16 operator-() const {
+    const __m256 neg = _mm256_set1_ps(-0.0f);
+    return FP32Vec16(_mm256_xor_ps(reg_low, neg), _mm256_xor_ps(reg_high, neg));
+  }
+
  FP32Vec16 operator/(const FP32Vec16& b) const {
    return FP32Vec16(_mm256_div_ps(reg_low, b.reg_low),
                     _mm256_div_ps(reg_high, b.reg_high));
@@ -739,6 +759,85 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
    _mm256_storeu_ps(ptr, reg_low);
    _mm256_storeu_ps(ptr + 8, reg_high);
  }
+
+  void save(float* ptr, const int elem_num) const {
+    // Partial store: cmpgt produces a sign-bit mask (0xFFFFFFFF/0 per lane)
+    // for the first elem_num lanes, applied across the two 8-wide halves.
+    if (elem_num <= 8) {
+      __m256i mask =
+          _mm256_cmpgt_epi32(_mm256_set1_epi32(elem_num),
+                             _mm256_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7));
+      _mm256_maskstore_ps(ptr, mask, reg_low);
+    } else {
+      _mm256_storeu_ps(ptr, reg_low);
+      __m256i mask =
+          _mm256_cmpgt_epi32(_mm256_set1_epi32(elem_num - 8),
+                             _mm256_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7));
+      _mm256_maskstore_ps(ptr + 8, mask, reg_high);
+    }
+  }
+
+  FP32Vec16 clamp(const FP32Vec16& min, const FP32Vec16& max) const {
+    return FP32Vec16(
+        _mm256_min_ps(max.reg_low, _mm256_max_ps(min.reg_low, reg_low)),
+        _mm256_min_ps(max.reg_high, _mm256_max_ps(min.reg_high, reg_high)));
+  }
+
+  FP32Vec16 abs() const {
+    const __m256 sign_mask = _mm256_set1_ps(-0.0f);
+    return FP32Vec16(_mm256_andnot_ps(sign_mask, reg_low),
+                     _mm256_andnot_ps(sign_mask, reg_high));
+  }
+
+  FP32Vec16 min(const FP32Vec16& b) const {
+    return FP32Vec16(_mm256_min_ps(reg_low, b.reg_low),
+                     _mm256_min_ps(reg_high, b.reg_high));
+  }
+
+  // Partial element-wise min over the first elem_num lanes only (tail path).
+  // Scalar via AliasReg: AVX2 has no masked vminps, so we spill, loop, reload.
+  FP32Vec16 min(const FP32Vec16& b, const int elem_num) const {
+    AliasReg ar_this_low, ar_this_high, ar_b_low, ar_b_high;
+    ar_this_low.reg = reg_low;
+    ar_this_high.reg = reg_high;
+    ar_b_low.reg = b.reg_low;
+    ar_b_high.reg = b.reg_high;
+    for (int i = 0; i < elem_num && i < 8; ++i)
+      ar_this_low.values[i] =
+          std::min(ar_this_low.values[i], ar_b_low.values[i]);
+    for (int i = 0; i < elem_num - 8 && i < 8; ++i)
+      ar_this_high.values[i] =
+          std::min(ar_this_high.values[i], ar_b_high.values[i]);
+    return FP32Vec16(ar_this_low.reg, ar_this_high.reg);
+  }
+
+  // Partial element-wise max over the first elem_num lanes only (tail path).
+  // Scalar via AliasReg: AVX2 has no masked vmaxps, so we spill, loop, reload.
+  FP32Vec16 max(const FP32Vec16& b, const int elem_num) const {
+    AliasReg ar_this_low, ar_this_high, ar_b_low, ar_b_high;
+    ar_this_low.reg = reg_low;
+    ar_this_high.reg = reg_high;
+    ar_b_low.reg = b.reg_low;
+    ar_b_high.reg = b.reg_high;
+    for (int i = 0; i < elem_num && i < 8; ++i)
+      ar_this_low.values[i] =
+          std::max(ar_this_low.values[i], ar_b_low.values[i]);
+    for (int i = 0; i < elem_num - 8 && i < 8; ++i)
+      ar_this_high.values[i] =
+          std::max(ar_this_high.values[i], ar_b_high.values[i]);
+    return FP32Vec16(ar_this_low.reg, ar_this_high.reg);
+  }
+
+  float reduce_min() const {
+    __m256 v = _mm256_min_ps(reg_low, reg_high);
+    __m256 v_shuffled = _mm256_permute_ps(v, 0b00001011);
+    __m256 v_min = _mm256_min_ps(v, v_shuffled);
+    v_shuffled = _mm256_permute_ps(v_min, 0b00000001);
+    v_min = _mm256_min_ps(v_min, v_shuffled);
+    v_shuffled = _mm256_permute2f128_ps(v_min, v_min, 0b00000001);
+    v_min = _mm256_min_ps(v_min, v_shuffled);
+    return _mm256_cvtss_f32(v_min);
+  }
 };
 #endif

@@ -791,6 +890,34 @@ struct INT8Vec64 : public Vec<INT8Vec64> {
  // non-temporal save
  void nt_save(int8_t* ptr) { _mm512_stream_si512((__m512i*)ptr, reg); }
 };
+#else
+struct INT8Vec16 : public Vec<INT8Vec16> {
+  constexpr static int VEC_ELEM_NUM = 16;
+  union AliasReg {
+    __m128i reg;
+    int8_t values[VEC_ELEM_NUM];
+  };
+
+  __m128i reg;
+
+  explicit INT8Vec16(const FP32Vec16& vec) {
+    __m256i lo_i32 = _mm256_cvtps_epi32(vec.reg_low);
+    __m256i hi_i32 = _mm256_cvtps_epi32(vec.reg_high);
+    __m256i packed16 = _mm256_packs_epi32(lo_i32, hi_i32);
+    packed16 = _mm256_permute4x64_epi64(packed16, 0xD8);
+    __m256i packed8 = _mm256_packs_epi16(packed16, _mm256_setzero_si256());
+    packed8 = _mm256_permute4x64_epi64(packed8, 0xD8);
+    reg = _mm256_castsi256_si128(packed8);
+  }
+
+  void save(int8_t* ptr) const { _mm_storeu_si128((__m128i*)ptr, reg); }
+
+  void save(int8_t* ptr, const int elem_num) const {
+    AliasReg ar;
+    ar.reg = reg;
+    for (int i = 0; i < elem_num; ++i) ptr[i] = ar.values[i];
+  }
+};
 #endif

 template <typename T>
@@ -215,7 +215,7 @@ void dynamic_quant_epilogue(const float* input, scalar_t* output,
        float zp_scale_val = a_scale[i] * static_cast<float>(azp[i]);
        token_zp_scale_vec = cvt_vec_t(zp_scale_val);
      }
-      for (; j < hidden_size - vec_elem_num; ++j) {
+      for (; j < hidden_size - vec_elem_num; j += vec_elem_num) {
        cvt_vec_t elems_fp32(input_ptr + j);
        elems_fp32 = elems_fp32 * token_scale_vec;
        if constexpr (AZP) {
@@ -20,6 +20,7 @@ ISA_TYPES = {
    "VEC16": 2,
    "NEON": 3,
    "VXE": 4,
+    "VSX": 5,
 }

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

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

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

+#ifdef __powerpc__
+  #include "cpu_attn_vsx.hpp"
+#endif
+
 """

    header += generate_helper_function()
@@ -207,6 +212,11 @@ def generate_header_file() -> str:
        ["VXE", "VEC", "VEC16"],
        fp8=False,
    )
+    header += _macro_block(
+        "#elif defined(__powerpc__)",
+        ["VSX", "VEC", "VEC16"],
+        fp8=False,
+    )
    header += _macro_block(
        "#elif defined(__AVX512F__)",
        ["VEC", "VEC16"],
@@ -223,7 +233,8 @@ def generate_header_file() -> str:
        fp8=False,
    )
    header += (
-        "#endif  /* CPU_CAPABILITY_AMXBF16 / __aarch64__ / __s390x__ */\n\n"
+        "#endif  /* CPU_CAPABILITY_AMXBF16 / __aarch64__ / "
+        "__s390x__ / __powerpc__ */\n\n"
        "#endif  // CPU_ATTN_DISPATCH_GENERATED_H\n"
    )

@@ -1,13 +1,12 @@
 // Adapted from
 // https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu

+// clang-format off
+
 #pragma once

 #include <ATen/ATen.h>
 #include <ATen/Parallel.h>
-#include <ATen/record_function.h>
-
-// clang-format off

 #if defined(_OPENMP)
 #include <omp.h>
@@ -16,40 +15,157 @@
 namespace {

 // dispatch bool
-#define AT_DISPATCH_BOOL(BOOL_V, BOOL_NAME, ...)                                 \
-  [&] {                                                                          \
-    if (BOOL_V) {                                                                \
-      constexpr bool BOOL_NAME = true;                                           \
-      return __VA_ARGS__();                                                      \
-    } else {                                                                     \
-      constexpr bool BOOL_NAME = false;                                          \
-      return __VA_ARGS__();                                                      \
-    }                                                                            \
+#define AT_DISPATCH_BOOL(BOOL_V, BOOL_NAME, ...) \
+  [&] {                                          \
+    if (BOOL_V) {                                \
+      constexpr bool BOOL_NAME = true;           \
+      return __VA_ARGS__();                      \
+    } else {                                     \
+      constexpr bool BOOL_NAME = false;          \
+      return __VA_ARGS__();                      \
+    }                                            \
  }()

-// dispatch: bfloat16, float16, int8_t, fp8_e4m3
-#define CPU_DISPATCH_PACKED_TYPES(TYPE, ...)                                    \
-  [&] {                                                                         \
-    switch (TYPE) {                                                             \
-      case at::ScalarType::BFloat16 : {                                         \
-        using packed_t = at::BFloat16;                                          \
-        return __VA_ARGS__();                                                   \
-      }                                                                         \
-      case at::ScalarType::Half: {                                              \
-        using packed_t = at::Half;                                              \
-        return __VA_ARGS__();                                                   \
-      }                                                                         \
-      case at::ScalarType::Char : {                                             \
-        using packed_t = int8_t;                                                \
-        return __VA_ARGS__();                                                   \
-      }                                                                         \
-      case at::ScalarType::Float8_e4m3fn : {                                    \
-        using packed_t = at::Float8_e4m3fn;                                     \
-        return __VA_ARGS__();                                                   \
-      }                                                                         \
-      default:                                                                  \
-        TORCH_CHECK(false, "Unsupported floating data type.\n");                \
-    }                                                                           \
+#define AT_DISPATCH_BOOL2(BOOL_V1, BOOL_NAME1, BOOL_V2, BOOL_NAME2, ...) \
+  [&] {                                                                  \
+    if (BOOL_V1) {                                                       \
+      constexpr bool BOOL_NAME1 = true;                                  \
+      if (BOOL_V2) {                                                     \
+        constexpr bool BOOL_NAME2 = true;                                \
+        return __VA_ARGS__();                                            \
+      } else {                                                           \
+        constexpr bool BOOL_NAME2 = false;                               \
+        return __VA_ARGS__();                                            \
+      }                                                                  \
+    } else {                                                             \
+      constexpr bool BOOL_NAME1 = false;                                 \
+      if (BOOL_V2) {                                                     \
+        constexpr bool BOOL_NAME2 = true;                                \
+        return __VA_ARGS__();                                            \
+      } else {                                                           \
+        constexpr bool BOOL_NAME2 = false;                               \
+        return __VA_ARGS__();                                            \
+      }                                                                  \
+    }                                                                    \
+  }()
+
+// dispatch: bfloat16, float16, int8_t, fp8_e4m3, uint8_t(mxfp4/int4)
+#define CPU_DISPATCH_PACKED_TYPES(TYPE, ...)                     \
+  [&] {                                                          \
+    switch (TYPE) {                                              \
+      case at::ScalarType::BFloat16: {                           \
+        using packed_t = at::BFloat16;                           \
+        return __VA_ARGS__();                                    \
+      }                                                          \
+      case at::ScalarType::Half: {                               \
+        using packed_t = at::Half;                               \
+        return __VA_ARGS__();                                    \
+      }                                                          \
+      case at::ScalarType::Char: {                               \
+        using packed_t = int8_t;                                 \
+        return __VA_ARGS__();                                    \
+      }                                                          \
+      case at::ScalarType::Float8_e4m3fn: {                      \
+        using packed_t = at::Float8_e4m3fn;                      \
+        return __VA_ARGS__();                                    \
+      }                                                          \
+      case at::ScalarType::Byte: {                               \
+        using packed_t = uint8_t;                                \
+        return __VA_ARGS__();                                    \
+      }                                                          \
+      default:                                                   \
+        TORCH_CHECK(false, "Unsupported floating data type.\n"); \
+    }                                                            \
+  }()
+
+// Helper MICRO for CPU_DISPATCH_FLOATING_TYPES_EXT:
+//   TYPE1: the primary dtype (input, output, weight);
+//   TYPE2: defined as PARAM_T input
+#define CPU_DISPATCH_TYPE1_WITH_PARAM(TYPE1, PARAM_T, ...)   \
+  switch (TYPE1) {                                           \
+    case at::ScalarType::BFloat16: {                         \
+      using scalar_t = at::BFloat16;                         \
+      using param_t = PARAM_T;                               \
+      return __VA_ARGS__();                                  \
+    }                                                        \
+    case at::ScalarType::Half: {                             \
+      using scalar_t = at::Half;                             \
+      using param_t = PARAM_T;                               \
+      return __VA_ARGS__();                                  \
+    }                                                        \
+    case at::ScalarType::Float: {                            \
+      using scalar_t = float;                                \
+      using param_t = PARAM_T;                               \
+      return __VA_ARGS__();                                  \
+    }                                                        \
+    default:                                                 \
+      TORCH_CHECK(false, "Unsupported floating data type."); \
+  }
+
+// Helper MICRO for CPU_DISPATCH_REDUCED_FLOATING_TYPES_EXT:
+//   TYPE1: the primary dtype (input, output, weight);
+//   TYPE2: defined as PARAM_T input
+#define CPU_DISPATCH_TYPE1_WITH_PARAM_REDUCED(TYPE1, PARAM_T, ...) \
+  switch (TYPE1) {                                                 \
+    case at::ScalarType::BFloat16: {                               \
+      using scalar_t = at::BFloat16;                               \
+      using param_t = PARAM_T;                                     \
+      return __VA_ARGS__();                                        \
+    }                                                              \
+    case at::ScalarType::Half: {                                   \
+      using scalar_t = at::Half;                                   \
+      using param_t = PARAM_T;                                     \
+      return __VA_ARGS__();                                        \
+    }                                                              \
+    default:                                                       \
+      TORCH_CHECK(false, "Unsupported floating data type.");       \
+  }
+
+// Helper MICRO for CPU_DISPATCH_REDUCED_FLOATING_TYPES_EXT:
+//   TYPE1: the dtype both for scalar_t and param_t
+#define CPU_DISPATCH_TYPE1_WITH_SAME_PARAM_REDUCED(TYPE1, ...)       \
+  switch (TYPE1) {                                                   \
+    case at::ScalarType::BFloat16: {                                 \
+      using scalar_t = at::BFloat16;                                 \
+      using param_t = at::BFloat16;                                  \
+      return __VA_ARGS__();                                          \
+    }                                                                \
+    case at::ScalarType::Half: {                                     \
+      using scalar_t = at::Half;                                     \
+      using param_t = at::Half;                                      \
+      return __VA_ARGS__();                                          \
+    }                                                                \
+    default:                                                         \
+      TORCH_CHECK(false, "Unsupported reduced floating data type."); \
+  }
+
+// dispatch with mixed dtypes (TYPE1, TYPE2):
+//   TYPE1: the primary dtype (input, output, weight);
+//   TYPE2: the secondary dtype (bias, etc.).
+#define CPU_DISPATCH_FLOATING_TYPES_EXT(TYPE1, TYPE2, ...)            \
+  [&] {                                                               \
+    if (TYPE2 == at::kFloat) {                                        \
+      CPU_DISPATCH_TYPE1_WITH_PARAM(TYPE1, float, __VA_ARGS__)        \
+    } else if (TYPE2 == at::ScalarType::BFloat16) {                   \
+      CPU_DISPATCH_TYPE1_WITH_PARAM(TYPE1, at::BFloat16, __VA_ARGS__) \
+    } else if (TYPE2 == at::ScalarType::Half) {                       \
+      CPU_DISPATCH_TYPE1_WITH_PARAM(TYPE1, at::Half, __VA_ARGS__)     \
+    } else {                                                          \
+      TORCH_CHECK(false, "Unsupported floating data type.");          \
+    }                                                                 \
+  }()
+
+// dispatch with mixed dtypes (reduced one, no float for TYPE1) (TYPE1, TYPE2):
+//   TYPE1: the primary dtype (input, output, weight);
+//   TYPE2: the secondary dtype (bias, etc.).
+#define CPU_DISPATCH_REDUCED_FLOATING_TYPES_EXT(TYPE1, TYPE2, ...)     \
+  [&] {                                                                \
+    if (TYPE2 == at::kFloat) {                                         \
+      CPU_DISPATCH_TYPE1_WITH_PARAM_REDUCED(TYPE1, float, __VA_ARGS__) \
+    } else {                                                           \
+      TORCH_CHECK(TYPE1 == TYPE2);                                     \
+      CPU_DISPATCH_TYPE1_WITH_SAME_PARAM_REDUCED(TYPE1, __VA_ARGS__)   \
+    }                                                                  \
  }()

 #define UNUSED(x) (void)(x)
@@ -70,13 +186,51 @@ namespace {
 #define CHECK_DIM(d, x) TORCH_CHECK(x.dim() == d, #x " must be a " #d "D tensor")

 #define CHECK_EQ(a, b) TORCH_CHECK((a) == (b), "CHECK_EQ(" #a ", " #b ") failed. ", a, " vs ", b)
+#define CHECK_GT(a, b) TORCH_CHECK((a) > (b), "CHECK_GT(" #a ", " #b ") failed. ", a, " vs ", b)
+#define CHECK_GE(a, b) TORCH_CHECK((a) >= (b), "CHECK_GE(" #a ", " #b ") failed. ", a, " vs ", b)

-// parallel routines
+template <bool is_only_lastdim_contiguous>
+static inline void CHECK_INPUT_SHAPE_DTYPE(const at::Tensor& tensor, const at::IntArrayRef sizes, at::ScalarType st) {
+  TORCH_CHECK(tensor.sizes() == sizes, "Input tensor shape mismatch: expected ", sizes, ", got ", tensor.sizes());
+  TORCH_CHECK(tensor.scalar_type() == st, "Input tensor dtype mismatch");
+  if constexpr (is_only_lastdim_contiguous) {
+    CHECK_LAST_DIM_CONTIGUOUS_INPUT(tensor);
+  } else {
+    CHECK_INPUT(tensor);
+  }
+}
+
+// [NB] Parallel Routines
+//
+//  * at::parallel_for - applies for most of generic use cases, this will be compiled
+//                       against openmp in default torch release.
+//
+//  * parallel_for     - same function as above, can choose payload partition scheme in
+//                       balance211.
+//
+//  * parallel_2d      - parallel for 2 dimensions, used in GEMM, etc.
+//                       this one will do payload balance across 2 dimensions.
+//
+
+// grain size for each thread
 constexpr int GRAIN_SIZE = 1024;

 template <typename T, typename std::enable_if<std::is_integral<T>::value, int>::type = 0>
-inline T div_up(T x, T y) { return (x + y - 1) / y; }
+inline T div_up(T x, T y) {
+  return (x + y - 1) / y;
+}

+// you can only use at::get_thread_num() with at::parallel_for()
+// as it is lazy initialized, otherwise it will always return 0.
+inline int get_thread_num() {
+#if defined(_OPENMP)
+  return omp_get_thread_num();
+#else
+  return 0;
+#endif
+}
+
+// balance payload across each thread
 template <typename T>
 inline void balance211(T n, T nth, T ith, T& n_start, T& n_end) {
 #if 0
@@ -94,10 +248,10 @@ inline void balance211(T n, T nth, T ith, T& n_start, T& n_end) {
    }
    n_end += n_start;
 #else
-    // pytorch aten partition pattern
-    T n_my = div_up(n, nth);
-    n_start = ith * n_my;
-    n_end = std::min(n_start + n_my, n);
+  // pytorch aten partition pattern
+  T n_my = div_up(n, nth);
+  n_start = ith * n_my;
+  n_end = std::min(n_start + n_my, n);
 #endif
 }

@@ -105,23 +259,15 @@ template <typename func_t>
 inline void parallel_for(int n, const func_t& f) {
 #if defined(_OPENMP)
 #pragma omp parallel
-{
+  {
    int nth = omp_get_num_threads();
    int ith = omp_get_thread_num();
    int tbegin, tend;
    balance211(n, nth, ith, tbegin, tend);
    f(tbegin, tend);
-}
+  }
 #else
-    f(0, n);
-#endif
-}
-
-inline int get_thread_num() {
-#if defined(_OPENMP)
-  return omp_get_thread_num();
-#else
-  return 0;
+  f(0, n);
 #endif
 }

@@ -137,7 +283,6 @@ int inline adjust_num_threads(int m) {

 template <typename func_t>
 inline void parallel_2d(int m, int n, const func_t& f) {
-
  // make sure we have even num_threads
  int nth = adjust_num_threads(m);

@@ -165,31 +310,59 @@ inline void parallel_2d(int m, int n, const func_t& f) {

 #if defined(_OPENMP)
 #pragma omp parallel num_threads(nth)
-{
-  int ith = omp_get_thread_num();
-  int ith_m = ith / nth_n;
-  int ith_n = ith % nth_n;
+  {
+    int ith = omp_get_thread_num();
+    int ith_m = ith / nth_n;
+    int ith_n = ith % nth_n;

-  int thread_block_m = div_up(m, nth_m);
-  int thread_block_n = div_up(n, nth_n);
+    int thread_block_m = div_up(m, nth_m);
+    int thread_block_n = div_up(n, nth_n);

-  int begin_m = ith_m * thread_block_m;
-  int end_m = std::min(m, begin_m + thread_block_m);
-  int begin_n = ith_n * thread_block_n;
-  int end_n = std::min(n, begin_n + thread_block_n);
+    int begin_m = ith_m * thread_block_m;
+    int end_m = std::min(m, begin_m + thread_block_m);
+    int begin_n = ith_n * thread_block_n;
+    int end_n = std::min(n, begin_n + thread_block_n);

-  f(begin_m, end_m, begin_n, end_n);
-}
+    f(begin_m, end_m, begin_n, end_n);
+  }
 #else
  f(0, m, 0, n);
 #endif
 }

+// limit max cache blocks
+// when we need to do pre-unpack for weights, e.g. fp8
+#define MAX_CACHE_BLOCK_SIZE 4
+
 template <typename T>
-int get_cache_blocks(int BLOCK_SIZE, int K) {
+inline int get_cache_blocks(int chunk_size) {
  // L2 2MB and ratio of 50%
  const int L2_size = 2048 * 1024 >> 1;
-  return std::max(1, int(L2_size / (BLOCK_SIZE * K * sizeof(T))));
+  return std::max(1, int(L2_size / (chunk_size * sizeof(T))));
+}
+
+template <>
+inline int get_cache_blocks<at::Float8_e4m3fn>(int chunk_size) {
+  // fp8 uses bf16 as accumulate type
+  int cache_block_size = get_cache_blocks<at::BFloat16>(chunk_size);
+  return std::min(MAX_CACHE_BLOCK_SIZE, cache_block_size);
+}
+
+// 2d sequential loop in range : [mb0, mb1), [nb0, nb1)
+template <typename T, typename func_t>
+inline void loop_2d(int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1, int64_t chunk_size, const func_t& f) {
+  // get number of blocks for L2 in most inner loop
+  int64_t cache_blocks_nb = get_cache_blocks<T>(chunk_size);
+
+  // loop order: [NB / cache_blocks_nb, MB, cache_blocks_nb]
+  // TODO: implement reverse order of [MB / cache_blocks_mb, NB, cache_blocks_mb]
+  for (int64_t nbb = nb0; nbb < nb1; nbb += cache_blocks_nb) {
+    for (int64_t mb = mb0; mb < mb1; ++mb) {
+      for (int64_t nb = nbb; nb < std::min(nbb + cache_blocks_nb, nb1); ++nb) {
+        f(mb, nb, nb - nbb);
+      }
+    }
+  }
 }

 // data indexing for dimension collapse
@@ -243,4 +416,10 @@ struct Unroll<1> {
  }
 };

-} // anonymous namespace
+// conditional data ptr for optional tensor
+template <typename T>
+inline T* conditional_data_ptr(const std::optional<at::Tensor>& opt) {
+  return opt.has_value() ? opt.value().data_ptr<T>() : nullptr;
+}
+
+}  // anonymous namespace
@@ -0,0 +1,709 @@
+// Adapted from
+// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu
+
+// clang-format off
+
+#include "common.h"
+#include "gemm.h"
+#include "vec.h"
+
+namespace {
+
+template <typename scalar_t>
+inline void copy_stub(scalar_t* __restrict__ y, const scalar_t* __restrict__ x, int64_t size) {
+  using Vec = at::vec::Vectorized<scalar_t>;
+  const bool is_padding = (x == nullptr);
+  for (int64_t d = 0; d < size; d += Vec::size()) {
+    Vec data_vec = is_padding ? Vec(0.f) : Vec::loadu(x + d);
+    data_vec.store(y + d);
+  }
+}
+
+// no remainder
+template <typename scalar_t>
+void inline update_conv_state(
+    scalar_t* __restrict__ conv_states,
+    const scalar_t* __restrict__ input,
+    int64_t width,
+    int64_t dim,
+    int64_t seqlen,
+    bool has_initial_states) {
+  // width for `conv_states`
+  int64_t width1 = width - 1;
+  int64_t w = 0;
+  for (; w < width1 - seqlen; ++w) {
+    scalar_t* y = conv_states + w * dim;
+    const scalar_t* x = has_initial_states ? conv_states + (w + seqlen) * dim : nullptr;
+    copy_stub(y, x, dim);
+  }
+  for (; w < width1; ++w) {
+    scalar_t* y = conv_states + w * dim;
+    const scalar_t* x = input + (w + seqlen - width1) * dim;
+    copy_stub(y, x, dim);
+  }
+}
+
+// A : [M, BLOCK_N]
+// B : [BLOCK_N, K], prepacked as [K/2, BLOCK_N, 2]
+// C : [M, BLOCK_N]
+// bias : [BLOCK_N]
+//
+// lda : leading dimension of `input` and `out`
+//
+template <typename scalar_t, int K, int BLOCK_N, bool has_bias, bool has_silu>
+struct tinygemm_kernel {
+  static inline void apply(
+      const scalar_t* __restrict__ A,
+      const scalar_t* __restrict__ B,
+      scalar_t* __restrict__ C,
+      const scalar_t* __restrict__ bias,
+      const scalar_t* __restrict__ conv_states,
+      bool has_initial_state,
+      int64_t M,
+      int64_t lda,
+      bool is_first_token) {
+    TORCH_CHECK(false, "tinygemm_kernel_nn: scalar path not implemented!");
+  }
+};
+
+#if defined(CPU_CAPABILITY_AVX512)
+template <int K, int BLOCK_N, bool has_bias, bool has_silu>
+struct tinygemm_kernel<at::BFloat16, K, BLOCK_N, has_bias, has_silu> {
+  static inline void apply(
+      const at::BFloat16* __restrict__ A,
+      const at::BFloat16* __restrict__ B,
+      at::BFloat16* __restrict__ C,
+      const at::BFloat16* __restrict__ bias,
+      const at::BFloat16* __restrict__ conv_states,
+      bool has_initial_state,
+      int64_t M,
+      int64_t lda,
+      bool is_first_token) {
+    assert(K == 4);
+    constexpr int ROWS = K;
+    constexpr int COLS = BLOCK_N / block_size_n();
+
+    // leading dimension size for b for next block [K/2, 32, 2]
+    constexpr int ldb = block_size_n() * K;
+
+    __m512bh va[ROWS * COLS];
+    __m512bh vb[ROWS * COLS];
+    __m512 vc[COLS * 2];
+
+    // k: {-3, -2, -1} -> {0, 1, 2}
+    auto set_conv_states = [&](int k, int col) -> __m512i {
+      return has_initial_state ? _mm512_loadu_si512(conv_states + (k + K - 1) * lda + col * 32)
+                               : _mm512_setzero_si512();
+    };
+
+#define MM512_LOAD_A(idx)                                                 \
+  ((idx) < 0 && is_first_token) ? (__m512bh)(set_conv_states((idx), col)) \
+                                : (__m512bh)(_mm512_loadu_si512(A + (idx) * lda + col * 32))
+
+#define MM512_PACK_A(ap, bp, a, b)                       \
+  do {                                                   \
+    __m512i r0 = (__m512i)(a);                           \
+    __m512i r1 = (__m512i)(b);                           \
+    __m512i d0 = _mm512_unpacklo_epi16(r0, r1);          \
+    __m512i d1 = _mm512_unpackhi_epi16(r0, r1);          \
+    r0 = _mm512_shuffle_i32x4(d0, d1, 0x88);             \
+    r1 = _mm512_shuffle_i32x4(d0, d1, 0xdd);             \
+    (ap) = (__m512bh)_mm512_shuffle_i32x4(r0, r1, 0x88); \
+    (bp) = (__m512bh)_mm512_shuffle_i32x4(r0, r1, 0xdd); \
+  } while (0)
+
+    // step 0 : preload a at time step [-3][-2][-1]
+    auto preloada = [&](auto i) {
+      constexpr int col = i;
+      int64_t m = 0;
+      va[1 * COLS + col] = MM512_LOAD_A(m - 3);
+      va[2 * COLS + col] = MM512_LOAD_A(m - 2);
+      va[3 * COLS + col] = MM512_LOAD_A(m - 1);
+    };
+    Unroll<COLS>{}(preloada);
+
+    auto loada = [&](auto i, int64_t m) {
+      constexpr int col = i;
+      // update previous time step
+      va[0 * COLS + col] = va[1 * COLS + col];
+      va[1 * COLS + col] = va[2 * COLS + col];
+      va[2 * COLS + col] = va[3 * COLS + col];
+      // load current time step
+      va[3 * COLS + col] = MM512_LOAD_A(m);
+    };
+
+    // step 1 : load weight for just once
+    auto loadb = [&](auto i) {
+      constexpr int row = i / COLS;
+      constexpr int col = i % COLS;
+      vb[row * COLS + col] = (__m512bh)(_mm512_loadu_si512(B + col * ldb + row * 32));
+    };
+    Unroll<ROWS * COLS>{}(loadb);
+
+    // [NB] accumulates 4x32 bfloat16 blocks
+    //
+    //   +------------+------------+
+    //   |    col0    |    col1    |
+    //   +------------+------------+
+    //   |  va0  va1  |  va0  va1  |
+    //   |  va2  va3  |  va2  va3  |
+    //   +------------+------------+
+    //   |  vc0  vc1  |  vc0  vc1  |
+    //   +------------+------------+
+    //
+    //  * va and vb shares the same memory layout
+    //  * block_n 32 with 4 rows equals to 4 registers
+    //  * 37 uops with avx512bf16 v.s. 57 uops with avx512f
+    //
+    auto compute = [&](auto i) {
+      constexpr int col = i;
+
+      // init accumulators
+      if constexpr (has_bias) {
+        __m512i b16 = _mm512_loadu_si512(reinterpret_cast<const __m512i*>(bias + col * 32));
+        vc[col * 2 + 0] = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(b16, 0));
+        vc[col * 2 + 1] = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(b16, 1));
+      } else {
+        vc[col * 2 + 0] = _mm512_set1_ps(0.f);
+        vc[col * 2 + 1] = _mm512_set1_ps(0.f);
+      }
+
+      // convert to vnni2 format
+      __m512bh va0, va1, va2, va3;
+      MM512_PACK_A(va0, va1, va[0 * COLS + col], va[1 * COLS + col]);
+      MM512_PACK_A(va2, va3, va[2 * COLS + col], va[3 * COLS + col]);
+
+      // accumulate
+      vc[col * 2 + 0] = _mm512_dpbf16_ps(vc[col * 2 + 0], va0, vb[0 * COLS + col]);
+      vc[col * 2 + 0] = _mm512_dpbf16_ps(vc[col * 2 + 0], va2, vb[2 * COLS + col]);
+      vc[col * 2 + 1] = _mm512_dpbf16_ps(vc[col * 2 + 1], va1, vb[1 * COLS + col]);
+      vc[col * 2 + 1] = _mm512_dpbf16_ps(vc[col * 2 + 1], va3, vb[3 * COLS + col]);
+    };
+
+    using fVec = at::vec::Vectorized<float>;
+    using bVec = at::vec::Vectorized<at::BFloat16>;
+    const fVec one = fVec(1.f);
+    auto storec = [&](auto i, int64_t m) {
+      constexpr int col = i;
+      fVec x0 = fVec(vc[col * 2 + 0]);
+      fVec x1 = fVec(vc[col * 2 + 1]);
+      if constexpr (has_silu) {
+        x0 = x0 / (one + x0.neg().exp_u20());
+        x1 = x1 / (one + x1.neg().exp_u20());
+      }
+      bVec out_vec = convert_from_float_ext<at::BFloat16>(x0, x1);
+      out_vec.store(C + m * lda + col * 32);
+    };
+
+    for (int64_t m = 0; m < M; ++m) {
+      // step 3.a : load a at current time step
+      Unroll<COLS>{}(loada, m);
+      // step 3.b : accumulate for window size (4)
+      Unroll<COLS>{}(compute);
+      // step 3.c : store c at current time step
+      Unroll<COLS>{}(storec, m);
+    }
+  }
+};
+#endif
+
+#define LAUNCH_TINYGEMM_KERNEL(K, NB_SIZE)                                                   \
+  tinygemm_kernel<scalar_t, K, NB_SIZE, has_bias, has_silu>::apply(                          \
+      input + bs * seqlen * dim + mb_start * dim + nb_start,                                 \
+      weight + nb_start * width,                                                             \
+      out + bs * seqlen * dim + mb_start * dim + nb_start,                                   \
+      has_bias ? bias + nb_start : nullptr,                                                  \
+      has_conv_states ? conv_states + conv_state_index * (K - 1) * dim + nb_start : nullptr, \
+      has_initial_states_value,                                                              \
+      mb_size,                                                                               \
+      dim,                                                                                   \
+      mb_start == 0);
+
+template <typename scalar_t>
+void causal_conv1d_fwd_kernel_impl(
+    scalar_t* __restrict__ out,
+    const scalar_t* __restrict__ input,
+    const scalar_t* __restrict__ weight,
+    const scalar_t* __restrict__ bias,
+    scalar_t* __restrict__ conv_states,
+    const int32_t* __restrict__ conv_indices,
+    const bool* __restrict__ has_initial_state,
+    bool silu_activation,
+    int64_t batch,
+    int64_t dim,
+    int64_t seqlen,
+    int64_t width,
+    int64_t num_seq_blocks) {
+  // handle 32 x 64 per block
+  constexpr int64_t BLOCK_M = block_size_m();
+  constexpr int64_t BLOCK_N = block_size_n() * 2;
+  const int64_t NB = div_up(dim, BLOCK_N);
+
+  const int64_t num_blocks_per_seq = div_up(seqlen, BLOCK_M);
+  const bool has_conv_states = conv_states != nullptr;
+  const bool has_conv_indices = conv_indices != nullptr;
+
+  // parallel on [batch, seq, NB]
+  AT_DISPATCH_BOOL2(bias != nullptr, has_bias, silu_activation, has_silu, [&] {
+    at::parallel_for(0, num_seq_blocks * NB, 0, [&](int64_t begin, int64_t end) {
+      int64_t mb{0}, nb{0};
+      data_index_init(begin, mb, num_seq_blocks, nb, NB);
+
+      for (int64_t i = begin; i < end; ++i) {
+        int64_t bs = mb / num_blocks_per_seq;
+
+        int64_t mb_start = (mb % num_blocks_per_seq) * BLOCK_M;
+        int64_t mb_size = std::min(seqlen - mb_start, BLOCK_M);
+        int64_t nb_start = nb * BLOCK_N;
+        int64_t nb_size = std::min(dim - nb_start, BLOCK_N);
+
+        const bool has_initial_states_value = has_conv_states ? has_initial_state[bs] : false;
+        int32_t conv_state_index = has_conv_indices ? conv_indices[bs] : bs;
+
+        switch (width << 4 | nb_size >> 4) {
+          case 0x42:
+            LAUNCH_TINYGEMM_KERNEL(4, 32);
+            break;
+          case 0x44:
+            LAUNCH_TINYGEMM_KERNEL(4, 64);
+            break;
+          default:
+            TORCH_CHECK(false, "Unexpected block size, ", width, " x ", nb_size);
+        }
+
+        // move to the next index
+        data_index_step(mb, num_seq_blocks, nb, NB);
+      }
+    });
+  });
+
+  // update conv_states if necessary
+  if (has_conv_states) {
+    at::parallel_for(0, batch, 0, [&](int64_t begin, int64_t end) {
+      for (int64_t bs = begin; bs < end; ++bs) {
+        update_conv_state(
+            conv_states + bs * (width - 1) * dim, input + bs * seqlen * dim, width, dim, seqlen, has_initial_state[bs]);
+      }
+    });
+  }
+}
+
+#define LAUNCH_TINYGEMM_VARLEN_KERNEL(K, NB_SIZE)                   \
+  tinygemm_kernel<scalar_t, K, NB_SIZE, has_bias, has_silu>::apply( \
+      input + batch_offset * dim + mb_start * dim + nb_start,       \
+      weight + nb_start * width,                                    \
+      out + batch_offset * dim + mb_start * dim + nb_start,         \
+      has_bias ? bias + nb_start : nullptr,                         \
+      nullptr,                                                      \
+      false,                                                        \
+      mb_size,                                                      \
+      dim,                                                          \
+      mb_start == 0);
+
+// TODO: add `has_initial_state` support for varlen kernel
+template <typename scalar_t>
+void causal_conv1d_fwd_varlen_kernel_impl(
+    scalar_t* __restrict__ out,
+    const scalar_t* __restrict__ input,
+    const scalar_t* __restrict__ weight,
+    const scalar_t* __restrict__ bias,
+    scalar_t* __restrict__ conv_states,
+    const int32_t* __restrict__ query_start_loc,
+    const int32_t* __restrict__ conv_indices,
+    const bool* __restrict__ has_initial_state,
+    const int32_t* __restrict__ block_indices,
+    bool silu_activation,
+    int64_t batch,
+    int64_t dim,
+    int64_t width,
+    int64_t num_seq_blocks) {
+  // handle 32 x 64 per block
+  constexpr int64_t BLOCK_M = block_size_m();
+  constexpr int64_t BLOCK_N = block_size_n() * 2;
+  const int64_t NB = div_up(dim, BLOCK_N);
+
+  const bool has_conv_states = conv_states != nullptr;
+  const bool has_conv_indices = conv_indices != nullptr;
+
+  // parallel on [batch, seq, NB]
+  AT_DISPATCH_BOOL2(bias != nullptr, has_bias, silu_activation, has_silu, [&] {
+    at::parallel_for(0, num_seq_blocks * NB, 0, [&](int64_t begin, int64_t end) {
+      int64_t mb{0}, nb{0};
+      data_index_init(begin, mb, num_seq_blocks, nb, NB);
+
+      for (int64_t i = begin; i < end; ++i) {
+        int32_t bs = block_indices[mb * 2 + 0];
+        int32_t batch_offset = query_start_loc[bs];
+        int32_t seqlen = query_start_loc[bs + 1] - query_start_loc[bs];
+
+        int64_t mb_start = block_indices[mb * 2 + 1] * BLOCK_M;
+        int64_t mb_size = std::min(seqlen - mb_start, BLOCK_M);
+        int64_t nb_start = nb * BLOCK_N;
+        int64_t nb_size = std::min(dim - nb_start, BLOCK_N);
+
+        switch (width << 4 | nb_size >> 4) {
+          case 0x42:
+            LAUNCH_TINYGEMM_VARLEN_KERNEL(4, 32);
+            break;
+          case 0x44:
+            LAUNCH_TINYGEMM_VARLEN_KERNEL(4, 64);
+            break;
+          default:
+            TORCH_CHECK(false, "Unexpected block size, ", width, " x ", nb_size);
+        }
+
+        // move to the next index
+        data_index_step(mb, num_seq_blocks, nb, NB);
+      }
+    });
+  });
+
+  // update conv_states if necessary
+  if (has_conv_states) {
+    at::parallel_for(0, batch, 0, [&](int64_t begin, int64_t end) {
+      for (int64_t bs = begin; bs < end; ++bs) {
+        int32_t conv_state_index = has_conv_indices ? conv_indices[bs] : bs;
+        int32_t seqlen = query_start_loc[bs + 1] - query_start_loc[bs];
+        int32_t batch_offset = query_start_loc[bs];
+        update_conv_state(
+            conv_states + conv_state_index * (width - 1) * dim,
+            input + batch_offset * dim,
+            width,
+            dim,
+            seqlen,
+            /* has_initial_state */ false);
+      }
+    });
+  }
+}
+
+template <typename scalar_t>
+void causal_conv1d_update_kernel_impl(
+    scalar_t* __restrict__ out,
+    const scalar_t* __restrict__ input,
+    scalar_t* __restrict__ conv_states,
+    const scalar_t* __restrict__ weight,
+    const scalar_t* __restrict__ bias,
+    const int32_t* __restrict__ conv_indices,
+    bool silu_activation,
+    int64_t batch,
+    int64_t dim,
+    int64_t seqlen,
+    int64_t width) {
+  // handle 32 x 64 per block
+  constexpr int64_t BLOCK_M = block_size_m();
+  constexpr int64_t BLOCK_N = block_size_n() * 2;
+  const int64_t NB = div_up(dim, BLOCK_N);
+
+  const bool has_conv_states = conv_states != nullptr;
+  const bool has_conv_indices = conv_indices != nullptr;
+
+  // parallel on [batch, NB]
+  AT_DISPATCH_BOOL2(bias != nullptr, has_bias, silu_activation, has_silu, [&] {
+    at::parallel_for(0, batch * NB, 0, [&](int64_t begin, int64_t end) {
+      int64_t bs{0}, nb{0};
+      data_index_init(begin, bs, batch, nb, NB);
+
+      for (int64_t i = begin; i < end; ++i) {
+        int64_t mb_start = 0;
+        int64_t mb_size = 1;
+        int64_t nb_start = nb * BLOCK_N;
+        int64_t nb_size = std::min(dim - nb_start, BLOCK_N);
+
+        const bool has_initial_states_value = true;
+        int32_t conv_state_index = has_conv_indices ? conv_indices[bs] : bs;
+
+        switch (width << 4 | nb_size >> 4) {
+          case 0x42:
+            LAUNCH_TINYGEMM_KERNEL(4, 32);
+            break;
+          case 0x44:
+            LAUNCH_TINYGEMM_KERNEL(4, 64);
+            break;
+          default:
+            TORCH_CHECK(false, "Unexpected block size, ", width, " x ", nb_size);
+        }
+
+        // move to the next index
+        data_index_step(bs, batch, nb, NB);
+      }
+    });
+  });
+
+#define CONV_STATE_INDEXR(w) conv_states + conv_state_index*(width - 1) * dim + (w) * dim
+
+  // update conv_states
+  at::parallel_for(0, batch, 0, [&](int64_t begin, int64_t end) {
+    for (int64_t bs = begin; bs < end; ++bs) {
+      // update old states, range [1, width - 1)
+      int32_t conv_state_index = has_conv_indices ? conv_indices[bs] : bs;
+      for (int64_t w = 1; w < width - 1; ++w) {
+        std::memcpy(CONV_STATE_INDEXR(w - 1), CONV_STATE_INDEXR(w), dim * sizeof(scalar_t));
+      }
+      // copy new states
+      std::memcpy(CONV_STATE_INDEXR(width - 2), input + bs * dim, dim * sizeof(scalar_t));
+    }
+  });
+}
+
+}  // anonymous namespace
+
+// from [dim, width] or [N, K]
+// to [N/BLOCK_N, K/2, BLOCK_N, 2]
+at::Tensor causal_conv1d_weight_pack(const at::Tensor& weight) {
+  CHECK_INPUT(weight);
+
+  int64_t dim = weight.size(0);
+  int64_t width = weight.size(1);
+  constexpr int64_t BLOCK_N = block_size_n();
+  TORCH_CHECK(width == 4, "causal_conv1d_weight_pack: support only width of 4");
+  TORCH_CHECK(dim % BLOCK_N == 0, "causal_conv1d_weight_pack: invalid dim size ", dim);
+
+  const int64_t N = dim, K2 = width >> 1;
+  const int64_t NB = div_up(N, BLOCK_N);
+
+  auto packed_weight = at::empty_like(weight);
+  AT_DISPATCH_REDUCED_FLOATING_TYPES(weight.scalar_type(), "causal_conv1d_fwd_kernel_impl", [&] {
+    // cast to float32 as vnni size is 2
+    const float* w_data = reinterpret_cast<float*>(weight.data_ptr<scalar_t>());
+    float* packed_data = reinterpret_cast<float*>(packed_weight.data_ptr<scalar_t>());
+
+    at::parallel_for(0, NB * K2 * BLOCK_N, 0, [&](int64_t begin, int64_t end) {
+      int64_t nb{0}, k2{0}, n{0};
+      data_index_init(begin, nb, NB, k2, K2, n, BLOCK_N);
+
+      // TODO: optimize this if we need to online prepacking.
+      for (int64_t i = begin; i < end; ++i) {
+        packed_data[i] = w_data[nb * BLOCK_N * K2 + n * K2 + k2];
+
+        // move to the next index
+        data_index_step(nb, NB, k2, K2, n, BLOCK_N);
+      }
+    });
+  });
+  return packed_weight;
+}
+
+#define CHECK_OPTIONAL_SHAPE_DTYPE(OPT, SIZE, DTYPE) \
+  if (OPT.has_value()) {                             \
+    const auto tensor = OPT.value();                 \
+    CHECK_CONTIGUOUS(tensor);                        \
+    CHECK_EQ(tensor.size(0), SIZE);                  \
+    CHECK_EQ(tensor.scalar_type(), DTYPE);           \
+  }
+
+template <int BLOCK_M>
+int64_t get_block_count(const std::optional<at::Tensor>& offsets, int64_t batch, int64_t seqlen) {
+  if (offsets.has_value()) {
+    const int32_t* offsets_data = offsets.value().data_ptr<int32_t>();
+    int32_t num_seq_blocks = 0;
+    for (int64_t row = 0; row < batch; ++row) {
+      num_seq_blocks += div_up(offsets_data[row + 1] - offsets_data[row], BLOCK_M);
+    }
+    return num_seq_blocks;
+  }
+  return batch * div_up(seqlen, int64_t(BLOCK_M));
+}
+
+template <int BLOCK_M>
+at::Tensor get_block_indices(const std::optional<at::Tensor>& offsets, int64_t num_seq_blocks) {
+  if (!offsets.has_value()) {
+    return at::Tensor();
+  }
+
+  const at::Tensor& offsets_ = offsets.value();
+  at::Tensor indices = at::empty({num_seq_blocks, 2}, offsets_.options());
+
+  int64_t batch = offsets_.size(0) - 1;
+
+  const int32_t* offsets_data = offsets_.data_ptr<int32_t>();
+  int32_t* indices_data = indices.data_ptr<int32_t>();
+
+  int64_t idx = 0;
+  for (int32_t row = 0; row < batch; ++row) {
+    int32_t blocks = div_up(offsets_data[row + 1] - offsets_data[row], BLOCK_M);
+
+    for (int32_t col = 0; col < blocks; ++col) {
+      indices_data[idx * 2 + 0] = row;
+      indices_data[idx * 2 + 1] = col;
+      idx++;
+    }
+  }
+  return indices;
+}
+
+// API aligned with GPUs
+//
+//   x: (batch, dim, seqlen) or (dim, cu_seq_len) for varlen
+//   weight: (dim, width)
+//   bias: (dim,)
+//   query_start_loc: (batch + 1) int32
+//   cache_indices: (batch)  int32
+//   has_initial_state: (batch) bool
+//   conv_states: (..., dim, width - 1) itype
+//   activation: either None or "silu" or "swish"
+//   pad_slot_id: int
+//
+at::Tensor causal_conv1d_fwd_cpu(
+    const at::Tensor& x,
+    const at::Tensor& weight,
+    const std::optional<at::Tensor>& bias,
+    const std::optional<at::Tensor>& conv_states,
+    const std::optional<at::Tensor>& query_start_loc,
+    const std::optional<at::Tensor>& conv_state_indices,
+    const std::optional<at::Tensor>& has_initial_state,
+    bool silu_activation,
+    int64_t pad_slot_id,
+    bool is_vnni) {
+  CHECK_CONTIGUOUS(weight);
+  auto packed_w = is_vnni ? weight : causal_conv1d_weight_pack(weight);
+
+  const bool is_var_seqlen = query_start_loc.has_value();
+  const int64_t input_ndim = is_var_seqlen ? 2 : 3;
+  TORCH_CHECK(x.dim() == input_ndim, "causal_conv1d_fwd_cpu: expect x to be ", input_ndim, "D tensor.");
+  TORCH_CHECK(x.stride(-2) == 1 && x.stride(-1) == x.size(-2), "causal_conv1d_fwd_cpu: expect x to be transposed.");
+
+  const int64_t batch = is_var_seqlen ? query_start_loc.value().size(0) - 1 : x.size(0);
+  const int64_t dim = x.size(-2);
+  const int64_t seqlen = x.size(-1);
+  const int64_t width = weight.size(-1);
+
+  const auto scalar_type = x.scalar_type();
+  CHECK_EQ(weight.scalar_type(), scalar_type);
+  CHECK_OPTIONAL_SHAPE_DTYPE(bias, dim, scalar_type);
+  CHECK_OPTIONAL_SHAPE_DTYPE(query_start_loc, batch + 1, at::kInt);
+  CHECK_OPTIONAL_SHAPE_DTYPE(conv_state_indices, batch, at::kInt);
+  CHECK_OPTIONAL_SHAPE_DTYPE(has_initial_state, batch, at::kBool);
+
+  if (conv_states.has_value()) {
+    auto& conv_states_val = conv_states.value();
+    int64_t padded_batch = conv_states_val.size(0);
+    CHECK_EQ(conv_states_val.scalar_type(), scalar_type);
+    CHECK_GE(padded_batch, batch);
+    CHECK_EQ(conv_states_val.size(1), dim);
+    CHECK_EQ(conv_states_val.size(2), width - 1);
+
+    // adjust `conv_states` to be contiguous on `dim`
+    // should happen only once
+    if (conv_states_val.stride(-2) != 1) {
+      auto conv_states_copy = conv_states_val.clone();
+      conv_states_val.as_strided_({padded_batch, dim, width - 1}, {(width - 1) * dim, 1, dim});
+      conv_states_val.copy_(conv_states_copy);
+    }
+  }
+
+  // block size for sequence blocks, 32
+  constexpr int64_t BLOCK_M = block_size_m();
+
+  // total number of sequence blocks
+  int64_t num_seq_blocks = get_block_count<BLOCK_M>(query_start_loc, batch, seqlen);
+
+  at::Tensor out = at::empty_like(x);
+  AT_DISPATCH_REDUCED_FLOATING_TYPES(scalar_type, "causal_conv1d_fwd_kernel_impl", [&] {
+    if (is_var_seqlen) {
+      // record seq blocks in Coordinate format, aka [num_seq_blocks, 2]
+      at::Tensor block_indices = get_block_indices<BLOCK_M>(query_start_loc, num_seq_blocks);
+
+      causal_conv1d_fwd_varlen_kernel_impl(
+          out.data_ptr<scalar_t>(),
+          x.data_ptr<scalar_t>(),
+          packed_w.data_ptr<scalar_t>(),
+          conditional_data_ptr<scalar_t>(bias),
+          conditional_data_ptr<scalar_t>(conv_states),
+          conditional_data_ptr<int32_t>(query_start_loc),
+          conditional_data_ptr<int32_t>(conv_state_indices),
+          conditional_data_ptr<bool>(has_initial_state),
+          block_indices.data_ptr<int32_t>(),
+          silu_activation,
+          batch,
+          dim,
+          width,
+          num_seq_blocks);
+    } else {
+      causal_conv1d_fwd_kernel_impl<scalar_t>(
+          out.data_ptr<scalar_t>(),
+          x.data_ptr<scalar_t>(),
+          packed_w.data_ptr<scalar_t>(),
+          conditional_data_ptr<scalar_t>(bias),
+          conditional_data_ptr<scalar_t>(conv_states),
+          conditional_data_ptr<int32_t>(conv_state_indices),
+          conditional_data_ptr<bool>(has_initial_state),
+          silu_activation,
+          batch,
+          dim,
+          seqlen,
+          width,
+          num_seq_blocks);
+    }
+  });
+  return out;
+}
+
+// API aligned with GPUs
+//
+//   x: (batch, dim) or (batch, dim, seqlen)
+//   conv_state: (..., dim, state_len), where state_len >= width - 1
+//   weight: (dim, width)
+//   bias: (dim,)
+//   cache_seqlens: (batch,), dtype int32.
+//   conv_state_indices: (batch,), dtype int32
+//   pad_slot_id: int
+//   out: (batch, dim) or (batch, dim, seqlen)
+//
+at::Tensor causal_conv1d_update_cpu(
+    const at::Tensor& x,
+    const at::Tensor& conv_states,
+    const at::Tensor& weight,
+    const std::optional<at::Tensor>& bias,
+    bool silu_activation,
+    const std::optional<at::Tensor>& cache_seqlens,
+    const std::optional<at::Tensor>& conv_state_indices,
+    int64_t pad_slot_id,
+    bool is_vnni) {
+  CHECK_CONTIGUOUS(x);
+  CHECK_CONTIGUOUS(weight);
+  auto packed_w = is_vnni ? weight : causal_conv1d_weight_pack(weight);
+
+  // TODO: add multi-token prediction support
+  TORCH_CHECK(x.dim() == 2, "causal_conv1d_update_cpu: expect x to be 2D tensor.");
+  TORCH_CHECK(!cache_seqlens.has_value(), "causal_conv1d_update_cpu: don't support cache_seqlens.");
+
+  int64_t batch = x.size(0);
+  int64_t dim = x.size(1);
+  int64_t seqlen = 1;
+  int64_t width = weight.size(-1);
+
+  const auto scalar_type = x.scalar_type();
+  CHECK_EQ(weight.scalar_type(), scalar_type);
+  CHECK_OPTIONAL_SHAPE_DTYPE(bias, dim, scalar_type);
+  CHECK_OPTIONAL_SHAPE_DTYPE(conv_state_indices, batch, at::kInt);
+
+  CHECK_EQ(conv_states.scalar_type(), scalar_type);
+  CHECK_EQ(conv_states.size(1), dim);
+  CHECK_EQ(conv_states.size(2), width - 1);
+
+  // adjust `conv_states` to be contiguous on `dim`
+  if (conv_states.stride(-2) != 1) {
+    int64_t num_cache_lines = conv_states.size(0);
+    auto conv_states_copy = conv_states.clone();
+    conv_states.as_strided_({num_cache_lines, dim, width - 1}, {(width - 1) * dim, 1, dim});
+    conv_states.copy_(conv_states_copy);
+  }
+
+  at::Tensor out = at::empty_like(x);
+  AT_DISPATCH_REDUCED_FLOATING_TYPES(scalar_type, "causal_conv1d_update_kernel_impl", [&] {
+    causal_conv1d_update_kernel_impl<scalar_t>(
+        out.data_ptr<scalar_t>(),
+        x.data_ptr<scalar_t>(),
+        conv_states.data_ptr<scalar_t>(),
+        packed_w.data_ptr<scalar_t>(),
+        conditional_data_ptr<scalar_t>(bias),
+        conditional_data_ptr<int32_t>(conv_state_indices),
+        silu_activation,
+        batch,
+        dim,
+        seqlen,
+        width);
+  });
+  return out;
+}
@@ -1,11 +1,12 @@
 // Adapted from
 // https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu

-#include "common.h"
-#include "vec.h"
+// clang-format off
+
 #include "gemm.h"

-// clang-format off
+#include "common.h"
+#include "vec.h"

 namespace {

@@ -26,13 +27,13 @@ inline void s8s8_compensation(int8_t* __restrict__ packed, int K) {
  const __m512i off = _mm512_set1_epi8(static_cast<char>(0x80));
  for (int k = 0; k < K / 4; ++k) {
    for (int col = 0; col < COLS; ++col) {
-      __m512i vb = _mm512_loadu_si512((const __m512i *)(packed + k * BLOCK_N * 4 + col * 64));
+      __m512i vb = _mm512_loadu_si512((const __m512i*)(packed + k * BLOCK_N * 4 + col * 64));
      vcomp[col] = _mm512_dpbusd_epi32(vcomp[col], off, vb);
    }
  }

  for (int col = 0; col < COLS; ++col) {
-    _mm512_storeu_si512((__m512i *)(packed + offset + col * 64), vcomp[col]);
+    _mm512_storeu_si512((__m512i*)(packed + offset + col * 64), vcomp[col]);
  }
 #else
  TORCH_CHECK(false, "s8s8_compensation not implemented!");
@@ -69,6 +70,43 @@ inline void pack_vnni<int8_t>(int8_t* __restrict__ packed, const int8_t* __restr
  s8s8_compensation<BLOCK_N>(packed, K);
 }

+// uint8_t: mxfp4 or int4
+// pack to vnni2 format as they are computed with bfloat16
+//
+// from [N, K'/2, 2] to [K'/2, N, 2], view 2x int4 as unit8:
+// from [N,    K   ] to [K,    N   ] where K = K'/2
+//
+template <>
+inline void pack_vnni<uint8_t>(uint8_t* __restrict__ packed, const uint8_t* __restrict__ weight, int N, int K) {
+  constexpr int BLOCK_N = block_size_n();
+
+  uint8_t unpacked[2 * BLOCK_N];
+
+  // 32-way pack (align with BLOCK_N), faster for avx512 unpacking
+  //
+  // for a range of (64):
+  //   {0, 1, 2, ..., 63}
+  //
+  // original format:
+  //   { 1|0,  3|2, ..., 63|62}
+  //
+  // packed format:
+  //   {32|0, 31|1, ..., 63|31}
+  //
+  for (int k = 0; k < K; ++k) {
+    // unpack first
+    for (int n = 0; n < N; ++n) {
+      uint8_t value = weight[n * K + k];
+      unpacked[n * 2 + 0] = value & 0xF;  // lower 4 bits
+      unpacked[n * 2 + 1] = value >> 4;   // higher 4 bits
+    }
+    // re-pack to 32-way
+    for (int n = 0; n < N; ++n) {
+      packed[k * N + n] = (unpacked[n + BLOCK_N] << 4) | unpacked[n];
+    }
+  }
+}
+
 template <typename scalar_t>
 inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ input, int64_t size) {
  using bVec = at::vec::Vectorized<scalar_t>;
@@ -76,7 +114,7 @@ inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ inpu
  constexpr int kVecSize = bVec::size();

  int64_t d;
-  #pragma GCC unroll 4
+#pragma GCC unroll 4
  for (d = 0; d <= size - kVecSize; d += kVecSize) {
    fVec data0 = fVec::loadu(input + d);
    fVec data1 = fVec::loadu(input + d + fVec::size());
@@ -89,13 +127,34 @@ inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ inpu
 }

 template <typename scalar_t>
-inline void copy_add_stub(scalar_t* __restrict__ out, const float* __restrict__ input, const float* __restrict__ bias, int64_t size) {
+inline void copy_stub(float* __restrict__ out, const scalar_t* __restrict__ input, int64_t size) {
  using bVec = at::vec::Vectorized<scalar_t>;
  using fVec = at::vec::Vectorized<float>;
  constexpr int kVecSize = bVec::size();

  int64_t d;
-  #pragma GCC unroll 4
+#pragma GCC unroll 4
+  for (d = 0; d <= size - kVecSize; d += kVecSize) {
+    fVec data0, data1;
+    bVec b_vec = bVec::loadu(input + d);
+    std::tie(data0, data1) = at::vec::convert_to_float(b_vec);
+    data0.store(out + d);
+    data1.store(out + d + fVec::size());
+  }
+  for (; d < size; ++d) {
+    out[d] = static_cast<float>(input[d]);
+  }
+}
+
+template <typename scalar_t>
+inline void copy_add_stub(
+    scalar_t* __restrict__ out, const float* __restrict__ input, const float* __restrict__ bias, int64_t size) {
+  using bVec = at::vec::Vectorized<scalar_t>;
+  using fVec = at::vec::Vectorized<float>;
+  constexpr int kVecSize = bVec::size();
+
+  int64_t d;
+#pragma GCC unroll 4
  for (d = 0; d <= size - kVecSize; d += kVecSize) {
    fVec data0 = fVec::loadu(input + d) + fVec::loadu(bias + d);
    fVec data1 = fVec::loadu(input + d + fVec::size()) + fVec::loadu(bias + d + fVec::size());
@@ -107,11 +166,51 @@ inline void copy_add_stub(scalar_t* __restrict__ out, const float* __restrict__
  }
 }

+template <typename scalar_t, bool has_bias>
+inline void scalar_sigmoid_and_mul(
+    scalar_t* __restrict__ out,
+    const float* __restrict__ input,
+    const float* __restrict__ bias,
+    const scalar_t* __restrict__ mul,
+    int SIZE) {
+  using bVec = at::vec::Vectorized<scalar_t>;
+  using fVec = at::vec::Vectorized<float>;
+  // scalar sigmoid
+  const fVec one = fVec(1.f);
+  fVec X;
+  if constexpr (has_bias) {
+    assert(bias != nullptr);
+    X = fVec(input[0] + bias[0]);
+  } else {
+    X = fVec(input[0]);
+  }
+  X = one / (one + X.neg().exp_u20());
+
+  // vec mul
+  constexpr int kVecSize = bVec::size();
+  for (int d = 0; d < SIZE; d += kVecSize) {
+    bVec m_bvec = bVec::loadu(mul + d);
+    fVec m_fvec0, m_fvec1;
+    std::tie(m_fvec0, m_fvec1) = at::vec::convert_to_float(m_bvec);
+    m_fvec0 = m_fvec0 * X;
+    m_fvec1 = m_fvec1 * X;
+
+    bVec out_vec = convert_from_float_ext<scalar_t>(m_fvec0, m_fvec1);
+    out_vec.store(out + d);
+  }
+}
+
 template <typename scalar_t, bool has_bias, int BLOCK_M, int BLOCK_N>
 struct tinygemm_kernel_nn {
  static inline void apply(
-      const scalar_t* __restrict__ A, const scalar_t* __restrict__ B, scalar_t* __restrict__ C,
-      const float* __restrict__ bias, int64_t K, int64_t lda, int64_t ldb, int64_t ldc) {
+      const scalar_t* __restrict__ A,
+      const scalar_t* __restrict__ B,
+      scalar_t* __restrict__ C,
+      const float* __restrict__ bias,
+      int64_t K,
+      int64_t lda,
+      int64_t ldb,
+      int64_t ldc) {
    TORCH_CHECK(false, "tinygemm_kernel_nn: scalar path not implemented!");
  }
 };
@@ -120,9 +219,14 @@ struct tinygemm_kernel_nn {
 template <bool has_bias, int BLOCK_M, int BLOCK_N>
 struct tinygemm_kernel_nn<at::BFloat16, has_bias, BLOCK_M, BLOCK_N> {
  static inline void apply(
-      const at::BFloat16* __restrict__ A, const at::BFloat16* __restrict__ B, at::BFloat16* __restrict__ C,
-      const float* __restrict__ bias, int64_t K, int64_t lda, int64_t ldb, int64_t ldc) {
-
+      const at::BFloat16* __restrict__ A,
+      const at::BFloat16* __restrict__ B,
+      at::BFloat16* __restrict__ C,
+      const float* __restrict__ bias,
+      int64_t K,
+      int64_t lda,
+      int64_t ldb,
+      int64_t ldc) {
    constexpr int ROWS = BLOCK_M;
    constexpr int COLS = BLOCK_N / 16;

@@ -145,7 +249,7 @@ struct tinygemm_kernel_nn<at::BFloat16, has_bias, BLOCK_M, BLOCK_N> {

    const int64_t K2 = K >> 1;
    const int64_t lda2 = lda >> 1;
-    const int64_t ldb2 = ldb; // ldb * 2 >> 1;
+    const int64_t ldb2 = ldb;  // ldb * 2 >> 1;
    const float* a_ptr = reinterpret_cast<const float*>(A);
    const float* b_ptr = reinterpret_cast<const float*>(B);

@@ -180,9 +284,7 @@ struct tinygemm_kernel_nn<at::BFloat16, has_bias, BLOCK_M, BLOCK_N> {
              (__m512i)(_mm512_cvtne2ps_pbh(vc[row * COLS + col + 1], vc[row * COLS + col])));
        }
      } else {
-        _mm256_storeu_si256(
-            reinterpret_cast<__m256i*>(C + row * ldc + col * 16),
-            (__m256i)(_mm512_cvtneps_pbh(vc[i])));
+        _mm256_storeu_si256(reinterpret_cast<__m256i*>(C + row * ldc + col * 16), (__m256i)(_mm512_cvtneps_pbh(vc[i])));
      }
    };
    Unroll<ROWS * COLS>{}(storec);
@@ -190,22 +292,33 @@ struct tinygemm_kernel_nn<at::BFloat16, has_bias, BLOCK_M, BLOCK_N> {
 };
 #endif

-#define LAUNCH_TINYGEMM_KERNEL_NN(MB_SIZE, NB_SIZE)                          \
-    tinygemm_kernel_nn<scalar_t, has_bias, MB_SIZE, NB_SIZE>::apply(         \
-        A + mb_start * lda, B + nb_start * 2, C + mb_start * ldc + nb_start, \
-        has_bias ? bias + nb_start : nullptr, K, lda, ldb, ldc);
+#define LAUNCH_TINYGEMM_KERNEL_NN(MB_SIZE, NB_SIZE)                \
+  tinygemm_kernel_nn<scalar_t, has_bias, MB_SIZE, NB_SIZE>::apply( \
+      A + mb_start * lda,                                          \
+      B + nb_start * 2,                                            \
+      C + mb_start * ldc + nb_start,                               \
+      has_bias ? bias + nb_start : nullptr,                        \
+      K,                                                           \
+      lda,                                                         \
+      ldb,                                                         \
+      ldc);

 template <typename scalar_t, bool has_bias>
 struct brgemm {
  static inline void apply(
-      const scalar_t* __restrict__ A, const scalar_t* __restrict__ B, scalar_t* __restrict__ C,
-      float* __restrict__ Ctmp, const float* __restrict__ bias,
-      int64_t M, int64_t N, int64_t K, int64_t lda, int64_t ldb, int64_t ldc) {
-
+      const scalar_t* __restrict__ A,
+      const scalar_t* __restrict__ B,
+      scalar_t* __restrict__ C,
+      float* __restrict__ Ctmp,
+      const float* __restrict__ bias,
+      int64_t M,
+      int64_t N,
+      int64_t K,
+      int64_t lda,
+      int64_t ldb,
+      int64_t ldc) {
    constexpr int BLOCK_N = block_size_n();
-    at::native::cpublas::brgemm(
-        M, N, K, lda, ldb, BLOCK_N, /* add_C */false,
-        A, B, Ctmp);
+    at::native::cpublas::brgemm(M, N, K, lda, ldb, BLOCK_N, /* add_C */ false, A, B, Ctmp);

    // copy from Ctmp to C
    for (int64_t m = 0; m < M; ++m) {
@@ -216,6 +329,21 @@ struct brgemm {
      }
    }
  }
+  static inline void apply(
+      const float* __restrict__ A,
+      const float* __restrict__ B,
+      scalar_t* __restrict__ C,
+      float* __restrict__ Ctmp,
+      const float* __restrict__ bias,
+      int64_t M,
+      int64_t N,
+      int64_t K,
+      int64_t lda,
+      int64_t ldb,
+      int64_t ldc) {
+    constexpr int BLOCK_N = block_size_n();
+    at::native::cpublas::brgemm(M, N, K, lda, ldb, BLOCK_N, /* add_C */ false, A, B, Ctmp);
+  }
 };

 template <typename scalar_t, bool has_bias>
@@ -232,15 +360,12 @@ void tinygemm_kernel(
    int64_t ldb,
    int64_t ldc,
    bool brg) {
-
  if (brg) {
-    brgemm<scalar_t, has_bias>::apply(
-        A, B, C, Ctmp, bias,
-        M, N, K, lda, ldb, ldc);
+    brgemm<scalar_t, has_bias>::apply(A, B, C, Ctmp, bias, M, N, K, lda, ldb, ldc);
    return;
  }

-  // pattern: 1-4-16
+  // pattern: 1-4-16, N = 16, 32, 48, 64
  constexpr int64_t BLOCK_M = 4;
  constexpr int64_t BLOCK_N = 64;
  const int64_t MB = div_up(M, BLOCK_M);
@@ -252,25 +377,88 @@ void tinygemm_kernel(
      int64_t nb_start = nb * BLOCK_N;
      int64_t nb_size = std::min(BLOCK_N, N - nb_start);

-      switch(mb_size << 4 | nb_size >> 4) {
+      switch (mb_size << 4 | nb_size >> 4) {
        // mb_size = 1
-        case 0x12: LAUNCH_TINYGEMM_KERNEL_NN(1, 32); break;
-        case 0x14: LAUNCH_TINYGEMM_KERNEL_NN(1, 64); break;
+        case 0x11:
+          LAUNCH_TINYGEMM_KERNEL_NN(1, 16);
+          break;
+        case 0x12:
+          LAUNCH_TINYGEMM_KERNEL_NN(1, 32);
+          break;
+        case 0x13:
+          LAUNCH_TINYGEMM_KERNEL_NN(1, 48);
+          break;
+        case 0x14:
+          LAUNCH_TINYGEMM_KERNEL_NN(1, 64);
+          break;
        // mb_size = 2
-        case 0x22: LAUNCH_TINYGEMM_KERNEL_NN(2, 32); break;
-        case 0x24: LAUNCH_TINYGEMM_KERNEL_NN(2, 64); break;
+        case 0x21:
+          LAUNCH_TINYGEMM_KERNEL_NN(2, 16);
+          break;
+        case 0x22:
+          LAUNCH_TINYGEMM_KERNEL_NN(2, 32);
+          break;
+        case 0x23:
+          LAUNCH_TINYGEMM_KERNEL_NN(2, 48);
+          break;
+        case 0x24:
+          LAUNCH_TINYGEMM_KERNEL_NN(2, 64);
+          break;
        // mb_size = 3
-        case 0x32: LAUNCH_TINYGEMM_KERNEL_NN(3, 32); break;
-        case 0x34: LAUNCH_TINYGEMM_KERNEL_NN(3, 64); break;
+        case 0x31:
+          LAUNCH_TINYGEMM_KERNEL_NN(3, 16);
+          break;
+        case 0x32:
+          LAUNCH_TINYGEMM_KERNEL_NN(3, 32);
+          break;
+        case 0x33:
+          LAUNCH_TINYGEMM_KERNEL_NN(3, 48);
+          break;
+        case 0x34:
+          LAUNCH_TINYGEMM_KERNEL_NN(3, 64);
+          break;
        // mb_size = 4
-        case 0x42: LAUNCH_TINYGEMM_KERNEL_NN(4, 32); break;
-        case 0x44: LAUNCH_TINYGEMM_KERNEL_NN(4, 64); break;
-        default: TORCH_CHECK(false, "Unexpected block size, ", mb_size, "x", nb_size);
+        case 0x41:
+          LAUNCH_TINYGEMM_KERNEL_NN(4, 16);
+          break;
+        case 0x42:
+          LAUNCH_TINYGEMM_KERNEL_NN(4, 32);
+          break;
+        case 0x43:
+          LAUNCH_TINYGEMM_KERNEL_NN(4, 48);
+          break;
+        case 0x44:
+          LAUNCH_TINYGEMM_KERNEL_NN(4, 64);
+          break;
+        default:
+          TORCH_CHECK(false, "Unexpected block size, ", mb_size, " x ", nb_size);
      }
    }
  }
 }

+template <typename scalar_t, bool has_bias>
+void tinygemm_kernel(
+    const float* __restrict__ A,
+    const float* __restrict__ B,
+    scalar_t* __restrict__ C,
+    float* __restrict__ Ctmp,
+    const float* __restrict__ bias,
+    int64_t M,
+    int64_t N,
+    int64_t K,
+    int64_t lda,
+    int64_t ldb,
+    int64_t ldc,
+    bool brg) {
+  TORCH_CHECK(brg, "Expected to use fp32 brgemm for small N GEMM");
+  if (brg) {
+    brgemm<scalar_t, has_bias>::apply(A, B, C, Ctmp, bias, M, N, K, lda, ldb, ldc);
+    return;
+  }
+  // TODO : add intrinsic path
+}
+
 template <typename scalar_t>
 void weight_packed_linear_kernel_impl(
    scalar_t* __restrict__ out,
@@ -282,29 +470,20 @@ void weight_packed_linear_kernel_impl(
    int64_t K,
    int64_t mat1_strideM,
    int64_t out_strideM) {
-
  constexpr int64_t BLOCK_M = block_size_m();
  constexpr int64_t BLOCK_N = block_size_n();
  const int64_t MB = div_up(M, BLOCK_M);
  const int64_t NB = div_up(N, BLOCK_N);

-  // use avx512-bf16 when a) M is small; b) dtype is bfloat16, otherwise use amx
-  const bool use_brgemm = (M > 4) || (!std::is_same_v<scalar_t, at::BFloat16>);
-
-  // l2 cache block for n
-  int64_t cache_blocks_nb = get_cache_blocks<scalar_t>(BLOCK_N, K);
+  const bool use_brgemm = can_use_brgemm<scalar_t>(M);

  // parallel on [MB, NB]
  AT_DISPATCH_BOOL(bias != nullptr, has_bias, [&] {
-    parallel_2d(MB, NB, [&](int64_t begin_mb, int64_t end_mb, int64_t begin_nb, int64_t end_nb) {
-
+    parallel_2d(MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
      // for brgemm, use float32 for accumulate
      alignas(64) float Ctmp[BLOCK_M * BLOCK_N];

-      for (int64_t nbb = begin_nb; nbb < end_nb; nbb += cache_blocks_nb) {
-      for (int64_t mb = begin_mb; mb < end_mb; ++mb) {
-      for (int64_t nb = nbb; nb < std::min(nbb + cache_blocks_nb, end_nb); ++nb) {
-
+      loop_2d<scalar_t>(mb0, mb1, nb0, nb1, BLOCK_N * K, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
        int64_t mb_start = mb * BLOCK_M;
        int64_t mb_size = std::min(M - mb_start, BLOCK_M);
        int64_t nb_start = nb * BLOCK_N;
@@ -323,7 +502,7 @@ void weight_packed_linear_kernel_impl(
            /* ldb */ nb_size,
            /* ldc */ out_strideM,
            /* brg */ use_brgemm);
-      }}}
+      });

      if (use_brgemm) {
        at::native::cpublas::brgemm_release();
@@ -332,20 +511,113 @@ void weight_packed_linear_kernel_impl(
  });
 }

-} // anonymous namespace
+template <typename scalar_t>
+void weight_packed_linear_kernel_impl(
+    scalar_t* __restrict__ out,
+    const scalar_t* __restrict__ mat1,
+    const float* __restrict__ mat2,
+    const float* __restrict__ bias,
+    const scalar_t* __restrict__ post_mul_mat,
+    int64_t M,
+    int64_t N,
+    int64_t K,
+    int64_t mat1_strideM,
+    int64_t out_strideM) {
+  constexpr int64_t BLOCK_M = block_size_m();
+  constexpr int64_t BLOCK_N = block_size_n();
+  const int64_t MB = div_up(M, BLOCK_M);
+  const int64_t NB = div_up(N, BLOCK_N);
+
+  const bool use_brgemm = true;  // TODO: add intrinsic path
+  // parallel on [MB, NB]
+  AT_DISPATCH_BOOL(bias != nullptr, has_bias, [&] {
+    parallel_2d(MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
+      // for brgemm, use float32 for accumulate
+      alignas(64) float Atmp[BLOCK_M * K];
+      alignas(64) float Ctmp[BLOCK_M * BLOCK_N];
+
+      loop_2d<float>(mb0, mb1, nb0, nb1, BLOCK_N * K, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
+        int64_t mb_start = mb * BLOCK_M;
+        int64_t mb_size = std::min(M - mb_start, BLOCK_M);
+        int64_t nb_start = nb * BLOCK_N;
+        int64_t nb_size = std::min(N - nb_start, BLOCK_N);
+        for (int64_t m = 0; m < mb_size; ++m) {
+          copy_stub<scalar_t>(Atmp + m * K, mat1 + mb_start * mat1_strideM + m * K, K);
+        }
+        tinygemm_kernel<scalar_t, has_bias>(
+            /*   A */ Atmp,
+            /*   B */ mat2 + nb_start * K /* nb * BLOCK_N * K */,
+            /*   C */ out + mb_start * out_strideM + nb_start,
+            /* Ctmp*/ Ctmp,
+            /* bias*/ bias + nb_start,
+            /*   M */ mb_size,
+            /*   N */ nb_size,
+            /*   K */ K,
+            /* lda */ mat1_strideM,
+            /* ldb */ nb_size,
+            /* ldc */ out_strideM,
+            /* brg */ use_brgemm);
+
+        if (post_mul_mat != nullptr) {
+          for (int64_t m = 0; m < mb_size; ++m) {
+            scalar_sigmoid_and_mul<scalar_t, has_bias>(
+                out + mb_start * out_strideM + nb_start + m * out_strideM,
+                Ctmp + m * BLOCK_N,
+                bias + nb_start,
+                post_mul_mat + mb_start * out_strideM + m * out_strideM,
+                out_strideM);
+          }
+        } else {
+          for (int64_t m = 0; m < mb_size; ++m) {
+            if constexpr (has_bias) {
+              copy_add_stub(
+                  out + mb_start * out_strideM + nb_start + m * out_strideM, Ctmp + m * BLOCK_N, bias + nb_start, N);
+            } else {
+              copy_stub(out + mb_start * out_strideM + nb_start + m * out_strideM, Ctmp + m * BLOCK_N, N);
+            }
+          }
+        }
+      });
+
+      if (use_brgemm) {
+        at::native::cpublas::brgemm_release();
+      }
+    });
+  });
+}
+
+}  // anonymous namespace

 // tinygemm interface
 template <typename scalar_t>
-void tinygemm_kernel(const scalar_t* __restrict__ A, const scalar_t* __restrict__ B, scalar_t* __restrict__ C,
-    float* __restrict__ Ctmp, int64_t M, int64_t N, int64_t K, int64_t lda, int64_t ldb, int64_t ldc, bool brg) {
+void tinygemm_kernel(
+    const scalar_t* __restrict__ A,
+    const scalar_t* __restrict__ B,
+    scalar_t* __restrict__ C,
+    float* __restrict__ Ctmp,
+    int64_t M,
+    int64_t N,
+    int64_t K,
+    int64_t lda,
+    int64_t ldb,
+    int64_t ldc,
+    bool brg) {
  tinygemm_kernel<scalar_t, false>(A, B, C, Ctmp, nullptr, M, N, K, lda, ldb, ldc, brg);
 }

-#define INSTANTIATE_TINYGEMM_TEMPLATE(TYPE)                                             \
-    template void tinygemm_kernel<TYPE>(                                                \
-        const TYPE* __restrict__ A, const TYPE* __restrict__ B, TYPE* __restrict__ C,   \
-        float* __restrict__ Ctmp, int64_t M, int64_t N, int64_t K, int64_t lda,         \
-        int64_t ldb, int64_t ldc, bool brg)
+#define INSTANTIATE_TINYGEMM_TEMPLATE(TYPE) \
+  template void tinygemm_kernel<TYPE>(      \
+      const TYPE* __restrict__ A,           \
+      const TYPE* __restrict__ B,           \
+      TYPE* __restrict__ C,                 \
+      float* __restrict__ Ctmp,             \
+      int64_t M,                            \
+      int64_t N,                            \
+      int64_t K,                            \
+      int64_t lda,                          \
+      int64_t ldb,                          \
+      int64_t ldc,                          \
+      bool brg)

 INSTANTIATE_TINYGEMM_TEMPLATE(at::BFloat16);
 INSTANTIATE_TINYGEMM_TEMPLATE(at::Half);
@@ -359,14 +631,23 @@ at::Tensor convert_weight_packed(at::Tensor& weight) {

  const int64_t ndim = weight.ndimension();
  TORCH_CHECK(ndim == 2 || ndim == 3, "expect weight to be 2d or 3d, got ", ndim, "d tensor.");
+
+  if (ndim == 2 && weight.size(0) < TILE_N) {
+    // for 2D weight and small OC shape, we use fma linear path, which needs transpose not pack
+    return weight.to(at::kFloat).t().contiguous();
+  }
+
  const auto st = weight.scalar_type();
  const int64_t E = ndim == 3 ? weight.size(0) : 1;
  const int64_t OC = ndim == 3 ? weight.size(1) : weight.size(0);
  const int64_t IC = ndim == 3 ? weight.size(2) : weight.size(1);

+  // mxfp4 or int4 are packed with uint8
+  const int64_t actual_IC = st == at::kByte ? IC * 2 : IC;
+
  // we handle 2 TILE_N at a time.
  TORCH_CHECK(OC % TILE_N == 0, "invalid weight out features ", OC);
-  TORCH_CHECK(IC % TILE_K == 0, "invalid weight input features ", IC);
+  TORCH_CHECK(actual_IC % TILE_K == 0, "invalid weight input features ", actual_IC);

  constexpr int64_t BLOCK_N = block_size_n();
  const int64_t NB = div_up(OC, BLOCK_N);
@@ -375,12 +656,14 @@ at::Tensor convert_weight_packed(at::Tensor& weight) {
  auto packed_weight = at::empty({}, weight.options());
  const int64_t stride = OC * IC;

-  TORCH_CHECK(st == at::kBFloat16 || st == at::kHalf || st == at::kChar || st == at::kFloat8_e4m3fn,
-      "expect weight to be bfloat16, float16, int8 or fp8_e4m3.");
+  // Note: for `kByte` (uint8), it represents either `mxfp4` or `int4`.
+  TORCH_CHECK(
+      st == at::kBFloat16 || st == at::kHalf || st == at::kChar || st == at::kFloat8_e4m3fn || st == at::kByte,
+      "expect weight to be bfloat16, float16, int8, fp8_e4m3 or uint8(mxfp4 or int4).");

  CPU_DISPATCH_PACKED_TYPES(st, [&] {
    // adjust most inner dimension size
-    const int packed_row_size = get_row_size<packed_t>(IC);
+    const int packed_row_size = get_row_size<packed_t>(actual_IC);
    auto sizes = weight.sizes().vec();
    sizes[ndim - 1] = packed_row_size;
    packed_weight.resize_(sizes);
@@ -399,10 +682,7 @@ at::Tensor convert_weight_packed(at::Tensor& weight) {
        int64_t n = nb * BLOCK_N;
        int64_t n_size = std::min(BLOCK_N, OC - n);
        pack_vnni<packed_t>(
-            packed_data + e * OC * packed_row_size + n * packed_row_size,
-            w_data + e * stride + n * IC,
-            n_size,
-            IC);
+            packed_data + e * OC * packed_row_size + n * packed_row_size, w_data + e * stride + n * IC, n_size, IC);

        // move to the next index
        data_index_step(e, E, nb, NB);
@@ -412,33 +692,71 @@ at::Tensor convert_weight_packed(at::Tensor& weight) {
  return packed_weight;
 }

+at::Tensor convert_scale_packed(at::Tensor& scale) {
+  CHECK_INPUT(scale);
+
+  const int64_t ndim = scale.ndimension();
+  TORCH_CHECK(ndim == 2 || ndim == 3, "expect scale to be 2d or 3d, got ", ndim, "d tensor.");
+  const auto st = scale.scalar_type();
+  const int64_t E = ndim == 3 ? scale.size(0) : 1;
+  const int64_t N = ndim == 3 ? scale.size(1) : scale.size(0);
+  // number of groups, e.g. K/32
+  const int64_t G = ndim == 3 ? scale.size(2) : scale.size(1);
+
+  constexpr int64_t BLOCK_N = block_size_n();
+  TORCH_CHECK(N % BLOCK_N == 0, "invalid weight out features ", N);
+  const int64_t NB = N / BLOCK_N;
+
+  auto packed_scale = at::empty_like(scale);
+  TORCH_CHECK(st == at::kByte, "expect scale to be uint8.");
+
+  const uint8_t* s_data = scale.data_ptr<uint8_t>();
+  uint8_t* packed_data = packed_scale.data_ptr<uint8_t>();
+
+  // parallel on src {E, NB, BLOCK_N, G}, dst {E, NB, G, BLOCK_N}
+  at::parallel_for(0, E * NB * BLOCK_N * G, 0, [&](int64_t begin, int64_t end) {
+    int64_t e{0}, nb{0}, n{0}, g{0};
+    data_index_init(begin, e, E, nb, NB, n, BLOCK_N, g, G);
+
+    for (int64_t i = begin; i < end; ++i) {
+      packed_data[e * N * G + nb * G * BLOCK_N + g * BLOCK_N + n] = s_data[i];
+      // move to the next index
+      data_index_step(e, E, nb, NB, n, BLOCK_N, g, G);
+    }
+  });
+  return packed_scale;
+}
+
 // mat1 : [M, K]
-// mat2 : [N, K]
+// mat2 : [N, K] ([K, N] if use_fma_gemm)
 // bias : [N]
 // out  : [M, N]
 //
-at::Tensor weight_packed_linear(at::Tensor& mat1, at::Tensor& mat2,
-    const std::optional<at::Tensor>& bias, bool is_vnni) {
-  RECORD_FUNCTION(
-    "sgl-kernel::weight_packed_linear", std::vector<c10::IValue>({mat1, mat2, bias}));
-
+at::Tensor
+weight_packed_linear(at::Tensor& mat1, at::Tensor& mat2, const std::optional<at::Tensor>& bias, bool is_vnni) {
  auto packed_w = is_vnni ? mat2 : convert_weight_packed(mat2);
+  bool use_fma_gemm = false;
+  if (packed_w.scalar_type() == at::kFloat) {
+    use_fma_gemm = true;
+  }
+
+  int64_t M = mat1.size(0);
+  int64_t K = mat1.size(1);
+  int64_t N = use_fma_gemm ? mat2.size(1) : mat2.size(0);

  CHECK_LAST_DIM_CONTIGUOUS_INPUT(mat1);
  CHECK_INPUT(mat2);
-
-  int64_t M = mat1.size(0);
-  int64_t N = mat2.size(0);
-  int64_t K = mat2.size(1);
-  CHECK_EQ(mat1.size(1), K);
  CHECK_DIM(2, mat1);
  CHECK_DIM(2, mat2);
+  if (!use_fma_gemm) {
+    CHECK_EQ(mat1.size(1), K);
+  }

+  auto dispatch_type = mat1.scalar_type();
  auto out = at::empty({M, N}, mat1.options());
-
  // strides
-  int64_t mat1_strideM = mat1.stride(0);
  int64_t out_strideM = out.stride(0);
+  int64_t mat1_strideM = mat1.stride(0);

  const bool has_bias = bias.has_value();
  const float* bias_data = nullptr;
@@ -447,12 +765,83 @@ at::Tensor weight_packed_linear(at::Tensor& mat1, at::Tensor& mat2,
    bias_data = bias.value().data_ptr<float>();
  }

-  AT_DISPATCH_REDUCED_FLOATING_TYPES(mat1.scalar_type(), "weight_packed_linear_kernel_impl", [&] {
+  AT_DISPATCH_REDUCED_FLOATING_TYPES(dispatch_type, "weight_packed_linear_kernel_impl", [&] {
+    if (use_fma_gemm) {
+      weight_packed_linear_kernel_impl<scalar_t>(
+          out.data_ptr<scalar_t>(),
+          mat1.data_ptr<scalar_t>(),
+          packed_w.data_ptr<float>(),
+          bias_data,
+          nullptr,
+          M,
+          N,
+          K,
+          mat1_strideM,
+          out_strideM);
+    } else {
+      weight_packed_linear_kernel_impl<scalar_t>(
+          out.data_ptr<scalar_t>(),
+          mat1.data_ptr<scalar_t>(),
+          packed_w.data_ptr<scalar_t>(),
+          bias_data,
+          M,
+          N,
+          K,
+          mat1_strideM,
+          out_strideM);
+    }
+  });
+
+  return out;
+}
+
+// mat1         : [M, K]
+// mat2         : [K, 1]
+// post_mul_mat : [M, K]
+// bias         : [N]
+// out          : [M, N]
+//
+at::Tensor fused_linear_sigmoid_mul(
+    at::Tensor& mat1,
+    at::Tensor& mat2,
+    const std::optional<at::Tensor>& bias,
+    bool is_vnni,
+    const at::Tensor& post_mul_mat) {
+  auto packed_w = is_vnni ? mat2 : convert_weight_packed(mat2);
+  TORCH_CHECK(packed_w.scalar_type() == at::kFloat, "fused_linear_sigmoid_mul requires packed float weight")
+
+  int64_t M = mat1.size(0);
+  int64_t K = mat1.size(1);
+  int64_t N = mat2.size(1);
+
+  CHECK_LAST_DIM_CONTIGUOUS_INPUT(mat1);
+  CHECK_INPUT(mat2);
+  CHECK_DIM(2, mat1);
+  CHECK_DIM(2, mat2);
+
+  int64_t out_strideM = post_mul_mat.size(1);
+  int64_t mat1_strideM = mat1.stride(0);
+  auto dispatch_type = mat1.scalar_type();
+  auto out = at::empty({M, out_strideM}, mat1.options());
+
+  TORCH_CHECK(
+      N == 1 && out_strideM % 32 == 0,
+      "post_mul_mat tensor size(1) should be 32 dividable, and the mat2 OC=1 (Mx1 as linear output shape)")
+
+  const bool has_bias = bias.has_value();
+  const float* bias_data = nullptr;
+  if (has_bias) {
+    CHECK_EQ(bias.value().size(0), N);
+    bias_data = bias.value().data_ptr<float>();
+  }
+
+  AT_DISPATCH_REDUCED_FLOATING_TYPES(dispatch_type, "fused_linear_sigmoid_mul", [&] {
    weight_packed_linear_kernel_impl<scalar_t>(
        out.data_ptr<scalar_t>(),
        mat1.data_ptr<scalar_t>(),
-        packed_w.data_ptr<scalar_t>(),
+        packed_w.data_ptr<float>(),
        bias_data,
+        post_mul_mat.data_ptr<scalar_t>(),
        M,
        N,
        K,
@@ -1,29 +1,55 @@
-#pragma once
-
-#include <ATen/native/CPUBlas.h>
+// Adapted from
+// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu

 // clang-format off

+#pragma once
+#include <ATen/native/CPUBlas.h>
+
+#include "common.h"
+
 // amx-bf16
 #define TILE_M 16
 #define TILE_N 16
 #define TILE_K 32

 // block size for AMX gemm
-constexpr int block_size_m() { return 2 * TILE_M; }
-constexpr int block_size_n() { return 2 * TILE_N; }
+constexpr int block_size_m() {
+  return 2 * TILE_M;
+}
+constexpr int block_size_n() {
+  return 2 * TILE_N;
+}

 // define threshold using brgemm (intel AMX)
-template <typename T> inline bool can_use_brgemm(int M);
-template <> inline bool can_use_brgemm<at::BFloat16>(int M) { return M > 4; }
-template <> inline bool can_use_brgemm<at::Half>(int M) { return true; }
-template <> inline bool can_use_brgemm<int8_t>(int M) { return M > 4; }
-template <> inline bool can_use_brgemm<uint8_t>(int M) { return M > 4; }
-template <> inline bool can_use_brgemm<at::Float8_e4m3fn>(int M) { return M > 4; }
-template <> inline bool can_use_brgemm<at::quint4x2>(int M) { return M > 4; }
+template <typename T>
+inline bool can_use_brgemm(int M);
+template <>
+inline bool can_use_brgemm<at::BFloat16>(int M) {
+  return M > 4;
+}
+template <>
+inline bool can_use_brgemm<at::Half>(int M) {
+  return true;
+}
+// this requires PyTorch 2.7 or above
+template <>
+inline bool can_use_brgemm<int8_t>(int M) {
+  return M > 4;
+}
+
+template <>
+inline bool can_use_brgemm<uint8_t>(int M) {
+  return M > 4;
+}
+
+template <>
+inline bool can_use_brgemm<at::Float8_e4m3fn>(int M) {
+  return M > 4;
+}

 // work around compiler internal error
-#define BLOCK_K 128 // 4 * TILE_K
+#define BLOCK_K 128  // 4 * TILE_K

 // adjust leading dimension size for K
 template <typename T>
@@ -36,18 +62,44 @@ inline int64_t get_row_size<int8_t>(int64_t K) {
  return K + sizeof(int32_t);
 }

+// uint8: mxfp4 or int4
+template <>
+inline int64_t get_row_size<uint8_t>(int64_t K) {
+  return K >> 1;
+}
+
 inline int64_t get_row_size(int64_t K, bool use_int8_w8a8) {
  return use_int8_w8a8 ? K + sizeof(int32_t) : K;
 }

+enum class CPUQuantMethod : int64_t { BF16 = 0, INT8_W8A8 = 1, FP8_W8A16 = 2, INT4_W4A8 = 3 };
+
+constexpr bool operator==(CPUQuantMethod a, int64_t b) {
+  return static_cast<int64_t>(a) == b;
+}
+
+constexpr bool operator==(int64_t a, CPUQuantMethod b) {
+  return a == static_cast<int64_t>(b);
+}
+
+enum class CPUQuantAlgo : int64_t { AWQ = 0, GPTQ = 1 };
+
+constexpr bool operator==(CPUQuantAlgo a, int64_t b) {
+  return static_cast<int64_t>(a) == b;
+}
+
+constexpr bool operator==(int64_t a, CPUQuantAlgo b) {
+  return a == static_cast<int64_t>(b);
+}
+
 inline int64_t get_4bit_block_k_size(int64_t group_size) {
  return group_size > 128 ? 128 : group_size;
 }

-// pack weight into vnni format
+// pack weight to vnni format
 at::Tensor convert_weight_packed(at::Tensor& weight);

-// pack weight to vnni format for int4 (adapted from sglang)
+// pack weight to vnni format for int4
 std::tuple<at::Tensor, at::Tensor, at::Tensor>
 convert_weight_packed_scale_zp(at::Tensor qweight, at::Tensor qzeros, at::Tensor scales);

@@ -105,35 +157,6 @@ void fused_experts_fp8_kernel_impl(
    int64_t topk,
    int64_t num_tokens_post_pad);

-// moe implementations for int4 w4a16
-template <typename scalar_t>
-void fused_experts_int4_w4a16_kernel_impl(
-    scalar_t* __restrict__ output,
-    scalar_t* __restrict__ ic0,
-    scalar_t* __restrict__ ic1,
-    scalar_t* __restrict__ ic2,
-    scalar_t* __restrict__ A_tmp,
-    scalar_t* __restrict__ B_tmp,
-    float* __restrict__ C_tmp,
-    const scalar_t* __restrict__ input,
-    const at::quint4x2* __restrict__ packed_w1,
-    const at::quint4x2* __restrict__ packed_w2,
-    const uint8_t* __restrict__ w1z,
-    const uint8_t* __restrict__ w2z,
-    const scalar_t* __restrict__ w1s,
-    const scalar_t* __restrict__ w2s,
-    int group_size,
-    const float* __restrict__ topk_weights,
-    const int32_t* __restrict__ sorted_ids,
-    const int32_t* __restrict__ expert_ids,
-    const int32_t* __restrict__ offsets,
-    int64_t M,
-    int64_t N,
-    int64_t K,
-    int64_t E,
-    int64_t topk,
-    int64_t num_tokens_post_pad);
-
 // shared expert implementation for int8 w8a8
 template <typename scalar_t>
 void shared_expert_int8_kernel_impl(
@@ -153,6 +176,37 @@ void shared_expert_int8_kernel_impl(
    int64_t N,
    int64_t K);

+template <typename scalar_t>
+void fused_experts_int4_w4a8_kernel_impl(
+    scalar_t* __restrict__ output,
+    scalar_t* __restrict__ ic0,
+    scalar_t* __restrict__ ic1,
+    scalar_t* __restrict__ ic2,
+    uint8_t* __restrict__ A_tmp,
+    uint8_t* __restrict__ Aq_tmp,
+    float* __restrict__ As_tmp,
+    int32_t* __restrict__ Azp_tmp,
+    float* __restrict__ C_tmp,
+    int8_t* __restrict__ dqB_tmp,
+    const scalar_t* __restrict__ input,
+    const uint8_t* __restrict__ packed_w1,
+    const uint8_t* __restrict__ packed_w2,
+    const int8_t* __restrict__ w1z,
+    const int8_t* __restrict__ w2z,
+    const float* __restrict__ w1s,
+    const float* __restrict__ w2s,
+    int group_size,
+    const float* __restrict__ topk_weights,
+    const int32_t* __restrict__ sorted_ids,
+    const int32_t* __restrict__ expert_ids,
+    const int32_t* __restrict__ offsets,
+    int64_t M,
+    int64_t N,
+    int64_t K,
+    int64_t E,
+    int64_t topk,
+    int64_t num_tokens_post_pad);
+
 template <typename scalar_t>
 void shared_expert_fp8_kernel_impl(
    scalar_t* __restrict__ output,
@@ -204,6 +258,7 @@ void tinygemm_kernel(
    int64_t ldc,
    bool brg);

+// block quantization
 template <typename scalar_t>
 void tinygemm_kernel(
    const scalar_t* __restrict__ A,
@@ -219,33 +274,26 @@ void tinygemm_kernel(
    int64_t ldb,
    int64_t ldc,
    bool brg,
-    int64_t block_size_K);
+    int64_t block_size_K,
+    bool do_unpack = true);

+// per tensor quantization
 template <typename scalar_t>
 void tinygemm_kernel(
    const scalar_t* __restrict__ A,
-    const at::quint4x2* __restrict__ B,
+    const at::Float8_e4m3fn* __restrict__ B,
    scalar_t* __restrict__ C,
-    const uint8_t* __restrict__ Bz,
-    const scalar_t* __restrict__ Bs,
    scalar_t* __restrict__ Btmp,
    float* __restrict__ Ctmp,
+    float scale,
    int64_t M,
    int64_t N,
    int64_t K,
-    int group_size,
    int64_t lda,
    int64_t ldb,
    int64_t ldc,
-    int64_t strideBz,
-    int64_t strideBs,
    bool brg);

-// int4 scaled GEMM (adapted from sglang)
-at::Tensor int4_scaled_mm_cpu(
-    at::Tensor& x, at::Tensor& w, at::Tensor& w_zeros, at::Tensor& w_scales, std::optional<at::Tensor> bias);
-
-// int4 tinygemm kernel interface(adapted from sglang)
 template <typename scalar_t>
 void tinygemm_kernel(
    scalar_t* C,
@@ -266,34 +314,21 @@ void tinygemm_kernel(
    bool store_out,
    bool use_brgemm);

-// TODO: debug print, remove me later
-inline void print_16x32i(const __m512i x) {
-  int32_t a[16];
-  _mm512_storeu_si512((__m512i *)a, x);
-
-  for (int i = 0; i < 16; i++){
-    std::cout << a[i] << " ";
-  }
-  std::cout << std::endl;
-}
-
-inline void print_16x32(const __m512 x) {
-  float a[16];
-  _mm512_storeu_ps((__m512 *)a, x);
-
-  for (int i = 0; i < 16; i++){
-    std::cout << a[i] << " ";
-  }
-  std::cout << std::endl;
-}
-
-
-inline void print_32x8u(const __m256i x) {
-  uint8_t a[32];
-  _mm256_storeu_si256((__m256i *)a, x);
-
-  for (int i = 0; i < 32; ++i) {
-    std::cout << int32_t(a[i]) << " ";
-  }
-  std::cout << std::endl;
-}
+// mxfp4
+template <typename scalar_t>
+void tinygemm_kernel(
+    const scalar_t* __restrict__ A,
+    const uint8_t* __restrict__ B,
+    scalar_t* __restrict__ C,
+    scalar_t* __restrict__ Btmp,
+    float* __restrict__ Ctmp,
+    const uint8_t* __restrict__ scale,
+    int64_t M,
+    int64_t N,
+    int64_t K,
+    int64_t lda,
+    int64_t ldb,
+    int64_t ldc,
+    bool brg,
+    int64_t block_size_K,
+    bool do_unpack = true);
@@ -1,10 +1,10 @@
-// SPDX-License-Identifier: Apache-2.0
-// Adapted from sgl-project/sglang
-// https://github.com/sgl-project/sglang/pull/8226
+// Adapted from
+// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu

-#include <ATen/ATen.h>
+// clang-format off
+
+#include <torch/all.h>

-#include "common.h"
 #include "gemm.h"
 #include "vec.h"

@@ -24,19 +24,22 @@ struct ActDtype<false> {
  using type = uint8_t;
 };

+#if defined(CPU_CAPABILITY_AVX512)
 struct alignas(32) m256i_wrapper {
  __m256i data;
 };

-#if defined(CPU_CAPABILITY_AVX512)
-inline std::array<m256i_wrapper, 2> load_zps_4vnni(
-    const int8_t* __restrict__ zps) {
-  __m256i vzps_low = _mm256_set1_epi64x(*reinterpret_cast<const int64_t*>(zps));
-  __m256i vzps_high =
-      _mm256_set1_epi64x(*reinterpret_cast<const int64_t*>(zps + 8));
+inline std::array<m256i_wrapper, 2> load_zps_4vnni(const int8_t* __restrict__ zps) {
+  // broadcast 01234567 to
+  // 01234567012345670123456701234567
+  __m256i vzps_low = _mm256_set1_epi64x(*reinterpret_cast<const long*>(zps));
+  __m256i vzps_high = _mm256_set1_epi64x(*reinterpret_cast<const long*>(zps + 8));
+  // shuffle from
+  // 01234567012345670123456701234567
+  // to
+  // 00001111222233334444555566667777
  __m256i shuffle_mask =
-      _mm256_set_epi8(7, 7, 7, 7, 6, 6, 6, 6, 5, 5, 5, 5, 4, 4, 4, 4, 3, 3, 3,
-                      3, 2, 2, 2, 2, 1, 1, 1, 1, 0, 0, 0, 0);
+      _mm256_set_epi8(7, 7, 7, 7, 6, 6, 6, 6, 5, 5, 5, 5, 4, 4, 4, 4, 3, 3, 3, 3, 2, 2, 2, 2, 1, 1, 1, 1, 0, 0, 0, 0);
  vzps_low = _mm256_shuffle_epi8(vzps_low, shuffle_mask);
  vzps_high = _mm256_shuffle_epi8(vzps_high, shuffle_mask);
  m256i_wrapper vzps_low_wp, vzps_high_wp;
@@ -45,8 +48,7 @@ inline std::array<m256i_wrapper, 2> load_zps_4vnni(
  return {vzps_low_wp, vzps_high_wp};
 }

-inline std::array<m256i_wrapper, 2> load_uint4_as_int8(
-    const uint8_t* __restrict__ qB) {
+inline std::array<m256i_wrapper, 2> load_uint4_as_int8(const uint8_t* __restrict__ qB) {
  __m256i packed = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(qB));
  const __m256i low_mask = _mm256_set1_epi8(0x0f);
  __m256i high = _mm256_srli_epi16(packed, 4);
@@ -58,37 +60,42 @@ inline std::array<m256i_wrapper, 2> load_uint4_as_int8(
  return {low_wp, high_wp};
 }

-template <int N, int ldb>
-void _dequant_weight_zp_only(const uint8_t* __restrict__ B, int8_t* dqB,
-                             const int8_t* __restrict__ qzeros, int64_t K) {
-  #pragma GCC unroll 2
+template <int64_t N, int64_t ldb>
+void _dequant_weight_zp_only(const uint8_t* __restrict__ B, int8_t* dqB, const int8_t* __restrict__ qzeros, int64_t K) {
+  // unpack weight int8 -> two int4
+  // subtract zero point
+  // B shape = [K, ldb] = [K, N / 2], actual shape = [K / 4, N / 2, 4]
+  // dqB shape = [K, N], actual shape = [K / 4, N, 4]
+#pragma GCC unroll 2
  for (int n = 0; n < N; n += 16) {
    auto [zps_low_wp, zps_high_wp] = load_zps_4vnni(&qzeros[n]);
    auto zps_low = zps_low_wp.data;
    auto zps_high = zps_high_wp.data;
    for (int k = 0; k < K; k += 4) {
-      auto [vb_low_wp, vb_high_wp] =
-          load_uint4_as_int8(B + ldb * k + n / 2 * 4);
+      auto [vb_low_wp, vb_high_wp] = load_uint4_as_int8(B + ldb * k + n / 2 * 4);
      auto vb_low = vb_low_wp.data;
      auto vb_high = vb_high_wp.data;
      vb_high = _mm256_sub_epi8(vb_high, zps_high);
      vb_low = _mm256_sub_epi8(vb_low, zps_low);
-      _mm256_storeu_si256(reinterpret_cast<__m256i_u*>(dqB + N * k + n * 4),
-                          vb_low);
-      _mm256_storeu_si256(
-          reinterpret_cast<__m256i_u*>(dqB + N * k + (n + 8) * 4), vb_high);
+      // store vb to B
+      _mm256_storeu_si256(reinterpret_cast<__m256i_u*>(dqB + N * k + n * 4), vb_low);
+      _mm256_storeu_si256(reinterpret_cast<__m256i_u*>(dqB + N * k + (n + 8) * 4), vb_high);
    }
  }
 }

-template <bool sym_quant_act, int N, bool accum>
-void _dequant_and_store(float* __restrict__ output,
-                        const int32_t* __restrict__ input,
-                        const float* __restrict__ scale_a,
-                        const int32_t* __restrict__ zp_a,
-                        const float* __restrict__ scale_b,
-                        const int32_t* __restrict__ comp_b, int M, int ldi,
-                        int ldo, int ldsa = 1) {
+template <bool accum, int64_t N, bool sym_quant_act>
+void _dequant_and_store(
+    float* __restrict__ output,
+    const int32_t* __restrict__ input,
+    const float* __restrict__ scale_a,
+    const int32_t* __restrict__ zp_a,
+    const float* __restrict__ scale_b,
+    const int32_t* __restrict__ comp_b,
+    int M,
+    int ldi,
+    int ldo,
+    int ldsa = 1) {
  for (int m = 0; m < M; ++m) {
    float a_scale = *(scale_a + m * ldsa);
    __m512 va_scale = _mm512_set1_ps(a_scale);
@@ -99,7 +106,7 @@ void _dequant_and_store(float* __restrict__ output,
      va_zp = _mm512_set1_epi32(a_zp);
    }
    int n = 0;
-  #pragma GCC unroll 2
+#pragma GCC unroll 2
    for (; n < N; n += 16) {
      __m512i vc = _mm512_loadu_si512(input + m * ldi + n);
      if constexpr (!sym_quant_act) {
@@ -122,8 +129,7 @@ void _dequant_and_store(float* __restrict__ output,
      if constexpr (sym_quant_act) {
        dq_val = (float)input[m * ldi + n] * a_scale * scale_b[n];
      } else {
-        dq_val = (float)(input[m * ldi + n] - a_zp * comp_b[n]) * a_scale *
-                 scale_b[n];
+        dq_val = (float)(input[m * ldi + n] - a_zp * comp_b[n]) * a_scale * scale_b[n];
      }
      if constexpr (accum) {
        output[m * ldo + n] += dq_val;
@@ -135,9 +141,10 @@ void _dequant_and_store(float* __restrict__ output,
 }

 #else
-template <int N, int ldb>
-void _dequant_weight_zp_only(const uint8_t* B, int8_t* dqB,
-                             const int8_t* qzeros, int64_t K) {
+template <int64_t N, int64_t ldb>
+void _dequant_weight_zp_only(const uint8_t* B, int8_t* dqB, const int8_t* qzeros, int64_t K) {
+  // B shape = [K, N / 2]
+  // dqB shape = [K, N]
  for (int k = 0; k < K; ++k) {
    for (int n = 0; n < N / 2; ++n) {
      int32_t b = (int32_t)B[k * ldb + n];
@@ -158,20 +165,31 @@ inline __m512i combine_m256i(std::array<m256i_wrapper, 2> two_256) {
  return combine_m256i(two_256[0].data, two_256[1].data);
 }

+// negate elements in a according to b's sign
 static inline __m512i _mm512_sign_epi8(__m512i a, __m512i b) {
  __m512i zero = _mm512_setzero_si512();
  __mmask64 blt0 = _mm512_movepi8_mask(b);
  return _mm512_mask_sub_epi8(a, blt0, zero, a);
 }

-template <bool sym_quant_act, int M, int N, int ldb>
-void _dequant_gemm_accum_small_M(float* __restrict__ C, const uint8_t* A,
-                                 const float* scales_a, const int32_t* qzeros_a,
-                                 const uint8_t* B, const float* scales_b,
-                                 const int8_t* qzeros_b, int64_t K, int64_t lda,
-                                 int64_t ldc) {
+template <int64_t M, int64_t N, int64_t ldb, bool sym_quant_act>
+void _dequant_gemm_accum_small_M(
+    float* __restrict__ C,
+    const uint8_t* A,
+    const float* scales_a,
+    const int32_t* qzeros_a,
+    const uint8_t* B,
+    const float* scales_b,
+    const int8_t* qzeros_b,
+    int64_t K,
+    int64_t lda,
+    int64_t ldc) {
+  // if sym_quant_act is true, A pointer type is passed in as uint8_t* but actually int8_t*.
+
  constexpr int COLS = N / 16;
-  __m512i ones = _mm512_set1_epi8(1);
+  // Computing compensation is faster than loading it for small M
+  // because it's memory bound.
+  __m512i ones = _mm512_set1_epi8(1);  // used for computing compensation
  __m512i va;
  __m512i vb[COLS];
  __m512i vc[M * COLS];
@@ -179,6 +197,7 @@ void _dequant_gemm_accum_small_M(float* __restrict__ C, const uint8_t* A,
  __m512i vzps[COLS];
  __m512i vcompensate[COLS];

+  // Load scales and zps
  Unroll<COLS>{}([&](auto i) {
    vscales[i] = _mm512_loadu_ps(scales_b + i * 16);
    vzps[i] = combine_m256i(load_zps_4vnni(qzeros_b + i * 16));
@@ -214,25 +233,25 @@ void _dequant_gemm_accum_small_M(float* __restrict__ C, const uint8_t* A,
    }
  };

+  // Accumulate along k
  constexpr const int unroll = 4;
  int k = 0;
  for (; k < K / 4 / unroll; k++) {
-    Unroll<unroll>{}(
-        [&](auto i) { Unroll<M * COLS>{}(compute, 4 * (k * unroll + i)); });
+    Unroll<unroll>{}([&](auto i) { Unroll<M * COLS>{}(compute, 4 * (k * unroll + i)); });
  }
  k *= 4 * unroll;
  for (; k < K; k += 4) {
    Unroll<M * COLS>{}(compute, k);
  }

+  // Store to C
  auto store = [&](auto i) {
    constexpr const int row = i / COLS;
    constexpr const int col = i % COLS;
+    // compute (qC - compensate * zp_a) * scale_a * scale_b
    __m512 vc_float;
    if constexpr (!sym_quant_act) {
-      vc[i] = _mm512_sub_epi32(
-          vc[i], _mm512_mullo_epi32(vcompensate[col],
-                                    _mm512_set1_epi32(*(qzeros_a + row))));
+      vc[i] = _mm512_sub_epi32(vc[i], _mm512_mullo_epi32(vcompensate[col], _mm512_set1_epi32(*(qzeros_a + row))));
    }
    vc_float = _mm512_cvtepi32_ps(vc[i]);
    vc_float = _mm512_mul_ps(vc_float, _mm512_set1_ps(*(scales_a + row)));
@@ -245,17 +264,28 @@ void _dequant_gemm_accum_small_M(float* __restrict__ C, const uint8_t* A,
  Unroll<M * COLS>{}(store);
 }

-  #define CALL_DEQUANT_GEMM_ACCUM_SMALL_M(M)               \
-    _dequant_gemm_accum_small_M<sym_quant_act, M, N, ldb>( \
-        C, A, scales_a, qzeros_a, B, scales_b, qzeros_b, K, lda, ldc);
+#define CALL_DEQUANT_GEMM_ACCUM_SMALL_M(M) \
+  _dequant_gemm_accum_small_M<M, N, ldb, sym_quant_act>(C, A, scales_a, qzeros_a, B, scales_b, qzeros_b, K, lda, ldc);
 #endif

-template <bool sym_quant_act, int N, int ldb>
-void _dequant_gemm_accum(float* C, const uint8_t* A, const float* scales_a,
-                         const int32_t* qzeros_a, const uint8_t* B,
-                         const float* scales_b, const int8_t* qzeros_b,
-                         const int32_t* compensation, int8_t* dqB, int64_t M,
-                         int64_t K, int64_t lda, int64_t ldc, bool use_brgemm) {
+template <int64_t N, int64_t ldb, bool sym_quant_act>
+void _dequant_gemm_accum(
+    float* C,
+    const uint8_t* A,
+    const float* scales_a,
+    const int32_t* qzeros_a,
+    const uint8_t* B,
+    const float* scales_b,
+    const int8_t* qzeros_b,
+    const int32_t* compensation,
+    int8_t* dqB,
+    int64_t M,
+    int64_t K,
+    int64_t lda,
+    int64_t ldc,
+    bool use_brgemm) {
+  // Compute GEMM int8 * int8 -> int32
+  // dequant result to float by applying scales/qzeros
 #if defined(CPU_CAPABILITY_AVX512)
  if (!use_brgemm) {
    switch (M) {
@@ -282,14 +312,12 @@ void _dequant_gemm_accum(float* C, const uint8_t* A, const float* scales_a,
  Tin* A_ptr = (Tin*)A;
  if (use_brgemm) {
    int32_t C_i32[M * N];
-    at::native::cpublas::brgemm(M, N, K, lda, N /*ldb*/, N /*ldc*/,
-                                false /* add_C */, A_ptr, dqB, C_i32,
-                                true /* is_vnni */);
+    at::native::cpublas::brgemm(
+        M, N, K, lda, N /*ldb*/, N /*ldc*/, false /* add_C */, A_ptr, dqB, C_i32, true /* is_vnni */);
    _mm_prefetch(B + N * K / 2, _MM_HINT_T0);
    _mm_prefetch(A + K, _MM_HINT_T0);
-    _dequant_and_store<sym_quant_act, N, true>(C, C_i32, scales_a, qzeros_a,
-                                               scales_b, compensation, M,
-                                               N /*ldi*/, ldc, 1 /*ldsa*/);
+    _dequant_and_store<true, N, sym_quant_act>(
+        C, C_i32, scales_a, qzeros_a, scales_b, compensation, M, N /*ldi*/, ldc, 1 /*ldsa*/);
  } else
 #endif
  {
@@ -297,13 +325,13 @@ void _dequant_gemm_accum(float* C, const uint8_t* A, const float* scales_a,
  }
 }

-template <int N>
+template <int64_t N>
 inline void copy_bias(const float* bias_ptr, float* y_buf, int64_t m) {
  if (bias_ptr) {
    for (int i = 0; i < m; ++i) {
      int j = 0;
 #if defined(CPU_CAPABILITY_AVX512)
-  #pragma GCC unroll 2
+#pragma GCC unroll 2
      for (; j < N; j += 16) {
        __m512 bias_vec = _mm512_loadu_ps(bias_ptr + j);
        _mm512_storeu_ps(y_buf + i * N + j, bias_vec);
@@ -313,11 +341,11 @@ inline void copy_bias(const float* bias_ptr, float* y_buf, int64_t m) {
        y_buf[i * N + j] = bias_ptr[j];
      }
    }
-  } else {
+  } else {  // initialize to zero
    for (int i = 0; i < m; ++i) {
      int j = 0;
 #if defined(CPU_CAPABILITY_AVX512)
-  #pragma GCC unroll 2
+#pragma GCC unroll 2
      for (; j < N; j += 16) {
        __m512 zero_vec = _mm512_setzero_ps();
        _mm512_storeu_ps(y_buf + i * N + j, zero_vec);
@@ -330,14 +358,13 @@ inline void copy_bias(const float* bias_ptr, float* y_buf, int64_t m) {
  }
 }

-template <int N, typename out_dtype>
-inline void store_out(const float* y_buf, out_dtype* c_ptr, int64_t m,
-                      int64_t lda) {
+template <typename out_dtype, int64_t N>
+inline void store_out(const float* y_buf, out_dtype* c_ptr, int64_t m, /* int64_t n, */ int64_t lda) {
  for (int i = 0; i < m; ++i) {
    int j = 0;
    if constexpr (std::is_same<out_dtype, float>::value) {
 #if defined(CPU_CAPABILITY_AVX512)
-  #pragma GCC unroll 2
+#pragma GCC unroll 2
      for (; j < N; j += 16) {
        __m512 y_vec = _mm512_loadu_ps(y_buf + i * N + j);
        _mm512_storeu_ps(c_ptr + i * lda + j, y_vec);
@@ -348,12 +375,11 @@ inline void store_out(const float* y_buf, out_dtype* c_ptr, int64_t m,
      }
    } else if constexpr (std::is_same<out_dtype, at::BFloat16>::value) {
 #if defined(CPU_CAPABILITY_AVX512)
-  #pragma GCC unroll 2
+#pragma GCC unroll 2
      for (; j < N; j += 16) {
        __m512 y_vec = _mm512_loadu_ps(y_buf + i * N + j);
        __m256i y_bf16_vec = at::vec::cvtfp32_bf16(y_vec);
-        _mm256_storeu_si256(reinterpret_cast<__m256i*>(c_ptr + i * lda + j),
-                            y_bf16_vec);
+        _mm256_storeu_si256(reinterpret_cast<__m256i*>(c_ptr + i * lda + j), y_bf16_vec);
      }
 #endif
      for (; j < N; ++j) {
@@ -361,12 +387,11 @@ inline void store_out(const float* y_buf, out_dtype* c_ptr, int64_t m,
      }
    } else if constexpr (std::is_same<out_dtype, at::Half>::value) {
 #if defined(CPU_CAPABILITY_AVX512)
-  #pragma GCC unroll 2
+#pragma GCC unroll 2
      for (; j < N; j += 16) {
        __m512 y_vec = _mm512_loadu_ps(y_buf + i * N + j);
        __m256i y_fp16_vec = at::vec::cvtfp32_fp16(y_vec);
-        _mm256_storeu_si256(reinterpret_cast<__m256i*>(c_ptr + i * lda + j),
-                            y_fp16_vec);
+        _mm256_storeu_si256(reinterpret_cast<__m256i*>(c_ptr + i * lda + j), y_fp16_vec);
      }
 #endif
      for (; j < N; ++j) {
@@ -392,16 +417,25 @@ void fill_val_stub(int32_t* __restrict__ output, int32_t value, int64_t size) {
  }
 }

-template <bool sym_quant_act, typename act_dtype, typename out_dtype>
+template <typename act_dtype, typename out_dtype, bool sym_quant_act>
 void _da8w4_linear_impl(
-    act_dtype* __restrict__ input, const float* __restrict__ input_scales,
+    act_dtype* __restrict__ input,
+    const float* __restrict__ input_scales,
    const int32_t* __restrict__ input_qzeros,
-    const uint8_t* __restrict__ weight, const float* __restrict__ weight_scales,
-    const int8_t* __restrict__ weight_qzeros, const float* __restrict__ bias,
-    out_dtype* __restrict__ output, float* __restrict__ output_temp,
-    int8_t* __restrict__ dequant_weight_temp, int64_t M, int64_t N, int64_t K,
+    const uint8_t* __restrict__ weight,
+    const float* __restrict__ weight_scales,
+    const int8_t* __restrict__ weight_qzeros,
+    const float* __restrict__ bias,
+    out_dtype* __restrict__ output,
+    float* __restrict__ output_temp,
+    int8_t* __restrict__ dequant_weight_temp,
+    int64_t M,
+    int64_t N,
+    int64_t K,
    int64_t num_groups) {
-  const bool use_brgemm = can_use_brgemm<act_dtype>(M);
+  // weight + compensation shape = [Nc, Kc, BLOCK_N * _block_k / 2 + BLOCK_N*sizeof(int32_t)]
+  // scales/qzeros shape = [Nc, G, BLOCK_N]
+  const bool use_brgemm = can_use_brgemm<int8_t>(M);
  int64_t block_m = [&]() -> long {
    if (M <= 48) {
      return M;
@@ -432,30 +466,24 @@ void _da8w4_linear_impl(
      int64_t mc_end = parallel_on_M ? mc + 1 : Mc;

      for (int mci = mc; mci < mc_end; ++mci) {
-        int64_t m_size =
-            mci * block_m + block_m > M ? M - mci * block_m : block_m;
+        int64_t m_size = mci * block_m + block_m > M ? M - mci * block_m : block_m;
+        // copy bias to y_buf if bias is not None
        auto bias_data = bias ? bias + nc * BLOCK_N : nullptr;
        copy_bias<BLOCK_N>(bias_data, C_tmp, m_size);
        for (int kci = 0; kci < Kc; ++kci) {
          int32_t* compensation_ptr =
              sym_quant_act
                  ? nullptr
-                  : (int32_t*)(void*)(weight +
-                                      (nc * Kc + kci) *
-                                          (BLOCK_N *
-                                           (_block_k / 2 + sizeof(int32_t))) +
-                                      _block_k * BLOCK_N / 2);
-          _dequant_gemm_accum<sym_quant_act, BLOCK_N, BLOCK_N / 2>(
+                  : (int32_t*)(void*)(weight + (nc * Kc + kci) * (BLOCK_N * (_block_k / 2 + sizeof(int32_t))) +
+                                      _block_k * BLOCK_N / 2) /*Bcomp*/;
+          _dequant_gemm_accum<BLOCK_N, BLOCK_N / 2, sym_quant_act>(
              /*C*/ C_tmp,
              /*A*/ (uint8_t*)input + mci * block_m * K + kci * _block_k,
              /*scales_a*/ input_scales + mci * block_m,
              /*qzeros_a*/ input_qzeros + mci * block_m,
-              /*B*/ weight + (nc * Kc + kci) *
-                                 (BLOCK_N * (_block_k / 2 + sizeof(int32_t))),
-              /*scales_b*/ weight_scales + nc * BLOCK_N * num_groups +
-                  kci / block_per_group * BLOCK_N,
-              /*qzeros_b*/ weight_qzeros + nc * BLOCK_N * num_groups +
-                  kci / block_per_group * BLOCK_N,
+              /*B*/ weight + (nc * Kc + kci) * (BLOCK_N * (_block_k / 2 + sizeof(int32_t))),
+              /*scales_b*/ weight_scales + nc * BLOCK_N * num_groups + kci / block_per_group * BLOCK_N,
+              /*qzeros_b*/ weight_qzeros + nc * BLOCK_N * num_groups + kci / block_per_group * BLOCK_N,
              /*Bcomp*/ compensation_ptr,
              /*dqB_tmp*/ dqB_tmp,
              /*M*/ m_size,
@@ -464,8 +492,8 @@ void _da8w4_linear_impl(
              /*ldc*/ BLOCK_N,
              /*use_brgemm*/ use_brgemm);
        }
-        store_out<BLOCK_N>(C_tmp, output + mci * block_m * N + nc * BLOCK_N,
-                           m_size, N /*lda*/);
+        // store y_buf to output with dtype conversion
+        store_out<out_dtype, BLOCK_N>(C_tmp, output + mci * block_m * N + nc * BLOCK_N, m_size, N /*lda*/);
      }
    }
    if (use_brgemm) {
@@ -476,15 +504,16 @@ void _da8w4_linear_impl(

 }  // anonymous namespace

-std::tuple<at::Tensor, at::Tensor, at::Tensor>
-convert_int4_weight_packed_with_compensation(const at::Tensor& weight,
-                                             const at::Tensor& scales,
-                                             const at::Tensor& qzeros) {
-  TORCH_CHECK(weight.dim() == 2,
-              "DA8W4 CPU: Weight should be a 2D tensor for packing");
-  TORCH_CHECK(
-      weight.size(1) % 2 == 0,
-      "DA8W4 CPU: Weight should have even number of columns for packing");
+/*
+return: packed_weight, packed_scales, packed_qzeros
+*/
+std::tuple<at::Tensor, at::Tensor, at::Tensor> convert_int4_weight_packed_with_compensation(
+    const at::Tensor& weight, const at::Tensor& scales, const at::Tensor& qzeros) {
+  // weight shape = [N, K]
+  // scales shape = [N, G]
+  // qzeros shape = [N, G]
+  TORCH_CHECK(weight.dim() == 2, "DA8W4 CPU: Weight should be a 2D tensor for packing");
+  TORCH_CHECK(weight.size(1) % 2 == 0, "DA8W4 CPU: Weight should have even number of columns for packing");

  auto new_scales = scales;
  auto new_qzeros = qzeros;
@@ -505,20 +534,21 @@ convert_int4_weight_packed_with_compensation(const at::Tensor& weight,
  int64_t Nc = N / block_n;
  int64_t Kc = K / _block_k;

+  // Reorder weight to [N/block_n, K/_block_k, _block_k, block_n]
+  // Reorder scales/qzeros to [N/block_n, G, block_n]
+  // weight + compensation shape = [Nc, Kc, block_n * _block_k / 2 + block_n*sizeof(int32_t)]
+  // scales/qzeros shape = [Nc, G, block_n]
  auto weight_view = weight.view({Nc, block_n, Kc, _block_k});
  at::Tensor weight_reordered = weight_view.permute({0, 2, 3, 1}).contiguous();
  at::Tensor blocked_weight;
-  at::Tensor blocked_scales =
-      new_scales.view({Nc, block_n, G}).permute({0, 2, 1}).contiguous();
-  at::Tensor blocked_qzeros =
-      new_qzeros.view({Nc, block_n, G}).permute({0, 2, 1}).contiguous();
-  auto weight_sub_qzero = weight.view({Nc, block_n, G, -1}).to(at::kInt) -
-                          new_qzeros.view({Nc, block_n, G, -1});
+  at::Tensor blocked_scales = new_scales.view({Nc, block_n, G}).permute({0, 2, 1}).contiguous();
+  at::Tensor blocked_qzeros = new_qzeros.view({Nc, block_n, G}).permute({0, 2, 1}).contiguous();
+  // Compensation = Σ(k)(W[k][n] - ZP[n]) for each block.
+  auto weight_sub_qzero = weight.view({Nc, block_n, G, -1}).to(at::kInt) - new_qzeros.view({Nc, block_n, G, -1});
  weight_sub_qzero = weight_sub_qzero.view({Nc, block_n, Kc, _block_k});
  at::Tensor compensation = weight_sub_qzero.sum(-1);
  compensation = compensation.permute({0, 2, 1}).contiguous().to(at::kInt);
-  int64_t buffer_size_nbytes =
-      _block_k * block_n / 2 + block_n * sizeof(int32_t);
+  int64_t buffer_size_nbytes = _block_k * block_n / 2 + block_n * sizeof(int32_t);
  blocked_weight = at::empty({Nc, Kc, buffer_size_nbytes}, weight.options());

  auto weight_ptr = weight_reordered.data_ptr<uint8_t>();
@@ -528,24 +558,25 @@ convert_int4_weight_packed_with_compensation(const at::Tensor& weight,
  at::parallel_for(0, num_blocks, 1, [&](int64_t begin, int64_t end) {
    for (const auto i : c10::irange(begin, end)) {
      auto in_ptr = weight_ptr + i * _block_k * block_n;
-      auto out_ptr =
-          blocked_weight_ptr + i * block_n * (_block_k / 2 + sizeof(int32_t));
+      auto out_ptr = blocked_weight_ptr + i * block_n * (_block_k / 2 + sizeof(int32_t));
      int32_t* comp_in_prt = compensation_ptr + i * block_n;
-      int32_t* comp_out_prt =
-          (int32_t*)(void*)(blocked_weight_ptr +
-                            i * block_n * (_block_k / 2 + sizeof(int32_t)) +
-                            _block_k * block_n / 2);
+      int32_t* comp_out_prt = (int32_t*)(void*)(blocked_weight_ptr + i * block_n * (_block_k / 2 + sizeof(int32_t)) +
+                                                _block_k * block_n / 2);
+      // Reorder weight block to VNNI4 and pack two lanes along N
+      // N=16 viewed as two lanes: a0, ...a7, b0, ...b7
+      // pack two lanes: [a0, b0], ..., [a7, b7]
+      // plain shape = [_block_k, block_n]
+      // packed shape = [_block_k / 4, block_n / 2, 4] viewed as [_block_k, block_n / 2]
      constexpr int n_group_size = 8;
      constexpr int vnni_size = 4;
-      constexpr int n_group = block_n / n_group_size;
+      constexpr int n_group = block_n / n_group_size;  // 4
      for (int nb = 0; nb < n_group; nb += 2) {
        for (int k = 0; k < _block_k; k += vnni_size) {
          for (int ni = 0; ni < n_group_size; ++ni) {
            for (int ki = 0; ki < vnni_size; ++ki) {
              int src_idx_1 = nb * n_group_size + ni + (k + ki) * block_n;
              int src_idx_2 = (nb + 1) * n_group_size + ni + (k + ki) * block_n;
-              int dst_idx = (nb / 2 * n_group_size + ni) * vnni_size +
-                            k * block_n / 2 + ki;
+              int dst_idx = (nb / 2 * n_group_size + ni) * vnni_size + k * block_n / 2 + ki;
              uint8_t src_1 = *(in_ptr + src_idx_1);
              uint8_t src_2 = *(in_ptr + src_idx_2);
              uint8_t dst = (src_1 & 0x0f) | ((src_2 & 0x0f) << 4);
@@ -554,39 +585,128 @@ convert_int4_weight_packed_with_compensation(const at::Tensor& weight,
          }
        }
      }
+      // compensation [block_n]
      for (int nb = 0; nb < block_n; nb++) {
        *(comp_out_prt + nb) = *(comp_in_prt + nb);
      }
    }
  });

-  return std::make_tuple(std::move(blocked_weight), std::move(blocked_scales),
-                         std::move(blocked_qzeros));
+  return std::make_tuple(std::move(blocked_weight), std::move(blocked_scales), std::move(blocked_qzeros));
 }

-std::tuple<at::Tensor, at::Tensor> autoawq_to_int4pack(at::Tensor qweight,
-                                                       at::Tensor qzeros) {
-  auto bitshifts = at::tensor({0, 4, 1, 5, 2, 6, 3, 7}, at::kInt) * 4;
-  auto qweight_unsq = qweight.unsqueeze(-1);
-  auto unpacked = at::bitwise_right_shift(qweight_unsq, bitshifts) & 0xF;
-  auto qweight_final = unpacked.flatten(-2).transpose(-1, -2).to(at::kByte);
+std::tuple<at::Tensor, at::Tensor> unpack_4bit_to_32bit_signed(const at::Tensor& qweight, const at::Tensor& qzeros) {
+  TORCH_CHECK(qweight.scalar_type() == at::kInt, "qweight must be int32");
+  TORCH_CHECK(qzeros.scalar_type() == at::kInt, "qzeros must be int32");
+  const auto W0 = qweight.size(0);
+  const auto W1 = qweight.size(1);
+  const auto Z0 = qzeros.size(0);
+  const auto Z1 = qzeros.size(1);

-  auto qzeros_unsq = qzeros.unsqueeze(-1);
-  auto qzeros_unpacked = at::bitwise_right_shift(qzeros_unsq, bitshifts) & 0xF;
-  auto qzeros_final = qzeros_unpacked.flatten(-2).to(at::kByte);
+  // unpacked_weights: (W0 * 8, W1), int8
+  auto unpacked_weights = at::zeros({W0 * 8, W1}, at::TensorOptions().dtype(at::kChar));
+  // unpacked_zeros: (Z0, Z1 * 8), int8
+  auto unpacked_zeros = at::zeros({Z0, Z1 * 8}, at::TensorOptions().dtype(at::kChar));

-  return std::make_tuple(qweight_final, qzeros_final);
+  const int32_t* qw_ptr = qweight.data_ptr<int32_t>();
+  const int32_t* qz_ptr = qzeros.data_ptr<int32_t>();
+  int8_t* uw_ptr = unpacked_weights.data_ptr<int8_t>();
+  int8_t* uz_ptr = unpacked_zeros.data_ptr<int8_t>();
+
+  // ---- unpack qweight ----
+  for (int64_t row = 0; row < W0 * 8; ++row) {
+    const int i = row & 7;         // row % 8
+    const int src_row = row >> 3;  // row // 8
+    const int shift = 4 * i;
+    for (int64_t col = 0; col < W1; ++col) {
+      int32_t v = qw_ptr[src_row * W1 + col];
+      uw_ptr[row * W1 + col] = static_cast<int8_t>((v >> shift) & 0xF);
+    }
+  }
+  // ---- unpack qzeros ----
+  for (int64_t col = 0; col < Z1 * 8; ++col) {
+    const int i = col & 7;
+    const int src_col = col >> 3;
+    const int shift = 4 * i;
+
+    for (int64_t row = 0; row < Z0; ++row) {
+      int32_t v = qz_ptr[row * Z1 + src_col];
+      uz_ptr[row * (Z1 * 8) + col] = static_cast<int8_t>((v >> shift) & 0xF);
+    }
+  }
+
+  return std::make_tuple(unpacked_weights, unpacked_zeros + 1);
+}
+
+std::tuple<at::Tensor, at::Tensor>
+autogptq_to_int4pack(const at::Tensor& qweight_tensor, const at::Tensor& qzeros_tensor) {
+  TORCH_CHECK(qweight_tensor.scalar_type() == at::kInt, "qweight_tensor must be int32");
+  TORCH_CHECK(qzeros_tensor.scalar_type() == at::kInt, "qzeros_tensor must be int32");
+  TORCH_CHECK(qweight_tensor.is_cpu(), "CPU only implementation");
+  if (qweight_tensor.dim() == 3) {
+    const int64_t B = qweight_tensor.size(0);
+    std::vector<at::Tensor> qweight_list;
+    std::vector<at::Tensor> qzeros_list;
+    qweight_list.reserve(B);
+    qzeros_list.reserve(B);
+    for (int64_t i = 0; i < B; ++i) {
+      auto outputs = unpack_4bit_to_32bit_signed(qweight_tensor[i], qzeros_tensor[i]);
+      at::Tensor unpacked_qweight = std::get<0>(outputs);
+      at::Tensor unpacked_qzeros = std::get<1>(outputs);
+      qweight_list.push_back(unpacked_qweight.transpose(0, 1).contiguous().to(at::kByte));
+      qzeros_list.push_back(unpacked_qzeros.contiguous().to(at::kByte));
+    }
+    return std::make_tuple(at::stack(qweight_list).detach(), at::stack(qzeros_list).detach());
+  }
+  auto outputs = unpack_4bit_to_32bit_signed(qweight_tensor, qzeros_tensor);
+  at::Tensor unpacked_qweight = std::get<0>(outputs);
+  at::Tensor unpacked_qzeros = std::get<1>(outputs);
+  at::Tensor return_qweight = unpacked_qweight.transpose(0, 1).contiguous().to(at::kByte);
+  at::Tensor return_qzeros = unpacked_qzeros.contiguous().to(at::kByte);
+  return std::make_tuple(return_qweight, return_qzeros);
+}
+
+std::tuple<at::Tensor, at::Tensor> int4pack(at::Tensor qweight, at::Tensor qzeros, int64_t quant_method_4bit) {
+  if (quant_method_4bit == CPUQuantAlgo::AWQ) {
+    // autoawq unpacking
+    qweight = qweight.contiguous();
+    qzeros = qzeros.contiguous();
+    // bitshifts: [0, 4, 1, 5, 2, 6, 3, 7] * 4
+    auto bitshifts = at::tensor({0, 4, 1, 5, 2, 6, 3, 7}, at::kInt) * 4;
+    auto qweight_unsq = qweight.unsqueeze(-1);  // [..., K, N/8, 1]
+    auto unpacked = (at::bitwise_right_shift(qweight_unsq, bitshifts) & 0xF).contiguous();
+    auto qweight_final = unpacked.flatten(-2).transpose(-1, -2).to(at::kByte).clone();
+    auto qzeros_unsq = qzeros.unsqueeze(-1);
+    auto qzeros_unpacked = (at::bitwise_right_shift(qzeros_unsq, bitshifts) & 0xF).contiguous();
+    auto qzeros_final = qzeros_unpacked.flatten(-2).to(at::kByte).clone();
+    return std::make_tuple(qweight_final, qzeros_final);
+  } else if (quant_method_4bit == CPUQuantAlgo::GPTQ) {
+    // autogptq unpacking
+    auto outputs = autogptq_to_int4pack(qweight, qzeros);
+    at::Tensor unpacked_qweight = std::get<0>(outputs);
+    at::Tensor unpacked_qzeros = std::get<1>(outputs);
+    return std::make_tuple(unpacked_qweight, unpacked_qzeros);
+  } else {
+    TORCH_CHECK(false, "CPU int4 pack only support AWQ or GPTQ...");
+  }
 }

 std::tuple<at::Tensor, at::Tensor, at::Tensor> convert_weight_packed_scale_zp(
-    at::Tensor qweight, at::Tensor qzeros, at::Tensor scales) {
-  auto res = autoawq_to_int4pack(qweight, qzeros);
-  auto _qweight = std::get<0>(res);
-  auto _qzeros = std::get<1>(res);
+    at::Tensor qweight,  // awq: (*, K, N / 8)  ||  gptq: (*, K / 8, N) , int32
+    at::Tensor qzeros,   // awq: (*, K / group_size, N / 8) ||  gptq: (*, K / group_size, N / 8) , int32
+    at::Tensor scales,   // awq: (*, K / group_size, N) ||  gptq: (*, K / group_size, N) , bfloat16
+    int64_t quant_method_4bit) {
+  at::Tensor _qweight;
+  at::Tensor _qzeros;
+
+  auto res = int4pack(qweight, qzeros, quant_method_4bit);
+  _qweight = std::get<0>(res);
+  _qzeros = std::get<1>(res);
+
  auto _scales = scales;
-  _qzeros = _qzeros.transpose(-2, -1).contiguous();
+  _qzeros = _qzeros.transpose(-2, -1).contiguous();  // .T
  _scales = _scales.transpose(-2, -1).contiguous();
-  if (_qweight.dim() == 3) {
+  if (_qweight.dim() == 3) {  // Dim=3 for MOE packing, TODO: refine a unified loop
    int64_t E = _qweight.size(0);
    int64_t K = _qweight.size(2);
    int64_t G = _scales.size(2);
@@ -595,17 +715,12 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> convert_weight_packed_scale_zp(
    int64_t block_n = block_size_n();
    int64_t Nc = _qweight.size(1) / block_n;
    int64_t Kc = K / _block_k;
-    int64_t buffer_size_nbytes =
-        _block_k * block_n / 2 + block_n * sizeof(int32_t);
-    auto blocked_weight =
-        at::empty({E, Nc, Kc, buffer_size_nbytes}, _qweight.options());
-    auto blocked_scales =
-        at::empty({E, Nc, G, block_n}, _scales.options()).to(at::kFloat);
-    auto blocked_qzeros =
-        at::empty({E, Nc, G, block_n}, _qzeros.options()).to(at::kChar);
+    int64_t buffer_size_nbytes = _block_k * block_n / 2 + block_n * sizeof(int32_t);
+    auto blocked_weight = at::empty({E, Nc, Kc, buffer_size_nbytes}, _qweight.options());
+    auto blocked_scales = at::empty({E, Nc, G, block_n}, _scales.options()).to(at::kFloat);
+    auto blocked_qzeros = at::empty({E, Nc, G, block_n}, _qzeros.options()).to(at::kChar);
    for (int i = 0; i < _qweight.size(0); i++) {
-      auto res_ = convert_int4_weight_packed_with_compensation(
-          _qweight[i], _scales[i], _qzeros[i]);
+      auto res_ = convert_int4_weight_packed_with_compensation(_qweight[i], _scales[i], _qzeros[i]);
      blocked_weight[i] = std::get<0>(res_);
      blocked_scales[i] = std::get<1>(res_);
      blocked_qzeros[i] = std::get<2>(res_);
@@ -614,8 +729,7 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> convert_weight_packed_scale_zp(
    _scales = blocked_scales;
    _qzeros = blocked_qzeros;
  } else {
-    auto res_ = convert_int4_weight_packed_with_compensation(_qweight, _scales,
-                                                             _qzeros);
+    auto res_ = convert_int4_weight_packed_with_compensation(_qweight, _scales, _qzeros);
    _qweight = std::get<0>(res_);
    _scales = std::get<1>(res_);
    _qzeros = std::get<2>(res_);
@@ -624,89 +738,93 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> convert_weight_packed_scale_zp(
  return std::make_tuple(_qweight, _qzeros, _scales);
 }

-at::Tensor int4_scaled_mm_cpu_with_quant(const at::Tensor& input,
-                                         const at::Tensor& weight,
-                                         const at::Tensor& weight_scales,
-                                         const at::Tensor& weight_qzeros,
-                                         const std::optional<at::Tensor>& bias,
-                                         at::ScalarType output_dtype) {
-  RECORD_FUNCTION("vllm::int4_scaled_mm_cpu_with_quant",
-                  std::vector<c10::IValue>({input, weight}));
-
+at::Tensor int4_scaled_mm_cpu_with_quant(
+    const at::Tensor& input,
+    const at::Tensor& weight,
+    const at::Tensor& weight_scales,
+    const at::Tensor& weight_qzeros,
+    const std::optional<at::Tensor>& bias,
+    at::ScalarType output_dtype) {
  int64_t M_a = input.size(0);
  int64_t K_a = input.size(1);
  int64_t lda = input.stride(0);

  const auto st = input.scalar_type();
  TORCH_CHECK(
-      st == at::kBFloat16 || st == at::kHalf,
-      "int4_scaled_mm_cpu_with_quant: expect A to be bfloat16 or half.");
+      st == at::kBFloat16 || st == at::kHalf, "int4_scaled_mm_cpu_with_quant: expect A to be bfloat16 or half.");

-  constexpr bool sym_quant_act = false;
+  constexpr bool sym_quant_act = false;  // TODO: add sym quant path
  using Tin = typename ActDtype<sym_quant_act>::type;
-  int64_t act_buffer_size =
-      M_a * K_a + M_a * sizeof(float) + M_a * sizeof(int32_t);
-  auto act_buffer =
-      at::empty({act_buffer_size}, input.options().dtype(at::kByte));
+  int64_t act_buffer_size = /* act quant */ M_a * K_a +
+                            /* act scale */ M_a * sizeof(float) +
+                            /* act zp */ M_a * sizeof(int32_t);
+  auto act_buffer = at::empty({act_buffer_size}, input.options().dtype(at::kByte));
+  // asym path, activation quants into uint8_t
  auto Aq_data = act_buffer.data_ptr<uint8_t>();
  auto As_data = reinterpret_cast<float*>(Aq_data + M_a * K_a);
  auto Azp_data = reinterpret_cast<int32_t*>(As_data + M_a);
-  fill_val_stub(Azp_data, 128, M_a);
+  fill_val_stub(Azp_data, 128, M_a);  // sym_a s8s8 is unified to u8s8 with compensation (128)

  auto out_sizes = input.sizes().vec();
  int64_t N = weight_scales.size(0) * weight_scales.size(-1);
  out_sizes.back() = N;
  auto output = at::empty(out_sizes, input.options());
+  // weight + compensation shape = [Nc, Kc, BLOCK_N * _block_k / 2 + BLOCK_N*sizeof(int32_t)]
+  // scales/qzeros shape = [Nc, G, BLOCK_N]
  int64_t Nc = weight.size(0);
  int64_t Kc = weight.size(1);
  int64_t _block_k = K_a / Kc;
  TORCH_CHECK(N == Nc * BLOCK_N, "DA8W4: weight and input shapes mismatch");
+  // scales/qzeros shape = [Nc, G, BLOCK_N]
  int64_t num_groups = weight_scales.size(1);

  const uint8_t* b_ptr = weight.data_ptr<uint8_t>();
  const float* b_scales_ptr = weight_scales.data_ptr<float>();
  const int8_t* b_qzeros_ptr = weight_qzeros.data_ptr<int8_t>();
-  const float* bias_ptr =
-      bias.has_value() ? bias.value().data_ptr<float>() : nullptr;
+  const float* bias_ptr = bias.has_value() ? bias.value().data_ptr<float>() : nullptr;
  int num_threads = at::get_num_threads();
-  int64_t temp_buffer_size = num_threads * BLOCK_M * BLOCK_N * sizeof(float) +
-                             num_threads * _block_k * BLOCK_N;
-  auto c_temp_buffer =
-      at::empty({temp_buffer_size}, input.options().dtype(at::kChar));
+  int64_t temp_buffer_size = /* output temp */ num_threads * BLOCK_M * BLOCK_N * sizeof(float) +
+                             /*  weight dequant temp */ num_threads * _block_k * BLOCK_N;
+  auto c_temp_buffer = at::empty({temp_buffer_size}, input.options().dtype(at::kChar));
  float* c_temp_ptr = (float*)((void*)(c_temp_buffer.data_ptr<int8_t>()));
-  int8_t* dqB_temp_ptr =
-      (int8_t*)((void*)(c_temp_ptr + num_threads * BLOCK_M * BLOCK_N));
+  int8_t* dqB_temp_ptr = (int8_t*)((void*)(c_temp_ptr + num_threads * BLOCK_M * BLOCK_N));

-#define LAUNCH_DA8W4_LINEAR_WITH_QUANT_IMPL(sym_quant_act)                 \
-  AT_DISPATCH_FLOATING_TYPES_AND2(                                         \
-      at::ScalarType::BFloat16, at::ScalarType::Half, output_dtype,        \
-      "int4_scaled_mm_cpu", [&] {                                          \
-        const scalar_t* __restrict__ A_data = input.data_ptr<scalar_t>();  \
-        scalar_t* __restrict__ c_ptr = output.data_ptr<scalar_t>();        \
-        at::parallel_for(0, M_a, 0, [&](int64_t begin, int64_t end) {      \
-          for (int64_t m = begin; m < end; ++m) {                          \
-            quantize_row_int8<scalar_t>(Aq_data + m * K_a, As_data[m],     \
-                                        A_data + m * lda, K_a);            \
-          }                                                                \
-        });                                                                \
-        _da8w4_linear_impl<sym_quant_act, Tin, scalar_t>(                  \
-            Aq_data, As_data, Azp_data, b_ptr, b_scales_ptr, b_qzeros_ptr, \
-            bias_ptr, c_ptr, c_temp_ptr, dqB_temp_ptr, M_a, N, K_a,        \
-            num_groups);                                                   \
+#define LAUNCH_DA8W4_LINEAR_WITH_QUANT_IMPL(sym_quant_act)                                                 \
+  AT_DISPATCH_FLOATING_TYPES_AND2(                                                                         \
+      at::ScalarType::BFloat16, at::ScalarType::Half, output_dtype, "int4_scaled_mm_cpu_with_quant", [&] { \
+        const scalar_t* __restrict__ A_data = input.data_ptr<scalar_t>();                                  \
+        scalar_t* __restrict__ c_ptr = output.data_ptr<scalar_t>();                                        \
+        at::parallel_for(0, M_a, 0, [&](int64_t begin, int64_t end) {                                      \
+          for (int64_t m = begin; m < end; ++m) {                                                          \
+            quantize_row_int8<scalar_t>(Aq_data + m * K_a, As_data[m], A_data + m * lda, K_a);             \
+          }                                                                                                \
+        });                                                                                                \
+        _da8w4_linear_impl<Tin, scalar_t, sym_quant_act>(                                                  \
+            Aq_data,                                                                                       \
+            As_data,                                                                                       \
+            Azp_data,                                                                                      \
+            b_ptr,                                                                                         \
+            b_scales_ptr,                                                                                  \
+            b_qzeros_ptr,                                                                                  \
+            bias_ptr,                                                                                      \
+            c_ptr,                                                                                         \
+            c_temp_ptr,                                                                                    \
+            dqB_temp_ptr,                                                                                  \
+            M_a,                                                                                           \
+            N,                                                                                             \
+            K_a,                                                                                           \
+            num_groups);                                                                                   \
      });

  LAUNCH_DA8W4_LINEAR_WITH_QUANT_IMPL(sym_quant_act);

  return output;
 }
-
-namespace {
-
 template <typename scalar_t>
-inline void copy_stub(scalar_t* __restrict__ out,
-                      const float* __restrict__ input, int64_t size) {
+inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ input, int64_t size) {
  using Vec = at::vec::Vectorized<scalar_t>;
  using fVec = at::vec::Vectorized<float>;
+// no remainder
 #pragma GCC unroll 4
  for (int64_t d = 0; d < size; d += Vec::size()) {
    fVec x0 = fVec::loadu(input + d);
@@ -716,40 +834,61 @@ inline void copy_stub(scalar_t* __restrict__ out,
  }
 }

-}  // anonymous namespace
-
 template <typename scalar_t>
-void tinygemm_kernel(scalar_t* C, float* C_temp, const uint8_t* A,
-                     const float* scales_a, const int32_t* qzeros_a,
-                     const uint8_t* B, const float* scales_b,
-                     const int8_t* qzeros_b, const int32_t* compensation,
-                     int8_t* dqB_tmp, int64_t M, int64_t K, int64_t lda,
-                     int64_t ldc_f, int64_t ldc_s, bool store_out,
-                     bool use_brgemm) {
-  _dequant_gemm_accum<false, BLOCK_N, BLOCK_N / 2>(
-      C_temp, A, scales_a, qzeros_a, B, scales_b, qzeros_b, compensation,
-      dqB_tmp, M, K, lda, ldc_f, use_brgemm);
+void tinygemm_kernel(
+    scalar_t* C,
+    float* C_temp,
+    const uint8_t* A,
+    const float* scales_a,
+    const int32_t* qzeros_a,
+    const uint8_t* B,
+    const float* scales_b,
+    const int8_t* qzeros_b,
+    const int32_t* compensation,
+    int8_t* dqB_tmp,
+    int64_t M,
+    int64_t K,
+    int64_t lda,
+    int64_t ldc_f,
+    int64_t ldc_s,
+    bool store_out,
+    bool use_brgemm) {
+  // TODO: add sym quant act, now only asym
+  _dequant_gemm_accum<BLOCK_N, BLOCK_N / 2, false>(
+      C_temp, A, scales_a, qzeros_a, B, scales_b, qzeros_b, compensation, dqB_tmp, M, K, lda, ldc_f, use_brgemm);
  if (store_out) {
+    // copy from Ctmp to C
    for (int64_t m = 0; m < M; ++m) {
      copy_stub<scalar_t>(C + m * ldc_s, C_temp + m * ldc_f, BLOCK_N);
    }
  }
 }

-#define INSTANTIATE_TINYGEMM_TEMPLATE(TYPE)                                 \
-  template void tinygemm_kernel<TYPE>(                                      \
-      TYPE * C, float* C_temp, const uint8_t* A, const float* scales_a,     \
-      const int32_t* qzeros_a, const uint8_t* B, const float* scales_b,     \
-      const int8_t* qzeros_b, const int32_t* compensation, int8_t* dqB_tmp, \
-      int64_t M, int64_t K, int64_t lda, int64_t ldc_f, int64_t ldc_s,      \
-      bool store_out, bool use_brgemm)
+#define INSTANTIATE_TINYGEMM_TEMPLATE(TYPE) \
+  template void tinygemm_kernel<TYPE>(      \
+      TYPE * C,                             \
+      float* C_temp,                        \
+      const uint8_t* A,                     \
+      const float* scales_a,                \
+      const int32_t* qzeros_a,              \
+      const uint8_t* B,                     \
+      const float* scales_b,                \
+      const int8_t* qzeros_b,               \
+      const int32_t* compensation,          \
+      int8_t* dqB_tmp,                      \
+      int64_t M,                            \
+      int64_t K,                            \
+      int64_t lda,                          \
+      int64_t ldc_f,                        \
+      int64_t ldc_s,                        \
+      bool store_out,                       \
+      bool use_brgemm)

 INSTANTIATE_TINYGEMM_TEMPLATE(at::BFloat16);
 INSTANTIATE_TINYGEMM_TEMPLATE(at::Half);

-at::Tensor int4_scaled_mm_cpu(at::Tensor& x, at::Tensor& w, at::Tensor& w_zeros,
-                              at::Tensor& w_scales,
-                              std::optional<at::Tensor> bias) {
-  return int4_scaled_mm_cpu_with_quant(x, w, w_scales, w_zeros, bias,
-                                       x.scalar_type());
+// int4 gemm dispatch api register
+at::Tensor int4_scaled_mm_cpu(
+    at::Tensor& x, at::Tensor& w, at::Tensor& w_zeros, at::Tensor& w_scales, std::optional<at::Tensor> bias) {
+  return int4_scaled_mm_cpu_with_quant(x, w, w_scales, w_zeros, bias, x.scalar_type());
 }
@@ -1,20 +1,83 @@
 // Adapted from
 // https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu

-#include "common.h"
-#include "vec.h"
-#include "gemm.h"
-
 // clang-format off

+#include "common.h"
+#include "gemm.h"
+#include "vec.h"
+
 namespace {

+template <typename scalar_t, bool has_bias, int BLOCK_N>
+struct scale_C {
+  static inline void apply(
+      scalar_t* __restrict__ C,
+      const int32_t* __restrict__ Ctmp,
+      const int32_t* __restrict__ Bcomp,
+      const float* __restrict__ bias,
+      float As,
+      const float* __restrict__ Bs) {
+    TORCH_CHECK(false, "scale_C: scalar path not implemented!");
+  }
+};
+
+#if defined(CPU_CAPABILITY_AVX512)
+template <bool has_bias, int BLOCK_N>
+struct scale_C<at::BFloat16, has_bias, BLOCK_N> {
+  static inline void apply(
+      at::BFloat16* __restrict__ C,
+      const int32_t* __restrict__ Ctmp,
+      const int32_t* __restrict__ Bcomp,
+      const float* __restrict__ bias,
+      float As,
+      const float* __restrict__ Bs) {
+    constexpr int COLS = BLOCK_N / 16;
+    static_assert(COLS % 2 == 0);
+
+    __m512 vc[COLS];
+    __m512 vd0 = _mm512_set1_ps(As);
+
+    auto compute = [&](auto col) {
+      __m512 vd1 = _mm512_loadu_ps(Bs + col * 16);
+      __m512i vcomp = _mm512_loadu_si512(Bcomp + col * 16);
+      __m512i vc32 = _mm512_loadu_si512(Ctmp + col * 16);
+      vc[col] = _mm512_cvtepi32_ps(_mm512_sub_epi32(vc32, vcomp));
+      if constexpr (has_bias) {
+        __m512 vbias = _mm512_loadu_ps(bias + col * 16);
+        vc[col] = _mm512_fmadd_ps(_mm512_mul_ps(vc[col], vd0), vd1, vbias);
+      } else {
+        vc[col] = _mm512_mul_ps(_mm512_mul_ps(vc[col], vd0), vd1);
+      }
+    };
+    Unroll<COLS>{}(compute);
+
+    auto storec = [&](auto col) {
+      // for COLS = 2, 4 use 512bit store
+      if constexpr (col % 2 == 0) {
+        _mm512_storeu_si512(
+            reinterpret_cast<__m512i*>((C + col * 16)), (__m512i)(_mm512_cvtne2ps_pbh(vc[col + 1], vc[col + 0])));
+      }
+    };
+    Unroll<COLS>{}(storec);
+  }
+};
+#endif
+
 template <typename scalar_t, bool has_bias, int BLOCK_M, int BLOCK_N>
 struct tinygemm_kernel_nn {
  static inline void apply(
-      const uint8_t* __restrict__ A, const int8_t* __restrict__ B, scalar_t* __restrict__ C,
-      const float* __restrict__ As, const float* __restrict__ Bs, const int32_t* __restrict__ Bcomp,
-      const float* __restrict__ bias, int64_t K, int64_t lda, int64_t ldb, int64_t ldc) {
+      const uint8_t* __restrict__ A,
+      const int8_t* __restrict__ B,
+      scalar_t* __restrict__ C,
+      const float* __restrict__ As,
+      const float* __restrict__ Bs,
+      const int32_t* __restrict__ Bcomp,
+      const float* __restrict__ bias,
+      int64_t K,
+      int64_t lda,
+      int64_t ldb,
+      int64_t ldc) {
    TORCH_CHECK(false, "tinygemm_kernel_nn: scalar path not implemented!");
  }
 };
@@ -23,10 +86,17 @@ struct tinygemm_kernel_nn {
 template <bool has_bias, int BLOCK_M, int BLOCK_N>
 struct tinygemm_kernel_nn<at::BFloat16, has_bias, BLOCK_M, BLOCK_N> {
  static inline void apply(
-      const uint8_t* __restrict__ A, const int8_t* __restrict__ B, at::BFloat16* __restrict__ C,
-      const float* __restrict__ As, const float* __restrict__ Bs, const int32_t* __restrict__ Bcomp,
-      const float* __restrict__ bias, int64_t K, int64_t lda, int64_t ldb, int64_t ldc) {
-
+      const uint8_t* __restrict__ A,
+      const int8_t* __restrict__ B,
+      at::BFloat16* __restrict__ C,
+      const float* __restrict__ As,
+      const float* __restrict__ Bs,
+      const int32_t* __restrict__ Bcomp,
+      const float* __restrict__ bias,
+      int64_t K,
+      int64_t lda,
+      int64_t ldb,
+      int64_t ldc) {
    constexpr int ROWS = BLOCK_M;
    constexpr int COLS = BLOCK_N / 16;
    static_assert(COLS % 2 == 0);
@@ -38,10 +108,10 @@ struct tinygemm_kernel_nn<at::BFloat16, has_bias, BLOCK_M, BLOCK_N> {
    __m512i vb[COLS];
    __m512i vc[ROWS * COLS];
    __m512i vcomp[COLS];
-    __m512  vd0;
-    __m512  vd1[COLS];
+    __m512 vd0;
+    __m512 vd1[COLS];

-    // oops! 4x4 spills but luckily we use 4x2
+    // oops! 4x4 spills but we use 4x2
    __m512 vbias[COLS];

    // [NOTE]: s8s8 igemm compensation in avx512-vnni
@@ -54,14 +124,12 @@ struct tinygemm_kernel_nn<at::BFloat16, has_bias, BLOCK_M, BLOCK_N> {
    // 1) 128 * b is pre-computed when packing B to vnni formats
    // 2) a + 128 is fused when dynamically quantize A
    //
-    auto loadc = [&](auto i) {
-      vc[i] = _mm512_set1_epi32(0);
-    };
+    auto loadc = [&](auto i) { vc[i] = _mm512_set1_epi32(0); };
    Unroll<ROWS * COLS>{}(loadc);

    const int64_t K4 = K >> 2;
    const int64_t lda4 = lda >> 2;
-    const int64_t ldb4 = ldb; // ldb * 4 >> 2;
+    const int64_t ldb4 = ldb;  // ldb * 4 >> 2;
    const int32_t* a_ptr = reinterpret_cast<const int32_t*>(A);
    const int32_t* b_ptr = reinterpret_cast<const int32_t*>(B);

@@ -89,7 +157,7 @@ struct tinygemm_kernel_nn<at::BFloat16, has_bias, BLOCK_M, BLOCK_N> {
      constexpr int col = i % COLS;

      // load a scale
-      if constexpr(col == 0) {
+      if constexpr (col == 0) {
        vd0 = _mm512_set1_ps(As[row]);
      }
      // load b scale and vcomp per 2 vectors
@@ -120,8 +188,7 @@ struct tinygemm_kernel_nn<at::BFloat16, has_bias, BLOCK_M, BLOCK_N> {
        }

        _mm512_storeu_si512(
-            reinterpret_cast<__m512i*>((C + row * ldc + col * 16)),
-            (__m512i)(_mm512_cvtne2ps_pbh(vc1, vc0)));
+            reinterpret_cast<__m512i*>((C + row * ldc + col * 16)), (__m512i)(_mm512_cvtne2ps_pbh(vc1, vc0)));
      }
    };
    Unroll<ROWS * COLS>{}(storec);
@@ -129,11 +196,19 @@ struct tinygemm_kernel_nn<at::BFloat16, has_bias, BLOCK_M, BLOCK_N> {
 };
 #endif

-#define LAUNCH_TINYGEMM_KERNEL_NN(MB_SIZE, NB_SIZE)                          \
-    tinygemm_kernel_nn<scalar_t, has_bias, MB_SIZE, NB_SIZE>::apply(         \
-        A + mb_start * lda, B + nb_start * 4, C + mb_start * ldc + nb_start, \
-        As + mb_start, Bs + nb_start, Bcomp + nb_start,                      \
-        has_bias ? bias + nb_start : nullptr, K, lda, ldb, ldc);
+#define LAUNCH_TINYGEMM_KERNEL_NN(MB_SIZE, NB_SIZE)                \
+  tinygemm_kernel_nn<scalar_t, has_bias, MB_SIZE, NB_SIZE>::apply( \
+      A + mb_start * lda,                                          \
+      B + nb_start * 4,                                            \
+      C + mb_start * ldc + nb_start,                               \
+      As + mb_start,                                               \
+      Bs + nb_start,                                               \
+      Bcomp + nb_start,                                            \
+      has_bias ? bias + nb_start : nullptr,                        \
+      K,                                                           \
+      lda,                                                         \
+      ldb,                                                         \
+      ldc);

 template <typename scalar_t, bool has_bias>
 void tinygemm_kernel(
@@ -151,10 +226,20 @@ void tinygemm_kernel(
    int64_t ldb,
    int64_t ldc,
    bool brg) {
-
  // B compensation
  const int32_t* Bcomp = reinterpret_cast<const int32_t*>(B + block_size_n() * K);

+  if (brg) {
+    constexpr int BLOCK_N = block_size_n();
+    at::native::cpublas::brgemm(M, N, K, lda, ldb, BLOCK_N, /* add_C */ false, A, B, Ctmp);
+
+    // apply compensation and scale
+    for (int64_t m = 0; m < M; ++m) {
+      scale_C<scalar_t, has_bias, BLOCK_N>::apply(C + m * ldc, Ctmp + m * BLOCK_N, Bcomp, bias, As[m], Bs);
+    }
+    return;
+  }
+
  // pattern: 1-4-16
  constexpr int64_t BLOCK_M = 4;
  constexpr int64_t BLOCK_N = 64;
@@ -167,26 +252,43 @@ void tinygemm_kernel(
      int64_t nb_start = nb * BLOCK_N;
      int64_t nb_size = std::min(BLOCK_N, N - nb_start);

-      switch(mb_size << 4 | nb_size >> 4) {
+      switch (mb_size << 4 | nb_size >> 4) {
        // mb_size = 1
-        case 0x12: LAUNCH_TINYGEMM_KERNEL_NN(1, 32); break;
-        case 0x14: LAUNCH_TINYGEMM_KERNEL_NN(1, 64); break;
+        case 0x12:
+          LAUNCH_TINYGEMM_KERNEL_NN(1, 32);
+          break;
+        case 0x14:
+          LAUNCH_TINYGEMM_KERNEL_NN(1, 64);
+          break;
        // mb_size = 2
-        case 0x22: LAUNCH_TINYGEMM_KERNEL_NN(2, 32); break;
-        case 0x24: LAUNCH_TINYGEMM_KERNEL_NN(2, 64); break;
+        case 0x22:
+          LAUNCH_TINYGEMM_KERNEL_NN(2, 32);
+          break;
+        case 0x24:
+          LAUNCH_TINYGEMM_KERNEL_NN(2, 64);
+          break;
        // mb_size = 3
-        case 0x32: LAUNCH_TINYGEMM_KERNEL_NN(3, 32); break;
-        case 0x34: LAUNCH_TINYGEMM_KERNEL_NN(3, 64); break;
+        case 0x32:
+          LAUNCH_TINYGEMM_KERNEL_NN(3, 32);
+          break;
+        case 0x34:
+          LAUNCH_TINYGEMM_KERNEL_NN(3, 64);
+          break;
        // mb_size = 4
-        case 0x42: LAUNCH_TINYGEMM_KERNEL_NN(4, 32); break;
-        case 0x44: LAUNCH_TINYGEMM_KERNEL_NN(4, 64); break;
-        default: TORCH_CHECK(false, "Unexpected block size, ", mb_size, "x", nb_size);
+        case 0x42:
+          LAUNCH_TINYGEMM_KERNEL_NN(4, 32);
+          break;
+        case 0x44:
+          LAUNCH_TINYGEMM_KERNEL_NN(4, 64);
+          break;
+        default:
+          TORCH_CHECK(false, "Unexpected block size, ", mb_size, "x", "nb_size");
      }
    }
  }
 }

-template<typename scalar_t>
+template <typename scalar_t>
 void int8_scaled_mm_kernel_impl(
    scalar_t* __restrict__ out,
    const uint8_t* __restrict__ mat1,
@@ -197,28 +299,22 @@ void int8_scaled_mm_kernel_impl(
    int64_t M,
    int64_t N,
    int64_t K) {
-
  constexpr int64_t BLOCK_M = block_size_m();
  constexpr int64_t BLOCK_N = block_size_n();
  const int64_t MB = div_up(M, BLOCK_M);
  const int64_t NB = div_up(N, BLOCK_N);

-  // TODO: brgemm u8s8 depends on PyTorch 2.7 release.
-  const bool use_brgemm = false;
+  const bool use_brgemm = can_use_brgemm<int8_t>(M);

  // K + 4 after compensation
  const int64_t packed_row_size = get_row_size<int8_t>(K);

  AT_DISPATCH_BOOL(bias != nullptr, has_bias, [&] {
-    at::parallel_for(0, MB * NB, 0, [&](int64_t begin, int64_t end) {
-      int64_t mb{0}, nb{0};
-      data_index_init(begin, mb, MB, nb, NB);
-
+    parallel_2d(MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
      // for brgemm, use int32_t for accumulate
      alignas(64) int32_t Ctmp[BLOCK_M * BLOCK_N];

-      for (int i = begin; i < end; ++i) {
-        UNUSED(i);
+      loop_2d<int8_t>(mb0, mb1, nb0, nb1, BLOCK_N * K, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
        int mb_start = mb * BLOCK_M;
        int mb_size = std::min(M - mb_start, BLOCK_M);
        int nb_start = nb * BLOCK_N;
@@ -239,10 +335,7 @@ void int8_scaled_mm_kernel_impl(
            /* ldb */ nb_size,
            /* ldc */ N,
            /* brg */ use_brgemm);
-
-        // move to the next index
-        data_index_step(mb, MB, nb, NB);
-      }
+      });

      if (use_brgemm) {
        at::native::cpublas::brgemm_release();
@@ -251,28 +344,47 @@ void int8_scaled_mm_kernel_impl(
  });
 }

-} // anonymous namespace
+}  // anonymous namespace

 // tinygemm interface
 template <typename scalar_t>
-void tinygemm_kernel(const uint8_t* __restrict__ A, const int8_t* __restrict__ B, scalar_t* __restrict__ C,
-    int32_t* __restrict__ Ctmp,  const float* __restrict__ As, const float* __restrict__ Bs,
-    int64_t M, int64_t N, int64_t K, int64_t lda, int64_t ldb, int64_t ldc, bool brg) {
+void tinygemm_kernel(
+    const uint8_t* __restrict__ A,
+    const int8_t* __restrict__ B,
+    scalar_t* __restrict__ C,
+    int32_t* __restrict__ Ctmp,
+    const float* __restrict__ As,
+    const float* __restrict__ Bs,
+    int64_t M,
+    int64_t N,
+    int64_t K,
+    int64_t lda,
+    int64_t ldb,
+    int64_t ldc,
+    bool brg) {
  tinygemm_kernel<scalar_t, false>(A, B, C, Ctmp, As, Bs, nullptr, M, N, K, lda, ldb, ldc, brg);
 }

-#define INSTANTIATE_TINYGEMM_TEMPLATE(TYPE)                                                     \
-    template void tinygemm_kernel<TYPE>(                                                        \
-        const uint8_t* __restrict__ A, const int8_t* __restrict__ B, TYPE* __restrict__ C,      \
-        int32_t* __restrict__ Ctmp, const float* __restrict__ As, const float* __restrict__ Bs, \
-        int64_t M, int64_t N, int64_t K, int64_t lda, int64_t ldb, int64_t ldc, bool brg)
+#define INSTANTIATE_TINYGEMM_TEMPLATE(TYPE) \
+  template void tinygemm_kernel<TYPE>(      \
+      const uint8_t* __restrict__ A,        \
+      const int8_t* __restrict__ B,         \
+      TYPE* __restrict__ C,                 \
+      int32_t* __restrict__ Ctmp,           \
+      const float* __restrict__ As,         \
+      const float* __restrict__ Bs,         \
+      int64_t M,                            \
+      int64_t N,                            \
+      int64_t K,                            \
+      int64_t lda,                          \
+      int64_t ldb,                          \
+      int64_t ldc,                          \
+      bool brg)

 INSTANTIATE_TINYGEMM_TEMPLATE(at::BFloat16);
 INSTANTIATE_TINYGEMM_TEMPLATE(at::Half);

 std::tuple<at::Tensor, at::Tensor> per_token_quant_int8_cpu(at::Tensor& A) {
-  RECORD_FUNCTION("sgl-kernel::per_token_quant_int8_cpu", std::vector<c10::IValue>({A}));
-
  CHECK_LAST_DIM_CONTIGUOUS_INPUT(A);
  CHECK_DIM(2, A);

@@ -281,8 +393,7 @@ std::tuple<at::Tensor, at::Tensor> per_token_quant_int8_cpu(at::Tensor& A) {
  int64_t lda = A.stride(0);

  const auto st = A.scalar_type();
-  TORCH_CHECK(st == at::kBFloat16 || st == at::kHalf,
-      "per_token_quant_int8: expect A to be bfloat16 or half.");
+  TORCH_CHECK(st == at::kBFloat16 || st == at::kHalf, "per_token_quant_int8: expect A to be bfloat16 or half.");

  auto Aq = at::empty({M, K}, A.options().dtype(at::kByte));
  auto As = at::empty({M}, A.options().dtype(at::kFloat));
@@ -292,13 +403,9 @@ std::tuple<at::Tensor, at::Tensor> per_token_quant_int8_cpu(at::Tensor& A) {
    float* __restrict__ As_data = As.data_ptr<float>();
    const scalar_t* __restrict__ A_data = A.data_ptr<scalar_t>();

-    at::parallel_for(0, M, 0, [&] (int64_t begin, int64_t end) {
+    at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
      for (int64_t m = begin; m < end; ++m) {
-        quantize_row_int8<scalar_t>(
-            Aq_data + m * K,
-            As_data[m],
-            A_data + m * lda,
-            K);
+        quantize_row_int8<scalar_t>(Aq_data + m * K, As_data[m], A_data + m * lda, K);
      }
    });
  });
@@ -315,11 +422,14 @@ std::tuple<at::Tensor, at::Tensor> per_token_quant_int8_cpu(at::Tensor& A) {
 // bias    : [N]
 // out     : [M, N]
 //
-at::Tensor int8_scaled_mm_cpu(at::Tensor& mat1, at::Tensor& mat2,
-    at::Tensor& scales1, at::Tensor& scales2,
-    std::optional<at::Tensor>& bias, at::ScalarType out_dtype, bool is_vnni) {
-  RECORD_FUNCTION("sgl-kernel::int8_scaled_mm_cpu", std::vector<c10::IValue>({mat1, mat2, scales1, scales2, bias}));
-
+at::Tensor int8_scaled_mm_cpu(
+    at::Tensor& mat1,
+    at::Tensor& mat2,
+    at::Tensor& scales1,
+    at::Tensor& scales2,
+    const std::optional<at::Tensor>& bias,
+    at::ScalarType out_dtype,
+    bool is_vnni) {
  auto packed_w = is_vnni ? mat2 : convert_weight_packed(mat2);

  CHECK_INPUT(mat1);
@@ -340,7 +450,8 @@ at::Tensor int8_scaled_mm_cpu(at::Tensor& mat1, at::Tensor& mat2,

  TORCH_CHECK(mat1.scalar_type() == at::kByte, "int8_scaled_mm: expect mat1 to be uint8.");
  TORCH_CHECK(mat2.scalar_type() == at::kChar, "int8_scaled_mm: expect mat2 to be int8.");
-  TORCH_CHECK(scales1.scalar_type() == at::kFloat && scales2.scalar_type() == at::kFloat,
+  TORCH_CHECK(
+      scales1.scalar_type() == at::kFloat && scales2.scalar_type() == at::kFloat,
      "int8_scaled_mm: expect scales to be float32.");

  auto out = at::empty({M, N}, mat1.options().dtype(out_dtype));
@@ -368,10 +479,13 @@ at::Tensor int8_scaled_mm_cpu(at::Tensor& mat1, at::Tensor& mat2,
 }

 // fused `per_token_quant_int8_cpu` and `int8_scaled_mm_cpu`
-at::Tensor int8_scaled_mm_with_quant(at::Tensor& mat1, at::Tensor& mat2, at::Tensor& scales2,
-    const std::optional<at::Tensor>& bias, at::ScalarType out_dtype, bool is_vnni) {
-  RECORD_FUNCTION("sgl-kernel::int8_scaled_mm_cpu", std::vector<c10::IValue>({mat1, mat2, scales2, bias}));
-
+at::Tensor int8_scaled_mm_with_quant(
+    at::Tensor& mat1,
+    at::Tensor& mat2,
+    at::Tensor& scales2,
+    const std::optional<at::Tensor>& bias,
+    at::ScalarType out_dtype,
+    bool is_vnni) {
  auto packed_w = is_vnni ? mat2 : convert_weight_packed(mat2);

  CHECK_LAST_DIM_CONTIGUOUS_INPUT(mat1);
@@ -390,14 +504,10 @@ at::Tensor int8_scaled_mm_with_quant(at::Tensor& mat1, at::Tensor& mat2, at::Ten
  CHECK_EQ(scales2.numel(), N);

  const auto st = mat1.scalar_type();
-  TORCH_CHECK(st == at::kBFloat16 || st == at::kHalf,
-      "int8_scaled_mm_with_quant: expect A to be bfloat16 or half.");
-  TORCH_CHECK(st == out_dtype,
-      "int8_scaled_mm_with_quant: expect A has same dtype with out_dtype.");
-  TORCH_CHECK(mat2.scalar_type() == at::kChar,
-      "int8_scaled_mm_with_quant: expect mat2 to be int8.");
-  TORCH_CHECK(scales2.scalar_type() == at::kFloat,
-      "int8_scaled_mm_with_quant: expect scales to be float32.");
+  TORCH_CHECK(st == at::kBFloat16 || st == at::kHalf, "int8_scaled_mm_with_quant: expect A to be bfloat16 or half.");
+  TORCH_CHECK(st == out_dtype, "int8_scaled_mm_with_quant: expect A has same dtype with out_dtype.");
+  TORCH_CHECK(mat2.scalar_type() == at::kChar, "int8_scaled_mm_with_quant: expect mat2 to be int8.");
+  TORCH_CHECK(scales2.scalar_type() == at::kFloat, "int8_scaled_mm_with_quant: expect scales to be float32.");

  const int64_t buffer_size = M * K + M * sizeof(float);
  auto buffer = at::empty({buffer_size}, mat1.options().dtype(at::kByte));
@@ -415,13 +525,9 @@ at::Tensor int8_scaled_mm_with_quant(at::Tensor& mat1, at::Tensor& mat2, at::Ten
    float* __restrict__ As_data = (float*)((void*)(Aq_data + M * K));
    const scalar_t* __restrict__ A_data = mat1.data_ptr<scalar_t>();

-    at::parallel_for(0, M, 0, [&] (int64_t begin, int64_t end) {
+    at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
      for (int64_t m = begin; m < end; ++m) {
-        quantize_row_int8<scalar_t>(
-            Aq_data + m * K,
-            As_data[m],
-            A_data + m * lda,
-            K);
+        quantize_row_int8<scalar_t>(Aq_data + m * K, As_data[m], A_data + m * lda, K);
      }
    });

@@ -1,12 +1,13 @@
 // Adapted from
 // https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu

-#include "common.h"
-#include "vec.h"
-#include "gemm.h"
-
 // clang-format off

+#include "moe.h"
+
+#include "common.h"
+#include "gemm.h"
+
 namespace {

 // [NOTE]: Fused MoE kernel with AMX
@@ -30,109 +31,6 @@ namespace {
 //     3. abstract at::native::cpublas::brgemm with WoQ gemm (M = 1 & M != 1)
 //

-template <typename scalar_t>
-inline void fill_stub(scalar_t* __restrict__ out, scalar_t val, int64_t size) {
-  using Vec = at::vec::Vectorized<scalar_t>;
-  const Vec data_vec(val);
-  at::vec::map<scalar_t>([data_vec](Vec out) { return out = data_vec; }, out, out, size);
-}
-
-template <typename scalar_t>
-inline void copy_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, int64_t size) {
-  using Vec = at::vec::Vectorized<scalar_t>;
-  // no remainder
-  #pragma GCC unroll 4
-  for (int64_t d = 0; d < size; d += Vec::size()) {
-    Vec data = Vec::loadu(input + d);
-    data.store(out + d);
-  }
-}
-
-template <typename scalar_t>
-inline void copy_mul_stub(scalar_t* __restrict__ out, const float* __restrict__ input, float weight, int64_t size) {
-  using bVec = at::vec::Vectorized<scalar_t>;
-  using fVec = at::vec::Vectorized<float>;
-  constexpr int kVecSize = bVec::size();
-  const fVec weight_vec = fVec(weight);
-  int64_t d;
-  #pragma GCC unroll 4
-  for (d = 0; d <= size - kVecSize; d += kVecSize) {
-    fVec data0 = fVec::loadu(input + d) * weight_vec;
-    fVec data1 = fVec::loadu(input + d + fVec::size()) * weight_vec;
-    bVec out_vec = convert_from_float_ext<scalar_t>(data0, data1);
-    out_vec.store(out + d);
-  }
-  for (; d < size; ++d) {
-    out[d] = static_cast<scalar_t>(input[d] * weight);
-  }
-}
-
-// acc from [topk, K] to [K]
-template <typename scalar_t>
-inline void sum_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, int64_t topk, int64_t K) {
-  using bVec = at::vec::Vectorized<scalar_t>;
-  using fVec = at::vec::Vectorized<float>;
-  constexpr int kVecSize = bVec::size();
-  if (topk == 1) {
-    // do copy for topk = 1
-    copy_stub(out, input, K);
-  } else {
-    // do sum for topk != 1
-    int64_t d;
-    #pragma GCC unroll 4
-    for (d = 0; d <= K - kVecSize; d += kVecSize) {
-      fVec sum_fvec0 = fVec(0.f);
-      fVec sum_fvec1 = fVec(0.f);
-      for (int t = 0; t < topk; ++t) {
-        bVec x_bvec = bVec::loadu(input + t * K + d);
-        fVec x_fvec0, x_fvec1;
-        std::tie(x_fvec0, x_fvec1) = at::vec::convert_to_float(x_bvec);
-
-        sum_fvec0 += x_fvec0;
-        sum_fvec1 += x_fvec1;
-      }
-      bVec out_bvec = convert_from_float_ext<scalar_t>(sum_fvec0, sum_fvec1);
-      out_bvec.store(out + d);
-    }
-    for (; d < K; ++d) {
-      float sum_val = 0.f;
-      for (int t = 0; t < topk; ++t) {
-        sum_val += static_cast<float>(input[t * K + d]);
-      }
-      out[d] = static_cast<scalar_t>(sum_val);
-    }
-  }
-}
-
-// out = input + input2 * scale
-template <typename scalar_t>
-inline void add_mul_stub(scalar_t* __restrict__ out, const float* __restrict__ input,
-    const scalar_t* __restrict__ input2, float scale, int64_t size) {
-
-  using bVec = at::vec::Vectorized<scalar_t>;
-  using fVec = at::vec::Vectorized<float>;
-  constexpr int kVecSize = bVec::size();
-  const fVec s_vec = fVec(scale);
-  int64_t d;
-  #pragma GCC unroll 4
-  for (d = 0; d <= size - kVecSize; d += kVecSize) {
-    fVec x0 = fVec::loadu(input + d);
-    fVec x1 = fVec::loadu(input + d + fVec::size());
-
-    bVec y_bvec = bVec::loadu(input2 + d);
-    fVec y0, y1;
-    std::tie(y0, y1) = at::vec::convert_to_float(y_bvec);
-
-    x0 = x0 + y0 * s_vec;
-    x1 = x1 + y1 * s_vec;
-    bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1);
-    out_vec.store(out + d);
-  }
-  for (; d < size; ++d) {
-    out[d] = static_cast<scalar_t>(input[d] + float(input2[d]) * scale);
-  }
-}
-
 template <int BLOCK_M>
 int moe_align_block_size(
    int32_t* __restrict__ sorted_ids,
@@ -144,8 +42,7 @@ int moe_align_block_size(
    int num_experts,
    int numel,
    int num_threads) {
-
-  #define T_INDEX(tt) total_cnts + (tt) * num_experts
+#define T_INDEX(tt) total_cnts + (tt) * num_experts

  // accumulate count of expert ids locally
  at::parallel_for(0, numel, 0, [&](int begin, int end) {
@@ -160,8 +57,7 @@ int moe_align_block_size(
  using iVec = at::vec::Vectorized<int32_t>;
  for (int t = 0; t < num_threads; ++t) {
    at::vec::map2<int32_t>(
-        [](iVec x, iVec y) { return x + y; },
-        T_INDEX(t + 1), T_INDEX(t + 1), T_INDEX(t), num_experts);
+        [](iVec x, iVec y) { return x + y; }, T_INDEX(t + 1), T_INDEX(t + 1), T_INDEX(t), num_experts);
  }

  // the last row holds sums of each experts
@@ -201,7 +97,9 @@ int moe_align_block_size(
  // padding value for sorted_ids: numel
  auto sorted_id_size = [=](const int32_t* sorted_ids_ptr) {
    for (int d = 0; d < BLOCK_M; ++d) {
-      if (sorted_ids_ptr[d] == numel) { return d; }
+      if (sorted_ids_ptr[d] == numel) {
+        return d;
+      }
    }
    return BLOCK_M;
  };
@@ -215,7 +113,7 @@ int moe_align_block_size(
      offsets[mb + 1] = sorted_id_size(sorted_ids + mb * BLOCK_M);
    }
  });
-  // TODO: do we need to vectorize this ?
+  // TODO: do we need to vecterize this ?
  for (int mb = 0; mb < num_token_blocks; ++mb) {
    offsets[mb + 1] += offsets[mb];
  }
@@ -236,7 +134,6 @@ inline void silu_and_mul(
    const float* __restrict__ input1,  // y: y0, y1
    int64_t m_size,
    int64_t N) {
-
  using bVec = at::vec::Vectorized<scalar_t>;
  using fVec = at::vec::Vectorized<float>;

@@ -269,8 +166,14 @@ inline void silu_and_mul(
 template <typename scalar_t, int BLOCK_M, int BLOCK_N>
 struct tinygemm_kernel_nn2 {
  static inline void apply(
-      const scalar_t* __restrict__ A, const scalar_t* __restrict__ B0, const scalar_t* __restrict__ B1,
-      scalar_t* __restrict__ C, int64_t K, int64_t lda, int64_t ldb, int64_t ldc) {
+      const scalar_t* __restrict__ A,
+      const scalar_t* __restrict__ B0,
+      const scalar_t* __restrict__ B1,
+      scalar_t* __restrict__ C,
+      int64_t K,
+      int64_t lda,
+      int64_t ldb,
+      int64_t ldc) {
    TORCH_CHECK(false, "tinygemm_kernel_nn: scalar path not implemented!");
  }
 };
@@ -279,9 +182,14 @@ struct tinygemm_kernel_nn2 {
 template <int BLOCK_M, int BLOCK_N>
 struct tinygemm_kernel_nn2<at::BFloat16, BLOCK_M, BLOCK_N> {
  static inline void apply(
-      const at::BFloat16* __restrict__ A, const at::BFloat16* __restrict__ B0, const at::BFloat16* __restrict__ B1,
-      at::BFloat16* __restrict__ C, int64_t K, int64_t lda, int64_t ldb, int64_t ldc) {
-
+      const at::BFloat16* __restrict__ A,
+      const at::BFloat16* __restrict__ B0,
+      const at::BFloat16* __restrict__ B1,
+      at::BFloat16* __restrict__ C,
+      int64_t K,
+      int64_t lda,
+      int64_t ldb,
+      int64_t ldc) {
    constexpr int ROWS = BLOCK_M;
    constexpr int COLS = BLOCK_N / 16;

@@ -304,7 +212,7 @@ struct tinygemm_kernel_nn2<at::BFloat16, BLOCK_M, BLOCK_N> {

    const int64_t K2 = K >> 1;
    const int64_t lda2 = lda >> 1;
-    const int64_t ldb2 = ldb; // ldb * 2 >> 1;
+    const int64_t ldb2 = ldb;  // ldb * 2 >> 1;
    const float* a_ptr = reinterpret_cast<const float*>(A);
    const float* b0_ptr = reinterpret_cast<const float*>(B0);
    const float* b1_ptr = reinterpret_cast<const float*>(B1);
@@ -352,17 +260,16 @@ struct tinygemm_kernel_nn2<at::BFloat16, BLOCK_M, BLOCK_N> {
        _mm512_storeu_si512(
            reinterpret_cast<__m512i*>((C + row * ldc + col * 16)),
            (__m512i)(_mm512_cvtne2ps_pbh(__m512(x1), __m512(x0))));
-        }
+      }
    };
    Unroll<ROWS * COLS>{}(storec);
  }
 };
 #endif

-#define LAUNCH_TINYGEMM_KERNEL_NN(MB_SIZE, NB_SIZE)                          \
-    tinygemm_kernel_nn2<scalar_t, MB_SIZE, NB_SIZE>::apply(                  \
-        A + mb_start * lda, B0 + nb_start * 2, B1 + nb_start * 2,            \
-        C + mb_start * ldc + nb_start, K, lda, ldb, ldc);
+#define LAUNCH_TINYGEMM_KERNEL_NN(MB_SIZE, NB_SIZE)       \
+  tinygemm_kernel_nn2<scalar_t, MB_SIZE, NB_SIZE>::apply( \
+      A + mb_start * lda, B0 + nb_start * 2, B1 + nb_start * 2, C + mb_start * ldc + nb_start, K, lda, ldb, ldc);

 template <typename scalar_t>
 void tinygemm_kernel(
@@ -376,7 +283,6 @@ void tinygemm_kernel(
    int64_t lda,
    int64_t ldb,
    int64_t ldc) {
-
  // pattern: 1-(2+2)-(8+8)
  constexpr int64_t BLOCK_M = 4;
  constexpr int64_t BLOCK_N = 32;
@@ -389,16 +295,25 @@ void tinygemm_kernel(
      int64_t nb_start = nb * BLOCK_N;
      int64_t nb_size = std::min(BLOCK_N, N - nb_start);

-      switch(mb_size << 4 | nb_size >> 4) {
+      switch (mb_size << 4 | nb_size >> 4) {
        // mb_size = 1
-        case 0x12: LAUNCH_TINYGEMM_KERNEL_NN(1, 32); break;
+        case 0x12:
+          LAUNCH_TINYGEMM_KERNEL_NN(1, 32);
+          break;
        // mb_size = 2
-        case 0x22: LAUNCH_TINYGEMM_KERNEL_NN(2, 32); break;
+        case 0x22:
+          LAUNCH_TINYGEMM_KERNEL_NN(2, 32);
+          break;
        // mb_size = 3
-        case 0x32: LAUNCH_TINYGEMM_KERNEL_NN(3, 32); break;
+        case 0x32:
+          LAUNCH_TINYGEMM_KERNEL_NN(3, 32);
+          break;
        // mb_size = 4
-        case 0x42: LAUNCH_TINYGEMM_KERNEL_NN(4, 32); break;
-        default: TORCH_CHECK(false, "Unexpected block size, ", mb_size, "x", nb_size);
+        case 0x42:
+          LAUNCH_TINYGEMM_KERNEL_NN(4, 32);
+          break;
+        default:
+          TORCH_CHECK(false, "Unexpected block size, ", mb_size, "x", "nb_size");
      }
    }
  }
@@ -407,8 +322,13 @@ void tinygemm_kernel(
 template <typename scalar_t, int BLOCK_M, int BLOCK_N>
 struct tinygemm_kernel_nn {
  static inline void apply(
-      const scalar_t* __restrict__ A, const scalar_t* __restrict__ B, float* __restrict__ C,
-      int64_t K, int64_t lda, int64_t ldb, int64_t ldc) {
+      const scalar_t* __restrict__ A,
+      const scalar_t* __restrict__ B,
+      float* __restrict__ C,
+      int64_t K,
+      int64_t lda,
+      int64_t ldb,
+      int64_t ldc) {
    TORCH_CHECK(false, "tinygemm_kernel_nn: scalar path not implemented!");
  }
 };
@@ -417,9 +337,13 @@ struct tinygemm_kernel_nn {
 template <int BLOCK_M, int BLOCK_N>
 struct tinygemm_kernel_nn<at::BFloat16, BLOCK_M, BLOCK_N> {
  static inline void apply(
-      const at::BFloat16* __restrict__ A, const at::BFloat16* __restrict__ B, float* __restrict__ C,
-      int64_t K, int64_t lda, int64_t ldb, int64_t ldc) {
-
+      const at::BFloat16* __restrict__ A,
+      const at::BFloat16* __restrict__ B,
+      float* __restrict__ C,
+      int64_t K,
+      int64_t lda,
+      int64_t ldb,
+      int64_t ldc) {
    constexpr int ROWS = BLOCK_M;
    constexpr int COLS = BLOCK_N / 16;

@@ -432,14 +356,12 @@ struct tinygemm_kernel_nn<at::BFloat16, BLOCK_M, BLOCK_N> {
    __m512bh vb[COLS];
    __m512 vc[ROWS * COLS];

-    auto loadc = [&](auto i) {
-      vc[i] = _mm512_set1_ps(0.f);
-    };
+    auto loadc = [&](auto i) { vc[i] = _mm512_set1_ps(0.f); };
    Unroll<ROWS * COLS>{}(loadc);

    const int64_t K2 = K >> 1;
    const int64_t lda2 = lda >> 1;
-    const int64_t ldb2 = ldb; // ldb * 2 >> 1;
+    const int64_t ldb2 = ldb;  // ldb * 2 >> 1;
    const float* a_ptr = reinterpret_cast<const float*>(A);
    const float* b_ptr = reinterpret_cast<const float*>(B);

@@ -466,17 +388,15 @@ struct tinygemm_kernel_nn<at::BFloat16, BLOCK_M, BLOCK_N> {
      constexpr int row = i / COLS;
      constexpr int col = i % COLS;
      _mm512_storeu_ps(reinterpret_cast<__m512*>(C + row * ldc + col * 16), vc[i]);
-
    };
    Unroll<ROWS * COLS>{}(storec);
  }
 };
 #endif

-#define LAUNCH_TINYGEMM_KERNEL_NN2(MB_SIZE, NB_SIZE)                         \
-    tinygemm_kernel_nn<scalar_t, MB_SIZE, NB_SIZE>::apply(                   \
-        A + mb_start * lda, B + nb_start * 2, C + mb_start * ldc + nb_start, \
-        K, lda, ldb, ldc);
+#define LAUNCH_TINYGEMM_KERNEL_NN2(MB_SIZE, NB_SIZE)     \
+  tinygemm_kernel_nn<scalar_t, MB_SIZE, NB_SIZE>::apply( \
+      A + mb_start * lda, B + nb_start * 2, C + mb_start * ldc + nb_start, K, lda, ldb, ldc);

 template <typename scalar_t>
 void tinygemm_kernel(
@@ -489,7 +409,6 @@ void tinygemm_kernel(
    int64_t lda,
    int64_t ldb,
    int64_t ldc) {
-
  // pattern: 1-2-8
  constexpr int64_t BLOCK_M = 4;
  constexpr int64_t BLOCK_N = 32;
@@ -502,16 +421,25 @@ void tinygemm_kernel(
      int64_t nb_start = nb * BLOCK_N;
      int64_t nb_size = std::min(BLOCK_N, N - nb_start);

-      switch(mb_size << 4 | nb_size >> 4) {
+      switch (mb_size << 4 | nb_size >> 4) {
        // mb_size = 1
-        case 0x12: LAUNCH_TINYGEMM_KERNEL_NN2(1, 32); break;
+        case 0x12:
+          LAUNCH_TINYGEMM_KERNEL_NN2(1, 32);
+          break;
        // mb_size = 2
-        case 0x22: LAUNCH_TINYGEMM_KERNEL_NN2(2, 32); break;
+        case 0x22:
+          LAUNCH_TINYGEMM_KERNEL_NN2(2, 32);
+          break;
        // mb_size = 3
-        case 0x32: LAUNCH_TINYGEMM_KERNEL_NN2(3, 32); break;
+        case 0x32:
+          LAUNCH_TINYGEMM_KERNEL_NN2(3, 32);
+          break;
        // mb_size = 4
-        case 0x42: LAUNCH_TINYGEMM_KERNEL_NN2(4, 32); break;
-        default: TORCH_CHECK(false, "Unexpected block size, ", mb_size, "x", nb_size);
+        case 0x42:
+          LAUNCH_TINYGEMM_KERNEL_NN2(4, 32);
+          break;
+        default:
+          TORCH_CHECK(false, "Unexpected block size, ", mb_size, "x", "nb_size");
      }
    }
  }
@@ -537,7 +465,6 @@ void fused_experts_kernel_impl(
    int64_t E,
    int64_t topk,
    int64_t num_tokens_post_pad) {
-
  // handle 2 tiles per block
  constexpr int64_t BLOCK_M = block_size_m();
  constexpr int64_t BLOCK_N = block_size_n();
@@ -552,39 +479,36 @@ void fused_experts_kernel_impl(
  const int64_t stride_e = 2 * N * K;
  const int64_t stride_n = K;

+  int64_t avg_M = std::max(int64_t(1), M * topk / E);
+  const bool use_brgemm = can_use_brgemm<scalar_t>(avg_M);
+
  // here we only parallel on half of 2N to fuse silu_and_mul with gemm
-  at::parallel_for(0, MB * NB, 0, [&](int64_t begin, int64_t end) {
+  parallel_2d(MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
    // get local pointers
-    int tid = at::get_thread_num();
+    int tid = get_thread_num();
    scalar_t* __restrict__ A = A_tmp + tid * BLOCK_M * K;
    float* __restrict__ C0 = C_tmp + tid * 2 * BLOCK_M * BLOCK_N;
    float* __restrict__ C1 = C0 + BLOCK_M * BLOCK_N;

-    bool is_brgemm_used = false;
-
-    for (int64_t i = begin; i < end; ++i) {
-      int64_t mb = i / NB;
-      int64_t nb = i % NB;
-
-      // nb0 from top half and nb1 from bottom half
-      int64_t nb0 = nb, nb1 = nb + NB;
-      int64_t n_size = std::min(N - nb0 * BLOCK_N, BLOCK_N);
+    loop_2d<scalar_t>(mb0, mb1, nb0, nb1, BLOCK_N * K * 2, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
+      // nb_upper from top half and nb_lower from bottom half
+      int64_t nb_upper = nb, nb_lower = nb + NB;
+      int64_t n_size = std::min(N - nb * BLOCK_N, BLOCK_N);

      // B shape [K, n_size] in vnni format
      int32_t expert_id = expert_ids[mb];
-      const scalar_t* __restrict__ B0 = packed_w1 + expert_id * stride_e + nb0 * BLOCK_N * stride_n;
-      const scalar_t* __restrict__ B1 = packed_w1 + expert_id * stride_e + nb1 * BLOCK_N * stride_n;
+      const scalar_t* __restrict__ B0 = packed_w1 + expert_id * stride_e + nb_upper * BLOCK_N * stride_n;
+      const scalar_t* __restrict__ B1 = packed_w1 + expert_id * stride_e + nb_lower * BLOCK_N * stride_n;

-      // 1.a load A
-      const int32_t* A_ids = sorted_ids + mb * BLOCK_M;
      int64_t m_size = offsets[mb + 1] - offsets[mb];

-      const bool use_brgemm = can_use_brgemm<scalar_t>(m_size);
-      is_brgemm_used = is_brgemm_used || use_brgemm;
-
-      for (int64_t m = 0; m < m_size; ++m) {
-        int32_t index = A_ids[m] / topk;
-        copy_stub(A + m * K, input + index * K, K);
+      if (nb_offset == 0) {
+        // 1.a load A
+        const int32_t* A_ids = sorted_ids + mb * BLOCK_M;
+        for (int64_t m = 0; m < m_size; ++m) {
+          int32_t index = A_ids[m] / topk;
+          copy_stub(A + m * K, input + index * K, K);
+        }
      }

      if (use_brgemm) {
@@ -616,12 +540,7 @@ void fused_experts_kernel_impl(

        // 1.d silu and mul
        const int64_t offset = offsets[mb];
-        silu_and_mul<scalar_t, BLOCK_N>(
-            ic1 + offset * N + nb * BLOCK_N,
-            C0,
-            C1,
-            m_size,
-            N);
+        silu_and_mul<scalar_t, BLOCK_N>(ic1 + offset * N + nb * BLOCK_N, C0, C1, m_size, N);
      } else {
        // fused 1.bcd: silu_and_mul(A @ B0, A @ B1)
        const int64_t offset = offsets[mb];
@@ -637,9 +556,9 @@ void fused_experts_kernel_impl(
            /* ldb   */ n_size,
            /* ldc   */ N);
      }
-    }
+    });

-    if (is_brgemm_used) {
+    if (use_brgemm) {
      at::native::cpublas::brgemm_release();
    }
  });
@@ -654,24 +573,16 @@ void fused_experts_kernel_impl(
  const int64_t stride_oc = IC;

  // parallel on [MB2, NB2]
-  at::parallel_for(0, MB2 * NB2, 0, [&](int64_t begin, int64_t end) {
+  parallel_2d(MB2, NB2, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
    // get local pointers
-    int tid = at::get_thread_num();
+    int tid = get_thread_num();
    // we won't be using C1 for gemm2
    float* __restrict__ C = C_tmp + tid * 2 * BLOCK_M * BLOCK_N;

-    bool is_brgemm_used = false;
-
-    for (int64_t i = begin; i < end; ++i) {
-      int64_t mb = i / NB2;
-      int64_t nb = i % NB2;
-
+    loop_2d<scalar_t>(mb0, mb1, nb0, nb1, BLOCK_N * IC, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
      int64_t m_size = offsets[mb + 1] - offsets[mb];
      int64_t n_size = std::min(OC - nb * BLOCK_N, BLOCK_N);

-      const bool use_brgemm = can_use_brgemm<scalar_t>(m_size);
-      is_brgemm_used = is_brgemm_used || use_brgemm;
-
      // A ptr from ic1 of [M * topk, N] in sorted order
      // so as to avoid copy A to tmp buffer again
      const scalar_t* __restrict__ A = ic1 + offsets[mb] * N;
@@ -714,9 +625,9 @@ void fused_experts_kernel_impl(
        float weight = topk_weights[index];
        copy_mul_stub(ic2 + index * K + nb * BLOCK_N, C + m * BLOCK_N, weight, n_size);
      }
-    }
+    });

-    if (is_brgemm_used) {
+    if (use_brgemm) {
      at::native::cpublas::brgemm_release();
    }
  });
@@ -743,7 +654,6 @@ void shared_expert_kernel_impl(
    int64_t M,
    int64_t N,
    int64_t K) {
-
  // handle 2 tiles per block
  constexpr int64_t BLOCK_M = block_size_m();
  constexpr int64_t BLOCK_N = block_size_n();
@@ -755,36 +665,29 @@ void shared_expert_kernel_impl(
  TORCH_CHECK(N % BLOCK_N == 0, "Fixme when N is not multiples of ", BLOCK_N);
  const int64_t stride_n = K;

+  const bool use_brgemm = can_use_brgemm<scalar_t>(M);
+
+  const bool apply_scaling_factor = fused_experts_out != nullptr;
+
  // here we only parallel on half of 2N to fuse silu_and_mul with gemm
-  at::parallel_for(0, MB * NB, 0, [&](int64_t begin, int64_t end) {
+  parallel_2d(MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
    // get local pointers
-    int tid = at::get_thread_num();
+    int tid = get_thread_num();
    float* __restrict__ C0 = C_tmp + tid * 2 * BLOCK_M * BLOCK_N;
    float* __restrict__ C1 = C0 + BLOCK_M * BLOCK_N;

-    bool is_brgemm_used = false;
-
-    for (int64_t i = begin; i < end; ++i) {
-      int64_t mb = i / NB;
-      int64_t nb = i % NB;
-
-      // nb0 from top half and nb1 from bottom half
-      int64_t nb0 = nb, nb1 = nb + NB;
-      int64_t n_size = std::min(N - nb0 * BLOCK_N, BLOCK_N);
+    loop_2d<scalar_t>(mb0, mb1, nb0, nb1, BLOCK_N * K * 2, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
+      // nb_upper from top half and nb_lower from bottom half
+      int64_t nb_upper = nb, nb_lower = nb + NB;
+      int64_t n_size = std::min(N - nb * BLOCK_N, BLOCK_N);
      int64_t m_size = std::min(M - mb * BLOCK_M, BLOCK_M);

-      //int64_t mb_start = mb * BLOCK_M;
-      //int64_t mb_size = std::min(M - mb_start, BLOCK_M);
-
      // A shape [m_size, K]
      const scalar_t* A = input + mb * BLOCK_M * K;

      // B shape [K, n_size] in vnni format
-      const scalar_t* __restrict__ B0 = packed_w1 + nb0 * BLOCK_N * stride_n;
-      const scalar_t* __restrict__ B1 = packed_w1 + nb1 * BLOCK_N * stride_n;
-
-      const bool use_brgemm = can_use_brgemm<scalar_t>(m_size);
-      is_brgemm_used = is_brgemm_used || use_brgemm;
+      const scalar_t* __restrict__ B0 = packed_w1 + nb_upper * BLOCK_N * stride_n;
+      const scalar_t* __restrict__ B1 = packed_w1 + nb_lower * BLOCK_N * stride_n;

      if (use_brgemm) {
        // 1.b gemm: C0 = A @ B0
@@ -814,12 +717,7 @@ void shared_expert_kernel_impl(
            /* C     */ C1);

        // 1.d silu and mul
-        silu_and_mul<scalar_t, BLOCK_N>(
-            ic1 + mb * BLOCK_M * N + nb * BLOCK_N,
-            C0,
-            C1,
-            m_size,
-            N);
+        silu_and_mul<scalar_t, BLOCK_N>(ic1 + mb * BLOCK_M * N + nb * BLOCK_N, C0, C1, m_size, N);
      } else {
        // fused 1.bcd: silu_and_mul(A @ B0, A @ B1)
        tinygemm_kernel(
@@ -834,9 +732,9 @@ void shared_expert_kernel_impl(
            /* ldb   */ n_size,
            /* ldc   */ N);
      }
-    }
+    });

-    if (is_brgemm_used) {
+    if (use_brgemm) {
      at::native::cpublas::brgemm_release();
    }
  });
@@ -850,24 +748,16 @@ void shared_expert_kernel_impl(
  const int64_t stride_oc = IC;

  // parallel on [MB2, NB2]
-  at::parallel_for(0, MB2 * NB2, 0, [&](int64_t begin, int64_t end) {
+  parallel_2d(MB2, NB2, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
    // get local pointers
-    int tid = at::get_thread_num();
+    int tid = get_thread_num();
    // we won't be using C1 for gemm2
    float* __restrict__ C = C_tmp + tid * 2 * BLOCK_M * BLOCK_N;

-    bool is_brgemm_used = false;
-
-    for (int64_t i = begin; i < end; ++i) {
-      int64_t mb = i / NB2;
-      int64_t nb = i % NB2;
-
+    loop_2d<scalar_t>(mb0, mb1, nb0, nb1, BLOCK_N * IC, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
      int64_t m_size = std::min(M - mb * BLOCK_M, BLOCK_M);
      int64_t n_size = std::min(OC - nb * BLOCK_N, BLOCK_N);

-      const bool use_brgemm = can_use_brgemm<scalar_t>(m_size);
-      is_brgemm_used = is_brgemm_used || use_brgemm;
-
      // A shape [m_size, IC]
      const scalar_t* __restrict__ A = ic1 + mb * BLOCK_M * N;

@@ -902,19 +792,21 @@ void shared_expert_kernel_impl(

      // 2.b copy from C to output and add fused_experts_out
      scalar_t* __restrict__ out = output + mb * BLOCK_M * K + nb * BLOCK_N;
-      const scalar_t* __restrict__ fused_out = fused_experts_out + mb * BLOCK_M * K + nb * BLOCK_N;
+      const scalar_t* __restrict__ fused_out =
+          apply_scaling_factor ? fused_experts_out + mb * BLOCK_M * K + nb * BLOCK_N : nullptr;
      for (int64_t m = 0; m < m_size; ++m) {
-        add_mul_stub(out + m * K, C + m * BLOCK_N, fused_out + m * K, routed_scaling_factor, n_size);
+        const scalar_t* __restrict__ fused_out_row = apply_scaling_factor ? (fused_out + m * K) : nullptr;
+        add_mul_stub(out + m * K, C + m * BLOCK_N, fused_out_row, routed_scaling_factor, n_size);
      }
-    }
+    });

-    if (is_brgemm_used) {
+    if (use_brgemm) {
      at::native::cpublas::brgemm_release();
    }
  });
 }

-} // anonymous namespace
+}  // anonymous namespace

 // common checks
 static inline void check_moe_scales(
@@ -922,14 +814,10 @@ static inline void check_moe_scales(
    bool use_fp8_w8a16,
    const std::optional<at::Tensor>& w1_scale,
    const std::optional<at::Tensor>& w2_scale,
-    const std::optional<std::vector<int64_t>> block_size,
-    const std::optional<at::Tensor>& a1_scale,
-    const std::optional<at::Tensor>& a2_scale) {
+    const std::optional<std::vector<int64_t>> block_size) {
  if (use_int8_w8a8) {
    TORCH_CHECK(w1_scale.has_value(), "missing w1_scale for int8 w8a8.");
    TORCH_CHECK(w2_scale.has_value(), "missing w2_scale for int8 w8a8.");
-    TORCH_CHECK(!a1_scale.has_value(), "static quantization for activation not supported.");
-    TORCH_CHECK(!a2_scale.has_value(), "static quantization for activation not supported.");
  }
  if (use_fp8_w8a16) {
    TORCH_CHECK(w1_scale.has_value(), "missing w1_scale for fp8 w8a16.");
@@ -939,16 +827,16 @@ static inline void check_moe_scales(
  }
 }

-#define CHECK_MOE_SCALES_FP8(DIM0, DIM1)                 \
-    auto w1s = w1_scale.value();                         \
-    auto w2s = w2_scale.value();                         \
-    auto block_size_val = block_size.value();            \
-    int64_t block_size_N = block_size_val[0];            \
-    int64_t block_size_K = block_size_val[1];            \
-    TORCH_CHECK(w1s.size(DIM0) == 2 * N / block_size_N); \
-    TORCH_CHECK(w1s.size(DIM1) == K / block_size_K);     \
-    TORCH_CHECK(w2s.size(DIM0) == K / block_size_N);     \
-    TORCH_CHECK(w2s.size(DIM1) == N / block_size_K)
+#define CHECK_MOE_SCALES_FP8(DIM0, DIM1)                      \
+  auto w1s = w1_scale.value();                                \
+  auto w2s = w2_scale.value();                                \
+  auto block_size_val = block_size.value();                   \
+  int64_t block_size_N = block_size_val[0];                   \
+  int64_t block_size_K = block_size_val[1];                   \
+  TORCH_CHECK(w1s.size(DIM0) == div_up(2 * N, block_size_N)); \
+  TORCH_CHECK(w1s.size(DIM1) == div_up(K, block_size_K));     \
+  TORCH_CHECK(w2s.size(DIM0) == div_up(K, block_size_N));     \
+  TORCH_CHECK(w2s.size(DIM1) == div_up(N, block_size_K))

 // hidden_states: [M, K]
 // w1: [E, 2N, K]
@@ -956,6 +844,7 @@ static inline void check_moe_scales(
 // topk_weights: [M, topk]
 // topk_ids: [M, topk] (int32_t)
 //
+
 at::Tensor fused_experts_cpu(
    at::Tensor& hidden_states,
    at::Tensor& w1,
@@ -963,16 +852,13 @@ at::Tensor fused_experts_cpu(
    at::Tensor& topk_weights,
    at::Tensor& topk_ids,
    bool inplace,
-    bool use_int8_w8a8,
-    bool use_fp8_w8a16,
+    int64_t moe_comp_method,
    const std::optional<at::Tensor>& w1_scale,
    const std::optional<at::Tensor>& w2_scale,
+    const std::optional<at::Tensor>& w1_zero,
+    const std::optional<at::Tensor>& w2_zero,
    const std::optional<std::vector<int64_t>> block_size,
-    const std::optional<at::Tensor>& a1_scale,
-    const std::optional<at::Tensor>& a2_scale,
    bool is_vnni) {
-  RECORD_FUNCTION("sgl-kernel::fused_experts_cpu", std::vector<c10::IValue>({hidden_states, w1, w2, topk_weights, topk_ids}));
-
  auto packed_w1 = is_vnni ? w1 : convert_weight_packed(w1);
  auto packed_w2 = is_vnni ? w2 : convert_weight_packed(w2);

@@ -985,32 +871,49 @@ at::Tensor fused_experts_cpu(
  CHECK_INPUT(w2);
  CHECK_EQ(topk_weights.sizes(), topk_ids.sizes());
  CHECK_DIM(2, hidden_states);
-  CHECK_DIM(3, w1);
-  CHECK_DIM(3, w2);
+  if (moe_comp_method == CPUQuantMethod::INT4_W4A8 && is_vnni) {
+    CHECK_DIM(4, w1);
+    CHECK_DIM(4, w2);
+  } else {
+    CHECK_DIM(3, w1);
+    CHECK_DIM(3, w2);
+  }
  CHECK_DIM(2, topk_weights);
  CHECK_DIM(2, topk_ids);

  CHECK_EQ(topk_ids.scalar_type(), at::kInt);
-  CHECK_EQ(topk_weights.scalar_type(), at::kFloat);
+
+  // TODO: support topk_weights to be bf16 or fp16 in the kernel.
+  // The topk_weights of llama4 is computed via Llama4MoE:custom_routing_function and is bf16/fp16
+  // while the kernel currently only supports it to be float32
+  auto topk_weights_ = topk_weights.to(at::kFloat);
+  CHECK_EQ(topk_weights_.scalar_type(), at::kFloat);

  int64_t M = hidden_states.size(0);
  int64_t K = hidden_states.size(1);
-  int64_t N = w1.size(1) / 2;
+  int64_t N = moe_comp_method == CPUQuantMethod::INT4_W4A8 ? w1_scale.value().size(1) * w1_scale.value().size(3) / 2
+                                                           : w1.size(1) / 2;
  int64_t E = w1.size(0);
-  int64_t topk = topk_weights.size(1);
+  int64_t topk = topk_weights_.size(1);

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

  // check weight shapes
  CHECK_EQ(w2.size(0), E);
-  CHECK_EQ(w2.size(1), K);
-  CHECK_EQ(packed_w1.size(2), packed_K);
-  CHECK_EQ(packed_w2.size(2), packed_N);
-
+  if (!(moe_comp_method == CPUQuantMethod::INT4_W4A8)) {
+    CHECK_EQ(w2.size(1), K);
+    CHECK_EQ(packed_w1.size(2), packed_K / (moe_comp_method == CPUQuantMethod::INT4_W4A8 ? 2 : 1));
+    CHECK_EQ(packed_w2.size(2), packed_N / (moe_comp_method == CPUQuantMethod::INT4_W4A8 ? 2 : 1));
+  }
  // check scales
-  check_moe_scales(use_int8_w8a8, use_fp8_w8a16, w1_scale, w2_scale, block_size, a1_scale, a2_scale);
+  check_moe_scales(
+      moe_comp_method == CPUQuantMethod::INT8_W8A8,
+      moe_comp_method == CPUQuantMethod::FP8_W8A16,
+      w1_scale,
+      w2_scale,
+      block_size);

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

@@ -1031,8 +934,8 @@ at::Tensor fused_experts_cpu(
  int32_t* __restrict__ sorted_ids = buffer.data_ptr<int32_t>();
  int32_t* __restrict__ expert_ids = sorted_ids + max_num_tokens_padded;
  int32_t* __restrict__ total_cnts = expert_ids + max_num_blocks;
-  int32_t* __restrict__ cumsums    = total_cnts + (num_threads + 1) * E;
-  int32_t* __restrict__ offsets    = cumsums    + (E + 1);
+  int32_t* __restrict__ cumsums = total_cnts + (num_threads + 1) * E;
+  int32_t* __restrict__ offsets = cumsums + (E + 1);

  // init sorted_ids with `numel` as the padding number
  // init expert_ids with `num_experts`
@@ -1064,26 +967,31 @@ at::Tensor fused_experts_cpu(
  //
  // for fp8 w8a16:
  //   7. intermediate_cache0 : [M * topk, 2N]
-  //   8. B_tmp : [T, BLOCK_N, std::max(K, N)]
+  //   8. B_tmp : [T, MAX_CACHE_BLOCK_SIZE, BLOCK_N, std::max(K, N)]
  //
-  int64_t buffer_size_nbytes = M * topk * N * 2 + M * topk * K * 2 +
-      num_threads * BLOCK_M * K * (use_int8_w8a8 ? 1 : 2) +
+  int64_t buffer_size_nbytes =
+      M * topk * N * 2 + M * topk * K * 2 +
+      num_threads * BLOCK_M * K *
+          (moe_comp_method == CPUQuantMethod::INT8_W8A8 | moe_comp_method == CPUQuantMethod::INT4_W4A8 ? 1 : 2) +
      num_threads * 2 * BLOCK_M * BLOCK_N * sizeof(float);

-  if (use_int8_w8a8) {
+  if (moe_comp_method == CPUQuantMethod::INT8_W8A8) {
    buffer_size_nbytes += std::max(M * K, M * topk * N) + M * topk * sizeof(float);
  }
-  if (use_fp8_w8a16) {
-    buffer_size_nbytes += M * topk * 2 * N * 2 + num_threads * BLOCK_N * std::max(K, N) * 2;
+  if (moe_comp_method == CPUQuantMethod::FP8_W8A16) {
+    buffer_size_nbytes += M * topk * 2 * N * 2 + num_threads * MAX_CACHE_BLOCK_SIZE * BLOCK_N * std::max(K, N) * 2;
+  }
+  if (moe_comp_method == CPUQuantMethod::INT4_W4A8) {
+    buffer_size_nbytes += M * topk * 2 * N * 2 + std::max(M * K, M * topk * N) + M * topk * sizeof(float) +
+                          num_threads * 2 * get_4bit_block_k_size(K / w1_scale.value().size(2)) * BLOCK_N;
  }
-
  auto buffer2 = at::empty({buffer_size_nbytes}, hidden_states.options().dtype(at::kChar));

  AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "fused_experts_kernel_impl", [&] {
    scalar_t* __restrict__ intermediate_cache1 = (scalar_t*)((void*)(buffer2.data_ptr<int8_t>()));
    scalar_t* __restrict__ intermediate_cache2 = intermediate_cache1 + M * topk * N;

-    if (use_int8_w8a8) {
+    if (moe_comp_method == CPUQuantMethod::INT8_W8A8) {
      uint8_t* __restrict__ A_tmp = (uint8_t*)((void*)(intermediate_cache2 + M * topk * K));
      float* __restrict__ C_tmp = (float*)((void*)(A_tmp + num_threads * BLOCK_M * K));
      uint8_t* __restrict__ Aq_tmp = (uint8_t*)((void*)(C_tmp + num_threads * 2 * BLOCK_M * BLOCK_N));
@@ -1107,7 +1015,7 @@ at::Tensor fused_experts_cpu(
          packed_w2.data_ptr<int8_t>(),
          w1s.data_ptr<float>(),
          w2s.data_ptr<float>(),
-          topk_weights.data_ptr<float>(),
+          topk_weights_.data_ptr<float>(),
          sorted_ids,
          expert_ids,
          offsets,
@@ -1117,7 +1025,7 @@ at::Tensor fused_experts_cpu(
          E,
          topk,
          num_tokens_post_pad);
-    } else if (use_fp8_w8a16) {
+    } else if (moe_comp_method == CPUQuantMethod::FP8_W8A16) {
      // here we just ignore C_tmp as it is not used
      scalar_t* __restrict__ A_tmp = (scalar_t*)((void*)(intermediate_cache2 + M * topk * K));
      float* __restrict__ C_tmp = (float*)((void*)(A_tmp + num_threads * BLOCK_M * K));
@@ -1140,6 +1048,48 @@ at::Tensor fused_experts_cpu(
          w2s.data_ptr<float>(),
          block_size_N,
          block_size_K,
+          topk_weights_.data_ptr<float>(),
+          sorted_ids,
+          expert_ids,
+          offsets,
+          M,
+          N,
+          K,
+          E,
+          topk,
+          num_tokens_post_pad);
+    } else if (moe_comp_method == CPUQuantMethod::INT4_W4A8) {
+      uint8_t* __restrict__ A_tmp = (uint8_t*)((void*)(intermediate_cache2 + M * topk * K));
+      float* __restrict__ C_tmp = (float*)((void*)(A_tmp + num_threads * BLOCK_M * K));
+      scalar_t* __restrict__ intermediate_cache0 = (scalar_t*)((void*)(C_tmp + num_threads * 2 * BLOCK_M * BLOCK_N));
+      uint8_t* __restrict__ Aq_tmp = (uint8_t*)((void*)(intermediate_cache0 + M * topk * 2 * N));
+      float* __restrict__ As_tmp = (float*)((void*)(Aq_tmp + std::max(M * K, M * topk * N)));
+      int8_t* __restrict__ dqB_tmp = (int8_t*)((void*)(As_tmp + M * topk));
+
+      // weight + compensation shape = [Nc, Kc, block_n * block_k / 2 + block_n*sizeof(int32_t)]
+      // scales/qzeros shape = [E, Nc, G, block_n]
+      int64_t num_groups = w1_scale.value().size(2);
+      const int group_size = K / num_groups;
+      // TODO: check scales and zeros
+      fused_experts_int4_w4a8_kernel_impl<scalar_t>(
+          out_hidden_states.data_ptr<scalar_t>(),
+          intermediate_cache0,
+          intermediate_cache1,
+          intermediate_cache2,
+          A_tmp,
+          Aq_tmp,
+          As_tmp,
+          nullptr,
+          C_tmp,
+          dqB_tmp,
+          hidden_states.data_ptr<scalar_t>(),
+          packed_w1.data_ptr<uint8_t>(),
+          packed_w2.data_ptr<uint8_t>(),
+          w1_zero.value().data_ptr<int8_t>(),
+          w2_zero.value().data_ptr<int8_t>(),
+          w1_scale.value().data_ptr<float>(),
+          w2_scale.value().data_ptr<float>(),
+          group_size,
          topk_weights.data_ptr<float>(),
          sorted_ids,
          expert_ids,
@@ -1163,7 +1113,7 @@ at::Tensor fused_experts_cpu(
          hidden_states.data_ptr<scalar_t>(),
          packed_w1.data_ptr<scalar_t>(),
          packed_w2.data_ptr<scalar_t>(),
-          topk_weights.data_ptr<float>(),
+          topk_weights_.data_ptr<float>(),
          sorted_ids,
          expert_ids,
          offsets,
@@ -1188,34 +1138,37 @@ at::Tensor shared_expert_cpu(
    at::Tensor& hidden_states,
    at::Tensor& w1,
    at::Tensor& w2,
-    at::Tensor& fused_experts_out,
-    double routed_scaling_factor,
+    const std::optional<at::Tensor>& fused_experts_out,
+    const std::optional<double> routed_scaling_factor,
    bool inplace,
    bool use_int8_w8a8,
    bool use_fp8_w8a16,
-    std::optional<at::Tensor>& w1_scale,
-    std::optional<at::Tensor>& w2_scale,
-    std::optional<std::vector<int64_t>> block_size,
-    std::optional<at::Tensor>& a1_scale,
-    std::optional<at::Tensor>& a2_scale,
+    const std::optional<at::Tensor>& w1_scale,
+    const std::optional<at::Tensor>& w2_scale,
+    const std::optional<std::vector<int64_t>> block_size,
    bool is_vnni) {
-  RECORD_FUNCTION("sgl-kernel::shared_expert_cpu", std::vector<c10::IValue>({hidden_states, w1, w2}));
-
  auto packed_w1 = is_vnni ? w1 : convert_weight_packed(w1);
  auto packed_w2 = is_vnni ? w2 : convert_weight_packed(w2);

  constexpr int64_t BLOCK_M = block_size_m();
  constexpr int64_t BLOCK_N = block_size_n();

+  double routed_scaling_factor_value = 0;
+  if (routed_scaling_factor.has_value()) {
+    TORCH_CHECK(fused_experts_out.has_value(), "shared_expert_cpu: expect fused_experts_out.");
+    const auto fused_experts_out_tensor = fused_experts_out.value();
+    routed_scaling_factor_value = routed_scaling_factor.value();
+    CHECK_INPUT(fused_experts_out_tensor);
+    CHECK_EQ(hidden_states.sizes(), fused_experts_out_tensor.sizes());
+  }
+
  const auto st = hidden_states.scalar_type();
  CHECK_INPUT(hidden_states);
-  CHECK_INPUT(fused_experts_out);
  CHECK_INPUT(w1);
  CHECK_INPUT(w2);
  CHECK_DIM(2, hidden_states);
  CHECK_DIM(2, w1);
  CHECK_DIM(2, w2);
-  CHECK_EQ(hidden_states.sizes(), fused_experts_out.sizes());
  CHECK_EQ(hidden_states.scalar_type(), st);

  int64_t M = hidden_states.size(0);
@@ -1232,7 +1185,7 @@ at::Tensor shared_expert_cpu(
  CHECK_EQ(packed_w2.size(1), packed_N);

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

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

@@ -1246,7 +1199,7 @@ at::Tensor shared_expert_cpu(
  //
  // for fp8 w8a16:
  //   5. intermediate_cache0 : [M, 2N]
-  //   6. B_tmp: [T, BLOCK_M, max(K, N)]
+  //   6. B_tmp: [T, MAX_CACHE_BLOCK_SIZE, BLOCK_M, max(K, N)]
  //
  int num_threads = at::get_num_threads();
  int64_t buffer_size_nbytes = M * N * 2 + num_threads * 2 * BLOCK_M * BLOCK_N * sizeof(float);
@@ -1255,7 +1208,7 @@ at::Tensor shared_expert_cpu(
    buffer_size_nbytes += std::max(M * K, M * N) + M * sizeof(float);
  }
  if (use_fp8_w8a16) {
-    buffer_size_nbytes += M * 2 * N * 2 + num_threads * BLOCK_M * std::max(K, N) * 2;
+    buffer_size_nbytes += M * 2 * N * 2 + num_threads * MAX_CACHE_BLOCK_SIZE * BLOCK_M * std::max(K, N) * 2;
  }

  auto buffer = at::empty({buffer_size_nbytes}, hidden_states.options().dtype(at::kChar));
@@ -1283,8 +1236,8 @@ at::Tensor shared_expert_cpu(
          packed_w2.data_ptr<int8_t>(),
          w1s.data_ptr<float>(),
          w2s.data_ptr<float>(),
-          fused_experts_out.data_ptr<scalar_t>(),
-          routed_scaling_factor,
+          conditional_data_ptr<scalar_t>(fused_experts_out),
+          routed_scaling_factor_value,
          M,
          N,
          K);
@@ -1306,8 +1259,8 @@ at::Tensor shared_expert_cpu(
          w2s.data_ptr<float>(),
          block_size_N,
          block_size_K,
-          fused_experts_out.data_ptr<scalar_t>(),
-          routed_scaling_factor,
+          conditional_data_ptr<scalar_t>(fused_experts_out),
+          routed_scaling_factor_value,
          M,
          N,
          K);
@@ -1319,8 +1272,8 @@ at::Tensor shared_expert_cpu(
          hidden_states.data_ptr<scalar_t>(),
          packed_w1.data_ptr<scalar_t>(),
          packed_w2.data_ptr<scalar_t>(),
-          fused_experts_out.data_ptr<scalar_t>(),
-          routed_scaling_factor,
+          conditional_data_ptr<scalar_t>(fused_experts_out),
+          routed_scaling_factor_value,
          M,
          N,
          K);
@@ -0,0 +1,178 @@
+// Adapted from
+// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu
+
+// clang-format off
+
+#pragma once
+#include "vec.h"
+
+template <typename scalar_t>
+inline void fill_stub(scalar_t* __restrict__ out, scalar_t val, int64_t size) {
+  using Vec = at::vec::Vectorized<scalar_t>;
+  const Vec data_vec(val);
+  at::vec::map<scalar_t>([data_vec](Vec out) { return out = data_vec; }, out, out, size);
+}
+
+template <typename scalar_t>
+inline void copy_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, int64_t size) {
+  using Vec = at::vec::Vectorized<scalar_t>;
+  constexpr int kVecSize = Vec::size();
+  int64_t d;
+#pragma GCC unroll 4
+  for (d = 0; d <= size - kVecSize; d += kVecSize) {
+    Vec data = Vec::loadu(input + d);
+    data.store(out + d);
+  }
+  for (; d < size; ++d) {
+    out[d] = input[d];
+  }
+}
+
+template <typename scalar_t>
+inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ input, int64_t size) {
+  using bVec = at::vec::Vectorized<scalar_t>;
+  using fVec = at::vec::Vectorized<float>;
+  constexpr int kVecSize = bVec::size();
+  int64_t d;
+#pragma GCC unroll 4
+  for (d = 0; d <= size - kVecSize; d += kVecSize) {
+    auto [x0, x1] = load_float_vec2(input + d);
+    bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1);
+    out_vec.store(out + d);
+  }
+  for (; d < size; ++d) {
+    out[d] = static_cast<scalar_t>(input[d]);
+  }
+}
+
+template <>
+inline void copy_stub<uint8_t>(uint8_t* __restrict__ out, const uint8_t* __restrict__ input, int64_t size) {
+  // size might be 64x + 32
+  std::memcpy(out, input, size * sizeof(uint8_t));
+}
+
+template <typename scalar_t, typename input_t>
+inline void copy_mul_stub(scalar_t* __restrict__ out, const input_t* __restrict__ input, float weight, int64_t size) {
+  static_assert(
+      std::is_same_v<input_t, float> || std::is_same_v<input_t, scalar_t>,
+      "copy_mul_stub only supports input_t == float or input_t == scalar_t");
+  using bVec = at::vec::Vectorized<scalar_t>;
+  using fVec = at::vec::Vectorized<float>;
+  constexpr int kVecSize = bVec::size();
+  const fVec weight_vec = fVec(weight);
+  int64_t d;
+#pragma GCC unroll 4
+  for (d = 0; d <= size - kVecSize; d += kVecSize) {
+    auto [x0, x1] = load_float_vec2(input + d);
+    x0 = x0 * weight_vec;
+    x1 = x1 * weight_vec;
+    bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1);
+    out_vec.store(out + d);
+  }
+  for (; d < size; ++d) {
+    out[d] = static_cast<scalar_t>(input[d] * weight);
+  }
+}
+
+// acc from [topk, K] to [K]
+template <typename scalar_t>
+inline void sum_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, int64_t topk, int64_t K) {
+  using bVec = at::vec::Vectorized<scalar_t>;
+  using fVec = at::vec::Vectorized<float>;
+  constexpr int kVecSize = bVec::size();
+  if (topk == 1) {
+    // do copy for topk = 1
+    copy_stub(out, input, K);
+  } else {
+    // do sum for topk != 1
+    int64_t d;
+#pragma GCC unroll 4
+    for (d = 0; d <= K - kVecSize; d += kVecSize) {
+      fVec sum_fvec0 = fVec(0.f);
+      fVec sum_fvec1 = fVec(0.f);
+      for (int t = 0; t < topk; ++t) {
+        bVec x_bvec = bVec::loadu(input + t * K + d);
+        fVec x_fvec0, x_fvec1;
+        std::tie(x_fvec0, x_fvec1) = at::vec::convert_to_float(x_bvec);
+
+        sum_fvec0 += x_fvec0;
+        sum_fvec1 += x_fvec1;
+      }
+      bVec out_bvec = convert_from_float_ext<scalar_t>(sum_fvec0, sum_fvec1);
+      out_bvec.store(out + d);
+    }
+    for (; d < K; ++d) {
+      float sum_val = 0.f;
+      for (int t = 0; t < topk; ++t) {
+        sum_val += static_cast<float>(input[t * K + d]);
+      }
+      out[d] = static_cast<scalar_t>(sum_val);
+    }
+  }
+}
+
+// out = input + input2 * scale
+template <typename scalar_t, typename input_t>
+inline void add_mul_stub(
+    scalar_t* __restrict__ out,
+    const input_t* __restrict__ input,
+    const scalar_t* __restrict__ input2,
+    float scale,
+    int64_t size) {
+  static_assert(
+      std::is_same_v<input_t, float> || std::is_same_v<input_t, scalar_t>,
+      "add_mul_stub only supports input_t == float or input_t == scalar_t");
+
+  // out = input (without scale factor)
+  if (input2 == nullptr) {
+    copy_stub(out, input, size);
+    return;
+  }
+
+  using bVec = at::vec::Vectorized<scalar_t>;
+  using fVec = at::vec::Vectorized<float>;
+  constexpr int kVecSize = bVec::size();
+  const fVec s_vec = fVec(scale);
+  int64_t d;
+#pragma GCC unroll 4
+  for (d = 0; d <= size - kVecSize; d += kVecSize) {
+    auto [x0, x1] = load_float_vec2(input + d);
+
+    bVec y_bvec = bVec::loadu(input2 + d);
+    fVec y0, y1;
+    std::tie(y0, y1) = at::vec::convert_to_float(y_bvec);
+
+    x0 = x0 + y0 * s_vec;
+    x1 = x1 + y1 * s_vec;
+    bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1);
+    out_vec.store(out + d);
+  }
+  for (; d < size; ++d) {
+    out[d] = static_cast<scalar_t>(input[d] + float(input2[d]) * scale);
+  }
+}
+
+template <typename scalar_t>
+inline void silu_and_mul_stub(
+    scalar_t* __restrict__ out, const scalar_t* __restrict__ input, const scalar_t* __restrict__ input2, int64_t size) {
+  using bVec = at::vec::Vectorized<scalar_t>;
+  using fVec = at::vec::Vectorized<float>;
+  const fVec one = fVec(1.f);
+
+  // no remainder
+#pragma GCC unroll 4
+  for (int64_t d = 0; d < size; d += bVec::size()) {
+    bVec x = bVec::loadu(input + d);
+    fVec x0, x1;
+    std::tie(x0, x1) = at::vec::convert_to_float(x);
+    bVec y = bVec::loadu(input2 + d);
+    fVec y0, y1;
+    std::tie(y0, y1) = at::vec::convert_to_float(y);
+    x0 = x0 / (one + x0.neg().exp_u20());
+    x1 = x1 / (one + x1.neg().exp_u20());
+    x0 = x0 * y0;
+    x1 = x1 * y1;
+    bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1);
+    out_vec.store(out + d);
+  }
+}
@@ -1,147 +1,11 @@
 // Adapted from
 // https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu

-#include "common.h"
-#include "gemm.h"
-#include "vec.h"
-
 // clang-format off

-namespace {
-
-template <typename scalar_t>
-inline void copy_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, int64_t size) {
-  using Vec = at::vec::Vectorized<scalar_t>;
-  // no remainder
-  #pragma GCC unroll 4
-  for (int64_t d = 0; d < size; d += Vec::size()) {
-    Vec data = Vec::loadu(input + d);
-    data.store(out + d);
-  }
-}
-
-template <typename scalar_t>
-inline void copy_mul_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, float weight, int64_t size) {
-  using bVec = at::vec::Vectorized<scalar_t>;
-  using fVec = at::vec::Vectorized<float>;
-  constexpr int kVecSize = bVec::size();
-  const fVec weight_vec = fVec(weight);
-  int64_t d;
-  #pragma GCC unroll 4
-  for (d = 0; d <= size - kVecSize; d += kVecSize) {
-    bVec x = bVec::loadu(input + d);
-    fVec x0, x1;
-    std::tie(x0, x1) = at::vec::convert_to_float(x);
-    x0 = x0 * weight_vec;
-    x1 = x1 * weight_vec;
-    bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1);
-    out_vec.store(out + d);
-  }
-  for (; d < size; ++d) {
-    out[d] = static_cast<scalar_t>(input[d] * weight);
-  }
-}
-
-// acc from [topk, K] to [K]
-template <typename scalar_t>
-inline void sum_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, int64_t topk, int64_t K) {
-  using bVec = at::vec::Vectorized<scalar_t>;
-  using fVec = at::vec::Vectorized<float>;
-  constexpr int kVecSize = bVec::size();
-  if (topk == 1) {
-    // do copy for topk = 1
-    copy_stub(out, input, K);
-  } else {
-    // do sum for topk != 1
-    int64_t d;
-    #pragma GCC unroll 4
-    for (d = 0; d <= K - kVecSize; d += kVecSize) {
-      fVec sum_fvec0 = fVec(0.f);
-      fVec sum_fvec1 = fVec(0.f);
-      for (int t = 0; t < topk; ++t) {
-        bVec x_bvec = bVec::loadu(input + t * K + d);
-        fVec x_fvec0, x_fvec1;
-        std::tie(x_fvec0, x_fvec1) = at::vec::convert_to_float(x_bvec);
-
-        sum_fvec0 += x_fvec0;
-        sum_fvec1 += x_fvec1;
-      }
-      bVec out_bvec = convert_from_float_ext<scalar_t>(sum_fvec0, sum_fvec1);
-      out_bvec.store(out + d);
-    }
-    for (; d < K; ++d) {
-      float sum_val = 0.f;
-      for (int t = 0; t < topk; ++t) {
-        sum_val += static_cast<float>(input[t * K + d]);
-      }
-      out[d] = static_cast<scalar_t>(sum_val);
-    }
-  }
-}
-
-// out = input + input2 * scale
-template <typename scalar_t>
-inline void add_mul_stub(
-    scalar_t* __restrict__ out,
-    const scalar_t* __restrict__ input,
-    const scalar_t* __restrict__ input2,
-    float scale,
-    int64_t size) {
-  using bVec = at::vec::Vectorized<scalar_t>;
-  using fVec = at::vec::Vectorized<float>;
-  constexpr int kVecSize = bVec::size();
-  const fVec s_vec = fVec(scale);
-
-  int64_t d;
-#pragma GCC unroll 4
-  for (d = 0; d <= size - kVecSize; d += kVecSize) {
-    bVec x_bvec = bVec::loadu(input + d);
-    fVec x0, x1;
-    std::tie(x0, x1) = at::vec::convert_to_float(x_bvec);
-
-    bVec y_bvec = bVec::loadu(input2 + d);
-    fVec y0, y1;
-    std::tie(y0, y1) = at::vec::convert_to_float(y_bvec);
-
-    x0 = x0 + y0 * s_vec;
-    x1 = x1 + y1 * s_vec;
-    bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1);
-    out_vec.store(out + d);
-  }
-  for (; d < size; ++d) {
-    out[d] = static_cast<scalar_t>(input[d] + float(input2[d]) * scale);
-  }
-}
-
-template <typename scalar_t>
-inline void silu_and_mul_stub(
-    scalar_t* __restrict__ out,
-    const scalar_t* __restrict__ input,
-    const scalar_t* __restrict__ input2,
-    int64_t size) {
-  using bVec = at::vec::Vectorized<scalar_t>;
-  using fVec = at::vec::Vectorized<float>;
-  const fVec one = fVec(1.f);
-
-  // no remainder
-#pragma GCC unroll 4
-  for (int64_t d = 0; d < size; d += bVec::size()) {
-    bVec x = bVec::loadu(input + d);
-    fVec x0, x1;
-    std::tie(x0, x1) = at::vec::convert_to_float(x);
-    bVec y = bVec::loadu(input2 + d);
-    fVec y0, y1;
-    std::tie(y0, y1) = at::vec::convert_to_float(y);
-    x0 = x0 / (one + x0.neg().exp_u20());
-    x1 = x1 / (one + x1.neg().exp_u20());
-    x0 = x0 * y0;
-    x1 = x1 * y1;
-    bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1);
-    out_vec.store(out + d);
-  }
-}
-
-} // anonymous namespace
+#include "common.h"
+#include "gemm.h"
+#include "moe.h"

 template <typename scalar_t>
 void fused_experts_fp8_kernel_impl(
@@ -169,7 +33,6 @@ void fused_experts_fp8_kernel_impl(
    int64_t E,
    int64_t topk,
    int64_t num_tokens_post_pad) {
-
  constexpr int64_t BLOCK_M = block_size_m();
  constexpr int64_t BLOCK_N = block_size_n();

@@ -183,35 +46,39 @@ void fused_experts_fp8_kernel_impl(
  const int64_t stride_e = 2 * N * K;
  const int64_t stride_n = K;

+  int64_t avg_M = std::max(int64_t(1), M * topk / E);
+  const bool use_brgemm = can_use_brgemm<at::Float8_e4m3fn>(avg_M);
+
+  int64_t B_tmp_size_per_thread = MAX_CACHE_BLOCK_SIZE * BLOCK_N * std::max(K, N);
+
  // here we only parallel on half of 2N to fuse silu_and_mul with gemm
-  at::parallel_for(0, MB * NB, 0, [&](int64_t begin, int64_t end) {
+  parallel_2d(MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
    // get local pointers
-    int tid = at::get_thread_num();
+    int tid = get_thread_num();
    scalar_t* __restrict__ A = A_tmp + tid * BLOCK_M * K;

-    bool is_brgemm_used = false;
-
-    for (int64_t i = begin; i < end; ++i) {
-      int64_t mb = i / NB;
-      int64_t nb = i % NB;
-
+    loop_2d<at::Float8_e4m3fn>(mb0, mb1, nb0, nb1, BLOCK_N * K, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
      int64_t n_size = std::min(2 * N - nb * BLOCK_N, BLOCK_N);

      // B shape [K, n_size] in vnni format
      int32_t expert_id = expert_ids[mb];
      const at::Float8_e4m3fn* __restrict__ B = packed_w1 + expert_id * stride_e + nb * BLOCK_N * stride_n;
-      const float* __restrict__ Bs = w1s + expert_id * scale_size_N * scale_size_K + (nb / blocks_n_per_group) * scale_size_K;
+      const float* __restrict__ Bs =
+          w1s + expert_id * scale_size_N * scale_size_K + (nb / blocks_n_per_group) * scale_size_K;
+
+      // do unpacking for the first row or a new expert
+      int32_t pre_expert_id = mb == 0 ? -1 : expert_ids[mb - 1];
+      bool do_unpack = (mb == mb0) || (expert_id != pre_expert_id);

-      // 1.a load A
-      const int32_t* A_ids = sorted_ids + mb * BLOCK_M;
      int64_t m_size = offsets[mb + 1] - offsets[mb];

-      const bool use_brgemm = can_use_brgemm<at::Float8_e4m3fn>(m_size);
-      is_brgemm_used = is_brgemm_used || use_brgemm;
-
-      for (int64_t m = 0; m < m_size; ++m) {
-        int32_t index = A_ids[m] / topk;
-        copy_stub(A + m * K, input + index * K, K);
+      if (nb_offset == 0) {
+        // 1.a load A
+        const int32_t* A_ids = sorted_ids + mb * BLOCK_M;
+        for (int64_t m = 0; m < m_size; ++m) {
+          int32_t index = A_ids[m] / topk;
+          copy_stub(A + m * K, input + index * K, K);
+        }
      }

      const int64_t offset = offsets[mb];
@@ -219,7 +86,7 @@ void fused_experts_fp8_kernel_impl(
          /*   A            */ A,
          /*   B            */ B,
          /*   C            */ ic0 + offset * 2 * N + nb * BLOCK_N,
-          /*   Btmp         */ B_tmp + tid * BLOCK_N * std::max(K, N),
+          /*   Btmp         */ B_tmp + tid * B_tmp_size_per_thread + nb_offset * BLOCK_N * K,
          /*   Ctmp         */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N,
          /*   scale        */ Bs,
          /*   M            */ m_size,
@@ -229,10 +96,11 @@ void fused_experts_fp8_kernel_impl(
          /*   ldb          */ n_size,
          /*   ldc          */ 2 * N,
          /*   brg          */ use_brgemm,
-          /*   block_size_K */ block_size_K);
-    }
+          /*   block_size_K */ block_size_K,
+          /*   do_unpack    */ do_unpack);
+    });

-    if (is_brgemm_used) {
+    if (use_brgemm) {
      at::native::cpublas::brgemm_release();
    }
  });
@@ -240,11 +108,7 @@ void fused_experts_fp8_kernel_impl(
  // stage 1.5: intermediate_cache1 = silu(intermediate_cache0)
  at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
    for (int64_t m = begin; m < end; ++m) {
-      silu_and_mul_stub(
-          ic1 + m * N,
-          ic0 + m * 2 * N,
-          ic0 + m * 2 * N + N,
-          N);
+      silu_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, ic0 + m * 2 * N + N, N);
    }
  });

@@ -260,22 +124,14 @@ void fused_experts_fp8_kernel_impl(
  const int64_t stride_oc = IC;

  // parallel on [MB2, NB2]
-  at::parallel_for(0, MB2 * NB2, 0, [&](int64_t begin, int64_t end) {
-    int tid = at::get_thread_num();
+  parallel_2d(MB2, NB2, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
+    int tid = get_thread_num();
    alignas(64) scalar_t C[BLOCK_M * BLOCK_K];

-    bool is_brgemm_used = false;
-
-    for (int64_t i = begin; i < end; ++i) {
-      int64_t mb = i / NB2;
-      int64_t nb = i % NB2;
-
+    loop_2d<at::Float8_e4m3fn>(mb0, mb1, nb0, nb1, BLOCK_N * IC, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
      int64_t m_size = offsets[mb + 1] - offsets[mb];
      int64_t n_size = std::min(OC - nb * BLOCK_N, BLOCK_N);

-      const bool use_brgemm = can_use_brgemm<at::Float8_e4m3fn>(m_size);
-      is_brgemm_used = is_brgemm_used || use_brgemm;
-
      // A ptr from ic1 of [M * topk, N] in sorted order
      // so as to avoid copy A to tmp buffer again
      const scalar_t* __restrict__ A = ic1 + offsets[mb] * N;
@@ -284,13 +140,18 @@ void fused_experts_fp8_kernel_impl(
      // B shape [IC, n_size] in vnni format
      int32_t expert_id = expert_ids[mb];
      const at::Float8_e4m3fn* __restrict__ B = packed_w2 + expert_id * stride_e2 + nb * BLOCK_N * stride_oc;
-      const float* __restrict__ Bs = w2s + expert_id * scale_size_N * scale_size_K + (nb / blocks_n_per_group) * scale_size_K;
+      const float* __restrict__ Bs =
+          w2s + expert_id * scale_size_N * scale_size_K + (nb / blocks_n_per_group) * scale_size_K;
+
+      // do unpacking for the first row or a new expert
+      int32_t pre_expert_id = mb == 0 ? -1 : expert_ids[mb - 1];
+      bool do_unpack = (mb == mb0) || (expert_id != pre_expert_id);

      tinygemm_kernel<scalar_t>(
          /*   A            */ A,
          /*   B            */ B,
          /*   C            */ C,
-          /*   Btmp         */ B_tmp + tid * BLOCK_N * std::max(K, N),
+          /*   Btmp         */ B_tmp + tid * B_tmp_size_per_thread + nb_offset * BLOCK_N * IC,
          /*   Ctmp         */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N,
          /*   scale        */ Bs,
          /*   M            */ m_size,
@@ -300,7 +161,8 @@ void fused_experts_fp8_kernel_impl(
          /*   ldb          */ n_size,
          /*   ldc          */ BLOCK_N,
          /*   brg          */ use_brgemm,
-          /*   block_size_K */ block_size_K);
+          /*   block_size_K */ block_size_K,
+          /*   do_unpack    */ do_unpack);

      // 2.b copy from C to ic2 in original order
      //   and also mul topk_weights in float32
@@ -309,9 +171,9 @@ void fused_experts_fp8_kernel_impl(
        float weight = topk_weights[index];
        copy_mul_stub(ic2 + index * K + nb * BLOCK_N, C + m * BLOCK_N, weight, n_size);
      }
-    }
+    });

-    if (is_brgemm_used) {
+    if (use_brgemm) {
      at::native::cpublas::brgemm_release();
    }
  });
@@ -374,7 +236,6 @@ void shared_expert_fp8_kernel_impl(
    int64_t M,
    int64_t N,
    int64_t K) {
-
  constexpr int64_t BLOCK_M = block_size_m();
  constexpr int64_t BLOCK_N = block_size_n();

@@ -385,21 +246,25 @@ void shared_expert_fp8_kernel_impl(
  int64_t blocks_n_per_group = block_size_N / BLOCK_N;

  const bool use_brgemm = can_use_brgemm<at::Float8_e4m3fn>(M);
+  const bool apply_scaling_factor = fused_experts_out != nullptr;

-  at::parallel_for(0, MB * NB, 0, [&](int64_t begin, int64_t end) {
-    int tid = at::get_thread_num();
+  int64_t B_tmp_size_per_thread = MAX_CACHE_BLOCK_SIZE * BLOCK_N * std::max(K, N);

-    for (int64_t i = begin; i < end; ++i) {
-      int64_t mb = i / NB;
-      int64_t nb = i % NB;
+  parallel_2d(MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
+    int tid = get_thread_num();
+
+    loop_2d<at::Float8_e4m3fn>(mb0, mb1, nb0, nb1, BLOCK_N * K, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
      int64_t m_size = std::min(M - mb * BLOCK_M, BLOCK_M);
      int64_t n_size = std::min(2 * N - nb * BLOCK_N, BLOCK_N);

+      // do unpacking for the first row
+      bool do_unpack = (mb == mb0);
+
      tinygemm_kernel<scalar_t>(
          /*   A            */ input + mb * BLOCK_M * K,
          /*   B            */ packed_w1 + nb * BLOCK_N * K,
          /*   C            */ ic0 + mb * BLOCK_M * 2 * N + nb * BLOCK_N,
-          /*   Btmp         */ B_tmp + tid * BLOCK_N * std::max(K, N),
+          /*   Btmp         */ B_tmp + tid * B_tmp_size_per_thread + nb_offset * BLOCK_N * K,
          /*   Ctmp         */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N,
          /*   scale        */ w1s + (nb / blocks_n_per_group) * scale_size_K,
          /*   M            */ m_size,
@@ -409,8 +274,9 @@ void shared_expert_fp8_kernel_impl(
          /*   ldb          */ n_size,
          /*   ldc          */ 2 * N,
          /*   brg          */ use_brgemm,
-          /*   block_size_K */ block_size_K);
-    }
+          /*   block_size_K */ block_size_K,
+          /*   do_unpack    */ do_unpack);
+    });

    if (use_brgemm) {
      at::native::cpublas::brgemm_release();
@@ -420,11 +286,7 @@ void shared_expert_fp8_kernel_impl(
  // stage 1.5: intermediate_cache1 = silu(intermediate_cache0)
  at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
    for (int64_t m = begin; m < end; ++m) {
-      silu_and_mul_stub(
-          ic1 + m * N,
-          ic0 + m * 2 * N,
-          ic0 + m * 2 * N + N,
-          N);
+      silu_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, ic0 + m * 2 * N + N, N);
    }
  });

@@ -437,22 +299,23 @@ void shared_expert_fp8_kernel_impl(
  scale_size_K = div_up(N, block_size_K);

  // parallel on [MB2, NB2]
-  at::parallel_for(0, MB2 * NB2, 0, [&](int64_t begin, int64_t end) {
-    int tid = at::get_thread_num();
+  parallel_2d(MB2, NB2, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
+    int tid = get_thread_num();
    alignas(64) scalar_t C[BLOCK_M * BLOCK_K];

-    for (int64_t i = begin; i < end; ++i) {
-      int64_t mb = i / NB2;
-      int64_t nb = i % NB2;
+    loop_2d<at::Float8_e4m3fn>(mb0, mb1, nb0, nb1, BLOCK_N * IC, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
      int64_t m_size = std::min(M - mb * BLOCK_M, BLOCK_M);
      int64_t n_size = std::min(OC - nb * BLOCK_N, BLOCK_N);

+      // do unpacking for the first row
+      bool do_unpack = (mb == mb0);
+
      // 2.a gemm: C = A @ B
      tinygemm_kernel<scalar_t>(
          /*   A            */ ic1 + mb * BLOCK_M * N,
          /*   B            */ packed_w2 + nb * BLOCK_N * N,
          /*   C            */ C,
-          /*   Btmp         */ B_tmp + tid * BLOCK_N * std::max(K, N),
+          /*   Btmp         */ B_tmp + tid * B_tmp_size_per_thread + nb_offset * BLOCK_N * IC,
          /*   Ctmp         */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N,
          /*   scale        */ w2s + (nb / blocks_n_per_group) * scale_size_K,
          /*   M            */ m_size,
@@ -462,15 +325,18 @@ void shared_expert_fp8_kernel_impl(
          /*   ldb          */ n_size,
          /*   ldc          */ BLOCK_N,
          /*   brg          */ use_brgemm,
-          /*   block_size_K */ block_size_K);
+          /*   block_size_K */ block_size_K,
+          /*   do_unpack    */ do_unpack);

      // 2.b copy from C to output and add fused_experts_out
      scalar_t* __restrict__ out = output + mb * BLOCK_M * K + nb * BLOCK_N;
-      const scalar_t* __restrict__ fused_out = fused_experts_out + mb * BLOCK_M * K + nb * BLOCK_N;
+      const scalar_t* __restrict__ fused_out =
+          apply_scaling_factor ? fused_experts_out + mb * BLOCK_M * K + nb * BLOCK_N : nullptr;
      for (int64_t m = 0; m < m_size; ++m) {
-        add_mul_stub(out + m * K, C + m * BLOCK_N, fused_out + m * K, routed_scaling_factor, n_size);
+        const scalar_t* __restrict__ fused_out_row = apply_scaling_factor ? (fused_out + m * K) : nullptr;
+        add_mul_stub(out + m * K, C + m * BLOCK_N, fused_out_row, routed_scaling_factor, n_size);
      }
-    }
+    });
  });

  if (use_brgemm) {
@@ -0,0 +1,323 @@
+// Adapted from
+// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu
+
+// clang-format off
+
+#include "common.h"
+#include "gemm.h"
+#include "moe.h"
+
+template <int64_t N>
+inline void copy_bias(const float* bias_ptr, float* y_buf, int64_t m, int64_t ldn) {
+  using Vec = at::vec::Vectorized<float>;
+  constexpr int kVecSize = Vec::size();
+  static_assert(N % kVecSize == 0, "copy_bias requires N to be a multiple of Vectorized<float>::size()");
+  const bool has_bias = bias_ptr != nullptr;
+  const Vec zero_vec(0.f);
+  for (int i = 0; i < m; ++i) {
+#pragma GCC unroll 2
+    for (int j = 0; j < N; j += kVecSize) {
+      Vec vec = has_bias ? Vec::loadu(bias_ptr + j) : zero_vec;
+      vec.store(y_buf + i * ldn + j);
+    }
+  }
+}
+
+template <typename scalar_t>
+void fused_experts_int4_w4a8_kernel_impl(
+    scalar_t* __restrict__ output,
+    scalar_t* __restrict__ ic0,
+    scalar_t* __restrict__ ic1,
+    scalar_t* __restrict__ ic2,
+    uint8_t* __restrict__ A_tmp,
+    uint8_t* __restrict__ Aq_tmp,
+    float* __restrict__ As_tmp,
+    int32_t* __restrict__ Azp_tmp,
+    float* __restrict__ C_tmp,
+    int8_t* __restrict__ dqB_tmp,
+    const scalar_t* __restrict__ input,
+    const uint8_t* __restrict__ packed_w1,
+    const uint8_t* __restrict__ packed_w2,
+    const int8_t* __restrict__ w1z,
+    const int8_t* __restrict__ w2z,
+    const float* __restrict__ w1s,
+    const float* __restrict__ w2s,
+    int group_size,
+    const float* __restrict__ topk_weights,
+    const int32_t* __restrict__ sorted_ids,
+    const int32_t* __restrict__ expert_ids,
+    const int32_t* __restrict__ offsets,
+    int64_t M,
+    int64_t N,
+    int64_t K,
+    int64_t E,
+    int64_t topk,
+    int64_t num_tokens_post_pad) {
+  constexpr int64_t BLOCK_M = block_size_m();
+  constexpr int64_t BLOCK_N = block_size_n();
+  int num_threads = at::get_num_threads();
+  // int64_t buffer_size_nbytes = M * topk * N * 2
+  //                              M * topk * K * 2 +
+  //                              num_threads * BLOCK_M * K +
+  //                              num_threads * 2 * BLOCK_M * BLOCK_N * sizeof(float)  +
+  //                              M * topk * 2 * N * 2 +
+  //                              max(M * K, M * topk * N)  +
+  //                              M * topk * sizeof(float);
+
+  // intermediate_cache1 (scalar_t):     START + M * topk * N
+  // intermediate_cache2 (scalar_t):     + M * topk * K
+  // A_tmp (uint8_t):                    + num_threads * BLOCK_M * K
+  // C_tmp (float):                      + num_threads * 2 * BLOCK_M * BLOCK_N
+  // intermediate_cache0 (scalar_t):     + M * topk * 2 * N
+  // Aq_tmp (uint8_t):                   + max(M * K, M * topk * N)
+  // As_tmp (float):                     + M * topk
+  // dqB_tmp (int8_t)                    + num_threads * _block_k * BlOCK_N
+
+  // stage 0: quantize input to uint8, [M, K]
+  at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
+    for (int64_t m = begin; m < end; ++m) {
+      quantize_row_int8<scalar_t>(Aq_tmp + m * K, As_tmp[m], input + m * K, K);
+    }
+  });
+  int64_t _block_k = get_4bit_block_k_size(group_size);
+  auto Azp = at::ones({M * topk}).to(at::kInt).mul(128);
+  auto Azp_ptr = Azp.data_ptr<int32_t>();
+  // stage 1: intermediate_cache0 = hidden_states @ w1
+  const int64_t MB = div_up(num_tokens_post_pad, BLOCK_M);
+  const int64_t NB = div_up(N, BLOCK_N);
+
+  int64_t block_per_group = group_size / _block_k;
+  int64_t Kc = K / _block_k;
+  int64_t num_groups = K / group_size;
+
+  const int64_t stride_e = 2 * NB * Kc * (BLOCK_N * (_block_k / 2 + sizeof(int32_t)));
+  const bool sym_quant_act = false;
+  // weight + compensation shape = [E, Nc, Kc, block_n * _block_k / 2 + block_n*sizeof(int32_t)]
+  // scales/qzeros shape = [E, Nc, G, block_n]
+
+  // here we only parallel on half of 2N to fuse silu_and_mul with gemm
+  at::parallel_for(0, MB * NB, 0, [&](int64_t begin, int64_t end) {
+    // get local pointers
+    int tid = at::get_thread_num();
+    int8_t* dqB_tmp1 = dqB_tmp + tid * 2 * _block_k * BLOCK_N;
+    int8_t* dqB_tmp2 = dqB_tmp1 + _block_k * BLOCK_N;
+    alignas(64) float As[BLOCK_M];
+    uint8_t* __restrict__ A = A_tmp + tid * BLOCK_M * K;
+    float* __restrict__ C0 = C_tmp + tid * 2 * BLOCK_M * BLOCK_N;
+    float* __restrict__ C1 = C0 + BLOCK_M * BLOCK_N;
+    bool is_brgemm_used = false;
+    for (int64_t i = begin; i < end; ++i) {
+      int64_t mb = i / NB;
+      int64_t nb = i % NB;
+      int64_t nb1 = nb + NB;
+      int64_t n_size = std::min(N - nb * BLOCK_N, BLOCK_N);
+      // B shape [K, n_size] in vnni format
+      int32_t expert_id = expert_ids[mb];
+      const uint8_t* __restrict__ B = packed_w1 + expert_id * stride_e;
+      // Bz and Bs: [E, K/gs, 2N]
+      const int8_t* __restrict__ Bz = w1z + expert_id * (num_groups) * (2 * N);
+      const float* __restrict__ Bs = w1s + expert_id * (num_groups) * (2 * N);
+
+      // 1.a load A
+      const int32_t* A_ids = sorted_ids + mb * BLOCK_M;
+      int64_t m_size = offsets[mb + 1] - offsets[mb];
+      const bool use_brgemm = can_use_brgemm<int8_t>(m_size);
+      is_brgemm_used = is_brgemm_used || use_brgemm;
+      // copy to A [BLOCK_M, K]
+      for (int64_t m = 0; m < m_size; ++m) {
+        int32_t index = A_ids[m] / topk;
+        copy_stub(A + m * K, Aq_tmp + index * K, K);
+        As[m] = As_tmp[index];
+      }
+      const int64_t offset = offsets[mb];
+      copy_bias<BLOCK_N>(nullptr, C0, m_size, BLOCK_N);
+      copy_bias<BLOCK_N>(nullptr, C1, m_size, BLOCK_N);
+      for (int kci = 0; kci < Kc; ++kci) {
+        int32_t* compensation_ptr =
+            sym_quant_act ? nullptr
+                          : (int32_t*)(void*)(B + (nb * Kc + kci) * (BLOCK_N * (_block_k / 2 + sizeof(int32_t))) +
+                                              _block_k * BLOCK_N / 2) /*Bcomp*/;
+        tinygemm_kernel<scalar_t>(
+            ic0 + offset * 2 * N + nb * BLOCK_N,
+            C0,
+            A + kci * _block_k,
+            As,
+            Azp_ptr,
+            B + (nb * Kc + kci) * (BLOCK_N * (_block_k / 2 + sizeof(int32_t))) /*B*/,
+            Bs + nb * BLOCK_N * num_groups + kci / block_per_group * BLOCK_N /*scales_b*/,
+            Bz + nb * BLOCK_N * num_groups + kci / block_per_group * BLOCK_N /*qzeros_b*/,
+            compensation_ptr,
+            dqB_tmp1,
+            m_size,
+            _block_k,
+            K,
+            BLOCK_N,
+            2 * N,
+            kci == Kc - 1,
+            use_brgemm);
+      }
+
+      for (int kci = 0; kci < Kc; ++kci) {
+        int32_t* compensation_ptr =
+            sym_quant_act ? nullptr
+                          : (int32_t*)(void*)(B + (nb1 * Kc + kci) * (BLOCK_N * (_block_k / 2 + sizeof(int32_t))) +
+                                              _block_k * BLOCK_N / 2) /*Bcomp*/;
+        tinygemm_kernel<scalar_t>(
+            ic0 + offset * 2 * N + nb1 * BLOCK_N,
+            C1,
+            A + kci * _block_k,
+            As,
+            Azp_ptr,
+            B + (nb1 * Kc + kci) * (BLOCK_N * (_block_k / 2 + sizeof(int32_t))) /*B*/,
+            Bs + nb1 * BLOCK_N * num_groups + kci / block_per_group * BLOCK_N /*scales_b*/,
+            Bz + nb1 * BLOCK_N * num_groups + kci / block_per_group * BLOCK_N /*qzeros_b*/,
+            compensation_ptr,
+            dqB_tmp2,
+            m_size,
+            _block_k,
+            K,
+            BLOCK_N,
+            2 * N,
+            kci == Kc - 1,
+            use_brgemm);
+      }
+    }
+
+    if (is_brgemm_used) {
+      at::native::cpublas::brgemm_release();
+    }
+  });
+
+  // stage 1.5: intermediate_cache1 = silu(intermediate_cache0)
+  at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
+    for (int64_t m = begin; m < end; ++m) {
+      silu_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, ic0 + m * 2 * N + N, N);
+    }
+  });
+
+  // stage 1.5: quantize ic1 to uint8, [M * topk, N]
+  at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
+    for (int64_t m = begin; m < end; ++m) {
+      quantize_row_int8<scalar_t>(Aq_tmp + m * N, As_tmp[m], ic1 + m * N, N);
+    }
+  });
+  // stage 2: intermediate_cache2 = intermediate_cache1 @ w2
+  //   w2 : [E, K, N] as [E, OC, IC]
+  const int64_t OC = K;  // rename K as OC
+  const int64_t IC = N;  // rename N as IC
+  const int64_t MB2 = MB;
+  const int64_t NB2 = div_up(OC, BLOCK_N);
+  const int64_t stride_oc = IC;
+  num_groups = IC / group_size;
+  Kc = IC / _block_k;
+  const int64_t stride_e2 = NB2 * Kc * (BLOCK_N * (_block_k / 2 + sizeof(int32_t)));
+  // parallel on [MB2, NB2]
+  at::parallel_for(0, MB2 * NB2, 0, [&](int64_t begin, int64_t end) {
+    int tid = at::get_thread_num();
+    int8_t* dqB_tmp1 = dqB_tmp + tid * 2 * _block_k * BLOCK_N;
+    float* __restrict__ C2 = C_tmp + tid * 2 * BLOCK_M * BLOCK_N;
+    bool is_brgemm_used = false;
+    for (int64_t i = begin; i < end; ++i) {
+      int64_t mb = i / NB2;
+      int64_t nb = i % NB2;
+
+      int64_t m_size = offsets[mb + 1] - offsets[mb];
+      int64_t n_size = std::min(OC - nb * BLOCK_N, BLOCK_N);
+      const bool use_brgemm = can_use_brgemm<int8_t>(m_size);
+      is_brgemm_used = is_brgemm_used || use_brgemm;
+      const int32_t* A_ids = sorted_ids + mb * BLOCK_M;
+
+      // B shape [IC, n_size] in vnni format
+      int32_t expert_id = expert_ids[mb];
+      const uint8_t* __restrict__ B = packed_w2 + expert_id * stride_e2;
+
+      // Bz and Bs: [E, IC/gs, OC]
+      const int8_t* __restrict__ Bz = w2z + expert_id * (num_groups)*OC;
+      const float* __restrict__ Bs = w2s + expert_id * (num_groups)*OC;
+
+      // A ptr from ic1 of [M * topk, N] in sorted order
+      // so as to avoid copy A to tmp buffer again
+      const uint8_t* __restrict__ A = Aq_tmp + offsets[mb] * IC;
+      const float* __restrict__ As = As_tmp + offsets[mb];
+      copy_bias<BLOCK_N>(nullptr, C2, m_size, BLOCK_N);
+      for (int kci = 0; kci < Kc; ++kci) {
+        int32_t* compensation_ptr =
+            sym_quant_act ? nullptr
+                          : (int32_t*)(void*)(B + (nb * Kc + kci) * (BLOCK_N * (_block_k / 2 + sizeof(int32_t))) +
+                                              _block_k * BLOCK_N / 2) /*Bcomp*/;
+        tinygemm_kernel<scalar_t>(
+            nullptr, /*store_out is false*/
+            C2,
+            A + kci * _block_k,
+            As,
+            Azp_ptr,
+            B + (nb * Kc + kci) * (BLOCK_N * (_block_k / 2 + sizeof(int32_t))),
+            Bs + nb * BLOCK_N * num_groups + kci / block_per_group * BLOCK_N /*scales_b*/,
+            Bz + nb * BLOCK_N * num_groups + kci / block_per_group * BLOCK_N /*zeros_b*/,
+            compensation_ptr,
+            dqB_tmp1,
+            m_size,
+            _block_k,
+            IC,
+            BLOCK_N,
+            BLOCK_N,
+            false,
+            use_brgemm);
+      }
+
+      // 2.b copy from C to ic2 in original order
+      //   and also mul topk_weights in float32
+      for (int64_t m = 0; m < m_size; ++m) {
+        int32_t index = A_ids[m];
+        float weight = topk_weights[index];
+        copy_mul_stub(ic2 + index * K + nb * BLOCK_N, C2 + m * BLOCK_N, weight, n_size);
+      }
+    }
+
+    if (is_brgemm_used) {
+      at::native::cpublas::brgemm_release();
+    }
+  });
+
+  // stage 3: out = intermediate_cache2.sum(dim=1)
+  //   from [M, topk, K] to [M, K]
+  at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
+    for (int64_t m = begin; m < end; ++m) {
+      sum_stub(output + m * K, ic2 + m * topk * K, topk, K);
+    }
+  });
+}
+
+#define INSTANTIATE_MOE_INT4_W4A8_TEMPLATE(TYPE)           \
+  template void fused_experts_int4_w4a8_kernel_impl<TYPE>( \
+      TYPE* __restrict__ output,                           \
+      TYPE* __restrict__ ic0,                              \
+      TYPE* __restrict__ ic1,                              \
+      TYPE* __restrict__ ic2,                              \
+      uint8_t* __restrict__ A_tmp,                         \
+      uint8_t* __restrict__ Aq_tmp,                        \
+      float* __restrict__ As_tmp,                          \
+      int32_t* __restrict__ Azp_tmp,                       \
+      float* __restrict__ C_tmp,                           \
+      int8_t* __restrict__ dqB_tmp,                        \
+      const TYPE* __restrict__ input,                      \
+      const uint8_t* __restrict__ packed_w1,               \
+      const uint8_t* __restrict__ packed_w2,               \
+      const int8_t* __restrict__ w1z,                      \
+      const int8_t* __restrict__ w2z,                      \
+      const float* __restrict__ w1s,                       \
+      const float* __restrict__ w2s,                       \
+      int group_size,                                      \
+      const float* __restrict__ topk_weights,              \
+      const int32_t* __restrict__ sorted_ids,              \
+      const int32_t* __restrict__ expert_ids,              \
+      const int32_t* __restrict__ offsets,                 \
+      int64_t M,                                           \
+      int64_t N,                                           \
+      int64_t K,                                           \
+      int64_t E,                                           \
+      int64_t topk,                                        \
+      int64_t num_tokens_post_pad)
+
+INSTANTIATE_MOE_INT4_W4A8_TEMPLATE(at::BFloat16);
+INSTANTIATE_MOE_INT4_W4A8_TEMPLATE(at::Half);
@@ -1,10 +1,10 @@
 // Adapted from
 // https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu

-#pragma once
-
 // clang-format off

+#pragma once
+
 #if defined(__AVX512F__) && defined(__AVX512BF16__) && defined(__AMX_BF16__)
 #define CPU_CAPABILITY_AVX512
 #endif
@@ -16,38 +16,51 @@ namespace {

 using namespace at::vec;

-template <typename scalar_t,
-          typename std::enable_if_t<is_reduced_floating_point_v<scalar_t>, int> = 0>
+template <typename scalar_t, typename std::enable_if_t<is_reduced_floating_point_v<scalar_t>, int> = 0>
 inline Vectorized<scalar_t> convert_from_float_ext(const Vectorized<float>& a, const Vectorized<float>& b) {
  return at::vec::convert_from_float<scalar_t>(a, b);
 }

+// allow f16, bf16
+template <typename scalar_t, typename std::enable_if_t<is_reduced_floating_point_v<scalar_t>, int> = 1>
+inline std::tuple<Vectorized<float>, Vectorized<float>> load_float_vec2(const scalar_t* __restrict__ data) {
+  using bVec = at::vec::Vectorized<scalar_t>;
+  using fVec = at::vec::Vectorized<float>;
+  bVec x_vec = bVec::loadu(data);
+  fVec x0, x1;
+  std::tie(x0, x1) = at::vec::convert_to_float(x_vec);
+  return std::make_tuple(x0, x1);
+}
+
+// allow  f32
+inline std::tuple<Vectorized<float>, Vectorized<float>> load_float_vec2(const float* __restrict__ data) {
+  using fVec = at::vec::Vectorized<float>;
+  fVec x0 = fVec::loadu(data);
+  fVec x1 = fVec::loadu(data + fVec::size());
+  return std::make_tuple(x0, x1);
+}
+
 #if defined(CPU_CAPABILITY_AVX512)

 // `at::vec::convert_from_float<>` from PyTorch doesn't have avx512-bf16 intrinsics
 // use native instruction for bfloat16->float32 conversion
 template <>
-inline Vectorized<at::BFloat16> convert_from_float_ext<at::BFloat16>(const Vectorized<float>& a, const Vectorized<float>& b) {
+inline Vectorized<at::BFloat16>
+convert_from_float_ext<at::BFloat16>(const Vectorized<float>& a, const Vectorized<float>& b) {
  return (__m512i)(_mm512_cvtne2ps_pbh(__m512(b), __m512(a)));
 }

-#define CVT_BF16_TO_FP32(a) \
-    _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(a), 16))
+#define CVT_BF16_TO_FP32(a) _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(a), 16))

-#define CVT_FP16_TO_FP32(a) \
-    _mm512_cvtps_ph(a, (_MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC))
+#define CVT_FP16_TO_FP32(a) _mm512_cvtph_ps(a)

 // this doesn't handle NaN.
 inline __m512bh cvt_e4m3_bf16_intrinsic_no_nan(__m256i fp8_vec) {
  const __m512i x = _mm512_cvtepu8_epi16(fp8_vec);
-
-  const __m512i mant = _mm512_slli_epi16(_mm512_and_si512(x, _mm512_set1_epi16(0x07)), 4);
-  const __m512i raw_exp = _mm512_srli_epi16(_mm512_and_si512(x, _mm512_set1_epi16(0x78)), 3);
-  const __m512i exp = _mm512_slli_epi16(_mm512_add_epi16(raw_exp, _mm512_set1_epi16(120)), 7);
-  const __m512i nonsign = _mm512_or_si512(exp, mant);
-
-  const __m512i sign = _mm512_slli_epi16(_mm512_and_si512(x, _mm512_set1_epi16(0x80)), 8);
-  const __m512i combined = _mm512_or_si512(nonsign, sign);
+  __m512i combined = _mm512_add_epi16(x, _mm512_set1_epi16(0x0780));
+  combined = _mm512_slli_epi16(combined, 4);
+  combined = _mm512_and_si512(combined, _mm512_set1_epi16(0x87f0));
+  combined = _mm512_add_epi16(combined, _mm512_set1_epi16(0x3c00));

  const __mmask32 is_nonzero = _mm512_cmpneq_epi16_mask(x, _mm512_setzero_si512());
  return (__m512bh)_mm512_maskz_mov_epi16(is_nonzero, combined);
@@ -113,10 +126,77 @@ inline __m512bh CVT_FP8_TO_BF16(__m256i a) {
 #endif
 }

+// faster version of float8_e4m3fn conversion to bfloat16
+//
+// we mapped cuda implementation from below link and vectorized with avx512:
+// https://github.com/thu-pacman/chitu/blob/1ed2078ec26581ebdca05b7306d4385f86edaa7c/csrc/cuda/marlin/marlin_gemm/dequant.h#L387
+//
+inline __attribute__((always_inline)) __m512bh CVT_FP8_TO_BF16_EXT(__m256i a) {
+  const __m512i mask0 = _mm512_set1_epi16(0x80);  // sign bit
+  const __m512i mask1 = _mm512_set1_epi16(0x7F);  // exponent and mantissa
+  const __m512i mask2 = _mm512_set1_epi16(0x4000);
+
+  __m512i x = _mm512_cvtepu8_epi16(a);
+  __m512i vsign = _mm512_and_si512(x, mask0);
+  vsign = _mm512_slli_epi16(vsign, 8);
+
+  __m512i vexp_and_mant = _mm512_and_si512(x, mask1);
+  vexp_and_mant = _mm512_slli_epi16(vexp_and_mant, 4);
+
+  // _MM_TERNLOG_A | _MM_TERNLOG_B | _MM_TERNLOG_C: 0b11111110
+  return (__m512bh)(_mm512_ternarylogic_epi32(vsign, mask2, vexp_and_mant, 0b11111110));
+}
+
+// bias for conversion of fp8 to bf16 1/256 in float32
+#define kFP8_BIAS 0x3b800000
+
+// remove warning: ignoring attributes on template argument ‘__m512bh’ [-Wignored-attributes]
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Wignored-attributes"
+
+#define MXFP4_VALUES \
+  -6.0f, -4.0f, -3.0f, -2.0f, -1.5f, -1.0f, -0.5f, -0.0f, 6.0f, 4.0f, 3.0f, 2.0f, 1.5f, 1.0f, 0.5f, 0.0f
+
+// convert 64 mxfp4 to 2x bf16 vectors, expect input 32-way packing
+inline std::tuple<__m512bh, __m512bh> cvt_mxfp4_e2m1_bf16_intrinsic_lut(__m256i a, __m512i s0, __m512i s1) {
+  // LUT
+  const __m512 values = _mm512_set_ps(MXFP4_VALUES);
+  const __m512i lut = (__m512i)(_mm512_cvtne2ps_pbh(values, values));
+
+  const __m512i abs_mask = _mm512_set1_epi16(0x7FFF);
+  const __m512i zero = _mm512_setzero_si512();
+
+  // expand values to 16-bit integers
+  __m512i x0 = _mm512_cvtepu8_epi16(a);
+  __m512i x1 = _mm512_srli_epi32(x0, 4);
+
+  // LUT to convert mxfp4 values to bf16
+  x0 = _mm512_permutexvar_epi16(x0, lut);
+  x1 = _mm512_permutexvar_epi16(x1, lut);
+
+  // check for zeros
+  __mmask32 mask0 = _mm512_cmp_epi16_mask(_mm512_and_si512(x0, abs_mask), zero, _MM_CMPINT_EQ);
+  __mmask32 mask1 = _mm512_cmp_epi16_mask(_mm512_and_si512(x1, abs_mask), zero, _MM_CMPINT_EQ);
+
+  // emulate bf16 mul with scale factor
+  x0 = _mm512_add_epi16(x0, s0);
+  x1 = _mm512_add_epi16(x1, s1);
+
+  // blend with zero
+  x0 = _mm512_mask_blend_epi16(mask0, x0, zero);
+  x1 = _mm512_mask_blend_epi16(mask1, x1, zero);
+
+  return std::make_tuple(__m512bh(x0), __m512bh(x1));
+}
+
+#define CVT_MXFP4_TO_BF16(a, s0, s1) cvt_mxfp4_e2m1_bf16_intrinsic_lut(a, s0, s1)
+
+#pragma GCC diagnostic pop
+
 #endif

 // vector to scalar reduction
-#if defined(CPU_CAPABILITY_AVX512) && 0
+#if defined(CPU_CAPABILITY_AVX512)
 inline float vec_reduce_sum(const Vectorized<float>& a) {
  return _mm512_reduce_add_ps(__m512(a));
 }
@@ -136,10 +216,9 @@ inline float vec_reduce_max(const Vectorized<float>& a) {

 // https://github.com/InternLM/lmdeploy/blob/086481ed84b59bee3b8e4274e5fc69620040c048/lmdeploy/pytorch/kernels/cuda/w8a8_triton_kernels.py#L282
 template <typename scalar_t>
-inline void quantize_row_int8(uint8_t* __restrict__ Aq, float& As,
-    const scalar_t* __restrict__ A, int64_t K, float eps = 1e-7) {
-
-  float amax = 0.f; // absolute max
+inline void
+quantize_row_int8(uint8_t* __restrict__ Aq, float& As, const scalar_t* __restrict__ A, int64_t K, float eps = 1e-7) {
+  float amax = 0.f;  // absolute max
  for (int64_t k = 0; k < K; ++k) {
    const float val = static_cast<float>(A[k]);
    amax = std::max(amax, std::abs(val));
@@ -158,9 +237,8 @@ inline void quantize_row_int8(uint8_t* __restrict__ Aq, float& As,

 #if defined(CPU_CAPABILITY_AVX512)
 template <>
-inline void quantize_row_int8<at::BFloat16>(uint8_t* __restrict__ Aq, float& As,
-    const at::BFloat16* __restrict__ A, int64_t K, float eps) {
-
+inline void quantize_row_int8<at::BFloat16>(
+    uint8_t* __restrict__ Aq, float& As, const at::BFloat16* __restrict__ A, int64_t K, float eps) {
  const __m512 signBit = _mm512_set1_ps(-0.0f);
  const __m512i off = _mm512_set1_epi32(128);

@@ -200,7 +278,7 @@ inline void quantize_row_int8<at::BFloat16>(uint8_t* __restrict__ Aq, float& As,
 // transpose utils
 // taken from my PR in ggml: https://github.com/ggml-org/llama.cpp/pull/8998
 #if defined(CPU_CAPABILITY_AVX512)
-inline void transpose_16x16_32bit(__m512i * v) {
+inline void transpose_16x16_32bit(__m512i* v) {
  __m512i v1[16];
  v1[0] = _mm512_unpacklo_epi32(v[0], v[1]);
  v1[1] = _mm512_unpackhi_epi32(v[0], v[1]);
@@ -293,16 +371,56 @@ inline std::tuple<__m512i, __m512i> transpose_2x32_16bit(__m512i r0, __m512i r1)
 }
 #pragma GCC diagnostic pop

+inline __attribute__((always_inline)) __m512 _mm512_fexp_u20_ps(const __m512 values) {
+  const __m512 vec_c0 = _mm512_set1_ps(0.00010703434948458272f);
+  const __m512 vec_c1 = _mm512_set1_ps(0.30354260500649682f);
+  const __m512 vec_c2 = _mm512_set1_ps(-0.22433836478672356);
+  const __m512 vec_c3 = _mm512_set1_ps(-0.079204240219773236);
+
+  const __m512 vec_exp_log2ef = _mm512_castsi512_ps(_mm512_set1_epi32(0x3fb8aa3b));  // log2(e)
+
+  const __m512 vec_a = _mm512_set1_ps(std::pow(2, 23) / std::log2(2));
+  const __m512 vec_b = _mm512_set1_ps(std::pow(2, 23) * 127.f);
+
+  const __m512 vec_ln_flt_min = _mm512_castsi512_ps(_mm512_set1_epi32(0xc2aeac50));
+  const __m512 vec_ln_flt_max = _mm512_castsi512_ps(_mm512_set1_epi32(0x42b17218));
+  __m512i vec_infinity = _mm512_set1_epi32(0x7F800000);
+  __m512i vec_zero = _mm512_setzero_epi32();
+
+  // Fast Exponential Computation on SIMD Architectures
+  // A. Cristiano I. Malossi, Yves Ineichen, Costas Bekas, and Alessandro
+  // Curioni exp(x) = 2**(x * log2(e))
+  //        = 2**xi * 2**xf   - TIPS we are using  the EEEE floating point
+  //        representation with identification to the exponent and the
+  //        mentissa
+  //  2**xf will be approximated to a polynomial of degree 3 computed with
+  //  Horner method
+  // mask for the boundary condition
+  auto min_mask = _mm512_cmp_ps_mask(values, vec_ln_flt_min, _CMP_LT_OS);
+  auto max_mask = _mm512_cmp_ps_mask(values, vec_ln_flt_max, _CMP_GT_OS);
+
+  // transformation with log2(e)
+  auto vec_src = _mm512_mul_ps(values, vec_exp_log2ef);
+  auto vec_fractional = _mm512_sub_ps(vec_src, _mm512_floor_ps(vec_src));
+
+  // compute polynomial using Horner Scheme, for superscalar processor
+  auto vec_res = _mm512_fmadd_ps(vec_fractional, vec_c3, vec_c2);
+  vec_res = _mm512_fmadd_ps(vec_fractional, vec_res, vec_c1);
+  vec_res = _mm512_fmadd_ps(vec_fractional, vec_res, vec_c0);
+
+  vec_src = _mm512_sub_ps(vec_src, vec_res);
+  // the tips is here, headache in perspective
+  auto tmp = _mm512_fmadd_ps(vec_a, vec_src, vec_b);
+  // headache bis - we loose precision with the cast but it "fits", but ok
+  // after f32 -> f16 later
+  __m512i casted_integer = _mm512_cvttps_epi32(tmp);
+  // boundary condition, lower than the min -> 0
+  casted_integer = _mm512_mask_mov_epi32(casted_integer, min_mask, vec_zero);
+  // boundary condition, larger than the max -> +oo
+  casted_integer = _mm512_mask_mov_epi32(casted_integer, max_mask, vec_infinity);
+  // final interpretation to float
+  return _mm512_castsi512_ps(casted_integer);
+}
 #endif

-// TODO: debug print, remove me later
-template<typename scalar_t>
-void print_array(scalar_t* ptr, int size) {
-  for (int d = 0; d < size; ++d) {
-    if (d % 16 == 0) { std::cout << std::endl; }
-    std::cout << ptr[d] << " ";
-  }
-  std::cout << std::endl;
-}
-
-} // anonymous namespace
+}  // anonymous namespace
@@ -0,0 +1,299 @@
+// Adapted from
+// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu
+
+// clang-format off
+
+// To use the transpose functions
+#include <ATen/native/cpu/utils.h>
+
+#include "vec.h"
+
+namespace {
+
+using namespace at::vec;
+
+template <typename index_t>
+inline index_t get_index(index_t* ind, int i) {
+  return (ind == nullptr) ? (index_t)i : ind[i];
+}
+
+#if defined(CPU_CAPABILITY_AVX512)
+// key: from [N, 32] to [32/2, N, 2]
+template <typename scalar_t, typename index_t>
+inline void pack_vnni_Nx32(
+    scalar_t* __restrict__ dst,
+    const scalar_t* __restrict__ src,
+    const index_t* __restrict__ ind,
+    int N,
+    int ld_src,
+    int ld_dst) {
+  __m512i vinputs[16];
+
+  int n = 0;
+  for (; n < N; ++n) {
+    index_t index = get_index(ind, n);
+    vinputs[n] = _mm512_loadu_si512(src + index * ld_src);
+  }
+  // padding with zero to avoid uninitialized vectors
+  for (; n < 16; ++n) {
+    vinputs[n] = _mm512_set1_epi32(0);
+  }
+
+  // pack key
+  transpose_16x16_32bit(vinputs);
+
+  const __mmask16 vmask = (1 << N) - 1;
+  for (int k = 0; k < 16; ++k) {
+    _mm512_mask_storeu_epi32(dst + k * ld_dst * 2, vmask, vinputs[k]);
+  }
+}
+
+template <typename scalar_t, typename index_t>
+inline void pack_vnni_N_remainder(
+    scalar_t* __restrict__ dst,
+    const scalar_t* __restrict__ src,
+    const index_t* __restrict__ ind,
+    int N,
+    int K,
+    int ld_src,
+    int ld_dst) {
+  __m512i vinputs[16];
+
+  int K2 = K >> 1;
+  const __mmask16 vmask = (1 << K2) - 1;
+
+  int n = 0;
+  for (; n < N; ++n) {
+    index_t index = get_index(ind, n);
+    vinputs[n] = _mm512_maskz_loadu_epi32(vmask, src + index * ld_src);
+  }
+  // padding with zero to avoid uninitialized vectors
+  for (; n < 16; ++n) {
+    vinputs[n] = _mm512_set1_epi32(0);
+  }
+
+  // pack key
+  transpose_16x16_32bit(vinputs);
+
+  const __mmask16 vmask2 = (1 << N) - 1;
+  for (int k = 0; k < K2; ++k) {
+    _mm512_mask_storeu_epi32(dst + k * ld_dst * 2, vmask2, vinputs[k]);
+  }
+}
+
+// value: from [K, 32] to [K/2, 32, 2]
+template <typename scalar_t, typename index_t>
+inline void pack_vnni_Kx32(
+    scalar_t* __restrict__ dst,
+    const scalar_t* __restrict__ src,
+    const index_t* __restrict__ ind,
+    int K,
+    int ld_src,
+    int ld_dst) {
+  __m512i vinputs[2];
+
+  int k = 0;
+  for (; k < K; ++k) {
+    index_t index = get_index(ind, k);
+    vinputs[k] = _mm512_loadu_si512(src + index * ld_src);
+  }
+  // padding with zero to avoid uninitialized vectors
+  for (; k < 2; ++k) {
+    vinputs[k] = _mm512_set1_epi32(0);
+  }
+
+  // pack value
+  __m512i d0, d1;
+  std::tie(d0, d1) = transpose_2x32_16bit(vinputs[0], vinputs[1]);
+  _mm512_storeu_si512(dst + 0 * ld_dst * 2, d0);
+  _mm512_storeu_si512(dst + 0 * ld_dst * 2 + 32, d1);
+}
+
+template <typename scalar_t, typename index_t>
+inline void pack_vnni_K_remainder(
+    scalar_t* __restrict__ dst,
+    const scalar_t* __restrict__ src,
+    const index_t* __restrict__ ind,
+    int K,
+    int N,
+    int ld_src,
+    int ld_dst) {
+  __m512i vinputs[2];
+
+  const __mmask32 vmask = (1 << N) - 1;
+
+  int k = 0;
+  for (; k < K; ++k) {
+    index_t index = get_index(ind, k);
+    vinputs[k] = _mm512_maskz_loadu_epi16(vmask, src + index * ld_src);
+  }
+  // padding with zero to avoid uninitialized vectors
+  for (; k < 2; ++k) {
+    vinputs[k] = _mm512_set1_epi32(0);
+  }
+
+  // pack value
+  __m512i d0, d1;
+  std::tie(d0, d1) = transpose_2x32_16bit(vinputs[0], vinputs[1]);
+
+  if (N <= 16) {
+    // 2N * 16bits: N * 32bits
+    const __mmask16 vmask2 = (1 << N) - 1;
+    _mm512_mask_storeu_epi32(dst + 0 * ld_dst * 2, vmask2, d0);
+  } else {
+    // 2(N-16) * 16bits: (N-16) * 32bits
+    const __mmask16 vmask2 = (1 << (N - 16)) - 1;
+    _mm512_storeu_epi32(dst + 0 * ld_dst * 2, d0);
+    _mm512_mask_storeu_epi32(dst + 0 * ld_dst * 2 + 32, vmask2, d1);
+  }
+}
+#endif
+
+// convert to vnni format
+// from [N, K/2, 2] to [K/2, N, 2] for bfloat16 and float16
+template <typename scalar_t, typename index_t, bool is_indexed>
+void pack_vnni(
+    scalar_t* __restrict__ dst,
+    const scalar_t* __restrict__ src,
+    const index_t* __restrict__ ind,
+    int N,
+    int K,
+    int ld_src,
+    int ld_dst) {
+#if defined(CPU_CAPABILITY_AVX512)
+  const int NB = div_up(N, 16);
+  const int KB = K / 32;
+  const int K_remainder = K - KB * 32;
+
+  for (int nb = 0; nb < NB; ++nb) {
+    int nb_size = std::min(N - nb * 16, 16);
+    for (int kb = 0; kb < KB; ++kb) {
+      // handle 16x512bits each block
+      pack_vnni_Nx32<scalar_t, index_t>(
+          /*    dst */ dst + ((kb * 32) >> 1) * ld_dst * 2 + nb * 16 * 2,
+          /*    src */ src + kb * 32 + (is_indexed ? 0 : nb * 16 * ld_src),
+          /*    ind */ is_indexed ? ind + nb * 16 : nullptr,
+          /*      N */ nb_size,
+          /* ld_src */ ld_src,
+          /* ld_dst */ ld_dst);
+    }
+    if (K_remainder > 0) {
+      pack_vnni_N_remainder<scalar_t, index_t>(
+          /*    dst */ dst + ((KB * 32) >> 1) * ld_dst * 2 + nb * 16 * 2,
+          /*    src */ src + KB * 32 + (is_indexed ? 0 : nb * 16 * ld_src),
+          /*    ind */ is_indexed ? ind + nb * 16 : nullptr,
+          /*      N */ nb_size,
+          /*      K */ K_remainder,
+          /* ld_src */ ld_src,
+          /* ld_dst */ ld_dst);
+    }
+  }
+#else
+  for (int n = 0; n < N; ++n) {
+    index_t index = get_index(ind, n);
+    for (int k = 0; k < K / 2; ++k) {
+      for (int d = 0; d < 2; ++d) {
+        dst[k * ld_dst * 2 + n * 2 + d] = src[index * ld_src + k * 2 + d];
+      }
+    }
+  }
+#endif
+}
+
+template <typename scalar_t>
+void pack_vnni(scalar_t* __restrict__ dst, const scalar_t* __restrict__ src, int N, int K, int ld_src, int ld_dst) {
+  pack_vnni<scalar_t, int32_t, false>(dst, src, nullptr, N, K, ld_src, ld_dst);
+}
+
+template <typename scalar_t, typename index_t>
+void pack_vnni(
+    scalar_t* __restrict__ dst,
+    const scalar_t* __restrict__ src,
+    const index_t* __restrict__ ind,
+    int N,
+    int K,
+    int ld_src,
+    int ld_dst) {
+  assert(ind != nullptr);
+  pack_vnni<scalar_t, index_t, true>(dst, src, ind, N, K, ld_src, ld_dst);
+}
+
+// convert to vnni format
+// from [K/2, 2, N] to [K/2, N, 2] for bfloat16 and float16
+template <typename scalar_t, typename index_t, bool is_indexed>
+void pack_vnni2(
+    scalar_t* __restrict__ dst,
+    const scalar_t* __restrict__ src,
+    const index_t* __restrict__ ind,
+    int K,
+    int N,
+    int ld_src,
+    int ld_dst) {
+#if defined(CPU_CAPABILITY_AVX512)
+  const int KB = div_up(K, 2);
+  const int NB = N / 32;
+  const int N_remainder = N - NB * 32;
+
+  for (int kb = 0; kb < KB; ++kb) {
+    int kb_size = std::min(K - kb * 2, 2);
+    for (int nb = 0; nb < NB; ++nb) {
+      // handle 2x512bits each block
+      pack_vnni_Kx32<scalar_t, index_t>(
+          /*    dst */ dst + ((kb * 2) >> 1) * ld_dst * 2 + nb * 32 * 2,
+          /*    src */ src + (is_indexed ? 0 : kb * 2 * ld_src) + nb * 32,
+          /*    ind */ is_indexed ? ind + kb * 2 : nullptr,
+          /*      K */ kb_size,
+          /* ld_src */ ld_src,
+          /* ld_dst */ ld_dst);
+    }
+    if (N_remainder > 0) {
+      pack_vnni_K_remainder(
+          /*    dst */ dst + ((kb * 2) >> 1) * ld_dst * 2 + NB * 32 * 2,
+          /*    src */ src + (is_indexed ? 0 : kb * 2 * ld_src) + NB * 32,
+          /*    ind */ is_indexed ? ind + kb * 2 : nullptr,
+          /*      K */ kb_size,
+          /*      N */ N_remainder,
+          /* ld_src */ ld_src,
+          /* ld_dst */ ld_dst);
+    }
+  }
+#else
+  int k = 0;
+  for (; k < (K >> 1) * 2; k += 2) {
+    index_t index0 = get_index(ind, k + 0);
+    index_t index1 = get_index(ind, k + 1);
+    for (int n = 0; n < N; ++n) {
+      dst[(k >> 1) * ld_dst * 2 + n * 2 + 0] = src[index0 * ld_src + n];
+      dst[(k >> 1) * ld_dst * 2 + n * 2 + 1] = src[index1 * ld_src + n];
+    }
+  }
+  if (K % 2 != 0) {
+    index_t index = get_index(ind, K - 1);
+    for (int n = 0; n < N; ++n) {
+      dst[(K >> 1) * ld_dst * 2 + n * 2 + 0] = src[index * ld_src + n];
+      dst[(K >> 1) * ld_dst * 2 + n * 2 + 1] = 0;
+    }
+    k += 2;
+  }
+#endif
+}
+
+template <typename scalar_t>
+void pack_vnni2(scalar_t* __restrict__ dst, const scalar_t* __restrict__ src, int K, int N, int ld_src, int ld_dst) {
+  pack_vnni2<scalar_t, int32_t, false>(dst, src, nullptr, K, N, ld_src, ld_dst);
+}
+
+template <typename scalar_t, typename index_t>
+void pack_vnni2(
+    scalar_t* __restrict__ dst,
+    const scalar_t* __restrict__ src,
+    const index_t* __restrict__ ind,
+    int K,
+    int N,
+    int ld_src,
+    int ld_dst) {
+  assert(ind != nullptr);
+  pack_vnni2<scalar_t, index_t, true>(dst, src, ind, K, N, ld_src, ld_dst);
+}
+
+}  // anonymous namespace
@@ -56,6 +56,8 @@ void shm_send_tensor_list(int64_t handle,

 std::vector<torch::Tensor> shm_recv_tensor_list(int64_t handle, int64_t src);

+// SGL CPU kernels
+
 at::Tensor weight_packed_linear(at::Tensor& mat1, at::Tensor& mat2,
                                const std::optional<at::Tensor>& bias,
                                bool is_vnni);
@@ -65,21 +67,32 @@ at::Tensor convert_weight_packed(at::Tensor& weight);
 at::Tensor fused_experts_cpu(
    at::Tensor& hidden_states, at::Tensor& w1, at::Tensor& w2,
    at::Tensor& topk_weights, at::Tensor& topk_ids, bool inplace,
-    bool use_int8_w8a8, bool use_fp8_w8a16,
-    const std::optional<at::Tensor>& w1_scale,
+    int64_t moe_comp_method, const std::optional<at::Tensor>& w1_scale,
    const std::optional<at::Tensor>& w2_scale,
-    const std::optional<std::vector<int64_t>> block_size,
-    const std::optional<at::Tensor>& a1_scale,
-    const std::optional<at::Tensor>& a2_scale, bool is_vnni);
+    const std::optional<at::Tensor>& w1_zero,
+    const std::optional<at::Tensor>& w2_zero,
+    const std::optional<std::vector<int64_t>> block_size, bool is_vnni);

 at::Tensor int8_scaled_mm_with_quant(at::Tensor& mat1, at::Tensor& mat2,
                                     at::Tensor& scales2,
                                     const std::optional<at::Tensor>& bias,
                                     at::ScalarType out_dtype, bool is_vnni);

+// Adapted from sglang: FP8 W8A16 kernel
+at::Tensor fp8_scaled_mm_cpu(at::Tensor& mat1, at::Tensor& mat2,
+                             at::Tensor& scales2,
+                             std::vector<int64_t> block_size,
+                             const std::optional<at::Tensor>& bias,
+                             at::ScalarType out_dtype, bool is_vnni);
+
 // Adapted from sglang: INT4 W4A8 kernels
 std::tuple<at::Tensor, at::Tensor, at::Tensor> convert_weight_packed_scale_zp(
-    at::Tensor qweight, at::Tensor qzeros, at::Tensor scales);
+    at::Tensor qweight,  // awq: (*, K, N / 8)  ||  gptq: (*, K / 8, N) , int32
+    at::Tensor qzeros,   // awq: (*, K / group_size, N / 8) ||  gptq: (*, K /
+                         // group_size, N / 8) , int32
+    at::Tensor scales,   // awq: (*, K / group_size, N) ||  gptq: (*, K /
+                         // group_size, N) , bfloat16
+    int64_t quant_method_4bit);

 at::Tensor int4_scaled_mm_cpu(at::Tensor& x, at::Tensor& w, at::Tensor& w_zeros,
                              at::Tensor& w_scales,
@@ -271,8 +284,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  ops.impl("rotary_embedding", torch::kCPU, &rotary_embedding);

  // Quantization
-#if defined(__AVX512F__) || (defined(__aarch64__) && !defined(__APPLE__)) || \
-    defined(__powerpc64__)
+#if defined(__AVX512F__) || defined(__AVX2__) || \
+    (defined(__aarch64__) && !defined(__APPLE__)) || defined(__powerpc64__)
  // Helper function to release oneDNN handlers
  ops.def("release_dnnl_matmul_handler(int handler) -> ()",
          &release_dnnl_matmul_handler);
@@ -353,10 +366,10 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  ops.def("convert_weight_packed(Tensor! weight) -> Tensor");
  ops.impl("convert_weight_packed", torch::kCPU, &convert_weight_packed);
  ops.def(
-      "fused_experts_cpu(Tensor! hidden_states, Tensor w1, Tensor w2, Tensor "
-      "topk_weights, Tensor topk_ids, bool inplace, bool use_int8_w8a8, bool "
-      "use_fp8_w8a16, Tensor? w1_scale, Tensor? w2_scale, SymInt[]? "
-      "block_size, Tensor? a1_scale, Tensor? a2_scale, bool is_vnni) -> "
+      "fused_experts_cpu(Tensor hidden_states, Tensor w1, Tensor w2, Tensor "
+      "topk_weights, Tensor topk_ids, bool "
+      "inplace, int moe_comp_method, Tensor? w1_scale, Tensor? w2_scale, "
+      "Tensor? w1_zero, Tensor? w2_zero, int[]? block_size, bool is_vnni) -> "
      "Tensor");
  ops.impl("fused_experts_cpu", torch::kCPU, &fused_experts_cpu);
  ops.def(
@@ -367,8 +380,9 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {

  // Adapted from sglang: INT4 W4A8 kernels
  ops.def(
-      "convert_weight_packed_scale_zp(Tensor qweight, Tensor qzeros, "
-      "Tensor scales) -> (Tensor, Tensor, Tensor)");
+      "convert_weight_packed_scale_zp(Tensor weight, Tensor qzeros, Tensor "
+      "scales, int quant_method_4bit) -> (Tensor, "
+      "Tensor, Tensor)");
  ops.impl("convert_weight_packed_scale_zp", torch::kCPU,
           &convert_weight_packed_scale_zp);

@@ -376,6 +390,13 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
      "int4_scaled_mm_cpu(Tensor(a0!) x, Tensor(a1!) w, Tensor(a2!) w_zeros, "
      "Tensor(a3!) w_scales, Tensor? bias) -> Tensor");
  ops.impl("int4_scaled_mm_cpu", torch::kCPU, &int4_scaled_mm_cpu);
+
+  // Adapted from sglang: FP8 W8A16 kernel
+  ops.def(
+      "fp8_scaled_mm_cpu(Tensor(a0!) mat1, Tensor(a1!) mat2, Tensor(a2!) "
+      "scales2, SymInt[] block_size, Tensor? bias, ScalarType out_dtype, "
+      "bool is_vnni) -> Tensor");
+  ops.impl("fp8_scaled_mm_cpu", torch::kCPU, &fp8_scaled_mm_cpu);
 #endif

  // CPU attention kernels
@@ -54,7 +54,7 @@ struct Counter {
 };

 inline int64_t get_available_l2_size() {
-#if defined(__s390x__)
+#if defined(__s390x__) || defined(__powerpc__)
  static int64_t size = []() {
    uint32_t l2_cache_size = 0;
    auto caps = at::cpu::get_cpu_capabilities();
@@ -232,28 +232,6 @@ void unmap_and_release(unsigned long long device, ssize_t size,
    }
  }

-  // ROCm workaround: hipMemRelease does not return physical VRAM to the
-  // free pool while the virtual-address reservation is still held.
-  // Cycling cuMemAddressFree → cuMemAddressReserve (at the same address)
-  // forces the driver to actually release the physical pages while keeping
-  // the same VA available for a later create_and_map.
-  if (first_error == no_error) {
-    first_error = cuMemAddressFree(d_mem, size);
-    if (first_error == no_error) {
-      CUdeviceptr d_mem_new = 0;
-      first_error = cuMemAddressReserve(&d_mem_new, size, 0, d_mem, 0);
-      if (first_error == no_error && d_mem_new != d_mem) {
-        cuMemAddressFree(d_mem_new, size);
-        snprintf(error_msg, sizeof(error_msg),
-                 "ROCm: VA re-reserve got %p instead of %p", (void*)d_mem_new,
-                 (void*)d_mem);
-        error_code = CUresult(1);
-        std::cerr << error_msg << std::endl;
-        return;
-      }
-    }
-  }
-
  if (first_error != no_error) {
    CUDA_CHECK(first_error);
  }
@@ -29,7 +29,11 @@
 */

 #include <cmath>
-#include <cuda_fp8.h>
+#ifndef USE_ROCM
+  #include <cuda_fp8.h>
+#else
+  #include <hip/hip_fp8.h>
+#endif
 #include <cuda_runtime.h>
 #include <type_traits>

@@ -42,7 +46,23 @@
 #include "type_convert.cuh"

 #ifndef FINAL_MASK
-  #define FINAL_MASK 0xffffffffu
+  #ifdef USE_ROCM
+    #define FINAL_MASK 0xffffffffffffffffULL
+  #else
+    #define FINAL_MASK 0xffffffffu
+  #endif
+#endif
+
+#ifdef USE_ROCM
+// ROCm-compatible FP8 conversion helpers
+__device__ __forceinline__ uint8_t rocm_cvt_float_to_fp8_e4m3(float val) {
+  #if defined(HIP_FP8_TYPE_OCP)
+  __hip_fp8_e4m3 fp8_val(val);
+  #else
+  __hip_fp8_e4m3_fnuz fp8_val(val);
+  #endif
+  return reinterpret_cast<uint8_t&>(fp8_val);
+}
 #endif

 namespace vllm {
@@ -314,9 +334,13 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
      for (int i = 0; i < kElemsPerLane; i++) {
        float scaled = elements[i] * inv_scale;
        scaled = fminf(fmaxf(scaled, -kFp8Max), kFp8Max);
+#ifndef USE_ROCM
        __nv_fp8_storage_t s =
            __nv_cvt_float_to_fp8(scaled, __NV_SATFINITE, __NV_E4M3);
        out_bytes[i] = static_cast<uint8_t>(s);
+#else
+      out_bytes[i] = rocm_cvt_float_to_fp8_e4m3(scaled);
+#endif
      }
      // One 16-byte STG per lane.
      *reinterpret_cast<uint4*>(token_fp8_ptr + dim_base) =
@@ -384,6 +408,7 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
  // PDL: enable programmatic stream serialization whenever the hardware
  // supports it (SM90+).  On pre-Hopper GPUs the attribute is unavailable,
  // so leave numAttrs = 0 and launch as a regular kernel.
+#ifndef USE_ROCM
  static int const sm_version = getSMVersion();
  // Host-side guard: the device kernel body is compiled as a no-op for
  // bf16 on pre-Ampere (sm_70/sm_75) because _typeConvert<BFloat16> is
@@ -410,6 +435,15 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
      q_inout, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, eps,
      num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size,
      kv_block_stride);
+#else
+  // ROCm: use standard kernel launch syntax (no PDL/stream serialization)
+  // clang-format off
+  fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<scalar_t_in>
+      <<<grid, kBlockSize, 0, stream>>>(
+          q_inout, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache,
+          eps, num_tokens_full, num_tokens_insert, num_heads_q,
+          cache_block_size, kv_block_stride);
+#endif
 }

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

 namespace vllm {

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

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

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

 namespace vllm {
@@ -14,9 +15,17 @@ void cutlass_scaled_mm_sm120_fp8(
    STD_TORCH_CHECK(bias->scalar_type() == out.scalar_type(),
                    "currently bias dtype must match output dtype ",
                    out.scalar_type());
+    if (vllm_is_batch_invariant()) {
+      return cutlass_scaled_mm_sm120_fp8_batch_invariant_epilogue<
+          c3x::ScaledEpilogueBias>(out, a, b, a_scales, b_scales, *bias);
+    }
    return cutlass_scaled_mm_sm120_fp8_epilogue<c3x::ScaledEpilogueBias>(
        out, a, b, a_scales, b_scales, *bias);
  } else {
+    if (vllm_is_batch_invariant()) {
+      return cutlass_scaled_mm_sm120_fp8_batch_invariant_epilogue<
+          c3x::ScaledEpilogue>(out, a, b, a_scales, b_scales);
+    }
    return cutlass_scaled_mm_sm120_fp8_epilogue<c3x::ScaledEpilogue>(
        out, a, b, a_scales, b_scales);
  }
@@ -179,6 +179,26 @@ inline void cutlass_gemm_sm120_fp8_dispatch(torch::stable::Tensor& out,
      out, a, b, std::forward<EpilogueArgs>(args)...);
 }

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

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

 namespace vllm {

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

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

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

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

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

+template <typename InType, typename OutType,
+          template <typename, typename> typename Epilogue,
+          typename... EpilogueArgs>
+inline void cutlass_gemm_sm89_fp8_batch_invariant_dispatch(
+    torch::stable::Tensor& out, torch::stable::Tensor const& a,
+    torch::stable::Tensor const& b, EpilogueArgs&&... args) {
+  static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
+  STD_TORCH_CHECK(a.scalar_type() ==
+                  torch::headeronly::ScalarType::Float8_e4m3fn);
+  STD_TORCH_CHECK(b.scalar_type() ==
+                  torch::headeronly::ScalarType::Float8_e4m3fn);
+
+  // keep the CUTLASS config independent of M for batch invariance
+  return sm89_fp8_config_M64::dispatch<InType, OutType, Epilogue>(
+      out, a, b, std::forward<EpilogueArgs>(args)...);
+}
+
+template <template <typename, typename> typename Epilogue,
+          typename... EpilogueArgs>
+void cutlass_scaled_mm_sm89_fp8_batch_invariant_epilogue(
+    torch::stable::Tensor& out, torch::stable::Tensor const& a,
+    torch::stable::Tensor const& b, EpilogueArgs&&... epilogue_args) {
+  STD_TORCH_CHECK(a.scalar_type() ==
+                  torch::headeronly::ScalarType::Float8_e4m3fn);
+  STD_TORCH_CHECK(b.scalar_type() ==
+                  torch::headeronly::ScalarType::Float8_e4m3fn);
+
+  if (out.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
+    return cutlass_gemm_sm89_fp8_batch_invariant_dispatch<
+        cutlass::float_e4m3_t, cutlass::bfloat16_t, Epilogue>(
+        out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
+  } else {
+    STD_TORCH_CHECK(out.scalar_type() == torch::headeronly::ScalarType::Half);
+    return cutlass_gemm_sm89_fp8_batch_invariant_dispatch<
+        cutlass::float_e4m3_t, cutlass::half_t, Epilogue>(
+        out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
+  }
+}
+
 }  // namespace vllm
@@ -236,17 +236,41 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input,
 #undef LAUNCH_KERNEL
 }

-template <typename T, typename DST_DTYPE>
-__global__ void per_token_group_quant_8bit_packed_kernel(
+// Register-resident fast path for group_size==128.
+//
+// Each thread holds 16 source elements (32 B = uint4 x 2) in registers across
+// the absmax reduce -> scale compute -> quantize pipeline. No shared memory.
+// UE8M0 scale extracted via bit math (bit-exact with exp2f(ceilf(log2f))).
+//
+// Loads two contiguous uint4s (16 B + 16 B = 32 B) per thread; on Blackwell
+// nvcc fuses these into a single 256-bit LDG.E.256.
+//
+// Constraints: GROUP_SIZE % (THREADS_PER_GROUP * VEC_SIZE) == 0; for
+// THREADS_PER_GROUP=8 and bf16/fp16 (VEC_SIZE=16), this means GROUP_SIZE=128.
+template <typename T, typename DST_DTYPE, int GROUP_SIZE>
+__global__ void per_token_group_quant_8bit_packed_register_kernel(
    const T* __restrict__ input, void* __restrict__ output_q,
-    unsigned int* __restrict__ output_s_packed, const int group_size,
-    const int num_groups_padded, const int groups_per_block,
-    const int padded_groups_per_row, const int groups_per_row, const int mn,
-    const int tma_aligned_mn, const int num_scale_elems, const float eps,
+    unsigned int* __restrict__ output_s_packed, const int64_t num_groups_padded,
+    const int groups_per_block, const int padded_groups_per_row,
+    const int groups_per_row, const int mn, const int output_q_mn_extent,
+    const int tma_aligned_mn, const int64_t num_scale_elems, const float eps,
    const float min_8bit, const float max_8bit) {
-  const int threads_per_group = 16;
-  const int64_t local_group_id = threadIdx.x / threads_per_group;
-  const int lane_id = threadIdx.x % threads_per_group;
+  static_assert(GROUP_SIZE == 128, "fast path supports GROUP_SIZE==128");
+  constexpr int THREADS_PER_GROUP = 8;
+  constexpr int VEC_SIZE = 32 / sizeof(T);  // 16 for bf16/fp16
+  static_assert(GROUP_SIZE == THREADS_PER_GROUP * VEC_SIZE,
+                "GROUP_SIZE must equal THREADS_PER_GROUP * VEC_SIZE");
+  // Each group's 8 threads must live in a single warp octet so the
+  // 0xffu << (threadIdx.x & 24u) shuffle mask selects exactly the lanes
+  // that share a group. Requires 32 % THREADS_PER_GROUP == 0 and the host
+  // to launch num_threads as a multiple of THREADS_PER_GROUP (which it does
+  // via num_threads = groups_per_block * THREADS_PER_GROUP).
+  static_assert(32 % THREADS_PER_GROUP == 0,
+                "THREADS_PER_GROUP must divide warp size for the shuffle "
+                "mask to be valid");
+
+  const int local_group_id = threadIdx.x / THREADS_PER_GROUP;
+  const int lane_id = threadIdx.x % THREADS_PER_GROUP;

  const int64_t block_group_id = blockIdx.x * groups_per_block;
  const int64_t global_group_id = block_group_id + local_group_id;
@@ -254,141 +278,207 @@ __global__ void per_token_group_quant_8bit_packed_kernel(
    return;
  }

-  // map flat group id to 2D indices (mn_idx, sf_k_idx)
  const int sf_k_idx =
      static_cast<int>(global_group_id % padded_groups_per_row);
  const int mn_idx = static_cast<int>(global_group_id / padded_groups_per_row);
-
-  // whether it is a valid group (not padding)
  const bool is_valid_group = (mn_idx < mn) && (sf_k_idx < groups_per_row);

-  // shared memory to cache each group's data to avoid double DRAM reads.
-  extern __shared__ __align__(16) char smem_raw[];
-  T* smem = reinterpret_cast<T*>(smem_raw);
-  T* smem_group = smem + local_group_id * group_size;
-
-  // compute scale for valid groups
-  float y_s = 0.f;
+  // Load 16 input elements (32 B) into registers as two adjacent uint4
+  // loads. nvcc keeps these as 2x LDG.E.128 on sm_100; the per-thread cost
+  // is dominated by HBM bandwidth at large MN, so a fused 256-bit load via
+  // inline PTX gave no measurable speedup.
+  // alignas(16) is required so the uint4* reinterpret_cast below is
+  // well-defined for T == bf16/fp16 (default alignof is 2).
+  alignas(16) T regs[VEC_SIZE];
+  float local_absmax = eps;
  if (is_valid_group) {
    const T* group_input =
-        input + static_cast<int64_t>(mn_idx) * groups_per_row * group_size +
-        sf_k_idx * group_size;
-    y_s = ComputeGroupScale<T, true>(group_input, smem_group, group_size,
-                                     lane_id, threads_per_group, eps, max_8bit);
+        input + static_cast<int64_t>(mn_idx) * groups_per_row * GROUP_SIZE +
+        sf_k_idx * GROUP_SIZE + lane_id * VEC_SIZE;
+    uint4* dst = reinterpret_cast<uint4*>(&regs[0]);
+    const uint4* src = reinterpret_cast<const uint4*>(group_input);
+    dst[0] = src[0];
+    dst[1] = src[1];
+#pragma unroll
+    for (int i = 0; i < VEC_SIZE; ++i) {
+      float v = fabsf(static_cast<float>(regs[i]));
+      local_absmax = fmaxf(local_absmax, v);
+    }
  }

-  // pack 4 scales into a uint32 exponent
+  // 8-lane subgroup shuffle reduce (octet of the warp). The mask selects the
+  // 8 lanes within the warp that share a group.
+  unsigned mask = 0xffu << (threadIdx.x & 24u);
+  local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 4));
+  local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 2));
+  local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 1));
+
+  float y_s = local_absmax / max_8bit;
+  y_s = fmaxf(y_s, 1e-10f);
+  uint32_t bits = __float_as_uint(y_s);
+  uint32_t exp_bits = (bits >> 23) & 0xffu;
+  uint32_t mant_bits = bits & 0x7fffffu;
+  uint8_t exp_byte =
+      static_cast<uint8_t>(exp_bits + (mant_bits != 0u ? 1u : 0u));
+
+  // Lane 0 writes the packed scale byte.
  if (lane_id == 0) {
-    // each uint32 in output_s_packed stores 4 packed scales
    const int sf_k_pack_idx = sf_k_idx / 4;
    const int pos = sf_k_idx % 4;
    const int out_idx = sf_k_pack_idx * tma_aligned_mn + mn_idx;
-
    if (is_valid_group) {
-      // reinterpret the UE8M0 scale y_s as IEEE bits, extract the 8-bit
-      // exponent, and place it into the correct byte of the 32-bit word.
-      const unsigned int bits = __float_as_uint(y_s);
-      const uint8_t exponent = static_cast<uint8_t>((bits >> 23u) & 0xffu);
-      reinterpret_cast<uint8_t*>(output_s_packed)[out_idx * 4 + pos] = exponent;
+      reinterpret_cast<uint8_t*>(output_s_packed)[out_idx * 4 + pos] = exp_byte;
    } else if (out_idx < num_scale_elems) {
-      // write zero for padding groups if within bounds of output_s_packed
      reinterpret_cast<uint8_t*>(output_s_packed)[out_idx * 4 + pos] = 0;
    }
  }

-  __syncthreads();
-
-  if (is_valid_group) {
-    DST_DTYPE* group_output =
-        static_cast<DST_DTYPE*>(output_q) +
-        static_cast<int64_t>(mn_idx) * groups_per_row * group_size +
-        sf_k_idx * group_size;
-    QuantizeGroup<T, DST_DTYPE>(smem_group, group_output, group_size, lane_id,
-                                threads_per_group, y_s, min_8bit, max_8bit);
+  // For padded mn rows that fall within output_q's allocated extent, write
+  // a uint4 of zeros to keep the buffer clean for downstream TMA loads.
+  // Skip writes for sf_k padding (those positions don't exist in output_q).
+  if (!is_valid_group) {
+    if (sf_k_idx < groups_per_row && mn_idx >= mn &&
+        mn_idx < output_q_mn_extent) {
+      DST_DTYPE* group_output =
+          static_cast<DST_DTYPE*>(output_q) +
+          static_cast<int64_t>(mn_idx) * groups_per_row * GROUP_SIZE +
+          sf_k_idx * GROUP_SIZE + lane_id * VEC_SIZE;
+      *reinterpret_cast<uint4*>(group_output) = make_uint4(0, 0, 0, 0);
+    }
+    return;
  }
+
+  // Reconstruct y_s as a power-of-2 float and use its reciprocal.
+  float y_s_q = __uint_as_float(static_cast<uint32_t>(exp_byte) << 23);
+  float inv_y = 1.0f / y_s_q;
+
+  // Quantize and pack into 16 fp8/int8 bytes (= uint4). VEC_SIZE==16 so we
+  // fill four 32-bit words, four bytes each.
+  uint32_t packed_lo = 0;
+  uint32_t packed_lo_hi = 0;
+  uint32_t packed_hi_lo = 0;
+  uint32_t packed_hi = 0;
+#pragma unroll
+  for (int i = 0; i < VEC_SIZE; ++i) {
+    float q =
+        fminf(fmaxf(static_cast<float>(regs[i]) * inv_y, min_8bit), max_8bit);
+    DST_DTYPE qb = DST_DTYPE(q);
+    uint8_t byte = *reinterpret_cast<uint8_t*>(&qb);
+    const int shift = (i & 3) * 8;
+    if (i < 4) {
+      packed_lo |= static_cast<uint32_t>(byte) << shift;
+    } else if (i < 8) {
+      packed_lo_hi |= static_cast<uint32_t>(byte) << shift;
+    } else if (i < 12) {
+      packed_hi_lo |= static_cast<uint32_t>(byte) << shift;
+    } else {
+      packed_hi |= static_cast<uint32_t>(byte) << shift;
+    }
+  }
+
+  uint4 packed_out =
+      make_uint4(packed_lo, packed_lo_hi, packed_hi_lo, packed_hi);
+  DST_DTYPE* group_output =
+      static_cast<DST_DTYPE*>(output_q) +
+      static_cast<int64_t>(mn_idx) * groups_per_row * GROUP_SIZE +
+      sf_k_idx * GROUP_SIZE + lane_id * VEC_SIZE;
+  *reinterpret_cast<uint4*>(group_output) = packed_out;
 }

+// Public entry point: register-resident packed quant kernel.
+// Constraints: group_size == 128 and bf16/fp16 input.
 void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
                                       torch::stable::Tensor& output_q,
                                       torch::stable::Tensor& output_s_packed,
                                       int64_t group_size, double eps,
                                       double min_8bit, double max_8bit) {
+  STD_TORCH_CHECK(group_size == 128,
+                  "per_token_group_quant_8bit_packed only supports "
+                  "group_size==128, got ",
+                  group_size, ".");
+  const auto in_dtype = input.scalar_type();
+  STD_TORCH_CHECK(
+      in_dtype == torch::headeronly::ScalarType::Half ||
+          in_dtype == torch::headeronly::ScalarType::BFloat16,
+      "per_token_group_quant_8bit_packed only supports bf16/fp16 input.");
+
  STD_TORCH_CHECK(input.is_contiguous());
  STD_TORCH_CHECK(output_q.is_contiguous());

  const int64_t k = input.size(-1);
-  STD_TORCH_CHECK(k % group_size == 0, "Last dimension (", k,
-                  ") must be divisible by group_size (", group_size, ").");
+  STD_TORCH_CHECK(k % group_size == 0, "input last dim k=", k,
+                  " is not divisible by group_size=", group_size, ".");

  const int64_t mn = input.numel() / k;
  const int64_t groups_per_row = k / group_size;
-
-  STD_TORCH_CHECK(output_s_packed.dim() == 2,
-                  "output_s_packed must be 2D, got dim=", output_s_packed.dim(),
-                  ".");
-
  const int64_t k_num_packed_sfk = (groups_per_row + 3) / 4;
  const int64_t tma_aligned_mn = ((mn + 3) / 4) * 4;

+  // output_q may be allocated with extra padded mn rows (e.g.,
+  // (tma_aligned_mn, k)) so the kernel can zero-fill them in-line and the
+  // caller can use torch.empty instead of torch.zeros. The grid only covers
+  // up to tma_aligned_mn, so we cap the extent there.
+  const int64_t output_q_mn_actual = output_q.numel() / k;
+  STD_TORCH_CHECK(output_q_mn_actual >= mn,
+                  "output_q must have at least mn rows; got ",
+                  output_q_mn_actual, " rows for mn=", mn, ".");
+  const int64_t output_q_mn_extent =
+      output_q_mn_actual < tma_aligned_mn ? output_q_mn_actual : tma_aligned_mn;
+
  STD_TORCH_CHECK(
      output_s_packed.scalar_type() == torch::headeronly::ScalarType::Int,
-      "output_s_packed must have dtype int32 for UE8M0-packed scales.");
-  // DeepGEMM expects SFA scales in MN-major form with shape
-  // [mn, ceil_div(K, 128 * 4)] and TMA-aligned stride on the last
-  // dimension.
+      "output_s_packed must be int32 for UE8M0-packed scales.");
  STD_TORCH_CHECK(output_s_packed.size(0) == mn &&
                      output_s_packed.size(1) == k_num_packed_sfk,
                  "output_s_packed shape must be [", mn, ", ", k_num_packed_sfk,
-                  "], but got [", output_s_packed.size(0), ", ",
+                  "]; got [", output_s_packed.size(0), ", ",
                  output_s_packed.size(1), "].");
-  // Verify column-major TMA-aligned layout
  STD_TORCH_CHECK(output_s_packed.stride(0) == 1 &&
                      output_s_packed.stride(1) == tma_aligned_mn,
-                  "output_s_packed must have strides [1, ", tma_aligned_mn,
-                  "], but got [", output_s_packed.stride(0), ", ",
+                  "output_s_packed strides must be [1, ", tma_aligned_mn,
+                  "]; got [", output_s_packed.stride(0), ", ",
                  output_s_packed.stride(1), "].");

  cudaStream_t stream = get_current_cuda_stream();

-  constexpr int THREADS_PER_GROUP = 16;
-
-  // Expand the grid to cover MN and K padding so every byte in
-  // output_s_packed is written (padding bytes get zeroed by the kernel).
+  constexpr int THREADS_PER_GROUP = 8;
  const int64_t padded_groups_per_row = k_num_packed_sfk * 4;
  const int64_t num_groups_padded = tma_aligned_mn * padded_groups_per_row;
-  // Number of elements in output_s_packed.
  const int64_t num_scale_elems = mn + (k_num_packed_sfk - 1) * tma_aligned_mn;
-
  const int groups_per_block = GetGroupsPerBlock(num_groups_padded);

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

-#define LAUNCH_PACKED_KERNEL(T, DST_DTYPE)                                     \
-  do {                                                                         \
-    dim3 grid(num_blocks);                                                     \
-    dim3 block(num_threads);                                                   \
-    size_t smem_bytes =                                                        \
-        static_cast<size_t>(groups_per_block) * group_size * sizeof(T);        \
-    per_token_group_quant_8bit_packed_kernel<T, DST_DTYPE>                     \
-        <<<grid, block, smem_bytes, stream>>>(                                 \
-            static_cast<const T*>(input.data_ptr()), output_q.data_ptr(),      \
-            reinterpret_cast<unsigned int*>(output_s_packed.data_ptr()),       \
-            static_cast<int>(group_size), static_cast<int>(num_groups_padded), \
-            groups_per_block, static_cast<int>(padded_groups_per_row),         \
-            static_cast<int>(groups_per_row), static_cast<int>(mn),            \
-            static_cast<int>(tma_aligned_mn),                                  \
-            static_cast<int>(num_scale_elems), static_cast<float>(eps),        \
-            static_cast<float>(min_8bit), static_cast<float>(max_8bit));       \
+#define LAUNCH_REG_KERNEL(T, DST_DTYPE)                                   \
+  do {                                                                    \
+    dim3 grid(static_cast<unsigned int>(num_blocks));                     \
+    dim3 block(num_threads);                                              \
+    per_token_group_quant_8bit_packed_register_kernel<T, DST_DTYPE, 128>  \
+        <<<grid, block, 0, stream>>>(                                     \
+            static_cast<const T*>(input.data_ptr()), output_q.data_ptr(), \
+            reinterpret_cast<unsigned int*>(output_s_packed.data_ptr()),  \
+            num_groups_padded, groups_per_block,                          \
+            static_cast<int>(padded_groups_per_row),                      \
+            static_cast<int>(groups_per_row), static_cast<int>(mn),       \
+            static_cast<int>(output_q_mn_extent),                         \
+            static_cast<int>(tma_aligned_mn), num_scale_elems,            \
+            static_cast<float>(eps), static_cast<float>(min_8bit),        \
+            static_cast<float>(max_8bit));                                \
  } while (0)

-  VLLM_STABLE_DISPATCH_FLOATING_TYPES(
-      input.scalar_type(), "per_token_group_quant_8bit_packed", ([&] {
+  VLLM_STABLE_DISPATCH_HALF_TYPES(
+      input.scalar_type(), "per_token_group_quant_8bit_packed_register", ([&] {
        if (dst_type == torch::headeronly::ScalarType::Float8_e4m3fn) {
-          LAUNCH_PACKED_KERNEL(scalar_t, __nv_fp8_e4m3);
+          LAUNCH_REG_KERNEL(scalar_t, __nv_fp8_e4m3);
        } else if (dst_type == torch::headeronly::ScalarType::Char) {
-          LAUNCH_PACKED_KERNEL(scalar_t, int8_t);
+          LAUNCH_REG_KERNEL(scalar_t, int8_t);
        } else {
          STD_TORCH_CHECK(
              false,
@@ -397,7 +487,7 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
        }
      }));

-#undef LAUNCH_PACKED_KERNEL
+#undef LAUNCH_REG_KERNEL
 }

 void per_token_group_quant_fp8(const torch::stable::Tensor& input,
@@ -8,3 +8,13 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input,
                                torch::stable::Tensor& output_s,
                                int64_t group_size, double eps, double min_8bit,
                                double max_8bit, bool scale_ue8m0 = false);
+
+// Public op: register-resident packed quant for the DeepGEMM Blackwell path.
+// Restricted to group_size == 128 and bf16/fp16 input; other configurations
+// raise STD_TORCH_CHECK. The legacy shared-memory fallback was removed because
+// no production caller (deep_gemm_moe / input_quant_fp8) uses other shapes.
+void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
+                                       torch::stable::Tensor& output_q,
+                                       torch::stable::Tensor& output_s_packed,
+                                       int64_t group_size, double eps,
+                                       double min_8bit, double max_8bit);
@@ -67,10 +67,6 @@ void shuffle_rows(const torch::Tensor& input_tensor,
                  torch::Tensor& output_tensor);

 #ifndef USE_ROCM
-// cuBLAS bf16 x bf16 -> fp32 router GEMM (fallback for non-SM90 / batch > 16)
-torch::Tensor router_gemm_bf16_fp32(torch::Tensor const& input,
-                                    torch::Tensor const& weight);
-
 // DeepSeek V3 optimized router GEMM kernel for SM90+
 // Computes output = mat_a @ mat_b.T where:
 //   mat_a: [num_tokens, hidden_dim] in bf16
@@ -1,52 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-// bf16 x bf16 -> fp32 router GEMM via cuBLAS.
-// Uses CUBLAS_COMPUTE_32F so bf16 operands accumulate into fp32,
-// matching TRT-LLM's cuBLAS fallback behaviour in dsv3RouterGemmOp.
-
-#include <torch/all.h>
-#include <ATen/cuda/CUDAContext.h>
-#include <cublas_v2.h>
-
-// cuBLAS column-major math for row-major PyTorch tensors:
-//   weight[N,K]_row  lda=K  -> cuBLAS sees (K,N) col-major; CUBLAS_OP_T ->
-//   (N,K) input[M,K]_row   ldb=K  -> cuBLAS sees (K,M) col-major; CUBLAS_OP_N
-//   -> (K,M) out[M,N]_row     ldc=N  -> cuBLAS sees (N,M) col-major (written as
-//   output^T)
-// cuBLAS: C(N,M) = weight(N,K) @ input(K,M)  =>  C^T = output[M,N]
-// params: m=N, n=M, k=K, lda=K (weight), ldb=K (input), ldc=N (output)
-
-torch::Tensor router_gemm_bf16_fp32(torch::Tensor const& input,
-                                    torch::Tensor const& weight) {
-  TORCH_CHECK(input.dtype() == torch::kBFloat16,
-              "router_gemm_bf16_fp32: input must be bfloat16");
-  TORCH_CHECK(weight.dtype() == torch::kBFloat16,
-              "router_gemm_bf16_fp32: weight must be bfloat16");
-  TORCH_CHECK(input.dim() == 2 && weight.dim() == 2,
-              "router_gemm_bf16_fp32: input and weight must be 2-D");
-  TORCH_CHECK(input.size(1) == weight.size(1),
-              "router_gemm_bf16_fp32: inner dimensions must match");
-
-  int64_t const M = input.size(0);
-  int64_t const N = weight.size(0);
-  int64_t const K = input.size(1);
-
-  auto out = torch::empty({M, N}, input.options().dtype(torch::kFloat32));
-
-  cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
-  TORCH_CUDABLAS_CHECK(
-      cublasSetStream(handle, at::cuda::getCurrentCUDAStream()));
-
-  float const alpha = 1.0f;
-  float const beta = 0.0f;
-
-  TORCH_CUDABLAS_CHECK(cublasGemmEx(
-      handle, CUBLAS_OP_T, CUBLAS_OP_N, static_cast<int>(N),
-      static_cast<int>(M), static_cast<int>(K), &alpha, weight.data_ptr(),
-      CUDA_R_16BF, static_cast<int>(K), input.data_ptr(), CUDA_R_16BF,
-      static_cast<int>(K), &beta, out.data_ptr(), CUDA_R_32F,
-      static_cast<int>(N), CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT));
-
-  return out;
-}
@@ -60,15 +60,6 @@ __device__ __forceinline__ float toFloat(T value) {
  }
 }

-#define FINAL_MASK 0xffffffff
-template <typename T>
-__inline__ __device__ T warpReduceSum(T val) {
-#pragma unroll
-  for (int mask = 16; mask > 0; mask >>= 1)
-    val += __shfl_xor_sync(FINAL_MASK, val, mask, 32);
-  return val;
-}
-
 // ====================== TopK softplus_sqrt things
 // ===============================

@@ -272,8 +263,14 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
      }
    }
    // Compute per-thread scale (using warp reduction when renormalizing).
+    // THREADS_PER_ROW-parameterized butterfly works for both warp sizes (32
+    // on CUDA, 64 on ROCm CDNA) and any THREADS_PER_ROW the dispatch picks.
    if (renormalize) {
-      selected_sum = warpReduceSum(selected_sum);
+#pragma unroll
+      for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) {
+        selected_sum +=
+            VLLM_SHFL_XOR_SYNC_WIDTH(selected_sum, mask, THREADS_PER_ROW);
+      }
    }
    float scale = static_cast<float>(routed_scaling_factor);
    if (renormalize) {
@@ -544,7 +541,6 @@ void topkGatingSoftplusSqrtKernelLauncher(
    const IndType* tid2eid, cudaStream_t stream) {
  static constexpr int WARPS_PER_TB = 4;
  static constexpr int BYTES_PER_LDG_POWER_OF_2 = 16;
-#ifndef USE_ROCM
  // for bfloat16 dtype, we need 4 bytes loading to make sure num_experts
  // elements can be loaded by a warp
  static constexpr int BYTES_PER_LDG_MULTIPLE_64 =
@@ -552,6 +548,19 @@ void topkGatingSoftplusSqrtKernelLauncher(
       std::is_same_v<InputType, __half>)
          ? 4
          : 8;
+  // Narrower LDG (ELTS_PER_LDG=1) used by 192/320/448/576 on ROCm WARP_SIZE=64
+  // where ELTS_PER_LDG=2 fails the EXPERTS%(ELTS_PER_LDG*WARP_SIZE)==0 check.
+  // On CUDA WARP_SIZE=32 the wider LDG already aligns, so the alias collapses
+  // back to BYTES_PER_LDG_MULTIPLE_64 — no behavioral change for CUDA.
+#ifdef USE_ROCM
+  static constexpr int BYTES_PER_LDG_MULTIPLE_64_NARROW =
+      (std::is_same_v<InputType, __nv_bfloat16> ||
+       std::is_same_v<InputType, __half>)
+          ? 2
+          : 4;
+#else
+  static constexpr int BYTES_PER_LDG_MULTIPLE_64_NARROW =
+      BYTES_PER_LDG_MULTIPLE_64;
 #endif
  switch (num_experts) {
    case 1:
@@ -584,27 +593,29 @@ void topkGatingSoftplusSqrtKernelLauncher(
    case 512:
      LAUNCH_SOFTPLUS_SQRT(512, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
      break;
-      // (CUDA only) support multiples of 64 when num_experts is not power of 2.
-      // ROCm uses WARP_SIZE 64 so 8 bytes loading won't fit for some of
-      // num_experts, alternatively we can test 4 bytes loading and enable it in
-      // future.
-#ifndef USE_ROCM
+      // Multiples of 64 that are not powers of 2. The kernel requires
+      // EXPERTS % (ELTS_PER_LDG * WARP_SIZE) == 0. With ELTS_PER_LDG=2
+      // (BYTES_PER_LDG_MULTIPLE_64), this holds for all five values on CUDA
+      // WARP_SIZE=32 but only for 384 on ROCm WARP_SIZE=64. The other four
+      // use BYTES_PER_LDG_MULTIPLE_64_NARROW (ELTS_PER_LDG=1), which
+      // satisfies the assertion for any multiple of 64 on either backend;
+      // on CUDA the narrow alias collapses back to the wider load, so CUDA
+      // behavior is unchanged.
    case 192:
-      LAUNCH_SOFTPLUS_SQRT(192, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
+      LAUNCH_SOFTPLUS_SQRT(192, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64_NARROW);
      break;
    case 320:
-      LAUNCH_SOFTPLUS_SQRT(320, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
+      LAUNCH_SOFTPLUS_SQRT(320, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64_NARROW);
      break;
    case 384:
      LAUNCH_SOFTPLUS_SQRT(384, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
      break;
    case 448:
-      LAUNCH_SOFTPLUS_SQRT(448, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
+      LAUNCH_SOFTPLUS_SQRT(448, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64_NARROW);
      break;
    case 576:
-      LAUNCH_SOFTPLUS_SQRT(576, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
+      LAUNCH_SOFTPLUS_SQRT(576, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64_NARROW);
      break;
-#endif
    default: {
      TORCH_CHECK(false, "Unsupported expert number: ", num_experts);
    }
@@ -16,14 +16,13 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
      "bias) -> ()");
  m.impl("topk_sigmoid", torch::kCUDA, &topk_sigmoid);

-#ifndef USE_ROCM
  m.def(
      "topk_softplus_sqrt(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
      "token_expert_indices, Tensor gating_output, bool renormalize, float "
      "routed_scaling_factor, Tensor? "
      "bias, Tensor? input_ids, Tensor? tid2eid) -> ()");
  m.impl("topk_softplus_sqrt", torch::kCUDA, &topk_softplus_sqrt);
-#endif
+
  // Calculate the result of moe by summing up the partial results
  // from all selected experts.
  m.def("moe_sum(Tensor input, Tensor! output) -> ()");
@@ -133,10 +132,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
      "Tensor)");
  m.impl("grouped_topk", torch::kCUDA, &grouped_topk);

-  // cuBLAS bf16 x bf16 -> fp32 router GEMM (fallback for non-SM90 / batch > 16)
-  m.def("router_gemm_bf16_fp32(Tensor input, Tensor weight) -> Tensor");
-  m.impl("router_gemm_bf16_fp32", torch::kCUDA, &router_gemm_bf16_fp32);
-
  // DeepSeek V3 optimized router GEMM for SM90+
  m.def("dsv3_router_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
  // conditionally compiled so impl registration is in source file
@@ -887,27 +887,14 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2)
  uint32_t* shared_ordered =
      reinterpret_cast<uint32_t*>(smem_raw + kFixedSmemLarge);

-  // RadixRowState for multi-CTA cooperative radix
+  // RadixRowState for multi-CTA cooperative radix.
+  // Zero-initialization is done host-side via cudaMemsetAsync in topk.cu
+  // before launch — that gives a stream-ordered happens-before edge for all
+  // CTAs, which the previous in-kernel init (CTA-0 only + intra-CTA
+  // __syncthreads) did not provide and which manifested as a race against
+  // CTA-1+'s first red_release on arrival_counter.
  RadixRowState* state = &params.row_states[group_id];

-  // -- Initialize RadixRowState (only needed if large rows exist) --
-  if (params.max_seq_len > RADIX_THRESHOLD) {
-    if (cta_in_group == 0) {
-      for (uint32_t buf = 0; buf < 3; buf++) {
-        for (uint32_t i = tx; i < RADIX; i += kThreadsPerBlock) {
-          state->histogram[buf][i] = 0;
-        }
-      }
-      if (tx == 0) {
-        state->remaining_k = 0;
-        state->prefix = 0;
-        state->arrival_counter = 0;
-        state->output_counter = 0;
-      }
-    }
-    __syncthreads();
-  }
-
  int barrier_phase = 0;
  const uint32_t total_iters = (params.num_rows + num_groups - 1) / num_groups;

@@ -1,7 +1,7 @@
 #include <cstdint>

 /* Adapted from ./csrc/quantization/gguf/mmq.cuh
-   based on ./vllm/model_executor/layers/fused_moe/fused_moe.py */
+   based on ./vllm/model_executor/layers/fused_moe/experts/triton_moe.py */
 template <typename scalar_t, int qk, int qr, int qi, bool need_sum,
          typename block_q_t, int mmq_x, int mmq_y, int nwarps,
          allocate_tiles_cuda_t allocate_tiles, load_tiles_cuda_t load_tiles,
@@ -42,8 +42,8 @@ __device__ __forceinline__ void atomicAdd_half2(half2* address, half2 val) {
 }

 //
-
-#if defined(__CUDA_ARCH__) || defined(USE_ROCM)
+#if defined(__CUDA_ARCH__) || \
+    (defined(USE_ROCM) && (HIP_VERSION_MAJOR * 100 + HIP_VERSION_MINOR) < 713)
  #if __CUDA_ARCH__ < 700 || defined(USE_ROCM)

 __device__ __forceinline__ void atomicAdd(half* address, half val) {
@@ -0,0 +1,204 @@
+#include <Python.h>
+
+extern "C" {
+
+#include <stdbool.h>
+#include <time.h>
+
+#if defined(__i386__) || defined(__x86_64__)
+  #include <cpuid.h>
+  #include <mwaitxintrin.h>
+#endif
+
+#if defined(CLOCK_MONOTONIC_RAW)
+  #define TIMEOUT_CLOCK CLOCK_MONOTONIC_RAW
+#else
+  #define TIMEOUT_CLOCK CLOCK_MONOTONIC
+#endif
+
+#define CPU_SUPPORT_NONE 0
+#define CPU_SUPPORT_MONITORX 1
+
+#define MWAITX_DEFAULT_TIMEOUT_CYCLES 1000000
+
+typedef struct {
+  unsigned int cpu_support;
+  unsigned int max_monitor_line_size;
+} spinloop_state_t;
+
+static void determine_cpu_support(spinloop_state_t* state) {
+  state->cpu_support = CPU_SUPPORT_NONE;
+  state->max_monitor_line_size = 0;
+
+#if defined(__i386__) || defined(__x86_64__)
+  unsigned int eax, ebx, ecx, edx;
+  if (__get_cpuid(0, &eax, &ebx, &ecx, &edx) == 1) {
+    // AMD CPU (possible monitorx/mwaitx support)
+    if (ebx == 0x68747541 && edx == 0x69746e65 && ecx == 0x444d4163) {
+      if (__get_cpuid(0x80000000, &eax, &ebx, &ecx, &edx) == 1 &&
+          eax >= 0x80000001 &&
+          __get_cpuid(0x80000001, &eax, &ebx, &ecx, &edx) == 1) {
+        if ((ecx & (1 << 29)) != 0) {
+          state->cpu_support = CPU_SUPPORT_MONITORX;
+        }
+      }
+    }
+  }
+
+  if (state->cpu_support == CPU_SUPPORT_MONITORX) {
+    if (__get_cpuid(5, &eax, &ebx, &ecx, &edx) == 1) {
+      state->max_monitor_line_size = ebx & 0xff;
+    }
+  }
+#endif
+}
+
+static PyObject* method_spinloop(PyObject* self, PyObject* args,
+                                 PyObject* kwargs) {
+  Py_buffer buffer;
+  PyObject* callback;
+  double timeout = 0.;
+
+  spinloop_state_t* state = (spinloop_state_t*)PyModule_GetState(self);
+  if (state == NULL) {
+    PyErr_SetString(PyExc_TypeError, "Failed to retrieve module state!");
+    return NULL;
+  }
+
+  static const char* keywords[] = {"buffer", "callback", "timeout", NULL};
+  if (!PyArg_ParseTupleAndKeywords(args, kwargs, "y*O|d", (char**)keywords,
+                                   &buffer, &callback, &timeout)) {
+    return NULL;
+  }
+
+  if (!PyCallable_Check(callback)) {
+    PyErr_SetString(PyExc_TypeError, "callback parameter must be callable!");
+    PyBuffer_Release(&buffer);
+    return NULL;
+  }
+
+  struct timespec t_start;
+  if (clock_gettime(TIMEOUT_CLOCK, &t_start) != 0) {
+    PyErr_SetString(PyExc_RuntimeError, "clock_gettime() failed!");
+    PyBuffer_Release(&buffer);
+    return NULL;
+  }
+
+  bool result = false;
+  bool error = false;
+  bool have_timeout = (timeout > 1e-9);
+  unsigned int iteration = 0;
+  const bool buffer_qualifies = (buffer.len <= state->max_monitor_line_size);
+
+  while (true) {
+    PyObject* res = PyObject_CallNoArgs(callback);
+    if (res == NULL) {
+      error = true;
+      break;
+    }
+    int ok = (res == Py_True);
+    Py_DECREF(res);
+
+    if (ok) {
+      result = true;
+      break;
+    }
+
+    // Check timeout at most every 16 iterations to avoid clock_gettime and
+    // comparison cost
+    if (have_timeout && (iteration & 15u) == 0) {
+      struct timespec t_now;
+      if (clock_gettime(TIMEOUT_CLOCK, &t_now) != 0) {
+        PyErr_SetString(PyExc_RuntimeError, "clock_gettime() failed!");
+        error = true;
+        break;
+      }
+
+      const double elapsed = (double)(t_now.tv_sec - t_start.tv_sec) +
+                             (t_now.tv_nsec - t_start.tv_nsec) * 1e-9;
+      if (elapsed >= timeout) {
+        result = false;
+        break;
+      }
+    }
+    ++iteration;
+
+#if defined(__i386__) || defined(__x86_64__)
+    // monitorx + mwaitx with qualified buffer
+    if (buffer_qualifies && state->cpu_support == CPU_SUPPORT_MONITORX) {
+      _mm_monitorx(buffer.buf, 0, 0);
+
+      // Check once more in case the buffer has been modified while we were
+      // arming the monitor hardware
+      res = PyObject_CallNoArgs(callback);
+      if (res == NULL) {
+        error = true;
+        break;
+      }
+      ok = (res == Py_True);
+      Py_DECREF(res);
+
+      if (ok) {
+        result = true;
+        break;
+      }
+
+      // Run mwaitx with enabled timeout (bit 1). The actual timeout value
+      // is not very important, we just want to ensure we don't lock up
+      // here for too long.
+      Py_BEGIN_ALLOW_THREADS _mm_mwaitx((1 << 1), 0,
+                                        MWAITX_DEFAULT_TIMEOUT_CYCLES);
+      Py_END_ALLOW_THREADS
+    }
+
+    // Fallback: Busy poll
+    else {
+#endif
+      // Give other threads a chance to be scheduled
+      Py_BEGIN_ALLOW_THREADS
+#if defined(__i386__) || defined(__x86_64__)
+      __builtin_ia32_pause();
+#elif defined(__aarch64__)
+        __asm__ volatile("yield" :: : "memory");
+#endif
+      Py_END_ALLOW_THREADS
+#if defined(__i386__) || defined(__x86_64__)
+    }
+#endif
+  }
+
+  PyBuffer_Release(&buffer);
+
+  if (error) {
+    return NULL;
+  }
+
+  if (result) {
+    Py_RETURN_TRUE;
+  }
+
+  Py_RETURN_FALSE;
+}
+
+static PyMethodDef spinloop_methods[] = {
+    {"spinloop", (PyCFunction)method_spinloop, METH_VARARGS | METH_KEYWORDS,
+     "Wait for store with callback"},
+    {NULL, NULL, 0, NULL}};
+
+static struct PyModuleDef spinloop_module = {
+    PyModuleDef_HEAD_INIT, "spinloop",
+    "Hardware-optimized spinloops for Python", sizeof(spinloop_state_t),
+    spinloop_methods};
+
+PyMODINIT_FUNC PyInit_spinloop(void) {
+  PyObject* m = PyModule_Create(&spinloop_module);
+  if (m != NULL) {
+    spinloop_state_t* state = (spinloop_state_t*)PyModule_GetState(m);
+    if (state != NULL) {
+      determine_cpu_support(state);
+    }
+  }
+  return m;
+}
+
+}  // extern "C"
@@ -20,7 +20,7 @@ void launch_persistent_topk(const torch::Tensor& logits,
  namespace P = vllm::persistent;

  const int64_t num_rows = logits.size(0);
-  const int64_t stride = logits.size(1);
+  const int64_t stride = logits.stride(0);
  cudaStream_t stream = at::cuda::getCurrentCUDAStream();

  static int num_sms = 0;
@@ -153,6 +153,38 @@ void launch_persistent_topk(const torch::Tensor& logits,
    TORCH_CHECK(workspace.size(0) >= static_cast<int64_t>(state_bytes),
                "workspace too small, need ", state_bytes, " bytes");

+    // Zero the per-group RadixRowState region before launch.
+    //
+    // Issued UNCONDITIONALLY so the memset is captured as its own node in
+    // the cudagraph (a separate cudaMemsetAsync node, sequenced before the
+    // persistent_topk_kernel launch on the same stream). The previous
+    // host-side guard `if (needs_cooperative)` was evaluated at capture time;
+    // when capture-time max_seq_len <= RADIX_THRESHOLD (always true under
+    // FULL_DECODE_ONLY with max_model_len < 32 K) the memset would NOT be
+    // captured, leaving the workspace state to accumulate across replays.
+    // That's a latent correctness bug if the runtime data ever takes the
+    // radix path, and removes one variable while debugging hangs in the
+    // decode/medium paths.
+    //
+    // Cost is sub-microsecond: state_bytes = num_groups * sizeof(RadixRowState)
+    // is ~3 KB per group, ~100 KB for the largest grids on this hardware.
+    //
+    // Why the memset is required (regardless of which path the kernel takes):
+    //   1. arrival_counter accumulates within a launch and is never reset,
+    //      so a prior call leaves it at a large positive value. Without this
+    //      reset, the very first wait_ge in the next call sees counter >>
+    //      target and returns instantly, breaking the barrier.
+    //   2. The previous in-kernel init only ran in CTA-0 with intra-CTA
+    //      __syncthreads(), so it had no happens-before edge to CTA-1+'s
+    //      first red_release. cudaMemsetAsync is stream-ordered: the zero
+    //      is globally visible before any CTA runs.
+    {
+      cudaError_t mz_err = cudaMemsetAsync(workspace.data_ptr<uint8_t>(), 0,
+                                           state_bytes, stream);
+      TORCH_CHECK(mz_err == cudaSuccess,
+                  "row_states memset failed: ", cudaGetErrorString(mz_err));
+    }
+
    P::PersistentTopKParams params;
    params.input = logits.data_ptr<float>();
    params.output = output.data_ptr<int32_t>();
@@ -211,7 +243,7 @@ void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,
  TORCH_CHECK(output.dim() == 2, "output must be 2D");

  const int64_t num_rows = logits.size(0);
-  const int64_t stride = logits.size(1);
+  const int64_t stride = logits.stride(0);

  TORCH_CHECK(lengths.numel() == num_rows, "lengths size mismatch");
  TORCH_CHECK(output.size(0) == num_rows && output.size(1) == k,
@@ -183,7 +183,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
      "int forced_token_heads_per_warp=-1) -> ()");
  ops.impl("fused_qk_norm_rope", torch::kCUDA, &fused_qk_norm_rope);

-#ifndef USE_ROCM
  // Horizontally-fused DeepseekV4-MLA: per-head RMSNorm + GPT-J RoPE for Q, and
  // GPT-J RoPE + UE8M0 FP8 quant + paged cache insert for KV, all in one
  // kernel launch.
@@ -194,7 +193,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
      "float eps, int cache_block_size) -> ()");
  ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert", torch::kCUDA,
           &fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert);
-#endif

  // Apply repetition penalties to logits in-place
  ops.def(
@@ -555,7 +553,8 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cache_ops), cache_ops) {
  // Batch swap: submit all block copies in a single driver call.
  cache_ops.def(
      "swap_blocks_batch(Tensor src_ptrs, Tensor dst_ptrs,"
-      "                  Tensor sizes) -> ()");
+      "                  Tensor sizes,"
+      "                  bool is_src_access_order_any=False) -> ()");
  cache_ops.impl("swap_blocks_batch", torch::kCPU, &swap_blocks_batch);

  // Reshape the key and value tensors and cache them.
@@ -41,6 +41,13 @@ ARG BUILD_BASE_IMAGE=nvidia/cuda:${CUDA_VERSION}-devel-ubuntu22.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-ubuntu${UBUNTU_VERSION}

+# OS family of BUILD_BASE_IMAGE. Controls package manager (apt vs dnf) and
+# Python bootstrap. Set to "manylinux" alongside a manylinux build base such
+# as pytorch/manylinux2_28-builder:cuda13.0 to produce wheels with a glibc
+# 2.28 floor (matches PyTorch's own published wheels). Default stays on
+# Ubuntu for backwards compatibility.
+ARG BUILD_OS=ubuntu
+
 # By parameterizing the Deadsnakes repository URL, we allow third-party to use
 # their own mirror. When doing so, we don't benefit from the transparent
 # installation of the GPG key of the PPA, as done by add-apt-repository, so we
@@ -94,35 +101,64 @@ FROM ${BUILD_BASE_IMAGE} AS base

 ARG CUDA_VERSION
 ARG PYTHON_VERSION
+ARG BUILD_OS

 ENV DEBIAN_FRONTEND=noninteractive

-# Install system dependencies including build tools
-RUN 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 \
-    # Install python dev headers if available (needed for cmake FindPython on Ubuntu 24.04
-    # which ships cmake 3.28 and requires Development.SABIModule; silently skipped on
-    # Ubuntu 20.04/22.04 where python3.x-dev is not available without a PPA)
-    && (apt-get install -y --no-install-recommends python${PYTHON_VERSION}-dev 2>/dev/null || true) \
-    && rm -rf /var/lib/apt/lists/* \
-    && curl -LsSf https://astral.sh/uv/install.sh | sh \
-    && $HOME/.local/bin/uv venv /opt/venv --python ${PYTHON_VERSION} \
+# Install system dependencies including build tools.
+# The Ubuntu path uses apt + deadsnakes-via-uv for Python; the manylinux path
+# (AlmaLinux 8, e.g. pytorch/manylinux2_28-builder) uses dnf and the Python
+# interpreters pre-installed at /opt/python/cpXY-cpXY/.
+RUN if [ "${BUILD_OS}" = "manylinux" ]; then \
+        # rdma-core-devel provides libibverbs headers; ccache lives in EPEL,
+        # which the pytorch manylinux image already enables. git/curl/sudo
+        # are typically pre-installed but listed defensively.
+        dnf install -y --setopt=install_weak_deps=False \
+            ccache \
+            git \
+            curl \
+            sudo \
+            rdma-core-devel \
+        && dnf clean all \
+        && rm -rf /var/cache/dnf; \
+    else \
+        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 \
+        # Install python dev headers if available (needed for cmake FindPython on Ubuntu 24.04
+        # which ships cmake 3.28 and requires Development.SABIModule; silently skipped on
+        # Ubuntu 20.04/22.04 where python3.x-dev is not available without a PPA)
+        && (apt-get install -y --no-install-recommends python${PYTHON_VERSION}-dev 2>/dev/null || true) \
+        && rm -rf /var/lib/apt/lists/*; \
+    fi
+
+# Install uv and bootstrap /opt/venv. Both paths converge on /opt/venv so all
+# downstream stages stay distro-agnostic.
+RUN curl -LsSf https://astral.sh/uv/install.sh | sh \
+    && if [ "${BUILD_OS}" = "manylinux" ]; then \
+           # manylinux images ship Python at /opt/python/cpXY-cpXY/; point uv
+           # at the matching interpreter rather than letting it fetch one.
+           PYV_NODOT=$(echo ${PYTHON_VERSION} | tr -d '.') \
+           && MANYLINUX_PY=/opt/python/cp${PYV_NODOT}-cp${PYV_NODOT}/bin/python${PYTHON_VERSION} \
+           && $HOME/.local/bin/uv venv /opt/venv --python "$MANYLINUX_PY"; \
+       else \
+           $HOME/.local/bin/uv venv /opt/venv --python ${PYTHON_VERSION}; \
+       fi \
    && rm -f /usr/bin/python3 /usr/bin/python3-config /usr/bin/pip \
-    && ln -s /opt/venv/bin/python3 /usr/bin/python3 \
-    && ln -s /opt/venv/bin/python3-config /usr/bin/python3-config \
-    && ln -s /opt/venv/bin/pip /usr/bin/pip \
+    && ln -sf /opt/venv/bin/python3 /usr/bin/python3 \
+    && ln -sf /opt/venv/bin/python3-config /usr/bin/python3-config \
+    && ln -sf /opt/venv/bin/pip /usr/bin/pip \
    && python3 --version && python3 -m pip --version

 # Activate virtual environment and add uv to PATH
@@ -433,6 +469,7 @@ FROM base AS dev
 ARG PIP_INDEX_URL UV_INDEX_URL
 ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
 ARG PYTORCH_CUDA_INDEX_BASE_URL
+ARG BUILD_OS

 # This timeout (in seconds) is necessary when installing some dependencies via uv since it's likely to time out
 # Reference: https://github.com/astral-sh/uv/pull/1694
@@ -442,7 +479,11 @@ 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 if [ "${BUILD_OS}" = "manylinux" ]; then \
+        dnf install -y numactl-devel && dnf clean all && rm -rf /var/cache/dnf; \
+    else \
+        apt-get update && apt-get install -y --no-install-recommends libnuma-dev && rm -rf /var/lib/apt/lists/*; \
+    fi


 # We can specify the standard or nightly build of PyTorch
@@ -819,7 +860,7 @@ LABEL org.opencontainers.image.source="https://github.com/vllm-project/vllm" \
 # define sagemaker first, so it is not default from `docker build`
 FROM vllm-openai-base AS vllm-sagemaker

-COPY examples/online_serving/sagemaker-entrypoint.sh .
+COPY examples/deployment/sagemaker-entrypoint.sh .
 RUN chmod +x sagemaker-entrypoint.sh
 ENTRYPOINT ["./sagemaker-entrypoint.sh"]

@@ -124,9 +124,9 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1

 # RIXL/UCX build stages
 FROM base AS build_rixl
-ARG RIXL_BRANCH="bf4a7214"
+ARG RIXL_BRANCH="39be1de8"
 ARG RIXL_REPO="https://github.com/ROCm/RIXL.git"
-ARG UCX_BRANCH="7009d7a1"
+ARG UCX_BRANCH="bfb51733"
 ARG UCX_REPO="https://github.com/openucx/ucx.git"
 ENV ROCM_PATH=/opt/rocm
 ENV UCX_HOME=/usr/local/ucx
@@ -192,6 +192,7 @@ RUN cd /opt/rixl && \
    sed -i "s/--exclude 'libamdhip64\*'/--exclude 'libamdhip64*' --exclude 'libcore*' --exclude 'libpull*'/" \
        contrib/build-wheel.sh && \
    mkdir -p /app/install && \
+    _ucx_install_dir=${UCX_HOME} \
    ./contrib/build-wheel.sh \
        --output-dir /app/install \
        --rocm-dir ${ROCM_PATH} \
@@ -506,6 +507,10 @@ RUN --mount=type=bind,from=export_vllm,src=/,target=/install \
    && pip uninstall -y vllm \
    && uv pip install --system *.whl

+# Install RIXL wheel
+RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \
+    uv pip install --system /rixl_install/*.whl
+
 ARG COMMON_WORKDIR
 ARG BASE_IMAGE
 ARG NIC_BACKEND
@@ -516,6 +521,10 @@ 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

+# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc
+# See: https://github.com/ROCm/rocm-libraries/issues/6266
+ENV HSA_ENABLE_IPC_MODE_LEGACY=1
+
 ENV TOKENIZERS_PARALLELISM=false

 # ENV that can improve safe tensor loading, and end-to-end time
--- a/Show More
+++ b/Show More