[Attention][TokenSpeed MLA] Also warm up prefill kernel from decode impl

The prefill backend may be paired with flash_attn / trtllm in production — in that case the prefill backend's __init__ never runs and the prefill kernel's first call pays a 1.5–2 minute JIT cost. Add the same idempotent `warmup_compile_prefill` invocation to TokenspeedMLAImpl.__init__ (the decode-side backend, always present when tokenspeed is selected). The function dedupes by config key, so the double call is a no-op when both backends are tokenspeed. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
update version
2026-05-06 07:45:16 +00:00 · 2026-05-06 07:28:20 +00:00 · 2026-05-06 07:20:48 +00:00 · 2026-05-06 03:14:49 +00:00 · 2026-05-06 02:53:46 +00:00 · 2026-05-06 02:49:09 +00:00
506 changed files with 36443 additions and 53536 deletions
@@ -12,19 +12,15 @@ 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 30m "
+      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 20m "
      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/quantization/test_cpu_fp8_scaled_mm.py"
+      pytest -x -v -s tests/kernels/test_awq_int4_to_int8.py"

 - label: CPU-Compatibility Tests
  depends_on: []
@@ -65,7 +61,6 @@ 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,18 +18,17 @@ 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 &&
+      (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_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  &&
-      pytest -v -s lora/test_worker.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)'

 - label: LoRA Fused/MoE Kernels
  timeout_in_minutes: 45
@@ -47,7 +46,6 @@ 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'

@@ -67,9 +65,8 @@ 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_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]"'
+      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]"'

 - label: LoRA Punica FP8/XPU Ops
  timeout_in_minutes: 45
@@ -87,7 +84,6 @@ 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'

@@ -107,12 +103,10 @@ 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_transformers_model.py &&
-      pytest -v -s lora/test_chatglm3_tp.py &&
-      pytest -s -v lora/test_minicpmv_tp.py'
+      pytest -v -s lora/test_qwen35_densemodel_lora.py &&
+      pytest -v -s lora/test_transformers_model.py'

 - label: LoRA Multimodal
  timeout_in_minutes: 45
@@ -130,6 +124,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 --ignore=v1/logits_processors/test_custom_online.py --ignore=v1/logits_processors/test_custom_offline.py &&
+      pytest -v -s v1/logits_processors &&
      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_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_tree_attention.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py &&
        pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --ignore=v1/kv_connector/unit/test_lmcache_integration.py --ignore=v1/kv_connector/unit/test_hf3fs_client.py --ignore=v1/kv_connector/unit/test_hf3fs_connector.py --ignore=v1/kv_connector/unit/test_hf3fs_metadata_server.py'
@@ -378,11 +378,9 @@ 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
@@ -399,15 +397,9 @@ fi

 echo "Raw commands: $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
+# Fix quoting before ROCm overrides (so overrides see correct structure)
+commands=$(re_quote_pytest_markers "$commands")
+echo "After re-quoting: $commands"

 commands=$(apply_rocm_test_overrides "$commands")
 echo "Final commands: $commands"
@@ -31,21 +31,6 @@ 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
@@ -61,6 +46,13 @@ 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 '
@@ -75,21 +67,6 @@ 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,6 +136,8 @@ 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" \
@@ -230,6 +230,7 @@ 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
@@ -271,6 +272,7 @@ 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
@@ -394,8 +396,8 @@ steps:
    - python3 pooling/embed/vision_embedding_offline.py --seed 0
    # Features demo
    - python3 features/automatic_prefix_caching/prefix_caching_offline.py
-    - 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 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 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

@@ -588,6 +590,7 @@ 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
@@ -618,7 +621,6 @@ 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/
@@ -862,6 +864,7 @@ 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
@@ -880,7 +883,7 @@ steps:
  - vllm/platforms/rocm.py
  commands:
  - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
-  - ATTENTION_BACKEND=ROCM_ATTN bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh
+  - ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh

 - label: V1 e2e (2 GPUs) # TBD
  timeout_in_minutes: 180
@@ -927,7 +930,6 @@ steps:
  - tests/renderers
  - tests/standalone_tests/lazy_imports.py
  - tests/tokenizers_
-  - tests/reasoning
  - tests/tool_parsers
  - tests/transformers_utils
  - tests/config
@@ -940,7 +942,7 @@ 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
+  - 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 tool_parsers
  - pytest -v -s transformers_utils
  - pytest -v -s config
@@ -1100,13 +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/fusions_e2e/test_tp2_ar_rms.py::test_tp2_ar_rms_fusions
+  - pytest -v -s tests/compile/passes/distributed/test_sequence_parallelism.py
+  - pytest -v -s tests/compile/passes/distributed/test_tp2_ar_rms.py::test_tp2_ar_rms_fusions

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

@@ -1171,6 +1173,7 @@ 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

@@ -1184,6 +1187,7 @@ 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)'
@@ -1203,6 +1207,7 @@ 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
@@ -1248,6 +1253,7 @@ 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"
@@ -1269,6 +1275,7 @@ 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  --------------------------------------------------------#
@@ -1649,8 +1656,8 @@ steps:
    - python3 pooling/embed/vision_embedding_offline.py --seed 0
    # Features demo
    - python3 features/automatic_prefix_caching/prefix_caching_offline.py
-    - 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 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 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

@@ -1796,17 +1803,15 @@ 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:
-  - 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
-
+  - 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

 - label: Multi-Modal Models (Extended Generation 2) # TBD
  timeout_in_minutes: 180
@@ -1818,10 +1823,8 @@ steps:
  - vllm/
  - tests/models/multimodal/generation
  commands:
-  - 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)'
-
+  - 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'

 - label: Multi-Modal Models (Extended Generation 3) # TBD
  timeout_in_minutes: 180
@@ -1841,7 +1844,6 @@ 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/
@@ -2202,6 +2204,7 @@ steps:
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
+  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -2278,6 +2281,7 @@ 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
@@ -2297,6 +2301,7 @@ 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
@@ -2359,6 +2364,7 @@ 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
@@ -2488,6 +2494,7 @@ 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
@@ -2512,6 +2519,7 @@ 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
@@ -2532,6 +2540,7 @@ 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
@@ -2619,7 +2628,6 @@ steps:
  agent_pool: mi325_1
  torch_nightly: true
  parallelism: 2
-  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -2645,7 +2653,6 @@ 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/
@@ -2711,6 +2718,7 @@ 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

@@ -2741,7 +2749,6 @@ steps:
  agent_pool: mi355_1
  fast_check: true
  torch_nightly: true
-  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -2757,7 +2764,6 @@ steps:
  agent_pool: mi355_1
  fast_check: true
  torch_nightly: true
-  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/
@@ -2933,8 +2939,8 @@ steps:
  - python3 pooling/embed/vision_embedding_offline.py --seed 0
  # Features demo
  - python3 features/automatic_prefix_caching/prefix_caching_offline.py
-  - 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 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 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

@@ -3046,7 +3052,7 @@ steps:
  - vllm/
  - tests/models/language/generation
  commands:
-  - 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/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)'

@@ -3054,7 +3060,6 @@ 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/
@@ -3254,7 +3259,6 @@ 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/
@@ -3280,7 +3284,6 @@ 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/
@@ -3321,7 +3324,7 @@ steps:
  - vllm/platforms/rocm.py
  commands:
  - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
-  - ATTENTION_BACKEND=ROCM_ATTN bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh
+  - ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh

 - label: Distributed NixlConnector PD accuracy (4 GPUs) # TBD
  timeout_in_minutes: 180
@@ -7,11 +7,7 @@ steps:
  timeout_in_minutes: 15
  device: h200_18gb
  source_file_dependencies:
-  - vllm/envs.py
-  - vllm/logger.py
-  - vllm/platforms/
-  - vllm/plugins/
-  - vllm/utils/
+  - vllm/
  - tests/cuda
  commands:
    - pytest -v -s cuda/test_cuda_context.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/
+    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
    - 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/
+    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
    - 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/
+    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
    - 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/
+    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
    - 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/
+    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
    - 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/
+    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
    - 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/
+    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
    - 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/
+    - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
    - 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/
@@ -7,17 +7,7 @@ steps:
  timeout_in_minutes: 15
  device: h200_18gb
  source_file_dependencies:
-  - 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/
+  - vllm/
  - tests/engine
  - tests/test_sequence
  - tests/test_config
@@ -62,27 +52,16 @@ steps:
  optional: true
  num_devices: 2
  source_file_dependencies:
-    - 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
+    - vllm/
+    - tests/v1/e2e
  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
@@ -90,27 +69,16 @@ steps:
  optional: true
  num_devices: 4
  source_file_dependencies:
-    - 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
+    - vllm/
+    - tests/v1/e2e
  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
@@ -26,11 +26,6 @@ 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
@@ -43,6 +38,11 @@ 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)
@@ -57,6 +57,11 @@ steps:
  - 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
@@ -19,7 +19,6 @@ 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
@@ -6,37 +6,24 @@ steps:
  key: v1-spec-decode
  timeout_in_minutes: 30
  source_file_dependencies:
-    - vllm/config/
-    - vllm/distributed/
-    - vllm/inputs/
-    - vllm/model_executor/
-    - vllm/platforms/
-    - vllm/sampling_params.py
-    - vllm/transformers_utils/
-    - vllm/utils/
-    - vllm/v1/
+    - vllm/
    - 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/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/
+    - vllm/
    - tests/v1/sample
    - tests/v1/logits_processors
    - tests/v1/test_oracle.py
@@ -51,7 +38,7 @@ steps:
    - pytest -v -s v1/test_outputs.py
  mirror:
    amd:
-      device: mi300_1
+      device: mi325_1
      depends_on:
      - image-build-amd

@@ -59,23 +46,7 @@ steps:
  key: v1-core-kv-metrics
  timeout_in_minutes: 30
  source_file_dependencies:
-    - 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/
+    - vllm/
    - tests/v1/core
    - tests/v1/executor
    - tests/v1/kv_offload
@@ -96,27 +67,18 @@ 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/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/
+    - vllm/
    - tests/v1
  device: cpu-small
  commands:
@@ -132,17 +94,7 @@ steps:
  timeout_in_minutes: 20
  device: h200_18gb
  source_file_dependencies:
-  - 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/
+  - vllm/
  - tests/test_regression
  commands:
  - pip install modelscope
@@ -175,8 +127,8 @@ steps:
    - python3 pooling/embed/vision_embedding_offline.py --seed 0
    # for features demo
    - python3 features/automatic_prefix_caching/prefix_caching_offline.py
-    - 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 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 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
@@ -186,18 +138,7 @@ steps:
  timeout_in_minutes: 20
  num_devices: 2
  source_file_dependencies:
-  - 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/
+  - vllm/
  - tests/v1/tracing
  commands:
  - "pip install \
@@ -221,19 +162,7 @@ steps:
  key: async-engine-inputs-utils-worker
  timeout_in_minutes: 50
  source_file_dependencies:
-  - 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/
+  - vllm/
  - tests/detokenizer
  - tests/multimodal
  - tests/utils_
@@ -248,30 +177,7 @@ steps:
  - image-build-cpu
  timeout_in_minutes: 30
  source_file_dependencies:
-  - 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/
+  - vllm/
  - tests/test_inputs.py
  - tests/test_outputs.py
  - tests/test_pooling_params.py
@@ -294,7 +200,7 @@ 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
+  - 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 tool_parsers
  - pytest -v -s transformers_utils
  - pytest -v -s config
@@ -37,7 +37,7 @@ steps:
    - examples/generate/multimodal/
    - examples/features/
    - examples/pooling/embed/vision_embedding_offline.py
-    - examples/features/tensorize_vllm_model.py
+    - examples/others/tensorize_vllm_model.py
  commands:
    - set -x
    - export VLLM_USE_V2_MODEL_RUNNER=1
@@ -55,8 +55,8 @@ steps:
    - python3 pooling/embed/vision_embedding_offline.py --seed 0
    # for features demo
    - python3 features/automatic_prefix_caching/prefix_caching_offline.py
-    - 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 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 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
@@ -7,21 +7,7 @@ steps:
  timeout_in_minutes: 45
  torch_nightly: true
  source_file_dependencies:
-  - vllm/config/
-  - vllm/distributed/
-  - vllm/engine/
-  - vllm/inputs/
-  - vllm/logging_utils/
-  - vllm/model_executor/
-  - vllm/multimodal/
-  - vllm/platforms/
-  - vllm/sampling_params.py
-  - vllm/sequence.py
-  - vllm/tasks.py
-  - vllm/transformers_utils/
-  - vllm/utils/
-  - vllm/v1/
-  - vllm/version.py
+  - vllm/
  - tests/models/test_initialization.py
  - tests/models/registry.py
  commands:
@@ -50,26 +36,18 @@ steps:
  key: basic-models-tests-other
  timeout_in_minutes: 45
  source_file_dependencies:
-  - vllm/config/
-  - vllm/distributed/
-  - vllm/engine/
-  - vllm/inputs/
-  - vllm/logging_utils/
-  - vllm/model_executor/
-  - vllm/multimodal/
-  - vllm/platforms/
-  - vllm/sampling_params.py
-  - vllm/sequence.py
-  - vllm/tasks.py
-  - vllm/transformers_utils/
-  - vllm/utils/
-  - vllm/v1/
-  - vllm/version.py
+  - vllm/
  - tests/models/test_terratorch.py
  - tests/models/test_transformers.py
  - 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
@@ -77,11 +55,7 @@ steps:
  - image-build-cpu
  timeout_in_minutes: 10
  source_file_dependencies:
-  - vllm/config/
-  - vllm/distributed/
-  - vllm/model_executor/
-  - vllm/platforms/
-  - vllm/utils/
+  - vllm/
  - tests/models/test_utils.py
  - tests/models/test_vision.py
  device: cpu-small
@@ -48,14 +48,6 @@ 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
@@ -70,6 +62,15 @@ 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
@@ -91,6 +92,11 @@ steps:
  - 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: mi300_1
+      device: mi325_1
      depends_on:
      - image-build-amd

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

@@ -49,7 +49,7 @@ steps:
    - pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model
  mirror:
    amd:
-      device: mi300_1
+      device: mi325_1
      depends_on:
      - image-build-amd

@@ -64,6 +64,11 @@ 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
@@ -115,7 +120,7 @@ steps:
    - pytest -v -s models/multimodal/test_mapping.py
  mirror:
    amd:
-      device: mi300_1
+      device: mi325_1
      depends_on:
      - image-build-amd

@@ -6,30 +6,7 @@ steps:
  key: pytorch-compilation-unit-tests
  timeout_in_minutes: 10
  source_file_dependencies:
-    - 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/
+    - vllm/
    - tests/compile
  commands:
  # Run unit tests defined directly under compile/,
@@ -47,30 +24,7 @@ steps:
  device: h100
  num_devices: 1
  source_file_dependencies:
-    - 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/
+    - vllm/
    - tests/compile/h100/
  commands:
  - "find compile/h100/ -name 'test_*.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"
@@ -79,30 +33,7 @@ steps:
  key: pytorch-compilation-passes-unit-tests
  timeout_in_minutes: 20
  source_file_dependencies:
-    - 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/
+    - vllm/
    - tests/compile/passes
  commands:
  - pytest -s -v compile/passes --ignore compile/passes/distributed
@@ -111,30 +42,7 @@ steps:
  key: pytorch-fullgraph-smoke-test
  timeout_in_minutes: 35
  source_file_dependencies:
-  - 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/
+  - vllm/
  - tests/compile
  commands:
  # Run smoke tests under fullgraph directory, except test_full_graph.py
@@ -148,30 +56,7 @@ steps:
  timeout_in_minutes: 30
  device: h200_18gb
  source_file_dependencies:
-  - 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/
+  - vllm/
  - tests/compile
  commands:
    # fp8 kv scales not supported on sm89, tested on Blackwell instead
@@ -17,7 +17,7 @@ steps:
    - VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers
  mirror:
    amd:
-      device: mi250_1
+      device: mi325_1
      depends_on:
      - image-build-amd
      commands:
@@ -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 @zyongye
-/vllm/model_executor/layers/quantization @mgoin @robertgshaw2-redhat @tlrmchlsmth @yewentao256 @pavanimajety @zyongye
+/vllm/model_executor/layers/fused_moe @mgoin @pavanimajety
+/vllm/model_executor/layers/quantization @mgoin @robertgshaw2-redhat @tlrmchlsmth @yewentao256 @pavanimajety
 /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,8 +18,7 @@
 /vllm/kernels/helion @ProExpertProg @zou3519
 /vllm/multimodal @DarkLight1337 @ywang96 @NickLucche @tjtanaa
 /vllm/vllm_flash_attn @LucasWilkinson @MatthewBonanni
-/CMakeLists.txt @tlrmchlsmth @LucasWilkinson @Harry-Chen
-/cmake @tlrmchlsmth @LucasWilkinson @Harry-Chen
+CMakeLists.txt @tlrmchlsmth @LucasWilkinson

 # Any change to the VllmConfig changes can have a large user-facing impact,
 # so spam a lot of people
@@ -71,10 +70,6 @@
 /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
@@ -82,11 +77,11 @@
 /tests/distributed/test_same_node.py @youkaichao
 /tests/entrypoints @DarkLight1337 @robertgshaw2-redhat @aarnphm @NickLucche
 /tests/evals @mgoin @vadiklyutiy
-/tests/kernels @mgoin @tlrmchlsmth @WoosukKwon @yewentao256 @zyongye
+/tests/kernels @mgoin @tlrmchlsmth @WoosukKwon @yewentao256
 /tests/kernels/ir @ProExpertProg @tjtanaa
 /tests/models @DarkLight1337 @ywang96
 /tests/multimodal @DarkLight1337 @ywang96 @NickLucche
-/tests/quantization @mgoin @robertgshaw2-redhat @yewentao256 @pavanimajety @zyongye
+/tests/quantization @mgoin @robertgshaw2-redhat @yewentao256 @pavanimajety
 /tests/test_inputs.py @DarkLight1337 @ywang96
 /tests/entrypoints/llm/test_struct_output_generate.py @mgoin @russellb @aarnphm
 /tests/v1/structured_output @mgoin @russellb @aarnphm
@@ -152,12 +147,6 @@ 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,7 +477,9 @@ pull_request_rules:
  conditions:
    - label != stale
    - or:
-      - files~=^examples/disaggregated/
+      - files~=^examples/online_serving/disaggregated[^/]*/.*
+      - files~=^examples/offline_inference/disaggregated[^/]*/.*
+      - files~=^examples/others/lmcache/
      - files~=^tests/v1/kv_connector/
      - files~=^vllm/distributed/kv_transfer/
      - title~=(?i)\bP/?D\b
@@ -131,16 +131,6 @@ 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,44 +9,7 @@ build:
    python: "3.12"
  jobs:
    post_checkout:
-      - |
-        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
+      # - bash docs/maybe_skip_pr_build.sh
      - git fetch origin main --unshallow --no-tags --filter=blob:none || true
    pre_create_environment:
      - pip install uv
@@ -13,12 +13,8 @@ cmake_minimum_required(VERSION 3.26)
 # cmake --install . --component _C
 project(vllm_extensions LANGUAGES CXX)

-set(CMAKE_CXX_STANDARD 20)
+set(CMAKE_CXX_STANDARD 17)
 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)
@@ -1,7 +1,7 @@
 include(FetchContent)

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

@@ -32,23 +32,18 @@ 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. 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.
+    # and create a local shim dir with it
    vllm_prepare_torch_gomp_shim(VLLM_TORCH_GOMP_SHIM_DIR)

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

@@ -326,6 +321,14 @@ 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)
@@ -338,14 +341,8 @@ 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)
-    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()
+    target_compile_options(dnnl_ext PRIVATE ${DNNL_COMPILE_FLAGS} -fPIC)
    list(APPEND LIBS dnnl_ext)
-
-
    set(USE_ONEDNN ON)
 else()
    set(USE_ONEDNN OFF)
@@ -409,15 +406,12 @@ 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
@@ -436,11 +430,10 @@ 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"
@@ -455,7 +448,7 @@ if (ENABLE_X86_ISA)

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

    # AMX + AVX512F + AVX512BF16 + AVX512VNNI
    define_extension_target(
@@ -76,6 +76,7 @@ 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,8 +12,7 @@ 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,
-                       bool is_src_access_order_any);
+                       const torch::Tensor& sizes);

 void reshape_and_cache(torch::Tensor& key, torch::Tensor& value,
                       torch::Tensor& key_cache, torch::Tensor& value_cache,
@@ -77,8 +77,7 @@ 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,
-                       bool is_src_access_order_any) {
+                       const torch::Tensor& sizes) {
  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");
@@ -98,13 +97,13 @@ void swap_blocks_batch(const torch::Tensor& src_ptrs,

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

-  // 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.
+  // 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));
  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).
@@ -125,12 +124,7 @@ void swap_blocks_batch(const torch::Tensor& src_ptrs,

  if (batch_fn != nullptr) {
    CUmemcpyAttributes attr = {};
-    // 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;
+    attr.srcAccessOrder = 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),
@@ -140,30 +134,12 @@ 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);
-    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;
-  }
+  } else
 #endif
  {
-    // Fallback for CUDA < 12.8, older CUDA drivers, and ROCm < 7.1:
-    // individual async copies. cudaMemcpyDefault lets the driver infer
-    // direction from pointer types.
+    // Fallback for CUDA < 12.8, older drivers, and ROCm:
+    // 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]),
@@ -486,10 +486,6 @@ 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,9 +6,6 @@

 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__) \
@@ -148,9 +145,6 @@ 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> {
@@ -308,10 +302,6 @@ 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>(
@@ -146,9 +146,6 @@ 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> {
@@ -411,10 +408,6 @@ 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,10 +688,6 @@ 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,17 +122,9 @@ 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
  }
 };

@@ -169,17 +161,9 @@ 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
  }
 };

@@ -263,12 +247,13 @@ 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(_mm256_broadcastsi128_si256((__m128i)vec8_data.reg)),
-        reg_high(_mm256_broadcastsi128_si256((__m128i)vec8_data.reg)) {}
+      : 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)) {}

  // E4M3 decode (AVX2 path) — same bit-layout trick as the AVX512 variant
  // above.  Result = true_E4M3 * 2^-8; caller applies scale * 2^8.
@@ -689,11 +674,6 @@ 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));
@@ -759,85 +739,6 @@ 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

@@ -890,34 +791,6 @@ 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 += vec_elem_num) {
+      for (; j < hidden_size - vec_elem_num; ++j) {
        cvt_vec_t elems_fp32(input_ptr + j);
        elems_fp32 = elems_fp32 * token_scale_vec;
        if constexpr (AZP) {
@@ -1,12 +1,13 @@
 // 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>
@@ -15,157 +16,40 @@
 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__();                                                      \
+    }                                                                            \
  }()

-#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__)   \
-    }                                                                  \
+// 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 UNUSED(x) (void)(x)
@@ -186,51 +70,13 @@ 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)

-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
+// parallel routines
 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
@@ -248,10 +94,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
 }

@@ -259,15 +105,23 @@ 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);
+    f(0, n);
+#endif
+}
+
+inline int get_thread_num() {
+#if defined(_OPENMP)
+  return omp_get_thread_num();
+#else
+  return 0;
 #endif
 }

@@ -283,6 +137,7 @@ 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);

@@ -310,59 +165,31 @@ 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>
-inline int get_cache_blocks(int chunk_size) {
+int get_cache_blocks(int BLOCK_SIZE, int K) {
  // L2 2MB and ratio of 50%
  const int L2_size = 2048 * 1024 >> 1;
-  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);
-      }
-    }
-  }
+  return std::max(1, int(L2_size / (BLOCK_SIZE * K * sizeof(T))));
 }

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

-// 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
+} // anonymous namespace
@@ -1,709 +0,0 @@
-// 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,12 +1,11 @@
 // Adapted from
 // https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu

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

 namespace {

@@ -27,13 +26,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!");
@@ -70,43 +69,6 @@ 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>;
@@ -114,7 +76,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());
@@ -127,34 +89,13 @@ inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ inpu
 }

 template <typename scalar_t>
-inline void copy_stub(float* __restrict__ out, const scalar_t* __restrict__ input, int64_t size) {
+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, 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
+  #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());
@@ -166,51 +107,11 @@ inline void copy_add_stub(
  }
 }

-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!");
  }
 };
@@ -219,14 +120,9 @@ 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;

@@ -249,7 +145,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);

@@ -284,7 +180,9 @@ 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);
@@ -292,33 +190,22 @@ 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) {
@@ -329,21 +216,6 @@ 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>
@@ -360,12 +232,15 @@ 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, N = 16, 32, 48, 64
+  // pattern: 1-4-16
  constexpr int64_t BLOCK_M = 4;
  constexpr int64_t BLOCK_N = 64;
  const int64_t MB = div_up(M, BLOCK_M);
@@ -377,88 +252,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 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;
+        case 0x12: LAUNCH_TINYGEMM_KERNEL_NN(1, 32); break;
+        case 0x14: LAUNCH_TINYGEMM_KERNEL_NN(1, 64); break;
        // mb_size = 2
-        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;
+        case 0x22: LAUNCH_TINYGEMM_KERNEL_NN(2, 32); break;
+        case 0x24: LAUNCH_TINYGEMM_KERNEL_NN(2, 64); break;
        // mb_size = 3
-        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;
+        case 0x32: LAUNCH_TINYGEMM_KERNEL_NN(3, 32); break;
+        case 0x34: LAUNCH_TINYGEMM_KERNEL_NN(3, 64); break;
        // mb_size = 4
-        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);
+        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, 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,
@@ -470,20 +282,29 @@ 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);

-  const bool use_brgemm = can_use_brgemm<scalar_t>(M);
+  // 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);

  // 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) {
+    parallel_2d(MB, NB, [&](int64_t begin_mb, int64_t end_mb, int64_t begin_nb, int64_t end_nb) {
+
      // for brgemm, use float32 for accumulate
      alignas(64) float Ctmp[BLOCK_M * BLOCK_N];

-      loop_2d<scalar_t>(mb0, mb1, nb0, nb1, BLOCK_N * K, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
+      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) {
+
        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;
@@ -502,7 +323,7 @@ void weight_packed_linear_kernel_impl(
            /* ldb */ nb_size,
            /* ldc */ out_strideM,
            /* brg */ use_brgemm);
-      });
+      }}}

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

-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
+} // 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);
@@ -631,23 +359,14 @@ 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(actual_IC % TILE_K == 0, "invalid weight input features ", actual_IC);
+  TORCH_CHECK(IC % TILE_K == 0, "invalid weight input features ", IC);

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

-  // 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).");
+  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.");

  CPU_DISPATCH_PACKED_TYPES(st, [&] {
    // adjust most inner dimension size
-    const int packed_row_size = get_row_size<packed_t>(actual_IC);
+    const int packed_row_size = get_row_size<packed_t>(IC);
    auto sizes = weight.sizes().vec();
    sizes[ndim - 1] = packed_row_size;
    packed_weight.resize_(sizes);
@@ -682,7 +399,10 @@ 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);
@@ -692,71 +412,33 @@ 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] ([K, N] if use_fma_gemm)
+// mat2 : [N, K]
 // 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) {
-  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;
-  }
+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}));

-  int64_t M = mat1.size(0);
-  int64_t K = mat1.size(1);
-  int64_t N = use_fma_gemm ? mat2.size(1) : mat2.size(0);
+  auto packed_w = is_vnni ? mat2 : convert_weight_packed(mat2);

  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;
@@ -765,83 +447,12 @@ weight_packed_linear(at::Tensor& mat1, at::Tensor& mat2, const std::optional<at:
    bias_data = bias.value().data_ptr<float>();
  }

-  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", [&] {
+  AT_DISPATCH_REDUCED_FLOATING_TYPES(mat1.scalar_type(), "weight_packed_linear_kernel_impl", [&] {
    weight_packed_linear_kernel_impl<scalar_t>(
        out.data_ptr<scalar_t>(),
        mat1.data_ptr<scalar_t>(),
-        packed_w.data_ptr<float>(),
+        packed_w.data_ptr<scalar_t>(),
        bias_data,
-        post_mul_mat.data_ptr<scalar_t>(),
        M,
        N,
        K,
@@ -1,12 +1,8 @@
-// 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"
+// clang-format off

 // amx-bf16
 #define TILE_M 16
@@ -14,42 +10,20 @@
 #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;
-}
-// 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;
-}
+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; }

 // 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>
@@ -62,44 +36,18 @@ 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 to vnni format
+// pack weight into vnni format
 at::Tensor convert_weight_packed(at::Tensor& weight);

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

@@ -157,6 +105,35 @@ 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(
@@ -176,37 +153,6 @@ 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,
@@ -258,7 +204,6 @@ void tinygemm_kernel(
    int64_t ldc,
    bool brg);

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

-// per tensor quantization
 template <typename scalar_t>
 void tinygemm_kernel(
    const scalar_t* __restrict__ A,
-    const at::Float8_e4m3fn* __restrict__ B,
+    const at::quint4x2* __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,
@@ -314,21 +266,34 @@ void tinygemm_kernel(
    bool store_out,
    bool use_brgemm);

-// 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);
+// 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;
+}
@@ -1,10 +1,10 @@
-// Adapted from
-// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu
+// SPDX-License-Identifier: Apache-2.0
+// Adapted from sgl-project/sglang
+// https://github.com/sgl-project/sglang/pull/8226

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

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

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

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

-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
+#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));
  __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;
@@ -48,7 +45,8 @@ inline std::array<m256i_wrapper, 2> load_zps_4vnni(const int8_t* __restrict__ zp
  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);
@@ -60,42 +58,37 @@ inline std::array<m256i_wrapper, 2> load_uint4_as_int8(const uint8_t* __restrict
  return {low_wp, high_wp};
 }

-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
+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
  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);
-      // 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);
+      _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 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) {
+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) {
  for (int m = 0; m < M; ++m) {
    float a_scale = *(scale_a + m * ldsa);
    __m512 va_scale = _mm512_set1_ps(a_scale);
@@ -106,7 +99,7 @@ void _dequant_and_store(
      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) {
@@ -129,7 +122,8 @@ void _dequant_and_store(
      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;
@@ -141,10 +135,9 @@ void _dequant_and_store(
 }

 #else
-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]
+template <int N, int ldb>
+void _dequant_weight_zp_only(const uint8_t* B, int8_t* dqB,
+                             const int8_t* qzeros, int64_t K) {
  for (int k = 0; k < K; ++k) {
    for (int n = 0; n < N / 2; ++n) {
      int32_t b = (int32_t)B[k * ldb + n];
@@ -165,31 +158,20 @@ 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 <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*.
-
+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) {
  constexpr int COLS = N / 16;
-  // 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 ones = _mm512_set1_epi8(1);
  __m512i va;
  __m512i vb[COLS];
  __m512i vc[M * COLS];
@@ -197,7 +179,6 @@ void _dequant_gemm_accum_small_M(
  __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));
@@ -233,25 +214,25 @@ void _dequant_gemm_accum_small_M(
    }
  };

-  // 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)));
@@ -264,28 +245,17 @@ void _dequant_gemm_accum_small_M(
  Unroll<M * COLS>{}(store);
 }

-#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);
+  #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);
 #endif

-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
+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) {
 #if defined(CPU_CAPABILITY_AVX512)
  if (!use_brgemm) {
    switch (M) {
@@ -312,12 +282,14 @@ void _dequant_gemm_accum(
  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<true, N, sym_quant_act>(
-        C, C_i32, scales_a, qzeros_a, scales_b, compensation, M, N /*ldi*/, ldc, 1 /*ldsa*/);
+    _dequant_and_store<sym_quant_act, N, true>(C, C_i32, scales_a, qzeros_a,
+                                               scales_b, compensation, M,
+                                               N /*ldi*/, ldc, 1 /*ldsa*/);
  } else
 #endif
  {
@@ -325,13 +297,13 @@ void _dequant_gemm_accum(
  }
 }

-template <int64_t N>
+template <int 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);
@@ -341,11 +313,11 @@ inline void copy_bias(const float* bias_ptr, float* y_buf, int64_t m) {
        y_buf[i * N + j] = bias_ptr[j];
      }
    }
-  } else {  // initialize to zero
+  } else {
    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);
@@ -358,13 +330,14 @@ inline void copy_bias(const float* bias_ptr, float* y_buf, int64_t m) {
  }
 }

-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) {
+template <int N, typename out_dtype>
+inline void store_out(const float* y_buf, out_dtype* c_ptr, int64_t m,
+                      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);
@@ -375,11 +348,12 @@ inline void store_out(const float* y_buf, out_dtype* c_ptr, int64_t m, /* int64_
      }
    } 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) {
@@ -387,11 +361,12 @@ inline void store_out(const float* y_buf, out_dtype* c_ptr, int64_t m, /* int64_
      }
    } 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) {
@@ -417,25 +392,16 @@ void fill_val_stub(int32_t* __restrict__ output, int32_t value, int64_t size) {
  }
 }

-template <typename act_dtype, typename out_dtype, bool sym_quant_act>
+template <bool sym_quant_act, typename act_dtype, typename out_dtype>
 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) {
-  // 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);
+  const bool use_brgemm = can_use_brgemm<act_dtype>(M);
  int64_t block_m = [&]() -> long {
    if (M <= 48) {
      return M;
@@ -466,24 +432,30 @@ 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;
-        // copy bias to y_buf if bias is not None
+        int64_t m_size =
+            mci * block_m + block_m > M ? M - mci * block_m : block_m;
        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) /*Bcomp*/;
-          _dequant_gemm_accum<BLOCK_N, BLOCK_N / 2, sym_quant_act>(
+                  : (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>(
              /*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,
@@ -492,8 +464,8 @@ void _da8w4_linear_impl(
              /*ldc*/ BLOCK_N,
              /*use_brgemm*/ use_brgemm);
        }
-        // 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*/);
+        store_out<BLOCK_N>(C_tmp, output + mci * block_m * N + nc * BLOCK_N,
+                           m_size, N /*lda*/);
      }
    }
    if (use_brgemm) {
@@ -504,16 +476,15 @@ void _da8w4_linear_impl(

 }  // anonymous namespace

-/*
-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");
+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");

  auto new_scales = scales;
  auto new_qzeros = qzeros;
@@ -534,21 +505,20 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> convert_int4_weight_packed_with_c
  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();
-  // 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});
+  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});
  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>();
@@ -558,25 +528,24 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> convert_int4_weight_packed_with_c
  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);
-      // 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]
+      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);
      constexpr int n_group_size = 8;
      constexpr int vnni_size = 4;
-      constexpr int n_group = block_n / n_group_size;  // 4
+      constexpr int n_group = block_n / n_group_size;
      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);
@@ -585,128 +554,39 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> convert_int4_weight_packed_with_c
          }
        }
      }
-      // 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> 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);
+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);

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

-  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...");
-  }
+  return std::make_tuple(qweight_final, qzeros_final);
 }

 std::tuple<at::Tensor, at::Tensor, at::Tensor> convert_weight_packed_scale_zp(
-    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);
-
+    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);
  auto _scales = scales;
-  _qzeros = _qzeros.transpose(-2, -1).contiguous();  // .T
+  _qzeros = _qzeros.transpose(-2, -1).contiguous();
  _scales = _scales.transpose(-2, -1).contiguous();
-  if (_qweight.dim() == 3) {  // Dim=3 for MOE packing, TODO: refine a unified loop
+  if (_qweight.dim() == 3) {
    int64_t E = _qweight.size(0);
    int64_t K = _qweight.size(2);
    int64_t G = _scales.size(2);
@@ -715,12 +595,17 @@ 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_);
@@ -729,7 +614,8 @@ 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_);
@@ -738,93 +624,89 @@ 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) {
+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}));
+
  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;  // TODO: add sym quant path
+  constexpr bool sym_quant_act = false;
  using Tin = typename ActDtype<sym_quant_act>::type;
-  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
+  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));
  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);  // sym_a s8s8 is unified to u8s8 with compensation (128)
+  fill_val_stub(Azp_data, 128, M_a);

  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 = /* 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));
+  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));
  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_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);                                                                                   \
+#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);                                                   \
      });

  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);
@@ -834,61 +716,40 @@ inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ inpu
  }
 }

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

-// 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());
+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,83 +1,20 @@
 // 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!");
  }
 };
@@ -86,17 +23,10 @@ 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);
@@ -108,10 +38,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 we use 4x2
+    // oops! 4x4 spills but luckily we use 4x2
    __m512 vbias[COLS];

    // [NOTE]: s8s8 igemm compensation in avx512-vnni
@@ -124,12 +54,14 @@ 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);

@@ -157,7 +89,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
@@ -188,7 +120,8 @@ 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);
@@ -196,19 +129,11 @@ 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(
@@ -226,20 +151,10 @@ 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;
@@ -252,43 +167,26 @@ 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,
@@ -299,22 +197,28 @@ 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);

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

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

  AT_DISPATCH_BOOL(bias != nullptr, has_bias, [&] {
-    parallel_2d(MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
+    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);
+
      // for brgemm, use int32_t for accumulate
      alignas(64) int32_t Ctmp[BLOCK_M * BLOCK_N];

-      loop_2d<int8_t>(mb0, mb1, nb0, nb1, BLOCK_N * K, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
+      for (int i = begin; i < end; ++i) {
+        UNUSED(i);
        int mb_start = mb * BLOCK_M;
        int mb_size = std::min(M - mb_start, BLOCK_M);
        int nb_start = nb * BLOCK_N;
@@ -335,7 +239,10 @@ 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();
@@ -344,47 +251,28 @@ 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);

@@ -393,7 +281,8 @@ 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));
@@ -403,9 +292,13 @@ 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);
      }
    });
  });
@@ -422,14 +315,11 @@ 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,
-    const std::optional<at::Tensor>& bias,
-    at::ScalarType out_dtype,
-    bool is_vnni) {
+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}));
+
  auto packed_w = is_vnni ? mat2 : convert_weight_packed(mat2);

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

  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));
@@ -479,13 +368,10 @@ at::Tensor int8_scaled_mm_cpu(
 }

 // 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) {
+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}));
+
  auto packed_w = is_vnni ? mat2 : convert_weight_packed(mat2);

  CHECK_LAST_DIM_CONTIGUOUS_INPUT(mat1);
@@ -504,10 +390,14 @@ at::Tensor int8_scaled_mm_with_quant(
  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));
@@ -525,9 +415,13 @@ at::Tensor int8_scaled_mm_with_quant(
    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,13 +1,12 @@
 // Adapted from
 // https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu

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

+// clang-format off
+
 namespace {

 // [NOTE]: Fused MoE kernel with AMX
@@ -31,6 +30,109 @@ 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,
@@ -42,7 +144,8 @@ 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) {
@@ -57,7 +160,8 @@ 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
@@ -97,9 +201,7 @@ 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;
  };
@@ -113,7 +215,7 @@ int moe_align_block_size(
      offsets[mb + 1] = sorted_id_size(sorted_ids + mb * BLOCK_M);
    }
  });
-  // TODO: do we need to vecterize this ?
+  // TODO: do we need to vectorize this ?
  for (int mb = 0; mb < num_token_blocks; ++mb) {
    offsets[mb + 1] += offsets[mb];
  }
@@ -134,6 +236,7 @@ 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>;

@@ -166,14 +269,8 @@ 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!");
  }
 };
@@ -182,14 +279,9 @@ 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;

@@ -212,7 +304,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);
@@ -260,16 +352,17 @@ 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(
@@ -283,6 +376,7 @@ 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;
@@ -295,25 +389,16 @@ 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);
      }
    }
  }
@@ -322,13 +407,8 @@ 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!");
  }
 };
@@ -337,13 +417,9 @@ 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;

@@ -356,12 +432,14 @@ 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);

@@ -388,15 +466,17 @@ 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(
@@ -409,6 +489,7 @@ 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;
@@ -421,25 +502,16 @@ 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);
      }
    }
  }
@@ -465,6 +537,7 @@ 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();
@@ -479,36 +552,39 @@ 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
-  parallel_2d(MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
+  at::parallel_for(0, MB * NB, 0, [&](int64_t begin, int64_t end) {
    // get local pointers
-    int tid = get_thread_num();
+    int tid = at::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;

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

      // 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 + nb_upper * BLOCK_N * stride_n;
-      const scalar_t* __restrict__ B1 = packed_w1 + expert_id * stride_e + nb_lower * BLOCK_N * stride_n;
+      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;

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

-      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 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 (use_brgemm) {
@@ -540,7 +616,12 @@ 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];
@@ -556,9 +637,9 @@ void fused_experts_kernel_impl(
            /* ldb   */ n_size,
            /* ldc   */ N);
      }
-    });
+    }

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

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

-    loop_2d<scalar_t>(mb0, mb1, nb0, nb1, BLOCK_N * IC, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
+    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<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;
@@ -625,9 +714,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 (use_brgemm) {
+    if (is_brgemm_used) {
      at::native::cpublas::brgemm_release();
    }
  });
@@ -654,6 +743,7 @@ 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();
@@ -665,29 +755,36 @@ 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
-  parallel_2d(MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
+  at::parallel_for(0, MB * NB, 0, [&](int64_t begin, int64_t end) {
    // get local pointers
-    int tid = get_thread_num();
+    int tid = at::get_thread_num();
    float* __restrict__ C0 = C_tmp + tid * 2 * BLOCK_M * BLOCK_N;
    float* __restrict__ C1 = C0 + BLOCK_M * 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);
+    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);
      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 + nb_upper * BLOCK_N * stride_n;
-      const scalar_t* __restrict__ B1 = packed_w1 + nb_lower * BLOCK_N * stride_n;
+      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;

      if (use_brgemm) {
        // 1.b gemm: C0 = A @ B0
@@ -717,7 +814,12 @@ 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(
@@ -732,9 +834,9 @@ void shared_expert_kernel_impl(
            /* ldb   */ n_size,
            /* ldc   */ N);
      }
-    });
+    }

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

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

-    loop_2d<scalar_t>(mb0, mb1, nb0, nb1, BLOCK_N * IC, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
+    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 = 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;

@@ -792,21 +902,19 @@ 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 =
-          apply_scaling_factor ? fused_experts_out + mb * BLOCK_M * K + nb * BLOCK_N : nullptr;
+      const scalar_t* __restrict__ fused_out = fused_experts_out + mb * BLOCK_M * K + nb * BLOCK_N;
      for (int64_t m = 0; m < m_size; ++m) {
-        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);
+        add_mul_stub(out + m * K, C + m * BLOCK_N, fused_out + m * K, routed_scaling_factor, n_size);
      }
-    });
+    }

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

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

 // common checks
 static inline void check_moe_scales(
@@ -814,10 +922,14 @@ 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<std::vector<int64_t>> block_size,
+    const std::optional<at::Tensor>& a1_scale,
+    const std::optional<at::Tensor>& a2_scale) {
  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.");
@@ -827,16 +939,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) == 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))
+#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)

 // hidden_states: [M, K]
 // w1: [E, 2N, K]
@@ -844,7 +956,6 @@ 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,
@@ -852,13 +963,16 @@ at::Tensor fused_experts_cpu(
    at::Tensor& topk_weights,
    at::Tensor& topk_ids,
    bool inplace,
-    int64_t moe_comp_method,
+    bool use_int8_w8a8,
+    bool use_fp8_w8a16,
    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);

@@ -871,49 +985,32 @@ at::Tensor fused_experts_cpu(
  CHECK_INPUT(w2);
  CHECK_EQ(topk_weights.sizes(), topk_ids.sizes());
  CHECK_DIM(2, hidden_states);
-  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(3, w1);
+  CHECK_DIM(3, w2);
  CHECK_DIM(2, topk_weights);
  CHECK_DIM(2, topk_ids);

  CHECK_EQ(topk_ids.scalar_type(), at::kInt);
-
-  // 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);
+  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 = moe_comp_method == CPUQuantMethod::INT4_W4A8 ? w1_scale.value().size(1) * w1_scale.value().size(3) / 2
-                                                           : w1.size(1) / 2;
+  int64_t N = 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, moe_comp_method == CPUQuantMethod::INT8_W8A8);
-  int64_t packed_N = get_row_size(N, moe_comp_method == CPUQuantMethod::INT8_W8A8);
+  int64_t packed_K = get_row_size(K, use_int8_w8a8);
+  int64_t packed_N = get_row_size(N, use_int8_w8a8);

  // check weight shapes
  CHECK_EQ(w2.size(0), E);
-  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_EQ(w2.size(1), K);
+  CHECK_EQ(packed_w1.size(2), packed_K);
+  CHECK_EQ(packed_w2.size(2), packed_N);
+
  // check scales
-  check_moe_scales(
-      moe_comp_method == CPUQuantMethod::INT8_W8A8,
-      moe_comp_method == CPUQuantMethod::FP8_W8A16,
-      w1_scale,
-      w2_scale,
-      block_size);
+  check_moe_scales(use_int8_w8a8, use_fp8_w8a16, w1_scale, w2_scale, block_size, a1_scale, a2_scale);

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

@@ -934,8 +1031,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`
@@ -967,31 +1064,26 @@ at::Tensor fused_experts_cpu(
  //
  // for fp8 w8a16:
  //   7. intermediate_cache0 : [M * topk, 2N]
-  //   8. B_tmp : [T, MAX_CACHE_BLOCK_SIZE, BLOCK_N, std::max(K, N)]
+  //   8. B_tmp : [T, BLOCK_N, std::max(K, N)]
  //
-  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) +
+  int64_t buffer_size_nbytes = M * topk * N * 2 + M * topk * K * 2 +
+      num_threads * BLOCK_M * K * (use_int8_w8a8 ? 1 : 2) +
      num_threads * 2 * BLOCK_M * BLOCK_N * sizeof(float);

-  if (moe_comp_method == CPUQuantMethod::INT8_W8A8) {
+  if (use_int8_w8a8) {
    buffer_size_nbytes += std::max(M * K, M * topk * N) + M * topk * sizeof(float);
  }
-  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;
+  if (use_fp8_w8a16) {
+    buffer_size_nbytes += M * topk * 2 * N * 2 + num_threads * BLOCK_N * std::max(K, N) * 2;
  }
+
  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 (moe_comp_method == CPUQuantMethod::INT8_W8A8) {
+    if (use_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));
@@ -1015,7 +1107,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,
@@ -1025,7 +1117,7 @@ at::Tensor fused_experts_cpu(
          E,
          topk,
          num_tokens_post_pad);
-    } else if (moe_comp_method == CPUQuantMethod::FP8_W8A16) {
+    } else if (use_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));
@@ -1048,48 +1140,6 @@ 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,
@@ -1113,7 +1163,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,
@@ -1138,37 +1188,34 @@ at::Tensor shared_expert_cpu(
    at::Tensor& hidden_states,
    at::Tensor& w1,
    at::Tensor& w2,
-    const std::optional<at::Tensor>& fused_experts_out,
-    const std::optional<double> routed_scaling_factor,
+    at::Tensor& fused_experts_out,
+    double routed_scaling_factor,
    bool inplace,
    bool use_int8_w8a8,
    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,
+    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,
    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);
@@ -1185,7 +1232,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);
+  check_moe_scales(use_int8_w8a8, use_fp8_w8a16, w1_scale, w2_scale, block_size, a1_scale, a2_scale);

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

@@ -1199,7 +1246,7 @@ at::Tensor shared_expert_cpu(
  //
  // for fp8 w8a16:
  //   5. intermediate_cache0 : [M, 2N]
-  //   6. B_tmp: [T, MAX_CACHE_BLOCK_SIZE, BLOCK_M, max(K, N)]
+  //   6. B_tmp: [T, 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);
@@ -1208,7 +1255,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 * MAX_CACHE_BLOCK_SIZE * BLOCK_M * std::max(K, N) * 2;
+    buffer_size_nbytes += M * 2 * N * 2 + num_threads * BLOCK_M * std::max(K, N) * 2;
  }

  auto buffer = at::empty({buffer_size_nbytes}, hidden_states.options().dtype(at::kChar));
@@ -1236,8 +1283,8 @@ at::Tensor shared_expert_cpu(
          packed_w2.data_ptr<int8_t>(),
          w1s.data_ptr<float>(),
          w2s.data_ptr<float>(),
-          conditional_data_ptr<scalar_t>(fused_experts_out),
-          routed_scaling_factor_value,
+          fused_experts_out.data_ptr<scalar_t>(),
+          routed_scaling_factor,
          M,
          N,
          K);
@@ -1259,8 +1306,8 @@ at::Tensor shared_expert_cpu(
          w2s.data_ptr<float>(),
          block_size_N,
          block_size_K,
-          conditional_data_ptr<scalar_t>(fused_experts_out),
-          routed_scaling_factor_value,
+          fused_experts_out.data_ptr<scalar_t>(),
+          routed_scaling_factor,
          M,
          N,
          K);
@@ -1272,8 +1319,8 @@ at::Tensor shared_expert_cpu(
          hidden_states.data_ptr<scalar_t>(),
          packed_w1.data_ptr<scalar_t>(),
          packed_w2.data_ptr<scalar_t>(),
-          conditional_data_ptr<scalar_t>(fused_experts_out),
-          routed_scaling_factor_value,
+          fused_experts_out.data_ptr<scalar_t>(),
+          routed_scaling_factor,
          M,
          N,
          K);
@@ -1,178 +0,0 @@
-// 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,11 +1,147 @@
 // 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"
+#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

 template <typename scalar_t>
 void fused_experts_fp8_kernel_impl(
@@ -33,6 +169,7 @@ 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();

@@ -46,39 +183,35 @@ 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
-  parallel_2d(MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
+  at::parallel_for(0, MB * NB, 0, [&](int64_t begin, int64_t end) {
    // get local pointers
-    int tid = get_thread_num();
+    int tid = at::get_thread_num();
    scalar_t* __restrict__ A = A_tmp + tid * BLOCK_M * K;

-    loop_2d<at::Float8_e4m3fn>(mb0, mb1, nb0, nb1, BLOCK_N * K, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
+    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 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;
-
-      // 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);
+      const float* __restrict__ Bs = w1s + expert_id * scale_size_N * scale_size_K + (nb / blocks_n_per_group) * scale_size_K;

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

-      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 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);
      }

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

-    if (use_brgemm) {
+    if (is_brgemm_used) {
      at::native::cpublas::brgemm_release();
    }
  });
@@ -108,7 +240,11 @@ 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);
    }
  });

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

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

-    loop_2d<at::Float8_e4m3fn>(mb0, mb1, nb0, nb1, BLOCK_N * IC, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
+    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<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;
@@ -140,18 +284,13 @@ 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;
-
-      // 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);
+      const float* __restrict__ Bs = w2s + expert_id * scale_size_N * scale_size_K + (nb / blocks_n_per_group) * scale_size_K;

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

      // 2.b copy from C to ic2 in original order
      //   and also mul topk_weights in float32
@@ -171,9 +309,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 (use_brgemm) {
+    if (is_brgemm_used) {
      at::native::cpublas::brgemm_release();
    }
  });
@@ -236,6 +374,7 @@ 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();

@@ -246,25 +385,21 @@ 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;

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

-  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) {
+    for (int64_t i = begin; i < end; ++i) {
+      int64_t mb = i / NB;
+      int64_t nb = i % NB;
      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 * B_tmp_size_per_thread + nb_offset * BLOCK_N * K,
+          /*   Btmp         */ B_tmp + tid * BLOCK_N * std::max(K, N),
          /*   Ctmp         */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N,
          /*   scale        */ w1s + (nb / blocks_n_per_group) * scale_size_K,
          /*   M            */ m_size,
@@ -274,9 +409,8 @@ void shared_expert_fp8_kernel_impl(
          /*   ldb          */ n_size,
          /*   ldc          */ 2 * N,
          /*   brg          */ use_brgemm,
-          /*   block_size_K */ block_size_K,
-          /*   do_unpack    */ do_unpack);
-    });
+          /*   block_size_K */ block_size_K);
+    }

    if (use_brgemm) {
      at::native::cpublas::brgemm_release();
@@ -286,7 +420,11 @@ 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);
    }
  });

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

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

-    loop_2d<at::Float8_e4m3fn>(mb0, mb1, nb0, nb1, BLOCK_N * IC, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
+    for (int64_t i = begin; i < end; ++i) {
+      int64_t mb = i / NB2;
+      int64_t nb = i % NB2;
      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 * B_tmp_size_per_thread + nb_offset * BLOCK_N * IC,
+          /*   Btmp         */ B_tmp + tid * BLOCK_N * std::max(K, N),
          /*   Ctmp         */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N,
          /*   scale        */ w2s + (nb / blocks_n_per_group) * scale_size_K,
          /*   M            */ m_size,
@@ -325,18 +462,15 @@ void shared_expert_fp8_kernel_impl(
          /*   ldb          */ n_size,
          /*   ldc          */ BLOCK_N,
          /*   brg          */ use_brgemm,
-          /*   block_size_K */ block_size_K,
-          /*   do_unpack    */ do_unpack);
+          /*   block_size_K */ block_size_K);

      // 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 =
-          apply_scaling_factor ? fused_experts_out + mb * BLOCK_M * K + nb * BLOCK_N : nullptr;
+      const scalar_t* __restrict__ fused_out = fused_experts_out + mb * BLOCK_M * K + nb * BLOCK_N;
      for (int64_t m = 0; m < m_size; ++m) {
-        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);
+        add_mul_stub(out + m * K, C + m * BLOCK_N, fused_out + m * K, routed_scaling_factor, n_size);
      }
-    });
+    }
  });

  if (use_brgemm) {
@@ -1,323 +0,0 @@
-// 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

-// clang-format off
-
 #pragma once

+// clang-format off
+
 #if defined(__AVX512F__) && defined(__AVX512BF16__) && defined(__AMX_BF16__)
 #define CPU_CAPABILITY_AVX512
 #endif
@@ -16,51 +16,38 @@ 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_cvtph_ps(a)
+#define CVT_FP16_TO_FP32(a) \
+    _mm512_cvtps_ph(a, (_MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC))

 // this doesn't handle NaN.
 inline __m512bh cvt_e4m3_bf16_intrinsic_no_nan(__m256i fp8_vec) {
  const __m512i x = _mm512_cvtepu8_epi16(fp8_vec);
-  __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 __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);

  const __mmask32 is_nonzero = _mm512_cmpneq_epi16_mask(x, _mm512_setzero_si512());
  return (__m512bh)_mm512_maskz_mov_epi16(is_nonzero, combined);
@@ -126,77 +113,10 @@ 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)
+#if defined(CPU_CAPABILITY_AVX512) && 0
 inline float vec_reduce_sum(const Vectorized<float>& a) {
  return _mm512_reduce_add_ps(__m512(a));
 }
@@ -216,9 +136,10 @@ 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));
@@ -237,8 +158,9 @@ quantize_row_int8(uint8_t* __restrict__ Aq, float& As, const scalar_t* __restric

 #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);

@@ -278,7 +200,7 @@ inline void quantize_row_int8<at::BFloat16>(
 // 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]);
@@ -371,56 +293,16 @@ 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

-}  // anonymous namespace
+// 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
@@ -1,299 +0,0 @@
-// 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,8 +56,6 @@ 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);
@@ -67,32 +65,21 @@ 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,
-    int64_t moe_comp_method, const std::optional<at::Tensor>& w1_scale,
+    bool use_int8_w8a8, bool use_fp8_w8a16,
+    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, bool is_vnni);
+    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);

 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,  // 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, at::Tensor scales);

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

  // Quantization
-#if defined(__AVX512F__) || defined(__AVX2__) || \
-    (defined(__aarch64__) && !defined(__APPLE__)) || defined(__powerpc64__)
+#if defined(__AVX512F__) || (defined(__aarch64__) && !defined(__APPLE__)) || \
+    defined(__powerpc64__)
  // Helper function to release oneDNN handlers
  ops.def("release_dnnl_matmul_handler(int handler) -> ()",
          &release_dnnl_matmul_handler);
@@ -366,10 +353,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, int moe_comp_method, Tensor? w1_scale, Tensor? w2_scale, "
-      "Tensor? w1_zero, Tensor? w2_zero, int[]? block_size, bool is_vnni) -> "
+      "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) -> "
      "Tensor");
  ops.impl("fused_experts_cpu", torch::kCPU, &fused_experts_cpu);
  ops.def(
@@ -380,9 +367,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {

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

@@ -390,13 +376,6 @@ 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
@@ -232,6 +232,28 @@ 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);
  }
@@ -42,8 +42,8 @@ __device__ __forceinline__ void atomicAdd_half2(half2* address, half2 val) {
 }

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

 __device__ __forceinline__ void atomicAdd(half* address, half val) {
@@ -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.stride(0);
+  const int64_t stride = logits.size(1);
  cudaStream_t stream = at::cuda::getCurrentCUDAStream();

  static int num_sms = 0;
@@ -153,23 +153,14 @@ 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.
+    // Zero the per-group RadixRowState region before launch — only when the
+    // radix path will actually run (max_seq_len > RADIX_THRESHOLD). The
+    // RadixRowState fields (arrival_counter, histograms) are only touched by
+    // radix_topk; the decode/medium paths inside the persistent kernel
+    // operate purely in shared memory and never read these globals, so a
+    // stale workspace is harmless for them.
    //
-    // 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):
+    // Why we need the memset (when needs_cooperative is true):
    //   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 >>
@@ -178,7 +169,7 @@ void launch_persistent_topk(const torch::Tensor& logits,
    //      __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.
-    {
+    if (needs_cooperative) {
      cudaError_t mz_err = cudaMemsetAsync(workspace.data_ptr<uint8_t>(), 0,
                                           state_bytes, stream);
      TORCH_CHECK(mz_err == cudaSuccess,
@@ -243,7 +234,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.stride(0);
+  const int64_t stride = logits.size(1);

  TORCH_CHECK(lengths.numel() == num_rows, "lengths size mismatch");
  TORCH_CHECK(output.size(0) == num_rows && output.size(1) == k,
@@ -553,8 +553,7 @@ 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,"
-      "                  bool is_src_access_order_any=False) -> ()");
+      "                  Tensor sizes) -> ()");
  cache_ops.impl("swap_blocks_batch", torch::kCPU, &swap_blocks_batch);

  // Reshape the key and value tensors and cache them.
@@ -860,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/deployment/sagemaker-entrypoint.sh .
+COPY examples/online_serving/sagemaker-entrypoint.sh .
 RUN chmod +x sagemaker-entrypoint.sh
 ENTRYPOINT ["./sagemaker-entrypoint.sh"]

@@ -507,10 +507,6 @@ 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
@@ -521,10 +517,6 @@ 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
@@ -9,7 +9,7 @@ ARG PYTORCH_AUDIO_BRANCH="v2.9.0"
 ARG PYTORCH_AUDIO_REPO="https://github.com/pytorch/audio.git"
 ARG FA_BRANCH="0e60e394"
 ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git"
-ARG AITER_BRANCH="v0.1.13"
+ARG AITER_BRANCH="v0.1.12.post2"
 ARG AITER_REPO="https://github.com/ROCm/aiter.git"
 ARG MORI_BRANCH="v1.1.0"
 ARG MORI_REPO="https://github.com/ROCm/mori.git"
@@ -132,7 +132,7 @@ The generated timeline is an interactive visualization in the form of an HTML fi

 Example output:

-<iframe src="../assets/contributing/vllm_bench_serve_timeline.html" width="100%" height="600" frameborder="0"></iframe>
+<iframe src="../../assets/contributing/vllm_bench_serve_timeline.html" width="100%" height="600" frameborder="0"></iframe>

 ##### Dataset statistics

@@ -270,36 +270,6 @@ Known supported models (with corresponding benchmarks):

 ## Input Processing

-### fastokens Tokenizer Mode
-
-By default vLLM uses the standard Hugging Face `tokenizers` library to power
-the fast tokenizer (`--tokenizer-mode hf`). For BPE tokenizers (Qwen, Llama,
-DeepSeek, GPT-OSS, etc.) you can switch to the
-[fastokens](https://github.com/crusoecloud/fastokens) Rust backend, a drop-in
-replacement that's substantially faster on encode/decode and on streaming
-detokenization:
-
-```console
-vllm serve Qwen/Qwen3-8B --tokenizer-mode fastokens
-```
-
-Equivalent in the offline API:
-
-```python
-from vllm import LLM
-llm = LLM(model="Qwen/Qwen3-8B", tokenizer_mode="fastokens")
-```
-
-The `fastokens` Python package must be installed; if it isn't, vLLM raises
-a clear `ImportError` at tokenizer load. `fastokens` loads a Hugging Face
-fast tokenizer with its inner Rust tokenizer replaced by the fastokens shim,
-so it is mutually exclusive with non-HF modes such as `mistral` or
-`deepseek_v32`.
-
-Tokenizer-bound workloads — long shared prefixes, bursty short prompts,
-batch detokenization — see the largest wins. If your bottleneck is GPU
-prefill/decode, the tokenizer change is unlikely to be visible end-to-end.
-
 ### Parallel Processing

 You can run input processing in parallel via [API server scale-out](../serving/data_parallel_deployment.md#internal-load-balancing).
@@ -278,7 +278,7 @@ Once your model implements `SupportsTranscription`, you can test the endpoints (
      http://localhost:8000/v1/audio/translations
    ```

-Or check out more examples in [examples/speech_to_text](../../../examples/speech_to_text).
+Or check out more examples in [examples/online_serving](../../../examples/online_serving).

 !!! note
    - If your model handles chunking internally (e.g., via its processor or encoder), set `min_energy_split_window_size=None` in the returned `SpeechToTextConfig` to disable server-side chunking.
@@ -3,7 +3,7 @@
 [Anyscale](https://www.anyscale.com) is a managed, multi-cloud platform developed by the creators of Ray.

 Anyscale automates the entire lifecycle of Ray clusters in your AWS, GCP, or Azure account, delivering the flexibility of open-source Ray
-without the operational overhead of maintaining Kubernetes control planes, configuring autoscalers, managing observability stacks, or manually managing head and worker nodes with helper scripts like [examples/ray_serving/run_cluster.sh](../../../examples/ray_serving/run_cluster.sh).
+without the operational overhead of maintaining Kubernetes control planes, configuring autoscalers, managing observability stacks, or manually managing head and worker nodes with helper scripts like [examples/online_serving/run_cluster.sh](../../../examples/online_serving/run_cluster.sh).

 When serving large language models with vLLM, Anyscale can rapidly provision [production-ready HTTPS endpoints](https://docs.anyscale.com/examples/deploy-ray-serve-llms) or [fault-tolerant batch inference jobs](https://docs.anyscale.com/examples/ray-data-llm).

@@ -17,7 +17,7 @@ Before you begin, ensure that you have the following:

 ## Installing the chart

-This guide uses the Helm chart at [examples/deployment/chart-helm](../../../examples/deployment/chart-helm).
+This guide uses the Helm chart at [examples/online_serving/chart-helm](../../../examples/online_serving/chart-helm).

 To install the chart with the release name `test-vllm`:

@@ -40,7 +40,7 @@ Deploy the following yaml file `lws.yaml`
                command:
                  - sh
                  - -c
-                  - "bash /vllm-workspace/examples/ray_serving/multi-node-serving.sh leader --ray_cluster_size=$(LWS_GROUP_SIZE); 
+                  - "bash /vllm-workspace/examples/online_serving/multi-node-serving.sh leader --ray_cluster_size=$(LWS_GROUP_SIZE); 
                    vllm serve meta-llama/Meta-Llama-3.1-405B-Instruct --port 8080 --tensor-parallel-size 8 --pipeline_parallel_size 2"
                resources:
                  limits:
@@ -73,7 +73,7 @@ Deploy the following yaml file `lws.yaml`
                command:
                  - sh
                  - -c
-                  - "bash /vllm-workspace/examples/ray_serving/multi-node-serving.sh worker --ray_address=$(LWS_LEADER_ADDRESS)"
+                  - "bash /vllm-workspace/examples/online_serving/multi-node-serving.sh worker --ray_address=$(LWS_LEADER_ADDRESS)"
                resources:
                  limits:
                    nvidia.com/gpu: "8"
@@ -36,7 +36,7 @@ pip install -U vllm \
    vllm serve qwen/Qwen1.5-0.5B-Chat --port 8001
    ```

-1. Use the script: [examples/applications/rag/retrieval_augmented_generation_with_langchain.py](../../../examples/applications/rag/retrieval_augmented_generation_with_langchain.py)
+1. Use the script: [examples/online_serving/retrieval_augmented_generation_with_langchain.py](../../../examples/online_serving/retrieval_augmented_generation_with_langchain.py)

 1. Run the script

@@ -74,7 +74,7 @@ pip install vllm \
    vllm serve qwen/Qwen1.5-0.5B-Chat --port 8001
    ```

-1. Use the script: [examples/applications/rag/retrieval_augmented_generation_with_llamaindex.py](../../../examples/applications/rag/retrieval_augmented_generation_with_llamaindex.py)
+1. Use the script: [examples/online_serving/retrieval_augmented_generation_with_llamaindex.py](../../../examples/online_serving/retrieval_augmented_generation_with_llamaindex.py)

 1. Run the script:

@@ -12,7 +12,8 @@ vLLM can be deployed on [RunPod](https://www.runpod.io/), a cloud GPU platform t
 SSH into your RunPod pod and launch the vLLM OpenAI-compatible server:

 ```bash
-vllm serve <model-name> \
+python -m vllm.entrypoints.openai.api_server \
+    --model <model-name> \
    --host 0.0.0.0 \
    --port 8000
 ```
@@ -59,7 +59,7 @@ See the vLLM SkyPilot YAML for serving, [serving.yaml](https://github.com/skypil

      echo 'Starting gradio server...'
      git clone https://github.com/vllm-project/vllm.git || true
-      python vllm/examples/applications/chatbot/gradio_openai_chatbot_webserver.py \
+      python vllm/examples/online_serving/gradio_openai_chatbot_webserver.py \
        -m $MODEL_NAME \
        --port 8811 \
        --model-url http://localhost:8081/v1 \
@@ -305,7 +305,7 @@ It is also possible to access the Llama-3 service with a separate GUI frontend,

      echo 'Starting gradio server...'
      git clone https://github.com/vllm-project/vllm.git || true
-      python vllm/examples/applications/api_client/gradio_openai_chatbot_webserver.py \
+      python vllm/examples/online_serving/gradio_openai_chatbot_webserver.py \
        -m $MODEL_NAME \
        --port 8811 \
        --model-url http://$ENDPOINT/v1 \
@@ -20,7 +20,7 @@ pip install vllm streamlit openai
    vllm serve Qwen/Qwen1.5-0.5B-Chat
    ```

-1. Use the script: [examples/applications/chatbot/streamlit_openai_chatbot_webserver.py](../../../examples/applications/chatbot/streamlit_openai_chatbot_webserver.py)
+1. Use the script: [examples/online_serving/streamlit_openai_chatbot_webserver.py](../../../examples/online_serving/streamlit_openai_chatbot_webserver.py)

 1. Start the streamlit web UI and start to chat:

@@ -78,9 +78,10 @@ Key points from the example YAML:
    - sh
    - -c
    - >
-      bash /vllm-workspace/examples/ray_serving/multi-node-serving.sh leader --ray_cluster_size=2;
-      vllm serve meta-llama/Llama-3.1-405B-Instruct
+      bash /vllm-workspace/examples/online_serving/multi-node-serving.sh leader --ray_cluster_size=2;
+      python3 -m vllm.entrypoints.openai.api_server
        --port 8080
+        --model meta-llama/Llama-3.1-405B-Instruct
        --tensor-parallel-size 8
        --pipeline-parallel-size 2
  ```
@@ -92,7 +93,7 @@ Key points from the example YAML:
    - sh
    - -c
    - >
-      bash /vllm-workspace/examples/ray_serving/multi-node-serving.sh worker --ray_address=$(ENTRY_ADDRESS)
+      bash /vllm-workspace/examples/online_serving/multi-node-serving.sh worker --ray_address=$(ENTRY_ADDRESS)
  ```

 ---
@@ -143,8 +144,8 @@ spec:
                command:
                  - sh
                  - -c
-                  - "bash /vllm-workspace/examples/ray_serving/multi-node-serving.sh leader --ray_cluster_size=2; 
-                    vllm serve meta-llama/Llama-3.1-405B-Instruct --port 8080 --tensor-parallel-size 8 --pipeline-parallel-size 2"
+                  - "bash /vllm-workspace/examples/online_serving/multi-node-serving.sh leader --ray_cluster_size=2; 
+                    python3 -m vllm.entrypoints.openai.api_server --port 8080 --model meta-llama/Llama-3.1-405B-Instruct --tensor-parallel-size 8 --pipeline-parallel-size 2"
                resources:
                  limits:
                    nvidia.com/gpu: "8"
@@ -177,7 +178,7 @@ spec:
                command:
                  - sh
                  - -c
-                  - "bash /vllm-workspace/examples/ray_serving/multi-node-serving.sh worker --ray_address=$(ENTRY_ADDRESS)"
+                  - "bash /vllm-workspace/examples/online_serving/multi-node-serving.sh worker --ray_address=$(ENTRY_ADDRESS)"
                resources:
                  limits:
                    nvidia.com/gpu: "8"
@@ -168,7 +168,7 @@ Priority is **1 = highest** (tried first).

 | Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | MM Prefix | DCP | Attention Types | Compute Cap. |
 | ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | --------- | --- | --------------- | ------------ |
-| `CPU_ATTN` | | fp16, bf16, fp32 | `auto`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ❌ | ❌ | ❌ | All | N/A |
+| `CPU_ATTN` | | fp16, bf16, fp32 | `auto`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | Any | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ❌ | ❌ | ❌ | All | N/A |
 | `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ❌ | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
 | `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ❌ | ✅ | Decoder | 10.x |
 | `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ❌ | ✅ | All | ≥8.0 |
@@ -179,6 +179,7 @@ Priority is **1 = highest** (tried first).
 | `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
 | `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | N/A |
 | `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder, Encoder Only | N/A |
+| `TREE_ATTN` | | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | ❌ | Decoder | Any |
 | `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `int8_per_token_head`, `fp8_per_token_head` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | Any |
 | `TURBOQUANT` | | fp16, bf16 | `turboquant_k8v4`, `turboquant_4bit_nc`, `turboquant_k3v4_nc`, `turboquant_3bit_nc` | 16, 32, 64, 128 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | Any |

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

 > **‡** TRT-LLM Ragged is the default on Blackwell (SM100).
 > On other GPUs, FlashAttention is used as the default.
@@ -222,5 +224,6 @@ MLA decode backends are selected using the standard
 | `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %1 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
 | `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 1, 64 | Any | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
 | `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
+| `TOKENSPEED_MLA` | fp16, bf16 | `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
 | `TRITON_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
 | `XPU_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | Any | 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | Any |
@@ -21,7 +21,7 @@ Let's say we want to serve the popular Qwen model by running `vllm serve Qwen/Qw

 Beyond that, there are two more things vLLM depends on Hugging Face for.

-1. **Tokenizer**: vLLM uses the tokenizer from Hugging Face to tokenize the input text. The tokenizer is loaded using [AutoTokenizer.from_pretrained](https://huggingface.co/docs/transformers/en/model_doc/auto#transformers.AutoTokenizer.from_pretrained) with the `model` argument as the model name and the `--revision` argument as the revision. It is also possible to use a tokenizer from another model by specifying the `--tokenizer` argument in the `vllm serve` command. Other relevant arguments are `--tokenizer-revision` and `--tokenizer-mode`. Setting `--tokenizer-mode fastokens` swaps in a drop-in Rust BPE backend for the HF fast tokenizer (see [fastokens Tokenizer Mode](../configuration/optimization.md#fastokens-tokenizer-mode)). Please check Hugging Face's documentation for the meaning of these arguments. This part of the logic can be found in the [get_tokenizer](https://github.com/vllm-project/vllm/blob/127c07480ecea15e4c2990820c457807ff78a057/vllm/transformers_utils/tokenizer.py#L87) function. After obtaining the tokenizer, notably, vLLM will cache some expensive attributes of the tokenizer in [vllm.tokenizers.hf.get_cached_tokenizer][].
+1. **Tokenizer**: vLLM uses the tokenizer from Hugging Face to tokenize the input text. The tokenizer is loaded using [AutoTokenizer.from_pretrained](https://huggingface.co/docs/transformers/en/model_doc/auto#transformers.AutoTokenizer.from_pretrained) with the `model` argument as the model name and the `--revision` argument as the revision. It is also possible to use a tokenizer from another model by specifying the `--tokenizer` argument in the `vllm serve` command. Other relevant arguments are `--tokenizer-revision` and `--tokenizer-mode`. Please check Hugging Face's documentation for the meaning of these arguments. This part of the logic can be found in the [get_tokenizer](https://github.com/vllm-project/vllm/blob/127c07480ecea15e4c2990820c457807ff78a057/vllm/transformers_utils/tokenizer.py#L87) function. After obtaining the tokenizer, notably, vLLM will cache some expensive attributes of the tokenizer in [vllm.tokenizers.hf.get_cached_tokenizer][].

 2. **Model weight**: vLLM downloads the model weight from the Hugging Face model hub using the `model` argument as the model name and the `--revision` argument as the revision. vLLM provides the argument `--load-format` to control what files to download from the model hub. By default, it will try to load the weights in the safetensors format and fall back to the PyTorch bin format if the safetensors format is not available. We can also pass `--load-format dummy` to skip downloading the weights.
    - It is recommended to use the safetensors format, as it is efficient for loading in distributed inference and also safe from arbitrary code execution. See the [documentation](https://huggingface.co/docs/safetensors/en/index) for more information on the safetensors format. This part of the logic can be found [here](https://github.com/vllm-project/vllm/blob/10b67d865d92e376956345becafc249d4c3c0ab7/vllm/model_executor/model_loader/loader.py#L385). Please note that:
@@ -62,7 +62,8 @@ The filesystem resolver is installed with vLLM by default and enables loading Lo
 3. **Start vLLM server**:
   Your base model can be `meta-llama/Llama-2-7b-hf`. Please make sure you set up the Hugging Face token in your env var `export HF_TOKEN=xxx235`.
   ```bash
-   vllm serve your-base-model \
+   python -m vllm.entrypoints.openai.api_server \
+       --model your-base-model \
       --enable-lora
   ```

@@ -16,7 +16,7 @@ User-set flags take precedence over optimization level defaults.

 ```bash
 # CLI usage
-vllm serve RedHatAI/Llama-3.2-1B-FP8 -O1
+python -m vllm.entrypoints.api_server --model RedHatAI/Llama-3.2-1B-FP8 -O1

 # Python API usage
 from vllm.entrypoints.llm import LLM
@@ -88,7 +88,7 @@ pip install "vllm>=0.9.2"
 #### Proxy (e.g. 10.0.1.1)

 ```shell
-cd {your vllm directory}/examples/disaggregated/p2p_nccl_xpyd/
+cd {your vllm directory}/examples/online_serving/disaggregated_serving_p2p_nccl_xpyd/
 python3 disagg_proxy_p2p_nccl_xpyd.py &
 ```

@@ -181,7 +181,7 @@ python3 disagg_proxy_p2p_nccl_xpyd.py &
 #### Proxy (e.g. 10.0.1.1)

 ```shell
-cd {your vllm directory}/examples/disaggregated/p2p_nccl_xpyd/
+cd {your vllm directory}/examples/online_serving/disaggregated_serving_p2p_nccl_xpyd/
 python3 disagg_proxy_p2p_nccl_xpyd.py &
 ```

@@ -1,17 +1,7 @@
 # Examples

-vLLM's examples are organized into the following categories:
+vLLM's examples are split into three categories:

- **[`basic/`](../../examples/basic)** – Minimal examples for offline inference and online serving.
- **[`generate/`](../../examples/generate)** – Text generation examples, including multimodal models.
- **[`pooling/`](../../examples/pooling)** – Examples for embedding, classification, scoring, reward, etc.
- **[`speech_to_text/`](../../examples/speech_to_text)** – Speech transcription, translation and real-time audio examples.
- **[`features/`](../../examples/features)** – Demonstrations of individual vLLM features: automatic prefix caching, speculative decoding, LoRA, structured outputs, prompt embedding, pause/resume, batch invariance, KV events, data parallelism, and more.
- **[`reasoning/`](../../examples/reasoning)** – Examples for reasoning with vLLM.
- **[`tool_calling/`](../../examples/tool_calling)** – Examples for function/tool calling with vLLM.
- **[`applications/`](../../examples/applications)** – Application examples such as chatbots and RAG (Retrieval-Augmented Generation).
- **[`rl/`](../../examples/rl)** – Reinforcement learning examples.
- **[`deployment/`](../../examples/deployment)** – Examples for deploying vLLM in production.
- **[`ray_serving/`](../../examples/ray_serving)** – Scalable serving using Ray.
- **[`disaggregated/`](../../examples/disaggregated)** – Examples for disaggregated serving (separate prefill and decode), including various kv cache connectors (LMCache, Mooncake, FlexKV, P2P NCCL) and failure recovery.
- **[`observability/`](../../examples/observability)** – Metrics, logging, tracing (OpenTelemetry), and dashboards (Grafana, Perses).
+- If you are using vLLM from within Python code, see the [Offline Inference](../../examples/offline_inference) section.
+- If you are using vLLM from an HTTP application or client, see the [Online Serving](../../examples/online_serving) section.
+- For examples of using some of vLLM's advanced features (e.g. LMCache or Tensorizer) which are not specific to either of the above use cases, see the [Others](../../examples/others) section.
@@ -36,10 +36,10 @@ The current reference pathway is **ExampleConnector**.
 Below ready-to-run scripts shows the workflow:

 1 Encoder instance + 1 PD instance:
-`examples/disaggregated/disaggregated_encoder/disagg_1e1pd_example.sh`
+`examples/online_serving/disaggregated_encoder/disagg_1e1pd_example.sh`

 1 Encoder instance + 1 Prefill instance + 1 Decode instance:
-`examples/disaggregated/disaggregated_encoder/disagg_1e1p1d_example.sh`
+`examples/online_serving/disaggregated_encoder/disagg_1e1p1d_example.sh`

 ---

@@ -17,15 +17,15 @@ Two main reasons:

 ## Usage example

-Please refer to [examples/disaggregated/disaggregated_prefill.sh](../../examples/disaggregated/disaggregated_prefill.sh) for the example usage of disaggregated prefilling.
+Please refer to [examples/online_serving/disaggregated_prefill.sh](../../examples/online_serving/disaggregated_prefill.sh) for the example usage of disaggregated prefilling.

 Now supports 6 types of connectors:

- **ExampleConnector**: refer to [examples/disaggregated/example_connector/run.sh](../../examples/disaggregated/example_connector/run.sh) for the example usage of ExampleConnector disaggregated prefilling.
- **LMCacheConnectorV1**: refer to [examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh](../../examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh) for the example usage of LMCacheConnectorV1 disaggregated prefilling which uses NIXL as the underlying KV transmission.
+- **ExampleConnector**: refer to [examples/offline_inference/disaggregated-prefill-v1/run.sh](../../examples/offline_inference/disaggregated-prefill-v1/run.sh) for the example usage of ExampleConnector disaggregated prefilling.
+- **LMCacheConnectorV1**: refer to [examples/others/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh](../../examples/others/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh) for the example usage of LMCacheConnectorV1 disaggregated prefilling which uses NIXL as the underlying KV transmission.
 - **NixlConnector**: refer to [tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh](../../tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh) for the example usage of NixlConnector disaggregated prefilling which support fully async send/recv. For detailed usage guide, see [NixlConnector Usage Guide](nixl_connector_usage.md). For feature compatibility details, see [NixlConnector Compatibility Matrix](nixl_connector_compatibility.md).
- **P2pNcclConnector**: refer to [examples/disaggregated/p2p_nccl_xpyd/disagg_example_p2p_nccl_xpyd.sh](../../examples/disaggregated/p2p_nccl_xpyd/disagg_example_p2p_nccl_xpyd.sh) for the example usage of P2pNcclConnector disaggregated prefilling.
- **MooncakeConnector**: refer to [examples/disaggregated/mooncake_connector/run_mooncake_connector.sh](../../examples/disaggregated/mooncake_connector/run_mooncake_connector.sh) for the example usage of MooncakeConnector disaggregated prefilling. For detailed usage guide, see [MooncakeConnector Usage Guide](mooncake_connector_usage.md).
+- **P2pNcclConnector**: refer to [examples/online_serving/disaggregated_serving_p2p_nccl_xpyd/disagg_example_p2p_nccl_xpyd.sh](../../examples/online_serving/disaggregated_serving_p2p_nccl_xpyd/disagg_example_p2p_nccl_xpyd.sh) for the example usage of P2pNcclConnector disaggregated prefilling.
+- **MooncakeConnector**: refer to [examples/online_serving/disaggregated_serving/mooncake_connector/run_mooncake_connector.sh](../../examples/online_serving/disaggregated_serving/mooncake_connector/run_mooncake_connector.sh) for the example usage of ExampleConnector disaggregated prefilling. For detailed usage guide, see [MooncakeConnector Usage Guide](mooncake_connector_usage.md).
 - **MultiConnector**: take advantage of the kv_connector_extra_config: dict[str, Any] already present in KVTransferConfig to stash all the connectors we want in an ordered list of kwargs.such as:

  ```bash
@@ -44,7 +44,7 @@ For NixlConnector, you may also specify one or multiple NIXL_Backend. Such as:
  --kv-transfer-config '{"kv_connector":"OffloadingConnector","kv_role":"kv_both","kv_connector_extra_config":{"block_size": 64, "cpu_bytes_to_use": 1000000000}}'
  ```

- **FlexKVConnectorV1**: refer to [examples/disaggregated/flexkv_connector/prefix_caching_flexkv.py](../../examples/disaggregated/flexkv_connector/prefix_caching_flexkv.py) for the example usage of FlexKVConnectorV1. FlexKV is a distributed KV Store and multi-level cache management system for ultra-large-scale LLM inference.
+- **FlexKVConnectorV1**: refer to [examples/offline_inference/prefix_caching_flexkv.py](../../examples/offline_inference/prefix_caching_flexkv.py) for the example usage of FlexKVConnectorV1. FlexKV is a distributed KV Store and multi-level cache management system for ultra-large-scale LLM inference.

  ```bash
  --kv-transfer-config '{"kv_connector":"FlexKVConnectorV1","kv_role":"kv_both"}'
@@ -31,7 +31,7 @@ vllm serve Qwen/Qwen2.5-7B-Instruct --port 8020 --kv-transfer-config '{"kv_conne
 ### Proxy

 ```bash
-python examples/disaggregated/disaggregated_serving/mooncake_connector/mooncake_connector_proxy.py --prefill http://192.168.0.2:8010 --decode http://192.168.0.3:8020
+python examples/online_serving/disaggregated_serving/mooncake_connector/mooncake_connector_proxy.py --prefill http://192.168.0.2:8010 --decode http://192.168.0.3:8020
 ```

 Now you can send requests to the proxy server through port 8000.
@@ -65,5 +65,5 @@ Now you can send requests to the proxy server through port 8000.

 Refer to these example scripts in the vLLM repository:

- [run_mooncake_connector.sh](../../examples/disaggregated/mooncake_connector/run_mooncake_connector.sh)
- [mooncake_connector_proxy.py](../../examples/disaggregated/mooncake_connector/mooncake_connector_proxy.py)
+- [run_mooncake_connector.sh](../../examples/online_serving/disaggregated_serving/mooncake_connector/run_mooncake_connector.sh)
+- [mooncake_connector_proxy.py](../../examples/online_serving/disaggregated_serving/mooncake_connector/mooncake_connector_proxy.py)
@@ -13,7 +13,7 @@ Install the NIXL library: `uv pip install nixl`, as a quick start on Nvidia plat
 - Refer to [NIXL official repository](https://github.com/ai-dynamo/nixl) for more installation instructions
 - The specified required NIXL version can be found in [requirements/kv_connectors.txt](../../requirements/kv_connectors.txt) and other relevant config files

-For ROCm platform, the [ROCm docker file](../../docker/Dockerfile.rocm) includes RIXL and ucx already.
+For ROCm platform, the [base ROCm docker file](../../docker/Dockerfile.rocm_base) includes RIXL and ucx already.

 - Refer to [RIXL official repository](https://github.com/rocm/rixl) for more information
 - The supportive libraries for RIXL can be found in [requirements/kv_connectors_rocm.txt](../../requirements/kv_connectors_rocm.txt)
@@ -47,7 +47,7 @@ After installing AutoAWQ, you are ready to quantize a model. Please refer to the
 To run an AWQ model with vLLM, you can use [TheBloke/Llama-2-7b-Chat-AWQ](https://huggingface.co/TheBloke/Llama-2-7b-Chat-AWQ) with the following command:

 ```bash
-python examples/deployment/llm_engine_example.py \
+python examples/offline_inference/llm_engine_example.py \
    --model TheBloke/Llama-2-7b-Chat-AWQ \
    --quantization awq
 ```
@@ -58,7 +58,7 @@ Here is an example of how to quantize `meta-llama/Llama-3.2-1B-Instruct`:
 To run an GPTQModel quantized model with vLLM, you can use [DeepSeek-R1-Distill-Qwen-7B-gptqmodel-4bit-vortex-v2](https://huggingface.co/ModelCloud/DeepSeek-R1-Distill-Qwen-7B-gptqmodel-4bit-vortex-v2) with the following command:

 ```bash
-python examples/deployment/llm_engine_example.py \
+python examples/offline_inference/llm_engine_example.py \
    --model ModelCloud/DeepSeek-R1-Distill-Qwen-7B-gptqmodel-4bit-vortex-v2
 ```

@@ -157,7 +157,7 @@ OpenAI Python client library does not officially support `reasoning` attribute f
            print(content, end="", flush=True)
    ```

-Remember to check whether the `reasoning` exists in the response before accessing it. You could check out the [example](https://github.com/vllm-project/vllm/blob/main/examples/reasoning/openai_chat_completion_with_reasoning_streaming.py).
+Remember to check whether the `reasoning` exists in the response before accessing it. You could check out the [example](https://github.com/vllm-project/vllm/blob/main/examples/online_serving/openai_chat_completion_with_reasoning_streaming.py).

 ## Tool Calling

@@ -72,17 +72,6 @@ only apply to model-based methods such as `draft_model`, `mtp`, `eagle3`, and
 | `rejection_sample_method` | `string` | `strict` | `strict`, `probabilistic`, or `synthetic`. |
 | `synthetic_acceptance_rate` | `float` | `None` | Average acceptance rate to target when `rejection_sample_method` is `synthetic`. Valid range is `[0, 1]`. |

-!!! note
-    Gemma 4 assistant checkpoints are handled as Gemma 4 MTP speculators, not
-    as generic draft models. Use `"method": "mtp"` with the assistant
-    checkpoint in `model`, as shown in the [MTP guide](mtp.md#gemma-4-assistant-models).
-
-    If startup logs show `SpeculativeConfig(method='draft_model', ...)` for a
-    Gemma 4 assistant checkpoint, the installed vLLM version does not include
-    Gemma 4 MTP support for that path. Upgrade to a version that includes
-    Gemma 4 MTP support instead of forcing the assistant checkpoint through
-    generic draft-model speculative decoding.
-
 ### Method-specific keys

 #### N-gram
@@ -9,31 +9,6 @@ MTP is useful when:
 - Your model natively supports MTP.
 - You want model-based speculative decoding with minimal extra configuration.

-## Gemma 4 Assistant Models
-
-Gemma 4 assistant checkpoints use vLLM's Gemma 4 MTP path. They are not generic
-draft models, even though they are passed through the `model` field in
-`--speculative-config`.
-
-Use `"method": "mtp"` when serving Gemma 4 with an assistant checkpoint:
-
-```bash
-vllm serve google/gemma-4-E2B-it \
-    --tensor-parallel-size 1 \
-    --max-model-len 8192 \
-    --speculative-config '{"method":"mtp","model":"gg-hf-am/gemma-4-E2B-it-assistant","num_speculative_tokens":1}'
-```
-
-The E2B, E4B, 26B-A4B, and 31B Gemma 4 IT assistant checkpoints are supported
-when their configuration uses `model_type: gemma4_assistant`. vLLM maps those
-checkpoints to `Gemma4MTPModel` internally and wires the assistant layers to
-share KV cache with the target model.
-
-If an older vLLM release logs `SpeculativeConfig(method='draft_model', ...)`
-for a Gemma 4 assistant checkpoint, that release is treating the assistant as a
-generic draft model and may fail during initialization for multimodal Gemma 4
-targets. Upgrade to a version with Gemma 4 MTP support instead.
-
 ## Offline Example

 ```python
@@ -88,7 +88,7 @@ vllm serve facebook/opt-13b \
     -tp=8
 ```

-By default, a ray instance will be launched automatically if no existing one is detected in the system, with `num-gpus` equals to `parallel_config.world_size`. We recommend properly starting a ray cluster before execution, referring to the [examples/ray_serving/run_cluster.sh](https://github.com/vllm-project/vllm/blob/main/examples/ray_serving/run_cluster.sh) helper script.
+By default, a ray instance will be launched automatically if no existing one is detected in the system, with `num-gpus` equals to `parallel_config.world_size`. We recommend properly starting a ray cluster before execution, referring to the [examples/online_serving/run_cluster.sh](https://github.com/vllm-project/vllm/blob/main/examples/online_serving/run_cluster.sh) helper script.

 --8<-- [end:supported-features]
 --8<-- [start:distributed-backend]
@@ -0,0 +1,24 @@
+#!/bin/bash
+# SPDX-License-Identifier: Apache-2.0
+# Skip PR builds unless the PR has the "documentation" or "ready" label.
+# Used by Read the Docs (see .readthedocs.yaml).
+
+if [[ "$READTHEDOCS_VERSION_TYPE" != "external" ]]; then
+  exit 0
+fi
+
+PR_URL="https://api.github.com/repos/vllm-project/vllm/pulls/${READTHEDOCS_VERSION}"
+CURL_ARGS=(-s -o /tmp/pr_response.json -w "%{http_code}")
+if [[ -n "$GITHUB_TOKEN" ]]; then
+  CURL_ARGS+=(-H "Authorization: token ${GITHUB_TOKEN}")
+fi
+HTTP_CODE=$(curl "${CURL_ARGS[@]}" "$PR_URL")
+
+if [[ "$HTTP_CODE" -ne 200 ]]; then
+  echo "GitHub API returned HTTP ${HTTP_CODE}, proceeding with build."
+elif grep -qE '"name": *"(documentation|ready)"' /tmp/pr_response.json; then
+  echo "Found required label, proceeding with build."
+else
+  echo "PR #${READTHEDOCS_VERSION} lacks 'documentation' or 'ready' label, cancelling build."
+  exit 1
+fi
@@ -14,7 +14,7 @@ To install `tensorizer`, run `pip install vllm[tensorizer]`.
 ## The basics

 To load a model using Tensorizer, the model first needs to be serialized by
-Tensorizer. [The example script](../../../examples/features/tensorize_vllm_model.py) takes care of this process.
+Tensorizer. [The example script](../../examples/others/tensorize_vllm_model.md) takes care of this process.

 Let's walk through a basic example by serializing `facebook/opt-125m` using the script, and then loading it for inference.

@@ -25,7 +25,7 @@ CLI arguments. The docstring for the script itself explains the CLI args
 and how to use it properly in great detail, and we'll use one of the examples from the docstring directly, assuming we want to serialize and save our model at our S3 bucket example `s3://my-bucket`:

 ```bash
-python examples/features/tensorize_vllm_model.py \
+python examples/others/tensorize_vllm_model.py \
   --model facebook/opt-125m \
   serialize \
   --serialized-directory s3://my-bucket \
@@ -35,7 +35,7 @@ python examples/features/tensorize_vllm_model.py \
 This saves the model tensors at `s3://my-bucket/vllm/facebook/opt-125m/v1`. If you intend on applying a LoRA adapter to your tensorized model, you can pass the HF id of the LoRA adapter in the above command, and the artifacts will be saved there too:

 ```bash
-python examples/features/tensorize_vllm_model.py \
+python examples/others/tensorize_vllm_model.py \
   --model facebook/opt-125m \
   --lora-path <lora_id> \
   serialize \
@@ -71,7 +71,7 @@ llm = LLM(
 As an example, CPU concurrency can be limited when serializing with `tensorizer` via the `limit_cpu_concurrency` parameter in the initializer for `TensorSerializer`. To set `limit_cpu_concurrency` to some arbitrary value, you would do so like this when serializing:

 ```bash
-python examples/features/tensorize_vllm_model.py \
+python examples/others/tensorize_vllm_model.py \
   --model facebook/opt-125m \
   --lora-path <lora_id> \
   serialize \
@@ -10,40 +10,39 @@

 ### Text-only Language Models

-| Model                                              | Architecture                                     | BF16/FP16/Dynamic FP8 | Compressed_tensors FP8 | MXFP4 |
-| -------------------------------------------------- | ------------------------------------------------ | --------------------- | ---------------------- | ----- |
-| openai/gpt-oss-20b                                 | GPTForCausalLM                                   |                       |                        | ✅    |
-| openai/gpt-oss-120b                                | GPTForCausalLM                                   |                       |                        | ✅    |
-| deepseek-ai/DeepSeek-R1-Distill-Llama-8B           | LlamaForCausalLM                                 | ✅                    |                        |       |
-| deepseek-ai/DeepSeek-R1-Distill-Qwen-14B           | QwenForCausalLM                                  | ✅                    |                        |       |
-| deepseek-ai/DeepSeek-R1-Distill-Qwen-32B           | QwenForCausalLM                                  | ✅                    |                        |       |
-| deepseek-ai/DeepSeek-R1-Distill-Llama-70B          | LlamaForCausalLM                                 | ✅                    |                        |       |
-| Qwen/Qwen2.5-72B-Instruct                          | Qwen2ForCausalLM                                 | ✅                    |                        |       |
-| Qwen/Qwen3-14B                                     | Qwen3ForCausalLM                                 | ✅                    |                        |       |
-| Qwen/Qwen3-32B                                     | Qwen3ForCausalLM                                 | ✅                    |                        |       |
-| Qwen/Qwen3-30B-A3B                                 | Qwen3ForCausalLM                                 | ✅                    |                        |       |
-| Qwen/Qwen3-30B-A3B-GPTQ-Int4                       | Qwen3ForCausalLM                                 | ✅                    |                        |       |
-| Qwen/Qwen3-coder-30B-A3B-Instruct                  | Qwen3ForCausalLM                                 | ✅                    |                        |       |
-| Qwen/QwQ-32B                                       | QwenForCausalLM                                  | ✅                    |                        |       |
-| deepseek-ai/DeepSeek-V2-Lite                       | DeepSeekForCausalLM                              | ✅                    |                        |       |
-| meta-llama/Llama-3.1-8B-Instruct                   | LlamaForCausalLM                                 | ✅                    |                        |       |
-| baichuan-inc/Baichuan2-13B-Chat                    | BaichuanForCausalLM                              | ✅                    |                        |       |
-| THUDM/GLM-4-9B-chat                                | GLMForCausalLM                                   | ✅                    |                        |       |
-| THUDM/CodeGeex4-All-9B                             | CodeGeexForCausalLM                              | ✅                    |                        |       |
-| chuhac/TeleChat2-35B                               | LlamaForCausalLM (TeleChat2 based on Llama arch) | ✅                    |                        |       |
-| 01-ai/Yi1.5-34B-Chat                               | YiForCausalLM                                    | ✅                    |                        |       |
-| THUDM/CodeGeex4-All-9B                             | CodeGeexForCausalLM                              | ✅                    |                        |       |
-| deepseek-ai/DeepSeek-Coder-33B-base                | DeepSeekCoderForCausalLM                         | ✅                    |                        |       |
-| baichuan-inc/Baichuan2-13B-Chat                    | BaichuanForCausalLM                              | ✅                    |                        |       |
-| meta-llama/Llama-2-13b-chat-hf                     | LlamaForCausalLM                                 | ✅                    |                        |       |
-| THUDM/CodeGeex4-All-9B                             | CodeGeexForCausalLM                              | ✅                    |                        |       |
-| Qwen/Qwen1.5-14B-Chat                              | QwenForCausalLM                                  | ✅                    |                        |       |
-| Qwen/Qwen1.5-32B-Chat                              | QwenForCausalLM                                  | ✅                    |                        |       |
-| RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8-dynamic    | LlamaForCausalLM                                 |                       | ✅                     |       |
+| Model                                     | Architecture                                         | FP16 | Dynamic FP8 | MXFP4 |
+| ----------------------------------------- | ---------------------------------------------------- | ---- | ----------- | ----- |
+| openai/gpt-oss-20b                        | GPTForCausalLM                                       |      |             | ✅    |
+| openai/gpt-oss-120b                       | GPTForCausalLM                                       |      |             | ✅    |
+| deepseek-ai/DeepSeek-R1-Distill-Llama-8B  | LlamaForCausalLM                                     | ✅   | ✅          |       |
+| deepseek-ai/DeepSeek-R1-Distill-Qwen-14B  | QwenForCausalLM                                      | ✅   | ✅          |       |
+| deepseek-ai/DeepSeek-R1-Distill-Qwen-32B  | QwenForCausalLM                                      | ✅   | ✅          |       |
+| deepseek-ai/DeepSeek-R1-Distill-Llama-70B | LlamaForCausalLM                                     | ✅   | ✅          |       |
+| Qwen/Qwen2.5-72B-Instruct                 | Qwen2ForCausalLM                                     | ✅   | ✅          |       |
+| Qwen/Qwen3-14B                            | Qwen3ForCausalLM                                     | ✅   | ✅          |       |
+| Qwen/Qwen3-32B                            | Qwen3ForCausalLM                                     | ✅   | ✅          |       |
+| Qwen/Qwen3-30B-A3B                        | Qwen3ForCausalLM                                     | ✅   | ✅          |       |
+| Qwen/Qwen3-30B-A3B-GPTQ-Int4              | Qwen3ForCausalLM                                     | ✅   | ✅          |       |
+| Qwen/Qwen3-coder-30B-A3B-Instruct         | Qwen3ForCausalLM                                     | ✅   | ✅          |       |
+| Qwen/QwQ-32B                              | QwenForCausalLM                                      | ✅   | ✅          |       |
+| deepseek-ai/DeepSeek-V2-Lite              | DeepSeekForCausalLM                                  | ✅   | ✅          |       |
+| meta-llama/Llama-3.1-8B-Instruct          | LlamaForCausalLM                                     | ✅   | ✅          |       |
+| baichuan-inc/Baichuan2-13B-Chat           | BaichuanForCausalLM                                  | ✅   | ✅          |       |
+| THUDM/GLM-4-9B-chat                       | GLMForCausalLM                                       | ✅   | ✅          |       |
+| THUDM/CodeGeex4-All-9B                    | CodeGeexForCausalLM                                  | ✅   | ✅          |       |
+| chuhac/TeleChat2-35B                      | LlamaForCausalLM (TeleChat2 based on Llama arch)     | ✅   | ✅          |       |
+| 01-ai/Yi1.5-34B-Chat                      | YiForCausalLM                                        | ✅   | ✅          |       |
+| THUDM/CodeGeex4-All-9B                    | CodeGeexForCausalLM                                  | ✅   | ✅          |       |
+| deepseek-ai/DeepSeek-Coder-33B-base       | DeepSeekCoderForCausalLM                             | ✅   | ✅          |       |
+| baichuan-inc/Baichuan2-13B-Chat           | BaichuanForCausalLM                                  | ✅   | ✅          |       |
+| meta-llama/Llama-2-13b-chat-hf            | LlamaForCausalLM                                     | ✅   | ✅          |       |
+| THUDM/CodeGeex4-All-9B                    | CodeGeexForCausalLM                                  | ✅   | ✅          |       |
+| Qwen/Qwen1.5-14B-Chat                     | QwenForCausalLM                                      | ✅   | ✅          |       |
+| Qwen/Qwen1.5-32B-Chat                     | QwenForCausalLM                                      | ✅   | ✅          |       |

 ### Multimodal Language Models

-| Model                        | Architecture                     | BF16 | Dynamic FP8 | MXFP4 |
+| Model                        | Architecture                     | FP16 | Dynamic FP8 | MXFP4 |
 | ---------------------------- | -------------------------------- | ---- | ----------- | ----- |
 | OpenGVLab/InternVL3_5-8B     | InternVLForConditionalGeneration | ✅   | ✅          |       |
 | OpenGVLab/InternVL3_5-14B    | InternVLForConditionalGeneration | ✅   | ✅          |       |
@@ -56,7 +55,7 @@

 ### Embedding and Reranker Language Models

-| Model                   | Architecture                   | BF16 | Dynamic FP8 | MXFP4 |
+| Model                   | Architecture                   | FP16 | Dynamic FP8 | MXFP4 |
 | ----------------------- | ------------------------------ | ---- | ----------- | ----- |
 | Qwen/Qwen3-Embedding-8B | Qwen3ForTextEmbedding          | ✅   | ✅          |       |
 | Qwen/Qwen3-Reranker-8B  | Qwen3ForSequenceClassification | ✅   | ✅          |       |
@@ -41,7 +41,6 @@ The score models is designed to compute similarity scores between two input prom
 | `GemmaForSequenceClassification` | Gemma-based | `BAAI/bge-reranker-v2-gemma`(see note), etc. | [bge-reranker-v2-gemma.jinja](../../../examples/pooling/score/template/bge-reranker-v2-gemma.jinja) | ✅︎ | ✅︎ |
 | `GteNewForSequenceClassification` | mGTE-TRM (see note) | `Alibaba-NLP/gte-multilingual-reranker-base`, etc. | N/A | | |
 | `LlamaBidirectionalForSequenceClassification`<sup>C</sup> | Llama-based with bidirectional attention | `nvidia/llama-nemotron-rerank-1b-v2`, etc. | [nemotron-rerank.jinja](../../../examples/pooling/score/template/nemotron-rerank.jinja) | ✅︎ | ✅︎ |
-| `ModernBertForSequenceClassification` | ModernBERT-based | `Alibaba-NLP/gte-reranker-modernbert-base`, etc. | N/A | | |
 | `Qwen2ForSequenceClassification`<sup>C</sup> | Qwen2-based | `mixedbread-ai/mxbai-rerank-base-v2`(see note), etc. | [mxbai_rerank_v2.jinja](../../../examples/pooling/score/template/mxbai_rerank_v2.jinja) | ✅︎ | ✅︎ |
 | `Qwen3ForSequenceClassification`<sup>C</sup> | Qwen3-based | `tomaarsen/Qwen3-Reranker-0.6B-seq-cls`, `Qwen/Qwen3-Reranker-0.6B`(see note), etc. | [qwen3_reranker.jinja](../../../examples/pooling/score/template/qwen3_reranker.jinja) | ✅︎ | ✅︎ |
 | `RobertaForSequenceClassification` | RoBERTa-based | `cross-encoder/quora-roberta-base`, etc. | N/A | | |
@@ -378,7 +378,7 @@ th {
 | `BloomForCausalLM` | BLOOM, BLOOMZ, BLOOMChat | `bigscience/bloom`, `bigscience/bloomz`, etc. | | ✅︎ |
 | `ChatGLMModel`, `ChatGLMForConditionalGeneration` | ChatGLM | `zai-org/chatglm2-6b`, `zai-org/chatglm3-6b`, `thu-coai/ShieldLM-6B-chatglm3`, etc. | ✅︎ | ✅︎ |
 | `CohereForCausalLM`, `Cohere2ForCausalLM` | Command-R, Command-A | `CohereLabs/c4ai-command-r-v01`, `CohereLabs/c4ai-command-r7b-12-2024`, `CohereLabs/c4ai-command-a-03-2025`, `CohereLabs/command-a-reasoning-08-2025`, etc. | ✅︎ | ✅︎ |
-| `Cohere2MoeForCausalLM` | Command (MoE) | (model checkpoints loaded with `trust_remote_code=True`) | ✅︎ | ✅︎ |
+| `CohereMoeForCausalLM` | Command (MoE) | (model checkpoints loaded with `trust_remote_code=True`) | ✅︎ | ✅︎ |
 | `CwmForCausalLM` | CWM | `facebook/cwm`, etc. | ✅︎ | ✅︎ |
 | `DbrxForCausalLM` | DBRX | `databricks/dbrx-base`, `databricks/dbrx-instruct`, etc. | | ✅︎ |
 | `DeciLMForCausalLM` | DeciLM | `nvidia/Llama-3_3-Nemotron-Super-49B-v1`, etc. | ✅︎ | ✅︎ |
@@ -659,9 +659,6 @@ Some models are supported only via the [Transformers modeling backend](#transfor
    For `Gemma4ForConditionalGeneration`:
    - audio input is only supported by the `gemma-4-E2B` and `gemma-4-E4B` variants.
    - The model does not ingest videos directly. However, vLLM’s Gemma 4 implementation supports video inputs by handling video processing internally. Users can send videos directly in the message structure to vLLM, where they are converted into text and image frames before being passed to the model.
-    - Gemma 4 assistant checkpoints for speculative decoding use vLLM's Gemma
-      4 MTP path, not generic draft-model speculative decoding. See the
-      [Gemma 4 assistant model MTP example](../features/speculative_decoding/mtp.md#gemma-4-assistant-models).

 !!! note
    For `InternVLChatModel`, only InternVL2.5 with Qwen2.5 text backbone (`OpenGVLab/InternVL2.5-1B` etc.), InternVL3 and InternVL3.5 have video inputs support currently.
@@ -4,11 +4,11 @@ For general troubleshooting, see [Troubleshooting](../usage/troubleshooting.md).

 ## Verify inter-node GPU communication

-After you start the Ray cluster, verify GPU-to-GPU communication across nodes. Proper configuration can be non-trivial. For more information, see [troubleshooting script](../usage/troubleshooting.md#incorrect-hardwaredriver). If you need additional environment variables for communication configuration, append them to [examples/ray_serving/run_cluster.sh](../../examples/ray_serving/run_cluster.sh), for example `-e NCCL_SOCKET_IFNAME=eth0`. Setting environment variables during cluster creation is recommended because the variables propagate to all nodes. In contrast, setting environment variables in the shell affects only the local node. For more information, see <https://github.com/vllm-project/vllm/issues/6803>.
+After you start the Ray cluster, verify GPU-to-GPU communication across nodes. Proper configuration can be non-trivial. For more information, see [troubleshooting script](../usage/troubleshooting.md#incorrect-hardwaredriver). If you need additional environment variables for communication configuration, append them to [examples/online_serving/run_cluster.sh](../../examples/online_serving/run_cluster.sh), for example `-e NCCL_SOCKET_IFNAME=eth0`. Setting environment variables during cluster creation is recommended because the variables propagate to all nodes. In contrast, setting environment variables in the shell affects only the local node. For more information, see <https://github.com/vllm-project/vllm/issues/6803>.

 ## No available node types can fulfill resource request

-The error message `Error: No available node types can fulfill resource request` can appear even when the cluster has enough GPUs. The issue often occurs when nodes have multiple IP addresses and vLLM can't select the correct one. Ensure that vLLM and Ray use the same IP address by setting `VLLM_HOST_IP` in [examples/ray_serving/run_cluster.sh](../../examples/ray_serving/run_cluster.sh) (with a different value on each node). Use `ray status` and `ray list nodes` to verify the chosen IP address. For more information, see <https://github.com/vllm-project/vllm/issues/7815>.
+The error message `Error: No available node types can fulfill resource request` can appear even when the cluster has enough GPUs. The issue often occurs when nodes have multiple IP addresses and vLLM can't select the correct one. Ensure that vLLM and Ray use the same IP address by setting `VLLM_HOST_IP` in [examples/online_serving/run_cluster.sh](../../examples/online_serving/run_cluster.sh) (with a different value on each node). Use `ray status` and `ray list nodes` to verify the chosen IP address. For more information, see <https://github.com/vllm-project/vllm/issues/7815>.

 ## Ray observability

@@ -251,7 +251,7 @@ The following extra parameters are supported:
 Our Responses API is compatible with [OpenAI's Responses API](https://platform.openai.com/docs/api-reference/responses);
 you can use the [official OpenAI Python client](https://github.com/openai/openai-python) to interact with it.

-Code example: [examples/tool_calling/openai_responses_client_with_tools.py](../../examples/tool_calling/openai_responses_client_with_tools.py)
+Code example: [examples/online_serving/openai_responses_client_with_tools.py](../../examples/tool_calling/openai_responses_client_with_tools.py)

 #### Extra parameters

@@ -456,8 +456,8 @@ Audio must be sent as base64-encoded PCM16 audio at 16kHz sample rate, mono chan

 #### Example Clients

- [openai_realtime_client.py](https://github.com/vllm-project/vllm/tree/main/examples/speech_to_text/realtime/openai_realtime_client.py) - Upload and transcribe an audio file
- [openai_realtime_microphone_client.py](https://github.com/vllm-project/vllm/tree/main/examples/speech_to_text/realtime/openai_realtime_microphone_client.py) - Gradio demo for live microphone transcription
+- [openai_realtime_client.py](https://github.com/vllm-project/vllm/tree/main/examples/online_serving/openai_realtime_client.py) - Upload and transcribe an audio file
+- [openai_realtime_microphone_client.py](https://github.com/vllm-project/vllm/tree/main/examples/online_serving/openai_realtime_microphone_client.py) - Gradio demo for live microphone transcription

 ### Tokenizer API

@@ -542,6 +542,6 @@ Key capabilities:
 - Scales from a single GPU to a multi-node cluster without code changes.
 - Provides observability and autoscaling policies through Ray dashboards and metrics.

-The following example shows how to deploy a large model like DeepSeek R1 with Ray Serve LLM: [examples/ray_serving/ray_serve_deepseek.py](../../examples/ray_serving/ray_serve_deepseek.py).
+The following example shows how to deploy a large model like DeepSeek R1 with Ray Serve LLM: [examples/online_serving/ray_serve_deepseek.py](../../examples/online_serving/ray_serve_deepseek.py).

 Learn more about Ray Serve LLM with the official [Ray Serve LLM documentation](https://docs.ray.io/en/latest/serve/llm/index.html).
@@ -78,7 +78,7 @@ For details, see the [Ray documentation](https://docs.ray.io/en/latest/index.htm

 ### Ray cluster setup with containers

-The helper script [examples/ray_serving/run_cluster.sh](../../examples/ray_serving/run_cluster.sh) starts containers across nodes and initializes Ray. By default, the script runs Docker without administrative privileges, which prevents access to the GPU performance counters when profiling or tracing. To enable admin privileges, add the `--cap-add=CAP_SYS_ADMIN` flag to the Docker command.
+The helper script [examples/online_serving/run_cluster.sh](../../examples/online_serving/run_cluster.sh) starts containers across nodes and initializes Ray. By default, the script runs Docker without administrative privileges, which prevents access to the GPU performance counters when profiling or tracing. To enable admin privileges, add the `--cap-add=CAP_SYS_ADMIN` flag to the Docker command.

 Choose one node as the head node and run:

@@ -162,7 +162,7 @@ vllm serve /path/to/the/model/in/the/container \

 Efficient tensor parallelism requires fast internode communication, preferably through high-speed network adapters such as InfiniBand.
 To set up the cluster to use InfiniBand, append additional arguments like `--privileged -e NCCL_IB_HCA=mlx5` to the
-[examples/ray_serving/run_cluster.sh](../../examples/ray_serving/run_cluster.sh) helper script.
+[examples/online_serving/run_cluster.sh](../../examples/online_serving/run_cluster.sh) helper script.
 Contact your system administrator for more information about the required flags.

 ## Enabling GPUDirect RDMA
@@ -60,4 +60,4 @@ The key insight is that requests paused with `mode="keep"` will produce tokens f

 ## Example

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

 ## Architecture

-The weight transfer system follows a **four-phase protocol** with a pluggable backend design:
+The weight transfer system follows a **two-phase protocol** with a pluggable backend design:

 1. **Initialization** (`init_weight_transfer_engine`): Establishes the communication channel between the trainer and inference workers. Called once before the training loop begins.
-2. **Start** (`start_weight_update`): Prepares the inference engine for a weight update.
-3. **Weight Update** (`update_weights`): Transfers updated weights from the trainer to the inference engine. May be called one or more times (e.g., for chunked transfers).
-4. **Finish** (`finish_weight_update`): Finalizes the weight update (e.g., runs post-processing for checkpoint-format weights). Called once after all weights have been transferred.
+2. **Weight Update** (`update_weights`): Transfers updated weights from the trainer to the inference engine. Called after each training step (or batch of steps).

 ## Available Backends

@@ -50,9 +48,7 @@ When running vLLM as an HTTP server, the following endpoints are available for w
 | Endpoint | Method | Description |
 | -------- | ------ | ----------- |
 | `/init_weight_transfer_engine` | POST | Initialize the weight transfer engine with backend-specific info |
-| `/start_weight_update` | POST | Start a weight update |
-| `/update_weights` | POST | Transfer a batch of weights with backend-specific metadata |
-| `/finish_weight_update` | POST | Finish the weight update and run post-processing |
+| `/update_weights` | POST | Trigger a weight update with backend-specific metadata |
 | `/pause` | POST | Pause generation before weight sync to handle inflight requests |
 | `/resume` | POST | Resume generation after weight sync |
 | `/get_world_size` | GET | Get the number of inference workers (useful for NCCL world size calculation) |
@@ -68,17 +64,11 @@ Both backends provide static methods that the trainer calls to send weights. The
 # 1. Initialize the transfer engine (backend-specific)
 EngineClass.trainer_init(init_info)

-# 2. Start weight update on inference side
-llm.start_weight_update(is_checkpoint_format=True)
-
-# 3. Send weights to inference workers
+# 2. Send weights to inference workers
 EngineClass.trainer_send_weights(
    iterator=model.named_parameters(),
    trainer_args=backend_specific_args,
 )
-
-# 4. Finish weight update on inference side
-llm.finish_weight_update()
 ```

 See the [NCCL](nccl.md) and [IPC](ipc.md) pages for backend-specific trainer APIs and full examples.
@@ -43,14 +43,16 @@ update_request = WeightTransferUpdateRequest(

 ### WeightTransferUpdateInfo

-The base `WeightTransferUpdateInfo` is a marker class for backend-specific update info:
+The base `WeightTransferUpdateInfo` includes an `is_checkpoint_format` flag:

 ```python
@dataclass
 class WeightTransferUpdateInfo(ABC):
-    pass
+    is_checkpoint_format: bool = True
 ```

+When `is_checkpoint_format=True` (the default), vLLM applies layerwise weight processing (repacking, renaming, etc.) on the received weights before loading them. Set to `False` if the trainer has already converted weights to the kernel format expected by the model.
+
 ## Implementing a Custom Engine

 To create a custom weight transfer backend:
@@ -38,15 +38,11 @@ trainer_args = IPCTrainerSendWeightsArgs(
    mode="ray",
    llm_handle=llm_actor_handle,
 )
-# start
-ray.get(llm_actor_handle.start_weight_update.remote(is_checkpoint_format=True))
-# send weights
+
 IPCWeightTransferEngine.trainer_send_weights(
    iterator=model.named_parameters(),
    trainer_args=trainer_args,
 )
-# finish
-ray.get(llm_actor_handle.finish_weight_update.remote())
 ```

 In Ray mode, the engine calls `llm_handle.update_weights.remote(...)` directly, passing the IPC handles via Ray's serialization.
@@ -61,27 +57,17 @@ trainer_args = IPCTrainerSendWeightsArgs(
    url="http://localhost:8000",
 )

-# start
-base_url = "http://localhost:8000"
-url = f"{base_url}/start_weight_update"
-response = requests.post(url, json={"is_checkpoint_format": True}, timeout=60)
-response.raise_for_status()
-# send weights
 IPCWeightTransferEngine.trainer_send_weights(
    iterator=model.named_parameters(),
    trainer_args=trainer_args,
 )
-# finish
-url = f"{base_url}/finish_weight_update"
-response = requests.post(url, json={}, timeout=60)
-response.raise_for_status()
 ```

-In HTTP mode, IPC handles are pickled, base64-encoded, and sent as JSON to the `/update_weights` endpoint. As with Ray mode, you must call `start_weight_update` before and `finish_weight_update` after.
+In HTTP mode, IPC handles are pickled, base64-encoded, and sent as JSON to the `/update_weights` endpoint.

 See [`IPCTrainerSendWeightsArgs`](https://github.com/vllm-project/vllm/blob/main/vllm/distributed/weight_transfer/ipc_engine.py) for the full list of configurable fields.

 ## Examples

- [RLHF with IPC weight syncing (offline, Ray)](../../../examples/rl/rlhf_ipc.py) - Colocated training and inference on a single GPU using Ray placement groups and CUDA IPC handles
- [RLHF with IPC weight syncing (online serving, HTTP)](../../../examples/rl/rlhf_http_ipc.py) - Weight transfer with a vLLM HTTP server where both server and trainer share the same GPU
+- [RLHF with IPC weight syncing (offline, Ray)](../../examples/rl/rlhf_ipc.md) - Colocated training and inference on a single GPU using Ray placement groups and CUDA IPC handles
+- [RLHF with IPC weight syncing (online serving, HTTP)](../../examples/rl/rlhf_http_ipc.md) - Weight transfer with a vLLM HTTP server where both server and trainer share the same GPU
@@ -84,15 +84,11 @@ Both the trainer (`NCCLTrainerSendWeightsArgs`) and inference side (`NCCLWeightT

 ## Receiving Weights (Inference Side)

-The inference side triggers weight reception using the four-phase protocol — `init_weight_transfer_engine`, `start_weight_update`, `update_weights`, `finish_weight_update`. The init phase is shown [above](#initialization). The remaining three steps are:
+The inference side triggers weight reception by calling `update_weights`:

 ```python
 from vllm.distributed.weight_transfer.base import WeightTransferUpdateRequest

-# 1. Start the weight update
-llm.start_weight_update(is_checkpoint_format=True)
-
-# 2. Receive weights (can be called multiple times for chunked transfers)
 llm.update_weights(
    WeightTransferUpdateRequest(
        update_info=dict(
@@ -103,17 +99,12 @@ llm.update_weights(
        )
    )
 )
-
-# 3. Finish the weight update
-llm.finish_weight_update()
 ```

 The `names`, `dtype_names`, and `shapes` lists describe each parameter. These must match the order in which the trainer iterates over its parameters.

-`start_weight_update` must be called before `update_weights`, and `finish_weight_update` must be called after all weight chunks have been transferred. The `is_checkpoint_format` flag controls whether layerwise reload processing is applied (`True` for checkpoint-format weights, `False` for pre-processed kernel-format weights).
-
 ## Examples

- [RLHF with NCCL weight syncing (offline, Ray)](../../../examples/rl/rlhf_nccl.py) - Trainer on one GPU, 2x tensor-parallel vLLM engine on two others, with packed NCCL weight broadcast
- [RLHF with async weight syncing (offline, Ray)](../../../examples/rl/rlhf_async_new_apis.py) - Async generation with mid-flight pause, weight sync, resume, and validation against a fresh model
- [RLHF with NCCL weight syncing (online serving, HTTP)](../../../examples/rl/rlhf_http_nccl.py) - Weight transfer with a running vLLM HTTP server using HTTP control plane and NCCL data plane
+- [RLHF with NCCL weight syncing (offline, Ray)](../../examples/rl/rlhf_nccl.md) - Trainer on one GPU, 2x tensor-parallel vLLM engine on two others, with packed NCCL weight broadcast
+- [RLHF with async weight syncing (offline, Ray)](../../examples/rl/rlhf_async_new_apis.md) - Async generation with mid-flight pause, weight sync, resume, and validation against a fresh model
+- [RLHF with NCCL weight syncing (online serving, HTTP)](../../examples/rl/rlhf_http_nccl.md) - Weight transfer with a running vLLM HTTP server using HTTP control plane and NCCL data plane
@@ -309,30 +309,6 @@ vLLM supports dynamically loading and unloading LoRA adapters at runtime via the

 **Warning:** Dynamic LoRA loading is not a secure operation and should not be enabled in deployments exposed to untrusted clients. If you must enable dynamic LoRA loading, restrict access to the `/v1/load_lora_adapter` and `/v1/unload_lora_adapter` endpoints to trusted administrators only, using a reverse proxy or network-level access controls. Do not expose these endpoints to end users. For details on configuring LoRA adapters, see the [LoRA Adapters documentation](../features/lora.md).

-## Cache Directory Security
-
-vLLM assumes that its cache directories are **private and trusted**. Cache contents are loaded without cryptographic integrity verification, including formats that support arbitrary code execution. If an untrusted user or process can write to vLLM's cache directories, they may be able to crash vLLM or cause it to execute arbitrary code.
-
-**Do not share vLLM cache directories with untrusted users or mount them from untrusted storage.** Treat the cache directory with the same care as the vLLM installation itself.
-
-### Cache Directory Configuration
-
-Most cache paths default to subdirectories under a single root. Changing `VLLM_CACHE_ROOT` changes the default location for all features that inherit from it. When `torch.compile` caching is enabled (the default), vLLM also redirects `TRITON_CACHE_DIR` into this tree. If compile caching is disabled, Triton falls back to its own default location (`~/.triton/cache`).
-
-| Environment Variable | Default | Description |
-| --- | --- | --- |
-| `VLLM_CACHE_ROOT` | `~/.cache/vllm` | Base cache directory. Respects `XDG_CACHE_HOME` if set. All paths below inherit from this unless explicitly overridden. |
-| *(torch.compile)* | `$VLLM_CACHE_ROOT/torch_compile_cache/` | Compilation cache for AOT-compiled models, Inductor graphs, and Triton kernels. Controlled by `VLLM_DISABLE_COMPILE_CACHE` (set to `1` to disable). |
-| `VLLM_ASSETS_CACHE` | `$VLLM_CACHE_ROOT/assets/` | Downloaded assets (e.g., tokenizer files). |
-| `VLLM_XLA_CACHE_PATH` | `$VLLM_CACHE_ROOT/xla_cache/` | XLA/TPU compilation cache. |
-| `VLLM_MEDIA_CACHE` | *(disabled)* | Optional cache for downloaded media (images, video, audio). Not enabled unless explicitly set. |
-
-### Recommendations
-
- **Restrict file permissions** on `VLLM_CACHE_ROOT` (and any other cache directories used by dependencies, such as `~/.triton` if compile caching is disabled) so that only the vLLM process owner can read and write to them.
- **Do not copy cache contents from untrusted sources.** If you distribute cache artifacts between environments, ensure they originate from a trusted build pipeline.
- **Container deployments:** If mounting cache directories into containers, ensure the volume source is trusted.
-
 ## Reporting Security Vulnerabilities

 If you believe you have found a security vulnerability in vLLM, please report it following the project's security policy. For more information on how to report security issues and the project's security policy, please see the [vLLM Security Policy](https://github.com/vllm-project/vllm/blob/main/SECURITY.md).
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
Yongye ZhuandClaude Opus 4.7	58c8a5eaa5	[Attention][TokenSpeed MLA] Also warm up prefill kernel from decode impl The prefill backend may be paired with flash_attn / trtllm in production — in that case the prefill backend's __init__ never runs and the prefill kernel's first call pays a 1.5–2 minute JIT cost. Add the same idempotent `warmup_compile_prefill` invocation to TokenspeedMLAImpl.__init__ (the decode-side backend, always present when tokenspeed is selected). The function dedupes by config key, so the double call is a no-op when both backends are tokenspeed. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>	2026-05-06 07:45:16 +00:00
Yongye Zhu	c4547482ca	update version Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>	2026-05-06 07:28:20 +00:00
Yongye Zhu	91ef0afcb2	adding to cuda.txt dependency Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>	2026-05-06 07:20:48 +00:00
Yongye Zhu	97cd2c41ad	fix precommit Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>	2026-05-06 03:14:49 +00:00
Yongye ZhuandClaude Opus 4.7	c001535038	[Attention][TokenSpeed MLA] Force prefill V tensor contiguous before kernel call `v` arrives at both `run_prefill_new_tokens` and `run_prefill_context_chunk` as the second half of `kv_nope.split([qk_nope_head_dim, v_head_dim], dim=-1)` in mla_attention.py — a non-contiguous view along the last dim. The kernel internally does `v.reshape(1, total_kv, h_k, 1, d_v)` which silently copies when the input is non-contiguous; pull that copy out so the layout is predictable at the kernel boundary. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>	2026-05-06 02:53:46 +00:00
Yongye ZhuandClaude Opus 4.7	de6bc297df	[Attention][TokenSpeed MLA] Surface install hint when package missing Previously a user explicitly selecting TOKENSPEED_MLA without `tokenspeed_mla` installed got either a generic "required dependencies not available" message (prefill backend) or a raw ModuleNotFoundError deep inside forward_mqa at the first request (decode backend). Now both backends fail at startup with the exact install command: `uv pip install tokenspeed-mla`. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>	2026-05-06 02:49:09 +00:00
Yongye ZhuandClaude Opus 4.7	0012818287	[Attention][TokenSpeed MLA] Fix trtllm LSE parity test: log2 → natural log trtllm_ragged_attention_deepseek returns LSE in log2; tokenspeed and merge_attn_states use natural log. Multiply the trtllm reference by ln 2 before comparison. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>	2026-05-06 02:49:09 +00:00
Yongye ZhuandClaude Opus 4.7	73cd7e25ae	[Attention][TokenSpeed MLA] Warm up BF16 prefill compile, drop seq_lens computation Pre-JIT both BF16 and FP8 prefill kernels at backend init since the dtype isn't visible from `__init__` — depends on `use_prefill_query_quantization`. Move the per-forward `seq_lens` computation into `prepare_metadata` and document the cuda-graph padding interaction with `query_start_loc`. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>	2026-05-06 02:49:09 +00:00
Yongye ZhuandClaude Opus 4.7	964c6eb485	[Attention][TokenSpeed MLA] Fix decode FP8 numerics: pass output_scale and assert FP8 Q The decode kernel needs both bmm scales to recover correct outputs from an FP8 KV cache: bmm1 (softmax_scale = scale * q_scale * k_scale) and bmm2 (output_scale = k_scale, since V is stored as V_real / k_scale). We were only passing bmm1, which left bmm2 = 1.0 and produced silently wrong output. Also assert query dtype is float8_e4m3fn on entry to forward_mqa. supports_quant_query_input=True (inherited from MLACommonImpl) tells the upstream pipeline to FP8-quantize Q via _decode_concat_quant_fp8_op; the kernel is shape-specialized for FP8 Q + FP8 KV, so any other dtype here means the upstream quant path didn't run and the kernel will produce garbage. Failing loud beats failing silent. Verified: gsm8k matches reference with TOKENSPEED_MLA decode + FLASH_ATTN prefill on Kimi-K2.5-NVFP4 / TP=4 / B200. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>	2026-05-06 02:49:09 +00:00
Yongye ZhuandClaude Opus 4.7	d0e6514bf8	[Attention] Add TOKENSPEED_MLA backend for DeepSeek R1 prefill + decode on Blackwell Wires the tokenspeed_mla CuTe DSL kernels into vLLM as a new MLA backend, covering both prefill (tokenspeed_mla_prefill) and decode (tokenspeed_mla_decode). Targets Blackwell (SM100) with FP8 KV cache and DeepSeek R1 MLA dimensions; users opt in via -ac '{"backend":"TOKENSPEED_MLA","mla_prefill_backend":"TOKENSPEED_MLA"}'. Includes numeric parity tests against the trtllm reference kernels. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>	2026-05-06 02:49:09 +00:00